#!/usr/bin/env python3
"""Stage A, preprocessing: one patient's synthetic timeline -> per-visit records (deterministic).

  export    visit table -> ALL_visits.txt (term glossary applied)
  parse     ALL_visits.txt -> visits.jsonl
  split     visits.jsonl -> visits/v*.json + visit_manifest.json (the extraction manifest)
  vitals    visit table -> vitals.json + vitals_timeline.json
  facts     patient.env: the per-patient values the prompts are filled with
"""
from __future__ import annotations

import argparse
import json
import re
import shlex
import sys
import unicodedata

import pandas as pd

from pipeline import common
from pipeline.common import BASE, VISIT_TABLE, patient_id

PATIENT = patient_id()

PID, DEPT, ISO, VT = "record_id", "department", "visit_date", "visit_type"
S, O, A, P = "subjective", "objective", "assessment", "plan"
ICDS = ["icd_1", "icd_2", "icd_3", "icd_4", "icd_5", "icd_6"]


def caseless_letters(text):
    """The number of letters without case (Unicode category Lo) in `text`."""
    return sum(1 for ch in text if unicodedata.category(ch) == "Lo")


def norm(x):
    """Collapse all whitespace (incl. embedded newlines) to single spaces; strip."""
    return " ".join(str(x).split())


def build_raw_text(df, patient):
    sub = df[df[PID] == patient].sort_values([ISO, "visit_no"])
    lines = []
    for _, r in sub.iterrows():
        icds = [str(r[c]).strip() for c in ICDS if str(r[c]).strip()]
        header = f"### {r[ISO]} | {norm(r[DEPT])} | {VT}:{norm(r[VT])} | ICD:[{','.join(icds)}]"
        lines.append(header)
        lines.append(f"S: {norm(r[S])}")
        lines.append(f"O: {norm(r[O])}")
        lines.append(f"A: {norm(r[A])}")
        lines.append(f"P: {norm(r[P])}")
    return "\n".join(lines) + "\n", len(sub)


def export_visits(patient: int, out_path: str) -> None:
    """Visit table rows of one patient -> all-visits file.

    One block per visit (header line plus S/O/A/P); the term glossary turns every term into
    "normalized form (term)".
    """
    df = pd.read_parquet(VISIT_TABLE)
    for c in [S, O, A, P, DEPT, VT] + ICDS:
        df[c] = df[c].fillna("")

    raw, n = build_raw_text(df, patient)
    g = common.load_glossary()
    pat = common.glossary_re(g)
    out = common.apply_glossary(raw, g, pat)

    left = caseless_letters(re.sub(r"\([^()]*\)", "", out))
    open(out_path, "w", encoding="utf-8").write(out)
    print(f"patient={patient}  visits={n}  lines={out.count(chr(10))}  caseless_letters_outside_parens={left}")
    print(f"wrote {out_path}")


HEADER_RE = re.compile(
    r"^###\s+(\d{4}-\d{2}-\d{2})\s+\|\s+(.+?)\s+\|\s+visit_type:([^|]*?)\s+\|\s+ICD:\[(.*?)\]\s*$"
)
SOAP_RE = re.compile(r"^([SOAP]):\s?(.*)$")


def dept_key(dept):
    base = re.sub(r"\s*\(.*?\)\s*", "", dept).strip().lower()
    return re.sub(r"[^a-z0-9]+", "_", base).strip("_")


def parse(path):
    visits = []
    cur = None
    field = None
    for line in open(path, encoding="utf-8"):
        line = line.rstrip("\n")
        m = HEADER_RE.match(line)
        if m:
            if cur:
                visits.append(cur)
            date, dept, vt, icd = m.groups()
            codes = [c.strip() for c in icd.split(",") if c.strip()]
            cur = {
                "visit_number": len(visits) + 1,
                "date": date,
                "department": dept,
                "department_key": dept_key(dept),
                "visit_type": vt,
                "icd": codes,
                "soap": {"S": "", "O": "", "A": "", "P": ""},
            }
            field = None
            continue
        sm = SOAP_RE.match(line)
        if sm and cur is not None:
            field = sm.group(1)
            cur["soap"][field] = sm.group(2)
        elif cur is not None and field:
            cur["soap"][field] += " " + line
    if cur:
        visits.append(cur)
    return visits


def parse_visits(inp: str, out: str) -> None:
    """ALL_visits.txt -> deterministic per-visit JSONL records."""
    visits = parse(inp)
    with open(out, "w", encoding="utf-8") as f:
        for v in visits:
            f.write(json.dumps(v, ensure_ascii=False) + "\n")
    depts = {}
    for v in visits:
        depts[v["department"]] = depts.get(v["department"], 0) + 1
    print(f"parsed {len(visits)} visits -> {out}")
    print("departments:", json.dumps(depts, ensure_ascii=False))
    print("first record:", json.dumps(visits[0], ensure_ascii=False)[:300])


def split() -> None:
    """Split visits.jsonl into one file per visit and write the extraction manifest (absolute paths)."""
    BASE = common.BASE.resolve()
    src = BASE / "visits.jsonl"
    if not src.exists():
        raise SystemExit(f"FATAL: {src} not found; run the parse step first")
    rows = [json.loads(l) for l in open(src, encoding="utf-8") if l.strip()]
    if not rows:
        raise SystemExit(f"FATAL: {src} is empty")

    out = BASE / "visits"
    out.mkdir(parents=True, exist_ok=True)

    vn = [r["visit_number"] for r in rows]
    if len(set(vn)) != len(vn):
        raise SystemExit(f"FATAL: duplicate visit_number in {src}; numbering is the join key for the "
                         f"whole pipeline (vitals, chain, questions all index by it)")

    manifest = []
    for r in sorted(rows, key=lambda x: x["visit_number"]):
        p = out / f"v{r['visit_number']}.json"
        p.write_text(json.dumps(r, ensure_ascii=False, indent=1), encoding="utf-8")
        manifest.append({"visit_number": r["visit_number"], "path": str(p)})

    (BASE / "visit_manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=1))
    print(f"patient {PATIENT}: {len(manifest)} visits -> {out}/v*.json")
    print(f"  v{min(vn)}..v{max(vn)}  ({min(r['date'] for r in rows)} .. {max(r['date'] for r in rows)})")
    print(f"-> {BASE/'visit_manifest.json'}")


COLS = {"weight": "weight_kg", "height": "height_cm", "blood_pressure": "BP", "pulse": "pulse"}


def clean(x):
    x = str(x).strip()
    return x if x not in ("", "-", "nan", "None", "NaN") else None


def vitals() -> None:
    """Structured vitals (height, weight, blood pressure, pulse) from the visit table.

    Visits are numbered by the same (date, visit_no) sort as visits.jsonl, and the row count is checked
    against visits.jsonl. Writes vitals.json and vitals_timeline.json.
    """
    df = pd.read_parquet(VISIT_TABLE)
    sub = (df[df["record_id"] == PATIENT]
           .sort_values(["visit_date", "visit_no"])
           .reset_index(drop=True))
    if sub.empty:
        sys.exit(f"FATAL: patient {PATIENT} not in {VISIT_TABLE}")
    sub["vn"] = range(1, len(sub) + 1)

    vpath = BASE / "visits.jsonl"
    if vpath.exists():
        visits = {json.loads(l)["visit_number"]: json.loads(l) for l in open(vpath)}
        if len(visits) != len(sub):
            sys.exit(f"FATAL: visit-count mismatch; the visit table has {len(sub)} rows, "
                     f"{vpath} has {len(visits)}. Numbering would drift; refusing to write.")
        bad = [int(r["vn"]) for _, r in sub.iterrows()
               if visits[int(r["vn"])]["date"] != str(r["visit_date"])[:10]]
        if bad:
            sys.exit(f"FATAL: date misalignment on {len(bad)} visits (first: {bad[:5]}). "
                     f"visit-table sort != visits.jsonl sort; refusing to write.")
    else:
        print(f"  note: {vpath} absent; skipping the alignment check")

    vit = {}
    for _, r in sub.iterrows():
        v = {name: clean(r[col]) for col, name in COLS.items() if clean(r[col])}
        if v:
            vit[int(r["vn"])] = v

    (BASE / "vitals.json").write_text(
        json.dumps({str(k): v for k, v in sorted(vit.items())}, ensure_ascii=False, indent=1))

    timelines: dict = {}
    for vn in sorted(vit):
        date = str(sub.loc[sub["vn"] == vn, "visit_date"].iloc[0])[:10]
        for measure, val in vit[vn].items():
            timelines.setdefault(measure, []).append(
                {"visit_number": vn, "date": date, "value": val})
    (BASE / "vitals_timeline.json").write_text(json.dumps(timelines, ensure_ascii=False, indent=1))

    print(f"patient {PATIENT}: {len(sub)} visits, {len(vit)} with >=1 structured vital")
    for m, series in timelines.items():
        vals = [s["value"] for s in series]
        print(f"  {m:<10} {len(series):>3} points  {series[0]['date']}..{series[-1]['date']}  "
              f"{vals[:6]}{'...' if len(vals) > 6 else ''}")
    if not timelines:
        print("  no structured vitals")
    print(f"-> {BASE/'vitals.json'}\n-> {BASE/'vitals_timeline.json'}")


def facts(pdesc: str) -> None:
    """Write patient.env: visit count, first and last visit date, and the clinical descriptor."""
    W = BASE
    v = [json.loads(l) for l in open(f"{W}/visits.jsonl")]
    env = {"MEMEBENCH_WORK": W, "MEMEBENCH_PATIENT": PATIENT, "MEMEBENCH_PDESC": pdesc or common.DEFAULT_PDESC,
           "MEMEBENCH_NVISITS": len(v), "MEMEBENCH_PFIRST": min(x["date"] for x in v),
           "MEMEBENCH_PLAST": max(x["date"] for x in v)}
    open(f"{W}/patient.env", "w").write("".join(f"export {k}={shlex.quote(str(x))}\n" for k, x in env.items()))
    print(f"  patient.env: {len(v)} visits")


def main():
    ap = argparse.ArgumentParser(description="Stage A: preprocess one patient's synthetic timeline.")
    sp = ap.add_subparsers(dest="cmd", required=True)
    sp.add_parser("export").add_argument("--out", required=True)
    p = sp.add_parser("parse")
    p.add_argument("--inp", default="ALL_visits.txt")
    p.add_argument("--out", default="visits.jsonl")
    sp.add_parser("split")
    sp.add_parser("vitals")
    sp.add_parser("facts").add_argument("--pdesc", default="")
    a = ap.parse_args()
    if a.cmd == "export":
        export_visits(PATIENT, a.out)
    elif a.cmd == "parse":
        parse_visits(a.inp, a.out)
    elif a.cmd == "split":
        split()
    elif a.cmd == "vitals":
        vitals()
    else:
        facts(a.pdesc)


if __name__ == "__main__":
    main()