import torch
import torch.nn as nn
import segmentation_models_pytorch as smp
import torchvision
import torchvision.transforms.functional as TVF
import PIL

class ImageNetPreprocess(nn.Module):
    def __init__(self, device, encoder_architecture="resnet18"):
        super().__init__()
        params = smp.encoders.get_preprocessing_params(encoder_architecture, "imagenet")
        self.normalize = torchvision.transforms.Normalize(mean=params["mean"], std=params["std"])
    def forward(self, x):
        return self.normalize(x)

class UNet(nn.Module):
    """A simple UNet wrapper that handles ImageNet normalization in-network"""
    def __init__(self, imagenet_pretrained: bool, encoder_architecture, input_shape, device, auto_resize_in: bool, auto_resize_out: bool):
        super().__init__()
        self.imagenet_pretrained = imagenet_pretrained
        self.input_shape = input_shape
        self.auto_resize_in = auto_resize_in
        self.auto_resize_out = auto_resize_out
        encoder_weights = "imagenet" if imagenet_pretrained else None
        self.underlying = smp.Unet(encoder_architecture, encoder_weights=encoder_weights, in_channels=3, activation='sigmoid').to(device)
        self.preprocess = ImageNetPreprocess(device, encoder_architecture).to(device)
    
    def forward(self, x):
        assert x.ndim == 4, f"Illegal input shape {x.shape}. Expected (B,C,H,W) shape."
        original_shape = x.shape[2:]
        if self.auto_resize_in and original_shape != self.input_shape:
            x = TVF.resize(x, self.input_shape, PIL.Image.NEAREST)
        if x.dtype == torch.uint8:
            x = x / 255.0
        if self.imagenet_pretrained:
            x = self.preprocess(x)
        y = self.underlying(x)
        if self.auto_resize_out:
            y = TVF.resize(y, original_shape, PIL.Image.NEAREST)
        y = torch.squeeze(y, dim=1)
        return y
