import hashlib
import json
import random
import statistics
import urllib.request
from pathlib import Path
import pyarrow.parquet as pq

OUT = Path("benchmark-recalculation")
OUT.mkdir(exist_ok=True)
GIT = "40f164d5b8f1d249bf95a6df8b74b577fd8e519d"
HF = "c104f840cc67f8b6eec6f759ebc8b2693d585d4a"
AA = "1f28d51b6e191e9f2bfdc22b09d92eecdab6654e"
SUBS = {
    "A": "20241029_OpenHands-CodeAct-2.1-sonnet-20241022",
    "B": "20241202_agentless-1.5_claude-3.5-sonnet-20241022",
}

def fetch(url, name, expected=None):
    with urllib.request.urlopen(url, timeout=60) as r:
        data = r.read()
    digest = hashlib.sha256(data).hexdigest()
    if expected is not None:
        assert digest == expected, (name, digest)
    p = OUT / name
    p.write_bytes(data)
    return p

manifest = fetch(
    f"https://huggingface.co/datasets/princeton-nlp/"
    f"SWE-bench_Verified/resolve/{HF}/data/test-00000-of-00001.parquet",
    "verified.parquet",
    "a45b1fe4e2f0c8390b2b2938ac83e92ed5979000856808f3679c07812e9e6dcd",
)
rows = pq.read_table(manifest).to_pylist()
ids = {r["instance_id"] for r in rows}
assert len(rows) == len(ids) == 500
resolved, by_repo = {}, {}
for label, sub in SUBS.items():
    base = f"https://raw.githubusercontent.com/SWE-bench/experiments/{GIT}"
    base += f"/evaluation/verified/{sub}/results/"
    result = json.loads(fetch(base + "results.json", label + ".json").read_text())
    repo = json.loads(fetch(base + "resolved_by_repo.json",
                            label + "-repo.json").read_text())
    success = set(result["resolved"])
    assert len(success) == len(result["resolved"])
    assert success <= ids
    assert sum(x["total"] for x in repo.values()) == 500
    for name, counts in repo.items():
        subset = [r for r in rows if r["repo"] == name]
        assert len(subset) == counts["total"]
        assert sum(r["instance_id"] in success for r in subset) == counts["resolved"]
    resolved[label], by_repo[label] = success, repo
    micro = len(success) / len(ids)
    macro = statistics.mean(x["resolved"] / x["total"] for x in repo.values())
    print(label, "resolved", len(success), "micro", round(100 * micro, 2),
          "macro", round(100 * macro, 2),
          "no_generation", len(result["no_generation"]),
          "no_logs", len(result["no_logs"]))

a, b = resolved["A"], resolved["B"]
print("both, A_only, B_only, neither:",
      len(a & b), len(a - b), len(b - a), len(ids - (a | b)))
d = [int(i in a) - int(i in b) for i in sorted(ids)]
rng = random.Random(20261023)
boot = sorted(statistics.mean(rng.choices(d, k=len(d))) for _ in range(10000))
print("paired difference pp:", round(100 * statistics.mean(d), 4))
print("naive percentile CI pp:", 100 * boot[249], 100 * boot[9749])

from collections import Counter
qpath = fetch(
    f"https://huggingface.co/datasets/ArtificialAnalysis/"
    f"SWE-Atlas-QnA-AA/resolve/{AA}/data/test-00000-of-00001.parquet",
    "qna.parquet",
    "c78b998726e9e7762807050adf79f18406f14dbf61a6fb87cc7dc0c0858f5d74",
)
qna = pq.read_table(qpath).to_pylist()
rubrics = [json.loads(r["rubric"]) for r in qna]
print("QnA tasks, criteria:", len(qna), sum(map(len, rubrics)))
for column in ("language", "category", "repository_url"):
    print(column, dict(Counter(r[column] for r in qna)))
print("changes", dict(Counter(c.get("change") for rs in rubrics for c in rs)))