# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Aggregation of the committee's final estimates. Everything is computed in log space and converted back.

Point rules (configs/treatments.yaml ``aggregation_rules``):
  mean, median, trimmed_mean_20 (20 % trimmed from each tail, scipy convention), inverse_variance (weights
  = 1 / var of each role's validation log errors), stacking (non-negative least squares on log points,
  weights normalised to sum 1), robust_bayes (Student-t, 4 df, location by IRLS; scale from validation).
Interval rules: median of each quantile across roles (default) and the equal-weight mixture quantile.

Weights are estimated ONCE on validation transcripts (``fit_weights``), saved to
``results/aggregation_weights.json`` and applied unchanged to test/holdout. When no weights are supplied,
inverse_variance and stacking fall back to equal weights and robust_bayes uses a MAD scale from the
committee itself; the runner records ``weights_source`` in the manifest so that this is never silent.
"""
from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Any

import numpy as np
from scipy.optimize import nnls
from scipy.stats import trim_mean

from ..config import RESULTS
from ..schemas import DebateTranscript

RULES: tuple[str, ...] = ("mean", "median", "trimmed_mean_20", "inverse_variance", "stacking", "robust_bayes")
QUANTILE_LEVELS: tuple[float, ...] = (0.05, 0.25, 0.50, 0.75, 0.95)
T_DF = 4
WEIGHTS_PATH = RESULTS / "aggregation_weights.json"


# ---------------------------------------------------------------------------------------------------
# point rules
# ---------------------------------------------------------------------------------------------------
def _logs(points: dict[str, float]) -> tuple[list[str], np.ndarray]:
    roles = [r for r, v in points.items() if v is not None and v > 0]
    return roles, np.array([math.log(points[r]) for r in roles], dtype=float)


def _normalise(w: dict[str, float], roles: list[str]) -> np.ndarray:
    arr = np.array([max(float(w.get(r, 0.0)), 0.0) for r in roles], dtype=float)
    if arr.sum() <= 0:
        arr = np.ones(len(roles))
    return arr / arr.sum()


def student_t_location(x: np.ndarray, scale: float, df: int = T_DF, iters: int = 100, tol: float = 1e-10) -> float:
    """IRLS location estimate under a Student-t likelihood with known scale (robust Bayes posterior mode under a flat prior)."""
    if len(x) == 0:
        return float("nan")
    mu = float(np.median(x))
    scale = max(float(scale), 1e-6)
    for _ in range(iters):
        z2 = ((x - mu) / scale) ** 2
        w = (df + 1.0) / (df + z2)
        new = float(np.sum(w * x) / np.sum(w))
        if abs(new - mu) < tol:
            mu = new
            break
        mu = new
    return mu


def mad_scale(x: np.ndarray, floor: float = 0.02) -> float:
    if len(x) < 2:
        return floor
    return max(1.4826 * float(np.median(np.abs(x - np.median(x)))), floor)


def aggregate_points(points: dict[str, float], weights: dict[str, Any] | None = None) -> dict[str, float]:
    """rule -> aggregated value (CAD). Empty dict when no valid point."""
    roles, lx = _logs(points)
    if len(roles) == 0:
        return {}
    out: dict[str, float] = {}
    out["mean"] = math.exp(float(lx.mean()))
    out["median"] = math.exp(float(np.median(lx)))
    out["trimmed_mean_20"] = math.exp(float(trim_mean(lx, 0.2))) if len(lx) >= 3 else out["mean"]
    w_iv = _normalise((weights or {}).get("inverse_variance", {}), roles)
    out["inverse_variance"] = math.exp(float(np.dot(w_iv, lx)))
    w_st = _normalise((weights or {}).get("stacking", {}), roles)
    out["stacking"] = math.exp(float(np.dot(w_st, lx)))
    rb = (weights or {}).get("robust_bayes", {}) or {}
    scale = float(rb.get("scale")) if rb.get("scale") else mad_scale(lx)
    out["robust_bayes"] = math.exp(student_t_location(lx, scale=scale, df=int(rb.get("df", T_DF))))
    return out


def aggregate(transcript: DebateTranscript, weights: dict[str, Any] | None = None) -> dict[str, float]:
    return aggregate_points(transcript.final_estimates, weights)


# ---------------------------------------------------------------------------------------------------
# interval rules
# ---------------------------------------------------------------------------------------------------
def median_of_quantiles(quantiles: dict[str, list[float]]) -> list[float]:
    """Element-wise median across roles of [q05, q25, q50, q75, q95] (in log space)."""
    rows = [q for q in quantiles.values() if q and len(q) == 5 and all(v > 0 for v in q)]
    if not rows:
        return []
    arr = np.log(np.array(rows, dtype=float))
    return [float(v) for v in np.exp(np.median(arr, axis=0))]


def mixture_quantiles(quantiles: dict[str, list[float]], levels: tuple[float, ...] = QUANTILE_LEVELS, grid: int = 2001) -> list[float]:
    """Quantiles of the equal-weight mixture of the roles' distributions. Each role's distribution is
    represented by linear interpolation of its quantile function in log space (tails extended linearly
    to 0.5 % / 99.5 %)."""
    rows = [np.log(np.array(q, dtype=float)) for q in quantiles.values() if q and len(q) == 5 and all(v > 0 for v in q)]
    if not rows:
        return []
    p = np.array(QUANTILE_LEVELS)
    ext_p = np.concatenate([[0.005], p, [0.995]])
    cdfs = []
    for r in rows:
        lo = r[0] - (r[1] - r[0]) * (0.05 - 0.005) / 0.20 if r[1] > r[0] else r[0] - 0.05
        hi = r[4] + (r[4] - r[3]) * (0.995 - 0.95) / 0.20 if r[4] > r[3] else r[4] + 0.05
        cdfs.append((np.concatenate([[lo], r, [hi]]), ext_p))
    xmin = min(c[0][0] for c in cdfs)
    xmax = max(c[0][-1] for c in cdfs)
    xs = np.linspace(xmin, xmax, grid)
    mix = np.zeros_like(xs)
    for xq, pq in cdfs:
        xq_m = np.maximum.accumulate(xq)  # guard against flat/non-monotone quantiles
        mix += np.interp(xs, xq_m, pq, left=0.0, right=1.0)
    mix /= len(cdfs)
    out = []
    for lv in levels:
        idx = int(np.searchsorted(mix, lv, side="left"))
        idx = min(max(idx, 0), grid - 1)
        out.append(float(math.exp(xs[idx])))
    return out


def aggregate_intervals(transcript_or_quantiles: DebateTranscript | dict[str, list[float]]) -> dict[str, list[float]]:
    """rule -> [q05, q25, q50, q75, q95]."""
    q = transcript_or_quantiles.final_quantiles if isinstance(transcript_or_quantiles, DebateTranscript) else transcript_or_quantiles
    out: dict[str, list[float]] = {}
    m = median_of_quantiles(q)
    if m:
        out["median_of_quantiles"] = m
    mx = mixture_quantiles(q)
    if mx:
        out["mixture"] = mx
    return out


def interval_90(iv: dict[str, list[float]]) -> dict[str, list[float]]:
    """DebateTranscript.aggregated_interval contract: rule -> [lo90, hi90]."""
    return {k: [v[0], v[4]] for k, v in iv.items() if len(v) == 5}


# ---------------------------------------------------------------------------------------------------
# weights estimated on validation
# ---------------------------------------------------------------------------------------------------
def fit_weights(valid_transcripts: list[DebateTranscript], truth: dict[str, float], *, min_obs: int = 3, df: int = T_DF) -> dict[str, Any]:
    """Estimate aggregation weights on validation transcripts (truth = realised sale price by tx_id).

    inverse_variance: w_r ∝ 1 / Var(log point_r − log P) over the properties where role r answered.
    stacking: NNLS of log P on the roles' log points (properties where every role answered; missing
    roles imputed by the row median), weights normalised to sum 1.
    robust_bayes: scale = robust pooled SD (1.4826·MAD) of the role log errors; df fixed at 4.
    """
    roles: list[str] = []
    for t in valid_transcripts:
        for r in t.final_estimates:
            if r not in roles:
                roles.append(r)
    errs: dict[str, list[float]] = {r: [] for r in roles}
    rows: list[list[float]] = []
    ys: list[float] = []
    for t in valid_transcripts:
        p = truth.get(t.tx_id)
        if p is None or p <= 0 or not t.final_estimates:
            continue
        lp = math.log(p)
        row = []
        for r in roles:
            v = t.final_estimates.get(r)
            if v is not None and v > 0:
                errs[r].append(math.log(v) - lp)
                row.append(math.log(v))
            else:
                row.append(float("nan"))
        if any(not math.isnan(x) for x in row):
            med = float(np.nanmedian(row))
            rows.append([med if math.isnan(x) else x for x in row])
            ys.append(lp)
    pooled = [e for r in roles for e in errs[r]]
    pooled_var = float(np.var(pooled, ddof=1)) if len(pooled) > 1 else 0.05
    inv_var: dict[str, float] = {}
    for r in roles:
        e = errs[r]
        var = float(np.var(e, ddof=1)) if len(e) >= min_obs else pooled_var
        inv_var[r] = 1.0 / max(var, 1e-6)
    s = sum(inv_var.values())
    inv_var = {r: v / s for r, v in inv_var.items()} if s > 0 else {r: 1.0 / len(roles) for r in roles}
    stacking: dict[str, float]
    if len(rows) >= max(len(roles), 2):
        X = np.array(rows, dtype=float)
        y = np.array(ys, dtype=float)
        w, _ = nnls(X, y)
        if w.sum() <= 0:
            w = np.ones(len(roles))
        w = w / w.sum()
        stacking = {r: float(v) for r, v in zip(roles, w)}
    else:
        stacking = {r: 1.0 / len(roles) for r in roles} if roles else {}
    scale = mad_scale(np.array(pooled)) if pooled else 0.15
    return {
        "roles": roles,
        "n_transcripts": len(rows),
        "inverse_variance": inv_var,
        "stacking": stacking,
        "robust_bayes": {"df": df, "scale": float(scale)},
        "role_error_sd": {r: (float(np.std(errs[r], ddof=1)) if len(errs[r]) > 1 else None) for r in roles},
        "role_n": {r: len(errs[r]) for r in roles},
    }


def save_weights(weights: dict[str, Any], path: Path | None = None) -> Path:
    path = path or WEIGHTS_PATH
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(weights, indent=1, ensure_ascii=False))
    return path


def load_weights(path: Path | None = None) -> dict[str, Any] | None:
    path = path or WEIGHTS_PATH
    if not path.exists():
        return None
    return json.loads(path.read_text())