import time
from datetime import timezone
import logging
import os

import pandas as pd
from obspy.clients.fdsn import Client
from obspy.geodetics.base import gps2dist_azimuth, degrees2kilometers, kilometers2degrees
from obspy.core.event import Origin, Magnitude, Event
from obspy import UTCDateTime, Catalog, read_events

logger = logging.getLogger(__name__)


def writeCatFile(conf, cat):
    '''
    Write a flat catalogue file as csv (quick utility)
    Args:
        conf: configuraiton dictionary
        cat: obspy Catalog
    '''
    evconf = conf['evconf']
    with open(evconf['savecatcsv'], 'w') as fout:
        for ev in cat:
            orig = ev.preferred_origin()
            mag = ev.preferred_magnitude()
            fout.write(
                f'{orig.time}, {orig.latitude}, {orig.longitude}, {orig.depth/1000.}, {mag.mag}\n')
    return


def to_dataframe(cat):
    '''
    Convert an obspy Catalog to a pandas DataFrame
    Args:
        cat: obspy Catalog
    Return:
        df: pandas DataFrame
    '''
    data = []
    for ev in cat:
        orig = ev.preferred_origin()
        mag = ev.preferred_magnitude()
        data.append({
            'time': orig.time.datetime,
            'latitude': orig.latitude,
            'longitude': orig.longitude,
            'depth': orig.depth / 1000.,
            'magnitude': mag.mag
        })
    df = pd.DataFrame(data)
    df.set_index('time', inplace=True)
    df.sort_index(inplace=True)
    try:
        df.index = df.index.tz_localize(timezone.utc)
    except TypeError:
        pass
    return df


def rdCatFile(conf):
    '''
    Read a flat catalogue file
    Args:
        conf: configuration dictionary
    Return:
        cat: obspy Catalog
    '''
    evconf = conf['evconf']
    fname = evconf['catalogue']['file']
    if not os.path.isfile(fname):
        logger.error(f'Specified catalogue file {fname} not found')
        exit()
    cat_list = []
    with open(fname, 'r') as fin:
        for l in fin:
            if l.startswith('#'):
                continue
            fs = l.split(',')
            # check the criteria
            if float(fs[4]) < evconf['minmag']:
                continue
            if float(fs[3]) > evconf['maxdepth']:
                continue
            # Time range to include
            time = UTCDateTime().strptime(fs[0], '%Y-%m-%dT%H:%M:%S.%fZ')
            if time < evconf['mintime'] or time > evconf['maxtime']:
                continue
            # Time periods to exclude
            bExclude = False
            for period in evconf['timeperiods']:
                if period['label'] == 'outage' and time > period['mintime'] and time < period['maxtime']:
                    bExclude = True
                    break
            if bExclude:
                continue
            # Region of interest
            if 'radius' in evconf['ROI']:
                # check for ROI circle
                if gps2dist_azimuth(evconf['ROI']['center'][0],
                                    evconf['ROI']['center'][1],
                                    float(fs[1]),
                                    float(fs[2]))[0]/1000. > evconf['ROI']['radius']:
                    continue
            else:
                # check for ROI rectangle
                if float(fs[1]) < evconf['ROI']['latrange'][0] or \
                        float(fs[1]) > evconf['ROI']['latrange'][1] or \
                        float(fs[2]) < evconf['ROI']['lonrange'][0] or \
                        float(fs[2]) > evconf['ROI']['lonrange'][1]:
                    continue
            orig = Origin(longitude=float(fs[2]),
                          latitude=float(fs[1]),
                          depth=float(fs[3]),
                          time=time)
            mag = Magnitude(mag=float(fs[4]))
            ev = Event(origins=[orig], magnitudes=[mag])
            ev.preferred_origin_id = orig.resource_id
            ev.preferred_magnitude_id = mag.resource_id
            cat_list.append(ev)
    cat_list_sorted = sorted(
        cat_list, key=lambda event: event.preferred_origin().time)
    cat = Catalog(cat_list_sorted)
    return cat


def excludeTimeperiods(evconf, cat_all):
    '''
    Apply the exclusion of time periods from the configuration
    Args:
        conf: configuration dictionary
        cat_all: input catalogue
    Return:
        cat: obspy Catalog
    '''
    cat_rem = Catalog()
    for period in evconf['timeperiods']:
        if period['label'] == 'outage':
            cat_rem += cat_all.filter(f"time > {period['mintime'].strftime('%Y-%m-%dT%H:%M:%S')}",
                                      f"time < {period['maxtime'].strftime('%Y-%m-%dT%H:%M:%S')}")
    cat = Catalog()
    for ev in cat_all:
        if ev not in cat_rem:
            cat += ev
    return cat


def rdObspyCatFile(conf):
    '''
    Read an Obspy catalog file
    Args:
        conf: configuration dictionary
    Return:
        cat: obspy Catalog
    '''
    evconf = conf['evconf']
    fmt = evconf['catalogue']['file'].split('.')[-1].upper()
    all_cat = read_events(evconf['catalogue']['file'], format=fmt)
    all_cat = all_cat.filter(f'magnitude >= {evconf["minmag"]}',
                             f'depth <= {evconf["maxdepth"]}',
                             f'time >= {evconf["mintime"].strftime("%Y-%m-%sT%H:%M:%S")}',
                             f'time <= {evconf["maxtime"].strftime("%Y-%m-%sT%H:%M:%S")}')
    cat = Catalog()
    for ev in cat:
        orig = ev.preferred_origin()
        if 'radius' in evconf['ROI']:
            # check for ROI circle
            if gps2dist_azimuth(evconf['ROI']['center'][0],
                                evconf['ROI']['center'][1],
                                origin.latitude,
                                origin.longitude) > evconf['ROI']['radius']:
                continue
        else:
            # check for ROI rectangle
            if orig.latitude < evconf['ROI']['latrange'][0] or \
                    orig.latitude > evconf['ROI']['latrange'][1] or \
                    orig.longitude < evconf['ROI']['lonrange'][0] or \
                    orig.longitude > evconf['ROI']['lonrange'][1]:
                continue
        cat += ev
    cat = excludeTimeperiods(evconf, cat)
    return cat


def getFDSN(conf):
    '''
    Retrieve a catalog using FDSN
    Args:
        conf: configuration dictionary
    Return:
        cat: obspy Catalog
    '''
    evconf = conf['evconf']
    try:
        client = Client(evconf['catalogue']['client'])
    except:
        logger.error(
            f'Cannot connect to FDSN server {evconf["catalogue"]["client"]} to retrieve events')
        exit()
    if 'radius' in evconf['ROI']:
        radius = kilometers2degrees(evconf['ROI']['radius'])
        cat_all = client.get_events(starttime=evconf['mintime'],
                                    endtime=evconf['maxtime'],
                                    minmagnitude=evconf['minmag'],
                                    maxdepth=evconf['maxdepth'],
                                    latitude=evconf['ROI']['center'][0],
                                    longitude=evconf['ROI']['center'][1],
                                    maxradius=radius)
    else:
        cat_all = client.get_events(starttime=evconf['mintime'],
                                    endtime=evconf['maxtime'],
                                    minmagnitude=evconf['minmag'],
                                    maxdepth=evconf['maxdepth'],
                                    minlatitude=evconf['ROI']['latrange'][0],
                                    maxlatitude=evconf['ROI']['latrange'][1],
                                    minlongitude=evconf['ROI']['lonrange'][0],
                                    maxlongitude=evconf['ROI']['lonrange'][1])
    cat = excludeTimeperiods(evconf, cat_all)
    return cat


def getEvents(conf):
    '''
    Wrapper to get events based on configuration options
    Args:
        conf: configuration dictionary
    Return:
        cat: obspy Catalog
    '''
    evconf = conf['evconf']
    if 'file' in evconf['catalogue']:
        if evconf['catalogue']['file'].find('csv') > -1:
            cat = rdCatFile(conf)
        else:
            cat = rdObspyCatFile(conf)
    else:
        cat = getFDSN(conf)
    logger.info(f'{cat.count()} event(s) in the catalogue')
    if evconf['savecatcsv'] is not None:
        writeCatFile(conf, cat)
    return cat


def cat2times(cat):
    '''
    Extract the origin times from an obspy Catalog
    Assume the Catalog is sorted by ascending origin time
    '''
    times = [ev.preferred_origin().time for ev in cat]
    return times


if __name__ == "__main__":
    '''
    Test the events.py script
    '''

    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': 'nz_testcat.csv'},
            'ROI': {
                'center': (-41.289, 174.777),
                'radius': 200.,
        },
        'minmag': 3.,
        'maxdepth': 20.,
        'mintime': UTCDateTime().strptime('2020-01-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
        'maxtime': UTCDateTime().strptime('2024-04-10T13:24:55', '%Y-%m-%dT%H:%M:%S'),
        'timeperiods': [{
            'mintime': UTCDateTime().strptime('2020-05-01T00:00:00', '%Y-%m-%dT%H:%M:%S'),
            'maxtime': UTCDateTime().strptime('2020-05-01T06: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': 'nz_testcat.csv'
        # 'savecatcsv': None
    },
    }

    tstart = UTCDateTime().now()
    cat = getEvents(conf)
    tend = UTCDateTime().now()
    logger.info(f'Time taken to get events: {tend-tstart}')
    times = cat2times(cat)
    logger.info(times)
