#!/usr/bin/env python3
from __future__ import annotations

import argparse
import re
from collections import Counter
from pathlib import Path

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

ROOT = Path(__file__).resolve().parents[1]


_TOK_RE = re.compile(
    r"""
    (A\s*\[\]|E\s*<>|A\s*<>|E\s*\[\])
    |(-->)
    |(&&|\|\||<=|>=|==|!=)
    |([(){}\[\],])
    |(\bforall\b|\bexists\b|\bimply\b|\bnot\b|\bdeadlock\b)
    |([A-Za-z_]\w*)
    |(\d+)
    """,
    re.VERBOSE | re.IGNORECASE,
)


def _norm_ws(s: str) -> str:
    return re.sub(r"\s+", " ", (s or "").strip())


def tokenize_tctl(s: str) -> list[str]:
    s = _norm_ws(s)
    return [t for tup in _TOK_RE.findall(s) for t in tup if t]


def token_f1_pair(gen: str, gold: str) -> float:
    cg = Counter(tokenize_tctl(gen))
    ch = Counter(tokenize_tctl(gold))
    inter = sum((cg & ch).values())
    tot_g = sum(cg.values())
    tot_h = sum(ch.values())
    if tot_g == 0 and tot_h == 0:
        return 1.0
    p = inter / tot_g if tot_g else 0.0
    r = inter / tot_h if tot_h else 0.0
    return 0.0 if (p + r) == 0 else (2 * p * r / (p + r))


def rm_anova_oneway(wide: pd.DataFrame, subject_col: str, condition_cols: list[str]) -> dict[str, float]:
    """
    One-way repeated-measures ANOVA (k conditions, n subjects).
    Assumes balanced data with one observation per subject x condition.
    """
    y = wide[condition_cols].to_numpy(dtype=float)
    n, k = y.shape
    grand = float(np.mean(y))
    cond_means = np.mean(y, axis=0)
    subj_means = np.mean(y, axis=1)

    ss_total = float(np.sum((y - grand) ** 2))
    ss_conditions = float(n * np.sum((cond_means - grand) ** 2))
    ss_subjects = float(k * np.sum((subj_means - grand) ** 2))
    ss_error = ss_total - ss_conditions - ss_subjects

    df_conditions = k - 1
    df_subjects = n - 1
    df_error = df_conditions * df_subjects

    ms_conditions = ss_conditions / df_conditions
    ms_error = ss_error / df_error
    f = ms_conditions / ms_error
    p = float(stats.f.sf(f, df_conditions, df_error))
    partial_eta2 = ss_conditions / (ss_conditions + ss_error) if (ss_conditions + ss_error) else float("nan")

    return {
        "n_subjects": n,
        "k_conditions": k,
        "SS_conditions": ss_conditions,
        "SS_error": ss_error,
        "df_conditions": df_conditions,
        "df_error": df_error,
        "F": float(f),
        "p": p,
        "partial_eta2": float(partial_eta2),
    }


def holm_adjust(pvals: list[float]) -> list[float]:
    m = len(pvals)
    order = np.argsort(pvals)
    adj = np.empty(m, dtype=float)
    for rank, idx in enumerate(order):
        adj[idx] = (m - rank) * pvals[idx]
    # enforce monotonicity in sorted order
    for i in range(1, m):
        prev_idx = order[i - 1]
        cur_idx = order[i]
        adj[cur_idx] = max(adj[cur_idx], adj[prev_idx])
    return [float(min(1.0, x)) for x in adj]


def main() -> int:
    ap = argparse.ArgumentParser(description="Repeated-measures ANOVA on per-query Token-F1 (RQ2).")
    ap.add_argument(
        "--sheet",
        type=Path,
        default=ROOT / "artifacts" / "rq2_query_translation_sheet.tsv",
        help="TSV produced by scripts/rq2_query_translation_eval.py",
    )
    args = ap.parse_args()

    df = pd.read_csv(args.sheet, sep="\t")
    # Use adapted gold so identifiers are in the same naming space as predictions.
    gold = df["ground_query_adapted_ref_xml"].fillna("").map(str).tolist()

    methods = {
        "ours": df["pred_ours"].fillna("").map(str).tolist(),
        "gpt": df["pred_gpt"].fillna("").map(str).tolist(),
        "grok": df["pred_grok"].fillna("").map(str).tolist(),
        "claude": df["pred_claude"].fillna("").map(str).tolist(),
    }

    rows = []
    for j in range(len(df)):
        for m, preds in methods.items():
            rows.append(
                {
                    "query_id": j,
                    "method": m,
                    "token_f1": token_f1_pair(preds[j], gold[j]),
                }
            )
    long = pd.DataFrame(rows)
    wide = long.pivot(index="query_id", columns="method", values="token_f1").reset_index()
    conds = ["ours", "gpt", "grok", "claude"]
    if any(c not in wide.columns for c in conds):
        missing = [c for c in conds if c not in wide.columns]
        raise SystemExit(f"Missing columns in pivot: {missing}")
    if wide[conds].isna().any().any():
        raise SystemExit("Found NaNs; expected full matrix for repeated-measures ANOVA.")

    # --- RM-ANOVA ---
    res = rm_anova_oneway(wide, subject_col="query_id", condition_cols=conds)
    means = {c: float(wide[c].mean()) for c in conds}
    sds = {c: float(wide[c].std(ddof=1)) for c in conds}

    print("RQ2 Token-F1 per query (adapted gold) — repeated-measures ANOVA (one-way)")
    print(f"n queries: {int(res['n_subjects'])}, k methods: {int(res['k_conditions'])}")
    print("\nDescriptives (mean ± sd):")
    for c in conds:
        print(f"  {c:6}  {means[c]:.3f} ± {sds[c]:.3f}")
    print("\nRM-ANOVA:")
    print(f"  F({int(res['df_conditions'])}, {int(res['df_error'])}) = {res['F']:.4f}")
    print(f"  p = {res['p']:.6g}")
    print(f"  partial eta^2 = {res['partial_eta2']:.4f}")

    # --- Post-hoc paired comparisons (t-test) ---
    pairs = [("ours", "gpt"), ("ours", "grok"), ("ours", "claude"), ("claude", "gpt"), ("claude", "grok"), ("gpt", "grok")]
    pvals = []
    stats_out = []
    for a, b in pairs:
        t, p = stats.ttest_rel(wide[a].to_numpy(), wide[b].to_numpy())
        pvals.append(float(p))
        stats_out.append((a, b, float(t), float(p), float((wide[a] - wide[b]).mean())))
    p_holm = holm_adjust(pvals)

    print("\nPost-hoc paired t-tests (Holm-adjusted p-values):")
    for (a, b, t, p, dmean), padj in zip(stats_out, p_holm):
        print(f"  {a:6} vs {b:6}  mean_diff={dmean:+.4f}  t={t:+.4f}  p={p:.6g}  p_holm={padj:.6g}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())