# Copyright (C) 2013 Ian W. Harry
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import copy
import numpy
import logging
from six.moves import range
from pycbc.tmpltbank import coord_utils

class PartitionedTmpltbank(object):
    """
    This class is used to hold a template bank partitioned into numerous bins
    based on position in the Cartesian parameter space where the axes are the
    principal components. It can also be used to hold intermediary
    products used while constructing (e.g.) a stochastic template bank.
    """
    def __init__(self, mass_range_params, metric_params, ref_freq,
                 bin_spacing, bin_range_check=1):
        """
        Set up the partitioned template bank class. The combination of the
        reference frequency, the bin spacing and the metric dictates how the
        parameter space will be partitioned.

        Parameters
        -----------
        mass_range_params : massRangeParameters object
            An initialized massRangeParameters object holding the details of
            the mass and spin ranges being considered.
        metric_params : metricParameters object
            An initialized metricParameters object holding the details of the
            parameter space metric that is being used.
        ref_freq : float
            The reference frequency to use as the upper frequency cutoff of
            the metric when partitioning the bank. In general this would be
            set to the *smallest* upper frequency cutoff that is possible in
            the given parameter space. However, in some cases this can lead
            to only a small number of partitions and the computational cost
            will increase dramatically. NOTE: when using the vary-fupper
            option this upper frequency cutoff is only used to determine which
            points should be matched against each other, it is *not* used in
            the actual metric-based calculation of the distance (which uses the
            frequency cutoffs of the points being considered).
        bin_spacing : float
            The metric distance to space the bins by. NOTE: If you want to
            place the bins to have a width corresponding to a minimal match of
            0.97 you would set this to :math:`(1 - 0.97)^{0.5}`.
            Note the square root,
            matches correspond to the square of parameter space distance.
        bin_range_check : int
            When computing matches consider points in the corresponding bin and
            all bins +/- this value in both chi_1 and chi_2 directions.
            DEFAULT = 1.
        """
        # Flags to be used in other methods of this class. Initialized here for
        # simplicity
        self.spin_warning_given = False

        # These will probably be used a lot, so add to object
        self.mass_range_params = mass_range_params
        self.metric_params = metric_params
        self.ref_freq = ref_freq
        self.bin_spacing = bin_spacing

        # Get parameter space extent
        vals = coord_utils.estimate_mass_range(1000000, mass_range_params,
                                          metric_params, ref_freq, covary=True)
        chi1_max = vals[0].max()
        chi1_min = vals[0].min()
        chi1_diff = chi1_max - chi1_min
        chi2_max = vals[1].max()
        chi2_min = vals[1].min()
        chi2_diff = chi2_max - chi2_min
        # Add a little bit extra as we may not have reached the edges.
        # FIXME: Maybe better to use the numerical code to find maxima here?
        chi1_min = chi1_min - 0.1*chi1_diff
        chi1_max = chi1_max + 0.1*chi1_diff
        chi2_min = chi2_min - 0.1*chi2_diff
        chi2_max = chi2_max + 0.1*chi2_diff

        massbank = {}
        bank = {}
        # Also add a little bit here
        for i in range(-2, int((chi1_max - chi1_min) // bin_spacing + 2)):
            bank[i] = {}
            massbank[i] = {}
            for j in range(-2, int((chi2_max - chi2_min) // bin_spacing + 2)):
                bank[i][j] = []
                massbank[i][j] = {}
                massbank[i][j]['mass1s'] = numpy.array([])

        self.massbank = massbank
        self.bank = bank
        # Record minimum and maximum bins
        self.min_chi1_bin = -2
        self.min_chi2_bin = -2
        self.max_chi1_bin = int((chi1_max - chi1_min) // bin_spacing + 1)
        self.max_chi2_bin = int((chi2_max - chi2_min) // bin_spacing + 1)
        self.chi1_min = chi1_min
        self.chi1_max = chi1_max
        self.chi2_min = chi2_min
        self.chi2_max = chi2_max

        # How many adjacent bins should we check?
        self.bin_range_check = 1
        self.bin_loop_order = coord_utils.outspiral_loop(self.bin_range_check)

    def get_point_from_bins_and_idx(self, chi1_bin, chi2_bin, idx):
        """Find masses and spins given bin numbers and index.

        Given the chi1 bin, chi2 bin and an index, return the masses and spins
        of the point at that index. Will fail if no point exists there.

        Parameters
        -----------
        chi1_bin : int
            The bin number for chi1.
        chi2_bin : int
            The bin number for chi2.
        idx : int
            The index within the chi1, chi2 bin.

        Returns
        --------
        mass1 : float
            Mass of heavier body.
        mass2 : float
            Mass of lighter body.
        spin1z : float
            Spin of heavier body.
        spin2z : float
            Spin of lighter body.
        """
        mass1 = self.massbank[chi1_bin][chi2_bin]['mass1s'][idx]
        mass2 = self.massbank[chi1_bin][chi2_bin]['mass2s'][idx]
        spin1z = self.massbank[chi1_bin][chi2_bin]['spin1s'][idx]
        spin2z = self.massbank[chi1_bin][chi2_bin]['spin2s'][idx]
        return mass1, mass2, spin1z, spin2z

    def get_freq_map_and_normalizations(self, frequency_list,
                                        upper_freq_formula):
        """
        If using the --vary-fupper capability we need to store the mapping
        between index and frequencies in the list. We also precalculate the
        normalization factor at every frequency, which is used when estimating
        overlaps to account for abrupt changes in termination frequency.

        Parameters
        -----------
        frequency_list : array of floats
            The frequencies for which the metric has been computed and lie
            within the parameter space being considered.
        upper_freq_formula : string
        """
        self.frequency_map = {}
        self.normalization_map = {}
        self.upper_freq_formula = upper_freq_formula
        # FIXME: Must this be sorted on input
        frequency_list.sort()

        for idx, frequency in enumerate(frequency_list):
            self.frequency_map[frequency] = idx
            self.normalization_map[frequency] = \
                             (self.metric_params.moments['I7'][frequency])**0.5

    def find_point_bin(self, chi_coords):
        """
        Given a set of coordinates in the chi parameter space, identify the
        indices of the chi1 and chi2 bins that the point occurs in. Returns
        these indices.

        Parameters
        -----------
        chi_coords : numpy.array
            The position of the point in the chi coordinates.

        Returns
        --------
        chi1_bin : int
            Index of the chi_1 bin.
        chi2_bin : int
            Index of the chi_2 bin.
        """
        # Identify bin
        chi1_bin = int((chi_coords[0] - self.chi1_min) // self.bin_spacing)
        chi2_bin = int((chi_coords[1] - self.chi2_min) // self.bin_spacing)
        self.check_bin_existence(chi1_bin, chi2_bin)
        return chi1_bin, chi2_bin


    def check_bin_existence(self, chi1_bin, chi2_bin):
        """
        Given indices for bins in chi1 and chi2 space check that the bin
        exists in the object. If not add it. Also check for the existence of
        all bins within +/- self.bin_range_check and add if not present.

        Parameters
        -----------
        chi1_bin : int
            The index of the chi1_bin to check
        chi2_bin : int
            The index of the chi2_bin to check
        """
        bin_range_check = self.bin_range_check
        # Check if this bin actually exists. If not add it
        if ( (chi1_bin < self.min_chi1_bin+bin_range_check) or
             (chi1_bin > self.max_chi1_bin-bin_range_check) or
             (chi2_bin < self.min_chi2_bin+bin_range_check) or
             (chi2_bin > self.max_chi2_bin-bin_range_check) ):
            for temp_chi1 in range(chi1_bin-bin_range_check,
                                                   chi1_bin+bin_range_check+1):
                if temp_chi1 not in self.massbank:
                    self.massbank[temp_chi1] = {}
                    self.bank[temp_chi1] = {}
                for temp_chi2 in range(chi2_bin-bin_range_check,
                                                   chi2_bin+bin_range_check+1):
                    if temp_chi2 not in self.massbank[temp_chi1]:
                        self.massbank[temp_chi1][temp_chi2] = {}
                        self.massbank[temp_chi1][temp_chi2]['mass1s'] =\
                                                                numpy.array([])
                        self.bank[temp_chi1][temp_chi2] = []

    def calc_point_distance(self, chi_coords):
        """
        Calculate distance between point and the bank. Return the closest
        distance.

        Parameters
        -----------
        chi_coords : numpy.array
            The position of the point in the chi coordinates.

        Returns
        --------
        min_dist : float
            The smallest **SQUARED** metric distance between the test point and
            the bank.
        indexes : The chi1_bin, chi2_bin and position within that bin at which
            the closest matching point lies.
        """
        chi1_bin, chi2_bin = self.find_point_bin(chi_coords)
        min_dist = 1000000000
        indexes = None
        for chi1_bin_offset, chi2_bin_offset in self.bin_loop_order:
            curr_chi1_bin = chi1_bin + chi1_bin_offset
            curr_chi2_bin = chi2_bin + chi2_bin_offset
            for idx, bank_chis in \
                            enumerate(self.bank[curr_chi1_bin][curr_chi2_bin]):
                dist = coord_utils.calc_point_dist(chi_coords, bank_chis)
                if dist < min_dist:
                    min_dist = dist
                    indexes = (curr_chi1_bin, curr_chi2_bin, idx)
        return min_dist, indexes

    def test_point_distance(self, chi_coords, distance_threshold):
        """
        Test if the distance between the supplied point and the bank is less
        than the supplied distance theshold.

        Parameters
        -----------
        chi_coords : numpy.array
            The position of the point in the chi coordinates.
        distance_threshold : float
            The **SQUARE ROOT** of the metric distance to test as threshold.
            E.g. if you want to test to a minimal match of 0.97 you would
            use 1 - 0.97 = 0.03 for this value.

        Returns
        --------
        Boolean
            True if point is within the distance threshold. False if not.

        """
        chi1_bin, chi2_bin = self.find_point_bin(chi_coords)
        for chi1_bin_offset, chi2_bin_offset in self.bin_loop_order:
            curr_chi1_bin = chi1_bin + chi1_bin_offset
            curr_chi2_bin = chi2_bin + chi2_bin_offset
            for bank_chis in self.bank[curr_chi1_bin][curr_chi2_bin]:
                dist = coord_utils.calc_point_dist(chi_coords, bank_chis)
                if dist < distance_threshold:
                    return True
        else:
            return False

    def calc_point_distance_vary(self, chi_coords, point_fupper, mus):
        """
        Calculate distance between point and the bank allowing the metric to
        vary based on varying upper frequency cutoff. Slower than
        calc_point_distance, but more reliable when upper frequency cutoff can
        change a lot.

        Parameters
        -----------
        chi_coords : numpy.array
            The position of the point in the chi coordinates.
        point_fupper : float
            The upper frequency cutoff to use for this point. This value must
            be one of the ones already calculated in the metric.
        mus : numpy.array
            A 2D array where idx 0 holds the upper frequency cutoff and idx 1
            holds the coordinates in the [not covaried] mu parameter space for
            each value of the upper frequency cutoff.

        Returns
        --------
        min_dist : float
            The smallest **SQUARED** metric distance between the test point and
            the bank.
        indexes : The chi1_bin, chi2_bin and position within that bin at which
            the closest matching point lies.
        """
        chi1_bin, chi2_bin = self.find_point_bin(chi_coords)
        min_dist = 1000000000
        indexes = None
        for chi1_bin_offset, chi2_bin_offset in self.bin_loop_order:
            curr_chi1_bin = chi1_bin + chi1_bin_offset
            curr_chi2_bin = chi2_bin + chi2_bin_offset
            # No points = Next iteration
            curr_bank = self.massbank[curr_chi1_bin][curr_chi2_bin]
            if not curr_bank['mass1s'].size:
                continue

            # *NOT* the same of .min and .max
            f_upper = numpy.minimum(point_fupper, curr_bank['freqcuts'])
            f_other = numpy.maximum(point_fupper, curr_bank['freqcuts'])
            # NOTE: freq_idxes is a vector!
            freq_idxes = numpy.array([self.frequency_map[f] for f in f_upper])
            # vecs1 gives a 2x2 vector: idx0 = stored index, idx1 = mu index
            vecs1 = mus[freq_idxes, :]
            # vecs2 gives a 2x2 vector: idx0 = stored index, idx1 = mu index
            range_idxes = numpy.arange(len(freq_idxes))
            vecs2 = curr_bank['mus'][range_idxes, freq_idxes, :]

            # Now do the sums
            dists = (vecs1 - vecs2)*(vecs1 - vecs2)
            # This reduces to 1D: idx = stored index
            dists = numpy.sum(dists, axis=1)
            norm_upper = numpy.array([self.normalization_map[f] \
                                      for f in f_upper])
            norm_other = numpy.array([self.normalization_map[f] \
                                      for f in f_other])
            norm_fac = norm_upper / norm_other
            renormed_dists = 1 - (1 - dists)*norm_fac
            curr_min_dist = renormed_dists.min()
            if curr_min_dist < min_dist:
                min_dist = curr_min_dist
                indexes = curr_chi1_bin, curr_chi2_bin, renormed_dists.argmin()

        return min_dist, indexes

    def test_point_distance_vary(self, chi_coords, point_fupper, mus,
                                 distance_threshold):
        """
        Test if distance between point and the bank is greater than distance
        threshold while allowing the metric to
        vary based on varying upper frequency cutoff. Slower than
        test_point_distance, but more reliable when upper frequency cutoff can
        change a lot.

        Parameters
        -----------
        chi_coords : numpy.array
            The position of the point in the chi coordinates.
        point_fupper : float
            The upper frequency cutoff to use for this point. This value must
            be one of the ones already calculated in the metric.
        mus : numpy.array
            A 2D array where idx 0 holds the upper frequency cutoff and idx 1
            holds the coordinates in the [not covaried] mu parameter space for
            each value of the upper frequency cutoff.
        distance_threshold : float
            The **SQUARE ROOT** of the metric distance to test as threshold.
            E.g. if you want to test to a minimal match of 0.97 you would
            use 1 - 0.97 = 0.03 for this value.

        Returns
        --------
        Boolean
            True if point is within the distance threshold. False if not.
        """
        chi1_bin, chi2_bin = self.find_point_bin(chi_coords)
        for chi1_bin_offset, chi2_bin_offset in self.bin_loop_order:
            curr_chi1_bin = chi1_bin + chi1_bin_offset
            curr_chi2_bin = chi2_bin + chi2_bin_offset
            # No points = Next iteration
            curr_bank = self.massbank[curr_chi1_bin][curr_chi2_bin]
            if not curr_bank['mass1s'].size:
                continue

            # *NOT* the same of .min and .max
            f_upper = numpy.minimum(point_fupper, curr_bank['freqcuts'])
            f_other = numpy.maximum(point_fupper, curr_bank['freqcuts'])
            # NOTE: freq_idxes is a vector!
            freq_idxes = numpy.array([self.frequency_map[f] for f in f_upper])
            # vecs1 gives a 2x2 vector: idx0 = stored index, idx1 = mu index
            vecs1 = mus[freq_idxes, :]
            # vecs2 gives a 2x2 vector: idx0 = stored index, idx1 = mu index
            range_idxes = numpy.arange(len(freq_idxes))
            vecs2 = curr_bank['mus'][range_idxes,freq_idxes,:]

            # Now do the sums
            dists = (vecs1 - vecs2)*(vecs1 - vecs2)
            # This reduces to 1D: idx = stored index
            dists = numpy.sum(dists, axis=1)
            # I wonder if this line actually speeds things up?
            if (dists > distance_threshold).all():
                continue
            # This is only needed for close templates, should we prune?
            norm_upper = numpy.array([self.normalization_map[f] \
                                      for f in f_upper])
            norm_other = numpy.array([self.normalization_map[f] \
                                      for f in f_other])
            norm_fac = norm_upper / norm_other
            renormed_dists = 1 - (1 - dists)*norm_fac
            if (renormed_dists < distance_threshold).any():
                return True
        else:
            return False

    def add_point_by_chi_coords(self, chi_coords, mass1, mass2, spin1z, spin2z,
                          point_fupper=None, mus=None):
        """
        Add a point to the partitioned template bank. The point_fupper and mus
        kwargs must be provided for all templates if the vary fupper capability
        is desired. This requires that the chi_coords, as well as mus and
        point_fupper if needed, to be precalculated. If you just have the
        masses and don't want to worry about translations see
        add_point_by_masses, which will do translations and then call this.

        Parameters
        -----------
        chi_coords : numpy.array
            The position of the point in the chi coordinates.
        mass1 : float
            The heavier mass of the point to add.
        mass2 : float
            The lighter mass of the point to add.
        spin1z: float
            The [aligned] spin on the heavier body.
        spin2z: float
            The [aligned] spin on the lighter body.
            The upper frequency cutoff to use for this point. This value must
            be one of the ones already calculated in the metric.
        mus : numpy.array
            A 2D array where idx 0 holds the upper frequency cutoff and idx 1
            holds the coordinates in the [not covaried] mu parameter space for
            each value of the upper frequency cutoff.
        """
        chi1_bin, chi2_bin = self.find_point_bin(chi_coords)
        self.bank[chi1_bin][chi2_bin].append(copy.deepcopy(chi_coords))
        curr_bank = self.massbank[chi1_bin][chi2_bin]

        if curr_bank['mass1s'].size:
            curr_bank['mass1s'] = numpy.append(curr_bank['mass1s'],
                                               numpy.array([mass1]))
            curr_bank['mass2s'] = numpy.append(curr_bank['mass2s'],
                                               numpy.array([mass2]))
            curr_bank['spin1s'] = numpy.append(curr_bank['spin1s'],
                                               numpy.array([spin1z]))
            curr_bank['spin2s'] = numpy.append(curr_bank['spin2s'],
                                               numpy.array([spin2z]))
            if point_fupper is not None:
                curr_bank['freqcuts'] = numpy.append(curr_bank['freqcuts'],
                                                 numpy.array([point_fupper]))
            # Mus needs to append onto axis 0. See below for contents of
            # the mus variable
            if mus is not None:
                curr_bank['mus'] = numpy.append(curr_bank['mus'],
                                            numpy.array([mus[:,:]]), axis=0)
        else:
            curr_bank['mass1s'] = numpy.array([mass1])
            curr_bank['mass2s'] = numpy.array([mass2])
            curr_bank['spin1s'] = numpy.array([spin1z])
            curr_bank['spin2s'] = numpy.array([spin2z])
            if point_fupper is not None:
                curr_bank['freqcuts'] = numpy.array([point_fupper])
            # curr_bank['mus'] is a 3D array
            # NOTE: mu relates to the non-covaried Cartesian coordinate system
            # Axis 0: Template index
            # Axis 1: Frequency cutoff index
            # Axis 2: Mu coordinate index
            if mus is not None:
                curr_bank['mus'] = numpy.array([mus[:,:]])

    def add_point_by_masses(self, mass1, mass2, spin1z, spin2z,
                            vary_fupper=False):
        """
        Add a point to the template bank. This differs from add point to bank
        as it assumes that the chi coordinates and the products needed to use
        vary_fupper have not already been calculated. This function calculates
        these products and then calls add_point_by_chi_coords.
        This function also
        carries out a number of sanity checks (eg. is the point within the
        ranges given by mass_range_params) that add_point_by_chi_coords does
        not do for speed concerns.

        Parameters
        -----------
        mass1 : float
            Mass of the heavier body
        mass2 : float
            Mass of the lighter body
        spin1z : float
            Spin of the heavier body
        spin2z : float
            Spin of the lighter body
        """
        # Test that masses are the expected way around (ie. mass1 > mass2)
        if mass2 > mass1:
            if not self.spin_warning_given:
                warn_msg = "Am adding a template where mass2 > mass1. The "
                warn_msg += "convention is that mass1 > mass2. Swapping mass1 "
                warn_msg += "and mass2 and adding point to bank. This message "
                warn_msg += "will not be repeated."
                logging.warn(warn_msg)
                self.spin_warning_given = True

        # These that masses obey the restrictions of mass_range_params
        if self.mass_range_params.is_outside_range(mass1, mass2, spin1z,
                                                                       spin2z):
            err_msg = "Point with masses given by "
            err_msg += "%f %f %f %f " %(mass1, mass2, spin1z, spin2z)
            err_msg += "(mass1, mass2, spin1z, spin2z) is not consistent "
            err_msg += "with the provided command-line restrictions on masses "
            err_msg += "and spins."
            raise ValueError(err_msg)

        # Get chi coordinates
        chi_coords = coord_utils.get_cov_params(mass1, mass2, spin1z, spin2z,
                                                self.metric_params,
                                                self.ref_freq)

        # Get mus and best fupper for this point, if needed
        if vary_fupper:
            mass_dict = {}
            mass_dict['m1'] = numpy.array([mass1])
            mass_dict['m2'] = numpy.array([mass2])
            mass_dict['s1z'] = numpy.array([spin1z])
            mass_dict['s2z'] = numpy.array([spin2z])
            freqs = numpy.array([self.frequency_map.keys()], dtype=float)
            freq_cutoff = coord_utils.return_nearest_cutoff(\
                                     self.upper_freq_formula, mass_dict, freqs)
            freq_cutoff = freq_cutoff[0]
            lambdas = coord_utils.get_chirp_params\
                (mass1, mass2, spin1z, spin2z, self.metric_params.f0,
                 self.metric_params.pnOrder)
            mus = []
            for freq in self.frequency_map:
                mus.append(coord_utils.get_mu_params(lambdas,
                                                    self.metric_params, freq) )
            mus = numpy.array(mus)
        else:
            freq_cutoff=None
            mus=None

        self.add_point_by_chi_coords(chi_coords, mass1, mass2, spin1z, spin2z,
                               point_fupper=freq_cutoff, mus=mus)


    def add_tmpltbank_from_xml_table(self, sngl_table, vary_fupper=False):
        """
        This function will take a sngl_inspiral_table of templates and add them
        into the partitioned template bank object.

        Parameters
        -----------
        sngl_table : sngl_inspiral_table
            List of sngl_inspiral templates.
        vary_fupper : False
            If given also include the additional information needed to compute
            distances with a varying upper frequency cutoff.
        """
        for sngl in sngl_table:
            self.add_point_by_masses(sngl.mass1, sngl.mass2, sngl.spin1z,
                                     sngl.spin2z, vary_fupper=vary_fupper)

    def add_tmpltbank_from_hdf_file(self, hdf_fp, vary_fupper=False):
        """
        This function will take a pointer to an open HDF File object containing
        a list of templates and add them into the partitioned template bank
        object.

        Parameters
        -----------
        hdf_fp : h5py.File object
            The template bank in HDF5 format.
        vary_fupper : False
            If given also include the additional information needed to compute
            distances with a varying upper frequency cutoff.
        """
        mass1s = hdf_fp['mass1'][:]
        mass2s = hdf_fp['mass2'][:]
        spin1zs = hdf_fp['spin1z'][:]
        spin2zs = hdf_fp['spin2z'][:]
        for idx in range(len(mass1s)):
            self.add_point_by_masses(mass1s[idx], mass2s[idx], spin1zs[idx],
                                     spin2zs[idx], vary_fupper=vary_fupper)

    def output_all_points(self):
        """Return all points in the bank.

        Return all points in the bank as lists of m1, m2, spin1z, spin2z.

        Returns
        -------
        mass1 : list
            List of mass1 values.
        mass2 : list
            List of mass2 values.
        spin1z : list
            List of spin1z values.
        spin2z : list
            List of spin2z values.
        """
        mass1 = []
        mass2 = []
        spin1z = []
        spin2z = []
        for i in self.massbank.keys():
            for j in self.massbank[i].keys():
                for k in range(len(self.massbank[i][j]['mass1s'])):
                    curr_bank = self.massbank[i][j]
                    mass1.append(curr_bank['mass1s'][k])
                    mass2.append(curr_bank['mass2s'][k])
                    spin1z.append(curr_bank['spin1s'][k])
                    spin2z.append(curr_bank['spin2s'][k])

        return mass1, mass2, spin1z, spin2z
