""" Gravitational Wave Surrogate classes for text and hdf5 files"""

from __future__ import division  # for py2

__copyright__ = "Copyright (C) 2014 Scott Field and Chad Galley"
__email__     = "sfield@astro.cornell.edu, crgalley@tapir.caltech.edu"
__status__    = "testing"
__author__    = "Jonathan Blackman, Scott Field, Chad Galley, Vijay Varma, Kevin Barkett"

__license__ = """
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

# adding "_" prefix to potentially unfamiliar module names
# so they won't show up in gws' tab completion
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline as _iuspline
from gwtools.harmonics import sYlm as _sYlm

if __package__ is "" or "None": # py2 and py3 compatible
  print("setting __package__ to gwsurrogate.new so relative imports work")
  __package__="gwsurrogate.new"

# assumes unique global names
from .saveH5Object import SimpleH5Object
from .saveH5Object import H5ObjectList
from .saveH5Object import H5ObjectDict
from .nodeFunction import NodeFunction
from .spline_evaluation import TensorSplineGrid, fast_complex_tensor_spline_eval
from gwsurrogate import spline_interp_Cwrapper
from .tidal_functions import UniversalRelationLambda2ToI, \
    UniversalRelationLambda2ToOmega2, UniversalRelationLambda2ToLambda3, \
    UniversalRelationLambda3ToOmega3, UniversalRelationLambda2ToAqm, \
    EffectiveDeformabilityFromDynamicalTides, PNT2Tidal, \
    EffectiveDissipativeDynamicalTides, StrainTidalEnhancementFactor

PARAM_NUDGE_TOL = 1.e-12 # Default relative tolerance for nudging edge cases

def _identity(r1, r2):
    return r1, r2
def _amp_phase(r1, r2):
    return r1['amp']*np.exp(1.j*r1['phase'])
def _re_im(r1, r2):
    return r1['re'] + 1.j*r1['im']

RECOMBINATION_FUNCS = {
    'identity': _identity,
    'amp_phase': _amp_phase,
    're_im': _re_im,
        }

def _mode_sum(modes, theta, phi):
    h = 0.
    for (ell, m), h_mode in modes.items(): # inefficient in py2
        h += _sYlm(-2, ell, m, theta, phi) * h_mode
    return h

def _splinterp(xout, xin, yin, k=3, ext='const'):
    """Uses InterpolatedUnivariateSpline to interpolate real or complex data"""
    if np.iscomplexobj(yin):
        re = _splinterp(xout, xin, np.real(yin), k=k, ext=ext)
        im = _splinterp(xout, xin, np.imag(yin), k=k, ext=ext)
        return re + 1.j*im
    else:
        return _iuspline(xin, yin, k=k, ext=ext)(xout)

def _splinterp_Cwrapper(xout, xin, yin):
    """Uses gsl splines with a wrapper to interpolate real or complex data.
    Uses natural boundary conditions instead of not-a-knot boundary conditions
    like InterpolatedUnivariateSpline."""
    if len(xin) != len(yin):
        raise Exception('Expected x and y input lengths to match.')
    if np.iscomplexobj(yin):
        re = _splinterp_Cwrapper(xout, xin, np.real(yin))
        im = _splinterp_Cwrapper(xout, xin, np.imag(yin))
        return re + 1.j*im
    else:
        return spline_interp_Cwrapper.interpolate(xout, xin, yin)

class ParamDim(SimpleH5Object):
    """
    A helper class containing the information and functions for a single
    parameter space dimension
    """

    def __init__(self, name='', min_val=0, max_val=1, rtol=PARAM_NUDGE_TOL):
        """
        name: A descriptive name for this parameter dimension
        min_val: The minimum allowed value for this parameter
        max_val: The maximum allowed value for this parameter
        rtol: A relative tolerance for nudging parameters lying outside
             [min_val, max_val]. Scaled by (max_val - min_val) for an absolute
             tolerance. Useful to avoid machine precision issues.
        """
        super(ParamDim, self).__init__()

        tol = rtol * (max_val - min_val)

        if min_val + 2*tol > max_val:
            raise Exception("tol %s is too large for %s with range [%s, %s]"%(
                            tol, name, min_val, max_val))

        self.name = name
        self.min_val = min_val
        self.max_val = max_val
        self.tol = tol
        self.tol_max = max_val - tol
        self.tol_min = min_val + tol

    def __str__(self):
        return self.name

    def __repr__(self):
        return '%s: [%s, %s] with tol %s'%(self.name, self.min_val,
                                           self.max_val, self.tol)

    def nudge(self, x):
        """
        Returns a nudged version of x lying within [min_val+tol, max_val-tol].
        x can have any shape. Returns x if it already lies in the interval.
        x must already lie within [min_val - tol, max_val + tol].
        """
        xmax = np.max(x)
        xmin = np.min(x)

        if xmax > self.tol_max:
            if xmax > self.max_val + self.tol:
                raise Exception("The maximum allowed %s is %s, got %s"%(
                                self.name, self.max_val, xmax))
            x = np.amin([x, np.ones(np.shape(x)) * self.tol_max], axis=0)

        if xmin < self.tol_min:
            if xmin < self.min_val - self.tol:
                raise Exception("The maximum allowed %s is %s, got %s"%(
                                self.name, self.min_val, xmin))
            x = np.amax([x, np.ones(np.shape(x)) * self.tol_min], axis=0)

        return x

class ParamSpace(SimpleH5Object):
    """
    A helper class for all the parameter domain related information and
    functions.
    """

    def __init__(self, name='', params=[]):
        """
        name: A descriptive name for this parameter space
        params: A list of ParamDim instances, one per parameter dimension
        """
        super(ParamSpace, self).__init__(['name', 'dim'], ['_params'])

        self.name = name
        self._params = H5ObjectList(params)
        self.dim = len(params)

    def __str__(self):
        return self.name

    def __repr__(self):
        return '%s: %s'%(self.name, [p.name for p in self._params])

    def param_names(self):
        """Returns the list of names for each parameter space dimension"""
        return [p.name for p in self._params]

    def min_vals(self):
        """Returns a list of minimum parameter space values"""
        return [p.min_val for p in self._params]

    def max_vals(self):
        """Returns a list of maximum parameter space values"""
        return [p.max_val for p in self._params]

    def nudge_params(self, x):
        """
        Nudges parameters lying slightly outside the valid domain to the
        boundary. x can be a single 1d parameter vector or a 2d vector with
        shape (n_params, self.dim).
        """

        xshape = np.shape(x)

        # It's convenient to be able to accept a float instead of a length-1
        # array for 1d parameter spaces.
        if len(xshape) == 0:
            x = np.array([x])
            xshape = np.shape(x)

        if len(xshape) == 1:
            if len(x) != self.dim:
                raise Exception("Parameter space has dimension %s, got %s."%(
                                self.dim, xshape))
            res = np.array([p.nudge(xi) for p, xi in zip(self._params, x)])

        elif len(xshape) == 2:
            if xshape[1] != self.dim:
                raise Exception("Expecting array with shape (n, %s), got %s"%(
                                self.dim, xshape))
            res = np.array([p.nudge(xi) for p, xi in zip(self._params, x.T)])

        else:
            raise Exception("x should be 1d or 2d, got shape {}".format(xshape))

        return res

    def h5_prepare_subs(self):
        """Setup dummy subordinates before loading them"""
        params = [ParamDim() for _ in range(self.dim)]
        self._params = H5ObjectList(params)

class _SingleFunctionSurrogate_NoChecks(SimpleH5Object):
    """
    A surrogate model for a single (real or complex) function on a 1d domain.
    Skips sanity/validity checks, so should be called with sanitized inputs.
    Use SingleFunctionSurrogate for actual surrogates of single functions.
    """

    def __init__(self, name=None, ei_basis=None, node_functions=[]):
        """
        name: A descriptive name for the function this surrogate models
        ei_basis: A basis with shape (n_nodes, len(domain)).
                  Interpolation is done via nodes.dot(ei_basis)
        node_functions: A list of evaluators for each node.
                        Each one takes a parameter space vector x and returns
                        the node evaluated at x.
        """
        super(_SingleFunctionSurrogate_NoChecks, self).__init__(
                data_keys=['name', 'ei_basis', 'n_nodes'],
                sub_keys=['node_functions'])

        self.name = name
        self.ei_basis = ei_basis
        self.n_nodes = len(node_functions)
        self.node_functions = H5ObjectList(node_functions)

    def __str__(self):
        return self.name

    def __repr__(self):
        return self.name

    def __call__(self, x):
        """
        Evaluates the surrogate at x, returning the result.
        """
        nodes = np.array([nf(x) for nf in self.node_functions])
        return nodes.dot(self.ei_basis)

    def h5_prepare_subs(self):
        """Setup NodeFunctions before loading them"""
        tmp_nodes = [NodeFunction() for _ in range(self.n_nodes)]
        self.node_functions = H5ObjectList(tmp_nodes)

class SingleFunctionSurrogate(_SingleFunctionSurrogate_NoChecks):
    """
    A surrogate model for a single (real or complex) function on a 1d domain.
    """

    def __init__(self, name=None, domain=None, param_space=None,
                 ei_basis=None, node_functions=[]):
        """
        name: A descriptive name for the function this surrogate models
        domain: A 1d array of the monotonically increasing domain
                (time/frequency) values
        param_space: A ParamSpace for this surrogate
        ei_basis: A basis with shape (n_nodes, len(domain)).
                  Interpolation is done via nodes.dot(ei_basis)
        node_functions: A list of evaluators for each node.
                        Each one takes a parameter space vector x and returns
                        the node evaluated at x.
        """
        super(SingleFunctionSurrogate, self).__init__(name=name,
                ei_basis=ei_basis, node_functions=node_functions)

        if domain is not None and np.min(np.diff(domain)) <= 0.0:
            raise Exception("domain should be monotonically increasing")

        self.domain = domain
        self.param_space = param_space

        self._h5_data_keys.append('domain')
        self._h5_subordinate_keys.append('param_space')

    def h5_prepare_subs(self):
        super(SingleFunctionSurrogate, self).h5_prepare_subs()
        self.param_space = ParamSpace()

    def __repr__(self):
        return '%s (%s)'%(self.name, self.param_space.name)

    def __call__(self, x, domain=None):
        """
        Evaluates the surrogate at x, returning the result.
        domain: An optional 1d array of domain values. If given, the result is
                evaluated at the domain values.
        """
        # Verify valid parameters and domain
        x = self.param_space.nudge_params(x)
        if domain is not None:
            if domain[0] < self.domain[0] or domain[-1] > self.domain[-1]:
                raise Exception("Domain must lie in [%s, %s]"%(
                                self.domain[0], self.domain[-1]))

        res = super(SingleFunctionSurrogate, self).__call__(x)

        if domain is not None:
            res = _splinterp(domain, self.domain, res)

        return res

class _ManyFunctionSurrogate_NoChecks(SimpleH5Object):
    """
    A container for surrogates sharing the same parameter space and domain.
    Skips checks, assuming they have already been performed.
    """

    def __init__(self, name='', single_function_components={},
                 many_function_components={}, combine_func='identity'):
        """
        name: A descriptive name for this surrogate.
        single_function_components: A dictionary of components, where the
                values are (ei_basis, node_functions) tuples.
        many_function_components: A dictionary of _ManyFunctionSurrogate_NoCheck
                components, where the values are
                (combine_func, single_function_components,
                 many_function_components) tuples.
        combine_func: A key of RECOMBINATION_FUNCS used to combine subordinate
                      results into a single result.
        """
        super(_ManyFunctionSurrogate_NoChecks, self).__init__(
                data_keys=['name', 'func_keys', 'sur_keys', 'combine_func'],
                sub_keys=['func_subs', 'sur_subs'])

        self.name = name
        self.combine_func = combine_func
        self.func_keys = list(single_function_components.keys())
        self.sur_keys = list(many_function_components.keys())
        if len(set(self.func_keys).intersection(self.sur_keys)) > 0:
            raise Exception("Component keys must be unique! Got %s, %s"%(
                            self.func_keys, self.sur_keys))

        func_sub_dict = {}
        for k, (ei, nf) in single_function_components.items(): # inefficient in py2
            func_sub_dict[k] = _SingleFunctionSurrogate_NoChecks(k, ei, nf)
        self.func_subs = H5ObjectDict(func_sub_dict)

        sur_sub_dict = {}
        for k, (cf, sfc, mfc) in many_function_components.items(): # inefficient in py2
            sur_sub_dict[k] = _ManyFunctionSurrogate_NoChecks(k, sfc, mfc, cf)
        self.sur_subs = H5ObjectDict(sur_sub_dict)

    def h5_prepare_subs(self):
        """
        Initialize subordinate surrogate class instances so they can
        load their own data from h5 files.
        """
        self.func_subs = H5ObjectDict({k: _SingleFunctionSurrogate_NoChecks()
                                       for k in self.func_keys})
        self.sur_subs = H5ObjectDict({k: _ManyFunctionSurrogate_NoChecks()
                                      for k in self.sur_keys})

    def __str__(self):
        return self.name

    def __call__(self, x):
        func_evals = {k: sur(x) for k, sur in self.func_subs.iteritems()} # inefficient in py2
        sur_evals = {k: sur(x) for k, sur in self.sur_subs.iteritems()} # inefficient in py2
        return RECOMBINATION_FUNCS[self.combine_func](func_evals, sur_evals)

    def _eval_func(self, x, key):
        return self.func_subs[key](x)

    def _eval_sur(self, x, key):
        return self.sur_subs[key](x)

class ManyFunctionSurrogate(_ManyFunctionSurrogate_NoChecks):
    """
    A container for surrogates sharing the same parameter space and domain.
    """

    def __init__(self, name='', domain=None, param_space=None,
                 single_function_components={}, many_function_components={},
                 combine_func='identity'):
        """
        name: A descriptive name for this surrogate.
        param_space: A ParamSpace for this surrogate.
        single_function_components: A dictionary of components, where the
                values are (ei_basis, node_functions) tuples.
        many_function_components: A dictionary of _ManyFunctionSurrogate_NoCheck
                components, where the values are
                (combine_func, single_function_components,
                 many_function_components) tuples.
        combine_func: A key of RECOMBINATION_FUNCS used to combine subordinate
                      results into a single result.
        """
        super(ManyFunctionSurrogate, self).__init__(
                name=name,
                single_function_components=single_function_components,
                many_function_components=many_function_components,
                combine_func=combine_func,
            )

        self.domain = domain
        self.param_space = param_space
        self._h5_data_keys.append('domain')
        self._h5_subordinate_keys.append('param_space')

    def h5_prepare_subs(self):
        super(ManyFunctionSurrogate, self).h5_prepare_subs()
        self.param_space = ParamSpace()

    def __call__(self, x):
        """
        Evaluates the surrogate at x, returning the result.
        """
        # Verify valid parameters
        x = self.param_space.nudge_params(x)
        return super(ManyFunctionSurrogate, self)(x)

class FastTensorSplineSurrogate(SimpleH5Object):
    """
    A special case of having a complex empirical interpolant combined with
    tensor splines for the real and imaginary parts of each empirical node,
    for each waveform mode. All tensor splines must use the same grid.
    Written for speed, minimizing python operations
    between tensor spline interpolations, which are done simultaneously for
    each mode to keep numpy busy. Obtained ~25ms evaluation time per waveform
    for 12 waveform modes. This was ~66ms when using the python class
    hierarchy, and using a separate call to numpy for each tensor spline
    interpolation. Note that similar C code written with gsl splines takes
    ~50ms, but should have room for optimization.
    """

    def __init__(self, name=None, domain=None, param_space=None,
                 knot_vecs=[], mode_data={}, modes=None):

        super(FastTensorSplineSurrogate, self).__init__(
                sub_keys=['param_space', 'ts_grid'],
                data_keys=['name', 'domain', 'ei', 'cre', 'cim', 'mode_list',
                           'mode_indices'])

        self.name = name
        self.domain = domain
        self.param_space = param_space
        if param_space is None:
            self.param_space = ParamSpace()
        if modes is None:
            modes = list(mode_data.keys())
        self.mode_list = modes
        self.mode_indices = {str(k): i for i, k in enumerate(modes)}
        self.ei = [mode_data[k][0] for k in modes]
        self.cre = [mode_data[k][1] for k in modes]
        self.cim = [mode_data[k][2] for k in modes]
        self.ts_grid = TensorSplineGrid(knot_vecs)

    def __call__(self, x, theta=None, phi=None, modes=None):
        """
        Return surrogate evaluation.
        Arguments:
            x : The intrinsic parameters (see self.param_space)
            theta/phi : polar and azimuthal angles of the direction of
                        gravitational wave emission. If given, sums up modes
                        and returns h_plus and h_cross (default returns modes)
            modes : A list of (ell, m) modes to be evaluated (default: all)
        Returns h:
            h : If theta and phi are None, h is a dictionary of waveform modes
                sampled at self.domain with (ell, m) keys.
                If theta and phi are given, h = h_plus - i * h_cross is a
                complex array given by the sum of the modes.
        """
        if (theta is None) != (phi is None):
            raise Exception("Either give theta and phi or neither")

        x = self.param_space.nudge_params(x)

        if modes is None:
            modes = self.mode_list

        h_modes = {}
        for k in modes:
            i = self.mode_indices[str(k)]

            h_eim = fast_complex_tensor_spline_eval(x,self.ts_grid,self.cre[i],self.cim[i])

            # Evaluate the empirical interpolant
            h_modes[k] = h_eim.dot(self.ei[i])

        if theta is not None:
            return _mode_sum(h_modes, theta, phi)

        return h_modes

class MultiModalSurrogate(ManyFunctionSurrogate):
    """
    A surrogate for multimodal waveforms, where each waveform mode has
    its own surrogate. Contains added functionality for evaluating on
    sphere.
    """

    def __init__(self, name=None, domain=None, param_space=None,
                 mode_data={}, mode_type='complex', modes=None):
        """
        name: A descriptive name for this surrogate.
        domain: A 1d array of the monotonically increasing domain
                (time/frequency) values
        param_space: A ParamSpace for this surrogate.
        mode_data: A dictionary of modes with (l, m) integer keys,
                   where the values are (ei_basis, node_functions) tuples.
        mode_type: Can be 'amp_phase', 're_im', or 'complex' depending on
                   how mode surrogates are built. If 'amp_phase' or 're_im',
                   mode_data values should instead be dictionaries with keys
                   ['amp', 'phase'] or ['re', 'im'] and
                   (ei_basis, node_functions) values.
        modes: A list of (ell, m) modes giving an ordering to mode_data.keys().
               If None, uses mode_data.keys().
        """

        if mode_type == 'complex':
            super(MultiModalSurrogate, self).__init__(name, domain, param_space,
                                                      mode_data, {}, 'identity')
        elif mode_type in ['amp_phase', 're_im']:
            mode_data = {k: (mode_type, v, {})
                         for k, v in mode_data.items()} # inefficient in py2
            super(MultiModalSurrogate, self).__init__(name, domain, param_space,
                                                      {}, mode_data, mode_type)
        else:
            raise ValueError("Invalid mode_type: %s"%(mode_type))

        self.mode_type = mode_type
        if modes is None:
            self.modes = list(mode_data.keys())
        else:
            self.modes = modes

        self._h5_data_keys.append('modes')
        self._h5_data_keys.append('mode_type')

    def __call__(self, x, theta=None, phi=None, modes=None):
        """
        Return surrogate evaluation.
        Arguments:
            x : The intrinsic parameters (see self.param_space)
            theta/phi : polar and azimuthal angles of the direction of
                        gravitational wave emission. If given, sums up modes
                        and returns h_plus and h_cross (default returns modes)
            modes : A list of (ell, m) modes to be evaluated (default: all)
        Returns h:
            h : If theta and phi are None, h is a dictionary of waveform modes
                sampled at self.domain with (ell, m) keys.
                If theta and phi are given, h = h_plus - i * h_cross is a
                complex array given by the sum of the modes.
        """
        if (theta is None) != (phi is None):
            raise Exception("Either give theta and phi or neither")

        x = self.param_space.nudge_params(x)

        if modes is None:
            modes = self.modes

        if self.mode_type == 'complex':
            h_modes = {k: self._eval_func(x, k) for k in modes}
        else:
            h_modes = {k: self._eval_sur(x, k) for k in modes}

        if theta is not None:
            return _mode_sum(h_modes, theta, phi)

        return h_modes

class AlignedSpinCoOrbitalFrameSurrogate(ManyFunctionSurrogate):
    """
    A surrogate for coorbital frame multimodal waveforms, where each waveform
    data piece has its own surrogate.

    The waveform data pieces are:
    Amplitude and phase of the (2,2) mode.
    Real and imaginary parts of coorbital frame waveform for other modes.
    """

    def __init__(self, name=None, domain=None, param_space=None, \
            phaseAlignIdx=None, TaylorT3_t_ref=None, \
            coorb_mode_data={(2, 2): {}}
            ):
        """
        name:               A descriptive name for this surrogate.

        domain:             A 1d array of the monotonically increasing time
                            values.

        param_space:        A ParamSpace for this surrogate.

        phaseAlignIdx:      This value should be loaded directly from the
            surrogate's h5 file. Index of domain at which the orbital phase is
            aligned. This is used when putting back the TaylorT3 contribution
            that was subtracted before modeling the phase.

        TaylorT3_t_ref:     This value should be loaded directly from the
            surrogate's h5 file. This is an arbitrary reference time used
            in the TaylorT3 contribution, but is fixed during the surrogate
            construction.

        coorb_mode_data: A dictionary of modes with (l, m) integer keys, where
            the values are themselves dictionaries containing the coorbital
            frame waveform for that mode. The coorbital frame is defined as:
            H_lm = h_lm*exp(i*m*phi_22/2), where h_lm is the inertial frame
            waveform and h_22 = A_22 * exp(-i phi_22). NOTE the minus sign.

            coorb_mode_data should be a dict with mode_key: mode_value pairs,
                where mode_value = AmpPhase_22 for mode_key = (2, 2)
                and mode_value = CoorbReIm_lm for other modes.

            Above, AmpPhase_22 and CoorbReIm_lm are themselved dictionaries:
                AmpPhase_22 = {'amp': Amp_22, 'phase': phi_22}
                CoorbReIm_lm = {'re': re_H_lm, 'im': im_H_lm}, where
                re_H_lm = Real(H_lm) and im_H_lm = Imag(H_lmc).
            Finally, all of Amp_22, phi_22, re_H_lm and im_H_lm should be
            (ei_basis, node_functions) tuples of that data piece.

            IMPORTANT NOTE: The phase of (2, 2) mode should be defined with
            a minus sign as shown above, this is opposite to what is done for
            MultiModalSurrogate.
        """

        # get list of modes, but move (2,2) mode to start of the list.
        # This is important because we need the phase of the 22 mode to
        # transform the other modes from coorbital frame to inertial frame.
        self.mode_list = list(coorb_mode_data.keys())
        mode22_idx = [i for i in range(len(self.mode_list)) \
            if self.mode_list[i] == tuple([2, 2])]
        if len(mode22_idx) != 1:
            raise Exception('Seems to have found multiple or no 22 mode!')
        mode22_idx = mode22_idx[0]

        # shift 22 mode to the first index
        self.mode_list.insert(0, self.mode_list.pop(mode22_idx))

        if self.mode_list[0] != tuple([2, 2]):
            raise Exception('Expected the first mode at this point to be the'\
                ' 22 mode.')
        # make sure shifting the 22 mode index did not delete or add a
        # mode by mistake
        if len(self.mode_list) != len(coorb_mode_data.keys()):
            raise Exception('Number of modes do not agree')

        self.mode_type = 'identity'
        many_function_components = {}
        for mode in self.mode_list:
            many_function_components[mode] = ('identity', \
                coorb_mode_data[mode], {})

        # required for TaylorT3
        self.phaseAlignIdx = phaseAlignIdx
        self.TaylorT3_t_ref = TaylorT3_t_ref
        self.TaylorT3_factor_without_eta = None

        super(AlignedSpinCoOrbitalFrameSurrogate, self).__init__(name,
                domain, param_space, {}, many_function_components,
                self.mode_type)

        self._h5_data_keys.append('mode_list')
        self._h5_data_keys.append('mode_type')
        self._h5_data_keys.append('phaseAlignIdx')
        self._h5_data_keys.append('TaylorT3_t_ref')

    def _search_omega(self, omega22, omega_val):
        """ Find closest index such taht omega22[index] = omega_val
        """
        # find first index where omega22 > omega_val
        idx = np.where(omega22 > omega_val)[0][0]
        # if idx-1 is closer to omega_val, pick that instead
        if abs(omega22[idx-1] - omega_val) < abs(omega22[idx] - omega_val):
            idx -= 1
        return idx

    def _coorbital_to_inertial_frame(self, h_coorb, h_22, mode_list, dtM,
        timesM, fM_low, fM_ref, do_not_align):
        """ Transforms a dict from Coorbital frame to inertial frame.

            The surrogate data is sparsely sampled, so upsamples to time
            step dtM if given. This is done in the coorbital frame since
            the waveform is slowly varying in that frame.

            If fM_low is given, only part of the waveform where frequency of
            the (2, 2) mode is greater than fM_low is retained.

            if do_not_align = False:
                Aligns the 22 mode phase to be 0 at fM_ref. This means
                that at this reference frequency, the heavier BH is roughly on
                the +ve x axis and the lighter BH is on the -ve x axis.
            do_not_align should be True only when converting from pySurrogate
            format to gwsurrogate format as we may want to do some checks that
            the waveform has not been modified
        """

        Amp_22 = h_22[0]['amp']
        phi_22 = h_22[0]['phase']
        domain = np.copy(self.domain)

        # Get omega22_sparse, the angular frequency of the 22 mode, from the
        # sparse surrogate domain.
        # Use np.diff instead of np.gradient to match the LAL version
        omega22_sparse = np.append(np.diff(phi_22)/np.diff(domain), 0)

        # t=0 is at the waveform peak for the surrogate
        peak22Idx = np.argmin(np.abs(domain))
        omega22_peak = omega22_sparse[peak22Idx]
        # We ignore the part after the peak.  This way we avoid the noisy part
        # at late times, which can randomly be at frequency = fM_low.
        omega22_sparse = omega22_sparse[domain <= domain[peak22Idx]]

        # Get initIdx such that the initial (2, 2) mode frequency ~ fM_low.
        # We will make this more precise below.
        if fM_low != 0:
            omega_low = 2*np.pi*fM_low
            if omega_low < omega22_sparse[0]:
                raise ValueError('f_low is lower than the minimum allowed'
                    ' frequency')
            if omega_low > omega22_peak:
                raise ValueError('f_low is higher than the peak frequency')

            # Choose 5 indices less, to ensure omega_low is included
            initIdx = self._search_omega(omega22_sparse, omega_low) - 5
            # But if initIdx < 0, we are at the start of the surrogate data
            # so just choose 0
            if initIdx < 0:
                initIdx = 0

        else:
            # If fM_low is 0, we use the entire waveform
            initIdx = 0

            # But, if fM_low = 0 and timesM is given, the output of the
            # interpolant depends very slightly on the length of the sparse
            # data used to construct the interpolant. So, to achieve machine
            # precision equivalence between using dtM and timesM options, we
            # need to do the following: truncate before interpolation to the
            # same index as the dtM option would have used above (initIdx).
            # Using 6 rather than 5 because of the greater than condition.
            if timesM is not None:
                initIdx = np.where(domain > timesM[0])[0][0] - 6

        Amp_22 = Amp_22[initIdx:]
        phi_22 = phi_22[initIdx:]
        domain = domain[initIdx:]

        if timesM is not None:
            if timesM[-1] > domain[-1]:
                raise Exception("'times' includes times larger than the"
                    " maximum time value in domain.")
            if timesM[0] < domain[0]:
                raise Exception("'times' starts before start of domain. Try"
                    " increasing initial value of times or reducing f_low.")

        if dtM is None and timesM is None:
            # Use the sparse domain
            timesM = domain
            omega22 = omega22_sparse[initIdx:]
            do_interp = False
        else:
            ## Interpolate onto uniform-domain/timesM if needed
            do_interp = True
            if dtM is not None:
                t0 = domain[0]
                tf = domain[-1]
                num_times = int(np.ceil((tf - t0)/dtM));
                timesM = t0 + dtM*np.arange(num_times)
            else:
                if timesM[0] < domain[0] or timesM[-1] > domain[-1]:
                    raise Exception('Trying to evaluate at times outside the'
                        ' domain.')

            Amp_22 = _splinterp_Cwrapper(timesM, domain, Amp_22)
            phi_22 = _splinterp_Cwrapper(timesM, domain, phi_22)

            # now recompute omega22 with the dense data, but retain only data
            # upto the peak to avoid the noisy part
            omega22 = np.append(np.diff(phi_22)/np.diff(timesM), 0)
            omega22 = omega22[timesM <= 0]

            # Truncate data so that only freqs above omega_low are retained
            # If timesM are already given, we don't need to truncate data
            if dtM is not None:
                if fM_low != 0:
                    startIdx = self._search_omega(omega22, omega_low)
                else:
                    # If fM_low is 0, we use the entire waveform
                    startIdx = 0

                Amp_22 = Amp_22[startIdx:]
                phi_22 = phi_22[startIdx:]
                omega22 = omega22[startIdx:]
                timesM = timesM[startIdx:]

        # Get reference index where waveform needs to be aligned.
        if (abs(fM_ref-fM_low) < 1e-13) and (dtM is not None):
            # This means that the data is already truncated at fM_low,
            # so we just need the first index for fM_ref=fM_low
            refIdx = 0
        else:
            omega_ref = 2*np.pi*fM_ref
            if omega_ref > omega22_peak:
                raise ValueError('f_ref is higher than the peak frequency')

            refIdx = self._search_omega(omega22, omega_ref)

        # do_not_align should be True only when converting from pySurrogate
        # format to gwsurrogate format as we may want to do some checks that
        # the waveform has not been modified
        if not do_not_align:
            # Set orbital phase to 0 refIdx. Note that the Coorbital
            # frame data is not affected by this constant phase shift.

            # The orbital phase is obtained as phi_22/2, so this leaves a pi
            # ambiguity.  But the surrogate data is already aligned such that
            # the heavier BH is on the +ve x-axis at t=-1000M. See Sec.VI.A.4
            # of arxiv:1812.07865, the resolves the pi ambiguity. This means
            # that the after the realignment, the orbital phase at reference
            # frequency is 0.
            phi_22 += -phi_22[refIdx]

        h_dict = {}
        for mode in mode_list:
            if mode == tuple([2, 2]):
                h_dict[mode] = Amp_22 * np.exp(-1j*phi_22)
            else:
                l,m = mode
                h_coorb_lm = 0
                if 're' in h_coorb[mode][0].keys():
                    h_coorb_lm += h_coorb[mode][0]['re'] + 1j * 0
                if 'im' in h_coorb[mode][0].keys():
                    h_coorb_lm += 1j*h_coorb[mode][0]['im']

                h_coorb_lm = h_coorb_lm[initIdx:]
                if do_interp:
                    h_coorb_lm = _splinterp_Cwrapper(timesM,domain,h_coorb_lm)

                h_dict[mode] = h_coorb_lm * np.exp(-1j*m*phi_22/2.)

        return timesM, h_dict, None     # None is for dynamics

    def _set_TaylorT3_factor(self):
        """ Sets a term used in the 0 PN TaylorT3 phase. See Eq.43 of
        arxiv.1812.07865.
        """
        # Set only once
        if self.TaylorT3_factor_without_eta is None:
            # TaylorT3_t_ref is arbitrary. This is where the phase diverges,
            # so we choose it much after ringdown. This matches what was used
            # in the construction of the surrogate. See discussion near Eq.43
            # of arxiv.1812.07865
            theta_without_eta = ((self.TaylorT3_t_ref -self.domain)/5)**(-1./8)
            self.TaylorT3_factor_without_eta = -2./theta_without_eta**5

    def _TaylorT3_phase_22(self, x):
        """ 0 PN TaylorT3 phase. See Eq.43 of arxiv.1812.07865
        """

        q, chi1z, chi2z = x
        eta = q/(1.+q)**2

        # 0PN TaylorT3 phase
        phi22_T3 = 1./eta**(3./8) * self.TaylorT3_factor_without_eta

        # Align at phaseAlignIdx
        phi22_T3 -= phi22_T3[self.phaseAlignIdx]

        return phi22_T3

    def __call__(self, x, fM_low=None, fM_ref=None, dtM=None,
            timesM=None, dfM=None, freqsM=None, mode_list=None, ellMax=None,
            precessing_opts=None, tidal_opts=None, par_dict=None,
            return_dynamics=False, do_not_align=False):
        """
    Return dimensionless surrogate modes.
    Arguments:
    x :             The intrinsic parameters EXCLUDING total Mass (see
                    self.param_space)

    fM_low :        Initial frequency of (2,2) mode in units of cycles/M.
                    If 0, will use the entire data of the surrogate.
                    Default None.

    fM_ref:         Frequency used to set the reference epoch at which
                    the reference frame is defined and the spins are specified.
                    See below for definition of the reference frame.
                    Default: None.

                    For time domain models, f_ref is used to determine a t_ref,
                    such that the frequency of the (2, 2) mode equals f_ref at
                    t=t_ref.

    dtM :           Uniform time step to use, in units of M. If None, the
                    returned time array will be the array used in the
                    construction of the surrogate, which can be nonuniformly
                    sampled.
                    Default None.

    timesM:         Time samples to evaluate the waveform at. Use either dtM or
                    timesM, not both.

    dfM :           This should always be None as for now we are assuming
                    a time domain model.

    freqsM:         Frequency samples to evaluate the waveform at. Use either
                    dfM or freqsM, not both.

    ellMax:         Maximum ell index for modes to include. All available m
                    indicies for each ell will be included automatically.
                    Default: None, in which case all available modes wll be
                    included.

    mode_list :     A list of (ell, m) modes to be evaluated.
                    Default None, which evaluates all avilable modes.
                    Will deduce the m<0 modes from m>0 modes.

    par_dict:       This should always be None for this model.

    do_not_align:   Ignore fM_ref and do not align the waveform. This should be
                    True only when converting from pySurrogate format to
                    gwsurrogate format as we may want to do some checks that
                    the waveform has not been modified.

    Returns
    timesM, h, dynamics:
        timesM : time array in units of M.
        h : A dictionary of waveform modes sampled at timesM with
            (ell, m) keys.
        dynamics: None, since this is a nonprecessing model.

    IMPORTANT NOTES:
    ===============

    The reference frame (or inertial frame) is defined as follows:
        The +ve z-axis is along the orbital angular momentum at the reference
        epoch. The separation vector from the lighter BH to the heavier BH at
        the reference epoch is along the +ve x-axis. The y-axis completes the
        right-handed triad. The reference epoch is set using f_ref.
        """

        if dfM is not None:
            raise ValueError('Expected dfM to be None for a Time domain model')
        if freqsM is not None:
            raise ValueError('Expected freqsM to be None for a Time domain'
                ' model')

        if mode_list is None:
            mode_list = self.mode_list
        if ellMax is not None:
            if ellMax > np.max(np.array(self.mode_list).T[0]):
                raise ValueError('ellMax is greater than max allowed ell.')
            include_modes = np.array(self.mode_list).T[0] <= ellMax
            mode_list = [self.mode_list[idx]
                    for idx in range(len(self.mode_list))
                    if include_modes[idx]]

        if par_dict is not None:
            raise ValueError('par_dict should be None for this model')

        # always evaluate the (2,2) mode, the other modes neeed this
        # for transformation from coorbital to inertial frame

        # At this stage the phase of the (2,2) mode is the residual after
        # removing the TaylorT3 part (see. Eq.44 of arxiv.1812.07865)
        h_22 = self._eval_sur(x, tuple([2, 2]))

        # Get the TaylorT3 part and add to get the actual phase
        self._set_TaylorT3_factor()
        h_22[0]['phase'] += self._TaylorT3_phase_22(x)

        h_coorb = {k: self._eval_sur(x, k) for k in mode_list \
                        if k != tuple([2,2])}

        return self._coorbital_to_inertial_frame(h_coorb, h_22, \
            mode_list, dtM, timesM, fM_low, fM_ref, do_not_align)

class AlignedSpinCoOrbitalFrameSurrogateTidal(AlignedSpinCoOrbitalFrameSurrogate):
    """
    A surrogate for coorbital frame multimodal waveforms, where each waveform
    data piece has its own surrogate.

    The waveform data pieces are:
    Amplitude and phase of the (2,2) mode.
    Real and imaginary parts of coorbital frame waveform for other modes.

    This generates tidal inspiral waveforms by taking the surrogate output tuned
    to BBH results and incorporates the PN tidal corrections according to the
    tidal splicing method

    NOTE: This returns the waveform only during the inspiral portion of the
    binary's evolution, where the PN expansion is still valid; additional work
    will need to be done in order to complete the merger/ringdown portion of the
    waveform

    NOTE: The waveform is output with the time set so that t=0 corresponds to
    the peak of the waveform for the BBH waveform from the underlying surrogate,
    and NOT the peak of the tidally spliced waveform
    """

    def _coorbital_to_inertial_frame(self, h_coorb, h_22, mode_list, dtM,
        timesM, fM_low, fM_ref, do_not_align, x):
        """ Transforms a dict from Coorbital frame to inertial frame.

            The surrogate data is sparsely sampled, so upsamples to time
            step dtM if given. This is done in the coorbital frame since
            the waveform is slowly varying in that frame.

            If fM_low must be specified. The option of fM_low == 0 has been
            turned off for this model because of its excessive computational
            cost to evaluate

            if do_not_align = False:
                Aligns the 22 mode phase to be 0 at fM_ref. This means
                that at this reference frequency, the heavier BH is roughly on
                the +ve x axis and the lighter BH is on the -ve x axis.
            do_not_align should be True only when converting from pySurrogate
            format to gwsurrogate format as we may want to do some checks that
            the waveform has not been modified
        """

        Amp_22 = h_22[0]['amp']
        phi_22 = h_22[0]['phase']
        domain = np.copy(self.domain)

        # Get omega22_sparse, the angular frequency of the 22 mode, from the
        # sparse surrogate domain.
        # Use np.gradient
        omega22_sparse = np.gradient(phi_22, domain)

        # t=0 is at the waveform peak for the surrogate
        peak22Idx = np.argmin(np.abs(domain))
        omega22_peak = omega22_sparse[peak22Idx]
        # We ignore the part after the peak.  This way we avoid the noisy part
        # at late times, which can randomly be at frequency = fM_low.
        omega22_sparse = omega22_sparse[domain <= domain[peak22Idx]]

        # Get initIdx such that the initial (2, 2) mode frequency ~ fM_low.
        # We will make this more precise below.
        if fM_low != 0:
            omega_low = 2*np.pi*fM_low
            if omega_low < omega22_sparse[0]:
                raise ValueError('f_low is lower than the minimum allowed'
                    ' frequency')
            if omega_low > omega22_peak:
                raise ValueError('f_low is higher than the peak frequency')

            # Choose 5 indices less, to ensure omega_low is included
            initIdx = self._search_omega(omega22_sparse, omega_low) - 5
            # But if initIdx < 0, we are at the start of the surrogate data
            # so just choose 0
            if initIdx < 0:
                initIdx = 0
        else:
            raise ValueError("The option of setting 'fM_low' to 0 is turned off"
                    " for this model; must specifiy a non-zero 'fM_low'")
            ## If fM_low is 0, we use the entire waveform where frequency is
            ## monotonic, uncomment if want to allow this option
            #freq_orbital = np.gradient(phi_22[:peak22Idx], domain[:peak22Idx])
            #if np.min(np.diff(freq_orbital))<=0:
            #  initIdx = len(freq_orbital)-np.argmin((np.diff(freq_orbital)>np.zeros(len(freq_orbital)-1))[::-1])-1
            #else:
            #  initIdx = 0

        Amp_22 = Amp_22[initIdx:peak22Idx]
        phi_22 = phi_22[initIdx:peak22Idx]
        domain = domain[initIdx:peak22Idx]
        v_domain = np.power(np.abs(np.gradient(phi_22, domain))/2,1./3.)
        if(np.min(np.diff(v_domain))<0):
            raise ValueError('frequency is not monotonic over the entire'
                ' considered here')

        if timesM is not None:
            # This check is performed after the tidal terms computed
            #if timesM[-1] > domain[-1]:
            #    raise Exception("'times' includes times larger than the"
            #        " maximum time value in domain.")
            if timesM[0] < domain[0]:
                raise Exception("'times' starts before start of domain. Try"
                    " increasing initial value of times or reducing f_low.")

        # For tidal splicing, always want to interpolate first to a dense domain
        # in order to compute an accurate orbital frequency for the PN equations
        # then interpolated to the desired times afterwards

        if dtM is None and timesM is None:
            raise ValueError("For this model, must specify either the 'dtM' or"
                " 'timesM' option")
        else:
            ## Interpolate onto uniform domain
            ## WARNING -- if the the time points are not sampled densely enough
            ## here, there is a potential for error due to inaccurate orbital
            ## freq being used for the PN tidal equations
            if dtM is not None:
                t0 = domain[0]
                tf = domain[-1]
                num_times = int(np.ceil((tf - t0)/dtM));
                timesM_tmp = t0 + dtM*np.arange(num_times)
            else:
                # Because the spliced waveform is shifted so the final time
                # is the peak of the final waveform, we must ensure the check
                # here is performed similarly
                if timesM[0] < (domain[0]-domain[-1]) or timesM[-1] > 0:
                    raise Exception('Trying to evaluate at times outside the'
                        ' domain.')
                min_dt = np.min(np.diff(timesM))
                t0 = domain[0] #timesM[0] - min_dt
                tf = domain[-1]
                num_times = int(np.ceil((tf - t0)/min_dt));
                timesM_tmp = t0 + min_dt*np.arange(num_times)

            Amp_22 = _splinterp_Cwrapper(timesM_tmp, domain, Amp_22)
            phi_22 = _splinterp_Cwrapper(timesM_tmp, domain, phi_22)

            # now recompute omega22 with the dense data, but retain only data
            # upto the peak to avoid the noisy part
            omega22 = np.gradient(phi_22, timesM_tmp)

            #omega22 = omega22[timesM_tmp <= timesM_tmp[np.argmax(Amp_22)]]

            # Truncate data so that only freqs above omega_low are retained
            # If timesM are already given, we don't need to truncate data
            if dtM is not None:
                if fM_low != 0:
                    startIdx = max(np.argmin(np.abs(omega22 - omega_low)) - 4,0)
                else:
                    raise ValueError("The option of setting 'fM_low' to 0"
                            " is turned off for this model; must specifiy a"
                            " non-zero 'fM_low'")
                    ## If fM_low is 0, we use the entire waveform that is monotonic
                    #startIdx = 0
                    #if np.min(np.diff(omega22))<=0:
                    #  startIdx = len(omega22)-np.argmin((np.diff(omega22)>np.zeros(len(omega22)-1))[::-1])-1
                    ## Because the splicing changes the frequencies slightly, to
                    ## ensure we have wiggle room for interpolation later, buffer
                    ## the altered initial frequency of the spliced waveform so
                    ## it is not less than the initial frequency of v_domain
                    #gap = int((domain[5]-domain[0])/(timesM_tmp[1]-timesM_tmp[0]))
                    #if startIdx<gap:
                    #  startIdx=gap

                Amp_22 = Amp_22[startIdx:]
                phi_22 = phi_22[startIdx:]
                omega22 = omega22[startIdx:]
                timesM_tmp = timesM_tmp[startIdx:]

        freq_orbital = np.abs(omega22)/2
        v = np.power(freq_orbital,1./3.)

        # Setup all of the tidal parameters
        # Use universal relations to compute parameters beyond the quad love num
        # NOTE: omega2AB and omega3AB are stored as M*omega{2,3}{A,B}, to use the
        #   dimensionless value set the total mass of the system and that the
        #   universal relations return M{A,B}*omega{2,3}{A,B}
        # Aqm is the dimensionless quadrupole moment, however for splicing the 2PN BBH
        #   (v^4) term, the qm of a BBH must be subtracted off (Aqm_BH = 1; see
        #   arXiv:gr-qc/9709032 just below eqn 8), which will be done in the tidal
        #   function itself
        # If the NS is spinning, the effective driving frequency that the NS sees,
        #   from its own reference frame, is shifted according to the NS spin by the
        #   dimensionless rotation = M omega_spin =  (M / m_NS) * chi_NS / \bar{I}
        # WARNING: This effect has been turned off (omega_spin = 0) for NS anti-
        #   aligned spins, b/c the resonance peak is shifted to early enough in the
        #   inspiral that the approximation being used to model it might be breaking
        #   down by the end of the late inspiral (only supposed to be good up until
        #   shortly after resonance)
        qqq = x[0]; chiAz = x[1]; chiBz = x[2]; lambda2A = x[3]; lambda2B = x[4]
        XA = qqq/(1.+qqq); XB = 1.-XA
        omega2A = lambda3A = omega3A = AqmA = 0.
        omega2B = lambda3B = omega3B = AqmB = 0.
        omegaSpinA = omegaSpinB = 0.
        ell2Adyn = ell2Adiss = ell2Bdyn = ell2Bdiss = np.zeros(len(timesM_tmp))
        ell3Adyn = ell3Bdyn = np.zeros(len(timesM_tmp))
        if(lambda2A>0):
            IbarA     = UniversalRelationLambda2ToI(lambda2A)
            omegaSpinA = max(chiAz,0) / IbarA / XA
            omega2A   = UniversalRelationLambda2ToOmega2(lambda2A)/XA
            lambda3A  = UniversalRelationLambda2ToLambda3(lambda2A)
            omega3A   = UniversalRelationLambda3ToOmega3(lambda3A)/XA
            AqmA      = UniversalRelationLambda2ToAqm(lambda2A)
            ell2Adyn  = EffectiveDeformabilityFromDynamicalTides \
                        (np.abs(freq_orbital-omegaSpinA),omega2A,2,qqq)
            ell3Adyn  = EffectiveDeformabilityFromDynamicalTides \
                        (np.abs(freq_orbital-omegaSpinA),omega3A,3,qqq)
        if(lambda2B>0):
            IbarB     = UniversalRelationLambda2ToI(lambda2B)
            omegaSpinB = max(chiBz,0) / IbarB / XB
            omega2B   = UniversalRelationLambda2ToOmega2(lambda2B)/XB
            lambda3B  = UniversalRelationLambda2ToLambda3(lambda2B)
            omega3B   = UniversalRelationLambda3ToOmega3(lambda3B)/XB
            AqmB      = UniversalRelationLambda2ToAqm(lambda2B)
            ell2Bdyn  = EffectiveDeformabilityFromDynamicalTides \
                        (np.abs(freq_orbital-omegaSpinB),omega2B,2,qqq)
            ell3Bdyn  = EffectiveDeformabilityFromDynamicalTides \
                        (np.abs(freq_orbital-omegaSpinB),omega3B,3,qqq)

        dt_tid, dp_tid = PNT2Tidal(v, qqq, lambda2A*ell2Adyn, \
                lambda3A*ell3Adyn, AqmA, chiAz, lambda2B*ell2Bdyn, \
                lambda3B*ell3Bdyn, AqmB, chiBz, order=5)

        timesM_tmp += dt_tid - dt_tid[0]

        # Limit the waveform to the last time in the array that is increasing
        find = np.argmin(np.diff(timesM_tmp)>0)
        if(find == 0):
            find = len(timesM_tmp)

        timesM_tmp = timesM_tmp[:find]

        # There is a small region of parameter space where the interpolation
        # behaves poorly at very late times due to oddly shaped steps, so we
        # need to check the final handful of steps for that and truncate as
        # needed to avoid interpolation failures
        numcheck = 500
        factorLimit = 2.
        tdiff = np.diff(timesM_tmp[-numcheck-1:])
        for i in range(len(tdiff)-1):
          if ((tdiff[i]>tdiff[i+1]*factorLimit) or (tdiff[i]<tdiff[i+1]/factorLimit)):
            find = len(timesM_tmp)-numcheck+i-1
            timesM_tmp = timesM_tmp[:find]
            break

        timesM_tmp -= timesM_tmp[-1]
        phi_22 = phi_22[:find] + 2.*(dp_tid[:find] - dp_tid[0])

        # Reinterpolate to the final time grid
        if dtM is not None:
            t0 = timesM_tmp[0]
            tf = timesM_tmp[-1]
            num_times = int(np.ceil((tf - t0)/dtM));
            timesM = t0 + dtM*np.arange(num_times)
            timesM -= timesM[-1] # Ensure peak amplitude at t=0
        else:
            if timesM[-1] > timesM_tmp[-1]:
                raise Exception("'times' includes times larger than the"
                    " maximum time value in domain after splicing. (Remember"
                    " that tidal effects cause the binary to merger earlier)")
            if timesM[0] < timesM_tmp[0]:
                raise Exception("'times' includes times smaller than the"
                    " initial time value in domain after splicing. (Remember"
                    " that tidal effects cause the binary to merger earlier)")

        # Find the 'v' corresponding to the final time array, then perform the
        # interpolation in the 'v' domain as that is where most of the PN
        # quantities are defined
        v_uniform = _splinterp_Cwrapper(timesM, timesM_tmp, v[:find])

        Amp_22 = _splinterp_Cwrapper(v_uniform, v[:find], Amp_22[:find])
        phi_22 = _splinterp_Cwrapper(v_uniform, v[:find], phi_22)
        freq_orbital = np.power(v_uniform,3.)

        # Dynamical Tidal deformability stuff on final array for strain
        # amplitudes
        ell2Adyn = ell2Adiss = ell2Bdyn = ell2Bdiss = np.zeros(len(timesM))
        if(lambda2A>0):
          ell2Adyn  = EffectiveDeformabilityFromDynamicalTides \
                      (np.abs(freq_orbital-omegaSpinA),omega2A,2,qqq)
          ell2Adiss = EffectiveDissipativeDynamicalTides \
                      (np.abs(freq_orbital-omegaSpinA),ell2Adyn,omega2A,XA)
        if(lambda2B>0):
          ell2Bdyn  = EffectiveDeformabilityFromDynamicalTides \
                      (np.abs(freq_orbital-omegaSpinB),omega2B,2,qqq)
          ell2Bdiss = EffectiveDissipativeDynamicalTides \
                      (np.abs(freq_orbital-omegaSpinB),ell2Bdyn,omega2B,XB)

        # Get reference index where waveform needs to be aligned.
        if (abs(fM_ref-fM_low) < 1e-13) and (dtM is not None):
            # This means that the data is already truncated at fM_low,
            # so we just need the first index for fM_ref=fM_low
            refIdx = 0
        else:
            omega_ref = 2*np.pi*fM_ref
            if omega_ref > omega22_peak:
                raise ValueError('f_ref is higher than the peak frequency')

            refIdx = np.argmin(np.abs(2.*freq_orbital - omega_ref))

        # do_not_align should be True only when converting from pySurrogate
        # format to gwsurrogate format as we may want to do some checks that
        # the waveform has not been modified
        if not do_not_align:
            # Set orbital phase to 0 refIdx. Note that the Coorbital
            # frame data is not affected by this constant phase shift.

            # The orbital phase is obtained as phi_22/2, so this leaves a pi
            # ambiguity.  But the surrogate data is already aligned such that
            # the heavier BH is on the +ve x-axis at t=-1000M. See Sec.VI.A.4
            # of arxiv:1812.07865, the resolves the pi ambiguity. This means
            # that the after the realignment, the orbital phase at reference
            # frequency is 0.
            phi_22 += -phi_22[refIdx]

        h_dict = {}
        for mode in mode_list:
            if mode == tuple([2, 2]):
                h_dict[mode] = (Amp_22+StrainTidalEnhancementFactor(2,2, \
                      qqq,(lambda2A*ell2Adiss),(lambda2B*ell2Bdiss),v_uniform)) \
                      * np.exp(-1j*phi_22)
            else:
                l,m = mode
                h_coorb_lm = 0
                if 're' in h_coorb[mode][0].keys():
                    h_coorb_lm += h_coorb[mode][0]['re'] + 1j * 0
                if 'im' in h_coorb[mode][0].keys():
                    h_coorb_lm += 1j*h_coorb[mode][0]['im']

                h_coorb_lm = h_coorb_lm[initIdx:peak22Idx]

                h_coorb_lm = _splinterp_Cwrapper(v_uniform,v_domain,h_coorb_lm)

                h_coorb_lm_amp = np.abs(h_coorb_lm)
                h_coorb_lm_phase = np.unwrap(np.angle(h_coorb_lm))
                h_coorb_lm_tid = StrainTidalEnhancementFactor(l,m,qqq, \
                        (lambda2A*ell2Adiss),(lambda2B*ell2Bdiss),v_uniform)
                h_dict[mode] = (h_coorb_lm_amp + h_coorb_lm_tid) \
                        * np.exp((-1j*m/2.)*phi_22+1j*h_coorb_lm_phase)

        return timesM, h_dict, None     # None is for dynamics

    def __call__(self, x, fM_low=None, fM_ref=None, dtM=None,
        timesM=None, dfM=None, freqsM=None, mode_list=None, ellMax=None,
        precessing_opts=None, tidal_opts=None, par_dict=None,
        do_not_align=False):
        """
    Return dimensionless surrogate modes.
    Arguments:
    x :             The intrinsic parameters EXCLUDING total Mass (see
                    self.param_space)

    fM_low :        Initial frequency of (2,2) mode in units of cycles/M.
                    If 0, will use the entire data of the surrogate.
                    Default None.

    fM_ref:         Frequency used to set the reference epoch at which
                    the reference frame is defined and the spins are specified.
                    See below for definition of the reference frame.
                    Default: None.

                    For time domain models, f_ref is used to determine a t_ref,
                    such that the frequency of the (2, 2) mode equals f_ref at
                    t=t_ref.

    dtM :           Uniform time step to use, in units of M. If None, the
                    returned time array will be the array used in the
                    construction of the surrogate, which can be nonuniformly
                    sampled.
                    Default None.

    timesM:         Time samples to evaluate the waveform at. Use either dtM or
                    timesM, not both.

    dfM :           This should always be None as for now we are assuming
                    a time domain model.

    freqsM:         Frequency samples to evaluate the waveform at. Use either
                    dfM or freqsM, not both.

    ellMax:         Maximum ell index for modes to include. All available m
                    indicies for each ell will be included automatically.
                    Default: None, in which case all available modes wll be
                    included.

    mode_list :     A list of (ell, m) modes to be evaluated.
                    Default None, which evaluates all avilable modes.
                    Will deduce the m<0 modes from m>0 modes.

    par_dict:       This should always be None for this model.

    do_not_align:   Ignore fM_ref and do not align the waveform. This should be
                    True only when converting from pySurrogate format to
                    gwsurrogate format as we may want to do some checks that
                    the waveform has not been modified.

    Returns
    timesM, h, dynamics:
        timesM : time array in units of M.
        h : A dictionary of waveform modes sampled at timesM with
            (ell, m) keys.
        dynamics: None, since this is a nonprecessing model.

    IMPORTANT NOTES:
    ===============

    The reference frame (or inertial frame) is defined as follows:
        The +ve z-axis is along the orbital angular momentum at the reference
        epoch. The separation vector from the lighter BH to the heavier BH at
        the reference epoch is along the +ve x-axis. The y-axis completes the
        right-handed triad. The reference epoch is set using f_ref.
        """

        if par_dict is not None:
            raise ValueError('Expected par_dict to be None.')
        if dfM is not None:
            raise ValueError('Expected dfM to be None for a Time domain model')
        if freqsM is not None:
            raise ValueError('Expected freqsM to be None for a Time domain'
                ' model')

        if mode_list is None:
            mode_list = self.mode_list
        if ellMax is not None:
            if ellMax > np.max(np.array(self.mode_list).T[0]):
                raise ValueError('ellMax is greater than max allowed ell.')
            include_modes = np.array(self.mode_list).T[0] <= ellMax
            mode_list = [self.mode_list[idx]
                    for idx in range(len(self.mode_list))
                    if include_modes[idx]]

        # The last to parameters are the tidal parameters and are not a part of
        # the base surrogate model
        x_sur = x[:-2]

        # always evaluate the (2,2) mode, the other modes neeed this
        # for transformation from coorbital to inertial frame

        # At this stage the phase of the (2,2) mode is the residual after
        # removing the TaylorT3 part (see. Eq.44 of arxiv.1812.07865)
        h_22 = self._eval_sur(x_sur, tuple([2, 2]))

        # Get the TaylorT3 part and add to get the actual phase
        self._set_TaylorT3_factor()
        h_22[0]['phase'] += self._TaylorT3_phase_22(x_sur)

        h_coorb = {k: self._eval_sur(x_sur, k) for k in mode_list \
                        if k != tuple([2,2])}

        return self._coorbital_to_inertial_frame(h_coorb, h_22, \
            mode_list, dtM, timesM, fM_low, fM_ref, do_not_align, x)

class SpEC_nonspinning_q10_surrogate(MultiModalSurrogate):
    """A special class for the SpEC nonspinning surrogate"""

    skip_2_0_mode = True

    def __call__(self, x, theta=None, phi=None, modes=None,
                 fake_neg_modes=True):
        """
        Return surrogate evaluation.
        Arguments:
            x : The intrinsic parameters (see self.param_space)
            theta/phi : polar and azimuthal angles of the direction of
                        gravitational wave emission. If given, sums up modes
                        and returns h_plus and h_cross (default returns modes)
            modes : A list of (ell, m) modes to be evaluated (default: all)
            fake_neg_modes: Deduce (ell, -m) modes from (ell, m) modes for m>0.
        Returns h:
            h : If theta and phi are None, h is a dictionary of waveform modes
                sampled at self.domain with (ell, m) keys.
                If theta and phi are given, h = h_plus - i * h_cross is a
                complex array given by the sum of the modes.
        """
        if modes is None:
            modes = self.modes

        if self.skip_2_0_mode:
            # removed by 2to3 tool
            #modes = filter(lambda mode: mode != (2, 0), modes)
            modes = [mode for mode in modes if mode != (2, 0)]

        h_modes = super(SpEC_nonspinning_q10_surrogate, self).__call__(
                x, modes=modes)

        if fake_neg_modes:
            new_modes = {}
            for (ell, m), h in h_modes.items(): # inefficient in py2
                if m > 0:
                    new_modes[(ell, -m)] = np.power(-1, ell) * h.conjugate()
            for k, v in new_modes.items(): # inefficient in py2
                h_modes[k] = v

        if theta is not None:
            return _mode_sum(h_modes, theta, phi)

        return h_modes


from __future__ import print_function, division

from sympy import Expr, sympify, Symbol, Matrix
from sympy.printing.pretty.stringpict import prettyForm
from sympy.core.containers import Tuple
from sympy.core.compatibility import is_sequence, string_types

from sympy.physics.quantum.dagger import Dagger
from sympy.physics.quantum.matrixutils import (
    numpy_ndarray, scipy_sparse_matrix,
    to_sympy, to_numpy, to_scipy_sparse
)

__all__ = [
    'QuantumError',
    'QExpr'
]

#-----------------------------------------------------------------------------
# Error handling
#-----------------------------------------------------------------------------

class QuantumError(Exception):
    pass

def _qsympify_sequence(seq):
    """Convert elements of a sequence to standard form.

    This is like sympify, but it performs special logic for arguments passed
    to QExpr. The following conversions are done:

    * (list, tuple, Tuple) => _qsympify_sequence each element and convert
      sequence to a Tuple.
    * basestring => Symbol
    * Matrix => Matrix
    * other => sympify

    Strings are passed to Symbol, not sympify to make sure that variables like
    'pi' are kept as Symbols, not the SymPy built-in number subclasses.

    Examples
    ========

    >>> from sympy.physics.quantum.qexpr import _qsympify_sequence
    >>> _qsympify_sequence((1,2,[3,4,[1,]]))
    (1, 2, (3, 4, (1,)))

    """

    return tuple(__qsympify_sequence_helper(seq))

def __qsympify_sequence_helper(seq):
    """
       Helper function for _qsympify_sequence
       This function does the actual work.
    """
    #base case. If not a list, do Sympification
    if not is_sequence(seq):
        if isinstance(seq, Matrix):
            return seq
        elif isinstance(seq, string_types):
            return Symbol(seq)
        else:
            return sympify(seq)

    # base condition, when seq is QExpr and also
    # is iterable.
    if isinstance(seq, QExpr):
        return seq

    #if list, recurse on each item in the list
    result = [__qsympify_sequence_helper(item) for item in seq]

    return Tuple(*result)

#-----------------------------------------------------------------------------
# Basic Quantum Expression from which all objects descend
#-----------------------------------------------------------------------------

class QExpr(Expr):
    """A base class for all quantum object like operators and states."""

    # In sympy, slots are for instance attributes that are computed
    # dynamically by the __new__ method. They are not part of args, but they
    # derive from args.

    # The Hilbert space a quantum Object belongs to.
    __slots__ = ['hilbert_space']

    is_commutative = False

    # The separator used in printing the label.
    _label_separator = u''

    @property
    def free_symbols(self):
        return {self}

    def __new__(cls, *args, **old_assumptions):
        """Construct a new quantum object.

        Parameters
        ==========

        args : tuple
            The list of numbers or parameters that uniquely specify the
            quantum object. For a state, this will be its symbol or its
            set of quantum numbers.

        Examples
        ========

        >>> from sympy.physics.quantum.qexpr import QExpr
        >>> q = QExpr(0)
        >>> q
        0
        >>> q.label
        (0,)
        >>> q.hilbert_space
        H
        >>> q.args
        (0,)
        >>> q.is_commutative
        False
        """

        # First compute args and call Expr.__new__ to create the instance
        args = cls._eval_args(args)
        if len(args) == 0:
            args = cls._eval_args(tuple(cls.default_args()))
        inst = Expr.__new__(cls, *args, **old_assumptions)
        # Now set the slots on the instance
        inst.hilbert_space = cls._eval_hilbert_space(args)
        return inst

    @classmethod
    def _new_rawargs(cls, hilbert_space, *args, **old_assumptions):
        """Create new instance of this class with hilbert_space and args.

        This is used to bypass the more complex logic in the ``__new__``
        method in cases where you already have the exact ``hilbert_space``
        and ``args``. This should be used when you are positive these
        arguments are valid, in their final, proper form and want to optimize
        the creation of the object.
        """

        obj = Expr.__new__(cls, *args, **old_assumptions)
        obj.hilbert_space = hilbert_space
        return obj

    #-------------------------------------------------------------------------
    # Properties
    #-------------------------------------------------------------------------

    @property
    def label(self):
        """The label is the unique set of identifiers for the object.

        Usually, this will include all of the information about the state
        *except* the time (in the case of time-dependent objects).

        This must be a tuple, rather than a Tuple.
        """
        if len(self.args) == 0:  # If there is no label specified, return the default
            return self._eval_args(list(self.default_args()))
        else:
            return self.args

    @property
    def is_symbolic(self):
        return True

    @classmethod
    def default_args(self):
        """If no arguments are specified, then this will return a default set
        of arguments to be run through the constructor.

        NOTE: Any classes that override this MUST return a tuple of arguments.
        Should be overridden by subclasses to specify the default arguments for kets and operators
        """
        raise NotImplementedError("No default arguments for this class!")

    #-------------------------------------------------------------------------
    # _eval_* methods
    #-------------------------------------------------------------------------

    def _eval_adjoint(self):
        obj = Expr._eval_adjoint(self)
        if obj is None:
            obj = Expr.__new__(Dagger, self)
        if isinstance(obj, QExpr):
            obj.hilbert_space = self.hilbert_space
        return obj

    @classmethod
    def _eval_args(cls, args):
        """Process the args passed to the __new__ method.

        This simply runs args through _qsympify_sequence.
        """
        return _qsympify_sequence(args)

    @classmethod
    def _eval_hilbert_space(cls, args):
        """Compute the Hilbert space instance from the args.
        """
        from sympy.physics.quantum.hilbert import HilbertSpace
        return HilbertSpace()

    #-------------------------------------------------------------------------
    # Printing
    #-------------------------------------------------------------------------

    # Utilities for printing: these operate on raw sympy objects

    def _print_sequence(self, seq, sep, printer, *args):
        result = []
        for item in seq:
            result.append(printer._print(item, *args))
        return sep.join(result)

    def _print_sequence_pretty(self, seq, sep, printer, *args):
        pform = printer._print(seq[0], *args)
        for item in seq[1:]:
            pform = prettyForm(*pform.right((sep)))
            pform = prettyForm(*pform.right((printer._print(item, *args))))
        return pform

    # Utilities for printing: these operate prettyForm objects

    def _print_subscript_pretty(self, a, b):
        top = prettyForm(*b.left(' '*a.width()))
        bot = prettyForm(*a.right(' '*b.width()))
        return prettyForm(binding=prettyForm.POW, *bot.below(top))

    def _print_superscript_pretty(self, a, b):
        return a**b

    def _print_parens_pretty(self, pform, left='(', right=')'):
        return prettyForm(*pform.parens(left=left, right=right))

    # Printing of labels (i.e. args)

    def _print_label(self, printer, *args):
        """Prints the label of the QExpr

        This method prints self.label, using self._label_separator to separate
        the elements. This method should not be overridden, instead, override
        _print_contents to change printing behavior.
        """
        return self._print_sequence(
            self.label, self._label_separator, printer, *args
        )

    def _print_label_repr(self, printer, *args):
        return self._print_sequence(
            self.label, ',', printer, *args
        )

    def _print_label_pretty(self, printer, *args):
        return self._print_sequence_pretty(
            self.label, self._label_separator, printer, *args
        )

    def _print_label_latex(self, printer, *args):
        return self._print_sequence(
            self.label, self._label_separator, printer, *args
        )

    # Printing of contents (default to label)

    def _print_contents(self, printer, *args):
        """Printer for contents of QExpr

        Handles the printing of any unique identifying contents of a QExpr to
        print as its contents, such as any variables or quantum numbers. The
        default is to print the label, which is almost always the args. This
        should not include printing of any brackets or parenteses.
        """
        return self._print_label(printer, *args)

    def _print_contents_pretty(self, printer, *args):
        return self._print_label_pretty(printer, *args)

    def _print_contents_latex(self, printer, *args):
        return self._print_label_latex(printer, *args)

    # Main printing methods

    def _sympystr(self, printer, *args):
        """Default printing behavior of QExpr objects

        Handles the default printing of a QExpr. To add other things to the
        printing of the object, such as an operator name to operators or
        brackets to states, the class should override the _print/_pretty/_latex
        functions directly and make calls to _print_contents where appropriate.
        This allows things like InnerProduct to easily control its printing the
        printing of contents.
        """
        return self._print_contents(printer, *args)

    def _sympyrepr(self, printer, *args):
        classname = self.__class__.__name__
        label = self._print_label_repr(printer, *args)
        return '%s(%s)' % (classname, label)

    def _pretty(self, printer, *args):
        pform = self._print_contents_pretty(printer, *args)
        return pform

    def _latex(self, printer, *args):
        return self._print_contents_latex(printer, *args)

    #-------------------------------------------------------------------------
    # Methods from Basic and Expr
    #-------------------------------------------------------------------------

    def doit(self, **kw_args):
        return self

    #-------------------------------------------------------------------------
    # Represent
    #-------------------------------------------------------------------------

    def _represent_default_basis(self, **options):
        raise NotImplementedError('This object does not have a default basis')

    def _represent(self, **options):
        """Represent this object in a given basis.

        This method dispatches to the actual methods that perform the
        representation. Subclases of QExpr should define various methods to
        determine how the object will be represented in various bases. The
        format of these methods is::

            def _represent_BasisName(self, basis, **options):

        Thus to define how a quantum object is represented in the basis of
        the operator Position, you would define::

            def _represent_Position(self, basis, **options):

        Usually, basis object will be instances of Operator subclasses, but
        there is a chance we will relax this in the future to accommodate other
        types of basis sets that are not associated with an operator.

        If the ``format`` option is given it can be ("sympy", "numpy",
        "scipy.sparse"). This will ensure that any matrices that result from
        representing the object are returned in the appropriate matrix format.

        Parameters
        ==========

        basis : Operator
            The Operator whose basis functions will be used as the basis for
            representation.
        options : dict
            A dictionary of key/value pairs that give options and hints for
            the representation, such as the number of basis functions to
            be used.
        """
        basis = options.pop('basis', None)
        if basis is None:
            result = self._represent_default_basis(**options)
        else:
            result = dispatch_method(self, '_represent', basis, **options)

        # If we get a matrix representation, convert it to the right format.
        format = options.get('format', 'sympy')
        result = self._format_represent(result, format)
        return result

    def _format_represent(self, result, format):
        if format == 'sympy' and not isinstance(result, Matrix):
            return to_sympy(result)
        elif format == 'numpy' and not isinstance(result, numpy_ndarray):
            return to_numpy(result)
        elif format == 'scipy.sparse' and \
                not isinstance(result, scipy_sparse_matrix):
            return to_scipy_sparse(result)

        return result

def split_commutative_parts(e):
    """Split into commutative and non-commutative parts."""
    c_part, nc_part = e.args_cnc()
    c_part = list(c_part)
    return c_part, nc_part

def split_qexpr_parts(e):
    """Split an expression into Expr and noncommutative QExpr parts."""
    expr_part = []
    qexpr_part = []
    for arg in e.args:
        if not isinstance(arg, QExpr):
            expr_part.append(arg)
        else:
            qexpr_part.append(arg)
    return expr_part, qexpr_part

def dispatch_method(self, basename, arg, **options):
    """Dispatch a method to the proper handlers."""
    method_name = '%s_%s' % (basename, arg.__class__.__name__)
    if hasattr(self, method_name):
        f = getattr(self, method_name)
        # This can raise and we will allow it to propagate.
        result = f(arg, **options)
        if result is not None:
            return result
    raise NotImplementedError(
        "%s.%s can't handle: %r" %
        (self.__class__.__name__, basename, arg)
    )


import numpy as np
import random
import statsmodels.api as sm
import itertools as it
import os
from statsmodels.sandbox.stats.multicomp import multipletests
from itertools import compress
from sys import argv

def read_gems(directory, file_name):
    """
    Read a GEM file of the form "PlinePgem".
    Args:
       directory (str): directory of the file location (ex: '/Users/kimm/')
       file_name (str): name of the file (ex: 'SHG0008H.Fragnum_PlinePgem')
    Returns:
       in_gems (list): tab-separated lists of lists
    """
    with open(directory + file_name) as f:
        next(f)
        for line in f:
            entries = line.strip().split("\t")
            frags = entries[4]
            gem_id = entries[0]
            in_gems.append([frags+'|'+gem_id])
    return in_gems

def read_chroms(directory, genome_name):
    """
    Read a tab-delimited text file with list of chromosomes and their sizes.
    Args:
       directory (str): directory of the file location (ex: '/Users/kimm/')
       genome_name (str): name of reference genome (ex: 'dm3', 'mm10', 'hg38', 'hg19', etc.)
                        Note: must have a text file <genome_name>.chrom.sizes in the directory.
    Returns:
       chrom_dict (dictionary): tab-separated lists of lists
    """
    chrom_dict = {}
    with open(directory + genome_name + '.chrom.sizes') as f:
        for line in f:
            tmp_list = line.strip().split("\t")
            chrom_dict[tmp_list[0]] = int(tmp_list[1])
    return chrom_dict

def extract_info(gem_list):
    """
    Extract necessary information from input gem.
    Args:
       gem_list (list): list with 10 items as encoded in PlinePgem file
    Returns:
       gem_id (str): GEM ID (ex: SHG0008H-1000-10000799-AAACACCCACTAGTAC-K00384-FA-1-0)
       chrom (str): chromosome name (ex: chr2L)
       frags (list of list): [start,end] for each fragment
                        (ex: [3251935, 3252561], [3413857, 3414485])
       span (int): start of leftmost fragment to end of rightmost fragment (ex: 162550)
       fraglen (list): length of each fragment (ex: [626, 628])
       f2fdist (list): distance between neighboring fragments (ex: [161296])
       fragnum (int): number of fragments
       fragstr (str): fragment coordinates string
       coord (str): gem coordinates
    """
    raw_frags = gem_list[0].split("|")[0].split(";")
    frags = []
    for j in range(len(raw_frags)):
        bedentry = raw_frags[j].split("(")[0].replace("-", ":").split(":")
        frags.append([int(x) for x in bedentry[1:]])
    if '|' in gem_list[0]:
        gem_id = gem_list[0].split("|")[1]
    else:
        gem_id = ''
    chrom = bedentry[0]
    span = frags[-1][1] - frags[0][0]
    fraglen = [x[1]-x[0] for x in frags]
    f2fdist = [frags[x][0]-frags[x-1][1] for x in range(1,len(frags))]
    fragnum = j + 1
    fragstr = ';'.join([chrom+':'+str(x[0])+'-'+str(x[1]) for x in frags])
    coord = chrom+':'+str(frags[0][0])+'-'+str(frags[-1][1])
    return gem_id, chrom, frags, span, fraglen, f2fdist, fragnum, fragstr, coord

def norm_shannon_ent(prob_list):
    """
    Compute the normalized Shannon entropy of a list of probabilities.
    Args:
       prob_list (list): list of probabilities (ex: [0.1, 0.4, 0.5])
    Returns:
       entropy (float): tab-separated lists of lists (ex: 0.85867)
    """
    if sum(prob_list) != 1:   ## input is count or frequency instead of probability
        prob_list = [i/sum(prob_list) for i in prob_list]
    entropy = sum([x*np.log2(1/x) for x in prob_list])/np.log2(len(prob_list))
    return float(format(entropy, '.5f'))

def create_pseudo_gems(f2f_bucket, fragnum_bucket, samp_size, file_name):
    """
    Create <samp_size> pseudoGEMs for each unique fragment number.
    Args:
       f2f_bucket (list): list of fragment-to-fragment distances, from which samples will be drawn
       fragnum_bucket (list): number of fragments in each GEM for all GEMs
       samp_size (int): number of pseudoGEMs to create
       file_name (str): file to write 1000 pseudoGEMs per frag_num
    Returns:
       pseudo_gem_tot_dist (dictionary): dictionary of a list of <samp_size> total distances for fragnum
       pseudo_gem_mean_ent (dictionary): dictionary of mean entropy for fragnum
    """
    unique_frag_num = sorted(list(set(fragnum_bucket)))
    f2f_bucket_size = len(f2f_bucket)
    bucket_array = np.array(f2f_bucket)
    pseudo_gem_tot_dist = {}
    pseudo_gem_mean_ent = {}
    np.random.seed(12345)
    with open(file_name, 'a') as file1:
        for i in unique_frag_num:
            rand_ind = list(np.random.choice(range(f2f_bucket_size), (i-1)*samp_size))
            count = 0
            sum_ent = 0
            tmp_tot_dist = []
            for j in range(0,len(rand_ind), (i-1)):
                tmp_dist = bucket_array[rand_ind[j:(j+i-1)]].tolist()
                # get normalized Shannon entropy of first 1000 pseudoGEMs if a given GEM has > 2 fragments
                if (count < 1000):
                    file1.write(str(i) + '\t')
                    file1.write(';'.join(map(str, tmp_dist)) + '\t')
                    if (i > 2):
                        tmp_ent = norm_shannon_ent(tmp_dist)
                        file1.write(str(tmp_ent) + '\n')
                        sum_ent = sum_ent + tmp_ent
                    else:
                        file1.write("NA" + '\n')
                    count += 1
                tmp_tot_dist.append(sum(tmp_dist))
            pseudo_gem_tot_dist[str(i)] = np.sort(tmp_tot_dist)
            pseudo_gem_mean_ent[str(i)] = float(format(sum_ent/1000, '.5f'))
        pseudo_gem_tot_dist['1'] = [] # empty for singleton
    file1.close()
    return pseudo_gem_tot_dist, pseudo_gem_mean_ent

def get_raw_pval(f2f_dist, pseudo_gem_dist, samp_size):
    """
    Get raw p-value by performing a distance test.
    Args:
       f2f_dist (list): fragment-to-fragment distances for a GEM (ex: [446,21231,249])
       pseudo_gem_dist (list): pseudo gem total distances of a given <fragnum>
       samp_size (int): number of pseudoGEMs to create (ex: 10000)
    Returns:
       raw_pval (float): raw p-value from a total distance test; '.' if singleton GEM
    """
    if not f2f_dist:
        pval = '.'
    else:
        tot_dist = sum(f2f_dist)
        count = max(1, np.searchsorted(pseudo_gem_dist, tot_dist))
        pval = count/samp_size
    return pval

def get_adj_pval(raw_pval_list, fdr_thresh, method):
    """
    Adjust raw pvalues by Benjamini Hochberg multiple testing adjustment.
    Args:
       raw_pval_list (list): list of raw p-values (ex: [0.1, 0.04, 0.1])
       fdr_thresh (float): false discovery rate (ex: 0.05)
       method (string): adjustment method (ex: 'fdr_bh')
    Returns:
       adj_pval_list (array): array of booleans and adjusted p-values (ex: [0.1, 0.1, 0.1])
    """
    adj_pval_list = multipletests(raw_pval_list, alpha = fdr_thresh, method = method)
    return(adj_pval_list)

def entropy_filter(in_gem, pseudo_gem_ent, cutoff_ent_filter):
    """
    Filter a GEM into subGEMs if potentially doublet or triplet (metric: norm shannon Entropy).
    Args:
       in_gem (string): input GEM with fragments and GEM ID (ex: 'chr2L:11-33;chr2L:44-55|tmp_gem_id')
       pseudo_gem_ent (float): mean of normalized Shannon entropies in 1000 pseudoGEM (ex: 0.12)
       cutoff_ent_filter (int): max_dist/sec_max_dist ratio cutoff threshold (ex: 2)
    Returns:
       num_cuts (int): number of cuts made in this filter (ex: 0,1, or 2)
       sub_gems (list of strings): [sub_gem1, sub_gem2, etc.]
    """
    gem_id, chrom, frags, span, fraglen, f2fdist, fragnum, frag_str, coord = list(extract_info(in_gem))
    num_cuts = 0
    sub_gems = []
    obs_ent = norm_shannon_ent(f2fdist)
    if obs_ent < pseudo_gem_ent[str(fragnum)]: # cut once
        max_dist = max(f2fdist)
        max_dist_ind = f2fdist.index(max_dist)
        sec_max_dist = sorted(set(f2fdist))[-2]
        sec_max_dist_ind = f2fdist.index(sec_max_dist)
        sub_frags = [frags[:(max_dist_ind+1)], frags[(max_dist_ind+1):]]
        if max_dist/sec_max_dist < cutoff_ent_filter: # cut twice
            sorted_ind = sorted([max_dist_ind, sec_max_dist_ind])
            sub_frags = [frags[:(sorted_ind[0]+1)], frags[(sorted_ind[0]+1):(sorted_ind[1]+1)], frags[(sorted_ind[1]+1):]]
        num_cuts = len(sub_frags)-1
        for i in range(len(sub_frags)):
            sub_gems.append(';'.join([chrom+':'+str(x[0])+'-'+str(x[1]) for x in sub_frags[i]])+'|'+gem_id+'-sub-'+str(len(sub_frags))+'-'+str(i+1))
    else:
        sub_gems = [frag_str+'|'+gem_id]
    return sub_gems, num_cuts

def write_master_result(out_gem_list, out_name):
    """
    Write out significance test results.
    Args:
       out_gem_list (list): list with 11 items including gem_id, frag_str, p-values, etc.
       out_name (string): output file name
    Returns:
       None
    """
    with open(out_name, 'a') as file1:
        for i in range(len(out_gem_list)):
            file1.write('\t'.join(map(str, out_gem_list[i])) + '\n')
    file1.close()

if __name__ == '__main__':
    """
    Final results with following columns:
    1) GEM ID 2) GEM coordinate 3) GEM span 4) fragnum 5) frag coord 6) orig/postEF 7) totdist
    8) pval1 9) adj.pval1 10) status1 11) pval2 12) adj.pval2 13) status2
    """
    ### Set parameters ###
    library_name = argv[1] ## Library name of our data ##
    genome_name = argv[2] ## Name of the reference genome ##
    fdr_thresh = float(argv[3])  ## Benjamini-Hochberg FDR; p-value cutoff ##
    cutoff_ent_filter = int(argv[4]) ## Ratio of largest to second largest F2F cutoff in our entropy filter ##
    samp_size = int(argv[5]) ## Number of pseudo-GEMs ##
    prefix = library_name + "_FDR_" + str(fdr_thresh) + "_ratiothresh_"+ str(cutoff_ent_filter) + "_pseudoGEM_"+str(samp_size)

    ### Set directory and input file name ###
    directory = argv[6]
    file_name = argv[7]
    chrom_dir = argv[8]
    out_directory = directory + library_name + "_DistTest_FDR_" + str(fdr_thresh) + '/'
    if not os.path.exists(out_directory):
        os.mkdir(out_directory)

    #### Log file ####
    out = open(out_directory + prefix + "_distTest_logFile.txt", "a")

    out.write("Software version: v0.2 (2019-10-08, Kim)" + "\n")
    out.write("Directory: " + directory + "\n")
    out.write("File name: " + file_name + "\n")
    out.write("Library name: " + library_name + "\n")
    out.write("Reference genome: " + genome_name + "\n")
    out.write("FDR threshold: " + str(fdr_thresh) + "\n")
    out.write("Ratio threshold in entropy filter: " + str(cutoff_ent_filter) + "\n")
    out.write("Number of pseudo-GEMs: " + str(samp_size) + "\n")
    out.write("Started processing domain-based distance test. \n")
    out.write("================================= \n")

    ### Read input GEM file ###
    in_gems = []
    in_gems = read_gems(directory, file_name)
    out.write("Finished reading the input GEM file. \n")
    out.write(str(len(in_gems)) + " total GEMs. \n")
    out.write("================================= \n")

    ### Read chromosome names and sizes ###
    chrom_dict = read_chroms(chrom_dir, genome_name)

    header = ['GEM_ID', 'GEM_coord', 'GEM_span', 'Frag_number', 'List_of_frag_coord', 'category', 'dist', 'rawpval1', 'adjpval1', 'decis1', 'rawpval2', 'adjpval2', 'decis2']
    with open(out_directory+prefix+'_distTest_master.txt', 'a') as file1:
        file1.write('\t'.join(map(str, header)) + '\n')
    file1.close()

    tot_pass = 0
    for chr_name in chrom_dict.keys():
        ### Subset input GEM files by chromosome ###
        out.write("================================= \n")
        out.write("===== Chromosome is: " + chr_name + ". ===== \n")
        out.write("================================= \n")
        subset_gems = [x for x in in_gems if chr_name == x[0].split(":")[0]]
        out.write("Finished subsetting GEMs. \n")
        out.write(str(len(subset_gems)) + " GEMs in " + chr_name + " (of " + str(len(in_gems)) + " total GEMs). \n")
        if len(subset_gems) < 100:
            out.write("Skipped: less than 100 GEMs. \n")
            out.write("================================= \n")
            continue

        ### Initialize output GEM list ###
        out_gems = []

        ### Create distance bucket ###
        f2f_bucket = []
        fragnum_bucket = []
        for i in range(len(subset_gems)):
            gem_id, chrom, frags, span, fraglen, f2fdist, fragnum, frag_str, coord = extract_info(subset_gems[i])
            f2f_bucket.extend(f2fdist)
            fragnum_bucket.append(fragnum)
        out.write("Maximum number of fragments in a GEM: "+str(max(fragnum_bucket))+"\n")

        ### Create pseudoGEMS ###i
        pseudo_file = out_directory + prefix + '_distTest_pseudo_' + chr_name + '.txt'
        pseudo_gem_tot_dist, pseudo_gem_mean_ent = create_pseudo_gems(f2f_bucket, fragnum_bucket, samp_size, pseudo_file)
        out.write("Finished creating "+str(samp_size)+" pseudoGEMs and total distance for each fragment number class. \n")
        del f2f_bucket

        ### Calculate raw p-values for subset_gems via distance test ###
        out.write("================================= \n")
        out.write(" 1) Distance test. \n")
        out.write("================================= \n")
        for k in range(len(subset_gems)):
            gem_id, chrom, frags, span, fraglen, f2fdist, fragnum, frag_str, coord = list(extract_info(subset_gems[k]))
            raw_pval = get_raw_pval(f2fdist, pseudo_gem_tot_dist[str(fragnum)], samp_size)
            out_gems.append([gem_id, coord, span, fragnum, frag_str, 'Orig', sum(f2fdist), raw_pval])
        out.write("Finished computing raw p-values for "+str(k+1)+" GEMs. \n")

        ### Adjust p-values by fragment number classes in each chromosome ###
        deferred_gems = []
        for fn in sorted(list(set(fragnum_bucket))):
            indx = [i for i,val in enumerate(fragnum_bucket) if val==fn]
            pvals1 = [out_gems[k][7] for k in indx]
            adj_pval1 = get_adj_pval(pvals1, fdr_thresh, 'fdr_bh')
            for k in range(len(indx)):
                out_gems[indx[k]].append(round(adj_pval1[1][k], 4))
                # Subcategories from first distance test #
                if adj_pval1[0][k]==True:
                    if fn < 101: # GEMs with <= 100 fragments pass; 20191008
                        out_gems[indx[k]].extend(['PASS', '.', '.', '.'])
                    else: # GEMs with > 100 fragments fail; 20191008
                        out_gems[indx[k]].extend(['FAIL', '.', '.', '.'])
                elif adj_pval1[0][k]==False:
                    if fn==2: # GEMs with 2 fragments fail
                        out_gems[indx[k]].extend(['FAIL', '.', '.', '.'])
                    else: # GEMs with > 2 fragments go to entropy filter
                        out_gems[indx[k]].extend(['DEFER', '.', '.', '.'])
                        deferred_gems.append(subset_gems[indx[k]])
        out.write("Finished adjusting p-values. \n")
        out.write("First test PASS: "+str(len([x for x in out_gems if x[9] =='PASS']))+"\n")
        out.write("First test FAIL: "+str(len([x for x in out_gems if x[9] =='FAIL']))+"\n")
        out.write("First test DEFER: "+str(len([x for x in out_gems if x[9] =='DEFER']))+"\n")
        pass1 = len([x for x in out_gems if x[9] =='PASS'])

        del pvals1
        del adj_pval1

        ### Entropy Filter for those in 'DEFER' category ###
        out.write("================================= \n")
        out.write(" 2) Entropy filter. \n")
        out.write("================================= \n")
        ef_gems = []
        num_frag_list = []
        num_cut_list = []
        for k in range(len(deferred_gems)):
            sub_gems, num_cuts = entropy_filter(deferred_gems[k], pseudo_gem_mean_ent, cutoff_ent_filter)
            num_cut_list.append(num_cuts)
            for i in range(len(sub_gems)):
                gem_id, chrom, frags, span, fraglen, f2fdist, fragnum, fragstr, coord = list(extract_info([sub_gems[i]]))
                num_frag_list.append(fragnum)
                raw_pval = get_raw_pval(f2fdist, pseudo_gem_tot_dist[str(fragnum)], samp_size)
                ef_gems.append([gem_id, coord, span, fragnum, fragstr, 'EF-'+str(num_cuts), sum(f2fdist),'.','.', 'DEFER', raw_pval])
        out.write("Finished entropy filter for removing doublets. \n")
        out.write("Number of GEMs not split: "+str(len([x for x in num_cut_list if x == 0]))+"\n")
        out.write("Number of GEMs split once: "+str(len([x for x in num_cut_list if x == 1]))+"\n")
        out.write("Number of GEMs split twice: "+str(len([x for x in num_cut_list if x == 2]))+"\n")
        out.write("Number of non-singleton GEMs after cutting: "+str(len([x for x in num_frag_list if x > 1]))+"\n")
        out.write("Number of singleton GEMs after cutting: "+str(len([x for x in num_frag_list if x == 1]))+"\n")

        ### Adjust p-values by fragment number classes in each chromosome ###
        out.write("================================= \n")
        out.write(" 3) Distance test.  \n")
        out.write("================================= \n")
        sub_fragnum_bucket = [x[3] for x in ef_gems]
        for fn in sorted(list(set(sub_fragnum_bucket))):
            indx = [i for i,val in enumerate(sub_fragnum_bucket) if val==fn]
            if fn==1:
                for k in range(len(indx)):
                    ef_gems[indx[k]].extend(['.','.'])
            else:
                pvals2 = [ef_gems[k][10] for k in indx]
                adj_pval2 = get_adj_pval(pvals2, fdr_thresh, 'fdr_bh')
                for k in range(len(indx)):
                    ef_gems[indx[k]].append(round(adj_pval2[1][k], 4))
                    if adj_pval2[0][k]==True:
                        ef_gems[indx[k]].append('PASS')
                    elif adj_pval2[0][k]==False:
                        ef_gems[indx[k]].append('FAIL')
        out.write("Finished adjusting p-values. \n")
        out.write("Second test PASS: "+str(len([x for x in ef_gems if x[12] =='PASS']))+"\n")
        out.write("Second test FAIL: "+str(len([x for x in ef_gems if x[12] =='FAIL']))+"\n")
        pass2 = len([x for x in ef_gems if x[12] =='PASS'])
        out.write("================================= \n")

        del pvals2
        del adj_pval2

        ### Write results ###
        write_master_result(out_gems, out_directory+prefix+'_distTest_master.txt')
        write_master_result(ef_gems, out_directory+prefix+'_distTest_master.txt')
        out.write("Finished writing files. \n")

        tot_pass += pass1+pass2

        out.write("Total " +  str(pass1+pass2) + " GEMs out of " + str(len(subset_gems))+ " GEMs passed. \n")
        out.write("i.e., " + str(round((pass1+pass2)*100/len(subset_gems), 3)) + "% passed in " +chr_name+ ". \n")
        out.write("================================= \n")
        del out_gems
        del ef_gems
        del deferred_gems
    out.write("================================= \n")
    out.write("======= Summary Statistics ======= \n")
    out.write("Total " +  str(tot_pass) + " GEMs out of " + str(len(in_gems))+ " GEMs passed. \n")
    out.write("i.e., " + str(round((tot_pass)*100/len(in_gems), 3)) + "% passed. \n")
    out.write("================================= \n")
    out.write("DONE. \n")

    out.close()

import numpy,import statsmodels,from statsmodels