config_schema.py

from pydantic import BaseModel, Field, model_validator
from typing import Dict, List, Any
from constants import DEFAULT_ADV

class ConfigError(Exception):
    pass

class BenchmarksConfig(BaseModel):
    equity: str = "SPY"
    volatility: str = "^VIX"
    risk_free: str = "^TNX"

    def get(self, key: str, default: Any = None) -> Any:
        return getattr(self, key, default)

class AppConfig(BaseModel):
    risk_free_rate: float = Field(0.04, ge=0.0, le=0.20)
    transaction_cost: float = Field(0.001, ge=0.0, le=0.05)
    trading_days_per_year: int = Field(252, ge=100, le=365)
    rolling_cov_days: int = 756
    currency_symbol: str = "$"
    default_adv_proxy: float = Field(DEFAULT_ADV, ge=0.0)
    
    benchmarks: BenchmarksConfig = Field(default_factory=BenchmarksConfig)
    
    single_asset_min: float = Field(-1.0, ge=-1.0, le=1.0)
    single_asset_max: float = Field(0.40, ge=0.01, le=1.0)
    sector_limit: float = Field(0.40, ge=0.01, le=1.0)
    gross_leverage_cap: float = Field(2.0, ge=1.0, le=5.0)
    short_borrow_cost: float = Field(0.015, ge=0.0, le=0.50)
    max_turnover: float = Field(3.0, ge=0.0)
    
    tax_rate_lt: float = Field(0.20, ge=0.0, le=1.0)
    tax_rate_st: float = Field(0.35, ge=0.0, le=1.0)
    lt_days: int = Field(366, ge=1)
    hrp_tax_lambda: float = Field(2.5, ge=0.0)
    
    cvar_alpha: float = Field(0.95, ge=0.50, le=0.999)
    cvar_lambda: float = Field(0.5, ge=0.0, le=20.0)
    baseline_risk_factor: float = Field(3.0, ge=0.1, le=25.0)
    monte_carlo_sims: int = Field(1500, ge=100, le=100_000)
    monte_carlo_years: float = Field(1.0, ge=0.1, le=50.0)
    
    garch_enabled: bool = True
    cvar_enabled: bool = True
    tax_enabled: bool = False
    dynamic_risk: bool = True
    hmm_regime: bool = True
    arima_enabled: bool = False
    anova_enabled: bool = False
    with_futures: bool = False
    overlay_mode: str = "beta_hedge"
    futures_universe: List[str] = Field(default_factory=lambda: ["MES", "ES"])
    futures_safety_multiplier: float = Field(3.0, ge=0.0, le=10.0)
    futures_target_beta: float = Field(0.0, ge=-5.0, le=5.0)
    futures_margin_headroom: float = Field(0.05, ge=0.0, le=1.0)
    return_frequency: str = "daily"
    
    # End-to-End Differentiable Optimization (Model 6)
    e2e_loss_type: str = "spo"       # "spo", "sharpe", or "calmar"
    e2e_epochs: int = Field(default=50)
    e2e_batch_size: int = Field(default=32)

    e2e_lr: float = 1e-3
    e2e_cache_dir: str = ".e2e_cache"
    
    universe_categories: Dict[str, List[str]] = Field(default_factory=lambda: {
        "Core Equities": ["SPY", "QQQ", "DIA", "IWM"],
        "Bonds & Rates": ["TLT", "IEF", "SHY", "AGG"],
        "Tech & Growth": ["AAPL", "MSFT", "NVDA", "TSLA"],
        "Defensive/Value": ["JNJ", "PG", "KO", "XLP"],
        "Commodities": ["GLD", "SLV", "USO", "PDBC"],
        "International": ["VEA", "VWO", "EFA", "EEM"],
        "Crypto Proxies": ["IBIT", "FBTC", "ETHE", "MSTR"]
    })
    
    # ─────────────────────────────────────────────
    # EXTENDED HISTORY & BOOTSTRAPPING
    # ─────────────────────────────────────────────
    extended_history: bool = False
    bootstrap_samples: int = 100
    stitch_overlap_days: int = 252
    proxy_mappings: Dict[str, Dict] = Field(default_factory=lambda: {
        'SPY': {'proxy': '^GSPC', 'proxy_start': '1950-01-03', 'overlap_days': 252},
        'TLT': {'proxy': '^TYX', 'proxy_start': '1977-01-03', 'is_yield': True},
        'GLD': {'proxy': 'GC=F', 'proxy_start': '1974-12-31'},
        'QQQ': {'proxy': '^IXIC', 'proxy_start': '1971-02-05'}
    })
    
    bond_metadata: Dict[str, Any] = Field(default_factory=dict)
    
    model_config = {"extra": "allow"}
    
    @model_validator(mode='after')
    def check_logic(self):
        if self.single_asset_min > self.single_asset_max:
            raise ConfigError("single_asset_min cannot be greater than single_asset_max.")
        if self.sector_limit < self.single_asset_max:
            raise ConfigError(f"sector_limit ({self.sector_limit}) cannot be smaller than single_asset_max ({self.single_asset_max}).")
        if self.tax_rate_lt > self.tax_rate_st:
            raise ConfigError("Long-term tax rate is mathematically expected to be <= short-term tax rate.")
        return self
        
    def get(self, key: str, default: Any = None) -> Any:
        """Backward compatibility for dict-like access."""
        if hasattr(self, key):
            return getattr(self, key)
        if self.model_extra is not None and key in self.model_extra:
            return self.model_extra[key]
        return default

    def setdefault(self, key: str, default: Any = None) -> Any:
        val = self.get(key, None)
        if val is None:
            self[key] = default
            return default
        return val

    def __getitem__(self, key: str) -> Any:
        if hasattr(self, key):
            return getattr(self, key)
        if self.model_extra is not None and key in self.model_extra:
            return self.model_extra[key]
        raise KeyError(key)

    def __setitem__(self, key: str, value: Any) -> None:
        if hasattr(self, key) or key in self.model_fields:
            setattr(self, key, value)
        elif key.startswith('_'):
            # Internal transient keys (e.g., _risk_input, _is_historical_backtest) are allowed silently
            if self.model_extra is None:
                self.__dict__['__pydantic_extra__'] = {}
            self.model_extra[key] = value
        else:
            import warnings
            warnings.warn(
                f"AppConfig: setting unknown key '{key}'. Did you mean one of "
                f"{sorted(list(self.model_fields.keys()))[:8]}...?",
                stacklevel=2
            )
            if self.model_extra is None:
                self.__dict__['__pydantic_extra__'] = {}
            self.model_extra[key] = value
        
    def update(self, other: dict) -> None:
        for k, v in other.items():
            setattr(self, k, v)

DEFAULT_CONFIG = AppConfig()