from __future__ import annotations
import io
import random
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter, ImageOps
import torch
from torchvision import transforms as T

IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def jpeg_recompress(img: Image.Image, quality: int) -> Image.Image:
    b = io.BytesIO()
    img.convert("RGB").save(b, format="JPEG", quality=int(quality), optimize=False)
    b.seek(0)
    return Image.open(b).convert("RGB")


def sensor_noise(img: Image.Image, sigma: float) -> Image.Image:
    a = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
    read = np.random.normal(0, sigma, a.shape).astype(np.float32)
    shot = np.random.normal(0, np.sqrt(np.clip(a, 0, 1) + 1e-4) * sigma * 0.35).astype(np.float32)
    a = np.clip(a + read + shot, 0, 1)
    return Image.fromarray((a * 255 + .5).astype(np.uint8), mode="RGB")


class SymmetricAugment:

    def __call__(self, image: Image.Image) -> Image.Image:
        im = ImageOps.exif_transpose(image).convert("RGB")
        if random.random() < .70:
            im = jpeg_recompress(im, random.randint(35, 100))
        if random.random() < .15:
            im = jpeg_recompress(im, random.randint(45, 92))
        if random.random() < .52:
            w, h = im.size; s = random.uniform(.45, 1.25)
            im = im.resize((max(64, int(w*s)), max(64, int(h*s))), random.choice([
                Image.Resampling.BICUBIC, Image.Resampling.BILINEAR, Image.Resampling.LANCZOS]))
        if random.random() < .20:
            im = im.filter(ImageFilter.GaussianBlur(random.uniform(.1, 1.25)))
        if random.random() < .12:
            im = im.filter(ImageFilter.UnsharpMask(radius=1, percent=random.randint(60, 130), threshold=2))
        if random.random() < .25:
            im = ImageEnhance.Brightness(im).enhance(random.uniform(.68, 1.30))
        if random.random() < .25:
            im = ImageEnhance.Contrast(im).enhance(random.uniform(.72, 1.28))
        if random.random() < .15:
            im = ImageEnhance.Color(im).enhance(random.uniform(.72, 1.28))
        if random.random() < .18:
            im = sensor_noise(im, random.uniform(.002, .016))
        if random.random() < .18:
            w, h = im.size; s = random.uniform(.60, .92)
            tmp = im.resize((max(64, int(w*s)), max(64, int(h*s))), Image.Resampling.BICUBIC)
            im = tmp.resize((w, h), Image.Resampling.BILINEAR)
            if random.random() < .7:
                im = jpeg_recompress(im, random.randint(55, 92))
        return im


def foundation_transform(size: int = 224):
    return T.Compose([
        T.Resize(size + 32, interpolation=T.InterpolationMode.BICUBIC),
        T.CenterCrop(size),
        T.ToTensor(),
        T.Normalize(IMAGENET_MEAN, IMAGENET_STD),
    ])