# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
from __future__ import annotations

import math

import numpy as np
import pytest

from agentic_valuation.debate.aggregation import RULES, aggregate, aggregate_intervals, aggregate_points, fit_weights, load_weights, median_of_quantiles, mixture_quantiles, save_weights
from agentic_valuation.schemas import DebateTranscript

ROLES = ["a", "b", "c", "d", "e"]


def _tr(tx, points, q=None):
    return DebateTranscript(tx_id=tx, arm="F", committee="c", protocol="p", rounds=0, packet_hash="h", turns=[], final_estimates=dict(zip(ROLES, points)), final_quantiles={r: (q or [0.8 * p, 0.9 * p, p, 1.1 * p, 1.25 * p]) for r, p in zip(ROLES, points)})


def test_point_rules_log_space():
    pts = {"a": 400000.0, "b": 500000.0, "c": 450000.0, "d": 480000.0, "e": 2000000.0}  # e is an outlier
    out = aggregate_points(pts)
    assert set(RULES) <= set(out)
    assert math.isclose(out["mean"], math.exp(np.mean(np.log(list(pts.values())))))
    assert out["median"] == pytest.approx(480000.0)
    assert out["trimmed_mean_20"] < out["mean"] and out["robust_bayes"] < out["mean"]  # robust to the outlier
    assert abs(out["robust_bayes"] - 470000) < 40000
    assert out["inverse_variance"] == pytest.approx(out["mean"]) and out["stacking"] == pytest.approx(out["mean"])  # equal weights when none supplied
    w = {"inverse_variance": {"a": 1.0, "b": 0, "c": 0, "d": 0, "e": 0}}
    assert math.isclose(aggregate_points(pts, w)["inverse_variance"], 400000.0)
    assert aggregate_points({}) == {}


def test_intervals():
    q = {"a": [380000, 420000, 450000, 480000, 530000], "b": [400000, 440000, 470000, 500000, 560000], "c": [300000, 350000, 400000, 450000, 500000]}
    m = median_of_quantiles(q)
    assert m == pytest.approx([380000, 420000, 450000, 480000, 530000])
    mx = mixture_quantiles(q)
    assert all(b >= a for a, b in zip(mx, mx[1:])) and 300000 <= mx[0] <= 400000 and 500000 <= mx[4] <= 560000
    iv = aggregate_intervals(q)
    assert set(iv) == {"median_of_quantiles", "mixture"}


def test_fit_weights_and_apply(tmp_path):
    rng = np.random.default_rng(0)
    sds = [0.05, 0.10, 0.20, 0.30, 0.40]
    truth, trs = {}, []
    for i in range(60):
        p = 400000.0 * math.exp(rng.normal(0, 0.3))
        truth[f"t{i}"] = p
        trs.append(_tr(f"t{i}", [p * math.exp(rng.normal(0, s)) for s in sds]))
    w = fit_weights(trs, truth)
    iv = w["inverse_variance"]
    assert math.isclose(sum(iv.values()), 1.0) and iv["a"] > iv["b"] > iv["e"]
    st = w["stacking"]
    assert all(v >= 0 for v in st.values()) and math.isclose(sum(st.values()), 1.0)
    assert 0.0 < w["robust_bayes"]["scale"] < 1.0 and w["robust_bayes"]["df"] == 4
    path = save_weights(w, tmp_path / "w.json")
    assert load_weights(path)["inverse_variance"] == pytest.approx(iv)
    # out of sample on fresh draws: inverse-variance beats the plain mean on average
    err_iv, err_mean = [], []
    for i in range(200):
        p = 400000.0 * math.exp(rng.normal(0, 0.3))
        t = _tr(f"v{i}", [p * math.exp(rng.normal(0, s)) for s in sds])
        agg = aggregate(t, w)
        assert set(RULES) <= set(agg)
        err_iv.append(abs(math.log(agg["inverse_variance"]) - math.log(p)))
        err_mean.append(abs(math.log(agg["mean"]) - math.log(p)))
    assert np.mean(err_iv) < np.mean(err_mean)