"""
halo finding

"""

#-----------------------------------------------------------------------------
# Copyright (c) yt Development Team. All rights reserved.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

import gc
import numpy as np

from yt.config import ytcfg
from yt.funcs import mylog
from yt.utilities.math_utils import \
    get_rotation_matrix, \
    periodic_dist
from yt.utilities.physical_constants import \
    mass_sun_cgs
from yt.utilities.physical_ratios import \
    rho_crit_g_cm3_h2, \
    TINY
from yt.utilities.parallel_tools.parallel_analysis_interface import \
    ParallelAnalysisInterface

from yt_astro_analysis.halo_analysis.halo_finding.hop.EnzoHop import RunHOP
from yt_astro_analysis.halo_analysis.halo_finding.fof.EnzoFOF import RunFOF

class Halo(object):
    """
    A data source that returns particle information about the members of a
    HOP-identified halo.
    """
    _distributed = False
    _processing = False
    _owner = 0
    indices = None
    extra_wrap = ["__getitem__"]

    def __init__(self, halo_list, id, indices=None, size=None, CoM=None,
        max_dens_point=None, group_total_mass=None, max_radius=None,
        bulk_vel=None, tasks=None, rms_vel=None, supp=None, ptype='all'):
        self.ptype = ptype
        self.halo_list = halo_list
        self._max_dens = halo_list._max_dens
        self.id = id
        self.data = halo_list._data_source
        self.ds = self.data.ds
        self.gridsize = (self.ds.domain_right_edge - \
                 self.ds.domain_left_edge)
        if indices is not None:
            self.indices = halo_list._base_indices[indices]
        else:
            self.indices = None
        # We assume that if indices = None, the instantiator has OTHER plans
        # for us -- i.e., setting it somehow else
        self.size = size
        self.CoM = CoM
        self.max_dens_point = max_dens_point
        self.group_total_mass = group_total_mass
        self.max_radius = max_radius
        self.bulk_vel = bulk_vel
        self.tasks = tasks
        self.rms_vel = rms_vel
        self.bin_count = None
        self.overdensity = None
        # A supplementary data dict.
        if supp is None:
            self.supp = {}
        else:
            self.supp = supp
        self._saved_fields = {}
        self._ds_sort = None
        self._particle_mask = None

    @property
    def particle_mask(self):
        # Dynamically create the masking array for particles, and get
        # the data using standard yt methods.
        if self._particle_mask is not None:
            return self._particle_mask
        # This is from disk.
        pid = self.__getitem__('particle_index')
        # This is from the sphere.
        if self._name == "RockstarHalo":
            ds = self.ds.sphere(self.CoM, self._radjust * self.max_radius)
        elif self._name == "LoadedHalo":
            ds = self.ds.sphere(self.CoM, np.maximum(self._radjust * \
            self.ds.quan(self.max_radius, 'code_length'), \
            self.ds.index.get_smallest_dx()))
        sp_pid = ds['particle_index']
        self._ds_sort = sp_pid.argsort()
        sp_pid = sp_pid[self._ds_sort]
        # This matches them up.
        self._particle_mask = np.in1d(sp_pid, pid)
        return self._particle_mask

    def center_of_mass(self):
        r"""Calculate and return the center of mass.

        The center of mass of the halo is directly calculated and returned.

        Examples
        --------
        >>> com = halos[0].center_of_mass()
        """
        if self.CoM is not None:
            return self.CoM
        pm = self["particle_mass"].in_units('Msun')
        c = {}
        # We shift into a box where the origin is the left edge
        c[0] = self["particle_position_x"] - self.ds.domain_left_edge[0]
        c[1] = self["particle_position_y"] - self.ds.domain_left_edge[1]
        c[2] = self["particle_position_z"] - self.ds.domain_left_edge[2]
        com = []
        for i in range(3):
            # A halo is likely periodic around a boundary if the distance
            # between the max and min particle
            # positions are larger than half the box.
            # So skip the rest if the converse is true.
            # Note we might make a change here when periodicity-handling is
            # fully implemented.
            if (c[i].max() - c[i].min()) < (self.ds.domain_width[i] / 2.):
                com.append(c[i])
                continue
            # Now we want to flip around only those close to the left boundary.
            sel = (c[i] <= (self.ds.domain_width[i]/2))
            c[i][sel] += self.ds.domain_width[i]
            com.append(c[i])

        c = (com * pm).sum(axis=1) / pm.sum()
        c = self.ds.arr(c, 'code_length')

        return c%self.ds.domain_width + self.ds.domain_left_edge

    def maximum_density(self):
        r"""Return the HOP-identified maximum density. Not applicable to
        FOF halos.

        Return the HOP-identified maximum density. Not applicable to FOF halos.

        Examples
        --------
        >>> max_dens = halos[0].maximum_density()
        """
        if self.max_dens_point is not None:
            return self.max_dens_point[0]
        return self._max_dens[self.id][0]

    def maximum_density_location(self):
        r"""Return the location HOP identified as maximally dense. Not
        applicable to FOF halos.

        Return the location HOP identified as maximally dense.

        Examples
        --------
        >>> max_dens_loc = halos[0].maximum_density_location()
        """
        if self.max_dens_point is not None:
            return self.max_dens_point[1:]
        return np.array([
                self._max_dens[self.id][1],
                self._max_dens[self.id][2],
                self._max_dens[self.id][3]])

    def total_mass(self):
        r"""Returns the total mass in solar masses of the halo.

        Returns the total mass in solar masses of just the particles in the
        halo.

        Examples
        --------
        >>> halos[0].total_mass()
        """
        if self.group_total_mass is not None:
            return self.group_total_mass
        return self["particle_mass"].in_units('Msun').sum()

    def bulk_velocity(self):
        r"""Returns the mass-weighted average velocity in cm/s.

        This calculates and returns the mass-weighted average velocity of just
        the particles in the halo in cm/s.

        Examples
        --------
        >>> bv = halos[0].bulk_velocity()
        """
        if self.bulk_vel is not None:
            return self.bulk_vel
        pm = self["particle_mass"].in_units('Msun')
        vx = (self["particle_velocity_x"] * pm).sum()
        vy = (self["particle_velocity_y"] * pm).sum()
        vz = (self["particle_velocity_z"] * pm).sum()
        return self.ds.arr([vx, vy, vz], vx.units) / pm.sum()

    def rms_velocity(self):
        r"""Returns the mass-weighted RMS velocity for the halo
        particles in cgs units.

        Calculate and return the mass-weighted RMS velocity for just the
        particles in the halo.  The bulk velocity of the halo is subtracted
        before computation.

        Examples
        --------
        >>> rms_vel = halos[0].rms_velocity()
        """
        if self.rms_vel is not None:
            return self.rms_vel
        bv = self.bulk_velocity()
        pm = self["particle_mass"].in_units('Msun')
        sm = pm.sum()
        vx = (self["particle_velocity_x"] - bv[0]) * pm / sm
        vy = (self["particle_velocity_y"] - bv[1]) * pm / sm
        vz = (self["particle_velocity_z"] - bv[2]) * pm / sm
        s = vx ** 2. + vy ** 2. + vz ** 2.
        ms = np.mean(s)
        return np.sqrt(ms) * pm.size

    def maximum_radius(self, center_of_mass=True):
        r"""Returns the maximum radius in the halo for all particles,
        either from the point of maximum density or from the
        center of mass.

        The maximum radius from the most dense point is calculated.  This
        accounts for periodicity.

        Parameters
        ----------
        center_of_mass : bool
            True chooses the center of mass when
            calculating the maximum radius.
            False chooses from the maximum density location for HOP halos
            (it has no effect for FOF halos).
            Default = True.

        Examples
        --------
        >>> radius = halos[0].maximum_radius()
        """
        if self.max_radius is not None:
            return self.max_radius
        if center_of_mass:
            center = self.center_of_mass()
        else:
            center = self.maximum_density_location()
        rx = np.abs(self["particle_position_x"] - center[0])
        ry = np.abs(self["particle_position_y"] - center[1])
        rz = np.abs(self["particle_position_z"] - center[2])
        DW = self.data.ds.domain_right_edge - self.data.ds.domain_left_edge
        r = np.sqrt(np.minimum(rx, DW[0] - rx) ** 2.0
                + np.minimum(ry, DW[1] - ry) ** 2.0
                + np.minimum(rz, DW[2] - rz) ** 2.0)
        return r.max()

    def __getitem__(self, key):
        return self.data[(self.ptype, key)][self.indices]

    def get_sphere(self, center_of_mass=True):
        r"""Returns a sphere source.

        This will generate a new, empty sphere source centered on this halo,
        with the maximum radius of the halo. This can be used like any other
        data container in yt.

        Parameters
        ----------
        center_of_mass : bool, optional
            True chooses the center of mass when
            calculating the maximum radius.
            False chooses from the maximum density location for HOP halos
            (it has no effect for FOF halos).
            Default = True.

        Returns
        -------
        sphere : `yt.data_objects.api.YTSphere`
            The empty data source.

        Examples
        --------
        >>> sp = halos[0].get_sphere()
        """
        if center_of_mass:
            center = self.center_of_mass()
        else:
            center = self.maximum_density_location()
        radius = self.maximum_radius()
        # A bit of a long-reach here...
        sphere = self.data.ds.sphere(center, radius=radius)
        return sphere

    def get_size(self):
        if self.size is not None:
            return self.size
        return self.indices.size

    def virial_mass(self, virial_overdensity=200., bins=300):
        r"""Return the virial mass of the halo in Msun,
        using only the particles
        in the halo (no baryonic information used).

        The virial mass is calculated, using the built in `Halo.virial_info`
        functionality.  The mass is then returned.

        Parameters
        ----------
        virial_overdensity : float
            The overdensity threshold compared to the universal average when
            calculating the virial mass. Default = 200.
        bins : int
            The number of spherical bins used to calculate overdensities.
            Default = 300.

        Returns
        -------
        mass : float
            The virial mass in solar masses of the particles in the halo.  -1
            if not virialized.

        Examples
        --------
        >>> vm = halos[0].virial_mass()
        """
        self.virial_info(bins=bins)
        vir_bin = self.virial_bin(virial_overdensity=virial_overdensity,
            bins=bins)
        if vir_bin != -1:
            return self.mass_bins[vir_bin]
        else:
            return -1

    def virial_radius(self, virial_overdensity=200., bins=300):
        r"""Return the virial radius of the halo in code units.

        The virial radius of the halo is calculated, using only the particles
        in the halo (no baryonic information used). Returns -1 if the halo is
        not virialized.

        Parameters
        ----------
        virial_overdensity : float
            The overdensity threshold compared to the universal average when
            calculating the virial radius. Default = 200.
        bins : integer
            The number of spherical bins used to calculate overdensities.
            Default = 300.

        Returns
        -------
        radius : float
            The virial radius in code units of the particles in the halo.  -1
            if not virialized.

        Examples
        --------
        >>> vr = halos[0].virial_radius()
        """
        self.virial_info(bins=bins)
        vir_bin = self.virial_bin(virial_overdensity=virial_overdensity,
            bins=bins)
        if vir_bin != -1:
            return self.radial_bins[vir_bin]
        else:
            return -1

    def virial_bin(self, virial_overdensity=200., bins=300):
        r"""Returns the bin index of the virial radius of the halo. Generally,
        it is better to call virial_radius instead, which calls this function
        automatically.
        """
        self.virial_info(bins=bins)
        over = (self.overdensity > virial_overdensity)
        if over.any():
            vir_bin = max(np.arange(bins + 1)[over])
            return vir_bin
        else:
            return -1

    def virial_info(self, bins=300):
        r"""Calculates the virial information for the halo. Generally, it is
        better to call virial_radius or virial_mass instead, which calls this
        function automatically.
        """
        # Skip if we've already calculated for this number of bins.
        if self.bin_count == bins and self.overdensity is not None:
            return None
        self.bin_count = bins
        # Cosmology
        h = self.ds.hubble_constant
        Om_matter = self.ds.omega_matter
        z = self.ds.current_redshift
        period = self.ds.domain_right_edge - \
            self.ds.domain_left_edge
        thissize = self.get_size()
        rho_crit = rho_crit_g_cm3_h2 * h ** 2.0 * Om_matter  # g cm^-3
        Msun2g = mass_sun_cgs
        rho_crit = rho_crit * ((1.0 + z) ** 3.0)
        # Get some pertinent information about the halo.
        self.mass_bins = self.ds.arr(np.zeros(self.bin_count + 1,
                                              dtype='float64'),'Msun')
        dist = np.empty(thissize, dtype='float64')
        cen = self.center_of_mass()
        mark = 0
        # Find the distances to the particles. I don't like this much, but I
        # can't see a way to eliminate a loop like this, either here or in yt.
        for pos in zip(self["particle_position_x"],
                self["particle_position_y"], self["particle_position_z"]):
            dist[mark] = periodic_dist(cen, pos, period)
            mark += 1
        # Set up the radial bins.
        # Multiply min and max to prevent issues with digitize below.
        self.radial_bins = np.logspace(np.log10(min(dist) * .99 + TINY),
            np.log10(max(dist) * 1.01 + 2 * TINY), num=self.bin_count + 1)
        self.radial_bins = self.ds.arr(self.radial_bins,'code_length')
        # Find out which bin each particle goes into, and add the particle
        # mass to that bin.
        inds = np.digitize(dist, self.radial_bins) - 1
        if self["particle_position_x"].size > 1:
            for index in np.unique(inds):
                self.mass_bins[index] += \
                np.sum(self["particle_mass"][inds == index]).in_units('Msun')
        # Now forward sum the masses in the bins.
        for i in range(self.bin_count):
            self.mass_bins[i + 1] += self.mass_bins[i]
        # Calculate the over densities in the bins.
        self.overdensity = self.mass_bins * Msun2g / \
            (4./3. * np.pi * rho_crit * \
            (self.radial_bins )**3.0)

    def _get_ellipsoid_parameters_basic(self):
        np.seterr(all='ignore')
        # check if there are 4 particles to form an ellipsoid
        # neglecting to check if the 4 particles in the same plane,
        # that is almost certainly never to occur,
        # will deal with it later if it ever comes up
        if np.size(self["particle_position_x"]) < 4:
            mylog.warning("Too few particles for ellipsoid parameters.")
            return (0, 0, 0, 0, 0, 0, 0)
        # Calculate the parameters that describe the ellipsoid of
        # the particles that constitute the halo. This function returns
        # all the parameters except for the center of mass.
        com = self.center_of_mass()
        position = [self["particle_position_x"],
                    self["particle_position_y"],
                    self["particle_position_z"]]
        # Locate the furthest particle from com, its vector length and index
        DW = np.array([self.gridsize[0],self.gridsize[1],self.gridsize[2]])
        position = [position[0] - com[0],
                    position[1] - com[1],
                    position[2] - com[2]]
        # different cases of particles being on other side of boundary
        for axis in range(np.size(DW)):
            cases = np.array([position[axis],
                                position[axis] + DW[axis],
                              position[axis] - DW[axis]])
            # pick out the smallest absolute distance from com
            position[axis] = np.choose(np.abs(cases).argmin(axis=0), cases)
        # find the furthest particle's index
        r = np.sqrt(position[0]**2 +
                    position[1]**2 +
                    position[2]**2)
        A_index = r.argmax()
        mag_A = r.max()
        # designate the A vector
        A_vector = (position[0][A_index],
                    position[1][A_index],
                    position[2][A_index])
        # designate the e0 unit vector
        e0_vector = A_vector / mag_A
        # locate the tB particle position by finding the max B
        e0_vector_copy = np.empty((np.size(position[0]), 3), dtype='float64')
        for i in range(3):
            e0_vector_copy[:, i] = e0_vector[i]
        rr = np.array([position[0],
                       position[1],
                       position[2]]).T # Similar to tB_vector in old code.
        tC_vector = np.cross(e0_vector_copy, rr)
        te2 = tC_vector.copy()
        for dim in range(3):
            te2[:,dim] *= np.sum(tC_vector**2., axis = 1)**(-0.5)
        te1 = np.cross(te2, e0_vector_copy)
        length = np.abs(-np.sum(rr * te1, axis = 1) * \
            (1. - np.sum(rr * e0_vector_copy, axis = 1)**2. * \
            mag_A**-2.)**(-0.5))
        # This problem apparently happens sometimes, that the NaNs are turned
        # into infs, which messes up the nanargmax below.
        length[length == np.inf] = 0.
        tB_index = np.nanargmax(length) # ignores NaNs created above.
        mag_B = length[tB_index]
        e1_vector = te1[tB_index]
        e2_vector = te2[tB_index]
        temp_e0 = rr.copy()
        temp_e1 = rr.copy()
        temp_e2 = rr.copy()
        for dim in range(3):
            temp_e0[:,dim] = e0_vector[dim]
            temp_e1[:,dim] = e1_vector[dim]
            temp_e2[:,dim] = e2_vector[dim]
        length = np.abs(np.sum(rr * temp_e2, axis = 1) * (1 - \
            np.sum(rr * temp_e0, axis = 1)**2. * mag_A**-2. - \
            np.sum(rr * temp_e1, axis = 1)**2. * mag_B**-2.)**(-0.5))
        length[length == np.inf] = 0.
        tC_index = np.nanargmax(length)
        mag_C = length[tC_index]
        # tilt is calculated from the rotation about x axis
        # needed to align e1 vector with the y axis
        # after e0 is aligned with x axis
        # find the t1 angle needed to rotate about z axis to align e0 onto x-z plane
        t1 = np.arctan(-e0_vector[1] / e0_vector[0])
        RZ = get_rotation_matrix(t1, (0, 0, 1))
        r1 = np.dot(RZ, e0_vector)
        # find the t2 angle needed to rotate about y axis to align e0 to x
        t2 = np.arctan(r1[2] / r1[0])
        RY = get_rotation_matrix(t2, (0, 1, 0))
        r2 = np.dot(RY, np.dot(RZ, e1_vector))
        # find the tilt angle needed to rotate about x axis to align e1 to y and e2 to z
        tilt = np.arctan(-r2[2] / r2[1])
        return (mag_A, mag_B, mag_C, e0_vector[0], e0_vector[1],
            e0_vector[2], tilt)

class HOPHalo(Halo):
    _name = "HOPHalo"
    pass

class FOFHalo(Halo):

    def maximum_density(self):
        r"""Not implemented."""
        return -1

    def maximum_density_location(self):
        r"""Not implemented."""
        return self.center_of_mass()

class HaloList(object):

    _fields = ["particle_position_%s" % ax for ax in 'xyz']

    def __init__(self, data_source, redshift=-1, ptype='all'):
        """
        Run hop on *data_source* with a given density *threshold*.
        Returns an iterable collection of
        *HopGroup* items.
        """
        self._data_source = data_source
        self.ptype = ptype
        self._groups = []
        self._max_dens = {}
        self.__obtain_particles()
        self._run_finder()
        mylog.info("Parsing outputs")
        self._parse_output()
        mylog.debug("Finished. (%s)", len(self))
        self.redshift = redshift

    def __obtain_particles(self):
        ii = slice(None)
        self.particle_fields = {}
        for field in self._fields:
            tot_part = self._data_source[(self.ptype, field)].size
            if field == "particle_index":
                self.particle_fields[field] = \
                    self._data_source[(self.ptype, field)][ii].astype('int64')
            else:
                self.particle_fields[field] = \
                    self._data_source[(self.ptype, field)][ii].astype('float64')
            del self._data_source[(self.ptype, field)]
        self._base_indices = np.arange(tot_part)[ii]
        gc.collect()

    def _parse_output(self):
        unique_ids = np.unique(self.tags)
        counts = np.bincount(self.tags + 1)
        sort_indices = np.argsort(self.tags)
        grab_indices = np.indices(self.tags.shape).ravel()[sort_indices]
        dens = self.densities[sort_indices]
        cp = 0
        for i in unique_ids:
            cp_c = cp + counts[i + 1]
            if i == -1:
                cp += counts[i + 1]
                continue
            group_indices = grab_indices[cp:cp_c]
            self._groups.append(self._halo_class(self, i, group_indices,
                                                 ptype=self.ptype))
            md_i = np.argmax(dens[cp:cp_c])
            px, py, pz = \
                [self.particle_fields['particle_position_%s' % ax][group_indices]
                 for ax in 'xyz']
            self._max_dens[i] = (dens[cp:cp_c][md_i], px[md_i],
                py[md_i], pz[md_i])
            cp += counts[i + 1]

    def __len__(self):
        return len(self._groups)

    def __iter__(self):
        for i in self._groups:
            yield i

    def __getitem__(self, key):
        return self._groups[key]

class HOPHaloList(HaloList):
    """
    Run hop on *data_source* with a given density *threshold*.
    Returns an iterable collection of *HopGroup* items.
    """
    _name = "HOP"
    _halo_class = HOPHalo
    _fields = ["particle_position_%s" % ax for ax in 'xyz'] + \
              ["particle_mass"]

    def __init__(self, data_source, threshold=160.0, ptype='all'):
        self.threshold = threshold
        mylog.info("Initializing HOP")
        HaloList.__init__(self, data_source, ptype=ptype)

    def _run_finder(self):
        self.densities, self.tags = \
            RunHOP(self.particle_fields["particle_position_x"] / self.period[0],
                self.particle_fields["particle_position_y"] / self.period[1],
                self.particle_fields["particle_position_z"] / self.period[2],
                self.particle_fields["particle_mass"].in_units('Msun'),
                self.threshold)
        self.particle_fields["densities"] = self.densities
        self.particle_fields["tags"] = self.tags

class FOFHaloList(HaloList):
    _name = "FOF"
    _halo_class = FOFHalo

    def __init__(self, data_source, link=0.2, redshift=-1, ptype='all'):
        self.link = link
        mylog.info("Initializing FOF")
        HaloList.__init__(self, data_source, redshift=redshift, ptype=ptype)

    def _run_finder(self):
        self.tags = \
        RunFOF(self.particle_fields["particle_position_x"] / self.period[0],
               self.particle_fields["particle_position_y"] / self.period[1],
               self.particle_fields["particle_position_z"] / self.period[2],
               self.link)
        self.densities = np.ones(self.tags.size, dtype='float64') * -1
        self.particle_fields["densities"] = self.densities
        self.particle_fields["tags"] = self.tags

class GenericHaloFinder(HaloList, ParallelAnalysisInterface):
    def __init__(self, ds, data_source, padding=0.0, ptype='all'):
        ParallelAnalysisInterface.__init__(self)
        self.ds = ds
        self.index = ds.index
        self.center = (np.array(data_source.right_edge) +
                       np.array(data_source.left_edge)) / 2.0
        self.ptype = ptype

    def _parse_halolist(self, threshold_adjustment):
        groups = []
        max_dens = {}
        hi = 0
        LE, RE = self.bounds
        for halo in self._groups:
            this_max_dens = halo.maximum_density_location()
            # if the most dense particle is in the box, keep it
            if np.all((this_max_dens >= LE) & (this_max_dens <= RE)):
                # Now we add the halo information to OURSELVES, taken from the
                # self.hop_list
                # We need to mock up the HOPHaloList thingie, so we need to
                #     set self._max_dens
                max_dens_temp = list(self._max_dens[halo.id])[0] / \
                    threshold_adjustment
                max_dens[hi] = [max_dens_temp] + \
                    list(self._max_dens[halo.id])[1:4]
                groups.append(self._halo_class(self, hi, ptype=self.ptype))
                groups[-1].indices = halo.indices
                self.comm.claim_object(groups[-1])
                hi += 1
        del self._groups, self._max_dens  # explicit >> implicit
        self._groups = groups
        self._max_dens = max_dens

    def _join_halolists(self):
        groups = {self.comm.rank: len(self)}
        groups = self.comm.par_combine_object(
            groups, datatype='dict', op='join')

        ngroups = np.array([groups[rank] for rank in sorted(groups)])
        offsets = ngroups.cumsum() - ngroups
        my_offset = offsets[self.comm.rank]
        for halo in self:
            halo.id += my_offset

    def _reposition_particles(self, bounds):
        # This only does periodicity.  We do NOT want to deal with anything
        # else.  The only reason we even do periodicity is the
        LE, RE = bounds
        dw = self.ds.domain_right_edge - self.ds.domain_left_edge
        for i, ax in enumerate('xyz'):
            arr = self._data_source[self.ptype, "particle_position_%s" % ax]
            arr[arr < LE[i] - self.padding] += dw[i]
            arr[arr > RE[i] + self.padding] -= dw[i]

class HOPHaloFinder(GenericHaloFinder, HOPHaloList):
    r"""HOP halo finder.

    Halos are built by:
    1. Calculating a density for each particle based on a smoothing kernel.
    2. Recursively linking particles to other particles from lower density
    particles to higher.
    3. Geometrically proximate chains are identified and
    4. merged into final halos following merging rules.

    Lower thresholds generally produce more halos, and the largest halos
    become larger. Also, halos become more filamentary and over-connected.

    Eisenstein and Hut. "HOP: A New Group-Finding Algorithm for N-Body
    Simulations." ApJ (1998) vol. 498 pp. 137-142

    Parameters
    ----------
    ds : `Dataset`
        The dataset on which halo finding will be conducted.
    subvolume : `yt.data_objects.data_containers.YTSelectionContainer`, optional
        A region over which HOP will be run, which can be used to run HOP
        on a subvolume of the full volume. Default = None, which defaults
        to the full volume automatically.
    threshold : float
        The density threshold used when building halos. Default = 160.0.
    ptype : string
        The particle type to be used for halo finding.
        Default: 'all'.
    padding : float
        When run in parallel, the finder needs to surround each subvolume
        with duplicated particles for halo finidng to work. This number
        must be no smaller than the radius of the largest halo in the box
        in code units. Default = 0.02.
    total_mass : float
        If HOP is run on the same dataset mulitple times, the total mass
        of particles in Msun units in the full volume can be supplied here
        to save time.
        This must correspond to the particles being operated on, meaning
        if stars are included in the halo finding, they must be included
        in this mass as well, and visa-versa.
        If halo finding on a subvolume, this still corresponds with the
        mass in the entire volume.
        Default = None, which means the total mass is automatically
        calculated.
    save_particles : bool
        If True, output member particles for each halo.
        Default: True.

    Examples
    --------
    >>> import yt
    >>> from yt.extensions.astro_analysis.halo_analysis import HaloCatalog
    >>> data_ds = yt.load('Enzo_64/RD0006/RedshiftOutput0006')
    >>> hc = HaloCatalog(data_ds=data_ds, finder_method='hop',
    ...                  finder_kwargs={"threshold": 160})
    >>> hc.create()
    """
    def __init__(self, ds, subvolume=None, threshold=160, dm_only=False,
                 ptype='all', padding=0.02, total_mass=None,
                 save_particles=True):
        if subvolume is not None:
            ds_LE = np.array(subvolume.left_edge)
            ds_RE = np.array(subvolume.right_edge)
        self.period = ds.domain_right_edge - ds.domain_left_edge
        self.save_particles = save_particles
        self._data_source = ds.all_data()
        GenericHaloFinder.__init__(self, ds, self._data_source, padding,
                                   ptype=ptype)
        # do it once with no padding so the total_mass is correct
        # (no duplicated particles), and on the entire volume, even if only
        # a small part is actually going to be used.
        self.padding = 0.0
        padded, LE, RE, self._data_source = \
            self.partition_index_3d(ds=self._data_source,
                padding=self.padding)

        if dm_only:
            raise RuntimeError(
                "dm_only has been removed. " +
                "Use ptype to specify a particle type, instead.")

        self.ptype = ptype

        # For scaling the threshold, note that it's a passthrough
        if total_mass is None:
            total_mass = self.comm.mpi_allreduce(
                self._data_source.quantities.total_quantity(
                    (self.ptype, "particle_mass")).in_units('Msun'), op='sum')
        # MJT: Note that instead of this, if we are assuming that the particles
        # are all on different processors, we should instead construct an
        # object representing the entire domain and sum it "lazily" with
        # Derived Quantities.
        if subvolume is not None:
            self._data_source = ds.region([0.] * 3, ds_LE, ds_RE)
        else:
            self._data_source = ds.all_data()
        self.padding = padding  # * ds["unitary"] # This should be clevererer
        padded, LE, RE, self._data_source = \
            self.partition_index_3d(ds=self._data_source,
            padding=self.padding)
        self.bounds = (LE, RE)
        # sub_mass can be skipped if subvolume is not used and this is not
        # parallel.
        if subvolume is None and \
                ytcfg.get("yt", "internals", "topcomm_parallel_size") == 1:
            sub_mass = total_mass
        else:
            sub_mass = \
                self._data_source.quantities.total_quantity(
                    (self.ptype, "particle_mass")).in_units('Msun')
        HOPHaloList.__init__(self, self._data_source,
            threshold * total_mass / sub_mass, ptype=self.ptype)
        self._parse_halolist(total_mass / sub_mass)
        self._join_halolists()

class FOFHaloFinder(GenericHaloFinder, FOFHaloList):
    r"""Friends-of-friends halo finder.

    Halos are found by linking together all pairs of particles closer than
    some distance from each other. Particles may have multiple links,
    and halos are found by recursively linking together all such pairs.

    Larger linking lengths produce more halos, and the largest halos
    become larger. Also, halos become more filamentary and over-connected.

    Davis et al. "The evolution of large-scale structure in a universe
    dominated by cold dark matter." ApJ (1985) vol. 292 pp. 371-394

    Parameters
    ----------
    ds : `Dataset`
        The dataset on which halo finding will be conducted.
    subvolume : `yt.data_objects.data_containers.YTSelectionContainer`, optional
        A region over which HOP will be run, which can be used to run HOP
        on a subvolume of the full volume. Default = None, which defaults
        to the full volume automatically.
    link : float
        If positive, the interparticle distance (compared to the overall
        average) used to build the halos. If negative, this is taken to be
        the *actual* linking length, and no other calculations will be
        applied.  Default = 0.2.
    ptype : string
        The type of particle to be used for halo finding.
        Default: 'all'.
    padding : float
        When run in parallel, the finder needs to surround each subvolume
        with duplicated particles for halo finidng to work. This number
        must be no smaller than the radius of the largest halo in the box
        in code units. Default = 0.02.
    save_particles : bool
        If True, output member particles for each halo.
        Default: True.

    Examples
    --------
    >>> import yt
    >>> from yt.extensions.astro_analysis.halo_analysis import HaloCatalog
    >>> data_ds = yt.load('Enzo_64/RD0006/RedshiftOutput0006')
    >>> hc = HaloCatalog(data_ds=data_ds, finder_method='fof',
    ...                  finder_kwargs={"link": 0.2})
    >>> hc.create()
    """
    def __init__(self, ds, subvolume=None, link=0.2, dm_only=False,
                 ptype='all', padding=0.02, save_particles=True):
        if subvolume is not None:
            ds_LE = np.array(subvolume.left_edge)
            ds_RE = np.array(subvolume.right_edge)
        self.period = ds.domain_right_edge - ds.domain_left_edge
        self.ds = ds
        self.index = ds.index
        self.redshift = ds.current_redshift
        self.save_particles = save_particles
        self._data_source = ds.all_data()
        GenericHaloFinder.__init__(self, ds, self._data_source, padding)
        self.padding = 0.0  # * ds["unitary"] # This should be clevererer
        # get the total number of particles across all procs, with no padding
        padded, LE, RE, self._data_source = \
            self.partition_index_3d(ds=self._data_source,
            padding=self.padding)

        if dm_only:
            raise RuntimeError(
                "dm_only has been removed. " +
                "Use ptype to specify a particle type, instead.")

        self.ptype = ptype

        if link > 0.0:
            n_parts = self.comm.mpi_allreduce(
                self._data_source["particle_position_x"].size, op='sum')
            # get the average spacing between particles
            #l = ds.domain_right_edge - ds.domain_left_edge
            #vol = l[0] * l[1] * l[2]
            # Because we are now allowing for datasets with non 1-periodicity,
            # but symmetric, vol is always 1.
            vol = 1.
            avg_spacing = (float(vol) / n_parts) ** (1. / 3.)
            linking_length = link * avg_spacing
        else:
            linking_length = np.abs(link)
        self.padding = padding
        if subvolume is not None:
            self._data_source = ds.region([0.] * 3, ds_LE,
                ds_RE)
        else:
            self._data_source = ds.all_data()
        padded, LE, RE, self._data_source = \
            self.partition_index_3d(ds=self._data_source,
            padding=self.padding)
        self.bounds = (LE, RE)
        # reflect particles around the periodic boundary
        #self._reposition_particles((LE, RE))
        # here is where the FOF halo finder is run
        mylog.info("Using a linking length of %0.3e", linking_length)
        FOFHaloList.__init__(self, self._data_source, linking_length,
                             redshift=self.redshift, ptype=self.ptype)
        self._parse_halolist(1.)
        self._join_halolists()


import numpy as np
import xarray as xr
import multiprocessing
from scipy.sparse import csr_matrix
from pathlib import Path
from apexpy import Apex
from datetime import datetime, timedelta
import logging
import warnings
try:
    import cvxpy as cp
    import bottleneck as bn
    import pandas
    from skimage import measure, morphology
    from sklearn.metrics.pairwise import rbf_kernel
except ImportError as imp_err:
    warnings.warn(f"Packages required for recreating dataset not installed: {imp_err}")

from trough import config, utils, _tec, _arb

logger = logging.getLogger(__name__)

def get_model(tec_data, hemisphere, omni_file):
    """Get magnetic latitudes of the trough according to the model in Deminov 2017
    for a specific time and set of magnetic local times.
    """
    logger.info("getting model")
    omni_data = xr.open_dataset(omni_file)
    logger.info(f"{omni_data.time.values[0]=} {omni_data.time.values[-1]=}")
    kp = _get_weighted_kp(tec_data.time, omni_data)
    logger.info(f"{kp.shape=}")
    apex = Apex(date=utils.datetime64_to_datetime(tec_data.time.values[0]))
    mlat = 65.5 * np.ones((tec_data.time.shape[0], tec_data.mlt.shape[0]))
    for i in range(10):
        glat, glon = apex.convert(mlat, tec_data.mlt.values[None, :], 'mlt', 'geo', 350, tec_data.time.values[:, None])
        mlat = _model_subroutine_lat(tec_data.mlt.values[None, :], glon, kp[:, None], hemisphere)
    if hemisphere == 'south':
        mlat = mlat * -1
    tec_data['model'] = xr.DataArray(
        mlat,
        coords={'time': tec_data.time, 'mlt': tec_data.mlt},
        dims=['time', 'mlt']
    )

def _model_subroutine_lat(mlt, glon, kp, hemisphere):
    """Get's model output mlat given MLT, geographic lon and weighted kp

    Parameters
    ----------
    mlt: numpy.ndarray (n_mlt, )
    glon: numpy.ndarray (n_mlt, )
    kp: float

    Returns
    -------
    mlat: numpy.ndarray (n_t, n_mlt)
    """
    phi_t = 3.16 - 5.6 * np.cos(np.deg2rad(15 * (mlt - 2.4))) + 1.4 * np.cos(np.deg2rad(15 * (2 * mlt - .8)))
    if hemisphere == 'north':
        phi_lon = .85 * np.cos(np.deg2rad(glon + 63)) - .52 * np.cos(np.deg2rad(2 * glon + 5))
    elif hemisphere == 'south':
        phi_lon = 1.5 * np.cos(np.deg2rad(glon - 119))
    else:
        raise ValueError(f"Invalid hemisphere: {hemisphere}, valid = ['north', 'south']")
    return 65.5 - 2.4 * kp + phi_t + phi_lon * np.exp(-.3 * kp)

def _get_weighted_kp(times, omni_data, tau=.6, T=10):
    """Get a weighed sum of kp values over time. See paper for details.
    """
    logger.info(f"_get_weighted_kp {times[0]=} {times[-1]=}")
    ap = omni_data.sel(time=slice(times[0] - np.timedelta64(T, 'h'), times[-1]))['ap'].values
    prehistory = np.column_stack([ap[T - i:ap.shape[0] - i] for i in range(T)])
    weight_factors = tau ** np.arange(T)
    ap_tau = np.sum((1 - tau) * prehistory * weight_factors, axis=1)
    return 2.1 * np.log(.2 * ap_tau + 1)

def estimate_background(tec, patch_shape):
    """Use a moving average filter to estimate the background TEC value. `patch_shape` must contain odd numbers

    Parameters
    ----------
    tec: numpy.ndarray[float]
    patch_shape: tuple

    Returns
    -------
    numpy.ndarray[float]
    """
    assert all([2 * (p // 2) + 1 == p for p in patch_shape]), "patch_shape must be all odd numbers"
    patches = utils.extract_patches(tec, patch_shape)
    return bn.nanmean(patches.reshape((tec.shape[0] - patch_shape[0] + 1,) + tec.shape[1:] + (-1,)), axis=-1)

def preprocess_interval(data, min_val=0, max_val=100, bg_est_shape=(1, 15, 15)):
    logger.info("preprocessing interval")
    tec = data['tec'].values
    # throw away outlier values
    tec[tec > max_val] = np.nan
    tec[tec < min_val] = np.nan
    # change to log
    log_tec = np.log10(tec + .001)
    # estimate background
    bg = estimate_background(log_tec, bg_est_shape)
    # subtract background
    x = log_tec - bg
    coords = {'time': data.time, 'mlat': data.mlat, 'mlt': data.mlt}
    data['x'] = xr.DataArray(x, coords=coords, dims=['time', 'mlat', 'mlt'])
    data['tec'] = xr.DataArray(tec, coords=coords, dims=['time', 'mlat', 'mlt'])

def fix_boundaries(labels):
    fixed = labels.copy()
    while True:
        boundary_pairs = np.unique(fixed[:, [0, -1]], axis=0)
        if np.all((boundary_pairs[:, 0] == boundary_pairs[:, 1]) | np.any(boundary_pairs == 0, axis=1)):
            break
        for i in range(boundary_pairs.shape[0]):
            if np.any(boundary_pairs[i] == 0) or boundary_pairs[i, 0] == boundary_pairs[i, 1]:
                continue
            fixed[fixed == boundary_pairs[i, 1]] = boundary_pairs[i, 0]
            break
    return fixed

def remove_auroral(data, hemisphere, offset=3):
    if hemisphere == 'north':
        data['labels'] *= data.mlat < (data['arb'] + offset)
    elif hemisphere == 'south':
        data['labels'] *= data.mlat > (data['arb'] - offset)
    else:
        raise ValueError(f"Invalid hemisphere: {hemisphere}, valid = ['north', 'south']")

def postprocess(data, hemisphere, perimeter_th=50, area_th=1, closing_r=0):
    if closing_r > 0:
        selem = morphology.disk(closing_r, dtype=bool)[:, :, None]
        data['labels'] = np.pad(data['labels'], ((0, 0), (0, 0), (closing_r, closing_r)), 'wrap')
        data['labels'] = morphology.binary_closing(data['labels'], selem)[:, :, closing_r:-closing_r]
    remove_auroral(data, hemisphere)
    for t in range(data.time.shape[0]):
        tmap = data['labels'][t].values
        labeled = measure.label(tmap, connectivity=2)
        labeled = fix_boundaries(labeled)
        props = pandas.DataFrame(measure.regionprops_table(labeled, properties=('label', 'area', 'perimeter')))
        error_mask = (props['perimeter'] < perimeter_th) | (props['area'] < area_th)
        for i, r in props[error_mask].iterrows():
            tmap[labeled == r['label']] = 0
        data['labels'][t] = tmap

def get_rbf_matrix(shape, bandwidth=1):
    X, Y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
    xy = np.column_stack((X.ravel(), Y.ravel()))
    gamma = np.log(2) / bandwidth ** 2
    basis = rbf_kernel(xy, gamma=gamma)
    basis[basis < .01] = 0
    return csr_matrix(basis)

def get_tv_matrix(im_shape, hw=1, vw=1):
    size = im_shape[0] * im_shape[1]

    right = np.eye(size)
    cols = im_shape[1] * (np.arange(size) // im_shape[1]) + (np.arange(size) + 1) % im_shape[1]
    right[np.arange(size), cols] = -1

    left = np.eye(size)
    cols = im_shape[1] * (np.arange(size) // im_shape[1]) + (np.arange(size) - 1) % im_shape[1]
    left[np.arange(size), cols] = -1

    up = np.eye(size)
    cols = np.arange(size) + im_shape[1]
    mask = cols < size
    cols = cols[mask]
    up[np.arange(size)[mask], cols] = -1

    down = np.eye(size)
    cols = np.arange(size) - im_shape[1]
    mask = cols >= 0
    cols = cols[mask]
    down[np.arange(size)[mask], cols] = -1
    return csr_matrix(hw * (right + left) + vw * (up + down))

def get_optimization_args(data, model_weight_max, rbf_bw, tv_hw, tv_vw, l2_weight, tv_weight):
    all_args = []
    # get rbf basis matrix
    basis = get_rbf_matrix((data.mlat.shape[0], data.mlt.shape[0]), rbf_bw)
    # get tv matrix
    tv = get_tv_matrix((data.mlat.shape[0], data.mlt.shape[0]), tv_hw, tv_vw) * tv_weight
    for i in range(data.time.shape[0]):
        # l2 norm cost away from model
        l2 = abs(data.mlat - data['model']).isel(time=i)
        l2 -= l2.min()
        l2 = (model_weight_max - 1) * l2 / l2.max() + 1
        l2 *= l2_weight
        fin_mask = np.isfinite(np.ravel(data['x'].isel(time=i)))
        if not fin_mask.any():
            raise Exception("WHY ALL NAN??")
        args = (cp.Variable(data.mlat.shape[0] * data.mlt.shape[0]), basis[fin_mask, :],
                np.ravel(data['x'].isel(time=i))[fin_mask], tv, np.ravel(l2))
        all_args.append(args)
    return all_args

def run_single(u, basis, x, tv, l2):
    main_cost = u.T @ basis.T @ x
    tv_cost = cp.norm1(tv @ u)
    l2_cost = l2 @ (u ** 2)
    total_cost = main_cost + tv_cost + l2_cost
    prob = cp.Problem(cp.Minimize(total_cost))
    try:
        prob.solve(solver=cp.GUROBI)
    except Exception as e:
        logger.info(f"FAILED, USING ECOS: {e}")
        prob.solve(solver=cp.ECOS)
    return u.value

def run_multiple(args, parallel=True):
    if parallel:
        with multiprocessing.Pool(processes=4) as p:
            results = p.starmap(run_single, args)
    else:
        results = []
        for arg in args:
            results.append(run_single(*arg))
    return np.stack(results, axis=0)

def label_trough_interval(start_date, end_date, params, hemisphere, tec_dir, arb_dir, omni_file):
    logger.info(f"labeling trough interval: {start_date=} {end_date=}")
    data = _tec.get_tec_data(start_date, end_date, hemisphere, tec_dir).to_dataset(name='tec')
    preprocess_interval(data, bg_est_shape=params.bg_est_shape)

    data['arb'] = _arb.get_arb_data(start_date, end_date, hemisphere, arb_dir)
    get_model(data, hemisphere, omni_file)
    args = get_optimization_args(data, params.model_weight_max, params.rbf_bw, params.tv_hw, params.tv_vw,
                                 params.l2_weight, params.tv_weight)
    logger.info("Running inversion optimization")
    data['score'] = xr.DataArray(
        run_multiple(args).reshape((data.time.shape[0], data.mlat.shape[0], data.mlt.shape[0])),
        coords={
            'time': data.time,
            'mlat': data.mlat,
            'mlt': data.mlt
        },
        dims=['time', 'mlat', 'mlt']
    )
    # threshold
    data['labels'] = data['score'] >= params.threshold
    # postprocess
    logger.info("Postprocessing inversion results")
    postprocess(data, hemisphere, params.perimeter_th, params.area_th, params.closing_rad)
    return data

def label_trough_dataset(start_date, end_date, params=None, tec_dir=None, arb_dir=None, omni_file=None,
                         output_dir=None):
    if params is None:
        params = config.trough_id_params
    if tec_dir is None:
        tec_dir = config.processed_tec_dir
    if arb_dir is None:
        arb_dir = config.processed_arb_dir
    if omni_file is None:
        omni_file = config.processed_omni_file
    if output_dir is None:
        output_dir = config.processed_labels_dir

    Path(output_dir).mkdir(exist_ok=True, parents=True)
    for year in range(start_date.year, end_date.year + 1):
        for hemisphere in ['north', 'south']:
            labels = []
            scores = []
            start = datetime(year, 1, 1, 0, 0)
            while start.year < year + 1:
                end = start + timedelta(days=1)
                if start >= end_date or end <= start_date:
                    start += timedelta(days=1)
                    continue
                start = max(start_date, start)
                end = min(end_date, end)
                data = label_trough_interval(start, end, params, hemisphere, tec_dir, arb_dir, omni_file)
                if end.year == start.year + 1:
                    data = data.isel(time=slice(0, -1))
                labels.append(data['labels'])
                scores.append(data['score'])
                start += timedelta(days=1)
            labels = xr.concat(labels, 'time')
            scores = xr.concat(scores, 'time')
            labels.to_netcdf(Path(output_dir) / f"labels_{hemisphere}_{year:04d}.nc")
            scores.to_netcdf(Path(output_dir) / f"scores_{hemisphere}_{year:04d}.nc")

def get_label_paths(start_date, end_date, hemisphere, processed_dir):
    file_dates = np.arange(
        np.datetime64(start_date, 'Y'),
        (np.datetime64(end_date, 's')).astype('datetime64[Y]') + 1,
        np.timedelta64(1, 'Y')
    )
    file_dates = utils.decompose_datetime64(file_dates)
    return [Path(processed_dir) / f"labels_{hemisphere}_{d[0]:04d}.nc" for d in file_dates]

def get_trough_labels(start_date, end_date, hemisphere, labels_dir=None):
    if labels_dir is None:
        labels_dir = config.processed_labels_dir
    data = utils.read_netcdfs(get_label_paths(start_date, end_date, hemisphere, labels_dir), 'time')
    return data.sel(time=slice(start_date, end_date))

def get_data(start_date, end_date, hemisphere, tec_dir=None, omni_file=None, labels_dir=None):
    if tec_dir is None:
        tec_dir = config.processed_tec_dir
    if omni_file is None:
        omni_file = config.processed_omni_file
    if labels_dir is None:
        labels_dir = config.processed_labels_dir
    data = xr.open_dataset(omni_file).sel(time=slice(start_date, end_date))
    data['tec'] = _tec.get_tec_data(start_date, end_date, hemisphere, tec_dir)
    data['labels'] = get_trough_labels(start_date, end_date, hemisphere, labels_dir)
    return data


import argparse
import h5sparse as h5py
from h5py import SoftLink
import os
import numpy as np
import blockmodel
from math import sqrt
from read_matrix import body_matrix, surface_matrix, revert_h_weighting, revert_c_weighting
from scipy.sparse import csr_matrix, diags, vstack
from scipy.sparse.linalg import svds
import pandas as pd
from time import time
from datetime import date

# kernel object
class Kernel():
    ## initialized with filegroup name & k value
    def __init__(self, params):
        self.filegroup = params.f
        self.num_of_vectors = params.k
        self.epsilon = params.epsilon
        self.regularization = params.regular
        self.is_coverage_to_weight = params.c
        self.is_thickness_to_weight = params.z
        self.is_smoother_to_apply = params.smooth
        self.is_loaded = False
        self.smoother = False

    ## load kernel results
    def load_kernel(self, db):
        if f"/{self.filegroup}" in db:
            for item in db[f"/{self.filegroup}"]:
                if f"/{self.filegroup}/{item}/kernel" in db:
                    attrs = db[f"/{self.filegroup}/{item}/kernel"].attrs
                    if attrs['thickness_weighted'] == self.is_thickness_to_weight:
                        if (not self.is_coverage_to_weight and attrs['coverage_weighted'] == -1) or (self.is_coverage_to_weight and attrs['coverage_weighted'] == self.regularization):
                            if not 'smoother_added' in attrs: attrs['smoother_added'] = False
                            if attrs['smoother_added'] == self.is_smoother_to_apply:
                                self.csr = db[f"/{self.filegroup}/{item}/kernel"][()]
                                self.d = db[f"/{self.filegroup}/{item}/residual"]
                                self.smoother = attrs['smoother_added']
                                if self.is_coverage_to_weight: self.c = db[f"/{self.filegroup}/{item}/weight"]
                                if not item == self.num_of_vectors:
                                    if not f"/{self.filegroup}/{self.num_of_vectors}/kernel" in db: db[f"/{self.filegroup}/{self.num_of_vectors}/kernel"] = SoftLink(f"/{self.filegroup}/{item}/kernel")
                                    if not f"/{self.filegroup}/{self.num_of_vectors}/residual" in db: db[f"/{self.filegroup}/{self.num_of_vectors}/residual"] = SoftLink(f"/{self.filegroup}/{item}/residual")
                                    if not f"/{self.filegroup}/{self.num_of_vectors}/weight" in db and self.is_coverage_to_weight: db[f"/{self.filegroup}/{self.num_of_vectors}/weight"] = SoftLink(f"/{self.filegroup}/{item}/weight")
                                self.is_loaded = True
                                break

    ## save kernel results
    def save_kernel(self, db):
        if not f"/{self.filegroup}" in db: db.create_group(f"/{self.filegroup}")
        if not f"/{self.filegroup}/{self.num_of_vectors}" in db: db.create_group(f"/{self.filegroup}/{self.num_of_vectors}")

        kernel_data = overwrite_data(db, f"/{self.filegroup}/{self.num_of_vectors}/kernel", self.csr)
        kernel_data.attrs['created'] = date.today().strftime('%Y-%m-%d')
        kernel_data.attrs['error_weighted'] = True
        kernel_data.attrs['smoother_added'] = self.smoother
        kernel_data.attrs['thickness_weighted'] = self.is_thickness_to_weight
        kernel_data.attrs['coverage_weighted'] = self.regularization if self.is_coverage_to_weight else -1

        if self.is_coverage_to_weight:
            weight_data = overwrite_data(db, f"/{self.filegroup}/{self.num_of_vectors}/weight", self.c)
            weight_data.attrs['created'] = date.today().strftime('%Y-%m-%d')
            weight_data.attrs['smoother_added'] = self.smoother

        residual_data = overwrite_data(db, f"/{self.filegroup}/{self.num_of_vectors}/residual", self.d)
        residual_data.attrs['created'] = date.today().strftime('%Y-%m-%d')
        residual_data.attrs['error_weighted'] = True
        residual_data.attrs['smoother_added'] = self.smoother

    ## Load existing data matrix
    def readmatrix(self):
        if self.filegroup == 'Sfast':
            filenames = ['ScS-S.4.05', 'SS-S.4']
            is_surface_to_add = True
        elif self.filegroup == 'Surf':
            filenames = []
            is_surface_to_add = True
        elif self.filegroup == 'Scomb':
            filenames = ['Scomb.4']
            is_surface_to_add = False
        elif self.filegroup == 'ScombSurf':
            filenames = ['Scomb.4']
            is_surface_to_add = True
        elif self.filegroup == 'Sonly':
            filenames = ['Scomb.4', 'SS.95-05.4', 'ScS-S.4.05', 'SS-S.4']
            is_surface_to_add = False
        elif self.filegroup == 'S':
            filenames = ['Scomb.4', 'SS.95-05.4', 'ScS-S.4.05', 'SS-S.4']
            is_surface_to_add = True
        elif self.filegroup == 'P':
            filenames = ['Pcomb.4', 'PP.95-05.4',  'PP-P.4']
            is_surface_to_add = False
        else:
            raise ValueError('unknown file combination %s' % self.filegroup)

        data = ()
        kernels = ()
        for filename in filenames:
            (d, K) = body_matrix(filename)
            data = data + (d,)
            kernels = kernels + (K,)
        if is_surface_to_add:
            fns = ['mat4.nd.r'+str(i)+'.swp' for i in range(4,15+1)]
            fns += ['mat4.nd.l'+str(i)+'.swp' for i in range(4,15+1)]
            (d, K) = surface_matrix(fns)
            data = data + (d,)
            kernels = kernels + (K,)

        return np.vstack(data), np.vstack(kernels)

    ## collect and weight kernel matrices
    def get_kernel(self):
        from config import H
        if not self.is_thickness_to_weight:
            H = np.ones(len(H))
        (self.d, K) = self.readmatrix()

        # Apply coverage weighting column by column (save memory)
        c = np.zeros(K.shape[1])
        for i in range(len(c)):
            c[i] = H[i//2578]
        if self.is_coverage_to_weight:
            for i in range(len(c)):
                if self.regularization == 0:
                    c[i] *= max(1, np.sum(K[:,i]!=0))
                elif self.regularization == 1:
                    c[i] *= max(1, np.sum(np.abs(K[:,i]), axis=0))
                elif self.regularization == 2:
                    c[i] *= max(1, np.sqrt(np.sum(K[:,i]**self.regularization, axis=0)))
                elif self.regularization > 2:
                    c[i] *= max(1, np.sqrt(np.sum(np.abs(K[:,i])**self.regularization, axis=0)))
                else:
                    raise ValueError('incorrect regularization norm: L-%s' % self.regularization)
                K[:,i] = K[:,i]/c[i]
        self.c = c

        # Compress K into scipy sparse matrix (csr format)
        self.csr = csr_matrix(K,dtype=np.float32)
        self.is_loaded = True

    ## Apply first differnce smoother
    def first_diffence(self, db, radial_lambda=5., lateral_lambda=5.):
        if self.smoother: return
        m = blockmodel.Model()
        dwt = self.c[...] ** 2

        # Insert extra rows to sensitivity matrix
        data = []
        row = []
        col = []

        # Insert extra elements to traveltime residual
        (mshell, nshell) = (2578, 18)

        # Radical smoother
        for i in range(mshell):
            for j in range(nshell-1):
                lower_block_num = i + j * mshell
                upper_block_num = lower_block_num + mshell

                # insert flam for lower block and -flam for upper to matrix
                newrow = row[-1] + 1 if len(row) > 0 else 0
                row.append(newrow)
                col.append(lower_block_num)
                data.append(radial_lambda / self.c[lower_block_num])
                dwt[lower_block_num] += (radial_lambda / self.c[lower_block_num]) ** 2

                row.append(newrow)
                col.append(upper_block_num)
                data.append(-radial_lambda / self.c[upper_block_num])
                dwt[upper_block_num] += (radial_lambda / self.c[upper_block_num]) ** 2

                # insert 0 residual to self.d
                self.d = np.append(self.d, 0)

        # Lateral smoother
        for this_block in range(nshell*mshell):
            # get surrounding blocks
            west_neigbor = m[this_block].neighbor('W').id
            east_neigbor = m[this_block].neighbor('E').id
            south_neighbor = m[this_block].neighbor('S').id
            north_neighbor = m[this_block].neighbor('N').id

            # insert 4*glam for center block and -glam for ambient to matrix
            newrow = row[-1] + 1 if len(row) > 0 else 0
            row.append(newrow)
            col.append(this_block)
            data.append(4 * lateral_lambda / self.c[this_block])
            dwt[this_block] += (4 * lateral_lambda / self.c[this_block]) ** 2

            row.append(newrow)
            col.append(north_neighbor)
            data.append(-lateral_lambda / self.c[north_neighbor])
            dwt[north_neighbor] += (lateral_lambda / self.c[north_neighbor]) ** 2

            row.append(newrow)
            col.append(south_neighbor)
            data.append(-lateral_lambda / self.c[south_neighbor])
            dwt[south_neighbor] += (lateral_lambda / self.c[south_neighbor]) ** 2

            row.append(newrow)
            col.append(east_neigbor)
            data.append(-lateral_lambda / self.c[east_neigbor])
            dwt[east_neigbor] += (lateral_lambda / self.c[east_neigbor]) ** 2

            row.append(newrow)
            col.append(west_neigbor)
            data.append(-lateral_lambda / self.c[west_neigbor])
            dwt[west_neigbor] += (lateral_lambda / self.c[west_neigbor]) ** 2

            # insert 0 residual to self.d
            self.d = np.append(self.d, 0)

        # compile smoother matrix and stack with csr
        print(f"inserting {len(row)} elements smoother in {row[-1]+1} rows...")
        smoother = csr_matrix((data, (row, col)), shape=(row[-1]+1, mshell*nshell))
        self.csr = vstack([self.csr, smoother])
        self.c = np.sqrt(dwt)

        # toggle on smoother status
        self.smoother = True

    ## decompose and reconstruct
    def svd(self, db):
        t = time()
        U, W, VT = svds(self.csr, k=self.num_of_vectors) if self.num_of_vectors < 46404 else svds(self.csr, k=46403)
        print('SVD with %s vectors finished in %.2f s' % (str(self.num_of_vectors) if self.num_of_vectors < 46404 else 'All', time()-t))

        ### Reshape and save SVD results with timing
        t = time()
        thres = self.epsilon * max(W)
        z_array = []
        for i in range(W.shape[0]):
            if W[i] < thres:
                break
            z_array.append(1/W[i])
        Z = diags(z_array)
        rank = len(z_array)

        V = np.transpose(VT[0:rank, :])
        UT = np.transpose(U[:, 0:rank])

        ### write to database
        svd_v_data = overwrite_data(db, f"/{self.filegroup}/{self.num_of_vectors}/svd_v", V)
        svd_z_data = overwrite_data(db, f"/{self.filegroup}/{self.num_of_vectors}/svd_z", z_array)
        svd_ut_data = overwrite_data(db, f"/{self.filegroup}/{self.num_of_vectors}/svd_ut", UT)
        svd_z_data.attrs['epsilon'] = self.epsilon

        t = time()
        m = V @ (Z @ (UT @ self.d))
        if self.is_thickness_to_weight: m = revert_h_weighting(m)
        if self.is_coverage_to_weight: m = revert_c_weighting(m, self.c)
        s = pd.Series(m.reshape(len(m)))
        model = overwrite_data(db, f"/{self.filegroup}/{self.num_of_vectors}/model", m)
        print(s.describe())
        print('inversion model calculated from SVD in %.2f s' % (time()-t))

# parse input parameters
def create_parser():
    parser = argparse.ArgumentParser(prog="Tomography Inversion Tools")
    parser.add_argument("f", help="filegroup")
    parser.add_argument("-k", help="truncated size", type=int, default=6000)
    parser.add_argument("--epsilon", "-e", help="neglecting threshold", type=float, default=1e-4)
    parser.add_argument("--regular", "-l", help="power magnitude for regularization", type=int, default=2)
    parser.add_argument("-c", help="coverage weighting", action='store_true')
    parser.add_argument("-z", help="thickness weighting", action='store_true')
    parser.add_argument("--smooth", help="apply smoother", action='store_true')
    parser.add_argument("--overwrite", "-o", help="overwrite kernel", action='store_true')
    parser.add_argument("--dispose", "-d", help="not to save kernel", action='store_true')
    return parser

def create_params(f, **kwargs):
    parser = create_parser()
    args = []
    args.append(f)
    for key in kwargs.keys():
        if not kwargs[key] is False:
            args.append(f"--{key}" if len(key) > 1 else f"-{key}") # note this is only compatible with Python >= 3.6
            if not type(kwargs[key]) is bool:args.append(str(kwargs[key]))
    return parser.parse_args(args)

# main program
def inversion(params, output_dir='.'):
    ## prepare kernel data
    db = open_hdf5(f"{output_dir}/database.hdf")
    data = Kernel(params)

    if not params.overwrite:
        print("try to load sensitivity matrix if possible...")
        data.load_kernel(db)

    if not data.is_loaded:
        if not params.overwrite: print("no sensitivity matrix availible")
        print("generating sensitivity matrix from data...")
        data.get_kernel()
        ## add smoother to kernel matrix
        if params.smooth: data.first_diffence(db, radial_lambda=5, lateral_lambda=5)
        if not params.dispose:
            data.save_kernel(db)
            print("sensitivity matrix saved.")
    else:
        print("sensitivity matrix loaded.")

    print(f"sensitivity matrix shape: {data.csr.shape}")
    ## implement svd
    data.svd(db)
    db.close()

# open database
def open_hdf5(filename):
    os.environ['HDF5_USE_FILE_LOCKING'] = 'FALSE'
    try:
        db = h5py.File(filename, mode='a')
    except:
        db = h5py.File(filename, mode='w')
    return db

# handle database write
def overwrite_data(db, path, newdata):
    if path in db: del db[path]
    data = db.create_dataset(path, data=newdata)
    return data

# run as main
if __name__ == '__main__':
    parser = create_parser()
    params = parser.parse_args()
    inversion(params)


# This file is part of PSL-Python.
# Copyright (c) 2021, Eijiro Shibusawa <phd_kimberlite@yahoo.co.jp>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import numpy as np
import cv2

from colormap import colormap as jet
from pixel_cost import absolute_difference as ad_cost
from pixel_cost import zero_mean_normalized_cross_correlation as zncc_cost
from pixel_cost import zero_mean_absolute_difference as zsad_cost
from pixel_cost import normalized_cross_correlation as ncc_cost

def depth_to_colormap(depth, min_z, max_z):
    mask = depth > 0
    delta = 1/min_z - 1/max_z
    index = np.round(np.maximum(0, np.minimum(1/(depth + 1E-6) - 1/max_z, delta) / delta) * 255).astype(np.int)
    img = (jet[index] * 255).astype(np.uint8)
    img[~mask] = (0,0,0)
    return img

class plane_sweep():
    def __init__(self):
        self.scale = 1.0
        self.near_z = 0.4
        self.far_z = 5000
        self.match_window_width = 7
        self.match_window_height = 7
        self.num_planes = 256
        self.near_y = 0.4
        self.far_y = 1.5
        self.num_ground = 16
        self.plane_generation_mode_ = {'uniform_depth':0, 'uniform_disparity':1}
        self.plane_generation_mode = self.plane_generation_mode_['uniform_disparity']
        self.matching_costs_ = {'sad':0, 'zncc':1, 'zsad':2, 'ncc':3}
        self.matching_costs = self.matching_costs_['sad']
        self.matching_gpu_batch_enabled = False
        self.sub_pixel_interpolation_mode_ = {'direct':0, 'inverse':1}
        self.sub_pixel_interpolation_mode = self.sub_pixel_interpolation_mode_['inverse']
        self.ouput_cost_volume_enabled = False
        self.sub_pixel_enabled = True
        self.box_filer_ad_enabled = True
        self.write_debug_warping_enabled = False
        self.ground_plane_enabled = False
        self.use_gpu = False
        self.clear_images()
        self.CV = None
        self.rays = None

    @staticmethod
    def generate_depth(near, far, num, mode):
        if mode == 1:
            minD = 1/far
            maxD = 1/near
            dstep = (maxD - minD)/(num - 1)
            depth = 1 / (np.arange(maxD, minD - dstep, -dstep))
        else:
            step = (far - near)/(num - 1)
            depth  = np.arange(near, far + step, step)
        return np.flip(depth)[:num]

    def generate_planes(self):
        # front parallel
        fp_planes = np.zeros((4, self.num_planes), dtype=np.float)
        fp_planes[2, :] = -1
        fp_planes_has_neighbor = np.full(self.num_planes, True, dtype=np.bool)
        fp_planes_has_neighbor[0] = fp_planes_has_neighbor[-1] = False
        fp_planes[3, :] = plane_sweep.generate_depth(self.near_z, self.far_z, self.num_planes, self.plane_generation_mode)
        self.planes = fp_planes
        self.planes_has_neighbor = fp_planes_has_neighbor

        # ground
        if self.ground_plane_enabled:
            g_planes = np.zeros((4, self.num_ground), dtype=np.float)
            g_planes[1, :] = -1
            g_planes[3, :] = plane_sweep.generate_depth(self.near_y, self.far_y, self.num_ground, self.plane_generation_mode)
            g_planes_has_neighbor = np.full(self.num_ground, True, dtype=np.bool)
            g_planes_has_neighbor[0] = g_planes_has_neighbor[-1] = False
            self.planes = np.hstack((self.planes, g_planes))
            self.planes_has_neighbor = np.hstack((self.planes_has_neighbor, g_planes_has_neighbor))

    def transform_planes(self, R, t):
        self.Hs = np.empty((self.planes.shape[1], 3, 3), dtype=self.planes.dtype)
        for k in range(0, self.Hs.shape[0]):
            p = self.planes[:,k]
            n = (p[0:3][:,np.newaxis])
            self.Hs[k] = R + np.dot(t, n.T)/p[3]

    def add_image(self, cam, R, t, img):
        # Xc = Rc*X - Tc
        self.cams.append(cam)
        self.Rs.append(R)
        self.Ts.append(t)
        self.imgs.append(img)

    def clear_images(self):
        self.imgs = list()
        self.cams = list()
        self.Rs = list()
        self.Ts = list()

    def get_cost_volume(self, ref):
        self.generate_planes()

        n_planes = self.planes.shape[1]
        n_imgs = len(self.imgs)
        sz_img = self.imgs[ref].shape

        accumlation_scale = 1.0 / (n_imgs - 1)
        if self.matching_costs == 1:
            cost_function = zncc_cost(self.imgs[ref], n_planes, (self.match_window_width, self.match_window_height), scale=accumlation_scale)
        elif self.matching_costs == 2:
            cost_function = zsad_cost(self.imgs[ref], n_planes, (self.match_window_width, self.match_window_height), self.box_filer_before_occlusion_enabled, scale=accumlation_scale)
        elif self.matching_costs == 3:
            cost_function = ncc_cost(self.imgs[ref], n_planes, (self.match_window_width, self.match_window_height), scale=accumlation_scale)
        else:
            cost_function = ad_cost(self.imgs[ref], n_planes, (self.match_window_width, self.match_window_height), self.box_filer_ad_enabled, scale=accumlation_scale)

        if n_imgs == 2: # stereo case
            self.mask = np.zeros((sz_img[0], sz_img[1], n_planes), np.bool)

        rays = self.cams[ref].unproject_rays(sz_img)
        for k in range(0, n_imgs):
            if k == ref:
                continue
            # relative pose
            R = np.dot(self.Rs[k], self.Rs[ref].T)
            t = np.dot(R, self.Ts[ref]) - self.Ts[k]
            self.transform_planes(R, t)
            img_other = self.imgs[k]
            cam_other = self.cams[k]
            for l in range(0, n_planes):
                xyz = np.dot(self.Hs[l], rays)
                xy = cam_other.project(xyz).astype(np.float32)
                warp = cv2.remap(img_other, xy[0].reshape(sz_img), xy[1].reshape(sz_img), cv2.INTER_LINEAR, None, cv2.BORDER_WRAP, 255)
                cost_function.accumulate(warp, l)
                if n_imgs == 2: # stereo case
                    mask = (xy[0,:] >= 0) & (xy[1,:] >= 0) & (xy[0,:] < sz_img[1]) & (xy[1,:] < sz_img[0])
                    mask = mask.reshape(sz_img)
                    self.mask[:,:,l] = mask

                if self.write_debug_warping_enabled:
                    cv2.imwrite('debug_warping_{:04d}_{:04d}.png'.format(k, l), warp)

        self.CV = cost_function.get_cost_volume()
        self.rays = rays

    def get_depth_from_planes(self, indices):
        if (self.planes is None) or (self.rays is None):
            return None

        xy = np.empty((2, self.rays.shape[1]), dtype=np.float)
        xy[0] = self.rays[0,:] / self.rays[2,:]
        xy[1] = self.rays[1,:] / self.rays[2,:]

        indices_0 = indices.reshape(-1)
        plane_depth = self.planes[3,indices_0]
        plane_dot = xy[0] * self.planes[0, indices_0] + xy[1] * self.planes[1, indices_0] + self.planes[2, indices_0]

        if self.sub_pixel_enabled:
            indices_mask = self.planes_has_neighbor[indices]
            indices_m = np.copy(indices)
            indices_m[indices_mask] += 1
            indices_p = np.copy(indices)
            indices_p[indices_mask] -= 1
            cost_0 = np.take_along_axis(self.CV, indices[:,:,np.newaxis], 2).reshape(-1)
            cost_p = np.take_along_axis(self.CV, indices_p[:,:,np.newaxis], 2).reshape(-1)
            cost_m = np.take_along_axis(self.CV, indices_m[:,:,np.newaxis], 2).reshape(-1)
            offset_denom = cost_p + cost_m - (cost_0 * 2)
            offset_denom_mask = np.abs(offset_denom) < 1E-5
            offset_denom[offset_denom_mask] = 1
            offset = (cost_m - cost_p)/(2 * offset_denom)
            offset[offset_denom_mask] = 0

            if self.sub_pixel_interpolation_mode == 0:
                step_p = self.planes[3, indices_p.reshape(-1)] - self.planes[3, indices_0]
                step_m = self.planes[3, indices_0] - self.planes[3, indices_m.reshape(-1)]
                depth_sub = np.copy(offset)
                depth_sub[offset < 0] = (offset*step_m)[offset < 0]
                depth_sub[offset > 0] = (offset*step_p)[offset > 0]
                D = -(plane_depth + depth_sub) / plane_dot
            else:
                step_p = 1/(self.planes[3, indices_0]) - 1/(self.planes[3, indices_p.reshape(-1)])
                step_m = 1/(self.planes[3, indices_m.reshape(-1)]) - 1/(self.planes[3, indices_0])

                depth_sub = np.copy(offset)
                depth_sub[offset < 0] = (offset*step_m)[offset < 0]
                depth_sub[offset > 0] = (offset*step_p)[offset > 0]
                D = -1/(1/plane_depth - depth_sub) / plane_dot
        else:
            D = -plane_depth / plane_dot

        D = D.reshape(indices.shape)
        return D

    def get_depth(self, ref = 0):
        try:
            import torch
            import psl_cuda as py_psl_cuda
        except ImportError:
            self.use_gpu = False

        if not self.use_gpu:
            if self.CV is None:
                self.get_cost_volume(ref)
            best_plane_indices = np.argmin(self.CV, axis=2)
            D = self.get_depth_from_planes(best_plane_indices)
        else:
            D = self.get_depth_gpu(ref)

        return D

    def get_depth_gpu(self, ref = 0):
        if self.sub_pixel_enabled:
            raise ValueError("Subpixel estimation is not supported!")

        try:
            import torch
            import psl_cuda as py_psl_cuda
            from pixel_cost_cuda import absolute_difference as ad_cost_cuda
            from pixel_cost_cuda import zero_mean_normalized_cross_correlation as zncc_cost_cuda
            from pixel_cost_cuda import zero_mean_absolute_difference as zsad_cost_cuda
            from pixel_cost_cuda import normalized_cross_correlation as ncc_cost_cuda
        except ImportError:
            return None

        self.generate_planes()
        n_imgs = len(self.imgs)

        # upload tensor
        if str(self.cams[ref].__class__.__name__) == 'unified_camera':
            # unified camera model has xi parameters
            Ks = [np.array([C.K[0, 0], C.K[1, 1], C.K[0, 2], C.K[1, 2], C.xi])  for C in self.cams]
            t_Ks = torch.tensor(np.array(Ks), dtype=torch.float32, device=torch.device('cpu'))

        t_imgs = torch.tensor(np.array(self.imgs), device=torch.device('cuda'))
        t_Ps = torch.tensor(self.planes.T, dtype=torch.float32, device=torch.device('cuda'))
        t_Rs = torch.tensor(np.array(self.Rs), dtype=torch.float32, device=torch.device('cpu'))
        t_Ts = torch.tensor(np.array(self.Ts), dtype=torch.float32, device=torch.device('cpu'))

        # cost function
        if self.matching_costs == 1:
            cost_function = zncc_cost_cuda(t_imgs[ref], (self.match_window_width, self.match_window_height),
                                self.matching_gpu_batch_enabled)
        elif self.matching_costs == 2:
            cost_function = zsad_cost_cuda(t_imgs[ref], (self.match_window_width, self.match_window_height),
                                self.box_filer_before_occlusion_enabled, self.matching_gpu_batch_enabled)
        elif self.matching_costs == 3:
            cost_function = ncc_cost_cuda(t_imgs[ref], (self.match_window_width, self.match_window_height),
                                self.matching_gpu_batch_enabled)
        else:
            cost_function = ad_cost_cuda(t_imgs[ref], (self.match_window_width, self.match_window_height),
                                self.box_filer_ad_enabled, self.matching_gpu_batch_enabled)

        # compute warped images
        rays_gpu = py_psl_cuda.get_ray_tensor(ref, t_imgs, t_Ks)
        warped_images_gpu = py_psl_cuda.get_warped_image_tensor(ref, t_imgs, t_Ks, t_Rs, t_Ts, rays_gpu, t_Ps)
        if self.write_debug_warping_enabled:
            warped_images = warped_images_gpu.to('cpu').numpy()
            for k in range(0, n_imgs - 1):
                for l in range(0, self.planes.shape[1]):
                    cv2.imwrite('debug_warping_{:04d}_{:04d}.png'.format(k, l), warped_images[k, l])

        # compute cost volume
        CV = cost_function.get_cost_volume(warped_images_gpu)
        del warped_images_gpu
        torch.cuda.empty_cache()

        # WTA
        indices_gpu = torch.argmin(CV, dim = 0)
        D_gpu = py_psl_cuda.get_depth_tensor(rays_gpu, indices_gpu, t_Ps)
        D = D_gpu.to('cpu').numpy()

        if self.ouput_cost_volume_enabled:
            self.CV = CV.to('cpu').numpy().transpose(1, 2, 0)

        del CV, rays_gpu, indices_gpu, D_gpu
        del t_imgs, t_Ps
        torch.cuda.empty_cache()

        return D


from core.macid_base import MACIDBase
from core.macid import MACID
import networkx as nx
from typing import Any, List, Dict, Union
from core.get_paths import directed_decision_free_path, find_all_dir_paths, find_all_undir_paths, get_motif, \
    is_active_indirect_frontdoor_trail, is_active_path
import copy

def _get_key_node(mb: MACIDBase, path: List[str]) -> str:
    """
    Returns the key node of a path (ie the first "fork" node in the path)
    """
    for _, b, _ in zip(path[:-2], path[1:-1], path[2:]):
        structure = get_motif(mb, path, path.index(b))
        if structure == "fork":
            return b
    return ""  # shouldn't happen

def _effective_dir_path_exists(mb: MACIDBase, start: str, finish: str, effective_set: List[str]) -> bool:
    """
    checks whether an effective directed path exists

    """
    start_finish_paths = find_all_dir_paths(mb, start, finish)
    for path in start_finish_paths:
        if _path_is_effective(mb, path, effective_set):
            return True
    else:
        return False

def _effective_undir_path_exists(mb: MACIDBase, start: str, finish: str, effective_set: List[str]) -> bool:
    """
    checks whether an effective undirected path exists
    """
    start_finish_paths = find_all_undir_paths(mb, start, finish)
    for path in start_finish_paths:
        if _path_is_effective(mb, path, effective_set):
            return True
    else:
        return False

def _path_is_effective(mb: MACIDBase, path: List[str], effective_set: List[str]) -> bool:
    """
    checks whether a path is effective
    """
    dec_nodes_in_path = set(mb.all_decision_nodes).intersection(set(path[1:]))  # exclude first node of the path
    all_dec_nodes_effective = all(dec_node in effective_set for dec_node in dec_nodes_in_path)
    # all([]) evaluates to true => this covers case where path has no decision nodes
    if all_dec_nodes_effective:
        return True
    else:
        return False

def _directed_effective_path_not_through_set_y(mb: MACIDBase, start: str, finish: str,
                                               effective_set: List[str], y: List[str] = []) -> bool:
    """
    checks whether a directed effective path exists that doesn't pass through any of the nodes in the set y.
    """
    start_finish_paths = find_all_dir_paths(mb, start, finish)
    for path in start_finish_paths:
        path_not_through_y = set(y).isdisjoint(set(path))
        if _path_is_effective(mb, path, effective_set) and path_not_through_y:
            return True
    else:
        return False

def _effective_backdoor_path_not_blocked_by_set_w(mb: MACIDBase, start: str, finish: str, effective_set: List[str],
                                                  w: List[str] = []) -> List[str]:
    """
    Returns the effective backdoor path not blocked if we condition on nodes in set w.
    If no such path exists, this returns None.
    """
    start_finish_paths: List[List[str]] = find_all_undir_paths(mb, start, finish)
    for path in start_finish_paths:
        is_backdoor_path = path[1] in mb.get_parents(path[0])
        not_blocked_by_w = is_active_path(mb, path, w)
        if is_backdoor_path and _path_is_effective(mb, path, effective_set) and not_blocked_by_w:
            return path
    return []

def _effective_undir_path_not_blocked_by_set_w(mb: MACIDBase, start: str, finish: str,
                                               effective_set: List[str], w: List[str] = []) -> Union[List[str], None]:
    """
    Returns an effective undirected path not blocked if we condition on nodes in set w.
    If no such path exists, this returns None.
    """
    start_finish_paths: List[List[str]] = find_all_undir_paths(mb, start, finish)
    for path in start_finish_paths:
        not_blocked_by_w = is_active_path(mb, path, w)
        if _path_is_effective(mb, path, effective_set) and not_blocked_by_w:
            return path
    else:
        return None

def direct_effect(macid: MACID, decision: str) -> bool:
    """checks to see whether this decision is motivated by a direct effect reasoning patter.
    Graphical Criterion:
    1) There is a directed decision free path from D_A to a utility node U_A
    """
    if decision not in macid.nodes:
        raise Exception(f"{decision} is not present in the macid")

    agent = macid.whose_node[decision]
    agent_utils = macid.utility_nodes_agent[agent]
    for u in agent_utils:
        if directed_decision_free_path(macid, decision, u):
            return True
    else:
        return False

def manipulation(macid: MACID, decision: str, effective_set: List[str]) -> bool:
    """checks to see whether this decision is motivated by an incentive for manipulation
    Graphical Criterion:
    1) There is a directed decision-free path from D_A to an effective decision node D_B.
    2) There is a directed, effective path from D_B to U_A (an effective path is a path in which all
    decision nodes, except possibly the initial node, and except fork nodes, are effective)
    3) There is a directed, effective path from D_A to U_B that does not pass through D_B.
    """
    if decision not in macid.nodes:
        raise Exception(f"{decision} is not present in the macid")

    if not all([node in macid.nodes for node in effective_set]):
        raise Exception("One or many of the nodes in the effective_set are not present in the macid.")

    agent = macid.whose_node[decision]
    agent_utils = macid.utility_nodes_agent[agent]
    reachable_decisions = []    # set of possible D_B
    list_decs = copy.deepcopy(macid.all_decision_nodes)
    list_decs.remove(decision)
    for dec_reach in list_decs:
        if dec_reach in effective_set:
            if directed_decision_free_path(macid, decision, dec_reach):
                reachable_decisions.append(dec_reach)

    for decision_b in reachable_decisions:
        agent_b = macid.whose_node[decision_b]
        agent_b_utils = macid.utility_nodes_agent[agent_b]

        for u in agent_utils:
            if _effective_dir_path_exists(macid, decision_b, u, effective_set):

                for u_b in agent_b_utils:
                    if _directed_effective_path_not_through_set_y(macid, decision, u_b, effective_set, [decision_b]):
                        return True
    else:
        return False

def signaling(macid: MACID, decision: str, effective_set: List[str]) -> bool:
    """checks to see whether this decision is motivated by an incentive for signaling

    Graphical Criterion:
    1) There is a directed decision-free path from D_A to an effective decision node D_B.
    2) There is a directed, effective path from D_B to U_A.
    3) There is an effective back-door path π from D_A to U_B that is not blocked by D_B U W^{D_A}_{D_B}.
    4) If C is the key node in π, there is an effective path from C to U_A that is not blocked by D_A U W^{C}_{D_A}

    """
    if decision not in macid.nodes:
        raise Exception(f"{decision} is not present in the macid")

    if not all([node in macid.nodes for node in effective_set]):
        raise Exception("One or many of the nodes in the effective_set are not present in the macid.")

    agent = macid.whose_node[decision]
    agent_utils = macid.utility_nodes_agent[agent]
    reachable_decisions = []    # set of possible D_B
    list_decs = copy.deepcopy(macid.all_decision_nodes)
    list_decs.remove(decision)
    for dec_reach in list_decs:
        if dec_reach in effective_set:
            if directed_decision_free_path(macid, decision, dec_reach):
                reachable_decisions.append(dec_reach)

    for decision_b in reachable_decisions:
        agent_b = macid.whose_node[decision_b]
        agent_b_utils = macid.utility_nodes_agent[agent_b]
        for u in agent_utils:
            if _effective_dir_path_exists(macid, decision_b, u, effective_set):
                for u_b in agent_b_utils:

                    decision_b_parents_not_desc_decision = [node for node in macid.get_parents(decision_b)
                                                            if node not in set(nx.descendants(macid, decision))]
                    cond_nodes = [decision_b] + decision_b_parents_not_desc_decision

                    if _effective_backdoor_path_not_blocked_by_set_w(macid, decision, u_b, effective_set, cond_nodes):
                        path = _effective_backdoor_path_not_blocked_by_set_w(macid, decision, u_b, effective_set,
                                                                             cond_nodes)
                        if _get_key_node(macid, path):
                            key_node = _get_key_node(macid, path)
                        else:
                            return False
                        decision_parents_not_desc_key_node = [node for node in macid.get_parents(decision)
                                                              if node not in set(nx.descendants(macid, key_node))]
                        cond_nodes2 = [decision] + decision_parents_not_desc_key_node

                        if _effective_undir_path_not_blocked_by_set_w(macid, key_node, u, effective_set, cond_nodes2):
                            return True
    else:
        return False

def revealing_or_denying(macid: MACID, decision: str, effective_set: List[str]) -> bool:
    """checks to see whether this decision is motivated by an incentive for revealing or denying

    Graphical Criterion:
    1) There is a directed decision-free path from D_A to an effective decision node D_B.
    2) There is a direced, effective path from D_B to U_A.
    3) There is an effective indirect front-door path π from D_A to U_B that is not blocked by D_B U W^{D_A}_{D_B}.
    """
    if decision not in macid.nodes:
        raise Exception(f"{decision} is not present in the macid")

    if not all([node in macid.nodes for node in effective_set]):
        raise Exception("One or many of the nodes in the effective_set are not present in the macid.")

    agent = macid.whose_node[decision]
    agent_utils = macid.utility_nodes_agent[agent]
    reachable_decisions = []    # set of possible D_B
    list_decs = copy.deepcopy(macid.all_decision_nodes)
    list_decs.remove(decision)
    for dec_reach in list_decs:
        if dec_reach in effective_set:
            if directed_decision_free_path(macid, decision, dec_reach):
                reachable_decisions.append(dec_reach)

    for decision_b in reachable_decisions:
        agent_b = macid.whose_node[decision_b]
        agent_b_utils = macid.utility_nodes_agent[agent_b]

        for u in agent_utils:
            if _effective_dir_path_exists(macid, decision_b, u, effective_set):

                for u_b in agent_b_utils:
                    decision_b_parents_not_desc_decision = [node for node in macid.get_parents(decision_b)
                                                            if node not in set(nx.descendants(macid, decision))]
                    cond_nodes = [decision_b] + decision_b_parents_not_desc_decision
                    if is_active_indirect_frontdoor_trail(macid, decision, u_b, cond_nodes):
                        return True
    else:
        return False

def get_reasoning_patterns(mb: MACID) -> Dict[str, List[Any]]:
    """ Return a dictionary matching each reasoning pattern with the decision nodes in the MAID which admit it.
    This finds all of the circumstances under which an agent in a MAID has a reason to prefer one strategy over
    another, when all other agents are playing WD strategies.
    (Pfeffer and Gal, 2007: On the Reasoning patterns of Agents in Games).
    """
    motivations: Dict[str, List[str]] = {'dir_effect': [], 'sig': [], 'manip': [], 'rev_den': []}
    effective_set = list(mb.all_decision_nodes)
    while True:
        new_set = [dec for dec in effective_set if direct_effect(mb, dec) or manipulation(mb, dec, effective_set)
                   or signaling(mb, dec, effective_set) or revealing_or_denying(mb, dec, effective_set)]

        if len(new_set) == len(effective_set):
            break
        effective_set = new_set

    for decision in effective_set:
        if direct_effect(mb, decision):
            motivations['dir_effect'].append(decision)
        elif signaling(mb, decision, effective_set):
            motivations['sig'].append(decision)
        elif manipulation(mb, decision, effective_set):
            motivations['manip'].append(decision)
        elif revealing_or_denying(mb, decision, effective_set):
            motivations['rev_den'].append(decision)

    return motivations