import os
from collections import defaultdict
from pathlib import Path
from typing import Union

import torch
import torch.nn as nn
from PIL import Image
from torchvision import transforms

ImageInput = Union[str, os.PathLike, Image.Image]

IMAGE_EXTENSIONS = {".bmp", ".jpg", ".jpeg", ".pgm", ".png", ".ppm", ".tif", ".tiff", ".webp"}

LPIPS_NET_CHOICES = ("alex", "vgg", "squeeze")


def _list_image_files(path: Union[str, os.PathLike]):
    path = os.fspath(path)
    if not os.path.isdir(path):
        raise ValueError(f"Expected a directory for LPIPS, got: {path}")
    files = []
    for entry in sorted(os.listdir(path)):
        full = os.path.join(path, entry)
        if os.path.isfile(full) and os.path.splitext(entry)[1].lower() in IMAGE_EXTENSIONS:
            files.append(full)
    if not files:
        raise ValueError(f"No images found under {path}.")
    return files


def _pair_directories_by_stem(dir_a, dir_b):
    files_a = _list_image_files(dir_a)
    files_b = _list_image_files(dir_b)
    by_stem_a = defaultdict(list)
    for f in files_a:
        by_stem_a[Path(f).stem].append(f)
    by_stem_b = defaultdict(list)
    for f in files_b:
        by_stem_b[Path(f).stem].append(f)
    common = sorted(set(by_stem_a.keys()) & set(by_stem_b.keys()))
    if not common:
        raise ValueError(f"No matching filename stems between {dir_a} and {dir_b}.")
    pairs = []
    for stem in common:
        pairs.append((sorted(by_stem_a[stem])[0], sorted(by_stem_b[stem])[0]))
    return pairs


def _open_rgb(image: ImageInput) -> Image.Image:
    if isinstance(image, (str, os.PathLike)):
        image = Image.open(image)
    if not isinstance(image, Image.Image):
        raise TypeError(f"LPIPS expects PIL images or image paths, got {type(image)}.")
    return image.convert("RGB")


class _AlexFeatures(nn.Module):
    def __init__(self):
        super().__init__()
        self.slice1 = nn.Sequential()
        self.slice2 = nn.Sequential()
        self.slice3 = nn.Sequential()
        self.slice4 = nn.Sequential()
        self.slice5 = nn.Sequential()
        self.slice1.add_module("0", nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2))
        self.slice1.add_module("1", nn.ReLU(inplace=True))
        self.slice2.add_module("2", nn.MaxPool2d(kernel_size=3, stride=2))
        self.slice2.add_module("3", nn.Conv2d(64, 192, kernel_size=5, padding=2))
        self.slice2.add_module("4", nn.ReLU(inplace=True))
        self.slice3.add_module("5", nn.MaxPool2d(kernel_size=3, stride=2))
        self.slice3.add_module("6", nn.Conv2d(192, 384, kernel_size=3, padding=1))
        self.slice3.add_module("7", nn.ReLU(inplace=True))
        self.slice4.add_module("8", nn.Conv2d(384, 256, kernel_size=3, padding=1))
        self.slice4.add_module("9", nn.ReLU(inplace=True))
        self.slice5.add_module("10", nn.Conv2d(256, 256, kernel_size=3, padding=1))
        self.slice5.add_module("11", nn.ReLU(inplace=True))

    def forward(self, x):
        h1 = self.slice1(x)
        h2 = self.slice2(h1)
        h3 = self.slice3(h2)
        h4 = self.slice4(h3)
        h5 = self.slice5(h4)
        return [h1, h2, h3, h4, h5]


class _VGG16Features(nn.Module):
    def __init__(self):
        super().__init__()
        self.slice1 = nn.Sequential()
        self.slice2 = nn.Sequential()
        self.slice3 = nn.Sequential()
        self.slice4 = nn.Sequential()
        self.slice5 = nn.Sequential()
        cfg = [
            (1, 0, nn.Conv2d(3, 64, 3, padding=1)),
            (1, 1, nn.ReLU(inplace=True)),
            (1, 2, nn.Conv2d(64, 64, 3, padding=1)),
            (1, 3, nn.ReLU(inplace=True)),
            (2, 4, nn.MaxPool2d(2, 2)),
            (2, 5, nn.Conv2d(64, 128, 3, padding=1)),
            (2, 6, nn.ReLU(inplace=True)),
            (2, 7, nn.Conv2d(128, 128, 3, padding=1)),
            (2, 8, nn.ReLU(inplace=True)),
            (3, 9, nn.MaxPool2d(2, 2)),
            (3, 10, nn.Conv2d(128, 256, 3, padding=1)),
            (3, 11, nn.ReLU(inplace=True)),
            (3, 12, nn.Conv2d(256, 256, 3, padding=1)),
            (3, 13, nn.ReLU(inplace=True)),
            (3, 14, nn.Conv2d(256, 256, 3, padding=1)),
            (3, 15, nn.ReLU(inplace=True)),
            (4, 16, nn.MaxPool2d(2, 2)),
            (4, 17, nn.Conv2d(256, 512, 3, padding=1)),
            (4, 18, nn.ReLU(inplace=True)),
            (4, 19, nn.Conv2d(512, 512, 3, padding=1)),
            (4, 20, nn.ReLU(inplace=True)),
            (4, 21, nn.Conv2d(512, 512, 3, padding=1)),
            (4, 22, nn.ReLU(inplace=True)),
            (5, 23, nn.MaxPool2d(2, 2)),
            (5, 24, nn.Conv2d(512, 512, 3, padding=1)),
            (5, 25, nn.ReLU(inplace=True)),
            (5, 26, nn.Conv2d(512, 512, 3, padding=1)),
            (5, 27, nn.ReLU(inplace=True)),
            (5, 28, nn.Conv2d(512, 512, 3, padding=1)),
            (5, 29, nn.ReLU(inplace=True)),
        ]
        for slice_idx, orig_idx, module in cfg:
            getattr(self, f"slice{slice_idx}").add_module(str(orig_idx), module)

    def forward(self, x):
        h1 = self.slice1(x)
        h2 = self.slice2(h1)
        h3 = self.slice3(h2)
        h4 = self.slice4(h3)
        h5 = self.slice5(h4)
        return [h1, h2, h3, h4, h5]


class _Fire(nn.Module):
    def __init__(self, in_channels, squeeze_channels, expand1x1_channels, expand3x3_channels):
        super().__init__()
        self.squeeze = nn.Conv2d(in_channels, squeeze_channels, kernel_size=1)
        self.squeeze_activation = nn.ReLU(inplace=True)
        self.expand1x1 = nn.Conv2d(squeeze_channels, expand1x1_channels, kernel_size=1)
        self.expand1x1_activation = nn.ReLU(inplace=True)
        self.expand3x3 = nn.Conv2d(squeeze_channels, expand3x3_channels, kernel_size=3, padding=1)
        self.expand3x3_activation = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.squeeze_activation(self.squeeze(x))
        return torch.cat(
            [
                self.expand1x1_activation(self.expand1x1(x)),
                self.expand3x3_activation(self.expand3x3(x)),
            ],
            dim=1,
        )


class _SqueezeNet11Features(nn.Module):
    def __init__(self):
        super().__init__()
        self.slice1 = nn.Sequential()
        self.slice2 = nn.Sequential()
        self.slice3 = nn.Sequential()
        self.slice4 = nn.Sequential()
        self.slice5 = nn.Sequential()
        self.slice6 = nn.Sequential()
        self.slice7 = nn.Sequential()
        self.slice1.add_module("0", nn.Conv2d(3, 64, kernel_size=3, stride=2))
        self.slice1.add_module("1", nn.ReLU(inplace=True))
        self.slice2.add_module("2", nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True))
        self.slice2.add_module("3", _Fire(64, 16, 64, 64))
        self.slice2.add_module("4", _Fire(128, 16, 64, 64))
        self.slice3.add_module("5", nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True))
        self.slice3.add_module("6", _Fire(128, 32, 128, 128))
        self.slice3.add_module("7", _Fire(256, 32, 128, 128))
        self.slice4.add_module("8", nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True))
        self.slice4.add_module("9", _Fire(256, 48, 192, 192))
        self.slice5.add_module("10", _Fire(384, 48, 192, 192))
        self.slice6.add_module("11", _Fire(384, 64, 256, 256))
        self.slice7.add_module("12", _Fire(512, 64, 256, 256))

    def forward(self, x):
        h1 = self.slice1(x)
        h2 = self.slice2(h1)
        h3 = self.slice3(h2)
        h4 = self.slice4(h3)
        h5 = self.slice5(h4)
        h6 = self.slice6(h5)
        h7 = self.slice7(h6)
        return [h1, h2, h3, h4, h5, h6, h7]


_NET_CONFIG = {
    "alex": {"factory": _AlexFeatures, "channels": (64, 192, 384, 256, 256)},
    "vgg": {"factory": _VGG16Features, "channels": (64, 128, 256, 512, 512)},
    "squeeze": {"factory": _SqueezeNet11Features, "channels": (64, 128, 256, 384, 384, 512, 512)},
}


class _ScalingLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("shift", torch.tensor([-0.030, -0.088, -0.188]).view(1, 3, 1, 1))
        self.register_buffer("scale", torch.tensor([0.458, 0.448, 0.450]).view(1, 3, 1, 1))

    def forward(self, x):
        return (x - self.shift) / self.scale


class _NetLinLayer(nn.Module):
    def __init__(self, chn_in, use_dropout=True):
        super().__init__()
        layers = []
        if use_dropout:
            layers.append(nn.Dropout())
        layers.append(nn.Conv2d(chn_in, 1, kernel_size=1, stride=1, padding=0, bias=False))
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)


def _normalize_tensor(x, eps=1e-10):
    norm = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True))
    return x / (norm + eps)


def _spatial_average(x):
    return x.mean(dim=(2, 3), keepdim=True)


class LPIPSModel(nn.Module):
    def __init__(self, net: str = "alex", use_dropout: bool = True):
        super().__init__()
        if net not in _NET_CONFIG:
            raise ValueError(f"net must be one of {LPIPS_NET_CHOICES}, got {net!r}")
        self.net_name = net
        self.scaling_layer = _ScalingLayer()
        self.net = _NET_CONFIG[net]["factory"]()
        chns = _NET_CONFIG[net]["channels"]
        for i, chn in enumerate(chns):
            setattr(self, f"lin{i}", _NetLinLayer(chn, use_dropout=use_dropout))
        self.num_layers = len(chns)
        for p in self.parameters():
            p.requires_grad = False

    def forward(self, in0, in1):
        in0 = self.scaling_layer(in0)
        in1 = self.scaling_layer(in1)
        feats0 = self.net(in0)
        feats1 = self.net(in1)
        val = 0
        for i in range(self.num_layers):
            diff = (_normalize_tensor(feats0[i]) - _normalize_tensor(feats1[i])) ** 2
            lin = getattr(self, f"lin{i}")
            val = val + _spatial_average(lin(diff))
        return val.view(-1)


class LPIPSCompute(nn.Module):
    def __init__(
        self,
        model: LPIPSModel,
        device: Union[str, torch.device] = "cpu",
        batch_size: int = 16,
        target_size: int = 512,
    ):
        super().__init__()
        self.model = model
        self.batch_size = batch_size
        self.target_size = target_size
        self._resize_transform = transforms.Compose(
            [
                transforms.Resize(target_size, interpolation=transforms.InterpolationMode.BICUBIC),
                transforms.CenterCrop(target_size),
                transforms.ToTensor(),
            ]
        )
        self._raw_transform = transforms.ToTensor()
        self.to(device)

    @property
    def device(self):
        try:
            return next(self.model.parameters()).device
        except StopIteration:
            return torch.device("cpu")

    def _to_tensor(self, image: Image.Image, do_resize: bool) -> torch.Tensor:
        transform = self._resize_transform if do_resize else self._raw_transform
        x = transform(image).clamp(0.0, 1.0) * 2.0 - 1.0
        return x

    @torch.no_grad()
    def _compute_pair(self, img_a: Image.Image, img_b: Image.Image, do_resize: bool) -> float:
        x0 = self._to_tensor(img_a, do_resize).unsqueeze(0).to(self.device)
        x1 = self._to_tensor(img_b, do_resize).unsqueeze(0).to(self.device)
        return float(self.model(x0, x1).item())

    @torch.no_grad()
    def _compute_pairs(self, pairs, do_resize: bool) -> float:
        scores = []
        batch_size = max(1, self.batch_size)
        for start in range(0, len(pairs), batch_size):
            chunk = pairs[start : start + batch_size]
            xs0 = torch.stack([self._to_tensor(_open_rgb(a), do_resize) for a, _ in chunk]).to(self.device)
            xs1 = torch.stack([self._to_tensor(_open_rgb(b), do_resize) for _, b in chunk]).to(self.device)
            scores.append(self.model(xs0, xs1).detach().cpu())
        merged = torch.cat(scores, dim=0)
        return float(merged.mean().item())

    @staticmethod
    def _is_dir(value) -> bool:
        return isinstance(value, (str, os.PathLike)) and os.path.isdir(os.fspath(value))

    @staticmethod
    def _is_image_input(value) -> bool:
        if isinstance(value, Image.Image):
            return True
        if isinstance(value, (str, os.PathLike)):
            return os.path.isfile(os.fspath(value))
        return False

    def compute(self, image_a, image_b) -> float:
        a_is_dir = self._is_dir(image_a)
        b_is_dir = self._is_dir(image_b)
        if a_is_dir != b_is_dir:
            raise ValueError("LPIPS.compute requires both inputs to be directories or both to be single images.")

        if a_is_dir:
            pairs = _pair_directories_by_stem(image_a, image_b)
            sizes = set()
            for path_a, path_b in pairs:
                with Image.open(path_a) as ia, Image.open(path_b) as ib:
                    sizes.add(ia.size)
                    sizes.add(ib.size)
            do_resize = len(sizes) > 1
            return self._compute_pairs(pairs, do_resize=do_resize)

        if not (self._is_image_input(image_a) and self._is_image_input(image_b)):
            raise ValueError("LPIPS.compute inputs must be image paths, PIL images, or directories.")
        img_a = _open_rgb(image_a)
        img_b = _open_rgb(image_b)
        do_resize = img_a.size != img_b.size
        return self._compute_pair(img_a, img_b, do_resize=do_resize)

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