"""
Reasoning-mode evaluation
=========================
Every run up to this point had thinking DISABLED (enable_thinking=False),
because the training targets contain no chain of thought. In other words the
very capability the LLM was chosen for was never engaged.

This script measures two things at once:

1) OPEN-ENDED QUESTIONS — questions with no rule-based answer. The training
   tasks could all be answered from a template. These cannot: differential
   diagnosis, pattern interpretation, synthesis of conflicting findings. Every
   answer has to come from the base model's pretraining.

2) THINKING ON vs OFF — the same question in both modes. Does the chain of
   thought improve the answer, or merely lengthen it?

With --compare-base the LoRA-free base model is run as well. If the base model
reasons well and the adapted one does not, that quantifies the damage
catastrophic forgetting did to reasoning. (The base model cannot read the soft
tokens; this comparison is about LANGUAGE and REASONING, not diagnostic
accuracy.)

Note: the questions are in Turkish because that is the language the adapter was
trained in and the language the assistant answers in.

Run:
    python eval_thinking.py \
        --features features.pt --text dataset.json \
        --projector projector.pt --n 2
"""
from __future__ import annotations
import argparse
import inspect
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from hfx_runtime import Projector, build_examples, SENTINEL

INSTR_MARKER_CANDIDATES = ["Yalnızca bu soruya cevap ver",
                           "Yukarıdaki MRI değerlendirmesine",
                           "Sınıflandırma (CN / MCI / AD) nedir"]

# Questions with NO rule-based answer — nothing in the training data covers them.
OPEN_QUESTIONS = [
    ("PATTERN",
     "Bu hastadaki bölgesel atrofi dağılımı hangi klinik tabloyu düşündürür? "
     "Tipik amnestik Alzheimer paterniyle uyumlu mu, değilse neden?"),
    ("DIFFERENTIAL",
     "Alzheimer dışında hangi tanılar düşünülmeli ve bunları ayırt etmek için "
     "ne gerekir?"),
    ("CONFLICT",
     "Bulgular arasında birbiriyle çelişen bir taraf var mı? Varsa bunu nasıl "
     "yorumlarsın?"),
]


def strip_body(prompt: str) -> str:
    """Keep the patient-data part of the prompt, drop the instruction paragraph."""
    idx = [prompt.find(m) for m in INSTR_MARKER_CANDIDATES]
    idx = [i for i in idx if i > 0]
    return prompt[:min(idx)].rstrip() if idx else prompt


def split_thought(text: str) -> tuple:
    """
    Separate the chain of thought from the final answer.

    Thinking is marked with channel control tokens whose exact spelling varies by
    release, so several patterns are tried; if none match, the whole text is
    treated as the answer.
    """
    for a, b in (("<|channel>thought", "<channel|>"),
                 ("<think>", "</think>"),
                 ("<|thought|>", "<|/thought|>")):
        if a in text:
            head, _, rest = text.partition(a)
            thought, _, answer = rest.partition(b)
            return thought.strip(), (head + answer).strip()
    return "", text.strip()


@torch.no_grad()
def ask(model, tok, embed_layer, projector, ex, device, supported, question,
        thinking: bool, max_new_tokens: int, use_soft: bool = True) -> str:
    body = strip_body(ex["prompt"])
    ptxt = body + "\n\n" + question + "\nYanıtını TÜRKÇE yaz."
    if SENTINEL not in ptxt:
        ptxt = SENTINEL + "\n" + ptxt

    msgs = [{"role": "user", "content": ptxt}]
    kw = dict(tokenize=False, add_generation_prompt=True)
    try:
        full = tok.apply_chat_template(msgs, enable_thinking=thinking, **kw)
    except TypeError:
        full = tok.apply_chat_template(msgs, **kw)

    pre_txt, post_txt = full.split(SENTINEL, 1)
    ids_pre = tok(pre_txt, add_special_tokens=False, return_tensors="pt").input_ids.to(device)
    ids_post = tok(post_txt, add_special_tokens=False, return_tensors="pt").input_ids.to(device)
    e_pre, e_post = embed_layer(ids_pre), embed_layer(ids_post)

    if use_soft and projector is not None:
        soft = projector(ex["feat"].unsqueeze(0).to(device)).to(e_pre.dtype)
        embeds = torch.cat([e_pre, soft, e_post], dim=1)
        n_soft, off = soft.size(1), e_pre.size(1)
    else:
        embeds = torch.cat([e_pre, e_post], dim=1)
        n_soft, off = 0, 0

    attn = torch.ones(embeds.shape[:2], dtype=torch.long, device=device)
    gk = {"inputs_embeds": embeds, "attention_mask": attn,
          "max_new_tokens": max_new_tokens, "do_sample": False,
          "repetition_penalty": 1.1}
    if n_soft and "mm_token_type_ids" in supported:
        mm = torch.zeros(embeds.shape[:2], dtype=torch.long, device=device)
        mm[0, off:off + n_soft] = 1
        gk["mm_token_type_ids"] = mm
    out = model.generate(**gk)
    return tok.decode(out[0], skip_special_tokens=False)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--features", required=True)
    ap.add_argument("--text", required=True)
    ap.add_argument("--projector", default="projector.pt",
                    help="projector checkpoint; LoRA adapter at <projector>_lora/")
    ap.add_argument("--model", default=None)
    ap.add_argument("--split", default="test")
    ap.add_argument("--n", type=int, default=2)
    # The chain of thought alone eats ~600 tokens; at a 600-token budget the
    # final answer was cut off before it began.
    ap.add_argument("--max-new-tokens", type=int, default=2000)
    ap.add_argument("--compare-base", action="store_true",
                    help="also run the base model without LoRA (reasoning comparison)")
    ap.add_argument("--no-4bit", dest="four_bit", action="store_false", default=True)
    args = ap.parse_args()

    import transformers
    from transformers import AutoTokenizer
    AutoCls = next(getattr(transformers, n) for n in
                   ("AutoModelForConditionalGeneration", "AutoModelForImageTextToText",
                    "AutoModelForCausalLM") if hasattr(transformers, n))

    device = "cuda" if torch.cuda.is_available() else "cpu"
    ck = torch.load(args.projector, map_location="cpu", weights_only=False)
    model_id = args.model or ck["model_id"]
    data, d = build_examples(args.features, args.text)
    class_names = list(d["class_names"])

    tok = AutoTokenizer.from_pretrained(model_id)
    load_kw = dict(device_map={"": 0} if device == "cuda" else None)
    if args.four_bit:
        from transformers import BitsAndBytesConfig
        load_kw["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
    try:
        base = AutoCls.from_pretrained(model_id, dtype=torch.bfloat16, **load_kw)
    except TypeError:
        base = AutoCls.from_pretrained(model_id, torch_dtype=torch.bfloat16, **load_kw)
    supported = set(inspect.signature(base.forward).parameters)
    embed_layer = base.get_input_embeddings()

    model = base
    lora_dir = args.projector + "_lora"
    if ck.get("lora") and os.path.isdir(lora_dir):
        from peft import PeftModel
        model = PeftModel.from_pretrained(base, lora_dir)
        print(f"[lora] adapter loaded: {lora_dir}")
    model.eval()

    projector = Projector(ck["in_dim"], ck["hidden"], ck["n_tokens"],
                          target_norm=ck.get("target_norm")).to(device)
    projector.load_state_dict(ck["projector"]); projector.eval()

    seen, items = set(), []
    for e in data[args.split]:
        if e["ptid"] in seen:
            continue
        seen.add(e["ptid"]); items.append(e)
        if len(items) >= args.n:
            break

    for ex in items:
        print("\n" + "=" * 74)
        print(f"PATIENT {ex['ptid']}  head={ex['head']}  true={class_names[ex['label']]}")
        print("=" * 74)
        for tag, q in OPEN_QUESTIONS:
            print(f"\n### [{tag}] {q}")
            for think in (False, True):
                raw = ask(model, tok, embed_layer, projector, ex, device,
                          supported, q, think, args.max_new_tokens)
                thought, answer = split_thought(raw)
                n_tok = len(tok(raw, add_special_tokens=False).input_ids)
                print(f"\n--- thinking={'ON' if think else 'OFF'} "
                      f"({n_tok} tokens"
                      f"{', thought ' + str(len(thought)) + ' chars' if thought else ''}) ---")
                if thought:
                    print("[thought] " + thought[:500])
                print(answer[:800])

            if args.compare_base:
                # LoRA disabled: how does the base model reason?
                # (no soft tokens — the base model cannot read them)
                with model.disable_adapter() if hasattr(model, "disable_adapter") \
                        else torch.no_grad():
                    raw = ask(model, tok, embed_layer, None, ex, device,
                              supported, q, True, args.max_new_tokens, use_soft=False)
                t2, a2 = split_thought(raw)
                print("\n--- BASE MODEL (LoRA off, no soft tokens, thinking ON) ---")
                if t2:
                    print("[thought] " + t2[:400])
                print(a2[:800])

    print("\n" + "=" * 74)
    print("  What to look for:")
    print("   · Is the thinking=ON answer better than thinking=OFF, or just")
    print("     longer? (length is not quality)")
    print("   · Does the chain of thought rest on the given measurements, or")
    print("     drift into general medical knowledge? The latter is ungrounded.")
    print("   · With --compare-base: if the base model reasons fluently and the")
    print("     adapted one cannot, the forgetting cost has been measured.")
    print("=" * 74)


if __name__ == "__main__":
    main()