import os
import time
import logging
from math import log10, floor, ceil
from numpy import arange, histogram, cumsum, percentile, histogram2d, where
from scipy.stats import percentileofscore
from obspy import UTCDateTime

logger = logging.getLogger(__name__)


def calcEOBS(conf, times):
    '''
    Compute elapsed observation times and statistics
    Args:
        calcconf: configuration dictionary for calculation
        times: list of event times
    Return:
        eobs: dictionary with elapsed observation time data
    '''
    calcconf = conf['calcconf']
    try:
        evconf = conf['evconf']
    except KeyError:
        pass
    plotconf = conf['plotconf']

    nEvAll = [x for x in conf['calcconf']['nevents']]
    if 'sequences' in conf['calcconf'] and conf['calcconf']['sequences']['do']:
        for nEv in conf['calcconf']['sequences']['nevents']:
            if nEv not in nEvAll:
                nEvAll.append(nEv)

    eobs = {'times': times}
    step = calcconf['eobsstep']
    dstep = calcconf['dstep']
    for nEv in nEvAll:
        eobs[nEv] = {}
        all_data = [x-y for x, y in zip(times[nEv-1:], times)]
        print(f'nEv: {nEv}, # all_data: {len(all_data)}')
        eobs[nEv]['eobs'] = all_data
        tmp = sorted(all_data)
        eobs[nEv]['median'] = tmp[len(tmp)//2]
        eobs[nEv]['mean'] = sum(tmp)/len(tmp)
        eobs[nEv]['pcs'] = percentile(
            all_data, calcconf['sequences']['sigthresh'])
        if min(all_data) == 0:
            dmin = 0
        else:
            dmin = floor(log10(min(all_data)))
        dmax = ceil(log10(max(all_data)))
        eobs[nEv]['bins'] = [pow(10, x)
                             for x in arange(dmin, dmax+(step/2), step)]
        count, bins_count = histogram(all_data, bins=eobs[nEv]['bins'])
        eobs[nEv]['pdf'] = count / sum(count)
        eobs[nEv]['cdf'] = cumsum(eobs[nEv]['pdf'])

        # Compute for background periods
        eobs[nEv]['bg'] = []
        for tp in calcconf['backgroundtime']:
            bgdata = {}
            bgdata['mintime'] = tp['mintime']
            bgdata['maxtime'] = tp['maxtime']
            tmp = []
            for x, t in zip(eobs[nEv]['eobs'], times[nEv-1:]):
                if t < tp['mintime'] or t > tp['maxtime']:
                    continue
                # Exclude any eruption periods from background statistics
                bEruption = False
                try:
                    for exc in [e for e in evconf['timeperiods'] if e['label'] == 'eruption']:
                        if t >= exc['mintime'] and t <= exc['maxtime']:
                            bEruption = True
                            break
                except KeyError:
                    pass
                if bEruption:
                    continue
                tmp.append([x, t])
            bgdata['eobs'], bgdata['times'] = zip(*tmp[nEv-1:])
            bg_data = sorted(bgdata['eobs'])
            print(f'nEv: {nEv}, # bg_data: {len(bg_data)}')
            bgdata['median'] = bg_data[len(bg_data)//2]
            bgdata['pcs'] = percentile(
                bgdata['eobs'], calcconf['sequences']['sigthresh'])
            count, bins_count = histogram(bg_data, bins=eobs[nEv]['bins'])
            bgdata['pdf'] = count / sum(count)
            bgdata['cdf'] = cumsum(bgdata['pdf'])
            bgdata['eobs_dist'] = [percentileofscore(
                bg_data, x) for x in all_data]
            bgdata['eobs_dist_diff'] = [y-x for x,
                                        y in zip(bgdata['eobs_dist'][1:], bgdata['eobs_dist'][:-1])]
            bgdata['eobs_dist_hist2d'] = histogram2d(bgdata['eobs_dist'][:-1],
                                                     bgdata['eobs_dist_diff'],
                                                     bins=(arange(0., 100+(dstep/2.), dstep),
                                                           arange(-100., 100.+(dstep/2.), dstep)))
            eobs[nEv]['bg'].append(bgdata)

    # Rate range
    with open(f'{plotconf["name"]}_rate_distribution.dat', 'w') as fout:
        bgn = 2
        if bgn in eobs:
            for nEv in calcconf['nevents']:
                if nEv < bgn:
                    continue
                for bind, bgdata in enumerate(eobs[nEv]['bg']):
                    stable_portion = 100. * \
                        len([x for x in bgdata['eobs_dist_diff'] if x >= -
                            10. and x <= 10.]) / len(bgdata['eobs_dist_diff'])
                    eobst_min = min(bgdata['eobs'])/(nEv-1)
                    raterangemin = percentileofscore(
                        eobs[bgn]['bg'][bind]['eobs'], eobst_min)
                    eobst_max = max(bgdata['eobs'])/(nEv-1)
                    raterangemax = percentileofscore(
                        eobs[bgn]['bg'][bind]['eobs'], eobst_max)
                    raterange = raterangemax - raterangemin
                    fout.write(
                        f'{plotconf["name"]}, Background: {bind}, nEv: {nEv}, stable portion: {stable_portion:.2f}, rate range min: {raterangemin:.2f}, rate range max: {raterangemax:.2f}, rate range: {raterange:.2f}\n')

    if calcconf['sequences']['do']:
        computeSequences(eobs, calcconf, plotconf)

    return eobs


def computeSequences(eobs, calcconf, plotconf):
    '''
    Identify sequences of interest based on elapsed time characteristics.
    The sequence is identified to start based on seeing N events occur within a configured time
    period T that all have elapsed time shorter than X% of the background distribution. The sequence
    is identified to end once the inter-event elapsed time exceeds Y% of the background
    distribution. N, T, X and Y are all specified in the configuration file.
    The nEv for this analysis is specified in the configuration file, but all background
    times will be used.
    Args:
        eobs: dictionary of elapsed observation time data
        calcconf: configuration parameters
    Return:
    '''

    # Detect sequences
    for nEv in calcconf['sequences']['nevents']:
        for bind, bgdata in enumerate(eobs[nEv]['bg']):
            bgdata['sequences'] = {}
            sig = eobs[nEv]['eobs'] < bgdata['pcs'][0]
            maxlimit = eobs[nEv]['eobs'] < bgdata['pcs'][1]
            for min_eq in [x-1 for x in calcconf['sequences']['min_sig_ev_seq_start']]:
                bgdata['sequences'][min_eq+1] = []
                sig_inds = where(sig)[0]
                ind = sig_inds[0]
                while ind < len(sig):
                    # Start sequence - N sig events must exist within time separation
                    if min_eq > 0:
                        next_sig_inds = sig_inds[sig_inds > ind]
                        if len(next_sig_inds) < min_eq:
                            ind = len(sig)
                            continue
                        tsep = eobs['times'][nEv-1:][next_sig_inds[min_eq-1]
                                                     ] - eobs['times'][nEv-1:][ind]
                        if tsep > calcconf['sequences']['max_t_seq_start']:
                            if sum(sig_inds > ind) > 0:
                                ind = next_sig_inds[0]
                            else:
                                ind = len(sig)
                            continue
                    endlimit = where(maxlimit[ind+1:] == False)[0]
                    if len(endlimit) == 0:
                        seq_end = ind
                    else:
                        seq_end = ind + endlimit[0]
                    bgdata['sequences'][min_eq+1].append({'starttime': eobs['times'][nEv-1:][ind],
                                                          'endtime': eobs['times'][nEv-1:][seq_end-nEv+2] + bgdata['pcs'][1],
                                                          'duration': eobs['times'][nEv-1:][seq_end-nEv+2] - eobs['times'][nEv-1:][ind] + bgdata['pcs'][1],
                                                          'events_in_seq': (seq_end - ind) + 1,
                                                          'sig_events_in_seq': sum(sig[ind:seq_end]) + 1})
                    if sum(sig_inds > seq_end) > 0:
                        ind = sig_inds[sig_inds > seq_end][0]
                    else:
                        ind = len(sig)

    # Write a summary file of the sequences
    with open(f'{plotconf["name"]}_sequences.dat', 'w') as fout:
        fout.write(
            f'Time separation for significant events {calcconf["sequences"]["max_t_seq_start"]}\n')
        for nEv in calcconf['sequences']['nevents']:
            fout.write(f'nEv: {nEv}\n')
            for bgdata in eobs[nEv]['bg']:
                fout.write(
                    f'Background times: {bgdata["mintime"]} - {bgdata["maxtime"]}\n')
                fout.write(
                    f'Background significance thresholds: {bgdata["pcs"][0]}, max threshold: {bgdata["pcs"][1]}\n')
                for min_eq in bgdata['sequences']:
                    fout.write(f'Min sig events in sequence: {min_eq}\n')
                    for seq in bgdata['sequences'][min_eq]:
                        fout.write(
                            f"starttime: {seq['starttime'].strftime('%Y-%m-%dT%H:%M:%S.%fZ')}, ")
                        fout.write(
                            f"endtime: {seq['endtime'].strftime('%Y-%m-%dT%H:%M:%S.%fZ')}, ")
                        fout.write(f"duration: {seq['duration']:.3f}, ")
                        fout.write(f"# events: {seq['events_in_seq']:d}, ")
                        fout.write(
                            f"# sig events: {seq['sig_events_in_seq']:d}\n")

    return


if __name__ == '__main__':
    '''
    Test the calcEOBS script
    '''
    import events as evs

    logging.basicConfig(filename=f'{os.path.basename(__file__).split(".")[0]}_{time.strftime("%y%m%dT%H%M%S", time.gmtime(time.time()))}.log',
                        level=logging.DEBUG)

    conf = {
        'evconf': {
            # 'catalogue': { 'client': 'GEONET' },
            'catalogue': {'file': 'redoubt_testdata.csv'},
            'ROI': {
                'center': (60.4852, -152.7438),
                'radius': 20.,
            },
            'minmag': 0.,
            'maxdepth': 50.,
            'mintime': UTCDateTime().strptime('1989-10-31T00:00:00', '%Y-%m-%dT%H:%M:%S'),
            'maxtime': UTCDateTime().strptime('2023-12-31T00:00:00', '%Y-%m-%dT%H:%M:%S'),
            'timeperiods': [{
                'mintime': UTCDateTime().strptime('2003-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'maxtime': UTCDateTime().strptime('2003-07-31T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'label': 'outage'
            },
                {
                'mintime': UTCDateTime().strptime('2004-04-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'maxtime': UTCDateTime().strptime('2004-08-31T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'label': 'outage'
            },
                {
                'mintime': UTCDateTime().strptime('2005-07-31T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'maxtime': UTCDateTime().strptime('2006-07-31T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'label': 'outage'
            },
                {
                'mintime': UTCDateTime().strptime('2008-06-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'maxtime': UTCDateTime().strptime('2008-09-30T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'label': 'outage'
            },
                {
                'mintime': UTCDateTime().strptime('2021-05-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'maxtime': UTCDateTime().strptime('2021-05-03T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'label': 'eruption'
            }],
            # 'savecatcsv': 'redoubt_testdata.csv'
            'savecatcsv': None
        },
        'calcconf': {
            'backgroundtime': [{
                'mintime': UTCDateTime().strptime('2003-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'maxtime': UTCDateTime().strptime('2016-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
            },
                {
                'mintime': UTCDateTime().strptime('2009-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'maxtime': UTCDateTime().strptime('2010-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
            },
                {
                'mintime': UTCDateTime().strptime('2010-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'maxtime': UTCDateTime().strptime('2016-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
            },
                {
                'mintime': UTCDateTime().strptime('2020-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
                'maxtime': UTCDateTime().strptime('2021-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
            },
            ],
            'nevents': [1, 2, 5, 10],
            'histstep': 0.1,
            'dstep': 5,
            'sigthresh': [5, 10, 50, 90, 95],
        },
    }

    cat = evs.getEvents(conf)
    times = evs.cat2times(cat)

    tstart = UTCDateTime().now()
    calcEOBS(conf, times)
    tend = UTCDateTime().now()
    logger.info(f'Time taken to compute EOBS: {tend-tstart}')
