simd-228-with-episodic-shocks.py

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import PercentFormatter
import seaborn as sns

# ------------------------------------------------------------------------------
# Time-Decay Function for the Static Inflation Rate
# ------------------------------------------------------------------------------
def time_decay_static_rate(epoch, r0, r_floor, dynamic_epochs):
    """
    Returns a linearly decayed static inflation rate for the dynamic period.
    """
    decay = (r0 - r_floor) / (dynamic_epochs * 1.5)
    return max(r0 - decay * epoch, r_floor)

# ------------------------------------------------------------------------------
# Current Solana Inflation Rate Formula
# ------------------------------------------------------------------------------
def current_inflation_rate(epoch, initial_inflation=0.08, dis_inflation=0.15, floor=0.015, epochs_per_year=182.5):
    """
    Compute the current Solana inflation rate with annual disinflation.
    """
    years = epoch / epochs_per_year
    annual_inflation = initial_inflation * ((1 - dis_inflation) ** years)
    return max(annual_inflation, floor) / epochs_per_year

# ------------------------------------------------------------------------------
# SIMD-228 Proposed Inflation Rate Formula with Time Decay
# ------------------------------------------------------------------------------
def simd228_inflation_rate(staking_rate, static_rate, c=None, alpha=1.0):
    """
    Compute the dynamic inflation rate using the proposed SIMD-228 formula.
    """
    if c is None:
        c = np.sqrt(1/3) / (1 - np.sqrt(2/3))  # Exact value ≈ 3.14626436994
    term1 = 1 - alpha * np.sqrt(staking_rate)
    term2 = alpha * c * max(1 - np.sqrt(2 * staking_rate), 0)
    inflation = static_rate * (term1 + term2)
    return max(inflation, 0.0)  # Ensure non-negative inflation

# ------------------------------------------------------------------------------
# Improved Staking Rate Adjustment Function
# ------------------------------------------------------------------------------
def adjust_staking_rate(current_staking, effective_yield, target_yield, 
                       adjustment_factor=0.05, max_change=0.01,
                       min_staking=0.1, dilution_impact=0.15, delay_epochs=0, epoch=0, total_epochs=182,
                       current_inflation=None):
    """
    Adjust staking rate based on yield deviation with realistic dilution and inflation-correlated decay.
    """
    EPOCHS_PER_YEAR = 182.5
    
    # Model cumulative dilution as an increase for low staking (< 0.5) over time
    if current_staking < 0.5:
        adjustment_factor = 0.1  # More aggressive adjustment for low staking
        max_change = 0.005  # Smaller max change for smoother increase
        max_dilution_increase = dilution_impact  
        dilution_progress = min(1.0, epoch / 182)  # Target 12 months for full dilution
        dilution_effect = (max_dilution_increase * dilution_progress) / EPOCHS_PER_YEAR
        target_staking = min(1.0, current_staking + dilution_effect) # Target 34.5% for 0.3 initial
    else:
        dilution_effect = -dilution_impact / EPOCHS_PER_YEAR  # Decrease for high staking (> 0.5)
    
    # Add delay-dependent decay correlated with current inflation
    if delay_epochs > 0 and current_inflation is not None:
        # Use current inflation as a factor for minimum decay (e.g., 1% decay per 1% inflation per year)
        inflation_impact = max(0.01 * current_inflation * EPOCHS_PER_YEAR, 0.001)  # Minimum 0.1% decay per year
        decay_factor = max(0.001, 1 - (delay_epochs / (2 * EPOCHS_PER_YEAR)) - inflation_impact)
        dilution_effect *= decay_factor
    else:
        decay_factor = max(0.001, 1 - (delay_epochs / (2 * EPOCHS_PER_YEAR)))  # Default decay
    
    decayed_staking = max(current_staking + dilution_effect, min_staking)
    
    yield_deviation = effective_yield - target_yield
    if yield_deviation < 0:
        response_strength = 1.5 * adjustment_factor  # Stronger unstaking for low yields
    else:
        response_strength = adjustment_factor  # Moderate staking for high yields
    
    change = response_strength * yield_deviation
    change = np.clip(change, -max_change, max_change)
    new_staking = np.clip(decayed_staking + change, min_staking, 1.0)
    
    # Ensure low staking rates increase toward the dilution target (e.g., 34% for 0.3 initial)
    if current_staking < 0.5 and new_staking < target_staking:
        new_staking = min(target_staking, new_staking + 0.001)  # Small incremental increase
    
    return new_staking

# ------------------------------------------------------------------------------
# Simulation with Implementation Delay and Episodic Shocks
# ------------------------------------------------------------------------------
def simulate_with_delay(initial_supply, initial_staking, r0, r_floor, c, total_epochs, delay_epochs,
                       adjustment_factor=0.05, shock_epochs=None, shock_factor=1.0,
                       penalty_threshold=None, extra_penalty=0.0, validator_commission=0.0412,
                       annual_to_epoch=True, min_staking=0.1, dilution_impact=0.15):
    """
    Simulate network evolution with delay, updated staking adjustment, and equilibrium dynamics.
    """
    EPOCHS_PER_YEAR = 182.5
    
    if annual_to_epoch:
        epoch_r0 = r0 / EPOCHS_PER_YEAR
        epoch_r_floor = r_floor / EPOCHS_PER_YEAR
        epoch_c = c
    else:
        epoch_r0 = r0
        epoch_r_floor = r_floor
        epoch_c = c
    
    supplies = [initial_supply]
    staking_rates = [max(initial_staking, min_staking)]
    inflation_rates = []
    yields = []
    static_rates = []
    
    current_staking = max(initial_staking, min_staking)
    
    for epoch in range(total_epochs):
        current_supply = supplies[-1]
        
        if epoch < delay_epochs:
            current_inflation = current_inflation_rate(epoch, initial_inflation=r0, dis_inflation=0.15)
            new_tokens = current_supply * current_inflation
            effective_yield = (new_tokens / (current_staking * current_supply)) * (1 - validator_commission) if current_staking > 0 else 0
            target_yield = epoch_r0 * 1.1 * (1 - validator_commission)
            current_staking = adjust_staking_rate(
                current_staking, effective_yield, target_yield, 
                adjustment_factor=0.01, max_change=0.001, min_staking=min_staking, 
                dilution_impact=dilution_impact, delay_epochs=delay_epochs, 
                epoch=epoch, total_epochs=total_epochs, current_inflation=current_inflation
            )
            current_static_rate = epoch_r0
        else:
            dynamic_epoch = epoch - delay_epochs
            current_static_rate = time_decay_static_rate(dynamic_epoch, epoch_r0, epoch_r_floor, total_epochs - delay_epochs)
            
            if dynamic_epoch < 50:
                alpha = 0.1 + (dynamic_epoch / 49) * 0.9  # Linear interpolation to 1.0 at 49 epochs
            else:
                alpha = 0.15
            
            current_inflation = simd228_inflation_rate(current_staking, current_static_rate, epoch_c, alpha=alpha)
            
            new_tokens = current_supply * current_inflation
            effective_yield = (new_tokens / (current_staking * current_supply)) * (1 - validator_commission) if current_staking > 0 else 0
            
            if shock_epochs and epoch in shock_epochs:
                current_staking *= shock_factor
                current_staking = np.clip(current_staking, min_staking, 1.0)
            
            target_yield = current_static_rate * (1 - validator_commission)
            current_staking = adjust_staking_rate(
                current_staking, effective_yield, target_yield, 
                adjustment_factor=adjustment_factor, max_change=0.005, min_staking=min_staking, 
                dilution_impact=dilution_impact, delay_epochs=delay_epochs, 
                epoch=epoch, total_epochs=total_epochs, current_inflation=current_inflation
            )
            
            if penalty_threshold and effective_yield < penalty_threshold / EPOCHS_PER_YEAR:
                current_staking = max(current_staking - extra_penalty, min_staking)
        
        supplies.append(current_supply + new_tokens)
        staking_rates.append(current_staking)
        inflation_rates.append(current_inflation)
        yields.append(effective_yield)
        static_rates.append(current_static_rate)
    
    return supplies, staking_rates, inflation_rates, yields, static_rates, EPOCHS_PER_YEAR

# ------------------------------------------------------------------------------
# Plot Scenario Function
# ------------------------------------------------------------------------------
def plot_scenario(scenario_name, scenario_params, horizons, delays, common_params):
    """
    Plot simulation results for a given scenario across different horizons and delays.
    """
    sns.set_style("whitegrid")
    for horizon_name, total_epochs in horizons.items():
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        fig.suptitle(f"{scenario_name} - {horizon_name}", fontsize=16)
        
        legend_entries = []
        
        for delay_name, delay_epochs in delays.items():
            if delay_epochs >= total_epochs:
                continue
                
            result = simulate_with_delay(
                common_params["initial_supply"],
                scenario_params["initial_staking"],
                common_params["r0"],
                common_params["r_floor"],
                common_params["c"],
                total_epochs,
                delay_epochs,
                adjustment_factor=common_params["adjustment_factor"],
                shock_epochs=scenario_params.get("shock_epochs", None),
                shock_factor=scenario_params.get("shock_factor", 1.0),
                penalty_threshold=scenario_params.get("penalty_threshold", None),
                extra_penalty=scenario_params.get("extra_penalty", 0.0),
                validator_commission=scenario_params.get("validator_commission", 0.0412),
                annual_to_epoch=True,
                min_staking=common_params.get("min_staking", 0.1),
                dilution_impact=common_params.get("dilution_impact", 0.15)
            )
            
            supplies, staking_rates, inflation_rates, yields, static_rates, epochs_per_year = result
            epochs_range = np.arange(total_epochs + 1)
            
            annual_inflation_rates = [rate * epochs_per_year for rate in inflation_rates]
            annual_yields = [yield_val * epochs_per_year for yield_val in yields]
            
            axes[0, 0].plot(epochs_range[:-1], supplies[:-1], label=f'Delay: {delay_name}')
            axes[0, 0].axvspan(0, delay_epochs, color='gray', alpha=0.1, label='Delay Period' if delay_name == list(delays.keys())[0] else None)
            axes[0, 0].set_xlabel('Epochs (2 days per epoch)')
            axes[0, 0].set_ylabel('Total Token Supply (SOL)')
            axes[0, 0].set_title('Total Supply')
            axes[0, 0].grid(True)
            
            ax_year = axes[0, 0].twiny()
            ax_year.set_xlim(axes[0, 0].get_xlim())
            ax_year.set_xticks([0, epochs_per_year/2, epochs_per_year, 3*epochs_per_year/2, 2*epochs_per_year])
            ax_year.set_xticklabels(['0', '0.5', '1', '1.5', '2'])
            ax_year.set_xlabel('Years')
            
            axes[0, 1].plot(epochs_range, staking_rates, label=f'Delay: {delay_name}')
            axes[0, 1].set_xlabel('Epochs (2 days per epoch)')
            axes[0, 1].set_ylabel('Staking Rate')
            axes[0, 1].set_title('Staking Rate')
            axes[0, 1].grid(True)
            axes[0, 1].yaxis.set_major_formatter(PercentFormatter(1.0))
            
            axes[1, 0].plot(epochs_range[:-1], annual_inflation_rates, label=f'Delay: {delay_name}')
            axes[1, 0].set_xlabel('Epochs (2 days per epoch)')
            axes[1, 0].set_ylabel('Annual Inflation Rate')
            axes[1, 0].set_title('Inflation Rate')
            axes[1, 0].grid(True)
            axes[1, 0].yaxis.set_major_formatter(PercentFormatter(1.0))
            
            axes[1, 1].plot(epochs_range[:-1], annual_yields, label=f'Delay: {delay_name}')
            axes[1, 1].set_xlabel('Epochs (2 days per epoch)')
            axes[1, 1].set_ylabel('Annual Effective Yield')
            axes[1, 1].set_title('Staking Yield')
            axes[1, 1].grid(True)
            axes[1, 1].yaxis.set_major_formatter(PercentFormatter(1.0))
            
            legend_entries.append(f'Delay: {delay_name}')
        
        # Add shock markers if applicable
        if "shock_epochs" in scenario_params and scenario_params["shock_epochs"]:
            for shock_epoch in scenario_params["shock_epochs"]:
                for ax in axes.flat:
                    ax.axvline(x=shock_epoch, color='red', linestyle='--', alpha=0.7, 
                              label='Shock' if shock_epoch == scenario_params["shock_epochs"][0] else None)
            if scenario_params["shock_epochs"]:
                legend_entries.append('Shock')
        
        fig.legend(legend_entries, loc='lower center', ncol=len(legend_entries), bbox_to_anchor=(0.5, 0.02))
        plt.tight_layout(rect=[0, 0.05, 1, 0.95])
        plt.subplots_adjust(top=0.9, bottom=0.1)
        plt.show()

# ------------------------------------------------------------------------------
# Compare Scenarios Function
# ------------------------------------------------------------------------------
def compare_scenarios(scenarios, horizons, common_params, delay_name="4 Months"):
    """
    Compare all scenarios for a fixed delay period.
    """
    delay_epochs = {"Now": 0, "4 Months": 60, "8 Months": 120, "1 Year": 182}[delay_name]
    
    for horizon_name, total_epochs in horizons.items():
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        fig.suptitle(f'Scenario Comparison - {horizon_name} (Delay: {delay_name})', fontsize=16)
        legend_entries = []
        
        for scenario_name, scenario_params in scenarios.items():
            result = simulate_with_delay(
                common_params["initial_supply"],
                scenario_params["initial_staking"],
                common_params["r0"],
                common_params["r_floor"],
                common_params["c"],
                total_epochs,
                delay_epochs,
                adjustment_factor=common_params["adjustment_factor"],
                shock_epochs=scenario_params.get("shock_epochs", None),
                shock_factor=scenario_params.get("shock_factor", 1.0),
                penalty_threshold=scenario_params.get("penalty_threshold", None),
                extra_penalty=scenario_params.get("extra_penalty", 0.0),
                validator_commission=scenario_params.get("validator_commission", 0.0412),
                annual_to_epoch=True,
                min_staking=common_params.get("min_staking", 0.1),
                dilution_impact=common_params.get("dilution_impact", 0.15)
            )
            
            supplies, staking_rates, inflation_rates, yields, static_rates, epochs_per_year = result
            epochs_range = np.arange(total_epochs + 1)
            
            annual_inflation_rates = [rate * epochs_per_year for rate in inflation_rates]
            annual_yields = [yield_val * epochs_per_year for yield_val in yields]
            
            axes[0, 0].plot(epochs_range[:-1], supplies[:-1], label=scenario_name)
            axes[0, 0].axvspan(0, delay_epochs, color='gray', alpha=0.1, label='Delay Period' if scenario_name == list(scenarios.keys())[0] else None)
            axes[0, 0].set_xlabel('Epochs (2 days per epoch)')
            axes[0, 0].set_ylabel('Total Token Supply (SOL)')
            axes[0, 0].set_title('Total Supply')
            axes[0, 0].grid(True)
            
            ax_year = axes[0, 0].twiny()
            ax_year.set_xlim(axes[0, 0].get_xlim())
            ax_year.set_xticks([0, epochs_per_year/2, epochs_per_year, 3*epochs_per_year/2, 2*epochs_per_year])
            ax_year.set_xticklabels(['0', '0.5', '1', '1.5', '2'])
            ax_year.set_xlabel('Years')
            
            axes[0, 1].plot(epochs_range, staking_rates, label=scenario_name)
            axes[0, 1].set_xlabel('Epochs (2 days per epoch)')
            axes[0, 1].set_ylabel('Staking Rate')
            axes[0, 1].set_title('Staking Rate')
            axes[0, 1].grid(True)
            axes[0, 1].yaxis.set_major_formatter(PercentFormatter(1.0))
            
            axes[1, 0].plot(epochs_range[:-1], annual_inflation_rates, label=scenario_name)
            axes[1, 0].set_xlabel('Epochs (2 days per epoch)')
            axes[1, 0].set_ylabel('Annual Inflation Rate')
            axes[1, 0].set_title('Inflation Rate')
            axes[1, 0].grid(True)
            axes[1, 0].yaxis.set_major_formatter(PercentFormatter(1.0))
            
            axes[1, 1].plot(epochs_range[:-1], annual_yields, label=scenario_name)
            axes[1, 1].set_xlabel('Epochs (2 days per epoch)')
            axes[1, 1].set_ylabel('Annual Effective Yield')
            axes[1, 1].set_title('Staking Yield')
            axes[1, 1].grid(True)
            axes[1, 1].yaxis.set_major_formatter(PercentFormatter(1.0))
            
            legend_entries.append(scenario_name)
            
            # Add shock markers if applicable
            if "shock_epochs" in scenario_params and scenario_params["shock_epochs"]:
                for shock_epoch in scenario_params["shock_epochs"]:
                    for ax in axes.flat:
                        ax.axvline(x=shock_epoch, color='red', linestyle='--', alpha=0.3,
                                 label=f'Shock ({scenario_name})' if shock_epoch == scenario_params["shock_epochs"][0] else None)
        
        fig.legend(legend_entries, loc='lower center', ncol=len(scenarios), bbox_to_anchor=(0.5, 0.02))
        plt.tight_layout(rect=[0, 0.05, 1, 0.95])
        plt.subplots_adjust(top=0.9, bottom=0.1)
        plt.show()

# ------------------------------------------------------------------------------
# Compare Worst-Case Scenarios Function
# ------------------------------------------------------------------------------
def compare_worst_cases(worst_cases, horizon="24 Months", delay_name="4 Months"):
    """
    Compare the worst-case scenarios to analyze the impact of penalties.
    """
    total_epochs = {"6 Months": 91, "12 Months": 182, "24 Months": 364}[horizon]
    delay_epochs = {"Now": 0, "4 Months": 60, "8 Months": 120, "1 Year": 182}[delay_name]
    
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    fig.suptitle(f'Worst-Case Comparison - {horizon} (Delay: {delay_name})', fontsize=16)
    legend_entries = []
    
    for scenario_name, scenario_params in worst_cases.items():
        result = simulate_with_delay(
            initial_supply=550e6,
            initial_staking=scenario_params["initial_staking"],
            r0=0.0468,
            r_floor=0.015,
            c=np.sqrt(1/3) / (1 - np.sqrt(2/3)),
            total_epochs=total_epochs,
            delay_epochs=delay_epochs,
            adjustment_factor=0.03,
            shock_epochs=scenario_params["shock_epochs"],
            shock_factor=scenario_params["shock_factor"],
            penalty_threshold=scenario_params.get("penalty_threshold", None),
            extra_penalty=scenario_params.get("extra_penalty", 0.0),
            validator_commission=scenario_params.get("validator_commission", 0.0412),
            annual_to_epoch=True,
            min_staking=0.1,
            dilution_impact=0.15
        )
        
        supplies, staking_rates, inflation_rates, yields, static_rates, epochs_per_year = result
        epochs_range = np.arange(total_epochs + 1)
        
        annual_inflation_rates = [rate * epochs_per_year for rate in inflation_rates]
        annual_yields = [yield_val * epochs_per_year for yield_val in yields]
        
        axes[0, 0].plot(epochs_range[:-1], supplies[:-1], label=scenario_name)
        axes[0, 0].axvspan(0, delay_epochs, color='gray', alpha=0.1, label='Delay Period' if scenario_name == "Worst-Case" else None)
        axes[0, 0].set_xlabel('Epochs (2 days per epoch)')
        axes[0, 0].set_ylabel('Total Token Supply (SOL)')
        axes[0, 0].set_title('Total Supply')
        axes[0, 0].grid(True)
        
        ax_year = axes[0, 0].twiny()
        ax_year.set_xlim(axes[0, 0].get_xlim())
        ax_year.set_xticks([0, epochs_per_year/2, epochs_per_year, 3*epochs_per_year/2, 2*epochs_per_year])
        ax_year.set_xticklabels(['0', '0.5', '1', '1.5', '2'])
        ax_year.set_xlabel('Years')
        
        axes[0, 1].plot(epochs_range, staking_rates, label=scenario_name)
        axes[0, 1].set_xlabel('Epochs (2 days per epoch)')
        axes[0, 1].set_ylabel('Staking Rate')
        axes[0, 1].set_title('Staking Rate')
        axes[0, 1].grid(True)
        axes[0, 1].yaxis.set_major_formatter(PercentFormatter(1.0))
        
        axes[1, 0].plot(epochs_range[:-1], annual_inflation_rates, label=scenario_name)
        axes[1, 0].set_xlabel('Epochs (2 days per epoch)')
        axes[1, 0].set_ylabel('Annual Inflation Rate')
        axes[1, 0].set_title('Inflation Rate')
        axes[1, 0].grid(True)
        axes[1, 0].yaxis.set_major_formatter(PercentFormatter(1.0))
        
        axes[1, 1].plot(epochs_range[:-1], annual_yields, label=scenario_name)
        axes[1, 1].set_xlabel('Epochs (2 days per epoch)')
        axes[1, 1].set_ylabel('Annual Effective Yield')
        axes[1, 1].set_title('Staking Yield')
        axes[1, 1].grid(True)
        axes[1, 1].yaxis.set_major_formatter(PercentFormatter(1.0))
        
        legend_entries.append(scenario_name)
        
        # Add shock markers
        for shock_epoch in scenario_params["shock_epochs"]:
            for ax in axes.flat:
                ax.axvline(x=shock_epoch, color='red', linestyle='--', alpha=0.3)
    
    fig.legend(legend_entries, loc='lower center', ncol=len(worst_cases), bbox_to_anchor=(0.5, 0.02))
    plt.tight_layout(rect=[0, 0.05, 1, 0.95])
    plt.subplots_adjust(top=0.9, bottom=0.1)
    plt.show()

# ------------------------------------------------------------------------------
# Sensitivity Analysis Function
# ------------------------------------------------------------------------------
def sensitivity_analysis(param_name, param_values, scenario_params, common_params, 
                        horizon_name="24 Months", delay_name="4 Months"):
    """
    Perform sensitivity analysis for a given parameter.
    """
    total_epochs = {"6 Months": 91, "12 Months": 182, "24 Months": 364}[horizon_name]
    delay_epochs = {"Now": 0, "4 Months": 60, "8 Months": 120, "1 Year": 182}[delay_name]
    
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    fig.suptitle(f'Sensitivity Analysis for {param_name} - {horizon_name} (Delay: {delay_name})', fontsize=16)
    legend_entries = []
    
    for param_value in param_values:
        if param_name == 'c':
            modified_common_params = common_params.copy()
            modified_common_params['c'] = param_value
            modified_scenario_params = scenario_params.copy()
        elif param_name == 'adjustment_factor':
            modified_common_params = common_params.copy()
            modified_common_params['adjustment_factor'] = param_value
            modified_scenario_params = scenario_params.copy()
        elif param_name == 'initial_staking':
            modified_common_params = common_params.copy()
            modified_scenario_params = scenario_params.copy()
            modified_scenario_params['initial_staking'] = param_value
        elif param_name == 'min_staking':
            modified_common_params = common_params.copy()
            modified_common_params['min_staking'] = param_value
            modified_scenario_params = scenario_params.copy()
        elif param_name == 'dilution_impact':
            modified_common_params = common_params.copy()
            modified_common_params['dilution_impact'] = param_value
            modified_scenario_params = scenario_params.copy()
        else:
            raise ValueError(f"Unknown parameter: {param_name}")
        
        result = simulate_with_delay(
            modified_common_params["initial_supply"],
            modified_scenario_params["initial_staking"],
            modified_common_params["r0"],
            modified_common_params["r_floor"],
            modified_common_params["c"],
            total_epochs,
            delay_epochs,
            adjustment_factor=modified_common_params["adjustment_factor"],
            shock_epochs=modified_scenario_params.get("shock_epochs", None),
            shock_factor=modified_scenario_params.get("shock_factor", 1.0),
            penalty_threshold=modified_scenario_params.get("penalty_threshold", None),
            extra_penalty=modified_scenario_params.get("extra_penalty", 0.0),
            validator_commission=modified_scenario_params.get("validator_commission", 0.0412),
            annual_to_epoch=True,
            min_staking=modified_common_params.get("min_staking", 0.1),
            dilution_impact=modified_common_params.get("dilution_impact", 0.15)
        )
        
        supplies, staking_rates, inflation_rates, yields, static_rates, epochs_per_year = result
        epochs_range = np.arange(total_epochs + 1)
        
        annual_inflation_rates = [rate * epochs_per_year for rate in inflation_rates]
        annual_yields = [yield_val * epochs_per_year for yield_val in yields]
        
        axes[0, 0].plot(epochs_range[:-1], supplies[:-1], label=f'{param_name}={param_value}')
        axes[0, 0].axvspan(0, delay_epochs, color='gray', alpha=0.1, label='Delay Period' if param_value == param_values[0] else None)
        axes[0, 0].set_xlabel('Epochs (2 days per epoch)')
        axes[0, 0].set_ylabel('Total Token Supply (SOL)')
        axes[0, 0].set_title('Total Supply')
        axes[0, 0].grid(True)
        
        axes[0, 1].plot(epochs_range, staking_rates, label=f'{param_name}={param_value}')
        axes[0, 1].set_xlabel('Epochs (2 days per epoch)')
        axes[0, 1].set_ylabel('Staking Rate')
        axes[0, 1].set_title('Staking Rate')
        axes[0, 1].grid(True)
        axes[0, 1].yaxis.set_major_formatter(PercentFormatter(1.0))
        
        axes[1, 0].plot(epochs_range[:-1], annual_inflation_rates, label=f'{param_name}={param_value}')
        axes[1, 0].set_xlabel('Epochs (2 days per epoch)')
        axes[1, 0].set_ylabel('Annual Inflation Rate')
        axes[1, 0].set_title('Inflation Rate')
        axes[1, 0].grid(True)
        axes[1, 0].yaxis.set_major_formatter(PercentFormatter(1.0))
        
        axes[1, 1].plot(epochs_range[:-1], annual_yields, label=f'{param_name}={param_value}')
        axes[1, 1].set_xlabel('Epochs (2 days per epoch)')
        axes[1, 1].set_ylabel('Annual Effective Yield')
        axes[1, 1].set_title('Staking Yield')
        axes[1, 1].grid(True)
        axes[1, 1].yaxis.set_major_formatter(PercentFormatter(1.0))
        
        legend_entries.append(f'{param_name}={param_value}')
    
    fig.legend(legend_entries, loc='lower center', ncol=len(param_values), bbox_to_anchor=(0.5, 0.02))
    plt.tight_layout(rect=[0, 0.05, 1, 0.95])
    plt.subplots_adjust(top=0.9, bottom=0.1)
    plt.show()

# ------------------------------------------------------------------------------
# Main Function
# ------------------------------------------------------------------------------
def main():
    horizons = {
        "6 Months": 91,
        "12 Months": 182,
        "24 Months": 364
    }
    
    delays = {
        "Now": 0,
        "4 Months": 60,
        "8 Months": 120,
        "1 Year": 182
    }
    
    scenarios = {
        "Best-Case": {
            "initial_staking": 0.70,
            "shock_epochs": [],  # No shocks
            "shock_factor": 1.0,
            "penalty_threshold": None,
            "extra_penalty": 0.0,
            "validator_commission": 0.0587  # Superminority (5.87%) for high-stake validators
        },
        "Median-Case": {
            "initial_staking": 0.65,
            "shock_epochs": [200],  # Shock at epoch 200
            "shock_factor": 0.8,  # 20% decrease in staking rate
            "penalty_threshold": None,
            "extra_penalty": 0.0,
            "validator_commission": 0.0438  # Supermajority (4.38%) for moderate-stake validators
        },
        "Worst-Case": {
            "initial_staking": 0.30,
            "shock_epochs": [100, 300],  # Shocks at epochs 100, 300 (within 24 months)
            "shock_factor": 0.7,  # 30% decrease in staking rate
            "penalty_threshold": 0.035,
            "extra_penalty": 0.03,  # Original penalty (3%)
            "validator_commission": 0.0366  # Long-tail (3.66%) for low-stake validators
        },
        "Worst-Case-No-Penalties": {
            "initial_staking": 0.30,
            "shock_epochs": [100, 300],  # Same shocks
            "shock_factor": 0.7,  # 30% decrease in staking rate
            "penalty_threshold": None,  # No penalties
            "extra_penalty": 0.0,
            "validator_commission": 0.0366  # Long-tail (3.66%) for low-stake validators
        },
        "Worst-Case-Small-Penalties": {
            "initial_staking": 0.30,
            "shock_epochs": [100, 300],  # Same shocks
            "shock_factor": 0.7,  # 30% decrease in staking rate
            "penalty_threshold": 0.035,  # Same threshold
            "extra_penalty": 0.01,  # Smaller penalty (1%)
            "validator_commission": 0.0366  # Long-tail (3.66%) for low-stake validators
        }
    }
    
    common_params = {
        "initial_supply": 550e6,
        "r0": 0.0468, 
        "r_floor": 0.015,
        "c": np.sqrt(1/3) / (1 - np.sqrt(2/3)),
        "adjustment_factor": 0.03,
        "min_staking": 0.1,  # Minimum staking rate representing core long-term holders
        "dilution_impact": 0.3
    }
    
    # 1. Run individual scenario simulations
    for scenario_name, scenario_params in scenarios.items():
        plot_scenario(scenario_name, scenario_params, horizons, delays, common_params)
    
    # 2. Compare all scenarios
    compare_scenarios(scenarios, horizons, common_params)
    
    # 3. Compare worst-case scenarios specifically
    worst_cases = {
        "Worst-Case": scenarios["Worst-Case"],
        "Worst-Case-No-Penalties": scenarios["Worst-Case-No-Penalties"],
        "Worst-Case-Small-Penalties": scenarios["Worst-Case-Small-Penalties"]
    }
    compare_worst_cases(worst_cases)
    
    # 4. Sensitivity analysis for various parameters
    # Parameter c (constant in SIMD-228 formula)
    exact_c = np.sqrt(1/3) / (1 - np.sqrt(2/3))
    sensitivity_analysis('c', [2.5, 3.0, 3.3, 3.5, exact_c, 4.0], 
                        scenarios["Median-Case"], common_params)
    
    # Adjustment factor (responsiveness of staking to yield changes)
    sensitivity_analysis('adjustment_factor', [0.01, 0.03, 0.05, 0.07, 0.10], 
                        scenarios["Median-Case"], common_params)
    
    # Initial staking rate
    sensitivity_analysis('initial_staking', [0.3, 0.4, 0.5, 0.6, 0.7], 
                        scenarios["Median-Case"], common_params)
    
    # Minimum staking level (core holders)
    sensitivity_analysis('min_staking', [0.05, 0.1, 0.15, 0.2, 0.25], 
                        scenarios["Median-Case"], common_params)
    
    # Dilution impact
    sensitivity_analysis('dilution_impact', [0.1, 0.15, 0.2, 0.25, 0.3], 
                        scenarios["Worst-Case"], common_params)

if __name__ == "__main__":
    main()