import numpy as np
import matplotlib.pyplot as plt
from functions import *
import socket
import matplotlib
import gc
import sys
import glob
from pathlib import Path

plt.rcParams.update({'font.size': 22})
mylightblue = '#a6cee3'
mydarkblue = '#1f78b4'
mylightgreen = '#b2df8a'
mydarkgreen = '#33a02c'
myred = '#e31a1c'
myyellow = 'ffff33ff'
mypurple = '984ea3ff'
myorange = 'ff7f00ff'
mybrown = 'a65628ff'
mygrey = '999999ff'
mypink = '#fb9a99'

TUD_cyan = '#00a6d6ff'
TUD_lightblue = '#6ebbd5ff'
TUD_darkblue = '#1d1c73ff'
TUD_orange = '#e64616ff'
TUD_red = '#e21a1aff'
TUD_yellow = '#ffd100ff'
TUD_green = '#008891ff'
TUD_lightgreen = '#a5ca1aff'
TUD_purple = '#6d177fff'
TUD_grey = '#6b8689ff'
TUD_brown = '#a65628ff'
TUD_pink = '#fb9a99'


analydir = '/path/to/Analyses/directory'
wdir = '/path/to/a/working/directory/'
changepointloc = 'location/of/changepoint.exe'
matplotlib.use('Agg')
cores = 4
   
os.chdir(wdir)



if len(sys.argv) == 2:
    thisdir = sys.argv[1]
else:
    thisdir = analydir

if thisdir [-4:] == 'hdf5':
    filelist = [thisdir]
else:
    filelist = glob.glob(thisdir + '**/Photon_HDF5/*.hdf5', recursive=True)

logger('starting analysis of %s'%thisdir, 'c', init=True)

for fullpath in filelist:
    photonhdf5 = str(Path(fullpath).relative_to(analydir))
    results_hdf_loc = init_results_hdf5(analydir, photonhdf5, overwrite_last_run=False,
                                        skip_existing=False, copyfrom='latest_run')
    if results_hdf_loc == 'None':
        logger('Analysis hdf not generated for ' + analydir + photonhdf5, 'w')
        continue
    elif results_hdf_loc == 'Skip':
        logger('Analysis hdf skipped for %s, because it exists already'%(analydir + photonhdf5), 'c')
        continue

    for chnl in ['chnl0', 'chnl1']:
        if not check_hdf_and_group_exist(results_hdf_loc, grp=chnl + '/spikes/exist_unique',
                                         latest_run=True):
            get_cps(results_hdf_loc, changepointloc, chnl, test=False, amnt_cores=cores, getfromrun='run_1')

    if (not check_hdf_and_group_exist(results_hdf_loc, grp='chnl0/spikes/eventrate_unique',
                                      latest_run=True) 
        and not check_hdf_and_group_exist(results_hdf_loc, grp='chnl1/spikes/eventrate_unique',
                                             latest_run=True)):
        generate_event_stats(results_hdf_loc, stage='unique')
    if (not check_hdf_and_group_exist(results_hdf_loc, grp='chnl0/spikes/exist_hsnr',
                                      latest_run=True) 
        and not check_hdf_and_group_exist(results_hdf_loc, grp='chnl1/spikes/exist_hsnr',
                                             latest_run=True)):
        sort_hsnr_states(results_hdf_loc)
    if (not check_hdf_and_group_exist(results_hdf_loc, grp='chnl0/spikes/eventrate_hsnr',
                                      latest_run=True) 
        and not check_hdf_and_group_exist(results_hdf_loc, grp='chnl1/spikes/eventrate_hsnr',
                                             latest_run=True)):
        generate_event_stats(results_hdf_loc, stage='hsnr')
    if (not check_hdf_and_group_exist(results_hdf_loc, grp='chnl0/spikes/exist_combined',
                                      latest_run=True) 
        and not check_hdf_and_group_exist(results_hdf_loc, grp='chnl1/spikes/exist_combined',
                                             latest_run=True)):
        combine_reentries(results_hdf_loc)

    if (not check_hdf_and_group_exist(results_hdf_loc, grp='chnl0/spikes/eventrate_combined',
                                      latest_run=True) 
        and not check_hdf_and_group_exist(results_hdf_loc, grp='chnl1/spikes/eventrate_combined',
                                             latest_run=True)):
        generate_event_stats(results_hdf_loc, stage='combined')

    logger('All data generated, now plotting', 'c')

    fig, ax = plt.subplots(figsize = (10,6))
    fcs_dict = {'type' : 'all'}
    plotted = plot_fcs(results_hdf_loc,ax, fcs_dict)
    fig.tight_layout()
    if plotted:
        savefig_with_params(results_hdf_loc, fig, fcs_dict, 'correl')

    fig, (ax0, ax1) = plt.subplots(nrows = 2, figsize = (10,10))
    for what in [('chnl0', ax0), ('chnl1', ax1)]:
        tt_dict={'chnl':what[0], 'ax':what[1], 'tt_binning':1e-4, 'stages':['unique', 'hsnr'], 'plt_cp':False}
        plotted = plot_cp_trace(results_hdf_loc, tt_dict)
    fig.tight_layout()
    if plotted:
        savefig_with_params(results_hdf_loc, fig, tt_dict, 'cptt')

    fig, (ax0, ax1) = plt.subplots(nrows = 2, figsize = (10,10))
    for what in [('chnl0', ax0), ('chnl1', ax1)]:
        tt_dict={'chnl':what[0], 'ax':what[1], 'tt_binning':1e-5, 'stages':['unique', 'hsnr'], 'tt_start':2, 'tt_end':4}
        plotted = plot_cp_trace(results_hdf_loc, tt_dict)
    fig.tight_layout()
    if plotted:
        savefig_with_params(results_hdf_loc, fig, tt_dict, 'cptt', mod='2s')

    

    for chnl in ['chnl0', 'chnl1']:
        fig, ax = plt.subplots(ncols = 2, nrows = 2, figsize = (10,10))
        evst_dict = {'chnl':chnl, 'stages':['unique', 'hsnr'], 'ax':ax, 'what':'photon', 'fig':fig, 'threshold':True}
        plotted = plot_event_stats(results_hdf_loc, evst_dict)
        fig.tight_layout(rect=[0, 0.05, 1, 0.95])
        if plotted:
            savefig_with_params(results_hdf_loc, fig, evst_dict, 'evstat')


        fig, ax = plt.subplots(ncols = 2, nrows = 2, figsize = (10,10))
        evst_dict = {'chnl':chnl, 'stages':['unique', 'hsnr'], 'ax':ax, 'what':'photon_rate', 'fig':fig, 'threshold':True}
        plotted = plot_event_stats(results_hdf_loc, evst_dict)
        fig.tight_layout(rect=[0, 0.05, 1, 0.95])
        if plotted:
            savefig_with_params(results_hdf_loc, fig, evst_dict, 'evstat')

        fig, ax = plt.subplots(ncols = 1, nrows = 1, figsize = (10,6))
        ovdict = {'chnl':chnl, 'ax':ax, 't_before':0.01, 't_after':0.01,
                  'stage':'hsnr', 'binning':0.0001, 'amnt':2000}
        plotted = plot_event_overlay(results_hdf_loc, ovdict)
        fig.tight_layout()
        if plotted:
            savefig_with_params(results_hdf_loc, fig, ovdict, 'evov')

        fig, ax = plt.subplots(ncols = 2, nrows = 2, sharey='row', figsize = (10,10))
        plt_int_dict = {'chnl':chnl, 'stage':'hsnr', 'ax':ax, 'threshold':'protein'}
        plot_interevents(results_hdf_loc, plt_int_dict)
        fig.tight_layout()
        if plotted:
            savefig_with_params(results_hdf_loc, fig, plt_int_dict, kind='int_ev')


    plt.close("all")
    gc.collect()

    logger('--------------------------\n\n\n\ndone with file %s \n\n\n\n----------------------------'%fullpath)


logger('done with the filelist', 'c')
