"""Evaluate the fine-tuned TinyBERT model on a held-out test set.

Standalone script — only needs `transformers`+`torch`. Expects a JSONL file
where each line is `{"messages": [system, user, assistant]}` (same chat
format used by the training data on GitHub), with the assistant content
being the gold JSON string. See:
https://github.com/innerkorehq/indian-address-parser

Usage:
    python evaluate_tinybert.py test.jsonl
    python evaluate_tinybert.py test.jsonl --n 100
"""
from __future__ import annotations

import argparse
import json
from collections import defaultdict

from inference_tinybert import ALL_FIELDS, extract_fields, load_model, MAX_LENGTH


def predict(model, tokenizer, raw_address: str) -> dict:
    import torch

    enc = tokenizer(
        raw_address, return_tensors="pt", return_offsets_mapping=True,
        truncation=True, max_length=MAX_LENGTH,
    )
    offsets = enc.pop("offset_mapping")[0].tolist()
    with torch.no_grad():
        out = model(**enc)
    pred_ids = out.logits[0].argmax(-1).tolist()
    return extract_fields(raw_address, offsets, pred_ids)


def compute_metrics(pairs: list[tuple[dict, dict]]) -> dict:
    n_total = len(pairs)
    field_correct = defaultdict(int)
    field_present_gold = defaultdict(int)
    field_recalled = defaultdict(int)
    overall_exact = 0

    for gold, pred in pairs:
        all_match = True
        for field in ALL_FIELDS:
            g = (gold.get(field) or "").strip().lower()
            p = (pred.get(field) or "").strip().lower()
            if g:
                field_present_gold[field] += 1
                if g == p:
                    field_recalled[field] += 1
                else:
                    all_match = False
            if g == p:
                field_correct[field] += 1
        if all_match:
            overall_exact += 1

    accs = {f: (field_correct[f] / n_total if n_total else 0) for f in ALL_FIELDS}
    return {
        "n_total": n_total,
        "overall_exact": overall_exact,
        "field_correct": field_correct,
        "field_present_gold": field_present_gold,
        "field_recalled": field_recalled,
        "field_accuracy": accs,
        "mean_field_accuracy": sum(accs.values()) / len(accs),
    }


def print_report(m: dict):
    n_total = m["n_total"]
    print(f"\nResults")
    print(f"Samples evaluated: {n_total}")
    print(f"Overall exact match (all present fields): {m['overall_exact']}/{n_total} ({100*m['overall_exact']/max(1,n_total):.1f}%)\n")
    print(f"{'Field':20s}  {'Accuracy':>10s}  {'Recall':>10s}  {'Gold presence':>14s}")
    print("-" * 60)
    for field in ALL_FIELDS:
        acc = m["field_accuracy"][field]
        n_gold = m["field_present_gold"][field]
        rec = m["field_recalled"][field] / n_gold if n_gold else float("nan")
        pres = n_gold / n_total if n_total else 0
        rec_str = f"{100*rec:.1f}%" if n_gold else "  n/a"
        print(f"{field:20s}  {100*acc:>9.1f}%  {rec_str:>10s}  {100*pres:>13.1f}%")
    print(f"\n{'Mean field accuracy':20s}  {100*m['mean_field_accuracy']:>9.1f}%")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("test_file", help="JSONL file with {'messages': [system, user, assistant]} per line")
    p.add_argument("--model", default=".", help="Model dir or HF repo id (default: current directory)")
    p.add_argument("--n", type=int, default=None)
    args = p.parse_args()

    with open(args.test_file, encoding="utf-8") as f:
        samples = [json.loads(line) for line in f]
    if args.n:
        samples = samples[: args.n]

    model, tokenizer = load_model(args.model)

    pairs = []
    for i, sample in enumerate(samples):
        msgs = sample["messages"]
        raw_address = msgs[1]["content"].replace("Parse this address:\n", "", 1)
        gold = json.loads(msgs[2]["content"])
        pred = predict(model, tokenizer, raw_address)
        pairs.append((gold, pred))
        if (i + 1) % 25 == 0:
            print(f"  ...{i + 1}/{len(samples)} evaluated")

    print_report(compute_metrics(pairs))


if __name__ == "__main__":
    main()