import torch

from ..core import ModelConfig
from ..core.device.npu_compatible_device import get_device_type
from ..models.lpips import LPIPSModel, LPIPS_NET_CHOICES, LPIPSCompute
from .base import Metric


_LPIPS_DEFAULT_FILES = {
    "alex": "LPIPS/alexnet.safetensors",
    "vgg": "LPIPS/vgg.safetensors",
    "squeeze": "LPIPS/squeezenet.safetensors",
}

_LPIPS_MODEL_NAMES = {
    "alex": "image_metrics_lpips_alex",
    "vgg": "image_metrics_lpips_vgg",
    "squeeze": "image_metrics_lpips_squeeze",
}


class LPIPSMetric(Metric):
    def __init__(self, model: LPIPSCompute):
        super().__init__()
        self.model = model

    @classmethod
    def from_pretrained(
        cls,
        net: str = "alex",
        model_config: ModelConfig = None,
        device: torch.device = get_device_type(),
        batch_size: int = 16,
        target_size: int = 512,
        vram_limit: float = None,
    ):
        if net not in LPIPS_NET_CHOICES:
            raise ValueError(f"net must be one of {LPIPS_NET_CHOICES}, got {net!r}")
        if model_config is None:
            model_config = ModelConfig(
                model_id="DiffSynth-Studio/ImageMetrics",
                origin_file_pattern=_LPIPS_DEFAULT_FILES[net],
            )
        model_pool = cls.download_and_load_models([model_config], torch_dtype=torch.float32, device=device, vram_limit=vram_limit)
        backbone = model_pool.fetch_model(_LPIPS_MODEL_NAMES[net])
        if backbone is None:
            raise RuntimeError(
                f"Failed to load LPIPS model for net={net!r}. The provided weights do not match the registered hash for {_LPIPS_MODEL_NAMES[net]}."
            )
        compute_model = LPIPSCompute(
            model=backbone,
            device=device,
            batch_size=batch_size,
            target_size=target_size,
        )
        return cls(compute_model)

    @torch.no_grad()
    def compute(self, image_a, image_b) -> float:
        return self.model.compute(image_a, image_b)

    def forward(self, image_a, image_b):
        return self.compute(image_a, image_b)