import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler as lrsch
from tqdm import tqdm
import numpy as np
from sklearn.mixture import GaussianMixture

from .estimator import TrainerBase
from . import block
from . import loss
from . import layers
from .utils import rsample_normal, rsample_categorical, temperature_func1


def gamma_func(e, max_value, minimum_epoch):
    return max((max_value - 1) * (1 - e / minimum_epoch) + 1, 1)


class VAEBase(TrainerBase):

    def get_optimizers(self):
        gen_optimizer = optim.Adam(self.parameters(), lr=self._lr)
        if self._kwargs.get("use_lrsch", False):
            scheduler = lrsch.CosineAnnealingWarmRestarts(gen_optimizer,
                                                          100, 2, 0.0001)
            return [gen_optimizer, scheduler]
        else:
            return [gen_optimizer]

    def _batch_end(self):
        if self._kwargs.get("use_lrsch", False):
            self._ops[-1].step(self._e + self._i / len(self._data_iter))

    def _epoch_start(self):
        self._alpha = 1.
        self._gamma = gamma_func(self._e,
                                 self._gamma_max_value,
                                 self._gamma_min_epoch)
        self._temp = temperature_func1(self._e, 0.01, 0.3, self._epoch)


class MClusterVAEsExact(VAEBase):

    def __init__(
        self, inpts, omic_oupts, n_clusters, omic_zdims, att_method="none",
        omic_hiddens=(), enc_hiddens=[200], enc_z_hiddens=[200],
        dec_hiddens=[200], gamma_max_value=5.0, gamma_min_epoch=300
    ):
        super().__init__()
        self._n_clusters = n_clusters
        self._inpts = inpts
        self._att_method = att_method
        self._gamma_max_value = gamma_max_value
        self._gamma_min_epoch = gamma_min_epoch

        self.encoder_y = block.Encoder(inpts, omic_oupts, [n_clusters],
                                       omic_hiddens, enc_hiddens, att_method)
        self.encoders_z = nn.ModuleList([
            block.VAEncoder(i+n_clusters, o, enc_z_hiddens)
            for i, o in zip(inpts, omic_zdims)
        ])
        self.decoders = nn.ModuleList([
            block.VADncoder(i+n_clusters, o, dec_hiddens)
            for i, o in zip(omic_zdims, inpts)
        ])
        self.z_maps = nn.ModuleList([
            nn.Sequential(
                nn.Linear(n_clusters, o * 2, bias=False),
                layers.SplitModule(o, o)
            ) for o in omic_zdims
        ])

    def get_criterions(self):
        fit_loss = loss.MOVBCExactLoss(self._kwargs["recon_ws"])
        return fit_loss

    def _batch_train(self, xs):
        # p(z_j|y), prepare
        iden = torch.eye(self._n_clusters,
                         device=self._device,
                         dtype=torch.float32)
        zy_mu, zy_logsigma2 = [], []
        for zmap in self.z_maps:
            zy_mu_j, zy_logsigma2_j = zmap(iden)   # ncluster x omic_dim
            zy_mu.append(zy_mu_j)
            zy_logsigma2.append(zy_logsigma2_j)

        # p(y|x_1,...,x_k), encode
        logits = self.encoder_y(xs)

        # p(z_ij|x_j,y=i), i=1,...,n_clusters, j=1,...,k
        z_mu, z_logsigma2, xs_rec = [], [], []
        for i in range(logits.size(1)):
            yi = torch.zeros_like(logits)
            yi[:, i] = 1.
            z_mu_i, z_logsigma2_i, xs_rec_i = [], [], []
            for x, enc, dec in zip(xs, self.encoders_z, self.decoders):
                x_y = torch.cat([x, yi], dim=1)
                z_mu_ij, z_logsigma2_ij = enc(x_y)
                z_ij = rsample_normal(z_mu_ij, z_logsigma2_ij)
                x_ij = dec(torch.cat([yi, z_ij], dim=1))
                z_mu_i.append(z_mu_ij)
                z_logsigma2_i.append(z_logsigma2_ij)
                xs_rec_i.append(x_ij)
            z_mu.append(z_mu_i)
            z_logsigma2.append(z_logsigma2_i)
            xs_rec.append(xs_rec_i)

        loss, losses = self._criterions(xs, xs_rec, logits,
                                        zy_mu, zy_logsigma2,
                                        z_mu, z_logsigma2,
                                        self._alpha, self._gamma)
        self._ops[0].zero_grad()
        loss.backward()
        self._ops[0].step()
        losses["total"] = loss
        self._accum.add(bs=self._bs, **losses)

    def transform(self, dat, bs, nw, device, disable=False):
        dataloader = dat.batch_iterator_loader(bs, nw, False, False, True)
        device = torch.device(device)
        self.encoder_y.to(device)
        self.encoder_y.eval()
        logits = []
        z_mu = [[] for _ in range(dat.n)]
        z_logsigma2 = [[] for _ in range(dat.n)]
        with torch.no_grad():
            for xs in tqdm(dataloader, "Embed: ", disable=disable):
                xs = [x.to(device) for x in xs]
                logit = self.encoder_y(xs)
                y = torch.eye(logit.size(1)).to(logit)[logit.argmax(dim=1)]
                for i, (x, enc) in enumerate(zip(xs, self.encoders_z)):
                    z_mu_j, z_logsigma2_j = enc(torch.cat([x, y], dim=1))
                    z_mu[i].append(z_mu_j)
                    z_logsigma2[i].append(z_logsigma2_j)
                logits.append(logit)

        logits = torch.cat(logits).detach().cpu().numpy()
        z_mu = [torch.cat(z_mu_j).detach().cpu().numpy() for z_mu_j in z_mu]
        z_lgs = [torch.cat(z_lgs_j).detach().cpu().numpy()
                 for z_lgs_j in z_logsigma2]

        return (logits, z_mu, z_lgs)

    def cluster(self, transformed_res):
        return np.argmax(transformed_res[0], axis=1)+1, transformed_res[0]


class MClusterVAEsGumbel(MClusterVAEsExact):

    def get_criterions(self):
        fit_loss = loss.MOVBCGumbelLoss(self._kwargs["recon_ws"])
        return fit_loss

    def _batch_train(self, xs):
        # p(y|x_1,...,x_k), encode
        logits = self.encoder_y(xs)

        y = rsample_categorical(logits, self._temp)

        # p(z_j|y), prepare
        zy_mu, zy_logsigma2 = [], []
        for zmap in self.z_maps:
            zy_mu_j, zy_logsigma2_j = zmap(y)   # ncluster x omic_dim
            zy_mu.append(zy_mu_j)
            zy_logsigma2.append(zy_logsigma2_j)

        # p(z_ij|x_j,y=i), i=1,...,n_clusters, j=1,...,k
        # z_mus, z_logsigma2s, zy_mus, zy_logsigma2s, xss = [], [], [], [], []
        z_mu, z_logsigma2, xs_rec = [], [], []
        for x, enc, dec in zip(xs, self.encoders_z, self.decoders):
            x_y = torch.cat([x, y], dim=1)
            z_mu_ij, z_logsigma2_ij = enc(x_y)
            z_ij = rsample_normal(z_mu_ij, z_logsigma2_ij)
            x_ij = dec(torch.cat([y, z_ij], dim=1))
            z_mu.append(z_mu_ij)
            z_logsigma2.append(z_logsigma2_ij)
            xs_rec.append(x_ij)

        loss, losses = self._criterions(xs, xs_rec, logits,
                                        zy_mu, zy_logsigma2,
                                        z_mu, z_logsigma2,
                                        self._alpha, self._gamma)
        self._ops[0].zero_grad()
        loss.backward()
        self._ops[0].step()
        losses["total"] = loss
        self._accum.add(bs=self._bs, **losses)

    def transform(self, dat, bs, nw, device, disable=False):
        dataloader = dat.batch_iterator_loader(bs, nw, False, False, True)
        device = torch.device(device)
        self.encoder_y.to(device)
        self.encoder_y.eval()
        logits = []
        z_mu = [[] for _ in range(dat.n)]
        z_logsigma2 = [[] for _ in range(dat.n)]
        with torch.no_grad():
            for xs in tqdm(dataloader, "Embed: ", disable=disable):
                xs = [x.to(device) for x in xs]
                logit = self.encoder_y(xs)
                y = torch.eye(logit.size(1)).to(logit)[logit.argmax(dim=1)]
                for i, (x, enc) in enumerate(zip(xs, self.encoders_z)):
                    z_mu_j, z_logsigma2_j = enc(torch.cat([x, y], dim=1))
                    z_mu[i].append(z_mu_j)
                    z_logsigma2[i].append(z_logsigma2_j)
                logits.append(logit)

        logits = torch.cat(logits).detach().cpu().numpy()
        z_mu = [torch.cat(z_mu_j).detach().cpu().numpy() for z_mu_j in z_mu]
        z_lgs = [torch.cat(z_lgs_j).detach().cpu().numpy()
                 for z_lgs_j in z_logsigma2]

        return (logits, z_mu, z_lgs)

    def cluster(self, transformed_res):
        return np.argmax(transformed_res[0], axis=1)+1, transformed_res[0]


class AE(VAEBase):

    def __init__(
        self, inpts, omic_oupts, n_clusters, con_dim, att_method="none",
        omic_hiddens=(), enc_hiddens=[200], dec_hiddens=[200]
    ):
        super().__init__()
        self._n_clusters = n_clusters
        self._inpts = inpts
        self._att_method = att_method

        self.encoder = block.Encoder(inpts, omic_oupts, [con_dim],
                                     omic_hiddens, enc_hiddens, att_method)
        self.decoder = block.Decoder(con_dim, inpts, dec_hiddens)

    def get_criterions(self):
        fit_loss = loss.MultiMSELoss(self._kwargs["recon_ws"], "sum")
        return fit_loss

    def _batch_train(self, xs):
        z = self.encoder(xs)
        xs_rec = self.decoder(z)

        loss, losses = self._criterions(xs, xs_rec)
        self._ops[0].zero_grad()
        loss.backward()
        self._ops[0].step()
        losses["total"] = loss
        self._accum.add(bs=self._bs, **losses)

    def transform(self, dat, bs, nw, device, disable=False):
        dataloader = dat.batch_iterator_loader(bs, nw, False, False, True)
        device = torch.device(device)
        self.encoder.to(device)
        self.encoder.eval()
        zs = []
        with torch.no_grad():
            for xs in tqdm(dataloader, "Embed: ", disable=disable):
                xs = [x.to(device) for x in xs]
                z = self.encoder(xs)
                zs.append(z)

        zs = torch.cat(zs).detach().cpu().numpy()

        return zs

    def cluster(self, transformed_res):
        gmm = GaussianMixture(n_components=self._n_clusters,
                              covariance_type="diag")
        cluster_res = gmm.fit_predict(transformed_res) + 1
        return cluster_res, None


class VAE(VAEBase):

    def __init__(
        self, inpts, omic_oupts, n_clusters, con_dim, att_method="none",
        omic_hiddens=(), enc_hiddens=[200], dec_hiddens=[200]
    ):
        super().__init__()
        self._n_clusters = n_clusters
        self._inpts = inpts
        self._att_method = att_method

        self.encoder = block.Encoder(inpts, omic_oupts, [con_dim, con_dim],
                                     omic_hiddens, enc_hiddens, att_method)
        self.decoder = block.Decoder(con_dim, inpts, dec_hiddens)

    def get_criterions(self):
        fit_loss = loss.VAELoss(self._kwargs["recon_ws"], "sum")
        return fit_loss

    def _batch_train(self, xs):
        z_mu, z_logsigma2 = self.encoder(xs)
        z = rsample_normal(z_mu, z_logsigma2)
        xs_rec = self.decoder(z)

        loss, losses = self._criterions(xs, xs_rec, z_mu, z_logsigma2)
        self._ops[0].zero_grad()
        loss.backward()
        self._ops[0].step()
        losses["total"] = loss
        self._accum.add(bs=self._bs, **losses)

    def transform(self, dat, bs, nw, device, disable=False):
        dataloader = dat.batch_iterator_loader(bs, nw, False, False, True)
        device = torch.device(device)
        self.encoder.to(device)
        self.encoder.eval()
        z_mus, z_logsigma2s = [], []
        with torch.no_grad():
            for xs in tqdm(dataloader, "Embed: ", disable=disable):
                xs = [x.to(device) for x in xs]
                z_mu, z_logsigma2 = self.encoder(xs)
                z_mus.append(z_mu)
                z_logsigma2s.append(z_logsigma2)

        z_mus = torch.cat(z_mus).detach().cpu().numpy()
        z_logsigma2s = torch.cat(z_logsigma2s).detach().cpu().numpy()

        return z_mus, z_logsigma2s

    def cluster(self, transformed_res):
        gmm = GaussianMixture(n_components=self._n_clusters,
                              covariance_type="diag")
        cluster_res = gmm.fit_predict(transformed_res[0]) + 1
        return cluster_res, None
