"""Prove the submission format end-to-end BEFORE the real run finishes.

Builds matching_results.tsv / candidate_pairs.tsv for the FULL test set using the
exact write_tsv() from run_test.py, with realistic content (mixed singletons and
multi-matches, S2 and S3 ids, the observed 5.58% singleton rate), then runs the
organisers' own validate_submission.py on them.

If this prints PASS, the format path is proven: the real run uses the same writer,
the same id source and the same header.
"""
import os, sys, subprocess, random
import numpy as np
import pyarrow.parquet as pq

sys.path.insert(0, "src")
from run_test import write_tsv               # the EXACT writer the pipeline uses

D = "/scratch/user4/dai/dataset/dataset"
OUT = "/scratch/user4/dai/model-3/work/format_check"
os.makedirs(OUT, exist_ok=True)

print("loading test ids ...")
s1 = pq.read_table(f"{D}/hf/test_p/test_source1.parquet", columns=["entity_id"]) \
       .column("entity_id").to_pylist()
pool = []
for s in (2, 3):
    pool += pq.read_table(f"{D}/hf/test_p/test_source{s}.parquet", columns=["entity_id"]) \
              .column("entity_id").to_pylist()
print(f"test S1 {len(s1):,}  pool {len(pool):,}")

rng = random.Random(42)
# mirror the training distribution: 5.58% singletons, mean 3.46 matches, max 11
sizes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
weights = [5.58, 5.40, 17.00, 24.05, 21.94, 14.59, 7.47, 2.90, 0.85, 0.19, 0.02, 0.01]

match_rows, cand_rows = [], []
used = set()                     # enforce one-owner-per-record, as the pipeline does
ptr = 0
for sid in s1:
    k = rng.choices(sizes, weights)[0]
    cands = []
    for _ in range(k + 3):       # candidates are a superset of matches
        if ptr >= len(pool):
            break
        c = pool[ptr]; ptr += 1
        if c not in used:
            cands.append(c)
    matched = cands[:k]
    used.update(matched)
    match_rows.append((sid, matched))
    cand_rows.append((sid, cands))

write_tsv(os.path.join(OUT, "matching_results.tsv"), match_rows, "matched_entity_ids")
write_tsv(os.path.join(OUT, "candidate_pairs.tsv"), cand_rows, "candidate_entity_ids")

print("\n--- first 3 lines of matching_results.tsv (tabs shown as <TAB>) ---")
with open(os.path.join(OUT, "matching_results.tsv"), encoding="utf-8") as f:
    for _ in range(3):
        print("   " + f.readline().rstrip("\n").replace("\t", "<TAB>"))

nz = sum(1 for _, m in match_rows if m)
print(f"\nrows: {len(match_rows):,} | with matches: {nz:,} | "
      f"singletons: {len(match_rows)-nz:,} ({1-nz/len(match_rows):.4f})")

print("\n=== running the organisers' validator ===", flush=True)
r = subprocess.run([sys.executable, f"{D}/../utils/validate_submission.py",
                    "--matching", os.path.join(OUT, "matching_results.tsv"),
                    "--candidate", os.path.join(OUT, "candidate_pairs.tsv"),
                    "--test-dir", f"{D}/test"],
                   capture_output=True, text=True)
print(r.stdout[-3000:])
if r.stderr.strip():
    print("STDERR:", r.stderr[-1500:])
print(f"exit code: {r.returncode}   (0 = safe to submit)")