"""Finite difference (PDE) valuation implementations.

This module follows the same structure as other valuation modules:
private implementation classes that plug into OptionValuation.

Current scope
-------------
PDE via finite differences for vanilla European and American call/put:
- time stepping: implicit, explicit, or Crank–Nicolson
- optional Rannacher smoothing for Crank–Nicolson
- spatial grids: spot or log-spot
- American handling: intrinsic projection or Gauss-Seidel/PSOR
"""

from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING

import logging
import math
import datetime as dt

import numpy as np

# optional acceleration via numba if available
try:
    from numba import njit as _njit
except ModuleNotFoundError:  # pragma: no cover

    def _njit(*args, **kwargs):  # type: ignore[misc]
        """Identity decorator when numba is not installed."""
        if args and callable(args[0]):
            return args[0]
        return lambda fn: fn


from ..enums import DayCountConvention, PDEEarlyExercise, PDEMethod, PDESpaceGrid, OptionType
from ..rates import DiscountCurve
from ..utils import calculate_year_fraction, log_timing
from ..exceptions import (
    ConfigurationError,
    StabilityError,
    UnsupportedFeatureError,
    ValidationError,
)
from .params import PDEParams

if TYPE_CHECKING:
    from .core import OptionValuation, UnderlyingData


logger = logging.getLogger(__name__)


def _solve_tridiagonal_thomas(
    lower: np.ndarray,
    diag: np.ndarray,
    upper: np.ndarray,
    rhs: np.ndarray,
) -> np.ndarray:
    """Solve a tridiagonal system Ax = rhs via the Thomas algorithm.

    A has:
      - lower: subdiagonal (length n-1)  -> A[i, i-1]
      - diag:  main diagonal (length n)  -> A[i, i]
      - upper: superdiagonal (length n-1)-> A[i, i+1]

    Parameters
    ----------
    lower
        Subdiagonal coefficients.
    diag
        Main diagonal coefficients.
    upper
        Superdiagonal coefficients.
    rhs
        Right-hand-side vector.

    Returns
    -------
    np.ndarray
        Solution vector of the tridiagonal system.
    """
    n = diag.size
    if rhs.size != n:
        raise ValidationError("rhs length must match diag length")
    if lower.size != n - 1 or upper.size != n - 1:
        raise ValidationError("lower/upper must have length n-1")

    # Copy to avoid mutating inputs
    a: np.ndarray = lower.astype(float, copy=True)
    d: np.ndarray = diag.astype(float, copy=True)
    c: np.ndarray = upper.astype(float, copy=True)
    y: np.ndarray = rhs.astype(float, copy=True)

    # Forward elimination
    for i in range(1, n):
        w = a[i - 1] / d[i - 1]
        d[i] -= w * c[i - 1]
        y[i] -= w * y[i - 1]

    # Back substitution
    x: np.ndarray = np.empty(n, dtype=float)
    x[-1] = y[-1] / d[-1]
    for i in range(n - 2, -1, -1):
        x[i] = (y[i] - c[i] * x[i + 1]) / d[i]
    return x


def _build_tau_grid(
    time_to_maturity: float,
    time_steps: int,
    dividend_taus: list[float],
) -> np.ndarray:
    """Build a tau (time remaining) grid that includes dividend dates."""
    base = np.linspace(0.0, time_to_maturity, time_steps + 1)
    if not dividend_taus:
        return base
    grid = np.unique(np.concatenate([base, np.array(dividend_taus, dtype=float)]))
    grid.sort()
    return grid


def _dividend_tau_schedule(
    *,
    discrete_dividends: Sequence[tuple[dt.datetime, float]],
    pricing_date: dt.datetime,
    maturity: dt.datetime,
    day_count_convention: DayCountConvention,
) -> list[tuple[float, float]]:
    """Return list of (tau, amount) for dividends between pricing_date and maturity.

    The range is closed: ``0.0 <= tau <= ttm``.  Boundary values (tau=0 for
    maturity-date dividends, tau=ttm for pricing-date dividends) are included
    so that ``_vanilla_fd_core`` can apply them as special-case jumps.
    """
    if not discrete_dividends:
        return []

    ttm = calculate_year_fraction(
        pricing_date,
        maturity,
        day_count_convention=day_count_convention,
    )
    schedule: dict[float, float] = {}
    for ex_date, amount in discrete_dividends:
        if pricing_date <= ex_date <= maturity:
            t = calculate_year_fraction(
                pricing_date,
                ex_date,
                day_count_convention=day_count_convention,
            )
            tau = ttm - t
            key = round(float(tau), 12)
            schedule[key] = schedule.get(key, 0.0) + float(amount)
    return sorted(schedule.items())


def _apply_dividend_jump(
    values: np.ndarray,
    grid: np.ndarray,
    amount: float,
    *,
    space_grid: PDESpaceGrid,
) -> None:
    """Apply the cash dividend jump condition V(S,t^-)=V(S-D,t^+)."""
    if amount == 0.0:
        return
    if space_grid is PDESpaceGrid.LOG_SPOT:
        spot_grid = np.exp(grid)
    else:
        spot_grid = grid
    shifted = np.interp(
        spot_grid - amount,
        spot_grid,
        values,
        left=values[0],
        right=values[-1],
    )
    values[:] = shifted


def _boundary_values(
    *,
    option_type: OptionType,
    strike: float,
    smax: float,
    df_tT: float,
    dq_tT: float,
    early_exercise: bool,
) -> tuple[float, float]:
    if option_type is OptionType.PUT:
        left = strike if early_exercise else strike * df_tT
        right = 0.0
    else:
        left = 0.0
        continuation = smax * dq_tT - strike * df_tT
        intrinsic = smax - strike
        right = max(continuation, intrinsic) if early_exercise else max(continuation, 0.0)
    return float(left), float(right)


def _build_log_grid(
    *,
    spot: float,
    strike: float,
    time_to_maturity: float,
    volatility: float,
    smax_mult: float,
    spot_steps: int,
    time_steps: int,
) -> tuple[np.ndarray, np.ndarray, float]:
    """Build log-spot grid using Hull's dz heuristic when possible."""
    d_tau = time_to_maturity / time_steps
    dz_hull = volatility * np.sqrt(3.0 * d_tau)

    smax = float(smax_mult * max(spot, strike))
    smin = float(max(max(spot, strike) / smax_mult, 1.0e-8))
    zmin_target = np.log(smin)
    zmax_target = np.log(smax)

    grid_width = spot_steps * dz_hull
    if (zmax_target - zmin_target) > grid_width:
        dz = (zmax_target - zmin_target) / spot_steps
        zmin = zmin_target
        zmax = zmax_target
    else:
        dz = dz_hull
        center = np.log(spot)
        zmin = center - 0.5 * grid_width
        zmax = center + 0.5 * grid_width
        if zmin > zmin_target:
            shift = zmin_target - zmin
            zmin += shift
            zmax += shift
        if zmax < zmax_target:
            shift = zmax_target - zmax
            zmin += shift
            zmax += shift

    Z = np.linspace(zmin, zmax, spot_steps + 1)
    S = np.exp(Z)
    return Z, S, dz


def _spot_operator_coeffs(
    *,
    spot_values: np.ndarray,
    dS: float,
    risk_free_rate: float,
    dividend_rate: float,
    volatility: float,
    hull_discounting: bool = False,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Spatial operator coefficients on the spot grid.

    When *hull_discounting* is True (Hull's explicit scheme), the rV
    term is excluded from beta and instead applied as an implicit
    divisor ``1 / (1 + r * dt)`` in the time-step function.
    """
    diffusion = (volatility**2) * (spot_values**2) / (dS**2)
    drift = (risk_free_rate - dividend_rate) * spot_values / dS
    gamma = 0.5 * (diffusion - drift)
    beta = -diffusion if hull_discounting else -(diffusion + risk_free_rate)
    alpha = 0.5 * (diffusion + drift)
    return gamma, beta, alpha


def _log_operator_coeffs(
    *,
    dz: float,
    risk_free_rate: float,
    dividend_rate: float,
    volatility: float,
    hull_discounting: bool = False,
    size: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Spatial operator coefficients on the log-spot grid.

    Returns constant (Toeplitz) arrays of length *size*, matching the
    signature of ``_spot_operator_coeffs`` so callers can treat both
    grids uniformly.

    When *hull_discounting* is True (Hull's explicit scheme), r is
    excluded from beta.
    """
    mu = risk_free_rate - dividend_rate - 0.5 * volatility**2
    diffusion = (volatility**2) / (dz**2)
    drift = mu / dz
    gamma = np.full(size, 0.5 * (diffusion - drift))
    beta = np.full(size, -diffusion if hull_discounting else -(diffusion + risk_free_rate))
    alpha = np.full(size, 0.5 * (diffusion + drift))
    return gamma, beta, alpha


def _scaled_operator_coeffs(
    *,
    gamma: np.ndarray,
    beta: np.ndarray,
    alpha: np.ndarray,
    d_tau: float,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    a = -d_tau * gamma
    b = -d_tau * beta
    c = -d_tau * alpha
    return a, b, c


# ---------------------------------------------------------------------------
# Time-step helpers (extracted from _vanilla_fd_core for readability)
# ---------------------------------------------------------------------------


def _explicit_step(
    V_prev: np.ndarray,
    j: np.ndarray,
    a: np.ndarray,
    b: np.ndarray,
    c: np.ndarray,
    left: float,
    right: float,
    intrinsic: np.ndarray | None,
    *,
    r_dt: float = 0.0,
) -> np.ndarray:
    """Explicit (forward-Euler) time step.

    When *r_dt* > 0 (Hull's explicit scheme), the interior update is
    divided by ``(1 + r_dt)`` to apply implicit discounting of the rV
    term.
    """
    V_new = V_prev.copy()
    interior = -a * V_prev[j - 1] + (1.0 - b) * V_prev[j] - c * V_prev[j + 1]
    V_new[j] = interior / (1.0 + r_dt)
    V_new[0] = left
    V_new[-1] = right
    if intrinsic is not None:
        V_new[:] = np.maximum(V_new, intrinsic)
    return V_new


@_njit(cache=True)
def _psor_core(
    x: np.ndarray,
    exercise_j: np.ndarray,
    rhs: np.ndarray,
    lower: np.ndarray,
    diag: np.ndarray,
    upper: np.ndarray,
    V_left: float,
    V_right: float,
    omega: float,
    tol: float,
    max_iter: int,
) -> tuple[np.ndarray, int]:
    """Numba-accelerated Projected SOR (Gauss-Seidel with overrelaxation)."""
    n = x.shape[0]
    for k in range(n):
        if x[k] < exercise_j[k]:
            x[k] = exercise_j[k]
    iter_used = max_iter
    for iter_idx in range(max_iter):
        max_diff = 0.0
        for k in range(n):
            left_val = x[k - 1] if k > 0 else V_left
            right_val = x[k + 1] if k < n - 1 else V_right
            gs = (rhs[k] - lower[k] * left_val - upper[k] * right_val) / diag[k]
            old = x[k]
            new = old + omega * (gs - old)
            if new < exercise_j[k]:
                new = exercise_j[k]
            x[k] = new
            diff = new - old if new > old else old - new
            if diff > max_diff:
                max_diff = diff
        if max_diff < tol:
            iter_used = iter_idx + 1
            break
    return x, iter_used


def _psor_solve(
    x: np.ndarray,
    exercise_j: np.ndarray,
    rhs: np.ndarray,
    lower: np.ndarray,
    diag: np.ndarray,
    upper: np.ndarray,
    V_left: float | np.floating,
    V_right: float | np.floating,
    omega: float,
    tol: float,
    max_iter: int,
) -> tuple[np.ndarray, int]:
    """Projected SOR (Gauss-Seidel with overrelaxation) for American exercise.

    Returns the updated interior values and the number of iterations used.
    Delegates to a Numba-JIT compiled inner loop when numba is available.
    """
    return _psor_core(
        x,
        exercise_j,
        rhs,
        lower,
        diag,
        upper,
        float(V_left),
        float(V_right),
        omega,
        tol,
        max_iter,
    )


def _implicit_cn_step(
    V_prev: np.ndarray,
    V: np.ndarray,
    j: np.ndarray,
    a: np.ndarray,
    b: np.ndarray,
    c: np.ndarray,
    left: float,
    right: float,
    method: PDEMethod,
    intrinsic: np.ndarray | None,
    american_solver: PDEEarlyExercise,
    omega: float | None,
    tol: float | None,
    max_iter: int | None,
) -> tuple[np.ndarray, int | None]:
    """One implicit or Crank-Nicolson time step (with optional early exercise).

    Returns the updated V array and the PSOR iteration count (None if PSOR was not used).
    """
    if method is PDEMethod.CRANK_NICOLSON:
        a = a * 0.5
        b = b * 0.5
        c = c * 0.5

    # Tridiagonal system (I - θ·dt·L)x = rhs, diagonals are (a, 1+b, c)
    diag = 1.0 + b

    if method is PDEMethod.IMPLICIT:
        rhs = V_prev[j].copy()
    else:
        rhs = -a * V_prev[j - 1] + (1.0 - b) * V_prev[j] - c * V_prev[j + 1]

    V[0] = left
    V[-1] = right

    rhs_adj = rhs.copy()
    rhs_adj[0] -= a[0] * V[0]
    rhs_adj[-1] -= c[-1] * V[-1]

    x = _solve_tridiagonal_thomas(a[1:], diag, c[:-1], rhs_adj)
    psor_iters: int | None = None

    if intrinsic is None:
        V[j] = x
    else:
        exercise_j = intrinsic[j]
        if american_solver is PDEEarlyExercise.INTRINSIC:
            V[j] = np.maximum(x, exercise_j)
        else:
            x, psor_iters = _psor_solve(
                x,
                exercise_j,
                rhs,
                a,
                diag,
                c,
                float(V[0]),
                float(V[-1]),
                float(omega),
                float(tol),
                int(max_iter),
            )
            V[j] = x

    return V, psor_iters


def _validate_fd_inputs(
    *,
    option_type: OptionType,
    time_to_maturity: float,
    spot_steps: int,
    time_steps: int,
    volatility: float,
    discount_curve: DiscountCurve,
    early_exercise: bool,
    method: PDEMethod,
    american_solver: PDEEarlyExercise,
    omega: float | None,
    tol: float | None,
    max_iter: int | None,
) -> None:
    """Validate FD PDE inputs before grid construction."""
    if option_type not in (OptionType.CALL, OptionType.PUT):
        raise UnsupportedFeatureError("FD PDE valuation supports only vanilla CALL/PUT.")
    if time_to_maturity <= 0:
        raise ValidationError("time_to_maturity must be positive")
    if spot_steps < 3:
        raise ValidationError("spot_steps must be >= 3")
    if time_steps < 1:
        raise ValidationError("time_steps must be >= 1")
    if volatility <= 0:
        raise ValidationError("volatility must be positive")
    if discount_curve is None:
        raise ValidationError("discount_curve is required for PDE valuation")
    if early_exercise and american_solver is PDEEarlyExercise.GAUSS_SEIDEL:
        if omega is None or tol is None or max_iter is None:
            raise ValidationError(
                "PSOR params (omega/tol/max_iter) are required for early exercise"
            )
    if method in (PDEMethod.EXPLICIT, PDEMethod.EXPLICIT_HULL) and (
        american_solver is PDEEarlyExercise.GAUSS_SEIDEL
    ):
        raise UnsupportedFeatureError("GAUSS_SEIDEL is not supported with explicit time stepping")


def _check_explicit_spot_stability(
    *,
    tau_grid: np.ndarray,
    volatility: float,
    smax: float,
    dS: float,
    time_to_maturity: float,
    discount_curve: DiscountCurve,
    dividend_curve: DiscountCurve | None,
    hull_discounting: bool,
) -> None:
    r"""CFL-style stability checks for an explicit scheme on a uniform spot grid.

    The dominant restriction comes from the diffusion term::

        dt <= dS² / (σ² S_max²)

    When *hull_discounting* is ``False`` (pure explicit), the reaction
    term :math:`-rV` is discretised explicitly too, giving the tighter
    bound::

        dt <= 1 / (σ² S_max² / dS² + r_max)

    With Hull implicit discounting, ``r`` is handled implicitly via the
    :math:`1/(1 + r \Delta t)` divisor so only diffusion matters.

    A secondary check enforces the central-difference stencil monotonicity
    condition :math:`\alpha, \gamma \ge 0`, i.e. diffusion dominates drift
    at the worst-case grid node (``S_max``).

    Raises ``StabilityError`` if any condition is violated.
    """
    if tau_grid.size < 2:
        return

    dt_steps = np.diff(tau_grid).astype(float)
    max_dt = float(np.max(dt_steps))
    if max_dt <= 0.0:
        return

    # Forward rates for each tau step (calendar times: t = T - tau)
    t_prev = time_to_maturity - tau_grid[:-1]
    t_next = time_to_maturity - tau_grid[1:]
    r_steps = np.array(
        [discount_curve.forward_rate(t1, t0) for t0, t1 in zip(t_prev, t_next)],
        dtype=float,
    )
    if dividend_curve is not None:
        q_steps = np.array(
            [dividend_curve.forward_rate(t1, t0) for t0, t1 in zip(t_prev, t_next)],
            dtype=float,
        )
    else:
        q_steps = np.zeros_like(r_steps)

    rq_abs_max = float(np.max(np.abs(r_steps - q_steps)))

    # (A) Diffusion CFL bound
    diffusion_max = (volatility**2) * (smax**2) / (dS**2)

    if hull_discounting:
        # r handled implicitly — only diffusion constrains dt
        dt_max = 1.0 / diffusion_max
        mode = "Hull implicit discounting"
    else:
        # Pure explicit: tighten for the reaction term -rV
        r_pos = max(float(np.max(r_steps)), 0.0)
        dt_max = (
            1.0 / (diffusion_max + r_pos) if (diffusion_max + r_pos) > 0.0 else 1.0 / diffusion_max
        )
        mode = "pure explicit discounting"

    if max_dt > dt_max:
        min_steps = int(math.ceil(time_to_maturity / dt_max))
        raise StabilityError(
            f"Explicit spot-grid scheme likely unstable (CFL violation, {mode}). "
            f"max_dt={max_dt:.4g} exceeds dt_max={dt_max:.4g}. "
            f"Increase time_steps to >= {min_steps}, or use log-spot/implicit/CN."
        )

    # (B) Drift monotonicity: central-difference stencil requires
    #     diffusion >= |drift| at worst-case S for alpha, gamma >= 0.
    drift_max = rq_abs_max * smax / dS
    if drift_max > diffusion_max:
        raise StabilityError(
            "Explicit spot-grid scheme likely unstable/oscillatory: drift dominates "
            "diffusion in the central-difference stencil. "
            f"diffusion_max={diffusion_max:.4g}, drift_max={drift_max:.4g}. "
            "Refine the spot grid (smaller dS), reduce smax, or use log-spot/implicit/CN."
        )


def _build_time_step_schedule(
    tau_grid: np.ndarray,
    method: PDEMethod,
    rannacher_steps: int,
) -> list[tuple[float, float, PDEMethod]]:
    """Build the time-step schedule from the tau grid.

    For Crank-Nicolson with Rannacher smoothing (Pooley-Vetzal-Forsyth 2003),
    the first *rannacher_steps* intervals are each replaced by two implicit
    (backward Euler) half-steps of size d_tau/2.  This damps payoff
    non-smoothness while preserving the overall time-grid structure.
    For all other methods the schedule is a straightforward pass-through.
    """
    steps: list[tuple[float, float, PDEMethod]] = []
    for n in range(1, tau_grid.size):
        tau_start = float(tau_grid[n - 1])
        tau_end = float(tau_grid[n])
        if method is PDEMethod.CRANK_NICOLSON and rannacher_steps > 0 and n <= rannacher_steps:
            tau_mid = 0.5 * (tau_start + tau_end)
            steps.append((tau_start, tau_mid, PDEMethod.IMPLICIT))
            steps.append((tau_mid, tau_end, PDEMethod.IMPLICIT))
        else:
            steps.append((tau_start, tau_end, method))
    return steps


def _vanilla_fd_core(
    *,
    spot: float,
    strike: float,
    time_to_maturity: float,
    volatility: float,
    discount_curve: DiscountCurve,
    dividend_curve: DiscountCurve | None,
    dividend_schedule: list[tuple[float, float]] | None,
    option_type: OptionType,
    smax_mult: float,
    spot_steps: int,
    time_steps: int,
    early_exercise: bool,
    method: PDEMethod,
    rannacher_steps: int,
    space_grid: PDESpaceGrid,
    american_solver: PDEEarlyExercise,
    omega: float | None = None,
    tol: float | None = None,
    max_iter: int | None = None,
) -> tuple[float, np.ndarray, np.ndarray, np.ndarray, float]:
    """Core finite-difference solver for vanilla option valuation.

    Returns
    -------
    tuple[float, np.ndarray, np.ndarray, np.ndarray, float]
        ``(price, spot_grid, V_final, V_prev, last_dtau)``
        where *V_prev* is the value slice one time step before pricing
        time and *last_dtau* is the size of that step (both needed for
        theta extraction from the grid).
    """
    _validate_fd_inputs(
        option_type=option_type,
        time_to_maturity=time_to_maturity,
        spot_steps=spot_steps,
        time_steps=time_steps,
        volatility=volatility,
        discount_curve=discount_curve,
        early_exercise=early_exercise,
        method=method,
        american_solver=american_solver,
        omega=omega,
        tol=tol,
        max_iter=max_iter,
    )

    smax = float(smax_mult * max(spot, strike))
    if space_grid is PDESpaceGrid.SPOT:
        grid = np.linspace(0.0, smax, spot_steps + 1)
        S = grid
        dS = smax / spot_steps
    else:
        grid, S, dz = _build_log_grid(
            spot=spot,
            strike=strike,
            time_to_maturity=time_to_maturity,
            volatility=volatility,
            smax_mult=smax_mult,
            spot_steps=spot_steps,
            time_steps=time_steps,
        )

    j = np.arange(1, spot_steps)  # interior indices 1..M-1

    # Standard terminal payoff at maturity
    if option_type is OptionType.PUT:
        payoff = np.maximum(strike - S, 0.0)
    else:
        payoff = np.maximum(S - strike, 0.0)

    V = payoff.copy()  # V at tau=0 (maturity)
    intrinsic = payoff if early_exercise else None

    schedule = dividend_schedule or []
    dividend_map = {round(tau, 12): amount for tau, amount in schedule}

    # Maturity-date dividend (tau=0): apply as an immediate jump
    # V(S,T⁻)=V(S-D,T⁺) right after setting the terminal condition.
    mat_div = dividend_map.pop(0.0, None)
    if mat_div is not None:
        _apply_dividend_jump(V, grid, mat_div, space_grid=space_grid)
        if early_exercise:
            V[:] = np.maximum(V, payoff)

    # Pricing-date dividend (tau=ttm): will be applied as a spot shift
    # at interpolation time.  Remove from the map so it's not applied
    # as a mid-grid jump during time-stepping.
    ttm_key = round(time_to_maturity, 12)
    pricing_div = dividend_map.pop(ttm_key, None)

    dividend_taus = list(dividend_map.keys())
    tau_grid = _build_tau_grid(time_to_maturity, time_steps, dividend_taus)

    if method in (PDEMethod.EXPLICIT, PDEMethod.EXPLICIT_HULL) and space_grid is PDESpaceGrid.SPOT:
        _check_explicit_spot_stability(
            tau_grid=tau_grid,
            volatility=volatility,
            smax=smax,
            dS=dS,
            time_to_maturity=time_to_maturity,
            discount_curve=discount_curve,
            dividend_curve=dividend_curve,
            hull_discounting=method is PDEMethod.EXPLICIT_HULL,
        )

    # March forward in tau: 0 -> T (equivalently backward in calendar time)
    df_0T = float(discount_curve.df(time_to_maturity))  # P(0,T)
    if dividend_curve is not None:
        dq_0T = float(dividend_curve.df(time_to_maturity))  # Dq(0,T)
    else:
        dq_0T = None

    psor_steps = 0
    psor_total_iters = 0
    psor_max_iters = 0
    psor_not_converged = 0

    steps = _build_time_step_schedule(tau_grid, method, rannacher_steps)

    V_prev = V.copy()
    last_dtau = 0.0

    for tau_prev, tau_curr, method_used in steps:
        d_tau = tau_curr - tau_prev
        t_prev = time_to_maturity - tau_prev
        t_curr = time_to_maturity - tau_curr

        r = float(discount_curve.forward_rate(t_curr, t_prev))
        if dividend_curve is not None:
            q = float(dividend_curve.forward_rate(t_curr, t_prev))
        else:
            q = 0.0

        hull_discounting = method_used is PDEMethod.EXPLICIT_HULL

        if space_grid is PDESpaceGrid.SPOT:
            gamma, beta, alpha = _spot_operator_coeffs(
                spot_values=S[1:-1],
                dS=dS,
                risk_free_rate=r,
                dividend_rate=q,
                volatility=volatility,
                hull_discounting=hull_discounting,
            )
        else:
            gamma, beta, alpha = _log_operator_coeffs(
                dz=dz,
                risk_free_rate=r,
                dividend_rate=q,
                volatility=volatility,
                hull_discounting=hull_discounting,
                size=spot_steps - 1,
            )

        df_0t = float(discount_curve.df(t_curr))
        df_tT: float = df_0T / df_0t
        if dividend_curve is not None:
            dq_0t = float(dividend_curve.df(t_curr))
            dq_tT: float = dq_0T / dq_0t  # type: ignore[operator]
        else:
            dq_tT = 1.0

        left, right = _boundary_values(
            option_type=option_type,
            strike=strike,
            smax=float(S[-1]),
            df_tT=df_tT,
            dq_tT=dq_tT,
            early_exercise=early_exercise,
        )

        V_prev = V.copy()
        last_dtau = d_tau

        a, b, c = _scaled_operator_coeffs(gamma=gamma, beta=beta, alpha=alpha, d_tau=d_tau)

        intrinsic_for_step = intrinsic if early_exercise else None

        if method_used in (PDEMethod.EXPLICIT, PDEMethod.EXPLICIT_HULL):
            V = _explicit_step(
                V_prev,
                j,
                a,
                b,
                c,
                left,
                right,
                intrinsic_for_step,
                r_dt=r * d_tau if hull_discounting else 0.0,
            )
        else:
            V, psor_iters = _implicit_cn_step(
                V_prev,
                V,
                j,
                a,
                b,
                c,
                left,
                right,
                method_used,
                intrinsic_for_step,
                american_solver,
                omega,
                tol,
                max_iter,
            )
            if psor_iters is not None:
                psor_steps += 1
                psor_total_iters += psor_iters
                psor_max_iters = max(psor_max_iters, psor_iters)
                if psor_iters == int(max_iter):
                    psor_not_converged += 1

        # Apply discrete dividend jump at tau if needed
        if dividend_map:
            amount = dividend_map.get(round(tau_curr, 12))
            if amount is not None:
                _apply_dividend_jump(V, grid, amount, space_grid=space_grid)
                if early_exercise:
                    V[:] = np.maximum(V, intrinsic)

    if psor_steps > 0:
        avg_iters = psor_total_iters / psor_steps
        logger.debug(
            "PDE PSOR steps=%d avg_iters=%.2f max_iters=%d not_converged=%d",
            psor_steps,
            avg_iters,
            psor_max_iters,
            psor_not_converged,
        )

    # Apply pricing-date dividend: the input spot is cum-dividend, so
    # interpolate at S₀ − D to get the ex-dividend option value.
    interp_spot = spot - pricing_div if pricing_div is not None else spot
    price = np.interp(interp_spot, S, V)
    return price, S, V, V_prev, last_dtau


class _FDGridGreeksMixin:
    """Mixin providing delta/gamma/theta extracted from the PDE solution grid.

    Subclasses must define ``_solve()`` returning
    ``(price, S, V, V_prev, last_dtau)``.
    """

    valuation_ctx: OptionValuation
    underlying: UnderlyingData

    def _solve(
        self,
    ) -> tuple[float, np.ndarray, np.ndarray, np.ndarray, float]: ...

    def _grid_greeks_data(
        self,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, float, int]:
        """Run the PDE solve and locate the spot node.

        Returns
        -------
        S, V, V_prev, last_dtau, j
            The spot grid, value vector, previous-step value vector,
            last time-step size, and the spot-grid index closest to
            the current spot.
        """
        _, S, V, V_prev, last_dtau = self._solve()
        spot = float(self.underlying.initial_value)
        j = int(np.searchsorted(S, spot))
        # Clamp to interior so the 3-point stencil is valid.
        j = max(1, min(j, len(S) - 2))
        return S, V, V_prev, last_dtau, j

    def delta(self) -> float:
        r"""Grid delta via central differences at the spot node.

        .. math::

            \Delta \approx \frac{V_{j+1} - V_{j-1}}{S_{j+1} - S_{j-1}}
        """
        S, V, _, _, j = self._grid_greeks_data()
        return float((V[j + 1] - V[j - 1]) / (S[j + 1] - S[j - 1]))

    def gamma(self) -> float:
        r"""Grid gamma via the standard second-order stencil.

        .. math::

            \Gamma \approx \frac{V_{j+1} - 2V_j + V_{j-1}}
                               {\tfrac12(S_{j+1} - S_{j-1}) \cdot (S_{j+1} - S_{j-1})/2}

        For a uniform grid this reduces to
        :math:`(V_{j+1} - 2V_j + V_{j-1}) / h^2`.
        """
        S, V, _, _, j = self._grid_greeks_data()
        h_up = S[j + 1] - S[j]
        h_dn = S[j] - S[j - 1]
        return float(
            2.0
            * (V[j + 1] * h_dn + V[j - 1] * h_up - V[j] * (h_up + h_dn))
            / (h_up * h_dn * (h_up + h_dn))
        )

    def theta(self) -> float:
        r"""Grid theta via backward difference between the last two time
        slices.

        .. math::

            \Theta \approx \frac{V^{n}(S_0) - V^{n-1}(S_0)}{\Delta t}

        Returned per **calendar day** (divided by 365).
        """
        S, V, V_prev, last_dtau, j = self._grid_greeks_data()
        if last_dtau <= 0.0:
            return 0.0
        theta_annual = (V_prev[j] - V[j]) / last_dtau
        return float(theta_annual / 365.0)


class _FDValuationBase(_FDGridGreeksMixin):
    """Base class for European/American FD valuation."""

    _early_exercise: bool = False

    def __init__(self, valuation_ctx: OptionValuation) -> None:
        self.valuation_ctx = valuation_ctx
        self.underlying = valuation_ctx.underlying  # type: ignore[assignment]

    def solve(self) -> tuple[float, np.ndarray, np.ndarray]:
        """Compute the full FD solution on the spot grid at pricing time."""
        pv, S, V, *_ = self._solve()
        return pv, S, V

    def _solve(self) -> tuple[float, np.ndarray, np.ndarray, np.ndarray, float]:
        """Run the PDE finite-difference solve."""
        params = self.valuation_ctx.params
        if not isinstance(params, PDEParams):
            raise ConfigurationError("PDE valuation requires PDEParams on OptionValuation")

        if self._early_exercise:
            logger.debug(
                "PDE American method=%s grid=%s solver=%s spot_steps=%d time_steps=%d",
                params.method.value,
                params.space_grid.value,
                params.american_solver.value,
                params.spot_steps,
                params.time_steps,
            )
        else:
            logger.debug(
                "PDE European method=%s grid=%s spot_steps=%d time_steps=%d",
                params.method.value,
                params.space_grid.value,
                params.spot_steps,
                params.time_steps,
            )

        spot = float(self.underlying.initial_value)
        strike = self.valuation_ctx.strike
        volatility = float(self.underlying.volatility)
        discount_curve = self.valuation_ctx.discount_curve
        dividend_curve = self.underlying.dividend_curve
        discrete_dividends = self.underlying.discrete_dividends

        time_to_maturity = calculate_year_fraction(
            self.valuation_ctx.pricing_date,
            self.valuation_ctx.maturity,
            day_count_convention=self.valuation_ctx.day_count_convention,
        )

        dividend_schedule = _dividend_tau_schedule(
            discrete_dividends=discrete_dividends,
            pricing_date=self.valuation_ctx.pricing_date,
            maturity=self.valuation_ctx.maturity,
            day_count_convention=self.valuation_ctx.day_count_convention,
        )

        smax_mult = float(params.smax_mult)
        spot_steps = int(params.spot_steps)
        time_steps = int(params.time_steps)

        return _vanilla_fd_core(
            spot=spot,
            strike=float(strike),
            time_to_maturity=float(time_to_maturity),
            volatility=volatility,
            discount_curve=discount_curve,
            dividend_curve=dividend_curve,
            dividend_schedule=dividend_schedule,
            option_type=self.valuation_ctx.option_type,
            smax_mult=smax_mult,
            spot_steps=spot_steps,
            time_steps=time_steps,
            early_exercise=self._early_exercise,
            method=params.method,
            rannacher_steps=int(params.rannacher_steps),
            space_grid=params.space_grid,
            american_solver=params.american_solver
            if self._early_exercise
            else PDEEarlyExercise.INTRINSIC,
            omega=float(params.omega) if self._early_exercise else None,
            tol=float(params.tol) if self._early_exercise else None,
            max_iter=int(params.max_iter) if self._early_exercise else None,
        )

    def present_value(self) -> float:
        """Return present value from the PDE solve."""
        params = self.valuation_ctx.params
        if not isinstance(params, PDEParams):
            raise ConfigurationError("PDE valuation requires PDEParams on OptionValuation")
        label = "PDE American" if self._early_exercise else "PDE European"
        with log_timing(logger, f"{label} present_value", params.log_timings):
            pv, *_ = self._solve()
        return float(pv)


class _FDEuropeanValuation(_FDValuationBase):
    """European option valuation using PDE finite differences."""

    _early_exercise = False


class _FDAmericanValuation(_FDValuationBase):
    """American option valuation using PDE finite differences."""

    _early_exercise = True
