import json, statistics, os
from collections import Counter
from datasets import load_dataset

OUT = os.environ.get("OUT_DIR", ".")
os.makedirs(OUT, exist_ok=True)

ds = load_dataset("ScaleAI/SWE-bench_Pro", split="test")
n = len(ds)
repos = sorted(set(ds["repo"]))
repo_counts = dict(Counter(ds["repo"]))

# Claim 1: public subset size
# Claim 2: public set = 11 repos (held-out 12 + commercial 18 not public)
# Claim 3: multi-file long-horizon patches
patch_files = []
for ex in ds:
    p = ex["patch"] or ""
    files = set()
    for line in p.splitlines():
        if line.startswith("diff --git"):
            files.add(line)
    patch_files.append(len(files))
multi = sum(1 for x in patch_files if x > 1)

# Claim 4: human-verified context adequacy
ps_lens = [len(ex["problem_statement"] or "") for ex in ds]
has_tp = sum(1 for ex in ds if ex["test_patch"] and len(ex["test_patch"].strip()) > 0)
has_req = sum(1 for ex in ds if ex["requirements"] and len(str(ex["requirements"]).strip()) > 0)
has_iface = sum(1 for ex in ds if ex["interface"] and len(str(ex["interface"]).strip()) > 0)
has_dh = sum(1 for ex in ds if ex["dockerhub_tag"])

# Claim 5: domain / language diversity + containerized env (contamination-resistant)
langs = sorted(set(ds["repo_language"]))

result = {
    "public_instances": n,
    "public_repos": len(repos),
    "repo_counts": repo_counts,
    "languages": langs,
    "patch_files_min": min(patch_files),
    "patch_files_max": max(patch_files),
    "patch_files_mean": round(statistics.mean(patch_files), 2),
    "patch_files_median": statistics.median(patch_files),
    "instances_multi_file": multi,
    "instances_multi_file_pct": round(100 * multi / n, 1),
    "ps_chars_mean": round(statistics.mean(ps_lens)),
    "ps_all_have_context": all(x > 200 for x in ps_lens),
    "has_test_patch": has_tp,
    "has_requirements": has_req,
    "has_interface": has_iface,
    "has_dockerhub_tag": has_dh,
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(result, f, indent=2)
print(json.dumps(result, indent=2))