import strax
import straxen
import numba
import numpy as np
from immutabledict import immutabledict

export, __all__ = strax.exporter()


# More info about the acquisition monitor can be found here:
# https://xe1t-wiki.lngs.infn.it/doku.php?id=xenon:xenon1t:alexelykov:acquisition_monitor

@export
@strax.takes_config(strax.Option('hit_min_amplitude_aqmon', default=50, track=False,
                                 help='Minimum hit threshold in ADC*counts above baseline'),
                    strax.Option('baseline_samples_aqmon', default=10, track=False,
                                 help='Number of samples to use at the start of the pulse to determine the baseline'))
class AqmonHits(strax.Plugin):
    """ Find hits in acquisition monitor data. These hits could be 
        then used by other plugins for deadtime calculations, 
        GPS SYNC analysis, etc.
    """
    __version__ = '0.0.3'

    depends_on = ('raw_records_aqmon')
    provides = ('aqmon_hits')
    data_kind = ('aqmon_hits')

    dtype = strax.hit_dtype

    save_when = strax.SaveWhen.TARGET

    def compute(self, raw_records_aqmon):
        rec = strax.raw_to_records(raw_records_aqmon)
        strax.zero_out_of_bounds(rec)
        strax.baseline(rec, baseline_samples=self.config['baseline_samples_aqmon'], flip=True)
        aqmon_hits = strax.find_hits(rec, min_amplitude=self.config['hit_min_amplitude_aqmon'])

        return aqmon_hits

#### Veto hardware ####:
# V1495 busy veto module:
# Generates a 25 ns NIM pulse whenever a veto begins and a 25 ns NIM signal when it ends.
# A new start signal can occur only after the previous busy instance ended.
# 1ms (1e6 ns) - minimum busy veto length, or until the board clears its memory

# DDC10 High Energy Veto:
# 10ms (1e7 ns) - fixed HE veto length in XENON1T DDC10,
# in XENONnT it will be calibrated based on the length of large S2 SE tails
# The start/stop signals for the HEV are generated by the V1495 board

@export
@strax.takes_config(
    strax.Option('min_veto_gap', default=int(1e6), type=int,
                 help='Minimum separation between veto stop and start pulses [ns]'),
    strax.Option('max_veto_gap', default=int(5e8), type=int,
                 help='Maximum separation between veto stop and start pulses [ns]'),
    strax.Option('channel_map', track=False, type=immutabledict,
                 help='immutabledict mapping subdetector to (min, max)'
                      'channel number.'))
class VetoIntervals(strax.OverlapWindowPlugin):
    """ Find pairs of veto start and veto stop signals and the veto duration between them
    busy_*  <= V1495 busy veto for tpc channels
    he_*    <= V1495 busy veto for high energy tpc channels
    hev_*   <= DDC10 hardware high energy veto
    """

    __version__ = '0.1.1'
    depends_on = ('aqmon_hits')
    provides = ('veto_intervals')
    data_kind = ('veto_intervals')

    def infer_dtype(self):
        dtype = [(('veto start time since unix epoch [ns]', 'time'), np.int64),
                 (('veto end time since unix epoch [ns]', 'endtime'), np.int64),
                 (('veto interval [ns]', 'veto_interval'), np.int64),
                 (('veto signal type', 'veto_type'), np.str('U9'))]
        return dtype

    def setup(self):
        aq_ch_start = straxen.n_hard_aqmon_start
        aqmon_channel_names = ('sum_wf', 'm_veto_sync',
                               'hev_stop', 'hev_start', 'he_stop', 'he_start', 'busy_stop', 'busy_start')
        channel_numbers = np.arange(aq_ch_start, aq_ch_start + len(aqmon_channel_names))

        # Keeping ch_range as class attribute for now, maybe will need it for other features
        self.channel_range = self.config['channel_map']['aqmon']
        self.channel_map = dict(zip(aqmon_channel_names, channel_numbers))
        self.veto_names = ['busy_', 'he_', 'hev_']

    def get_window_size(self):
        # Give a very wide window
        return (self.config['max_veto_gap'] * 100)

    def compute(self, aqmon_hits):
        hits = aqmon_hits

        # Make a dict to populate later
        res_temp = dict()

        for i, veto in enumerate(self.veto_names):
            veto_hits = channel_select(hits, self.channel_map[veto + 'stop'], self.channel_map[veto + 'start'])
            vetos = self.merge_vetos(veto_hits, gap=self.config['min_veto_gap'], \
                                     dtype=strax.time_fields, t=0)

            # Check that we found a veto interval and update the resulting dict
            if len(vetos):
                res_temp.setdefault("time", []).extend(vetos['time'])
                res_temp.setdefault("endtime", []).extend(vetos['endtime'])
                res_temp.setdefault("veto_interval", []).extend((vetos['endtime'] - vetos['time']))
                res_temp.setdefault("veto_type", []).extend([veto + 'veto'] * len(vetos))

        res = strax.dict_to_rec(res_temp, dtype=self.dtype)
        res.sort(order='time')

        return res

    @staticmethod
    def merge_vetos(channels, gap, dtype, t):
        if len(channels):
            start, stop = strax.find_peak_groups(channels, gap_threshold=gap)
            result = np.zeros(len(start), dtype=dtype)
            result['time'] = start
            result['endtime'] = stop
        else:
            result = np.zeros(0, dtype=dtype)
        return result


@numba.njit
def channel_select(rr, ch_stop, ch_start):
    """Return data in between stop and start channels of the acquisition monitor (AM)"""
    return rr[((rr['channel'] == ch_stop) | (rr['channel'] == ch_start))] 
