"""CCAR/DFAST stress testing report engine.

Produces regulatory-formatted stress test reports by composing:
- ScenarioEngine (3 Fed scenarios: baseline, adverse, severely_adverse)
- CECLEngine (expected credit loss computation)
- PortfolioAnalyzer / MigrationMatrix (concentration, migration)

Per ADR-007: stdlib-only at module level; numpy/scipy used lazily if available.
"""

from __future__ import annotations

import datetime
import json
from dataclasses import dataclass, field

from orxaq.credit.causal_pd import CausalPDModel
from orxaq.credit.cecl import CECLEngine
from orxaq.credit.portfolio import LoanExposure, MigrationMatrix
from orxaq.credit.scenarios import ADVERSE, BASELINE, SEVERELY_ADVERSE, ScenarioEngine

# Canonical Fed scenario names used in CCAR/DFAST submissions
FED_SCENARIO_NAMES = ("baseline", "adverse", "severely_adverse")

# Map name -> macro variable dict
FED_SCENARIOS: dict[str, dict[str, float]] = {
    "baseline": BASELINE,
    "adverse": ADVERSE,
    "severely_adverse": SEVERELY_ADVERSE,
}

# CCAR/DFAST mandated horizon (9 quarters = ~2.25 years)
CCAR_HORIZON_QUARTERS = 9

# Basel III Tier-1 capital ratio proxy for capital impact estimation
_BASEL_TIER1_RATIO = 0.08

__all__ = [
    "CCAR_HORIZON_QUARTERS",
    "FED_SCENARIO_NAMES",
    "FED_SCENARIOS",
    "StressTestReport",
    "run_stress_test",
]


@dataclass(frozen=True)
class StressTestReport:
    """CCAR/DFAST regulatory stress test report for a single scenario.

    Attributes:
        scenario_name: One of "baseline", "adverse", "severely_adverse".
        quarterly_ecl: List of 9 projected ECL values (one per CCAR quarter).
        portfolio_loss_rate: Total 9Q ECL / total EAD (dimensionless [0,1]).
        capital_impact: Estimated Tier-1 capital ratio impact (Basel proxy).
        migration_matrices: List of per-quarter rating migration probability dicts.
        confidence_intervals: Mapping of metric name -> (low, high) 90% bands.
        generated_at: ISO-8601 UTC timestamp of report generation.
    """

    scenario_name: str
    quarterly_ecl: list[float]
    portfolio_loss_rate: float
    capital_impact: float
    migration_matrices: list[dict[str, dict[str, float]]]
    confidence_intervals: dict[str, tuple[float, float]]
    generated_at: str = field(
        default_factory=lambda: datetime.datetime.now(datetime.UTC).isoformat()
    )

    def to_dict(self) -> dict[str, object]:
        """Return a JSON-serializable dict representation."""
        return {
            "scenario_name": self.scenario_name,
            "quarterly_ecl": self.quarterly_ecl,
            "portfolio_loss_rate": self.portfolio_loss_rate,
            "capital_impact": self.capital_impact,
            "migration_matrices": self.migration_matrices,
            "confidence_intervals": {k: list(v) for k, v in self.confidence_intervals.items()},
            "generated_at": self.generated_at,
        }

    def to_json(self, indent: int = 2) -> str:
        """Serialize report to a JSON string."""
        return json.dumps(self.to_dict(), indent=indent)


def run_stress_test(
    portfolio: list[LoanExposure],
    scenarios: dict[str, dict[str, float]],
    pd_model: CausalPDModel,
    cecl_engine: CECLEngine,
    *,
    horizon_quarters: int = CCAR_HORIZON_QUARTERS,
    scenario_name: str | None = None,
) -> StressTestReport:
    """Run a CCAR/DFAST stress test for a single scenario.

    Composes ScenarioEngine, CECLEngine, PortfolioAnalyzer, and
    MigrationMatrix to produce a regulatory-grade stress report.

    Args:
        portfolio: List of LoanExposure objects representing the credit book.
        scenarios: Dict mapping scenario names to macro variable overrides.
            Typically pass FED_SCENARIOS or a single-key subset.
        pd_model: Fitted CausalPDModel for PD estimation under stress.
            May be None; stress is applied via ScenarioEngine propagation.
        cecl_engine: Configured CECLEngine for ECL computation.
        horizon_quarters: CCAR horizon length (default 9).
        scenario_name: Which scenario key in ``scenarios`` to run.
            If None, uses the first key found.

    Returns:
        StressTestReport with all regulatory metrics populated.
    """
    if not scenarios:
        msg = "scenarios dict must not be empty"
        raise ValueError(msg)

    name = scenario_name if scenario_name is not None else next(iter(scenarios))
    if name not in scenarios:
        msg = f"scenario '{name}' not found in scenarios dict"
        raise KeyError(msg)

    macro_overrides = scenarios[name]

    # Build raw borrower dicts from LoanExposure for scenario propagation
    raw_portfolio = _exposures_to_dicts(portfolio)

    # Apply scenario: stress macro variables through the causal DAG
    scenario_engine = ScenarioEngine(cecl_engine._dag)  # noqa: SLF001
    stressed_portfolio = scenario_engine.apply_scenario(macro_overrides, raw_portfolio)

    # Compute per-quarter ECL over the CCAR 9-quarter horizon
    quarterly_ecl = _compute_quarterly_ecl(
        stressed_portfolio,
        cecl_engine,
        horizon_quarters,
    )

    # Aggregate portfolio metrics
    total_ead = sum(e.ead for e in portfolio)
    total_ecl = sum(quarterly_ecl)
    portfolio_loss_rate = (total_ecl / total_ead) if total_ead > 0.0 else 0.0
    capital_impact = _BASEL_TIER1_RATIO * portfolio_loss_rate

    # Build migration matrices per quarter (record stressed rating shifts)
    migration_matrices = _build_migration_matrices(portfolio, stressed_portfolio, horizon_quarters)

    # Confidence intervals from CECL confidence bands across all quarters
    confidence_intervals = _compute_confidence_intervals(
        stressed_portfolio,
        cecl_engine,
        quarterly_ecl,
        total_ead,
    )

    return StressTestReport(
        scenario_name=name,
        quarterly_ecl=quarterly_ecl,
        portfolio_loss_rate=portfolio_loss_rate,
        capital_impact=capital_impact,
        migration_matrices=migration_matrices,
        confidence_intervals=confidence_intervals,
    )


# ---------------------------------------------------------------------------
# Private helpers
# ---------------------------------------------------------------------------


def _exposures_to_dicts(portfolio: list[LoanExposure]) -> list[dict[str, float]]:
    """Convert LoanExposure objects to borrower dicts for ScenarioEngine."""
    return [
        {
            "pd": e.pd,
            "lgd": e.lgd,
            "ead": e.ead,
            "loan_id_hash": float(hash(e.loan_id) % 100_000),
        }
        for e in portfolio
    ]


def _compute_quarterly_ecl(
    stressed_portfolio: list[dict[str, float]],
    cecl_engine: CECLEngine,
    horizon_quarters: int,
) -> list[float]:
    """Compute per-quarter ECL using CECL engine with 1-quarter horizons.

    Each quarter's ECL is estimated independently by passing horizon=1
    and scaling by the quarter's survival probability (geometric decay).
    """
    mean_pd = _portfolio_mean_pd(stressed_portfolio)
    ecl_results = []
    for q in range(1, horizon_quarters + 1):
        # Survival factor: probability of not having defaulted by quarter q
        survival = (1.0 - mean_pd) ** (q - 1)
        result = cecl_engine.expected_credit_loss(stressed_portfolio, horizon_quarters=1)
        quarterly_ecl = result.total_ecl * survival
        ecl_results.append(quarterly_ecl)
    return ecl_results


def _portfolio_mean_pd(portfolio: list[dict[str, float]]) -> float:
    """Compute mean PD across stressed portfolio."""
    if not portfolio:
        return 0.0
    return sum(b.get("pd", 0.0) for b in portfolio) / len(portfolio)


def _build_migration_matrices(
    original: list[LoanExposure],
    stressed: list[dict[str, float]],
    horizon_quarters: int,
) -> list[dict[str, dict[str, float]]]:
    """Build simplified per-quarter migration probability matrices.

    Maps each loan from its original rating (derived from PD bucket)
    to a stressed rating under the scenario, repeated across quarters
    with increasing severity.
    """

    def _pd_to_rating(pd: float) -> str:
        if pd < 0.01:
            return "AAA"
        if pd < 0.03:
            return "AA"
        if pd < 0.06:
            return "A"
        if pd < 0.10:
            return "BBB"
        if pd < 0.15:
            return "BB"
        if pd < 0.25:
            return "B"
        return "CCC"

    matrices = []
    for q in range(horizon_quarters):
        # Severity increases linearly over the horizon
        severity = (q + 1) / horizon_quarters
        mm = MigrationMatrix()
        for i, exposure in enumerate(original):
            from_rating = _pd_to_rating(exposure.pd)
            # Under stress, PD increases by severity factor
            stressed_pd_raw = stressed[i].get("pd", exposure.pd)
            stressed_pd = min(1.0, stressed_pd_raw * (1.0 + severity))
            to_rating = _pd_to_rating(stressed_pd)
            mm.record(f"loan_{i}", from_rating, to_rating)
        matrices.append(mm.probability_matrix())

    return matrices


def _compute_confidence_intervals(
    stressed_portfolio: list[dict[str, float]],
    cecl_engine: CECLEngine,
    quarterly_ecl: list[float],
    total_ead: float,
) -> dict[str, tuple[float, float]]:
    """Derive 90% confidence intervals from CECL confidence bands.

    Uses the ECLResult.confidence_band from the CECL engine and
    propagates it to loss rate and capital impact metrics.
    """
    result = cecl_engine.expected_credit_loss(stressed_portfolio, horizon_quarters=1)
    ecl_low, ecl_high = result.confidence_band

    # Scale to 9-quarter horizon
    horizon_scale = len(quarterly_ecl)
    total_ecl_low = ecl_low * horizon_scale
    total_ecl_high = ecl_high * horizon_scale

    loss_rate_low = (total_ecl_low / total_ead) if total_ead > 0.0 else 0.0
    loss_rate_high = (total_ecl_high / total_ead) if total_ead > 0.0 else 0.0

    return {
        "total_ecl": (total_ecl_low, total_ecl_high),
        "portfolio_loss_rate": (loss_rate_low, loss_rate_high),
        "capital_impact": (
            _BASEL_TIER1_RATIO * loss_rate_low,
            _BASEL_TIER1_RATIO * loss_rate_high,
        ),
    }
