#!/usr/bin/env python3
"""Stage B, deterministic parts of the event chain and the visit-level change annotation.

  event-chain RECORDS     extraction records -> event_chain.jsonl
  windows                 change-annotation windows for the tags step -> _tag_manifest.json
  tags-post TAGS          change tags of the tags step -> tags.jsonl
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import re
import sys
from collections import Counter

from pipeline.common import BASE, FIELDS, patient_id

PATIENT = patient_id()


def event_chain(records_path: str) -> None:
    """Merge the per-visit extractions with visit metadata, ICD notes and vitals, check grounding and
    cross-bucket duplication, and write event_chain.jsonl."""
    visits = {json.loads(l)["visit_number"]: json.loads(l) for l in open(f"{BASE}/visits.jsonl")}

    VISITS = sorted(visits)

    if os.path.exists(f"{BASE}/vitals.json"):
        orig_vitals = {int(k): v for k, v in json.load(open(f"{BASE}/vitals.json")).items()}
    else:
        sys.exit(f"FATAL: no vitals source in {BASE}. Run the vitals step first.")

    raw = json.load(open(records_path))
    records = raw["records"] if isinstance(raw, dict) else raw
    by_v = {r["visit_number"]: r for r in records}

    WORD = re.compile(r"[a-z0-9]+")

    def toks(s):
        return [w for w in WORD.findall(str(s).lower()) if len(w) >= 3]

    def gfrac(a, src):
        t = toks(a)
        return 1.0 if not t else sum(1 for w in t if w in src) / len(t)

    chain, gs, weak, dups = [], [], [], []
    for v in VISITS:
        meta = visits[v]
        rec = by_v.get(v, {})
        src = set(toks(" ".join(meta["soap"].values())))
        out = {"patient_id": PATIENT, "visit_number": v, "date": meta["date"], "department": meta["department"],
               "icd_notes": meta.get("icd") or []}
        for f in FIELDS:
            out[f] = rec.get(f, []) or []
            for a in out[f]:
                g = gfrac(a, src)
                gs.append(g)
                if g < 0.5:
                    weak.append((v, f, a[:50], round(g, 2)))
        out["vitals"] = orig_vitals.get(v, {})
        seen = {}
        for f in FIELDS:
            for a in out[f]:
                k = re.sub(r"\s+", " ", str(a).lower()).strip()
                if k in seen and seen[k] != f:
                    dups.append((v, k[:40], seen[k], f))
                seen[k] = f
        chain.append(out)

    with open(f"{BASE}/event_chain.jsonl", "w", encoding="utf-8") as fh:
        for c in chain:
            fh.write(json.dumps(c, ensure_ascii=False) + "\n")

    n = len(gs)
    miss = [v for v in VISITS if v not in by_v]
    print(f"patient: {PATIENT}  visits: {min(VISITS)}..{max(VISITS)}")
    print(f"records: {len(chain)}  events: {n}  mean_grounding: {sum(gs)/n if n else 0:.3f}")
    print(f"vitals: {sum(1 for c in chain if c['vitals'])} visits carry >=1")
    print(f"missing visits: {len(miss)} {miss[:20]}")
    print(f"weak events (<0.5): {len(weak)} ({100*len(weak)/max(1,n):.1f}%)")
    print(f"cross-bucket duplicates: {len(dups)}")
    for w in weak[:12]:
        print("   WEAK", w)
    for d in dups[:12]:
        print("   DUP", d)
    for f in FIELDS:
        print(f"  {f:26} visits_with={sum(1 for c in chain if c[f]):3}  events={sum(len(c[f]) for c in chain)}")
    print("wrote event_chain.jsonl")


def windows() -> None:
    """Tags-step preparation: per visit, up to two most recent prior visits in the same department and the
    most recent visit in a different department. Writes batch files and a manifest."""
    chain = [json.loads(l) for l in open(f"{BASE}/event_chain.jsonl")]
    chain.sort(key=lambda r: r["visit_number"])
    by_v = {r["visit_number"]: r for r in chain}

    def facts(r):
        return {f: r.get(f, []) for f in FIELDS}

    prev_by_dept = {}
    items = []
    for r in chain:
        v = r["visit_number"]
        dept = r["department"]
        prevs = prev_by_dept.get(dept, [])
        tprev = prevs[-2:]
        xdept = None
        for pr in reversed(chain[:chain.index(r)]):
            if pr["department"] != dept:
                xdept = pr["visit_number"]
                break
        items.append({
            "visit_number": v, "date": r["date"], "department": dept,
            "is_first_in_dept": len(prevs) == 0,
            "prev_visits": [p for p in tprev],
            "prev_cross_dept_visit": xdept,
            "current_facts": facts(r),
            "prev_facts": [{"visit_number": p, "date": by_v[p]["date"], "facts": facts(by_v[p])} for p in tprev],
        })
        prev_by_dept.setdefault(dept, []).append(v)

    bdir = f"{BASE}/_tag_batches"
    os.makedirs(bdir, exist_ok=True)
    for f in os.listdir(bdir):
        os.remove(os.path.join(bdir, f))
    SZ = 20
    manifest = []
    for i in range(0, len(items), SZ):
        ch = items[i:i+SZ]
        p = os.path.join(bdir, f"tag_{i//SZ:02d}.json")
        json.dump(ch, open(p, "w"), ensure_ascii=False)
        manifest.append({"batch_id": i//SZ, "path": p, "n": len(ch),
                         "first_v": ch[0]["visit_number"], "last_v": ch[-1]["visit_number"]})
    json.dump(manifest, open(f"{BASE}/_tag_manifest.json", "w"))
    n_first = sum(1 for it in items if it["is_first_in_dept"])
    print(f"visits: {len(items)} | first-in-dept: {n_first} | batches: {len(manifest)} (size {SZ})")
    print("example item:", json.dumps(items[0], ensure_ascii=False)[:500])
    print("manifest:", json.dumps(manifest))


def tags_post(tags_path: str) -> None:
    """Merge the change tags with the window metadata into tags.jsonl, with statistics."""
    items = {}
    for p in glob.glob(f"{BASE}/_tag_batches/*.json"):
        for it in json.load(open(p)):
            items[it["visit_number"]] = it

    raw = json.load(open(tags_path))
    tagrows = raw["tags"] if isinstance(raw, dict) else raw
    tagmap = {t["visit_number"]: t for t in tagrows}

    out = []
    for v in sorted(items):
        it = items[v]
        t = tagmap.get(v, {"tag": "MISSING", "noise_score": -1, "note": ""})
        out.append({
            "visit_number": v, "date": it["date"], "department": it["department"],
            "prev_visits": it["prev_visits"], "prev_cross_dept_visit": it["prev_cross_dept_visit"],
            "tag": t["tag"], "noise_score": t["noise_score"], "note": t["note"],
        })
    with open(f"{BASE}/tags.jsonl", "w", encoding="utf-8") as fh:
        for o in out:
            fh.write(json.dumps(o, ensure_ascii=False) + "\n")

    tagc = Counter(o["tag"] for o in out)
    scorec = Counter(o["noise_score"] for o in out)
    print(f"tags: {len(out)}/{len(items)}  | tag dist: {dict(tagc)}  "
          f"| noise_score dist: {dict(sorted(scorec.items()))}")
    first = [o["visit_number"] for o in out if items[o["visit_number"]]["is_first_in_dept"]
             and not (o["tag"] == "change" and o["noise_score"] == 0)]
    print(f"first-in-dept visits not tagged change/0: {len(first)} {first[:15]}")
    by_dept = {}
    for o in out:
        d = o["department"].split(" (")[0]
        by_dept.setdefault(d, []).append(o)
    print("per-dept no_change share:")
    for d, rows in sorted(by_dept.items(), key=lambda x: -len(x[1])):
        nc = sum(1 for r in rows if r["tag"] == "no_change")
        mean = sum(r["noise_score"] for r in rows) / len(rows)
        print(f"   {d:24s} {len(rows):4d} visits  no_change={nc:4d} ({100*nc/len(rows):.0f}%)  meanScore={mean:.2f}")
    print("wrote tags.jsonl")


def main():
    ap = argparse.ArgumentParser(description="Stage B: event chain and change annotation.")
    sp = ap.add_subparsers(dest="cmd", required=True)
    sp.add_parser("event-chain").add_argument("records")
    sp.add_parser("windows")
    sp.add_parser("tags-post").add_argument("tags")
    a = ap.parse_args()
    if a.cmd == "event-chain":
        event_chain(a.records)
    elif a.cmd == "windows":
        windows()
    else:
        tags_post(a.tags)


if __name__ == "__main__":
    main()