#!/usr/bin/env python3
"""Axis 6 (temporal localization and trend): LLM extraction with deterministic computation.

The LLM reads each candidate timeline and returns its numeric series. classify() computes the trend
shape from the series; extremum questions take the highest or lowest reading. Golds are formatted from
PROSE_TEMPLATE_* with the computed operands. synthesize() authors synthetic questions.

  python -m pipeline.questions.axis6     discovery -> _discover_axis6.json
"""
from __future__ import annotations

import json
import random
from datetime import date, timedelta
from pathlib import Path

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

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

AXIS = 6
SHAPE_VOCAB = ["steadily increased", "steadily decreased", "increased then decreased",
               "decreased then increased", "remained essentially unchanged"]

PROSE_TEMPLATE_EXTREMUM_Q = ("Across the whole record, what was this patient's {extremum_kind} charted "
                             "{series_entity} value, and when was it measured?")
PROSE_TEMPLATE_EXTREMUM_A = ("The {extremum_kind} charted {series_entity} was {extremum_num:g}{unit_sp} "
                             "on {extremum_date}.")
PROSE_TEMPLATE_MC_Q = ("Between {range_from} and {range_to}, how did this patient's {series_entity} "
                       "change over the readings taken in that period?")
PROSE_TEMPLATE_MC_A = "Over {range_from} to {range_to}, {series_entity} {correct_option_text}."


def sign(x):
    return 1 if x > 0 else (-1 if x < 0 else 0)


def classify(ns):
    """Trend shape of a series, computed deterministically; None rejects the series."""
    lo, hi, med = min(ns), max(ns), sorted(ns)[len(ns) // 2]
    if med != 0 and (hi - lo) / abs(med) <= 0.05:
        return "remained essentially unchanged"
    d = [b - a for a, b in zip(ns, ns[1:])]
    nz = [sign(x) for x in d if x != 0]
    if not nz:
        return "remained essentially unchanged"
    flips = sum(1 for a, b in zip(nz, nz[1:]) if a != b)
    if flips == 0:
        return "steadily increased" if nz[0] > 0 else "steadily decreased"
    if flips == 1:
        return "increased then decreased" if nz[0] > 0 else "decreased then increased"
    return None


def best_window(points):
    nums = [p["num"] for p in points]
    for length in range(len(nums), 2, -1):
        for i in range(0, len(nums) - length + 1):
            sh = classify(nums[i:i + length])
            if sh is not None:
                return (i, i + length, sh)
    return None


def mid(d1, d2):
    a, b = date.fromisoformat(d1), date.fromisoformat(d2)
    return (a + (b - a) / 2).isoformat()


EXTRACT_SYS = (
    "You read one entity's timeline from a synthetic longitudinal medical chart and extract its\n"
    "numeric readings. Charts are terse, and their format varies across patients and departments.\n"
    "Return only JSON.")

EXTRACT_SCHEMA = """Return JSON in this form:
{"is_numeric_series": <true|false>,
 "measurand": "<what is actually measured, normalized, e.g. 'Total Cholesterol' or 'Body weight'>",
 "unit": "<unit, or '' if none>",
 "points": [{"date": "YYYY-MM-DD", "visit": <int>, "num": <number>, "raw": "<the source text>"}],
 "aliases": ["<every name or abbreviation this chart uses for this measurand>"],
 "reject_reason": "<why, if is_numeric_series is false>"}"""


def extract_series(cli, t):
    """LLM extraction of the numeric readings of one timeline.

    Trailing parentheticals are dropped; dates, censored values ("< 6") and readings that mix two measurands
    are rejected.
    """
    segs = [{"value": s.get("value"), "date": (s.get("valid_from") or {}).get("date"),
             "visit": (s.get("valid_from") or {}).get("visit_number")} for s in t.get("segments") or []]
    entity, department, aspect = t.get("entity"), t.get("department"), t.get("aspect")
    segment_lines = "\n".join(f"{s['date']} v{s['visit']}: {s['value']!r}" for s in segs)
    usr = (
        f"Entity: {entity}\nDepartment: {department}\nAspect: {aspect}\n\n"
        "Segments, each one value and the date it was first recorded:\n" + segment_lines + "\n\n"
        "Extract the numeric readings of one measurand.\n\n"
        "Reject the series (is_numeric_series false) if any of these holds:\n"
        "- The values are not numeric measurements but impressions, findings or narrative.\n"
        "- More than one measurand is charted here, for example '<lab1>/<lab2> <value1>/<value2>', which\n"
        "  is ambiguous.\n"
        "- The values are censored, such as '< 6' or '> 60', and cannot be ordered.\n"
        "- Fewer than 3 readings carry a real number on distinct dates.\n"
        "- A number is a date or a lab accession number, not a reading.\n\n"
        "Take the number that is the reading, and ignore trailing text in parentheses, lab codes, embedded\n"
        "dates, reference ranges and targets such as 'goal LDL < 100'. Read a value in scientific notation,\n"
        "such as '1.2x10^3', as a number. Give one reading per date, and use the segment's own date and\n"
        "visit.\n\n"
        + EXTRACT_SCHEMA)
    r = cli.complete_json(EXTRACT_SYS, usr)
    if not r.get("is_numeric_series"):
        return None, r.get("reject_reason", "not a numeric series")
    pts = [p for p in (r.get("points") or []) if isinstance(p.get("num"), (int, float)) and p.get("date")]
    seen, uniq = set(), []
    for p in sorted(pts, key=lambda x: x["date"]):
        if p["date"] in seen:
            continue
        seen.add(p["date"])
        uniq.append(p)
    if len(uniq) < 3:
        return None, f"only {len(uniq)} usable readings"
    return {"entity": r.get("measurand") or t.get("entity"), "department": t.get("department"),
            "unit": (r.get("unit") or "").strip(), "points": uniq,
            "aliases": [a.lower() for a in (r.get("aliases") or [])]}, None


UNIQ_SYS = ("You check whether a claimed extreme value is unique in a patient's chart. Return only JSON.")
LINES_PER_CHECK = 70


def extremum_is_unique(cli, s, kind, ex_num, ex_date, hay_rows):
    """Whether the extreme reading of a series is unique, checked against every chart line naming the
    measurand, LINES_PER_CHECK lines per call; the first group that finds a higher (or lower) value or an
    equal value on another date decides."""
    names = ", ".join(sorted({s["entity"], *s.get("aliases", [])}))
    rel = [f"  {d} v{v}: {tx}" for v, d, tx in hay_rows
           if any(n and n.lower() in tx.lower() for n in {s["entity"], *s.get("aliases", [])})]
    for i in range(0, len(rel), LINES_PER_CHECK):
        unique, violations = _unique_in(cli, names, kind, ex_num, ex_date, rel[i:i + LINES_PER_CHECK])
        if not unique:
            return False, violations
    return True, []


def _unique_in(cli, names, kind, ex_num, ex_date, rel):
    """One uniqueness call over a group of chart lines."""
    chart_lines = "\n".join(rel)
    beyond = "higher" if kind == "highest" else "lower"
    r = cli.complete_json(
        UNIQ_SYS,
        f"Measurand: {names}\nClaim: the {kind} charted value is {ex_num:g} on {ex_date}.\n\n"
        "Chart lines that mention it:\n" + chart_lines + "\n\n"
        "Each line starts with the date of the visit that recorded it. Ignore reference ranges and targets.\n"
        "Is the claim true, and is that date unique?\n"
        f"- A line showing a {beyond} value anywhere makes the claim false.\n"
        "- A separate measurement with an equal value on another date makes it not unique. A reading\n"
        "  copied forward into a later visit, with the same value and its original test date or no date\n"
        "  of its own, is the same reading and is fine, because charts copy readings forward.\n\n"
        "Return JSON in this form:\n"
        '{"unique": <true|false>,\n'
        ' "violations": [{"date": "...", "value": <n>, "why": "beats|ties", "line": "..."}]}')
    return bool(r.get("unique")), r.get("violations") or []


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

    ts = [json.loads(l) for l in open(f"{BASE}/evidence_chain.jsonl") if l.strip()]
    chain = [json.loads(l) for l in open(f"{BASE}/event_chain.jsonl") if l.strip()]
    LAST_V = max(r["visit_number"] for r in chain)
    hay_rows = [(r["visit_number"], r["date"], str(a))
                for r in chain for f in FIELDS + ["icd_notes"] for a in (r.get(f) or [])]
    hay_rows += [(r["visit_number"], r["date"], f"{k}={val}")
                 for r in chain for k, val in (r.get("vitals") or {}).items()]

    cand = [t for t in ts if len(t.get("segments") or []) >= 3]
    series, rejects = [], []
    with Progress("axis6 extract (LLM)", total=len(cand), out=f"{BASE}/progress_axis6_discover.json") as pg:
        for t in cand:
            try:
                s, why = extract_series(cli, t)
            except Exception as e:
                s, why = None, f"error: {str(e)[:80]}"
            if s:
                series.append(s)
            else:
                rejects.append((t.get("entity"), why))
            pg.tick(ok=bool(s))

    candidates = []
    for s in series:
        pts, nums = s["points"], [p["num"] for p in s["points"]]
        ent, unit_sp = s["entity"], (f" {s['unit']}" if s["unit"] else "")
        evidence = [{"visit": p.get("visit"), "date": p["date"],
                     "field": "evidence_chain", "atom": p.get("raw")} for p in pts]

        for kind in ("highest", "lowest"):
            ex = max(nums) if kind == "highest" else min(nums)
            if nums.count(ex) != 1:
                continue
            k = nums.index(ex)
            ex_date, ex_visit = pts[k]["date"], pts[k].get("visit")
            try:
                ok, _ = extremum_is_unique(cli, s, kind, ex, ex_date, hay_rows)
            except Exception:
                ok = False
            if not ok:
                continue
            ops = {"extremum_kind": kind, "series_entity": ent, "extremum_num": ex,
                   "unit_sp": unit_sp, "extremum_date": ex_date, "n_points": len(pts),
                   "series_department": s["department"], "unit": s["unit"]}
            answer = PROSE_TEMPLATE_EXTREMUM_A.format(**ops)
            candidates.append({
                "qid": f"{PATIENT}/axis{AXIS}/{len(candidates):03d}", "patient_id": PATIENT,
                "axis": f"{AXIS}_{kind}_value_and_date",
                "question": PROSE_TEMPLATE_EXTREMUM_Q.format(**ops), "answer": answer,
                "answer_type": "open_exact", "grounding": "record",
                "evidence_visits": [p.get("visit") for p in pts],
                "gt_visit": ex_visit, "gt_date": ex_date,
                "question_appearance_window": {"from": ex_date, "to": max(r["date"] for r in chain)},
                "derivation": {"evidence": evidence,
                               "steps": [f"{ent} has {len(pts)} readings: {[(p['date'], p['num']) for p in pts]}.",
                                         f"{kind} = {ex:g}{unit_sp}, at {ex_date}.",
                                         f"Uniqueness confirmed against every chart line naming {ent}."],
                               "gold": answer, "rivals": [], "operands": ops, "derived": True},
                "provenance": f"LLM-extracted series ({ENGINE}); {kind} unique across the chart",
                "axis_specific": {"template": f"{kind}_value_and_date", "n_points": len(pts)},
            })

        w = best_window(pts)
        if w:
            i, j, shape = w
            wp = pts[i:j]
            rf = mid(pts[i - 1]["date"], wp[0]["date"]) if i > 0 else \
                (date.fromisoformat(wp[0]["date"]) - timedelta(days=15)).isoformat()
            rt = mid(wp[-1]["date"], pts[j]["date"]) if j < len(pts) else \
                (date.fromisoformat(wp[-1]["date"]) + timedelta(days=15)).isoformat()
            rnd = random.Random(f"{PATIENT}|{ent}|trend")
            opts = SHAPE_VOCAB[:]
            rnd.shuffle(opts)
            opts = opts[:4]
            if shape not in opts:
                opts[rnd.randrange(4)] = shape
            ops = {"range_from": rf, "range_to": rt, "series_entity": ent,
                   "correct_option_text": shape, "shape_class": shape,
                   "series_department": s["department"], "n_window": len(wp), "n_points": len(pts)}
            answer = PROSE_TEMPLATE_MC_A.format(**ops)
            nums_w = [p["num"] for p in wp]
            candidates.append({
                "qid": f"{PATIENT}/axis{AXIS}/{len(candidates):03d}", "patient_id": PATIENT,
                "axis": f"{AXIS}_trend_direction_mc",
                "question": PROSE_TEMPLATE_MC_Q.format(**ops), "answer": answer,
                "answer_type": "multiple_choice", "grounding": "record",
                "options": [{"label": "ABCD"[k], "content": o, "is_correct": o == shape}
                            for k, o in enumerate(opts)],
                "evidence_visits": [p.get("visit") for p in wp],
                "gt_visit": wp[-1].get("visit"), "gt_date": wp[-1]["date"],
                "question_appearance_window": {"from": rt, "to": max(r["date"] for r in chain)},
                "derivation": {"evidence": evidence[i:j],
                               "steps": [f"{ent} in {rf}..{rt}: {nums_w}.",
                                         f"Deltas: {[round(b-a,2) for a,b in zip(nums_w, nums_w[1:])]}.",
                                         f"classify() -> '{shape}'."],
                               "gold": answer, "rivals": [], "operands": ops, "derived": True},
                "provenance": f"LLM-extracted series ({ENGINE}); shape COMPUTED by classify()",
                "axis_specific": {"template": "trend_direction_mc", "shape_class": shape},
            })

    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}  LAST_V={LAST_V}  engine={ENGINE}")
    print(f"  timelines with >=3 segments : {len(cand)}")
    print(f"  LLM-extracted numeric series: {len(series)}")
    print(f"  rejected                    : {len(rejects)}")
    for e, why in rejects[:5]:
        print(f"      {str(e)[:38]:<38} {str(why)[:56]}")
    print(f"  CANDIDATES: {len(candidates)}  -> {BASE}/_discover_axis{AXIS}.json")


def synthesize(cli, chain, by_v, forbidden, n, batch, existing_series):
    """Author `n` numeric series whose values follow a fixed profile for the intended shape; one LLM call picks the
    measurands and baselines."""
    out, vs = [], sorted(by_v)
    SHAPES = ["increased then decreased", "decreased then increased", "steadily increased", "steadily decreased"]
    FRAC = {"increased then decreased": [0.2, 0.5, 0.85, 0.5, 0.25],
            "decreased then increased": [0.85, 0.5, 0.15, 0.5, 0.8],
            "steadily increased": [0.15, 0.4, 0.65, 0.9],
            "steadily decreased": [0.9, 0.65, 0.4, 0.15]}
    patient, context = patient_line(), patient_context(chain)
    tracked_list = "\n".join(f"- {e}" for e in existing_series[:50])
    forbidden_inline = ", ".join(forbidden[:25])
    r = cli.complete_json(
        "You choose recurring numeric lab tests for a trend question of a medical-memory benchmark.\n"
        "Return only JSON.",
        f"{patient}\n{context}\n\n"
        f"Choose {n} different recurring numeric lab tests that this patient's care might track over\n"
        "years. Give a plausible baseline value and unit for each; the trajectory is computed from it.\n\n"
        "Avoid lab tests that the chart already tracks, because a shared word or abbreviation makes both\n"
        "unusable:\n" + tracked_list +
        "\nAlso avoid these, which answer keys already track: " + forbidden_inline + "\n\n"
        "Return JSON in this form:\n"
        '{"labs": [{"label": "<chart name>", "unit": "<unit or \'\'>", "baseline": <number>,\n'
        '           "decimals": <0|1|2>, "entity": "<measurand>", "topic": "<slug>",\n'
        '           "department": "<one of the departments listed above>"}]}')
    labs = r.get("labs") or []
    if not isinstance(labs, list):
        print("    axis6: the reply's labs are not a list")
        labs = []
    depts = {x["department"] for x in chain}
    for i, lab in enumerate(labs[:n]):
        if not (isinstance(lab, dict) and all(isinstance(lab.get(k), str) and lab[k].strip()
                                              for k in ("label", "entity"))):
            print(f"    axis6[{i}] rejected: the lab lacks a label or an entity")
            continue
        try:
            base_v, dec = float(lab["baseline"]), int(lab.get("decimals", 1))
        except (KeyError, TypeError, ValueError):
            print(f"    axis6[{i}] rejected: the lab lacks a numeric baseline")
            continue
        shape = SHAPES[i % len(SHAPES)]
        frac = FRAC[shape]
        dept = next((d for d in depts if lab.get("department", "") and lab["department"] in d),
                    by_v[vs[i % len(vs)]]["department"])
        hosts = [v for v in vs if by_v[v]["department"] == dept][:len(frac)]
        if len(hosts) < 3:
            hosts = vs[i::max(1, len(vs) // len(frac))][:len(frac)]
        if len(hosts) < 3:
            continue
        unit = str(lab.get("unit") or "").strip()
        pts, atoms = [], []
        for k, v in enumerate(hosts):
            val = round(base_v * (0.5 + frac[k]), dec)
            txt = f"{lab['label']} {val:g} {unit}".strip()
            pts.append((v, by_v[v]["date"], val))
            atoms.append({"atom_id": f"{batch}/{i:02d}_{k:02d}", "batch_id": batch, "axis": "6",
                          "visit_number": v, "date": by_v[v]["date"], "department": by_v[v]["department"],
                          "field": "tests_and_medication", "text": txt, "entities": [lab["entity"]],
                          "topic": lab.get("topic"), "anchors_axis": "6", "anchors_on": []})
        nums = [p[2] for p in pts]
        gold = f"Over {pts[0][1]} to {pts[-1][1]}, {lab['entity']} {shape}."
        out.append({"atom": atoms[0], "atoms": atoms, "q": {
            "qid": f"{PATIENT}/axis6/S{i:03d}", "patient_id": PATIENT, "axis": "6_trend_direction",
            "question": (f"Between {pts[0][1]} and {pts[-1][1]}, how did this patient's {lab['entity']} "
                         f"change over the readings taken in that period?"),
            "answer": gold, "answer_type": "open_exact", "grounding": f"synthetic:{batch}",
            "evidence_visits": [p[0] for p in pts], "gt_visit": pts[-1][0], "gt_date": pts[-1][1],
            "question_appearance_window": {"from": pts[-1][1], "to": by_v[vs[-1]]["date"]},
            "derivation": {"evidence": [{"visit": p[0], "date": p[1], "field": "tests_and_medication",
                                         "atom": f"{lab['label']} {p[2]:g} {unit}".strip()} for p in pts],
                           "steps": [f"{lab['entity']} measured {len(pts)} times: {nums}.",
                                     f"Deltas: {[round(b-a,2) for a,b in zip(nums,nums[1:])]}.",
                                     f"Therefore the trajectory {shape}."],
                           "gold": gold, "rivals": [],
                           "operands": {"series_entity": lab["entity"], "series_department": dept,
                                        "unit": unit, "points": [{"date": p[1], "num": p[2]} for p in pts],
                                        "n_points": len(pts), "shape_class": shape},
                           "derived": True},
            "provenance": f"synthetic trend ({batch}); value sequence computed so the shape is exact "
                          f"(absolute values immaterial for a memory test)",
            "axis_specific": {"shape_class": shape, "depends_on_atoms": [x["atom_id"] for x in atoms]},
        }})
    return out


if __name__ == "__main__":
    discover()