import os
import sys
import numpy as np
import warnings

from sph import sph
from smart_dict import SmartDict
from vector import Vector, KVector
from safe_open import safe_open

__author__ = 'Johan G. McQuillan'
__email__ = 'johan.mcquillan.13@ucl.ac.uk'

class Cell(object):
    """Simulation cell which holds Atom objects in a 3D mesh.

    All lengths measured in Bohr radii (a0).
    All energies measured in electron volts (eV).

    Attributes:
        name (string): Name of simulation; used for plot titles
        fermi_level (float): Fermi Level of simulation
        vector (Vector): Vector from origin to far corner of cell.
            Components are maximum lengths of cell in each dimension
        grid_spacing (float): Resolution of mesh points
        real_mesh (array(float)): Mesh of real space; [i, j, k] returns np.array([x, y, z])
        atoms ({int : Atom}): Atom objects of simulation indexed by atom number
        bands ({Vector : [float]}): List of band energies indexed by k-point vector
        support_mesh (array(SmartDict)): Mesh of support function values.
            Indexed by [i, j, k][atom_key][l][zeta][m]
    """

    BOLTZMANN = 8.6173303E-5        # Boltzmann's Constant in eV/K
    ELECTRON_MASS = 9.10938E-31     # Electron Mass in kg
    H_BAR = 4.135667662E-15         # Reduced Planck's Constant in eV.s
    DELTA_S_FACTOR = 2              # Default delta S given by minimum delta S scaled by this factor

    MESH_FOLDER = 'meshes'          # Folder to save mesh files
    SUPPORT_FNAME = 'supp'          # Prefix for support mesh files
    LDOS_FNAME = 'ldos'             # Prefix for summed LDOS files
    PSI_FNAME = 'psi'               # Prefix for wavefunction files
    PROP_PSI_FNAME = 'prop'         # Prefix for propagated wavefunction files
    CURRENT_FNAME = 'current'       # Prefix for current files
    EXT = 'dat'                     # Mesh file extension

    PRINT_RELATIVE_TO_EF = True     # Print energies as absolute or relative to Fermi level
    PROG_BAR_INTERVALS = 20         # Number of intervals in debug progress bar
    PROG_BAR_CHARACTER = '>'

    def __init__(self, name, fermi_level, x_length, y_length, z_length,
                 grid_spacing=0.5, group_size=400):
        """Constructs 3D cell with given dimensional.

        All lengths measured in Bohr radii (a0);
        All energies measured in Hartrees (Ha).

        Args:
            name (string): Name of simulation; used for plot titles.
            fermi_level (float): Fermi Level of simulation.
            x_length (float): Length of cell along x.
            y_length (float): Length of cell along y.
            z_length (float): Length of cell along z.
            grid_spacing (float, opt): Resolution of mesh points.
            group_size (int, opt): Maximum number of atoms to be saved to same support file.
        """

        self.name = name
        self.fermi_level = fermi_level
        self.grid_spacing = grid_spacing
        self.atom_group_size = group_size

        vector_x = int(x_length / grid_spacing) * grid_spacing
        vector_y = int(y_length / grid_spacing) * grid_spacing
        vector_z = int(z_length / grid_spacing) * grid_spacing
        self.vector = Vector(vector_x, vector_y, vector_z)

        # Initialise Cartesian meshes
        self.real_mesh = np.transpose(np.mgrid[0:x_length:grid_spacing,
                                               0: y_length: grid_spacing, 0:
                                               z_length: grid_spacing],
                                      (1, 2, 3, 0))

        # Total number of mesh points
        self.mesh_points = self.real_mesh.size

        # Initialise atoms and bands
        self.atoms = {}
        self.bands = {}

        # Currently support function mesh
        self.support_mesh = None
        # Support meshes are split into groups, each with self.group_size atoms
        # For number of atoms > 400, storing mesh for all atoms takes huge amounts of RAM
        # First group is 0, second is 1 etc.
        # -1 indicates no group stored
        self.current_group = -1  # Current atom group stored

        self.default_delta_s = self.delta_s()  # Default delta_s

        # Vectorised method to calculate wavefunction
        self.psi_vec = np.vectorize(self.calculate_psi_grid_vec)

    def energy_list(self):
        """Return sorted list of energies from all k-points."""

        energies = []
        for K in self.bands:
            energies.extend(self.bands[K])
        return sorted(energies)

    def has_band(self, K, E):
        """Check if cell stores specified band.

        Args:
            K (Vector): 3D Cartesian k-space vector.
            E (float): Band energy.
        """

        return K in self.bands and E in self.bands[K]

    def add_atom(self, atom, atom_key):
        """Add atom to self.atoms, indexed by atom_key.

        Args:
            atom (Atom): Atom object.
            atom_key (int): Atom number, as given in Conquest_out.
        """

        # Add atom to dict
        self.atoms[atom_key] = atom

        # Add band energies and k-points to self.bands
        for K in atom.bands:
            # If cell does not have k-point, create empty band energy list
            if K not in self.bands:
                self.bands[K] = []

            # Add band energies to k-point
            for E in atom.bands[K]:
                if E not in self.bands[K]:
                    self.bands[K].append(E)

            # Sort energy list
            self.bands[K] = sorted(self.bands[K])

    def fermi_dirac(self, energy, temperature):
        """Calculate Fermi-Dirac occupation factor.

        Args:
            energy (float): Energy in eV.
            temperature (float): Absolute temperature in K.

        Returns:
            float: Occupation factor, between 0 and 1.
        """

        with warnings.catch_warnings():
            # Suppress RuntimeWarning from overflow and underflow in np.exp
            warnings.simplefilter('ignore', RuntimeWarning)
            f = 1.0 / (np.exp((energy - self.fermi_level) / (self.BOLTZMANN * temperature)) + 1)
        return f

    def bias_to_energy_range(self, V):
        """Return absolute energy range from the Fermi level to the bias voltage."""

        if V > 0:
            min_E = self.fermi_level
            max_E = self.fermi_level + V
        else:
            min_E = self.fermi_level + V
            max_E = self.fermi_level
        return min_E, max_E

    def get_nearest_mesh_value(self, x, points=None):
        """Return nearest mesh point to x. Not constrained to lie within simulation cell.

        Args:
            x (float): Value in a0 to find nearest point; Works for x, y, and z dimensions.
            points (float, opt): Number of points in cell dimension.
                If given, will return mesh index of the found point

        Returns:
            float: Nearest mesh point value
            int: Index of mesh point; only given if points argument is specified
        """

        # Get quotient and remainder wrt grid spacing
        div_x, mod_x = divmod(x, self.grid_spacing)

        # Check if x should be rounded up or down
        if mod_x >= self.grid_spacing / 2:
            div_x += 1

        # Get new point
        new_x = div_x * self.grid_spacing
        if points is not None:
            i = div_x - 1
            while i < 0:
                i += points
            while i >= points:
                i -= points
            return new_x, int(i)
        else:
            return new_x

    def constrain_relative_vector(self, vector):
        """Using periodic boundaries, return smallest Vector that is equivalent to input Vector."""

        x, y, z = vector.components

        # Check if vector components are greater than half of cell sides
        # If greater, add or subtract cell length

        while x > self.vector.x / 2:
            x -= self.vector.x
        while x <= -self.vector.x / 2:
            x += self.vector.x

        while y > self.vector.y / 2:
            y -= self.vector.y
        while y <= -self.vector.y / 2:
            y += self.vector.y

        while z > self.vector.z / 2:
            z -= self.vector.z
        while z <= -self.vector.z / 2:
            z += self.vector.z

        return Vector(x, y, z)

    def calculate_support_group(self, group, interpolation='cubic', debug=False):
        """Evaluate support function for each PAO for a given atom group.

        Args:
            group (int): Atom group
            interpolation (string, opt): Method of interpolation.
                Possible arguments are 'linear', 'quadratic', 'cubic'.
            debug (bool, opt): If true, print extra information during runtime.

        Returns:
            array(SmartDict): Mesh of support function values.
                Indexed by [x, y, z][atom_key][l][zeta][m]
        """

        if debug:
            sys.stdout.write('Calculating support mesh for atom group {}\n'.format(group))
            sys.stdout.flush()

        # Initialise support grid
        support_grid = np.empty(self.real_mesh.shape[:3], dtype=SmartDict)

        # Get first and last atom numbers for this group
        lower_bound = group * self.atom_group_size
        upper_bound = (group + 1) * self.atom_group_size

        # Get atom keys for group
        atom_list = sorted([a for a in self.atoms.keys() if lower_bound <= a < upper_bound])

        # Iterate over atoms
        for atom_key in atom_list:
            atom = self.atoms[atom_key]

            debug_str = '  Support grid for atom {:3} - {!s}: '.format(atom_key, atom.ion_name)
            if debug:
                sys.stdout.write(debug_str)
                sys.stdout.flush()

            # Get atom cutoff radius
            cut = atom.get_max_cutoff()

            # Get mesh points of maximum range of atoms orbitals in each direction
            x_lower_lim, i_start = self.get_nearest_mesh_value(atom.atom_pos.x - cut,
                                                               points=self.real_mesh.shape[0])
            x_upper_lim = self.get_nearest_mesh_value(atom.atom_pos.x + cut) + self.grid_spacing
            y_lower_lim, j_start = self.get_nearest_mesh_value(atom.atom_pos.y - cut,
                                                               points=self.real_mesh.shape[1])
            y_upper_lim = self.get_nearest_mesh_value(atom.atom_pos.y + cut) + self.grid_spacing
            z_lower_lim, k_start = self.get_nearest_mesh_value(atom.atom_pos.z - cut,
                                                               points=self.real_mesh.shape[2])
            z_upper_lim = self.get_nearest_mesh_value(atom.atom_pos.z + cut) + self.grid_spacing

            # Get array of mesh points within cutoff
            local_mesh = np.transpose(np.mgrid[x_lower_lim:x_upper_lim:self.grid_spacing,
                                               y_lower_lim:y_upper_lim:self.grid_spacing,
                                               z_lower_lim:z_upper_lim:self.grid_spacing],
                                      (1, 2, 3, 0))
            lm_shape = local_mesh.shape[:3]

            # Progress bar initialisation
            points_done = 0
            bars_done = 0
            total_points = local_mesh.shape[0]*local_mesh.shape[1]*local_mesh.shape[2]

            # The local mesh may exist over the periodic boundaries
            # Roll the full support mesh such the [0, 0, 0] entry corresponds
            #   physically to the same point as the [0, 0, 0] point on the local mesh
            rolled_mesh = np.roll(support_grid, -i_start, 0)
            rolled_mesh = np.roll(rolled_mesh, -j_start, 1)
            rolled_mesh = np.roll(rolled_mesh, -k_start, 2)

            # Extract the part of the full support mesh that corresponds to
            #   the same physical space as the local mesh
            partial_mesh = rolled_mesh[0:lm_shape[0], 0:lm_shape[1], 0:lm_shape[2]]

            # Iterate over the local mesh
            for local_ijk in np.ndindex(lm_shape):
                position = local_mesh[local_ijk]
                r = Vector(*position)

                # Find the shortest Vector between r and atom_pos using periodic boundaries
                relative_position = self.constrain_relative_vector(r - atom.atom_pos)

                # Convert indices of local mesh point from tuple into a list
                partial_ijk_list = list(local_ijk)
                for i in range(len(partial_ijk_list)):
                    # If point is outside boundary, reduce to lie within cell
                    if partial_ijk_list[i] >= self.real_mesh.shape[i]:
                        partial_ijk_list[i] -= self.real_mesh.shape[i]

                # Convert back to tuple
                partial_ijk = tuple(partial_ijk_list)

                # Iterate over orbitals
                for l in atom.radials:
                    for zeta in atom.radials[l]:
                        # Get radial part of wavefunction
                        R = atom.get_radial_value_relative(l, zeta, relative_position,
                                                           interpolation=interpolation)

                        # If R == 0, do not store
                        if R != 0.0:
                            for m in range(-l, l + 1):
                                # Get spherical harmonic
                                Y = sph(l, m, relative_position)

                                # Initialise support mesh entry
                                if partial_mesh[partial_ijk] is None:
                                    partial_mesh[partial_ijk] = SmartDict()

                                if m not in partial_mesh[partial_ijk][atom_key][l][zeta]:
                                    partial_mesh[partial_ijk][atom_key][l][zeta][m] = 0.0

                                # Store support function value
                                partial_mesh[partial_ijk][atom_key][l][zeta][m] += R * Y
                points_done += 1

                # Update progress bar
                prog = float(points_done) / total_points
                if debug and prog * self.PROG_BAR_INTERVALS >= bars_done:
                    percent = prog * 100
                    sys.stdout.write('\r')
                    sys.stdout.write(debug_str)
                    sys.stdout.write(
                        ' [{:<{}}]'.format(self.PROG_BAR_CHARACTER * bars_done,
                                           self.PROG_BAR_INTERVALS))
                    sys.stdout.write(' {:3.0f}%'.format(percent))
                    sys.stdout.flush()
                    bars_done += 1

            # Copy partial support mesh into full support mesh
            rolled_mesh[0:lm_shape[0], 0:lm_shape[1], 0:lm_shape[2]] = partial_mesh

            # Roll mesh back to original position
            rolled_mesh = np.roll(rolled_mesh, i_start, 0)
            rolled_mesh = np.roll(rolled_mesh, j_start, 1)
            support_grid = np.roll(rolled_mesh, k_start, 2)

            if debug:
                sys.stdout.write('\n')
                sys.stdout.flush()

        return support_grid

    def support_group_filename(self, group):
        """Return standardised filename for relevant support function file."""

        return os.path.join(self.MESH_FOLDER, '{}_{}_{}_{}_{}.{}'.format(
                self.SUPPORT_FNAME, self.name, self.grid_spacing, self.atom_group_size, group,
                self.EXT))

    def write_support_group(self, group, support_mesh, debug=False):
        """Write support function group to file."""

        filename = self.support_group_filename(group)
        support_file = safe_open(filename, 'w')

        if debug:
            sys.stdout.write('Writing support group to {}: '.format(filename))
            sys.stdout.flush()
        points_done = 0
        bars_done = 0

        # Iterate over mesh points
        for ijk in np.ndindex(self.real_mesh.shape[:3]):
            i, j, k = ijk

            # If support function values exist at mesh point
            if support_mesh[ijk]:
                support_file.write('{} {} {}\n'.format(i, j, k))

                # Iterate over atoms
                for atom_key in support_mesh[ijk]:
                    # Write atom index
                    support_file.write('{}\n'.format(atom_key))

                    # Iterate over orbitals
                    for l in support_mesh[ijk][atom_key]:
                        for zeta in support_mesh[ijk][atom_key][l]:
                            for m in support_mesh[ijk][atom_key][l][zeta]:
                                # Write orbital data
                                support_file.write('{} {} {} {}\n'.format(
                                    l, zeta, m, support_mesh[ijk][atom_key][l][zeta][m]))

            # Update progress bar
            points_done += 1
            if debug and float(points_done) / self.mesh_points*self.PROG_BAR_INTERVALS > bars_done:
                sys.stdout.write('\r')
                sys.stdout.write(' [{:<{}}]'.format(self.PROG_BAR_CHARACTER * bars_done,
                                                    self.PROG_BAR_INTERVALS))
                sys.stdout.flush()
                bars_done += 1

        if debug:
            sys.stdout.write('\n')
            sys.stdout.flush()
        support_file.close()

    def read_support_group(self, group, debug=False):
        """Read support function group from file."""

        filename = self.support_group_filename(group)
        support_file = open(filename, 'r')
        support_mesh = np.empty(self.real_mesh.shape[:3], dtype=SmartDict)

        if debug:
            sys.stdout.write('Reading support group from {}\n'.format(filename))
            sys.stdout.flush()

        # Iterate over file lines
        end_of_file = False
        line = support_file.next()
        line_split = line.split()
        while not end_of_file:
            try:
                # Get mesh indices
                i, j, k = [int(a) for a in line_split[:3]]

                # Read atom data
                reading_atoms = True
                line = support_file.next()
                line_split = line.split()
                while reading_atoms:
                    if len(line_split) == 1:
                        # Get atom index
                        atom_key = int(line)
                        # Read orbital data
                        reading_orbitals = True
                        while reading_orbitals:
                            line = support_file.next()
                            line_split = line.split()
                            if len(line_split) != 4:
                                reading_orbitals = False
                            elif len(line_split) == 1:
                                reading_atoms = True
                            else:
                                l = int(line_split[0])
                                zeta = int(line_split[1])
                                m = int(line_split[2])
                                value = float(line_split[3])
                                if support_mesh[i, j, k] is None:
                                    support_mesh[i, j, k] = SmartDict()
                                support_mesh[i, j, k][atom_key][l][zeta][m] = value
                    else:
                        reading_atoms = False
            except StopIteration:
                end_of_file = True
        if debug:
            sys.stdout.write('Support grid successfully read\n')
            sys.stdout.flush()
        return support_mesh

    def get_support_group(self, group, recalculate=False, interpolation='cubic', debug=False):
        """Get support function mesh for given atom group.

        If it is already stored in a file and recalculate is False, then it will read from file.
        Otherwise, it will calculate support functions.

        Args:
            group (int): Atom group
            recalculate (bool, opt): Force recalculation, even if already stored.
            interpolation (string, opt): Method of interpolation.
                Possible arguments are 'linear', 'quadratic', 'cubic'.
            debug (bool, opt): Print extra information during runtime.

        Returns:
            array(SmartDict): Mesh of support function values.
                Indexed by [i, j, k][atom_key][l][zeta][m].
        """

        # Check if support function group is already stored in field
        if group == self.current_group and self.support_mesh is not None:
            return self.support_mesh
        else:
            if not recalculate and os.path.isfile(self.support_group_filename(group)):
                # Read support grid from file
                support_mesh = self.read_support_group(group, debug=debug)
            else:
                # Recalculate support grid
                support_mesh = self.calculate_support_group(group, interpolation=interpolation,
                                                            debug=debug)
                # Write to file
                self.write_support_group(group, support_mesh, debug=debug)
            self.current_group = group
            self.support_mesh = support_mesh
            return support_mesh

    def calculate_psi_grid_vec(self, support_dict, atom_key, l, zeta, m, coefficient):
        """Calculate wavefunction contribution from a point of the support function mesh
        for a given orbital and atom."""

        if (support_dict is not None and atom_key in support_dict and l in support_dict[atom_key]
                and zeta in support_dict[atom_key][l] and m in support_dict[atom_key][l][zeta]):
            psi = coefficient * support_dict[atom_key][l][zeta][m]
        else:
            psi = complex(0, 0)
        return psi

    def calculate_psi_grid(self, K, E, recalculate=False, vectorised=True, interpolation='cubic',
                           debug=False):
        """Evaluate wavefunction over the mesh for a given k-point and energy.

        Args:
            K (KVector): K-point Vector.
            E (float): Band energy.
            recalculate (bool, opt): Force recalculation, even if already stored.
            vectorised (bool, opt): If true, use NumPy vectorisation.
            interpolation (string, opt): Method of interpolation.
                Possible arguments are 'linear', 'quadratic', 'cubic'
            debug (bool, opt): Print extra information during runtime.

        Returns:
            array(complex): Mesh of complex wavefunction values.
        """

        # Initialise progress bar
        atoms_done = 0
        total_atoms = len(self.atoms)
        bars_done = 0

        # Initialise wavefunction mesh
        psi_grid = np.zeros_like(self.real_mesh[..., 0], dtype=complex)

        # Print debug info
        if self.PRINT_RELATIVE_TO_EF:
            E_str = '{} eV'.format(E - self.fermi_level)
        else:
            E_str = '{} eV'.format(E)
        debug_str = 'Calculating psi(r) at k = {}, E = {}:\n'.format(K, E_str)
        if debug:
            sys.stdout.write(debug_str)
            sys.stdout.flush()

        # Iterate over all atoms
        previous_group = 0
        support_grid = self.get_support_group(0, recalculate=recalculate, debug=debug)
        for atom_key in sorted(self.atoms.iterkeys()):
            atom = self.atoms[atom_key]
            group = atom_key / self.atom_group_size

            # Check if current atom group is the one currently read from file
            if group != previous_group:
                support_grid = self.get_support_group(group, recalculate=recalculate,
                                                      interpolation=interpolation, debug=debug)
                previous_group = group

            # Iterate over orbitals
            for l in atom.bands[K][E]:
                for zeta in atom.bands[K][E][l]:
                    for m in atom.bands[K][E][l][zeta]:
                        # Evaluate wavefunction contribution over mesh
                        coefficient = atom.get_coefficient(K, E, l, zeta, m)
                        if vectorised:
                            psi_grid += self.psi_vec(support_grid, atom_key, l, zeta, m, coefficient)
                        else:
                            for ijk in np.ndindex(self.real_mesh.shape[:3]):
                                if support_grid[ijk]:
                                    if atom_key in support_grid[ijk]:
                                        psi_grid[ijk] += coefficient*support_grid[ijk][atom_key][l][zeta][m]
            # Update progress bar
            atoms_done += 1
            prog = float(atoms_done) / total_atoms
            if debug and prog * self.PROG_BAR_INTERVALS >= bars_done:
                percent = prog * 100
                sys.stdout.write('\r')
                sys.stdout.write(' [{:<{}}]'.format(self.PROG_BAR_CHARACTER * bars_done,
                                                    self.PROG_BAR_INTERVALS))
                sys.stdout.write(' {:3.0f}%'.format(percent))
                sys.stdout.flush()
                bars_done = int(prog * self.PROG_BAR_INTERVALS)
        if debug:
            sys.stdout.write('\n')
            sys.stdout.flush()
        return psi_grid

    def psi_filename(self, K, E):
        """Return standardised filename for relevant wavefunction file."""

        return os.path.join(self.MESH_FOLDER, '{}_{}_{}_{}_{}_{}_{}.{}'.format(
                self.PSI_FNAME, self.name, self.grid_spacing, K.x, K.y, K.z, E, self.EXT))

    def write_psi_grid(self, psi_grid, K, E):
        """Write wavefunction function mesh to file."""

        filename = self.psi_filename(K, E)
        psi_file = safe_open(filename, 'w')
        for ijk in np.ndindex(self.real_mesh.shape[:3]):
            i, j, k = ijk
            if psi_grid[ijk] != 0:
                psi = psi_grid[ijk]
                psi_file.write('{} {} {} {} {}\n'.format(i, j, k, psi.real, psi.imag))
        psi_file.close()

    def read_psi_grid(self, K, E, debug=False):
        """Read wavefunction mesh from file."""

        filename = self.psi_filename(K, E)
        psi_file = open(filename, 'r')
        psi_grid = np.zeros_like(self.real_mesh[..., 0], dtype=complex)

        if debug:
            if self.PRINT_RELATIVE_TO_EF:
                E_str = '{} eV'.format(E - self.fermi_level)
            else:
                E_str = '{} eV'.format(E)
            sys.stdout.write('Reading psi(r) at k = {!s}, E = {}\n'.format(K, E_str))

        for line in psi_file:
            line_split = line.split()
            i = int(line_split[0])
            j = int(line_split[1])
            k = int(line_split[2])
            real = float(line_split[3])
            imag = float(line_split[4])
            psi_grid[i, j, k] = complex(real, imag)
        return psi_grid

    def get_psi_grid(self, K, E, recalculate=False, vectorised=True, interpolation='cubic',
                     debug=False):
        """Get mesh of complex wavefunction values.

        If it is already stored in a file and recalculate is False, then it will read from file.
        Otherwise, it will calculate wavefunction.

        Args:
            K (KVector): K-point Vector
            E (float): Band energy
            recalculate (bool, opt): Force recalculation, even if already stored
            vectorised (bool, opt): If true, use NumPy vectorisation
            interpolation (string, opt): Method of interpolation.
                Possible arguments are 'linear', 'quadratic', 'cubic'
            debug (bool, opt): Print extra information during runtime

        Returns:
            array(complex): Mesh of complex wavefunction values
        """

        if not recalculate and os.path.isfile(self.psi_filename(K, E)):
            # Read data from file
            psi_grid = self.read_psi_grid(K, E, debug=debug)
        else:
            psi_grid = self.calculate_psi_grid(K, E, recalculate=recalculate, vectorised=vectorised,
                                               interpolation=interpolation, debug=debug)
            self.write_psi_grid(psi_grid, K, E)
        return psi_grid

    def calculate_ldos_grid(self, min_E, max_E, T, recalculate=False, vectorised=True,
                            interpolation='cubic', debug=False):
        """Calculate summed LDOS.

        Args:
            min_E (float): Minimum absolute energy
            max_E (float): Maximum absolute energy
            T (float): Absolute temperature in Kelvin
            recalculate (bool, opt): Force recalculation, even if already stored
            vectorised (bool, opt): If true, use NumPy vectorisation
            interpolation (string, opt): Method of interpolation.
                Possible arguments are 'linear', 'quadratic', 'cubic'
            debug (bool, opt): Print extra information during runtime

        Returns:
            array(float): LDoS mesh
        """

        # Print debug info
        if debug:
            sys.stdout.write('Calculating local density of states grid\n')
            sys.stdout.flush()

        # Initialise mesh
        ldos_grid = np.zeros_like(self.real_mesh[..., 0], dtype=float)

        total_k_weight = sum([K.weight for K in self.bands])

        # Iterate over energies
        for K in self.bands:
            for E in self.bands[K]:
                if min_E <= E <= max_E:
                    psi_grid = self.get_psi_grid(K, E, recalculate=recalculate,
                                                 vectorised=vectorised, interpolation=interpolation,
                                                 debug=debug)
                    fd = self.fermi_dirac(E, T)
                    if E > self.fermi_level:
                        fd = 1 - fd
                    if vectorised:
                        ldos_grid += (K.weight / total_k_weight) * fd * (abs(psi_grid))**2
                    else:
                        for ijk in np.ndindex(self.real_mesh.shape[:3]):
                            ldos_grid[ijk] += (K.weight/total_k_weight)*fd*(abs(psi_grid[ijk]))**2
        return ldos_grid

    def ldos_filename(self, min_E, max_E, T):
        """Return standardised filename for relevant LDOS file"""

        return os.path.join(self.MESH_FOLDER, '{}_{}_{}_{}_{}_{}.{}'.format(
                self.LDOS_FNAME, self.name, self.grid_spacing, min_E, max_E, T, self.EXT))

    def write_ldos_grid(self, ldos_grid, min_E, max_E, T, debug=False):
        """Write LDoS mesh to file.

        Args:
            ldos_grid (array(float)): 3D array of LDOS values.
            min_E (float): Minimum energy.
            max_E (float): Maximum energy.
            T (float): Absolute temperature in K.
            debug (bool, opt): Print extra information during runtime.
        """

        filename = self.ldos_filename(min_E, max_E, T)

        # Get LDOS mesh
        ldos_file = safe_open(filename, 'w')
        if debug:
            sys.stdout.write('Writing LDOS grid to {}\n'.format(filename))
            sys.stdout.flush()

        # Iterate over mesh points
        for i, j, k in np.ndindex(self.real_mesh.shape[:3]):
            # If LDOS is non-zero at mesh point, write data to file
            if ldos_grid[i, j, k]:
                ldos_file.write('{} {} {} {}\n'.format(i, j, k, ldos_grid[i, j, k]))

        ldos_file.close()

    def read_ldos_grid(self, min_E, max_E, T, debug=False):
        """Read LDOS mesh from file."""

        filename = self.ldos_filename(min_E, max_E, T)
        ldos_file = open(filename, 'r')
        ldos_grid = np.zeros_like(self.real_mesh[..., 0])

        if debug:
            sys.stdout.write('Reading LDoS grid from {}\n'.format(filename))
            sys.stdout.flush()

        for line in ldos_file:
            line_split = line.split()
            # Get mesh indices
            i = int(line_split[0])
            j = int(line_split[1])
            k = int(line_split[2])
            # Get LDoS value
            value = float(line_split[3])
            ldos_grid[i, j, k] = value

        if debug:
            sys.stdout.write('LDoS grid successfully read\n')
            sys.stdout.flush()
        return ldos_grid

    def get_ldos_grid(self, min_E, max_E, T, recalculate=False, vectorised=True,
                      interpolation='cubic', debug=False):
        """Get LDOS mesh by calculating or reading from file.

        Args:
            min_E (float): Minimum absolute energy
            max_E (float): Maximum absolute energy
            T (float): Absolute temperature in K
            recalculate (bool, opt): Force recalculation of meshes, even if already stored
            vectorised (bool, opt): If true, use NumPy vectorisation
            interpolation (string, opt): Method of interpolation.
                Possible arguments are 'linear', 'quadratic', 'cubic'
            debug (bool, opt): Print extra information during runtime

        Returns:
            array(float): 3D mesh of LDOS values
        """

        # Read ldos grid from file if not stored by cell
        if not recalculate and os.path.isfile(self.ldos_filename(min_E, max_E, T)):
            ldos_grid = self.read_ldos_grid(min_E, max_E, T, debug=debug)
        else:
            # Calculate LDOS on mesh
            ldos_grid = self.calculate_ldos_grid(min_E, max_E, T, recalculate=recalculate,
                                                 vectorised=vectorised, interpolation=interpolation,
                                                 debug=debug)
            self.write_ldos_grid(ldos_grid, min_E, max_E, T, debug=debug)
        return ldos_grid

    def periodic_gradient(self, mesh):
        """Calculate gradient of mesh with periodic boundary conditions enforced.

        This is done by padding the mesh with extra columns in each dimension, then copying the
        columns from the other side of the mesh into the pads. This makes the columns corresponding
        to a boundary of the simulation cell effectively next to the columns near the opposing
        boundary, such that numerical calculation of the gradient should be continuous over the
        boundary.
        """

        # Width of padding
        pad = 3
        # Get shape of padded array
        padded_shape = (mesh.shape[0] + 2*pad, mesh.shape[1] + 2*pad, mesh.shape[2] + 2*pad)

        # Copy mesh into padded mesh
        padded_mesh = np.zeros(padded_shape, dtype=mesh.dtype)
        padded_mesh[pad:-pad, pad:-pad, pad:-pad] = mesh

        # Copy boundary regions into padding
        padded_mesh[:pad, pad:-pad, pad:-pad] = mesh[-pad:, :, :]
        padded_mesh[pad:-pad, :pad, pad:-pad] = mesh[:, -pad:, :]
        padded_mesh[pad:-pad, pad:-pad, :pad] = mesh[:, :, -pad:]
        padded_mesh[-pad:, pad:-pad, pad:-pad] = mesh[:pad, :, :]
        padded_mesh[pad:-pad, -pad:, pad:-pad] = mesh[:, :pad, :]
        padded_mesh[pad:-pad, pad:-pad, -pad:] = mesh[:, :, :pad]

        # Get gradient
        padded_gradient = np.transpose(np.array(np.gradient(padded_mesh, self.grid_spacing)),
                                       (1, 2, 3, 0))

        # Return unpadded gradient
        return padded_gradient[pad:-pad, pad:-pad, pad:-pad]

    def delta_s(self):
        """Calculate the default value for delta S.

        This is the minimum delta S from Paz and Soler, multiplied by DELTA_S_FACTOR."""

        min_delta_s = 2.0 * self.grid_spacing / self.H_BAR * np.sqrt(4.85*2.0*self.ELECTRON_MASS)
        return self.DELTA_S_FACTOR * min_delta_s

    def kappa_squared(self, tip_work_func, tip_fermi_level):
        """Calculate decay constant of tip wavefunction"""

        return 2*self.ELECTRON_MASS/(self.H_BAR**2)*(tip_work_func - tip_fermi_level)

    def greens_function(self, distance, tip_work_func, tip_energy):
        """Evaluate Tersoff-Hamann Green's Function"""

        if distance == 0:
            return 0
        else:
            kappa2 = self.kappa_squared(tip_work_func, tip_energy)
            return np.exp(- kappa2 * distance) / (4*np.pi*distance)

    def broadened_surface(self, charge_density_mesh, fraction, max_height_index, delta_s=None):
        """Calculate the magnitude of the c mesh"""

        if delta_s is None:
            delta_s = self.default_delta_s

        # Get isovalue
        max_value = np.max(charge_density_mesh)
        isovalue = fraction * max_value

        # Evaluate logarithm on all non-zero entrie
        log_mesh = np.empty(charge_density_mesh.shape, dtype=float)
        zero_points = charge_density_mesh == 0
        log_mesh[~zero_points] = np.log(charge_density_mesh[charge_density_mesh != 0] / isovalue)
        log_mesh[zero_points] = np.inf

        # Apply broadening to surface
        broadened_mesh = np.where(abs(log_mesh) < delta_s,
                                  15.0 / (16.0 * delta_s) * (1.0 - (log_mesh / delta_s) ** 2) ** 2, 0)

        # Remove overlapping layers of surface
        # Iterate over x and y
        for i, j in np.ndindex(broadened_mesh.shape[:2]):
            on_surface = False
            past_surface = False
            # Iterate over z, starting from top
            for k in reversed(range(broadened_mesh.shape[2])):
                if k < max_height_index:
                    if past_surface:
                        # First surface has been traversed, replace subsequent elements with zeros
                        broadened_mesh[i, j, k] = 0
                    elif broadened_mesh[i, j, k] != 0 and not on_surface:
                        # Highest surface has been reached
                        on_surface = True
                    elif on_surface and broadened_mesh[i, j, k] == 0:
                        # Was on surface, now just below
                        on_surface = False
                        past_surface = True
                else:
                    broadened_mesh[i, j, k] = 0
        return broadened_mesh

    def get_c(self, charge_density_mesh, fraction, tip_height_index, delta_s=None):
        """Return c mesh for surface integration.

        Args:
            charge_density_mesh (array(float)): Charge density or LDOS mesh.
            fraction (float): Isovalue given by fraction of maximum mesh value.
            tip_height_index (int): Z-index of tip height.
            delta_s (float, opt): Surface broadening parameter; If None, uses default value.
        """

        if delta_s is None:
            delta_s = self.default_delta_s

        # Find all points of zero density
        zero_points = charge_density_mesh == 0

        # Get magnitude of c mesh
        broadened_mesh = self.broadened_surface(charge_density_mesh, fraction, tip_height_index,
                                                delta_s=delta_s)

        # Get gradient of density mesh
        gradient_surface = self.periodic_gradient(charge_density_mesh)

        # Calculate final mesh
        vector_surface = np.zeros(gradient_surface.shape, dtype=float)
        for i in range(3):
            gradient_surface[~zero_points, i] = gradient_surface[~zero_points, i] / charge_density_mesh[~zero_points]
            gradient_surface[zero_points, i] = 0
            vector_surface[..., i] = np.multiply(broadened_mesh, gradient_surface[..., i])

        return vector_surface

    def get_A_mesh(self, c, wavefunction_mesh):
        """Calculate A mesh."""

        grad_wavefunction = self.periodic_gradient(wavefunction_mesh)
        return self.mesh_dot_product(c, grad_wavefunction)

    @staticmethod
    def mesh_dot_product(vector_mesh_A, vector_mesh_B):
        """Return dot/scalar product of two vector meshes."""

        # Check if resulting mesh will be complex
        if vector_mesh_A.dtype == complex or vector_mesh_B.dtype == complex:
            dtype = complex
        else:
            dtype = float

        scalar_mesh = np.zeros(vector_mesh_A.shape[:3], dtype=dtype)
        for i in range(3):
            scalar_mesh += vector_mesh_A[..., i]*vector_mesh_B[..., i]
        return scalar_mesh

    def get_B_mesh(self, c, wavefunction_mesh):
        """Calculate B mesh."""

        B = np.zeros_like(c, dtype=complex)
        for i in range(3):
            B[..., i] = c[..., i] * wavefunction_mesh
        return B

    def greens_function_mesh(self, z_index, tip_work_func, tip_energy, debug=False):
        """Calculate Tersoff-Hamann tip Greens function over entire mesh.

        Uses symmetry of the function to reduce number of calculations.

        Args:
            z_index (float): z-index of tip plane.
            tip_work_func (float): Work function of tip.
            tip_energy (float): Fermi-level of tip.
            debug (bool, opt): Print extra information during runtime.
        """

        # Create square grid with sides equal to the size of the x or y=axis of the real mesh
        #   whichever is larger
        plane_length_half = - (- (max(self.real_mesh.shape[:2]) + 1) / 2)
        plane_length_full = plane_length_half * 2 - 1
        plane_shape_half = (plane_length_half,) * 2
        plane_UR = np.zeros(plane_shape_half, dtype=float)

        # Initialise Green's function mesh
        G_shape = (plane_length_full, plane_length_full, self.real_mesh.shape[2])
        G_mesh = np.zeros(G_shape, dtype=float)

        # Print debug info
        debug_str = 'Calculating G(r - R): '
        if debug:
            sys.stdout.write(debug_str)
            sys.stdout.flush()
        points_done = 0
        bars_done = 0

        # Iterate over z values
        for k in range(G_shape[2]):
            # Iterate over right-angle triangle in x and y
            # The tip is considered to be in corner of the mesh at i = j = 0
            for i in range(plane_length_half):
                for j in range(i + 1):
                    # If above tip, define as 0
                    if k >= z_index:
                        G = 0
                    else:
                        # Calculate Green's function
                        distance = self.grid_spacing * np.sqrt(i**2 + j**2 + (z_index - k)**2)
                        G = self.greens_function(distance, tip_work_func, tip_energy)

                    # Copy values from this side of triangle into other
                    # In effect, reflect values along x = y
                    for x, y in [i, j], [j, i]:
                        plane_UR[x, y] = G
            # Flip square to create four subsquares
            plane_UL = np.flipud(plane_UR[1:, :])
            plane_LR = np.fliplr(plane_UR[:, 1:])
            plane_LL = np.flipud(plane_LR[1:, :])

            # Combine subsquares
            # Tip is now defined above the centre of mesh
            plane_U = np.concatenate((plane_UL, plane_UR), axis=0)
            plane_L = np.concatenate((plane_LL, plane_LR), axis=0)
            plane = np.concatenate((plane_L, plane_U), axis=1)

            G_mesh[..., k] = plane

            # Update progress bar
            points_done += 1
            prog = float(points_done) / self.real_mesh.shape[2]
            if debug and prog * self.PROG_BAR_INTERVALS >= bars_done:
                percent = prog * 100
                sys.stdout.write('\r')
                sys.stdout.write(debug_str)
                sys.stdout.write(' [{:<{}}]'.format(self.PROG_BAR_CHARACTER * bars_done,
                                                    self.PROG_BAR_INTERVALS))
                sys.stdout.write(' {:3.0f}%'.format(percent))
                sys.stdout.flush()
                bars_done += 1

        if debug:
            sys.stdout.write('\n')
            sys.stdout.flush()
        return G_mesh

    def propagated_psi_filename(self, K, E, T, fraction, z, delta_s=None):
        """Return standardised filename for relevant propagated wavefunction file."""

        if delta_s is None:
            delta_s = self.default_delta_s

        return os.path.join(self.MESH_FOLDER, '{}_{}_{}_{}_{}_{}_{}_{}_{}_{}_{}.{}'.format(
                self.PROP_PSI_FNAME, self.name, self.grid_spacing, K.x, K.y, K.z, E, T, fraction, z,
                delta_s, self.EXT))

    def write_prop_psi(self, psi, K, E, T, fraction, z, delta_s=None):
        """Write propagated wavefunction function mesh to file."""

        if delta_s is None:
            delta_s = self.default_delta_s

        filename = self.propagated_psi_filename(K, E, T, fraction, z, delta_s=delta_s)
        psi_file = safe_open(filename, 'w')
        for i, j in np.ndindex(psi.shape):
            if psi[i, j] != 0:
                p = psi[i, j]
                psi_file.write('{} {} {} {}\n'.format(i, j, p.real, p.imag))

        psi_file.close()

    def read_prop_psi(self, K, E, T, fraction, z, delta_s=None, debug=False):
        """Read propagated wavefunction mesh from file."""

        if delta_s is None:
            delta_s = self.default_delta_s

        filename = self.propagated_psi_filename(K, E, T, fraction, z, delta_s=delta_s)
        psi_file = open(filename, 'r')
        psi = np.zeros(self.real_mesh.shape[:2], dtype=complex)

        if debug:
            if self.PRINT_RELATIVE_TO_EF:
                E_str = '{} eV'.format(E - self.fermi_level)
            else:
                E_str = '{} eV'.format(E)
            sys.stdout.write('Reading psi(R) at k = {!s}, E = {}\n'.format(K, E_str))
            sys.stdout.flush()

        for line in psi_file:
            line_split = line.split()
            i = int(line_split[0])
            j = int(line_split[1])
            real = float(line_split[2])
            imag = float(line_split[3])
            psi[i, j] = complex(real, imag)

        return psi

    def calculate_current_scan(self, z, V, T, tip_work_func, tip_energy, delta_s=None,
                               fraction=0.025, recalculate=False, vectorised=True,
                               interpolation='cubic', debug=False):
        """Calculate tunnelling current across a plane ie. in constant-height mode.

        Args:
            z (float): z-value of plane in Bohr radii; Uses nearest mesh point to given value
            V (float): Bias voltage
            T (float): Absolute temperature
            tip_work_func (float): Work function of tip
            tip_energy (float): Fermi-level of tip
            delta_s (float, opt): Surface broadening parameter; If None, uses default value
            fraction (float, opt): Fraction of max charge density to use as value for isosurface
            recalculate (bool, opt): Force recalculation, even if already stored.
            vectorised (bool, opt): If true, use NumPy vectorisation.
            interpolation (str, opt): Method of interpolation.
                Possible arguments are 'linear', 'quadratic', 'cubic'.
            debug (bool, opt): Print extra information during runtime.

        Returns:
            array(float): 2D array of current values.
        """

        if delta_s is None:
            delta_s = self.default_delta_s

        if debug:
            sys.stdout.write('Calculating I(R) at V = {}V\n'.format(V))
            sys.stdout.flush()

        min_E, max_E = self.bias_to_energy_range(V)

        total_k_weight = 0
        total_energies = 0
        for K in self.bands:
            total_k_weight += K.weight
            for E in self.bands[K]:
                if min_E <= E <= max_E:
                    total_energies += 1
        energies_done = 0

        # Find nearest mesh z-value to given
        z, k = self.get_nearest_mesh_value(z, points=self.real_mesh.shape[2])

        # Initialise meshes
        current = np.zeros(self.real_mesh.shape[:2], dtype=float)
        psi = np.zeros_like(current, dtype=complex)
        elements = current.shape[0]*current.shape[1]

        # Calculate meshes
        ldos = self.get_ldos_grid(min_E, max_E, T, recalculate=recalculate, vectorised=vectorised,
                                  interpolation=interpolation, debug=debug)
        c = self.get_c(ldos, fraction, k, delta_s=delta_s)
        G_conjugate = np.conjugate(self.greens_function_mesh(k, tip_work_func,
                                                             tip_energy, debug=debug))

        # Print debug info
        if debug:
            sys.stdout.write('Calculating grad(G)\n')
            sys.stdout.flush()

        G_conjugate_gradient = self.periodic_gradient(G_conjugate)
        G_centre = G_conjugate.shape[0] / 2

        # Iterate over energies
        for K in self.bands:
            w = K.weight
            for E in self.bands[K]:
                if min_E <= E <= max_E:
                    fd = self.fermi_dirac(E - V, T)
                    if V < 0:
                        fd = 1 - fd

                    # Get propagated wavefunction
                    if not recalculate and os.path.isfile(
                            self.propagated_psi_filename(K, E, T, fraction, z, delta_s=delta_s)):
                        # Read data from file
                        psi = self.read_prop_psi(K, E, T, fraction, z, delta_s=delta_s, debug=debug)
                        read = True
                    else:
                        read = False
                        points_done = 0
                        bars_done = 0

                        # Get unpropagated wavefunction
                        raw_psi = self.get_psi_grid(K, E, recalculate=recalculate,
                                                    vectorised=vectorised,
                                                    interpolation=interpolation, debug=debug)

                        # Print debug info
                        prog = float(energies_done) / total_energies * 100
                        if self.PRINT_RELATIVE_TO_EF:
                            E_str = '{} eV'.format(E - self.fermi_level)
                        else:
                            E_str = '{} eV'.format(E)
                        debug_str = 'Calculating psi(R) at k = {!s}, E = {}:\n  {:5.1f}%'.format(
                                K, E_str, prog)
                        if debug:
                            sys.stdout.write(debug_str)
                            sys.stdout.flush()

                        # Get meshes
                        A = self.get_A_mesh(c, raw_psi)
                        B = self.get_B_mesh(c, raw_psi)

                        # Iterate over tip positions
                        for i, j in np.ndindex(current.shape):
                            # By rolling the tip Green's function, we move the position of the tip
                            G_conjugate_rolled = np.roll(G_conjugate, (i - G_centre), 0)
                            G_conjugate_rolled = np.roll(G_conjugate_rolled, (j - G_centre), 1)[
                                                 :self.real_mesh.shape[0], :self.real_mesh.shape[1]]
                            G_conjugate_gradient_rolled = np.roll(G_conjugate_gradient, (i - G_centre), 0)
                            G_conjugate_gradient_rolled = np.roll(G_conjugate_gradient_rolled, (j - G_centre), 1)[
                                                          :self.real_mesh.shape[0], :self.real_mesh.shape[1]]

                            # Perform volume integral
                            integrand = G_conjugate_rolled * A - self.mesh_dot_product(B, G_conjugate_gradient_rolled)
                            psi[i, j] = np.sum(integrand) * self.grid_spacing ** 3

                            # Update progress bar
                            points_done += 1
                            prog = float(points_done) / elements
                            if debug and prog * self.PROG_BAR_INTERVALS >= bars_done:
                                percent = prog * 100
                                sys.stdout.write('\r')
                                sys.stdout.write(' [{:<{}}] {:3.0f}%'.format(
                                        self.PROG_BAR_CHARACTER * bars_done,
                                        self.PROG_BAR_INTERVALS, percent))
                                sys.stdout.flush()
                                bars_done += 1

                        # Save propagated wavefunction
                        self.write_prop_psi(psi, K, E, T, fraction, z, delta_s=delta_s)

                    # Add to current total
                    current += fd * (w / total_k_weight) * abs(psi)**2

                    energies_done += 1
                    if debug and not read:
                        sys.stdout.write('\n')
                        sys.stdout.flush()
        return current

    def current_filename(self, z, V, T, fraction, delta_s=None):
        """Return standardised filename for relevant current file."""

        if delta_s is None:
            delta_s = self.default_delta_s
        return os.path.join(self.MESH_FOLDER, '{}_{}_{}_{}_{}_{}_{}_{}.{}'.format(
                self.CURRENT_FNAME, self.name, self.grid_spacing, z, V, T, fraction, delta_s,
                self.EXT))

    def write_current(self, current, z, V, T, fraction, delta_s=None, debug=False):
        """Write current to file."""

        filename = self.current_filename(z, V, T, fraction, delta_s)

        current_file = safe_open(filename, 'w')
        if debug:
            sys.stdout.write('Writing current grid to {}\n'.format(filename))

        # Iterate over mesh points
        for i, j in np.ndindex(current.shape):
            # If LDOS is non-zero at mesh point, write data to file
            if current[i, j] != 0:
                current_file.write('{} {} {}\n'.format(i, j, current[i, j]))

        current_file.close()

    def read_current(self, z, V, T, fraction, delta_s=None, debug=False):
        """Read current grid from file."""

        filename = self.current_filename(z, V, T, fraction, delta_s=delta_s)
        current_file = open(filename, 'r')
        current = np.zeros(self.real_mesh.shape[:2], dtype=float)

        if debug:
            sys.stdout.write('Reading I(R) from {}\n'.format(filename))
            sys.stdout.flush()

        for line in current_file:
            line_split = line.split()
            # Get mesh indices
            i = int(line_split[0])
            j = int(line_split[1])

            # Get current value
            value = float(line_split[2])
            current[i, j] = value

        if debug:
            sys.stdout.write('I(R) successfully read\n')
            sys.stdout.flush()

        return current

    def get_current_scan(self, z, V, T, tip_work_func, tip_energy, delta_s=None, fraction=0.025,
                         recalculate=False, vectorised=True, interpolation='cubic', debug=False):
        """Get constant-height tunnelling current as 2d array.

        If it is already stored in a file and recalculate is False, then it will read from file.
        Otherwise, it will calculate current.

        Args:
            z (float): z-value of plane in Bohr radii; Uses nearest mesh point to given value.
            V (float): Bias voltage.
            T (float): Absolute temperature.
            tip_work_func (float): Work function of tip.
            tip_energy (float): Fermi-level of tip.
            delta_s (float, opt): Surface broadening parameter; If None, uses default value.
            fraction (float, opt): Fraction of max charge density to use as value for isosurface.
            recalculate (bool, opt): Force recalculation, even if already stored.
            vectorised (bool, opt): If true, use NumPy vectorisation.
            debug (bool, opt): Print extra information during runtime.

        Returns:
            array(float): 2D array of current values.
        """

        if delta_s is None:
            delta_s = self.default_delta_s
        if not recalculate and os.path.isfile(self.current_filename(z, V, T, fraction)):
            # Read data from file
            current = self.read_current(z, V, T, fraction, debug=debug)
        else:
            current = self.calculate_current_scan(z, V, T, tip_work_func, tip_energy,
                                                  delta_s=delta_s, fraction=fraction,
                                                  recalculate=recalculate,
                                                  interpolation=interpolation,
                                                  vectorised=vectorised, debug=debug)
            self.write_current(current, z, V, T, fraction)

        return current

    def get_spectrum_th(self, xy, min_V, max_V, sigma, T, fraction, z, delta_s=None, dE=0.005, debug=False):
        """Get spectroscopic data from a list of specific tip positions.

        Args:
            xy (list(list(float)): x-y points of tip in a0; Given as [[x1, y1], [x2, y2], ...].
                Uses nearest mesh point.
            min_V (float): Lower bound for voltage range.
            max_V (float): Upper bound for voltage range.
            sigma (float): State smearing parameter in eV.
            T (float): Absolute temperature in K.
            z (float): z-value of plane in a0; Uses nearest mesh point to given value.
            fraction (float): Fraction of maximum charge density to use as value for isosurface.
            delta_s (float, opt): Surface broadening parameter; If None, uses default value.
            dE (float, opt): Energy resolution of data points.
            debug (bool, opt): Print extra information during runtime.

        Returns:
            array(float): Range of voltage values with resolution dE.
            array(float): Tunnelling conductance values for each tip position.
        """

        # Convert voltage range into absolute energies
        min_E = min_V + self.fermi_level
        max_E = max_V + self.fermi_level

        # Get positions of tip
        mesh_positions = []
        mesh_indices = []
        for l in range(len(xy)):
            x, i = self.get_nearest_mesh_value(xy[l][0], points=self.real_mesh.shape[0])
            y, j = self.get_nearest_mesh_value(xy[l][1], points=self.real_mesh.shape[1])
            mesh_positions.append((x, y))
            mesh_indices.append((i, j))

        Es = []
        psis = []
        weights = []
        l = 0
        # Iterate over energies in voltage range
        for K in self.bands:
            for E in self.bands[K]:
                # Assume states 3*sigma away from the edges contribute negligibly
                if min_E - 3*sigma < E < max_E + 3*sigma:
                    Es.append(E)
                    weights.append(K.weight)
                    psi = self.read_prop_psi(K, E, T, fraction, z, delta_s=delta_s, debug=debug)
                    psis.append([])

                    # Get wavefunction for each tip position
                    for m in range(len(mesh_positions)):
                        psis[l].append(abs(psi[mesh_indices[m]])**2)
                    l += 1

        # Convert lists to arrays
        Es = np.array(Es)
        psis = np.array(psis)
        weights = np.array(weights)

        # Generate energy points for plot
        E_range = np.arange(min_E, max_E, dE)
        # Initialise LDOS
        LDOS = np.zeros((E_range.shape[0], len(xy)))

        for u in range(len(E_range)):
            LDOS[u] = np.sum(weights[..., None] * psis *
                             np.exp(-(((E_range[u] - Es[..., None]) / sigma)**2) / 2), axis=0)
        V_range = E_range - self.fermi_level

        return V_range, LDOS

    def get_cits(self, V, T, fraction, sigma, delta_s=None, debug=False):
        """Calculate Current Imaging Tunnelling Spectroscopy scan - UNFINISHED and UNRELIABLE."""

        if delta_s is None:
            delta_s = self.default_delta_s

        min_E, max_E = self.bias_to_energy_range(V)

        ldos = self.get_ldos_grid(min_E, max_E, T, debug=debug)
        b = self.broadened_surface(ldos, fraction, self.real_mesh[..., -1, 2], delta_s=delta_s)

        scan = np.zeros(b.shape[:2])

        for K in self.bands:
            w = K.weight
            for E in self.bands[K]:
                if -3*sigma < E - V - self.fermi_level < 3*sigma:
                    exp = np.exp(-(((E - V - self.fermi_level) / sigma)**2) / 2)
                    psi = self.get_psi_grid(K, E, debug=debug)
                    b_psi = b * abs(psi)**2

                    for i, j in np.ndindex(scan.shape):
                        scan[i, j] += w * exp * np.sum(b_psi[i, j])
        scan = scan.reshape(b.shape[:2])

        return scan


import torch
import torch.nn as nn

class FCOSLoss(nn.Module):
    def __init__(self, limit_range, central_sampling=True, central_sampling_radius=1.5, \
                        strides=[8, 16, 32, 64, 128], regression_loss_type='giou', num_classes=81):
        super(FCOSLoss, self).__init__()
        self.limit_range = limit_range
        self.strides = strides
        self.num_classes = num_classes
        self.central_sampling = central_sampling
        self.regression_loss_type = regression_loss_type
        self.central_sampling_radius = central_sampling_radius

    def forward(self, prediction, batch_bboxes):
        # prediction   : list of 3 tensors cls_probs, cnt_logits, reg_values
        # batch_bboxes : [B x M x 5] --> [x1, y1, x2, y2, cls_id]

        cls_probs_pred, cnt_logits_pred, reg_values_pred = prediction
        # cls_probs_pred  : [[B x 81 x H x W], [B x 81 x H x W], ...]
        # cnt_logits_pred : [[B x 1 x H x W], [B x 1 x H x W], ...]
        # reg_values_pred : [[B x 4 x H x W], [B x 4 x H x W], ...]

        comb_cls_probs_pred = []
        comb_cnt_logits_pred = []
        comb_reg_values_pred = []

        comb_cls_probs_target = []
        comb_cnt_logits_target = []
        comb_reg_values_target = []

        for i, (cls_p, cnt_p, reg_p) in enumerate(zip(cls_probs_pred, cnt_logits_pred, reg_values_pred)):
            num_classes, feat_h, feat_w = cls_p.shape[1:4]
            cls_target, cnt_target, reg_target = \
                self.generate_targets(feat_h, feat_w, batch_bboxes, \
                                        self.strides[i], self.limit_range[i])
            # cls_target : [B x H*W x 81]
            # cnt_target : [B x H*W x 1]
            # reg_target : [B x H*W x 4]

            # TODO : Cant we just calculate loss for each scale inside this loop rather than
            #        Combining these tensors and computing together.
            cls_p = torch.reshape(cls_p, (-1, num_classes, feat_h*feat_w)).permute([0, 2, 1])      # [B x H*W x 81]
            cnt_p = torch.reshape(cnt_p, (-1, 1, feat_h*feat_w)).permute([0, 2, 1])                # [B x H*W x 1]
            reg_p = torch.reshape(reg_p, (-1, 4, feat_h*feat_w)).permute([0, 2, 1])                # [B x H*W x 4]

            comb_cls_probs_pred.append(cls_p)
            comb_cnt_logits_pred.append(cnt_p)
            comb_reg_values_pred.append(reg_p)

            comb_cls_probs_target.append(cls_target)
            comb_cnt_logits_target.append(cnt_target)
            comb_reg_values_target.append(reg_target)

        comb_cls_probs_target = torch.cat(comb_cls_probs_target, dim=1)      # [B x sum of all H*W x 81]
        comb_cnt_logits_target = torch.cat(comb_cnt_logits_target, dim=1)    # [B x sum of all H*W x 1]
        comb_reg_values_target = torch.cat(comb_reg_values_target, dim=1)    # [B x sum of all H*W x 4]

        comb_cls_probs_pred = torch.cat(comb_cls_probs_pred, dim=1)    # [B x sum of all H*W x 81]
        comb_cnt_logits_pred = torch.cat(comb_cnt_logits_pred, dim=1)  # [B x sum of all H*W x 1]
        comb_reg_values_pred = torch.cat(comb_reg_values_pred, dim=1)  # [B x sum of all H*W x 4]

        mask_pos = comb_cnt_logits_target > -1                         # [B x sum of all H*W x 1]
        cls_loss = self.compute_cls_loss(comb_cls_probs_pred, comb_cls_probs_target, mask_pos)
        cnt_loss = self.compute_centerness_loss(comb_cnt_logits_pred, comb_cnt_logits_target, mask_pos)
        reg_loss = self.compute_coordinate_reg_loss(comb_reg_values_pred, comb_reg_values_target, mask_pos)
        return cls_loss, cnt_loss, reg_loss

    def generate_targets(self, feat_h, feat_w, batch_bbox, \
                                stride, limit_range):
        # feat_h     : Height of the feature map
        # feat_w     : Width of the feature map
        # batch_bbox : [B x M x 5]  --> [x1, y1, x2, y2, cls_id]
        #            : Not all images in B are having M bboxes but the are padded to make them M
        #            : So that they can be stacked in a tensor.
        # stride : value of stride for this feature-map
        # limit_range : [min_val, max_val]. To filter most relevant bboxes out of these feature-map

        target_device = batch_bbox.device
        grid_y = torch.arange(0, feat_h * stride, stride, dtype=torch.float32)
        grid_x = torch.arange(0, feat_w * stride, stride, dtype=torch.float32)

        grid_y, grid_x = torch.meshgrid(grid_y, grid_x)
        coords = torch.stack([torch.reshape(grid_x, [-1]), \
                                torch.reshape(grid_y, [-1])], -1) + stride // 2
        coords = torch.unsqueeze(coords, 0).to(target_device)
        # coords : [1 x H*W x 2] : center-points of each grid-cell : [x, y] at last dim.
        #                        : This coords are now in original image space i.e. 800x1024
        coords_x, coords_y = coords[:, :, 0], coords[:, :, 1]   # [1 x H*W]

        batch_bb_x1, batch_bb_y1, batch_bb_x2, batch_bb_y2 = \
            batch_bbox[:, :, 0], batch_bbox[:, :, 1], batch_bbox[:, :, 2], batch_bbox[:, :, 3]
        # [B x M]

        # coords_x   : [1 x H*W] --> [1 x H*W x 1]
        # batch_bb_x1: [B x M] --> [B x 1 x M]
        l_off = torch.unsqueeze(coords_x, 2) - torch.unsqueeze(batch_bb_x1, 1)
        r_off = torch.unsqueeze(batch_bb_x2, 1) - torch.unsqueeze(coords_x, 2)
        t_off = torch.unsqueeze(coords_y, 2) - torch.unsqueeze(batch_bb_y1, 1)
        b_off = torch.unsqueeze(batch_bb_y2, 1) - torch.unsqueeze(coords_y, 2)
        ltrb_off = torch.stack([l_off, t_off, r_off, b_off], -1)    # [B x H*W x M x 4]
        # [left_x, top_y, right_x, bottom_y]

        areas = (l_off + r_off) * (t_off + b_off)            # [B x H*W x M]
        # For each location we have relationg with All M bboxes.
        # But we will select only one bbox out of these M bboxes for which
        # the area is lowest.

        off_min = torch.min(ltrb_off, axis=-1)[0]           # [B x H*W x M]
        off_max = torch.max(ltrb_off, axis=-1)[0]           # [B x H*W x M]

        mask_feat_map_limit = (off_max > limit_range[0]) & (off_max <= limit_range[1])
        # [B x H*W x M]

        mask_in_gtbbox = off_min > 0        # [B x H*W x M]
        # The bbox out of M which are padded, for such bbox this area will return
        # -ve and it will be the minimum area out of all the M bboxes. This would
        # break the next logic. So to prevent it, this mask is introduced.

        if self.central_sampling:
            # Ref. : https://github.com/yqyao/FCOS_PLUS/issues/13#issuecomment-564823086
            img_level_radius = self.central_sampling_radius * stride

            gtbbox_xc = (batch_bb_x1 + batch_bb_x2) / 2     # [B x M]
            gtbbox_yc = (batch_bb_y1 + batch_bb_y2) / 2     # [B x M]

            # coords_x   : [1 x H*W] --> [1 x H*W x 1]
            # batch_bb_x1: [B x M] --> [B x 1 x M]
            center_x_off = torch.abs(torch.unsqueeze(coords_x, 2) - torch.unsqueeze(gtbbox_xc, 1))
            center_y_off = torch.abs(torch.unsqueeze(coords_y, 2) - torch.unsqueeze(gtbbox_yc, 1))
            center_xy_off = torch.stack([center_x_off, center_y_off], dim=-1)   # [B x H*W x M x 2]
            center_off_max = torch.max(center_xy_off, dim=-1)[0]                # [B x H*W x M]
            mask_central_sampling = center_off_max < img_level_radius           # [B x H*W X M]

        else:
            mask_central_sampling = torch.ones_like(mask_in_gtbbox)     # [B x H*W X M]

        mask_positive_bbox = mask_feat_map_limit & mask_in_gtbbox & mask_central_sampling
        # [B x H*W x M]

        areas[~mask_positive_bbox] = 9999999            # [B x H*W x M]
        # We would make area for negative bboxes infinite so that the minimum area
        # will be computed from positive bboxes only.

        area_min_idx = torch.min(areas, dim=-1)[1]      # [B x H*W]

        # areas : [B x H*W x M]
        ltrb_off_mask = torch.zeros_like(areas, dtype=torch.bool).scatter_(\
                                -1, area_min_idx.unsqueeze(dim=-1), 1)
        # This is the binary mask of shape [B x H*W x M] with value 1 for
        # all the indexes mentioned by area_min_idx at the last dimension which is having
        # M values. So out of those M values at last dimension, only one value which is
        # specified in area_min_idx will be made to 1.0 and rest will be kept as it is at 0.0

        reg_targets = ltrb_off[ltrb_off_mask]       # [B*H*W x 4]
        reg_targets = torch.reshape(reg_targets, (-1, feat_h * feat_w, 4))  # [B x H*W x 4]

        classes = torch.unsqueeze(batch_bbox[:, :, 4], dim=1)           # [B x 1 x M]
        classes = torch.broadcast_tensors(classes, areas.long())[0]     # [B x H*W x M]

        cls_targets = classes[ltrb_off_mask]                                # [B*H*W x 1]
        cls_targets = torch.reshape(cls_targets, (-1, feat_h * feat_w, 1))  # [B x H*W x 1]

        mask_positive_bbox_2 = mask_positive_bbox.long().sum(dim=-1)        # [B x H*W]
        mask_positive_bbox_2 = mask_positive_bbox_2 >= 1

        left_right_min = torch.min(reg_targets[:, :, 0], reg_targets[:, :, 2])      # [B x H*W]
        left_right_max = torch.max(reg_targets[:, :, 0], reg_targets[:, :, 2])      # [B x H*W]
        top_bottom_min = torch.min(reg_targets[:, :, 1], reg_targets[:, :, 3])      # [B x H*W]
        top_bottom_max = torch.max(reg_targets[:, :, 1], reg_targets[:, :, 3])      # [B x H*W]
        cnt_targets = torch.unsqueeze(torch.sqrt((left_right_min * top_bottom_min) /
                            (left_right_max * top_bottom_max + 1e-9)), dim=-1)     # [B x H*W x 1]

        cls_targets[~mask_positive_bbox_2] = 0              # [B x H*W x 1]
        # Assigning label=0 which is background class, for all such locations
        # which are negative i.e. locations which dont get associated to any bbox
        # print(cls_targets.shape)

        cls_targets = torch.nn.functional.one_hot(cls_targets[:, :, 0].long(), num_classes=self.num_classes)
        # [B x H*W x 1] ---> [B x H*W x 81]

        cnt_targets[~mask_positive_bbox_2] = -1             # [B x H*W x 1]
        reg_targets[~mask_positive_bbox_2] = -1             # [B x H*W x 4]
        return cls_targets, cnt_targets, reg_targets

    def compute_cls_loss(self, comb_cls_probs_pred, comb_cls_probs_target, mask_pos,
                            alpha=0.25, gamma=2.0):
        # comb_cls_probs_pred   : [B x sum of all H*W x 81]
        # comb_cls_probs_target : [B x sum of all H*W x 81]
        # mask_pos              : [B x sum of all H*W x 1]

        assert comb_cls_probs_pred.shape[:2] == comb_cls_probs_target.shape[:2]

        num_pos = torch.sum(mask_pos, dim=[1, 2]).clamp_(min=1).float()     # [B]

        # Focal loss to be computed for all grid points
        # Positive as well as negative
        pt = comb_cls_probs_pred * comb_cls_probs_target + \
            (1 - comb_cls_probs_pred) * (1 - comb_cls_probs_target)
        w = alpha * comb_cls_probs_target + (1 - alpha) * (1 - comb_cls_probs_target)
        focal_loss = (-1) * w * torch.pow((1 - pt), gamma) * torch.log(pt + 1e-10)
        focal_loss = focal_loss.sum(dim=[1, 2])         # [B]

        focal_loss = focal_loss / num_pos               # [B]
        return focal_loss.mean()                        # [1]

    def compute_centerness_loss(self, comb_cnt_logits_pred, comb_cnt_logits_target, mask_pos):
        # comb_cnt_logits_pred   : [B x sum of all H*W x 1]
        # comb_cnt_logits_target : [B x sum of all H*W x 1]
        # mask_pos               : [B x sum of all H*W x 1]

        num_pos = torch.sum(mask_pos, dim=[1, 2]).clamp_(min=1).float()     # [B]

        cnt_loss_batch = []
        batch_size = comb_cnt_logits_pred.shape[0]
        for i in range(batch_size):
            preds = comb_cnt_logits_pred[i][mask_pos[i]]          # [n]
            targets = comb_cnt_logits_target[i][mask_pos[i]]      # [n]

            cnt_loss = nn.functional.binary_cross_entropy_with_logits(input=preds, \
                                target=targets, reduction='none') # [n]
            cnt_loss = torch.sum(cnt_loss)                        # [1]
            cnt_loss_batch.append(cnt_loss)

        cnt_loss_batch = torch.stack(cnt_loss_batch, dim=0)       # [B]
        cnt_loss_batch = cnt_loss_batch / num_pos                 # [B]
        return cnt_loss_batch.mean()                              # [1]

    def compute_coordinate_reg_loss(self, comb_reg_values_pred, comb_reg_values_target, mask_pos):
        # comb_reg_values_pred   : [B x sum of all H*W x 4]
        # comb_reg_values_target : [B x sum of all H*W x 4]
        # mask_pos               : [B x sum of all H*W x 1]

        num_pos = torch.sum(mask_pos, dim=[1, 2]).clamp_(min=1).float()     # [B]

        reg_loss_batch = []
        batch_size = comb_reg_values_pred.shape[0]
        for i in range(batch_size):
            pred_pos = comb_reg_values_pred[i][mask_pos[i, :, 0]]           # [n, 4]
            target_pos = comb_reg_values_target[i][mask_pos[i, :, 0]]       # [n, 4]

            if self.regression_loss_type == 'iou':
                reg_loss = self.iou_loss(pred_pos, target_pos)  # [1]
            elif self.regression_loss_type == 'giou':
                reg_loss = self.giou_loss(pred_pos, target_pos) # [1]
            else:
                raise NotImplementedError(f"Regression Loss type: {self.regression_loss_type} is not supported.")
            reg_loss_batch.append(reg_loss)

        reg_loss_batch = torch.stack(reg_loss_batch, dim=0)     # [B]
        reg_loss_batch = reg_loss_batch / num_pos
        return reg_loss_batch.mean()                            # [1]

    def iou_loss(self, pred_pos, target_pos):
        # pred_pos      : [n, 4]    : [left_x, top_y, right_x, bottom_y]
        # target_pos    : [n, 4]    : [left_x, top_y, right_x, bottom_y]

        x1y1_intersection = torch.min(pred_pos[:, :2], target_pos[:, :2])   # [n, 2]
        x2y2_intersection = torch.min(pred_pos[:, 2:], target_pos[:, 2:])   # [n, 2]
        wh = (x1y1_intersection + x2y2_intersection).clamp(min=0)           # [n, 2]
        overlap_area = wh[:, 0] * wh[:, 1]                                  # [n]
        area1 = (pred_pos[:, 2] + pred_pos[:, 0]) \
                    * (pred_pos[:, 3] + pred_pos[:, 1])             # [n]
        area2 = (target_pos[:, 2] + target_pos[:, 0]) \
                    * (target_pos[:, 3] + target_pos[:, 1])         # [n]

        iou = overlap_area / (area1 + area2 - overlap_area + 1e-9)  # [n]

        iou_loss = -torch.log(iou.clamp(1e-9))                      # [n]

        return iou_loss.sum()                                       # [1]

    def giou_loss(self, pred_pos, target_pos):
        # pred_pos      : [n, 4]    : [left_x, top_y, right_x, bottom_y]
        # target_pos    : [n, 4]    : [left_x, top_y, right_x, bottom_y]

        x1y1_outer = torch.max(pred_pos[:, :2], target_pos[:, :2])   # [n, 2]
        x2y2_outer = torch.max(pred_pos[:, 2:], target_pos[:, 2:])   # [n, 2]
        wh_outer = (x1y1_outer + x2y2_outer).clamp(min=0)            # [n, 2]
        outer_area = wh_outer[:, 0] * wh_outer[:, 1]                 # [n]

        x1y1_intersection = torch.min(pred_pos[:, :2], target_pos[:, :2])   # [n, 2]
        x2y2_intersection = torch.min(pred_pos[:, 2:], target_pos[:, 2:])   # [n, 2]
        wh = (x1y1_intersection + x2y2_intersection).clamp(min=0)           # [n, 2]
        overlap_area = wh[:, 0] * wh[:, 1]                                  # [n]
        area1 = (pred_pos[:, 2] + pred_pos[:, 0]) \
                    * (pred_pos[:, 3] + pred_pos[:, 1])             # [n]
        area2 = (target_pos[:, 2] + target_pos[:, 0]) \
                    * (target_pos[:, 3] + target_pos[:, 1])         # [n]

        union = area1 + area2 - overlap_area                        # [n]
        iou = overlap_area / (union + 1e-9)                         # [n]

        giou_loss = 1 - iou + \
                        ((outer_area - union) / (outer_area + 1e-9))    # [n]

        return giou_loss.sum()                                      # [1]

if __name__ == '__main__':
    import sys
    import numpy as np
    sys.path.append("../")
    import config_temp as config
    loss_fn = FCOSLoss(limit_range=config.limit_range, central_sampling=config.central_sampling, \
                        central_sampling_radius=config.central_sampling_radius, strides=config.strides, \
                        regression_loss_type=config.regression_loss_type, num_classes=config.num_classes)

    cls_probs_pred = []
    cnt_logits_pred = []
    reg_values_pred = []
    batch_bboxes = []
    for s in config.strides:
        f_h, f_w = config.input_size[0] // s, config.input_size[1] // s

        cls_probs_pred.append(torch.randn(1, config.num_classes, f_h, f_w))
        cnt_logits_pred.append(torch.randn(1, 1, f_h, f_w))
        reg_values_pred.append(torch.randn(1, 4, f_h, f_w))

        rnd_num = np.random.randint(low=1, high=10)
        pad_len = 10 - rnd_num
        bbox = torch.randn(1, rnd_num, 5)
        bbox = torch.nn.functional.pad(bbox, pad=(0, 0, 0, pad_len))
        batch_bboxes.append(bbox)

    prediction = [cls_probs_pred, cnt_logits_pred, reg_values_pred]

    batch_bboxes = torch.cat(batch_bboxes, dim=1)

    loss_fn(prediction, batch_bboxes)


import time, os, math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from body_model.body_model import BodyModel
from body_model.utils import SMPLH_PATH, SMPL_JOINTS, VPOSER_PATH, SMPL_PARENTS, KEYPT_VERTS
from utils.transforms import rotation_matrix_to_angle_axis
from datasets.amass_utils import CONTACT_INDS

NUM_BODY_JOINTS = len(SMPL_JOINTS) - 1
BETA_SIZE = 16
CONTACT_THRESH = 0.5

class HumorLoss(nn.Module):

    def __init__(self,
                    kl_loss=1.0,
                    kl_loss_anneal_start=0,
                    kl_loss_anneal_end=0,
                    kl_loss_cycle_len=-1, # if > 0 will anneal KL loss cyclicly
                    regr_trans_loss=1.0,
                    regr_trans_vel_loss=1.0,
                    regr_root_orient_loss=1.0,
                    regr_root_orient_vel_loss=1.0,
                    regr_pose_loss=1.0,
                    regr_pose_vel_loss=1.0,
                    regr_joint_loss=1.0,
                    regr_joint_vel_loss=1.0,
                    regr_joint_orient_vel_loss=1.0,
                    regr_vert_loss=1.0,
                    regr_vert_vel_loss=1.0,
                    contacts_loss=0.0, # classification loss on binary contact prediction
                    contacts_vel_loss=0.0, # velocity near 0 at predicted contacts
                    smpl_joint_loss=0.0,
                    smpl_mesh_loss=0.0,
                    smpl_joint_consistency_loss=0.0,
                    smpl_vert_consistency_loss=0.0,
                    smpl_batch_size=480):
        super(HumorLoss, self).__init__()
        '''
        All loss inputs are weights for that loss term. If the weight is 0, the loss is not used.

        - regr_*_loss :                 L2 regression losses on various state terms (root trans/orient, body pose, joint positions, and joint velocities)
        - smpl_joint_loss :             L2 between GT joints and joint locations resulting from SMPL model (parameterized by trans/orient/body poase)
        - smpl_mesh_loss :              L2 between GT and predicted vertex locations resulting from SMPL model (parameterized by trans/orient/body poase)
        - smpl_joint_consistency_loss : L2 between regressed joints and predicted joint locations from SMPL model (ensures consistency between
                                        state joint locations and joint angle predictions)
        - kl_loss :                     divergence between predicted posterior and prior

        - smpl_batch_size : the size of batches that will be given to smpl. if less than this is passed in, will be padded accordingly. however, passed
                            in batches CANNOT be greater than this number.
        '''
        self.kl_loss_weight = kl_loss
        self.kl_loss_anneal_start = kl_loss_anneal_start
        self.kl_loss_anneal_end = kl_loss_anneal_end
        self.use_kl_anneal = self.kl_loss_anneal_end > self.kl_loss_anneal_start

        self.kl_loss_cycle_len = kl_loss_cycle_len
        self.use_kl_cycle = False
        if self.kl_loss_cycle_len > 0:
            self.use_kl_cycle = True
            self.use_kl_anneal = False

        self.contacts_loss_weight = contacts_loss
        self.contacts_vel_loss_weight = contacts_vel_loss
        self.bce_loss = nn.BCEWithLogitsLoss(reduction='none')

        # build dict of all possible regression losses based on inputs
        # keys must be the same as we expect from the pred/gt data
        self.regr_loss_weight_dict = {
            'trans' : regr_trans_loss,
            'trans_vel' : regr_trans_vel_loss,
            'root_orient' : regr_root_orient_loss,
            'root_orient_vel' : regr_root_orient_vel_loss,
            'pose_body' : regr_pose_loss,
            'pose_body_vel' : regr_pose_vel_loss,
            'joints' : regr_joint_loss,
            'joints_vel' : regr_joint_vel_loss,
            'joints_orient_vel' : regr_joint_orient_vel_loss,
            'verts' : regr_vert_loss,
            'verts_vel' : regr_vert_vel_loss
        }

        self.smpl_joint_loss_weight = smpl_joint_loss
        self.smpl_mesh_loss_weight = smpl_mesh_loss
        self.smpl_joint_consistency_loss_weight = smpl_joint_consistency_loss
        self.smpl_vert_consistency_loss_weight = smpl_vert_consistency_loss

        self.l2_loss = nn.MSELoss(reduction='none')
        self.regr_loss = nn.MSELoss(reduction='none')

        smpl_losses = [self.smpl_joint_loss_weight, self.smpl_mesh_loss_weight, self.smpl_joint_consistency_loss_weight, self.smpl_vert_consistency_loss_weight]
        self.smpl_batch_size = smpl_batch_size
        self.use_smpl_losses = False
        if sum(smpl_losses) > 0.0:
            self.use_smpl_losses = True
            # need a body model to compute the losses
            male_bm_path = os.path.join(SMPLH_PATH, 'male/model.npz')
            self.male_bm = BodyModel(bm_path=male_bm_path, num_betas=16, batch_size=self.smpl_batch_size)
            female_bm_path = os.path.join(SMPLH_PATH, 'female/model.npz')
            self.female_bm = BodyModel(bm_path=female_bm_path, num_betas=16, batch_size=self.smpl_batch_size)

    def forward(self, pred_dict, gt_dict, cur_epoch, gender=None, betas=None):
        '''
        Compute the loss.

        All data in the dictionaries should be of size B x D.

        group_regr_losses will be used to aggregate every group_regr_lossth batch idx together into
        the stats_dict. This can be useful when there are multiple output steps and you want to track
        each separately.
        '''
        loss = 0.0
        stats_dict = dict()

        #
        # KL divergence
        #
        if self.kl_loss_weight > 0.0:
            qm, qv = pred_dict['posterior_distrib']
            pm, pv = pred_dict['prior_distrib']
            kl_loss = self.kl_normal(qm, qv, pm, pv)
            kl_stat_loss = kl_loss.mean()
            kl_loss = kl_stat_loss
            # print(kl_loss.size())
            stats_dict['kl_loss'] = kl_stat_loss
            anneal_weight = 1.0
            if self.use_kl_anneal or self.use_kl_cycle:
                anneal_epoch = cur_epoch
                anneal_start = self.kl_loss_anneal_start
                anneal_end = self.kl_loss_anneal_end
                if self.use_kl_cycle:
                    anneal_epoch = cur_epoch % self.kl_loss_cycle_len
                    anneal_start = 0
                    anneal_end = self.kl_loss_cycle_len // 2 # optimize full weight for second half of cycle
                if anneal_epoch >= anneal_start:
                    anneal_weight = (anneal_epoch - anneal_start) / (anneal_end - anneal_start)
                else:
                    anneal_weight = 0.0
                anneal_weight = 1.0 if anneal_weight > 1.0 else anneal_weight

            loss = loss + anneal_weight*self.kl_loss_weight*kl_loss

            stats_dict['kl_anneal_weight'] = anneal_weight
            stats_dict['kl_weighted_loss'] = loss

        #
        # Reconstruction
        #

        # regression terms
        for cur_key in gt_dict.keys():
            # print(cur_key)
            if cur_key not in self.regr_loss_weight_dict:
                continue
            cur_regr_weight = self.regr_loss_weight_dict[cur_key]
            if cur_regr_weight > 0.0:
                pred_val = pred_dict[cur_key]
                gt_val = gt_dict[cur_key]

                if cur_key == 'root_orient' or cur_key == 'pose_body':
                    # rotations use L2 for matrices
                    cur_regr_loss = self.l2_loss(pred_val, gt_val)
                else:
                    cur_regr_loss = self.regr_loss(pred_val, gt_val)
                agg_cur_regr_loss = cur_regr_loss
                cur_regr_stat_loss = agg_cur_regr_loss.mean()
                agg_cur_regr_loss = cur_regr_stat_loss
                stats_dict[cur_key + '_loss'] = cur_regr_stat_loss
                loss = loss + cur_regr_weight*agg_cur_regr_loss

        if self.contacts_loss_weight > 0.0:
            if 'contacts' in gt_dict.keys() and 'contacts' in pred_dict.keys():
                gt_contacts = gt_dict['contacts']
                pred_contacts = pred_dict['contacts']
                # pred is assumed to be logits from network (i.e. sigmoid has not been applied yet)
                cur_contacts_loss = self.bce_loss(pred_contacts, gt_contacts)
                cur_stat_contacts_loss = cur_contacts_loss.mean()
                cur_contacts_loss = cur_stat_contacts_loss
                stats_dict['contacts_loss'] = cur_stat_contacts_loss
                loss = loss + self.contacts_loss_weight*cur_contacts_loss

                # other accuracy statistics
                pred_contacts = (torch.sigmoid(pred_contacts) > CONTACT_THRESH).to(torch.bool)
                gt_contacts = gt_contacts.to(torch.bool)
                # counts for confusion matrix
                # true positive (pred contact, labeled contact)
                true_pos = pred_contacts & gt_contacts
                true_pos_cnt = torch.sum(true_pos).to(torch.float)
                # false positive (pred contact, not lebeled contact)
                false_pos = pred_contacts & ~(gt_contacts)
                false_pos_cnt = torch.sum(false_pos).to(torch.float)
                # false negative (pred no contact, labeled contact)
                false_neg = ~(pred_contacts) & gt_contacts
                false_neg_cnt = torch.sum(false_neg).to(torch.float)
                # true negative (pred no contact, no labeled contact)
                true_neg = (~pred_contacts) & (~gt_contacts)
                true_neg_cnt = torch.sum(true_neg).to(torch.float)

                acc = (true_pos_cnt + true_neg_cnt) / (true_pos_cnt + false_pos_cnt + false_neg_cnt + true_neg_cnt)
                pos_acc = true_pos_cnt / (true_pos_cnt + false_neg_cnt)
                neg_acc = true_neg_cnt / (true_neg_cnt + false_pos_cnt)
                stats_dict['contacts_acc'] = acc
                stats_dict['contacts_pos_acc'] = pos_acc
                stats_dict['contacts_neg_acc'] = neg_acc
            else:
                print('Cannot compute contact loss without contact pred/gt! Skipping...')

        if self.contacts_vel_loss_weight > 0.0:
            if 'contacts' in pred_dict.keys() and 'joints_vel' in pred_dict.keys():
                pred_contacts = torch.sigmoid(pred_dict['contacts'])
                pred_joints_vel = pred_dict['joints_vel'].reshape((-1, len(SMPL_JOINTS), 3))
                contact_joints_vel = pred_joints_vel[:,CONTACT_INDS,:]
                # use predicted contact probability to weight regularization on joint velocity
                vel_mag = torch.norm(contact_joints_vel, dim=-1)
                cur_contact_vel_loss = pred_contacts*(vel_mag**2)
                cur_stat_contact_vel_loss = cur_contact_vel_loss.mean()
                cur_contact_vel_loss = cur_stat_contact_vel_loss
                stats_dict['contacts_vel_loss'] = cur_stat_contact_vel_loss
                loss = loss + self.contacts_vel_loss_weight*cur_contact_vel_loss
            else:
                print('Cannot compute contact vel loss without contact and joints_vel pred! Skipping...')

        # terms requiring SMPL reconstruction
        if self.use_smpl_losses:
            if gender is None or betas is None:
                raise Exception('Must pass gender and betas to MotionVAE loss to use SMPL losses!')

            try:
                pred_trans = pred_dict['trans']
                pred_orient = pred_dict['root_orient']
                pred_pose = pred_dict['pose_body']
                gt_trans = gt_dict['trans']
                gt_orient = gt_dict['root_orient']
                gt_pose = gt_dict['pose_body']
            except KeyError:
                print('ERROR: In order to use SMPL losses must have trans, root_orient, and pose_body in pred and gt dicts!')
                exit()

            # need to transform rotation matrices to aa for SMPL model
            B = pred_trans.size(0)
            pred_orient = rotation_matrix_to_angle_axis(pred_orient.reshape((B, 3, 3)))
            gt_orient = rotation_matrix_to_angle_axis(gt_orient.reshape((B, 3, 3)))
            pred_pose = rotation_matrix_to_angle_axis(pred_pose.reshape((B*NUM_BODY_JOINTS, 3, 3))).reshape((B, NUM_BODY_JOINTS*3))
            gt_pose = rotation_matrix_to_angle_axis(gt_pose.reshape((B*NUM_BODY_JOINTS, 3, 3))).reshape((B, NUM_BODY_JOINTS*3))

            pred_vals = [pred_trans, pred_orient, pred_pose]
            gt_vals = [gt_trans, gt_orient, gt_pose, betas]

            # have to split by gender to make sure we use the correct body model
            gender_names = ['male', 'female']
            mask_list = []
            pred_joints = []
            pred_mesh = []
            gt_joints = []
            gt_mesh = []
            for gender_name in gender_names:
                # print(gender_name)
                gender_idx = gender[:, 0] == gender_name
                mask_list.append(gender_idx)
                cur_pred_vals = [val[gender_idx] for val in pred_vals]
                cur_gt_vals = [val[gender_idx] for val in gt_vals]

                # need to pad extra frames with zeros in case not as long as expected
                pad_size = self.smpl_batch_size - cur_pred_vals[0].size(0)
                if pad_size == self.smpl_batch_size:
                    # skip if no frames for this gender
                    continue
                pad_list = cur_pred_vals + cur_gt_vals
                if pad_size < 0:
                    raise Exception('SMPL model batch size not large enough to accomodate!')
                elif pad_size > 0:
                    pad_list = self.zero_pad_tensors(pad_list, pad_size)

                # reconstruct SMPL
                cur_pred_trans, cur_pred_orient, cur_pred_pose, cur_gt_trans, cur_gt_orient, cur_gt_pose, cur_betas = pad_list
                bm = self.male_bm if gender_name == 'male' else self.female_bm
                pred_body = bm(pose_body=cur_pred_pose, betas=cur_betas, root_orient=cur_pred_orient, trans=cur_pred_trans)
                gt_body = bm(pose_body=cur_gt_pose, betas=cur_betas, root_orient=cur_gt_orient, trans=cur_gt_trans)
                if pad_size > 0:
                    pred_joints.append(pred_body.Jtr[:-pad_size])
                    pred_mesh.append(pred_body.v[:-pad_size])
                    gt_joints.append(gt_body.Jtr[:-pad_size])
                    gt_mesh.append(gt_body.v[:-pad_size])
                else:
                    pred_joints.append(pred_body.Jtr)
                    pred_mesh.append(pred_body.v)
                    gt_joints.append(gt_body.Jtr)
                    gt_mesh.append(gt_body.v)

            pred_joints = torch.cat(pred_joints, axis=0)[:,:len(SMPL_JOINTS),:]
            pred_mesh = torch.cat(pred_mesh, axis=0)
            gt_joints = torch.cat(gt_joints, axis=0)[:,:len(SMPL_JOINTS),:]
            gt_mesh = torch.cat(gt_mesh, axis=0)

            pred_verts = pred_mesh[:,KEYPT_VERTS,:]
            gt_verts = gt_mesh[:,KEYPT_VERTS,:]

            # now compute SMPL-related losses
            if self.smpl_joint_loss_weight > 0.0:
                smpl_joint_loss = self.regr_loss(pred_joints, gt_joints)
                smpl_joint_stat_loss = smpl_joint_loss.mean()
                smpl_joint_loss = smpl_joint_stat_loss
                stats_dict['smpl_joint_loss'] = smpl_joint_stat_loss
                loss = loss + self.smpl_joint_loss_weight*smpl_joint_loss
            if self.smpl_mesh_loss_weight > 0.0:
                smpl_mesh_loss = self.regr_loss(pred_mesh, gt_mesh)
                smpl_mesh_stat_loss = smpl_mesh_loss.mean()
                smpl_mesh_loss = smpl_mesh_stat_loss
                stats_dict['smpl_mesh_loss'] = smpl_mesh_stat_loss
                loss = loss + self.smpl_mesh_loss_weight*smpl_mesh_loss
            if self.smpl_joint_consistency_loss_weight > 0.0:
                if not 'joints' in pred_dict.keys():
                    print('Must regress joints in order to use smpl joint consistency loss!')
                    exit()
                regressed_joints = pred_dict['joints'].reshape((B, len(SMPL_JOINTS), -1))
                # need to reorder regressed joints with mask_list to ensure consistency with smpl joints
                regressed_joints = torch.cat([regressed_joints[mask_list[i]] for i in range(len(mask_list))], axis=0)

                smpl_joint_consistency_loss = self.regr_loss(pred_joints, regressed_joints)
                smpl_joint_consistency_stat_loss = smpl_joint_consistency_loss.mean()
                smpl_joint_consistency_loss = smpl_joint_consistency_stat_loss
                stats_dict['smpl_joint_consistency_loss'] = smpl_joint_consistency_stat_loss
                loss = loss + self.smpl_joint_consistency_loss_weight*smpl_joint_consistency_loss
            if self.smpl_vert_consistency_loss_weight > 0.0:
                if not 'verts' in pred_dict.keys():
                    print('Must regress verts in order to use smpl vert consistency loss!')
                    exit()
                regressed_verts = pred_dict['verts'].reshape((B, len(KEYPT_VERTS), -1))
                # need to reorder regressed verts with mask_list to ensure consistency with smpl verts
                regressed_verts = torch.cat([regressed_verts[mask_list[i]] for i in range(len(mask_list))], axis=0)

                smpl_vert_consistency_loss = self.regr_loss(pred_verts, regressed_verts)
                smpl_vert_consistency_stat_loss = smpl_vert_consistency_loss.mean()
                smpl_vert_consistency_loss = smpl_vert_consistency_stat_loss
                stats_dict['smpl_vert_consistency_loss'] = smpl_vert_consistency_stat_loss
                loss = loss + self.smpl_vert_consistency_loss_weight*smpl_vert_consistency_loss

        if self.kl_loss_weight > 0.0:
            stats_dict['reconstr_weighted_loss'] = loss - stats_dict['kl_weighted_loss']

        return loss, stats_dict

    def zero_pad_tensors(self, pad_list, pad_size):
        '''
        Assumes tensors in pad_list are B x D
        '''
        new_pad_list = []
        for pad_idx, pad_tensor in enumerate(pad_list):
            padding = torch.zeros((pad_size, pad_tensor.size(1))).to(pad_tensor)
            new_pad_list.append(torch.cat([pad_tensor, padding], dim=0))
        return new_pad_list

    def kl_normal(self, qm, qv, pm, pv):
        """
        Computes the elem-wise KL divergence between two normal distributions KL(q || p) and
        sum over the last dimension

        Args:
            qm: tensor: (batch, dim): q mean
            qv: tensor: (batch, dim): q variance
            pm: tensor: (batch, dim): p mean
            pv: tensor: (batch, dim): p variance

        Return:
            kl: tensor: (batch,): kl between each sample
        """
        element_wise = 0.5 * (torch.log(pv) - torch.log(qv) + qv / pv + (qm - pm).pow(2) / pv - 1)
        kl = element_wise.sum(-1)
        return kl

    def log_normal(self, x, m, v):
        """
        Computes the elem-wise log probability of a Gaussian and then sum over the
        last dim. Basically we're assuming all dims are batch dims except for the
        last dim.    Args:
            x: tensor: (batch_1, batch_2, ..., batch_k, dim): Observation
            m: tensor: (batch_1, batch_2, ..., batch_k, dim): Mean
            v: tensor: (batch_1, batch_2, ..., batch_k, dim): Variance    Return:
            log_prob: tensor: (batch_1, batch_2, ..., batch_k): log probability of
                each sample. Note that the summation dimension is not kept
        """
        log_prob = -torch.log(torch.sqrt(v)) - math.log(math.sqrt(2*math.pi)) \
                        - ((x - m)**2 / (2*v))
        log_prob = torch.sum(log_prob, dim=-1)
        return log_prob