"""PerilVine: exposure-weighted vine copula for multi-peril insurance losses.

The core idea is simple: perils in a UK home insurance portfolio are
correlated. Flood and escape of water tend to spike together. Storm
damage co-moves with flooding. If you price each peril independently and
then add them up, you underestimate the tail risk — the scenarios where
multiple perils hit at once.

A vine copula separates the dependence structure from the marginal
distributions. We fit the vine on pseudo-observations (empirical CDF
ranks) and optionally weight each observation by the policy's earned
exposure. This means policies that were at risk for a full year have
more influence on the fitted structure than short-period covers.

Typical usage::

    from insurance_copula import PerilVine
    from scipy.stats import gamma

    vine = PerilVine(perils=['flood', 'storm', 'eow', 'subsidence'])
    vine.fit(losses_df, exposure=earned_premium_series)

    samples = vine.simulate(10_000, marginals={
        'flood': gamma(a=0.8, scale=5_000),
        'storm': gamma(a=1.2, scale=800),
        'eow':   gamma(a=1.5, scale=600),
        'subsidence': gamma(a=0.6, scale=3_000),
    })

    diag = vine.diagnostics()
    print(diag.kendall_tau_matrix)
"""

from __future__ import annotations

import warnings
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd

if TYPE_CHECKING:
    from scipy.stats import rv_continuous
    from ._results import AggregateResult, PerilVineDiagnostics

# Attempt to import pyvinecopulib. If unavailable (ARM64 / no wheels),
# all methods that need it will raise ImportError with a clear message.
try:
    import pyvinecopulib as pv  # type: ignore[import-untyped]
    HAS_PYVINECOPULIB = True
except ImportError:
    HAS_PYVINECOPULIB = False

_FAMILY_SET_MAP: dict[str, object] = {}
_VINE_TYPE_MAP: dict[str, str] = {
    "rvine": "rvine",
    "cvine": "cvine",
    "dvine": "dvine",
}

_VALID_SELECTION_CRITERIA = {"aic", "bic", "loglik"}
_VALID_TREE_CRITERIA = {"tau", "hoeffd", "rho", "mcor"}


def _build_family_set_map() -> None:
    """Populate _FAMILY_SET_MAP once pyvinecopulib is available."""
    global _FAMILY_SET_MAP
    if not HAS_PYVINECOPULIB:
        return
    _FAMILY_SET_MAP = {
        "parametric": pv.parametric,
        "archimedean": pv.archimedean,
        "all": pv.all,
        "tll": pv.tll,
    }


if HAS_PYVINECOPULIB:
    _build_family_set_map()


class PerilVine:
    """Vine copula for multi-peril insurance loss dependence.

    Wraps ``pyvinecopulib.Vinecop`` with an insurance-specific workflow:
    exposure-weighted fitting, automatic pseudo-observation transform,
    conditional expectation estimation, and aggregate loss simulation.

    Parameters
    ----------
    perils:
        Ordered list of peril names, e.g.
        ``['flood', 'subsidence', 'storm', 'eow']``.
    family_set:
        Pair-copula families to consider during structure selection.
        One of ``'parametric'``, ``'archimedean'``, ``'all'``, ``'tll'``.
        Defaults to ``'parametric'``. Use ``'all'`` to include
        nonparametric TLL copulas (slower but more flexible).
    selection_criterion:
        Information criterion for pair-copula family selection.
        One of ``'aic'``, ``'bic'``, ``'loglik'``. Defaults to
        ``'bic'``. BIC penalises complexity harder and reduces
        over-fitting on small insurance datasets.
    tree_criterion:
        Criterion for selecting the spanning tree at each level of the
        vine. One of ``'tau'``, ``'hoeffd'``, ``'rho'``, ``'mcor'``.
        Defaults to ``'tau'`` (Kendall's tau — fast and robust).
    trunc_lvl:
        Truncation level: only fit pair copulas up to this tree level.
        ``None`` means fit a full vine (all ``d-1`` trees). For 5-6
        perils, ``trunc_lvl=3`` is typically sufficient and avoids
        over-fitting sparse higher-order conditional dependencies.
    vine_type:
        Vine structure type. One of ``'rvine'``, ``'cvine'``,
        ``'dvine'``. Defaults to ``'rvine'`` (most flexible). C-vine
        and D-vine impose structure useful when one peril dominates
        (C-vine) or losses are ordered in time (D-vine).
    n_threads:
        Number of threads for parallel fitting. Defaults to 1.

    Raises
    ------
    ValueError
        If any argument is invalid.
    """

    def __init__(
        self,
        perils: list[str],
        family_set: str = "parametric",
        selection_criterion: str = "bic",
        tree_criterion: str = "tau",
        trunc_lvl: int | None = 3,
        vine_type: str = "rvine",
        n_threads: int = 1,
    ) -> None:
        self._validate_init(
            perils, family_set, selection_criterion, tree_criterion,
            trunc_lvl, vine_type, n_threads,
        )
        self.perils = list(perils)
        self.family_set = family_set
        self.selection_criterion = selection_criterion
        self.tree_criterion = tree_criterion
        self.trunc_lvl = trunc_lvl
        self.vine_type = vine_type
        self.n_threads = n_threads

        self._vine: object | None = None
        self._u: np.ndarray | None = None  # pseudo-obs used for fitting
        self._n_obs: int | None = None
        self._fitted: bool = False

    # ------------------------------------------------------------------
    # Validation
    # ------------------------------------------------------------------

    @staticmethod
    def _validate_init(
        perils: list[str],
        family_set: str,
        selection_criterion: str,
        tree_criterion: str,
        trunc_lvl: int | None,
        vine_type: str,
        n_threads: int,
    ) -> None:
        if not perils:
            raise ValueError("perils must be a non-empty list")
        if len(perils) < 2:
            raise ValueError(
                f"Need at least 2 perils to fit a vine copula, got {len(perils)}"
            )
        if len(set(perils)) != len(perils):
            raise ValueError(f"Peril names must be unique, got: {perils}")

        valid_family_sets = {"parametric", "archimedean", "all", "tll"}
        if family_set not in valid_family_sets:
            raise ValueError(
                f"family_set must be one of {valid_family_sets}, "
                f"got '{family_set}'"
            )
        if selection_criterion not in _VALID_SELECTION_CRITERIA:
            raise ValueError(
                f"selection_criterion must be one of "
                f"{_VALID_SELECTION_CRITERIA}, got '{selection_criterion}'"
            )
        if tree_criterion not in _VALID_TREE_CRITERIA:
            raise ValueError(
                f"tree_criterion must be one of {_VALID_TREE_CRITERIA}, "
                f"got '{tree_criterion}'"
            )
        if trunc_lvl is not None:
            if not isinstance(trunc_lvl, int) or trunc_lvl < 1:
                raise ValueError(
                    f"trunc_lvl must be a positive integer or None, "
                    f"got {trunc_lvl!r}"
                )
        if vine_type not in _VINE_TYPE_MAP:
            raise ValueError(
                f"vine_type must be one of {set(_VINE_TYPE_MAP)}, "
                f"got '{vine_type}'"
            )
        if not isinstance(n_threads, int) or n_threads < 1:
            raise ValueError(
                f"n_threads must be a positive integer, got {n_threads!r}"
            )

    # ------------------------------------------------------------------
    # Fitting
    # ------------------------------------------------------------------

    def fit(
        self,
        losses: pd.DataFrame,
        exposure: pd.Series | None = None,
    ) -> "PerilVine":
        """Fit the vine copula to peril loss data.

        Transforms losses to pseudo-observations via nonparametric PIT,
        then fits the vine using pyvinecopulib's structure selection.

        Parameters
        ----------
        losses:
            DataFrame with one column per peril (columns must match
            ``self.perils``). Rows are observations (typically one per
            policy-year). Values should be non-negative loss amounts.
        exposure:
            Optional earned exposure per observation (e.g. days at risk,
            earned premium units). Passed as weights to
            ``FitControlsVinecop``. Must have the same index as
            ``losses``.

        Returns
        -------
        PerilVine
            Self, to allow method chaining.

        Raises
        ------
        ImportError
            If pyvinecopulib is not installed.
        ValueError
            If column names don't match perils, or shapes are
            inconsistent.
        """
        self._validate_losses(losses, exposure)
        self._require_pyvinecopulib()

        from ._marginals import to_pseudo_obs

        # Reorder columns to match self.perils ordering
        losses_ordered = losses[self.perils]
        u = to_pseudo_obs(losses_ordered)

        controls_kwargs: dict = {
            "family_set": _FAMILY_SET_MAP[self.family_set],
            "selection_criterion": self.selection_criterion,
            "tree_criterion": self.tree_criterion,
            "num_threads": self.n_threads,
        }

        if self.trunc_lvl is not None:
            controls_kwargs["trunc_lvl"] = self.trunc_lvl

        if exposure is not None:
            weights = exposure.reindex(losses.index).fillna(1.0).values
            weights = weights.astype(float)
            if np.any(weights <= 0):
                raise ValueError(
                    "Exposure weights must all be positive. "
                    "Found non-positive values."
                )
            controls_kwargs["weights"] = weights

        controls = pv.FitControlsVinecop(**controls_kwargs)

        self._vine = pv.Vinecop.from_data(u, controls=controls)
        self._u = u
        self._n_obs = len(losses)
        self._fitted = True

        return self

    def _validate_losses(
        self,
        losses: pd.DataFrame,
        exposure: pd.Series | None,
    ) -> None:
        """Validate loss DataFrame and optional exposure series."""
        if not isinstance(losses, pd.DataFrame):
            raise TypeError(
                f"losses must be a pandas DataFrame, got {type(losses).__name__}"
            )
        missing_cols = [p for p in self.perils if p not in losses.columns]
        if missing_cols:
            raise ValueError(
                f"Loss DataFrame is missing columns for perils: {missing_cols}. "
                f"Available columns: {list(losses.columns)}"
            )
        if len(losses) < 2:
            raise ValueError(
                f"Need at least 2 observations to fit a vine copula, "
                f"got {len(losses)}"
            )
        if exposure is not None:
            if len(exposure) != len(losses):
                raise ValueError(
                    f"exposure has {len(exposure)} rows but losses has "
                    f"{len(losses)} rows. They must match."
                )

    # ------------------------------------------------------------------
    # Simulation
    # ------------------------------------------------------------------

    def simulate(
        self,
        n: int,
        marginals: dict[str, "rv_continuous"] | None = None,
        seed: int | None = None,
    ) -> pd.DataFrame:
        """Simulate from the fitted vine copula.

        Parameters
        ----------
        n:
            Number of samples to draw.
        marginals:
            Optional dict mapping peril name to a fitted scipy
            ``rv_continuous`` instance. If provided, samples are
            back-transformed from ``[0,1]`` to loss space. If None,
            returns copula samples on ``[0,1]``.
        seed:
            Random seed for reproducibility.

        Returns
        -------
        pd.DataFrame
            DataFrame of shape ``(n, d)`` with one column per peril.
            Values are on ``[0,1]`` if no marginals, else in loss space.

        Raises
        ------
        RuntimeError
            If the model has not been fitted yet.
        """
        self._require_fitted()

        if n < 1:
            raise ValueError(f"n must be positive, got {n}")

        seeds: list[int] = [seed] if seed is not None else []
        u = self._vine.simulate(n, seeds=seeds)  # type: ignore[union-attr]

        if marginals is not None:
            from ._marginals import back_transform, validate_marginals
            validate_marginals(self.perils, marginals)
            return back_transform(u, self.perils, marginals)

        return pd.DataFrame(u, columns=self.perils)

    # ------------------------------------------------------------------
    # Conditional expectation
    # ------------------------------------------------------------------

    def conditional_expectation(
        self,
        target_peril: str,
        condition_peril: str,
        threshold: float,
        n_samples: int = 50_000,
        marginals: dict[str, "rv_continuous"] | None = None,
    ) -> float:
        """Estimate E[target_peril | condition_peril > threshold].

        Uses accept-reject Monte Carlo. The threshold is interpreted in
        loss space if marginals are provided, otherwise on ``[0,1]``.

        Parameters
        ----------
        target_peril:
            Peril whose conditional expected loss we want.
        condition_peril:
            Conditioning peril.
        threshold:
            Threshold for the conditioning peril.
        n_samples:
            Number of unconditional draws for accept-reject. Increase
            for thresholds in the far tail.
        marginals:
            Optional scipy distributions for back-transformation.

        Returns
        -------
        float
            Estimated conditional expectation.
        """
        self._require_fitted()

        from ._conditional import conditional_expectation as _ce

        return _ce(
            vine=self._vine,
            target_peril=target_peril,
            condition_peril=condition_peril,
            threshold=threshold,
            perils=self.perils,
            n_samples=n_samples,
            marginals=marginals,
            seed=None,
        )

    # ------------------------------------------------------------------
    # Aggregate loss
    # ------------------------------------------------------------------

    def aggregate_loss(
        self,
        n_sim: int,
        marginals: dict[str, "rv_continuous"],
        n_policies: int,
        seed: int | None = None,
    ) -> "AggregateResult":
        """Simulate aggregate portfolio loss distribution.

        For each of ``n_sim`` scenarios, draws ``n_policies`` independent
        samples from the vine and sums all peril losses. Returns the
        empirical distribution of total portfolio loss.

        Parameters
        ----------
        n_sim:
            Number of scenarios.
        marginals:
            Fitted scipy distributions for each peril.
        n_policies:
            Number of policies in the portfolio per scenario.
        seed:
            Base random seed.

        Returns
        -------
        AggregateResult
            Contains the simulated distribution and PML quantiles.
        """
        self._require_fitted()

        from ._aggregate import simulate_aggregate_loss

        return simulate_aggregate_loss(
            vine=self._vine,
            n_sim=n_sim,
            marginals=marginals,
            perils=self.perils,
            n_policies=n_policies,
            seed=seed,
        )

    # ------------------------------------------------------------------
    # Diagnostics
    # ------------------------------------------------------------------

    def diagnostics(self) -> "PerilVineDiagnostics":
        """Compute model diagnostics for the fitted vine.

        Returns Kendall tau matrix, tail dependence coefficients,
        AIC/BIC, and a flat list of pair-copula summaries.

        Returns
        -------
        PerilVineDiagnostics
            Populated diagnostics object.

        Raises
        ------
        RuntimeError
            If the model has not been fitted yet.
        """
        self._require_fitted()

        from ._diagnostics import compute_diagnostics

        return compute_diagnostics(
            vine=self._vine,
            u=self._u,  # type: ignore[arg-type]
            perils=self.perils,
        )

    # ------------------------------------------------------------------
    # Serialisation
    # ------------------------------------------------------------------

    def to_json(self) -> str:
        """Serialise the fitted vine to a JSON string.

        Wraps pyvinecopulib's native serialisation with a metadata
        envelope containing peril names and hyperparameters.

        Returns
        -------
        str
            JSON string.

        Raises
        ------
        RuntimeError
            If the model has not been fitted yet.
        """
        self._require_fitted()

        from ._serialisation import serialise_vine

        return serialise_vine(
            vine=self._vine,
            perils=self.perils,
            family_set=self.family_set,
            selection_criterion=self.selection_criterion,
            tree_criterion=self.tree_criterion,
            trunc_lvl=self.trunc_lvl,
            vine_type=self.vine_type,
        )

    @classmethod
    def from_json(cls, s: str) -> "PerilVine":
        """Restore a PerilVine from a JSON string.

        Parameters
        ----------
        s:
            JSON string produced by :meth:`to_json`.

        Returns
        -------
        PerilVine
            Restored model. The ``_vine`` attribute is populated and
            ``_fitted`` is True.

        Raises
        ------
        ImportError
            If pyvinecopulib is not installed.
        ValueError
            If the JSON is malformed.
        """
        from ._serialisation import deserialise_vine

        vine_obj, metadata = deserialise_vine(s)

        instance = cls(
            perils=metadata["perils"],
            family_set=metadata["family_set"],
            selection_criterion=metadata["selection_criterion"],
            tree_criterion=metadata["tree_criterion"],
            trunc_lvl=metadata["trunc_lvl"],
            vine_type=metadata["vine_type"],
        )
        instance._vine = vine_obj
        instance._fitted = True
        return instance

    # ------------------------------------------------------------------
    # Guards
    # ------------------------------------------------------------------

    @staticmethod
    def _require_pyvinecopulib() -> None:
        if not HAS_PYVINECOPULIB:
            raise ImportError(
                "pyvinecopulib is required for vine copula fitting. "
                "Install it with: pip install 'insurance-copula[vine]'\n"
                "Note: pre-built wheels are available for x86_64 Linux, "
                "macOS (Intel/Apple Silicon), and Windows. "
                "ARM64 Linux (e.g. Raspberry Pi) requires building from "
                "source — use Databricks for testing on x86_64."
            )

    def _require_fitted(self) -> None:
        if not self._fitted or self._vine is None:
            raise RuntimeError(
                "PerilVine has not been fitted yet. "
                "Call .fit(losses) before using this method."
            )

    # ------------------------------------------------------------------
    # Properties and dunder methods
    # ------------------------------------------------------------------

    @property
    def is_fitted(self) -> bool:
        """True if the model has been fitted."""
        return self._fitted

    @property
    def n_perils(self) -> int:
        """Number of perils in the model."""
        return len(self.perils)

    @property
    def n_obs(self) -> int | None:
        """Number of observations used for fitting (None if not fitted)."""
        return self._n_obs

    def __repr__(self) -> str:
        status = "fitted" if self._fitted else "unfitted"
        return (
            f"PerilVine("
            f"perils={self.perils}, "
            f"family_set='{self.family_set}', "
            f"trunc_lvl={self.trunc_lvl}, "
            f"vine_type='{self.vine_type}', "
            f"status={status}"
            f")"
        )
