# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02_gti.ipynb (unless otherwise specified).

__all__ = ['get_gti']

# Cell
from astropy.io import fits
import numpy as np
import pandas as pd

from .config import MJD, Config, Cache


# Cell
def _get(config):
    """Combine the GTI intervals that fall within the gti range
    return an array of alternate start/stop times
    """

    gti_files = list(config.files.gti.glob('*.fits'))
    if config.verbose>1:
        print(f'Processing {len(gti_files)} FITS files with GTI information ... ', end='')
    starts=[]
    stops=[]
    for i, ft1 in enumerate(gti_files):
        with fits.open(ft1) as hdu:
            gti_data = hdu['GTI'].data
            start = gti_data.START
            if i>0:
                assert start[0]>= stops[-1][-1], f'file {ft1} has start time not following preceding file'
            starts.append(start)
            stops.append( gti_data.STOP)
    start = MJD(np.concatenate(starts))
    stop  = MJD(np.concatenate(stops))

    if config.verbose>1:
        livetime = (stop-start).sum()
        print( f' {len(gti_files)} files, {len(start)} intervals with'\
               f' {int(livetime):,} days live time')

    sel = slice(None)
    if config.mjd_range is not None:
        a,b =  np.searchsorted(start, config.mjd_range)
        if a>0 or b<len(start):
            if config.verbose>1:
                print(f'\tcut from {len(start):,} to {a} - {b}, or {b-a:,} entries after MJD range selection')
            sel = slice(a,b)


    a,b =start[sel], stop[sel]
    fraction = np.sum(b-a)/(b[-1]-a[0])
    assert len(a)==len(b)
    gtimes = np.array([a,b]).T.flatten()
    assert np.sum(np.diff(gtimes)<0)==0, 'Bad GTI ordering'

    return gtimes

# Cell
def get_gti(config, key='gti'):
    """
    Return a GTI functor that checks a list of times

    - key -- the cache key to use to save/restore interval times
    """

    class GTI():

        def __init__(self, config, key=key):
            assert config.cache is not None
            self.gtimes = config.cache(key, _get, config)

        def __call__(self, time):

            x = np.digitize(time, self.gtimes)
            return np.bitwise_and(x,1).astype(bool)
        def __repr__(self):
            g = self.gtimes
            return f'GTI: {len(g)/2} intervals from {g[0]:.1f} to {g[1]:.1f}'

    return GTI(config, key)
