#!/usr/bin/env python3
"""Axis 7 (cross-department synthesis): LLM discovery, then templates.

The LLM reads the chain in overlapping windows of chart facts, each grouped by department, and returns pairs
of facts in different departments that must be joined to answer, where neither department's record alone suffices (e.g. a referral whose
outcome is recorded in another department). Questions and golds use Q_CLOSURE / A_CLOSURE.
synthesize() authors synthetic questions.

  python -m pipeline.questions.axis7     discovery -> _discover_axis7.json
"""
from __future__ import annotations

import json
import re
from collections import defaultdict
from pathlib import Path

from pipeline.common import BASE, FIELDS, Progress, has_text, model, patient_context, patient_id, patient_line
from pipeline.llm import ChatClient

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

AXIS = 7
SCAN_FIELDS = ["doctor_plans", "doctor_extra_notes", "doctor_notes", "health_advice", "tests_and_medication",
               "diagnosis"]
WINDOW, OVERLAP = 400, 100

Q_CLOSURE = ("On {ref_date} my {ref_dept} doctor said {ref_atom_verbatim}. Did that ever happen, and "
             "what actually happened at that follow-up — for example, anything found or done?")
A_CLOSURE = ('Yes. {ref_dept} recorded "{ref_atom_verbatim}" on {ref_date}, and {closer_dept} recorded '
             '"{closer_atom_verbatim}" on {closer_date}: {outcome_summary}')

def _sentence(text):
    """The text with a closing full stop."""
    text = str(text or "").strip()
    return text if not text or text[-1] in ".!?" else text + "."


DISCOVER_SYS = (
    "You read one patient's chart, grouped by department, and find cross-department joins. A join is a\n"
    "question that can be answered only by combining a fact from department A with a fact from\n"
    "department B: neither department's record alone is enough, and B's record does not cite A's. The\n"
    "typical case: department A records a referral or a plan, and its outcome is recorded later in\n"
    "department B. Return only JSON.")


def quoted_field(text, rec):
    """The field of the visit's atom that contains the quoted text (case and spacing ignored), or None."""
    def norm(t):
        return re.sub(r"\s+", " ", str(t or "")).strip().lower().rstrip(".;, ")
    q = norm(text)
    return next((f for f in FIELDS for a in rec.get(f) or [] if q and q in norm(a)), None)


def discover_joins(cli, chain, by_v):
    """Find cross-department joins over windows of WINDOW chart facts in visit order; consecutive windows share
    OVERLAP facts, and each window is shown grouped by department."""
    rows = [(r["visit_number"], r["date"], r["department"], f, str(a)[:120])
            for r in sorted(chain, key=lambda x: x["visit_number"]) for f in SCAN_FIELDS for a in r.get(f) or []]
    found, seen = [], set()
    for start in range(0, max(len(rows) - OVERLAP, 1), WINDOW - OVERLAP):
        by_dept = defaultdict(list)
        for v, d, dept, f, t in rows[start:start + WINDOW]:
            by_dept[dept].append(f"  v{v} {d} {f}: {t}")
        chart_by_department = "\n\n".join(f"### {dept}\n" + "\n".join(lines) for dept, lines in by_dept.items())
        joins = _discover_window(cli, chart_by_department)
        if not isinstance(joins, list):
            joins = [joins]
        n_bad = sum(1 for j in joins if not isinstance(j, dict))
        if n_bad:
            print(f"  skipped {n_bad} join(s) that are not objects")
        for j in (j for j in joins if isinstance(j, dict)):
            k = (j.get("ref_visit"), j.get("closer_visit"), str(j.get("ref_atom_verbatim")).lower())
            if k not in seen:
                seen.add(k)
                found.append(j)
    return found


def _discover_window(cli, chart_by_department):
    """One discovery call over a window of chart facts."""
    usr = (
        "Chart by department:\n" + chart_by_department + "\n\n"
        "Find cross-department joins. Each join must:\n"
        "- involve two different departments, A and B;\n"
        "- have a fact in A, such as a referral, a plan or a request, whose outcome or resolution is a fact\n"
        "  in B;\n"
        "- be impossible to answer from either department alone: B's note does not cite A's referral or\n"
        "  plan.\n"
        "Do not invent links; report only joins that the chart facts actually support. Copy both facts\n"
        "exactly as they appear above, without the visit, date and field.\n\n"
        "Return JSON in this form:\n"
        '{"joins": [{\n'
        '  "ref_visit": <int>,\n'
        '  "ref_dept": "<A>",\n'
        '  "ref_atom_verbatim": "<A\'s fact>",\n'
        '  "closer_visit": <int>,\n'
        '  "closer_dept": "<B>",\n'
        '  "closer_atom_verbatim": "<B\'s fact>",\n'
        '  "outcome_summary": "<one line: what came of it>",\n'
        '  "topic": "<slug>"\n'
        "}]}")
    try:
        return cli.complete_json(DISCOVER_SYS, usr).get("joins") or []
    except Exception as e:
        print(f"  discover window FAILED: {str(e)[:120]}")
        return []


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

    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("axis7 discover (LLM)", total=1, out=f"{BASE}/progress_axis7_discover.json") as pg:
        joins = discover_joins(cli, chain, by_v)
        pg.tick(ok=bool(joins))

    candidates = []
    for j in joins:
        rv, cv = j.get("ref_visit"), j.get("closer_visit")
        if (rv not in by_v or cv not in by_v or cv <= rv
                or by_v[rv]["department"] == by_v[cv]["department"]):
            continue
        ref_field = quoted_field(j.get("ref_atom_verbatim"), by_v[rv])
        closer_field = quoted_field(j.get("closer_atom_verbatim"), by_v[cv])
        if not (ref_field and closer_field):
            print(f"  skip v{rv}->v{cv}: a quoted fact is not in its visit")
            continue
        ops = {"ref_date": by_v[rv]["date"], "ref_dept": by_v[rv]["department"].split(" (")[0],
               "ref_atom_verbatim": j.get("ref_atom_verbatim"),
               "closer_date": by_v[cv]["date"], "closer_dept": by_v[cv]["department"].split(" (")[0],
               "closer_atom_verbatim": j.get("closer_atom_verbatim"),
               "outcome_summary": _sentence(j.get("outcome_summary", ""))}
        candidates.append({
            "qid": f"{PATIENT}/axis{AXIS}/{len(candidates):03d}", "patient_id": PATIENT,
            "axis": "7_cross_dept_synthesis",
            "question": Q_CLOSURE.format(**ops), "answer": A_CLOSURE.format(**ops),
            "answer_type": "open_exact", "grounding": "record",
            "evidence_visits": [rv, cv], "gt_visit": cv, "gt_date": by_v[cv]["date"],
            "question_appearance_window": {"from": by_v[cv]["date"], "to": plast},
            "derivation": {
                "evidence": [{"visit": rv, "date": by_v[rv]["date"], "field": ref_field,
                              "atom": j.get("ref_atom_verbatim")},
                             {"visit": cv, "date": by_v[cv]["date"], "field": closer_field,
                              "atom": j.get("closer_atom_verbatim")}],
                "steps": [f"{ops['ref_dept']} v{rv}: {j.get('ref_atom_verbatim')!r} — a request, "
                          f"no outcome here.",
                          f"{ops['closer_dept']} v{cv}: {j.get('closer_atom_verbatim')!r} — the outcome, "
                          f"another dept.",
                          "The outcome note does not cite the referral; answering requires joining both."],
                "gold": A_CLOSURE.format(**ops), "rivals": [], "operands": ops, "derived": True},
            "provenance": f"LLM-discovered cross-dept join ({ENGINE}); A={ops['ref_dept']} B={ops['closer_dept']}",
            "axis_specific": {"A_dept": ops["ref_dept"], "B_dept": ops["closer_dept"], "topic": j.get("topic")},
        })

    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"  cross-dept joins discovered (LLM): {len(joins)}")
    print(f"  CANDIDATES (record pairs)        : {len(candidates)}")
    for c in candidates[:3]:
        print(f"      {c['axis_specific']['A_dept']} -> {c['axis_specific']['B_dept']}: {c['question'][:55]}")
    print(f"  -> {BASE}/_discover_axis{AXIS}.json")


def synthesize(cli, chain, by_v, forbidden, n, batch):
    """Author department A's referral and department B's outcome as two atoms; the join is the question."""
    out, vs = [], sorted(by_v)
    for i in range(n):
        a_host = vs[int(len(vs) * (0.15 + 0.5 * i / max(1, n)))]
        a_dept = by_v[a_host]["department"]
        later = [v for v in vs if v > a_host and by_v[v]["department"] != a_dept]
        if not later:
            continue
        b_host = later[min(len(later) - 1, 12)]
        b_dept = by_v[b_host]["department"]
        a_date, b_date = by_v[a_host]["date"], by_v[b_host]["date"]
        patient, context = patient_line(), patient_context(chain)
        forbidden_list = "\n".join(f"- {e}" for e in forbidden[:50])
        r = cli.complete_json(
            "You write two terse chart facts for one referral story in a health record. Return only JSON.",
            f"{patient}\n{context}\n"
            f"A: v{a_host} · {a_date} · {a_dept} — records the referral\n"
            f"B: v{b_host} · {b_date} · {b_dept} — records the outcome\n\n"
            "A refers the patient to B about a topic, and B later records what came of it. Neither fact may\n"
            "mention the other department's note: the point is that the join itself is not recorded. Write\n"
            "in terse chart style, such as 'refer GI for dyspepsia', not in prose.\n\n"
            "Do not use these, which existing answer keys already track:\n" + forbidden_list +
            "\n\nReturn JSON in this form:\n"
            '{"ref_text": "<A\'s referral fact>",\n'
            ' "outcome_text": "<B\'s outcome fact>",\n'
            ' "outcome_summary": "<one line: what came of it>",\n'
            ' "entity": "<topic>",\n'
            ' "topic": "<slug>"}')
        missing = [k for k in ("ref_text", "outcome_text", "outcome_summary") if not has_text(r, k)]
        if missing:
            print(f"    axis7[{i}] rejected: the reply lacks {', '.join(missing)}")
            continue
        a1 = {"atom_id": f"{batch}/{len(out)*2:03d}", "batch_id": batch, "axis": "7",
              "visit_number": a_host, "date": by_v[a_host]["date"], "department": a_dept,
              "field": "doctor_plans", "text": r["ref_text"], "entities": [r.get("entity", "")],
              "topic": r.get("topic"), "anchors_axis": "7", "anchors_on": []}
        a2 = {**a1, "atom_id": f"{batch}/{len(out)*2+1:03d}", "visit_number": b_host,
              "date": by_v[b_host]["date"], "department": b_dept, "field": "doctor_notes",
              "text": r["outcome_text"]}
        gold = (f"Yes. {a_dept.split(' (')[0]} recorded \"{r['ref_text']}\" on {by_v[a_host]['date']}, and "
                f"{b_dept.split(' (')[0]} recorded \"{r['outcome_text']}\" on {by_v[b_host]['date']}: "
                f"{_sentence(r['outcome_summary'])}")
        out.append({"atom": a1, "atom2": a2, "q": {
            "qid": f"{PATIENT}/axis7/S{len(out):03d}", "patient_id": PATIENT, "axis": "7_cross_dept_synthesis",
            "question": (f"On {by_v[a_host]['date']} my {a_dept.split(' (')[0]} doctor said "
                         f"{r['ref_text']}. Did that ever happen, and what actually happened at that follow-up "
                         f"— for example, anything found or done?"),
            "answer": gold, "answer_type": "open_exact", "grounding": f"synthetic:{batch}",
            "evidence_visits": [a_host, b_host], "gt_visit": b_host, "gt_date": by_v[b_host]["date"],
            "question_appearance_window": {"from": by_v[b_host]["date"], "to": by_v[vs[-1]]["date"]},
            "derivation": {"evidence": [
                {"visit": a_host, "date": by_v[a_host]["date"], "field": "doctor_plans", "atom": r["ref_text"]},
                {"visit": b_host, "date": by_v[b_host]["date"], "field": "doctor_notes", "atom": r["outcome_text"]}],
                "steps": [f"{a_dept} v{a_host}: {r['ref_text']!r} — a referral with no outcome in this dept.",
                          f"{b_dept} v{b_host}: {r['outcome_text']!r} — the outcome, in a DIFFERENT dept.",
                          "The outcome note does not cite the referral; answering requires joining both."],
                "gold": gold, "rivals": [],
                "operands": {"ref_visit": a_host, "ref_date": by_v[a_host]["date"], "ref_dept": a_dept,
                             "ref_atom_verbatim": r["ref_text"], "target_dept": b_dept,
                             "closer_date": by_v[b_host]["date"], "closer_dept": b_dept,
                             "closer_plans_verbatim": r["outcome_text"]},
                "derived": True},
            "provenance": f"synthetic cross-dept join ({batch}); A={a_dept} B={b_dept}, join unrecorded",
            "axis_specific": {"A_dept": a_dept, "B_dept": b_dept,
                              "depends_on_atoms": [a1["atom_id"], a2["atom_id"]]},
        }})
    return out


if __name__ == "__main__":
    discover()