"""Pinned model inference and official test split; no training or remote code."""
import hashlib
import os
from functools import lru_cache
from pathlib import Path
from urllib.request import urlopen

ROOT = Path(__file__).resolve().parent
CACHE = Path(os.environ.get("APP_CACHE", ROOT / ".cache"))
os.environ.setdefault("HF_HOME", str(CACHE / "huggingface"))
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")

import numpy as np
import torch
from PIL import Image, ImageOps

MODEL_ID = "therayelab/pneumoniamnist-resnet18"
REVISION = "a5e0ebebfcf20bee6701e6d528e33c5bf7a33b94"
DATA_URL = "https://zenodo.org/records/10519652/files/pneumoniamnist.npz?download=1"
DATA_MD5 = "28209eda62fecd6e6a2d98b1501bb15f"
LABELS = ["Normal", "Pneumonia"]
torch.set_num_threads(min(4, os.cpu_count() or 1))

@lru_cache(maxsize=1)
def model_bundle():
    from transformers import AutoImageProcessor, AutoModelForImageClassification
    processor = AutoImageProcessor.from_pretrained(MODEL_ID, revision=REVISION, use_fast=False)
    model = AutoModelForImageClassification.from_pretrained(
        MODEL_ID, revision=REVISION, use_safetensors=True, trust_remote_code=False
    ).eval()
    if model.config.id2label != {0: "0", 1: "1"}:
        raise ValueError("Unexpected model label mapping; review before using this demo.")
    return processor, model

def clean_image(image):
    if image is None:
        raise ValueError("Upload an image first.")
    if image.width * image.height > 20_000_000:
        raise ValueError("Please use an image smaller than 20 megapixels.")
    return ImageOps.exif_transpose(image).convert("RGB")

def predict(images):
    processor, model = model_bundle()
    inputs = processor(images=[clean_image(i) for i in images], return_tensors="pt")
    with torch.inference_mode():
        probabilities = model(**inputs).logits.softmax(dim=-1).cpu().numpy()
    return probabilities

@lru_cache(maxsize=1)
def test_data():
    CACHE.mkdir(parents=True, exist_ok=True)
    path = CACHE / "pneumoniamnist.npz"
    if not path.exists():
        with urlopen(DATA_URL, timeout=120) as response:
            data = response.read(20_000_000)
        if hashlib.md5(data).hexdigest() != DATA_MD5:
            raise ValueError("Dataset checksum failed; please retry later.")
        path.write_bytes(data)
    if hashlib.md5(path.read_bytes()).hexdigest() != DATA_MD5:
        raise ValueError("Dataset checksum failed; remove the cached dataset and retry.")
    with np.load(path, allow_pickle=False) as data:
        images = data["test_images"].copy()
        labels = data["test_labels"].reshape(-1).copy()
    assert images.shape == (624, 28, 28) and set(labels) == {0, 1}
    return images, labels

@lru_cache(maxsize=1)
def test_predictions():
    images, _ = test_data()
    return np.concatenate([predict([Image.fromarray(i) for i in images[s:s+16]])
                           for s in range(0, len(images), 16)])

def metrics(labels, probabilities, threshold):
    predicted = (np.asarray(probabilities)[:, 1] >= threshold).astype(int)
    labels = np.asarray(labels)
    tn = int(np.sum((labels == 0) & (predicted == 0)))
    fp = int(np.sum((labels == 0) & (predicted == 1)))
    fn = int(np.sum((labels == 1) & (predicted == 0)))
    tp = int(np.sum((labels == 1) & (predicted == 1)))
    divide = lambda a, b: a / b if b else None
    return {"n": len(labels), "accuracy": divide(tp+tn, len(labels)),
            "sensitivity": divide(tp, tp+fn), "specificity": divide(tn, tn+fp),
            "precision": divide(tp, tp+fp), "tn": tn, "fp": fp, "fn": fn, "tp": tp}