"""PixelPrism v0.1 — DeepfakeDetector subclass for the BitMind DFD Arena.

This is a sanitized single-detector submission representing the most
informative component of PixelPrism's production V9 16-detector ensemble:
the Swin V2 transformer head (haywoodsloan/ai-image-detector-deploy),
which V9's permutation-importance audit ranked at 0.271 (38% of total
discriminative power, the strongest single feature in our ensemble).

The full PixelPrism V9 ensemble fuses 16 detectors via a
HistGradientBoostingClassifier meta-classifier and runs against a live
production API at https://pixelprism.ai. Some V9 components depend on
non-MIT-redistributable weights (FLUX-schnell for DIRE-FLUX, SDXL for
DIRE-SDXL) so this submission is the open-redistributable subset.

For the full 16-detector live numbers see https://pixelprism.ai/leaderboard
(refreshed monthly with each retrain).
"""
from __future__ import annotations

from pathlib import Path

import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import AutoImageProcessor, Swinv2ForImageClassification

from arena.detectors.deepfake_detectors import DeepfakeDetector
from arena.detectors import DETECTOR_REGISTRY


# Module-level cache so repeated instantiations don't reload the model.
_PROCESSOR = None
_MODEL = None


@DETECTOR_REGISTRY.register_module(module_name='PixelPrism')
class PixelPrismDetector(DeepfakeDetector):
    """Single-detector submission: Swin V2 wrapped with PixelPrism's production
    inference path. Outputs P(AI) ∈ [0, 1] per image.

    Attributes set from the YAML config:
        hf_repo (str):     "pixelprism-ai/dfd-arena-mini"
        backbone_repo (str): "haywoodsloan/ai-image-detector-deploy"
        ai_label_idx (int):  the class index that means "AI" in the Swin head (0)
    """

    def __init__(
        self,
        model_name: str = 'PixelPrism',
        config: str = 'pixelprism_config.yaml',
        cuda: bool = True,
    ):
        super().__init__(model_name, config, cuda)

    # ───────────────────────── model load ─────────────────────────

    def load_model(self):
        global _PROCESSOR, _MODEL

        if _MODEL is not None and _PROCESSOR is not None:
            self.processor = _PROCESSOR
            self.model = _MODEL
            return

        # The Swin V2 model + preprocessor come from haywoodsloan's MIT-licensed
        # repo; we re-host the same files at pixelprism-ai/dfd-arena-mini so the
        # arena evaluator can pull everything from one place.
        repo = self.hf_repo

        config_path = hf_hub_download(repo, 'config.json')
        weights_path = hf_hub_download(repo, 'model.safetensors')
        preproc_path = hf_hub_download(repo, 'preprocessor_config.json')

        self.processor = AutoImageProcessor.from_pretrained(Path(preproc_path).parent)
        self.model = Swinv2ForImageClassification.from_pretrained(Path(weights_path).parent)
        self.model.to(self.device).eval()

        _PROCESSOR = self.processor
        _MODEL = self.model

    # ───────────────────────── inference ─────────────────────────

    def preprocess(self, image: Image.Image) -> torch.Tensor:
        if image.mode != 'RGB':
            image = image.convert('RGB')
        inputs = self.processor(images=image, return_tensors='pt')
        return inputs['pixel_values'].to(self.device)

    def __call__(self, image: Image.Image) -> float:
        """Returns P(AI) ∈ [0, 1] for a single PIL image."""
        pixel_values = self.preprocess(image)
        with torch.no_grad():
            logits = self.model(pixel_values=pixel_values).logits
            probs = torch.softmax(logits, dim=-1).cpu().numpy().ravel()
        # ai_label_idx is set from the YAML config (typically 0 for haywoodsloan)
        ai_idx = int(getattr(self, 'ai_label_idx', 0))
        return float(probs[ai_idx])