import sys

import torch
import torch.nn as nn

sys.path.insert(0, "MobileStyleGAN.pytorch")

from core.models.mapping_network import MappingNetwork
from core.models.mobile_synthesis_network import MobileSynthesisNetwork
from core.models.synthesis_network import SynthesisNetwork


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        # teacher model
        mapping_net_params = {"style_dim": 512, "n_layers": 8, "lr_mlp": 0.01}
        synthesis_net_params = {
            "size": 1024,
            "style_dim": 512,
            "blur_kernel": [1, 3, 3, 1],
            "channels": [512, 512, 512, 512, 512, 256, 128, 64, 32],
        }
        self.mapping_net = MappingNetwork(**mapping_net_params).eval()
        self.synthesis_net = SynthesisNetwork(**synthesis_net_params).eval()
        # student network
        self.student = MobileSynthesisNetwork(
            style_dim=self.mapping_net.style_dim, channels=synthesis_net_params["channels"][:-1]
        )

        self.style_mean = nn.Parameter(torch.zeros((1, 512)), requires_grad=False)

    def forward(self, var: torch.Tensor, truncation_psi: float = 0.5, generator: str = "student") -> torch.Tensor:
        style = self.mapping_net(var)
        style = self.style_mean + truncation_psi * (style - self.style_mean)
        if generator == "student":
            img = self.student(style)["img"]
        elif generator == "teacher":
            img = self.synthesis_net(style)["img"]
        else:
            raise ValueError
        return img