import torch
import torch.nn as nn
import torchvision.transforms as tr
from PIL import Image, ImageOps
import numpy as np


class LetterBoxResize:
    """Изменение размера с сохранением пропорций (letter box)"""
    def __init__(self, target_size, fill=0):
        self.target_size = target_size
        self.fill = fill
        
    def __call__(self, img):
        return ImageOps.pad(img, self.target_size, color=(self.fill,)*3)

def get_val_transform(mean, std, image_size=256, crop_size=224):
    """Трансформация валидации: ресайз, центрированное обрезание, нормализация"""
    return tr.Compose([
        LetterBoxResize((image_size, image_size)),
        tr.CenterCrop(crop_size),
        tr.ToTensor(),
        tr.Normalize(mean=mean, std=std)
    ])

def de_normalize(tensor, mean, std):
    """
    Де-нормализация изображения: преобразование тензора обратно в PIL.Image.
    tensor: torch.Tensor размером (C, H, W)
    """
    tensor = tensor.clone().detach().cpu()
    for t, m, s in zip(tensor, mean, std):
        t.mul_(s).add_(m)
    tensor = torch.clamp(tensor, 0, 1)
    np_img = tensor.permute(1,2,0).numpy()
    np_img = (np_img * 255).astype(np.uint8)
    return Image.fromarray(np_img)