# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Faithfulness audit of debate transcripts (docs/architecture.md §3, docs/evaluation_protocol.md §7).

Inputs: `results/runs/<run_id>/transcripts/<tx_id>.json` (DebateTranscript) and `results/runs/<run_id>/packets/<tx_id>.json`
(EvidencePacket). Every check is deterministic and operates on the stored artefacts only.

Turn level (one row per estimate/revision turn):
    cited_share_valid      share of `evidence_ids_used` that exist in the packet (NaN if nothing cited)
    adj_share_within_40    share of adjustments whose adjusted price is within ±40 % of the raw comparable price
    point_in_hull          point inside [min adj, max adj] × [0.75, 1.25] (NaN if no adjustments)
    quantile_sane          quantiles ordered, positive, point within [q05, q95], q95/q05 within [1.01, 4]
    summary_words          word count of `reasoning_summary` (contract: ≤ 120)
Committee level (one row per transcript × round):
    iqr_log, convergence (iqr_r / iqr_0), anchoring_index (corr across roles of revision r−0 with the gap to the
    round-0 median), share_toward_median.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

from ..schemas import AgentEstimate, DebateTranscript, EvidencePacket

ADJ_BAND = 0.40
HULL_BAND = 0.25
MAX_SUMMARY_WORDS = 120


def _word_count(text: str) -> int:
    return len([w for w in (text or "").split() if w.strip()])


def quantile_sanity(est: AgentEstimate) -> bool:
    qs = [est.q05, est.q25, est.q50, est.q75, est.q95]
    if any(q <= 0 for q in qs) or est.point_estimate <= 0:
        return False
    if any(b < a for a, b in zip(qs, qs[1:])):
        return False
    if not (est.q05 <= est.point_estimate <= est.q95):
        return False
    ratio = est.q95 / est.q05
    return 1.01 <= ratio <= 4.0


def audit_estimate(est: AgentEstimate, packet: EvidencePacket | None) -> dict[str, Any]:
    """Deterministic checks of one AgentEstimate against its packet (packet may be None → id checks are NaN)."""
    ids = packet.evidence_ids() if packet is not None else None
    comp_price = {c.evidence_id: c.price for c in packet.comparables} if packet is not None else {}
    cited = list(est.evidence_ids_used)
    cited_valid = float(np.mean([c in ids for c in cited])) if (cited and ids is not None) else np.nan
    n_invalid = int(sum(c not in ids for c in cited)) if ids is not None else 0
    adj = est.adjustments
    within, adj_prices, n_adj_unknown = [], [], 0
    for a in adj:
        raw = comp_price.get(a.comparable_id)
        if raw is None or raw <= 0:
            n_adj_unknown += 1
            continue
        within.append(abs(a.adjusted_price / raw - 1.0) <= ADJ_BAND)
        adj_prices.append(a.adjusted_price)
    adj_share = float(np.mean(within)) if within else np.nan
    if adj_prices:
        lo, hi = min(adj_prices) * (1 - HULL_BAND), max(adj_prices) * (1 + HULL_BAND)
        in_hull = bool(lo <= est.point_estimate <= hi)
        hull_dev = 0.0 if in_hull else float(min(abs(np.log(est.point_estimate / lo)), abs(np.log(est.point_estimate / hi))))
        justified = (not in_hull) and len(est.key_assumptions) > 0
    else:
        in_hull, hull_dev, justified = np.nan, np.nan, np.nan
    return {
        "n_cited": len(cited), "n_cited_invalid": n_invalid, "cited_share_valid": cited_valid,
        "n_adjustments": len(adj), "n_adjustments_unknown_comp": n_adj_unknown, "adj_share_within_40": adj_share,
        "point_in_hull": in_hull, "hull_log_deviation": hull_dev, "hull_deviation_justified": justified,
        "quantile_sane": quantile_sanity(est), "summary_words": _word_count(est.reasoning_summary),
        "summary_within_limit": _word_count(est.reasoning_summary) <= MAX_SUMMARY_WORDS,
        "n_missing_data": len(est.missing_data), "n_counterarguments": len(est.counterarguments),
    }


def audit_transcript(tr: DebateTranscript, packet: EvidencePacket | None) -> tuple[pd.DataFrame, pd.DataFrame]:
    """(turn-level table, committee-level table) for one transcript."""
    turn_rows = []
    for t in tr.turns:
        if t.estimate is None:
            continue
        row = {"tx_id": tr.tx_id, "arm": tr.arm, "committee": tr.committee, "protocol": tr.protocol, "role": t.role,
               "round": t.round, "kind": t.kind, "model": t.model, "provider": t.provider,
               "validation_retries": t.validation_retries, "error": t.error is not None}
        row.update(audit_estimate(t.estimate, packet))
        row["packet_hash_match"] = (packet.content_hash() == tr.packet_hash) if packet is not None else np.nan
        turn_rows.append(row)
    turns = pd.DataFrame(turn_rows)

    # committee dynamics: last estimate of each role per round
    est: dict[int, dict[str, float]] = {}
    for t in tr.turns:
        if t.estimate is not None:
            est.setdefault(t.round, {})[t.role] = float(np.log(t.estimate.point_estimate))
    comm_rows = []
    if est:
        r0 = est[min(est)]
        med0 = float(np.median(list(r0.values())))
        iqr0 = float(np.subtract(*np.percentile(list(r0.values()), [75, 25]))) if len(r0) >= 2 else np.nan
        for r in sorted(est):
            vals = est[r]
            v = np.array(list(vals.values()))
            iqr = float(np.subtract(*np.percentile(v, [75, 25]))) if len(v) >= 2 else np.nan
            roles = [k for k in vals if k in r0]
            if r > min(est) and len(roles) >= 3:
                rev = np.array([vals[k] - r0[k] for k in roles])
                gap = np.array([med0 - r0[k] for k in roles])
                anchoring = float(np.corrcoef(rev, gap)[0, 1]) if rev.std() > 0 and gap.std() > 0 else np.nan
                toward = float(np.mean([abs(vals[k] - med0) < abs(r0[k] - med0) for k in roles]))
                mean_abs_rev = float(np.mean(np.abs(rev)))
            else:
                anchoring, toward, mean_abs_rev = np.nan, np.nan, (0.0 if r == min(est) else np.nan)
            comm_rows.append({"tx_id": tr.tx_id, "arm": tr.arm, "committee": tr.committee, "protocol": tr.protocol, "round": r,
                              "n_roles": len(v), "median_log": float(np.median(v)), "iqr_log": iqr,
                              "convergence": (iqr / iqr0) if (iqr0 and np.isfinite(iqr0) and iqr0 > 0) else np.nan,
                              "anchoring_index": anchoring, "share_toward_median": toward, "mean_abs_revision": mean_abs_rev,
                              "total_cost_usd": tr.total_cost_usd, "total_tokens": tr.total_tokens})
    return turns, pd.DataFrame(comm_rows)


def load_packet(path: Path) -> EvidencePacket | None:
    if not path.exists():
        return None
    return EvidencePacket.model_validate_json(path.read_text(encoding="utf-8"))


def audit_run(run_dir: Path) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Audit every transcript of a run directory. Returns (turn table, committee table), both with run_id."""
    run_dir = Path(run_dir)
    tdir, pdir = run_dir / "transcripts", run_dir / "packets"
    turns, comms = [], []
    if not tdir.exists():
        return pd.DataFrame(), pd.DataFrame()
    for f in sorted(tdir.glob("*.json")):
        tr = DebateTranscript.model_validate_json(f.read_text(encoding="utf-8"))
        packet = load_packet(pdir / f"{tr.tx_id}.json")
        t, c = audit_transcript(tr, packet)
        turns.append(t)
        comms.append(c)
    turns_df = pd.concat([t for t in turns if not t.empty], ignore_index=True) if any(not t.empty for t in turns) else pd.DataFrame()
    comm_df = pd.concat([c for c in comms if not c.empty], ignore_index=True) if any(not c.empty for c in comms) else pd.DataFrame()
    for df in (turns_df, comm_df):
        if not df.empty:
            df.insert(0, "run_id", run_dir.name)
    return turns_df, comm_df


def summarise_faithfulness(turns: pd.DataFrame, comms: pd.DataFrame, by: list[str] | None = None) -> pd.DataFrame:
    """Aggregate the turn and committee tables into one row per (run_id, arm, committee[, round])."""
    by = by or ["run_id", "arm", "committee"]
    if turns.empty:
        return pd.DataFrame()
    agg = turns.groupby(by, dropna=False).agg(
        n_turns=("tx_id", "size"), n_properties=("tx_id", "nunique"),
        cited_share_valid=("cited_share_valid", "mean"), share_turns_all_ids_valid=("cited_share_valid", lambda s: float(np.nanmean(s >= 1.0)) if s.notna().any() else np.nan),
        adj_share_within_40=("adj_share_within_40", "mean"), share_turns_adj_ok=("adj_share_within_40", lambda s: float(np.nanmean(s >= 1.0)) if s.notna().any() else np.nan),
        point_in_hull=("point_in_hull", "mean"), hull_deviation_justified=("hull_deviation_justified", "mean"),
        quantile_sane=("quantile_sane", "mean"), summary_words=("summary_words", "mean"), summary_within_limit=("summary_within_limit", "mean"),
        validation_retries=("validation_retries", "mean"), packet_hash_match=("packet_hash_match", "mean"),
    ).reset_index()
    if not comms.empty:
        last = comms.sort_values("round").groupby(["run_id", "tx_id"], as_index=False).tail(1)
        c = last.groupby([b for b in by if b in last.columns], dropna=False).agg(
            iqr_log_final=("iqr_log", "mean"), convergence=("convergence", "mean"),
            anchoring_index=("anchoring_index", "mean"), share_toward_median=("share_toward_median", "mean"),
        ).reset_index()
        first = comms[comms["round"] == comms.groupby(["run_id", "tx_id"])["round"].transform("min")]
        c0 = first.groupby([b for b in by if b in first.columns], dropna=False).agg(iqr_log_round0=("iqr_log", "mean")).reset_index()
        agg = agg.merge(c0, how="left").merge(c, how="left")
    return agg