import numpy as np
from tqdm import trange, tqdm
import torch
import math
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F

def matrix_log_density_gaussian(x, mu, logvar):
	batch_size, dim = x.shape
	x = x.view(batch_size, 1, dim)
	mu = mu.view(1, batch_size, dim)
	logvar = logvar.view(1, batch_size, dim)
	return log_density_gaussian(x, mu, logvar)
	
def log_density_gaussian(x, mu, logvar):
	normalization = - 0.5 * (math.log(2 * math.pi) + logvar)
	inv_var = torch.exp(-logvar)
	log_density = normalization - 0.5 * ((x - mu)**2 * inv_var)
	return log_density

def log_importance_weight_matrix(batch_size, dataset_size):
	N = dataset_size
	M = batch_size - 1
	strat_weight = (N - M) / (N * M)
	W = torch.Tensor(batch_size, batch_size).fill_(1 / M)
	W.view(-1)[::M + 1] = 1 / N
	W.view(-1)[1::M + 1] = strat_weight
	W[M - 1, 0] = strat_weight
	return W.log()

rand_tensor = torch.empty((3,2))
mu_tensor = torch.empty((3,2))
logvar_tensor = torch.empty((3,2))

rand_tensor[0,0]=5
rand_tensor[0,1]=3
rand_tensor[1,0]=7
rand_tensor[1,1]=9
rand_tensor[2,0]=8
rand_tensor[2,1]=6

mu_tensor[0,0]=0.6
mu_tensor[0,1]=0.2
mu_tensor[1,0]=0.3
mu_tensor[1,1]=0.8
mu_tensor[2,0]=0.7
mu_tensor[2,1]=0.1

logvar_tensor[0,0]=0.4
logvar_tensor[0,1]=0.8
logvar_tensor[1,0]=0.4
logvar_tensor[1,1]=0.3
logvar_tensor[2,0]=0.1
logvar_tensor[2,1]=0.4

zeros = torch.zeros_like(rand_tensor)
log_pz = log_density_gaussian(rand_tensor,zeros, zeros).sum(1)
mu = Variable(torch.Tensor([0]))
logsigma = Variable(torch.Tensor([math.log(1)]))
size = rand_tensor.size()
mu = mu.expand(size)
logsigma = logsigma.expand(size)
mu = mu.type_as(rand_tensor)
logsigma = logsigma.type_as(rand_tensor)
c = Variable(torch.Tensor([np.log(2 * np.pi)])).type_as(rand_tensor.data)
inv_sigma = torch.exp(-logsigma)
tmp = (rand_tensor - mu) * inv_sigma
logpz_ricky = (-0.5 * (tmp * tmp + 2 * logsigma + c)).view(3, -1).sum(1)

log_q_zCx = log_density_gaussian(rand_tensor,mu_tensor,logvar_tensor).sum(dim=1)

mu = mu_tensor
logsigma = 0.5*logvar_tensor
c = Variable(torch.Tensor([np.log(2 * np.pi)])).type_as(rand_tensor.data)
inv_sigma = torch.exp(-logsigma)
tmp = (rand_tensor - mu) * inv_sigma
logqz_condx_ricky = (-0.5 * (tmp * tmp + 2 * logsigma + c)).view(3, -1).sum(1)

mat_log_qz = matrix_log_density_gaussian(rand_tensor,mu_tensor,logvar_tensor)

rand_tensor = rand_tensor.view(3, 1, 2)

mu_logsigma_tensor = torch.empty((1,3,2,2))
mu_logsigma_tensor[0,0,0,0]=mu_tensor[0,0]
mu_logsigma_tensor[0,0,0,1]=0.5*logvar_tensor[0,0]
mu_logsigma_tensor[0,1,0,0]=mu_tensor[1,0]
mu_logsigma_tensor[0,1,0,1]=0.5*logvar_tensor[1,0]
mu_logsigma_tensor[0,2,0,0]=mu_tensor[2,0]
mu_logsigma_tensor[0,2,0,1]=0.5*logvar_tensor[2,0]

mu_logsigma_tensor[0,0,1,0]=mu_tensor[0,1]
mu_logsigma_tensor[0,0,1,1]=0.5*logvar_tensor[0,1]
mu_logsigma_tensor[0,1,1,0]=mu_tensor[1,1]
mu_logsigma_tensor[0,1,1,1]=0.5*logvar_tensor[1,1]
mu_logsigma_tensor[0,2,1,0]=mu_tensor[2,1]
mu_logsigma_tensor[0,2,1,1]=0.5*logvar_tensor[2,1]

mu = mu_logsigma_tensor .select(-1, 0)
logsigma = mu_logsigma_tensor .select(-1, 1)
c = Variable(torch.Tensor([np.log(2 * np.pi)])).type_as(rand_tensor.data)
inv_sigma = torch.exp(-logsigma)
tmp = (rand_tensor - mu) * inv_sigma
_logqz_ricky = (-0.5 * (tmp * tmp + 2 * logsigma + c))

log_qz = torch.logsumexp(mat_log_qz.sum(2), dim=1, keepdim=False)
log_prod_qzi = torch.logsumexp(mat_log_qz, dim=1, keepdim=False).sum(1)

logqz_ricky = (torch.logsumexp(_logqz_ricky.sum(2), dim=1, keepdim=False) - math.log(3 * 12))
logqz_prodmarginals_ricky = (torch.logsumexp(_logqz_ricky, dim=1, keepdim=False) - math.log(3 * 12)).sum(1)

log_pz
logpz_ricky

log_q_zCx
logqz_condx_ricky

mat_log_qz
_logqz_ricky

log_qz
logqz_ricky

log_prod_qzi
logqz_prodmarginals_ricky

log_iw_mat = log_importance_weight_matrix(3, 12).to(rand_tensor.device)
logiw_matrix = Variable(log_importance_weight_matrix(3, 12).type_as(_logqz_ricky.data))

logqz_ricky_mss = torch.logsumexp(logiw_matrix + _logqz_ricky.sum(2), dim=1, keepdim=False)
logqz_prodmarginals_ricky_mss = torch.logsumexp(logiw_matrix.view(3, 3, 1) + _logqz_ricky, dim=1, keepdim=False).sum(1)

mat_log_qz_mss = mat_log_qz + log_iw_mat.view(3, 3, 1)

log_qz_mss = torch.logsumexp(mat_log_qz_mss.sum(2), dim=1, keepdim=False)
log_prod_qzi_mss = torch.logsumexp(mat_log_qz_mss, dim=1, keepdim=False).sum(1)



Add AE
Map to PCA
Nrmazlizarion - Scaling for10 dims - same unit
running example	
Losses Graphs