import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm

# -----------------------------
# Black-Scholes European Put
# -----------------------------
def european_put_black_scholes(S0, K, r, sigma, T):
    d1 = (np.log(S0/K) + (r + 0.5*sigma**2)*T) / (sigma*np.sqrt(T))
    d2 = d1 - sigma*np.sqrt(T)
    return K*np.exp(-r*T)*norm.cdf(-d2) - S0*norm.cdf(-d1)

# -----------------------------
# GBM Simulation
# -----------------------------
def simulate_gbm(S0, r, sigma, T, M, N, seed=None):
    if seed is not None:
        np.random.seed(seed)
    dt = T / N
    # Vectorized simulation for speed
    S = np.zeros((M, N + 1))
    S[:, 0] = S0
    for t in range(1, N + 1):
        z = np.random.standard_normal(M)
        S[:, t] = S[:, t-1] * np.exp((r - 0.5 * sigma**2) * dt + sigma * np.sqrt(dt) * z)
    return S

# -----------------------------
# Longstaff-Schwartz Monte Carlo (Fixed)
# -----------------------------
def lsm_american_put(S0, K, r, sigma, T, M=100_000, N=50, basis_deg=2, seed=42, return_boundary=False):
    S = simulate_gbm(S0, r, sigma, T, M, N, seed)
    dt = T / N
    df = np.exp(-r * dt)

    # Immediate exercise value (payoff) at each step
    payoff = np.maximum(K - S, 0)
    
    # V stores the value of the option at each path
    # Initialize with the payoff at maturity
    V = payoff[:, -1]

    boundary = []

    # Backward induction
    for t in range(N - 1, 0, -1):
        # Identify In-The-Money paths
        itm = payoff[:, t] > 0
        
        # Default: option value at t is just the discounted value from t+1
        V = V * df
        
        if np.any(itm):
            X = S[itm, t]
            Y = V[itm] # These are already discounted from t+1
            
            # Regression: Estimate Continuation Value
            coeffs = np.polyfit(X, Y, basis_deg)
            continuation_value = np.polyval(coeffs, X)
            
            # Exercise if payoff > estimated continuation value
            exercise = payoff[itm, t] > continuation_value
            
            # Update V for paths where we exercise
            # Get the indices of the ITM paths that should exercise
            itm_indices = np.where(itm)[0]
            exercise_indices = itm_indices[exercise]
            V[exercise_indices] = payoff[exercise_indices, t]
            
            # Boundary: The highest stock price at which we still exercise
            if len(exercise_indices) > 0:
                boundary.append((t * dt, np.max(S[exercise_indices, t])))
            else:
                boundary.append((t * dt, np.nan))
        else:
            boundary.append((t * dt, np.nan))

    # Final price is the average of discounted values at t=1
    price = np.mean(V * df)
    
    if return_boundary:
        return price, boundary[::-1]
    return price

# -----------------------------
# QuantLib American Put (v1.18)
# -----------------------------

def quantlib_american_put(S0, K, r, sigma, T):
    import QuantLib as ql

    today = ql.Date().todaysDate()
    ql.Settings.instance().evaluationDate = today
    maturity = today + int(T*365)
    payoff = ql.PlainVanillaPayoff(ql.Option.Put, K)
    exercise = ql.AmericanExercise(today, maturity)

    spot = ql.QuoteHandle(ql.SimpleQuote(S0))
    dividend_curve = ql.YieldTermStructureHandle(ql.FlatForward(today, 0.0, ql.Actual365Fixed()))
    riskfree_curve = ql.YieldTermStructureHandle(ql.FlatForward(today, r, ql.Actual365Fixed()))
    vol_curve = ql.BlackVolTermStructureHandle(ql.BlackConstantVol(today, ql.NullCalendar(), sigma, ql.Actual365Fixed()))

    process = ql.BlackScholesMertonProcess(spot, dividend_curve, riskfree_curve, vol_curve)
    engine = ql.FdBlackScholesVanillaEngine(process, 200, 400)
    option = ql.VanillaOption(payoff, exercise)
    option.setPricingEngine(engine)
    return option.NPV()

# -----------------------------
# Main script
# -----------------------------
if __name__ == "__main__":
    # Parameters
    S0, K, r, sigma, T = 100, 100, 0.05, 0.2, 1.0
    M, N = 200_000, 50

    # European Put
    eur_bs = european_put_black_scholes(S0, K, r, sigma, T)
    print(f"European Put (BS): {eur_bs:.4f}")

    # LSM American Put
    lsm_price, boundary = lsm_american_put(S0, K, r, sigma, T, M, N, return_boundary=True)
    print(f"LSM American Put (M={M}): {lsm_price:.4f}")

    # QuantLib American Put
    ql_price = quantlib_american_put(S0, K, r, sigma, T)
    print(f"QuantLib American Put: {ql_price:.4f}")

    print(f"Lower bound (immediate exercise): {max(K-S0,0):.4f}")

    # Plot Exercise Boundary
    times = [b[0] for b in boundary]
    boundaries = [b[1] for b in boundary]

    plt.figure(figsize=(8,5))
    plt.plot(times, boundaries, color='orange', label="LSM Exercise Boundary")
    plt.axhline(K, color='red', linestyle='--', label="Strike Price")
    plt.xlabel("Time to Maturity")
    plt.ylabel("Stock Price for Exercise")
    plt.title("American Put LSM Exercise Boundary")
    plt.gca().invert_xaxis()
    plt.legend()
    plt.grid(True)
    plt.show()