import psycopg
import sqlite3
import pandas as pd
import numpy as np
import json
import random

def get_available_pricing_dates():
    """Returns a list of available historical snapshot dates."""
    try:
        conn_deriva = psycopg.connect('postgresql://deriva:deriva_dev_2026@localhost:15400/deriva')
        query = "SELECT DISTINCT DATE(created_at) as c_date FROM yield_curve_points ORDER BY c_date DESC"
        df = pd.read_sql(query, conn_deriva)
        conn_deriva.close()
        return df['c_date'].astype(str).tolist()
    except Exception as e:
        print(f"Failed to fetch dates: {e}")
        return ['2026-03-16']

def load_base_market_data(target_date=None):
    print(f"Connecting to Deriva & Climate DB for ISDA Baseline on Date: {target_date or 'LATEST'}...")
    
    # Extract Base JSON from Deriva
    conn_deriva = psycopg.connect('postgresql://deriva:deriva_dev_2026@localhost:15400/deriva')
    try:
        if target_date:
            query_eval = f"SELECT dataset::text FROM eval_datasets WHERE DATE(created_at) <= '{target_date}' ORDER BY created_at DESC LIMIT 1"
        else:
            query_eval = "SELECT dataset::text FROM eval_datasets ORDER BY created_at DESC LIMIT 1"
            
        cur = conn_deriva.cursor()
        cur.execute(query_eval)
        res = cur.fetchone()
        eval_ds = json.loads(res[0]) if res else {'spots': {}, 'vols': {}}
        
        # HW Vols mapping
        df_rate_vols = pd.read_sql("SELECT currency, tenor_days, vol_1f as hw_sigma, hw_kappa_1f FROM rate_vol_configs", conn_deriva)
        
        # Curves (Fetch snapshot on target_date)
        date_filter = f"AND DATE(p.created_at) = '{target_date}'" if target_date else ""
        
        df_yc = pd.read_sql(f"""
            SELECT DISTINCT ON (c.curve_name, p.tenor_days)
                c.curve_name, c.currency, p.tenor_days, p.zero_rate 
            FROM yield_curves c 
            JOIN yield_curve_points p ON c.id = p.curve_id 
            WHERE c.curve_role = 'irs' {date_filter}
            ORDER BY c.curve_name, p.tenor_days ASC, p.created_at DESC
        """, conn_deriva)
    except Exception as e:
        print(f"Deriva Fetch Error: {e}")
        eval_ds = {'spots': {}, 'vols': {}}
        df_rate_vols = pd.DataFrame()
        df_yc = pd.DataFrame()
    finally:
        conn_deriva.close()
        
    # Extract Firm Reference mapping from Local DB
    conn_climate = sqlite3.connect('C:/Users/Variet-Worker/Desktop/climate_risk/data/climate_risk.db')
    try:
        df_assets = pd.read_sql("SELECT asset_code, ksic_code as sector, gics_sector, credit_rating FROM firm_reference_data", conn_climate)
        df_assets['asset_type'] = df_assets.apply(lambda x: 'FX' if x['gics_sector'] == 'Currency' else ('Bond' if x['gics_sector'] in ['Sovereign', 'Corporate Bond'] else ('Index' if x['gics_sector'] == 'Index' else ('Rate' if x['gics_sector'] == 'Rate' else 'Equity'))), axis=1)
    except Exception as e:
        print(f"Climate DB Error: {e}")
        df_assets = pd.DataFrame(columns=['asset_code', 'asset_type', 'sector', 'gics_sector', 'credit_rating'])
    
    conn_climate.close()
    
    # Load Live Bond Yields
    import os
    live_bond_yields = {}
    bond_path = 'C:/Users/Variet-Worker/Desktop/climate_risk/data/live_bond_yields.json'
    if os.path.exists(bond_path):
        try:
            with open(bond_path, 'r') as f:
                live_bond_yields = json.load(f).get('yields', {})
        except Exception:
            pass

    # Merge exact Spot/Vol limits
    def get_spot(code, asset_type):
        if code in live_bond_yields:
            return float(live_bond_yields[code])
        s = eval_ds['spots'].get(code)
        if s is not None: return float(s)
        # Mock randomized scaling using target_date seed to keep it consistent
        seed = int(target_date.replace('-', '')) if target_date else 0
        rand = random.Random(hash(code) + seed)
        return rand.uniform(100.0, 50000.0) if asset_type == 'Equity' else 100.0
        
    def get_vol(code, asset_type):
        v = eval_ds['vols'].get(code)
        if v is not None: return float(v)
        seed = int(target_date.replace('-', '')) if target_date else 0
        rand = random.Random(hash(code) + seed * 2)
        return rand.uniform(0.15, 0.40)
        
    df_assets['base_price'] = df_assets.apply(lambda r: get_spot(r['asset_code'], r['asset_type']), axis=1)
    df_assets['base_vol'] = df_assets.apply(lambda r: get_vol(r['asset_code'], r['asset_type']), axis=1)

    # Missing Curve check
    if len(df_yc) == 0:
        df_yc = pd.DataFrame({
            'curve_name': ['IRS-KRW-CD91D-CD91D']*7,
            'currency': ['KRW']*7,
            'tenor_days': [30, 90, 180, 365, 365*3, 365*5, 365*10],
            'zero_rate': [0.035, 0.0352, 0.0355, 0.036, 0.038, 0.040, 0.042]
        })
        
    if len(df_rate_vols) == 0:
        df_rate_vols = pd.DataFrame({'currency': ['KRW', 'USD'], 'tenor_days': [365, 365], 'hw_sigma': [0.01, 0.01], 'hw_kappa_1f': [0.05, 0.05]})

    # Do not inject fake mock data to inflate universe
    # Only return real assets present in SQLite firm_reference_data

    return {
        'yield_curve': df_yc,
        'assets': df_assets,
        'rate_vols': df_rate_vols
    }

if __name__ == '__main__':
    dates = get_available_pricing_dates()
    print("Available Dates:", dates)
    data = load_base_market_data(dates[0] if dates else None)
    print("\n[Base Yield Curve Points]")
    print(data['yield_curve'].head(3))
    print("\n[Base Asset Universe from Security Master & Eval DS]")
    print(data['assets'].head())
