#!/usr/bin/env python3
"""Step 4: show base + LoRA CAN perform capability C.

Loads the frozen base with the trained adapter attached and runs the SAME
held-out eval set from step 2. Accuracy jumps from near-zero to high. This is
the before/after that proves the adapter, and only the adapter, carries the
capability.
"""

import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from demo.common import ADAPTER_DIR, STATE_DIR, evaluate, load_base_plus_adapter, load_tokenizer
from demo.dataset import build_dataset


def main() -> int:
    print("=== Step 4: base + LoRA on capability C (expected: high) ===\n")
    if not (ADAPTER_DIR / "adapter_config.json").exists():
        print(f"  no adapter at {ADAPTER_DIR}; run step 3 first.")
        return 1

    _, held_out = build_dataset()
    tokenizer = load_tokenizer()
    model = load_base_plus_adapter(ADAPTER_DIR)
    accuracy, correct, total, samples = evaluate(model, tokenizer, held_out, show=4)

    print("  sample base+LoRA outputs:")
    for request, gold, out, ok in samples:
        print(f"    request : {request}")
        print(f"    expected: {gold}")
        print(f"    model   : {out!r}  [{'ok' if ok else 'wrong'}]\n")

    base_acc = None
    base_path = STATE_DIR / "eval_base.json"
    if base_path.exists():
        base_acc = json.loads(base_path.read_text())["accuracy"]

    print(f"  BASE+LoRA ACCURACY ON C: {accuracy:.3f}  ({correct}/{total})")
    if base_acc is not None:
        print(f"  (base alone was {base_acc:.3f} on the same held-out set)")

    (STATE_DIR / "eval_lora.json").write_text(
        json.dumps({"accuracy": accuracy, "correct": correct, "total": total}, indent=2)
    )

    if accuracy < 0.80:
        print("\nRESULT: WEAK - adapter accuracy below 0.80; consider more epochs/data.")
        return 1
    print("\nRESULT: PASS - the adapter carries capability C.")
    return 0


if __name__ == "__main__":
    sys.exit(main())