"""Run the GLiNER2.5-Decide Core ML classifier without loading the original model weights."""
import argparse
import json
from pathlib import Path

import coremltools as ct
import numpy as np

from preprocessing import load_processor, prepare_decision, task_labels


def package_name(precision: str, length: int, max_heads: int, max_options: int) -> str:
    return f"gliner2_decide_classification_{precision}_L{length}_H{max_heads}_K{max_options}.mlpackage"


def decode(tasks: dict, logits: np.ndarray) -> dict:
    """Apply the native activation and label rules to per-head logits of shape (heads, options)."""
    results = {}
    for head, (name, labels) in enumerate(task_labels(tasks).items()):
        config = tasks[name] if isinstance(tasks[name], dict) else {}
        values = logits[head, : len(labels)].astype(np.float64)
        activation = config.get("class_act", "auto")
        multi = config.get("multi_label", False)
        if activation == "sigmoid" or (activation == "auto" and multi):
            probs = 1.0 / (1.0 + np.exp(-values))
        else:
            probs = np.exp(values - values.max())
            probs /= probs.sum()
        if multi:
            threshold = config.get("cls_threshold", 0.5)
            chosen = [{"label": labels[j], "confidence": float(probs[j])} for j in range(len(labels))
                      if probs[j] >= threshold]
            best = int(probs.argmax())
            results[name] = chosen or [{"label": labels[best], "confidence": float(probs[best])}]
        else:
            best = int(probs.argmax())
            results[name] = {"label": labels[best], "confidence": float(probs[best])}
    return results


class CoreMLDecide:
    def __init__(self, model_dir: str, precision: str = "fp16", length: int = 128, max_heads: int = 4,
                 max_options: int | None = None, compute_units=ct.ComputeUnit.ALL):
        model_dir = Path(model_dir)
        # Published buckets: L128 holds 8 labels per head, L256 and L512 hold 32.
        max_options = max_options or (8 if length == 128 else 32)
        self.length, self.max_heads, self.max_options = length, max_heads, max_options
        self.processor = load_processor(str(model_dir))
        package = model_dir / package_name(precision, length, max_heads, max_options)
        self.model = ct.models.MLModel(str(package), compute_units=compute_units)

    def classify(self, text: str, tasks: dict) -> dict:
        arrays = prepare_decision(self.processor, text, tasks, self.length, self.max_heads, self.max_options)
        logits = np.asarray(self.model.predict(arrays)["logits"])[0]
        return decode(tasks, logits)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-dir", required=True)
    parser.add_argument("--text", required=True)
    parser.add_argument("--tasks", required=True, help='JSON object, e.g. {"intent": ["a", "b"]}')
    parser.add_argument("--precision", default="fp16")
    parser.add_argument("--length", type=int, default=128)
    args = parser.parse_args()
    model = CoreMLDecide(args.model_dir, args.precision, args.length)
    print(json.dumps(model.classify(args.text, json.loads(args.tasks)), indent=2))


if __name__ == "__main__":
    main()