"""Score GLiNER2.5-Decide on Fastino's Fast Decisions dev split with the dataset card's protocol.

One ``classify_text`` call per head with the head's labels as given; single-label and multi-label
predictions are compared to ``true_label`` as sets; the reported average is the mean of the 17
per-domain accuracies. The card's published numbers use the held-out test split, not this one.
"""
import argparse
import json
import time
from pathlib import Path

import coremltools as ct
import numpy as np
import torch
from gliner2 import AutoExtractor
from huggingface_hub import snapshot_download

from convert_names import MODEL_ID, MODEL_REVISION
from preprocessing import prepare_decision
from runtime import decode

DATASET_ID = "fastino/fast-decisions"
DATASET_REVISION = "1a33070cabf94ce2e29105482dd2ef6c157ad7f2"
DOMAINS = [
    "support_intent", "support_topic", "document_type", "review_sentiment", "agent_handoff", "email_triage",
    "ticket_route", "product_feedback", "banking_intent", "clinic_request", "travel_request", "news_topic",
    "paper_field", "sports_recap", "restaurant_review", "benefits_request", "screen_tags",
]
UNITS = {"all": ct.ComputeUnit.ALL, "cpu_ne": ct.ComputeUnit.CPU_AND_NE, "cpu": ct.ComputeUnit.CPU_ONLY,
         "cpu_gpu": ct.ComputeUnit.CPU_AND_GPU}


def labels_of(prediction) -> list[str]:
    """Normalize a native or Core ML head result the way the dataset card does."""
    if isinstance(prediction, str):
        return [prediction]
    if isinstance(prediction, dict):
        return [prediction["label"]]
    return [entry["label"] if isinstance(entry, dict) else entry for entry in prediction]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--backend", choices=["native", "coreml"], required=True)
    parser.add_argument("--package", default="build/gliner2_decide_classification_fp16_L512_H4_K32.mlpackage")
    parser.add_argument("--length", type=int, default=512)
    parser.add_argument("--max-heads", type=int, default=4)
    parser.add_argument("--max-options", type=int, default=32)
    parser.add_argument("--compute-units", choices=list(UNITS), default="all")
    parser.add_argument("--out", required=True)
    args = parser.parse_args()
    torch.set_num_threads(8)
    data = Path(snapshot_download(DATASET_ID, repo_type="dataset", revision=DATASET_REVISION))
    native = AutoExtractor.from_pretrained(snapshot_download(MODEL_ID, revision=MODEL_REVISION),
                                           map_location="cpu").eval()
    model = None
    if args.backend == "coreml":
        model = ct.models.MLModel(args.package, compute_units=UNITS[args.compute_units])
    bucket = (args.length, args.max_heads, args.max_options)
    domains, rows, latencies = {}, [], []
    start = time.perf_counter()
    for domain in DOMAINS:
        hits = total = 0
        for index, line in enumerate((data / f"{domain}.jsonl").open()):
            row = json.loads(line)
            for head in row["output"]["classifications"]:
                tasks = {head["task"]: head["labels"]}
                if model is None:
                    with torch.no_grad():
                        got = native.classify_text(row["input"], tasks)[head["task"]]
                else:
                    arrays = prepare_decision(native.processor, row["input"], tasks, *bucket)
                    call = time.perf_counter()
                    logits = np.asarray(model.predict(arrays)["logits"])[0]
                    latencies.append((time.perf_counter() - call) * 1000)
                    got = decode(tasks, logits)[head["task"]]
                predicted = labels_of(got)
                hit = sorted(predicted) == sorted(head["true_label"])
                hits += hit
                total += 1
                rows.append({"domain": domain, "index": index, "task": head["task"], "predicted": predicted,
                             "gold": head["true_label"], "hit": hit})
        domains[domain] = {"heads": total, "hits": hits, "accuracy": hits / total}
        print(f"{domain:18s} {hits / total:.3f}", flush=True)
    ordered = sorted(latencies)
    result = {
        "model": MODEL_ID, "revision": MODEL_REVISION, "dataset": DATASET_ID, "dataset_revision": DATASET_REVISION,
        "split": "development (public, 100 rows per domain); published card scores use the held-out test split",
        "backend": args.backend, "package": args.package if model else None,
        "compute_units": args.compute_units if model else None,
        "average_accuracy": sum(value["accuracy"] for value in domains.values()) / len(domains),
        "pooled_accuracy": sum(v["hits"] for v in domains.values()) / sum(v["heads"] for v in domains.values()),
        "model_call_p50_ms": ordered[len(ordered) // 2] if ordered else None,
        "model_call_p95_ms": ordered[int(0.95 * (len(ordered) - 1))] if ordered else None,
        "seconds": time.perf_counter() - start, "domains": domains,
    }
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(result, indent=2) + "\n")
    out.with_suffix(".rows.jsonl").write_text("".join(json.dumps(r) + "\n" for r in rows))
    print(json.dumps({k: v for k, v in result.items() if k != "domains"}, indent=2))


if __name__ == "__main__":
    main()