"""Scenario analysis engine for stress testing.

Custom shock vectors, cascading effects, historical scenarios,
scenario comparison, and probabilistic scenario analysis.

All functions are pure: data in, results out. No DB access, no API calls.
Requires numpy (core dependency).

Usage::

    from cpz.risk.scenario import custom_scenario, cascade_effects

    result = custom_scenario(shocks, weights, cov_matrix)
    cascaded = cascade_effects(primary_shock, transmission)
"""

from __future__ import annotations

import math
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field

import numpy as np


# ── Models ────────────────────────────────────────────────────────────


class ScenarioImpact(BaseModel):
    """Impact of a scenario on a portfolio."""
    total_impact_pct: float = 0.0
    per_position: Dict[str, float] = Field(default_factory=dict)
    expected_shortfall: float = 0.0
    max_drawdown_estimate: float = 0.0
    recovery_estimate_days: Optional[int] = None


class CascadeResult(BaseModel):
    """Cascading effects from a primary shock."""
    primary_shock: Dict[str, float] = Field(default_factory=dict)
    secondary_effects: Dict[str, float] = Field(default_factory=dict)
    total_shock: Dict[str, float] = Field(default_factory=dict)
    amplification_factor: float = 1.0


class ScenarioComparison(BaseModel):
    """Side-by-side comparison of multiple scenarios."""
    scenarios: Dict[str, ScenarioImpact] = Field(default_factory=dict)
    worst_case: str = ""
    best_case: str = ""


# ── Historical scenario definitions ──────────────────────────────────

HISTORICAL_SCENARIOS: Dict[str, Dict[str, float]] = {
    "covid_2020": {"equities": -0.34, "bonds": 0.05, "commodities": -0.20, "crypto": -0.50, "volatility": 3.0},
    "gfc_2008": {"equities": -0.57, "bonds": 0.15, "commodities": -0.35, "crypto": 0.0, "volatility": 4.0},
    "dotcom_2000": {"equities": -0.49, "bonds": 0.10, "commodities": -0.10, "crypto": 0.0, "volatility": 2.5},
    "black_monday_1987": {"equities": -0.22, "bonds": 0.04, "commodities": 0.0, "crypto": 0.0, "volatility": 5.0},
    "rate_shock_2022": {"equities": -0.24, "bonds": -0.10, "commodities": 0.25, "crypto": -0.55, "volatility": 1.8},
    "euro_crisis_2011": {"equities": -0.20, "bonds": 0.08, "commodities": -0.05, "crypto": 0.0, "volatility": 2.0},
    "flash_crash_2010": {"equities": -0.09, "bonds": 0.01, "commodities": -0.02, "crypto": 0.0, "volatility": 3.5},
    "rate_hike_200bp": {"equities": -0.15, "bonds": -0.12, "commodities": 0.05, "crypto": -0.20, "volatility": 1.5},
    "crypto_winter": {"equities": -0.05, "bonds": 0.02, "commodities": 0.0, "crypto": -0.70, "volatility": 1.2},
    "stagflation": {"equities": -0.25, "bonds": -0.08, "commodities": 0.30, "crypto": -0.30, "volatility": 2.0},
}

# Default transmission matrix for cascade effects
DEFAULT_TRANSMISSION = {
    "equities": {"bonds": -0.3, "commodities": 0.4, "crypto": 0.6, "volatility": -2.0},
    "bonds": {"equities": -0.2, "commodities": 0.1, "crypto": 0.1, "volatility": -0.5},
    "commodities": {"equities": -0.3, "bonds": -0.1, "crypto": 0.2, "volatility": -1.0},
    "crypto": {"equities": -0.1, "bonds": 0.0, "commodities": 0.0, "volatility": -0.3},
}

# Asset classification
_CRYPTO = frozenset(["BTC", "ETH", "SOL", "DOGE", "AVAX", "MATIC", "ADA", "XRP", "DOT", "LINK", "SHIB", "UNI", "AAVE"])
_BONDS = frozenset(["TLT", "IEF", "SHY", "AGG", "BND", "LQD", "HYG", "TIP", "GOVT"])
_COMMODITIES = frozenset(["GLD", "SLV", "USO", "UNG", "DBA", "DBC", "PDBC", "GSG"])


def _classify_asset(symbol: str) -> str:
    s = symbol.upper().replace("/", "").replace("-", "")
    if any(c in s for c in _CRYPTO):
        return "crypto"
    if s in _BONDS:
        return "bonds"
    if s in _COMMODITIES:
        return "commodities"
    return "equities"


# ── Public API ────────────────────────────────────────────────────────


def custom_scenario(
    shock_vector: Dict[str, float],
    position_weights: Dict[str, float],
    cov_matrix: Optional[Dict[str, Dict[str, float]]] = None,
    correlation_stress: float = 0.0,
) -> ScenarioImpact:
    """Apply a custom shock vector to portfolio positions.

    Args:
        shock_vector: Shocks by asset class {"equities": -0.20, "crypto": -0.50, ...}.
        position_weights: Position weights {symbol: weight}.
        cov_matrix: Optional covariance matrix for correlation-adjusted impact.
        correlation_stress: Increase correlations by this amount under stress (0-1).

    Returns:
        ScenarioImpact with total and per-position impact.
    """
    per_position: Dict[str, float] = {}
    total = 0.0

    for sym, weight in position_weights.items():
        asset_class = _classify_asset(sym)
        shock = shock_vector.get(asset_class, shock_vector.get("equities", 0.0))

        # Under stress, correlations increase — amplify impact
        if correlation_stress > 0:
            shock *= (1 + correlation_stress * 0.5)

        impact = abs(weight) * shock
        per_position[sym] = round(impact * 100, 2)
        total += impact

    es = total * 1.15  # Expected shortfall approximation
    recovery = max(5, int(abs(total) * 200)) if abs(total) > 0.01 else None

    return ScenarioImpact(
        total_impact_pct=round(total * 100, 2),
        per_position=per_position,
        expected_shortfall=round(es * 100, 2),
        max_drawdown_estimate=round(abs(total) * 120, 2),
        recovery_estimate_days=recovery,
    )


def cascade_effects(
    primary_shock: Dict[str, float],
    transmission_matrix: Optional[Dict[str, Dict[str, float]]] = None,
    iterations: int = 3,
    decay: float = 0.5,
) -> CascadeResult:
    """Model second-order cascading effects from a primary shock.

    Simulates how a shock in one asset class propagates to others through
    economic linkages (e.g., rate shock -> equity multiple compression).

    Args:
        primary_shock: Initial shocks by asset class.
        transmission_matrix: How shocks transmit between classes.
            Default uses empirical cross-asset transmission coefficients.
        iterations: Number of cascade rounds (default 3).
        decay: How quickly cascade effects diminish (default 0.5).

    Returns:
        CascadeResult with primary, secondary, and total effects.
    """
    transmission = transmission_matrix or DEFAULT_TRANSMISSION
    secondary: Dict[str, float] = {k: 0.0 for k in set(list(primary_shock.keys()) + list(transmission.keys()))}

    current = dict(primary_shock)
    for iteration in range(iterations):
        round_effects: Dict[str, float] = {}
        for source, shock_val in current.items():
            if source in transmission:
                for target, coeff in transmission[source].items():
                    effect = shock_val * coeff * (decay ** iteration)
                    round_effects[target] = round_effects.get(target, 0.0) + effect

        for target, effect in round_effects.items():
            secondary[target] = secondary.get(target, 0.0) + effect
        current = round_effects

    total: Dict[str, float] = {}
    all_keys = set(list(primary_shock.keys()) + list(secondary.keys()))
    for k in all_keys:
        total[k] = round(primary_shock.get(k, 0.0) + secondary.get(k, 0.0), 4)

    primary_magnitude = sum(abs(v) for v in primary_shock.values())
    total_magnitude = sum(abs(v) for v in total.values())
    amplification = total_magnitude / primary_magnitude if primary_magnitude > 0 else 1.0

    return CascadeResult(
        primary_shock={k: round(v, 4) for k, v in primary_shock.items()},
        secondary_effects={k: round(v, 4) for k, v in secondary.items()},
        total_shock=total,
        amplification_factor=round(amplification, 2),
    )


def historical_scenario(
    scenario_name: str,
    position_weights: Dict[str, float],
) -> ScenarioImpact:
    """Apply a named historical scenario to the portfolio.

    Args:
        scenario_name: One of the predefined scenarios (e.g., "covid_2020", "gfc_2008").
        position_weights: Current position weights.

    Returns:
        ScenarioImpact with the historical scenario applied.

    Raises:
        ValueError: If scenario_name is not recognized.
    """
    if scenario_name not in HISTORICAL_SCENARIOS:
        available = ", ".join(sorted(HISTORICAL_SCENARIOS.keys()))
        raise ValueError(f"Unknown scenario '{scenario_name}'. Available: {available}")

    shocks = HISTORICAL_SCENARIOS[scenario_name]
    return custom_scenario(shocks, position_weights)


def scenario_comparison(
    scenarios: Dict[str, Dict[str, float]],
    position_weights: Dict[str, float],
) -> ScenarioComparison:
    """Run multiple scenarios and compare side-by-side.

    Args:
        scenarios: Named scenarios {name: {asset_class: shock}}.
        position_weights: Current position weights.

    Returns:
        ScenarioComparison with all impacts and worst/best case.
    """
    results: Dict[str, ScenarioImpact] = {}
    worst_name, worst_impact = "", 0.0
    best_name, best_impact = "", float("-inf")

    for name, shocks in scenarios.items():
        impact = custom_scenario(shocks, position_weights)
        results[name] = impact

        if impact.total_impact_pct < worst_impact:
            worst_impact = impact.total_impact_pct
            worst_name = name
        if impact.total_impact_pct > best_impact:
            best_impact = impact.total_impact_pct
            best_name = name

    return ScenarioComparison(
        scenarios=results,
        worst_case=worst_name,
        best_case=best_name,
    )
