"""
Parametric volatility models for testing and simulation.

Implements common non-stationary volatility models as described in
Cavaliere & Taylor (2007, 2008).
"""

import numpy as np
from typing import Callable, Optional


def single_break(
    tau: float = 0.5,
    delta: float = 3.0,
    sigma0: float = 1.0,
) -> Callable[[np.ndarray], np.ndarray]:
    """
    Single abrupt volatility break.
    
    Parameters
    ----------
    tau : float, default=0.5
        Break fraction (0 < tau < 1)
    delta : float, default=3.0
        Volatility ratio: delta = sigma0 / sigma1
    sigma0 : float, default=1.0
        Initial volatility level
        
    Returns
    -------
    omega : callable
        Volatility function ω(s) where:
        ω(s) = sigma0 for s < tau
        ω(s) = sigma1 for s >= tau
        
    Notes
    -----
    Model A from Cavaliere & Taylor (2008, Section 5):
    ω₁(s)² = σ₀² + (σ₁² - σ₀²)I(s ≥ τ)
    
    Common parameterizations:
    - Early negative break: tau=0.2, delta=3 (sigma increases)
    - Late positive break: tau=0.8, delta=1/3 (sigma decreases)
    """
    sigma1 = sigma0 / delta
    
    def omega(s: np.ndarray) -> np.ndarray:
        s = np.atleast_1d(s)
        return np.where(s < tau, sigma0, sigma1)
    
    return omega


def double_break(
    tau: float = 0.3,
    delta: float = 3.0,
    sigma0: float = 1.0,
    symmetric: bool = True,
) -> Callable[[np.ndarray], np.ndarray]:
    """
    Two volatility breaks.
    
    Parameters
    ----------
    tau : float, default=0.3
        First break fraction
    delta : float, default=3.0
        Volatility ratio
    sigma0 : float, default=1.0
        Initial and final volatility level
    symmetric : bool, default=True
        If True, second break at 1-tau (symmetric breaks)
        
    Returns
    -------
    omega : callable
        Volatility function with two breaks
        
    Notes
    -----
    Model B from Cavaliere & Taylor (2008, Section 5):
    ω₂(s)² = σ₀² + (σ₁² - σ₀²)I(τ ≤ s < (1-τ))
    """
    sigma1 = sigma0 / delta
    tau2 = 1 - tau if symmetric else 0.7
    
    def omega(s: np.ndarray) -> np.ndarray:
        s = np.atleast_1d(s)
        return np.where((s >= tau) & (s < tau2), sigma1, sigma0)
    
    return omega


def trending_volatility(
    delta: float = 3.0,
    sigma0: float = 1.0,
) -> Callable[[np.ndarray], np.ndarray]:
    """
    Linear trending volatility.
    
    Parameters
    ----------
    delta : float, default=3.0
        Volatility ratio: delta = sigma0 / sigma1
    sigma0 : float, default=1.0
        Initial volatility at s=0
        
    Returns
    -------
    omega : callable
        Volatility function ω(s) = sigma0 + (sigma1 - sigma0)s
        
    Notes
    -----
    Model C from Cavaliere & Taylor (2008, Section 5):
    ω₃(s) = σ₀ + (σ₁ - σ₀)s
    
    Generates linearly trending variance:
    - Upward trend: delta < 1 (sigma increases)
    - Downward trend: delta > 1 (sigma decreases)
    """
    sigma1 = sigma0 / delta
    
    def omega(s: np.ndarray) -> np.ndarray:
        s = np.atleast_1d(s)
        return sigma0 + (sigma1 - sigma0) * s
    
    return omega


def smooth_transition(
    tau: float = 0.5,
    gamma: float = 10.0,
    delta: float = 3.0,
    sigma0: float = 1.0,
) -> Callable[[np.ndarray], np.ndarray]:
    """
    Smooth (logistic) transition volatility break.
    
    Parameters
    ----------
    tau : float, default=0.5
        Transition midpoint
    gamma : float, default=10.0
        Transition speed (larger = faster transition)
    delta : float, default=3.0
        Volatility ratio
    sigma0 : float, default=1.0
        Initial volatility level
        
    Returns
    -------
    omega : callable
        Volatility function with smooth transition
        
    Notes
    -----
    Logistic smooth transition model:
    ω(s)² = σ₀² + (σ₁² - σ₀²)S(s)
    S(s) = (1 + exp(-γ(s - τ)))⁻¹
    
    As γ → ∞, converges to abrupt break at tau
    As γ → 0, converges to constant σ = 0.5(σ₀ + σ₁)
    """
    sigma1 = sigma0 / delta
    
    def S(s):
        """Logistic transition function."""
        return 1.0 / (1.0 + np.exp(-gamma * (s - tau)))
    
    def omega(s: np.ndarray) -> np.ndarray:
        s = np.atleast_1d(s)
        var0_sq = sigma0 ** 2
        var1_sq = sigma1 ** 2
        var_sq = var0_sq + (var1_sq - var0_sq) * S(s)
        return np.sqrt(var_sq)
    
    return omega


def integrated_sv(
    nu: float = 4.0,
    c: float = 0.0,
    sigma0: float = 1.0,
    T: int = 1000,
    seed: Optional[int] = None,
) -> Callable[[np.ndarray], np.ndarray]:
    """
    Exponential (near-)integrated stochastic volatility.
    
    Parameters
    ----------
    nu : float, default=4.0
        Volatility of volatility parameter
    c : float, default=0.0
        Mean reversion parameter (c=0 for integrated, c>0 for stationary)
    sigma0 : float, default=1.0
        Initial volatility level
    T : int, default=1000
        Number of time steps
    seed : int, optional
        Random seed
        
    Returns
    -------
    omega : callable
        Stochastic volatility function
        
    Notes
    -----
    Model D from Cavaliere & Taylor (2008, Section 5):
    ω₄(s)² = σ₀² exp(ν J_c(s))
    
    where J_c(s) is:
    - Brownian motion if c=0 (integrated)
    - Ornstein-Uhlenbeck process if c>0 (stationary)
    
    Note: This generates one realization of the stochastic process.
    """
    if seed is not None:
        np.random.seed(seed)
    
    # Generate driving process J_c(s)
    dt = 1.0 / T
    dW = np.random.normal(0, np.sqrt(dt), T)
    
    if c == 0:
        # Integrated: J(s) is Brownian motion
        J = np.concatenate([[0], np.cumsum(dW)])
    else:
        # Stationary: J(s) is Ornstein-Uhlenbeck
        J = np.zeros(T + 1)
        for t in range(T):
            J[t + 1] = J[t] * np.exp(-c * dt) + dW[t]
    
    # Construct volatility path
    s_grid = np.linspace(0, 1, T + 1)
    vol_path = sigma0 * np.exp(nu * J)
    
    # Create interpolated function
    from scipy.interpolate import interp1d
    omega_func = interp1d(s_grid, vol_path, kind='linear', bounds_error=False, 
                          fill_value=(vol_path[0], vol_path[-1]))
    
    def omega(s: np.ndarray) -> np.ndarray:
        return omega_func(s)
    
    return omega


def garch_11(
    alpha0: float = 0.1,
    alpha1: float = 0.1,
    beta1: float = 0.8,
    T: int = 1000,
    seed: Optional[int] = None,
) -> Callable[[np.ndarray], np.ndarray]:
    """
    Stationary GARCH(1,1) volatility.
    
    Parameters
    ----------
    alpha0 : float, default=0.1
        Constant term
    alpha1 : float, default=0.1
        ARCH parameter
    beta1 : float, default=0.8
        GARCH parameter
    T : int, default=1000
        Number of time steps
    seed : int, optional
        Random seed
        
    Returns
    -------
    omega : callable
        GARCH volatility function
        
    Notes
    -----
    Model 5 from Cavaliere & Taylor (2008, Section 5):
    σ²ₜ = α₀ + α₁ε²ₜ₋₁ + β₁σ²ₜ₋₁
    
    Stationarity requires α₁ + β₁ < 1
    
    Note: Returns one simulated path from the GARCH process.
    """
    if alpha1 + beta1 >= 1:
        raise ValueError("GARCH(1,1) requires alpha1 + beta1 < 1 for stationarity")
    
    if seed is not None:
        np.random.seed(seed)
    
    # Simulate GARCH process
    sigma2 = np.zeros(T + 1)
    epsilon = np.random.normal(0, 1, T + 1)
    
    # Initialize at unconditional variance
    sigma2[0] = alpha0 / (1 - alpha1 - beta1)
    
    for t in range(T):
        sigma2[t + 1] = alpha0 + alpha1 * (epsilon[t] * np.sqrt(sigma2[t]))**2 + beta1 * sigma2[t]
    
    sigma = np.sqrt(sigma2)
    
    # Create interpolated function
    s_grid = np.linspace(0, 1, T + 1)
    from scipy.interpolate import interp1d
    omega_func = interp1d(s_grid, sigma, kind='linear', bounds_error=False,
                          fill_value=(sigma[0], sigma[-1]))
    
    def omega(s: np.ndarray) -> np.ndarray:
        return omega_func(s)
    
    return omega
