"""Stress testing for prediction market portfolios.

Provides scenario-based stress tests using Monte Carlo simulation
with user-defined correlation and price shock overrides.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from horizon._horizon import SimPosition, SimulationResult, monte_carlo


@dataclass
class StressScenario:
    """A stress test scenario definition.

    Args:
        name: Human-readable scenario label.
        correlation_override: If set, overrides the correlation matrix with
            a uniform correlation at this level (e.g. 0.95 for spike).
        price_shocks: Dict mapping market_id -> shocked current_price.
            Markets not in the dict keep their original price.
    """

    name: str
    correlation_override: float | None = None
    price_shocks: dict[str, float] = field(default_factory=dict)


@dataclass
class ScenarioResult:
    """Result of a single stress scenario."""

    scenario: StressScenario
    sim: SimulationResult


@dataclass
class StressTestResult:
    """Aggregated stress test results across all scenarios."""

    scenarios: list[ScenarioResult]

    @property
    def worst_scenario(self) -> str:
        """Name of the scenario with the worst (lowest) mean PnL."""
        if not self.scenarios:
            return ""
        return min(self.scenarios, key=lambda s: s.sim.mean_pnl).scenario.name

    @property
    def worst_pnl(self) -> float:
        """Worst mean PnL across all scenarios."""
        if not self.scenarios:
            return 0.0
        return min(s.sim.mean_pnl for s in self.scenarios)

    def summary(self) -> dict[str, Any]:
        """Return a summary dict of all scenario results."""
        return {
            "n_scenarios": len(self.scenarios),
            "worst_scenario": self.worst_scenario,
            "worst_pnl": self.worst_pnl,
            "scenarios": {
                s.scenario.name: {
                    "mean_pnl": s.sim.mean_pnl,
                    "var_95": s.sim.var_95,
                    "cvar_95": s.sim.cvar_95,
                    "max_loss": s.sim.max_loss,
                    "win_probability": s.sim.win_probability,
                }
                for s in self.scenarios
            },
        }


# ===========================================================================
# Built-in scenarios
# ===========================================================================

CORRELATION_SPIKE = StressScenario(
    name="correlation_spike",
    correlation_override=0.95,
)

ALL_RESOLVE_NO = StressScenario(
    name="all_resolve_no",
    price_shocks={},  # Applied dynamically: all prices → 0.01
)

LIQUIDITY_SHOCK = StressScenario(
    name="liquidity_shock",
    price_shocks={},  # Applied dynamically: all prices → 0.50
)

TAIL_RISK = StressScenario(
    name="tail_risk",
    correlation_override=0.99,
)


# ===========================================================================
# stress_test
# ===========================================================================


def stress_test(
    positions: list[SimPosition],
    scenarios: list[StressScenario] | None = None,
    n_simulations: int = 10000,
    seed: int | None = None,
) -> StressTestResult:
    """Run stress tests across multiple scenarios.

    Args:
        positions: Portfolio positions as SimPosition objects.
        scenarios: List of stress scenarios to test.
            Defaults to [CORRELATION_SPIKE, ALL_RESOLVE_NO, LIQUIDITY_SHOCK, TAIL_RISK].
        n_simulations: Monte Carlo simulations per scenario.
        seed: PRNG seed for reproducibility.

    Returns:
        StressTestResult with per-scenario simulation results.
    """
    if scenarios is None:
        scenarios = [CORRELATION_SPIKE, ALL_RESOLVE_NO, LIQUIDITY_SHOCK, TAIL_RISK]

    if not positions:
        empty_results = [
            ScenarioResult(
                scenario=s,
                sim=monte_carlo([], n_simulations, None, seed),
            )
            for s in scenarios
        ]
        return StressTestResult(scenarios=empty_results)

    n = len(positions)
    results = []

    for scenario in scenarios:
        # Apply price shocks
        shocked_positions = []
        for pos in positions:
            new_price = pos.current_price

            # Dynamic built-in scenarios
            if scenario.name == "all_resolve_no":
                new_price = 0.01
            elif scenario.name == "liquidity_shock":
                new_price = 0.50

            # Explicit price shocks override
            if pos.market_id in scenario.price_shocks:
                new_price = scenario.price_shocks[pos.market_id]

            shocked_positions.append(
                SimPosition(
                    market_id=pos.market_id,
                    side=pos.side,
                    size=pos.size,
                    entry_price=pos.entry_price,
                    current_price=new_price,
                )
            )

        # Build correlation matrix
        corr = None
        if scenario.correlation_override is not None and n > 1:
            rho = max(-1.0, min(1.0, scenario.correlation_override))
            corr = [[rho if i != j else 1.0 for j in range(n)] for i in range(n)]

        sim = monte_carlo(shocked_positions, n_simulations, corr, seed)
        results.append(ScenarioResult(scenario=scenario, sim=sim))

    return StressTestResult(scenarios=results)
