"""Phase 3 Step 5 — re-measure a trained LoRA on the validation set and compare to base-8B.

Builds the SAME validation prompts used for the base measurement (apples-to-apples), runs them through
the Modal base+adapter batch function, parses, and computes Cohen's κ vs the Kim ground truth.

Run (from a Modal-authed machine):  python -m eval.measure_lora interaction
                                     python -m eval.measure_lora critical
"""
from __future__ import annotations

import json
import os
import sys

from eval import kappa as K
from eval.prompts_v3 import build_interaction_prompt
from eval.step_critical import build_critical_prompt, CE
from eval.step_c import goal_prompt, decomp_prompt
from prompt_card.scoring import observable_axes as OA

# base-8B κ per axis/category (from the cascade) for the lock-vs-keep comparison
BASE_KAPPA = {"interaction": 0.320, "goal_stated": 0.226, "decomposition": 0.261, "re_questioning": 0.058}
# single-feature axes: (prompt builder over the scored turns, the feature name, which GT axis holds it)
SINGLE = {"goal_stated": (goal_prompt, "goal_stated", "input_quality"),
          "decomposition": (decomp_prompt, "decomposition", "technique")}


def build(axis, gt, convs):
    prompts, truth = [], []
    for r in gt:
        conv = convs[r["id"]]
        ut = OA._user_turns(conv)
        if axis == "interaction":
            for row in r["interaction"]:
                i = int(row["turn"][1:]) - 1
                prompts.append(build_interaction_prompt(ut[i - 1], ut[i]))
                truth.append(set(["refinement_attempt"]) if row["refinement"] else set())
        elif axis == "critical":
            for row in r["critical"]:
                i = int(row["turn"][1:]) - 1
                prompts.append(build_critical_prompt(OA._prev_assistant(conv, i), ut[i]))
                truth.append(set(row["types"]))
        elif axis in SINGLE:
            builder, feat, gt_axis = SINGLE[axis]
            key = "features" if gt_axis == "input_quality" else "types"
            for row in r[gt_axis]:
                i = int(row["turn"][1:]) - 1
                prompts.append(builder(ut[i]))
                truth.append({feat} if feat in row[key] else set())
    return prompts, truth


def main(axis, adapter=""):
    from modal_eval_lora import app, evaluate
    gt = K.load_gt(); convs = K.load_convs()
    prompts, truth = build(axis, gt, convs)
    tag = f"{axis} (adapter={adapter or axis})"
    print(f"[lora-eval] {tag}: {len(prompts)} prompts -> Modal base+adapter ...", flush=True)
    with app.run():
        resp = evaluate.remote(axis, prompts, adapter=adapter)

    if axis == "interaction":
        fields = ("refinement_attempt",)
    elif axis in SINGLE:
        fields = (SINGLE[axis][1],)
    else:
        fields = CE
    fail = 0
    preds = []
    for rtext in resp:
        d = OA.parse(rtext, fields)
        if d is None:
            fail += 1; d = {}
        preds.append({f for f in fields if d.get(f)})

    print(f"[lora-eval] {axis}: parse_fail {fail}/{len(prompts)}")
    if axis == "interaction" or axis in SINGLE:
        feat = "refinement_attempt" if axis == "interaction" else SINGLE[axis][1]
        yt = [int(feat in s) for s in truth]
        yp = [int(feat in s) for s in preds]
        k = K.cohen_kappa(yt, yp)
        base_k = BASE_KAPPA[axis]
        print(f"  {axis} κ: base {base_k:+.3f} -> LoRA {k:+.3f}  [{K.binary_counts(yt, yp)}]")
        result = {"axis": axis, "lora_kappa": k, "base_kappa": base_k}
    else:
        base = json.load(open(os.path.join(os.path.dirname(__file__), "_cache", "critical.json")))
        per = {}
        for t in CE:
            yt = [int(t in s) for s in truth]; yp = [int(t in s) for s in preds]
            per[t] = K.cohen_kappa(yt, yp)
            print(f"  {t:24} base {base['per_type'][t]:+.3f} -> LoRA {(per[t] if per[t] is not None else float('nan')):+.3f}")
        valid = [v for v in per.values() if v is not None]
        head = sum(valid) / len(valid)
        print(f"  per-type mean: base {base['headline']:+.3f} -> LoRA {head:+.3f}")
        result = {"axis": axis, "lora_per_type": per, "lora_headline": head, "base_headline": base["headline"]}

    result["adapter"] = adapter or axis
    json.dump(result, open(os.path.join(os.path.dirname(__file__), "_cache", f"lora_{adapter or axis}.json"), "w"),
              indent=1, default=str)
    print(f"  decision: {'LOCK LoRA' if _wins(result, axis) else 'KEEP base'}")


def _wins(result, axis):
    # Phase-5 rule: LoRA κ ≥ base + 0.1 → use LoRA.
    if "lora_kappa" in result:
        return (result["lora_kappa"] or 0) >= (result["base_kappa"] or 0) + 0.1
    return (result["lora_headline"] or 0) >= (result["base_headline"] or 0) + 0.1


if __name__ == "__main__":
    # usage: python -m eval.measure_lora <axis> [adapter_dir_name]
    main(sys.argv[1] if len(sys.argv) > 1 else "interaction",
         sys.argv[2] if len(sys.argv) > 2 else "")