import torch
import torchvision.transforms.functional as TF
import torchvision.transforms as T
import random
from PIL import Image


ROTATION_DEGREES  = 15       
NOISE_STD         = 0.02     
BRIGHTNESS_FACTOR = 0.3      
CONTRAST_FACTOR   = 0.3      



def _add_gaussian_noise(tensor: torch.Tensor, std: float = NOISE_STD) -> torch.Tensor:
    noise = torch.randn_like(tensor) * std
    return (tensor + noise).clamp(0.0, 1.0)



def augment_image(img: Image.Image) -> Image.Image:
    img = img.convert("RGB")

    # 1 — Random rotation
    angle = random.uniform(-ROTATION_DEGREES, ROTATION_DEGREES)
    img = TF.rotate(
        img,
        angle=angle,
        interpolation=TF.InterpolationMode.BICUBIC,
        fill=255,           
    )

    # 2 — Brightness + contrast jitter
    jitter = T.ColorJitter(
        brightness=BRIGHTNESS_FACTOR,
        contrast=CONTRAST_FACTOR,
    )
    img = jitter(img)

    # 3 — Gaussian noise (operate in float-tensor space)
    tensor = TF.to_tensor(img)              
    tensor = _add_gaussian_noise(tensor)
    img    = TF.to_pil_image(tensor)        
    
    return img