#!/usr/bin/env python3
"""Evaluate GLiNER checkpoints and external models on an eval.json file.

This script expects GLiNER-style records with:
  - tokenized_text: list[str]
  - ner: list[[start_token_idx, end_token_idx, label], ...]

It rebuilds each sample text from tokenized tokens (" ".join(tokens)),
runs batched inference, and computes exact-match entity metrics
(label + start_char + end_char):
  - overall precision / recall / F1
  - per-entity precision / recall / F1
  - counts (tp/fp/fn, predicted/true totals)

Examples:
  python scripts/eval_models_on_eval_json.py \
    --eval-json /mnt/doscratch/attempt_vanilla/data/eval.json \
    --model /mnt/doscratch/attempt_vanilla/gliner_logs/gliner-pii-large-full-v1/checkpoint-5500

  python scripts/eval_models_on_eval_json.py \
    --eval-json /mnt/doscratch/attempt_vanilla/data/eval.json \
    --checkpoints-root /mnt/doscratch/attempt_vanilla/gliner_logs/gliner-pii-large-full-v1 \
    --checkpoint-glob 'checkpoint-*'
"""

from __future__ import annotations

import argparse
import json
import re
import time
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import torch
from tqdm import tqdm

from gliner import GLiNER


Span = tuple[str, int, int]


@dataclass(frozen=True)
class Example:
    text: str
    gold_spans: frozenset[Span]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Evaluate GLiNER models on eval.json with overall and per-entity PRF."
    )
    parser.add_argument(
        "--eval-json",
        required=True,
        type=Path,
        help="Path to eval.json in GLiNER format.",
    )
    parser.add_argument(
        "--model",
        action="append",
        default=[],
        help=(
            "Model reference(s): local checkpoint/model directory or HF id. "
            "Repeatable."
        ),
    )
    parser.add_argument(
        "--checkpoints-root",
        type=Path,
        default=None,
        help="Directory that contains checkpoint-* folders to evaluate.",
    )
    parser.add_argument(
        "--checkpoint-glob",
        default="checkpoint-*",
        help="Glob for checkpoint discovery under --checkpoints-root.",
    )
    parser.add_argument(
        "--base-model",
        default="nvidia/gliner-PII",
        help="External baseline model id to include unless --skip-base-model is set.",
    )
    parser.add_argument(
        "--skip-base-model",
        action="store_true",
        help="Do not include --base-model in evaluation.",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=12,
        help="Inference batch size.",
    )
    parser.add_argument(
        "--threshold",
        type=float,
        default=0.5,
        help="Prediction threshold for GLiNER inference.",
    )
    parser.add_argument(
        "--flat-ner",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Use flat NER decoding (default: true).",
    )
    parser.add_argument(
        "--multi-label",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Allow multi-label decoding per span/token.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=0,
        help="Evaluate only first N examples (0 = all).",
    )
    parser.add_argument(
        "--device",
        default="auto",
        help="Device for model inference: auto|cpu|cuda|cuda:0 etc.",
    )
    parser.add_argument(
        "--case-sensitive-labels",
        action="store_true",
        help="Keep label casing as-is. Default normalizes labels to lowercase.",
    )
    parser.add_argument(
        "--output-json",
        type=Path,
        default=None,
        help="Optional path to write full metrics JSON.",
    )
    return parser.parse_args()


def normalize_label(label: str, case_sensitive: bool) -> str:
    value = str(label).strip()
    return value if case_sensitive else value.lower()


def token_offsets(tokens: list[str]) -> list[tuple[int, int]]:
    offsets: list[tuple[int, int]] = []
    cursor = 0
    for token in tokens:
        start = cursor
        end = cursor + len(token)
        offsets.append((start, end))
        cursor = end + 1
    return offsets


def parse_examples(
    data: list[dict],
    case_sensitive_labels: bool,
) -> tuple[list[Example], list[str], dict[str, int]]:
    examples: list[Example] = []
    all_labels: set[str] = set()
    skipped_bad_rows = 0
    skipped_bad_spans = 0

    for idx, row in enumerate(data):
        tokens = row.get("tokenized_text")
        if not isinstance(tokens, list) or not all(isinstance(t, str) for t in tokens):
            skipped_bad_rows += 1
            continue

        text = " ".join(tokens)
        offsets = token_offsets(tokens)
        spans: set[Span] = set()

        for ner_item in row.get("ner", []):
            if not isinstance(ner_item, (list, tuple)) or len(ner_item) < 3:
                skipped_bad_spans += 1
                continue

            start_tok, end_tok, label = ner_item[0], ner_item[1], ner_item[2]
            if not isinstance(start_tok, int) or not isinstance(end_tok, int):
                skipped_bad_spans += 1
                continue
            if start_tok < 0 or end_tok < start_tok or end_tok >= len(tokens):
                skipped_bad_spans += 1
                continue

            norm_label = normalize_label(label, case_sensitive_labels)
            start_char = offsets[start_tok][0]
            end_char = offsets[end_tok][1]
            spans.add((norm_label, start_char, end_char))
            all_labels.add(norm_label)

        examples.append(Example(text=text, gold_spans=frozenset(spans)))

    labels = sorted(all_labels)
    stats = {
        "input_rows": len(data),
        "used_rows": len(examples),
        "skipped_bad_rows": skipped_bad_rows,
        "skipped_bad_spans": skipped_bad_spans,
        "num_labels": len(labels),
    }
    return examples, labels, stats


def checkpoint_sort_key(path: Path) -> tuple[int, str]:
    m = re.search(r"(\d+)(?!.*\d)", path.name)
    if m is None:
        return (10**18, path.name)
    return (int(m.group(1)), path.name)


def collect_models(
    model_args: list[str],
    checkpoints_root: Path | None,
    checkpoint_glob: str,
    base_model: str,
    skip_base_model: bool,
) -> list[str]:
    model_refs: list[str] = []

    for ref in model_args:
        if ref:
            model_refs.append(ref)

    if checkpoints_root is not None:
        ckpts = [p for p in checkpoints_root.glob(checkpoint_glob) if p.is_dir()]
        for ckpt in sorted(ckpts, key=checkpoint_sort_key):
            model_refs.append(str(ckpt))

    if not skip_base_model and base_model:
        model_refs.append(base_model)

    deduped: list[str] = []
    seen: set[str] = set()
    for ref in model_refs:
        if ref not in seen:
            deduped.append(ref)
            seen.add(ref)
    return deduped


def resolve_device(device_arg: str) -> str:
    if device_arg != "auto":
        return device_arg
    return "cuda" if torch.cuda.is_available() else "cpu"


def safe_prf(tp: int, fp: int, fn: int) -> dict[str, float]:
    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0
    return {"precision": precision, "recall": recall, "f1": f1}


def compute_metrics(
    gold_by_sample: list[set[Span]],
    pred_by_sample: list[set[Span]],
) -> tuple[dict, dict[str, dict]]:
    overall_tp = 0
    overall_fp = 0
    overall_fn = 0

    true_by_label: dict[str, set[tuple[int, int, int]]] = defaultdict(set)
    pred_by_label: dict[str, set[tuple[int, int, int]]] = defaultdict(set)

    for sample_idx, (gold, pred) in enumerate(zip(gold_by_sample, pred_by_sample)):
        tp_items = gold & pred
        overall_tp += len(tp_items)
        overall_fp += len(pred - gold)
        overall_fn += len(gold - pred)

        for label, start, end in gold:
            true_by_label[label].add((sample_idx, start, end))
        for label, start, end in pred:
            pred_by_label[label].add((sample_idx, start, end))

    overall = {
        **safe_prf(overall_tp, overall_fp, overall_fn),
        "tp": overall_tp,
        "fp": overall_fp,
        "fn": overall_fn,
        "pred_total": overall_tp + overall_fp,
        "true_total": overall_tp + overall_fn,
    }

    per_entity: dict[str, dict] = {}
    all_labels = sorted(set(true_by_label) | set(pred_by_label))
    for label in all_labels:
        true_items = true_by_label.get(label, set())
        pred_items = pred_by_label.get(label, set())
        tp = len(true_items & pred_items)
        fp = len(pred_items - true_items)
        fn = len(true_items - pred_items)
        per_entity[label] = {
            **safe_prf(tp, fp, fn),
            "tp": tp,
            "fp": fp,
            "fn": fn,
            "pred_total": tp + fp,
            "true_total": tp + fn,
        }

    return overall, per_entity


def batched(iterable: list[Example], n: int) -> Iterable[list[Example]]:
    for i in range(0, len(iterable), n):
        yield iterable[i : i + n]


@torch.no_grad()
def evaluate_model(
    model_ref: str,
    examples: list[Example],
    labels: list[str],
    batch_size: int,
    threshold: float,
    flat_ner: bool,
    multi_label: bool,
    case_sensitive_labels: bool,
    device: str,
) -> dict:
    model = GLiNER.from_pretrained(model_ref)
    model.to(device)
    model.eval()

    gold_by_sample: list[set[Span]] = [set(ex.gold_spans) for ex in examples]
    pred_by_sample: list[set[Span]] = []

    for batch in tqdm(
        list(batched(examples, batch_size)),
        desc=f"eval:{model_ref}",
        leave=False,
    ):
        texts = [ex.text for ex in batch]
        batch_preds = model.inference(
            texts,
            labels=labels,
            threshold=threshold,
            flat_ner=flat_ner,
            multi_label=multi_label,
            batch_size=len(texts),
        )

        for sample_preds in batch_preds:
            pred_spans: set[Span] = set()
            for item in sample_preds:
                label = normalize_label(item.get("label", ""), case_sensitive_labels)
                start = item.get("start")
                end = item.get("end")
                if not isinstance(start, int) or not isinstance(end, int):
                    continue
                pred_spans.add((label, start, end))
            pred_by_sample.append(pred_spans)

    overall, per_entity = compute_metrics(gold_by_sample, pred_by_sample)
    return {
        "model": model_ref,
        "overall": overall,
        "per_entity": per_entity,
    }


def print_result(result: dict) -> None:
    model_ref = result["model"]
    overall = result["overall"]
    per_entity = result["per_entity"]

    print(f"\nModel: {model_ref}")
    print(
        "Overall: "
        f"P={overall['precision']:.4f} "
        f"R={overall['recall']:.4f} "
        f"F1={overall['f1']:.4f} "
        f"(tp={overall['tp']} fp={overall['fp']} fn={overall['fn']} "
        f"pred={overall['pred_total']} true={overall['true_total']})"
    )

    print("Per-entity:")
    print("  label\tprecision\trecall\tf1\ttp\tfp\tfn\tpred\ttrue")
    for label in sorted(per_entity):
        m = per_entity[label]
        print(
            f"  {label}\t"
            f"{m['precision']:.4f}\t"
            f"{m['recall']:.4f}\t"
            f"{m['f1']:.4f}\t"
            f"{m['tp']}\t{m['fp']}\t{m['fn']}\t"
            f"{m['pred_total']}\t{m['true_total']}"
        )


def main() -> None:
    args = parse_args()

    with args.eval_json.open("r", encoding="utf-8") as f:
        raw = json.load(f)
    if not isinstance(raw, list):
        raise ValueError(f"{args.eval_json} must contain a JSON list of samples.")

    if args.limit and args.limit > 0:
        raw = raw[: args.limit]

    examples, labels, parse_stats = parse_examples(
        raw,
        case_sensitive_labels=args.case_sensitive_labels,
    )
    if not examples:
        raise ValueError("No valid examples found in eval JSON.")
    if not labels:
        raise ValueError("No labels found in eval JSON.")

    model_refs = collect_models(
        model_args=args.model,
        checkpoints_root=args.checkpoints_root,
        checkpoint_glob=args.checkpoint_glob,
        base_model=args.base_model,
        skip_base_model=args.skip_base_model,
    )
    if not model_refs:
        raise ValueError("No models to evaluate. Provide --model or --checkpoints-root.")

    device = resolve_device(args.device)

    print(f"Eval file: {args.eval_json}")
    print(f"Examples: {parse_stats['used_rows']} (from {parse_stats['input_rows']})")
    print(f"Labels: {parse_stats['num_labels']}")
    print(f"Device: {device}")
    print(f"Models: {len(model_refs)}")

    all_results: list[dict] = []
    for model_ref in model_refs:
        started = time.time()
        result = evaluate_model(
            model_ref=model_ref,
            examples=examples,
            labels=labels,
            batch_size=args.batch_size,
            threshold=args.threshold,
            flat_ner=args.flat_ner,
            multi_label=args.multi_label,
            case_sensitive_labels=args.case_sensitive_labels,
            device=device,
        )
        result["elapsed_sec"] = round(time.time() - started, 3)
        print_result(result)
        all_results.append(result)

    payload = {
        "eval_json": str(args.eval_json),
        "parse_stats": parse_stats,
        "labels": labels,
        "settings": {
            "batch_size": args.batch_size,
            "threshold": args.threshold,
            "flat_ner": args.flat_ner,
            "multi_label": args.multi_label,
            "case_sensitive_labels": args.case_sensitive_labels,
            "device": device,
        },
        "results": all_results,
    }

    if args.output_json is not None:
        args.output_json.parent.mkdir(parents=True, exist_ok=True)
        with args.output_json.open("w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)
        print(f"\nSaved metrics JSON: {args.output_json}")


if __name__ == "__main__":
    main()