import torch
from torch import nn

# import torch.nn.functional as F
from .blocks import AdaOctConv, Oct_Conv_aftup, Oct_conv_lreLU, Oct_conv_up, OctConv

# from blocks import *


def define_network(
    net_type, alpha_in, alpha_out, style_kernel, input_nc, nf, output_nc, freq_ratio
):
    net = None
    sk = style_kernel

    if net_type == "Encoder":
        net = Encoder(
            in_dim=input_nc,
            nf=nf,
            style_kernel=[sk, sk],
            alpha_in=alpha_in,
            alpha_out=alpha_out,
        )
    elif net_type == "Generator":
        net = Decoder(
            nf=nf,
            out_dim=output_nc,
            style_channel=256,
            style_kernel=[sk, sk, 3],
            alpha_in=alpha_in,
            freq_ratio=freq_ratio,
            alpha_out=alpha_out,
        )
    return net


class Encoder(nn.Module):
    def __init__(self, in_dim, nf=64, style_kernel=[3, 3], alpha_in=0.5, alpha_out=0.5):
        super(Encoder, self).__init__()

        self.conv = nn.Conv2d(
            in_channels=in_dim, out_channels=nf, kernel_size=7, stride=1, padding=3
        )

        self.OctConv1_1 = OctConv(
            in_channels=nf,
            out_channels=nf,
            kernel_size=3,
            stride=2,
            padding=1,
            groups=64,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="first",
        )
        self.OctConv1_2 = OctConv(
            in_channels=nf,
            out_channels=2 * nf,
            kernel_size=1,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="normal",
        )
        self.OctConv1_3 = OctConv(
            in_channels=2 * nf,
            out_channels=2 * nf,
            kernel_size=3,
            stride=1,
            padding=1,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="normal",
        )

        self.OctConv2_1 = OctConv(
            in_channels=2 * nf,
            out_channels=2 * nf,
            kernel_size=3,
            stride=2,
            padding=1,
            groups=128,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="normal",
        )
        self.OctConv2_2 = OctConv(
            in_channels=2 * nf,
            out_channels=4 * nf,
            kernel_size=1,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="normal",
        )
        self.OctConv2_3 = OctConv(
            in_channels=4 * nf,
            out_channels=4 * nf,
            kernel_size=3,
            stride=1,
            padding=1,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="normal",
        )

        self.pool_h = nn.AdaptiveAvgPool2d((style_kernel[0], style_kernel[0]))
        self.pool_l = nn.AdaptiveAvgPool2d((style_kernel[1], style_kernel[1]))

        self.relu = Oct_conv_lreLU()

    def forward(self, x):
        enc_feat = []
        out = self.conv(x)

        out = self.OctConv1_1(out)
        out = self.relu(out)
        out = self.OctConv1_2(out)
        out = self.relu(out)
        out = self.OctConv1_3(out)
        out = self.relu(out)
        enc_feat.append(out)

        out = self.OctConv2_1(out)
        out = self.relu(out)
        out = self.OctConv2_2(out)
        out = self.relu(out)
        out = self.OctConv2_3(out)
        out = self.relu(out)
        enc_feat.append(out)

        out_high, out_low = out
        out_sty_h = self.pool_h(out_high)
        out_sty_l = self.pool_l(out_low)
        out_sty = out_sty_h, out_sty_l

        return out, out_sty, enc_feat

    def forward_test(self, x, cond):
        out = self.conv(x)

        out = self.OctConv1_1(out)
        out = self.relu(out)
        out = self.OctConv1_2(out)
        out = self.relu(out)
        out = self.OctConv1_3(out)
        out = self.relu(out)

        out = self.OctConv2_1(out)
        out = self.relu(out)
        out = self.OctConv2_2(out)
        out = self.relu(out)
        out = self.OctConv2_3(out)
        out = self.relu(out)

        if cond == "style":
            out_high, out_low = out
            out_sty_h = self.pool_h(out_high)
            out_sty_l = self.pool_l(out_low)
            return out_sty_h, out_sty_l
        else:
            return out


class Decoder(nn.Module):
    def __init__(
        self,
        nf=64,
        out_dim=3,
        style_channel=512,
        style_kernel=[3, 3, 3],
        alpha_in=0.5,
        alpha_out=0.5,
        freq_ratio=[1, 1],
        pad_type="reflect",
    ):
        super(Decoder, self).__init__()

        group_div = [1, 2, 4, 8]
        self.up_oct = Oct_conv_up(scale_factor=2)

        self.AdaOctConv1_1 = AdaOctConv(
            in_channels=4 * nf,
            out_channels=4 * nf,
            group_div=group_div[0],
            style_channels=style_channel,
            kernel_size=style_kernel,
            stride=1,
            padding=1,
            oct_groups=4 * nf,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="normal",
        )
        self.OctConv1_2 = OctConv(
            in_channels=4 * nf,
            out_channels=2 * nf,
            kernel_size=1,
            stride=1,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="normal",
        )
        self.oct_conv_aftup_1 = Oct_Conv_aftup(
            in_channels=2 * nf,
            out_channels=2 * nf,
            kernel_size=3,
            stride=1,
            padding=1,
            pad_type=pad_type,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
        )

        self.AdaOctConv2_1 = AdaOctConv(
            in_channels=2 * nf,
            out_channels=2 * nf,
            group_div=group_div[1],
            style_channels=style_channel,
            kernel_size=style_kernel,
            stride=1,
            padding=1,
            oct_groups=2 * nf,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="normal",
        )
        self.OctConv2_2 = OctConv(
            in_channels=2 * nf,
            out_channels=nf,
            kernel_size=1,
            stride=1,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="normal",
        )
        self.oct_conv_aftup_2 = Oct_Conv_aftup(
            nf, nf, 3, 1, 1, pad_type, alpha_in, alpha_out
        )

        self.AdaOctConv3_1 = AdaOctConv(
            in_channels=nf,
            out_channels=nf,
            group_div=group_div[2],
            style_channels=style_channel,
            kernel_size=style_kernel,
            stride=1,
            padding=1,
            oct_groups=nf,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="normal",
        )
        self.OctConv3_2 = OctConv(
            in_channels=nf,
            out_channels=nf // 2,
            kernel_size=1,
            stride=1,
            alpha_in=alpha_in,
            alpha_out=alpha_out,
            type="last",
            freq_ratio=freq_ratio,
        )

        self.conv4 = nn.Conv2d(in_channels=nf // 2, out_channels=out_dim, kernel_size=1)

    def forward(self, content, style):
        out = self.AdaOctConv1_1(content, style)
        out = self.OctConv1_2(out)
        out = self.up_oct(out)
        out = self.oct_conv_aftup_1(out)

        out = self.AdaOctConv2_1(out, style)
        out = self.OctConv2_2(out)
        out = self.up_oct(out)
        out = self.oct_conv_aftup_2(out)

        out = self.AdaOctConv3_1(out, style)
        out = self.OctConv3_2(out)
        out, out_high, out_low = out

        out = self.conv4(out)
        out_high = self.conv4(out_high)
        out_low = self.conv4(out_low)

        return out, out_high, out_low

    def forward_test(self, content, style):
        out = self.AdaOctConv1_1(content, style, "test")
        out = self.OctConv1_2(out)
        out = self.up_oct(out)
        out = self.oct_conv_aftup_1(out)

        out = self.AdaOctConv2_1(out, style, "test")
        out = self.OctConv2_2(out)
        out = self.up_oct(out)
        out = self.oct_conv_aftup_2(out)

        out = self.AdaOctConv3_1(out, style, "test")
        out = self.OctConv3_2(out)

        out = self.conv4(out[0])
        return out


############## Contrastive Loss function ##############
def calc_mean_std(feat, eps=1e-5):
    # eps is a small value added to the variance to avoid divide-by-zero.
    size = feat.size()
    assert len(size) == 4
    N, C = size[:2]
    feat_var = feat.view(N, C, -1).var(dim=2) + eps
    feat_std = feat_var.sqrt().view(N, C, 1, 1)
    feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
    return feat_mean, feat_std


def calc_content_loss(input, target):
    assert input.size() == target.size()
    mse_loss = nn.MSELoss()
    return mse_loss(input, target)


def calc_style_loss(input, target):
    assert input.size() == target.size()
    mse_loss = nn.MSELoss()
    input_mean, input_std = calc_mean_std(input)
    target_mean, target_std = calc_mean_std(target)

    loss = mse_loss(input_mean, target_mean) + mse_loss(input_std, target_std)
    return loss


class EFDM_loss(nn.Module):
    def __init__(self):
        super(EFDM_loss, self).__init__()
        self.mse_loss = nn.MSELoss()

    def efdm_single(self, style, trans):
        B, C, W, H = style.size(0), style.size(1), style.size(2), style.size(3)

        value_style, index_style = torch.sort(style.view(B, C, -1))
        value_trans, index_trans = torch.sort(trans.view(B, C, -1))
        inverse_index = index_trans.argsort(-1)

        return self.mse_loss(
            trans.view(B, C, -1), value_style.gather(-1, inverse_index)
        )

    def forward(self, style_E, style_S, translate_E, translate_S, neg_idx):
        loss = 0.0
        batch = style_E[0][0].shape[0]
        for b in range(batch):
            poss_loss = 0.0
            neg_loss = 0.0

            # Positive loss
            for i in range(len(style_E)):
                poss_loss += self.efdm_single(
                    style_E[i][0][b].unsqueeze(0), translate_E[i][0][b].unsqueeze(0)
                ) + self.efdm_single(
                    style_E[i][1][b].unsqueeze(0), translate_E[i][1][b].unsqueeze(0)
                )
            for i in range(len(style_S)):
                poss_loss += self.efdm_single(
                    style_S[i][0][b].unsqueeze(0), translate_S[i][0][b].unsqueeze(0)
                ) + self.efdm_single(
                    style_S[i][1][b].unsqueeze(0), translate_S[i][1][b].unsqueeze(0)
                )

            # Negative loss
            for nb in neg_idx[b]:
                for i in range(len(style_E)):
                    neg_loss += self.efdm_single(
                        style_E[i][0][nb].unsqueeze(0),
                        translate_E[i][0][b].unsqueeze(0),
                    ) + self.efdm_single(
                        style_E[i][1][nb].unsqueeze(0),
                        translate_E[i][1][b].unsqueeze(0),
                    )
                for i in range(len(style_S)):
                    neg_loss += self.efdm_single(
                        style_S[i][0][nb].unsqueeze(0),
                        translate_S[i][0][b].unsqueeze(0),
                    ) + self.efdm_single(
                        style_S[i][1][nb].unsqueeze(0),
                        translate_S[i][1][b].unsqueeze(0),
                    )

            loss += poss_loss / neg_loss

        return loss