"""Claim 5 audit: recompute the reported average rank of 1.24 from Table 1 itself.

Table 1 is parsed straight out of the arXiv HTML text dump (paper/paper.txt) so the
numbers are not re-typed by hand.  We then rank the 12 methods on every
(dataset, metric) cell and report the mean rank of "Ours" under several readings of
the claim, since the paper is ambiguous about which metrics enter the average
(Sec 4.1 says "across all metrics", the challenge statement says accuracy/F1/AUROC).
"""

import json
import re

import numpy as np

METHODS = ["Autoformer", "Crossformer", "FEDformer", "Informer", "iTransformer",
           "MTST", "Nonformer", "PatchTST", "Reformer", "Transformer",
           "Medformer", "Ours"]
DATASETS = ["ADFTD", "PTB", "PTB-XL", "APAVA", "SleepEDF", "FLAAP", "UCI-HAR"]
METRICS = ["Accuracy", "Precision", "Recall", "F1-score", "AUROC"]


def parse(path="paper/paper.txt"):
    txt = open(path, encoding="utf-8").read()
    # Table 1 spans from the header block to the "Table 1 :" caption
    start = txt.index(" Dataset \n Metrics \n Autoformer")
    end = txt.index(" Table 1 : ")
    body = txt[start:end]
    # every "value ± std" pair, in reading order
    vals = [float(m) for m in re.findall(r"(\d+\.\d+)\s*±", body)]
    assert len(vals) == 7 * 5 * 12, f"expected 420 numbers, parsed {len(vals)}"
    arr = np.array(vals).reshape(7, 5, 12)
    return arr


def ranks_for_ours(arr, metric_idx):
    """Rank (1 = best, higher value is better for every metric here) of 'Ours'."""
    out = {}
    for di, ds in enumerate(DATASETS):
        for mi in metric_idx:
            col = arr[di, mi]
            # competition ranking with ties sharing the average rank
            order = (-col).argsort()
            r = np.empty(12)
            r[order] = np.arange(1, 13)
            # handle exact ties
            for v in np.unique(col):
                m = col == v
                if m.sum() > 1:
                    r[m] = r[m].mean()
            out[(ds, METRICS[mi])] = float(r[METHODS.index("Ours")])
    return out


def main():
    arr = parse()
    report = {"table_shape": list(arr.shape)}
    # spot-check a few known cells against the paper text
    checks = {
        ("ADFTD", "Accuracy", "Ours"): 53.91,
        ("ADFTD", "F1-score", "Ours"): 51.79,
        ("ADFTD", "F1-score", "Medformer"): 50.65,
        ("PTB-XL", "F1-score", "Ours"): 63.51,
        ("UCI-HAR", "Accuracy", "Medformer"): 91.65,
        ("SleepEDF", "AUROC", "Ours"): 95.37,
    }
    ok = {}
    for (ds, me, mo), want in checks.items():
        got = arr[DATASETS.index(ds), METRICS.index(me), METHODS.index(mo)]
        ok[f"{ds}/{me}/{mo}"] = {"parsed": float(got), "paper": want,
                                 "match": bool(abs(got - want) < 1e-9)}
    report["parse_spot_checks"] = ok

    variants = {
        "all_5_metrics": [0, 1, 2, 3, 4],
        "acc_f1_auroc (challenge wording)": [0, 3, 4],
        "acc_f1": [0, 3],
    }
    for name, idx in variants.items():
        r = ranks_for_ours(arr, idx)
        vals = np.array(list(r.values()))
        report[name] = {
            "n_cells": int(vals.size),
            "avg_rank_ours": float(vals.mean()),
            "n_rank1": int((vals == 1).sum()),
            "worst_rank": float(vals.max()),
            "cells_not_rank1": {f"{k[0]}/{k[1]}": v for k, v in r.items() if v != 1},
        }

    # average rank of every method under the 5-metric reading, for context
    allranks = {}
    for mo in METHODS:
        rs = []
        for di in range(7):
            for mi in range(5):
                col = arr[di, mi]
                order = (-col).argsort()
                rr = np.empty(12)
                rr[order] = np.arange(1, 13)
                for v in np.unique(col):
                    m = col == v
                    if m.sum() > 1:
                        rr[m] = rr[m].mean()
                rs.append(rr[METHODS.index(mo)])
        allranks[mo] = float(np.mean(rs))
    report["avg_rank_all_methods_5metrics"] = allranks

    # one-tailed Wilcoxon signed-rank, Ours vs Medformer (paper's stated test)
    from scipy.stats import wilcoxon
    ours = arr[:, :, METHODS.index("Ours")].ravel()
    med = arr[:, :, METHODS.index("Medformer")].ravel()
    st, p = wilcoxon(ours, med, alternative="greater")
    report["wilcoxon_ours_gt_medformer"] = {"statistic": float(st), "p_value": float(p),
                                            "n_pairs": int(ours.size),
                                            "ours_wins": int((ours > med).sum())}
    print(json.dumps(report, indent=2))
    with open("results/table1_rank_audit.json", "w") as f:
        json.dump(report, f, indent=2)


if __name__ == "__main__":
    main()