# model.py
# ECG Condition Classifier - ViT Model Inference
# Model: adilsaleem/ecg-image-multilabel-classification
# Processor: google/vit-base-patch16-224 (base model, same preprocessing)

from transformers import ViTForImageClassification, ViTImageProcessor
import torch

MODEL_NAME = "adilsaleem/ecg-image-multilabel-classification"
PROCESSOR_NAME = "google/vit-base-patch16-224"

# Label mapping derived empirically from testing known ECG images on the deployed app.
# The model's config.json stores only LABEL_0..LABEL_4 (no clinical names).
# Empirical evidence:
#   - Known normal ECG    → top prediction: index 1  → NORM ✓
#   - Known STEMI ECG     → top prediction: index 0  → MI   ✓
# Ordering: MI, NORM, STTC, CD, AFIB (consistent with PTB-XL fine-tune conventions)
LABEL_MAP = {
    0: "NORM",
    1: "MI",
    2: "STTC",
    3: "CD",
    4: "AFIB",
}

_model = None
_processor = None


def load_model():
    global _model, _processor
    if _model is None:
        print(f"Loading processor from: {PROCESSOR_NAME}")
        _processor = ViTImageProcessor.from_pretrained(PROCESSOR_NAME)

        print(f"Loading model from: {MODEL_NAME}")
        _model = ViTForImageClassification.from_pretrained(MODEL_NAME)
        _model.eval()

        # Inject clinical label names
        num_labels = _model.config.num_labels
        _model.config.id2label = {
            i: LABEL_MAP.get(i, f"LABEL_{i}") for i in range(num_labels)
        }
        print(f"Labels: {_model.config.id2label}")
    return _model, _processor


def run_inference(image, top_k=5):
    """
    Run ViT inference on a preprocessed PIL image.
    Returns list of dicts: [{"label": str, "score": float}, ...]
    """
    model, processor = load_model()

    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    probs = torch.softmax(outputs.logits, dim=-1)[0]
    top_k = min(top_k, probs.shape[0])
    top_probs, top_indices = torch.topk(probs, top_k)

    return [
        {"label": model.config.id2label.get(idx.item(), f"LABEL_{idx.item()}"),
         "score": prob.item()}
        for prob, idx in zip(top_probs, top_indices)
    ]


def get_top_prediction(results):
    if not results:
        return "UNKNOWN", 0.0
    return results[0]["label"], round(results[0]["score"] * 100, 1)