import numpy as np
from pybullet_utils import bullet_client
import math
from .math_utils import *

class HumanoidPoseInterpolator(object):

  # REPRESENTATION_MODE_CHECKPOINT
  # def __init__(self):
  def __init__(self, arg_parser = None):

    # REPRESENTATION_MODE_CHECKPOINT
    # self.state_representation_mode = "Quaternion"
    # self.action_representation_mode = "AxisAngle"
    # self.state_representation_mode = "6D"
    # self.action_representation_mode = "6D"

    # REPRESENTATION_MODE_CHECKPOINT
    self._arg_parser = arg_parser
    self.state_representation_mode = self._arg_parser.parse_string('state_repr', default="Quaternion")
    self.action_representation_mode = self._arg_parser.parse_string('action_repr', default="AxisAngle")

    if self.action_representation_mode == "Quaternion":
      self.action_dim = 4
    elif self.action_representation_mode == "Euler":
      self.action_dim = 3
    elif self.action_representation_mode == "AxisAngle":
      self.action_dim = 4
    elif self.action_representation_mode == "RotVec":
      self.action_dim = 3
    elif self.action_representation_mode == "RotMat":
      self.action_dim = 9
    elif self.action_representation_mode == "6D":
      self.action_dim = 6

    pass

  def Reset(self,

            basePos=[0, 0, 0],
            baseOrn=[0, 0, 0, 1],

            chestRot=[0, 0, 0, 1],
            neckRot=[0, 0, 0, 1],
            rightHipRot=[0, 0, 0, 1],
            rightKneeRot=[0],
            rightAnkleRot=[0, 0, 0, 1],
            rightShoulderRot=[0, 0, 0, 1],
            rightElbowRot=[0],
            leftHipRot=[0, 0, 0, 1],
            leftKneeRot=[0],
            leftAnkleRot=[0, 0, 0, 1],
            leftShoulderRot=[0, 0, 0, 1],
            leftElbowRot=[0],

            baseLinVel=[0, 0, 0],
            baseAngVel=[0, 0, 0],

            chestVel=[0, 0, 0],
            neckVel=[0, 0, 0],
            rightHipVel=[0, 0, 0],
            rightKneeVel=[0],
            rightAnkleVel=[0, 0, 0],
            rightShoulderVel=[0, 0, 0],
            rightElbowVel=[0],
            leftHipVel=[0, 0, 0],
            leftKneeVel=[0],
            leftAnkleVel=[0, 0, 0],
            leftShoulderVel=[0, 0, 0],
            leftElbowVel=[0]):

    self._basePos = basePos
    self._baseLinVel = baseLinVel
    #print("HumanoidPoseInterpolator.Reset: baseLinVel = ", baseLinVel)
    self._baseOrn = baseOrn
    self._baseAngVel = baseAngVel

    self._chestRot = chestRot
    self._chestVel = chestVel
    self._neckRot = neckRot
    self._neckVel = neckVel

    self._rightHipRot = rightHipRot
    self._rightHipVel = rightHipVel
    self._rightKneeRot = rightKneeRot
    self._rightKneeVel = rightKneeVel
    self._rightAnkleRot = rightAnkleRot
    self._rightAnkleVel = rightAnkleVel

    self._rightShoulderRot = rightShoulderRot
    self._rightShoulderVel = rightShoulderVel
    self._rightElbowRot = rightElbowRot
    self._rightElbowVel = rightElbowVel

    self._leftHipRot = leftHipRot
    self._leftHipVel = leftHipVel
    self._leftKneeRot = leftKneeRot
    self._leftKneeVel = leftKneeVel
    self._leftAnkleRot = leftAnkleRot
    self._leftAnkleVel = leftAnkleVel

    self._leftShoulderRot = leftShoulderRot
    self._leftShoulderVel = leftShoulderVel
    self._leftElbowRot = leftElbowRot
    self._leftElbowVel = leftElbowVel

  def ComputeLinVel(self, posStart, posEnd, deltaTime):
    vel = [(posEnd[0] - posStart[0]) / deltaTime, (posEnd[1] - posStart[1]) / deltaTime,
           (posEnd[2] - posStart[2]) / deltaTime]
    return vel

  def ComputeAngVel(self, ornStart, ornEnd, deltaTime, bullet_client):
    dorn = bullet_client.getDifferenceQuaternion(ornStart, ornEnd)
    axis, angle = bullet_client.getAxisAngleFromQuaternion(dorn)
    angVel = [(axis[0] * angle) / deltaTime, (axis[1] * angle) / deltaTime,
              (axis[2] * angle) / deltaTime]
    return angVel

  def ComputeAngVelRel(self, ornStart, ornEnd, deltaTime, bullet_client):
    ornStartConjugate = [-ornStart[0], -ornStart[1], -ornStart[2], ornStart[3]]
    pos_diff, q_diff = bullet_client.multiplyTransforms([0, 0, 0], ornStartConjugate, [0, 0, 0],
                                                        ornEnd)
    axis, angle = bullet_client.getAxisAngleFromQuaternion(q_diff)
    angVel = [(axis[0] * angle) / deltaTime, (axis[1] * angle) / deltaTime,
              (axis[2] * angle) / deltaTime]
    return angVel

  def NormalizeVector(self, vec):
    # length2 = vec[0] * vec[0] + vec[1] * vec[1] + vec[2] * vec[2]
    length2 = self.DotProduct(vec, vec)
    if (length2 > 0):
      length = math.sqrt(length2)
      vec[0] /= length
      vec[1] /= length
      vec[2] /= length
      return vec

  def NormalizeQuaternion(self, orn):
    # length2 = orn[0] * orn[0] + orn[1] * orn[1] + orn[2] * orn[2] + orn[3] * orn[3]
    length2 = self.DotProduct(orn, orn)
    if (length2 > 0):
      length = math.sqrt(length2)
      orn[0] /= length
      orn[1] /= length
      orn[2] /= length
      orn[3] /= length
      return orn

    #print("Normalize? length=",length)

  def PostProcessMotionData(self, frameData):
    baseOrn1Start = [frameData[5], frameData[6], frameData[7], frameData[4]]

    chestRotStart = [frameData[9], frameData[10], frameData[11], frameData[8]]

    neckRotStart = [frameData[13], frameData[14], frameData[15], frameData[12]]
    rightHipRotStart = [frameData[17], frameData[18], frameData[19], frameData[16]]
    rightAnkleRotStart = [frameData[22], frameData[23], frameData[24], frameData[21]]
    rightShoulderRotStart = [frameData[26], frameData[27], frameData[28], frameData[25]]
    leftHipRotStart = [frameData[31], frameData[32], frameData[33], frameData[30]]
    leftAnkleRotStart = [frameData[36], frameData[37], frameData[38], frameData[35]]
    leftShoulderRotStart = [frameData[40], frameData[41], frameData[42], frameData[39]]

  def GetPose(self):
    pose = [
        # these 7 elements will be zero-ed out in pybullet_deep_mimic_env.set_action()
        self._basePos[0], self._basePos[1], self._basePos[2],
        self._baseOrn[0], self._baseOrn[1], self._baseOrn[2], self._baseOrn[3],
        # these values will be given in ConvertFromAction()
        self._chestRot[0], self._chestRot[1], self._chestRot[2], self._chestRot[3],
        self._neckRot[0], self._neckRot[1], self._neckRot[2], self._neckRot[3],
        self._rightHipRot[0], self._rightHipRot[1], self._rightHipRot[2], self._rightHipRot[3],
        self._rightKneeRot[0],
        self._rightAnkleRot[0], self._rightAnkleRot[1], self._rightAnkleRot[2], self._rightAnkleRot[3],
        self._rightShoulderRot[0], self._rightShoulderRot[1], self._rightShoulderRot[2], self._rightShoulderRot[3],
        self._rightElbowRot[0],
        self._leftHipRot[0], self._leftHipRot[1], self._leftHipRot[2], self._leftHipRot[3],
        self._leftKneeRot[0],
        self._leftAnkleRot[0], self._leftAnkleRot[1], self._leftAnkleRot[2], self._leftAnkleRot[3],
        self._leftShoulderRot[0], self._leftShoulderRot[1], self._leftShoulderRot[2], self._leftShoulderRot[3],
        self._leftElbowRot[0]
    ]
    return pose

  def Slerp(self, frameFraction, frameData, frameDataNext, bullet_client):
    keyFrameDuration = frameData[0]
    basePos1Start = [frameData[1], frameData[2], frameData[3]]
    basePos1End = [frameDataNext[1], frameDataNext[2], frameDataNext[3]]
    self._basePos = [
        basePos1Start[0] + frameFraction * (basePos1End[0] - basePos1Start[0]),
        basePos1Start[1] + frameFraction * (basePos1End[1] - basePos1Start[1]),
        basePos1Start[2] + frameFraction * (basePos1End[2] - basePos1Start[2])
    ]
    self._baseLinVel = self.ComputeLinVel(basePos1Start, basePos1End, keyFrameDuration)
    baseOrn1Start = [frameData[5], frameData[6], frameData[7], frameData[4]]
    baseOrn1Next = [frameDataNext[5], frameDataNext[6], frameDataNext[7], frameDataNext[4]]
    self._baseOrn = bullet_client.getQuaternionSlerp(baseOrn1Start, baseOrn1Next, frameFraction)
    self._baseAngVel = self.ComputeAngVel(baseOrn1Start, baseOrn1Next, keyFrameDuration,
                                          bullet_client)

    ##pre-rotate to make z-up
    #y2zPos=[0,0,0.0]
    #y2zOrn = p.getQuaternionFromEuler([1.57,0,0])
    #basePos,baseOrn = p.multiplyTransforms(y2zPos, y2zOrn,basePos1,baseOrn1)

    chestRotStart = [frameData[9], frameData[10], frameData[11], frameData[8]]
    chestRotEnd = [frameDataNext[9], frameDataNext[10], frameDataNext[11], frameDataNext[8]]
    self._chestRot = bullet_client.getQuaternionSlerp(chestRotStart, chestRotEnd, frameFraction)
    self._chestVel = self.ComputeAngVelRel(chestRotStart, chestRotEnd, keyFrameDuration,
                                           bullet_client)

    neckRotStart = [frameData[13], frameData[14], frameData[15], frameData[12]]
    neckRotEnd = [frameDataNext[13], frameDataNext[14], frameDataNext[15], frameDataNext[12]]
    self._neckRot = bullet_client.getQuaternionSlerp(neckRotStart, neckRotEnd, frameFraction)
    self._neckVel = self.ComputeAngVelRel(neckRotStart, neckRotEnd, keyFrameDuration,
                                          bullet_client)

    rightHipRotStart = [frameData[17], frameData[18], frameData[19], frameData[16]]
    rightHipRotEnd = [frameDataNext[17], frameDataNext[18], frameDataNext[19], frameDataNext[16]]
    self._rightHipRot = bullet_client.getQuaternionSlerp(rightHipRotStart, rightHipRotEnd,
                                                         frameFraction)
    self._rightHipVel = self.ComputeAngVelRel(rightHipRotStart, rightHipRotEnd, keyFrameDuration,
                                              bullet_client)

    rightKneeRotStart = [frameData[20]]
    rightKneeRotEnd = [frameDataNext[20]]
    self._rightKneeRot = [
        rightKneeRotStart[0] + frameFraction * (rightKneeRotEnd[0] - rightKneeRotStart[0])
    ]
    self._rightKneeVel = [(rightKneeRotEnd[0] - rightKneeRotStart[0]) / keyFrameDuration]

    rightAnkleRotStart = [frameData[22], frameData[23], frameData[24], frameData[21]]
    rightAnkleRotEnd = [frameDataNext[22], frameDataNext[23], frameDataNext[24], frameDataNext[21]]
    self._rightAnkleRot = bullet_client.getQuaternionSlerp(rightAnkleRotStart, rightAnkleRotEnd,
                                                           frameFraction)
    self._rightAnkleVel = self.ComputeAngVelRel(rightAnkleRotStart, rightAnkleRotEnd,
                                                keyFrameDuration, bullet_client)

    rightShoulderRotStart = [frameData[26], frameData[27], frameData[28], frameData[25]]
    rightShoulderRotEnd = [
        frameDataNext[26], frameDataNext[27], frameDataNext[28], frameDataNext[25]
    ]
    self._rightShoulderRot = bullet_client.getQuaternionSlerp(rightShoulderRotStart,
                                                              rightShoulderRotEnd, frameFraction)
    self._rightShoulderVel = self.ComputeAngVelRel(rightShoulderRotStart, rightShoulderRotEnd,
                                                   keyFrameDuration, bullet_client)

    rightElbowRotStart = [frameData[29]]
    rightElbowRotEnd = [frameDataNext[29]]
    self._rightElbowRot = [
        rightElbowRotStart[0] + frameFraction * (rightElbowRotEnd[0] - rightElbowRotStart[0])
    ]
    self._rightElbowVel = [(rightElbowRotEnd[0] - rightElbowRotStart[0]) / keyFrameDuration]

    leftHipRotStart = [frameData[31], frameData[32], frameData[33], frameData[30]]
    leftHipRotEnd = [frameDataNext[31], frameDataNext[32], frameDataNext[33], frameDataNext[30]]
    self._leftHipRot = bullet_client.getQuaternionSlerp(leftHipRotStart, leftHipRotEnd,
                                                        frameFraction)
    self._leftHipVel = self.ComputeAngVelRel(leftHipRotStart, leftHipRotEnd, keyFrameDuration,
                                             bullet_client)

    leftKneeRotStart = [frameData[34]]
    leftKneeRotEnd = [frameDataNext[34]]
    self._leftKneeRot = [
        leftKneeRotStart[0] + frameFraction * (leftKneeRotEnd[0] - leftKneeRotStart[0])
    ]
    self._leftKneeVel = [(leftKneeRotEnd[0] - leftKneeRotStart[0]) / keyFrameDuration]

    leftAnkleRotStart = [frameData[36], frameData[37], frameData[38], frameData[35]]
    leftAnkleRotEnd = [frameDataNext[36], frameDataNext[37], frameDataNext[38], frameDataNext[35]]
    self._leftAnkleRot = bullet_client.getQuaternionSlerp(leftAnkleRotStart, leftAnkleRotEnd,
                                                          frameFraction)
    self._leftAnkleVel = self.ComputeAngVelRel(leftAnkleRotStart, leftAnkleRotEnd,
                                               keyFrameDuration, bullet_client)

    leftShoulderRotStart = [frameData[40], frameData[41], frameData[42], frameData[39]]
    leftShoulderRotEnd = [
        frameDataNext[40], frameDataNext[41], frameDataNext[42], frameDataNext[39]
    ]
    self._leftShoulderRot = bullet_client.getQuaternionSlerp(leftShoulderRotStart,
                                                             leftShoulderRotEnd, frameFraction)
    self._leftShoulderVel = self.ComputeAngVelRel(leftShoulderRotStart, leftShoulderRotEnd,
                                                  keyFrameDuration, bullet_client)

    leftElbowRotStart = [frameData[43]]
    leftElbowRotEnd = [frameDataNext[43]]
    self._leftElbowRot = [
        leftElbowRotStart[0] + frameFraction * (leftElbowRotEnd[0] - leftElbowRotStart[0])
    ]
    self._leftElbowVel = [(leftElbowRotEnd[0] - leftElbowRotStart[0]) / keyFrameDuration]

    pose = self.GetPose()
    return pose

  def ConvertFromAction(self, pybullet_client, action):
    #turn action into pose

    self.Reset()  #?? needed?
    index = 0

    self._chestRot = getQuaternionFromAction(action[index:index+self.action_dim], self.action_representation_mode)
    index += self.action_dim
    self._neckRot = getQuaternionFromAction(action[index:index+self.action_dim], self.action_representation_mode)
    index += self.action_dim
    self._rightHipRot = getQuaternionFromAction(action[index:index+self.action_dim], self.action_representation_mode)
    index += self.action_dim
    self._rightKneeRot = [action[index]]
    index += 1
    self._rightAnkleRot = getQuaternionFromAction(action[index:index+self.action_dim], self.action_representation_mode)
    index += self.action_dim
    self._rightShoulderRot = getQuaternionFromAction(action[index:index+self.action_dim], self.action_representation_mode)
    index += self.action_dim
    self._rightElbowRot = [action[index]]
    index += 1
    self._leftHipRot = getQuaternionFromAction(action[index:index+self.action_dim], self.action_representation_mode)
    index += self.action_dim
    self._leftKneeRot = [action[index]]
    index += 1
    self._leftAnkleRot = getQuaternionFromAction(action[index:index+self.action_dim], self.action_representation_mode)
    index += self.action_dim
    self._leftShoulderRot = getQuaternionFromAction(action[index:index+self.action_dim], self.action_representation_mode)
    index += self.action_dim
    self._leftElbowRot = [action[index]]
    index += 1

    # if self.action_representation_mode == "6D":
    #   sixdim = action[index:index + 6]
    #   axis, angle = getAxisAngleFromSixDim(sixdim)
    #   self._chestRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
    #   index += 6

    #   sixdim = action[index:index + 6]
    #   axis, angle = getAxisAngleFromSixDim(sixdim)
    #   self._neckRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
    #   index += 6

    #   sixdim = action[index:index + 6]
    #   axis, angle = getAxisAngleFromSixDim(sixdim)
    #   self._rightHipRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
    #   index += 6

    #   angle = action[index]
    #   self._rightKneeRot = [angle]
    #   index += 1

    #   sixdim = action[index:index + 6]
    #   axis, angle = getAxisAngleFromSixDim(sixdim)
    #   self._rightAnkleRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
    #   index += 6

    #   sixdim = action[index:index + 6]
    #   axis, angle = getAxisAngleFromSixDim(sixdim)
    #   self._rightShoulderRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
    #   index += 6

    #   angle = action[index]
    #   self._rightElbowRot = [angle]
    #   index += 1

    #   sixdim = action[index:index + 6]
    #   axis, angle = getAxisAngleFromSixDim(sixdim)
    #   self._leftHipRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
    #   index += 6

    #   angle = action[index]
    #   self._leftKneeRot = [angle]
    #   index += 1

    #   sixdim = action[index:index + 6]
    #   axis, angle = getAxisAngleFromSixDim(sixdim)
    #   self._leftAnkleRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
    #   index += 6

    #   sixdim = action[index:index + 6]
    #   axis, angle = getAxisAngleFromSixDim(sixdim)
    #   self._leftShoulderRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
    #   index += 6

    #   angle = action[index]
    #   self._leftElbowRot = [angle]
    #   index += 1

    # elif self.action_representation_mode == "AxisAngle":
    #   angle = action[index]
    #   axis = [action[index + 1], action[index + 2], action[index + 3]]
    #   index += 4
    #   self._chestRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
    #   #print("pose._chestRot=",pose._chestRot)

    #   angle = action[index]
    #   axis = [action[index + 1], action[index + 2], action[index + 3]]
    #   index += 4
    #   self._neckRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)

    #   angle = action[index]
    #   axis = [action[index + 1], action[index + 2], action[index + 3]]
    #   index += 4
    #   self._rightHipRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)

    #   angle = action[index]
    #   index += 1
    #   self._rightKneeRot = [angle]

    #   angle = action[index]
    #   axis = [action[index + 1], action[index + 2], action[index + 3]]
    #   index += 4
    #   self._rightAnkleRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)

    #   angle = action[index]
    #   axis = [action[index + 1], action[index + 2], action[index + 3]]
    #   index += 4
    #   self._rightShoulderRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)

    #   angle = action[index]
    #   index += 1
    #   self._rightElbowRot = [angle]

    #   angle = action[index]
    #   axis = [action[index + 1], action[index + 2], action[index + 3]]
    #   index += 4
    #   self._leftHipRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)

    #   angle = action[index]
    #   index += 1
    #   self._leftKneeRot = [angle]

    #   angle = action[index]
    #   axis = [action[index + 1], action[index + 2], action[index + 3]]
    #   index += 4
    #   self._leftAnkleRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)

    #   angle = action[index]
    #   axis = [action[index + 1], action[index + 2], action[index + 3]]
    #   index += 4
    #   self._leftShoulderRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)

    #   angle = action[index]
    #   index += 1
    #   self._leftElbowRot = [angle]

    pose = self.GetPose()

    return pose


from mpi4py import MPI
from neuron import h
from Simulation import Simulation
from cells import AfferentFiber
import random as rnd
import time
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import pickle
from tools import seed_handler as sh
sh.set_seed()

comm = MPI.COMM_WORLD
sizeComm = comm.Get_size()
rank = comm.Get_rank()

class CollisionEesNatural(Simulation):
	""" Simulation to evaluate the effect of EES on the natural firing rate in function of the fiber
	legnth, of its firing rate and of the frequency of stimulation.
	"""

	def __init__(self, parallelContext, eesFrequencies, fiberDelays, fiberFiringRates, segmentToRecord = None, tstop = 5000):
		""" Object initialization.

		Keyword arguments:
		parallelContext -- Neuron parallelContext object.
		eesFrequencies -- List of stimulation frequencies to test.
		fiberDelays -- List of fiber delays to test.
		fiberFiringRates -- List of fiber firing rates to test.
		"""

		Simulation.__init__(self,parallelContext)

		if rank==1:
			print "\nMPI execution: the different processes have different stimulation starting time."
			print "The final result is the mean results between each process\n"

		# Variables initializations
		self._eesFrequencies = eesFrequencies
		self._fiberDelays = fiberDelays
		self._fiberFiringRates = fiberFiringRates
		self._segmentToRecord = segmentToRecord

		self._init_lists()
		self._results = np.zeros([len(self._eesFrequencies),len(self._fiberDelays),len(self._fiberFiringRates)])

		self._create_fibers()
		self._create_ees_objects()
		self._connect_ees_to_fibers()

		self._set_tstop(tstop)
		self._set_integration_step(AfferentFiber.get_update_period())

	"""
	Redefinition of inherited methods
	"""

	def _update(self):
		""" Update simulation parameters. """
		self._update_afferents()

	def _end_integration(self):
		""" Print the total simulation time and extract the results. """
		Simulation._end_integration(self)
		self._extract_results()

	def save_results(self,name=""):
		""" Save the simulation results.

		Keyword arguments:
		name -- string to add at predefined file name (default = "").
		"""
		fileName = time.strftime("%Y_%m_%d_resultsCollisionEesNatural"+name+".p")
		with open(self._resultsFolder+fileName, 'w') as pickle_file:
			pickle.dump(self._results, pickle_file)
			pickle.dump(self._eesFrequencies, pickle_file)
			pickle.dump(self._fiberDelays, pickle_file)
			pickle.dump(self._fiberFiringRates, pickle_file)

	def plot(self,delay,nColorLevels=None,name=""):
		""" Plot the simulation results.

		Plot the percantage of collisions for a given delay in fucntion of the afferent
		firing rate and of the stimulation frequency.
		Keyword arguments:
		delay -- fiber delay for which we want the plot.
		Threshold -- threshold to plot binary simulation results (default = None).
		"""
		if rank == 0:
			fig, ax = plt.subplots(figsize=(16,9))

			dataToPlot = self._results[:,delay,:]
			title = "Percentage of sensory information erased by the stimulation\n (delay "+str(self._fiberDelays[delay])+" ms)"
			if nColorLevels is not None:
				dataToPlot=np.round(dataToPlot/100*nColorLevels)*100/nColorLevels
				print dataToPlot
				title += "\nnColorLevels = "+str(nColorLevels)

			# cmap = plt.cm.gray
			cmap = plt.cm.bone_r
			im = ax.imshow(dataToPlot, cmap=cmap, interpolation='nearest',origin="lower",vmin = 0, vmax = 100)
			ax.set_title(title)

			# Move left and bottom spines outward by 10 points
			ax.spines['left'].set_position(('outward', 10))
			ax.spines['bottom'].set_position(('outward', 10))
			# Hide the right and top spines
			ax.spines['right'].set_visible(False)
			ax.spines['top'].set_visible(False)
			# Only show ticks on the left and bottom spines
			ax.yaxis.set_ticks_position('left')
			ax.xaxis.set_ticks_position('bottom')
			fig.colorbar(im, orientation='vertical',label='% Erased APs')

			plt.yticks(range(len(self._eesFrequencies)),self._eesFrequencies)
			plt.xticks(range(len(self._fiberFiringRates)),self._fiberFiringRates)
			plt.ylabel('EES eesFrequency (Hz)')
			plt.xlabel('Natural afferent firing rate (hz)')

			fileName = time.strftime("%Y_%m_%d_CollisionEesNatural_Delay_"+str(self._fiberDelays[delay])+name+".pdf")
			plt.savefig(self._resultsFolder+fileName, format="pdf",transparent=True)
			plt.show(block=False)

	"""
	Specific Methods of this class
	"""
	def _init_lists(self):
		""" Initialize lists containg the fibers, netcon objects and ees objects. """
		self._fiberList = [[[] for i in range(len(self._fiberDelays))] for j in range(len(self._eesFrequencies))]
		self._netconList = [[[] for i in range(len(self._fiberDelays))] for j in range(len(self._eesFrequencies))]
		self._eesList = []

	def _create_fibers(self):
		""" Create the fibers with the defined different delays. """
		for i in range(len(self._eesFrequencies)):
			for j in range(len(self._fiberDelays)):
				for k in range(len(self._fiberFiringRates)):
					self._fiberList[i][j].append(AfferentFiber(self._fiberDelays[j]))
					if self._segmentToRecord is None:
						self._fiberList[i][j][k].set_firing_rate(self._fiberFiringRates[k])
					if self._segmentToRecord is not None:
						self._fiberList[i][j][k].set_firing_rate(self._fiberFiringRates[k],False)
						self._fiberList[i][j][k].set_recording(True,self._segmentToRecord)

	def _create_ees_objects(self):
		""" Create different ees objects with the defined stimulation frequencies. """
		scale = rnd.random()
		for i in range(len(self._eesFrequencies)):
			self._eesList.append(h.NetStim())
			self._eesList[i].interval = 1000.0/self._eesFrequencies[i]
			self._eesList[i].number = 10000
			self._eesList[i].start = 10.0*scale
			self._eesList[i].noise = 0

	def _connect_ees_to_fibers(self):
		""" Connect fibers ojects to ees objects to make the stimulation activate these fibers. """
		for i in range(len(self._eesFrequencies)):
			for j in range(len(self._fiberDelays)):
				for k in range(len(self._fiberFiringRates)):
					self._netconList[i][j].append(h.NetCon(self._eesList[i],self._fiberList[i][j][k].cell))
					self._netconList[i][j][k].delay = 1
					self._netconList[i][j][k].weight[0] = AfferentFiber.get_ees_weight()

	def _update_afferents(self):
		""" Update the afferents fiber state. """
		for i in range(len(self._eesFrequencies)):
			for j in range(len(self._fiberDelays)):
				for k in range(len(self._fiberFiringRates)):
					self._fiberList[i][j][k].update(h.t)

	def _extract_results(self):
		""" Extract the simulation results. """
		for i in range(len(self._eesFrequencies)):
			for j in range(len(self._fiberDelays)):
				for k in range(len(self._fiberFiringRates)):
					sent,arr,coll,perc=self._fiberList[i][j][k].get_stats()
					self._results[i,j,k]=perc
		comm.Barrier()
		if sizeComm>1:
			temp = comm.gather(self._results, root=0)
			if rank==0:
				for i in range(1,sizeComm):
					self._results += temp[i]
				self._results/=sizeComm

	def plot_isoinformation_surface(self,percentage=50):
		""" Plot a surface where the number of AP erased by the stimulation is equal. """
		if rank==0:
			Z = np.zeros([len(self._fiberFiringRates),len(self._fiberDelays)])
			temp = (self._results - percentage)**2
			for x in xrange(len(self._fiberFiringRates)):
				for y in xrange(len(self._fiberDelays)):
					Z[x,y]=self._eesFrequencies[temp[:,y,x].argmin()]

			fig, ax = plt.subplots(figsize=(16,9))
			im = ax.imshow(Z, cmap=plt.cm.bone, interpolation='nearest',origin="lower")
			ax.set_title("Isoinformation surface - "+str(percentage)+"% of APs erased")

			# Move left and bottom spines outward by 10 points
			ax.spines['left'].set_position(('outward', 10))
			ax.spines['bottom'].set_position(('outward', 10))
			# Hide the right and top spines
			ax.spines['right'].set_visible(False)
			ax.spines['top'].set_visible(False)
			# Only show ticks on the left and bottom spines
			ax.yaxis.set_ticks_position('left')
			ax.xaxis.set_ticks_position('bottom')
			fig.colorbar(im, orientation='vertical', label="Stimulation frequency (Hz)")

			plt.xticks(range(len(self._fiberDelays)),self._fiberDelays)
			plt.yticks(range(len(self._fiberFiringRates)),self._fiberFiringRates)
			plt.xlabel('Fiber delay (ms)')
			plt.ylabel('Natural afferents firing rate (hz)')

			fileName = time.strftime("%Y_%m_%d_CollisionEesNatural_Isoinfo_"+str(percentage)+"perc.pdf")
			plt.savefig(self._resultsFolder+fileName, format="pdf",transparent=True)
			plt.show(block=False)

	def plot_recorded_segment(self, freqInd = 0, delInd = 0, firInd = 0):
		""" Plot recorded spikes from a fiber.

			Keyword arguments:
			freqInd -- index of the stimulation frequencies.
			delInd -- index of the fiber delay.
			firInd -- index of the fiber natural firing rate.
		"""

		if self._segmentToRecord == None: return
		naturalSignals, eesInducedSignals, trigger, time = self._fiberList[freqInd][delInd][firInd].get_recording()
		nNaturalSent,nNaturalArrived,nCollisions,percErasedAp = self._fiberList[freqInd][delInd][firInd].get_stats()

		fig1, ax1 = plt.subplots(1, 1, figsize=(8,4.5))
		msToRec = 20
		sumTrigNaturalSignal = np.zeros(msToRec)
		sumTrigEesSignal = np.zeros(msToRec)
		nPeripheralStims = nNaturalSent*np.ones(msToRec)

		for i,val in enumerate(trigger):
			if val and i+msToRec<len(time):
				sumTrigNaturalSignal += naturalSignals[i:i+msToRec]
				sumTrigEesSignal += eesInducedSignals[i:i+msToRec]
		ax1.plot(sumTrigNaturalSignal,color='b',label='peripheral')
		ax1.plot(sumTrigEesSignal,color='r',label='spinal')
		ax1.plot(nPeripheralStims,color='g',ls = '--', label='n peripheral stims')
		ax1.set_ylim([-10,nNaturalSent+10])

		collisionPerc = 100*(1-sumTrigNaturalSignal.max()/float(nNaturalSent))
		ax1.set_title("Spikes in segment {0}, collision perc: {1:.1f}%".format(self._segmentToRecord,collisionPerc))
		ax1.legend()

		fig2, ax2 = plt.subplots(1, 1, figsize=(16,4))
		ax2.plot(time,naturalSignals,color='b',label='peripheral')
		ax2.plot(time,eesInducedSignals,color='r',label='spinal')
		ax2.plot(time,trigger,color='g',label='trigger',ls='--')
		ax2.set_ylim([-0.5,1.5])
		ax2.legend()
		plt.show()


__author__ = 'sebastian'

import numpy as np
import math
import os
from photogrammetry_importer.utility.blender_logging_utility import log_report

class Camera:
    """
    This class represents a reconstructed camera and provides functionality to manage
    intrinsic and extrinsic camera parameters as well as image information.
    """
    panoramic_type_equirectangular = "EQUIRECTANGULAR"

    IMAGE_FP_TYPE_NAME = "NAME"
    IMAGE_FP_TYPE_RELATIVE = "RELATIVE"
    IMAGE_FP_TYPE_ABSOLUTE = "ABSOLUTE"

    DEPTH_MAP_WRT_UNIT_VECTORS = "DEPTH_MAP_WRT_UNIT_VECTORS"
    DEPTH_MAP_WRT_CANONICAL_VECTORS = "DEPTH_MAP_WRT_CANONICAL_VECTORS"

    def __init__(self):
        self._center = np.array([0, 0, 0], dtype=float)              # C = -R^T t
        self._translation_vec = np.array([0, 0, 0], dtype=float)     # t = -R C
        self.normal = np.array([0, 0, 0], dtype=float)
        self.color = np.array([255, 255, 255], dtype=int)

        # use for these attributes the getter and setter methods
        self._quaternion = np.array([0, 0, 0, 0], dtype=float)
        self._rotation_mat = np.zeros((3, 3), dtype=float)

        self._calibration_mat = np.zeros((3, 3), dtype=float)

        self.image_fp_type = None
        self.image_dp = None
        self._relative_fp = None
        self._absolute_fp = None
        self._undistorted_relative_fp = None
        self._undistorted_absolute_fp = None
        self.width = None
        self.height = None
        self.panoramic_type = None

        self.depth_map_fp = None
        self.depth_map_callback = None
        self.depth_map_semantic = None

        self.id = None  # an unique identifier (natural number)

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return str('Camera: ' + self._relative_fp + ' ' + str(self._center) + ' ' + str(self.normal))

    def get_file_name(self):
        return os.path.basename(self.get_absolute_fp())

    def set_relative_fp(self, relative_fp, image_fp_type):
        self._relative_fp = relative_fp
        self.image_fp_type = image_fp_type

    def get_relative_fp(self):
        return self._get_relative_fp(
            self._relative_fp, self._absolute_fp)

    def get_undistorted_relative_fp(self):
        return self._get_relative_fp(
            self._undistorted_relative_fp, self._undistorted_absolute_fp)

    def _get_relative_fp(self, relative_fp, absolute_fp):
        if self.image_fp_type == Camera.IMAGE_FP_TYPE_NAME:
            assert relative_fp is not None
            return os.path.basename(relative_fp)
        elif self.image_fp_type == Camera.IMAGE_FP_TYPE_RELATIVE:
            assert relative_fp is not None
            return relative_fp
        elif self.image_fp_type == Camera.IMAGE_FP_TYPE_ABSOLUTE:
            assert absolute_fp is not None
            return absolute_fp
        else:
            assert False

    def set_absolute_fp(self, absolute_fp):
        self._absolute_fp = absolute_fp

    def get_absolute_fp(self):
        return self._get_absolute_fp(
            self._relative_fp, self._absolute_fp)

    def get_undistored_absolute_fp(self):
        if self.image_fp_type == Camera.IMAGE_FP_TYPE_ABSOLUTE:
            assert False # Not supported for undistorted images
        return self._get_absolute_fp(
            self._undistorted_relative_fp, self._undistorted_absolute_fp)

    def _get_absolute_fp(self, relative_fp, absolute_fp):
        if self.image_fp_type == Camera.IMAGE_FP_TYPE_NAME:
            assert self.image_dp is not None
            assert relative_fp is not None
            return os.path.join(self.image_dp, os.path.basename(relative_fp))
        elif self.image_fp_type == Camera.IMAGE_FP_TYPE_RELATIVE:
            assert self.image_dp is not None
            assert relative_fp is not None
            return os.path.join(self.image_dp, relative_fp)
        elif self.image_fp_type == Camera.IMAGE_FP_TYPE_ABSOLUTE:
            assert absolute_fp is not None
            return absolute_fp
        else:
            assert False

    def has_undistorted_absolute_fp(self):
        requirements = False
        if self.image_fp_type == Camera.IMAGE_FP_TYPE_NAME:
            requirements = (self.image_dp is not None) and (self._undistorted_relative_fp is not None)
        elif self.image_fp_type == Camera.IMAGE_FP_TYPE_RELATIVE:
            requirements = (self.image_dp is not None) and (self._undistorted_relative_fp is not None)
        elif self.image_fp_type == Camera.IMAGE_FP_TYPE_ABSOLUTE:
            requirements = (self._undistorted_absolute_fp is not None)

        has_fp = False
        if requirements:
            fp = self._get_absolute_fp(
                self._undistorted_relative_fp,
                self._undistorted_absolute_fp)
            if os.path.isfile(fp):
                has_fp = True
        return has_fp

    def get_blender_obj_gui_str(self):
        # Replace special characters
        #image_fp_clean = image_fp.replace("/", "_").replace("\\", "_").replace(":", "_")
        image_fp_stem = os.path.splitext(self.get_relative_fp())[0]
        # Blender supports only object names with length 63
        # However, we need also space for additional suffixes
        image_fp_suffix = image_fp_stem[-40:]
        return image_fp_suffix

    def set_calibration(self, calibration_mat, radial_distortion):
        self._calibration_mat = np.asarray(calibration_mat, dtype=float)
        self._radial_distortion = radial_distortion
        assert self._radial_distortion is not None

    def has_focal_length(self):
        return self._calibration_mat[0][0] > 0

    def get_focal_length(self):
        return self._calibration_mat[0][0]

    def get_field_of_view(self):
        assert self.width is not None and self.height is not None
        angle = math.atan(max(self.width, self.height) / (self.get_focal_length() * 2.0)) * 2.0
        return angle

    def has_intrinsics(self):
        return self.has_focal_length() and self.is_principal_point_initialized()

    def check_calibration_mat(self):
        assert self.has_focal_length() and self.is_principal_point_initialized()

    def get_calibration_mat(self):
        self.check_calibration_mat()
        return self._calibration_mat

    def set_calibration_mat(self, calibration_mat):
        self._calibration_mat = calibration_mat

    def set_principal_point(self, principal_point):
        self._calibration_mat[0][2] = principal_point[0]
        self._calibration_mat[1][2] = principal_point[1]

    def get_principal_point(self):
        calibration_mat = self.get_calibration_mat()
        cx = calibration_mat[0][2]
        cy = calibration_mat[1][2]
        return np.asarray([cx,cy], dtype=float)

    def is_principal_point_initialized(self):
        cx_zero = np.isclose(self._calibration_mat[0][2], 0.0)
        cy_zero = np.isclose(self._calibration_mat[1][2], 0.0)
        initialized = (not cx_zero) and (not cy_zero)
        return initialized

    def is_panoramic(self):
        return self.panoramic_type is not None

    def set_panoramic_type(self, panoramic_type):
        self.panoramic_type = panoramic_type

    def get_panoramic_type(self):
        return self.panoramic_type

    @staticmethod
    def compute_calibration_mat(focal_length, cx, cy):
        return np.array([[focal_length, 0, cx], [0, focal_length, cy], [0,0,1]], dtype=float)

    def set_quaternion(self, quaternion):
        self._quaternion = quaternion
        # we must change the rotation matrixes as well
        self._rotation_mat = Camera.quaternion_to_rotation_matrix(quaternion)

    def set_rotation_mat(self, rotation_mat, check_rotation=True):
        if check_rotation:
            assert Camera.is_rotation_mat_valid(rotation_mat)
        self._rotation_mat = rotation_mat
        # we must change the quaternion as well
        self._quaternion = Camera.rotation_matrix_to_quaternion(rotation_mat)

    def set_camera_center_after_rotation(self, center, check_rotation=True):
        if check_rotation:
            assert Camera.is_rotation_mat_valid(self._rotation_mat)
        self._center = center
        self._translation_vec = - np.dot(self._rotation_mat, center)    # t = -R C

    def set_camera_translation_vector_after_rotation(self, translation_vector, check_rotation=True):
        if check_rotation:
            assert Camera.is_rotation_mat_valid(self._rotation_mat)
        self._translation_vec = translation_vector
        self._center = - np.dot(self._rotation_mat.transpose(), translation_vector) # C = -R^T t

    def get_quaternion(self):
        return self._quaternion

    def get_rotation_mat(self):
        return self._rotation_mat

    def get_translation_vec(self):
        return self._translation_vec

    def get_camera_center(self):
        return self._center

    def set_4x4_cam_to_world_mat(self, cam_to_world_mat, check_rotation=True):
        self.set_rotation_mat(
            cam_to_world_mat[0:3, 0:3].transpose(), check_rotation=check_rotation)
        self.set_camera_center_after_rotation(
            cam_to_world_mat[0:3, 3], check_rotation=check_rotation)

    @staticmethod
    def is_rotation_mat_valid(some_mat):
        # Test if rotation_mat is really a rotation matrix (i.e. det = -1 or det = 1)
        det = np.linalg.det(some_mat)
        res = np.isclose(det, 1) or np.isclose(det, -1)
        return res

    @staticmethod
    def quaternion_to_rotation_matrix(q):
        """
        Original C++ Method ('SetQuaternionRotation()') defined in  pba/src/pba/DataInterface.h
        Parallel bundle adjustment (pba) code (used by visualsfm) is provided here:
        http://grail.cs.washington.edu/projects/mcba/
        """
        qq = math.sqrt(q[0]*q[0]+q[1]*q[1]+q[2]*q[2]+q[3]*q[3])
        if qq > 0:  # Normalize the quaternion
            qw = q[0]/qq
            qx = q[1]/qq
            qy = q[2]/qq
            qz = q[3]/qq
        else:
            qw = 1
            qx = qy = qz = 0
        m = np.zeros((3, 3), dtype=float)
        m[0][0] = float(qw*qw + qx*qx- qz*qz- qy*qy )
        m[0][1] = float(2*qx*qy -2*qz*qw )
        m[0][2] = float(2*qy*qw + 2*qz*qx)
        m[1][0] = float(2*qx*qy+ 2*qw*qz)
        m[1][1] = float(qy*qy+ qw*qw - qz*qz- qx*qx)
        m[1][2] = float(2*qz*qy- 2*qx*qw)
        m[2][0] = float(2*qx*qz- 2*qy*qw)
        m[2][1] = float(2*qy*qz + 2*qw*qx )
        m[2][2] = float(qz*qz+ qw*qw- qy*qy- qx*qx)
        return m

    @staticmethod
    def rotation_matrix_to_quaternion(m):
        """
        Original C++ Method ('GetQuaternionRotation()') defined in  pba/src/pba/DataInterface.h
        Parallel bundle adjustment (pba) code (used by visualsfm) is provided here:
        http://grail.cs.washington.edu/projects/mcba/
        """
        q = np.array([0, 0, 0, 0], dtype=float)
        q[0] = 1 + m[0][0] + m[1][1] + m[2][2]
        if q[0] > 0.000000001:
            q[0] = math.sqrt(q[0]) / 2.0
            q[1] = (m[2][1] - m[1][2]) / ( 4.0 * q[0])
            q[2] = (m[0][2] - m[2][0]) / ( 4.0 * q[0])
            q[3] = (m[1][0] - m[0][1]) / ( 4.0 * q[0])
        else:
            if m[0][0] > m[1][1] and m[0][0] > m[2][2]:
                s = 2.0 * math.sqrt(1.0 + m[0][0] - m[1][1] - m[2][2])
                q[1] = 0.25 * s
                q[2] = (m[0][1] + m[1][0]) / s
                q[3] = (m[0][2] + m[2][0]) / s
                q[0] = (m[1][2] - m[2][1]) / s
            elif m[1][1] > m[2][2]:
                s = 2.0 * math.sqrt(1.0 + m[1][1] - m[0][0] - m[2][2])
                q[1] = (m[0][1] + m[1][0]) / s
                q[2] = 0.25 * s
                q[3] = (m[1][2] + m[2][1]) / s
                q[0] = (m[0][2] - m[2][0]) / s
            else:
                s = 2.0 * math.sqrt(1.0 + m[2][2] - m[0][0] - m[1][1])
                q[1] = (m[0][2] + m[2][0]) / s
                q[2] = (m[1][2] + m[2][1]) / s
                q[3] = 0.25 * s
                q[0] = (m[0][1] - m[1][0]) / s
        return q

    def set_depth_map(self, depth_map_ifp, depth_map_callback, depth_map_semantic):
        self.depth_map_fp = depth_map_ifp
        self.depth_map_callback = depth_map_callback
        self.depth_map_semantic = depth_map_semantic

    def get_depth_map(self):
        if os.path.isfile(self.depth_map_fp):
            return self.depth_map_callback(self.depth_map_fp)
        else:
            return None

    def get_4x4_cam_to_world_mat(self):
        """
        This matrix can be used to convert points given in camera coordinates
        into points given in world coordinates.
        M = [R^T    c]
            [0      1]
        :return:
        """
        homogeneous_mat = np.identity(4, dtype=float)
        homogeneous_mat[0:3, 0:3] = self.get_rotation_mat().transpose()
        homogeneous_mat[0:3, 3] = self.get_camera_center()
        return homogeneous_mat

    def convert_depth_map_to_world_coords(  self,
                                            depth_map_display_sparsity=100):
        """
        Do not confuse z_buffer with depth_buffer!
        z_buffer contains values in [0,1]
        depth_buffer contains the actual distance values

        :param depth_buffer_matrix:
        :param n_th_result_point:
        :return:
        """
        assert 0 < depth_map_display_sparsity

        depth_map = self.get_depth_map()
        height, width = depth_map.shape

        if self.height == height and self.width == width:
            x_step_size = 1.0
            y_step_size = 1.0
        else:
            x_step_size = self.width / width
            y_step_size = self.height / height

        fx = self.get_calibration_mat()[0][0]
        fy = self.get_calibration_mat()[1][1]
        cx, cy = self.get_principal_point()

        indices = np.indices((height, width))
        y_index_list = indices[0].flatten()
        x_index_list = indices[1].flatten()

        # Use the local coordinate system of the camera to analyze its viewing directions
        # The Blender camera coordinate system looks along the negative z axis (blue),
        # the up axis points along the y axis (green).

        y_index_list = y_index_list[::-1]   # Reverse order of indices
        x_index_list = x_index_list[::-1]   # Reverse order of indices
        fx = -fx
        fy = -fy
        depth_values = depth_map.flatten()

        assert len(x_index_list) == len(y_index_list) == len(depth_values)

        # https://github.com/colmap/colmap/blob/dev/src/base/reconstruction.cc
        #   // COLMAP assumes that the upper left pixel center is (0.5, 0.5)
        # https://github.com/simonfuhrmann/mve/blob/master/libs/mve/depthmap.cc
        #  math::Vec3f v = invproj * math::Vec3f(
        #       (float)x + 0.5f, (float)y + 0.5f, 1.0f);
        x_index_coord_list = x_step_size * x_index_list + 0.5
        y_index_coord_list = y_step_size * y_index_list + 0.5

        # The cannoncial vectors are defined according to p.155 of
        # "Multiple View Geometry" by Hartley and Zisserman using a canonical
        # focal length of 1 , i.e. vec = [(x - cx) / fx, (y - cy) / fy, 1]
        x_coords_canonical = (x_index_coord_list - cx) / fx
        y_coords_canonical = (y_index_coord_list - cy) / fy
        z_coords_canonical = np.ones(len(depth_values), dtype=float)

        # Determine non-background data
        non_background_flags = depth_values > 0
        x_coords_canonical_filtered = x_coords_canonical[non_background_flags]
        y_coords_canonical_filtered = y_coords_canonical[non_background_flags]
        z_coords_canonical_filtered = z_coords_canonical[non_background_flags]
        depth_values_filtered = depth_values[non_background_flags]

        if depth_map_display_sparsity != 100:
            x_coords_canonical_filtered = x_coords_canonical_filtered[::depth_map_display_sparsity]
            y_coords_canonical_filtered = y_coords_canonical_filtered[::depth_map_display_sparsity]
            z_coords_canonical_filtered = z_coords_canonical_filtered[::depth_map_display_sparsity]
            depth_values_filtered = depth_values_filtered[::depth_map_display_sparsity]

        if self.depth_map_semantic == Camera.DEPTH_MAP_WRT_CANONICAL_VECTORS:
            # In this case, the depth values are defined w.r.t. the canonical
            # vectors. This kind of depth data is used by Colmap.
            x_coords_filtered = x_coords_canonical_filtered * depth_values_filtered
            y_coords_filtered = y_coords_canonical_filtered * depth_values_filtered
            z_coords_filtered = z_coords_canonical_filtered * depth_values_filtered

        elif self.depth_map_semantic == Camera.DEPTH_MAP_WRT_UNIT_VECTORS:
            # In this case the depth values are defined w.r.t. the normalized
            # canonical vectors. This kind of depth data is used by MVE.
            cannonical_norms_filtered = np.linalg.norm(
                np.array(
                    [x_coords_canonical_filtered,
                    y_coords_canonical_filtered,
                    z_coords_canonical_filtered],
                    dtype=float),
                axis=0)
            # Instead of normalizing the x,y and z component, we divide the
            # depth values by the corresponding norm.
            normalized_depth_values_filtered = depth_values_filtered / cannonical_norms_filtered
            x_coords_filtered = x_coords_canonical_filtered * normalized_depth_values_filtered
            y_coords_filtered = y_coords_canonical_filtered * normalized_depth_values_filtered
            z_coords_filtered = z_coords_canonical_filtered * normalized_depth_values_filtered

        else:
            assert False

        hom_entries = np.ones_like(z_coords_filtered)
        cam_coords_hom = np.dstack(
            (x_coords_filtered,
            y_coords_filtered,
            z_coords_filtered,
            hom_entries))[0]

        world_coords_hom = self.get_4x4_cam_to_world_mat().dot(cam_coords_hom.T).T
        world_coords = np.delete(world_coords_hom, 3, 1)

        return world_coords


from torchvision import transforms
from torch.utils.data import DataLoader
from torch.utils.data.sampler import BatchSampler

import numpy as np
import torch
import os

from elbo_functions import deviance_upper_bound, elbo, KL_closed, minibatch_KLD_upper_bound, minibatch_KLD_upper_bound_iter
from model_test import MSE_test_GPapprox, MSE_test
from utils import SubjectSampler, VaryingLengthSubjectSampler, VaryingLengthBatchSampler, HensmanDataLoader
from predict_HealthMNIST import recon_complete_gen, gen_rotated_mnist_plot, variational_complete_gen
from validation import validate

def hensman_training(nnet_model, type_nnet, epochs, dataset, optimiser, type_KL, num_samples, latent_dim, covar_module0,
                     covar_module1, likelihoods, m, H, zt_list, P, T, varying_T, Q, weight, id_covariate, loss_function,
                     natural_gradient=False, natural_gradient_lr=0.01, subjects_per_batch=20, memory_dbg=False,
                     eps=1e-6, results_path=None, validation_dataset=None, generation_dataset=None,
                     prediction_dataset=None, gp_model=None, csv_file_test_data=None, csv_file_test_label=None,
                     test_mask_file=None, data_source_path=None):

    """
    Perform training with minibatching and Stochastic Variational Inference [Hensman et. al, 2013]. See L-VAE supplementary
    materials

    :param nnet_model: encoder/decoder neural network model
    :param type_nnet: type of encoder/decoder
    :param epochs: numner of epochs
    :param dataset: dataset to use in training
    :param optimiser: optimiser to be used
    :param type_KL: type of KL divergenve computation to use
    :param num_samples: number of samples to use
    :param latent_dim: number of latent dimensions
    :param covar_module0: additive kernel (sum of cross-covariances) without id covariate
    :param covar_module1: additive kernel (sum of cross-covariances) with id covariate
    :param likelihoods: GPyTorch likelihood model
    :param m: variational mean
    :param H: variational variance
    :param zt_list: list of inducing points
    :param P: number of unique instances
    :param T: number of longitudinal samples per individual
    :param Q: number of covariates
    :param weight: value for the weight
    :param id_covariate: covariate number of the id
    :param loss_function: selected loss function
    :param natural_gradient: use of natural gradients
    :param natural_gradient_lr: natural gradients learning rate
    :param subject_per_batch; number of subjects per batch (vectorisation)
    :param memory_dbg: enable debugging
    :param eps: jitter
    :param results_path: path to results
    :param validation_dataset: dataset for vaildation set
    :param generation_dataset: dataset to help with sample image generation
    :param prediction_dataset; dataset with subjects for prediction
    :param gp_mode: GPyTorch gp model
    :param csv_file_test_data: path to test data
    :param csv_file_test_label: path to test label
    :param test_mask_file: path to test mask
    :param data_source_path: path to data source

    :return trained models and resulting losses

    """

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    N = len(dataset)
    assert type_KL == 'GPapprox_closed'

    if varying_T:
        n_batches = (P + subjects_per_batch - 1)//subjects_per_batch
        dataloader = HensmanDataLoader(dataset, batch_sampler=VaryingLengthBatchSampler(VaryingLengthSubjectSampler(dataset, id_covariate), subjects_per_batch), num_workers=4)
    else:
        batch_size = subjects_per_batch*T
        n_batches = (P*T + batch_size - 1)//(batch_size)
        dataloader = HensmanDataLoader(dataset, batch_sampler=BatchSampler(SubjectSampler(dataset, P, T), batch_size, drop_last=False), num_workers=4)

    net_train_loss_arr = np.empty((0, 1))
    recon_loss_arr = np.empty((0, 1))
    nll_loss_arr = np.empty((0, 1))
    kld_loss_arr = np.empty((0, 1))
    penalty_term_arr = np.empty((0, 1))
    best_val_pred_mse = np.Inf
    best_epoch = 0
    for epoch in range(1, epochs + 1):
        recon_loss_sum = 0
        nll_loss_sum = 0
        kld_loss_sum = 0
        net_loss_sum = 0
        iid_kld_sum = 0
        for batch_idx, sample_batched in enumerate(dataloader):
            optimiser.zero_grad()
            nnet_model.train()
            covar_module0.train()
            covar_module1.train()
            indices = sample_batched['idx']
            data = sample_batched['digit'].double().to(device)
            train_x = sample_batched['label'].double().to(device)
            mask = sample_batched['mask'].double().to(device)
            N_batch = data.shape[0]

            covariates = torch.cat((train_x[:, :id_covariate], train_x[:, id_covariate+1:]), dim=1)

            recon_batch, mu, log_var = nnet_model(data)
            [recon_loss, nll] = nnet_model.loss_function(recon_batch, data, mask)
            recon_loss = torch.sum(recon_loss)
            nll_loss = torch.sum(nll)

            PSD_H = H if natural_gradient else torch.matmul(H, H.transpose(-1, -2))

            if varying_T:
                P_in_current_batch = torch.unique(train_x[:, id_covariate]).shape[0]
                kld_loss, grad_m, grad_H = minibatch_KLD_upper_bound_iter(covar_module0, covar_module1, likelihoods, latent_dim, m, PSD_H, train_x, mu, log_var, zt_list, P, P_in_current_batch, N, natural_gradient, id_covariate, eps)
            else:
                P_in_current_batch = N_batch // T
                kld_loss, grad_m, grad_H = minibatch_KLD_upper_bound(covar_module0, covar_module1, likelihoods, latent_dim, m, PSD_H, train_x, mu, log_var, zt_list, P, P_in_current_batch, T, natural_gradient, eps)

            recon_loss = recon_loss * P/P_in_current_batch
            nll_loss = nll_loss * P/P_in_current_batch

            if loss_function == 'nll':
                net_loss = nll_loss + kld_loss
            elif loss_function == 'mse':
                kld_loss = kld_loss / latent_dim
                net_loss = recon_loss + weight * kld_loss

            net_loss.backward()
            optimiser.step()

            if natural_gradient:
                LH = torch.cholesky(H)
                iH = torch.cholesky_solve(torch.eye(H.shape[-1], dtype=torch.double).to(device), LH)
                iH_new = iH + natural_gradient_lr*(grad_H + grad_H.transpose(-1,-2))
                LiH_new = torch.cholesky(iH_new)
                H = torch.cholesky_solve(torch.eye(H.shape[-1], dtype=torch.double).to(device), LiH_new).detach()
                m = torch.matmul(H, torch.matmul(iH, m) - natural_gradient_lr*(grad_m - 2*torch.matmul(grad_H, m))).detach()

            net_loss_sum += net_loss.item() / n_batches
            recon_loss_sum += recon_loss.item() / n_batches
            nll_loss_sum += nll_loss.item() / n_batches
            kld_loss_sum += kld_loss.item() / n_batches

        print('Iter %d/%d - Loss: %.3f  - GP loss: %.3f  - NLL Loss: %.3f  - Recon Loss: %.3f' % (
            epoch, epochs, net_loss_sum, kld_loss_sum, nll_loss_sum, recon_loss_sum), flush=True)
        penalty_term_arr = np.append(penalty_term_arr, 0.0)
        net_train_loss_arr = np.append(net_train_loss_arr,  net_loss_sum)
        recon_loss_arr = np.append(recon_loss_arr, recon_loss_sum)
        nll_loss_arr = np.append(nll_loss_arr, nll_loss_sum)
        kld_loss_arr = np.append(kld_loss_arr, kld_loss_sum)

        if (not epoch % 25) and epoch != epochs:
            with torch.no_grad():
                nnet_model.eval()
                covar_module0.eval()
                covar_module1.eval()
                if validation_dataset is not None:
                    full_mu = torch.zeros(len(dataset), latent_dim, dtype=torch.double).to(device)
                    prediction_x = torch.zeros(len(dataset), Q, dtype=torch.double).to(device)
                    for batch_idx, sample_batched in enumerate(dataloader):
                        label_id = sample_batched['idx']
                        prediction_x[label_id] = sample_batched['label'].double().to(device)
                        data = sample_batched['digit'].double().to(device)
                        covariates = torch.cat((prediction_x[label_id, :id_covariate], prediction_x[label_id, id_covariate+1:]), dim=1)

                        mu, log_var = nnet_model.encode(data)
                        full_mu[label_id] = mu
                    val_pred_mse = validate(nnet_model, type_nnet, validation_dataset, type_KL, num_samples, latent_dim, covar_module0, covar_module1, likelihoods, zt_list, T, weight, full_mu, prediction_x, id_covariate, loss_function, eps=1e-6)
                    if val_pred_mse < best_val_pred_mse:
                        best_val_pred_mse = val_pred_mse
                        best_epoch = epoch

                        prediction_dataloader = DataLoader(prediction_dataset, batch_sampler=VaryingLengthBatchSampler(
                            VaryingLengthSubjectSampler(prediction_dataset, id_covariate), subjects_per_batch),
                                                           num_workers=4)
                        full_mu = torch.zeros(len(prediction_dataset), latent_dim, dtype=torch.double).to(device)
                        prediction_x = torch.zeros(len(prediction_dataset), Q, dtype=torch.double).to(device)

                        with torch.no_grad():
                            for batch_idx, sample_batched in enumerate(prediction_dataloader):
                                label_id = sample_batched['idx']
                                prediction_x[label_id] = sample_batched['label'].double().to(device)
                                data = sample_batched['digit'].double().to(device)
                                covariates = torch.cat(
                                    (prediction_x[label_id, :id_covariate], prediction_x[label_id, id_covariate + 1:]),
                                    dim=1)

                                mu, log_var = nnet_model.encode(data)
                                full_mu[label_id] = mu
                            covar_module0.eval()
                            covar_module1.eval()
                            if type_KL == 'GPapprox' or type_KL == 'GPapprox_closed':
                                MSE_test_GPapprox(csv_file_test_data, csv_file_test_label, test_mask_file,
                                                  data_source_path, type_nnet,
                                                  nnet_model, covar_module0, covar_module1, likelihoods, results_path,
                                                  latent_dim, prediction_x,
                                                  full_mu, zt_list, P, T, id_covariate, varying_T,
                                                  save_file='result_error_best.csv')

                        print('Saving better model')
                        try:
                            torch.save(nnet_model.state_dict(), os.path.join(results_path, 'nnet_model_best.pth'))
                            torch.save(gp_model.state_dict(), os.path.join(results_path, 'gp_model_best.pth'))
                            torch.save(zt_list, os.path.join(results_path, 'zt_list_best.pth'))
                            torch.save(m, os.path.join(results_path, 'm_best.pth'))
                            torch.save(H, os.path.join(results_path, 'H_best.pth'))

                            if results_path and generation_dataset:
                                prediction_dataloader = DataLoader(prediction_dataset,
                                                                   batch_sampler=VaryingLengthBatchSampler(
                                                                       VaryingLengthSubjectSampler(prediction_dataset,
                                                                                id_covariate),
                                                                       subjects_per_batch), num_workers=4)
                                full_mu = torch.zeros(len(prediction_dataset), latent_dim, dtype=torch.double).to(
                                    device)
                                prediction_x = torch.zeros(len(prediction_dataset), Q, dtype=torch.double).to(device)
                                for batch_idx, sample_batched in enumerate(prediction_dataloader):
                                    label_id = sample_batched['idx']
                                    prediction_x[label_id] = sample_batched['label'].double().to(device)
                                    data = sample_batched['digit'].double().to(device)
                                    covariates = torch.cat((prediction_x[label_id, :id_covariate],
                                                            prediction_x[label_id, id_covariate + 1:]), dim=1)

                                    mu, log_var = nnet_model.encode(data)
                                    full_mu[label_id] = mu

                                recon_complete_gen(generation_dataset, nnet_model, type_nnet,
                                                   results_path, covar_module0,
                                                   covar_module1, likelihoods, latent_dim,
                                                   './data', prediction_x, full_mu, epoch,
                                                   zt_list, P, T, id_covariate, varying_T)
                        except e:
                            print(e)
                            print('Saving intermediate model failed!')
                            pass
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()

    return penalty_term_arr, net_train_loss_arr, nll_loss_arr, recon_loss_arr, kld_loss_arr, m, H, best_epoch

def minibatch_training(nnet_model, type_nnet, epochs, dataset, optimiser, type_KL, num_samples, latent_dim,
                       covar_module0, covar_module1, likelihoods, zt_list, P, T, Q, weight, id_covariate,
                       loss_function, memory_dbg=False, eps=1e-6, results_path=None, validation_dataset=None,
                       generation_dataset=None, prediction_dataset=None):

    """
    Perform training with minibatching (psuedo-minibatching) similar to GPPVAE [Casale el. al, 2018]. See L-VAE supplementary
    materials

    :param nnet_model: encoder/decoder neural network model
    :param type_nnet: type of encoder/decoder
    :param epochs: numner of epochs
    :param dataset: dataset to use in training
    :param optimiser: optimiser to be used
    :param type_KL: type of KL divergenve computation to use
    :param num_samples: number of samples to use
    :param latent_dim: number of latent dimensions
    :param covar_module0: additive kernel (sum of cross-covariances) without id covariate
    :param covar_module1: additive kernel (sum of cross-covariances) with id covariate
    :param likelihoods: GPyTorch likelihood model
    :param zt_list: list of inducing points
    :param P: number of unique instances
    :param T: number of longitudinal samples per individual
    :param Q: number of covariates
    :param weight: value for the weight
    :param id_covariate: covariate number of the id
    :param loss_function: selected loss function
    :param memory_dbg: enable debugging
    :param eps: jitter
    :param results_path: path to results
    :param validation_dataset: dataset for vaildation set
    :param generation_dataset: dataset to help with sample image generation
    :param prediction_dataset; dataset with subjects for prediction

    :return trained models and resulting losses

    """

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

    batch_size = T
    assert (type_KL == 'GPapprox_closed' or type_KL == 'GPapprox')

    # set up Data Loader for training
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    net_train_loss_arr = np.empty((0, 1))
    recon_loss_arr = np.empty((0, 1))
    nll_loss_arr = np.empty((0, 1))
    gp_loss_arr = np.empty((0, 1))
    penalty_term_arr = np.empty((0, 1))

    for epoch in range(1, epochs + 1):

        optimiser.zero_grad()

        full_mu = torch.zeros(len(dataset), latent_dim, dtype=torch.double, requires_grad=True).to(device)
        full_log_var = torch.zeros(len(dataset), latent_dim, dtype=torch.double, requires_grad=True).to(device)
        train_x = torch.zeros(len(dataset), Q, dtype=torch.double, requires_grad=False).to(device)

        #Step 1: Encode the sample data to obtain \bar{\mu} and diag(W)
        with torch.no_grad():
            for batch_idx, sample_batched in enumerate(dataloader):
                indices = sample_batched['idx']
                data = sample_batched['digit'].double().to(device)
                train_x[indices] = sample_batched['label'].double().to(device)

                covariates = torch.cat((train_x[indices, :id_covariate], train_x[indices, id_covariate+1:]), dim=1)
                mu, log_var = nnet_model.encode(data)

                full_mu[indices] = mu
                full_log_var[indices] = log_var

        mu_grads = torch.zeros(len(dataset), latent_dim, dtype=torch.double, requires_grad=True).to(device)
        log_var_grads = torch.zeros(len(dataset), latent_dim, dtype=torch.double, requires_grad=True).to(device)

        gp_losses = 0
        gp_loss_sum = 0
        param_list = []

        #Steps 2 & 3: compute d and E, compute gradients of KLD w.r.t S and theta
        if type_KL == 'GPapprox':
            for sample in range(0, num_samples):
                Z = nnet_model.sample_latent(full_mu, full_log_var)
                for i in range(0, latent_dim):
                    Z_dim = Z[:, i]
                    gp_loss = -elbo(covar_module0[i], covar_module1[i], likelihoods[i], train_x, Z_dim,
                                    zt_list[i].to(device), P, T, eps)
                    gp_loss_sum = gp_loss.item() + gp_loss_sum
                    gp_losses = gp_losses + gp_loss
            gp_losses = gp_losses / num_samples
            gp_loss_sum /= num_samples

        elif type_KL == 'GPapprox_closed':
            for i in range(0, latent_dim):
                mu_sliced = full_mu[:, i]
                log_var_sliced = full_log_var[:, i]
                gp_loss = deviance_upper_bound(covar_module0[i], covar_module1[i],
                                               likelihoods[i], train_x,
                                               mu_sliced, log_var_sliced,
                                               zt_list[i].to(device), P,
                                               T, eps)
                gp_loss_sum = gp_loss.item() + gp_loss_sum
                gp_losses = gp_losses + gp_loss

        for i in range(0, latent_dim):
            param_list += list(covar_module0[i].parameters())
            param_list += list(covar_module1[i].parameters())
#            param_list.append(zt_list[i])

        if loss_function == 'mse':
            gp_losses = weight*gp_losses/latent_dim
            gp_loss_sum /= latent_dim

        mu_grads = torch.autograd.grad(gp_losses, full_mu, retain_graph=True)[0]
        log_var_grads = torch.autograd.grad(gp_losses, full_log_var, retain_graph=True)[0]
        grads = torch.autograd.grad(gp_losses, param_list)

        for ind, p in enumerate(param_list):
            p.grad = grads[ind]

        recon_loss_sum = 0
        nll_loss_sum = 0
        #Step 4: compute reconstruction losses w.r.t phi and psi, add dKLD/dphi to the gradients
        for batch_idx, sample_batched in enumerate(dataloader):
            data = sample_batched['digit'].double().to(device)
            mask = sample_batched['mask'].double().to(device)
            indices = sample_batched['idx']

            label = sample_batched['label'].double().to(device)
            covariates = torch.cat((label[:, :id_covariate], label[:, id_covariate+1:]), dim=1)
            recon_batch, mu, log_var = nnet_model(data)

            [recon_loss, nll] = nnet_model.loss_function(recon_batch, data, mask)
            recon_loss = torch.sum(recon_loss)
            nll = torch.sum(nll)

            mu.backward(mu_grads[indices], retain_graph = True)
            log_var.backward(log_var_grads[indices], retain_graph = True)

            if loss_function == 'mse':
                recon_loss.backward()
            elif loss_function == 'nll':
                nll.backward()

            recon_loss_sum = recon_loss_sum + recon_loss.item()
            nll_loss_sum = nll_loss_sum + nll.item()

        #Do logging
        print('Iter %d/%d - Loss: %.3f  - GP loss: %.3f  - NLL loss: %.3f  - Recon Loss: %.3f' % (
            epoch, epochs, recon_loss_sum + weight*gp_loss_sum, gp_loss_sum, nll_loss_sum, recon_loss_sum))
        penalty_term_arr = np.append(penalty_term_arr, 0.0)
        net_train_loss_arr = np.append(net_train_loss_arr,  recon_loss_sum + weight*gp_loss_sum)
        nll_loss_arr = np.append(nll_loss_arr, nll_loss_sum)
        recon_loss_arr = np.append(recon_loss_arr, recon_loss_sum)
        gp_loss_arr = np.append(gp_loss_arr, gp_loss_sum)

        #Step 5: apply gradients using an Adam optimiser
        optimiser.step()

        if (not epoch % 100) and epoch != epochs:
            if validation_dataset is not None:
                validate(nnet_model, type_nnet, validation_dataset, type_KL, num_samples, latent_dim, covar_module0, covar_module1, likelihoods, zt_list, T, weight, full_mu, train_x, id_covariate, loss_function, eps=1e-6)
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

            if results_path and generation_dataset:
                prediction_dataloader = DataLoader(prediction_dataset, batch_size=1000, shuffle=False, num_workers=4)
                full_mu = torch.zeros(len(prediction_dataset), latent_dim, dtype=torch.double).to(device)
                prediction_x = torch.zeros(len(prediction_dataset), Q, dtype=torch.double).to(device)
                with torch.no_grad():
                    for batch_idx, sample_batched in enumerate(prediction_dataloader):
                        # no mini-batching. Instead get a batch of dataset size
                        label_id = sample_batched['idx']
                        prediction_x[label_id] = sample_batched['label'].double().to(device)
                        data = sample_batched['digit'].double().to(device)
                        covariates = torch.cat((prediction_x[label_id, :id_covariate], prediction_x[label_id, id_covariate+1:]), dim=1)

                        mu, log_var = nnet_model.encode(data)

                        full_mu[label_id] = mu

                    recon_complete_gen(generation_dataset, nnet_model, type_nnet,
                                       results_path, covar_module0,
                                       covar_module1, likelihoods, latent_dim,
                                       './data', prediction_x, full_mu, epoch,
                                       zt_list, P, T, id_covariate)

    return penalty_term_arr, net_train_loss_arr, nll_loss_arr, recon_loss_arr, gp_loss_arr

def standard_training(nnet_model, type_nnet, epochs, dataset, optimiser, type_KL, num_samples,
    latent_dim, covar_modules, likelihoods, zt_list, id_covariate, P, T, Q, weight, constrain_scales,
    loss_function, memory_dbg=False, eps=1e-6, validation_dataset=None, generation_dataset=None, prediction_dataset=None):

    """
    Perform training without minibatching.

    :param nnet_model: encoder/decoder neural network model
    :param type_nnet: type of encoder/decoder
    :param epochs: numner of epochs
    :param dataset: dataset to use in training
    :param optimiser: optimiser to be used
    :param type_KL: type of KL divergenve computation to use
    :param num_samples: number of samples to use
    :param latent_dim: number of latent dimensions
    :param covar_modules: additive kernel (sum of cross-covariances)
    :param likelihoods: GPyTorch likelihood model
    :param zt_list: list of inducing points
    :param id_covariate: covariate number of the id
    :param P: number of unique instances
    :param T: number of longitudinal samples per individual
    :param Q: number of covariates
    :param weight: value for the weight
    :param constrain_scales: boolean to constrain scales to 1
    :param loss_function: selected loss function
    :param memory_dbg: enable debugging
    :param eps: jitter
    :param validation_dataset: dataset for vaildation set
    :param generation_dataset: dataset to help with sample image generation
    :param prediction_dataset; dataset with subjects for prediction

    :return trained models and resulting losses

    """
    if type_KL == 'closed':
        covar_module = covar_modules[0]
    elif type_KL == 'GPapprox' or type_KL == 'GPapprox_closed':
        covar_module0 = covar_modules[0]
        covar_module1 = covar_modules[1]

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

    # set up Data Loader for training
    dataloader = DataLoader(dataset, batch_size=len(dataset), shuffle=False, num_workers=4)

    net_train_loss_arr = np.empty((0, 1))
    recon_loss_arr = np.empty((0, 1))
    nll_loss_arr = np.empty((0, 1))
    gp_loss_arr = np.empty((0, 1))
    penalty_term_arr = np.empty((0, 1))

    for epoch in range(1, epochs + 1):
        for batch_idx, sample_batched in enumerate(dataloader):

            # no mini-batching. Instead get a batch of dataset size.
            optimiser.zero_grad()                                       # clear gradients
            label_id = sample_batched['idx']
            label = sample_batched['label']
            data = sample_batched['digit']
            data = data.double().to(device)
            mask = sample_batched['mask']
            mask = mask.to(device)

            train_x = label.double().to(device)
            covariates = torch.cat((train_x[:, :id_covariate], train_x[:, id_covariate+1:]), dim=1)

            # encode data
            recon_batch, mu, log_var = nnet_model(data)

            [recon_loss, nll] = nnet_model.loss_function(recon_batch, data, mask)
            recon_loss = torch.sum(recon_loss)
            nll_loss = torch.sum(nll)

            gp_loss_avg = torch.tensor([0.0]).to(device)
            net_loss = torch.tensor([0.0]).to(device)
            penalty_term = torch.tensor([0.0]).to(device)

            for sample_iter in range(0, num_samples):

                # Iterate over specified number of samples. Default: num_samples = 1.
                Z = nnet_model.sample_latent(mu, log_var)
                gp_loss = torch.tensor([0.0]).to(device)

                for i in range(0, latent_dim):
                    Z_dim = Z[:, i].view(-1).type(torch.DoubleTensor).to(device)

                    if type_KL == 'closed':

                        # Closed-form KL divergence formula
                        kld1 = KL_closed(covar_module[i], train_x, likelihoods[i], data,  mu[:, i], log_var[:, i])
                        gp_loss = gp_loss + kld1
                    elif type_KL == 'conj_gradient':

                        # GPyTorch default: use modified batch conjugate gradients
                        # See: https://arxiv.org/abs/1809.11165
                        gp_models[i].set_train_data(train_x.to(device), Z_dim.to(device))
                        gp_loss = gp_loss - mlls[i](gp_models[i](train_x.to(device)), Z_dim)
                    elif type_KL == 'GPapprox':

                        # Our proposed efficient approximate GP inference scheme
                        # See: http://arxiv.org/abs/2006.09763
                        loss = -elbo(covar_module0[i], covar_module1[i], likelihoods[i], train_x, Z_dim,
                                     zt_list[i].to(device), P, T, eps)
                        gp_loss = gp_loss + loss

                    elif type_KL == 'GPapprox_closed':

                        # A variant of our proposed efficient approximate GP inference scheme.
                        # The key difference with GPapprox is the direct use of the variational mean and variance,
                        # instead of a sample from Z. We can call this a deviance upper bound.
                        # See the L-VAE supplement for more details: http://arxiv.org/abs/2006.09763
                        loss = deviance_upper_bound(covar_module0[i], covar_module1[i], likelihoods[i], train_x,
                                                    mu[:, i].view(-1), log_var[:, i].view(-1), zt_list[i].to(device), P,
                                                    T, eps)
                        gp_loss = gp_loss + loss

                if type_KL == 'closed' or type_KL == 'GPapprox' or type_KL == 'GPapprox_closed':
                    if loss_function == 'mse':
                        gp_loss_avg = gp_loss_avg + (gp_loss / latent_dim)
                    elif loss_function == 'nll':
                        gp_loss_avg = gp_loss_avg + gp_loss
                elif type_KL == 'conj_gradient':
                    if loss_function == 'mse':
                        gp_loss = gp_loss * data.shape[0] / latent_dim
                    elif loss_function == 'nll':
                        gp_loss = gp_loss * data.shape[0]
                    gp_loss_avg = gp_loss_avg + gp_loss

            if type_KL == 'closed' or type_KL == 'GPapprox' or type_KL == 'GPapprox_closed':
                gp_loss_avg = gp_loss_avg / num_samples
                if loss_function == 'mse':
                    net_loss = recon_loss + weight * gp_loss_avg
                elif loss_function == 'nll':
                    net_loss = nll_loss + gp_loss_avg
            elif type_KL == 'conj_gradient':
                gp_loss_avg = gp_loss_avg / num_samples
                penalty_term = -0.5 * log_var.sum() / latent_dim
                if loss_function == 'mse':
                    net_loss = recon_loss + weight * (gp_loss_avg + penalty_term)
                elif loss_function == 'nll':
                    net_loss = nll_loss + gp_loss_avg + penalty_term

            net_loss.backward()

            if type_KL == 'closed' or type_KL == 'GPapprox' or type_KL == 'GPapprox_closed':
                print('Iter %d/%d - Loss: %.3f  - GP loss: %.3f  - NLL Loss: %.3f  - Recon Loss: %.3f' % (
                    epoch, epochs, net_loss.item(), gp_loss_avg.item(), nll_loss.item(), recon_loss.item()))
            elif type_KL == 'conj_gradient':
                print('Iter %d/%d - Loss: %.3f  - GP loss: %.3f  - Penalty: %.3f  - NLL Loss: %.3f  - Recon Loss: %.3f' % (
                    epoch, epochs, net_loss.item(), gp_loss_avg.item(), penalty_term.item(), nll_loss.item(), recon_loss.item()))

            penalty_term_arr = np.append(penalty_term_arr, penalty_term.cpu().item())
            net_train_loss_arr = np.append(net_train_loss_arr, net_loss.cpu().item())
            recon_loss_arr = np.append(recon_loss_arr, recon_loss.cpu().item())
            nll_loss_arr = np.append(nll_loss_arr, nll_loss.cpu().item())
            gp_loss_arr = np.append(gp_loss_arr, gp_loss_avg.cpu().item())
            optimiser.step()
            if constrain_scales:
                for i in range(0, latent_dim):
                    likelihoods[i].noise = torch.tensor([1], dtype=torch.float).to(device)

            if (not epoch % 100) and epoch != epochs:
                if validation_dataset is not None:
                    standard_validate(nnet_model, type_nnet, validation_dataset, type_KL, num_samples, latent_dim, covar_module0, covar_module1, likelihoods, zt_list, T, weight, mu, train_x, id_covariate, loss_function, eps=1e-6)
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()

    return penalty_term_arr, net_train_loss_arr, nll_loss_arr, recon_loss_arr, gp_loss_arr

def variational_inference_optimization(nnet_model, type_nnet, epochs, dataset, prediction_dataset, optimiser,
    latent_dim, covar_module0, covar_module1, likelihoods, zt_list, P, T, Q, weight, constrain_scales,
    id_covariate, loss_function, memory_dbg=False, eps=1e-6, results_path=None, save_path=None, gp_model_folder=None,
    generation_dataset=None):

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

    # set up Data Loader for training
    dataloader = DataLoader(dataset, batch_size=len(dataset), shuffle=False, num_workers=4)

    net_train_loss_arr = np.empty((0, 1))
    recon_loss_arr = np.empty((0, 1))
    nll_loss_arr = np.empty((0, 1))
    gp_loss_arr = np.empty((0, 1))
    penalty_term_arr = np.empty((0, 1))

    for batch_idx, sample_batched in enumerate(dataloader):
        label_id = sample_batched['idx']
        label = sample_batched['label'].double().to(device)
        data = sample_batched['digit'].double().to(device)
        mask = sample_batched['mask'].double().to(device)

        covariates = torch.cat((label[:, :id_covariate], label[:, id_covariate+1:]), dim=1)

        # encode data
        mu, log_var = nnet_model.encode(data)

    mu = torch.nn.Parameter(mu.clone().detach(), requires_grad=True)
    log_var = torch.nn.Parameter(log_var.clone().detach(), requires_grad=True)

    try:
        mu = torch.load(os.path.join(gp_model_folder, 'mu.pth'), map_location=torch.device(device)).detach().to(device).requires_grad_(True)
        log_var = torch.load(os.path.join(gp_model_foder, 'log_var.pth'), map_location=torch.device(device)).detach().to(device).requires_grad_(True)
    except:
        pass

    optimiser.add_param_group({'params': mu})
    optimiser.add_param_group({'params': log_var})

    for epoch in range(1, epochs + 1):
        optimiser.zero_grad()
        Z = nnet_model.sample_latent(mu, log_var)
        recon_batch = nnet_model.decode(Z)
        [recon_loss, nll] = nnet_model.loss_function(recon_batch, data, mask)
        recon_loss = torch.sum(recon_loss)
        nll_loss = torch.sum(nll)

        gp_loss_avg = torch.tensor([0.0]).to(device)
        net_loss = torch.tensor([0.0]).to(device)
        penalty_term = torch.tensor([0.0]).to(device)

        for i in range(0, latent_dim):
            loss = deviance_upper_bound(covar_module0[i], covar_module1[i], likelihoods[i], label,
                                        mu[:, i].view(-1), log_var[:, i].view(-1), zt_list[i].to(device), P,
                                        T, eps)
            gp_loss_avg = gp_loss_avg + loss / latent_dim

        if loss_function == 'mse':
            net_loss = recon_loss + weight * gp_loss_avg
        elif loss_function == 'nll':
            net_loss = nll_loss + gp_loss_avg

        net_loss.backward()

        print('Iter %d/%d - Loss: %.3f  - GP loss: %.3f  - NLL Loss: %.3f  - Recon Loss: %.3f' % (
              epoch, epochs, net_loss.item(), gp_loss_avg.item(), nll_loss.item(), recon_loss.item()),
              flush=True)

        penalty_term_arr = np.append(penalty_term_arr, penalty_term.cpu().item())
        net_train_loss_arr = np.append(net_train_loss_arr, net_loss.cpu().item())
        recon_loss_arr = np.append(recon_loss_arr, recon_loss.cpu().item())
        nll_loss_arr = np.append(nll_loss_arr, nll_loss.cpu().item())
        gp_loss_arr = np.append(gp_loss_arr, gp_loss_avg.cpu().item())
        optimiser.step()

        if not epoch % 100:
            sv_pth = os.path.join(save_path, 'recon_' + str(epoch) + '.pdf')
            gen_rotated_mnist_plot(data[1920:2080].cpu().detach(), recon_batch[1920:2080].cpu().detach(), label[1920:2080].cpu().detach(), seq_length=20, num_sets=8, save_file=sv_pth)

    torch.save(nnet_model.state_dict(), os.path.join(save_path, 'final-vae_model.pth'))
    torch.save(mu, os.path.join(save_path, 'mu.pth'))
    torch.save(log_var, os.path.join(save_path, 'log_var.pth'))
    for i in range(0, latent_dim):
        torch.save(covar_module0[i].state_dict(), os.path.join(save_path, 'cov_module0_' + str(i) + '.pth'))
        torch.save(covar_module1[i].state_dict(), os.path.join(save_path, 'cov_module1_' + str(i) + '.pth'))

    prediction_dataloader = DataLoader(prediction_dataset, batch_size=len(prediction_dataset), shuffle=False, num_workers=1)
    for batch_idx, sample_batched in enumerate(prediction_dataloader):
        label_pred = sample_batched['label'].double().to(device)
        data_pred = sample_batched['digit'].double().to(device)
        mask_pred = sample_batched['mask'].double().to(device)
        covariates = torch.cat((label_pred[:, :id_covariate], label_pred[:, id_covariate+1:]), dim=1)
        # encode data
        mu_pred, log_var_pred = nnet_model.encode(data_pred)
        break

    try:
        mu_pred = torch.load(os.path.join(gp_model_folder, 'mu_pred.pth'), map_location=torch.device(device)).detach().to(device).requires_grad_(True)
        log_var_pred = torch.load(os.path.join(gp_model_folder, 'log_var_pred.pth'), map_location=torch.device(device)).detach().to(device).requires_grad_(True)
    except:
        pass

    mu_pred = torch.nn.Parameter(mu_pred.clone().detach(), requires_grad=True)
    log_var_pred = torch.nn.Parameter(log_var_pred.clone().detach(), requires_grad=True)
    adam_param_list = []
    adam_param_list.append({'params': mu_pred})
    adam_param_list.append({'params': log_var_pred})
    optimiser_pred = torch.optim.Adam(adam_param_list, lr=1e-3)
    for epoch in range(1, 1001):
        optimiser_pred.zero_grad()

        Z = nnet_model.sample_latent(mu_pred, log_var_pred)

        recon_batch = nnet_model.decode(Z)
        [recon_loss, nll] = nnet_model.loss_function(recon_batch,
                                                     data_pred,
                                                     mask_pred)

        recon_loss = torch.sum(recon_loss)
        nll_loss = torch.sum(nll)

        gp_loss_avg = torch.tensor([0.0]).to(device)

        prediction_mu = torch.cat((mu_pred, mu), dim=0)
        prediction_log_var = torch.cat((log_var_pred, log_var), dim=0)
        prediction_x = torch.cat((label_pred, label), dim=0)

        for i in range(0, latent_dim):
            loss = deviance_upper_bound(covar_module0[i], covar_module1[i], likelihoods[i], prediction_x,
                                        prediction_mu[:, i].view(-1), prediction_log_var[:, i].view(-1),
                                        zt_list[i].to(device), P+8, T, eps)
            gp_loss_avg = gp_loss_avg + loss / latent_dim

        if loss_function == 'mse':
            net_loss = recon_loss + weight * gp_loss_avg
        elif loss_function == 'nll':
            net_loss = nll_loss + gp_loss_avg

        net_loss.backward()

        print('Iter %d/1000 - Total Loss: %.3f  - GP Loss: %.3f  - Recon Loss: %.3f' % (
              epoch, net_loss.item(), gp_loss_avg.item(), recon_loss.item()),
              flush=True)

        optimiser_pred.step()

    torch.save(mu_pred, os.path.join(save_path, 'mu_pred.pth'))
    torch.save(log_var_pred, os.path.join(save_path, 'log_var_pred.pth'))

    l = [i*20 + k for i in range(0,8) for k in range(0,5)]
    prediction_x = torch.cat((label_pred[l],
                               label))
    prediction_mu = torch.cat((mu_pred[l],
                               mu))

    if generation_dataset:
        variational_complete_gen(generation_dataset, nnet_model, type_nnet,
                                 results_path, covar_module0,
                                 covar_module1, likelihoods, latent_dim,
                                 './data', prediction_x, prediction_mu, 'final',
                                 zt_list, P, T, id_covariate)

    exit(0)


# -*- coding: utf-8 -*-
# pylint: disable=wrong-import-position

"""
Generate single-DOM Retro tables binned in (t,r,theta), with each bin
containing a survival probability and average directionality vector. The length
of this vector indicates "how directional" light is, from 0 (isotropic) to 1
(perfectly directional). The information provided comes from the single-DOM
Retro tables generated by CLSim, which are binned in
(theta,r,t,theta_dir,deltaphi_dir).
"""

from __future__ import absolute_import, division, print_function

__all__ = ['generate_t_r_theta_table']

__author__ = 'P. Eller, J.L. Lanfranchi'
__license__ = '''Copyright 2017 Philipp Eller and Justin L. Lanfranchi

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.'''

from os.path import abspath, dirname
import sys

import numpy as np

if __name__ == '__main__' and __package__ is None:
    PARENT_DIR = dirname(dirname(abspath(__file__)))
    if PARENT_DIR not in sys.path:
        sys.path.append(PARENT_DIR)
from retro import DFLT_NUMBA_JIT_KWARGS, numba_jit
from retro.const import SPEED_OF_LIGHT_M_PER_NS
from retro.utils.stats import weighted_average

@numba_jit(**DFLT_NUMBA_JIT_KWARGS)
def generate_t_r_theta_table(
        table,
        n_photons,
        group_refractive_index,
        t_bin_width,
        angular_acceptance_fract,
        thetadir_centers,
        deltaphidir_centers,
        theta_bin_edges
    ):
    """Transform information from a raw single-DOM table (as output from CLSim)
    that is binned in (r, theta, t, theta_dir, deltaphi_dir) into a more
    compact representation, with a probability and an average direction vector
    (represented by a single theta and phi) in each (t, r, theta) bin.

    Parameters
    ----------
    table
    n_photons
    group_refractive_index
    t_bin_width
    angular_acceptance_fract
    thetadir_centers
    deltaphidir_centers
    theta_bin_edges

    Returns
    -------
    survival_probs
    average_thetas
    average_phis
    lengths

    """
    # Source tables are photon counts binned in
    # (r, theta, t, dir_theta, dir_phi)
    n_r_bins = table.shape[0]
    n_theta_bins = table.shape[1]
    n_t_bins = table.shape[2]

    # (Base) survival probability (that will be modified by directionality):
    # We can either
    # 1. Sum over the directionality dimensions, which means "probability that
    #    photon is going in any one of these directions is P_{det, tot}. This
    #    is like the total area under a curve, but then we have to distribute
    #    this total area among all the directions via the distribution we
    #    parameterize as fn of dir vector length."
    # 2. Max over directionality dimensions, which means "the best P_det occurs
    #    if the photon is moving in this particular direction, whereupon it
    #    will be detected with P_{det, max}. It gets worse from there according
    #    to the distribution we parameterize as fn of dir vector length."

    norm = (
        1
        / n_photons
        / (SPEED_OF_LIGHT_M_PER_NS / group_refractive_index)
        / t_bin_width
        * angular_acceptance_fract
        * n_theta_bins
    )

    # Destination tables are to be binned in (t, r, costheta) (there are as
    # many costheta bins as theta bins in the original tables)
    dest_shape = (n_t_bins, n_r_bins, n_theta_bins)

    survival_probs = np.empty(dest_shape, dtype=np.float32)
    average_thetas = np.empty(dest_shape, dtype=np.float32)
    average_phis = np.empty(dest_shape, dtype=np.float32)
    lengths = np.empty(dest_shape, dtype=np.float32)

    for r_i in range(n_r_bins):
        for theta_j in range(n_theta_bins):
            for t_k in range(n_t_bins):
                # flip coszen_dir (photon direction)
                weights = table[r_i, theta_j, t_k, ::-1, :].astype(np.float64)
                weights_tot = weights.sum()
                if weights_tot == 0:
                    # If no photons, just set the average direction to the
                    # theta of the bin center...
                    average_theta = 0.5 * (theta_bin_edges[theta_j]
                                           + theta_bin_edges[theta_j + 1])
                    # ... and lengths to 0
                    length = 0.0
                    average_phi = 0.0
                else:
                    # Average theta
                    weights_theta = weights.sum(axis=1)
                    average_theta = weighted_average(thetadir_centers,
                                                     weights_theta)

                    # Average delta phi
                    projected_survival_prob = (
                        (weights.T * np.sin(thetadir_centers)).T
                    )
                    weights_phi = projected_survival_prob.sum(axis=0)
                    average_phi = weighted_average(deltaphidir_centers,
                                                   weights_phi)

                    # Length of vector (using projections from all vectors
                    # onto average vector cos(angle) between average vector
                    # and all angles)
                    coscos = np.cos(thetadir_centers)*np.cos(average_theta)
                    sinsin = np.sin(thetadir_centers)*np.sin(average_theta)
                    cosphi = np.cos(deltaphidir_centers - average_phi)
                    # Other half of sphere
                    cospsi = (coscos + np.outer(sinsin, cosphi).T).T
                    cospsi_avg = (cospsi * weights).sum() / weights_tot
                    length = max(0.0, 2 * (cospsi_avg - 0.5))

                # Output tables are expected to be in (flip(t), r, costheta).
                # In addition to time being flipped, coszen is expected to be
                # ascending, and therefore its binning is also flipped as
                # compared to the theta binning in the original table.
                dest_bin = (
                    n_t_bins - 1 - t_k,
                    r_i,
                    n_theta_bins - 1 - theta_j
                )

                survival_probs[dest_bin] = norm * table[r_i, theta_j, t_k]
                average_thetas[dest_bin] = average_theta
                average_phis[dest_bin] = average_phi
                lengths[dest_bin] = length

    return survival_probs, average_thetas, average_phis, lengths