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

export, __all__ = strax.exporter()


@export
class OnlineMonitorNV(strax.Plugin):
    """
    Plugin to write data of nVeto detector to the online-monitor. 
    Data that is written by this plugin should be small (~MB/chunk) 
    to not overload the runs-database.

    This plugin takes 'hitlets_nv' and 'events_nv'. Although they are
    not strictly related, they are aggregated into a single data_type
    in order to minimize the number of documents in the online monitor.

    Produces 'online_monitor_nv' with info on the hitlets_nv and events_nv
    """

    __version__ = '0.0.4'

    depends_on = ('hitlets_nv', 'events_nv')
    provides = 'online_monitor_nv'
    data_kind = 'online_monitor_nv'

    # Needed in case we make again an muVETO child.
    ends_with = '_nv'

    channel_map = straxen.URLConfig(
        track=False,
        type=immutabledict,
        help='immutabledict mapping subdetector to (min, max) '
             'channel number.')

    events_area_bounds = straxen.URLConfig(
        type=tuple, default=(-0.5, 130.5),
        help='Boundaries area histogram of events_nv_area_per_chunk [PE]')

    events_area_nbins = straxen.URLConfig(
        type=int, default=131,
        help='Number of bins of histogram of events_nv_area_per_chunk, '
             'defined value 1 PE/bin')

    def infer_dtype(self):
        self.channel_range = self.channel_map['nveto']
        self.n_channel = (self.channel_range[1] - self.channel_range[0]) + 1
        return veto_monitor_dtype(self.ends_with, self.n_channel, self.events_area_nbins)

    def compute(self, hitlets_nv, events_nv, start, end):
        # General setup
        res = np.zeros(1, dtype=self.dtype)
        res['time'] = start
        res['endtime'] = end

        # Count number of hitlets_nv per PMT
        hitlets_channel_count, _ = np.histogram(hitlets_nv['channel'],
                                                bins=self.n_channel,
                                                range=[self.channel_range[0],
                                                       self.channel_range[1] + 1])
        res[f'hitlets{self.ends_with}_per_channel'] = hitlets_channel_count

        # Count number of events_nv with coincidence cut
        res[f'events{self.ends_with}_per_chunk'] = len(events_nv)
        sel = events_nv['n_contributing_pmt'] >= 4
        res[f'events{self.ends_with}_4coinc_per_chunk'] = np.sum(sel)
        sel = events_nv['n_contributing_pmt'] >= 5
        res[f'events{self.ends_with}_5coinc_per_chunk'] = np.sum(sel)
        sel = events_nv['n_contributing_pmt'] >= 8
        res[f'events{self.ends_with}_8coinc_per_chunk'] = np.sum(sel)
        sel = events_nv['n_contributing_pmt'] >= 10
        res[f'events{self.ends_with}_10coinc_per_chunk'] = np.sum(sel)

        # Get histogram of events_nv_area per chunk
        events_area, bins_ = np.histogram(events_nv['area'],
                                          bins=self.events_area_nbins,
                                          range=self.events_area_bounds)
        res[f'events{self.ends_with}_area_per_chunk'] = events_area
        return res


def veto_monitor_dtype(veto_name: str = '_nv',
                       n_pmts: int = 120,
                       n_bins: int = 131) -> list:
    dtype = []
    dtype += strax.time_fields  # because mutable
    dtype += [((f'hitlets{veto_name} per channel', f'hitlets{veto_name}_per_channel'), (np.int64, n_pmts)),
              ((f'events{veto_name}_area per chunk', f'events{veto_name}_area_per_chunk'), np.int64, n_bins),
              ((f'events{veto_name} per chunk', f'events{veto_name}_per_chunk'), np.int64),
              ((f'events{veto_name} 4-coincidence per chunk', f'events{veto_name}_4coinc_per_chunk'), np.int64),
              ((f'events{veto_name} 5-coincidence per chunk', f'events{veto_name}_5coinc_per_chunk'), np.int64),
              ((f'events{veto_name} 8-coincidence per chunk', f'events{veto_name}_8coinc_per_chunk'), np.int64),
              ((f'events{veto_name} 10-coincidence per chunk', f'events{veto_name}_10coinc_per_chunk'), np.int64)
              ]
    return dtype
