import sqlite3
import pandas as pd
import numpy as np
from scipy.stats import norm
from base_market_data_loader import load_base_market_data, get_available_pricing_dates

def rating_to_debt_ratio(rating):
    mapping = {'AAA': 0.40, 'AA': 0.50, 'A': 0.60, 'BBB': 0.70, 'BB': 0.85, 'B': 0.95}
    return mapping.get(str(rating).upper(), 0.65)

def merton_model(V, D, r, T, sigma_V):
    d1 = (np.log(V / D) + (r + 0.5 * sigma_V**2) * T) / (sigma_V * np.sqrt(T))
    d2 = d1 - sigma_V * np.sqrt(T)
    E = V * norm.cdf(d1) - D * np.exp(-r * T) * norm.cdf(d2)
    sigma_E = (V / E) * norm.cdf(d1) * sigma_V if E > 0 else 0
    
    term = norm.cdf(d2) + (V / D) * np.exp(r * T) * norm.cdf(-d1)
    if term <= 0:
        spread_bps = 10000.0
    else:
        spread_bps = (-1.0 / T) * np.log(term) * 10000.0
        
    return E, sigma_E, spread_bps, norm.cdf(-d2)

def compute_market_shocks():
    print("--- Starting ISDA Merton Full Revaluation Pre-Calc Engine (Multi-Date) ---")
    
    target_dates = get_available_pricing_dates()
    # For UI testing speed, limit to recent 3 dates.
    target_dates = target_dates[:3]
    print(f"Discovered Dates for Shock Generation: {target_dates}")
    
    conn = sqlite3.connect('C:/Users/Variet-Worker/Desktop/climate_risk/data/climate_risk.db')
    ngfs_df = pd.read_sql('SELECT * FROM ngfs_macro_time_series', conn)
    ksic_df = pd.read_sql('SELECT * FROM ksic_carbon_beta', conn)
    
    results_credit = []
    results_equity = []
    results_curve = []
    results_fx = []
    
    scenarios = ngfs_df['scenario'].unique()
    years = sorted(ngfs_df['year'].unique())
    r_base = 0.03
    T_mat = 5.0
    
    for t_date in target_dates:
        print(f"Processing baseline date: {t_date}")
        base_data = load_base_market_data(t_date)
        assets_df = base_data['assets']
        df_yc = base_data['yield_curve']
        rate_vols = base_data.get('rate_vols', pd.DataFrame())
        
        for sc in scenarios:
            sc_data = ngfs_df[ngfs_df['scenario'] == sc]
            
            base_gdp = sc_data[(sc_data['variable'] == 'GDP|MER') & (sc_data['year'] == '2025')]['value'].values[0]
            base_co2 = sc_data[(sc_data['variable'] == 'Emissions|CO2|Price') & (sc_data['year'] == '2025')]['value'].values[0]
            base_rate = sc_data[(sc_data['variable'] == 'Policy Rate|Short-term') & (sc_data['year'] == '2025')]['value'].values[0]
            
            for yr in years:
                val_gdp = sc_data[(sc_data['variable'] == 'GDP|MER') & (sc_data['year'] == yr)]['value'].values[0]
                val_co2 = sc_data[(sc_data['variable'] == 'Emissions|CO2|Price') & (sc_data['year'] == yr)]['value'].values[0]
                val_rate = sc_data[(sc_data['variable'] == 'Policy Rate|Short-term') & (sc_data['year'] == yr)]['value'].values[0]
                
                gdp_shock = (val_gdp - base_gdp) / base_gdp 
                co2_jump = max(0, val_co2 - base_co2) / 1000.0 
                rate_delta_decimal = (val_rate - base_rate) / 100.0
                r_stress_proxy = r_base + rate_delta_decimal
                
                for index, asset in assets_df.iterrows():
                    sector = asset['sector']
                    beta_row = ksic_df[ksic_df['ksic_code'] == sector]
                    c_beta = beta_row['carbon_beta'].values[0] if len(beta_row) > 0 else 1.0
                    
                    V_base = float(asset['base_price'])
                    sigma_V_base = float(asset['base_vol'])
                    
                    if asset['asset_type'].lower() == 'fx':
                        fx_beta = 1.5 if sector == 'EXPORT' else -0.8
                        dx = gdp_shock * fx_beta
                        base_spot = V_base
                        stress_spot = base_spot * (1 + dx)
                        base_vol = sigma_V_base
                        stress_vol = base_vol * (1 + abs(dx) * 2.5) 
                        results_fx.append({
                            'pricing_date': t_date, 'scenario': sc, 'year': yr, 'asset_code': asset['asset_code'],
                            'base_spot': base_spot, 'stress_spot': stress_spot, 
                            'base_vol': base_vol, 'stress_vol': stress_vol
                        })
                    elif asset['asset_type'].lower() in ['bond', 'equity', 'index'] or 'KS200' in str(asset['asset_code']):
                        D_ratio = rating_to_debt_ratio(asset['credit_rating'])
                        D = V_base * D_ratio
                        
                        E0, vol0, spread0, pd0 = merton_model(V_base, D, r_base, T_mat, sigma_V_base)
                        
                        dV_ratio = (gdp_shock * c_beta) - (co2_jump * c_beta)
                        V_stress = max(V_base * (1 + dV_ratio), V_base * 0.01)
                        sigma_V_stress = sigma_V_base
                        
                        E1, vol1, spread1, pd1 = merton_model(V_stress, D, r_stress_proxy, T_mat, sigma_V_stress)
                        
                        if asset['asset_type'].lower() in ['equity', 'index'] or 'KS200' in str(asset['asset_code']):
                            results_equity.append({
                                'pricing_date': t_date, 'scenario': sc, 'year': yr, 'asset_code': asset['asset_code'],
                                'base_price': E0, 'stress_price': E1,
                                'base_implied_vol': vol0, 'stress_implied_vol': vol1
                            })
                        else:
                            results_credit.append({
                                'pricing_date': t_date, 'scenario': sc, 'year': yr, 'asset_code': asset['asset_code'],
                                'base_spread_bps': max(spread0, 0), 'stress_spread_bps': max(spread1, 0),
                                'base_pd': pd0, 'stress_pd': pd1
                            })
                
                for index, point in df_yc.iterrows():
                    currency = point.get('currency', 'KRW')
                    tenor_days = point['tenor_days']
                    base_zero = point['zero_rate']
                    curve_name = point.get('curve_name', f'YIELD_CURVE_{tenor_days}D')
                    
                    if not rate_vols.empty:
                        curr_vols = rate_vols[rate_vols['currency'] == currency].copy()
                        if not curr_vols.empty:
                            curr_vols['tenor_diff'] = abs(curr_vols['tenor_days'] - tenor_days)
                            best_match = curr_vols.sort_values('tenor_diff').iloc[0]
                            base_hw_sigma = float(best_match['hw_sigma'])
                            hw_kappa = float(best_match['hw_kappa_1f'])
                        else:
                            base_hw_sigma = 0.01
                            hw_kappa = 0.05
                    else:
                        base_hw_sigma = 0.01
                        hw_kappa = 0.05
                    
                    T_yrs = tenor_days / 365.0
                    B_factor = (1 - np.exp(-hw_kappa * T_yrs)) / hw_kappa if hw_kappa > 0 else T_yrs
                    zero_rate_shift = B_factor / T_yrs if T_yrs > 0 else 1.0
                    stress_zero = base_zero + (rate_delta_decimal * zero_rate_shift)
                    stress_hw_sigma = base_hw_sigma * (1.30 if rate_delta_decimal > 0.02 else 1.0)
                    
                    results_curve.append({
                        'pricing_date': t_date, 'scenario': sc, 'year': yr, 'curve_name': curve_name, 'tenor_days': tenor_days,
                        'base_zero': base_zero, 'stress_zero': stress_zero, 
                        'base_hw_sigma': base_hw_sigma, 'stress_hw_sigma': stress_hw_sigma
                    })
                        
    if len(results_equity) > 0:
        print("Saving equity...")
        pd.DataFrame(results_equity).to_sql('shocks_equity', conn, if_exists='replace', index=False)
    if len(results_credit) > 0:
        print("Saving credit...")
        pd.DataFrame(results_credit).to_sql('shocks_credit', conn, if_exists='replace', index=False)
    if len(results_fx) > 0:
        print("Saving fx...")
        pd.DataFrame(results_fx).to_sql('shocks_fx', conn, if_exists='replace', index=False)
    if len(results_curve) > 0:
        print("Saving curve...")
        pd.DataFrame(results_curve).to_sql('shocks_curve', conn, if_exists='replace', index=False)
    
    conn.close()
    print("Merton Stressed Parameters successfully built for ALL Pricing Dates.")

if __name__ == '__main__':
    compute_market_shocks()
