import torch
from utils import *


class AdaIN(torch.nn.Module):
    def __init__(self):
        super(AdaIN, self).__init__()
        # initialize instance normalization function
        # this is the basis of our AdaIN layer, it follows an equation similar to a z-score
        # (x - mu)/sigma
        self.instance_norm = torch.nn.InstanceNorm2d(3)
    
    # forward method for our layer
    # x would be the content input and y would be the style input
    # both x and y are tensors
    def forward(self, x, y):
        # size is shaped (N, num_channels, Height, Width)
        x_size = x.size()
        
        # we do not need these since they will be calculated by the instance normalization function
        #x_mean, x_std = mean_and_std_of_image(x)
        y_mean, y_std = mean_and_std_of_image(y)

        x_norm = self.instance_norm(x)

        
        print(x_norm.size())
        # expand size of tensors so that there are no shape errors when performing AdaIN operation
        # if not self.training:
        #     x_norm = x_norm.view(*x_norm.shape, 1)

        x_size = x_norm.size()
        print(x_size)
        return y_std.expand(x_size) * x_norm + y_mean.expand(x_size)