"""Marginal transform utilities for insurance loss data.

The vine copula operates in the unit hypercube [0,1]^d. Before fitting,
observed losses must be transformed to pseudo-observations on [0,1] via
the probability integral transform (PIT). After simulation, uniform
samples must be back-transformed to loss space using fitted marginals.

We use the nonparametric PIT (empirical CDF ranks) rather than assuming
a parametric family for marginals. This is standard practice in
copula modelling: fit the dependence structure separately from the
marginal distributions.
"""

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


# Numerical stability bounds for pseudo-observations.
# Exact 0 and 1 cause problems in copula log-likelihood evaluation.
_PSEUDO_OBS_CLIP = (1e-10, 1.0 - 1e-10)


def to_pseudo_obs(data: pd.DataFrame) -> np.ndarray:
    """Transform loss data to pseudo-observations via nonparametric PIT.

    Uses pyvinecopulib's ``to_pseudo_obs`` if available, otherwise falls
    back to a scipy-based rank transform. The result is clipped to
    ``(1e-10, 1-1e-10)`` for numerical stability during vine fitting.

    Parameters
    ----------
    data:
        DataFrame of shape ``(n_obs, d)`` containing non-negative loss
        values. Columns are perils.

    Returns
    -------
    np.ndarray
        Array of shape ``(n_obs, d)`` with values in ``(0, 1)``.

    Raises
    ------
    ValueError
        If ``data`` contains NaN values.
    """
    if data.isnull().any().any():
        raise ValueError(
            "Loss data contains NaN values. Fill or remove them before "
            "fitting the vine copula."
        )
    if len(data) < 2:
        raise ValueError(
            f"Need at least 2 observations for pseudo-obs transform, "
            f"got {len(data)}."
        )

    try:
        import pyvinecopulib as pv  # type: ignore[import-untyped]

        u = pv.to_pseudo_obs(data.values.astype(float))
    except ImportError:
        u = _pseudo_obs_fallback(data.values.astype(float))

    u = np.clip(u, *_PSEUDO_OBS_CLIP)
    return u


def _pseudo_obs_fallback(x: np.ndarray) -> np.ndarray:
    """Fallback pseudo-obs using scipy rankdata.

    Produces the same result as pv.to_pseudo_obs: i/(n+1) ranks.

    Parameters
    ----------
    x:
        Array of shape ``(n, d)``.

    Returns
    -------
    np.ndarray
        Pseudo-observations in ``(0, 1)``.
    """
    from scipy.stats import rankdata

    n = x.shape[0]
    u = np.empty_like(x, dtype=float)
    for j in range(x.shape[1]):
        u[:, j] = rankdata(x[:, j]) / (n + 1)
    return u


def back_transform(
    u: np.ndarray,
    perils: list[str],
    marginals: dict[str, "rv_continuous"],
) -> pd.DataFrame:
    """Back-transform uniform copula samples to loss space via marginal PPFs.

    Parameters
    ----------
    u:
        Array of shape ``(n, d)`` with values in ``[0, 1]``.
    perils:
        List of peril names corresponding to columns of ``u``.
    marginals:
        Mapping from peril name to a fitted scipy ``rv_continuous``
        instance (or any object with a ``.ppf(q)`` method).

    Returns
    -------
    pd.DataFrame
        DataFrame of shape ``(n, d)`` with loss values, one column per
        peril.

    Raises
    ------
    ValueError
        If any peril in ``perils`` is missing from ``marginals``.
    """
    missing = [p for p in perils if p not in marginals]
    if missing:
        raise ValueError(
            f"Marginals not provided for perils: {missing}. "
            f"All perils must have a corresponding scipy distribution."
        )

    u_arr = np.asarray(u, dtype=float)
    if u_arr.shape[1] != len(perils):
        raise ValueError(
            f"u has {u_arr.shape[1]} columns but {len(perils)} perils were "
            f"specified."
        )

    result = {}
    for j, peril in enumerate(perils):
        ppf = marginals[peril].ppf
        result[peril] = ppf(u_arr[:, j])

    return pd.DataFrame(result)


def validate_marginals(
    perils: list[str],
    marginals: dict[str, "rv_continuous"],
) -> None:
    """Validate that all perils have a marginal with a working PPF.

    Parameters
    ----------
    perils:
        Expected peril names.
    marginals:
        Candidate marginals dict.

    Raises
    ------
    ValueError
        If perils are missing or the PPF raises on a test call.
    """
    missing = [p for p in perils if p not in marginals]
    if missing:
        raise ValueError(f"Missing marginals for perils: {missing}")

    # Smoke-test each PPF at the median.
    for peril in perils:
        m = marginals[peril]
        if not hasattr(m, "ppf"):
            raise TypeError(
                f"Marginal for '{peril}' must have a .ppf() method. "
                f"Pass a scipy.stats distribution or a frozen distribution."
            )
        try:
            val = float(m.ppf(0.5))
        except Exception as exc:
            raise ValueError(
                f"PPF of marginal for '{peril}' failed at p=0.5: {exc}"
            ) from exc
        if not np.isfinite(val):
            warnings.warn(
                f"PPF of marginal for '{peril}' returned non-finite value "
                f"{val} at p=0.5. Check your marginal parameters.",
                stacklevel=2,
            )
