# -*- coding: utf-8 -*-
"""
Created on Tue Jan  6 10:23:49 2026

@author: S.T.Hwang
"""

# HYOTC/Equity/American.py
# -*- coding: utf-8 -*-

from __future__ import annotations

import numpy as np
from math import log, sqrt, exp, isfinite
from scipy.stats import norm, multivariate_normal
from scipy.optimize import brentq


# ============================================================
# 0) Helpers: 입력/옵션 플래그
# ============================================================

def _opt_flag(option: str | None = None, is_call: bool | None = None) -> str:
    """
    내부 통일용: 'call'/'put' 반환
    - option="call"/"put" 우선
    - is_call=True/False도 지원
    """
    if option is not None:
        opt = option.lower().strip()
        if opt not in ("call", "put"):
            raise ValueError("option must be 'call' or 'put'")
        return opt
    if is_call is None:
        return "call"
    return "call" if bool(is_call) else "put"


def _assert_pos(name: str, x: float):
    if not (x > 0 and isfinite(x)):
        raise ValueError(f"{name} must be positive and finite. got {x}")


# ============================================================
# 1) Black–Scholes (European) price/delta
# ============================================================

def bs_price(S: float, K: float, T: float, r: float, q: float, sigma: float, option: str = "call") -> float:
    """
    Black-Scholes European option price with continuous dividend yield q.
    """
    option = _opt_flag(option=option)

    if T <= 0:
        intrinsic = max(S - K, 0.0) if option == "call" else max(K - S, 0.0)
        return intrinsic

    _assert_pos("S", S)
    _assert_pos("K", K)
    _assert_pos("sigma", sigma)

    vsqrt = sigma * sqrt(T)
    d1 = (log(S / K) + (r - q + 0.5 * sigma * sigma) * T) / vsqrt
    d2 = d1 - vsqrt

    if option == "call":
        return S * exp(-q * T) * norm.cdf(d1) - K * exp(-r * T) * norm.cdf(d2)
    else:
        return K * exp(-r * T) * norm.cdf(-d2) - S * exp(-q * T) * norm.cdf(-d1)


def bs_delta(S: float, K: float, T: float, r: float, q: float, sigma: float, option: str = "call") -> float:
    """
    Black-Scholes European delta with continuous dividend yield q.
    """
    option = _opt_flag(option=option)

    if T <= 0:
        if option == "call":
            return 1.0 if S > K else 0.0
        else:
            return -1.0 if S < K else 0.0

    _assert_pos("S", S)
    _assert_pos("K", K)
    _assert_pos("sigma", sigma)

    d1 = (log(S / K) + (r - q + 0.5 * sigma * sigma) * T) / (sigma * sqrt(T))

    if option == "call":
        return exp(-q * T) * norm.cdf(d1)
    else:
        return exp(-q * T) * (norm.cdf(d1) - 1.0)


# ============================================================
# 2) Barone-Adesi & Whaley (1987) Approximation
# ============================================================

def _baw_params(T: float, r: float, q: float, sigma: float):
    """
    BAW common parameters.
    b = r - q  (cost of carry)
    """
    b = r - q
    sig2 = sigma * sigma
    M = 2.0 * r / sig2
    N = 2.0 * b / sig2
    kappa = 1.0 - exp(-r * T)  # 1 - e^{-rT}
    return b, M, N, kappa


def _baw_q1_q2(T: float, r: float, q: float, sigma: float):
    """
    Exponents used in BAW.
    q2 (call) > 0, q1 (put) < 0
    """
    b, M, N, kappa = _baw_params(T, r, q, sigma)
    # 보호: r≈0이면 kappa≈0가 되어 발산 가능 → 작은 값 바닥
    kappa = max(kappa, 1e-12)

    sqrt_term = sqrt((N - 1.0) ** 2 + 4.0 * M / kappa)
    q2 = 0.5 * (-(N - 1.0) + sqrt_term)  # call exponent
    q1 = 0.5 * (-(N - 1.0) - sqrt_term)  # put exponent
    return q1, q2


def american_baw(S: float, K: float, T: float, r: float, q: float, sigma: float, option: str = "call") -> float:
    """
    American option price by Barone-Adesi & Whaley approximation.

    Notes:
    - Call with dividend (q>0) may have early exercise.
    - For non-dividend call (q=0), American call = European call (no early exercise).
    """
    option = _opt_flag(option=option)

    if T <= 0:
        return max(S - K, 0.0) if option == "call" else max(K - S, 0.0)

    _assert_pos("S", S)
    _assert_pos("K", K)
    _assert_pos("sigma", sigma)

    # Non-dividend call: no early exercise in BSM world
    if option == "call" and q <= 0.0:
        return bs_price(S, K, T, r, q, sigma, option="call")

    b, M, N, kappa = _baw_params(T, r, q, sigma)
    q1, q2 = _baw_q1_q2(T, r, q, sigma)

    # choose exponent and sign by option
    if option == "call":
        qn = q2
        # 초기 guess: K 근처에서 시작
        S_guess = max(K, 1e-8)
        # S* 찾기 방정식
        def f(Sstar: float) -> float:
            c = bs_price(Sstar, K, T, r, q, sigma, option="call")
            d1 = (log(Sstar / K) + (r - q + 0.5 * sigma * sigma) * T) / (sigma * sqrt(T))
            A2 = (Sstar / qn) * (1.0 - exp((b - r) * T) * norm.cdf(d1))
            return (Sstar - K) - (c + A2)

        # Root bracket (실무 안전범위)
        lo = 1e-12
        hi = max(5.0 * K, 5.0 * S)
        # brentq가 실패하면 bracket 확장
        for _ in range(10):
            if f(lo) * f(hi) < 0:
                break
            hi *= 2.0

        S_star = brentq(f, lo, hi, maxiter=200)

        # Price
        if S >= S_star:
            return S - K
        c = bs_price(S, K, T, r, q, sigma, option="call")
        d1 = (log(S_star / K) + (r - q + 0.5 * sigma * sigma) * T) / (sigma * sqrt(T))
        A2 = (S_star / qn) * (1.0 - exp((b - r) * T) * norm.cdf(d1))
        return c + A2 * (S / S_star) ** qn

    else:
        qn = q1  # negative
        def f(Sstar: float) -> float:
            p = bs_price(Sstar, K, T, r, q, sigma, option="put")
            d1 = (log(Sstar / K) + (r - q + 0.5 * sigma * sigma) * T) / (sigma * sqrt(T))
            A1 = -(Sstar / qn) * (1.0 - exp((b - r) * T) * norm.cdf(-d1))
            return (K - Sstar) - (p + A1)

        lo = 1e-12
        hi = max(5.0 * K, 5.0 * S)
        for _ in range(10):
            if f(lo) * f(hi) < 0:
                break
            hi *= 2.0

        S_star = brentq(f, lo, hi, maxiter=200)

        if S <= S_star:
            return K - S
        p = bs_price(S, K, T, r, q, sigma, option="put")
        d1 = (log(S_star / K) + (r - q + 0.5 * sigma * sigma) * T) / (sigma * sqrt(T))
        A1 = -(S_star / qn) * (1.0 - exp((b - r) * T) * norm.cdf(-d1))
        return p + A1 * (S / S_star) ** qn


# ============================================================
# 3) Bjerksund–Stensland (2002) Approximation
#    (현업에서 BAW 대체로 많이 쓰는 근사 중 하나)
# ============================================================

def _phi(S, T, gamma, H, I, r, b, sigma):
    """
    φ(S, T, γ, H, I) = e^{λT} S^γ [ N(d) - (I/S)^k N(d - 2 ln(I/S)/(σ√T)) ]
    (BS2002 내부 구성요소)
    """
    if T <= 0:
        return 0.0

    sig2 = sigma * sigma
    sig_sqrtT = sigma * sqrt(T)
    lamb = -r + gamma * b + 0.5 * gamma * (gamma - 1.0) * sig2
    k = (2.0 * b / sig2) + (2.0 * gamma - 1.0)

    # d
    d = -(log(S / H) + (b + (gamma - 0.5) * sig2) * T) / sig_sqrtT
    d2 = d - 2.0 * log(I / S) / sig_sqrtT

    return exp(lamb * T) * (S ** gamma) * (norm.cdf(d) - (I / S) ** k * norm.cdf(d2))


def _psi(S, T, gamma, H, I2, I1, t1, r, b, sigma):
    """
    Ψ function in BS2002:
    e^{λT} S^γ [ M(-e1, -f1, ρ) - (I2/S)^k M(-e2, -f2, ρ)
               - (I1/S)^k M(-e3, -f3, -ρ) + (I1/I2)^k M(-e4, -f4, -ρ) ]
    """
    if T <= 0:
        return 0.0
    if t1 <= 0:
        return 0.0

    sig2 = sigma * sigma
    lamb = -r + gamma * b + 0.5 * gamma * (gamma - 1.0) * sig2
    k = (2.0 * b / sig2) + (2.0 * gamma - 1.0)

    sqrt_t1 = sqrt(t1)
    sqrt_T = sqrt(T)

    rho = sqrt(t1 / T)
    # e's (t1)
    e1 = (log(S / I1) + (b + (gamma - 0.5) * sig2) * t1) / (sigma * sqrt_t1)
    e2 = (log((I2 * I2) / (S * I1)) + (b + (gamma - 0.5) * sig2) * t1) / (sigma * sqrt_t1)
    e3 = (log(S / I1) - (b + (gamma - 0.5) * sig2) * t1) / (sigma * sqrt_t1)
    e4 = (log((I2 * I2) / (S * I1)) - (b + (gamma - 0.5) * sig2) * t1) / (sigma * sqrt_t1)

    # f's (T)
    f1 = (log(S / H) + (b + (gamma - 0.5) * sig2) * T) / (sigma * sqrt_T)
    f2 = (log((I2 * I2) / (S * H)) + (b + (gamma - 0.5) * sig2) * T) / (sigma * sqrt_T)
    f3 = (log((I1 * I1) / (S * H)) + (b + (gamma - 0.5) * sig2) * T) / (sigma * sqrt_T)
    f4 = (log((S * I1 * I1) / (H * I2 * I2)) + (b + (gamma - 0.5) * sig2) * T) / (sigma * sqrt_T)

    # bivariate normal CDF
    # scipy.stats.multivariate_normal.cdf expects mean and cov
    mean = np.array([0.0, 0.0])
    cov_pos = np.array([[1.0, rho], [rho, 1.0]])
    cov_neg = np.array([[1.0, -rho], [-rho, 1.0]])

    def M(a, b_, cov_):
        # P(Z1<=a, Z2<=b)
        return multivariate_normal.cdf([a, b_], mean=mean, cov=cov_)

    term1 = M(-e1, -f1, cov_pos)
    term2 = (I2 / S) ** k * M(-e2, -f2, cov_pos)
    term3 = (I1 / S) ** k * M(-e3, -f3, cov_neg)
    term4 = (I1 / I2) ** k * M(-e4, -f4, cov_neg)

    return exp(lamb * T) * (S ** gamma) * (term1 - term2 - term3 + term4)


def american_bs2002(S, K, T, r, q, sigma, is_call=True):
    """
    Bjerksund–Stensland (2002) American approximation.

    Parameters
    - S, K, T, r, q, sigma
    - is_call : True for call, False for put

    Returns
    - approximate American option price
    """
    opt = _opt_flag(is_call=is_call)

    if T <= 0:
        return max(S - K, 0.0) if opt == "call" else max(K - S, 0.0)

    _assert_pos("S", S)
    _assert_pos("K", K)
    _assert_pos("sigma", sigma)

    b = r - q

    # Put은 call로 변환(put-call symmetry)
    # P(S,K,r,q) = C(K,S,r,b?) 형태로 구현할 수도 있으나
    # 여기서는 안정적으로 "call 공식 + 패리티"를 사용
    if opt == "put":
        # American put ≈ American call on swapped params via parity-like relation은 위험
        # 실무에선 보통 put용 BS2002 별도식을 쓰거나, 여기서는 call로 뒤집은 뒤 패리티 적용
        # 가장 안전: call로 계산 후 put-call parity (European) + early exercise 차이 무시(근사)
        # → 근사를 근사로 유지하되, 사용자는 FDM/LSM로 검증 권장.
        call_val = american_bs2002(S, K, T, r, q, sigma, is_call=True)
        euro_call = bs_price(S, K, T, r, q, sigma, option="call")
        euro_put = bs_price(S, K, T, r, q, sigma, option="put")
        return call_val - euro_call + euro_put

    # Non-dividend call: American = European
    if q <= 0.0:
        return bs_price(S, K, T, r, q, sigma, option="call")

    sig2 = sigma * sigma
    beta = (0.5 - b / sig2) + sqrt((b / sig2 - 0.5) ** 2 + 2.0 * r / sig2)

    # t1, t2 (two-stage)
    t1 = 0.5 * T

    # boundaries (I1, I2)
    B_infty = (beta / (beta - 1.0)) * K
    B0 = max(K, (r / (r - b)) * K) if (r != b) else K

    h = -(b * t1 + 2.0 * sigma * sqrt(t1)) * (K * K) / ((B_infty - B0) * B0)
    I1 = B0 + (B_infty - B0) * (1.0 - exp(h))
    h = -(b * T + 2.0 * sigma * sqrt(T)) * (K * K) / ((B_infty - B0) * B0)
    I2 = B0 + (B_infty - B0) * (1.0 - exp(h))

    # α parameter
    alpha1 = (I1 - K) * (I1 ** (-beta))
    alpha2 = (I2 - K) * (I2 ** (-beta))

    if S >= I2:
        return S - K

    # European call component
    euro = bs_price(S, K, T, r, q, sigma, option="call")

    # Approx formula (call)
    # C ≈ φ(S,T,1,I2,I2) - φ(S,T,1,K,I2) + α2*S^β
    #   - α2*φ(S,T,β,I2,I2) + α2*φ(S,T,β,K,I2)
    #   + α1*ψ(S,T,β,I2,I1,I1,t1,...) - α1*ψ(S,T,β,K,I1,I1,t1,...)
    #
    # (표준 형태를 그대로 구현)
    term = (
        _phi(S, T, 1.0, I2, I2, r, b, sigma)
        - _phi(S, T, 1.0, K,  I2, r, b, sigma)
        + alpha2 * (S ** beta)
        - alpha2 * _phi(S, T, beta, I2, I2, r, b, sigma)
        + alpha2 * _phi(S, T, beta, K,  I2, r, b, sigma)
        + alpha1 * _psi(S, T, beta, I2, I1, I1, t1, r, b, sigma)
        - alpha1 * _psi(S, T, beta, K,  I1, I1, t1, r, b, sigma)
    )

    # term 자체가 이미 call 근사값
    return float(term)



