data_repository.py

from dataclasses import dataclass
from typing import List, Dict, Tuple, Optional
import pandas as pd
import logging

from data import fetch_fama_french_factors, fetch_data, fetch_risk_free_rate, fetch_risk_free_series, build_monthly_returns
from database import get_pg_engine
from core_engine import build_spread_map, load_portfolio_state_dict
from core_types import PortfolioState
from config import AppConfig

logger = logging.getLogger('portfolio_engine')

@dataclass
class DataSnapshot:
    opt_tickers: List[str]
    opt_returns_df: pd.DataFrame
    bench_rets_monthly: pd.Series
    opt_ff_df: Optional[pd.DataFrame]
    rfr: float
    vol_raw: Optional[pd.Series]
    display_df: pd.DataFrame
    bench_display: pd.Series
    master_state: PortfolioState
    spread_map: Dict[str, float]
    train_yrs: float
    test_yrs: float
    # Full-resolution data needed by the pipeline orchestrator
    returns_df: pd.DataFrame = None
    bench_rets: pd.Series = None
    raw: Dict = None
    prices: Dict = None
    eq_bench: str = "SPY"
    vol_bench: str = "^VIX"
    rfr_bench: str = "^TNX"


def guess_currency(ticker: str) -> str:
    if ticker.endswith('.AT') or ticker.endswith('.DE') or ticker.endswith('.PA') or ticker.endswith('.MI') or ticker.endswith('.AS') or ticker.endswith('.MC'): return 'EUR'
    if ticker.endswith('.L'): return 'GBP'
    if ticker.endswith('.T'): return 'JPY'
    if ticker.endswith('.AX'): return 'AUD'
    if ticker.endswith('.TO'): return 'CAD'
    if ticker.endswith('.SW'): return 'CHF'
    return 'USD'

class DataRepository:

    """
    Repository layer responsible for fetching, cleaning, and assembling 
    all market data, benchmarks, and portfolio state required by the engine.
    """
    def __init__(self, cfg: AppConfig):
        self.cfg = cfg
        self.trading_days = self.cfg.get("trading_days_per_year", 252)

    def fetch_all(self, input_tickers: List[str], model_id: int) -> DataSnapshot:
        b = self.cfg.get("benchmarks", {})
        eq_bench = b.get("equity", "SPY")
        vol_bench = b.get("volatility", "^VIX")
        rfr_bench = b.get("risk_free", "^TNX")
        
        # --- Options Parsing Logic ---
        original_input_tickers = list(input_tickers)
        option_tickers = [t for t in input_tickers if t.startswith('OPT:')]
        standard_tickers = [t for t in input_tickers if not t.startswith('OPT:')]
        for ot in option_tickers:
            parts = ot.split(':')
            if len(parts) >= 2:
                underlying = parts[1]
                if underlying not in standard_tickers:
                    standard_tickers.append(underlying)
        input_tickers = standard_tickers
        # ----------------------
        
        # --- Currency Logic ---
        curr_sym = self.cfg.get("currency_symbol", "$")
        currency_map = {'$': 'USD', '€': 'EUR', '£': 'GBP', '¥': 'JPY', 'CHF': 'CHF'}
        base_currency = currency_map.get(curr_sym, 'USD')
        
        fx_tickers_needed = set()
        ticker_to_currency = {}
        for t in input_tickers:
            c = guess_currency(t)
            ticker_to_currency[t] = c
            if c != base_currency and c != 'USD':
                fx_tickers_needed.add(f"{c}{base_currency}=X")
            elif c == 'USD' and base_currency != 'USD':
                # e.g. base is EUR, ticker is AAPL(USD), need USDEUR=X
                fx_tickers_needed.add(f"USD{base_currency}=X")
        
        augmented_tickers = list(input_tickers) + list(fx_tickers_needed)
        # ----------------------
        
        ff_df = fetch_fama_french_factors() if model_id in [4, 5] else None
        
        years_to_fetch = self.cfg.get('data_history_years', 15.0)
        valid_tickers = fetch_data(augmented_tickers, b, years=years_to_fetch, cfg=self.cfg.model_dump() if hasattr(self.cfg, 'model_dump') else dict(self.cfg))
        
        self.cfg["risk_free_rate"] = fetch_risk_free_rate(rfr_bench, self.cfg.get("risk_free_rate", 0.05))
        rfr_series = fetch_risk_free_series(rfr_bench)
        final_rfr = rfr_series if not rfr_series.empty else self.cfg["risk_free_rate"]
        
        macro_tickers = {eq_bench, vol_bench, rfr_bench, "^IRX"}
        opt_tickers = [t for t in valid_tickers if t not in macro_tickers or t in input_tickers]
        if not opt_tickers:
            raise SystemExit("No Usable Tickers found.")
            
        legacy_state_dict = load_portfolio_state_dict() if self.cfg.get('_use_saved_basis', False) else {}
        
        pg_engine = get_pg_engine()
        raw, prices = {}, {}
        all_tks = list(dict.fromkeys(opt_tickers + list(macro_tickers)))
        
        if all_tks:
            ph = ",".join(["?"] * len(all_tks)) if pg_engine.name == 'sqlite' else ",".join(["%s"] * len(all_tks))
            query = f"SELECT ticker, date, close_price FROM daily_prices WHERE ticker IN ({ph}) ORDER BY date"
            try:
                df_all = pd.read_sql(query, pg_engine, params=tuple(all_tks))
                if not df_all.empty:
                    df_all['date'] = pd.to_datetime(df_all['date'])
                    for t, group in df_all.groupby('ticker'):
                        group = group.drop_duplicates(subset=['date'], keep='last')
                        prices[t] = group.iloc[-1]['close_price']
                        raw[t] = group.set_index('date')['close_price']
            except Exception as e:
                logger.warning(f"DB read failed, returning empty context: {e}")
                
        # --- Apply Currency Conversion ---
        for t in input_tickers:
            if t not in raw: continue
            c = ticker_to_currency.get(t, 'USD')
            if c != base_currency:
                fx_pair = f"{c}{base_currency}=X"
                if c == 'USD':
                    fx_pair = f"USD{base_currency}=X"
                if fx_pair in raw:
                    # Align indices and multiply
                    fx_series = raw[fx_pair]
                    aligned_fx = fx_series.reindex(raw[t].index).ffill().bfill()
                    raw[t] = raw[t] * aligned_fx
                    prices[t] = prices[t] * aligned_fx.iloc[-1]
                else:
                    logger.warning(f"FX pair {fx_pair} missing for {t}. Using unadjusted {c} prices.")
        # ---------------------------------
        
        # Enforce forward-fill to bridge mismatched international trading holidays (US vs EU).
        for k in raw:
            raw[k] = raw[k].ffill()
        
        bench_rets = raw[eq_bench].pct_change().dropna() if eq_bench in raw else pd.Series(dtype=float)
        vol_raw = raw.get(vol_bench, None)
        
        MIN_DAYS_SHORT = self.trading_days * 2
        all_rets = {}
        for t in opt_tickers:
            if t not in raw: continue
            s = raw[t].pct_change().dropna(how='all')
            if len(s) >= MIN_DAYS_SHORT: all_rets[t] = s
            
        # --- Options Synthetic Generation ---
        import numpy as np
        from options_pricing import black_scholes_call, _d1_d2
        from scipy.stats import norm
        for ot in option_tickers:
            parts = ot.split(':')
            if len(parts) >= 5:
                # OPT:AAPL:C:150:0.5
                underlying = parts[1]
                op_type = parts[2].upper()
                try:
                    strike = float(parts[3])
                    ttm = float(parts[4])
                except:
                    continue
                if underlying in raw:
                    s_underlying = raw[underlying]
                    r_series = final_rfr.reindex(s_underlying.index).ffill().fillna(0.04) if isinstance(final_rfr, pd.Series) else pd.Series(final_rfr, index=s_underlying.index)
                    
                    synthetic_prices = []
                    for date, S in s_underlying.items():
                        rate = r_series.loc[date]
                        if op_type == 'C':
                            price = black_scholes_call(S, strike, ttm, rate, 0.20)
                        else:
                            # Put price
                            d1, d2 = _d1_d2(S, strike, ttm, rate, 0.20, 0.0)
                            price = strike * np.exp(-rate * ttm) * norm.cdf(-d2) - S * norm.cdf(-d1)
                        synthetic_prices.append(max(price, 1e-6)) # Avoid exact zero
                    
                    opt_series = pd.Series(synthetic_prices, index=s_underlying.index)
                    opt_rets = opt_series.pct_change().dropna()
                    opt_rets = opt_rets.replace([np.inf, -np.inf], 0.0).clip(-0.99, 5.0)
                    if len(opt_rets) >= MIN_DAYS_SHORT:
                        all_rets[ot] = opt_rets
                        if ot not in opt_tickers:
                            opt_tickers.append(ot)
        # ------------------------------------
              
        if not all_rets: raise SystemExit("No usable tickers with enough history.")
        
        returns_df = pd.DataFrame(all_rets)
        # Pad delisted assets
        for col in returns_df.columns:
            if returns_df[col].min() <= -0.99:
                dead_indices = returns_df[col][returns_df[col] <= -0.99].index
                if len(dead_indices) > 0:
                    dead_idx = dead_indices[0]
                    returns_df.loc[dead_idx:, col] = returns_df.loc[dead_idx:, col].fillna(0.0)
                    
        returns_df = returns_df.dropna()
        final_opt_tickers = [t for t in opt_tickers if t in returns_df.columns]
        spread_map = build_spread_map(final_opt_tickers, self.cfg.get("sector_map", {}))
        
        if self.cfg.get('return_frequency', 'daily') == 'monthly':
            opt_returns_df = build_monthly_returns(returns_df)
            bench_rets_monthly = build_monthly_returns(pd.DataFrame({eq_bench: bench_rets}))[eq_bench] if not bench_rets.empty else bench_rets
            opt_ff_df = build_monthly_returns(ff_df) if ff_df is not None else None
            self.cfg['_trading_periods'] = 12
        else:
            opt_returns_df, bench_rets_monthly, opt_ff_df = returns_df, bench_rets, ff_df
            self.cfg['_trading_periods'] = self.trading_days
            
        DISPLAY_DAYS = self.trading_days * 6
        display_df = returns_df.iloc[-DISPLAY_DAYS:] if len(returns_df) > DISPLAY_DAYS else returns_df
        bench_display = bench_rets.reindex(display_df.index).dropna()
        
        final_tickers = list(returns_df.columns)
        master_state = PortfolioState.build(final_tickers, prices, legacy_state_dict, self.cfg)
        
        total_days = len(returns_df)
        OOS_TEST_DAYS = int(min(total_days * 0.2, 504)) # Dynamic 20% test up to 2 years
        OOS_TRAIN_DAYS = max(100, total_days - OOS_TEST_DAYS)
        train_yrs = OOS_TRAIN_DAYS / self.trading_days
        test_yrs = OOS_TEST_DAYS / self.trading_days

        return DataSnapshot(
            opt_tickers=final_opt_tickers,
            opt_returns_df=opt_returns_df,
            bench_rets_monthly=bench_rets_monthly,
            opt_ff_df=opt_ff_df,
            rfr=final_rfr,
            vol_raw=vol_raw,
            display_df=display_df,
            bench_display=bench_display,
            master_state=master_state,
            spread_map=spread_map,
            train_yrs=train_yrs,
            test_yrs=test_yrs,
            returns_df=returns_df,
            bench_rets=bench_rets,
            raw=raw,
            prices=prices,
            eq_bench=eq_bench,
            vol_bench=vol_bench,
            rfr_bench=rfr_bench
        )