#!/usr/bin/env python3
"""Check a LiteRT GLiNER2 port against the original checkpoint on the bundled parity corpus.

    pip install gliner2==2.0.0 torch numpy ai-edge-litert
    python runtime/verify_parity.py --bundle .            # from the repository folder
    python runtime/verify_parity.py --bundle . --span gliner25_decide_span_wfp16.tflite --count gliner25_decide_count_wfp16.tflite

Runs every corpus request twice, through gliner2 with the original weights (downloaded from Fastino at the pinned
revision) and through this repository's LiteRT runtime, and prints per call type how many structured results are
identical, the lowest span F1 and the largest confidence difference. The corpus (runtime/parity_corpus.json) is short
synthetic texts with schema templates for every call type gliner2 2.0.0 offers; it has no gold labels.
"""
from __future__ import annotations

import argparse
import collections
import json
import pathlib
import sys

HERE = pathlib.Path(__file__).resolve().parent
SOURCES = {"boundary": ("fastino/gliner2.5-base-v1", "1a8bc24e00dc7300b9017c81d63e3dcdabb26596"),
           "span": ("fastino/GLiNER2.5-Decide", "bbe10ff77ebb238777c17d3a8ac9260e30929057")}


def build_schema(t: dict):
    from gliner2 import AttributeGroup, Schema
    s = Schema()
    if "entities" in t:
        s.entities(t["entities"], dtype=t.get("dtype", "list"))
    if "attributes" in t:
        s.entity_attributes({k: AttributeGroup(labels=v["labels"], multi_label=v.get("multi_label", False)) for k, v in t["attributes"].items()})
    if "relations" in t:
        s.relations(t["relations"])
    for c in t.get("classifications", []):
        s.classification(c["task"], c["labels"], multi_label=c.get("multi_label", False))
    for st in t.get("structures", []):
        builder = s.structure(st["name"], mode=st.get("mode"), anchor=st.get("anchor"))
        for f in st["fields"]:
            builder.field(f["name"], dtype=f.get("dtype", "list"), choices=f.get("choices"))
    return s


def without_confidence(value):
    if isinstance(value, dict):
        return {k: without_confidence(v) for k, v in value.items() if k != "confidence"}
    if isinstance(value, list):
        return [without_confidence(v) for v in value]
    return value


def confidence_errors(reference, actual):
    if isinstance(reference, dict) and isinstance(actual, dict):
        out = []
        for key, value in reference.items():
            if key == "confidence" and key in actual:
                out.append(abs(float(value) - float(actual[key])))
            elif key in actual:
                out.extend(confidence_errors(value, actual[key]))
        return out
    if isinstance(reference, (list, tuple)) and isinstance(actual, (list, tuple)):
        return [e for a, b in zip(reference, actual) for e in confidence_errors(a, b)]
    return []


def answers(value, path: str = "") -> list[str]:
    out = []
    if isinstance(value, dict):
        if {"text", "start", "end"} <= set(value):
            return [f"{path}|{value['text']}|{value['start']}|{value['end']}"]
        if "head" in value and "tail" in value:
            return [f"{path}|H:{'/'.join(answers(value['head']))}|T:{'/'.join(answers(value['tail']))}"]
        if "label" in value:
            return [f"{path}|{value['label']}"]
        for key, item in value.items():
            if key != "confidence":
                out += answers(item, f"{path}.{key}" if path else key)
    elif isinstance(value, (list, tuple)):
        if len(value) == 2 and isinstance(value[0], str) and isinstance(value[1], float):
            return [f"{path}|{value[0]}"]
        for item in value:
            out += answers(item, path)
    elif value is not None:
        out.append(f"{path}|{value}")
    return out


def f1(reference: list[str], actual: list[str]) -> float:
    a, b = collections.Counter(reference), collections.Counter(actual)
    if not a and not b:
        return 1.0
    return 2 * sum((a & b).values()) / (sum(a.values()) + sum(b.values()))


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--bundle", default=str(HERE.parent))
    parser.add_argument("--corpus", default=str(HERE / "parity_corpus.json"))
    parser.add_argument("--span", help="span architecture: the span file in place of the heads file's span_rep, count and classifier")
    parser.add_argument("--count", help="span architecture: the count file in place of the heads file's count_embed")
    args = parser.parse_args()
    sys.path.insert(0, str(HERE))
    import torch
    from gliner2 import AutoExtractor
    from huggingface_hub import snapshot_download
    config = json.loads((pathlib.Path(args.bundle) / "source" / "config.json").read_text())
    architecture = config.get("architecture", "span")
    if architecture == "boundary":
        from litert_boundary_runtime import LiteRTBoundaryExtractor as Runtime
    else:
        from litert_span_runtime import LiteRTSpanExtractor as Runtime
    repo, revision = SOURCES[architecture]
    native = AutoExtractor.from_pretrained(snapshot_download(repo, revision=revision), map_location="cpu").eval()
    torch.set_num_threads(4)
    if (args.span or args.count) and architecture != "span":
        parser.error("--span and --count apply to the span architecture only")
    runtime = Runtime(args.bundle, span=args.span, count=args.count) if architecture == "span" else Runtime(args.bundle)
    corpus = json.loads(pathlib.Path(args.corpus).read_text())
    by_kind = {}
    for group in corpus["groups"].values():
        for text in group["texts"]:
            for kind in group["templates"]:
                schema = build_schema(corpus["templates"][kind])
                expected = native.extract(text, schema, include_confidence=True, include_spans=True)
                k = by_kind.setdefault(kind, {"calls": 0, "identical": 0, "refused": 0, "minSpanF1": 1.0, "maxConfidenceDifference": 0.0})
                k["calls"] += 1
                try:
                    actual = runtime.extract(text, schema, include_confidence=True, include_spans=True)
                except ValueError:  # larger than the graphs' fixed sizes: refused, never truncated
                    k["refused"] += 1
                    continue
                k["identical"] += without_confidence(expected) == without_confidence(actual)
                k["minSpanF1"] = round(min(k["minSpanF1"], f1(answers(expected), answers(actual))), 4)
                k["maxConfidenceDifference"] = round(max([k["maxConfidenceDifference"], *confidence_errors(expected, actual)]), 6)
    for kind, k in by_kind.items():
        print(f"{kind:26s} {json.dumps(k)}")
    total = sum(k["calls"] - k["refused"] for k in by_kind.values())
    same = sum(k["identical"] for k in by_kind.values())
    print(f"identical structured results: {same} / {total}")
    return 0 if same == total else 1


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