from . import ccllib as lib
from . import constants as const
from .core import check
import numpy as np

import collections

# Same mapping for non-Limber integration methods
nonlimber_methods = {
    'native': const.CCL_NONLIMBER_METHOD_NATIVE,
    'angpow': const.CCL_NONLIMBER_METHOD_ANGPOW,
}

function_types = {
    'dndz': const.CCL_CLT_NZ,
    'bias': const.CCL_CLT_BZ,
    'mag_bias': const.CCL_CLT_SZ,
    'red_frac': const.CCL_CLT_RF,
    'ia_bias': const.CCL_CLT_BA,
    'lensing_win': const.CCL_CLT_WL,
    'mag_win': const.CCL_CLT_WM,
}

# Define symbolic 'None' type for arrays, to allow proper handling by swig
# wrapper
NoneArr = np.array([])

class Tracer(object):
    """A tracer of the matter density field.

    .. note:: This class cannot be used directly. Use one of
              :obj:`NumberCountsTracer`, :obj:`WeakLensingTracer`
              or :obj:`CMBLensingTracer` instead.

    This class contains all information describing the transfer functon of
    a tracer (e.g., galaxy density, lensing shear) of the matter distribution.
    """
    def __init__(self, *args, **kwargs):
        raise NotImplementedError(
            "A `Tracer` object cannot be used directly. Use one of "
            "`NumberCountsTracer`, `WeakLensingTracer` or `CMBLensingTracer` "
            "instead.")

    def _build_tracer(
            self, cosmo, tracer_type, has_rsd=False,
            dndz=None, bias=None, mag_bias=None, ia_bias=None,
            red_frac=None, z_source=1100.):
        """Build the CCL_ClTracer.

        Args:
            cosmo (:obj:`Cosmology`): Cosmology object.
            tracer_type (:obj:): Specifies the type of tracer. Must be one of
                    const.CL_TRACER_NC: number count tracer
                    const.CL_TRACER_WL: lensing tracer
                    const.CL_TRACER_CL: CMB lensing tracer
            has_rsd (bool, optional): Flag for whether the tracer has a
                redshift-space distortion term. Defaults to False.
            dndz (tuple of arrays, optional): A tuple of arrays (z, N(z))
                giving the redshift distribution of the objects. The units are
                arbitrary; N(z) will be normalized to unity. If `None`, the
                tracer is assumed to not have a redshift distribution (e.g.,
                it has a single source source redshift like the CMB). Defaults
                to None.
            bias (tuple of arrays, optional): A tuple of arrays (z, b(z))
                giving the galaxy bias. If `None`, the tracer is assumbed to
                not have a bias parameter. Defaults to None.
            mag_bias (tuple of arrays, optional): A tuple of arrays (z, s(z))
                giving the magnification bias as a function of redshift. If
                `None`, the tracer is assumed to not have magnification bias
                terms. Defaults to None.
            ia_bias (tuple of arrays, optional): A tuple of arrays
                (z, b_IA(z)) giving the intrinsic alignment amplitude b_IA(z).
                If `None`, the tracer is assumped to not have intrinsic
                alignments. Defaults to None.
            red_frac (tuple of arrays,, optional): A tuple of arrays
                (z, f_red(z)) givng the red fraction of galaxies as a function
                of redshift. If `None`, then the tracer is assumed to not have
                a red fraction. Defaults to None.
            z_source (float, optional): Redshift of source plane for CMB
                lensing. Defaults to 1100.
        """

        # Verify cosmo object
        cosmo = cosmo.cosmo

        has_magnification = mag_bias is not None
        if (red_frac is None) != (ia_bias is None):
            raise ValueError(
                "Either both or none of `red_frac` and `ia_bias` "
                "must be specified.")
        has_intrinsic_alignment = red_frac is not None

        # Passing None for certain arguments causes segmentation faults at the
        # moment. The following checks try to guard against these instances
        # but this should probably be checked for at the C level.
        if tracer_type in [const.CL_TRACER_WL,
                           const.CL_TRACER_NC]:
            if not isinstance(dndz, collections.Iterable) \
               or len(dndz) != 2 \
               or not (isinstance(dndz[0], collections.Iterable)
                       and isinstance(dndz[1], collections.Iterable)):
                raise ValueError("dndz needs to be a tuple of two arrays.")
        if tracer_type in [const.CL_TRACER_NC]:
            if not isinstance(bias, collections.Iterable) \
               or len(bias) != 2 \
               or not (isinstance(bias[0], collections.Iterable)
                       and isinstance(bias[1], collections.Iterable)):
                raise ValueError("bias needs to be a tuple of two arrays.")

        # Convert array arguments that are 'None' into 'NoneArr' type and
        # check whether arrays were specified as tuples
        self.z_n, self.n = _check_array_params(dndz)
        self.z_b, self.b = _check_array_params(bias)
        self.z_s, self.s = _check_array_params(mag_bias)
        self.z_ba, self.ba = _check_array_params(ia_bias)
        self.z_rf, self.rf = _check_array_params(red_frac)
        self.z_source = z_source

        # Construct new ccl_cl_tracer
        status = 0
        return_val = lib.cl_tracer_new_wrapper(
                            cosmo,
                            tracer_type,
                            int(has_rsd),
                            int(has_magnification),
                            int(has_intrinsic_alignment),
                            self.z_n, self.n,
                            self.z_b, self.b,
                            self.z_s, self.s,
                            self.z_ba, self.ba,
                            self.z_rf, self.rf,
                            float(self.z_source),
                            status)

        if (isinstance(return_val, int)):
            self.has_cltracer = False
            check(return_val)
        else:
            self.has_cltracer = True
            self.cltracer, status = return_val

    def get_internal_function(self, cosmo, function, a):
        """
        Method to evaluate any internal function of redshift for this tracer.

        Args:
            cosmo (:obj:`Cosmology`): Cosmology object.
            function (:obj:`str`): Specifies which function to evaluate. Must
                be one of
                    'dndz': number density
                    'bias': bias
                    'mag_bias': magnification bias
                    'red_frac': red fraction
                    'ia_bias': intrinsic alignment bias
                    'lensing_win': weak lensing window function
                    'mag_win': magnification window function
            a (:obj: float or array-like): list of scale factors at which to
                evaluate the function.

        Returns:
            Array of function values at the input scale factors.
        """
        # Access ccl_cosmology object
        cosmo_in = cosmo
        cosmo = cosmo.cosmo

        # Check that specified function type exists
        if function not in function_types.keys():
            raise ValueError(
                "Internal function type '%s' not recognized."
                % function)

        # Check input types
        status = 0
        is_scalar = False
        if isinstance(a, float):
            is_scalar = True
            aarr = np.array([a])
            na = 1
        elif isinstance(a, np.ndarray):
            aarr = a
            na = a.size
        else:
            aarr = a
            na = len(a)

        # Evaluate function
        farr, status = lib.clt_fa_vec(cosmo, self.cltracer,
                                      function_types[function],
                                      aarr, na, status)
        check(status, cosmo_in)
        if is_scalar:
            return farr[0]
        else:
            return farr

    def __del__(self):
        """Free memory associated with CCL_ClTracer object.
        """
        if hasattr(self, 'has_cltracer'):
            if self.has_cltracer:
                lib.cl_tracer_free(self.cltracer)

class NumberCountsTracer(Tracer):
    """A Tracer for galaxy number counts (galaxy clustering).

    Args:
        cosmo (:obj:`Cosmology`): Cosmology object.
        has_rsd (bool): Flag for whether the tracer has a
            redshift-space distortion term.
        dndz (tuple of arrays): A tuple of arrays (z, N(z))
            giving the redshift distribution of the objects. The units are
            arbitrary; N(z) will be normalized to unity.
        bias (tuple of arrays): A tuple of arrays (z, b(z))
            giving the galaxy bias.
        mag_bias (tuple of arrays, optional): A tuple of arrays (z, s(z))
            giving the magnification bias as a function of redshift. If
            `None`, the tracer is assumed to not have magnification bias
            terms. Defaults to None.
    """

    def __init__(self, cosmo, has_rsd, dndz, bias, mag_bias=None):
        # Call Tracer constructor with appropriate arguments
        self._build_tracer(
            cosmo=cosmo, tracer_type=const.CL_TRACER_NC,
            has_rsd=has_rsd,
            dndz=dndz, bias=bias, mag_bias=mag_bias,
            ia_bias=None, red_frac=None)

class WeakLensingTracer(Tracer):
    """A Tracer for weak lensing shear (galaxy shapes).

    Args:
        cosmo (:obj:`Cosmology`): Cosmology object.
        dndz (tuple of arrays): A tuple of arrays (z, N(z))
            giving the redshift distribution of the objects. The units are
            arbitrary; N(z) will be normalized to unity.
        ia_bias (tuple of arrays, optional): A tuple of arrays
            (z, b_IA(z)) giving the intrinsic alignment amplitude b_IA(z).
            If `None`, the tracer is assumped to not have intrinsic
            alignments. Defaults to None.
        red_frac (tuple of arrays,, optional): A tuple of arrays
            (z, f_red(z)) givng the red fraction of galaxies as a function
            of redshift. If `None`, then the tracer is assumed to not have
            a red fraction. Defaults to None.
    """

    def __init__(self, cosmo, dndz, ia_bias=None, red_frac=None):
        # Call Tracer constructor with appropriate arguments
        self._build_tracer(
            cosmo=cosmo, tracer_type=const.CL_TRACER_WL,
            has_rsd=False,
            dndz=dndz, bias=None, mag_bias=None,
            ia_bias=ia_bias, red_frac=red_frac)

class CMBLensingTracer(Tracer):
    """A Tracer for CMB lensing.

    Args:
        cosmo (:obj:`Cosmology`): Cosmology object.
        z_source (float): Redshift of source plane for CMB lensing.
    """

    def __init__(self, cosmo, z_source):
        # Call Tracer constructor with appropriate arguments
        self._build_tracer(
            cosmo=cosmo, tracer_type=const.CL_TRACER_CL,
            has_rsd=False,
            dndz=None, bias=None, mag_bias=None,
            ia_bias=None, red_frac=None, z_source=z_source)

def _check_array_params(f_arg):
    """Check whether an argument `f_arg` passed into the constructor of
    Tracer() is valid.

    If the argument is set to `None`, it will be replaced with a special array
    that signals to the CCL wrapper that this argument is NULL.
    """
    if f_arg is None:
        # Return empty array if argument is None
        f = NoneArr
        z_f = NoneArr
    else:
        z_f = np.atleast_1d(np.array(f_arg[0], dtype=float))
        f = np.atleast_1d(np.array(f_arg[1], dtype=float))
    return z_f, f

def angular_cl(cosmo, cltracer1, cltracer2, ell,
               l_limber=-1., l_logstep=1.05, l_linstep=20., dchi=3.,
               dlk=0.003, zmin=0.05, non_limber_method="native"):
    """Calculate the angular (cross-)power spectrum for a pair of tracers.

    Args:
        cosmo (:obj:`Cosmology`): A Cosmology object.
        cltracer1, cltracer2 (:obj:`Tracer`): Tracer objects, of any kind.
        ell (float or array_like): Angular wavenumber(s) at which to evaluate
            the angular power spectrum.
        l_limber (float) : Angular wavenumber beyond which Limber's
            approximation will be used. Defaults to 1.
        l_logstep (float) : logarithmic step in ell at low multipoles.
            Defaults to 1.05.
        l_linstep (float) : linear step in ell at high multipoles.
            Defaults to 20.
        dchi (float) : comoving distance step size in non-limber native
            integrals. Defaults to 3.
        dlk (float) : logarithmic step for the k non-limber native integral.
            Defaults to 0.003.
        zmin (float) : minimal redshift for the integrals. Defualts to 0.05.
        non_limber_method (str) : non-Limber integration method. Supported:
            "native" and "angpow". Defaults to 'native'.

    Returns:
        float or array_like: Angular (cross-)power spectrum values,
            :math:`C_\\ell`, for the pair of tracers, as a function of
            :math:`\\ell`.
    """
    # Access ccl_cosmology object
    cosmo = cosmo.cosmo

    if non_limber_method not in nonlimber_methods.keys():
        raise ValueError(
            "'%s' is not a valid non-Limber integration method." %
            non_limber_method)

    # Access CCL_ClTracer objects
    clt1 = cltracer1.cltracer
    clt2 = cltracer2.cltracer

    status = 0
    # Return Cl values, according to whether ell is an array or not
    if isinstance(ell, float) or isinstance(ell, int):
        # Use single-value function
        cl_one, status = lib.angular_cl_vec(
            cosmo, clt1, clt2, l_limber, l_logstep, l_linstep, dchi, dlk, zmin,
            nonlimber_methods[non_limber_method], [ell], 1, status)
        cl = cl_one[0]
    elif isinstance(ell, np.ndarray):
        # Use vectorised function
        cl, status = lib.angular_cl_vec(
            cosmo, clt1, clt2, l_limber, l_logstep, l_linstep, dchi, dlk, zmin,
            nonlimber_methods[non_limber_method], ell, ell.size, status)
    else:
        # Use vectorised function
        cl, status = lib.angular_cl_vec(
            cosmo, clt1, clt2, l_limber, l_logstep, l_linstep, dchi, dlk, zmin,
            nonlimber_methods[non_limber_method], ell, len(ell), status)
    check(status)
    return cl


import numba
import numpy as np
import strax
from immutabledict import immutabledict
from strax.processing.general import _touching_windows
import straxen
from .pulse_processing import HITFINDER_OPTIONS, HITFINDER_OPTIONS_he, HE_PREAMBLE
from straxen.get_corrections import is_cmt_option
from .peaklet_classification import PeakletClassification

export, __all__ = strax.exporter()
FAKE_MERGED_S2_TYPE = -42

@export
@strax.takes_config(
    strax.Option('peaklet_gap_threshold', default=700, infer_type=False,
                 help="No hits for this many ns triggers a new peak"),
    strax.Option('peak_left_extension', default=30, infer_type=False,
                 help="Include this many ns left of hits in peaks"),
    strax.Option('peak_right_extension', default=200, infer_type=False,
                 help="Include this many ns right of hits in peaks"),
    strax.Option('peak_min_pmts', default=2, infer_type=False,
                 help="Minimum number of contributing PMTs needed to define a peak"),
    strax.Option('peak_split_gof_threshold',
                 # See https://xe1t-wiki.lngs.infn.it/doku.php?id=
                 # xenon:xenonnt:analysis:strax_clustering_classification
                 # #natural_breaks_splitting
                 # for more information
                 default=(
                     None,  # Reserved
                     ((0.5, 1.0), (6.0, 0.4)),
                     ((2.5, 1.0), (5.625, 0.4))), infer_type=False,
                 help='Natural breaks goodness of fit/split threshold to split '
                      'a peak. Specify as tuples of (log10(area), threshold).'),
    strax.Option('peak_split_filter_wing_width', default=70, infer_type=False,
                 help='Wing width of moving average filter for '
                      'low-split natural breaks'),
    strax.Option('peak_split_min_area', default=40., infer_type=False,
                 help='Minimum area to evaluate natural breaks criterion. '
                      'Smaller peaks are not split.'),
    strax.Option('peak_split_iterations', default=20, infer_type=False,
                 help='Maximum number of recursive peak splits to do.'),
    strax.Option('diagnose_sorting', track=False, default=False, infer_type=False,
                 help="Enable runtime checks for sorting and disjointness"),
    strax.Option('gain_model', infer_type=False,
                 help='PMT gain model. Specify as '
                 '(str(model_config), str(version), nT-->boolean'),
    strax.Option('tight_coincidence_window_left', default=50, infer_type=False,
                 help="Time range left of peak center to call "
                      "a hit a tight coincidence (ns)"),
    strax.Option('tight_coincidence_window_right', default=50, infer_type=False,
                 help="Time range right of peak center to call "
                      "a hit a tight coincidence (ns)"),
    strax.Option('n_tpc_pmts', type=int,
                 help='Number of TPC PMTs'),
    strax.Option('saturation_correction_on', default=True, infer_type=False,
                 help='On off switch for saturation correction'),
    strax.Option('saturation_reference_length', default=100, infer_type=False,
                 help="Maximum number of reference sample used "
                      "to correct saturated samples"),
    strax.Option('saturation_min_reference_length', default=20, infer_type=False,
                 help="Minimum number of reference sample used "
                      "to correct saturated samples"),
    strax.Option('peaklet_max_duration', default=int(10e6), infer_type=False,
                 help="Maximum duration [ns] of a peaklet"),
    strax.Option('channel_map', track=False, type=immutabledict,
                 help="immutabledict mapping subdetector to (min, max) "
                      "channel number."),
    *HITFINDER_OPTIONS,
)
class Peaklets(strax.Plugin):
    """
    Split records into:
        -peaklets
        -lone_hits

    Peaklets are very aggressively split peaks such that we are able
    to find S1-S2s even if they are close to each other. (S2) Peaks
    that are split into too many peaklets will be merged later on.

    To get Peaklets from records apply/do:
        1. Hit finding
        2. Peak finding
        3. Peak splitting using the natural breaks algorithm
        4. Compute the digital sum waveform

    Lone hits are all hits which are outside of any peak. The area of
    lone_hits includes the left and right hit extension, except the
    extension overlaps with any peaks or other hits.
    """
    depends_on = ('records',)
    provides = ('peaklets', 'lone_hits')
    data_kind = dict(peaklets='peaklets',
                     lone_hits='lone_hits')
    parallel = 'process'
    compressor = 'zstd'

    __version__ = '0.5.0'

    def infer_dtype(self):
        return dict(peaklets=strax.peak_dtype(
                        n_channels=self.config['n_tpc_pmts']),
                    lone_hits=strax.hit_dtype)

    def setup(self):
        if self.config['peak_min_pmts'] > 2:
            # Can fix by re-splitting,
            raise NotImplementedError(
                f"Raising the peak_min_pmts to {self.config['peak_min_pmts']} "
                f"interferes with lone_hit definition. "
                f"See github.com/XENONnT/straxen/issues/295")

        self.to_pe = straxen.get_correction_from_cmt(self.run_id,
                                       self.config['gain_model'])

        # Check config of `hit_min_amplitude` and define hit thresholds
        # if cmt config
        if is_cmt_option(self.config['hit_min_amplitude']):
            self.hit_thresholds = straxen.get_correction_from_cmt(self.run_id,
                self.config['hit_min_amplitude'])
        # if hitfinder_thresholds config
        elif isinstance(self.config['hit_min_amplitude'], str):
            self.hit_thresholds = straxen.hit_min_amplitude(
                self.config['hit_min_amplitude'])
        else: # int or array
            self.hit_thresholds = self.config['hit_min_amplitude']

        self.channel_range = self.config['channel_map']['tpc']

    def compute(self, records, start, end):
        r = records

        hits = strax.find_hits(r, min_amplitude=self.hit_thresholds)

        # Remove hits in zero-gain channels
        # they should not affect the clustering!
        hits = hits[self.to_pe[hits['channel']] != 0]

        hits = strax.sort_by_time(hits)

        # Use peaklet gap threshold for initial clustering
        # based on gaps between hits
        peaklets = strax.find_peaks(
            hits, self.to_pe,
            gap_threshold=self.config['peaklet_gap_threshold'],
            left_extension=self.config['peak_left_extension'],
            right_extension=self.config['peak_right_extension'],
            min_channels=self.config['peak_min_pmts'],
            result_dtype=self.dtype_for('peaklets'),
            max_duration=self.config['peaklet_max_duration'],
        )

        # Make sure peaklets don't extend out of the chunk boundary
        # This should be very rare in normal data due to the ADC pretrigger
        # window.
        self.clip_peaklet_times(peaklets, start, end)

        # Get hits outside peaklets, and store them separately.
        # fully_contained is OK provided gap_threshold > extension,
        # which is asserted inside strax.find_peaks.
        is_lone_hit = strax.fully_contained_in(hits, peaklets) == -1
        lone_hits = hits[is_lone_hit]
        strax.integrate_lone_hits(
            lone_hits, records, peaklets,
            save_outside_hits=(self.config['peak_left_extension'],
                               self.config['peak_right_extension']),
            n_channels=len(self.to_pe))

        # Compute basic peak properties -- needed before natural breaks
        hits = hits[~is_lone_hit]
        # Define regions outside of peaks such that _find_hit_integration_bounds
        # is not extended beyond a peak.
        outside_peaks = self.create_outside_peaks_region(peaklets, start, end)
        strax.find_hit_integration_bounds(
            hits, outside_peaks, records,
            save_outside_hits=(self.config['peak_left_extension'],
                               self.config['peak_right_extension']),
            n_channels=len(self.to_pe),
            allow_bounds_beyond_records=True,
        )

        # Transform hits to hitlets for naming conventions. A hit refers
        # to the central part above threshold a hitlet to the entire signal
        # including the left and right extension.
        # (We are not going to use the actual hitlet data_type here.)
        hitlets = hits
        del hits

        hitlet_time_shift = (hitlets['left'] - hitlets['left_integration']) * hitlets['dt']
        hitlets['time'] = hitlets['time'] - hitlet_time_shift
        hitlets['length'] = (hitlets['right_integration'] - hitlets['left_integration'])
        hitlets = strax.sort_by_time(hitlets)
        rlinks = strax.record_links(records)

        strax.sum_waveform(peaklets, hitlets, r, rlinks, self.to_pe)

        strax.compute_widths(peaklets)

        # Split peaks using low-split natural breaks;
        # see https://github.com/XENONnT/straxen/pull/45
        # and https://github.com/AxFoundation/strax/pull/225
        peaklets = strax.split_peaks(
            peaklets, hitlets, r, rlinks, self.to_pe,
            algorithm='natural_breaks',
            threshold=self.natural_breaks_threshold,
            split_low=True,
            filter_wing_width=self.config['peak_split_filter_wing_width'],
            min_area=self.config['peak_split_min_area'],
            do_iterations=self.config['peak_split_iterations'])

        # Saturation correction using non-saturated channels
        # similar method used in pax
        # see https://github.com/XENON1T/pax/pull/712
        # Cases when records is not writeable for unclear reason
        # only see this when loading 1T test data
        # more details on https://numpy.org/doc/stable/reference/generated/numpy.ndarray.flags.html
        if not r['data'].flags.writeable:
            r = r.copy()

        if self.config['saturation_correction_on']:
            peak_list = peak_saturation_correction(
                r, rlinks, peaklets, hitlets, self.to_pe,
                reference_length=self.config['saturation_reference_length'],
                min_reference_length=self.config['saturation_min_reference_length'])

            # Compute the width again for corrected peaks
            strax.compute_widths(peaklets, select_peaks_indices=peak_list)

        # Compute tight coincidence level.
        # Making this a separate plugin would
        # (a) doing hitfinding yet again (or storing hits)
        # (b) increase strax memory usage / max_messages,
        #     possibly due to its currently primitive scheduling.
        hit_max_times = np.sort(
            hitlets['time']
            + hitlets['dt'] * hit_max_sample(records, hitlets)
            + hitlet_time_shift  # add time shift again to get correct maximum
        )
        peaklet_max_times = (
                peaklets['time']
                + np.argmax(peaklets['data'], axis=1) * peaklets['dt'])
        tight_coincidence, tight_coincidence_channel = get_tight_coin(
            hit_max_times,
            hitlets['channel'],
            peaklet_max_times,
            self.config['tight_coincidence_window_left'],
            self.config['tight_coincidence_window_right'],
            self.channel_range)

        peaklets['tight_coincidence'] = tight_coincidence
        peaklets['tight_coincidence_channel'] = tight_coincidence_channel

        if self.config['diagnose_sorting'] and len(r):
            assert np.diff(r['time']).min(initial=1) >= 0, "Records not sorted"
            assert np.diff(hitlets['time']).min(initial=1) >= 0, "Hits/Hitlets not sorted"
            assert np.all(peaklets['time'][1:]
                          >= strax.endtime(peaklets)[:-1]), "Peaks not disjoint"

        # Update nhits of peaklets:
        counts = strax.touching_windows(hitlets, peaklets)
        counts = np.diff(counts, axis=1).flatten()
        peaklets['n_hits'] = counts

        return dict(peaklets=peaklets,
                    lone_hits=lone_hits)

    def natural_breaks_threshold(self, peaks):
        rise_time = -peaks['area_decile_from_midpoint'][:, 1]

        # This is ~1 for an clean S2, ~0 for a clean S1,
        # and transitions gradually in between.
        f_s2 = 8 * np.log10(rise_time.clip(1, 1e5) / 100)
        f_s2 = 1 / (1 + np.exp(-f_s2))

        log_area = np.log10(peaks['area'].clip(1, 1e7))
        thresholds = self.config['peak_split_gof_threshold']
        return (
            f_s2 * np.interp(
                log_area,
                *np.transpose(thresholds[2]))
            + (1 - f_s2) * np.interp(
                log_area,
                *np.transpose(thresholds[1])))

    @staticmethod
    @numba.njit(nogil=True, cache=True)
    def clip_peaklet_times(peaklets, start, end):
        for p in peaklets:
            if p['time'] < start:
                p['time'] = start
            if strax.endtime(p) > end:
                p['length'] = (end - p['time']) // p['dt']

    @staticmethod
    def create_outside_peaks_region(peaklets, start, end):
        """
        Creates time intervals which are outside peaks.

        :param peaklets: Peaklets for which intervals should be computed.
        :param start: Chunk start
        :param end: Chunk end
        :return: array of strax.time_fields dtype.
        """
        if not len(peaklets):
            return np.zeros(0, dtype=strax.time_fields)

        outside_peaks = np.zeros(len(peaklets) + 1,
                                 dtype=strax.time_fields)

        outside_peaks[0]['time'] = start
        outside_peaks[0]['endtime'] = peaklets[0]['time']
        outside_peaks[1:-1]['time'] = strax.endtime(peaklets[:-1])
        outside_peaks[1:-1]['endtime'] = peaklets['time'][1:]
        outside_peaks[-1]['time'] = strax.endtime(peaklets[-1])
        outside_peaks[-1]['endtime'] = end
        return outside_peaks

@numba.jit(nopython=True, nogil=True, cache=True)
def peak_saturation_correction(records, rlinks, peaks, hitlets, to_pe,
                               reference_length=100,
                               min_reference_length=20,
                               use_classification=False,
                               ):
    """Correct the area and per pmt area of peaks from saturation
    :param records: Records
    :param rlinks: strax.record_links of corresponding records.
    :param peaks: Peaklets / Peaks
    :param hitlets: Hitlets found in records to build peaks.
        (Hitlets are hits including the left/right extension)
    :param to_pe: adc to PE conversion (length should equal number of PMTs)
    :param reference_length: Maximum number of reference sample used
    to correct saturated samples
    :param min_reference_length: Minimum number of reference sample used
    to correct saturated samples
    :param use_classification: Option of using classification to pick only S2
    """

    if not len(records):
        return
    if not len(peaks):
        return

    # Search for peaks with saturated channels
    mask = peaks['n_saturated_channels'] > 0
    if use_classification:
        mask &= peaks['type'] == 2
    peak_list = np.where(mask)[0]
    # Look up records that touch each peak
    record_ranges = _touching_windows(
        records['time'],
        strax.endtime(records),
        peaks[peak_list]['time'],
        strax.endtime(peaks[peak_list]))

    # Create temporary arrays for calculation
    dt = records[0]['dt']
    n_channels = len(peaks[0]['saturated_channel'])
    len_buffer = np.max(peaks['length'] * peaks['dt']) // dt + 1
    max_nrecord = len_buffer // len(records[0]['data']) + 1

    # Buff the sum wf [pe] of non-saturated channels
    b_sumwf = np.zeros(len_buffer, dtype=np.float32)
    # Buff the records 'data' [ADC] in saturated channels
    b_pulse = np.zeros((n_channels, len_buffer), dtype=np.int16)
    # Buff the corresponding record index of saturated channels
    b_index = np.zeros((n_channels, max_nrecord), dtype=np.int64)

    # Main
    for ix, peak_i in enumerate(peak_list):
        # reset buffers
        b_sumwf[:] = 0
        b_pulse[:] = 0
        b_index[:] = -1

        p = peaks[peak_i]
        channel_saturated = p['saturated_channel'] > 0

        for record_i in range(record_ranges[ix][0], record_ranges[ix][1]):
            r = records[record_i]
            r_slice, b_slice = strax.overlap_indices(
                r['time'] // dt, r['length'],
                p['time'] // dt, p['length'] * p['dt'] // dt)

            ch = r['channel']
            if channel_saturated[ch]:
                b_pulse[ch, slice(*b_slice)] += r['data'][slice(*r_slice)]
                b_index[ch, np.argmin(b_index[ch])] = record_i
            else:
                b_sumwf[slice(*b_slice)] += r['data'][slice(*r_slice)] \
                    * to_pe[ch]

        _peak_saturation_correction_inner(
            channel_saturated, records, p,
            to_pe, b_sumwf, b_pulse, b_index,
            reference_length, min_reference_length)

        # Back track sum wf downsampling
        peaks[peak_i]['length'] = p['length'] * p['dt'] / dt
        peaks[peak_i]['dt'] = dt

    strax.sum_waveform(peaks, hitlets, records, rlinks, to_pe, peak_list)
    return peak_list

@numba.jit(nopython=True, nogil=True, cache=True)
def _peak_saturation_correction_inner(channel_saturated, records, p,
                                      to_pe, b_sumwf, b_pulse, b_index,
                                      reference_length=100,
                                      min_reference_length=20,
                                      ):
    """Would add a third level loop in peak_saturation_correction
    Which is not ideal for numba, thus this function is written
    :param channel_saturated: (bool, n_channels)
    :param p: One peak/peaklet
    :param to_pe: adc to PE conversion (length should equal number of PMTs)
    :param b_sumwf, b_pulse, b_index: Filled buffers
    """
    dt = records['dt'][0]
    n_channels = len(channel_saturated)

    for ch in range(n_channels):
        if not channel_saturated[ch]:
            continue
        b = b_pulse[ch]
        r0 = records[b_index[ch][0]]

        # Define the reference region as reference_length before the first saturation point
        # unless there are not enough samples
        bl = np.inf
        for record_i in b_index[ch]:
            if record_i == -1:
                break
            bl = min(bl, records['baseline'][record_i])

        s0 = np.argmax(b >= np.int16(bl))
        ref = slice(max(0, s0-reference_length), s0)

        if (b[ref] * to_pe[ch] > 1).sum() < min_reference_length:
            # the pulse is saturated, but there are not enough reference samples to get a good ratio
            # This actually distinguished between S1 and S2 and will only correct S2 signals
            continue
        if (b_sumwf[ref] > 1).sum() < min_reference_length:
            # the same condition applies to the waveform model
            continue
        if np.sum(b[ref]) * to_pe[ch] / np.sum(b_sumwf[ref]) > 1:
            # The pulse is saturated, but insufficient information is available in the other channels
            # to reliably reconstruct it
            continue

        scale = np.sum(b[ref]) / np.sum(b_sumwf[ref])

        # Loop over the record indices of the saturated channel (saved in b_index buffer)
        for record_i in b_index[ch]:
            if record_i == -1:
                break
            r = records[record_i]
            r_slice, b_slice = strax.overlap_indices(
                r['time'] // dt, r['length'],
                p['time'] // dt + s0,  p['length'] * p['dt'] // dt - s0)

            if r_slice[1] == r_slice[0]:  # This record proceeds saturation
                continue
            b_slice = b_slice[0] + s0, b_slice[1] + s0

            # First is finding the highest point in the desaturated record
            # because we need to bit shift the whole record if it exceeds int16 range
            apax = scale * max(b_sumwf[slice(*b_slice)])

            if np.int32(apax) >= 2**15:  # int16(2**15) is -2**15
                bshift = int(np.floor(np.log2(apax) - 14))

                tmp = r['data'].astype(np.int32)
                tmp[slice(*r_slice)] = b_sumwf[slice(*b_slice)] * scale

                r['area'] = np.sum(tmp)  # Auto covert to int64
                r['data'][:] = np.right_shift(tmp, bshift)
                r['amplitude_bit_shift'] += bshift
            else:
                r['data'][slice(*r_slice)] = b_sumwf[slice(*b_slice)] * scale
                r['area'] = np.sum(r['data'])

@export
@strax.takes_config(
    strax.Option('n_he_pmts', track=False, default=752, infer_type=False,
                 help="Maximum channel of the he channels"),
    strax.Option('he_channel_offset', track=False, default=500, infer_type=False,
                 help="Minimum channel number of the he channels"),
    strax.Option('le_to_he_amplification', default=20, track=True, infer_type=False,
                 help="Difference in amplification between low energy and high "
                      "energy channels"),
    strax.Option('peak_min_pmts_he', default=2, infer_type=False,
                 child_option=True, parent_option_name='peak_min_pmts',
                 track=True,
                 help="Minimum number of contributing PMTs needed to define a peak"),
    strax.Option('saturation_correction_on_he', default=False, infer_type=False,
                 child_option=True, parent_option_name='saturation_correction_on',
                 track=True,
                 help='On off switch for saturation correction for High Energy'
                      ' channels'),
    *HITFINDER_OPTIONS_he
)
class PeakletsHighEnergy(Peaklets):
    __doc__ = HE_PREAMBLE + Peaklets.__doc__
    depends_on = 'records_he'
    provides = 'peaklets_he'
    data_kind = 'peaklets_he'
    __version__ = '0.0.2'
    child_plugin = True
    save_when = strax.SaveWhen.TARGET

    def infer_dtype(self):
        return strax.peak_dtype(n_channels=self.config['n_he_pmts'])

    def setup(self):
        self.to_pe = straxen.get_correction_from_cmt(self.run_id,
                                       self.config['gain_model'])

        buffer_pmts = np.zeros(self.config['he_channel_offset'])
        self.to_pe = np.concatenate((buffer_pmts, self.to_pe))
        self.to_pe *= self.config['le_to_he_amplification']

        # Check config of `hit_min_amplitude_he` and define hit thresholds
        # if cmt config
        if is_cmt_option(self.config['hit_min_amplitude_he']):
            self.hit_thresholds = straxen.get_correction_from_cmt(self.run_id,
                self.config['hit_min_amplitude_he'])
        # if hitfinder_thresholds config
        elif isinstance(self.config['hit_min_amplitude_he'], str):
            self.hit_thresholds = straxen.hit_min_amplitude(
                self.config['hit_min_amplitude_he'])
        else: # int or array
            self.hit_thresholds = self.config['hit_min_amplitude_he']

        self.channel_range = self.config['channel_map']['he']

    def compute(self, records_he, start, end):
        result = super().compute(records_he, start, end)
        return result['peaklets']

@export
class PeakletClassificationHighEnergy(PeakletClassification):
    __doc__ = HE_PREAMBLE + PeakletClassification.__doc__
    provides = 'peaklet_classification_he'
    depends_on = ('peaklets_he',)
    __version__ = '0.0.2'
    child_plugin = True

    def compute(self, peaklets_he):
        return super().compute(peaklets_he)

@export
@strax.takes_config(
    strax.Option('s2_merge_max_duration', default=50_000, infer_type=False,
                 help="Do not merge peaklets at all if the result would be a peak "
                      "longer than this [ns]"),
    strax.Option('s2_merge_gap_thresholds', default=((1.7, 2.65e4), (4.0, 2.6e3), (5.0, 0.)),
                 infer_type=False,
                 help="Points to define maximum separation between peaklets to allow "
                      "merging [ns] depending on log10 area of the merged peak\n"
                      "where the gap size of the first point is the maximum gap to allow merging"
                      "and the area of the last point is the maximum area to allow merging. "
                      "The format is ((log10(area), max_gap), (..., ...), (..., ...))"
                 ),
    strax.Option('gain_model', infer_type=False,
                 help='PMT gain model. Specify as '
                      '(str(model_config), str(version), nT-->boolean'),
    strax.Option('merge_without_s1', default=True, infer_type=False,
                 help="If true, S1s will be igored during the merging. "
                      "It's now possible for a S1 to be inside a S2 post merging"),
)
class MergedS2s(strax.OverlapWindowPlugin):
    """
    Merge together peaklets if peak finding favours that they would
    form a single peak instead.
    """
    depends_on = ('peaklets', 'peaklet_classification', 'lone_hits')
    data_kind = 'merged_s2s'
    provides = 'merged_s2s'
    __version__ = '0.4.2'

    def setup(self):
        self.to_pe = straxen.get_correction_from_cmt(self.run_id,
                                                     self.config['gain_model'])

    def infer_dtype(self):
        # wrong order
        dtype = (strax.unpack_dtype(self.deps['peaklets'].dtype_for('peaklets'))
                 +[(('Bayes peak classification type', 'type_bayes'), np.dtype('i1'))]
                 +[(('S1 ln probability', 's1_prob'), np.dtype('<f4'))]
                 +[(('S2 ln probability', 's2_prob'), np.dtype('<f4'))]
                )
        tocount = len(dtype)
        # shamless hack
        order = [0, 1, 2, 3, 4, tocount-3, tocount-2, tocount-1, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
        dtype = [dtype[i] for i in order]
        return  dtype

    def get_window_size(self):
        return 5 * (int(self.config['s2_merge_gap_thresholds'][0][1])
                    + self.config['s2_merge_max_duration'])

    def compute(self, peaklets, lone_hits):
        if self.config['merge_without_s1']:
            peaklets = peaklets[peaklets['type'] != 1]

        if len(peaklets) <= 1:
            return np.zeros(0, dtype=peaklets.dtype)

        gap_thresholds = self.config['s2_merge_gap_thresholds']
        max_gap = gap_thresholds[0][1]
        max_area = 10 ** gap_thresholds[-1][0]

        if max_gap < 0:
            # Do not merge at all
            return np.zeros(0, dtype=peaklets.dtype)
        else:
            # Max gap and area should be set by the gap thresholds
            # to avoid contradictions
            start_merge_at, end_merge_at = self.get_merge_instructions(
                peaklets['time'], strax.endtime(peaklets),
                areas=peaklets['area'],
                types=peaklets['type'],
                gap_thresholds=gap_thresholds,
                max_duration=self.config['s2_merge_max_duration'],
                max_gap=max_gap,
                max_area=max_area,
            )
            merged_s2s = strax.merge_peaks(
                peaklets,
                start_merge_at, end_merge_at,
                max_buffer=int(self.config['s2_merge_max_duration']//np.gcd.reduce(peaklets['dt'])),
            )
            merged_s2s['type'] = 2

            # Updated time and length of lone_hits and sort again:
            lh = np.copy(lone_hits)
            del lone_hits
            lh_time_shift = (lh['left'] - lh['left_integration']) *lh['dt']
            lh['time'] = lh['time'] - lh_time_shift
            lh['length'] = (lh['right_integration'] - lh['left_integration'])
            lh = strax.sort_by_time(lh)
            strax.add_lone_hits(merged_s2s, lh, self.to_pe)

            strax.compute_widths(merged_s2s)

        return merged_s2s

    @staticmethod
    @numba.njit(cache=True, nogil=True)
    def get_merge_instructions(
            peaklet_starts, peaklet_ends, areas, types,
            gap_thresholds, max_duration, max_gap, max_area):
        """
        Finding the group of peaklets to merge. To do this start with the
        smallest gaps and keep merging until the new, merged S2 has such a
        large area or gap to adjacent peaks that merging is not required
        anymore.
        see https://github.com/XENONnT/straxen/pull/548 and https://github.com/XENONnT/straxen/pull/568

        :returns: list of the first index of peaklet to be merged and
        list of the exclusive last index of peaklet to be merged
        """

        peaklet_gaps = peaklet_starts[1:] - peaklet_ends[:-1]
        peaklet_start_index = np.arange(len(peaklet_starts))
        peaklet_end_index = np.arange(len(peaklet_starts))

        for gap_i in np.argsort(peaklet_gaps):
            start_idx = peaklet_start_index[gap_i]
            inclusive_end_idx = peaklet_end_index[gap_i + 1]
            sum_area = np.sum(areas[start_idx:inclusive_end_idx + 1])
            this_gap = peaklet_gaps[gap_i]

            if inclusive_end_idx < start_idx:
                raise ValueError('Something went wrong, left is bigger then right?!')

            if this_gap > max_gap:
                break
            if sum_area > max_area:
                # For very large S2s, we assume that natural breaks is taking care
                continue
            if (sum_area > 0) and (
                    this_gap > merge_s2_threshold(np.log10(sum_area),
                                                  gap_thresholds)):
                # The merged peak would be too large
                continue

            peak_duration = (peaklet_ends[inclusive_end_idx] - peaklet_starts[start_idx])
            if peak_duration >= max_duration:
                continue

            # Merge gap in other words this means p @ gap_i and p @gap_i + 1 share the same
            # start, end and area:
            peaklet_start_index[start_idx:inclusive_end_idx + 1] = peaklet_start_index[start_idx]
            peaklet_end_index[start_idx:inclusive_end_idx + 1] = peaklet_end_index[inclusive_end_idx]

        start_merge_at = np.unique(peaklet_start_index)
        end_merge_at = np.unique(peaklet_end_index)
        if not len(start_merge_at) == len(end_merge_at):
            raise ValueError('inconsistent start and end merge instructions')

        merge_start, merge_stop_exclusive = _filter_s1_starts(
            start_merge_at, types, end_merge_at)

        return merge_start, merge_stop_exclusive

@numba.njit(cache=True, nogil=True)
def _filter_s1_starts(start_merge_at, types, end_merge_at):
    for start_merge_idx, _ in enumerate(start_merge_at):
        while types[start_merge_at[start_merge_idx]] != 2:
            if end_merge_at[start_merge_idx] - start_merge_at[start_merge_idx] <= 1:
                break
            start_merge_at[start_merge_idx] += 1

    start_merge_with_s2 = types[start_merge_at] == 2
    merges_at_least_two_peaks = end_merge_at - start_merge_at >= 1

    keep_merges = start_merge_with_s2 & merges_at_least_two_peaks
    return start_merge_at[keep_merges], end_merge_at[keep_merges] + 1

@numba.njit(cache=True, nogil=True)
def merge_s2_threshold(log_area, gap_thresholds):
    """Return gap threshold for log_area of the merged S2
    with linear interpolation given the points in gap_thresholds
    :param log_area: Log 10 area of the merged S2
    :param gap_thresholds: tuple (n, 2) of fix points for interpolation
    """
    for i, (a1, g1) in enumerate(gap_thresholds):
        if log_area < a1:
            if i == 0:
                return g1
            a0, g0 = gap_thresholds[i - 1]
            return (log_area - a0) * (g1 - g0) / (a1 - a0) + g0
    return gap_thresholds[-1][1]

@export
class MergedS2sHighEnergy(MergedS2s):
    __doc__ = HE_PREAMBLE + MergedS2s.__doc__
    depends_on = ('peaklets_he', 'peaklet_classification_he')
    data_kind = 'merged_s2s_he'
    provides = 'merged_s2s_he'
    __version__ = '0.0.1'
    child_plugin = True

    def infer_dtype(self):
        return strax.unpack_dtype(self.deps['peaklets_he'].dtype_for('peaklets_he'))

    def compute(self, peaklets_he):
        # There are not any lone hits for the high energy channel,
        #  so create a dummy for the compute method.
        lone_hits = np.zeros(0, dtype=strax.hit_dtype)
        return super().compute(peaklets_he, lone_hits)

@export
@strax.takes_config(
    strax.Option('diagnose_sorting', track=False, default=False, infer_type=False,
                 help="Enable runtime checks for sorting and disjointness"),
    strax.Option('merge_without_s1', default=True, infer_type=False,
                 help="If true, S1s will be igored during the merging. "
                      "It's now possible for a S1 to be inside a S2 post merging"),
)
class Peaks(strax.Plugin):
    """
    Merge peaklets and merged S2s such that we obtain our peaks
    (replacing all peaklets that were later re-merged as S2s). As this
    step is computationally trivial, never save this plugin.
    """
    depends_on = ('peaklets', 'peaklet_classification', 'merged_s2s')
    data_kind = 'peaks'
    provides = 'peaks'
    parallel = True
    save_when = strax.SaveWhen.EXPLICIT

    __version__ = '0.1.2'

    def infer_dtype(self):
        # wrong order
        dtype = (strax.unpack_dtype(self.deps['peaklets'].dtype_for('peaklets'))
                 +[(('Bayes peak classification type', 'type_bayes'), np.dtype('i1'))]
                 +[(('S1 ln probability', 's1_prob'), np.dtype('<f4'))]
                 +[(('S2 ln probability', 's2_prob'), np.dtype('<f4'))]
                )
        tocount = len(dtype)
        # shamless hack
        order = [0, 1, 2, 3, 4, tocount-3, tocount-2, tocount-1, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
        dtype = [dtype[i] for i in order]
        return  dtype

    def compute(self, peaklets, merged_s2s):
        # Remove fake merged S2s from dirty hack, see above
        merged_s2s = merged_s2s[merged_s2s['type'] != FAKE_MERGED_S2_TYPE]

        if self.config['merge_without_s1']:
            is_s1 = peaklets['type'] == 1
            peaks = strax.replace_merged(peaklets[~is_s1], merged_s2s)
            peaks = strax.sort_by_time(np.concatenate([peaklets[is_s1],
                                                       peaks]))
        else:
            peaks = strax.replace_merged(peaklets, merged_s2s)

        if self.config['diagnose_sorting']:
            assert np.all(np.diff(peaks['time']) >= 0), "Peaks not sorted"
            if self.config['merge_without_s1']:
                to_check = peaks['type'] != 1
            else:
                to_check = peaks['type'] != FAKE_MERGED_S2_TYPE

            assert np.all(peaks['time'][to_check][1:]
                            >= strax.endtime(peaks)[to_check][:-1]), "Peaks not disjoint"
        return peaks

@export
class PeaksHighEnergy(Peaks):
    __doc__ = HE_PREAMBLE + Peaks.__doc__
    depends_on = ('peaklets_he', 'peaklet_classification_he', 'merged_s2s_he')
    data_kind = 'peaks_he'
    provides = 'peaks_he'
    __version__ = '0.0.1'
    child_ends_with = '_he'

    def infer_dtype(self):
        return self.deps['peaklets_he'].dtype_for('peaklets')

    def compute(self, peaklets_he, merged_s2s_he):
        return super().compute(peaklets_he, merged_s2s_he)

@numba.jit(nopython=True, nogil=True, cache=True)
def get_tight_coin(hit_max_times, hit_channel, peak_max_times, left, right,
                   channels=(0, 493)):
    """Calculates the tight coincidence based on hits and PMT channels.

    Defined by number of hits within a specified time range of the
    the peak's maximum amplitude.
    Imitates tight_coincidence variable in pax:
    github.com/XENON1T/pax/blob/master/pax/plugins/peak_processing/BasicProperties.py

    :param hit_max_times: Time of the hit amplitude in ns.
    :param hit_channel: PMT channels of the hits
    :param peak_max_times: Time of the peaks maximum in ns.
    :param left: Left boundary in which we search for the tight
        coincidence in ns.
    :param right: Right boundary in which we search for the tight
        coincidence in ns.
    :param channel_range: (min/max) channel for the corresponding detector.

    :returns: n_coin_hit, n_coin_channel of length peaks containing the
        tight coincidence.
    """
    left_hit_i = 0
    n_coin_hit = np.zeros(len(peak_max_times), dtype=np.int16)
    n_coin_channel = np.zeros(len(peak_max_times), dtype=np.int16)
    start_ch, end_ch = channels
    channels_seen = np.zeros(end_ch-start_ch+1, dtype=np.bool_)

    # loop over peaks
    for p_i, p_t in enumerate(peak_max_times):
        channels_seen[:] = 0
        # loop over hits starting from the last one we left at
        for left_hit_i in range(left_hit_i, len(hit_max_times)):

            # if the hit is in the window, its a tight coin
            d = hit_max_times[left_hit_i] - p_t
            if (-left <= d) & (d <= right):
                n_coin_hit[p_i] += 1
                channels_seen[hit_channel[left_hit_i]-start_ch] = 1

            # stop the loop when we know we're outside the range
            if d > right:
                n_coin_channel[p_i] = np.sum(channels_seen)
                break

        # Add channel information in case there are no hits beyond
        # the last peak:
        n_coin_channel[p_i] = np.sum(channels_seen)

    return n_coin_hit, n_coin_channel

@numba.njit(cache=True, nogil=True)
def hit_max_sample(records, hits):
    """Return the index of the maximum sample for hits"""
    result = np.zeros(len(hits), dtype=np.int16)
    for i, h in enumerate(hits):
        r = records[h['record_i']]
        w = r['data'][h['left']:h['right']]
        result[i] = np.argmax(w)
    return result


"""Compute positions from an ephemeris installed as a Python package."""

import os
import numpy as np

class DateError(ValueError):
    """Date input is outside the range covered by the ephemeris."""

class Ephemeris(object):
    """A JPL planetary ephemeris that, given dates, computes positions."""

    def __init__(self, module):
        self.name = module.__name__.upper()
        self.dirpath = os.path.dirname(module.__file__)
        self.names = tuple(sorted(
            name.split('-')[-1].split('.')[0]
            for name in os.listdir(self.dirpath)
            if not name.startswith('constants') and name.endswith('.npy')
            ))
        path = self.path('constants.npy')
        self.__dict__.update((k.decode('ascii'), v) for k, v in np.load(path))
        self.earth_share = 1.0 / (1.0 + self.EMRAT)
        self.moon_share = self.EMRAT / (1.0 + self.EMRAT)
        self.sets = {}

    def path(self, filename):
        """Compute the path to a particular file in the ephemeris."""
        return os.path.join(self.dirpath, filename)

    def load(self, name):
        """Load the polynomial series for `name` and return it."""
        s = self.sets.get(name)
        if s is None:
            self.sets[name] = s = np.load(self.path('jpl-%s.npy' % name))
        return s

    def position(self, name, tdb, tdb2=0.0):
        """Compute the position of `name` at time ``tdb [+ tdb2]``.

        The position is returned as a NumPy array ``[x y z]``.

        The barycentric dynamical time `tdb` argument should be a float.
        If there are many dates you want computed, then make `tdb` an
        array, which is more efficient than calling this method multiple
        times; the return value will be a two-dimensional array giving a
        row of values for each coordinate.

        For extra precision, the time can be split into two floats; a
        popular choice is to use `tdb` for the integer or half-integer
        date, and `tdb2` to hold the remaining fraction.

        Consult the `names` attribute of this ephemeris for the values
        of `name` it supports, such as ``'mars'`` or ``'earthmoon'``.

        """
        bundle = self.compute_bundle(name, tdb, tdb2)
        return self.position_from_bundle(bundle)

    def position_and_velocity(self, name, tdb, tdb2=0.0):
        """Compute the position and velocity of `name` at ``tdb [+ tdb2]``.

        The position and velocity are returned in a 2-tuple::

            ([x y z], [xdot ydot zdot])

        The barycentric dynamical time `tdb` argument should be a float.
        If there are many dates you want computed, then make `tdb` an
        array, which is more efficient than calling this method multiple
        times; the return values will be two-dimensional arrays giving a
        row of values for each coordinate.

        For extra precision, the time can be split into two floats; a
        popular choice is to use `tdb` for the integer or half-integer
        date, and `tdb2` to hold the remaining fraction.

        Consult the `names` attribute of this ephemeris for the values
        of `name` it supports, such as ``'mars'`` or ``'earthmoon'``.

        """
        bundle = self.compute_bundle(name, tdb, tdb2)
        position = self.position_from_bundle(bundle)
        velocity = self.velocity_from_bundle(bundle)
        return position, velocity

    def compute(self, name, tdb):
        """Legacy routine that concatenates position and velocity vectors.

        This routine is deprecated.  Use the methods `position()` and
        `position_and_velocity()` instead.  This method follows the same
        calling convention, but incurs extra copy operations in order to
        return a single NumPy array::

            [x y z xdot ydot zdot]

        """
        bundle = self.compute_bundle(name, tdb, 0.0)
        position = self.position_from_bundle(bundle)
        velocity = self.velocity_from_bundle(bundle)
        return np.concatenate((position, velocity))

    def compute_bundle(self, name, tdb, tdb2=0.0):
        """Return a tuple of coefficients and parameters for `tdb`.

        The return value is a tuple that bundles together the
        coefficients and other Chebyshev intermediate values that are
        needed for the computation of either the position or velocity.
        The bundle can then be passed to either `position_from_bundle()`
        or `velocity_from_bundle()` to finish the computation.  See the
        package-level documentation for details; most users will simply
        call `position()` or `position_and_velocity()` instead.

        The barycentric dynamical time `tdb` argument should be a float.
        If there are many dates you want computed, then make `tdb` an
        array, which is more efficient than calling this method multiple
        times; the return values will be arrays providing a value for
        each time in `tdb`.

        For extra precision, the time can be split into two floats; a
        popular choice is to use `tdb` for the integer or half-integer
        date, and `tdb2` to hold the remaining fraction.

        Consult the `names` attribute of this ephemeris for the values
        of `name` it supports, such as ``'mars'`` or ``'earthmoon'``.

        """
        input_was_scalar = getattr(tdb, 'shape', ()) == ()
        if input_was_scalar:
            tdb = np.array((tdb,))
        # no need to deal with tdb2; numpy broadcast will add fine below.

        coefficient_sets = self.load(name)
        number_of_sets, axis_count, coefficient_count = coefficient_sets.shape

        jalpha, jomega = self.jalpha, self.jomega
        days_per_set = (jomega - jalpha) / number_of_sets
        # to keep precision, first subtract, then add
        index, offset = divmod((tdb - jalpha) + tdb2, days_per_set)
        index = index.astype(int)

        if (index < 0).any() or (number_of_sets < index).any():
            raise DateError('ephemeris %s only covers dates %.1f through %.1f'
                            % (self.name, jalpha, jomega))

        omegas = (index == number_of_sets)
        index[omegas] -= 1
        offset[omegas] += days_per_set

        coefficients = np.rollaxis(coefficient_sets[index], 1)

        # Chebyshev recurrence:

        T = np.empty((coefficient_count, len(index)))
        T[0] = 1.0
        T[1] = t1 = 2.0 * offset / days_per_set - 1.0
        twot1 = t1 + t1
        for i in range(2, coefficient_count):
            T[i] = twot1 * T[i-1] - T[i-2]

        bundle = coefficients, days_per_set, T, twot1
        return bundle

    def position_from_bundle(self, bundle):
        """Return position, given the `coefficient_bundle()` return value."""

        coefficients, days_per_set, T, twot1 = bundle
        return (T.T * coefficients).sum(axis=2)

    def velocity_from_bundle(self, bundle):
        """Return velocity, given the `coefficient_bundle()` return value."""

        coefficients, days_per_set, T, twot1 = bundle
        coefficient_count = coefficients.shape[2]

        # Chebyshev derivative:

        dT = np.empty_like(T)
        dT[0] = 0.0
        dT[1] = 1.0
        dT[2] = twot1 + twot1
        for i in range(3, coefficient_count):
            dT[i] = twot1 * dT[i-1] - dT[i-2] + T[i-1] + T[i-1]
        dT *= 2.0
        dT /= days_per_set

        return (dT.T * coefficients).sum(axis=2)


#!/usr/bin/env python
# -*- coding: utf-8 -*
''' Functions to compute MESE weights
'''
from __future__ import print_function, division

import numpy as np
from scipy.spatial import ConvexHull

from icecube import dataclasses
from icecube import icetray
from icecube import NewNuFlux
from icecube import AtmosphericSelfVeto
from icecube.icetray.i3logging import log_info, log_warn

from ic3_labels.labels.utils import muon as mu_utils
from ic3_labels.labels.utils import tau as tau_utils
from ic3_labels.labels.utils.cascade import get_cascade_of_primary_nu

def atmosphericFlux(
        neutrinoEnergy,
        neutrinoZenith,
        neutrinoType,
        atmFluxConv,
        atmFluxPrompt):
    """
    Excessively complicated flux_atm calculation.

    This bundle is adapted from Nancy's adaptation-for-mrichman of her dataset
    wrangling scripts.  Each part is more flexible than it really needs to be
    for this application, but I considered it safer to keep it all rather than
    go through the potentially error-prone process of streamlining the code.
    """
    if isinstance(neutrinoEnergy, float):
        neutrinoEnergy = np.atleast_1d(neutrinoEnergy)
        neutrinoZenith = np.atleast_1d(neutrinoZenith)
        neutrinoType = np.atleast_1d(neutrinoType)

    p_types = [dataclasses.I3Particle.ParticleType(t) for t in neutrinoType]
    atmflux = np.zeros(len(neutrinoEnergy))

    badmask = neutrinoEnergy < 10.
    if atmFluxConv is not None:
        conv = atmFluxConv.getFlux(
            p_types,
            neutrinoEnergy, np.cos(neutrinoZenith))
        conv[badmask] = np.nan
        atmflux += conv
    if atmFluxPrompt is not None:
        prompt = atmFluxPrompt.getFlux(
            p_types,
            neutrinoEnergy, np.cos(neutrinoZenith))
        prompt[badmask] = np.nan
        atmflux += prompt
    # return atmflux

    if ((isinstance(neutrinoZenith, float))
            and (isinstance(neutrinoType, str))):
        for i in range(len(neutrinoEnergy)):
            if neutrinoEnergy[i] < 10.:
                atmflux[i] = np.nan
                continue

            conv = 0.
            prompt = 0.
            if atmFluxConv is not None:
                conv = atmFluxConv.getFlux(
                    dataclasses.I3Particle.ParticleType(neutrinoType),
                    neutrinoEnergy[i],
                    np.cos(neutrinoZenith))
            if atmFluxPrompt is not None:
                prompt = atmFluxPrompt.getFlux(
                    dataclasses.I3Particle.ParticleType(neutrinoType),
                    neutrinoEnergy[i],
                    np.cos(neutrinoZenith))
            atmflux[i] = conv+prompt
    else:
        for i in range(len(neutrinoEnergy)):
            if neutrinoEnergy[i] < 10.:
                atmflux[i] = np.nan
                continue
            conv = 0.
            prompt = 0.
            if atmFluxConv is not None:
                conv = atmFluxConv.getFlux(
                    dataclasses.I3Particle.ParticleType(neutrinoType[i]),
                    neutrinoEnergy[i],
                    np.cos(neutrinoZenith[i]))
            if atmFluxPrompt is not None:
                prompt = atmFluxPrompt.getFlux(
                    dataclasses.I3Particle.ParticleType(neutrinoType[i]),
                    neutrinoEnergy[i],
                    np.cos(neutrinoZenith[i]))
            atmflux[i] = conv+prompt

    if len(atmflux) == 1:
        atmflux = atmflux[0]

    return atmflux

class MESEWeights(icetray.I3ConditionalModule):

    """Calculate weights for MESE 7yr cascade ps paper.
    The returned weights are rates in Hz. To obtain number of events, this
    still has to be multiplied by livetime.
    """

    def __init__(self, context):
        icetray.I3ConditionalModule.__init__(self, context)
        self.AddParameter("DatasetType",
                          "Type of dataset. Must be one of: "
                          "'muongun', 'nugen', 'genie'")
        self.AddParameter("DatasetNFiles", "Number of files")
        self.AddParameter("DatasetNEventsPerRun",
                          "Number of generated events per file")
        self.AddParameter("OutputKey", "Save weights to this frame key.",
                          'MESE_weights')

    def Configure(self):
        self._dataset_type = self.GetParameter("DatasetType")
        self._n_files = self.GetParameter("DatasetNFiles")
        self._n_events_per_run = self.GetParameter("DatasetNEventsPerRun")
        self._ngen = self._n_events_per_run * self._n_files
        self._output_key = self.GetParameter("OutputKey")

        self._dataset_type = self._dataset_type.lower()

        if self._dataset_type not in ['muongun', 'nugen', 'genie']:
            raise ValueError('Unkown dataset_type: {!r}'.format(dataset_type))

        # get Honda2006
        self.honda = NewNuFlux.makeFlux("honda2006")
        self.honda.knee_reweighting_model = 'gaisserH3a_elbert'
        self.honda.relative_kaon_contribution = .91
        # get self-veto
        self.af = AtmosphericSelfVeto.AnalyticPassingFraction
        self.honda_veto_hese = self.af('conventional', veto_threshold=1.25e3)
        self.honda_veto_mese = self.af('conventional', veto_threshold=1e2)
        # get the sarcevic model for prompt neutrinos
        self.enberg = NewNuFlux.makeFlux("sarcevic_std")
        self.enberg_veto_hese = self.af('charm', veto_threshold=1.25e3)
        self.enberg_veto_mese = self.af('charm', veto_threshold=1e2)
        self.conv_flux_multiplier = 1.07
        self.prompt_flux_multiplier = .2

    def Geometry(self, frame):
        geoMap = frame['I3Geometry'].omgeo
        domPosDict = {(i[0][0], i[0][1]): (i[1].position.x,
                                           i[1].position.y,
                                           i[1].position.z)
                      for i in geoMap if i[1].omtype.name == 'IceCube'}
        points = [
            domPosDict[(31, 1)], domPosDict[(1, 1)],
            domPosDict[(6, 1)], domPosDict[(50, 1)],
            domPosDict[(74, 1)], domPosDict[(72, 1)],
            domPosDict[(78, 1)], domPosDict[(75, 1)],

            domPosDict[(31, 60)], domPosDict[(1, 60)],
            domPosDict[(6, 60)], domPosDict[(50, 60)],
            domPosDict[(74, 60)], domPosDict[(72, 60)],
            domPosDict[(78, 60)], domPosDict[(75, 60)]
            ]
        self._convex_hull = ConvexHull(points)
        self._dom_pos_dict = domPosDict
        self.PushFrame(frame)

    def Physics(self, frame):

        mese_dict = {
            'n_files': self._n_files,
            'n_events_per_run': self._n_events_per_run,
        }

        # get MC info
        energy_true = frame['MCPrimary'].energy
        zenith_true = frame['MCPrimary'].dir.zenith
        azimuth_true = frame['MCPrimary'].dir.azimuth

        # -------
        # NuGen
        # -------
        if self._dataset_type in ['nugen', 'genie']:
            # get oneweight / n_gen
            oneweight = frame['I3MCWeightDict']['OneWeight'] / self._ngen
            true_type = frame['I3MCWeightDict']['PrimaryNeutrinoType']
            is_tau = (np.abs(true_type) == 16).all()

            # calculate astrophysical weights
            mese_dict['weight_E269'] = 2.09e-18 * oneweight * (
                                                    energy_true / 1e5)**-2.69
            mese_dict['weight_E250'] = 2.23e-18 * oneweight * (
                                                    energy_true / 1e5)**-2.5

            # calculate atmospheric weights
            if is_tau:
                mese_dict['weight_conv'] = oneweight * atmosphericFlux(
                        neutrinoEnergy=energy_true,
                        neutrinoZenith=zenith_true,
                        neutrinoType=true_type,
                        atmFluxConv=None,
                        atmFluxPrompt=None,) * 2. * self.conv_flux_multiplier

                mese_dict['weight_prompt'] = oneweight * atmosphericFlux(
                        neutrinoEnergy=energy_true,
                        neutrinoZenith=zenith_true,
                        neutrinoType=true_type,
                        atmFluxConv=None,
                        atmFluxPrompt=None,) * 2. * self.prompt_flux_multiplier
            else:
                mese_dict['weight_conv'] = oneweight * atmosphericFlux(
                        neutrinoEnergy=energy_true,
                        neutrinoZenith=zenith_true,
                        neutrinoType=true_type,
                        atmFluxConv=self.honda,
                        atmFluxPrompt=None,) * 2. * self.conv_flux_multiplier
                mese_dict['weight_prompt'] = oneweight * atmosphericFlux(
                    neutrinoEnergy=energy_true,
                    neutrinoZenith=zenith_true,
                    neutrinoType=true_type,
                    atmFluxConv=None,
                    atmFluxPrompt=self.enberg,)*2.*self.prompt_flux_multiplier

            # ---------------------
            # Atmospheric Self Veto
            # ---------------------
            # get true_depth
            if 'IntersectionPoint' in frame:
                true_depth = frame['IntersectionPoint'].z
            else:
                muon = mu_utils.get_muon_of_inice_neutrino(frame)
                tau = tau_utils.get_tau_of_inice_neutrino(frame)

                if muon is not None:
                    # found a muon
                    entry = self._get_muon_entry(frame, muon)
                    true_depth = entry.z

                elif tau is not None:
                    # found a tau
                    entry = self._get_particle_entry(tau)
                    true_depth = entry.z
                else:

                    # no muon or tau exists: cascade
                    cascade = get_cascade_of_primary_nu(frame,
                                                        frame['MCPrimary'],
                                                        convex_hull=None,
                                                        extend_boundary=800)[0]

                    if cascade is not None:
                        true_depth = cascade.pos.z
                    else:
                        cascade = get_cascade_of_primary_nu(
                            frame,
                            frame['MCPrimary'],
                            convex_hull=None,
                            extend_boundary=float('inf'))[0]

                        # Muon coming out of hadronic shower?
                        daughters = frame['I3MCTree'].get_daughters(cascade)

                        # collect possible muons from daughters of daughters
                        # e.g. Nu -> Nu + Hadrons -> Mu
                        muons = []
                        for d in daughters:
                            muons.extend([
                                m for m in frame['I3MCTree'].get_daughters(d)
                                if mu_utils.is_muon(m)])
                        if muons:
                            # pick highest energy muon
                            indices = np.argsort([m.energy for m in muons])
                            muon = muons[indices[-1]]
                            entry = self._get_muon_entry(frame, muon)
                            true_depth = entry.z
                        else:
                            true_depth = cascade.pos.z

            # apply self veto
            veto_args = (true_type, energy_true,
                         np.cos(zenith_true),
                         1950. - true_depth
                         )

            if 'IsHese' in frame:
                if frame['IsHese'].value:
                    mese_dict['veto_conv'] = self.honda_veto_hese(*veto_args)
                    mese_dict['veto_prompt'] = self.enberg_veto_hese(
                                                                    *veto_args)
                else:
                    mese_dict['veto_conv'] = self.honda_veto_mese(*veto_args)
                    mese_dict['veto_prompt'] = self.enberg_veto_mese(
                                                                    *veto_args)

            else:
                log_warn('WARNING: IsHese does not exist. Using MESE veto')
                mese_dict['veto_conv'] = self.honda_veto_mese(*veto_args)
                mese_dict['veto_prompt'] = self.honda_veto_mese(*veto_args)

            mese_dict['weight_conv'] *= mese_dict['veto_conv']
            mese_dict['weight_prompt'] *= mese_dict['veto_prompt']
            # ---------------------

        # -------
        # MuonGun
        # -------
        elif self._dataset_type == 'muongun':
            if 'MuonWeight_GaisserH4a' in frame:
                # --- Where does magic number of 1.6 come from? MuonMultiplier
                mese_dict['muon_weight'] = \
                    frame['MuonWeight_GaisserH4a'].value * 1.6 / self._ngen

        # -----------------
        # Experimental Data
        # -----------------
        elif self._dataset_type == 'data':
            mjd = frame['I3EventHeader'].start_time.mod_julian_day_double

        # -----------------------------------------------------
        # final track cut:
        # drop low energy downgoing tracks and duplicate events
        # -----------------------------------------------------
        try:
            # get TrackFit_zenith
            TrackFit_zenith = frame['TrackFit'].dir.zenith

            # get energy_millipede
            energy_millipede = frame['MillipedeDepositedEnergy'].value

            # mask events
            track_mask = data_dict['is_cascade_reco'] | \
                ~((np.cos(TrackFit_zenith) > 0.3) & (energy_millipede < 10e3))

            if self._dataset_type in ['muongun', 'nugen', 'genie']:
                uniq_mask = np.r_[True, np.diff(energy_true) != 0]
            else:
                uniq_mask = np.r_[True, np.diff(mjd) != 0]
            mese_dict['passed_final_track_cut'] = track_mask & uniq_mask
        except Exception as e:
            # log_warn(e)
            pass
        # -----------------------------------------------------
        for k, item in mese_dict.items():
            mese_dict[k] = float(item)
        frame[self._output_key] = dataclasses.I3MapStringDouble(mese_dict)

        self.PushFrame(frame)

    def _get_particle_entry(self, particle):

        entry = mu_utils.get_muon_initial_point_inside(
                                        particle, self._convex_hull)
        if entry is None:
            # get closest approach point as entry approximation
            entry = mu_utils.get_particle_closest_approach_to_position(
                                    particle, dataclasses.I3Position(0, 0, 0))
        return entry

    def _get_muon_entry(self, frame, muon):

        entry = mu_utils.get_muon_initial_point_inside(muon, self._convex_hull)
        if entry is None:
            # get closest approach point as entry approximation
            entry = mu_utils.get_muon_closest_approach_to_center(frame, muon)

        return entry


'''
说明：
没有文件写操作，完全将数据集加载到内存中进行操作
对于较大的数据集会很耗内存，但训练速度快！！！
'''

import os
import time
import numpy as np
import pandas as pd
import pickle
import yaml
import jieba.analyse
import warnings
# import logging
import multiprocessing
import matplotlib.pyplot as plt
import tensorflow as tf
# from collections import Counter
import keras
from keras import losses
from keras.preprocessing import sequence
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle
from keras.layers import Dense, Dropout, SpatialDropout1D, Flatten
from keras.layers.embeddings import Embedding
from keras.layers.wrappers import Bidirectional
from keras.layers.recurrent import LSTM
from keras.layers.convolutional import Conv1D, Convolution1D
from keras.layers.pooling import GlobalMaxPool1D, MaxPooling1D
from keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard
from keras.models import Sequential, model_from_yaml
from keras import backend as K
warnings.filterwarnings(action='ignore', category=UserWarning, module='gensim')
from gensim.models.word2vec import Word2Vec
from gensim.corpora.dictionary import Dictionary
from log.logger import MyLogger
from models.text_utils import TextUtils

#设置日志级别，默认是logging.WARNING，低于该级别的不会输出，级别排序:CRITICAL>ERROR>WARNING>INFO>DEBUG
# logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)

# logging.basicConfig(filename='log/new.log', #日志文件
#                     filemode='a', #模式，有w和a，w就是写模式，每次都会重新写日志，覆盖之前的日志，a是追加模式，默认是追加模式
#                     format='%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s', #日志格式
#                     level=logging.DEBUG,#控制台打印的日志级别
#                 ) #只会保存log到文件，不会输出到控制台

#随机数和种子之间的关系遵从以下两个规则：种子不同，产生不同的随机数；种子相同，即使实例不同也产生相同的随机数
np.random.seed(3347)
LABEL_POS, LABEL_NEG = 1, 0 #正样例标签1、负样例标签0
N_ROWS = 260000 #每个类别的评论数
N_ITER = 8  #词向量训练的迭代次数
WIN_SIZE = 8  #词向量上下文最大距离，一般取[5-10]
MIN_COUNT = 5 #需要计算词向量的最小词频。这个值可以去掉一些很生僻的低频词
EMBEDDING_SIZE = 300 #词向量的维度，数据集越大该值越大
HIDDEN_LAYER_SIZE = 64 #隐层的大小
BATCH_SIZE = 64 #批的大小
NUM_EPOCHS = 10 #LSTM样本训练的轮数
TEST_SIZE = 0.2 #总样本中，测试样本的占比
CPU_COUNT = multiprocessing.cpu_count() #cpu线程数量

COMMENT_DIR = '../comment' #源评论文件目录
USER_DIR = '../user' #用户自定义文件目录
MODEL_DIR = '../log' #保存模型目录
USER_DICT_PATH = os.path.join(USER_DIR, 'user_dict.txt') #自定义用户表
STOP_WORDS_PATH = os.path.join(USER_DIR, 'stop_words.txt') #停用词表
LSTM_MODEL_PATH = os.path.join(MODEL_DIR, 'lstm.yml') #保存LSTM模型结构的路径
LSTM_WEIGHT_PATH = os.path.join(MODEL_DIR, 'lstm.h5') #保存LSTM权重的路径
# WORD2VEC_PATH = os.path.join(MODEL_DIR, 'word2vec_model.pkl') #保存word2vec词向量模型的路径
WORD2VEC_PATH = os.path.join('../model_lstm', 'word2vec.model') #保存word2vec词向量模型的路径
HIST_PATH = os.path.join(MODEL_DIR, 'hist.pkl') #保存训练过程中的历史记录
jieba.load_userdict(USER_DICT_PATH)
jieba.analyse.set_stop_words(STOP_WORDS_PATH)
pos_tags = ['n', 'vn', 'v', 'ad', 'a', 'e', 'y'] #是名词、形容词、动词、副词、叹词、语气词
BEST_MODEL_PATH = os.path.join(MODEL_DIR, 'model-{epoch:02d}-{val_acc:.2f}.h5') #保存训练最好的模型的权值

class SentimentModel:
    def __init__(self):
        self.max_len = 100  # 文本保留的最大长度

        self.logging = MyLogger()  # 自定义日志器
        self.margin = 0.6  # 阈值
        self.theta = lambda t: (K.sign(t)+1)/2
        # self.theta = lambda t: K.sigmoid(100 * t)  # 软化
        self.logging.info('###该模型为情感二分类模型###')
        self.logging.info('CPU线程数：{}'.format(CPU_COUNT))

    # 二分类：自定义损失函数（hinge loss+triplet loss）
    def loss_new(self, y_true, y_pred):
        return -(1 - self.theta(y_true - self.margin) * self.theta(y_pred - self.margin) - self.theta(
            1 - self.margin - y_true) * self.theta(1 - self.margin - y_pred)) * (
                           y_true * K.log(y_pred + K.epsilon()) + (1 - y_true) * K.log(1 - y_pred + K.epsilon()))

    # 二分类：focal_loss损失函数，解决样本不均衡分布问题
    def focal_loss(self, y_true, y_pred, alpha=0.25, gamma=2.0):
        pt_1 = tf.where(tf.equal(y_true, LABEL_POS), y_pred, tf.ones_like(y_pred))
        pt_0 = tf.where(tf.equal(y_true, LABEL_NEG), y_pred, tf.zeros_like(y_pred))
        return -K.sum(alpha * K.pow(1. - pt_1, gamma) * K.log(pt_1)) \
               - K.sum((1. - alpha) * K.pow(pt_0, gamma) * K.log(1. - pt_0))

    '''
     input_csv_dir：csv格式的评论数据集目录
     返回值：DataFrame对象
     '''
    def load_train_data(self, input_csv_dir):
        self.logging.info('开始加载数据......')
        t1 = time.time()
        df_data = pd.DataFrame()
        if os.path.isdir(input_csv_dir):
            for f in os.listdir(input_csv_dir):
                csv_file = os.path.join(input_csv_dir, f)
                csv_data = pd.read_csv(csv_file, usecols=['label', 'segs'], nrows=N_ROWS)
                df_data = df_data.append(csv_data, ignore_index=True)
        else:
            df_data = pd.read_csv(input_csv_dir, usecols=['label', 'segs'], nrows=N_ROWS)

        t2 = time.time()
        self.logging.info('数据加载完成！总用时：{}s'.format(t2 - t1))
        # self.max_len = max(df['segs'].apply(lambda x: len(x)))  # 序列的最大长度

        self.logging.info('数据集总记录数：{}'.format(len(df_data)))
        nb_pos = len(df_data[df_data['label'] == LABEL_POS])  # 正样例数
        nb_neg = len(df_data[df_data['label'] == LABEL_NEG])  # 负样例数
        self.logging.info('正面例数：{} 负面样例数：{}'.format(nb_pos, nb_neg))
        return shuffle(df_data)  # 随机打乱数据集

    #训练词向量模型
    def train_wd2vect(self, sentences): #分词列表
        if os.path.exists(WORD2VEC_PATH):
            word2vec_model = Word2Vec.load(WORD2VEC_PATH)
        else:
            self.logging.info('开始训练词向量模型......')
            t1 = time.time()
            word2vec_model = Word2Vec(sentences,
                                      size=EMBEDDING_SIZE, #词向量的维度
                                      window=WIN_SIZE,  #在一个句子中，当前词和预测词的最大距离(词向量上下文最大距离)
                                      min_count=MIN_COUNT, #词频少于min_count次数的单词会被丢弃掉
                                      sg=0, #训练算法：sg=0 使用cbow训练, sg=1 使用skip-gram 对低频词较为敏感
                                      workers=CPU_COUNT, #设置多线程训练模型，机器的核数越多，训练越快
                                      iter=N_ITER  #随机梯度下降法中迭代的最大次数，默认是5。对于大语料，可以增大这个值
                                    )
            t2 = time.time()
            self.logging.info('词向量训练结束！总用时：{}min'.format((t2 - t1) / 60.0))

            if not os.path.exists(MODEL_DIR):
                os.mkdir(MODEL_DIR)
            word2vec_model.save(WORD2VEC_PATH) #保存词向量模型
            self.logging.info('词向量模型已保存......')
            # word2vec_model.save_word2vec_format(out.model, binary=False)
        return word2vec_model

    # 根据词向量模型得到词索引{词: 索引} 和 词向量{词: 词向量}
    def create_dicts(self, wd2vec_model):
        if wd2vec_model is not None:
            gensim_dict = Dictionary()  # {索引: 词}
            # 实现词袋模型
            gensim_dict.doc2bow(wd2vec_model.wv.vocab.keys(), allow_update=True)  # (token_id, token_count)
            word2index = {wd: idx + 1 for idx, wd in gensim_dict.items()}  # 词索引字典 {词: 索引}，索引从1开始计数
            word_vectors = {wd: wd2vec_model.wv[wd] for wd in word2index.keys()}  # 词向量 {词: 词向量}
            return word2index, word_vectors
        else:
            return None

    # 获取字典长度和权重矩阵
    def get_embedding_weights(self, word2index, word_vectors):
        vocab_size = len(word2index) + 1  # 字段大小(索引数字的个数)，因为有的词语索引为0，所以+1
        embedding_weights = np.zeros((vocab_size, EMBEDDING_SIZE))  # vocab_size * EMBEDDING_SIZE的0矩阵
        for wd, idx in word2index.items():  # 从索引为1的词语开始，用词向量填充矩阵
            embedding_weights[idx, :] = word_vectors[wd]  # 词向量矩阵，第一行是0向量（没有索引为0的词语）
        return embedding_weights

    #构建CNN模型
    def build_cnn_model(self, embedding_weights):
        model = Sequential()
        vocab_size = len(embedding_weights)
        # 嵌入层将正整数（下标）转换为具有固定大小的向量. eg. [[4], [20]] -> [[0.25, 0.1], [0.6, -0.2]]
        model.add(Embedding(input_dim=vocab_size,  # 字典长度
                            output_dim=EMBEDDING_SIZE,
                            input_length=self.max_len,
                            weights=[embedding_weights]))  # (None, MAX_SENTENCE_LENGTH, EMBEDDING_SIZE), where None is the batch dimension
        model.add(SpatialDropout1D(0.2))
        model.add(Convolution1D(filters=128, kernel_size=3, activation='relu'))
        model.add(GlobalMaxPool1D())  # 对于时间信号的全局最大池化,MaxPooling1D限制每一步的池化大小
        model.add(Dense(1, activation="sigmoid"))
        # model.add(Dense(2, activation='softmax'))

        model.summary()
        return model

    # def build_cnn_model(self, embedding_weights):
    #     model = Sequential()
    #     vocab_size = len(embedding_weights)
    #     # 嵌入层将正整数（下标）转换为具有固定大小的向量. eg. [[4], [20]] -> [[0.25, 0.1], [0.6, -0.2]]
    #     model.add(Embedding(input_dim=vocab_size,  # 字典长度
    #                         output_dim=EMBEDDING_SIZE,
    #                         input_length=self.max_len,
    #                         weights=[embedding_weights]))  # (None, MAX_SENTENCE_LENGTH, EMBEDDING_SIZE), where None is the batch dimension
    #     model.add(SpatialDropout1D(0.2))
    #     model.add(Conv1D(filters=128, kernel_size=3, activation='relu'))
    #     model.add(MaxPooling1D(pool_size=3)) #对于时间信号的全局最大池化,MaxPooling1D限制每一步的池化大小
    #     model.add(Conv1D(filters=64, kernel_size=3, activation='relu'))
    #     model.add(MaxPooling1D(pool_size=3))  # 对于时间信号的全局最大池化,MaxPooling1D限制每一步的池化大小
    #     model.add(Conv1D(filters=32, kernel_size=3, activation='relu'))
    #     model.add(Flatten())
    #     model.add(Dense(1, activation="sigmoid"))
    #     # model.add(Dense(2, activation='softmax'))
    #     model.summary()
    #     return model

    #构建LSTM模型
    def build_lstm_model(self, embedding_weights):
        model = Sequential()
        vocab_size = len(embedding_weights)
        # 嵌入层将正整数（下标）转换为具有固定大小的向量. eg. [[4], [20]] -> [[0.25, 0.1], [0.6, -0.2]]
        model.add(Embedding(input_dim=vocab_size,  # 字典长度
                            output_dim=EMBEDDING_SIZE,
                            # mask_zero=True,
                            input_length=self.max_len,
                            weights=[embedding_weights]))  # (None, MAX_SENTENCE_LENGTH, EMBEDDING_SIZE), where None is the batch dimension
        model.add(SpatialDropout1D(0.2))
        model.add(LSTM(HIDDEN_LAYER_SIZE, dropout=0.2, recurrent_dropout=0.2, return_sequences=False))
        #     model.add(GRU(HIDDEN_LAYER_SIZE, dropout=0.2, recurrent_dropout=0.2, return_sequences=False))
        model.add(Dense(HIDDEN_LAYER_SIZE//2, activation='relu'))
        model.add(Dropout(0.3))
        # model.add(Dense(1, activation="sigmoid"))
        model.add(Dense(2, activation="softmax"))

        model.summary()
        return model

    #构建CNN-LSTM模型
    def build_cnn_lstm_model(self, embedding_weights):
        model = Sequential()
        vocab_size = len(embedding_weights)
        # 嵌入层将正整数（下标）转换为具有固定大小的向量. eg. [[4], [20]] -> [[0.25, 0.1], [0.6, -0.2]]
        model.add(Embedding(input_dim=vocab_size,  # 字典长度
                            output_dim=EMBEDDING_SIZE,
                            input_length=self.max_len,
                            weights=[embedding_weights]))  # (None, MAX_SENTENCE_LENGTH, EMBEDDING_SIZE), where None is the batch dimension
        model.add(SpatialDropout1D(0.2))
        model.add(Conv1D(filters=128, kernel_size=3, strides=1, padding='same', activation='relu'))
        model.add(MaxPooling1D(pool_size=2, strides=2))
        # model.add(Conv1D(filters=64, kernel_size=3, strides=1, padding='same', activation='relu'))
        # model.add(MaxPooling1D(pool_size=2, strides=2))
        model.add(LSTM(HIDDEN_LAYER_SIZE, dropout=0.2, recurrent_dropout=0.2))
        #     model.add(GRU(HIDDEN_LAYER_SIZE, dropout=0.2, recurrent_dropout=0.2))
        model.add(Dense(1, activation="sigmoid"))

        model.summary()
        return model

    #构建LSTM-CNN模型
    def build_lstm_cnn_model(self, embedding_weights):
        model = Sequential()
        vocab_size = len(embedding_weights)
        # 嵌入层将正整数（下标）转换为具有固定大小的向量. eg. [[4], [20]] -> [[0.25, 0.1], [0.6, -0.2]]
        model.add(Embedding(input_dim=vocab_size,  # 字典长度
                            output_dim=EMBEDDING_SIZE,
                            # mask_zero=True,
                            input_length=self.max_len,
                            weights=[embedding_weights]))  # (None, MAX_SENTENCE_LENGTH, EMBEDDING_SIZE), where None is the batch dimension
        model.add(SpatialDropout1D(0.2))
        model.add(LSTM(HIDDEN_LAYER_SIZE, dropout=0.2, recurrent_dropout=0.2))
        #     model.add(GRU(HIDDEN_LAYER_SIZE, dropout=0.2, recurrent_dropout=0.2))
        model.add(Conv1D(filters=32, kernel_size=3, strides=1, padding='same', activation='relu'))
        model.add(MaxPooling1D(pool_size=2, strides=2))
        model.add(Dense(1, activation="sigmoid"))

        model.summary()
        return model

    #构建Bi-LSTM模型
    def build_bilstm_model(self, embedding_weights):
        vocab_size = len(embedding_weights)
        # 嵌入层将正整数（下标）转换为具有固定大小的向量. eg. [[4], [20]] -> [[0.25, 0.1], [0.6, -0.2]]
        model = Sequential()
        model.add(Embedding(input_dim=vocab_size,
                            output_dim=EMBEDDING_SIZE,
                            input_length=self.max_len,
                            weights=[embedding_weights]))  # (None, MAX_SENTENCE_LENGTH, EMBEDDING_SIZE), where None is the batch dimension
        model.add(SpatialDropout1D(0.2))
        model.add(Bidirectional(LSTM(HIDDEN_LAYER_SIZE, dropout=0.2, recurrent_dropout=0.2, return_sequences=True)))
        model.add(Bidirectional(LSTM(HIDDEN_LAYER_SIZE, dropout=0.2)))
        # model.add(Dense(128, activation='relu'))
        # model.add(Dropout(0.5))
        # model.add(Flatten())
        model.add(Dense(1, activation="sigmoid"))

        # inputs = Input(shape=(None,))
        # embedded = Embedding(vocab_size, EMBEDDING_SIZE, input_length=self.max_len)(inputs)
        # lstm_out = LSTM(HIDDEN_LAYER_SIZE, dropout=0.2, recurrent_dropout=0.2, return_sequences=False)(embedded)
        # predict = Dense(1, activation='softmax')(lstm_out)
        # model = Model(inputs=inputs, outputs=predict)
        model.summary()
        return model

    def F1_score(self, y_true, y_pred):
        # Only computes a batch-wise average of recall.
        def recall(y_true, y_pred):
            true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
            possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
            recall = true_positives / (possible_positives + K.epsilon())
            return recall

        # Only computes a batch-wise average of precision.
        def precision(y_true, y_pred):
            true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
            predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
            precision = true_positives / (predicted_positives + K.epsilon())
            return precision

        precision = precision(y_true, y_pred)
        recall = recall(y_true, y_pred)
        return 2 * ((precision * recall) / (precision + recall + K.epsilon()))

    #构建并训练LSTM模型
    def train_lstm_model(self, embedding_weights, X, y):
        lstm_model = self.build_lstm_model(embedding_weights)
        # lstm_model.compile(loss=losses.binary_crossentropy, optimizer="adam", metrics=["acc", self.F1_score])
        lstm_model.compile(loss=self.loss_new, optimizer="adam", metrics=["acc", self.F1_score])
        Xtrain, Xtest, ytrain, ytest = train_test_split(X, y, test_size=TEST_SIZE, random_state=42, shuffle=True)
        ytrain = keras.utils.to_categorical(ytrain, num_classes=2)
        ytest = keras.utils.to_categorical(ytest, num_classes=2)

        self.logging.info('开始训练LSTM模型......')
        self.logging.info('训练集大小：{} 测试集大小：{}'.format(round((1-TEST_SIZE)*len(X)), round(len(X)*TEST_SIZE)))
        t1 = time.time()
        # 保存训练最好的模型（每3轮训练检测一次）
        # cp = ModelCheckpoint(filepath=BEST_MODEL_PATH, monitor='val_acc', mode='max', save_best_only=True, verbose=1,
        #                      period=2)
        # tb = TensorBoard(log_dir='log')
        es = EarlyStopping(monitor='val_loss', patience=NUM_EPOCHS / 2)  # 经过NUM_EPOCHS/2轮训练，当val_acc不再提升，则停止训练
        callbacks_lst = [es]
        hist = lstm_model.fit(Xtrain,
                        ytrain,
                        batch_size=BATCH_SIZE,
                        epochs=NUM_EPOCHS,
                        callbacks=callbacks_lst,
                        validation_data=(Xtest, ytest))
        t2 = time.time()
        self.logging.info('LSTM模型训练结束！总用时：{}min'.format((t2 - t1) / 60.0))

        with open(HIST_PATH, 'wb') as output: #保存history对象
            pickle.dump(hist.history, output)

        # yaml_string = lstm_model.to_yaml()  # 保存模型结构为YAML字符串
        # with open(LSTM_MODEL_PATH, 'w') as fout:
        #     fout.write(yaml.dump(yaml_string, default_flow_style=True))
        # lstm_model.save_weights(LSTM_WEIGHT_PATH)  # 保存模型权重
        # self.logging.info('LSTM模型已经保存......')

        self.logging.info(lstm_model.metrics_names)
        loss, acc, F1_score = lstm_model.evaluate(Xtest, ytest, batch_size=BATCH_SIZE//2)
        self.logging.info('模型评估结果：loss:{} acc:{} F1:{}'.format(loss, acc, F1_score))

    def plot_hist(self):
        # plot loss and accuracy
        with open(HIST_PATH, 'rb') as pkl_hist:
            history = pickle.load(pkl_hist)
        plt.subplot(211)
        plt.title("Accuracy")
        # plt.plot(history["acc"], color="g", label="Train")
        plt.plot(history["val_acc"], color="b", label="Validation")
        plt.legend(loc="best")
        plt.subplot(212)
        plt.title("Loss")
        # plt.plot(history["loss"], color="g", label="Train")
        plt.plot(history["val_loss"], color="r", label="Validation")
        plt.legend(loc="best")
        plt.tight_layout() #自动调整子图间的间距
        plt.show()

    #加载LSTM模型
    def load_lstm_model(self):
        self.logging.info('loading models......')
        with open(LSTM_MODEL_PATH, 'r') as f:
            yaml_string = yaml.load(f)
        lstm_model = model_from_yaml(yaml_string)
        self.logging.info('loading weights......')
        lstm_model.load_weights(LSTM_WEIGHT_PATH)
        lstm_model.compile(loss=self.focal_loss, optimizer="adam", metrics=["accuracy"])
        return lstm_model

    #将分词序列转换成索引序列（分词序列由list列表传入）
    def text2index_from_lst(self, word2index, comm_seqs):
        data = []
        for seqs in comm_seqs:
            wd_idxs = []
            for wd in seqs:
                if wd in word2index.keys():
                    wd_idxs.append(word2index[wd])  # 单词转索引数字
                else:
                    wd_idxs.append(0)  # 索引字典里没有的词转为数字0
            data.append(wd_idxs)

        return sequence.pad_sequences(data, self.max_len)  # 对齐序列

    #情感预测（传入评论列表）
    def predict_by_lst(self, new_comms):
        lstm_model = self.load_lstm_model() #LSTM模型
        word2vec_model = Word2Vec.load(WORD2VEC_PATH) #词向量模型
        word2index, _ = self.create_dicts(word2vec_model)
        new_comms = [TextUtils.process(com) for com in new_comms if not TextUtils.is_blank(com)]
        wd_seqs = [TextUtils.tokenize(com) for com in new_comms if not TextUtils.is_blank(com)]
        X = self.text2index_from_lst(word2index, wd_seqs)
        for i, x in enumerate(X):
            x = np.array(x).reshape(1, -1)
            res = lstm_model.predict_classes(x)
            if res[0][0] == LABEL_POS:
                print(new_comms[i], '\n', 'positive!')
            else:
                print(new_comms[i], '\n', 'negative!')

    #情感预测（传入文件路径）
    def predict_by_file(self, file_path):
        with open(file_path, 'r', encoding='utf-8') as fin:
            # new_comms = fin.read().splitlines()
            new_coms = [line.strip() for line in fin if not TextUtils.is_blank(line)]
            self.predict_by_lst(new_coms)

    #训练
    def train(self, comment_dir=COMMENT_DIR):
        #加载数据
        df_data = self.load_train_data(comment_dir)
        #训练词向量模型
        word2vec_model = self.train_wd2vect(df_data['segs'])
        #创建索引词典和词向量
        word2index, word_vectors = self.create_dicts(word2vec_model)
        # 获取权值矩阵
        embedding_weights = self.get_embedding_weights(word2index, word_vectors)
        #文本序列索引化
        comms_seqs = self.text2index_from_lst(word2index, df_data['segs'])
        #序列LSTM
        self.train_lstm_model(embedding_weights, comms_seqs, df_data['label'])

if __name__ == "__main__":
    model = SentimentModel()
    model.train('data/test.csv')
    model.plot_hist()
    # new_comms = ['性价比非常高，显示屏幕高清，处理速度快！不带鼠标。很轻，没有想象中的薄。自带的office365需要激活',
    #              '非常好看！很满意，这个价钱买到质量这么好的鞋子简直太感动了，客服推荐的码数刚好，很合脚，穿起来不磨脚，防滑，耐脏，绝对值！',
    #              '什么破商家。我申请换货，不让我换，硬要求我退货，我说不退，直接把我电话挂掉，服务差评',
    #              '冰箱收到了，空间挺大的，冷冻室有点小，能效是三级，冷冻，冷藏，保鲜三室可以调温，送来的时候冰箱里面有点味道！送电以后就没了！总得来说这次购物不错！',
    #              '发申通物流，很不靠谱，太慢了，本来就是冲着京东物流才买的，结果发的申通，下次再不买它家的了',
    #              '鞋子已是第二双，朋友看到好看帮忙买的，质量杠杠的，又轻又舒服，春天到了配啥衣服都好看，出去玩也不怕累脚了，值得推荐！！！！',
    #              '用了两天才过来评论的，首先电脑颜值高，四面窄观影简直不要太舒服，背光键盘也很nice，尺寸大小也很合适，开机使用反应超快的。刚拿到手就插上电@了，语音助手指导小白也能上手。值得购买！',
    #              '很差，给老爸买的靴子，本来想着尽一份孝心，结果给我的感触很大，家里在县城，那么多快递可以到，偏偏发一个到不了的快递，也不跟买家核实地址，也不打招呼，哪怕跟我们核实一下，我们可以自己选择快递啊。那么大老远，还下着大雨，直接让老人去取。这次购物体验很差，很差。卖家服务态度也很差，很差']
    # model.predict_by_lst(new_comms)

    # model.predict_by_file("user/samples.txt")