from typing import Optional

import torch
import torch.nn as nn

from models.stylegan2 import FullyConnectedLayer, normalize_2nd_moment
from models.rtm_core import RTMMappingNetwork


class RTMMappingNetworkStyleGAN2(nn.Module):
    def __init__(
        self,
        z_dim: int,
        c_dim: int,
        w_dim: int,
        num_ws: Optional[int],
        num_layers: int = 8,
        embed_features: Optional[int] = None,
        layer_features: Optional[int] = None,
        activation: str = "lrelu",
        lr_multiplier: float = 0.01,
        w_avg_beta: float = 0.998,
        rtm_num_tokens: int = 4,
        rtm_H_cycles: int = 4,
        rtm_L_cycles: int = 1,
        rtm_H_layers: int = 2,
        rtm_L_layers: int = 2,
        rtm_hidden_size: int = 128,
        rtm_expansion: float = 4.0,
        rtm_refinement_steps: int = 4,
        rtm_with_grad: bool = False,
        rtm_cycle_noise_std: float = 0.0,
        use_rtm_equalized: bool = True,
        rtm_lr_multiplier: float = 0.01,
    ):
        super().__init__()
        del layer_features
        self.z_dim = z_dim
        self.c_dim = c_dim
        self.w_dim = w_dim
        self.num_ws = num_ws
        self.num_layers = num_layers
        self.w_avg_beta = w_avg_beta

        if embed_features is None:
            embed_features = w_dim
        if c_dim == 0:
            embed_features = 0

        if c_dim > 0:
            self.embed = FullyConnectedLayer(c_dim, embed_features)
        else:
            self.embed = None

        self.fuse = FullyConnectedLayer(
            z_dim + embed_features, w_dim,
            activation=activation, lr_multiplier=lr_multiplier,
        )

        self.rtm = RTMMappingNetwork(
            code_dim=w_dim,
            num_tokens=rtm_num_tokens,
            H_cycles=rtm_H_cycles,
            L_cycles=rtm_L_cycles,
            H_layers=rtm_H_layers,
            L_layers=rtm_L_layers,
            hidden_size=rtm_hidden_size,
            expansion=rtm_expansion,
            refinement_steps=rtm_refinement_steps,
            with_grad=rtm_with_grad,
            cycle_noise_std=rtm_cycle_noise_std,
            use_equalized=use_rtm_equalized,
            lr_multiplier=rtm_lr_multiplier,
        )

        if num_ws is not None and w_avg_beta is not None:
            self.register_buffer("w_avg", torch.zeros([w_dim]))

    def forward(
        self,
        z: torch.Tensor,
        c: Optional[torch.Tensor],
        truncation_psi: float = 1.0,
        truncation_cutoff: Optional[int] = None,
        update_emas: bool = False,
    ) -> torch.Tensor:
        x = None
        if self.z_dim > 0:
            x = normalize_2nd_moment(z.to(torch.float32))
        if self.c_dim > 0:
            assert c is not None and self.embed is not None
            y = normalize_2nd_moment(self.embed(c.to(torch.float32)))
            x = torch.cat([x, y], dim=1) if x is not None else y

        x = self.fuse(x)
        x = self.rtm(x)
        if isinstance(x, (list, tuple)):
            x = x[-1]

        if update_emas and self.w_avg_beta is not None and hasattr(self, "w_avg"):
            self.w_avg.copy_(
                x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta)
            )

        if self.num_ws is not None:
            x = x.unsqueeze(1).repeat([1, self.num_ws, 1])

        if truncation_psi != 1:
            assert self.w_avg_beta is not None
            if self.num_ws is None or truncation_cutoff is None:
                x = self.w_avg.lerp(x, truncation_psi)
            else:
                x[:, :truncation_cutoff] = self.w_avg.lerp(
                    x[:, :truncation_cutoff], truncation_psi
                )
        return x

    def forward_ws_trajectory(
        self,
        z: torch.Tensor,
        c: Optional[torch.Tensor],
        truncation_psi: float = 1.0,
        truncation_cutoff: Optional[int] = None,
    ):
        x = None
        if self.z_dim > 0:
            x = normalize_2nd_moment(z.to(torch.float32))
        if self.c_dim > 0:
            assert c is not None and self.embed is not None
            y = normalize_2nd_moment(self.embed(c.to(torch.float32)))
            x = torch.cat([x, y], dim=1) if x is not None else y

        x = self.fuse(x)
        flat_list = self.rtm.forward_w_trajectory(x)

        ws_out = []
        for xf in flat_list:
            if isinstance(xf, (list, tuple)):
                xf = xf[-1]
            xcur = xf
            if self.num_ws is not None:
                xcur = xcur.unsqueeze(1).repeat([1, self.num_ws, 1])
            if truncation_psi != 1:
                assert self.w_avg_beta is not None
                if self.num_ws is None or truncation_cutoff is None:
                    xcur = self.w_avg.lerp(xcur, truncation_psi)
                else:
                    xcur[:, :truncation_cutoff] = self.w_avg.lerp(
                        xcur[:, :truncation_cutoff], truncation_psi
                    )
            ws_out.append(xcur)
        return ws_out