"""同时作用于图像与标注框的几何/像素变换。

样本约定为 dict:
    image:  PIL.Image 或 torch.Tensor(CHW)
    boxes:  np.ndarray (N, 4) xyxy，原图像素坐标
    labels: np.ndarray (N,)
"""

from __future__ import annotations

import math
import random
from collections.abc import Callable, Iterable, Mapping, Sequence
from typing import Any

import numpy as np
import torch
from PIL import Image

Sample = dict


class Compose:
    def __init__(self, transforms: Sequence[Callable[[Sample], Sample]]) -> None:
        self.transforms = list(transforms)

    def __call__(self, sample: Sample) -> Sample:
        for transform in self.transforms:
            sample = transform(sample)
        return sample


class Resize:
    """按目标尺寸缩放（不保持长宽比），同步缩放框。"""

    def __init__(self, size: Sequence[int]) -> None:
        self.height, self.width = int(size[0]), int(size[1])

    def __call__(self, sample: Sample) -> Sample:
        image: Image.Image = sample["image"]
        src_w, src_h = image.size
        sample["image"] = image.resize((self.width, self.height), Image.BILINEAR)
        boxes = sample.get("boxes")
        if boxes is not None and len(boxes):
            scale = np.array([self.width / src_w, self.height / src_h] * 2, dtype=np.float32)
            sample["boxes"] = boxes.astype(np.float32) * scale
        return sample


class Letterbox:
    """保长宽比缩放到目标尺寸，不足的边用 `fill` 填充（居中）。

    薄膜缺陷图的宽高比跨度极大（长条划痕 vs 团块），直接拉伸会把长条压扁到不可辨认；
    letterbox 保住形状，代价是丢掉绝对尺度——需要尺度信息时用 shape features 补回来。
    """

    def __init__(self, size: Sequence[int], fill: int = 114) -> None:
        self.height, self.width = int(size[0]), int(size[1])
        self.fill = fill

    def __call__(self, sample: Sample) -> Sample:
        image: Image.Image = sample["image"]
        src_w, src_h = image.size
        scale = min(self.width / src_w, self.height / src_h)
        new_w, new_h = max(1, round(src_w * scale)), max(1, round(src_h * scale))
        pad_x, pad_y = (self.width - new_w) // 2, (self.height - new_h) // 2

        fill = self.fill if image.mode != "L" else int(self.fill)
        canvas = Image.new(image.mode, (self.width, self.height), fill)
        canvas.paste(image.resize((new_w, new_h), Image.BILINEAR), (pad_x, pad_y))
        sample["image"] = canvas

        boxes = sample.get("boxes")
        if boxes is not None and len(boxes):
            sample["boxes"] = boxes.astype(np.float32) * scale + np.array(
                [pad_x, pad_y, pad_x, pad_y], dtype=np.float32
            )
        sample["letterbox"] = {"scale": scale, "pad": (pad_x, pad_y), "src_size": (src_w, src_h)}
        return sample


class RandomHorizontalFlip:
    def __init__(self, p: float = 0.5) -> None:
        self.p = p

    def __call__(self, sample: Sample) -> Sample:
        if random.random() >= self.p:
            return sample
        image: Image.Image = sample["image"]
        width = image.size[0]
        sample["image"] = image.transpose(Image.FLIP_LEFT_RIGHT)
        boxes = sample.get("boxes")
        if boxes is not None and len(boxes):
            flipped = boxes.copy()
            flipped[:, 0] = width - boxes[:, 2]
            flipped[:, 2] = width - boxes[:, 0]
            sample["boxes"] = flipped
        return sample


class RandomVerticalFlip:
    def __init__(self, p: float = 0.5) -> None:
        self.p = p

    def __call__(self, sample: Sample) -> Sample:
        if random.random() >= self.p:
            return sample
        image: Image.Image = sample["image"]
        height = image.size[1]
        sample["image"] = image.transpose(Image.FLIP_TOP_BOTTOM)
        boxes = sample.get("boxes")
        if boxes is not None and len(boxes):
            flipped = boxes.copy()
            flipped[:, 1] = height - boxes[:, 3]
            flipped[:, 3] = height - boxes[:, 1]
            sample["boxes"] = flipped
        return sample


class CenterCrop:
    """在原始分辨率上取中心区域，不缩放。

    Resize 的备选：完全不重采样，纹理保持原始尺度，也没有填充边可供网络识别宽高比。
    图比目标小时用 symmetric 方式补齐（镜像延拓，比灰边少引入人造边缘）。
    """

    def __init__(self, size: Sequence[int]) -> None:
        self.height, self.width = int(size[0]), int(size[1])

    def __call__(self, sample: Sample) -> Sample:
        array = np.asarray(sample["image"])
        squeeze = array.ndim == 2
        if squeeze:
            array = array[:, :, None]
        src_h, src_w = array.shape[:2]

        pad_y = max(0, self.height - src_h)
        pad_x = max(0, self.width - src_w)
        if pad_y or pad_x:
            array = np.pad(
                array,
                ((pad_y // 2, pad_y - pad_y // 2), (pad_x // 2, pad_x - pad_x // 2), (0, 0)),
                mode="symmetric",
            )
        src_h, src_w = array.shape[:2]
        y0 = (src_h - self.height) // 2
        x0 = (src_w - self.width) // 2
        cropped = array[y0 : y0 + self.height, x0 : x0 + self.width]

        sample["image"] = Image.fromarray(cropped[:, :, 0] if squeeze else cropped)
        boxes = sample.get("boxes")
        if boxes is not None and len(boxes):
            shifted = boxes.astype(np.float32) - np.array(
                [x0 - pad_x // 2, y0 - pad_y // 2] * 2, dtype=np.float32
            )
            shifted[:, 0::2] = shifted[:, 0::2].clip(0, self.width)
            shifted[:, 1::2] = shifted[:, 1::2].clip(0, self.height)
            keep = (shifted[:, 2] > shifted[:, 0]) & (shifted[:, 3] > shifted[:, 1])
            sample["boxes"] = shifted[keep]
            labels = sample.get("labels")
            if labels is not None and len(labels):
                sample["labels"] = np.asarray(labels)[keep]
        return sample


class RandomResizedCrop:
    """随机裁一块再缩放回目标尺寸；SimCLR 生成正样本对的主力增强。

    对缺陷小图不能裁太狠——`scale` 下限别低于 0.4，否则容易裁到纯背景，
    把"背景 vs 背景"当成正样本对，学出来的特征就废了。
    """

    def __init__(
        self,
        size: Sequence[int],
        scale: Sequence[float] = (0.5, 1.0),
        ratio: Sequence[float] = (0.75, 1.333),
    ) -> None:
        self.height, self.width = int(size[0]), int(size[1])
        self.scale = (float(scale[0]), float(scale[1]))
        self.ratio = (float(ratio[0]), float(ratio[1]))

    def __call__(self, sample: Sample) -> Sample:
        image: Image.Image = sample["image"]
        src_w, src_h = image.size

        for _ in range(10):
            area = src_w * src_h * random.uniform(*self.scale)
            aspect = math.exp(random.uniform(math.log(self.ratio[0]), math.log(self.ratio[1])))
            crop_w, crop_h = round(math.sqrt(area * aspect)), round(math.sqrt(area / aspect))
            if 0 < crop_w <= src_w and 0 < crop_h <= src_h:
                break
        else:  # 退化为中心裁剪
            crop_w, crop_h = min(src_w, src_h), min(src_w, src_h)

        x0 = random.randint(0, src_w - crop_w)
        y0 = random.randint(0, src_h - crop_h)
        sample["image"] = image.crop((x0, y0, x0 + crop_w, y0 + crop_h)).resize(
            (self.width, self.height), Image.BILINEAR
        )

        boxes = sample.get("boxes")
        if boxes is not None and len(boxes):
            shifted = boxes.astype(np.float32) - np.array([x0, y0, x0, y0], dtype=np.float32)
            shifted[:, 0::2] = shifted[:, 0::2].clip(0, crop_w) * (self.width / crop_w)
            shifted[:, 1::2] = shifted[:, 1::2].clip(0, crop_h) * (self.height / crop_h)
            keep = (shifted[:, 2] > shifted[:, 0]) & (shifted[:, 3] > shifted[:, 1])
            sample["boxes"] = shifted[keep]
            labels = sample.get("labels")
            if labels is not None and len(labels):
                sample["labels"] = np.asarray(labels)[keep]
        return sample


class GaussianBlur:
    """SimCLR 的标准增强之一，防止网络靠高频噪声区分两个视图。"""

    def __init__(self, p: float = 0.5, radius: Sequence[float] = (0.1, 1.5)) -> None:
        self.p = p
        self.radius = (float(radius[0]), float(radius[1]))

    def __call__(self, sample: Sample) -> Sample:
        if random.random() < self.p:
            from PIL import ImageFilter

            sample["image"] = sample["image"].filter(
                ImageFilter.GaussianBlur(random.uniform(*self.radius))
            )
        return sample


class ColorJitter:
    """亮度/对比度扰动；薄膜图像的光照波动是主要域偏移来源。"""

    def __init__(self, brightness: float = 0.2, contrast: float = 0.2) -> None:
        self.brightness = brightness
        self.contrast = contrast

    def __call__(self, sample: Sample) -> Sample:
        from PIL import ImageEnhance

        image = sample["image"]
        if self.brightness > 0:
            factor = 1.0 + random.uniform(-self.brightness, self.brightness)
            image = ImageEnhance.Brightness(image).enhance(factor)
        if self.contrast > 0:
            factor = 1.0 + random.uniform(-self.contrast, self.contrast)
            image = ImageEnhance.Contrast(image).enhance(factor)
        sample["image"] = image
        return sample


class ToTensor:
    def __call__(self, sample: Sample) -> Sample:
        array = np.asarray(sample["image"], dtype=np.float32) / 255.0
        if array.ndim == 2:
            array = array[:, :, None]
        sample["image"] = torch.from_numpy(array.transpose(2, 0, 1)).contiguous()
        return sample


class Normalize:
    def __init__(self, mean: Sequence[float], std: Sequence[float]) -> None:
        self.mean = torch.tensor(mean, dtype=torch.float32).view(-1, 1, 1)
        self.std = torch.tensor(std, dtype=torch.float32).view(-1, 1, 1)

    def __call__(self, sample: Sample) -> Sample:
        sample["image"] = (sample["image"] - self.mean) / self.std
        return sample


class PerImageStandardize:
    """逐图 z-score：减自身均值、除自身标准差。

    薄膜图之间的整体亮度差异来自光照和相机增益，不是缺陷属性。用固定的
    ImageNet mean/std 归一化保不住这一点——实测聚类会把大半的簇花在
    "背景偏亮/偏暗"上，只剩一个簇装真正有缺陷的图。逐图标准化把亮度和
    对比度这两个全局因子直接除掉，逼特征去看局部结构。
    """

    def __init__(self, eps: float = 1e-5) -> None:
        self.eps = eps

    def __call__(self, sample: Sample) -> Sample:
        image = sample["image"]
        sample["image"] = (image - image.mean()) / (image.std() + self.eps)
        return sample


_TRANSFORMS: dict[str, Callable[..., Any]] = {
    "Resize": Resize,
    "Letterbox": Letterbox,
    "CenterCrop": CenterCrop,
    "RandomResizedCrop": RandomResizedCrop,
    "GaussianBlur": GaussianBlur,
    "RandomHorizontalFlip": RandomHorizontalFlip,
    "RandomVerticalFlip": RandomVerticalFlip,
    "ColorJitter": ColorJitter,
    "ToTensor": ToTensor,
    "Normalize": Normalize,
    "PerImageStandardize": PerImageStandardize,
}


def build_transforms(cfgs: Iterable[Mapping[str, Any]]) -> Compose:
    """由 yaml 列表构建，如 `- {name: Resize, size: [512, 512]}`。"""
    transforms = []
    for cfg in cfgs:
        cfg = dict(cfg)
        name = cfg.pop("name")
        if name not in _TRANSFORMS:
            raise KeyError(f"未知 transform: {name}，可选 {sorted(_TRANSFORMS)}")
        transforms.append(_TRANSFORMS[name](**cfg))
    return Compose(transforms)