import pickle
import numpy as np
import torch
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
from quaternion import convert_dictionary
from helperFunctions import eps


"""
Loss function
"""


class SimpleLoss(nn.Module):
	def __init__(self, alpha):
		super().__init__()
		self.alpha = alpha
		self.mse = nn.MSELoss().cuda()
		self.ce = nn.CrossEntropyLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = [ydata_bin, ydata_res]
		# ypred = [score, residual]
		l1 = self.ce(ypred[0], ytrue[0])
		l2 = self.mse(ypred[1], ytrue[1])
		return torch.add(l1, self.alpha, l2)


class GeodesicLoss(nn.Module):
	def __init__(self, alpha, kmeans_file, my_loss=None):
		super().__init__()
		self.alpha = alpha
		kmeans = pickle.load(open(kmeans_file, 'rb'))
		self.cluster_centers_ = Variable(torch.from_numpy(kmeans.cluster_centers_).float()).cuda()
		if my_loss is None:
			self.mse = nn.MSELoss().cuda()
		else:
			self.mse = my_loss
		self.ce = nn.CrossEntropyLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_label, ydata)
		# ypred = (score, residual)
		l1 = self.ce(ypred[0], ytrue[0])
		_, ind = torch.max(ypred[0], dim=1)
		y = torch.index_select(self.cluster_centers_, 0, ind)
		l2 = self.mse(y+ypred[1], ytrue[1])
		return torch.add(l1, self.alpha, l2)


class GeodesicLossQ(nn.Module):
	def __init__(self, alpha, kmeans_file, my_loss=None):
		super().__init__()
		self.alpha = alpha
		kmeans = pickle.load(open(kmeans_file, 'rb'))
		self.cluster_centers_ = Variable(torch.from_numpy(convert_dictionary(kmeans.cluster_centers_)).float()).cuda()
		if my_loss is None:
			self.mse = nn.MSELoss().cuda()
		else:
			self.mse = my_loss
		self.ce = nn.CrossEntropyLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_label, ydata)
		# ypred = (score, residual)
		l1 = self.ce(ypred[0], ytrue[0])
		_, ind = torch.max(ypred[0], dim=1)
		y = torch.index_select(self.cluster_centers_, 0, ind)
		l2 = self.mse(y+ypred[1], ytrue[1])
		return torch.add(l1, self.alpha, l2)


class SimpleRelaXedLoss(nn.Module):
	def __init__(self, alpha):
		super().__init__()
		self.alpha = alpha
		self.mse = nn.MSELoss().cuda()
		self.kl = nn.KLDivLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = [ydata_bin, ydata_res]
		# ypred = [score, residual]
		l1 = self.kl(F.log_softmax(ypred[0], dim=1), ytrue[0])
		l2 = self.mse(ypred[1], ytrue[1])
		return torch.add(l1, self.alpha, l2)


class RelaXedLoss(nn.Module):
	def __init__(self, alpha, kmeans_file, my_loss):
		super().__init__()
		self.alpha = alpha
		kmeans = pickle.load(open(kmeans_file, 'rb'))
		self.cluster_centers_ = Variable(torch.from_numpy(kmeans.cluster_centers_).float()).cuda()
		self.my_loss = my_loss
		self.kl = nn.KLDivLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_label, ydata)
		# ypred = (score, residual)
		l1 = self.kl(F.log_softmax(ypred[0], dim=1), ytrue[0])
		_, ind = torch.max(ypred[0], dim=1)
		y = torch.index_select(self.cluster_centers_, 0, ind)
		l2 = self.my_loss(y+ypred[1], ytrue[1])
		return torch.add(l1, self.alpha, l2)


class RelaXedProbabilisticLoss(nn.Module):
	def __init__(self, alpha, gmm_file, my_loss):
		super().__init__()
		self.alpha = alpha
		gmm = pickle.load(open(gmm_file, 'rb'))
		self.cluster_centers = Variable(torch.from_numpy(gmm.means_).float()).cuda()
		self.n_clusters = gmm.n_components
		self.my_loss = my_loss
		self.kl = nn.KLDivLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_prob, ydata)
		# ypred = (score, residual)
		l1 = self.kl(F.log_softmax(ypred[0], dim=1), ytrue[0])
		l2 = torch.stack([self.my_loss(ytrue[1], torch.add(ypred[1], 1.0, self.cluster_centers.index_select(0, Variable(i*torch.ones(1).long().cuda()))))
		                  for i in range(self.n_clusters)])
		l2 = torch.mean(torch.sum(torch.mul(F.softmax(ypred[0], dim=1), torch.t(l2)), dim=1))
		return torch.add(l1, self.alpha, l2)


class ProbabilisticLoss(nn.Module):
	def __init__(self, alpha, kmeans_file, my_loss):
		super().__init__()
		self.alpha = alpha
		kmeans = pickle.load(open(kmeans_file, 'rb'))
		self.cluster_centers = Variable(torch.from_numpy(kmeans.cluster_centers_).float()).cuda()
		self.n_clusters = kmeans.n_clusters
		self.my_loss = my_loss
		self.ce = nn.CrossEntropyLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_prob, ydata)
		# ypred = (score, residual)
		l1 = self.ce(ypred[0], ytrue[0])
		l2 = torch.stack([self.my_loss(ytrue[1], torch.add(ypred[1], 1.0, self.cluster_centers.index_select(0, Variable(i*torch.ones(1).long().cuda()))))
		                  for i in range(self.n_clusters)])
		l2 = torch.mean(torch.sum(torch.mul(F.softmax(ypred[0], dim=1), torch.t(l2)), dim=1))
		return torch.add(l1, self.alpha, l2)


class RelaXedProbabilisticLossQ(nn.Module):
	def __init__(self, alpha, kmeans_file, my_loss):
		super().__init__()
		self.alpha = alpha
		kmeans = pickle.load(open(kmeans_file, 'rb'))
		self.cluster_centers = Variable(torch.from_numpy(convert_dictionary(kmeans.cluster_centers_)).float()).cuda()
		self.n_clusters = kmeans.n_clusters
		self.my_loss = my_loss
		self.kl = nn.KLDivLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_prob, ydata)
		# ypred = (score, residual)
		l1 = self.kl(F.log_softmax(ypred[0], dim=1), ytrue[0])
		l2 = torch.stack([self.my_loss(ytrue[1], torch.add(ypred[1], 1.0, self.cluster_centers.index_select(0, Variable(i*torch.ones(1).long().cuda()))))
		                  for i in range(self.n_clusters)])
		l2 = torch.mean(torch.sum(torch.mul(F.softmax(ypred[0], dim=1), torch.t(l2)), dim=1))
		return torch.add(l1, self.alpha, l2)


class RelaXedProbabilisticMultiresLoss(RelaXedProbabilisticLoss):
	def __init__(self, alpha, kmeans_file, my_loss):
		super().__init__(alpha, kmeans_file, my_loss)

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_prob, ydata)
		# ypred = (score, residual)
		l1 = self.kl(F.log_softmax(ypred[0], dim=1), ytrue[0])
		y = self.cluster_centers + ypred[1]
		l2 = torch.stack([self.my_loss(ytrue[1], torch.squeeze(y.index_select(1, Variable(i*torch.ones(1).long().cuda())))) for i in range(self.n_clusters)])
		l2 = torch.mean(torch.sum(torch.mul(F.softmax(ypred[0], dim=1), torch.t(l2)), dim=1))
		return torch.add(l1, self.alpha, l2)


class ProbabilisticMultiresLoss(ProbabilisticLoss):
	def __init__(self, alpha, kmeans_file, my_loss):
		super().__init__(alpha, kmeans_file, my_loss)

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_prob, ydata)
		# ypred = (score, residual)
		l1 = self.ce(ypred[0], ytrue[0])
		y = self.cluster_centers + ypred[1]
		l2 = torch.stack([self.my_loss(ytrue[1], torch.squeeze(y.index_select(1, Variable(i*torch.ones(1).long().cuda())))) for i in range(self.n_clusters)])
		l2 = torch.mean(torch.sum(torch.mul(F.softmax(ypred[0], dim=1), torch.t(l2)), dim=1))
		return torch.add(l1, self.alpha, l2)


class RelaXedProbabilisticMultiresLossQ(RelaXedProbabilisticLossQ):
	def __init__(self, alpha, kmeans_file, my_loss):
		super().__init__(alpha, kmeans_file, my_loss)

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_prob, ydata)
		# ypred = (score, residual)
		l1 = self.kl(F.log_softmax(ypred[0], dim=1), ytrue[0])
		y = self.cluster_centers + ypred[1]
		l2 = torch.stack([self.my_loss(ytrue[1], torch.squeeze(y.index_select(1, Variable(i*torch.ones(1).long().cuda())))) for i in range(self.n_clusters)])
		l2 = torch.mean(torch.sum(torch.mul(F.softmax(ypred[0], dim=1), torch.t(l2)), dim=1))
		return torch.add(l1, self.alpha, l2)


class RiemannianLoss(nn.Module):
	def __init__(self, alpha, pose_dict):
		super().__init__()
		self.alpha = alpha
		self.key_poses = torch.from_numpy(pose_dict).float().cuda()
		self.ce = nn.CrossEntropyLoss().cuda()
		proj = np.array([[0,0,0,0,0,-1,0,1,0], [0,0,1,0,0,0,-1,0,0], [0,-1,0,1,0,0,0,0,0]])
		self.proj = torch.from_numpy(proj).float().cuda()
		self.Id = torch.eye(3).float().cuda()

	def my_loss(self, ypred, ytrue):
		# geodesic loss between predicted and gt rotations
		tmp = torch.stack([torch.trace(torch.mm(ypred[i].t(), ytrue[i])) for i in range(ytrue.size(0))])
		angle = torch.acos(torch.clamp((tmp - 1.0) / 2, -1 + eps, 1 - eps))
		return torch.mean(angle)

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_bin, ydata_rot)
		# ypred = (score, residual)
		l1 = self.ce(ypred[0], ytrue[0])
		_, ind = torch.max(ypred[0], dim=1)
		yres = ypred[1]
		angle = torch.norm(yres, 2, 1)
		axis = F.normalize(yres)
		axis = torch.mm(axis, self.proj).view(-1, 3, 3)
		y = torch.stack([self.Id + torch.sin(angle[i])*axis[i] + (1.0 - torch.cos(angle[i]))*torch.mm(axis[i], axis[i]) for i in range(angle.size(0))])
		y = torch.bmm(torch.index_select(self.key_poses, 0, ind), y)
		l2 = self.my_loss(y, ytrue[1])
		return torch.add(l1, self.alpha, l2)


# loss
class loss_m0(nn.Module):
	def __init__(self, alpha):
		super().__init__()
		self.alpha = alpha
		self.mse = nn.MSELoss().cuda()
		self.ce = nn.CrossEntropyLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = [ydata_bin, ydata_res]
		# ypred = [score, residual]
		l1 = self.ce(ypred[0], ytrue[0])
		l2 = self.mse(ypred[1], ytrue[1])
		return torch.add(l1, self.alpha, l2)


class loss_m1(nn.Module):
	def __init__(self, alpha, kmeans_file, my_loss=None):
		super().__init__()
		self.alpha = alpha
		kmeans = pickle.load(open(kmeans_file, 'rb'))
		self.cluster_centers_ = Variable(torch.from_numpy(kmeans.cluster_centers_).float()).cuda()
		if my_loss is None:
			self.mse = nn.MSELoss().cuda()
		else:
			self.mse = my_loss
		self.ce = nn.CrossEntropyLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_label, ydata)
		# ypred = (score, residual)
		l1 = self.ce(ypred[0], ytrue[0])
		_, ind = torch.max(ypred[0], dim=1)
		y = torch.index_select(self.cluster_centers_, 0, ind)
		l2 = self.mse(y+ypred[1], ytrue[1])
		return torch.add(l1, self.alpha, l2)


class loss_m2(nn.Module):
	def __init__(self, alpha, num_clusters):
		super().__init__()
		self.alpha = alpha
		self.num_clusters = num_clusters
		self.mse = nn.MSELoss().cuda()
		self.ce = nn.CrossEntropyLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_label, ydata_res)
		# ypred = (score, residual)
		l1 = self.ce(ypred[0], ytrue[0])
		_, label = torch.max(ypred[0], dim=1)
		label = torch.zeros(label.size(0), self.num_clusters).scatter_(1, label.unsqueeze(1).data.cpu(), 1.0)
		label = Variable(label.unsqueeze(2).cuda())
		yres = torch.squeeze(torch.bmm(ytrue[1].permute(0, 2, 1), label), 2)
		l2 = self.mse(ypred[1], yres)
		return torch.add(l1, self.alpha, l2)


class loss_m3(nn.Module):
	def __init__(self, alpha, kmeans_file, my_loss=None):
		super().__init__()
		self.alpha = alpha
		kmeans = pickle.load(open(kmeans_file, 'rb'))
		self.cluster_centers = Variable(torch.from_numpy(kmeans.cluster_centers_).float()).cuda()
		self.n_clusters = kmeans.n_clusters
		if my_loss is None:
			self.mse = nn.MSELoss(reduce=False).cuda()
		else:
			self.mse = my_loss
		self.kl = nn.KLDivLoss().cuda()

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_prob, ydata)
		# ypred = (score, residual)
		l1 = self.kl(F.log_softmax(ypred[0], dim=1), ytrue[0])
		l2 = torch.stack([self.mse(ytrue[1], torch.add(ypred[1], 1.0, self.cluster_centers.index_select(0, Variable(i*torch.ones(1).long().cuda()))))
		                  for i in range(self.n_clusters)])
		l2 = torch.mean(torch.sum(torch.mul(F.softmax(ypred[0], dim=1), torch.t(l2)), dim=1))
		return torch.add(l1, self.alpha, l2)


class loss_m4(loss_m3):
	def __init__(self, alpha, kmeans_file, my_loss=None):
		super().__init__(alpha, kmeans_file, my_loss)

	def forward(self, ypred, ytrue):
		# ytrue = (ydata_prob, ydata)
		# ypred = (score, residual)
		l1 = self.kl(F.log_softmax(ypred[0], dim=1), ytrue[0])
		y = self.cluster_centers + ypred[1]
		l2 = torch.stack([self.mse(ytrue[1], torch.squeeze(y.index_select(1, Variable(i*torch.ones(1).long().cuda())))) for i in range(self.n_clusters)])
		l2 = torch.mean(torch.sum(torch.mul(F.softmax(ypred[0], dim=1), torch.t(l2)), dim=1))
		return torch.add(l1, self.alpha, l2)
