# Open and analyze UCN data for a whole run
# Derek Fujimoto
# June 2024

import ROOT
ROOT.gROOT.SetBatch(1)
from rootloader import tfile, ttree, attrdict
from .exceptions import *
from .applylist import applylist
from .ucnbase import ucnbase
from .ttreeslow import ttreeslow
from .ucncycle import ucncycle
from .datetime import to_datetime
import itertools, warnings, os

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from collections.abc import Iterable
from tqdm import tqdm

# patch warnings
def new_format(message, category, filename, lineno, line):
    filename = os.path.basename(filename)
    return f'\n{filename}:{lineno}: {category.__name__}: {message}\n'

warnings.formatwarning = new_format

class ucnrun(ucnbase):
    """UCN run data. Cleans data and performs analysis

    Args:
        run (int|str): if int, generate filename with self.datadir
            elif str then run is the path to the file
        ucn_only (bool): if true set filter tIsUCN==1 on all hit trees

    Attributes:
        comment (str): comment input by users
        cycle (int|none): cycle number, none if no cycle selected
        cycle_param (attrdict): cycle parameters from sequencer settings
        experiment_number (str): experiment number input by users
        month (int): month of run start
        run_number (int): run number
        run_title (str): run title input by users
        shifter (str): experimenters on shift at time of run
        start_time (str): start time of the run
        stop_time (str): stop time of the run
        supercycle (int|none): supercycle number, none if no cycle selected
        tfile (tfile): stores tfile raw readback
        year (int): year of run start

    Notes:
        * Can access attributes of tfile directly from top-level object
        * Need to define the values if you want non-default behaviour
        * Object is indexed as [cycle, period] for easy access to sub time frames

        Cycle param contents

        `nperiods`: Number of periods in each cycle
        `nsupercyc`: Number of supercycles contained in the run
        `enable`: Enable status of the sequencer
        `inf_cyc_enable`: Enable status of infinite cycles
        `cycle`: Cycle ID numbers
        `supercycle`: Supercycle ID numbers
        `valve_states`: Valve states in each period and cycle
        `period_end_times`: End time of each period in each cycle in epoch time
        `period_durations_s`: Duration in sections of each period in each cycle
        `ncycles`: Number of total cycles contained in the run
        `filter`: A list indicating how we should filter cycles. More on that in [filters](filters.md)
        `cycle_time`: The start and end times of each cycle

    Examples:

        Loading runs
        ```python
        from ucndata import ucnrun, settings
        # load from filename
        ucnrun('/path/to/file/ucn_run_00002684.root')
        # load from run number
        ucnrun.datadir = '/path/to/file/'
        ucnrun(2684)
        ```

        Slicing
        ```python
        from ucndata import ucnrun
        run = ucnrun(2684)
        run[0, 0]   # cycle 0, period 0
        run[:]      # all cycles, no distinction on period
        run[:, 0]   # get all cycles, period 0
        run[0, :]   # cycle 0, all periods
        run[4:7, 2] # cycles 4, 5, 6, period 2
        ```

        Get beam properties
        ```python
        from ucndata import ucnrun
        run = ucnrun(2684)
        run.beam_current_uA # beam current in uA
        run.beam_on_s       # beam duration on in s
        run.beam_off_s      # beam duration off in s
        ```

        Draw hits
        ```python
        import matplotlib.pyplot as plt
        from ucndata import ucnrun
        run = ucnrun(2684)

        # draw all hits in the file
        run.get_hits_histogram('Li6').plot()

        # draw hits in each cycle
        for cycle in run:
            cycle.get_hits_histogram('Li6').plot(label=cycle.cycle)

        # adjust the timing of cycles and periods
        run._modify_ptiming(cycle=0, period=0, dt_start_s=1, dt_start_s=0)

        # inspect the data: draw a figure with hits histogram, beam current, and optional slow control data
        run.inspect('Li6', bin_ms=100, xmode='dur')
        ```
    """

    def __init__(self, run, ucn_only=True):

        # check if copying
        if run is None:
            return

        # make filename from defaults
        elif isinstance(run, (int, float, np.float64, np.int64)):
            filename = os.path.join(self.datadir, f'ucn_run_{int(run):0>8d}.root')

        # fetch from specified path
        elif isinstance(run, str):
            filename = run

        else:
            raise TypeError(f'Unknown input type "{type(run)}" for run')

        self.path = os.path.abspath(filename)

        # read
        self.tfile = tfile(filename, empty_ok=False, quiet=True,
                            key_filter=self.keyfilter)
        head = self.tfile['header'].to_dataframe()

        # reformat header and move to top level
        for k, val in head.items():
            setattr(self, k.replace(' ', '_').lower(), val.loc[0])

        if type(self.run_number) is pd.Series:
            self.run_number = int(float(self.run_number[0]))
        else:
            self.run_number = int(float(self.run_number))

        # set other header items
        date = pd.to_datetime(self.start_time)
        self.year = date.year
        self.month = date.month

        # reformat tfile branch names to remove spaces
        for key in tuple(self.tfile.keys()):
            if ' ' in key or '-' in key:
                self.tfile[key.replace(' ', '_').replace('-','')] = self.tfile[key]
                del self.tfile[key]

        # some trees should be dataframes by default
        for name in self.DATAFRAME_TREES:
            if name in self.tfile.keys():
                self.tfile[name] = self.tfile[name].to_dataframe()
                self.tfile[name].reset_index(inplace=True, drop=True)

        # set cycle parameters
        self._get_cycle_param()

        # get cycle times
        if hasattr(self, 'cycle_param'):
            cycle_failed = False
            if isinstance(self.cycle_times_mode, str):
                self.cycle_times_mode = [self.cycle_times_mode]

            for mode in self.cycle_times_mode:
                try:
                    self.set_cycle_times(mode=mode)
                except (AttributeError, IndexError):
                    warnings.warn(f'Run {self.run_number}: Unable to set cycle times. SequencerTree must exist and have entries.',
                                MissingDataWarning)
                    break
                except (KeyError, CycleError):
                    print(f'Run {self.run_number}: Cycle time detection mode {mode} failed')
                    cycle_failed = True
                else:
                    break

            if cycle_failed:
                print(f'Run {self.run_number}: Set cycle times based on {mode} detection mode')

        # setup tree filters
        for detector in self.DET_NAMES.keys():
            try:
                tree = self.tfile[self.DET_NAMES[detector]['hits']]
            except KeyError:
                continue

            # purge bad timestamps
            tree.set_filter('tUnixTimePrecise > 15e8', inplace=True)

            # filter on ucn hits
            if ucn_only:
                tree.set_filter('tIsUCN>0', inplace=True)


        # make slow control tree
        self.epics = ttreeslow((self.tfile[name] for name in self.EPICS_TREES if name in self.tfile.keys()),
                               parent=self)

        # store fetched cycles
        self._cycledict = dict()

        # store fetched histogram
        self._hits_hist = {'detector':None,
                           'bin_ms': None,
                           'hist':None}

        # histograms with edges set by period and cycle timings
        self._nhits = {}

        # pointer to self for cycles and periods
        self._run = self

    def __next__(self):
        # permit iteration over object like it was a list

        if self._iter_current < self.cycle_param.ncycles:

            # skip elements that should be filtered
            if self.cycle_param.filter is not None:

                # skip
                while not self.cycle_param.filter[self._iter_current]:
                    self._iter_current += 1

                    # end condition
                    if self._iter_current >= self.cycle_param.ncycles:
                        raise StopIteration()

            # iterate
            cyc = self[self._iter_current]
            self._iter_current += 1
            return cyc

        # end of iteration
        else:
            raise StopIteration()

    def __repr__(self):
        klist = [d for d in self.__dict__.keys() if d[0] != '_']
        if klist:

            # sort without caps
            klist.sort(key=lambda x: x.lower())

            # get number of columns based on terminal size
            maxsize = max((len(k) for k in klist)) + 2
            terminal_width = os.get_terminal_size().columns
            ncolumns = int(np.floor(terminal_width / maxsize))
            ncolumns = min(ncolumns, len(klist))

            # split into chunks
            needed_len = int(np.ceil(len(klist) / ncolumns)*ncolumns) - len(klist)
            klist = np.concatenate((klist, np.full(needed_len, '')))
            klist = np.array_split(klist, ncolumns)

            # print
            s = f'run {self.run_number}:\n'
            for key in zip(*klist):
                s += '  '
                s += ''.join(['{0: <{1}}'.format(k, maxsize) for k in key])
                s += '\n'
            return s
        else:
            return self.__class__.__name__ + "()"

    def __getitem__(self, key):
        # get cycle or period based on slicing indexes

        # get a single key
        if isinstance(key, (np.integer, int)):

            if key < 0:
                key = self.cycle_param.ncycles + key

            if key > self.cycle_param.ncycles:
                raise IndexError(f'Run {self.run_number}: Index larger than number of cycles ({self.cycle_param.ncycles})')

            return self.get_cycle(key)

        # slice on cycles
        if isinstance(key, slice):
            cycles = self.get_cycle()[:self.cycle_param.ncycles]

            # no filter
            if self.cycle_param.filter is None or all(self.cycle_param.filter):
                cyc = cycles[key]

            # yes filter
            else:

                # fetch the filter and slice in the same way as the return value
                cfilter = self.cycle_param.filter[key]

                # fetch cycles and slice, then apply filter
                cyc = np.array(cycles[key])
                cyc = cyc[cfilter]

            return applylist(cyc)

        # slice on periods
        if isinstance(key, tuple):
            cycles = self[key[0]]
            if isinstance(cycles, (np.ndarray, applylist, list)):
                return applylist([c[key[1]] for c in cycles])
            else:
                return cycles[key[1]]

        raise IndexError(f'Run {self.run_number} given an unknown index type ({type(key)})')

    def __len__(self):
        return self.cycle_param.ncycles

    def _get_cycle_param(self):
        # set self.cycle_param dict

        # reformat cycle param tree
        paramtree = self.tfile.CycleParamTree
        self.cycle_param = {'nperiods': paramtree.nPeriods[0],
                            'nsupercyc': paramtree.nSuperCyc[0],
                            'enable': bool(paramtree.enable[0]),
                            'inf_cyc_enable': bool(paramtree.infCyclesEnable[0]),
                            'ncycles': 1
                            }
        self.cycle_param = attrdict(self.cycle_param)

        # cycle filter
        self.cycle_param['filter'] = None

        # setup cycle paramtree array outputs from transition trees
        treelist = []
        for detector in self.DET_NAMES.values():
            if detector['transitions'] in self.tfile.keys():
                treelist.append(self.tfile[detector['transitions']])

        # check if trees exist
        if not treelist:
            warnings.warn(f'Run {self.run_number}: no detector transition tree, cannot fully set up cycle_param', MissingDataWarning)

            # number of cycles, periods
            self.cycle_param['ncycles'] = 1
            self.cycle_param['nperiods'] = 1

            # cycle and supercycle indices
            self.cycle_param['cycle'] = pd.Series([0], index=[0])
            self.cycle_param['supercycle'] = pd.Series([0], index=[0])

            self.cycle_param['cycle'].name = 'cycleIndex'
            self.cycle_param['supercycle'].name = 'superCycleIndex'

            return

        # get longest
        treelen = [len(t) for t in treelist]
        tree = treelist[np.argmax(treelen)]

        # cycle and supercycle indices
        self.cycle_param['cycle'] = tree['cycleIndex'].astype(int)
        self.cycle_param['supercycle'] = tree['superCycleIndex'].astype(int)

        # convert the array in each cell into a dataframe
        def item_to_df(x):
            s = pd.DataFrame(np.array(x).copy(), index=np.arange(len(x))+1)
            return s

        # valve states -------------------------------------------------------
        df = tree[[col for col in tree.columns if 'valveState' in col]]
        col_map = {col:int(col.replace("valveStatePeriod", "")) for col in df.columns}
        df = df.rename(columns=col_map)
        df.columns.name = 'period'
        df.reset_index(inplace=True, drop=True)
        df.index.name = 'cycle_idx'

        # valve states should not change across cycles
        df2 = df.loc[0]
        df2 = pd.concat([item_to_df(df2[period]) for period in df2.index], axis='columns')

        # rename columns and index
        df2.columns = np.arange(len(df2.columns))
        df2.columns.name = 'period'
        df2.index.name = 'valve'
        self.cycle_param['valve_states'] = df2.transpose()

        # period end times ---------------------------------------------------
        df = tree[[col for col in tree.columns if 'cyclePeriod' in col]]
        col_map = {col:int(col.replace("cyclePeriod", "").replace("EndTime", "")) for col in df.columns}

        # rename columns and index
        df = df.rename(columns=col_map)
        df.columns.name = 'period'
        df.index.name = 'cycle'
        self.cycle_param['period_end_times'] = df.transpose()

        # period durations ---------------------------------------------------
        cycle_start = tree.cycleStartTime
        df_diff = df.diff(axis='columns')
        df_diff[0] = df[0] - cycle_start

        self.cycle_param['period_durations_s'] = df_diff.transpose()

        # number of cycles
        self.cycle_param['ncycles'] = len(df.index)

    def _get_nhits(self, detector, cycle=None, period=None):
        # Get number ucn hits

        # Args:
        #     detector (str): Li6|He3
        #     cycle (int): cycle number, if None compute for whole run
        #     period (int): period number, if None compute for whole cycle, if cycle is not also None

        # Notes:
        #     Because of how RDataFrame works it is better to compute a histogram whose bins correspond to the period or cycle start/end times than to set a filter and get the resulting tree size.
        #     The histogram bin quantities is saved as self._nhits or self._nhits_cycle
        #     Both ucncycle and ucnperiod classes call this function to get the counts

        # get hits tree
        tree = self.tfile[self.DET_NAMES[detector]['hits']]

        # get hits for run
        if cycle is None and period is None:
            return tree.size

        else:

            # make hits histogram
            if detector not in self._nhits.keys():
                
                # use period and cycle start and end times as bin edges
                edges = []
                for cyclei in range(self.cycle_param.ncycles):

                    # get cycle start and end
                    cycle_start = self.cycle_param.cycle_times.start[cyclei].copy()
                    period_ends = self.cycle_param.period_end_times[cyclei].copy()

                    # shorten periods that extend past the end of the cycle (edge case)
                    if cyclei < self.cycle_param.ncycles-1:
                        next_cycle_start = self.cycle_param.cycle_times.start[cyclei+1]
                    else:
                        next_cycle_start = self.cycle_param.cycle_times.stop.iloc[-1]

                    period_ends[period_ends > next_cycle_start] = next_cycle_start

                    # save to list
                    edges.append(cycle_start)
                    edges.extend(list(period_ends))

                edges = np.append(edges, self.cycle_param.cycle_times.stop.iloc[-1])

                # discard duplicate edges - this drop counts for zero length periods, we re-insert these after
                # first bin is hits before first edge - ie before cycle start
                hits = tree.hist1d('tUnixTimePrecise', edges=np.unique(edges)).y[1:]

                # rebuild hits array with zero counts for zero length periods
                dur = self.cycle_param.period_durations_s.values.transpose() # durations[cycle,period]
                idx = 0 # index in the hits array to copy for non-zero durations
                self._nhits[detector] = []
                for cyci in range(dur.shape[0]):
                    for peri in range(dur.shape[1]):

                        # non-zero duration, add hits from historgram
                        # index exceeds hits len if run ended before last cycle finished
                        if dur[cyci,peri] > 0 and idx < len(hits):
                            self._nhits[detector].append(hits[idx])
                            idx += 1
                        
                        # zero duration add zero counts
                        else:
                            self._nhits[detector].append(0)
                        
                    # end of cycle hits after last period
                    if idx < len(hits):
                        self._nhits[detector].append(hits[idx])
                        idx += 1

                # convert to np array
                self._nhits[detector] = np.array(self._nhits[detector])

            # get hits for cycle
            if period is None:
                nperiodsp1 = len(self.cycle_param.period_end_times)+1
                return int(self._nhits[detector][cycle*nperiodsp1:(cycle+1)*nperiodsp1].sum())

            # get hits for period
            else:
                return int(self._nhits[detector][cycle*(len(self.cycle_param.period_end_times)+1)+period])

    def _modify_ptiming(self, cycle, period, dt_start_s=0, dt_stop_s=0, update_duration=True):
        # Change start and end times of periods

        # Args:
        #     cycle (int): cycle id number
        #     period (int): period id number
        #     dt_start_s (float): change to the period start time
        #     dt_stop_s (float): change to the period stop time
        #     update_duration (bool): if true, update period durations dataframe

        # Notes:
        #     * as a result of this, cycles may overlap or have gaps
        #     * periods are forced to not overlap and have no gaps
        #     * cannot change cycle end time, but can change cycle start time
        #     * this function resets all saved histgrams and hits

        # get cycle parameters
        cycpar = self.cycle_param

        # start times - prevent unnecessary recursion
        if dt_start_s != 0:

            # adjust cycle start
            if period == 0:
                cycpar.cycle_times.loc[cycle, 'start'] += dt_start_s

            # period start time adjustment
            else:
                cycpar.period_end_times.loc[period-1, cycle] += dt_start_s

            # recursively adjust adjacent size zero periods
            if period-1 > 0 and cycpar.period_durations_s.loc[period-1, cycle] == 0:
                self._modify_ptiming(cycle, period-1,
                                    dt_start_s = dt_start_s,
                                    dt_stop_s  = 0,
                                    update_duration = False)

        # stop times - prevent unnecessary recursion
        if dt_stop_s != 0:

            # period end time adjustment
            cycpar.period_end_times.loc[period, cycle] += dt_stop_s

            # force periods to stay within cycle bounds
            cycend = cycpar.cycle_times.loc[cycle, 'stop']
            perend = cycpar.period_end_times.loc[period, cycle]
            cycpar.period_end_times.loc[period, cycle] = min(perend, cycend)

            # recursively adjust adjacent size zero periods
            try:
                if cycpar.period_durations_s.loc[period+1, cycle] == 0:
                    self._modify_ptiming(cycle, period+1,
                                        dt_start_s = 0,
                                        dt_stop_s  = dt_stop_s,
                                        update_duration = False)

            # period not in durations dataframe
            except KeyError:
                pass

        # adjust period and cycle durations
        if update_duration:

            start = cycpar.cycle_times['start']
            stop = cycpar.cycle_times['stop']

            df = cycpar.period_end_times
            df_diff = df.diff()
            df_diff.loc[0] = df.loc[0] - start
            cycpar.period_durations_s = df_diff

            cycpar.cycle_times.loc[cycle, 'duration (s)'] = (stop-start).loc[0]

        # set in memory
        self.cycle_param = cycpar

        # remove saved cycles to account for updated limits
        if cycle in self._cycledict.keys():
            del self._cycledict[cycle]

        # reset histogram for number of hits
        self._nhits = {}

        # reset all saved histograms
        for key in self._hits_hist.keys():
            self._hits_hist[key] = None

    def check_data(self, raise_error=False):
        """Run some checks to determine if the data is ok.

        Args:
            raise_error (bool): if true, raise an error if check fails, else return false

        Returns:
            bool: true if check passes, else false.

        Notes:
            * Do the self.SLOW_TREES exist and have entries?
            * Are there nonzero counts in UCNHits?

        Example:
            ```python
            >>> run.check_data()
            True
            ```
        """

        # check some necessary data trees
        for tree in self.SLOW_TREES:

            msg = None

            # does tree exist?
            if tree not in self.tfile.keys():
                msg = f'Missing ttree "{tree}" in run {self.run_number}'

            # does tree have entries?
            elif len(self.tfile[tree]) == 0:
                msg = f'Zero entries found in "{tree}" ttree in run {self.run_number}'

            # raise error or return
            if msg is not None:
                if raise_error:
                    raise MissingDataError(msg)
                else:
                    print(msg)
                    return False

        for name, det in self.DET_NAMES.items():

            # check for nonzero counts
            if not self.tfile[det['hits']].tIsUCN.any():
                msg = f'No UCN hits in "{name}" ttree in run {self.run_number}'

            # check if sequencer was enabled but no run transitions
            elif any(self.tfile.SequencerTree.sequencerEnabled):

                if len(self.tfile[det['transitions']]) == 0:
                    msg = 'No cycles found in run {self.run_number}, although sequencer was active'

            # raise error or return
            if msg is not None:
                if raise_error:
                    raise MissingDataError(msg)
                else:
                    tqdm.write(msg)
                    return False

        return True

    def draw_cycle_times(self, ax=None, xmode='datetime', do_legend=False):
        """Draw cycle start times as thick black lines, period end times as dashed lines

        Args:
            ax (plt.Axes): axis to draw in, if None, draw in current axes
            xmode (str): datetime|duration|epoch
            do_legend (bool): if true draw legend colors for period numbers

        Notes:
            Assumed periods:    0 - irradiation
                                1 - storage
                                2 - count
        """

        # check input
        if all((xmode not in i for i in ('datetime', 'duration', 'epoch'))):
            raise RuntimeError('xmode must be one of datetime|duration|epoch')

        # get axis to draw in
        if ax is None:
            ax = plt.gca()

        # run start time
        if xmode in 'duration':
            run_start = self.cycle_param.cycle_times.loc[0, 'start']
        else:
            run_start = 0

        # draw lines
        non_zero_periods = []
        for cyc in self.get_cycle():

            # get x value
            if xmode in 'datetime':
                start = pd.to_datetime(cyc.cycle_start, unit='s')
            else:
                start = cyc.cycle_start - run_start

            # draw
            ax.axvline(start, color='k', ls='-', lw=2)

            for i, per in enumerate(cyc):

                # skip zero length periods
                if per.period_start == per.period_stop:
                    continue

                # get x value
                if xmode in 'datetime':
                    x = pd.to_datetime(per.period_stop, unit='s')
                else:
                    x = per.period_stop - run_start

                # draw
                ax.axvline(x, color=f'C{i}', ls=':', lw=1)
                non_zero_periods.append(i)

            # get cycle text - strikeout if not good
            text = f'Cycle {cyc.cycle}'

            # check if filtered
            if cyc.cycle_param.filter is None or cyc.cycle_param.filter:
                color = 'k'
            else:
                text = '\u0336'.join(text) + '\u0336'
                color = 'r'

            text += ' $\\downarrow$'

            ypos = ax.get_ylim()[1]
            ax.text(start, ypos, text,
                    va='top',
                    ha='left',
                    fontsize='xx-small',
                    color=color,
                    rotation='vertical',
                    clip_on=True,)

        # add periods to legend
        if do_legend:
            handles = [mpatches.Patch(color=f'C{period}', label=f'Period {period}') \
                    for period in np.unique(non_zero_periods)]
            ax.legend(handles=handles)

    def gen_cycle_filter(self, quiet=False):
        """Generate filter array for cycles. Use with self.set_cycle_filter to filter cycles.

        Args:
            quiet (bool): if true don't print or raise exception

        Returns:
            np.array(bool): true if keep cycle, false if discard

        Notes:
            calls `ucncycle.check_data` on each cycle

        Example:
            ```python
            >>> run = ucnrun(2575)
            >>> run.gen_cycle_filter()
            Run 2575, cycle 0: 1A current dropped below 0.1 uA
            Run 2575, cycle 1: 1A current dropped below 0.1 uA

            array([False, False,  True,  True,  True,  True,  True,  True,  True,
                    True,  True,  True])
            ```
        """

        cycles = self.get_cycle()
        iterator = tqdm(cycles, desc=f'Run {self.run_number}: Scanning cycles',
                                leave=False,
                                total=self.cycle_param.ncycles)
        cfilter = [c.check_data(quiet=quiet, raise_error=False) for c in iterator]
        return np.array(cfilter)

    def get_cycle(self, cycle=None):
        """Return a copy of this object, but trees are trimmed to only one cycle.

        Note that this process converts all objects to dataframes

        Args:
            cycle (int): cycle number, if None, get all cycles

        Returns:
            ucncycle:
                if cycle > 0:  ucncycle object
                if cycle < 0 | None: a list ucncycle objects for all cycles

        Examples:
            ```python
            # get single cycle
            >>> run.get_cycle(0)
            run 1846 (cycle 0):
                comment            cycle_start        month              shifters           supercycle
                cycle              cycle_stop         run_number         start_time         tfile
                cycle_param        experiment_number  run_title          stop_time          year

            # get all cycles
            >>> len(run.get_cycle())
            17
            ```
        """

        if cycle is None or cycle < 0:
            ncycles = len(self.cycle_param.cycle_times.index)
            return applylist(map(self.get_cycle, range(ncycles)))
        elif cycle in self._cycledict.keys():
            return self._cycledict[cycle]
        else:
            self._cycledict[cycle] = ucncycle(self, cycle)
            return self._cycledict[cycle]

    def inspect(self, detector='Li6', bin_ms=100, xmode='duration', slow=None):
        """Draw counts and BL1A current with indicated periods to determine data quality

        Args:
            detector (str): detector from which to get the counts from. Li6|He3
            bin_ms (int): histogram bin size in ms
            xmode (str): datetime|duration|epoch
            slow (list|str): name of slow control tree to add in a separate axis, can be a list of names
        """

        # check input
        if all((xmode not in i for i in ('datetime', 'duration', 'epoch'))):
            raise RuntimeError('xmode must be one of datetime|duration|epoch')

        # number of rows in figure
        nrows = 2

        # check input
        if slow is not None:
            if isinstance(slow, str):
                if slow not in self.epics.columns:
                    raise KeyError(f'Input slow ("{slow}") not found in tree "epics"')
                slow = [slow]
            elif isinstance(slow, Iterable):
                for s in slow:
                    if s not in self.epics.columns:
                        raise KeyError(f'Input slow ("{slow}") not found in tree "epics"')
            else:
                raise RuntimeError('Input "slow" must be a string or iterable')

            # extra row for slow control
            nrows += len(slow)

        # make figure
        _, axes = plt.subplots(nrows=nrows, ncols=1, sharex=True,
                        layout='constrained', figsize=(8,10))

        # get current and histogram
        current = self.beam1a_current_uA
        current.sort_index(inplace=True)
        hist = self.get_hits_histogram(detector, bin_ms=bin_ms).to_dataframe()
        hist.set_index('tUnixTimePrecise', inplace=True)
        hist = hist['Count']

        # get slow control
        if slow is not None:
            slow = {s:self.epics[s].to_dataframe() for s in slow}

        # run start time
        run_start = self.cycle_param.cycle_times.loc[0, 'start']

        for cyc in self.get_cycle():
            for i, per in enumerate(cyc):

                # draw current
                cur = current.loc[per.period_start:per.period_stop]

                if len(cur) > 0:
                    if xmode in 'datetime':
                        cur.index = pd.to_datetime(cur.index, unit='s')
                    elif xmode in 'duration':
                        cur.index -= run_start

                    cur.plot(ax=axes[0], color=f'C{i}')

                # draw histogram
                hi = hist.loc[per.period_start:per.period_stop]

                if len(hi) > 0:
                    if xmode in 'datetime':
                        hi.index = pd.to_datetime(hi.index, unit='s')
                    elif xmode in 'duration':
                        hi.index -= run_start

                    hi.plot(ax=axes[1], color=f'C{i}')

                # draw slow control
                if slow is not None:
                    for j, (key, val) in enumerate(slow.items()):
                        v = val.loc[per.period_start:per.period_stop]
                        if len(v) > 0:
                            if xmode in 'datetime':
                                v.index = pd.to_datetime(v.index, unit='s')
                            elif xmode in 'duration':
                                v.index -= run_start
                        v.plot(ax=axes[j+2], color=f'C{i}')

            # draw the rest of the run - current
            cur = current.loc[per.period_stop:cyc.cycle_stop]

            if len(cur) > 0:
                if xmode in 'datetime':
                    cur.index = pd.to_datetime(cur.index, unit='s')
                elif xmode in 'duration':
                    cur.index -= run_start

                cur.plot(ax=axes[0], color=f'k')

            # draw the rest of the run - histogram
            hi = hist.loc[per.period_stop:cyc.cycle_stop]

            if len(hi) > 0:
                if xmode in 'datetime':
                    hi.index = pd.to_datetime(hi.index, unit='s')
                elif xmode in 'duration':
                    hi.index -= run_start

                hi.plot(ax=axes[1], color=f'k')

            # draw slow control
            if slow is not None:
                for i, (key, val) in enumerate(slow.items()):
                    v = val.loc[per.period_stop:cyc.cycle_stop]
                    if len(v) > 0:
                        if xmode in 'datetime':
                            v.index = pd.to_datetime(v.index, unit='s')
                        elif xmode in 'duration':
                            v.index -= run_start
                    v.plot(ax=axes[i+2], color=f'k')

        # draw vertical markers
        for i, ax in enumerate(axes):
            self.draw_cycle_times(ax=ax, xmode=xmode, do_legend=(i==0))

        # plot elements
        axes[0].set_title(self.run_number,fontsize='x-small')
        axes[0].set_ylabel('BL1A Current (uA)')
        axes[1].set_ylabel(f'UCN Counts/{bin_ms/1000:g}s')

        if slow is not None:
            slow_keys = tuple(slow.keys())
            for i, ax in enumerate(axes[2:]):
                ax.set_ylabel(slow_keys[i].split('_')[-2])

        if xmode in 'datetime':
            axes[-1].set_xlabel('')
        elif xmode in 'duration':
            axes[-1].set_xlabel('Time Since Run Start (s)')
        else:
            axes[-1].set_xlabel('Epoch Time')
        axes[1].set_yscale('log')

        return axes

    def keyfilter(self, name):
        """Don't load all the data in each file, only that which is needed"""

        name = name.replace(' ', '_').lower()

        # reject some keys based on partial matches
        reject_keys = ('v1725', 'v1720', 'v792', 'tv1725', 'charge', 'edge_diff',
                       'pulse_widths', 'iv2', 'iv3', 'rate')
        for key in reject_keys:
            if key in name:
                return False

        return True

    def set_cycle_filter(self, cfilter=None):
        """Set filter for which cycles to fetch when slicing or iterating

        Args:
            cfilter (None|iterable): list of bool, True if keep cycle, False if reject.
                if None then same as if all True

        Returns:
            None: sets self.cycle_param.filter

        Notes:
            Filter is ONLY applied when fetching cycles as a slice or as an iterator. ucnrun.get_cycle() always returns unfiltered cycles.

            Examples where the filter is applied:
                * run[:]
                * run[3:10]
                * run[:3]
                * for c in run: print(c)

            Examples where the filter is not applied:
                * run[2]
                * run.get_cycle()
                * run.get_cycle(2)

        Example:

            ```python
            # check how many cycles are fetched without filter
            >>> len(run[:])
            17

            # apply a filter
            >>> filter = np.full(17, True)
            >>> filter[2] = False
            >>> run.set_cycle_filter(filter)

            # check that cycle 2 is filtered out
            >>> len(run[:])
            16
            >>> for c in run:
                    print(c.cycle)
            0
            1
            3
            4
            5
            6
            7
            8
            9
            10
            11
            12
            13
            14
            15
            16
            ```
        """

        # check input
        cfilter = np.array(cfilter).astype(bool)

        if len(cfilter) != self.cycle_param.ncycles:
            raise RuntimeError(f'Run {self.run_number}: Length of cycle filter ({len(cfilter)}) does not match expected number of cycles ({self.cycle_param.ncycles})')

        # set
        self.cycle_param.filter = cfilter

        # set for already fetched cycles
        for key, cyc in self._cycledict.items():
            cyc.cycle_param.filter = self.cycle_param.filter[key]

    def set_cycle_times(self, mode):
        """Get start and end times of each cycle from the sequencer and save
        into self.cycle_param.cycle_times

        Run this if you want to change how cycle start times are calculated

        Args:
            mode (str): default|matched|sequencer|he3|li6
                if matched: look for identical timestamps in RunTransitions from detectors
                if sequencer: look for inCycle timestamps in SequencerTree
                if he3: use He3 detector cycle start times
                if li6: use Li6 detector cycle start times

        Returns:
            pd.DataFrame: with columns "start", "stop", "offset" and "duration (s)". Values are in epoch time. Indexed by cycle id. Offset is the difference in detector start times: he3_start-li6_start


        Notes:
            - If run ends before sequencer stop is called, a stop is set to final timestamp.
            - If the sequencer is disabled mid-run, a stop is set when disable ocurrs.
            - If sequencer is not enabled, then make the entire run one cycle
            - For matched mode,
                - set run stops as start of next transition
                - set offset as start_He3 - start_Li6
                - set start/stop/duration based on start_He3
            - If the object reflects a single cycle, return from cycle_start, cycle_stop

        Example:
            ```python
            # this calculates new cycle start and end times based on the selected method
            >>> run.set_cycle_times('li6')
            ```
        """

        # check if single cycle
        if hasattr(self, 'cycle'):
            return pd.DataFrame({'start':[self.cycle_start],
                                 'stop':[self.cycle_stop],
                                 'duration (s)': [self.cycle_stop-self.cycle_start],
                                 'offset (s)': [0.0],
                                 'supercycle': [self.supercycle]},
                                 index=[self.cycle])

        # get mode
        mode = mode.lower()

        # get run end time from control trees - used in matched and detector cycles times
        run_stop = -np.inf
        for treename in self.SLOW_TREES:
            try:
                max_time = self.tfile[treename].index.max()

            # tree not found, skip
            except KeyError:
                continue

            run_stop = max((max_time, run_stop))

        # bad end time
        if np.isinf(run_stop):
            raise MissingDataError(f"Run {self.run_number}: Missing slow control trees, cannot find run end time.")

        # get squencer data
        try:
            df = self.tfile.SequencerTree
        except AttributeError:
            df = None
        else:
            if type(df) == ttree:
                df = df.to_dataframe()

        ## if no sequencer, make the whole run a single cycle
        if df is None or not any(df.sequencerEnabled):

            times = {'start': np.inf,
                     'stop':  -np.inf,
                     'supercycle': [0]}

            # use timestamps from slow control trees to determine timestamps
            for treename in self.SLOW_TREES:

                # check for tree
                if treename not in self.tfile.keys():
                    continue

                # get times
                if isinstance(self.tfile[treename], pd.DataFrame):
                    idx = self.tfile[treename].index
                else:
                    idx = self.tfile[treename].to_dataframe().index

                # find min and max
                times['start'] = [min((idx.min(), times['start']))]
                times['stop']  = [max((idx.max(), times['stop']))]
                break

        ## get matched timesteps from He3 and Li6 RunTransitions
        elif mode in 'matched':
            hestart = self.tfile[self.DET_NAMES['He3']['transitions']].cycleStartTime
            listart = self.tfile[self.DET_NAMES['Li6']['transitions']].cycleStartTime
            scycle = self.tfile[self.DET_NAMES['He3']['transitions']].superCycleIndex

            # drop duplicate timestamps
            hestart = hestart.drop_duplicates()
            listart = listart.drop_duplicates()

            # get all possible pairs and sort by time difference
            pairs = sorted(itertools.product(hestart, listart),
                            key = lambda t: abs(t[0] - t[1]))

            # save output
            matchedhe3 = []
            matchedli6 = []

            # go through all possible pairs
            for pair in pairs:

                # if none of the two start times are already in the matched list, add the pair to the matched list
                if pair[0] not in matchedhe3 and pair[1] not in matchedli6:
                    offset = pair[0] - pair[1]

                    # discard if time difference is too large
                    if abs(offset) > 20:
                        raise CycleError(f'He3 cycle start time ({pair[0]}) too distant from Li6 start ({pair[1]}) in run {self.run_number}')
                    else:
                        matchedhe3.append(pair[0])
                        matchedli6.append(pair[1])

            matchedhe3 = np.sort(matchedhe3)
            matchedli6 = np.sort(matchedli6)

            # warnings for unmatched cycles
            unmatched = [t not in matchedhe3 for t in hestart.values]
            if any(unmatched):
                raise CycleError(f'Found no match to He3 cycles at {hestart.values[unmatched]} in run {self.run_number}')

            unmatched = [t not in matchedhe3 for t in listart.values]
            if any(unmatched):
                raise CycleError(f'Found no match to Li6 cycles at {listart.values[unmatched]} in run {self.run_number}')

            # setup output
            times = {'start': matchedhe3,
                     'duration (s)': np.diff(np.concatenate((matchedhe3, [run_stop]))),
                     'offset (s)': matchedhe3-matchedli6}
            times['stop'] = times['start'] + times['duration (s)']
            times['supercycle'] = scycle

        ## get timestamps from sequencer
        elif mode in 'sequencer':

            # if sequencer is not enabled cause a stop transition
            df.inCycle *= df.sequencerEnabled

            # start counting only after first start flag
            df = df.loc[df.loc[df.cycleStarted > 0].index[0]:]

            # get start and end times
            df = df.diff()
            times = {'start': df.index[df.inCycle == 1],
                    'stop': df.index[df.inCycle == -1],
                    'supercycle': 0}

            # check lengths
            if len(times['start']) > len(times['stop']):
                times['stop'].append(df.index[-1])

        ## detector start times
        elif mode in 'he3':

            start = self.tfile[self.DET_NAMES['He3']['transitions']].cycleStartTime
            start = start.drop_duplicates()

            # setup output
            times = {'start': start,
                     'duration (s)': np.diff(np.concatenate((start, [run_stop])))
                    }
            times['stop'] = times['start'] + times['duration (s)']
            times['supercycle'] = self.tfile[self.DET_NAMES['He3']['transitions']].superCycleIndex

        ## detector start times
        elif mode in 'li6':

            start = self.tfile[self.DET_NAMES['Li6']['transitions']].cycleStartTime
            start = start.drop_duplicates()

            # setup output
            times = {'start': start,
                     'duration (s)': np.diff(np.concatenate((start, [run_stop]))),
                    }
            times['stop'] = times['start'] + times['duration (s)']
            times['supercycle'] = 0

        ## bad mode
        else:
            raise RuntimeError(f"Bad mode input {mode}")

        # convert to dataframe
        times = pd.DataFrame(times)
        times['duration (s)'] = times.stop - times.start
        times.index.name = 'cycle'

        # save
        self.cycle_param['cycle_times'] = times
        self.cycle_param['ncycles'] = len(times.index)

        # finish setting up cycle_param
        if 'period_end_times' not in self.cycle_param.keys():
            self.cycle_param['period_end_times'] = pd.DataFrame({0:times.loc[0, 'stop']}, index=[0])
            self.cycle_param['period_end_times'].index.name = 'period'
            self.cycle_param['period_end_times'].columns.name = 'cycle'

        if 'period_durations_s' not in self.cycle_param.keys():
            self.cycle_param['period_durations_s'] = pd.DataFrame({0:times.loc[0, 'duration (s)']}, index=[0])
            self.cycle_param['period_durations_s'].index.name = 'period'
            self.cycle_param['period_durations_s'].columns.name = 'cycle'

        return times
