"""Black-Scholes <-> Heat equation transforms via Wilmott substitution.

The substitution:
    x = ln(S/K),  tau = sigma^2 * T / 2
    k = 2(r-q) / sigma^2
    alpha = -(k-1)/2,  beta = -(k+1)^2 / 4
    V(S, t) = K * exp(alpha*x + beta*tau) * u(x, tau)

transforms the Black-Scholes PDE into the heat equation u_tau = u_xx.
"""

import numpy as np
from scipy.stats import norm

from .validation import GridConfig, MarketParams

# Maximum safe exponent to prevent overflow (exp(709) ~ float64 max)
_MAX_EXP = 700.0


def compute_transform_params(market: MarketParams) -> tuple[float, float, float, float]:
    """Compute Wilmott transform parameters from market data.

    Returns:
        (k, alpha, beta, t_eff) where t_eff = sigma^2 * T / 2 is the
        effective time for the heat equation.
    """
    q = getattr(market, "q", 0.0)
    r_eff = market.r - q
    k = 2 * r_eff / market.sigma**2
    alpha = -(k - 1) / 2
    beta = -((k - 1) ** 2) / 4 - 2 * market.r / (market.sigma**2)
    t_eff = market.sigma**2 * market.T / 2
    return k, alpha, beta, t_eff


def make_grid(config: GridConfig) -> np.ndarray:
    """Build spatial grid with optional snapping/stretching extensions."""
    if config.snap_points:
        from .adaptive_grid import make_snapped_grid

        return make_snapped_grid(config)
    if config.stretch > 0:
        from .adaptive_grid import make_stretched_grid

        return make_stretched_grid(config)
    return np.linspace(-config.L, config.L, config.N, endpoint=False)


def make_taper(x_grid: np.ndarray, config: GridConfig) -> np.ndarray:
    """Create cosine taper that smoothly zeros the function near domain boundaries.

    taper(x) = 1                                        if |x| <= L - taper_width
    taper(x) = 0.5*(1 + cos(pi*(|x| - edge)/width))    if edge < |x| < L
    """
    taper = np.ones_like(x_grid)
    edge = config.L - config.taper_width
    abs_x = np.abs(x_grid)
    mask = abs_x > edge
    t = (abs_x[mask] - edge) / config.taper_width
    taper[mask] = 0.5 * (1 + np.cos(np.pi * t))
    return taper


def bs_to_heat_initial(
    x_grid: np.ndarray,
    market: MarketParams,
    config: GridConfig,
    option_type: str = "call",
) -> np.ndarray:
    """Convert BS payoff to heat equation initial condition u(x, 0) with cosine taper.

    u(x, 0) = exp(-alpha * x) * payoff_normalized(x) * taper(x)

    where payoff_normalized = max(e^x - 1, 0) for call, max(1 - e^x, 0) for put.
    """
    _, alpha, _, _ = compute_transform_params(market)

    if option_type == "call":
        payoff = np.maximum(np.exp(x_grid) - 1, 0)
    elif option_type == "put":
        payoff = np.maximum(1 - np.exp(x_grid), 0)
    else:
        raise ValueError(f"option_type must be 'call' or 'put', got '{option_type}'")

    raw_exponent = -alpha * x_grid
    exponent = np.clip(raw_exponent, -_MAX_EXP, _MAX_EXP)
    with np.errstate(over="ignore", invalid="ignore"):
        u0 = np.exp(exponent) * payoff
    u0 = np.nan_to_num(u0, nan=0.0, posinf=0.0, neginf=0.0)

    # Overflow taper: when |alpha*x| > 10, the initial condition starts
    # losing significance; by |alpha*x| = 20 it is fully tapered to zero.
    # These thresholds are conservative (exp(700) ~ float64 max = exp(709)).
    _OVERFLOW_ONSET = 10.0
    _OVERFLOW_FULL = 2.0 * _OVERFLOW_ONSET
    abs_exp = np.abs(raw_exponent)
    overflow_mask = abs_exp > _OVERFLOW_ONSET
    if np.any(overflow_mask):
        overflow_taper = np.ones_like(x_grid)
        t = np.minimum(
            (abs_exp[overflow_mask] - _OVERFLOW_ONSET)
            / (_OVERFLOW_FULL - _OVERFLOW_ONSET),
            1.0,
        )
        overflow_taper[overflow_mask] = 0.5 * (1 + np.cos(np.pi * t))
        with np.errstate(over="ignore", invalid="ignore"):
            u0 = np.where(overflow_taper > 0, u0 * overflow_taper, 0.0)
        u0 = np.nan_to_num(u0, nan=0.0, posinf=0.0, neginf=0.0)

    taper = make_taper(x_grid, config)
    return np.where(taper > 0, u0 * taper, 0.0)


def heat_to_bs_price(
    u: np.ndarray, x_grid: np.ndarray, market: MarketParams
) -> np.ndarray:
    """Convert heat solution u(x, t_eff) back to BS prices V(S) on the full grid."""
    _, alpha, beta, t_eff = compute_transform_params(market)
    exponent = np.clip(alpha * x_grid + beta * t_eff, -_MAX_EXP, _MAX_EXP)
    return market.K * np.exp(exponent) * u


def extract_price_at_spot(
    u: np.ndarray, x_grid: np.ndarray, market: MarketParams
) -> float:
    """Extract option price at spot by interpolating the heat solution."""
    _, alpha, beta, t_eff = compute_transform_params(market)
    x0 = np.log(market.S / market.K)
    u_at_x0 = float(np.interp(x0, x_grid, u))
    exponent = np.clip(alpha * x0 + beta * t_eff, -_MAX_EXP, _MAX_EXP)
    return market.K * np.exp(exponent) * u_at_x0


def bs_exact_price(market: MarketParams, option_type: str = "call") -> float:
    """Exact Black-Scholes-Merton formula with continuous dividend yield q."""
    sqrt_t = np.sqrt(market.T)
    d1 = (
        np.log(market.S / market.K)
        + (market.r - market.q + market.sigma**2 / 2) * market.T
    ) / (market.sigma * sqrt_t)
    d2 = d1 - market.sigma * sqrt_t

    if option_type == "call":
        return float(
            market.S * np.exp(-market.q * market.T) * norm.cdf(d1)
            - market.K * np.exp(-market.r * market.T) * norm.cdf(d2)
        )
    if option_type == "put":
        return float(
            market.K * np.exp(-market.r * market.T) * norm.cdf(-d2)
            - market.S * np.exp(-market.q * market.T) * norm.cdf(-d1)
        )
    raise ValueError(f"option_type must be 'call' or 'put', got '{option_type}'")

