from model.AttentiveRNN import AttentiveRNN
from model.Autoencoder import AutoEncoder
from torch import nn
from torch import sum, pow, abs


class Generator(nn.Module):
    def __init__(self, repetition, blocks=3, layers=1, input_ch=3, out_ch=32, kernel_size=None, stride=1, padding=1, groups=1,
                 dilation=1):
        super(Generator, self).__init__()
        if kernel_size is None:
            kernel_size = [3, 3]
        self.attentiveRNN = AttentiveRNN(repetition,
            blocks=blocks, layers=layers, input_ch=input_ch, out_ch=out_ch,
            kernel_size=None, stride=stride, padding=padding, groups=groups, dilation=dilation
        )
        self.autoencoder = AutoEncoder()
        self.blocks = blocks
        self.layers = layers
        self.input_ch = input_ch
        self.out_ch = out_ch
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.groups = groups
        self.dilation = dilation
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        attentiveRNNresults = self.attentiveRNN(x)
        x = self.autoencoder(attentiveRNNresults['x'] * attentiveRNNresults['attention_map_list'][-1])
        ret = {
            'x' : x,
            'attention_maps' : attentiveRNNresults['attention_map_list']
        }
        return ret

    def binary_diff_mask(self, clean, dirty, thresold=0.1):
        # this parts corrects gamma, and always remember, sRGB values are not in linear scale with lights intensity,
        clean = pow(clean, 0.45)
        dirty = pow(dirty, 0.45)
        diff = abs(clean - dirty)
        diff = sum(diff, dim=1)

        bin_diff = (diff < thresold).to(clean.dtype)

        return bin_diff

    def loss(self, clean, dirty, thresold=0.1):
        # check diff if they are working as intended
        diff_mask = self.binary_diff_mask(clean, dirty, thresold)

        attentive_rnn_loss = self.attentiveRNN.loss(clean, diff_mask)
        autoencoder_loss = self.autoencoder.loss(clean, dirty)
        ret = {
            "attentive_rnn_loss" : attentive_rnn_loss,
            "autoencoder_loss" : autoencoder_loss,
        }
        return ret

if __name__ == "__main__":
    import torch
    from torchinfo import summary

    torch.set_default_tensor_type(torch.FloatTensor)
    generator = Generator(3, blocks=2)
    batch_size = 2
    summary(generator, input_size=(batch_size, 3, 720,720))