"""CCAR/DFAST scenario engine.

Applies macroeconomic stress scenarios to credit portfolios using
do-calculus (interventional queries) via the WorldDAG.
"""

from __future__ import annotations

from orxaq.credit.causal_pd import CausalPDModel
from orxaq.kernel import WorldDAG

# Fed CCAR scenario definitions (macro variable overrides)
BASELINE: dict[str, float] = {
    "unemployment_rate": 0.04,
    "fed_funds_rate": 0.03,
    "gdp_growth": 0.02,
    "hpi": 220.0,
    "cpi": 260.0,
}

ADVERSE: dict[str, float] = {
    "unemployment_rate": 0.07,
    "fed_funds_rate": 0.01,
    "gdp_growth": -0.01,
    "hpi": 180.0,
    "cpi": 275.0,
}

SEVERELY_ADVERSE: dict[str, float] = {
    "unemployment_rate": 0.10,
    "fed_funds_rate": 0.005,
    "gdp_growth": -0.04,
    "hpi": 150.0,
    "cpi": 280.0,
}


class ScenarioEngine:
    """CCAR/DFAST scenario analysis engine.

    Applies macroeconomic stress scenarios to credit portfolios by
    propagating interventions through the causal DAG.
    """

    def __init__(self, dag: WorldDAG) -> None:
        self._dag = dag

    def apply_scenario(
        self,
        scenario: dict[str, float],
        data: list[dict[str, float]],
    ) -> list[dict[str, float]]:
        """Apply macro scenario to portfolio data using do-calculus.

        For each macro variable in the scenario, sets it to the stressed
        value (do-operator) and propagates effects through the DAG to
        downstream borrower and performance variables.
        """
        result = []
        for borrower in data:
            modified = dict(borrower)
            # Set macro variables to scenario values
            for var_id, value in scenario.items():
                modified[var_id] = value
            # Propagate effects through DAG
            modified = self._propagate(modified, scenario)
            result.append(modified)
        return result

    def stress_pd(
        self,
        pd_model: CausalPDModel,
        scenario: dict[str, float],
        portfolio: list[dict[str, float]],
    ) -> float:
        """Calculate portfolio-level stressed PD under a scenario.

        Uses interventional queries for each macro variable in the
        scenario to estimate the causal effect of the stress.
        """
        if not pd_model.is_fitted:
            msg = "PD model not fitted. Call pd_model.fit() first."
            raise RuntimeError(msg)

        if not portfolio:
            return 0.0

        pd_sum = 0.0
        for borrower in portfolio:
            # Apply each macro intervention and take the max PD
            # (stress scenario = worst-case across macro shocks)
            pds = []
            for var_id, value in scenario.items():
                pd_val = pd_model.interventional_pd(borrower, var_id, value)
                pds.append(pd_val)
            # Use the mean of interventional PDs across macro variables
            # as the stressed PD for this borrower
            pd_sum += sum(pds) / len(pds) if pds else pd_model.predict_pd(borrower)

        return pd_sum / len(portfolio)

    def _propagate(
        self,
        borrower: dict[str, float],
        scenario: dict[str, float],
    ) -> dict[str, float]:
        """Propagate scenario effects through DAG in topological order.

        For variables downstream of scenario variables, estimate the
        shifted value based on causal edge weights and scenario deltas.
        """
        order = self._dag.topological_sort()
        result = dict(borrower)
        scenario_vars = set(scenario.keys())

        for var_id in order:
            if var_id in scenario_vars:
                continue  # Already set by scenario

            parents = self._dag.parents(var_id)
            affected_parents = parents & scenario_vars

            if not affected_parents:
                # Check if any parent was already modified by propagation
                continue

            # Estimate shift: for each affected parent, apply proportional
            # shift based on the edge confidence (simplified propagation)
            shift = 0.0
            for parent_id in parents:
                if parent_id not in result:
                    continue
                # Get edge confidence as proxy for effect size
                children = self._dag.children(parent_id)
                if var_id in children:
                    # Use a simplified linear propagation model
                    # Direction: positive confidence means positive relationship
                    shift += result.get(parent_id, 0.0) * 0.01

            if shift != 0.0 and var_id in result:
                result[var_id] = result[var_id] + shift

        return result
