import os
import logging
from numpy import meshgrid, transpose, loadtxt, hstack, ones
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

logger = logging.getLogger(__name__)

def addCmap():
    '''
    Create custom colormap for density plots
    '''
    cmap = loadtxt('xnvar.cmap') 
    return ListedColormap(hstack((cmap, ones((cmap.shape[0],1)))))

def addExclusionTimes(ax, conf):
    '''
    Plot exclusion times as filled coloured polygons
    Args:
        ax: Plot axis object
        conf: configuration dictionary
    '''
    evconf = conf['evconf']
    cols = ['orange', 'grey', 'green', 'red', 'yellow', 'blue']
    cats = list(sorted(set([t['label'] for t in evconf['timeperiods']])))
    ylim = ax.get_ylim()
    for tp in evconf['timeperiods']:
        xs = [tp['mintime'].datetime, 
              tp['mintime'].datetime, 
              tp['maxtime'].datetime, 
              tp['maxtime'].datetime]
        ys = [ylim[0], ylim[1], ylim[1], ylim[0]]
        ax.fill(xs, ys, c=cols[cats.index(tp['label'])], alpha=0.5)
    ax.set_ylim(ylim)
    return ax

def plotTimeEOBS(eobs, conf):
    '''
    Plot eobs against time
    Args:
        eobs: dictionary of elapsed times, keys are number of events for elapse
        conf: configuration dictionary
    '''
    evconf = conf['evconf']
    plotconf = conf['plotconf']
    cols = ['r', 'g', 'b', 'cyan', 'magenta', 'orange', 'yellow', 'grey']
    for nEv in [n for n in eobs if n != 'times']:
        fig, ax = plt.subplots(1,1)
        data = eobs[nEv]['eobs']
        plott = [x.datetime for x in eobs['times'][nEv-1:]]
        ax.scatter(plott, data, marker='x', c='k', s=10)
        ax.axhline(eobs[nEv]['median'], c='k')
        for tp,col in zip(eobs[nEv]['bg'], cols):
            xs,ys = zip(*[(x.datetime,y) for x,y in zip(tp['times'], tp['eobs'])])
            ax.scatter(xs, ys, marker='x', s=10, c=col)
            ax.axhline(tp['median'], c=col)
        ax.set_xlabel('Time')
        ax.set_ylabel('Elapsed observation time (s)')
        ax.set_yscale('log')
        ax.set_title(f'{plotconf["name"]}, N={nEv}')
        if 'mintime' in plotconf['TimeEOBS']:
            ax.set_xlim(xmin = plotconf['TimeEOBS']['mintime'].datetime)
        if 'maxtime' in plotconf['TimeEOBS']:
            ax.set_xlim(xmax = plotconf['TimeEOBS']['maxtime'].datetime)
        ax = addExclusionTimes(ax, conf)
        fig.autofmt_xdate()
        plt.savefig(f'{plotconf["name"]}_t-eobs_{nEv}.png')
    return

def plotEOBScdf(eobs, conf):
    '''
    Plot the CDF of the elapsed observation times
    Args:
        eobs: dictionary of elapsed times, keys are number of events for elapse
        conf: configuration dictionary
    '''
    plotconf = conf['plotconf']
    cols = ['r', 'g', 'b', 'cyan', 'magenta', 'orange', 'yellow', 'grey']
    for nEv in [n for n in eobs if n != 'times']:
        fig, ax = plt.subplots(1,1)
        ax.plot(eobs[nEv]['bins'][:-1], eobs[nEv]['cdf'], lw=1, label=f'all', c='k')
        for sig, pc in zip(conf['calcconf']['sequences']['sigthresh'], eobs[nEv]['pcs']):
            ax.axvline(pc, c='k', lw=0.5)
            ax.text(pc, (sig/100.)+0.05, f'{sig}%: {pc:.2e}', ha='center')
        for tp,col in zip(eobs[nEv]['bg'], cols):
            ax.plot(eobs[nEv]['bins'][:-1], tp['cdf'], lw=1,
                    label=f'{tp["mintime"].strftime("%Y-%m-%d")}-{tp["maxtime"].strftime("%Y-%m-%d")}', 
                    c=col)
            for sig, pc in zip(conf['calcconf']['sequences']['sigthresh'], tp['pcs']):
                ax.axvline(pc, c=col, lw=0.5)
                ax.text(pc, (sig/100.)-0.05, f'{sig}%: {pc:.2e}', ha='center', c=col)
        plt.legend()
        ax.set_xlabel('Elapsed observation time (s)')
        ax.set_ylabel('Empirical CDF')
        ax.set_ylim(0., 1.)
        ax.set_xscale('log')
        ax.grid(which='both', ls=':')
        ax.yaxis.tick_left()
        ax.set_title(f'{plotconf["name"]}, N={nEv}')
        plt.savefig(f'{plotconf["name"]}_cdf_{nEv}.png')
    return

def plotTimeXn(eobs, conf):
    '''
    Plot eobs against time for different Xn
    Args:
        eobs: dictionary of elapsed times, keys are number of events for elapse
        conf: configuration dictionary
    '''
    evconf = conf['evconf']
    plotconf = conf['plotconf']
    fig, ax = plt.subplots(1,1)
    for nEv in [n for n in eobs if n != 'times']:
        data = eobs[nEv]['eobs']
        xs = [x.datetime for x in eobs['times'][nEv-1:]]
        ax.plot(xs, data, label=f'{nEv}')
    if 'mintime' in plotconf['TimeXn']:
        ax.set_xlim(xmin = plotconf['TimeXn']['mintime'].datetime)
    if 'maxtime' in plotconf['TimeXn']:
        ax.set_xlim(xmax = plotconf['TimeXn']['maxtime'].datetime)
    ax.set_xlabel('Time')
    ax.set_ylabel('Elapsed observation time (s)')
    ax.set_yscale('log')
    plt.legend()
    ax = addExclusionTimes(ax, conf)
    fig.autofmt_xdate()
    ax.set_title(f'{plotconf["name"]}, N={nEv}')
    plt.savefig(f'{plotconf["name"]}_t-eobs-xn.png')
    return

def plotEOBSdistribution(eobs, conf):
    '''
    Plot the elapsed time distributions
    Args:
        eobs: dictionary of elapsed times, keys are number of events for elapse
        conf: configuration dictionary
    '''
    axn_cmap = addCmap()
    plotconf = conf['plotconf']
    for nEv in [n for n in eobs if n != 'times']:
        for i,bgdata in enumerate(eobs[nEv]['bg']):
            fig, ax = plt.subplots(1,1)
            hist, xedges, yedges = bgdata['eobs_dist_hist2d']
            X, Y = meshgrid(xedges, yedges)
            #cb = ax.pcolormesh(X, Y, transpose(hist), cmap='gist_ncar_r')
            cb = ax.pcolormesh(X, Y, transpose(hist), cmap=axn_cmap,
                               vmax = plotconf['EOBSdistribution']['cbarmax'])
            ax.set_xlabel('Distribution Position')
            ax.set_ylabel('Distribution Change Between Earthquakes')
            ax.set_title(f'{plotconf["name"]}, N={nEv}\nBackground: {bgdata["mintime"].strftime("%Y/%m/%dT%H:%M:%S")}-{bgdata["maxtime"].strftime("%Y/%m/%dT%H:%M:%S")}')
            plt.colorbar(cb, ax = ax)
            plt.savefig(f'{plotconf["name"]}_eobs-distribution_{nEv}_bg{i}.png')
    return

def plotMapEOBS(cat, conf):
    '''
    Plot a map of the events
    Args:
        cat: catalog of events
        conf: configuration dictionary
    '''
    from matplotlib.figure import Figure
    import cartopy.crs as ccrs
    from cartopy.io.img_tiles import GoogleTiles
    tiler = GoogleTiles(style='satellite')
    proj = tiler.crs

    plotconf = conf['plotconf']

    # Data parameters
    lats = [ev.preferred_origin().latitude for ev in cat]
    lons = [ev.preferred_origin().longitude for ev in cat]
    times = [ev.preferred_origin().time for ev in cat]

    # Figure
    fig = Figure()
    ax = fig.add_subplot(111, projection = proj)
    margin = 1.
    if 'margin' in plotconf['MapEOBS'] and plotconf['MapEOBS']['margin'] > 0:
        margin = plotconf['MapEOBS']['margin']
    bounds = [min(lons) - margin,
              max(lons) + margin,
              min(lats) - margin,
              max(lats) + margin]
    ax.set_extent(bounds, crs=ccrs.PlateCarree())
    ax.add_image(tiler, 8, alpha=0.7, zorder=0.)
    ax.scatter(lons, lats, c=times, transform=ccrs.PlateCarree(), cmap='jet')
    gl = ax.gridlines(draw_labels=True, linestyle='--')
    gl.top_labels = False
    gl.right_labels = False
    ax.set_title(f'{plotconf["name"]}')
    fig.savefig(f'{plotconf["name"]}_map_cat.png')
    return

def makePlots(eobs, cat, conf):
    '''
    Use plotconf to determine which plots to make
    Args:
        eobs: dictionary of elapsed times, keys are number of events for elapse
        conf: configuration dictionary
    '''
    plotconf = conf['plotconf']
    if 'TimeEOBS' in plotconf and plotconf['TimeEOBS']['do']:
        logger.info(f'Making TimeEOBS plot ...')
        plotTimeEOBS(eobs, conf)
    if 'EOBScdf' in plotconf and plotconf['EOBScdf']['do']:
        logger.info(f'Making EOBScdf plot ...')
        plotEOBScdf(eobs, conf)
    if 'TimeXn' in plotconf and plotconf['TimeXn']['do']:
        logger.info(f'Making TimeXn plot ...')
        plotTimeXn(eobs, conf)
    if 'EOBSdistribution' in plotconf and plotconf['EOBSdistribution']['do']:
        logger.info(f'Making EOBSdistribution plot ...')
        plotEOBSdistribution(eobs, conf)
    if 'MapEOBS' in plotconf and plotconf['MapEOBS']['do']:
        logger.info(f'Making MapEOBS plot ...')
        plotMapEOBS(cat, conf)
    return

if __name__ == '__main__':
    '''
    Test the plotting 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)

