"""
simulation.py — Monte Carlo simulation engine for Open FAIR scenarios.

The primary entry point is :class:`FAIRSimulator`, which accepts a
:class:`~pycrq.ontology.FAIRScenario` (or a
:class:`~pycrq.ontology.RiskRegister`) and returns rich result objects
containing per-iteration arrays and summary statistics.
"""

from __future__ import annotations

import warnings
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import numpy as np
from scipy.special import expit  # logistic sigmoid

from .ontology import FAIRScenario, FrequencyFactors, LossMagnitudeFactors, RiskRegister


# ---------------------------------------------------------------------------
# SimulationResult
# ---------------------------------------------------------------------------

@dataclass
class SimulationResult:
    """Stores the per-iteration output of a single-scenario Monte Carlo run.

    Attributes
    ----------
    scenario_name:         Name of the simulated scenario.
    n_simulations:         Number of Monte Carlo iterations.
    loss_event_frequency:  Iterations × Annual loss event frequency (LEF).
    loss_magnitude:        Iterations × Per-event loss magnitude (LM).
    annual_loss_exposure:  Iterations × ALE = LEF × LM.
    threat_event_frequency: Optional per-iteration TEF array.
    vulnerability:         Optional per-iteration vulnerability array.
    primary_loss:          Optional per-iteration primary loss array.
    secondary_loss:        Optional per-iteration secondary loss array.
    loss_by_category:      Per-category loss arrays keyed by category name.
    """

    scenario_name: str
    n_simulations: int
    loss_event_frequency: np.ndarray
    loss_magnitude: np.ndarray
    annual_loss_exposure: np.ndarray
    threat_event_frequency: Optional[np.ndarray] = None
    vulnerability: Optional[np.ndarray] = None
    primary_loss: Optional[np.ndarray] = None
    secondary_loss: Optional[np.ndarray] = None
    loss_by_category: Dict[str, np.ndarray] = field(default_factory=dict)

    # -- Derived statistics --

    @property
    def mean_ale(self) -> float:
        """Mean Annual Loss Exposure across all simulations."""
        return float(np.mean(self.annual_loss_exposure))

    @property
    def std_ale(self) -> float:
        """Standard deviation of Annual Loss Exposure."""
        return float(np.std(self.annual_loss_exposure))

    def percentile(self, p: float) -> float:
        """Return the *p*-th percentile of the ALE distribution (0 ≤ p ≤ 100)."""
        if not 0 <= p <= 100:
            raise ValueError(f"Percentile p must be in [0, 100]; got {p}")
        return float(np.percentile(self.annual_loss_exposure, p))

    def var(self, confidence: float = 0.95) -> float:
        """Value-at-Risk: the ALE not exceeded with probability *confidence*.

        Parameters
        ----------
        confidence: Confidence level in (0, 1).  Default 0.95.
        """
        if not 0 < confidence < 1:
            raise ValueError(f"confidence must be in (0, 1); got {confidence}")
        return float(np.percentile(self.annual_loss_exposure, confidence * 100))

    def cvar(self, confidence: float = 0.95) -> float:
        """Conditional Value-at-Risk (Expected Shortfall) at *confidence*.

        Returns the mean ALE in the tail beyond the VaR threshold.

        Parameters
        ----------
        confidence: Confidence level in (0, 1).  Default 0.95.
        """
        threshold = self.var(confidence)
        tail = self.annual_loss_exposure[self.annual_loss_exposure >= threshold]
        if len(tail) == 0:
            return threshold
        return float(np.mean(tail))

    def summary(self) -> Dict:
        """Return a dictionary of key summary statistics."""
        ale = self.annual_loss_exposure
        summary = {
            "scenario_name": self.scenario_name,
            "n_simulations": self.n_simulations,
            "mean_ale": self.mean_ale,
            "std_ale": self.std_ale,
            "min_ale": float(np.min(ale)),
            "p10_ale": self.percentile(10),
            "median_ale": self.percentile(50),
            "p90_ale": self.percentile(90),
            "p95_ale": self.percentile(95),
            "p99_ale": self.percentile(99),
            "max_ale": float(np.max(ale)),
            "var_95": self.var(0.95),
            "cvar_95": self.cvar(0.95),
        }
        # Add frequency/magnitude summaries if available
        if self.loss_event_frequency is not None:
            lef = self.loss_event_frequency
            summary["mean_lef"] = float(np.mean(lef))
            summary["p50_lef"] = float(np.percentile(lef, 50))
        if self.loss_magnitude is not None:
            lm = self.loss_magnitude
            summary["mean_lm"] = float(np.mean(lm))
            summary["p50_lm"] = float(np.percentile(lm, 50))
        if self.loss_by_category:
            summary["loss_by_category"] = {
                k: float(np.mean(v)) for k, v in self.loss_by_category.items()
            }
        return summary


# ---------------------------------------------------------------------------
# AggregateSimulationResult
# ---------------------------------------------------------------------------

@dataclass
class AggregateSimulationResult:
    """Aggregated results across all scenarios in a :class:`~pycrq.ontology.RiskRegister`.

    The ``aggregate_ale`` array is computed as the element-wise sum of ALE
    arrays from all scenarios (i.e., correlated total risk).

    Attributes
    ----------
    register_name:    Name of the simulated risk register.
    n_simulations:    Number of Monte Carlo iterations.
    scenario_results: List of per-scenario :class:`SimulationResult` objects.
    aggregate_ale:    Combined ALE array (computed in ``__post_init__``).
    """

    register_name: str
    n_simulations: int
    scenario_results: List[SimulationResult]
    aggregate_ale: np.ndarray = field(init=False)

    def __post_init__(self) -> None:
        if not self.scenario_results:
            self.aggregate_ale = np.zeros(self.n_simulations)
        else:
            self.aggregate_ale = sum(
                r.annual_loss_exposure for r in self.scenario_results
            )

    @property
    def mean_ale(self) -> float:
        """Mean aggregate Annual Loss Exposure."""
        return float(np.mean(self.aggregate_ale))

    def percentile(self, p: float) -> float:
        """Return the *p*-th percentile of the aggregate ALE distribution."""
        if not 0 <= p <= 100:
            raise ValueError(f"Percentile p must be in [0, 100]; got {p}")
        return float(np.percentile(self.aggregate_ale, p))

    def var(self, confidence: float = 0.95) -> float:
        """Value-at-Risk for the aggregate portfolio at *confidence*."""
        return float(np.percentile(self.aggregate_ale, confidence * 100))

    def cvar(self, confidence: float = 0.95) -> float:
        """Conditional Value-at-Risk for the aggregate portfolio."""
        threshold = self.var(confidence)
        tail = self.aggregate_ale[self.aggregate_ale >= threshold]
        if len(tail) == 0:
            return threshold
        return float(np.mean(tail))

    def summary(self) -> Dict:
        """Return a dictionary of aggregate and per-scenario summary statistics."""
        agg = self.aggregate_ale
        per_scenario = [r.summary() for r in self.scenario_results]
        return {
            "register_name": self.register_name,
            "n_simulations": self.n_simulations,
            "aggregate": {
                "mean_ale": self.mean_ale,
                "median_ale": self.percentile(50),
                "p90_ale": self.percentile(90),
                "p95_ale": self.percentile(95),
                "p99_ale": self.percentile(99),
                "max_ale": float(np.max(agg)),
                "var_95": self.var(0.95),
                "cvar_95": self.cvar(0.95),
            },
            "scenarios": per_scenario,
        }


# ---------------------------------------------------------------------------
# FAIRSimulator
# ---------------------------------------------------------------------------

class FAIRSimulator:
    """Monte Carlo simulator implementing the Open FAIR ontology.

    Parameters
    ----------
    n_simulations: Number of Monte Carlo iterations per run (default 10 000).
    seed:          Optional random seed for reproducibility.
    """

    def __init__(self, n_simulations: int = 10_000, seed: Optional[int] = None) -> None:
        if n_simulations < 1:
            raise ValueError(f"n_simulations must be >= 1; got {n_simulations}")
        self.n_simulations = int(n_simulations)
        self.seed = seed
        self._rng = np.random.default_rng(seed)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _sample(self, dist, n: int) -> np.ndarray:
        """Sample *n* values from *dist*, seeding its internal RNG."""
        # Temporarily set numpy random seed for scipy-backed distributions
        if self.seed is not None:
            np.random.seed(self._rng.integers(0, 2**31 - 1))
        return dist.sample(n)

    def _compute_tef(self, ff: FrequencyFactors) -> np.ndarray:
        """Derive Threat Event Frequency array from frequency factors.

        Priority order:
        1. Direct ``tef`` distribution.
        2. ``contact_frequency × probability_of_action``.
        3. ``contact_frequency`` alone (probability_of_action assumed 1.0).
        4. Raise :class:`ValueError`.
        """
        n = self.n_simulations

        if ff.tef is not None:
            return np.maximum(self._sample(ff.tef, n), 0.0)

        if ff.contact_frequency is not None:
            cf = np.maximum(self._sample(ff.contact_frequency, n), 0.0)
            if ff.probability_of_action is not None:
                poa = np.clip(self._sample(ff.probability_of_action, n), 0.0, 1.0)
                return cf * poa
            return cf

        raise ValueError(
            "FrequencyFactors must supply at least one of: tef, contact_frequency."
        )

    def _compute_vulnerability(self, ff: FrequencyFactors) -> np.ndarray:
        """Derive Vulnerability array from frequency factors.

        Priority order:
        1. Direct ``vulnerability`` distribution.
        2. Logistic function of threat_capability and control_strength scores.
           vuln = sigmoid(6 * (tcap - cs) / 100)
        3. threat_capability / 100  (if only capability is provided).
        4. 1 - control_strength / 100  (if only strength is provided).
        5. Warn and return Uniform(0, 1).
        """
        n = self.n_simulations

        if ff.vulnerability is not None:
            return np.clip(self._sample(ff.vulnerability, n), 0.0, 1.0)

        if ff.threat_capability is not None and ff.control_strength is not None:
            tcap = self._sample(ff.threat_capability, n)
            cs = self._sample(ff.control_strength, n)
            # Logistic: vuln = 1 / (1 + exp(-6 * (tcap - cs) / 100))
            vuln = expit(6.0 * (tcap - cs) / 100.0)
            return np.clip(vuln, 0.0, 1.0)

        if ff.threat_capability is not None:
            tcap = self._sample(ff.threat_capability, n)
            return np.clip(tcap / 100.0, 0.0, 1.0)

        if ff.control_strength is not None:
            cs = self._sample(ff.control_strength, n)
            return np.clip(1.0 - cs / 100.0, 0.0, 1.0)

        warnings.warn(
            "No vulnerability inputs found in FrequencyFactors. "
            "Defaulting to Uniform(0, 1) vulnerability.",
            stacklevel=3,
        )
        return self._rng.uniform(0.0, 1.0, size=n)

    def _compute_primary_loss(
        self, lmf: LossMagnitudeFactors
    ) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
        """Compute primary loss magnitude and per-category breakdown.

        Priority order:
        1. Direct ``primary_loss`` distribution.
        2. Sum of individual loss category distributions.
        3. Raise :class:`ValueError`.

        Returns
        -------
        (primary_loss_array, category_dict)
        """
        n = self.n_simulations
        category_map = {
            "productivity": lmf.productivity,
            "response": lmf.response,
            "replacement": lmf.replacement,
            "competitive_advantage": lmf.competitive_advantage,
            "fines_judgments": lmf.fines_judgments,
            "reputation": lmf.reputation,
        }
        active_cats = {k: v for k, v in category_map.items() if v is not None}

        if lmf.primary_loss is not None:
            pl = np.maximum(self._sample(lmf.primary_loss, n), 0.0)
            return pl, {}

        if active_cats:
            cat_arrays: Dict[str, np.ndarray] = {}
            for cat_name, dist in active_cats.items():
                cat_arrays[cat_name] = np.maximum(self._sample(dist, n), 0.0)
            pl = sum(cat_arrays.values())
            return pl, cat_arrays

        raise ValueError(
            "LossMagnitudeFactors must supply at least one of: "
            "primary_loss, or one or more loss category distributions "
            "(productivity, response, replacement, competitive_advantage, "
            "fines_judgments, reputation)."
        )

    def _compute_secondary_loss(
        self, lmf: LossMagnitudeFactors
    ) -> Optional[np.ndarray]:
        """Compute secondary loss array, or return None if not specified.

        Secondary loss is modelled as:
          secondary_loss = Binomial(1, SLEF) * SLM
        where SLEF is the per-iteration secondary loss event probability/frequency
        and SLM is the secondary loss magnitude per occurrence.
        """
        n = self.n_simulations

        if lmf.secondary_loss_event_frequency is None or lmf.secondary_loss_magnitude is None:
            return None

        slef = np.clip(self._sample(lmf.secondary_loss_event_frequency, n), 0.0, 1.0)
        slm = np.maximum(self._sample(lmf.secondary_loss_magnitude, n), 0.0)

        # Bernoulli occurrence based on SLEF probability
        occurrence = self._rng.binomial(1, slef).astype(float)
        return occurrence * slm

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def simulate(self, scenario: FAIRScenario) -> SimulationResult:
        """Run a Monte Carlo simulation for a single :class:`~pycrq.ontology.FAIRScenario`.

        Parameters
        ----------
        scenario: The scenario to simulate.

        Returns
        -------
        :class:`SimulationResult` containing per-iteration arrays and
        summary statistics.
        """
        n = self.n_simulations
        ff = scenario.frequency_factors
        lmf = scenario.loss_magnitude_factors
        T = scenario.time_horizon_years

        # --- Frequency ---
        tef = self._compute_tef(ff)
        vuln = self._compute_vulnerability(ff)
        # Loss Event Frequency = TEF × Vulnerability × time horizon
        lef = tef * vuln * T

        # --- Loss Magnitude ---
        primary_loss, loss_by_category = self._compute_primary_loss(lmf)
        secondary_loss = self._compute_secondary_loss(lmf)

        if secondary_loss is not None:
            lm = primary_loss + secondary_loss
        else:
            lm = primary_loss.copy()

        # --- Annual Loss Exposure ---
        ale = lef * lm

        return SimulationResult(
            scenario_name=scenario.name,
            n_simulations=n,
            loss_event_frequency=lef,
            loss_magnitude=lm,
            annual_loss_exposure=ale,
            threat_event_frequency=tef,
            vulnerability=vuln,
            primary_loss=primary_loss,
            secondary_loss=secondary_loss,
            loss_by_category=loss_by_category,
        )

    def simulate_register(self, register: RiskRegister) -> AggregateSimulationResult:
        """Simulate all scenarios in a :class:`~pycrq.ontology.RiskRegister`.

        Parameters
        ----------
        register: The risk register to simulate.

        Returns
        -------
        :class:`AggregateSimulationResult` with per-scenario and aggregate
        statistics.
        """
        if not register.scenarios:
            warnings.warn(
                f"RiskRegister '{register.name}' contains no scenarios.",
                stacklevel=2,
            )

        results = [self.simulate(s) for s in register.scenarios]

        return AggregateSimulationResult(
            register_name=register.name,
            n_simulations=self.n_simulations,
            scenario_results=results,
        )
