#!/usr/bin/env python3
"""CCAR/DFAST Stress Testing — Reference Implementation.

Multi-scenario stress test: baseline → adverse → severely adverse.
Demonstrates causal propagation of macroeconomic shocks through the
credit DAG to compute stressed PD and portfolio-level ECL impact.

Usage:
    python -m examples.ccar_stress_test
    # or
    python examples/ccar_stress_test.py
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

from orxaq.credit.causal_pd import CausalPDModel
from orxaq.credit.cecl import CECLEngine
from orxaq.credit.ontology import build_credit_ontology
from orxaq.credit.scenarios import (
    ADVERSE,
    BASELINE,
    SEVERELY_ADVERSE,
    ScenarioEngine,
)
from orxaq.credit.synthetic import CreditDataGenerator
from orxaq.kernel import WorldDAG

SCENARIOS = {
    "baseline": BASELINE,
    "adverse": ADVERSE,
    "severely_adverse": SEVERELY_ADVERSE,
}


def setup() -> tuple[WorldDAG, list[dict], CausalPDModel]:
    """Build ontology, generate data, and fit PD model."""
    dag = build_credit_ontology()
    gen = CreditDataGenerator(dag, seed=42)
    data = gen.generate(500)
    model = CausalPDModel(dag)
    model.fit(data)
    return dag, data, model


def run_stress_scenarios(
    dag: WorldDAG,
    data: list[dict],
    model: CausalPDModel,
) -> dict[str, dict]:
    """Run all three CCAR scenarios and collect results."""
    engine = ScenarioEngine(dag)
    results = {}

    for name, scenario in SCENARIOS.items():
        stressed_pd = engine.stress_pd(model, scenario, data)
        base_pd = engine.stress_pd(model, BASELINE, data)
        delta = stressed_pd - base_pd

        results[name] = {
            "pd": round(stressed_pd, 6),
            "base_pd": round(base_pd, 6),
            "delta": round(delta, 6),
            "scenario_params": {k: v for k, v in scenario.items()},
        }

    return results


def compute_stressed_ecl(
    dag: WorldDAG,
    data: list[dict],
    model: CausalPDModel,
) -> dict[str, float]:
    """Compute ECL under each scenario for comparison."""
    cecl = CECLEngine(model, dag)
    scenario_engine = ScenarioEngine(dag)
    ecl_results = {}

    for name, scenario in SCENARIOS.items():
        stressed_data = scenario_engine.apply_scenario(scenario, data)
        result = cecl.expected_credit_loss(stressed_data)
        ecl_results[name] = round(result.total_ecl, 4)

    return ecl_results


def print_scenario_report(
    results: dict[str, dict],
    ecl_results: dict[str, float],
) -> None:
    """Print formatted CCAR report."""
    print(f"{'='*70}")
    print("CCAR/DFAST STRESS TEST RESULTS")
    print(f"{'='*70}")

    # Macro parameters table
    print("\nMacroeconomic Scenarios:")
    print(f"  {'Variable':<20} {'Baseline':>12} {'Adverse':>12} {'Sev. Adverse':>12}")
    print(f"  {'-'*20} {'-'*12} {'-'*12} {'-'*12}")

    macro_vars = ["unemployment_rate", "fed_funds_rate", "gdp_growth", "hpi", "cpi"]
    for var in macro_vars:
        vals = [SCENARIOS[s].get(var, 0.0) for s in ["baseline", "adverse", "severely_adverse"]]
        if var in ("unemployment_rate", "fed_funds_rate", "gdp_growth"):
            fmt = "{:>11.1%}"
        else:
            fmt = "{:>11.0f}"
        print(f"  {var:<20} " + " ".join(fmt.format(v) for v in vals))

    # PD impact table
    print("\nPortfolio PD Impact:")
    print(f"  {'Scenario':<20} {'PD':>10} {'Delta':>10}")
    print(f"  {'-'*20} {'-'*10} {'-'*10}")
    for name in ["baseline", "adverse", "severely_adverse"]:
        r = results[name]
        print(f"  {name:<20} {r['pd']:>10.4%} {r['delta']:>+10.4%}")

    # ECL impact table
    print("\nPortfolio ECL Impact:")
    print(f"  {'Scenario':<20} {'ECL':>15}")
    print(f"  {'-'*20} {'-'*15}")
    base_ecl = ecl_results["baseline"]
    for name in ["baseline", "adverse", "severely_adverse"]:
        ecl = ecl_results[name]
        delta = ecl - base_ecl
        sign = "+" if delta > 0 else ""
        print(f"  {name:<20} ${ecl:>13,.2f}  ({sign}{delta:,.2f})")


def export_json(results: dict[str, dict], ecl_results: dict[str, float]) -> str:
    """Export full stress test results as JSON."""
    output = {
        "scenarios": results,
        "ecl_by_scenario": ecl_results,
    }
    return json.dumps(output, indent=2)


def main() -> None:
    """Run the full CCAR stress test pipeline."""
    # 1. Setup
    dag, data, model = setup()
    print(f"Portfolio: {len(data)} loans, PD model fitted")

    # 2. Run stress scenarios
    results = run_stress_scenarios(dag, data, model)

    # 3. Compute stressed ECL
    ecl_results = compute_stressed_ecl(dag, data, model)

    # 4. Print report
    print_scenario_report(results, ecl_results)

    # 5. JSON export
    print(f"\n{'='*70}")
    print("JSON OUTPUT")
    print(f"{'='*70}")
    print(export_json(results, ecl_results))


if __name__ == "__main__":
    main()
