"""
Simplified Stochastic Process Framework
Supports array dt for varying time steps
"""

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.math.gamma(nu)) * x**nu * np.exp(-x)
        case "custom" if custom:
            return custom(tau)
        case _:
            return np.exp(-np.abs(t))


def time_from_dt(dt):
    """
    Convert dt to time array.
    
    Parameters
    ----------
    dt : float or array
        If float: constant time step
        If array: varying time steps between samples
        
    Returns
    -------
    t : array
        Cumulative time array
    """
    if np.isscalar(dt):
        return None  # Signal to use linear indexing
    return np.concatenate([[0], np.cumsum(dt)])


def build_covariance(n, dt, variance_array, timescale, corr_type="exp", custom=None):
    """
    Build covariance matrix from autocorrelation.
    Now supports varying dt (dt as array).
    """
    # Get actual time differences
    if np.isscalar(dt):
        dts = np.full(n - 1, dt)
    else:
        dts = np.asarray(dt)
        if len(dts) != n - 1:
            raise ValueError(f"dt array length ({len(dts)}) must be n-1 ({n-1})")
    
    # Build time array: t[0]=0, t[i+1] = t[i] + dt[i]
    t = np.zeros(n)
    for i in range(1, n):
        t[i] = t[i-1] + dts[i-1]
    
    # Variance at each time point
    V = np.diag(variance_array)
    
    # Correlation matrix using actual time differences
    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 or array of shape (n-1,)
        Time step(s) between samples
    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":
        std = np.sqrt(variance_array)
        return mean_array + std * rng.standard_normal(n)
    
    # Build covariance matrix (now handles varying dt)
    cov = build_covariance(n, dt, variance_array, timescale, corr_type, custom)
    cov += np.eye(n) * 1e-10
    
    try:
        L = np.linalg.cholesky(cov)
        z = rng.standard_normal(n)
        return mean_array + L @ z
    except np.linalg.LinAlgError:
        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."""
    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]
        )
    
    # Return actual time array
    if np.isscalar(dt):
        time = np.arange(n_samples) * dt
    else:
        time = np.zeros(n_samples)
        for i in range(1, n_samples):
            time[i] = time[i-1] + dt[i-1]
    
    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]
        time = None
    else:
        mean = np.mean(samples, axis=0)
        var = np.var(samples, axis=0)
        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)
        time = None
    
    # Frequencies based on average dt
    if np.isscalar(dt):
        avg_dt = dt
    else:
        avg_dt = np.mean(dt)
    
    lags = np.arange(n) * avg_dt
    freqs = fftfreq(n, avg_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: Constant dt ---
    print("=" * 50)
    print("Example 1: Constant dt")
    print("=" * 50)
    
    n = 500
    dt_const = 0.01
    
    mean_arr = np.zeros(n)
    var_arr = np.ones(n)
    
    samples = generate_samples(n, mean_arr, var_arr, dt_const, seed=42)
    stats = compute_stats(samples, dt_const)
    
    print(f"dt = {dt_const} (scalar)")
    print(f"Total time: {n * dt_const:.2f} seconds")
    print(f"Mean: {stats['mean']:.4f}, Variance: {stats['variance']:.4f}")
    
    
    # --- Example 2: Varying dt with linspace ---
    print("\n" + "=" * 50)
    print("Example 2: Varying dt (linspace)")
    print("=" * 50)
    
    dt_var = np.linspace(0.001, 0.05, n - 1)  # dt increases from 1ms to 50ms
    
    mean_arr = np.zeros(n)
    var_arr = np.ones(n)
    
    samples = generate_samples(n, mean_arr, var_arr, dt_var, corr_type="exp", 
                               timescale=0.1, seed=42)
    
    # Get actual time
    time = np.zeros(n)
    for i in range(1, n):
        time[i] = time[i-1] + dt_var[i-1]
    
    stats = compute_stats(samples, dt_var)
    
    print(f"dt range: [{dt_var[0]:.4f}, {dt_var[-1]:.4f}] seconds")
    print(f"Total time: {time[-1]:.2f} seconds (vs {n*dt_const:.2f} for constant)")
    print(f"Mean: {stats['mean']:.4f}, Variance: {stats['variance']:.4f}")
    print(f"Actual time array: [{time[0]:.4f}, {time[1]:.4f}, ..., {time[-1]:.4f}]")
    
    
    # --- Example 3: Sparse then dense sampling ---
    print("\n" + "=" * 50)
    print("Example 3: Sparse → Dense Sampling")
    print("=" * 50)
    
    # Sparse first half, dense second half
    dt_sparse_dense = np.concatenate([
        np.full((n//2 - 1), 0.1),    # Sparse: 100ms steps
        np.full((n//2), 0.001)       # Dense: 1ms steps
    ])
    
    mean_arr = np.zeros(n)
    var_arr = np.ones(n)
    
    samples = generate_samples(n, mean_arr, var_arr, dt_sparse_dense, 
                               corr_type="gaussian", timescale=0.5, seed=42)
    
    time = np.zeros(n)
    for i in range(1, n):
        time[i] = time[i-1] + dt_sparse_dense[i-1]
    
    print(f"First half: dt = {dt_sparse_dense[0]:.3f} (sparse)")
    print(f"Second half: dt = {dt_sparse_dense[-1]:.3f} (dense)")
    print(f"Total time: {time[-1]:.2f} seconds")
    print(f"Time at sample {n//2}: {time[n//2]:.2f} seconds (90% of total!)")
    
    
    # --- Example 4: Time-varying mean with varying dt ---
    print("\n" + "=" * 50)
    print("Example 4: Time-Varying Mean + Varying dt")
    print("=" * 50)
    
    dt_var = np.linspace(0.01, 0.1, n - 1)
    time = np.zeros(n)
    for i in range(1, n):
        time[i] = time[i-1] + dt_var[i-1]
    
    mean_arr = 5 * np.sin(2 * np.pi * 0.1 * time)  # 0.1 Hz sine in actual time
    var_arr = np.ones(n)
    
    samples = generate_samples(n, mean_arr, var_arr, dt_var, seed=42)
    
    stats = compute_stats(samples, dt_var)
    
    print(f"Mean tracks actual time, not sample index")
    print(f"Target mean range: [{mean_arr.min():.2f}, {mean_arr.max():.2f}]")
    print(f"Sample mean range: [{samples.min():.2f}, {samples.max():.2f}]")
    print(f"Mean correlation: {np.corrcoef(mean_arr, samples)[0,1]:.4f}")
    
    
    # --- Example 5: Compare autocorrelation with varying dt ---
    print("\n" + "=" * 50)
    print("Example 5: Autocorrelation Shape Difference")
    print("=" * 50)
    
    dt_const = 0.01
    dt_var = np.linspace(0.001, 0.05, n - 1)
    
    timescale = 0.1
    
    # Generate with constant dt
    samples_const = generate_samples(n, np.zeros(n), np.ones(n), dt_const,
                                     corr_type="exp", timescale=timescale, seed=42)
    stats_const = compute_stats(samples_const, dt_const)
    
    # Generate with varying dt
    samples_var = generate_samples(n, np.zeros(n), np.ones(n), dt_var,
                                   corr_type="exp", timescale=timescale, seed=43)
    stats_var = compute_stats(samples_var, dt_var)
    
    print("With varying dt, autocorrelation lags represent actual time, not indices")
    print(f"Constant dt - lag 10 = {stats_const['autocorrelation'][10]:.4f}")
    print(f"Varying dt  - lag 10 = {stats_var['autocorrelation'][10]:.4f}")
    print(f"(Lag 10 means ~{0.1:.2f}s for constant, but actual time varies)")