import torch
import torch.nn as nn
import pdb
from dl.models.layers.base import CNNLayer

class Encoder1C(nn.Module):
    def __init__(self, config):
        super(Encoder1C, self).__init__()
        modules = []
        self.latent_dim = config.latent_dim
        self.hidden_states = config.hidden_states
        self.num_sampling = config.num_sampling

        # Design Encoder Factor-VAE ref
        modules.append(CNNLayer(in_channels=1, out_channels=32))
        modules.append(CNNLayer(in_channels=32, out_channels=32))
        modules.append(CNNLayer(in_channels=32, out_channels=64))
        modules.append(CNNLayer(in_channels=64, out_channels=64))
        self.hidden_layers = nn.ModuleList(modules)

        self.dense = nn.Linear(config.dense_dim[0], config.dense_dim[1])
        self.mu = nn.Linear(config.dense_dim[1], self.latent_dim)
        self.logvar = nn.Linear(config.dense_dim[1], self.latent_dim)

    def forward(self, input):
        all_hidden_states = ()

        output = input
        if self.hidden_states:
            all_hidden_states = all_hidden_states + (output,)
        for i, hidden_layer in enumerate(self.hidden_layers):
            output = hidden_layer(output)
            if self.hidden_states:
                all_hidden_states = all_hidden_states + (output,)
        # output = torch.flatten(output, start_dim=1)
        output = self.dense(
            output.contiguous().view(output.size(0), -1)
        )  # 4-D tensor: [Batch, *] --> 2-D tensor: [Batch, latent dim]
        # pdb.set_trace()
        mu = self.mu(output)  # [Batch, latent dim]
        logvar = self.logvar(output)  # [Batch, latent dim]

        z = self.reparameterization(mu, logvar)
        outputs = (
                      z,
                      mu,
                      logvar,
                  ) + (
                      all_hidden_states,
                  )  # (z, mu, logvar, (outputs))
        return outputs

    def reparameterization(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        if self.num_sampling == 1:
            eps = torch.randn_like(logvar)
            z = mu + std * eps
            return z
        else:
            batch, dim = std.size()
            eps = torch.randn(size=(self.num_sampling, batch, dim))
            z = mu + std * eps

            return z


class Encoder3C(Encoder1C):
    def __init__(self, config):
        super(Encoder3C, self).__init__(config)
        modules = []
        self.latent_dim = config.latent_dim
        self.hidden_states = config.hidden_states
        self.num_sampling = config.num_sampling

        # Design Encoder Factor-VAE ref
        modules.append(CNNLayer(in_channels=3, out_channels=32))
        modules.append(CNNLayer(in_channels=32, out_channels=32))
        modules.append(CNNLayer(in_channels=32, out_channels=64))
        modules.append(CNNLayer(in_channels=64, out_channels=64))
        self.hidden_layers = nn.ModuleList(modules)

        self.dense = nn.Linear(config.dense_dim[0], config.dense_dim[1])
        self.mu = nn.Linear(config.dense_dim[1], self.latent_dim)
        self.logvar = nn.Linear(config.dense_dim[1], self.latent_dim)


class MNISTEncoder1C(Encoder1C):
    def __init__(self, config):
        super(MNISTEncoder1C, self).__init__(config)
        modules = []
        self.latent_dim = config.latent_dim
        self.hidden_states = config.hidden_states
        self.num_sampling = config.num_sampling

        # Design Encoder Factor-VAE ref
        modules.append(CNNLayer(in_channels=1, out_channels=16))
        modules.append(CNNLayer(in_channels=16, out_channels=32))
        modules.append(CNNLayer(in_channels=32, out_channels=32))
        self.hidden_layers = nn.ModuleList(modules)

        self.dense = nn.Linear(32, config.dense_dim[1])
        self.mu = nn.Linear(config.dense_dim[1], self.latent_dim)
        self.logvar = nn.Linear(config.dense_dim[1], self.latent_dim)
