"""
Simplified Stochastic Process Framework
Using arrays for mean, variance over time
"""

import numpy as np
from scipy.fft import fft, fftfreq
from enum import Enum


class CorrType(Enum):
    WHITE = "white"
    EXPONENTIAL = "exp"
    GAUSSIAN = "gaussian"
    SINUSOIDAL = "sinusoidal"
    MATERN = "matern"
    CUSTOM = "custom"


def autocorr(tau, timescale, corr_type="exp", custom=None):
    """Compute autocorrelation for time lags tau."""
    t = tau / timescale
    
    match corr_type:
        case "white":
            return np.where(tau == 0, 1.0, 0.0)
        case "exp":
            return np.exp(-np.abs(t))
        case "gaussian":
            return np.exp(-t**2)
        case "sinusoidal":
            return np.exp(-np.abs(t)) * np.cos(np.pi * t)
        case "matern":
            nu = 1.5
            x = np.sqrt(2 * nu) * np.abs(t)
            return (2**(1-nu) / np.random.gamma(nu)) * x**nu * np.exp(-x)
        case "custom" if custom:
            return custom(tau)
        case _:
            return np.exp(-np.abs(t))


def build_covariance(n, dt, variance_array, timescale, corr_type="exp", custom=None):
    """Build covariance matrix from autocorrelation."""
    # Time grid
    t = np.arange(n) * dt
    
    # Variance at each time point (diagonal)
    V = np.diag(variance_array)
    
    # Correlation matrix
    tau = np.abs(t[:, None] - t[None, :])
    R = autocorr(tau, timescale, corr_type, custom)
    
    # Combined: Cov = D^(1/2) * R * D^(1/2)
    D_half = np.diag(np.sqrt(variance_array))
    return D_half @ R @ D_half


def generate_samples(n, mean_array, variance_array, dt, 
                     corr_type="white", timescale=1.0, custom=None, seed=None):
    """
    Generate correlated samples.
    
    Parameters
    ----------
    n : int
        Number of samples
    mean_array : array of shape (n,)
        Mean at each time point
    variance_array : array of shape (n,)
        Variance at each time point
    dt : float
        Time step
    corr_type : str
        Autocorrelation type
    timescale : float
        Correlation timescale
    seed : int, optional
        Random seed
        
    Returns
    -------
    samples : array of shape (n,)
    """
    rng = np.random.default_rng(seed)
    
    if corr_type == "white":
        # Simple: independent samples
        std = np.sqrt(variance_array)
        return mean_array + std * rng.standard_normal(n)
    
    # Build covariance matrix
    cov = build_covariance(n, dt, variance_array, timescale, corr_type, custom)
    
    # Ensure positive definite
    cov += np.eye(n) * 1e-10
    
    # Cholesky decomposition
    try:
        L = np.linalg.cholesky(cov)
        z = rng.standard_normal(n)
        return mean_array + L @ z
    except np.linalg.LinAlgError:
        # Fallback: eigendecomposition
        eigenvalues, eigenvectors = np.linalg.eigh(cov)
        eigenvalues = np.maximum(eigenvalues, 1e-10)
        L = eigenvectors @ np.diag(np.sqrt(eigenvalues))
        z = rng.standard_normal(n)
        return mean_array + L @ z


def generate_paths(n_samples, n_paths, mean_array, variance_array, dt,
                   corr_type="white", timescale=1.0, custom=None, seed=None):
    """
    Generate multiple independent paths.
    
    Parameters
    ----------
    n_samples : int
        Samples per path
    n_paths : int
        Number of paths
    mean_array, variance_array : array of shape (n_samples,)
        Parameters for each time point
    dt : float
        Time step
    seed : int, optional
        Random seed
        
    Returns
    -------
    paths : array of shape (n_paths, n_samples)
    time : array of shape (n_samples,)
    """
    rng = np.random.default_rng(seed)
    seeds = rng.integers(0, 2**31, size=n_paths)
    
    paths = np.zeros((n_paths, n_samples))
    
    for i in range(n_paths):
        paths[i] = generate_samples(
            n_samples, mean_array, variance_array, dt,
            corr_type, timescale, custom, seed=seeds[i]
        )
    
    time = np.arange(n_samples) * dt
    return paths, time


def compute_stats(samples, dt=1.0):
    """Compute empirical statistics."""
    n = samples.shape[-1]
    
    if samples.ndim == 1:
        mean = np.mean(samples)
        var = np.var(samples)
        acorr = np.correlate(samples - mean, samples - mean, mode='full')
        acorr = acorr[n-1:] / acorr[n-1]
    else:
        mean = np.mean(samples, axis=0)
        var = np.var(samples, axis=0)
        # Average autocorrelation across paths
        acorr = np.zeros(n)
        for s in samples:
            m = np.mean(s)
            a = np.correlate(s - m, s - m, mode='full')
            a = a[n-1:] / a[n-1]
            acorr += a
        acorr /= len(samples)
    
    lags = np.arange(n) * dt
    freqs = fftfreq(n, dt)
    psd = np.abs(fft(samples - np.mean(samples)))**2 / n
    
    return {
        'mean': mean,
        'variance': var,
        'autocorrelation': acorr,
        'lags': lags,
        'frequencies': freqs,
        'psd': psd
    }


# =============================================================================
# Usage Examples
# =============================================================================

if __name__ == "__main__":
    
    # --- Example 1: Simple white noise ---
    print("=" * 50)
    print("Example 1: White Noise")
    print("=" * 50)
    
    n = 500
    dt = 0.01
    time = np.arange(n) * dt
    
    mean_arr = np.zeros(n)
    var_arr = np.ones(n)
    
    samples = generate_samples(n, mean_arr, var_arr, dt, corr_type="white", seed=42)
    
    stats = compute_stats(samples)
    print(f"Mean: {stats['mean']:.4f} (target: 0)")
    print(f"Variance: {stats['variance']:.4f} (target: 1)")
    
    
    # --- Example 2: Colored noise ---
    print("\n" + "=" * 50)
    print("Example 2: Colored Noise (Exponential)")
    print("=" * 50)
    
    mean_arr = np.zeros(n)
    var_arr = np.full(n, 4.0)
    
    samples = generate_samples(n, mean_arr, var_arr, dt, 
                               corr_type="exp", timescale=0.05, seed=42)
    
    stats = compute_stats(samples, dt)
    print(f"Mean: {stats['mean']:.4f}")
    print(f"Variance: {stats['variance']:.4f}")
    print(f"Autocorrelation at lag 0: {stats['autocorrelation'][0]:.4f}")
    print(f"Autocorrelation at lag 10: {stats['autocorrelation'][10]:.4f}")
    
    
    # --- Example 3: Time-varying mean ---
    print("\n" + "=" * 50)
    print("Example 3: Time-Varying Mean")
    print("=" * 50)
    
    mean_arr = 5 * np.sin(2 * np.pi * 0.5 * time)  # 0.5 Hz sine
    var_arr = np.full(n, 1.0)
    
    samples = generate_samples(n, mean_arr, var_arr, dt, corr_type="gaussian", 
                               timescale=0.02, seed=42)
    
    stats = compute_stats(samples, dt)
    print(f"Mean range: [{samples.min():.2f}, {samples.max():.2f}]")
    print(f"Target mean range: [{mean_arr.min():.2f}, {mean_arr.max():.2f}]")
    
    
    # --- Example 4: Time-varying variance ---
    print("\n" + "=" * 50)
    print("Example 4: Time-Varying Variance")
    print("=" * 50)
    
    mean_arr = np.zeros(n)
    var_arr = 1.0 + 0.8 * np.sin(2 * np.pi * 2.0 * time)  # Modulated variance
    var_arr = np.maximum(var_arr, 0.1)  # Ensure positive
    
    samples = generate_samples(n, mean_arr, var_arr, dt, corr_type="exp",
                               timescale=0.01, seed=42)
    
    stats = compute_stats(samples, dt)
    print(f"Empirical mean: {stats['mean']:.4f}")
    print(f"Empirical variance (avg): {np.mean(stats['variance']):.4f}")
    print(f"Variance range: [{var_arr.min():.2f}, {var_arr.max():.2f}]")
    
    
    # --- Example 5: Multiple paths ---
    print("\n" + "=" * 50)
    print("Example 5: Multiple Paths")
    print("=" * 50)
    
    mean_arr = np.zeros(n)
    var_arr = np.ones(n)
    
    paths, time = generate_paths(n_samples=n, n_paths=50, 
                                  mean_array=mean_arr, 
                                  variance_array=var_arr, 
                                  dt=dt, 
                                  corr_type="matern", 
                                  timescale=0.1,
                                  seed=123)
    
    print(f"Shape: {paths.shape} (50 paths, 500 samples each)")
    print(f"Ensemble mean at t=0: {np.mean(paths[:, 0]):.4f}")
    print(f"Ensemble mean at t=2.5: {np.mean(paths[:, 250]):.4f}")
    
    
    # --- Example 6: Custom autocorrelation ---
    print("\n" + "=" * 50)
    print("Example 6: Custom Autocorrelation")
    print("=" * 50)
    
    def damped_oscillator(tau):
        return np.exp(-tau / 0.03) * np.cos(2 * np.pi * 10 * tau)
    
    mean_arr = np.zeros(n)
    var_arr = np.full(n, 2.0)
    
    samples = generate_samples(n, mean_arr, var_arr, dt, 
                               corr_type="custom", custom=damped_oscillator, seed=42)
    
    stats = compute_stats(samples, dt)
    print("Custom: exp(-t/0.03) * cos(2π * 10 * t)")
    print(f"Mean: {stats['mean']:.4f}")
    print(f"Variance: {stats['variance']:.4f}")