import io
import random

import numpy as np
import torch
from PIL import Image, ImageFilter
from torchvision.transforms import InterpolationMode
from torchvision.transforms import functional as F

# Per-channel mean and std of the ImageNet training set (RGB order)
# Required when using torchvision pretrained weights — they were trained with
# this exact normalisation and expect it at inference time
_IMAGENET_MEAN = (0.485, 0.456, 0.406)
_IMAGENET_STD  = (0.229, 0.224, 0.225)


# Computes per-channel mean and std from real (label=0) training samples only
# Used for the real-norm experiment to test whether the model relies on
# colour/brightness differences between real and fake rather than identity cues
def compute_real_stats(dataset, indices, max_samples=1000, seed=42):
    real_indices = [i for i in indices if dataset.samples[i][1] == 0]
    if not real_indices:
        return _IMAGENET_MEAN, _IMAGENET_STD

    if len(real_indices) > max_samples:
        rng = np.random.RandomState(seed)
        real_indices = rng.choice(real_indices, max_samples, replace=False).tolist()

    means, vars_ = [], []
    for i in real_indices:
        path, _ = dataset.samples[i]
        img = np.array(Image.open(path).convert("RGB"), dtype=np.float32) / 255.0
        means.append(img.mean(axis=(0, 1)))
        vars_.append(img.var(axis=(0, 1)))

    mean = tuple(float(x) for x in np.mean(means, axis=0))
    std  = tuple(float(x) for x in np.sqrt(np.mean(vars_, axis=0)))
    return mean, std


# Single-image preprocessing pipeline for training and evaluation
# Square-crops first to remove the real-rectangular / fake-square geometry cue,
# then resizes, augments (train only), and normalizes
# augment=None uses DEFAULTS; augment=dict overrides specific keys; NO_AUGMENT disables all stochastic ops
class DFFImagePipeline:

    DEFAULTS = {
        "crop_scale": [0.85, 1.0],
        "center_jitter": 0.1,
        "hflip_p": 0.5,
        "rotation_degrees": 15,
        "brightness": 0.4,
        "contrast": 0.4,
        "saturation": 0.3,
        "hue": 0.05,
        "grayscale_p": 0.2,
        "blur_p": 0.3,
        "blur_radius": [0.1, 1.5],
        "jpeg_p": 0.3,
        "jpeg_quality": [65, 95],
        "erase_p": 0.3,
        "erase_scale": [0.02, 0.15],
        "noise_p": 0.2,
        "noise_std": 0.05,
    }

    # Pass as augment=DFFImagePipeline.NO_AUGMENT to keep crop+resize+normalize but skip all randomness
    NO_AUGMENT: dict = {
        "crop_scale": [1.0, 1.0],
        "center_jitter": 0.0,
        "hflip_p": 0.0,
        "rotation_degrees": 0,
        "brightness": 0.0,
        "contrast": 0.0,
        "saturation": 0.0,
        "hue": 0.0,
        "grayscale_p": 0.0,
        "blur_p": 0.0,
        "jpeg_p": 0.0,
        "erase_p": 0.0,
        "noise_p": 0.0,
    }

    def __init__(self, *, image_size: int, train: bool, augment: dict | None = None,
                 normalize_mean=None, normalize_std=None):
        self.image_size = image_size
        self.train = train
        self.normalize_mean = normalize_mean or _IMAGENET_MEAN
        self.normalize_std  = normalize_std  or _IMAGENET_STD

        cfg = {**self.DEFAULTS, **(augment or {})}
        self.crop_scale = tuple(cfg["crop_scale"])
        self.center_jitter = cfg["center_jitter"]
        self.hflip_p = cfg["hflip_p"]
        self.rotation_degrees = cfg["rotation_degrees"]
        self.brightness = cfg["brightness"]
        self.contrast = cfg["contrast"]
        self.saturation = cfg["saturation"]
        self.hue = cfg["hue"]
        self.grayscale_p = cfg["grayscale_p"]
        self.blur_p = cfg["blur_p"]
        self.blur_radius = tuple(cfg["blur_radius"])
        self.jpeg_p = cfg["jpeg_p"]
        self.jpeg_quality = tuple(cfg["jpeg_quality"])
        self.erase_p = cfg["erase_p"]
        self.erase_scale = tuple(cfg["erase_scale"])
        self.noise_p = cfg["noise_p"]
        self.noise_std = cfg["noise_std"]

    # Geometry transforms
    def _crop_square(self, img: Image.Image) -> Image.Image:
        width, height = img.size
        short_side = min(width, height)
        center_top = max((height - short_side) // 2, 0)
        center_left = max((width - short_side) // 2, 0)

        if self.train:
            min_scale, max_scale = self.crop_scale
            scale = random.uniform(min_scale, max_scale)
            crop_size = max(1, int(short_side * scale))
            top = max((height - crop_size) // 2, 0)
            left = max((width - crop_size) // 2, 0)

            jitter_y = int((height - crop_size) * self.center_jitter)
            jitter_x = int((width - crop_size) * self.center_jitter)
            if jitter_y > 0:
                top += random.randint(-jitter_y, jitter_y)
            if jitter_x > 0:
                left += random.randint(-jitter_x, jitter_x)

            top = max(0, min(top, height - crop_size))
            left = max(0, min(left, width - crop_size))
        else:
            crop_size = short_side
            top = center_top
            left = center_left

        return F.crop(img, top=top, left=left, height=crop_size, width=crop_size)

    def _maybe_flip(self, img: Image.Image) -> Image.Image:
        if self.train and random.random() < self.hflip_p:
            return F.hflip(img)
        return img

    def _maybe_rotate(self, img: Image.Image) -> Image.Image:
        if self.train and self.rotation_degrees > 0:
            angle = random.uniform(-self.rotation_degrees, self.rotation_degrees)
            return F.rotate(img, angle, interpolation=InterpolationMode.BILINEAR, fill=0)
        return img

    # Photometric transforms
    def _jitter_factor(self, amount: float) -> float:
        return random.uniform(max(0.0, 1.0 - amount), 1.0 + amount)

    def _maybe_color_jitter(self, img: Image.Image) -> Image.Image:
        if not self.train:
            return img
        img = F.adjust_brightness(img, self._jitter_factor(self.brightness))
        img = F.adjust_contrast(img, self._jitter_factor(self.contrast))
        img = F.adjust_saturation(img, self._jitter_factor(self.saturation))
        img = F.adjust_hue(img, random.uniform(-self.hue, self.hue))
        return img

    def _maybe_grayscale(self, img: Image.Image) -> Image.Image:
        if self.train and random.random() < self.grayscale_p:
            return F.to_grayscale(img, num_output_channels=3)
        return img

    def _maybe_blur(self, img: Image.Image) -> Image.Image:
        if self.train and random.random() < self.blur_p:
            radius = random.uniform(*self.blur_radius)
            return img.filter(ImageFilter.GaussianBlur(radius=radius))
        return img

    # Bias-reduction transforms
    # JPEG recompression removes high-frequency GAN artifacts that survive other augmentations
    def _maybe_jpeg(self, img: Image.Image) -> Image.Image:
        if self.train and random.random() < self.jpeg_p:
            quality = random.randint(*self.jpeg_quality)
            buf = io.BytesIO()
            img.save(buf, format="JPEG", quality=quality)
            buf.seek(0)
            img = Image.open(buf).convert("RGB")
            img.load()  # decode pixels while buf is still in scope
        return img

    # Random erasing forces the model to use multiple regions rather than a single discriminative patch
    def _maybe_erase(self, tensor: torch.Tensor) -> torch.Tensor:
        if self.train and random.random() < self.erase_p:
            c, h, w = tensor.shape
            area = h * w
            min_scale, max_scale = self.erase_scale
            erase_area = area * random.uniform(min_scale, max_scale)
            aspect = random.uniform(0.3, 1.0 / 0.3)
            eh = int(round((erase_area * aspect) ** 0.5))
            ew = int(round((erase_area / aspect) ** 0.5))
            eh, ew = min(eh, h), min(ew, w)
            top = random.randint(0, h - eh)
            left = random.randint(0, w - ew)
            tensor = tensor.clone()
            tensor[:, top:top + eh, left:left + ew] = torch.rand(c, eh, ew)
        return tensor

    # Gaussian noise improves robustness to sensor noise vs GAN noise patterns
    def _maybe_noise(self, tensor: torch.Tensor) -> torch.Tensor:
        if self.train and random.random() < self.noise_p:
            noise = torch.randn_like(tensor) * self.noise_std
            tensor = tensor + noise
        return tensor

    # Pipeline entrypoint
    def __call__(self, img: Image.Image) -> torch.Tensor:
        img = self._crop_square(img)
        img = F.resize(img, [self.image_size, self.image_size], interpolation=InterpolationMode.BILINEAR)
        img = self._maybe_flip(img)
        img = self._maybe_rotate(img)
        img = self._maybe_color_jitter(img)
        img = self._maybe_grayscale(img)
        img = self._maybe_blur(img)
        img = self._maybe_jpeg(img)
        tensor = F.to_tensor(img)
        tensor = self._maybe_erase(tensor)
        tensor = self._maybe_noise(tensor)
        return F.normalize(tensor, self.normalize_mean, self.normalize_std)


# Convenience wrapper used by splits.py and evaluate.py
def get_transforms(train=True, image_size=224, augment=None,
                   normalize_mean=None, normalize_std=None):
    return DFFImagePipeline(image_size=image_size, train=train, augment=augment,
                            normalize_mean=normalize_mean, normalize_std=normalize_std)
