import numpy as np
import matplotlib.pyplot as plt

# ---------------------------------------------------------
# 1. Analytical Hull-White Bond Pricing
# ---------------------------------------------------------
def bond_price(r_t, t, T, a, sigma, r0):
    """Zero-coupon bond price P(t, T) in Hull-White model for flat curve r0."""
    B = (1 - np.exp(-a * (T - t))) / a
    A = np.exp((B - (T - t)) * (a**2 * r0 - sigma**2/2) / a**2 - (sigma**2 * B**2 / (4 * a)))
    return A * np.exp(-B * r_t)

def get_swap_npv(r_t, t, T_end, strike, a, sigma, r0):
    """NPV of a Payer Swap: Receives Floating, Pays Fixed strike."""
    payment_dates = np.arange(t + 1, T_end + 1, 1.0)
    if len(payment_dates) == 0: 
        return np.zeros_like(r_t)
    
    # NPV = 1 - P(t, T_end) - strike * Sum[P(t, T_i)]
    p_end = bond_price(r_t, t, T_end, a, sigma, r0)
    fixed_leg = sum(bond_price(r_t, t, pd, a, sigma, r0) for pd in payment_dates)
    
    return np.maximum(1 - p_end - strike * fixed_leg, 0)

# ---------------------------------------------------------
# 2. Joint Exact Simulator (Short Rate & Stochastic Integral)
# ---------------------------------------------------------
def simulate_hw_exact_joint(r0, a, sigma, exercise_dates, M):
    """Samples (r_t, Integral[r]) jointly to get exact discount factors."""
    M = int(M)
    num_dates = len(exercise_dates)
    r_matrix = np.zeros((M, num_dates + 1))
    d_matrix = np.zeros((M, num_dates)) # d[i] = exp(-integral from t_i to t_{i+1})
    
    r_matrix[:, 0] = r0
    t_steps = np.insert(exercise_dates, 0, 0.0)

    for i in range(len(t_steps) - 1):
        t, T = t_steps[i], t_steps[i+1]
        dt = T - t
        
        # Drift adjustments for flat initial curve r0
        alpha_t = r0 + (sigma**2 / (2 * a**2)) * (1 - np.exp(-a * t))**2
        alpha_T = r0 + (sigma**2 / (2 * a**2)) * (1 - np.exp(-a * T))**2
        
        # Means
        mean_r = r_matrix[:, i] * np.exp(-a * dt) + alpha_T - alpha_t * np.exp(-a * dt)
        # The expected value of the integral is derived from the bond price: E[exp(-I)] = P(t,T)
        mean_I = -np.log(bond_price(r_matrix[:, i], t, T, a, sigma, r0))
        
        # Covariance Matrix Components
        var_r = (sigma**2 / (2 * a)) * (1 - np.exp(-2 * a * dt))
        B = (1 - np.exp(-a * dt)) / a
        var_I = (sigma**2 / a**2) * (dt - 2*B + (1 - np.exp(-2*a*dt))/(2*a))
        cov_rI = (sigma**2 / (2 * a**2)) * (1 - np.exp(-a * dt))**2
        
        cov_matrix = [[var_r, cov_rI], [cov_rI, var_I]]
        
        # Sample joint normal innovations
        Z = np.random.multivariate_normal([0, 0], cov_matrix, M)
        
        r_matrix[:, i+1] = mean_r + Z[:, 0]
        # Important: The variance of the integral affects the mean of the exponent (convexity)
        # mean_I here is already the log of the bond price (risk-neutral expectation)
        d_matrix[:, i] = np.exp(-(mean_I + Z[:, 1] - 0.5 * var_I)) 
        
    return r_matrix[:, 1:], d_matrix

# ---------------------------------------------------------
# 3. LSM Pricing Logic
# ---------------------------------------------------------
def price_bermudan_swaption(r0, a, sigma, strike, exercise_dates, T_end, M):
    # 1. Generate exact paths
    r_at_ex, discounts = simulate_hw_exact_joint(r0, a, sigma, exercise_dates, M)
    
    # 2. Final exercise date payoff
    T_last = exercise_dates[-1]
    cash_flows = get_swap_npv(r_at_ex[:, -1], T_last, T_end, strike, a, sigma, r0)
    
    # 3. Backward Induction
    # exercise_dates[:-1] because we already handled the last date
    for i in reversed(range(len(exercise_dates) - 1)):
        t_current = exercise_dates[i]
        
        # Pull cash flows back to current time using stochastic discount
        # Note: If path was exercised later, this is the discounted value of that exercise.
        cash_flows = cash_flows * discounts[:, i]
        
        # Current intrinsic value
        X = r_at_ex[:, i]
        immediate_payoff = get_swap_npv(X, t_current, T_end, strike, a, sigma, r0)
        
        # Only regress In-The-Money paths
        itm = immediate_payoff > 0
        if np.any(itm):
            # Regression: Basis functions [1, r, r^2]
            A = np.column_stack([np.ones_like(X[itm]), X[itm], X[itm]**2])
            coeffs = np.linalg.lstsq(A, cash_flows[itm], rcond=None)[0]
            continuation_value = A @ coeffs
            
            # Exercise decision
            exercise = immediate_payoff[itm] > continuation_value
            
            itm_indices = np.where(itm)[0]
            exercise_indices = itm_indices[exercise]
            
            # Update cash flows for paths where we exercise early
            cash_flows[exercise_indices] = immediate_payoff[exercise_indices]

    # Final discount to t=0
    # The first discount factor in 'discounts' is from t1 to t0
    # But wait, our 'discounts' matrix is (M, num_dates). 
    # Let's just use the analytical P(0, t1) for the very last step to t=0.
    final_price = np.mean(cash_flows * bond_price(r0, 0, exercise_dates[0], a, sigma, r0))
    return final_price

# ---------------------------------------------------------
# Execution
# ---------------------------------------------------------
if __name__ == "__main__":
    # Params
    r0_val, a_val, sigma_val = 0.05, 0.1, 0.01
    strike_val = 0.05
    ex_dates = np.array([1.0, 2.0, 3.0, 4.0])
    maturity = 5.0
    num_paths = 100_000

    price = price_bermudan_swaption(r0_val, a_val, sigma_val, strike_val, ex_dates, maturity, num_paths)
    
    print(f"--- 100% Exact LSM Bermudan Swaption ---")
    print(f"Parameters: a={a_val}, sigma={sigma_val}, strike={strike_val}")
    print(f"Exercise Dates: {ex_dates}")
    print(f"Calculated Price: {price:.6f}")