import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from scipy.stats import norm
import warnings
warnings.filterwarnings('ignore')

class HullWhiteModel:
    """Hull-White Interest Rate Model Implementation"""
    
    def __init__(self, scen_size=1000, time_len=30, step_size=360, a=0.1, sigma=0.1, r0=0.05):
        self.scen_size = scen_size
        self.time_len = time_len
        self.step_size = step_size
        self.a = a
        self.sigma = sigma
        self.r0 = r0
        self.dt = time_len / step_size
        
        # Generate time grid
        self.t = np.linspace(0, time_len, step_size + 1)
        
        # Market forward rates (constant for simplicity)
        self.mkt_fwd = np.full(step_size + 1, r0)
        
        # Market zero-coupon bond prices
        self.mkt_zcb = np.exp(-self.mkt_fwd * self.t)
        
        # Alpha function
        self.alpha = self._calculate_alpha()
        
        # Generate random numbers
        np.random.seed(42)  # For reproducibility
        self.random_normal = np.random.standard_normal((scen_size, step_size))
        
    def _calculate_alpha(self):
        """Calculate alpha(t) = f^M(0,t) + sigma^2/(2*a^2) * (1-exp(-a*t))^2"""
        return self.mkt_fwd + (self.sigma**2 / (2 * self.a**2)) * (1 - np.exp(-self.a * self.t))**2
    
    def simulate_short_rates(self):
        """Simulate short rate paths using Hull-White model"""
        r_paths = np.zeros((self.scen_size, self.step_size + 1))
        r_paths[:, 0] = self.r0
        
        for i in range(1, self.step_size + 1):
            # Calculate conditional mean
            exp_factor = np.exp(-self.a * self.dt)
            mean_r = r_paths[:, i-1] * exp_factor + self.alpha[i] - self.alpha[i-1] * exp_factor
            
            # Calculate conditional variance
            var_r = (self.sigma**2 / (2 * self.a)) * (1 - np.exp(-2 * self.a * self.dt))
            std_r = np.sqrt(var_r)
            
            # Generate next step
            r_paths[:, i] = mean_r + std_r * self.random_normal[:, i-1]
            
        return r_paths
    
    def calculate_discount_factors(self, r_paths):
        """Calculate discount factors from short rate paths"""
        # Accumulate short rates (discrete approximation of integral)
        accum_rates = np.zeros_like(r_paths)
        for i in range(1, self.step_size + 1):
            accum_rates[:, i] = accum_rates[:, i-1] + r_paths[:, i-1] * self.dt
        
        # Calculate discount factors
        discount_factors = np.exp(-accum_rates)
        return discount_factors
    
    def theoretical_mean_short_rate(self):
        """Calculate theoretical mean of short rates E[r(t)|F_0]"""
        return self.r0 * np.exp(-self.a * self.t) + self.alpha - self.alpha[0] * np.exp(-self.a * self.t)
    
    def theoretical_var_short_rate(self):
        """Calculate theoretical variance of short rates Var[r(t)|F_0]"""
        return (self.sigma**2 / (2 * self.a)) * (1 - np.exp(-2 * self.a * self.t))

def create_short_rate_plot(scen_size, time_len, step_size, a, sigma, r0, num_paths):
    """Create short rate simulation plot"""
    model = HullWhiteModel(scen_size, time_len, step_size, a, sigma, r0)
    r_paths = model.simulate_short_rates()
    
    fig, ax = plt.subplots(figsize=(12, 8))
    
    # Plot first num_paths scenarios
    for i in range(min(num_paths, scen_size)):
        ax.plot(model.t, r_paths[i], alpha=0.7, linewidth=1)
    
    ax.set_xlabel('Time (years)')
    ax.set_ylabel('Short Rate')
    ax.set_title(f'Hull-White Short Rate Simulation ({num_paths} paths)\na={a}, σ={sigma}, scenarios={scen_size}')
    ax.grid(True, alpha=0.3)
    
    return fig

def create_convergence_plot(scen_size, time_len, step_size, a, sigma, r0):
    """Create mean convergence plot"""
    model = HullWhiteModel(scen_size, time_len, step_size, a, sigma, r0)
    r_paths = model.simulate_short_rates()
    
    # Calculate simulated means and theoretical expectations
    simulated_mean = np.mean(r_paths, axis=0)
    theoretical_mean = model.theoretical_mean_short_rate()
    
    fig, ax = plt.subplots(figsize=(12, 8))
    
    ax.plot(model.t, theoretical_mean, 'b-', linewidth=2, label='Theoretical E[r(t)]')
    ax.plot(model.t, simulated_mean, 'r--', linewidth=2, label='Simulated Mean')
    
    ax.set_xlabel('Time (years)')
    ax.set_ylabel('Short Rate')
    ax.set_title(f'Mean Convergence Analysis\na={a}, σ={sigma}, scenarios={scen_size}')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    return fig

def create_variance_plot(scen_size, time_len, step_size, a, sigma, r0):
    """Create variance convergence plot"""
    model = HullWhiteModel(scen_size, time_len, step_size, a, sigma, r0)
    r_paths = model.simulate_short_rates()
    
    # Calculate simulated variance and theoretical variance
    simulated_var = np.var(r_paths, axis=0)
    theoretical_var = model.theoretical_var_short_rate()
    
    fig, ax = plt.subplots(figsize=(12, 8))
    
    ax.plot(model.t, theoretical_var, 'b-', linewidth=2, label='Theoretical Var[r(t)]')
    ax.plot(model.t, simulated_var, 'r--', linewidth=2, label='Simulated Variance')
    
    ax.set_xlabel('Time (years)')
    ax.set_ylabel('Variance')
    ax.set_title(f'Variance Convergence Analysis\na={a}, σ={sigma}, scenarios={scen_size}')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    return fig

def create_discount_factor_plot(scen_size, time_len, step_size, a, sigma, r0):
    """Create discount factor convergence plot"""
    model = HullWhiteModel(scen_size, time_len, step_size, a, sigma, r0)
    r_paths = model.simulate_short_rates()
    discount_factors = model.calculate_discount_factors(r_paths)
    
    # Calculate mean discount factors
    mean_discount = np.mean(discount_factors, axis=0)
    
    fig, ax = plt.subplots(figsize=(12, 8))
    
    ax.plot(model.t, model.mkt_zcb, 'b-', linewidth=2, label='Market Zero-Coupon Bonds')
    ax.plot(model.t, mean_discount, 'r--', linewidth=2, label='Simulated Mean Discount Factor')
    
    ax.set_xlabel('Time (years)')
    ax.set_ylabel('Discount Factor')
    ax.set_title(f'Discount Factor Convergence\na={a}, σ={sigma}, σ/a={sigma/a:.2f}, scenarios={scen_size}')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    return fig

def create_parameter_sensitivity_plot(base_scen_size, time_len, step_size, base_a, base_sigma, r0, vary_param):
    """Create parameter sensitivity analysis"""
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle(f'Parameter Sensitivity Analysis - Varying {vary_param}', fontsize=16)
    
    if vary_param == "sigma":
        param_values = [0.05, 0.075, 0.1, 0.125]
        base_param = base_a
        param_label = "σ"
        base_label = f"a={base_a}"
    else:  # vary a
        param_values = [0.05, 0.1, 0.15, 0.2]
        base_param = base_sigma
        param_label = "a"
        base_label = f"σ={base_sigma}"
    
    axes = [ax1, ax2, ax3, ax4]
    
    for i, param_val in enumerate(param_values):
        if vary_param == "sigma":
            model = HullWhiteModel(base_scen_size, time_len, step_size, base_a, param_val, r0)
            ratio = param_val / base_a
        else:
            model = HullWhiteModel(base_scen_size, time_len, step_size, param_val, base_sigma, r0)
            ratio = base_sigma / param_val
        
        r_paths = model.simulate_short_rates()
        discount_factors = model.calculate_discount_factors(r_paths)
        mean_discount = np.mean(discount_factors, axis=0)
        
        axes[i].plot(model.t, model.mkt_zcb, 'b-', linewidth=2, label='Market ZCB')
        axes[i].plot(model.t, mean_discount, 'r--', linewidth=2, label='Simulated Mean')
        axes[i].set_title(f'{param_label}={param_val}, σ/a={ratio:.2f}')
        axes[i].grid(True, alpha=0.3)
        axes[i].legend()
    
    return fig

def generate_statistics_table(scen_size, time_len, step_size, a, sigma, r0):
    """Generate summary statistics table"""
    model = HullWhiteModel(scen_size, time_len, step_size, a, sigma, r0)
    r_paths = model.simulate_short_rates()
    
    # Calculate statistics at key time points
    time_points = [0, int(step_size*0.25), int(step_size*0.5), int(step_size*0.75), step_size]
    times = [model.t[i] for i in time_points]
    
    stats_data = []
    for i, t_idx in enumerate(time_points):
        rates_at_t = r_paths[:, t_idx]
        theoretical_mean = model.theoretical_mean_short_rate()[t_idx]
        theoretical_var = model.theoretical_var_short_rate()[t_idx]
        
        stats_data.append({
            'Time': f'{times[i]:.1f}',
            'Simulated Mean': f'{np.mean(rates_at_t):.4f}',
            'Theoretical Mean': f'{theoretical_mean:.4f}',
            'Mean Error': f'{abs(np.mean(rates_at_t) - theoretical_mean):.4f}',
            'Simulated Std': f'{np.std(rates_at_t):.4f}',
            'Theoretical Std': f'{np.sqrt(theoretical_var):.4f}',
            'Std Error': f'{abs(np.std(rates_at_t) - np.sqrt(theoretical_var)):.4f}'
        })
    
    return pd.DataFrame(stats_data)

# Create Gradio interface
with gr.Blocks(title="Hull-White Interest Rate Model Dashboard") as demo:
    gr.Markdown("""
    # 📊 Hull-White Interest Rate Model Dashboard
    
    This interactive dashboard allows actuaries and financial professionals to explore the Hull-White short rate model:
    
    **$$dr(t) = (θ(t) - ar(t))dt + σdW$$**
    
    Adjust the parameters below to see how they affect the interest rate simulations and convergence properties.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Model Parameters")
            scen_size = gr.Slider(100, 10000, value=1000, step=100, label="Number of Scenarios")
            time_len = gr.Slider(5, 50, value=30, step=5, label="Time Horizon (years)")
            step_size = gr.Slider(100, 500, value=360, step=60, label="Number of Time Steps")
            a = gr.Slider(0.01, 0.5, value=0.1, step=0.01, label="Mean Reversion Speed (a)")
            sigma = gr.Slider(0.01, 0.3, value=0.1, step=0.01, label="Volatility (σ)")
            r0 = gr.Slider(0.01, 0.15, value=0.05, step=0.01, label="Initial Rate (r₀)")
            
            gr.Markdown("### Display Options")
            num_paths = gr.Slider(1, 50, value=10, step=1, label="Number of Paths to Display")
            
            with gr.Row():
                vary_param = gr.Radio(["sigma", "a"], value="sigma", label="Parameter Sensitivity Analysis")
        
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.TabItem("Short Rate Paths"):
                    short_rate_plot = gr.Plot(label="Short Rate Simulation")
                
                with gr.TabItem("Mean Convergence"):
                    convergence_plot = gr.Plot(label="Mean Convergence Analysis")
                
                with gr.TabItem("Variance Convergence"):
                    variance_plot = gr.Plot(label="Variance Convergence Analysis")
                
                with gr.TabItem("Discount Factors"):
                    discount_plot = gr.Plot(label="Discount Factor Analysis")
                
                with gr.TabItem("Parameter Sensitivity"):
                    sensitivity_plot = gr.Plot(label="Parameter Sensitivity Analysis")
                
                with gr.TabItem("Statistics"):
                    stats_table = gr.Dataframe(label="Summary Statistics")
    
    gr.Markdown("""
    ### About the Hull-White Model
    
    - **Mean Reversion Speed (a)**: Controls how quickly rates revert to the long-term mean
    - **Volatility (σ)**: Controls the randomness in rate movements  
    - **σ/a Ratio**: Key parameter for convergence - ratios > 1 show poor convergence
    - **Scenarios**: More scenarios improve Monte Carlo convergence but increase computation time
    
    **Model Features:**
    - Gaussian short rate process
    - Analytical formulas for conditional moments
    - Market-consistent calibration capability
    - Monte Carlo simulation for complex derivatives
    """)
    
    # Update all plots when parameters change
    inputs = [scen_size, time_len, step_size, a, sigma, r0]
    
    # Connect inputs to outputs
    for inp in inputs + [num_paths]:
        inp.change(
            fn=create_short_rate_plot,
            inputs=inputs + [num_paths],
            outputs=short_rate_plot
        )
    
    for inp in inputs:
        inp.change(
            fn=create_convergence_plot,
            inputs=inputs,
            outputs=convergence_plot
        )
        
        inp.change(
            fn=create_variance_plot,
            inputs=inputs,
            outputs=variance_plot
        )
        
        inp.change(
            fn=create_discount_factor_plot,
            inputs=inputs,
            outputs=discount_plot
        )
        
        inp.change(
            fn=generate_statistics_table,
            inputs=inputs,
            outputs=stats_table
        )
    
    # Parameter sensitivity updates
    for inp in inputs[:-1] + [vary_param]:  # Exclude r0 from base params for sensitivity
        inp.change(
            fn=create_parameter_sensitivity_plot,
            inputs=[scen_size, time_len, step_size, a, sigma, r0, vary_param],
            outputs=sensitivity_plot
        )
    
    # Initialize plots on load
    demo.load(
        fn=create_short_rate_plot,
        inputs=inputs + [num_paths],
        outputs=short_rate_plot
    )
    demo.load(
        fn=create_convergence_plot,
        inputs=inputs,
        outputs=convergence_plot
    )
    demo.load(
        fn=create_variance_plot,
        inputs=inputs,
        outputs=variance_plot
    )
    demo.load(
        fn=create_discount_factor_plot,
        inputs=inputs,
        outputs=discount_plot
    )
    demo.load(
        fn=create_parameter_sensitivity_plot,
        inputs=[scen_size, time_len, step_size, a, sigma, r0, vary_param],
        outputs=sensitivity_plot
    )
    demo.load(
        fn=generate_statistics_table,
        inputs=inputs,
        outputs=stats_table
    )

if __name__ == "__main__":
    demo.launch()