#!/usr/bin/env python3
# -*- encodinng: uft-8 -*-

import torch
import torch.nn as nn

'''MI Estimator network is seperated from the main network'''


class MIEstimator(nn.Module):
    def __init__(self, dim=128):
        super(MIEstimator, self).__init__()
        self.dim = dim
        self.mimin_glob = CLUBMean(self.dim * 2, self.dim)  # Can also use CLUBEstimator, but CLUBMean is more stable
        self.mimin = CLUBMean(self.dim, self.dim)
        self.mimin_rna = CLUBMean(self.dim, self.dim)
        self.mimin_atac = CLUBMean(self.dim, self.dim)

    def forward(self, rna_private, atac_private, global_embed, rna_share, atac_share):
        mimin = self.mimin(rna_private, atac_private)
        mimin += self.mimin_glob(torch.cat((rna_private, atac_private), dim=1), global_embed).mean()
        mimin += self.mimin_rna(rna_private, rna_share)
        mimin += self.mimin_atac(atac_private, atac_share)
        return mimin

    def learning_loss(self, rna_private, atac_private, global_embed, rna_share, atac_share):
        mimin_loss = self.mimin.learning_loss(rna_private, atac_private)
        mimin_loss += self.mimin_glob.learning_loss(torch.cat((rna_private, atac_private), dim=1), global_embed).mean()
        mimin_loss += self.mimin_rna.learning_loss(rna_private, rna_share)
        mimin_loss += self.mimin_atac.learning_loss(atac_private, atac_share)
        return mimin_loss


'''From the original code (https://github.com/Linear95/CLUB), CLUBEstimator is used to estimate the mutual information between histology and pathways features'''


class CLUBEstimator(nn.Module):  # CLUB: Mutual Information Contrastive Learning Upper Bound
    '''
        This class provides the CLUB estimation to I(X,Y)
        Method:
            forward() :      provides the estimation with input samples
            loglikeli() :   provides the log-likelihood of the approximation q(Y|X) with input samples
        Arguments:
            x_dim, y_dim :         the dimensions of samples from X, Y respectively
            hidden_size :          the dimension of the hidden layer of the approximation network q(Y|X)
            x_samples, y_samples : samples from X and Y, having shape [sample_size, x_dim/y_dim]
    '''

    def __init__(self, x_dim, y_dim, hidden_size=512):
        super(CLUBEstimator, self).__init__()
        # p_mu outputs mean of q(Y|X)
        # print("create CLUB with dim {}, {}, hiddensize {}".format(x_dim, y_dim, hidden_size))
        self.p_mu = nn.Sequential(nn.Linear(x_dim, hidden_size),
                                  nn.ReLU(),
                                  nn.Linear(hidden_size, y_dim))
        # p_logvar outputs log of variance of q(Y|X)
        self.p_logvar = nn.Sequential(nn.Linear(x_dim, hidden_size),
                                      nn.ReLU(),
                                      nn.Linear(hidden_size, y_dim),
                                      nn.Tanh())

    def get_mu_logvar(self, x_samples):
        mu = self.p_mu(x_samples)
        logvar = self.p_logvar(x_samples)
        return mu, logvar  #q(y|x)

    def forward(self, x_samples, y_samples):
        mu, logvar = self.get_mu_logvar(x_samples)

        # log of conditional probability of positive sample pairs
        positive = - (mu - y_samples) ** 2 / 2. / logvar.exp()

        prediction_1 = mu.unsqueeze(1)  # shape [nsample,1,dim]
        y_samples_1 = y_samples.unsqueeze(0)  # shape [1,nsample,dim]

        # log of conditional probability of negative sample pairs
        negative = - ((y_samples_1 - prediction_1) ** 2).mean(dim=1) / 2. / logvar.exp()

        return (positive.sum(dim=-1) - negative.sum(dim=-1)).mean()

    def loglikeli(self, x_samples, y_samples):  # unnormalized loglikelihood
        mu, logvar = self.get_mu_logvar(x_samples)
        return (-(mu - y_samples) ** 2 / logvar.exp() - logvar).sum(dim=1).mean(dim=0)

    def learning_loss(self, x_samples, y_samples):
        return - self.loglikeli(x_samples, y_samples)


class CLUBMean(nn.Module):  # Set variance of q(y|x) to 1, logvar = 0. Update 11/26/2022
    def __init__(self, x_dim, y_dim, hidden_size=512):
        # p_mu outputs mean of q(Y|X)
        # print("create CLUB with dim {}, {}, hiddensize {}".format(x_dim, y_dim, hidden_size))

        super(CLUBMean, self).__init__()

        if hidden_size is None:
            self.p_mu = nn.Linear(x_dim, y_dim)
        else:
            self.p_mu = nn.Sequential(nn.Linear(x_dim, int(hidden_size)),
                                      nn.ReLU(),
                                      nn.Linear(int(hidden_size), y_dim))

    def get_mu_logvar(self, x_samples):
        # variance is set to 1, which means logvar=0
        mu = self.p_mu(x_samples)
        return mu, 0

    def forward(self, x_samples, y_samples):

        mu, logvar = self.get_mu_logvar(x_samples)

        # log of conditional probability of positive sample pairs
        positive = - (mu - y_samples) ** 2 / 2.

        prediction_1 = mu.unsqueeze(1)  # shape [nsample,1,dim]
        y_samples_1 = y_samples.unsqueeze(0)  # shape [1,nsample,dim]

        # log of conditional probability of negative sample pairs
        negative = - ((y_samples_1 - prediction_1) ** 2).mean(dim=1) / 2.

        return (positive.sum(dim=-1) - negative.sum(dim=-1)).mean()

    def loglikeli(self, x_samples, y_samples):  # unnormalized loglikelihood
        mu, logvar = self.get_mu_logvar(x_samples)
        return (-(mu - y_samples) ** 2).sum(dim=1).mean(dim=0)

    def learning_loss(self, x_samples, y_samples):
        return - self.loglikeli(x_samples, y_samples)
