import torch; torch.manual_seed(0)
import torch.nn as nn
import torch.nn.functional as F
import torch.utils
import torch.distributions
import matplotlib.pyplot as plt; plt.rcParams['figure.dpi'] = 200
from src.cocktails.representation_learning.simple_model import SimpleNet

device = 'cuda' if torch.cuda.is_available() else 'cpu'

def get_activation(activation):
    if activation == 'tanh':
        activ = F.tanh
    elif activation == 'relu':
        activ = F.relu
    elif activation == 'mish':
        activ = F.mish
    elif activation == 'sigmoid':
        activ = F.sigmoid
    elif activation == 'leakyrelu':
        activ = F.leaky_relu
    elif activation == 'exp':
        activ = torch.exp
    else:
        raise ValueError
    return activ

class IngredientEncoder(nn.Module):
    def __init__(self, input_dim, deepset_latent_dim, hidden_dims, activation, dropout):
        super(IngredientEncoder, self).__init__()
        self.linears = nn.ModuleList()
        self.dropouts = nn.ModuleList()
        dims = [input_dim] + hidden_dims + [deepset_latent_dim]
        for d_in, d_out in zip(dims[:-1], dims[1:]):
            self.linears.append(nn.Linear(d_in, d_out))
            self.dropouts.append(nn.Dropout(dropout))
        self.activation = get_activation(activation)
        self.n_layers = len(self.linears)
        self.layer_range = range(self.n_layers)

    def forward(self, x):
        for i_layer, layer, dropout in zip(self.layer_range, self.linears, self.dropouts):
            x = layer(x)
            if i_layer != self.n_layers - 1:
                x = self.activation(dropout(x))
        return x  # do not use dropout on last layer?

class DeepsetCocktailEncoder(nn.Module):
    def __init__(self, input_dim, deepset_latent_dim, hidden_dims_ing, activation,
                 hidden_dims_cocktail, latent_dim, aggregation, dropout):
        super(DeepsetCocktailEncoder, self).__init__()
        self.input_dim = input_dim  # dimension of ingredient representation + quantity
        self.ingredient_encoder = IngredientEncoder(input_dim, deepset_latent_dim, hidden_dims_ing, activation, dropout)  # encode each ingredient separately
        self.deepset_latent_dim = deepset_latent_dim  # dimension of the deepset aggregation
        self.aggregation = aggregation
        self.latent_dim = latent_dim
        # post aggregation network
        self.linears = nn.ModuleList()
        self.dropouts = nn.ModuleList()
        dims = [deepset_latent_dim] + hidden_dims_cocktail
        for d_in, d_out in zip(dims[:-1], dims[1:]):
            self.linears.append(nn.Linear(d_in, d_out))
            self.dropouts.append(nn.Dropout(dropout))
        self.FC_mean  = nn.Linear(hidden_dims_cocktail[-1], latent_dim)
        self.FC_logvar   = nn.Linear(hidden_dims_cocktail[-1], latent_dim)
        self.softplus = nn.Softplus()

        self.activation = get_activation(activation)
        self.n_layers = len(self.linears)
        self.layer_range = range(self.n_layers)

    def forward(self, nb_ingredients, x):

        # reshape x in (batch size * nb ingredients, dim_ing_rep)
        batch_size = x.shape[0]
        all_ingredients = []
        for i in range(batch_size):
            for j in range(nb_ingredients[i]):
                all_ingredients.append(x[i, self.input_dim * j: self.input_dim * (j + 1)].reshape(1, -1))
        x = torch.cat(all_ingredients, dim=0)
        # encode ingredients in parallel
        ingredients_encodings = self.ingredient_encoder(x)
        assert ingredients_encodings.shape == (torch.sum(nb_ingredients), self.deepset_latent_dim)

        # aggregate
        x = []
        index_first = 0
        for i in range(batch_size):
            index_last = index_first + nb_ingredients[i]
            # aggregate
            if self.aggregation == 'sum':
                x.append(torch.sum(ingredients_encodings[index_first:index_last], dim=0).reshape(1, -1))
            elif self.aggregation == 'mean':
                x.append(torch.mean(ingredients_encodings[index_first:index_last], dim=0).reshape(1, -1))
            else:
                raise ValueError
            index_first = index_last
        x = torch.cat(x, dim=0)
        assert x.shape[0] == batch_size

        for i_layer, layer, dropout in zip(self.layer_range, self.linears, self.dropouts):
            x = self.activation(dropout(layer(x)))
        mean = self.FC_mean(x)
        logvar = self.FC_logvar(x)
        return mean, logvar


class MultiHeadModel(nn.Module):
    def __init__(self, encoder, auxiliaries_dict, activation, hidden_dims_decoder):
        super(MultiHeadModel, self).__init__()
        self.encoder = encoder
        self.latent_dim = self.encoder.output_dim
        self.auxiliaries_str = []
        self.auxiliaries = nn.ModuleList()
        for aux_str in sorted(auxiliaries_dict.keys()):
            if aux_str == 'taste_reps':
                self.taste_reps_decoder = SimpleNet(input_dim=self.latent_dim, hidden_dims=[], output_dim=auxiliaries_dict[aux_str]['dim_output'],
                                                    activation=activation, dropout=0.0, final_activ=auxiliaries_dict[aux_str]['final_activ'])
            else:
                self.auxiliaries_str.append(aux_str)
                if aux_str == 'ingredients_quantities':
                    hd = hidden_dims_decoder
                else:
                    hd = []
                self.auxiliaries.append(SimpleNet(input_dim=self.latent_dim, hidden_dims=hd, output_dim=auxiliaries_dict[aux_str]['dim_output'],
                                                  activation=activation, dropout=0.0, final_activ=auxiliaries_dict[aux_str]['final_activ']))

    def get_all_auxiliaries(self, x):
        return [aux(x) for aux in self.auxiliaries]

    def get_auxiliary(self, z, aux_str):
        if aux_str == 'taste_reps':
            return self.taste_reps_decoder(z)
        else:
            index = self.auxiliaries_str.index(aux_str)
            return self.auxiliaries[index](z)

    def forward(self, x, aux_str=None):
        z = self.encoder(x)
        if aux_str is not None:
            return z, self.get_auxiliary(z, aux_str), [aux_str]
        else:
            return z, self.get_all_auxiliaries(z), self.auxiliaries_str

def get_multihead_model(input_dim, activation, hidden_dims_cocktail, latent_dim, dropout, auxiliaries_dict, hidden_dims_decoder):
    encoder = SimpleNet(input_dim, hidden_dims_cocktail, latent_dim, activation, dropout)
    model = MultiHeadModel(encoder, auxiliaries_dict, activation, hidden_dims_decoder)
    return model