"""Scenario simulator — propagate shocks through the causal DAG.

Defines macroeconomic scenarios, runs multi-step propagation through
the World DAG, and produces impact analysis with confidence intervals.
Stdlib-only (ADR-007).
"""

from __future__ import annotations

import json
import math
import random
from dataclasses import dataclass, field

from orxaq.kernel.world_dag import WorldDAG


@dataclass(frozen=True)
class Shock:
    """A single variable shock in a scenario."""

    variable: str
    value: float
    shock_type: str = "absolute"  # absolute, relative, delta


@dataclass(frozen=True)
class Scenario:
    """A named macroeconomic scenario with variable shocks."""

    name: str
    description: str
    shocks: tuple[Shock, ...]
    severity: str = "baseline"  # baseline, adverse, severely_adverse, custom
    horizon_quarters: int = 4

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "description": self.description,
            "shocks": [
                {"variable": s.variable, "value": s.value, "shock_type": s.shock_type}
                for s in self.shocks
            ],
            "severity": self.severity,
            "horizon_quarters": self.horizon_quarters,
        }

    @classmethod
    def from_dict(cls, data: dict) -> Scenario:
        shocks = tuple(
            Shock(
                variable=s["variable"],
                value=s["value"],
                shock_type=s.get("shock_type", "absolute"),
            )
            for s in data.get("shocks", [])
        )
        return cls(
            name=data["name"],
            description=data.get("description", ""),
            shocks=shocks,
            severity=data.get("severity", "custom"),
            horizon_quarters=data.get("horizon_quarters", 4),
        )

    @classmethod
    def from_json(cls, path: str) -> Scenario:
        with open(path) as f:
            return cls.from_dict(json.loads(f.read()))


@dataclass
class PropagationStep:
    """One step of propagation through the DAG."""

    variable: str
    original_value: float
    new_value: float
    delta: float
    source_variables: list[str] = field(default_factory=list)
    depth: int = 0


@dataclass
class SimulationResult:
    """Full simulation output."""

    scenario: Scenario
    steps: list[PropagationStep]
    affected_variables: dict[str, float]  # variable -> delta
    total_affected: int
    max_depth: int
    impact_score: float  # normalized 0-1 impact metric

    def to_dict(self) -> dict:
        return {
            "scenario": self.scenario.to_dict(),
            "steps": [
                {
                    "variable": s.variable,
                    "original": s.original_value,
                    "new_value": s.new_value,
                    "delta": s.delta,
                    "sources": s.source_variables,
                    "depth": s.depth,
                }
                for s in self.steps
            ],
            "affected_variables": self.affected_variables,
            "total_affected": self.total_affected,
            "max_depth": self.max_depth,
            "impact_score": self.impact_score,
        }


class ScenarioSimulator:
    """Propagate scenario shocks through the World DAG.

    Uses topological ordering and edge confidence to estimate
    downstream impact of macroeconomic interventions.
    """

    def __init__(self, dag: WorldDAG, seed: int = 42) -> None:
        self._dag = dag
        self._rng = random.Random(seed)

    def simulate(
        self,
        scenario: Scenario,
        baseline: dict[str, float] | None = None,
    ) -> SimulationResult:
        """Run simulation: apply shocks, propagate through DAG."""
        if baseline is None:
            baseline = {}

        # Apply initial shocks
        state = dict(baseline)
        steps: list[PropagationStep] = []
        shocked_vars: set[str] = set()

        for shock in scenario.shocks:
            original = state.get(shock.variable, 0.0)
            if shock.shock_type == "absolute":
                new_val = shock.value
            elif shock.shock_type == "relative":
                new_val = original * shock.value
            else:  # delta
                new_val = original + shock.value

            state[shock.variable] = new_val
            shocked_vars.add(shock.variable)
            steps.append(
                PropagationStep(
                    variable=shock.variable,
                    original_value=original,
                    new_value=new_val,
                    delta=new_val - original,
                    source_variables=[],
                    depth=0,
                )
            )

        # Propagate through DAG in topological order
        order = self._dag.topological_sort()
        affected: dict[str, float] = {}
        max_depth = 0

        for var_id in order:
            if var_id in shocked_vars:
                affected[var_id] = state[var_id] - baseline.get(var_id, 0.0)
                continue

            parents = self._dag.parents(var_id)
            if not parents:
                continue

            # Check if any parent has changed
            parent_deltas = {}
            for p in parents:
                if p in affected:
                    parent_deltas[p] = affected[p]

            if not parent_deltas:
                continue

            # Compute propagated effect
            original = state.get(var_id, baseline.get(var_id, 0.0))
            delta = self._compute_propagated_delta(var_id, parent_deltas, original)

            if abs(delta) < 1e-10:
                continue

            new_val = original + delta
            state[var_id] = new_val
            affected[var_id] = delta

            # Compute depth: max parent depth + 1
            depth = 0
            for step in steps:
                if step.variable in parent_deltas:
                    depth = max(depth, step.depth + 1)
            max_depth = max(max_depth, depth)

            steps.append(
                PropagationStep(
                    variable=var_id,
                    original_value=original,
                    new_value=new_val,
                    delta=delta,
                    source_variables=list(parent_deltas.keys()),
                    depth=depth,
                )
            )

        # Compute impact score: normalized magnitude of changes
        impact = self._compute_impact_score(affected, baseline)

        return SimulationResult(
            scenario=scenario,
            steps=steps,
            affected_variables=affected,
            total_affected=len(affected),
            max_depth=max_depth,
            impact_score=impact,
        )

    def _compute_propagated_delta(
        self,
        var_id: str,
        parent_deltas: dict[str, float],
        original: float,
    ) -> float:
        """Compute the propagated delta for a variable based on parent changes."""
        total_delta = 0.0
        for parent_id, p_delta in parent_deltas.items():
            # Use edge confidence as effect strength
            edge_id = f"{parent_id}\u2192{var_id}"
            edge = self._dag._edges.get(edge_id)
            confidence = edge.confidence if edge else 0.5

            # Linear propagation with confidence scaling
            # Add small noise for realism
            noise = self._rng.gauss(0, 0.01)
            total_delta += p_delta * confidence * (1.0 + noise)

        return total_delta

    def _compute_impact_score(
        self,
        affected: dict[str, float],
        baseline: dict[str, float],
    ) -> float:
        """Compute normalized impact score (0-1)."""
        if not affected:
            return 0.0

        magnitudes = []
        for var_id, delta in affected.items():
            base = abs(baseline.get(var_id, 1.0))
            if base > 0:
                magnitudes.append(abs(delta) / base)
            else:
                magnitudes.append(abs(delta))

        if not magnitudes:  # pragma: no cover
            return 0.0  # pragma: no cover

        # Use sigmoid-like normalization
        avg_magnitude = sum(magnitudes) / len(magnitudes)
        score = 2.0 / (1.0 + math.exp(-avg_magnitude * 5.0)) - 1.0
        return min(max(score, 0.0), 1.0)

    def monte_carlo(
        self,
        scenario: Scenario,
        baseline: dict[str, float] | None = None,
        n_simulations: int = 100,
    ) -> dict:
        """Run Monte Carlo simulation with random perturbations."""
        results: list[SimulationResult] = []
        for i in range(n_simulations):
            sim = ScenarioSimulator(self._dag, seed=42 + i)
            results.append(sim.simulate(scenario, baseline))

        # Aggregate affected variables across simulations
        all_vars: set[str] = set()
        for r in results:
            all_vars.update(r.affected_variables.keys())

        var_stats: dict[str, dict] = {}
        for var_id in sorted(all_vars):
            deltas = [r.affected_variables.get(var_id, 0.0) for r in results]
            deltas.sort()
            n = len(deltas)
            mean = sum(deltas) / n
            p5_idx = max(0, int(n * 0.05))
            p95_idx = min(n - 1, int(n * 0.95))
            var_stats[var_id] = {
                "mean_delta": mean,
                "p5": deltas[p5_idx],
                "p95": deltas[p95_idx],
                "min": deltas[0],
                "max": deltas[-1],
            }

        impact_scores = [r.impact_score for r in results]
        return {
            "scenario": scenario.name,
            "n_simulations": n_simulations,
            "variable_stats": var_stats,
            "impact_score_mean": sum(impact_scores) / len(impact_scores),
            "impact_score_p5": sorted(impact_scores)[max(0, int(len(impact_scores) * 0.05))],
            "impact_score_p95": sorted(impact_scores)[
                min(len(impact_scores) - 1, int(len(impact_scores) * 0.95))
            ],
        }


# ── Built-in scenarios ──────────────────────────────────────

BASELINE_SCENARIO = Scenario(
    name="baseline",
    description="Fed CCAR Baseline — normal economic conditions",
    severity="baseline",
    horizon_quarters=4,
    shocks=(
        Shock("unemployment_rate", 0.04),
        Shock("fed_funds_rate", 0.03),
        Shock("gdp_growth", 0.02),
    ),
)

ADVERSE_SCENARIO = Scenario(
    name="adverse",
    description="Fed CCAR Adverse — moderate recession",
    severity="adverse",
    horizon_quarters=4,
    shocks=(
        Shock("unemployment_rate", 0.07),
        Shock("fed_funds_rate", 0.01),
        Shock("gdp_growth", -0.01),
    ),
)

SEVERE_SCENARIO = Scenario(
    name="severely_adverse",
    description="Fed CCAR Severely Adverse — deep recession",
    severity="severely_adverse",
    horizon_quarters=4,
    shocks=(
        Shock("unemployment_rate", 0.10),
        Shock("fed_funds_rate", 0.005),
        Shock("gdp_growth", -0.04),
    ),
)

BUILT_IN_SCENARIOS = {
    "baseline": BASELINE_SCENARIO,
    "adverse": ADVERSE_SCENARIO,
    "severely_adverse": SEVERE_SCENARIO,
}
