"""Classifies host_assets/example.json with LiteRT CompiledModel on the CPU.

Run from the downloaded repository with the packages of requirements-lock.txt:

    python examples/run_example.py
    python examples/run_example.py --model gliner25_decide_s128_wfp16.tflite

The script builds the three graph inputs with the gliner2 2.0.0 host runtime
(host_assets/runtime), runs the graph through the LiteRT CompiledModel Python
API, turns the 32 logits into per-task decisions, and compares them with the
official gliner2 result stored in example.json. No checkpoint is downloaded.
"""

import argparse
import json
import os
import sys
from pathlib import Path

import numpy as np

os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

from ai_edge_litert.compiled_model import (  # noqa: E402
    CompiledModel,
    CpuOptions,
    HardwareAccelerator,
    Options,
)

SIGNATURE = "serving_default"
CPU_THREADS = 4
LABEL_SLOTS = 32


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--assets", type=Path, default=Path("host_assets"))
    parser.add_argument(
        "--model",
        type=Path,
        default=Path("gliner25_decide_s256_wfp16.tflite"),
    )
    parser.add_argument("--table", choices=("fp16", "fp32"), default="fp16")
    parser.add_argument("--result", type=Path, help="optional JSON output")
    return parser.parse_args()


def graph_window(model):
    """Returns N from the attention-mask input [1, N]."""
    details = model.get_input_tensor_details(SIGNATURE).values()
    return next(d["shape"][1] for d in details if len(d["shape"]) == 2)


def run_graph(model, arrays):
    """Runs one request and returns the 32 logits.

    The converter names the inputs args_0..args_2, so each buffer is matched
    to its array by shape: [1,N,1024], [1,N] and [1,32,N] are distinct.
    """
    details = model.get_input_tensor_details(SIGNATURE)
    names = model.get_signature_by_index(0)["inputs"]
    by_shape = {tuple(array.shape): array for array in arrays}
    input_buffers = model.create_input_buffers(0)
    output_buffers = model.create_output_buffers(0)
    try:
        for buffer, name in zip(input_buffers, names):
            array = by_shape[tuple(details[name]["shape"])]
            buffer.write(np.ascontiguousarray(array, dtype=np.float32))
        model.run_by_index(0, input_buffers, output_buffers)
        return output_buffers[0].read(LABEL_SLOTS, np.float32).copy()
    finally:
        for buffer in input_buffers + output_buffers:
            buffer.destroy()


def report(example, prepared, logits, probabilities):
    """Prints every head and returns the per-head record."""
    heads, offset = [], 0
    for task, probs, official in zip(
        prepared.tasks, probabilities, example["task_results"]
    ):
        count = len(task["labels"])
        head_logits = logits[offset:offset + count]
        offset += count
        print(f"{task['task']}:")
        for label, prob in zip(task["labels"], probs):
            print(f"  {label:<20} {float(prob):.4f}")
        heads.append({
            "task": task["task"],
            "multi_label": task["multi_label"],
            "labels": task["labels"],
            "probabilities": [float(p) for p in probs],
            "logits": [float(x) for x in head_logits],
            "official_probabilities": official["probs"],
            "official_logits": official["logits"],
        })
    return heads


def max_abs_difference(heads, key, official_key):
    return max(
        float(np.max(np.abs(np.subtract(h[key], h[official_key]))))
        for h in heads
    )


def main():
    args = parse_args()
    sys.path.insert(0, str((args.assets / "runtime").resolve()))
    from decide_inputs import DecideHost  # pylint: disable=import-outside-toplevel
    from host_decide import decide  # pylint: disable=import-outside-toplevel

    example = json.loads((args.assets / "example.json").read_text())
    host = DecideHost(args.assets, table=args.table)
    model = CompiledModel.from_file(
        str(args.model),
        options=Options(
            hardware_accelerators=HardwareAccelerator.CPU,
            cpu_options=CpuOptions(num_threads=CPU_THREADS),
        ),
    )
    try:
        seq = graph_window(model)
        prepared = host.prepare_classification(
            example["text"], example["tasks"], seq=seq
        )
        logits = run_graph(model, prepared.ordered_inputs())
    finally:
        model.close()
    if not np.isfinite(logits).all():
        raise RuntimeError("the graph returned non-finite logits")
    decisions, probabilities, label_count = decide(logits, prepared.tasks)

    print(example["text"])
    print(f"window s{seq}, {prepared.encoded_length} encoded tokens, "
          f"{label_count} labels, {args.table} table")
    heads = report(example, prepared, logits, probabilities)
    print(json.dumps(decisions, indent=2))
    same = decisions == example["official_result"]
    max_dlogit = max_abs_difference(heads, "logits", "official_logits")
    max_dprob = max_abs_difference(
        heads, "probabilities", "official_probabilities"
    )
    print(f"equal to the official gliner2 result: {same} "
          f"(max |dlogit| {max_dlogit:.2e}, max |dprob| {max_dprob:.2e})")
    if args.result:
        args.result.write_text(json.dumps({
            "model": args.model.name,
            "window": seq,
            "table": args.table,
            "backend": f"LiteRT CompiledModel CPU, {CPU_THREADS} threads",
            "text": example["text"],
            "encoded_length": prepared.encoded_length,
            "label_count": label_count,
            "decisions": decisions,
            "official_result": example["official_result"],
            "decisions_equal_official": same,
            "max_abs_dlogit_vs_official": max_dlogit,
            "max_abs_dprob_vs_official": max_dprob,
            "all_logits_finite": True,
            "heads": heads,
        }, indent=2) + "\n")


if __name__ == "__main__":
    main()