#!/usr/bin/env python3
"""Axis 3 (drug safety, multiple choice): LLM discovery, then question generation.

Stage 1: the LLM reads the drug and constraint timelines of the evidence chain and returns standing drug
constraints (class, unsafe member, safe equivalents). Stage 2: the axis-3 rules (pipeline.prompts) turn
each constraint into a multiple-choice question with exactly one correct option. The gold is
options[].is_correct. The options are put in an order seeded by the question id and labelled A-D in that order;
the question text does not list them. synthesize() authors synthetic questions.

  python -m pipeline.questions.axis3     discovery -> _discover_axis3.json
"""
from __future__ import annotations

import json
import random
from pathlib import Path

from pipeline import prompts
from pipeline.common import (BASE, Progress, has_text, lists_options, model, patient_context, patient_id,
                             patient_line)
from pipeline.llm import ChatClient, subst

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

AXIS = 3
DRUG_ASPECTS = {"med_rationale_referral", "medication", "constraint_or_safety", "care_plan"}

DISCOVER_SYS = (
    "You read one patient's longitudinal record of medications and constraints and find standing\n"
    "drug-safety constraints. A constraint is a documented reason, specific to this patient, why some\n"
    "drugs a clinician might choose are inappropriate for this patient while a clinically reasonable\n"
    "alternative remains fine: an adverse reaction, an allergy, an intolerance, or an interacting\n"
    "medication that the patient takes. Return only JSON.")


def discover_constraints(cli, timelines):
    """Read the drug-bearing timelines and return the patient's standing drug constraints."""
    rows = []
    for t in timelines:
        segs = t.get("segments") or []
        rows.append(f"- entity={t.get('entity')!r} aspect={t.get('aspect')} dept={t.get('department')} "
                    f"status={t.get('current_status')!r} "
                    f"values={[str(s.get('value'))[:60] for s in segs[:4]]}")
    record = "\n".join(rows)
    usr = (
        "Medication and constraint record, from the evidence chain:\n" + record + "\n\n"
        "Find every standing drug-safety constraint documented for this patient. A constraint qualifies\n"
        "only if all of these hold:\n"
        "- It is specific to this patient: an allergy, intolerance or adverse reaction on file, or a drug\n"
        "  the patient takes that interacts with a drug class. General guideline caution does not count.\n"
        "- It makes at least three drugs a clinician might reach for inappropriate for this patient, while\n"
        "  at least one clinically reasonable alternative stays appropriate, so that a four-option question\n"
        "  with exactly one appropriate answer is possible.\n"
        "- It is standing, that is, still true at later visits, and not a one-off event.\n"
        "Exclude negated screening entries such as 'no known allergy', 'denied' or '(-)'.\n\n"
        "For each constraint give:\n"
        "- the drug class;\n"
        "- at least three drugs that are unsafe for this patient;\n"
        "- one to four safe alternatives;\n"
        "- a realistic scenario in which the class would be prescribed to the patient;\n"
        "- the first-line trap drugs that a clinician, or a model without memory, would reach for first.\n"
        "  These must be among the unsafe drugs, or the question is pointless.\n\n"
        "Return JSON in this form:\n"
        '{"constraints": [{\n'
        '  "constraint_id": "<slug>",\n'
        '  "implicated_class": "<class>",\n'
        '  "constraint_text": "<the documented fact, close to its wording>",\n'
        '  "source_entity": "<the entity exactly as listed above>",\n'
        '  "unsafe_examples": ["<drug>", ...],\n'
        '  "safe_examples": ["<drug>", ...],\n'
        '  "scenario_hint": "<why the class would be prescribed>",\n'
        '  "required_patient_info": ["<the exact memory fact needed>"],\n'
        '  "first_line_trap": "<the drug a model without memory picks by default>"\n'
        "}]}")
    r = cli.complete_json(DISCOVER_SYS, usr)
    return r.get("constraints") or []


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

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

    drug_tl = [t for t in tl if t.get("aspect") in DRUG_ASPECTS]

    with Progress("axis3 discover (LLM)", total=2, out=f"{BASE}/progress_axis3_discover.json") as pg:
        try:
            seeds = discover_constraints(cli, drug_tl)
        except Exception as e:
            print(f"  discover FAILED: {str(e)[:160]}")
            seeds = []
        pg.tick(ok=bool(seeds))

        AX3 = subst(prompts.AXIS3_RULES)
        COMMON = subst(prompts.AXIS3_COMMON)
        system = subst(f"{COMMON}\n\n{AX3}")
        candidates = []
        for i, s in enumerate(seeds):
            if len(s.get("unsafe_examples") or []) < 3:
                print(f"  skip {s.get('constraint_id')}: fewer than three unsafe drugs for the traps")
                continue
            gt_visit, gt_date = _anchor(s, tl, chain)
            if gt_visit is None:
                print(f"  skip {s.get('constraint_id')}: its source entity is not in the evidence chain")
                continue
            seed_json = json.dumps({
                "constraint_id": s.get("constraint_id"), "implicated_class": s.get("implicated_class"),
                "scenario_hint": s.get("scenario_hint"), "constraint_text": s.get("constraint_text"),
                "unsafe_examples": s.get("unsafe_examples"), "safe_examples": s.get("safe_examples"),
                "required_patient_info": s.get("required_patient_info"),
                "first_line_trap": s.get("first_line_trap"),
                "gt_date": gt_date, "grounding": "record"}, ensure_ascii=False)
            usr = (
                "Write the axis-3 multiple-choice drug-safety question for this seed, following every rule\n"
                "above. Return only a JSON object of the form {\"questions\": [ ... ]} with one question, or\n"
                "with an empty list if the rules cannot be met.\n\nSeed:\n" + seed_json)
            try:
                out = cli.complete_json(system, usr)
                qs = out.get("questions") or ([out] if out.get("options") else [])
            except Exception as e:
                print(f"  gen[{i}] FAILED: {str(e)[:120]}")
                continue
            for q in qs:
                if not has_text(q):
                    print(f"  reject {s.get('constraint_id')}: the reply has no question text")
                    continue
                opts = q.get("options") or []
                if not (isinstance(opts, list) and all(isinstance(o, dict) for o in opts)):
                    print(f"  reject {s.get('constraint_id')}: the options are not a list of objects")
                    continue
                for o in opts:
                    if not o.get("content") and o.get("text"):
                        o["content"] = o["text"]
                if lists_options(q.get("question"), opts):
                    print(f"  reject {s.get('constraint_id')}: the question text lists the options")
                    continue
                if sum(1 for o in opts if o.get("is_correct")) != 1:
                    print(f"  reject {s.get('constraint_id')}: not exactly 1 correct option")
                    continue
                if not all(_names(o, s.get("safe_examples") if o.get("is_correct") else s.get("unsafe_examples"))
                           for o in opts):
                    print(f"  reject {s.get('constraint_id')}: an option is not one of the seed's safe or "
                          f"unsafe drugs")
                    continue
                qid = f"{PATIENT}/axis{AXIS}/{len(candidates):03d}"
                opts = shuffled(opts, qid)
                safe = next(o for o in opts if o.get("is_correct"))
                candidates.append({
                    "qid": qid, "patient_id": PATIENT,
                    "axis": "3_drug_safety", "question": q.get("question"), "options": opts,
                    "answer": safe.get("content"), "answer_type": "multiple_choice", "grounding": "record",
                    "constraint_id": s.get("constraint_id"),
                    "evidence_visits": [gt_visit] if gt_visit else [], "gt_visit": gt_visit, "gt_date": gt_date,
                    "question_appearance_window": {"from": gt_date, "to": plast},
                    "derivation": {
                        "evidence": [{"visit": gt_visit, "date": gt_date, "field": "evidence_chain",
                                      "atom": s.get("constraint_text")}],
                        "steps": [f"Standing constraint: {s.get('constraint_text')!r}.",
                                  f"Of the 4 options only {safe.get('content')!r} is appropriate given it.",
                                  safe.get("explanation", "")],
                        "gold": safe.get("content"),
                        "rivals": [{"value": o.get("content"), "why_wrong": o.get("explanation")}
                                   for o in opts if not o.get("is_correct")],
                        "operands": {"implicated_class": s.get("implicated_class"),
                                     "constraint_text": s.get("constraint_text"),
                                     "first_line_trap": s.get("first_line_trap")},
                        "derived": False},
                    "provenance": f"LLM-discovered constraint ({ENGINE}); appropriate/inappropriate split "
                                  f"turns on recalling {s.get('constraint_id')}",
                    "trap_design": q.get("trap_design"),
                    "axis_specific": {"constraint_id": s.get("constraint_id"),
                                      "required_patient_info": s.get("required_patient_info")},
                })
        pg.tick(ok=True)

    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"  drug/constraint timelines scanned : {len(drug_tl)}")
    print(f"  constraints discovered (LLM)      : {len(seeds)}")
    print(f"  MC questions built                : {len(candidates)}")
    for c in candidates[:3]:
        print(f"      [{c['constraint_id']}] {str(c['question'])[:70]}  -> {c['answer']}")
    print(f"  CANDIDATES: {len(candidates)}  -> {BASE}/_discover_axis{AXIS}.json")


def shuffled(options, qid):
    """The options in an order seeded by the question id, relabelled A, B, C, ... in that order."""
    out = [dict(o) for o in options]
    random.Random(qid).shuffle(out)
    for k, o in enumerate(out):
        o["label"] = chr(ord("A") + k)
    return out


def _names(option, drugs):
    """Whether an option names one of the drugs (case-insensitive; either text may contain the other)."""
    text = str(option.get("content") or "").strip().lower()
    return bool(text) and any(d and (d in text or text in d) for d in (str(x).strip().lower() for x in drugs or []))


def _anchor(seed, tl, chain):
    """Anchor a constraint at the earliest visit of its source entity's timeline; (None, None) when the
    entity is not in the evidence chain."""
    ent = seed.get("source_entity")
    for t in tl:
        if t.get("entity") == ent and t.get("segments"):
            vf = (t["segments"][0].get("valid_from") or {})
            if vf.get("visit_number") in chain:
                return vf["visit_number"], vf.get("date")
    return None, None


def synthesize(cli, chain, by_v, forbidden, n, batch):
    """Author a standing drug constraint with four options; the gold is the safe member (options[].is_correct)."""
    out, vs = [], sorted(by_v)
    for i in range(n):
        host = vs[int(len(vs) * (0.15 + 0.4 * i / max(1, n)))]
        dept = by_v[host]["department"]
        host_date = by_v[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 one drug-safety multiple-choice question for a medical-memory benchmark. Return only JSON.",
            f"{patient}\n{context}\n"
            f"Visit: v{host} · {host_date} · {dept}\n\n"
            "Write a standing drug constraint, such as an allergy, an intolerance or an interaction, that\n"
            "makes three drugs a clinician might choose unsafe for this patient while one clinically\n"
            "reasonable alternative stays safe. Then write four options: the three unsafe drugs and the safe\n"
            "one. A model without memory must be drawn to an unsafe first-line choice.\n\n"
            "Do not name, describe or hint at the constraint in the question or the options; the model must\n"
            "recall it from memory. Never use the words \"safe\", \"unsafe\", \"risk\", \"dangerous\" or \"avoid\" in\n"
            "the question or the option text; the options are plain drug names.\n\n"
            "Do not use these drugs or classes, which existing answer keys already track:\n"
            + forbidden_list +
            "\n\nReturn JSON in this form:\n"
            '{"atom_text": "<terse chart fact recording the constraint>",\n'
            ' "entity": "<the constraint>",\n'
            ' "topic": "<slug>",\n'
            ' "question": "<the patient asks which drug to start, without listing the options>",\n'
            ' "options": [{"label": "A", "content": "<drug>", "is_correct": false,\n'
            '              "explanation": "<why it is or is not appropriate>"},\n'
            "             ...]}")
        if not (has_text(r) and has_text(r, "atom_text") and has_text(r, "entity") and r.get("options")):
            print(f"    axis3[{i}] rejected: the reply lacks the question, its options, the fact or its entity")
            continue
        opts = r["options"]
        if not (isinstance(opts, list) and all(isinstance(o, dict) and has_text(o, "content") for o in opts)):
            print(f"    axis3[{i}] rejected: an option is not an object with a drug name")
            continue
        if lists_options(r["question"], opts):
            print(f"    axis3[{i}] rejected: the question text lists the options")
            continue
        n_correct = sum(1 for o in opts if o.get("is_correct"))
        if n_correct != 1:
            print(f"    axis3[{i}] rejected: {n_correct} correct options, need exactly 1")
            continue
        qid = f"{PATIENT}/axis3/S{len(out):03d}"
        opts = shuffled(opts, qid)
        safe = next(o for o in opts if o.get("is_correct"))
        atom = {"atom_id": f"{batch}/{len(out):03d}", "batch_id": batch, "axis": "3",
                "visit_number": host, "date": by_v[host]["date"], "department": dept,
                "field": "doctor_extra_notes", "text": r["atom_text"], "entities": [r["entity"]],
                "topic": r.get("topic"), "anchors_axis": "3", "anchors_on": []}
        out.append({"atom": atom, "q": {
            "qid": qid, "patient_id": PATIENT, "axis": "3_drug_safety",
            "question": r["question"], "options": opts, "answer": safe["content"],
            "answer_type": "multiple_choice", "grounding": f"synthetic:{batch}",
            "evidence_visits": [host], "gt_visit": host, "gt_date": by_v[host]["date"],
            "question_appearance_window": {"from": by_v[host]["date"], "to": by_v[vs[-1]]["date"]},
            "derivation": {"evidence": [{"visit": host, "date": by_v[host]["date"],
                                         "field": "doctor_extra_notes", "atom": r["atom_text"]}],
                           "steps": [f"Standing constraint recorded v{host}: {r['atom_text']!r}.",
                                     f"Of the four options only {safe['content']} is compatible with it.",
                                     f"{safe.get('explanation') or ''}"],
                           "gold": safe["content"],
                           "rivals": [{"value": o.get("content"), "why_wrong": o.get("explanation")}
                                      for o in opts if not o.get("is_correct")],
                           "operands": {"constraint_text": r["atom_text"], "constraint_id": r.get("topic")},
                           "derived": False},
            "provenance": f"synthetic drug-safety constraint ({batch}); only the recalled constraint "
                          f"distinguishes the safe member",
            "axis_specific": {"constraint_id": r.get("topic"), "depends_on_atoms": [atom["atom_id"]]},
        }})
    return out


if __name__ == "__main__":
    discover()