#!/usr/bin/env python3
"""Axis 4 (entity scoping): LLM discovery, then question generation.

Stage 1: the LLM reads candidate atoms and resolves who a condition belongs to (the patient, a named
relative or a household contact) and whether a model without memory could answer. Stage 2: the axis-4
rules (pipeline.prompts) write the question and gold.

  python -m pipeline.questions.axis4     discovery -> _discover_axis4.json
"""
from __future__ import annotations

import json
import re
from collections import Counter
from pathlib import Path

from pipeline import prompts
from pipeline.common import BASE, Progress, model, patient_id
from pipeline.llm import ChatClient, subst

PATIENT = patient_id()
ENGINE = model("generation")

AXIS = 4
SCAN_FIELDS = ["diagnosis", "patient_self_description", "doctor_notes", "doctor_extra_notes"]

DISCOVER_SYS = (
    "You read one patient's chart and find facts whose owner is ambiguous or not obvious: a condition or\n"
    "event that belongs to a relative, such as a mother, father, son, daughter, wife, husband or\n"
    "sibling; a household-contact or contagion event; or the patient's own condition that a reader could\n"
    "take for someone else's. These facts test whether a memory method attributes a fact to the right\n"
    "person. Return only JSON.")


RELATIVE = re.compile(r"(?i)\b(?:mother|father|mom|dad|parents?|sisters?|brothers?|siblings?|sons?|daughters?"
                      r"|child(?:ren)?|wife|husband|spouse|family|fhx|aunt|uncle|grand\w*|niece|nephew|cousin"
                      r"|household|contact)\b")


HISTORY = re.compile(r"(?i)^\s*(?:history of|hx of|h/o|known)\b")
ABBREVIATIONS = {
    "htn": "hypertension", "dm": "diabetes", "t2dm": "type 2 diabetes", "t1dm": "type 1 diabetes",
    "cad": "coronary artery disease", "chf": "heart failure", "hf": "heart failure", "mi": "myocardial infarction",
    "af": "atrial fibrillation", "afib": "atrial fibrillation", "cva": "stroke", "ckd": "chronic kidney disease",
    "esrd": "end-stage renal disease", "copd": "chronic obstructive pulmonary disease", "tb": "tuberculosis",
    "hbv": "hepatitis b", "hcv": "hepatitis c", "gerd": "reflux", "oa": "osteoarthritis", "ra": "rheumatoid arthritis",
    "bph": "benign prostatic hyperplasia", "uti": "urinary tract infection", "pud": "peptic ulcer",
}


def expand(text):
    """The lower-cased text followed by the full names of the abbreviations it contains."""
    t = str(text or "").lower()
    extra = [full for ab, full in ABBREVIATIONS.items() if re.search(rf"(?<![a-z0-9]){ab}(?![a-z0-9])", t)]
    return " ".join([t] + extra)


def own_diagnoses(chain):
    """The patient's own conditions as (text, atom, visit, field), text lower-cased with abbreviations expanded:
    diagnosis atoms and the patient_self_description atoms that state a history ("history of X"); atoms that
    mention a relative or a household contact are left out."""
    atoms = [(a, r.get("visit_number"), "diagnosis") for r in chain for a in (r.get("diagnosis") or [])]
    atoms += [(a, r.get("visit_number"), "patient_self_description") for r in chain
              for a in (r.get("patient_self_description") or []) if HISTORY.search(str(a))]
    return [(expand(a), a, v, f) for a, v, f in atoms if not RELATIVE.search(str(a))]


def patient_has(entity, own):
    """(atom, visit, field) of the first of the patient's own conditions that names the entity (abbreviations
    expanded), or None."""
    e = str(entity or "").strip().lower()
    names = [e] if len(e) >= 4 else []
    names += [full for ab, full in ABBREVIATIONS.items() if re.search(rf"(?<![a-z0-9]){ab}(?![a-z0-9])", e)]
    return next(((a, v, f) for text, a, v, f in own if any(n in text for n in names)), None)


def field_of(rec, text):
    """The scanned field of a visit whose atom is the text, or else contains it (case and spacing ignored), or
    None."""
    def norm(t):
        return re.sub(r"\s+", " ", str(t or "")).strip().lower()
    t = norm(text)
    atoms = [(f, norm(a)) for f in SCAN_FIELDS for a in rec.get(f) or []]
    return (next((f for f, a in atoms if t and a == t), None)
            or next((f for f, a in atoms if t and t in a), None))


def patient_owned(scoping):
    """Whether discovery resolved the fact to the patient."""
    owner = str(scoping.get("owner") or "").strip().lower()
    return scoping.get("kind") == "patient_owned" or bool(
        re.fullmatch(r"(?:the )?patient(?:'s)?(?: (?:him|her)self)?(?: own)?", owner))


def discover_scopings(cli, by_v):
    rows = []
    for v in sorted(by_v):
        for f in SCAN_FIELDS:
            for a in by_v[v].get(f) or []:
                rows.append(f"v{v} {by_v[v]['date']} {f}: {str(a)[:120]}")
    chunks = ["\n".join(rows[i:i + 140]) for i in range(0, len(rows), 140)] or [""]
    found = []
    for chart_facts in chunks:
        usr = (
            "Chart facts, one per line (visit, date, field, text):\n" + chart_facts + "\n\n"
            "Find facts that belong to someone other than the patient, or whose owner a reader could get\n"
            "wrong. Classify each one:\n"
            "- relative_owned: a condition or event of a named relative, such as family history,\n"
            "  'mother's ...' or 'father had ...'.\n"
            "- contagion: a household contact or exposure event, such as a family member's infection.\n"
            "- patient_owned: the patient's own condition or event that a reader could take for someone\n"
            "  else's, for example because it is mentioned together with a relative's.\n"
            "Skip the patient's own diagnoses that nothing links to another person. Skip negated family\n"
            "history such as 'denied family history of X': there the entity does not exist to be owned.\n\n"
            "For each fact give the owner (the relative, 'household contact' or 'the patient'), the entity\n"
            "(the condition or event), the visit, and a one-line note on why a model without memory would\n"
            "attribute it to the wrong person.\n\n"
            "Return JSON in this form:\n"
            '{"scopings": [{\n'
            '  "visit": <int>,\n'
            '  "kind": "<relative_owned|contagion|patient_owned>",\n'
            '  "owner": "<who it belongs to, e.g. \'the patient\'s mother\' or \'the patient\'>",\n'
            '  "entity": "<the condition or event>",\n'
            '  "atom_verbatim": "<the source text>",\n'
            '  "why_memory_required": "<one line>"\n'
            "}]}")
        try:
            found += cli.complete_json(DISCOVER_SYS, usr).get("scopings") or []
        except Exception as e:
            print(f"  discover chunk FAILED: {str(e)[:100]}")
    return found


def discover() -> None:
    """Discover the axis-4 candidates of this patient -> _discover_axis4.json."""
    cli = ChatClient(model=ENGINE, reasoning_effort="medium", max_completion_tokens=6000)

    chain = [json.loads(l) for l in open(f"{BASE}/event_chain.jsonl") if l.strip()]
    by_v = {r["visit_number"]: r for r in chain}
    plast = max(r["date"] for r in chain)

    with Progress("axis4 discover (LLM)", total=1, out=f"{BASE}/progress_axis4_discover.json") as pg:
        scopings = discover_scopings(cli, by_v)
        pg.tick(ok=bool(scopings))

    AX4 = subst(prompts.AXIS4_RULES)
    COMMON = subst(prompts.AXIS4_COMMON)
    system = subst(f"{COMMON}\n\n{AX4}")

    candidates = []
    seen = set()
    own = own_diagnoses(chain)
    for s in scopings:
        rv = s.get("visit")
        if rv not in by_v:
            continue
        k = (str(s.get("owner")).lower().strip(), str(s.get("entity")).lower().strip()[:40])
        if k in seen:
            continue
        seen.add(k)
        own_case = patient_owned(s)
        also = None if own_case else patient_has(s.get("entity"), own)
        seed = {"kind": s.get("kind"), "owner": s.get("owner"), "entity": s.get("entity"),
                "atom_verbatim": s.get("atom_verbatim"), "why_memory_required": s.get("why_memory_required")}
        if also:
            seed["patient_also_has"] = also[0]
        fact_json = json.dumps(seed, ensure_ascii=False)
        usr = (
            "Write one entity-attribution question for this seed, following the axis-4 rules above. The\n"
            "question must not reveal the owner; the gold answer states it. Return only JSON:\n"
            '{"question": "<question>", "answer": "<gold naming the correct owner>"}\n\n'
            "Seed:\n" + fact_json)
        try:
            q = cli.complete_json(system, usr)
        except Exception as e:
            print(f"  gen FAILED: {str(e)[:100]}")
            continue
        if not q.get("question") or not q.get("answer"):
            continue
        if str(q["answer"]).lstrip().startswith("Yes") != (own_case or bool(also)):
            print(f"  reject {s.get('entity')!r}: the gold's yes/no does not match the resolved owner")
            continue
        evidence = [{"visit": rv, "date": by_v[rv]["date"], "field": field_of(by_v[rv], s.get("atom_verbatim")),
                     "atom": s.get("atom_verbatim")}]
        owner_step = f"Owner resolved: {s.get('owner')} (kind={s.get('kind')})."
        ev_visits = [rv]
        if also:
            patom, pv, pfield = also
            if pv != rv:
                ev_visits.append(pv)
            evidence.append({"visit": pv, "date": by_v[pv]["date"] if pv in by_v else None, "field": pfield,
                             "atom": patom})
            owner_step = (f"The fact at v{rv} is {s.get('owner')}'s; the patient's own {s.get('entity')} is "
                          f"documented at v{pv}.")
        gv = max(v for v in ev_visits if v in by_v)
        candidates.append({
            "qid": f"{PATIENT}/axis{AXIS}/{len(candidates):03d}", "patient_id": PATIENT,
            "axis": "4_entity_scoping", "question": q["question"], "answer": q["answer"],
            "answer_type": "open_exact", "grounding": "record",
            "evidence_visits": ev_visits, "gt_visit": gv, "gt_date": by_v[gv]["date"],
            "question_appearance_window": {"from": by_v[gv]["date"], "to": plast},
            "derivation": {
                "evidence": evidence,
                "steps": [f"Atom at v{rv}: {s.get('atom_verbatim')!r}.", owner_step,
                          f"{s.get('why_memory_required')}"],
                "gold": q["answer"], "rivals": [],
                "operands": {"owner": s.get("owner"), "entity": s.get("entity"), "kind": s.get("kind"),
                             "patient_also_has": also[0] if also else None},
                "derived": False},
            "provenance": f"LLM-resolved scoping ({ENGINE}); owner={s.get('owner')}",
            "axis_specific": {"kind": s.get("kind"), "owner": s.get("owner"), "entity": s.get("entity")},
        })

    out = {"axis": AXIS, "patient_id": PATIENT, "n_record": len(candidates), "candidates": candidates}
    Path(f"{BASE}/_discover_axis{AXIS}.json").write_text(json.dumps(out, ensure_ascii=False, indent=1))
    print(f"\npatient {PATIENT}  engine={ENGINE}")
    print(f"  scopings discovered (LLM): {len(scopings)}  by kind: {dict(Counter(s.get('kind') for s in scopings))}")
    print(f"  CANDIDATES: {len(candidates)}  -> {BASE}/_discover_axis{AXIS}.json")
    for c in candidates[:3]:
        print(f"      Q: {c['question'][:70]}  A: {str(c['answer'])[:40]}")


if __name__ == "__main__":
    discover()