# -*- coding: utf-8 -*-
"""
Created on Tue Jan  6 09:46:52 2026

@author: S.T.Hwang
"""

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

def bs_price(S, K, T, r, q, sigma, option="call"):
    """Black-Scholes European option price (continuous dividend yield q)."""
    if T <= 0:
        intrinsic = max(S - K, 0.0) if option == "call" else max(K - S, 0.0)
        return intrinsic

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

    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, K, T, r, q, sigma, option="call"):
    """Black-Scholes European delta (continuous dividend yield q)."""
    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

    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)

def _baw_params(T, r, q, sigma):
    """
    BAW에서 쓰는 공통 파라미터.
    b = r - q (cost of carry)
    """
    b = r - q
    sig2 = sigma * sigma
    # M = 2r/sigma^2, N = 2b/sigma^2
    M = 2.0 * r / sig2
    N = 2.0 * b / sig2
    # Kappa = 1 - exp(-rT)
    kappa = 1.0 - exp(-r * T)
    return b, M, N, kappa

def _baw_q1_q2(T, r, q, sigma):
    """
    q1, q2는 BAW 논문/실무 표기에서 임계가격 방정식에 들어가는 지수.
    (문헌마다 기호가 조금 다르지만, 아래 형태가 표준적으로 쓰입니다.)
    """
    b, M, N, kappa = _baw_params(T, r, q, sigma)
    sig2 = sigma * sigma

    # sqrt term
    sqrt_term = sqrt((N - 1.0) ** 2 + 4.0 * M / kappa)

    # Call 쪽 지수(양수)
    q2 = 0.5 * (-(N - 1.0) + sqrt_term)
    # Put 쪽 지수(음수)
    q1 = 0.5 * (-(N - 1.0) - sqrt_term)
    return q1, q2

def american_baw(S, K, T, r, q, sigma, option="call"):
  
    """
    Barone-Adesi & Whaley approximation for American option (continuous dividend yield).

    Parameters
    ----------
    S : float
    K : float
    T : float (years)
    r : float
    q : float (dividend yield)
    sigma : float
    option : "call" or "put"

    Returns
    -------
    price : float
    """

    # 만기면 내재가치
    if T <= 0:
        return max(S - K, 0.0) if option == "call" else max(K - S, 0.0)

    # 배당이 없으면(또는 cost-of-carry가 r에 가까우면) Call은 보통 조기행사 없음 → 유럽형과 동일
    # (엄밀히는 q==0이면 American call = European call)
    if option == "call" and q <= 1e-14:
        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)

    # 유럽형 가격/델타
    euro = bs_price(S, K, T, r, q, sigma, option=option)

    # --------- 임계가격 S* 찾기 ---------
    # 임계가격에서:
    # option = call : V(S*) = S* - K  and  V'(S*) = 1
    # option = put  : V(S*) = K - S*  and  V'(S*) = -1
    #
    # BAW에서는 (S/S*)^q 형태의 premium term을 붙여서 위 조건으로 S*를 푼다.

    if option == "call":
        # A2(S*) = (S*/q2) * (1 - e^{-qT} N(d1*))  (문헌 형태 중 가장 흔한 형태)
        # 여기선 smooth pasting을 이용해 f(S*)=0을 구성해 root 찾음.
        def f(S_star):
            # 유럽형 call at S_star
            C_e = bs_price(S_star, K, T, r, q, sigma, option="call")
            delta_e = bs_delta(S_star, K, T, r, q, sigma, option="call")
            # A2 from smooth pasting: A2 = (1 - delta_e) * S_star / q2
            A2 = (1.0 - delta_e) * S_star / q2
            # value matching: C_e + A2 = S_star - K
            return C_e + A2 - (S_star - K)

        # root bracket: S*는 보통 K보다 큼
        low = K
        high = K * 50.0
        # 브래킷이 안 잡히면 더 확장
        fl, fh = f(low), f(high)
        if fl * fh > 0:
            high = K * 200.0
            fh = f(high)
            if fl * fh > 0:
                # 안전장치: 그래도 안 되면 유럽형으로 fallback
                return euro

        S_star = brentq(f, low, high, maxiter=200)

        # S >= S*이면 즉시행사(내재가)
        if S >= S_star:
            return S - K

        # 아니면: American = European + A2*(S/S*)^q2
        delta_star = bs_delta(S_star, K, T, r, q, sigma, option="call")
        A2 = (1.0 - delta_star) * S_star / q2
        return euro + A2 * (S / S_star) ** q2

    else:  # put
        def f(S_star):
            P_e = bs_price(S_star, K, T, r, q, sigma, option="put")
            delta_e = bs_delta(S_star, K, T, r, q, sigma, option="put")
            # A1 from smooth pasting: A1 = (-1 - delta_e) * S_star / q1
            A1 = (-1.0 - delta_e) * S_star / q1
            # value matching: P_e + A1 = K - S_star
            return P_e + A1 - (K - S_star)

        # put의 S*는 보통 K보다 작음
        low = 1e-12
        high = K

        fl, fh = f(low), f(high)
        if fl * fh > 0:
            # 확장 시도
            low = 1e-10
            fl = f(low)
            if fl * fh > 0:
                return euro

        S_star = brentq(f, low, high, maxiter=200)

        # S <= S*이면 즉시행사(내재가)
        if S <= S_star:
            return K - S

        # 아니면: American = European + A1*(S/S*)^q1
        delta_star = bs_delta(S_star, K, T, r, q, sigma, option="put")
        A1 = (-1.0 - delta_star) * S_star / q1
        return euro + A1 * (S / S_star) ** q1
    



#############################################
########### Bjerksund와 Stensland(2002) 근사해
#################################################
# -----------------------------
# 1) Black-Scholes European
# -----------------------------
def bs_european(S, K, T, r, q, sigma, is_call=True):
    if T <= 0:
        payoff = max(S - K, 0.0) if is_call else max(K - S, 0.0)
        return payoff

    if sigma <= 0:
        # deterministic
        fwd = S * math.exp((r - q) * T)
        disc = math.exp(-r * T)
        payoff = max(fwd - K, 0.0) if is_call else max(K - fwd, 0.0)
        return disc * payoff

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

    if is_call:
        return S * math.exp(-q * T) * norm.cdf(d1) - K * math.exp(-r * T) * norm.cdf(d2)
    else:
        return K * math.exp(-r * T) * norm.cdf(-d2) - S * math.exp(-q * T) * norm.cdf(-d1)


# -----------------------------
# 2) BS2002 helper: Phi, Psi
#    (Bjerksund-Stensland 2002)
# -----------------------------
def _phi(S, T, gamma, H, I, r, q, sigma):
    """
    Phi(S, T; gamma, H, I) = e^{lambda*T} * S^gamma * N(d) - e^{lambda*T} * I^gamma * (S/I)^{kappa} * N(d - 2*ln(I/S)/(sigma*sqrtT))
    구현체마다 표기가 다른데, 아래는 BS2002에서 많이 쓰는 안정형 형태로 구성.
    """
    if T <= 0:
        return 0.0

    b = r - q
    sig2 = sigma * sigma
    sqrtT = math.sqrt(T)

    # kappa(= 2b/sigma^2 + (2gamma-1)) 관련 항
    # 아래는 표준 구현에서 쓰이는 kappa 형태
    kappa = (2.0 * b / sig2) + (2.0 * gamma - 1.0)

    # lambda = -r + gamma*b + 0.5*gamma*(gamma-1)*sigma^2
    lam = -r + gamma * b + 0.5 * gamma * (gamma - 1.0) * sig2

    # d = [ln(S/H) + (b + (gamma-0.5)*sigma^2)T] / (sigma*sqrtT)
    d = (math.log(S / H) + (b + (gamma - 0.5) * sig2) * T) / (sigma * sqrtT)

    # d2 = d - 2 ln(I/S)/(sigma*sqrtT)
    d2 = d - (2.0 * math.log(I / S)) / (sigma * sqrtT)

    term1 = math.exp(lam * T) * (S ** gamma) * norm.cdf(d)
    term2 = math.exp(lam * T) * (I ** gamma) * ((S / I) ** kappa) * norm.cdf(d2)
    return term1 - term2


def _psi(S, T, gamma, H, I2, I1, r, q, sigma):
    """
    Psi는 Phi의 조합으로 정의되는 보조함수.
    BS2002 표준 구현에서 아래 형태를 많이 사용합니다.
    """
    if T <= 0:
        return 0.0

    return _phi(S, T, gamma, H, I2, r, q, sigma) - _phi(S, T, gamma, I1, I2, r, q, sigma)


# -----------------------------
# 3) BS2002: American Call/Put Approx
# -----------------------------
def american_bs2002(S, K, T, r, q, sigma, is_call=True):
    """
    Bjerksund-Stensland (2002) approximation.

    입력:
      S: spot
      K: strike
      T: maturity (year)
      r: risk-free
      q: dividend yield (continuous)
      sigma: vol
      is_call: True=Call, False=Put
    """
    # 만기 처리
    if T <= 0:
        return max(S - K, 0.0) if is_call else max(K - S, 0.0)

    b = r - q  # cost of carry

    # 배당이 없고(b=r) 콜이면 조기행사 이점이 거의 없어서 (정확히는 q=0이면)
    # 미국형 콜 = 유럽형 콜 (BS)로 처리하는 게 일반적
    if is_call and q <= 1e-14:
        return bs_european(S, K, T, r, q, sigma, is_call=True)

    # sigma=0 특이 케이스
    if sigma <= 0:
        # 조기행사 포함이면 사실상 즉시행사 vs 만기행사 비교
        euro = bs_european(S, K, T, r, q, sigma, is_call=is_call)
        intrinsic = max(S - K, 0.0) if is_call else max(K - S, 0.0)
        return max(euro, intrinsic)

    sig2 = sigma * sigma

    # ----- 공통 파라미터 -----
    # beta(=q2 역할) : 조기행사 경계의 지수
    # (표준 BS2002 구현에서 아래를 사용)
    # beta = (1/2 - b/sig^2) + sqrt((b/sig^2 - 1/2)^2 + 2r/sig^2)
    tmp = (b / sig2 - 0.5)
    beta = 0.5 - (b / sig2) + math.sqrt(tmp * tmp + 2.0 * r / sig2)

    if is_call:
        # Call: beta > 1 이어야 정상
        # B_inf, B0, I(임계값) 계산
        B_inf = (beta / (beta - 1.0)) * K
        B0 = max(K, (r / (r - b)) * K) if (r - b) > 1e-14 else B_inf

        # h(T) (BS2002 time-shape)
        h = -(b * T + 2.0 * sigma * math.sqrt(T)) * (B0 / (B_inf - B0))
        I = B0 + (B_inf - B0) * (1.0 - math.exp(h))

        # alpha
        alpha = (I - K) * (I ** (-beta))

        # 가격
        if S >= I:
            return S - K

        # 보조로 쓰는 t1, t2 (BS2002: two time points)
        t1 = 0.5 * (math.sqrt(5.0) - 1.0) * T
        t2 = T

        # 구성 (표준 구현)
        price = (
            alpha * (S ** beta)
            - alpha * _phi(S, t1, beta, I, I, r, q, sigma)
            + _phi(S, t1, 1.0, I, I, r, q, sigma)
            - _phi(S, t1, 1.0, K, I, r, q, sigma)
            - K * _phi(S, t1, 0.0, I, I, r, q, sigma)
            + K * _phi(S, t1, 0.0, K, I, r, q, sigma)
        )

        return float(price)

    else:
        # Put은 put-call parity 기반으로 직접 BS2002 put 공식을 쓰거나,
        # call 근사 + put-call parity를 쓰는 경우가 많습니다.
        # 여기서는 "직접 put"에 흔히 쓰이는 변환(대칭성) 방식으로 구현합니다.

        # 대칭 변환:
        # Put(S,K,r,q) = Call(K,S,q,r) 를 (적절한 스케일)로 변환하는 형태가 있음.
        # 가장 안정적인 실무 구현은:
        #   P_Am = C_Am(S,K,T,r,q,σ; putflag) 로 별도 공식을 적용
        # 여기서는 표준 BS2002 put 구현을 작성합니다.

        # put용 beta (보통 음수쪽 근)
        tmp = (b / sig2 - 0.5)
        beta_p = 0.5 - (b / sig2) - math.sqrt(tmp * tmp + 2.0 * r / sig2)  # 음수

        # 임계값 계산
        B_inf = (beta_p / (beta_p - 1.0)) * K
        B0 = min(K, (r / (r - b)) * K) if (r - b) > 1e-14 else B_inf

        h = (b * T - 2.0 * sigma * math.sqrt(T)) * (B0 / (B0 - B_inf))
        I = B_inf + (B0 - B_inf) * math.exp(h)

        alpha = (K - I) * (I ** (-beta_p))

        if S <= I:
            return K - S

        t1 = 0.5 * (math.sqrt(5.0) - 1.0) * T

        price = (
            alpha * (S ** beta_p)
            - alpha * _phi(S, t1, beta_p, I, I, r, q, sigma)
            - _phi(S, t1, 1.0, I, I, r, q, sigma)
            + _phi(S, t1, 1.0, K, I, r, q, sigma)
            + K * _phi(S, t1, 0.0, I, I, r, q, sigma)
            - K * _phi(S, t1, 0.0, K, I, r, q, sigma)
        )
        return float(price)
    
 # -----------------------------
 # Longstaff-Schwartz simulation
 # -----------------------------
   
    def lsm_american_gbm(
    S0: float,
    K: float,
    T: float,          # years
    r: float,
    q: float,
    sigma: float,
    n_steps: int,      # time steps
    n_paths: int,      # MC paths
    callput: str = "put",
    basis_deg: int = 2,   # polynomial degree: 2면 (1,S,S^2)
    seed: int = 111,
    antithetic: bool = True,
    return_paths: bool = False
):
    """
    Longstaff-Schwartz (LSM) American option pricer under GBM.

    Returns:
      price (float)
      (optional) S_paths (n_paths, n_steps+1)
    """

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

    rng = np.random.default_rng(seed)
    dt = T / n_steps
    disc = np.exp(-r * dt)

    # --- 1) 경로 생성 (GBM) ---
    # d ln S = (r - q - 0.5σ^2)dt + σ sqrt(dt) Z
    mu = (r - q - 0.5 * sigma**2) * dt
    vol = sigma * np.sqrt(dt)

    if antithetic:
        half = (n_paths + 1) // 2
        Z_half = rng.standard_normal(size=(half, n_steps))
        Z = np.vstack([Z_half, -Z_half])[:n_paths, :]
    else:
        Z = rng.standard_normal(size=(n_paths, n_steps))

    ln_incr = mu + vol * Z  # (n_paths, n_steps)
    lnS = np.cumsum(ln_incr, axis=1)
    lnS = np.hstack([np.zeros((n_paths, 1)), lnS])  # t=0 포함
    S = S0 * np.exp(lnS)  # (n_paths, n_steps+1)

    # --- 2) 내재가치(payoff) ---
    if callput == "call":
        intrinsic = np.maximum(S - K, 0.0)
    else:
        intrinsic = np.maximum(K - S, 0.0)

    # --- 3) LSM: 뒤에서 앞으로 회귀하며 exercise 결정 ---
    # cashflow[t] = "t에서 exercise 했다면 받는 금액" (그 외 0)
    cashflow = intrinsic[:, -1].copy()  # 만기에는 무조건 intrinsic (미국형도 만기는 동일)

    # exercise_time: 언제 exercise 했는지 추적(디버깅/분석용)
    exercise_time = np.full(n_paths, n_steps, dtype=int)

    # t = n_steps-1 ... 1 (보통 t=0은 exercise 안한다고 가정)
    for t in range(n_steps - 1, 0, -1):
        # 지금 시점까지 디스카운트(한 스텝씩 뒤로 오면서 disc 곱해줌)
        cashflow *= disc

        # ITM(in-the-money) path만 회귀
        itm = intrinsic[:, t] > 0
        if not np.any(itm):
            continue

        X = S[itm, t]  # 상태변수(여기서는 S만 사용)
        Y = cashflow[itm]  # 계속 보유했을 때의 "할인된 미래 현금흐름"

        # --- 회귀: E[continuation | S_t] ≈ beta0 + beta1*S + beta2*S^2 + ...
        # basis: [1, X, X^2, ..., X^basis_deg]
        A = np.vstack([X**d for d in range(basis_deg + 1)]).T  # (n_itm, deg+1)
        beta, *_ = np.linalg.lstsq(A, Y, rcond=None)
        continuation = A @ beta  # (n_itm,)

        # exercise if intrinsic >= continuation
        ex_now = intrinsic[itm, t] >= continuation

        # exercise한 경로는 cashflow를 intrinsic로 교체, time 업데이트
        idx_itm = np.where(itm)[0]
        ex_idx = idx_itm[ex_now]

        cashflow[ex_idx] = intrinsic[ex_idx, t]
        exercise_time[ex_idx] = t

        # exercise한 경로는 이후 시점 cashflow가 더 이상 없다고 보는 효과:
        # 이미 cashflow를 "t 시점 지급액"으로 갈아끼웠기 때문에,
        # 뒤로 더 갈 때는 disc만 적용되며 유지됨(LSM 표준 방식)

    # 마지막으로 t=0로 한 번 더 디스카운트 (t=1 -> t=0)
    price = np.mean(cashflow) * disc

    if return_paths:
        return float(price), S, exercise_time
    return float(price)
