#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Run-to-run variation of a hosted-LLM arm: compare a repeated run with the original run on the same
properties (same packets, same protocol, same models; only the providers' sampling differs).

Usage: python3 scripts/36_repeat_variation.py --orig final_G_debate_r2_mixed_random_mid --repeat repeat_G_debate_r2_mixed_random_mid
Output: results/evaluation/repeat_variation.json (+ .csv per property)
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import polars as pl

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
from agentic_valuation.config import PROCESSED, RESULTS  # noqa: E402


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--orig", default="final_G_debate_r2_mixed_random_mid")
    ap.add_argument("--repeat", default="_repeat_G_debate_r2_mixed_random_mid")
    ap.add_argument("--aggregation", default="median")
    a = ap.parse_args(argv)
    o = pl.read_parquet(RESULTS / "runs" / a.orig / "predictions.parquet").filter(pl.col("aggregation") == a.aggregation)
    r = pl.read_parquet(RESULTS / "runs" / a.repeat / "predictions.parquet").filter(pl.col("aggregation") == a.aggregation)
    truth = pl.read_parquet(PROCESSED / "analytic.parquet", columns=["tx_id", "price"])
    j = o.join(r, on="tx_id", suffix="_rep").join(truth, on="tx_id")
    if j.height == 0:
        sys.exit("no common properties")
    d = (j["point"].log() - j["point_rep"].log()).to_numpy()
    e1 = (j["point"].log() - j["price"].log()).abs().to_numpy()
    e2 = (j["point_rep"].log() - j["price"].log()).abs().to_numpy()
    w1 = (j["q95"].log() - j["q05"].log()).to_numpy()
    w2 = (j["q95_rep"].log() - j["q05_rep"].log()).to_numpy()
    out = {
        "orig": a.orig, "repeat": a.repeat, "aggregation": a.aggregation, "n": int(j.height),
        "mean_abs_log_diff_points": float(np.mean(np.abs(d))), "median_abs_log_diff_points": float(np.median(np.abs(d))), "max_abs_log_diff_points": float(np.max(np.abs(d))),
        "sd_log_diff": float(np.std(d, ddof=1)) if j.height > 1 else None,
        "mean_abs_log_error_orig": float(e1.mean()), "mean_abs_log_error_repeat": float(e2.mean()),
        "corr_abs_errors": float(np.corrcoef(e1, e2)[0, 1]) if j.height > 2 else None,
        "mean_log_width90_orig": float(w1.mean()), "mean_log_width90_repeat": float(w2.mean()),
        "share_repeat_inside_orig_90": float(((j["point_rep"] >= j["q05"]) & (j["point_rep"] <= j["q95"])).mean()),
    }
    (RESULTS / "evaluation").mkdir(parents=True, exist_ok=True)
    (RESULTS / "evaluation" / "repeat_variation.json").write_text(json.dumps(out, indent=1))
    j.select(["tx_id", "price", "point", "point_rep", "q05", "q95", "q05_rep", "q95_rep"]).write_csv(RESULTS / "evaluation" / "repeat_variation.csv")
    print(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()