"""Score the model on a test file, with or without your LoRA adapter.

usage:
  python evaluate.py none                       # base model, prompt only
  python evaluate.py adapters                   # base model + your adapter
  python evaluate.py adapters data/wild.jsonl   # any test file
"""
import json
import sys

from mlx_lm import generate, load

MODEL = "mlx-community/Qwen2.5-1.5B-Instruct-4bit"
FIELDS = ["title", "day", "time", "duration_min", "people"]
SHOW_MISSES = 5


def main():
    adapter = sys.argv[1] if len(sys.argv) > 1 else "none"
    path = sys.argv[2] if len(sys.argv) > 2 else "data/test.jsonl"
    model, tok = load(MODEL, adapter_path=None if adapter == "none" else adapter)
    rows = [json.loads(line) for line in open(path)]

    valid = exact = shown = 0
    field_hits = dict.fromkeys(FIELDS, 0)
    for row in rows:
        # system + user only: the assistant turn is the answer key
        prompt = tok.apply_chat_template(
            row["messages"][:2], add_generation_prompt=True, tokenize=False)
        out = generate(model, tok, prompt=prompt, max_tokens=120).strip()
        gold = json.loads(row["messages"][2]["content"])
        try:
            pred = json.loads(out)
            valid += 1
        except json.JSONDecodeError:
            pred = {}
        if not isinstance(pred, dict):
            pred = {}

        hits = {f: pred.get(f) == gold[f] for f in FIELDS}
        for f in FIELDS:
            field_hits[f] += hits[f]
        if all(hits.values()):
            exact += 1
        elif shown < SHOW_MISSES:
            shown += 1
            print(f"\nmsg : {row['messages'][1]['content']}")
            print(f"got : {out}")
            print(f"want: {json.dumps(gold)}")

    n = len(rows)
    print(f"\nvalid JSON : {valid}/{n}")
    print(f"exact match: {exact}/{n}")
    for f in FIELDS:
        print(f"  {f:<13} {field_hits[f]}/{n}")


if __name__ == "__main__":
    main()