import torch
import torch.nn as nn

from models.init_weight import init_net
from models.model_blocks import AdaInResBlock
from models.model_blocks import ResBlock
from models.semantic_face_fusion_model import SemanticFaceFusionModule
from models.shape_aware_identity_model import ShapeAwareIdentityExtractor


class Encoder(nn.Module):
    """
    Hififace encoder part
    """

    def __init__(self):
        super(Encoder, self).__init__()
        self.conv_first = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)

        self.channel_list = [64, 128, 256, 512, 512, 512, 512, 512]
        self.down_sample = [True, True, True, True, True, False, False]

        self.block_list = nn.ModuleList()

        for i in range(7):
            self.block_list.append(
                ResBlock(self.channel_list[i], self.channel_list[i + 1], down_sample=self.down_sample[i])
            )

    def forward(self, x):
        x = self.conv_first(x)
        z_enc = None

        for i in range(7):
            x = self.block_list[i](x)
            if i == 1:
                z_enc = x
        return z_enc, x


class Decoder(nn.Module):
    """
    Hififace decoder part
    """

    def __init__(self):
        super(Decoder, self).__init__()
        self.block_list = nn.ModuleList()
        self.channel_list = [512, 512, 512, 512, 512, 256]
        self.up_sample = [False, False, True, True, True]

        for i in range(5):
            self.block_list.append(
                AdaInResBlock(self.channel_list[i], self.channel_list[i + 1], up_sample=self.up_sample[i])
            )

    def forward(self, x, id_vector):
        """
        Parameters:
        -----------
        x: encoder encoded feature map
        id_vector: 3d shape aware identity vector

        Returns:
        --------
        z_dec
        """
        for i in range(5):
            x = self.block_list[i](x, id_vector)
        return x


class Generator(nn.Module):
    """
    Hififace Generator
    """

    def __init__(self, identity_extractor_config):
        super(Generator, self).__init__()
        self.id_extractor = ShapeAwareIdentityExtractor(identity_extractor_config)
        self.id_extractor.requires_grad_(False)
        self.encoder = init_net(Encoder())
        self.decoder = init_net(Decoder())
        self.sff_module = init_net(SemanticFaceFusionModule())

    @torch.no_grad()
    def interp(self, i_source, i_target, shape_rate=1.0, id_rate=1.0):
        shape_aware_id_vector = self.id_extractor.interp(i_source, i_target, shape_rate, id_rate)
        z_enc, x = self.encoder(i_target)
        z_dec = self.decoder(x, shape_aware_id_vector)

        i_r, i_low, m_r, m_low = self.sff_module(i_target, z_enc, z_dec, shape_aware_id_vector)

        return i_r, i_low, m_r, m_low

    def forward(self, i_source, i_target, need_id_grad=False):
        """
        Parameters:
        -----------
        i_source: torch.Tensor, shape (B, 3, H, W), in range [0, 1], source face image
        i_target: torch.Tensor, shape (B, 3, H, W), in range [0, 1], target face image
        need_id_grad: bool, whether to calculate id extractor module's gradient

        Returns:
        --------
        i_r:    torch.Tensor
        i_low:  torch.Tensor
        m_r:    torch.Tensor
        m_low:  torch.Tensor
        """
        if need_id_grad:
            shape_aware_id_vector = self.id_extractor(i_source, i_target)
        else:
            with torch.no_grad():
                shape_aware_id_vector = self.id_extractor(i_source, i_target)
        z_enc, x = self.encoder(i_target)
        z_dec = self.decoder(x, shape_aware_id_vector)

        i_r, i_low, m_r, m_low = self.sff_module(i_target, z_enc, z_dec, shape_aware_id_vector)

        return i_r, i_low, m_r, m_low