# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Tests of the evaluation layer: hand-computed metrics, bootstrap coverage, Holm, interval coverage,
conformal rescaling, faithfulness audit on a Pydantic-built transcript, LaTeX fragment hygiene."""
from __future__ import annotations

import math
from datetime import date, timedelta
from pathlib import Path

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

from agentic_valuation.evaluation import (apply_conformal, block_bootstrap_contrast, compare_arms, conformalise, coverage_by,
                                           diebold_mariano, holm, per_record, pit_from_quantiles, reliability, stratified_metrics,
                                           summarise, summarise_by)
from agentic_valuation.evaluation.faithfulness import audit_estimate, audit_run, audit_transcript, summarise_faithfulness
from agentic_valuation.evaluation.tables import arm_pretty, escape, num, write_contrasts, write_coverage, write_main_results, write_strata
from agentic_valuation.schemas import (AgentEstimate, AgentTurn, Adjustment, Comparable, DebateTranscript, EvidencePacket,
                                       MarketContext, SubjectProperty)

LOG125 = math.log(1.25)


# ----------------------------------------------------------------------------------------------- helpers

def _truth(n: int, seed: int = 1) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    d0 = date(2025, 4, 1)
    return pd.DataFrame({
        "tx_id": [f"T{i:05d}" for i in range(n)],
        "price": np.exp(rng.normal(math.log(400_000), 0.4, n)).round(-3),
        "date": [d0 + timedelta(days=int(x)) for x in rng.integers(0, 270, n)],
        "mun_code": rng.choice([f"M{k}" for k in range(12)], n),
        "region_name": rng.choice(["Montréal", "Capitale-Nationale", "Outaouais", "Estrie"], n),
        "property_class": rng.choice(["house", "condo", "plex"], n),
        "urban_class": rng.choice(["metro", "suburban", "rural"], n),
        "year_built": rng.integers(1900, 2024, n),
        "floor_area_m2": rng.uniform(60, 250, n),
    })


def _preds(truth: pd.DataFrame, arm: str, sigma: float, bias: float = 0.0, seed: int = 2, intervals: bool = True,
           shared: np.ndarray | None = None, iscale: float = 1.0) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    e = rng.normal(bias, sigma, len(truth)) + (shared if shared is not None else 0.0)
    point = truth["price"].to_numpy() * np.exp(e)
    z = stats.norm.ppf([0.05, 0.25, 0.5, 0.75, 0.95])
    s = sigma * iscale
    q = np.exp(np.log(point)[:, None] + s * z[None, :]) if intervals else np.full((len(truth), 5), np.nan)
    return pd.DataFrame({"tx_id": truth["tx_id"], "split": "test", "arm": arm, "committee": "mixed_random", "aggregation": "median",
                         "point": point, "q05": q[:, 0], "q25": q[:, 1], "q50": q[:, 2], "q75": q[:, 3], "q95": q[:, 4],
                         "cost_usd": 0.5, "latency_s": 10.0, "tokens": 1000, "run_id": "r1"})


# ----------------------------------------------------------------------------------------------- metrics

def test_metrics_hand_computed():
    truth = pd.DataFrame({"tx_id": ["a", "b", "c", "d"], "price": [100.0, 200.0, 300.0, 400.0]})
    preds = pd.DataFrame({
        "tx_id": ["a", "b", "c", "d"], "arm": "X", "point": [110.0, 180.0, 300.0, 600.0],
        "q05": [90.0, 150.0, 250.0, 500.0], "q25": [100.0, 170.0, 280.0, 550.0], "q50": [110.0, 180.0, 300.0, 600.0],
        "q75": [120.0, 190.0, 320.0, 650.0], "q95": [130.0, 210.0, 350.0, 700.0], "cost_usd": [1, 1, 1, 1],
    })
    e = np.log(np.array([110, 180, 300, 600]) / np.array([100, 200, 300, 400]))
    s = summarise(preds, truth)
    assert s["n"] == 4
    assert s["mean_abs_log_error"] == pytest.approx(np.mean(np.abs(e)))
    assert s["mae_cad"] == pytest.approx((10 + 20 + 0 + 200) / 4)
    assert s["mdape"] == pytest.approx(np.median([0.1, 0.1, 0.0, 0.5]))
    assert s["rmse_log"] == pytest.approx(np.sqrt(np.mean(e**2)))
    assert s["bias_log"] == pytest.approx(np.mean(e))
    assert s["severe_rate"] == pytest.approx(0.25)          # only d: log(1.5) > log(1.25)
    assert s["over_rate"] == pytest.approx(0.25)            # d: +50 %
    assert s["under_rate"] == pytest.approx(0.0)            # b: −10 % only
    # raw coverage: 50 % → price in [q25,q75]: a yes (100∈[100,120]), b no (200∉[170,190]), c yes, d no → 0.5
    assert s["coverage_50"] == pytest.approx(0.5)
    # 90 %: a yes, b no (200∉[150,210]? yes it is), c yes (300∈[250,350]), d no (400∉[500,700]) → 0.75
    assert s["coverage_90"] == pytest.approx(0.75)
    assert s["mean_log_width_90"] == pytest.approx(np.mean(np.log([130 / 90, 210 / 150, 350 / 250, 700 / 500])))
    assert s["cost_usd_total"] == pytest.approx(4.0)
    # pinball for record c (price 300): q=[250,280,300,320,350] → τ·(y−q) for y≥q, (τ−1)(y−q) otherwise, in log space
    y = math.log(300)
    lq = np.log([250, 280, 300, 320, 350])
    taus = [0.05, 0.25, 0.5, 0.75, 0.95]
    pb = np.mean([t * (y - q) if y >= q else (t - 1) * (y - q) for t, q in zip(taus, lq)])
    rec = per_record(preds, truth)
    assert rec.loc[rec["tx_id"] == "c", "pinball_log"].iloc[0] == pytest.approx(pb)


def test_interval_metrics_nan_without_quantiles():
    truth = _truth(50)
    p = _preds(truth, "A", 0.1, intervals=False)
    s = summarise(p, truth)
    assert np.isnan(s["coverage_90"]) and np.isnan(s["pinball_log"]) and np.isnan(s["crps_log"])
    assert np.isfinite(s["mean_abs_log_error"])


def test_normal_fit_coverage_matches_raw_on_normal_quantiles():
    truth = _truth(4000)
    p = _preds(truth, "D", 0.15)
    s = summarise(p, truth)
    # quantiles are exactly log-normal → fitted 90 % band coincides with [q05,q95] and coverage ≈ nominal
    assert s["coverage_90_normal_fit"] == pytest.approx(s["coverage_90"], abs=1e-9)
    assert abs(s["coverage_90"] - 0.90) < 0.03
    assert abs(s["coverage_50"] - 0.50) < 0.04
    assert abs(s["coverage_80"] - 0.80) < 0.03
    assert abs(s["coverage_95"] - 0.95) < 0.03


def test_summarise_by_groups():
    truth = _truth(100)
    p = pd.concat([_preds(truth, "A", 0.1), _preds(truth, "B", 0.2, seed=3)])
    out = summarise_by(p, truth)
    assert set(out["arm"]) == {"A", "B"} and (out["n"] == 100).all()


# ----------------------------------------------------------------------------------------------- inference

def test_bootstrap_ci_covers_true_difference():
    """Simulation: arm B has |e| smaller by a known amount on average; CI should cover the population difference."""
    truth = _truth(600, seed=11)
    rng = np.random.default_rng(5)
    shared = rng.normal(0, 0.05, len(truth))
    a = _preds(truth, "A", 0.16, shared=shared, seed=21)
    b = _preds(truth, "B", 0.10, shared=shared, seed=22)
    # population mean |e| for N(0, s²+0.05²): s_tot·sqrt(2/π)
    true_diff = (math.hypot(0.10, 0.05) - math.hypot(0.16, 0.05)) * math.sqrt(2 / math.pi)
    c = compare_arms(pd.concat([a, b]), truth, [("A", "B")], B=400, seed=1)
    row = c.iloc[0]
    assert row["ci_lo"] < true_diff < row["ci_hi"]
    assert row["diff"] < 0 and row["p_boot"] < 0.05 and row["dm_p"] < 0.05
    assert row["n_blocks"] > 10


def test_bootstrap_no_difference_under_null():
    truth = _truth(300, seed=3)
    a = _preds(truth, "A", 0.15, seed=1)
    b = a.assign(arm="B")
    rec_a, rec_b = per_record(a, truth), per_record(b, truth)
    r = block_bootstrap_contrast(rec_a, rec_b, B=200, seed=0)
    assert r["diff"] == 0.0 and r["ci_lo"] <= 0 <= r["ci_hi"]


def test_holm_matches_statsmodels_and_is_monotone():
    from statsmodels.stats.multitest import multipletests
    p = [0.01, 0.04, 0.03, 0.20, 0.005]
    ours = holm(p)
    ref = multipletests(p, method="holm")[1]
    assert np.allclose(ours, ref)
    assert holm([0.5])[0] == 0.5
    assert np.isnan(holm([0.1, np.nan])[1]) and holm([0.1, np.nan])[0] == 0.1


def test_diebold_mariano_detects_difference_and_respects_order():
    rng = np.random.default_rng(0)
    la = np.abs(rng.normal(0, 0.2, 500))
    lb = np.abs(rng.normal(0, 0.1, 500))
    r = diebold_mariano(la, lb)
    assert r["dm_stat"] < 0 and r["dm_p"] < 0.01
    same = diebold_mariano(la, la)
    assert np.isnan(same["dm_stat"]) or abs(same["dm_mean_diff"]) < 1e-12


def test_compare_arms_requires_unambiguous_spec():
    truth = _truth(40)
    a = _preds(truth, "A", 0.1)
    a2 = a.assign(aggregation="mean")
    with pytest.raises(ValueError, match="ambiguous"):
        compare_arms(pd.concat([a, a2, _preds(truth, "B", 0.1)]), truth, [("A", "B")], B=20)
    c = compare_arms(pd.concat([a, a2, _preds(truth, "B", 0.1)]), truth, [({"arm": "A", "aggregation": "mean"}, "B")], B=20)
    assert len(c) == 1


# ----------------------------------------------------------------------------------------------- calibration

def test_coverage_by_suppresses_small_cells():
    truth = _truth(200)
    rec = per_record(_preds(truth, "D", 0.15), truth)
    rec["grp"] = ["big"] * 180 + ["tiny"] * 20
    cb = coverage_by(rec, ["grp"], min_n=30)
    assert cb.loc[cb["grp"] == "tiny", "suppressed"].all()
    assert cb.loc[cb["grp"] == "tiny", "coverage"].isna().all()
    assert cb.loc[cb["grp"] == "big", "coverage"].notna().all()


def test_reliability_and_pit_uniform_for_calibrated_forecast():
    truth = _truth(3000, seed=9)
    rec = per_record(_preds(truth, "D", 0.15, seed=4), truth)
    r = reliability(rec)
    q = r[r["kind"] == "quantile"]
    assert (q["gap"].abs() < 0.03).all()
    pit = pit_from_quantiles(rec)
    assert np.isfinite(pit).all() and ((pit > 0) & (pit < 1)).all()
    # Kolmogorov–Smirnov against U(0,1) should not reject at 1 %
    assert stats.kstest(pit, "uniform").pvalue > 0.01


def test_conformal_scale_achieves_nominal_on_validation_and_improves_test():
    truth_v, truth_t = _truth(1000, seed=31), _truth(1000, seed=32)
    # too-narrow intervals (half the true spread) → factor must be ≈ 2 and coverage ≈ 0.90 on validation
    pv = _preds(truth_v, "D", 0.15, seed=41, iscale=0.5)
    pt = _preds(truth_t, "D", 0.15, seed=42, iscale=0.5)
    sc = conformalise(pv, truth_v)
    assert 1.6 < sc.scale_90 < 2.5
    assert sc.achieved_valid_90 >= 0.90
    adj_v = apply_conformal(pv, sc)
    assert summarise(adj_v, truth_v)["coverage_90"] >= 0.90 - 1e-9
    before, after = summarise(pt, truth_t)["coverage_90"], summarise(adj_v := apply_conformal(pt, sc), truth_t)["coverage_90"]
    assert before < 0.75 and abs(after - 0.90) < 0.04
    # point and q50 untouched, ordering preserved
    assert np.allclose(adj_v["point"], pt["point"]) and np.allclose(adj_v["q50"], pt["q50"])
    assert (adj_v["q05"] <= adj_v["q25"]).all() and (adj_v["q75"] <= adj_v["q95"]).all()
    # too-wide intervals → factor < 1 (narrowing)
    sc2 = conformalise(_preds(truth_v, "D", 0.15, seed=43, iscale=2.0), truth_v)
    assert sc2.scale_90 < 0.7


# ----------------------------------------------------------------------------------------------- strata

def test_stratified_metrics_long_table_and_suppression():
    truth = _truth(400)
    st = stratified_metrics(_preds(truth, "A", 0.1), truth, strata=["region_group", "year_built_decade"], metrics=["mean_abs_log_error"])
    assert set(st.columns) >= {"arm", "stratum", "level", "n", "metric", "value", "suppressed"}
    assert set(st["stratum"]) == {"region_group", "year_built_decade"}
    small = st[st["n"] < 30]
    assert small["suppressed"].all() and small["value"].isna().all()
    big = st[st["n"] >= 30]
    assert (~big["suppressed"]).all() and big["value"].notna().all()
    assert set(st[st["stratum"] == "region_group"]["level"]) <= {"Montréal", "Québec", "Outaouais", "other"}


# ----------------------------------------------------------------------------------------------- faithfulness

def _packet(tx_id: str = "T1") -> EvidencePacket:
    d = date(2025, 6, 1)
    comps = [Comparable(evidence_id=f"C{k:02d}", sale_date=d - timedelta(days=60 * k), price=p, distance_km=0.5 * k, property_class="house",
                        municipality="Gatineau", months_before_valuation=2.0 * k) for k, p in ((1, 400_000.0), (2, 440_000.0), (3, 380_000.0))]
    subj = SubjectProperty(tx_id=tx_id, valuation_date=d, information_cutoff=d - timedelta(days=1), property_class="house", region="Outaouais",
                           municipality="Gatineau", lat=45.477, lng=-75.701)
    return EvidencePacket(packet_id="pk1", subject=subj, comparables=comps, market=MarketContext(as_of=d - timedelta(days=1)))


def _estimate(p: float, cited, adj, words: int = 50, assumptions=()) -> AgentEstimate:
    return AgentEstimate(point_estimate=p, q05=p * 0.85, q25=p * 0.95, q50=p, q75=p * 1.05, q95=p * 1.15, evidence_ids_used=cited,
                         adjustments=adj, key_assumptions=list(assumptions), reasoning_summary=" ".join(["w"] * words))


def test_audit_estimate_hand_built():
    pk = _packet()
    good = _estimate(410_000, ["P", "C01", "C02", "M1"], [Adjustment(comparable_id="C01", adjusted_price=410_000, reason="r"),
                                                         Adjustment(comparable_id="C02", adjusted_price=420_000, reason="r")])
    r = audit_estimate(good, pk)
    assert r["cited_share_valid"] == 1.0 and r["adj_share_within_40"] == 1.0 and r["point_in_hull"] is True
    assert r["quantile_sane"] is True and r["summary_words"] == 50 and r["summary_within_limit"]
    bad = _estimate(700_000, ["P", "C01", "C77"], [Adjustment(comparable_id="C01", adjusted_price=200_000, reason="r"),
                                                  Adjustment(comparable_id="C03", adjusted_price=390_000, reason="r")], words=150, assumptions=["unique view"])
    r = audit_estimate(bad, pk)
    assert r["cited_share_valid"] == pytest.approx(2 / 3) and r["n_cited_invalid"] == 1
    assert r["adj_share_within_40"] == pytest.approx(0.5)            # 200k is −50 % of 400k
    assert r["point_in_hull"] is False and r["hull_deviation_justified"] is True   # 700k > 390k × 1.25
    assert r["hull_log_deviation"] == pytest.approx(math.log(700_000 / (390_000 * 1.25)))
    assert not r["summary_within_limit"]
    # without a packet the id checks are NaN
    assert np.isnan(audit_estimate(good, None)["cited_share_valid"])


def test_audit_transcript_committee_dynamics(tmp_path: Path):
    pk = _packet()
    roles = ["r1", "r2", "r3", "r4"]
    r0 = {"r1": 380_000.0, "r2": 400_000.0, "r3": 420_000.0, "r4": 500_000.0}
    med0 = float(np.median(np.log(list(r0.values()))))
    turns = []
    for role in roles:
        turns.append(AgentTurn(role=role, round=0, kind="estimate", model="m", provider="mock", estimate=_estimate(r0[role], ["P", "C01"], [])))
    # round 1: everyone moves halfway (in log) toward the round-0 median → anchoring index 1, share toward median 1
    r1 = {role: float(np.exp(np.log(v) + 0.5 * (med0 - np.log(v)))) for role, v in r0.items()}
    for role in roles:
        turns.append(AgentTurn(role=role, round=1, kind="revision", model="m", provider="mock", estimate=_estimate(r1[role], ["P", "C01"], [])))
    tr = DebateTranscript(tx_id="T1", arm="G_debate_r1", committee="all_anthropic", protocol="structured_debate", rounds=1,
                          packet_hash=pk.content_hash(), turns=turns, final_estimates=r1)
    t, c = audit_transcript(tr, pk)
    assert len(t) == 8 and t["packet_hash_match"].all()
    c1 = c[c["round"] == 1].iloc[0]
    assert c1["anchoring_index"] == pytest.approx(1.0)
    assert c1["share_toward_median"] == 1.0
    assert c1["convergence"] == pytest.approx(0.5, abs=1e-6)
    # round-trip through files + audit_run + summary
    run = tmp_path / "run_x"
    (run / "transcripts").mkdir(parents=True)
    (run / "packets").mkdir()
    (run / "transcripts" / "T1.json").write_text(tr.to_json())
    (run / "packets" / "T1.json").write_text(pk.model_dump_json())
    turns_df, comm_df = audit_run(run)
    assert len(turns_df) == 8 and (turns_df["run_id"] == "run_x").all()
    summ = summarise_faithfulness(turns_df, comm_df)
    assert summ.iloc[0]["cited_share_valid"] == 1.0 and summ.iloc[0]["anchoring_index"] == pytest.approx(1.0)


# ----------------------------------------------------------------------------------------------- tables

def test_latex_fragments_have_no_trailing_backslashes(tmp_path: Path):
    truth = _truth(80)
    preds = pd.concat([_preds(truth, "A_hedonic", 0.1, intervals=False), _preds(truth, "G_debate_r2", 0.1, seed=7)])
    summ = summarise_by(preds, truth)
    tex = write_main_results(summ, tmp_path / "main.tex")
    text = (tmp_path / "main.tex").read_text()
    assert not text.rstrip().endswith("\\\\")
    assert text.rstrip().endswith(r"\end{tabular}")
    assert r"\toprule" in tex and r"\bottomrule" in tex and r"\num{" in tex
    assert "A (hedonic)" in tex and "G (debate r2)" in tex
    assert (tmp_path / "main.csv").exists()
    cov = write_coverage(summ, tmp_path / "coverage.tex")
    assert "{--}" in cov and not cov.rstrip().endswith("\\\\")   # arm without intervals prints --
    c = compare_arms(preds, truth, [("A_hedonic", "G_debate_r2")], B=30)
    write_contrasts(c, tmp_path / "contrasts.tex")
    assert not (tmp_path / "contrasts.tex").read_text().rstrip().endswith("\\\\")
    st = stratified_metrics(preds, truth, strata=["region_group"], metrics=["mean_abs_log_error"])
    write_strata(st, tmp_path / "strata.tex")
    s = (tmp_path / "strata.tex").read_text()
    assert not s.rstrip().endswith("\\\\") and "Region:" in s


def test_escape_and_num():
    assert escape("a_b & 50%") == r"a\_b \& 50\%"
    assert num(float("nan")) == "{--}" and num(0.12345, 2) == r"\num{0.12}" and num(None) == "{--}"
    assert arm_pretty("J_cost_matched") == "J (cost matched)" and arm_pretty("human") == "human"