"""greeks_package.fdm – American option pricing via Finite Difference Method (Crank–Nicolson + PSOR)

This submodule prices **American call and put options** (early exercise allowed). Use
``option_type="call"`` or ``option_type="put"`` in ``FDM(...)``. It is **not vectorized**:
it prices a single option per call, returning a result object. It does not operate on
pandas DataFrames like the main package Greeks.

**Recommended:** Use ``gp.FDM(...)`` after ``import greeks_package as gp``, or
``from greeks_package import FDM``.

Usage::

    import greeks_package as gp
    res = gp.FDM(685.74, 685.0, 0.1743, 0.0402, 321/365)
    res.price
    res.delta
    res.gamma
    res.theta_day
    res.S   # full asset price grid
    res.V   # full option value vector

    res = gp.FDM(685.74, 685.0, 0.1743, 0.0402, 321/365, option_type="put")  # American put
    res = gp.FDM(685.74, 685.0, 0.1743, 0.0402, 321/365, bs=True)              # with BS benchmark
"""

import numpy as np
from scipy.stats import norm


class _FDMResult:
    """
    Return object for FDM(). Access results as attributes:
        res = FDM(...)
        res.price
        res.delta
        res.gamma
        res.theta_day
        res.S
        res.V
    """
    def __init__(self, price, delta, gamma, theta_day, S, V):
        self.price     = price
        self.delta     = delta
        self.gamma     = gamma
        self.theta_day = theta_day
        self.S         = S
        self.V         = V

    def __repr__(self):
        return (f"FDMResult(price={self.price:.4f}, delta={self.delta:.6f}, "
                f"gamma={self.gamma:.6f}, theta_day={self.theta_day:.6f})")


def FDM(S0, K, sigma, r, T,
        NAS=200, NTS=2000, S_max_mult=4.0, rannacher=4,
        omega=1.2, tol=1e-8, max_iter=5000,
        bs=False, option_type="call", silent=False):
    """
    Price an American call or put option using Crank-Nicolson + PSOR.

    This function is for **American options only** (early exercise handled via
    PSOR). Use ``option_type="call"`` for calls, ``option_type="put"`` for puts.
    It does not support European options.

    Prints results automatically on call unless silent=True. Returns an FDMResult object.

    **Note:** This function is not vectorized. It prices one option per call.
    For vectorized European Greeks on option chains, use the main package API
    (e.g. ``greeks``, ``first_order``, etc.).

    Required inputs
    ---------------
    S0        : float — current spot price
    K         : float — strike price
    sigma     : float — annualised volatility (e.g. 0.1743 for 17.43%)
    r         : float — risk-free rate (continuously compounded)
    T         : float — time to expiry in years

    Optional inputs
    ---------------
    NAS       : int   — asset steps             (default 200)
    NTS       : int   — time steps              (default 2000)
    S_max_mult: float — S_max = mult × K        (default 4.0)
    rannacher : int   — fully implicit warmup   (default 4)
    omega     : float — PSOR relaxation factor  (default 1.2)
    tol       : float — PSOR convergence tol    (default 1e-8)
    max_iter  : int   — PSOR max iterations     (default 5000)
    bs        : bool  — show Black-Scholes benchmark (default False)
    option_type: str  — "call" or "put"        (default "call")
    silent    : bool  — suppress all printed output (default False)

    Returns
    -------
    FDMResult with attributes: price, delta, gamma, theta_day, S, V

    Example
    -------
    res = FDM(685.74, 685.0, 0.1743, 0.0402, 321/365)
    res = FDM(685.74, 685.0, 0.1743, 0.0402, 321/365, option_type="put")
    res = FDM(685.74, 685.0, 0.1743, 0.0402, 321/365, bs=True)
    res = FDM(685.74, 685.0, 0.1743, 0.0402, 321/365, silent=True)
    res.price
    res.delta
    """

    option_type = option_type.strip().lower()
    if option_type not in ("call", "put"):
        raise ValueError("option_type must be 'call' or 'put'")

    # ── Grid ──────────────────────────────────────────────────────────────────
    S_max  = S_max_mult * K
    dS     = S_max / NAS
    dt     = T / NTS
    S      = np.linspace(0, S_max, NAS + 1)
    if option_type == "call":
        Payoff = np.maximum(S - K, 0)
    else:
        Payoff = np.maximum(K - S, 0)

    # ── Tridiagonal Coefficients ───────────────────────────────────────────────
    S_int = S[1:NAS]
    a     = 0.5 * sigma**2 * S_int**2 / dS**2
    d     = r * S_int / (2 * dS)

    def _build(theta):
        A_lo = -theta     * dt * (a - d)
        A_di =  1 + theta * dt * (2*a + r)
        A_up = -theta     * dt * (a + d)
        B_lo =  (1-theta) * dt * (a - d)
        B_di =  1 - (1-theta) * dt * (2*a + r)
        B_up =  (1-theta) * dt * (a + d)
        return A_lo, A_di, A_up, B_lo, B_di, B_up

    A_lo_cn, A_di_cn, A_up_cn, B_lo_cn, B_di_cn, B_up_cn = _build(0.5)
    A_lo_im, A_di_im, A_up_im, B_lo_im, B_di_im, B_up_im = _build(1.0)

    # ── PSOR ──────────────────────────────────────────────────────────────────
    def _psor(A_lo, A_di, A_up, rhs, V_init, g):
        V = V_init.copy()
        for _ in range(max_iter):
            V_prev = V.copy()
            # Forward sweep: update each node using already-updated left neighbour
            # and previous-iteration right neighbour, then project onto constraint.
            # Equivalent to the scalar loop but avoids Python overhead via
            # a Thomas-style cumulative pass written in numpy where possible.
            # The sequential dependency prevents full vectorisation, so we use
            # a compiled-friendly loop via np.nditer for cache efficiency.
            for j in range(len(rhs)):
                left  = V[j-1] if j > 0          else 0.0
                right = V[j+1] if j < len(rhs)-1 else 0.0
                V_gs  = (rhs[j] - A_lo[j]*left - A_up[j]*right) / A_di[j]
                V_sor = V[j] + omega * (V_gs - V[j])
                V[j]  = V_sor if V_sor > g[j] else g[j]
            if np.max(np.abs(V - V_prev)) < tol:
                break
        return V

    # ── Time March ────────────────────────────────────────────────────────────
    V = Payoff.copy()

    for k in range(NTS):
        tau_new = (k + 1) * dt

        if k < rannacher:
            A_lo, A_di, A_up = A_lo_im, A_di_im, A_up_im
            B_lo, B_di, B_up = B_lo_im, B_di_im, B_up_im
        else:
            A_lo, A_di, A_up = A_lo_cn, A_di_cn, A_up_cn
            B_lo, B_di, B_up = B_lo_cn, B_di_cn, B_up_cn

        if option_type == "call":
            V_lower = 0.0
            V_upper = S_max - K * np.exp(-r * tau_new)
        else:
            V_lower = K * np.exp(-r * tau_new)
            V_upper = 0.0
        rhs        = B_lo * V[0:NAS-1] + B_di * V[1:NAS] + B_up * V[2:NAS+1]
        rhs[0]    -= A_lo[0]  * V_lower
        rhs[-1]   -= A_up[-1] * V_upper
        V[1:NAS]   = _psor(A_lo, A_di, A_up, rhs, V[1:NAS], Payoff[1:NAS])
        V[0]       = V_lower
        V[NAS]     = V_upper

    # ── Greeks ────────────────────────────────────────────────────────────────
    idx       = int(np.clip(np.argmin(np.abs(S - S0)), 1, NAS - 1))
    price     = V[idx]
    delta     = (V[idx+1] - V[idx-1]) / (2 * dS)
    gamma     = (V[idx+1] - 2*V[idx] + V[idx-1]) / dS**2
    theta_yr  = -(0.5*sigma**2*S[idx]**2*gamma + r*S[idx]*delta - r*price)
    theta_day = theta_yr / 365

    # ── Black-Scholes (nested, computed only if bs=True) ──────────────────────
    def _black_scholes():
        d1   = (np.log(S0/K) + (r + 0.5*sigma**2)*T) / (sigma*np.sqrt(T))
        d2   = d1 - sigma*np.sqrt(T)
        if option_type == "call":
            bs_p = S0*norm.cdf(d1) - K*np.exp(-r*T)*norm.cdf(d2)
            bs_d = norm.cdf(d1)
            bs_t = (-(S0*norm.pdf(d1)*sigma)/(2*np.sqrt(T))
                    - r*K*np.exp(-r*T)*norm.cdf(d2)) / 365
        else:
            bs_p = K*np.exp(-r*T)*norm.cdf(-d2) - S0*norm.cdf(-d1)
            bs_d = norm.cdf(d1) - 1.0
            bs_t = (-(S0*norm.pdf(d1)*sigma)/(2*np.sqrt(T))
                    + r*K*np.exp(-r*T)*norm.cdf(-d2)) / 365
        bs_g = norm.pdf(d1) / (S0*sigma*np.sqrt(T))
        return bs_p, bs_d, bs_g, bs_t

    # ── Print ─────────────────────────────────────────────────────────────────
    if not silent:
        days = round(T * 365)
        label = "AMERICAN CALL" if option_type == "call" else "AMERICAN PUT"
        print("=" * 62)
        print(f"  {label} — Crank-Nicolson + PSOR")
        print("=" * 62)
        print(f"\n  Option Parameters:")
        print(f"    Spot (S₀):           ${S0:.2f}")
        print(f"    Strike (K):          ${K:.2f}")
        print(f"    Volatility (σ):       {sigma*100:.2f}%")
        print(f"    Risk-free rate:       {r*100:.2f}%")
        print(f"    Time to expiry:       {T:.4f} yrs ({days} days)")
        print(f"\n{'─'*62}")

        if bs:
            bs_p, bs_d, bs_g, bs_t = _black_scholes()
            print(f"  {'':28s} {'FDM':>10s}  {'Black-Scholes':>12s}")
            print(f"{'─'*62}")
            print(f"  Option Price:          {'${:.4f}'.format(price):>10s}  {'${:.4f}'.format(bs_p):>12s}")
            print(f"  Delta:                 {delta:>+10.6f}  {bs_d:>+12.6f}")
            print(f"  Gamma:                 {gamma:>+10.6f}  {bs_g:>+12.6f}")
            print(f"  Theta (per day):       {theta_day:>+10.6f}  {bs_t:>+12.6f}")
            print(f"{'─'*62}")
            print(f"\n  Error vs Black-Scholes:")
            print(f"    Price: ${abs(price-bs_p):.4f}  ({abs(price-bs_p)/bs_p*100:.3f}%)")
        else:
            print(f"  Option Price:    ${price:.4f}")
            print(f"{'─'*62}")
            print(f"  Delta:           {delta:>+.6f}")
            print(f"  Gamma:           {gamma:>+.6f}")
            print(f"  Theta (per day): {theta_day:>+.6f}")

        print("=" * 62)

    return _FDMResult(price, delta, gamma, theta_day, S, V)


__all__ = [
    "FDM",
    "_FDMResult",
]