#!/usr/bin/env python3
"""Join tree tip order + divergence + taxonomy into one JSON payload for the Moat plot."""
import csv, json, re, os
import os
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

D = os.path.join(ROOT, "data")
REF = "MG008164"  # the sihek sequence all distances are measured from

tips = re.findall(r'[(,]([A-Za-z][A-Za-z0-9_.\-]*):', open(f"{D}/nd2_ml.treefile").read())
assert len(tips) == 156, len(tips)

taxa = {r["label"]: r for r in csv.DictReader(open(f"{D}/taxa_accessions.csv"))}
missing = [t for t in tips if t not in taxa]
assert not missing, f"tips absent from taxa_accessions: {missing[:5]}"

dist = {r["accession"]: r for r in csv.DictReader(open(f"{D}/divergence_from_sihek.csv"))}

recs = []
for i, label in enumerate(tips):
    t = taxa[label]
    acc = t["accession"]
    if acc == REF:
        d, sites = 0.0, 1041
    else:
        row = dist.get(acc)
        assert row, f"no distance for {acc} ({label})"
        d, sites = float(row["uncorrected_p_distance_pct"]), int(row["comparable_sites"])
    recs.append({
        "i": i,                       # angular position = ML tree tip order
        "sp": t["species"],
        "acc": acc,
        "sub": t["subfamily"],
        "d": round(d, 4),             # radial position = ND2 p-distance from the sihek
        "n": sites,
        "sihek": "cinnamominus" in label,
        "gen": t["species"].split()[0],
    })

# Empty annuli: sort distances, find runs with no sequence in them.
ds = sorted(r["d"] for r in recs)
gaps = [(ds[k], ds[k + 1]) for k in range(len(ds) - 1) if ds[k + 1] - ds[k] >= 0.9]
payload = {
    "ref": REF,
    "n": len(recs),
    "maxd": max(ds),
    "gaps": [{"lo": round(a, 4), "hi": round(b, 4), "w": round(b - a, 3)} for a, b in gaps],
    "pts": recs,
}

out = os.path.join(ROOT, "figures", "moat_data.json")
json.dump(payload, open(out, "w"), separators=(",", ":"))
print("wrote", out, os.path.getsize(out), "bytes")
print("points:", len(recs), " max distance:", round(max(ds), 3))
print("subfamily counts:", {s: sum(1 for r in recs if r["sub"] == s) for s in {r["sub"] for r in recs}})
print("\nempty annuli (>=0.9 pp wide):")
for g in payload["gaps"]:
    print(f"  {g['lo']:7.3f} -> {g['hi']:7.3f}   width {g['w']:6.3f} pp")