"""Run multi-label inference with a published PyTorch or CPU INT8 ONNX model."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer

DEFAULT_MODEL_ID = "NajahUniv/AraUni-MARBERTv2-Intent-Classifier"


def choose_device(requested: str) -> str:
    if requested != "auto":
        return requested
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def resolve_runtime(backend: str, precision: str, device: str) -> tuple[str, str]:
    resolved_backend = "pytorch" if backend == "auto" else backend
    resolved_precision = precision
    if precision == "auto":
        if resolved_backend == "onnx":
            resolved_precision = "int8"
        elif device == "cuda":
            resolved_precision = "bf16" if torch.cuda.is_bf16_supported() else "fp16"
        else:
            resolved_precision = "fp32"
    if resolved_backend == "onnx" and (device != "cpu" or resolved_precision != "int8"):
        raise ValueError("the published ONNX artifact supports CPU INT8 only")
    if resolved_backend == "pytorch" and resolved_precision == "int8":
        raise ValueError("precision=int8 requires backend=onnx")
    if resolved_backend == "pytorch" and resolved_precision in {"bf16", "fp16"} and device != "cuda":
        raise ValueError("the example enables bf16/fp16 only on CUDA")
    return resolved_backend, resolved_precision


def hub_or_local_file(model_id: str, filename: str, revision: str | None) -> str:
    local = Path(model_id) / filename
    if local.is_file():
        return str(local)
    return hf_hub_download(model_id, filename, revision=revision)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-id", default=DEFAULT_MODEL_ID)
    parser.add_argument("--revision", help="Pin a release tag or reviewed commit SHA")
    parser.add_argument("--text", required=True)
    parser.add_argument("--device", default="auto", choices=("auto", "cpu", "mps", "cuda"))
    parser.add_argument("--backend", default="auto", choices=("auto", "pytorch", "onnx"))
    parser.add_argument(
        "--precision",
        default="auto",
        choices=("auto", "fp32", "bf16", "fp16", "int8"),
    )
    parser.add_argument("--top-k", type=int, default=5)
    args = parser.parse_args()

    device = "cpu" if args.backend == "onnx" and args.device == "auto" else choose_device(args.device)
    backend, precision = resolve_runtime(args.backend, args.precision, device)
    load_kwargs = {"revision": args.revision} if args.revision else {}
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_id,
        trust_remote_code=True,
        **load_kwargs,
    )
    config = AutoConfig.from_pretrained(
        args.model_id,
        trust_remote_code=True,
        **load_kwargs,
    )
    thresholds = config.thresholds
    if backend == "onnx":
        import onnxruntime as ort

        onnx_config_path = hub_or_local_file(
            args.model_id,
            "onnx/onnx_config.json",
            args.revision,
        )
        onnx_config = json.loads(Path(onnx_config_path).read_text(encoding="utf-8"))
        thresholds = onnx_config["thresholds"]
        model_path = hub_or_local_file(
            args.model_id,
            "onnx/model_int8.onnx",
            args.revision,
        )
        session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
        encoded = tokenizer(
            args.text,
            return_tensors="np",
            truncation=True,
            max_length=config.max_length,
        )
        logits = session.run(
            ["logits"],
            {
                "input_ids": encoded["input_ids"].astype(np.int64),
                "attention_mask": encoded["attention_mask"].astype(np.int64),
            },
        )[0][0]
        probabilities = 1.0 / (1.0 + np.exp(-logits))
    else:
        dtype = {
            "fp32": torch.float32,
            "bf16": torch.bfloat16,
            "fp16": torch.float16,
        }[precision]
        model = AutoModelForSequenceClassification.from_pretrained(
            args.model_id,
            trust_remote_code=True,
            torch_dtype=dtype,
            **load_kwargs,
        ).to(device).eval()
        encoded = tokenizer(
            args.text,
            return_tensors="pt",
            truncation=True,
            max_length=model.config.max_length,
        ).to(device)
        with torch.inference_mode():
            probabilities = torch.sigmoid(model(**encoded).logits[0]).cpu().float().numpy()

    labels = [config.id2label[index] for index in range(config.num_labels)]
    ranked = sorted(
        (
            {
                "label": label,
                "probability": float(probabilities[index]),
                "threshold": float(thresholds[label]),
                "selected": float(probabilities[index]) >= float(thresholds[label]),
            }
            for index, label in enumerate(labels)
        ),
        key=lambda item: item["probability"],
        reverse=True,
    )
    print(
        json.dumps(
            {
                "text": args.text,
                "backend": backend,
                "precision": precision,
                "device": device,
                "selected_labels": [item["label"] for item in ranked if item["selected"]],
                "scores": ranked[: max(1, args.top_k)],
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()