# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Interval calibration: PIT, reliability of the five quantiles, conditional coverage, sharpness and a
split-conformal rescaling of the log-interval (docs/evaluation_protocol.md §5)."""
from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Any

import numpy as np
import pandas as pd
from scipy import stats

from ._common import QUANTILE_COLUMNS, TAUS, to_pandas, validate_predictions
from .metrics import NOMINAL_LEVELS, fit_lognormal_from_quantiles, per_record


# ----------------------------------------------------------------------------------------------- PIT

def pit_from_quantiles(records: pd.DataFrame) -> np.ndarray:
    """Approximate probability integral transform F_i(log P_i) from the five quantiles.

    The predictive CDF is interpolated in *probit* space: z = Φ⁻¹(F) is piecewise linear in log price between the
    knots (log q_τ, Φ⁻¹(τ)), τ ∈ {.05,.25,.5,.75,.95}. This is exact when the forecast is log-normal and reduces to
    the usual piecewise-linear CDF otherwise up to the probit map. Below q05 / above q95 the tails of the normal fitted to
    the five quantiles are used, so PIT ∈ (0, 1). NaN when the record has no interval.
    """
    q = records[QUANTILE_COLUMNS].to_numpy(dtype=float)
    y = np.log(records["price"].to_numpy(dtype=float))
    out = np.full(len(records), np.nan)
    ok = np.all(np.isfinite(q), axis=1) & np.all(q > 0, axis=1)
    if not ok.any():
        return out
    lq = np.log(q[ok])
    yy = y[ok]
    mu, sigma = fit_lognormal_from_quantiles(lq)
    z_knots = stats.norm.ppf(TAUS)
    z = np.empty(len(yy))
    for i in range(len(yy)):
        xs = lq[i]
        if yy[i] < xs[0] or yy[i] > xs[-1] or not np.all(np.diff(xs) > 0):
            s = sigma[i] if sigma[i] > 0 else 1e-6
            z[i] = (yy[i] - mu[i]) / s
        else:
            z[i] = np.interp(yy[i], xs, z_knots)
    out[ok] = np.clip(stats.norm.cdf(z), 1e-9, 1 - 1e-9)
    return out


def pit_histogram(pit: np.ndarray, bins: int = 10) -> pd.DataFrame:
    """Histogram of PIT values with the uniform reference (`expected` = 1/bins)."""
    p = np.asarray(pit, dtype=float)
    p = p[np.isfinite(p)]
    counts, edges = np.histogram(p, bins=bins, range=(0.0, 1.0))
    share = counts / counts.sum() if counts.sum() else np.full(bins, np.nan)
    return pd.DataFrame({"bin_lo": edges[:-1], "bin_hi": edges[1:], "count": counts, "share": share, "expected": 1.0 / bins})


# ----------------------------------------------------------------------------------------------- reliability

def reliability(records: pd.DataFrame) -> pd.DataFrame:
    """Empirical frequency P(price ≤ q_τ) for each nominal τ, and interval coverage at each nominal level."""
    q = records[QUANTILE_COLUMNS].to_numpy(dtype=float)
    y = records["price"].to_numpy(dtype=float)
    ok = np.all(np.isfinite(q), axis=1)
    rows = []
    for k, tau in enumerate(TAUS):
        rows.append({"kind": "quantile", "nominal": float(tau), "empirical": float(np.mean(y[ok] <= q[ok, k])) if ok.any() else np.nan, "n": int(ok.sum())})
    for level in NOMINAL_LEVELS:
        col = {0.5: "cov_50", 0.8: "cov_80", 0.9: "cov_90", 0.95: "cov_95"}[level]
        v = records[col].to_numpy(dtype=float) if col in records.columns else np.full(len(records), np.nan)
        rows.append({"kind": "interval", "nominal": level, "empirical": float(np.nanmean(v)) if np.isfinite(v).any() else np.nan, "n": int(np.isfinite(v).sum())})
    df = pd.DataFrame(rows)
    df["gap"] = df["empirical"] - df["nominal"]
    return df


def coverage_by(records: pd.DataFrame, by: Sequence[str], levels: Sequence[float] = NOMINAL_LEVELS, min_n: int = 30) -> pd.DataFrame:
    """Conditional coverage per subgroup (long table). Cells with n < min_n are suppressed (NaN, suppressed=True)."""
    by = list(by)
    cols = {0.5: "cov_50", 0.8: "cov_80", 0.9: "cov_90", 0.95: "cov_95"}
    rows = []
    for key, g in records.groupby(by, sort=True, dropna=False):
        key = key if isinstance(key, tuple) else (key,)
        base = dict(zip(by, key))
        for level in levels:
            v = g[cols[level]].to_numpy(dtype=float)
            n = int(np.isfinite(v).sum())
            row = dict(base)
            row.update({"nominal": level, "n": n, "coverage": float(np.nanmean(v)) if n >= min_n else np.nan,
                        "mean_log_width": float(np.nanmean(g["width_90" if level == 0.9 else "width_50"])) if (n >= min_n and level in (0.5, 0.9)) else np.nan,
                        "suppressed": n < min_n})
            rows.append(row)
    return pd.DataFrame(rows)


def sharpness(records: pd.DataFrame) -> dict[str, float]:
    """Width statistics of the intervals (log scale and relative to the median)."""
    w90, w50 = records["width_90"].to_numpy(dtype=float), records["width_50"].to_numpy(dtype=float)
    q05, q95, q50 = (records[c].to_numpy(dtype=float) for c in ("q05", "q95", "q50"))
    rel = (q95 - q05) / q50
    f = lambda v: float(np.nanmean(v)) if np.isfinite(v).any() else float("nan")
    return {"mean_log_width_90": f(w90), "median_log_width_90": float(np.nanmedian(w90)) if np.isfinite(w90).any() else float("nan"),
            "mean_log_width_50": f(w50), "mean_relative_width_90": f(rel), "n": int(np.isfinite(w90).sum())}


# ----------------------------------------------------------------------------------------------- conformal

@dataclass(frozen=True)
class ConformalScales:
    """Multiplicative factors applied to the log half-widths around log q50 (one per interval)."""

    scale_90: float
    scale_50: float
    n_valid: int
    level_90: float = 0.90
    level_50: float = 0.50
    achieved_valid_90: float = float("nan")
    achieved_valid_50: float = float("nan")
    note: str = field(default="split conformal, asymmetric scaled residual score, centre = log q50")

    def as_dict(self) -> dict[str, float]:
        return {"scale_90": self.scale_90, "scale_50": self.scale_50, "n_valid": self.n_valid,
                "achieved_valid_90": self.achieved_valid_90, "achieved_valid_50": self.achieved_valid_50}


def _scores(logy: np.ndarray, c: np.ndarray, lo: np.ndarray, hi: np.ndarray) -> np.ndarray:
    eps = 1e-6
    lo, hi = np.maximum(lo, eps), np.maximum(hi, eps)
    return np.maximum((c - logy) / lo, (logy - c) / hi)


def _conformal_quantile(s: np.ndarray, level: float) -> float:
    n = len(s)
    k = int(np.ceil((n + 1) * level))
    if k > n:
        return float(np.max(s))
    return float(np.sort(s)[k - 1])


def conformalise(valid_preds: Any, valid_truth: Any, level_90: float = 0.90, level_50: float = 0.50) -> ConformalScales:
    """Split-conformal scale factors estimated on the validation table (pure function).

    Score s_i = max((c_i − y_i)/ℓ_i, (y_i − c_i)/u_i) with c = log q50, ℓ = log q50 − log q05, u = log q95 − log q50
    (and the analogous q25/q75 pair for the 50 % interval). The factor is the ⌈(n+1)(1−α)⌉-th smallest score,
    which guarantees marginal coverage ≥ 1−α on exchangeable test data (Vovk et al. 2005; Romano et al. 2019).
    Factor > 1 widens, < 1 narrows.
    """
    d = per_record(valid_preds, valid_truth)
    d = d[d["has_interval"]]
    if d.empty:
        raise ValueError("no validation records with intervals")
    y = np.log(d["price"].to_numpy(dtype=float))
    lq = np.log(d[QUANTILE_COLUMNS].to_numpy(dtype=float))
    c = lq[:, 2]
    s90 = _scores(y, c, c - lq[:, 0], lq[:, 4] - c)
    s50 = _scores(y, c, c - lq[:, 1], lq[:, 3] - c)
    k90, k50 = _conformal_quantile(s90, level_90), _conformal_quantile(s50, level_50)
    return ConformalScales(scale_90=k90, scale_50=k50, n_valid=len(d),
                           achieved_valid_90=float(np.mean(s90 <= k90)), achieved_valid_50=float(np.mean(s50 <= k50)))


def apply_conformal(preds: Any, scales: ConformalScales) -> pd.DataFrame:
    """Return a copy of a prediction table with q05/q95 (and q25/q75) rescaled around log q50; q50 and point unchanged."""
    p = validate_predictions(to_pandas(preds))
    lq = np.log(p[QUANTILE_COLUMNS].to_numpy(dtype=float))
    c = lq[:, 2]
    new = lq.copy()
    new[:, 0] = c - scales.scale_90 * np.maximum(c - lq[:, 0], 1e-6)
    new[:, 4] = c + scales.scale_90 * np.maximum(lq[:, 4] - c, 1e-6)
    new[:, 1] = c - scales.scale_50 * np.maximum(c - lq[:, 1], 1e-6)
    new[:, 3] = c + scales.scale_50 * np.maximum(lq[:, 3] - c, 1e-6)
    # keep the point inside [q05, q95] and the quantiles ordered
    new = np.sort(new, axis=1)
    out = p.copy()
    for k, col in enumerate(QUANTILE_COLUMNS):
        out[col] = np.exp(new[:, k])
    lp = np.log(out["point"].to_numpy(dtype=float))
    out["q05"] = np.minimum(out["q05"], np.exp(lp))
    out["q95"] = np.maximum(out["q95"], np.exp(lp))
    return out