import torch
from piqa import SSIM
import torchvision


class DecoderBlock(torch.nn.Module):
    def __init__(self, in_channels, inter, out_channels, dropout_rate, bias, last=False):
        super().__init__()
        self.last = last
        self.conv1 = torch.nn.Conv2d(in_channels=in_channels, out_channels=inter, kernel_size=3, padding=1, bias=bias)
        self.d1 = torch.nn.Dropout(p=dropout_rate)
        self.b1 = torch.nn.BatchNorm2d(inter)
        self.relu1 = torch.nn.LeakyReLU(0.1)
        self.conv2 = torch.nn.Conv2d(inter, inter, kernel_size=3, padding=1, bias=bias)
        self.d2 = torch.nn.Dropout(p=dropout_rate)
        self.b2 = torch.nn.BatchNorm2d(inter)
        self.relu2 = torch.nn.LeakyReLU(0.1)
        if not self.last:
            self.convT1 = torch.nn.ConvTranspose2d(inter, out_channels, kernel_size=2, stride=2)
        else:
            self.conv3 = torch.nn.Conv2d(in_channels=inter, out_channels=out_channels, kernel_size=1, bias=bias)

    def forward(self, x):
        x = self.conv1(x)
        x = self.b1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.b2(x)
        x = self.relu2(x)
        x = self.d1(x)
        if not self.last:
            x = self.convT1(x)
        else:
            x = self.conv3(x)
        return x


class EncoderBlock(torch.nn.Module):
    def __init__(self, in_channels, out_channels, dropout_rate, bias):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=bias)
        self.d1 = torch.nn.Dropout(p=dropout_rate)
        self.b1 = torch.nn.BatchNorm2d(out_channels)
        self.relu1 = torch.nn.LeakyReLU(0.1)
        self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=bias)
        self.d2 = torch.nn.Dropout(p=dropout_rate)
        self.b2 = torch.nn.BatchNorm2d(out_channels)
        self.relu2 = torch.nn.LeakyReLU(0.1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.d1(x)
        x = self.b1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.d2(x)
        x = self.b2(x)
        x = self.relu2(x)
        return x


class UNet(torch.nn.Module):
    def __init__(self, in_channels=3, n_class=3, dropout_rate=0.0, magnitude=64):
        super().__init__()
        self.encoder_block1 = EncoderBlock(in_channels=in_channels, out_channels=magnitude, dropout_rate=dropout_rate, bias=True)
        self.encoder_block2 = EncoderBlock(in_channels=magnitude, out_channels=2*magnitude, dropout_rate=dropout_rate, bias=True)
        self.encoder_block3 = EncoderBlock(in_channels=2*magnitude, out_channels=4*magnitude, dropout_rate=dropout_rate, bias=True)
        self.encoder_block4 = EncoderBlock(in_channels=4*magnitude, out_channels=8*magnitude, dropout_rate=dropout_rate, bias=True)
        self.encoder_block5 = EncoderBlock(in_channels=8*magnitude, out_channels=16*magnitude, dropout_rate=dropout_rate, bias=True)
        self.bridge = torch.nn.ConvTranspose2d(16 * magnitude, 8 * magnitude, kernel_size=2, stride=2)
        self.pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.decoder_block1 = DecoderBlock(in_channels=16*magnitude, inter=8*magnitude, out_channels=4*magnitude,
                                           dropout_rate=dropout_rate, bias=True)
        self.decoder_block2 = DecoderBlock(in_channels=8*magnitude, inter=4*magnitude, out_channels=2*magnitude,
                                           dropout_rate=dropout_rate, bias=True)
        self.decoder_block3 = DecoderBlock(in_channels=4*magnitude, inter=2*magnitude, out_channels=magnitude,
                                           dropout_rate=dropout_rate, bias=True)
        self.decoder_block4 = DecoderBlock(in_channels=2*magnitude, inter=magnitude, out_channels=n_class,
                                           dropout_rate=dropout_rate, bias=True, last=True)
        self.diff = torch.tensor(0.0)

    def forward(self, x):
        # Encoder
        residual = x
        e_block1 = self.encoder_block1(x)
        x = self.pool(e_block1)
        e_block2 = self.encoder_block2(x)
        x = self.pool(e_block2)
        e_block3 = self.encoder_block3(x)
        x = self.pool(e_block3)
        e_block4 = self.encoder_block4(x)
        x = self.pool(e_block4)
        x = self.encoder_block5(x)
        x = self.bridge(x)
        # Decoder
        x = torch.cat([x, e_block4], dim=1)
        x = self.decoder_block1(x)
        x = torch.cat([x, e_block3], dim=1)
        x = self.decoder_block2(x)
        x = torch.cat([x, e_block2], dim=1)
        x = self.decoder_block3(x)
        x = torch.cat([x, e_block1], dim=1)
        x = self.decoder_block4(x)
        self._regularize(x)
        x += residual
        return x

    def _regularize(self, x):
        self.diff = torch.norm(x)


class SSIMLoss(SSIM):
    def __init__(self):
        super().__init__()
        self.relu6 = torch.nn.ReLU6()

    def forward(self, x, y):
        return 1. - super().forward(self.relu6(x * 6.0) / 6.0, y)


class ReLU1(torch.nn.Module):
    def __init__(self, size):
        super().__init__()
        self.relu6 = torch.nn.ReLU6()
        self.size = size

    def forward(self, x):
        return self.relu6(x * 6.0) / 6.0


class L1_SSIM(torch.nn.Module):
    def __init__(self, alpha=0.5):
        super().__init__()
        self.alpha = alpha
        self.ssim_loss = SSIMLoss()
        self.l1_loss = torch.nn.L1Loss()

    def forward(self, x, target):
        ssim = self.ssim_loss(x, target)
        l1 = self.l1_loss(x, target)
        loss = self.alpha * ssim + (1 - self.alpha) * l1
        return loss.mean(), ssim.mean(), l1.mean()


class L1_DICE(torch.nn.Module):
    def __init__(self, alpha=0.5):
        super().__init__()
        self.alpha = alpha
        self.l1_loss = torch.nn.L1Loss()

    def forward(self, inputs, targets, smooth=1.):
        l1 = self.l1_loss(inputs, targets)
        inputs = inputs.view(-1)
        targets = targets.view(-1)
        intersection = (inputs * targets).sum()
        union = torch.sum(inputs * inputs) + torch.sum(targets * targets)
        dice_loss = 1 - (2.*intersection + smooth)/(union + smooth)
        loss = (self.alpha * (l1) + (1 - self.alpha) * dice_loss)
        return loss.mean(), dice_loss.mean(), l1.mean()


class VGGPerceptualLoss(torch.nn.Module):
    def __init__(self, resize=True):
        super(VGGPerceptualLoss, self).__init__()
        self.device = torch.device('cuda')
        blocks = []
        blocks.append(torchvision.models.vgg16(pretrained=True).features[:4].eval().to(self.device))
        blocks.append(torchvision.models.vgg16(pretrained=True).features[4:9].eval().to(self.device))
        blocks.append(torchvision.models.vgg16(pretrained=True).features[9:16].eval().to(self.device))
        blocks.append(torchvision.models.vgg16(pretrained=True).features[16:23].eval().to(self.device))
        for bl in blocks:
            for p in bl.parameters():
                p.requires_grad = False
        self.blocks = torch.nn.ModuleList(blocks).to(self.device)
        self.transform = torch.nn.functional.interpolate
        self.resize = resize
        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(self.device))
        self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(self.device))

    def forward(self, input, target, feature_layers=[0, 1, 2, 3], style_layers=[]):
        if input.shape[1] != 3:
            input = input.repeat(1, 3, 1, 1)
            target = target.repeat(1, 3, 1, 1)
        input = (input-self.mean) / self.std
        target = (target-self.mean) / self.std
        if self.resize:
            input = self.transform(input, mode='bilinear', size=(224, 224), align_corners=False)
            target = self.transform(target, mode='bilinear', size=(224, 224), align_corners=False)
        loss = 0.0
        x = input
        y = target
        for i, block in enumerate(self.blocks):
            x = block(x)
            y = block(y)
            if i in feature_layers:
                loss += torch.nn.functional.l1_loss(x, y)
            if i in style_layers:
                act_x = x.reshape(x.shape[0], x.shape[1], -1)
                act_y = y.reshape(y.shape[0], y.shape[1], -1)
                gram_x = act_x @ act_x.permute(0, 2, 1)
                gram_y = act_y @ act_y.permute(0, 2, 1)
                loss += torch.nn.functional.l1_loss(gram_x, gram_y)
        return loss, 0, 0


class BatchNormConv2d(torch.nn.Module):
    def __init__(self, input_channels, output_channels, resample_mode):
        super(BatchNormConv2d, self).__init__()
        self.input_channels = input_channels
        self.output_channels = output_channels

        if resample_mode == "down":
            self.kernel_size = (5, 5)
            self.stride = (2, 2)
            self.padding = (2, 2)
            self.conv = torch.nn.Conv2d(input_channels, output_channels, kernel_size=self.kernel_size,
                                        stride=self.stride, padding=self.padding)
        elif resample_mode == "up":
            self.kernel_size = (5, 5)
            self.stride = (2, 2)
            self.padding = (2, 2)
            self.out_padding = (1, 1)
            self.conv = torch.nn.ConvTranspose2d(input_channels, output_channels, kernel_size=self.kernel_size,
                                                 stride=self.stride, padding=self.padding, output_padding=self.out_padding)
        else:
            self.kernel_size = (3, 3)
            self.stride = (1, 1)
            self.padding = (1, 1)
            self.conv = torch.nn.Conv2d(input_channels, output_channels, kernel_size=self.kernel_size,
                                        stride=self.stride, padding=self.padding)

        self.norm = torch.nn.BatchNorm2d(output_channels)
        self.activation = torch.nn.SiLU()

    def forward(self, inputs):
        conv_out = self.conv(inputs)
        norm_out = self.norm(conv_out)
        return self.activation(norm_out)


class InvertedResidual(torch.nn.Module):
    def __init__(self, input_channels, expand_channels, squeeze_channels, downsample_flag):
        super(InvertedResidual, self).__init__()
        self.input_channels = input_channels
        self.expand_channels = expand_channels
        self.conv_expand = torch.nn.Conv2d(input_channels, expand_channels, kernel_size=(1, 1), stride=(1, 1))
        self.conv_contract = SeparableConv2d(expand_channels, input_channels, squeeze_channels, False)
        self.norm = torch.nn.BatchNorm2d(expand_channels)
        self.activation = torch.nn.SiLU()

    def forward(self, inputs):
        expand_out = self.conv_expand(inputs)
        expand_out = self.activation(self.norm(expand_out))
        return inputs + self.conv_contract(expand_out)


class SeparableConv2d(torch.nn.Module):
    def __init__(self, input_channels, output_channels, squeeze_channels, downsample_flag):
        super(SeparableConv2d, self).__init__()
        self.input_channels = input_channels
        self.conv_depthwise = torch.nn.Conv2d(input_channels, input_channels, kernel_size=(3, 3),
                                              stride=(1, 1), padding=(1, 1), groups=input_channels)
        self.norm_depthwise = torch.nn.BatchNorm2d(input_channels)
        self.activation = torch.nn.SiLU()
        self.squeeze_excite_flag = squeeze_channels > 0
        if self.squeeze_excite_flag:
            self.conv_reduce = torch.nn.Conv2d(input_channels, squeeze_channels, kernel_size=(1, 1), stride=(1, 1))
            self.conv_expand = torch.nn.Conv2d(squeeze_channels, input_channels, kernel_size=(1, 1), stride=(1, 1))
        if downsample_flag:
            self.conv_pointwise = torch.nn.Conv2d(input_channels, output_channels, kernel_size=(1, 1), stride=(2, 2))
        else:
            self.conv_pointwise = torch.nn.Conv2d(input_channels, output_channels, kernel_size=(1, 1), stride=(1, 1))
        self.norm_pointwise = torch.nn.BatchNorm2d(output_channels)
        self.downsample_flag = downsample_flag

    def forward(self, inputs):
        depthwise_out = self.conv_depthwise(inputs)
        depthwise_out = self.norm_depthwise(depthwise_out)
        depthwise_out = self.activation(depthwise_out)
        if self.squeeze_excite_flag is True:
            squeezed_out = torch.nn.functional.adaptive_avg_pool2d(depthwise_out, 1)
            squeezed_out = self.conv_reduce(squeezed_out)
            squeezed_out = self.activation(squeezed_out)
            squeezed_out = self.conv_expand(squeezed_out)
            squeezed_out = torch.sigmoid(squeezed_out)
            depthwise_out = squeezed_out * depthwise_out

        pointwise_out = self.conv_pointwise(depthwise_out)
        pointwise_out = self.norm_pointwise(pointwise_out)
        pointwise_out = self.activation(pointwise_out)
        return pointwise_out


class BarcodeDecoder(torch.nn.Module):
    def __init__(self):
        super(BarcodeDecoder, self).__init__()
        self.input_channels = 1
        self.activation = torch.nn.SiLU()
        self.init_num_channels = 64
        self.layers = torch.nn.Sequential()
        # self.layers.append(torch.nn.Linear(25600,51200))
        current_num_channels = self.init_num_channels * 2
        self.layers.append(torch.nn.Unflatten(1, (current_num_channels, 40, 160)))
        self.layers.append(InvertedResidual(current_num_channels, current_num_channels*2, 0, False))
        self.layers.append(BatchNormConv2d(current_num_channels, self.init_num_channels, "up"))
        self.layers.append(InvertedResidual(self.init_num_channels, self.init_num_channels*2, 0, False))
        self.layers.append(torch.nn.ConvTranspose2d(self.init_num_channels, self.input_channels,
                                                    kernel_size=(5, 5), stride=(2, 2), padding=(2, 2), output_padding=(1, 1)))
        self.layers.append(torch.nn.Tanh())

    def forward(self, inputs):
        return self.layers(inputs)


class BarcodeEncoder(torch.nn.Module):
    def __init__(self):
        super(BarcodeEncoder, self).__init__()
        self.input_channels = 1
        self.activation = torch.nn.SiLU()
        self.init_num_channels = 64
        self.layers = torch.nn.Sequential()
        current_num_channels = self.init_num_channels
        self.layers.append(BatchNormConv2d(self.input_channels, current_num_channels, "down"))
        self.layers.append(InvertedResidual(current_num_channels, current_num_channels*2, 0, False))
        self.layers.append(BatchNormConv2d(current_num_channels, current_num_channels*2, "down"))
        current_num_channels *= 2
        self.layers.append(InvertedResidual(current_num_channels, current_num_channels*2, 0, False))
        self.flatten = torch.nn.Flatten(start_dim=1)
        self.layers.append(self.flatten)
        # self.layers.append(torch.nn.Linear(51200,25600))

    def forward(self, inputs):
        return self.layers(inputs)
