import torch
import torch.nn as nn

from cg.models.encoders.base import Encoder1C, Encoder3C

class LieEncoder1C(Encoder1C):
    def __init__(self, config):
        super(LieEncoder1C, self).__init__(config)

        # modules = []
        self.to_means = nn.ModuleList([])
        self.to_logvar = nn.ModuleList([])

        self.subgroup_sizes_ls = config.subgroup_sizes_ls
        self.subspace_sizes_ls = config.subspace_sizes_ls
        self.latent_dim = config.latent_dim
        self.hidden_states = config.hidden_states
        self.num_sampling = config.num_sampling

        self.dense1 = nn.Linear(256, 256)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(256, sum(self.subgroup_sizes_ls))
        self.active = nn.Sigmoid()

        for i, subgroup_size_i in enumerate(
            self.subgroup_sizes_ls
        ):  # subgroup_size int = latent_dim
            self.to_means.append(
                nn.Sequential(
                    nn.Linear(subgroup_size_i, subgroup_size_i * 4),
                    nn.ReLU(True),
                    nn.Linear(subgroup_size_i * 4, self.subspace_sizes_ls[i]),
                )
            )
            self.to_logvar.append(
                nn.Sequential(
                    nn.Linear(subgroup_size_i, subgroup_size_i * 4),
                    nn.ReLU(True),
                    nn.Linear(subgroup_size_i * 4, self.subspace_sizes_ls[i]),
                )
            )
        # self.hidden_layers = nn.ModuleList(modules)

    def forward(self, input):
        group_feats = input
        for i, hidden_layer in enumerate(self.hidden_layers):
            group_feats = hidden_layer(group_feats)

        group_feats = self.dense1(
            group_feats.view(group_feats.size(0), -1)
        )  # 4-D tensor: [Batch, *] --> 2-D tensor: [Batch, latent dim]
        group_feats = self.relu(group_feats)
        group_feats = self.dense2(group_feats)
        group_feats = self.active(group_feats)

        b_idx = 0
        means_ls, logvars_ls = [], []

        for i, subgroup_size_i in enumerate(self.subgroup_sizes_ls):
            e_idx = b_idx + subgroup_size_i
            means_ls.append(self.to_means[i](group_feats[:, b_idx:e_idx]))
            logvars_ls.append(self.to_logvar[i](group_feats[:, b_idx:e_idx]))
            b_idx = e_idx
        mu = torch.cat(means_ls, dim=-1)
        logvar = torch.cat(logvars_ls, dim=-1)
        z = self.reparameterization(mu, logvar)
        # z: after main encoder (E_{img} + E_{group}),
        # group_feats: (E_{img})
        outputs = (z, mu, logvar, group_feats)

        return outputs


class LieEncoder3C(Encoder3C):
    def __init__(self, config):
        super(LieEncoder3C, self).__init__(config)

        # modules = []
        self.to_means = nn.ModuleList([])
        self.to_logvar = nn.ModuleList([])

        self.subgroup_sizes_ls = config.subgroup_sizes_ls
        self.subspace_sizes_ls = config.subspace_sizes_ls
        self.latent_dim = config.latent_dim
        self.hidden_states = config.hidden_states
        self.num_sampling = config.num_sampling

        self.dense1 = nn.Linear(256, 256)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(256, sum(self.subgroup_sizes_ls))
        self.active = nn.Sigmoid()

        for i, subgroup_size_i in enumerate(
            self.subgroup_sizes_ls
        ):  # subgroup_size int = latent_dim
            self.to_means.append(
                nn.Sequential(
                    nn.Linear(subgroup_size_i, subgroup_size_i * 4),
                    nn.ReLU(True),
                    nn.Linear(subgroup_size_i * 4, self.subspace_sizes_ls[i]),
                )
            )
            self.to_logvar.append(
                nn.Sequential(
                    nn.Linear(subgroup_size_i, subgroup_size_i * 4),
                    nn.ReLU(True),
                    nn.Linear(subgroup_size_i * 4, self.subspace_sizes_ls[i]),
                )
            )
        # self.hidden_layers = nn.ModuleList(modules)

    def forward(self, input):
        group_feats = input
        for i, hidden_layer in enumerate(self.hidden_layers):
            group_feats = hidden_layer(group_feats)

        group_feats = self.dense1(
            group_feats.reshape(group_feats.size(0), -1)
        )  # 4-D tensor: [Batch, *] --> 2-D tensor: [Batch, latent dim]
        group_feats = self.relu(group_feats)
        group_feats = self.dense2(group_feats)
        group_feats = self.active(group_feats)

        b_idx = 0
        means_ls, logvars_ls = [], []

        for i, subgroup_size_i in enumerate(self.subgroup_sizes_ls):
            e_idx = b_idx + subgroup_size_i
            means_ls.append(self.to_means[i](group_feats[:, b_idx:e_idx]))
            logvars_ls.append(self.to_logvar[i](group_feats[:, b_idx:e_idx]))
            b_idx = e_idx
        mu = torch.cat(means_ls, dim=-1)
        logvar = torch.cat(logvars_ls, dim=-1)
        z = self.reparameterization(mu, logvar)
        outputs = (z, mu, logvar, group_feats)  # z: after main encoder

        return outputs