import matplotlib.pyplot as plt
import numpy as np
import time
from os import sep
from multiprocessing import Pool, cpu_count

from constants import OUTPUT_DIR, PARAMS_ID, PARAMS_ID_INV, COLOURS_DARK_BLUE
from data_io import read_varying_mean_sd_fpt_and_params, collect_fpt_mean_stats_and_params, read_fpt_and_params,\
                    write_fpt_and_params
from formulae import stoch_gillespie, stoch_tauleap_lowmem, stoch_tauleap, get_physical_fp_stable_and_not, map_init_name_to_init_cond
from params import Params
from presets import presets
from plotting import plot_table_params


def get_fpt(ensemble, init_cond, params, num_steps=1000000, establish_switch=False, brief=True):
    # TODO could pass simmethod tau or gillespie to params and parse here
    if establish_switch:
        fpt_flag = False
        establish_flag = True
    else:
        fpt_flag = True
        establish_flag = False
    fp_times = np.zeros(ensemble)
    for i in xrange(ensemble):
        if brief:
            species_end, times_end = stoch_tauleap_lowmem(init_cond, num_steps, params, fpt_flag=fpt_flag,
                                                          establish_flag=establish_flag)
        else:
            species, times = stoch_gillespie(init_cond, num_steps, params, fpt_flag=fpt_flag,
                                             establish_flag=establish_flag)
            times_end = times[-1]
            # plotting
            #plt.plot(times, species)
            #plt.show()
        fp_times[i] = times_end
        if establish_switch:
            print "establish time is", fp_times[i]
    return fp_times


def get_mean_fpt(init_cond, params, samplesize=32, establish_switch=False):
    fpt = get_fpt(samplesize, init_cond, params, establish_switch=establish_switch)
    return np.mean(fpt)


def wrapper_get_fpt(fn_args_dict):
    np.random.seed()  # TODO double check that this fixes cluster RNG issues
    if fn_args_dict['kwargs'] is not None:
        return get_fpt(*fn_args_dict['args'], **fn_args_dict['kwargs'])
    else:
        return get_fpt(*fn_args_dict['args'])


def fast_fp_times(ensemble, init_cond, params, num_processes, num_steps='default', establish_switch=False):
    if num_steps == 'default':
        kwargs_dict = {'num_steps': 1000000, 'establish_switch': establish_switch}
    else:
        kwargs_dict = {'num_steps': num_steps, 'establish_switch': establish_switch}

    fn_args_dict = [0]*num_processes
    print "NUM_PROCESSES:", num_processes
    assert ensemble % num_processes == 0
    for i in xrange(num_processes):
        subensemble = ensemble / num_processes
        print "process:", i, "job size:", subensemble, "runs"
        fn_args_dict[i] = {'args': (subensemble, init_cond, params),
                           'kwargs': kwargs_dict}
    t0 = time.time()
    pool = Pool(num_processes)
    results = pool.map(wrapper_get_fpt, fn_args_dict)
    pool.close()
    pool.join()
    print "TIMER:", time.time() - t0

    fp_times = np.zeros(ensemble)
    for i, result in enumerate(results):
        fp_times[i*subensemble:(i+1)*subensemble] = result
    return fp_times


def fast_mean_fpt_varying(param_vary_name, param_vary_values, params, num_processes, init_name="x_all", samplesize=30, establish_switch=False):
    assert samplesize % num_processes == 0
    mean_fpt_varying = [0]*len(param_vary_values)
    sd_fpt_varying = [0]*len(param_vary_values)
    for idx, pv in enumerate(param_vary_values):
        params_step = params.mod_copy( {param_vary_name: pv} )
        init_cond = map_init_name_to_init_cond(params, init_name)
        fp_times = fast_fp_times(samplesize, init_cond, params_step, num_processes, establish_switch=establish_switch)
        mean_fpt_varying[idx] = np.mean(fp_times)
        sd_fpt_varying[idx] = np.std(fp_times)
    return mean_fpt_varying, sd_fpt_varying


def fpt_histogram(fpt_list, params, figname_mod="", flag_show=False, flag_norm=True, flag_xlog10=False, flag_ylog10=False, fs=12):
    ensemble_size = len(fpt_list)
    bins = np.linspace(np.min(fpt_list), np.max(fpt_list), 50) #50)
    #bins = np.arange(0, 3*1e4, 50)  # to plot against FSP

    # normalize
    if flag_norm:
        y_label = 'Probability'
        weights = np.ones_like(fpt_list) / ensemble_size
    else:
        y_label = 'Frequency'
        weights = np.ones_like(fpt_list)

    # prep fig before axes mod
    fig = plt.figure(figsize=(8,6), dpi=120)
    ax = plt.gca()

    # mod axes (log)
    if flag_xlog10:
        ax.set_xscale("log", nonposx='clip')
        max_log = np.ceil(np.max(np.log10(fpt_list)))  # TODO check this matches multihist
        bins = np.logspace(0.1, max_log, 100)
    if flag_ylog10:
        ax.set_yscale("log", nonposx='clip')

    # plot
    plt.hist(fpt_list, bins=bins, alpha=0.6, weights=weights)
    plt.hist(fpt_list, histtype='step', bins=bins, alpha=0.6, label=None, weights=weights, edgecolor='k', linewidth=0.5,
             fill=False)

    # draw mean line
    #plt.axvline(np.mean(fpt_list), color='k', linestyle='dashed', linewidth=2)

    # labels
    plt.title('First-passage time histogram (%d runs) - %s' % (ensemble_size, params.system), fontsize=fs)
    ax.set_xlabel('First-passage time (cell division timescale)', fontsize=fs)
    ax.set_ylabel(y_label, fontsize=fs)
    ax.tick_params(labelsize=fs)
    # plt.locator_params(axis='x', nbins=4)
    #plt.legend(loc='upper right', fontsize=fs)
    # create table of params
    plot_table_params(ax, params)
    # save and show
    plt_save = "fpt_histogram" + figname_mod
    plt.savefig(OUTPUT_DIR + sep + plt_save + '.pdf', bbox_inches='tight')
    if flag_show:
        plt.show()
    return ax


def fpt_histogram_multi(multi_fpt_list, labels, figname_mod="", fs=12, bin_linspace=80, colours=COLOURS_DARK_BLUE,
                        figsize=(8,6), ec='k', lw=0.5, flag_norm=False, flag_show=False, flag_xlog10=False,
                        flag_ylog10=False, flag_disjoint=False):

    # resize fpt lists if not all same size (to the min size)
    fpt_lengths = [len(fpt) for fpt in multi_fpt_list]
    ensemble_size = np.min(fpt_lengths)

    # cleanup data to same size
    if sum(fpt_lengths - ensemble_size) > 0:
        print "Resizing multi_fpt_list elements:", fpt_lengths, "to the min size of:", ensemble_size
        for idx in xrange(len(fpt_lengths)):
            multi_fpt_list[idx] = multi_fpt_list[idx][:ensemble_size]
    bins = np.linspace(np.min(multi_fpt_list), np.max(multi_fpt_list), bin_linspace)

    # normalize
    if flag_norm:
        y_label = 'Probability'
        weights = np.ones_like(multi_fpt_list) / ensemble_size
    else:
        y_label = 'Frequency'
        weights = np.ones_like(multi_fpt_list)

    # prep fig before axes mod
    fig = plt.figure(figsize=figsize, dpi=120)
    ax = plt.gca()

    # mod axes (log)
    if flag_xlog10:
        ax.set_xscale("log", nonposx='clip')
        max_log = np.ceil(np.max(np.log10(multi_fpt_list)))
        bins = np.logspace(0.1, max_log, 100)
    if flag_ylog10:
        ax.set_yscale("log", nonposx='clip')

    # plot calls
    if flag_disjoint:
        plt.hist(multi_fpt_list, bins=bins, color=colours, label=labels, weights=weights, edgecolor=ec, linewidth=lw)
    else:
        for idx, fpt_list in enumerate(multi_fpt_list):
            plt.hist(fpt_list, bins=bins, alpha=0.6, color=colours[idx], label=labels[idx],
                     weights=weights[idx,:])
            plt.hist(fpt_list, histtype='step', bins=bins, alpha=0.6, color=colours[idx],
                     label=None,weights=weights[idx,:], edgecolor=ec, linewidth=lw, fill=False)

    # labels
    plt.title('First-passage time histogram (%d runs)' % (ensemble_size), fontsize=fs)
    ax.set_xlabel('First-passage time (cell division timescale)', fontsize=fs)
    ax.set_ylabel(y_label, fontsize=fs)
    plt.legend(loc='upper right', fontsize=fs)
    ax.tick_params(labelsize=fs)
    # plt.locator_params(axis='x', nbins=4)

    # save and show
    plt_save = "fpt_multihistogram" + figname_mod
    fig.savefig(OUTPUT_DIR + sep + plt_save + '.pdf', bbox_inches='tight')
    if flag_show:
        plt.show()


def plot_mean_fpt_varying(mean_fpt_varying, sd_fpt_varying, param_vary_name, param_set, params, samplesize, SEM_flag=True, show_flag=False, figname_mod=""):
    if SEM_flag:
        sd_fpt_varying = sd_fpt_varying / np.sqrt(samplesize)  # s.d. from CLT since sample mean is approx N(mu, sd**2/n)
    plt.errorbar(param_set, mean_fpt_varying, yerr=sd_fpt_varying, label="sim")
    plt.title("Mean FP Time, %s varying (sample=%d)" % (param_vary_name, samplesize))
    ax = plt.gca()
    ax.set_xlabel(param_vary_name)
    ax.set_ylabel('Mean FP time')

    # log options
    for i in xrange(len(mean_fpt_varying)):
        print i, param_set[i], mean_fpt_varying[i], sd_fpt_varying[i]
    flag_xlog10 = True
    flag_ylog10 = True
    if flag_xlog10:
        ax.set_xscale("log", nonposx='clip')
        #ax.set_xlim([0.8*1e2, 1*1e7])
    if flag_ylog10:
        ax.set_yscale("log", nonposx='clip')
        #ax.set_ylim([0.8*1e2, 3*1e5])

    # create table of params
    plot_table_params(ax, params)
    plt_save = "mean_fpt_varying" + figname_mod
    plt.savefig(OUTPUT_DIR + sep + plt_save + '.png', bbox_inches='tight')
    if show_flag:
        plt.show()
    return ax


if __name__ == "__main__":
    # SCRIPT FLAGS
    run_compute_fpt = False
    run_read_fpt = False
    run_generate_hist_multi = False
    run_load_hist_multi = False
    run_collect = False
    run_means_read_and_plot = False
    run_means_collect_and_plot = True

    # SCRIPT PARAMETERS
    establish_switch = True
    brief = True
    num_steps = 1000000  # default 1000000
    ensemble = 1  # default 100

    # DYNAMICS PARAMETERS
    params = presets('preset_xyz_constant')  # preset_xyz_constant, preset_xyz_constant_fast, valley_2hit

    # OTHER PARAMETERS
    init_cond = np.zeros(params.numstates, dtype=int)
    init_cond[0] = int(params.N)

    # PLOTTING
    FS = 16
    EC = 'k'
    LW = 0.5
    FIGSIZE=(8,6)

    if run_compute_fpt:
        fp_times = get_fpt(ensemble, init_cond, params, num_steps=num_steps, establish_switch=establish_switch, brief=brief)
        write_fpt_and_params(fp_times, params)
        fpt_histogram(fp_times, params, flag_show=True, figname_mod="XZ_model_withFeedback_mu1e-1")

    if run_read_fpt:
        dbdir = OUTPUT_DIR
        dbdir_100 = dbdir + sep + "fpt_mean" + sep + "100_c95"
        fp_times_xyz_100, params_a = read_fpt_and_params(dbdir_100)
        dbdir_10k = dbdir + sep + "fpt_mean" + sep + "10k_c95"
        fp_times_xyz_10k, params_b = read_fpt_and_params(dbdir_10k)

    if run_generate_hist_multi:
        ensemble = 21
        num_proc = cpu_count() - 1
        param_vary_id = "N"
        param_idx = PARAMS_ID_INV[param_vary_id]
        param_vary_values = [1e2, 1e3, 1e4]
        param_vary_labels = ['A', 'B', 'C']
        params_ensemble = [params.params_list[:] for _ in param_vary_values]
        multi_fpt = np.zeros((len(param_vary_values), ensemble))
        multi_fpt_labels = ['label' for _ in param_vary_values]
        for idx, param_val in enumerate(param_vary_values):
            param_val_string = "%s=%.3f" % (param_vary_id, param_val)
            params_step = params.mod_copy({param_vary_id: param_val})
            #fp_times = get_fpt(ensemble, init_cond, params_set[idx], num_steps=num_steps)
            fp_times = fast_fp_times(ensemble, init_cond, params_step, num_proc, establish_switch=establish_switch)
            write_fpt_and_params(fp_times, params_step, filename="fpt_multi", filename_mod=param_val_string)
            multi_fpt[idx,:] = np.array(fp_times)
            multi_fpt_labels[idx] = "%s (%s)" % (param_vary_labels[idx], param_val_string)
        fpt_histogram_multi(multi_fpt, multi_fpt_labels, flag_show=True, flag_ylog10=False)

    if run_load_hist_multi:
        flag_norm = True
        dbdir = OUTPUT_DIR + sep + "may25_100"
        #dbdir_c80 = dbdir + "fpt_feedback_z_ens1040_c0.80_params"
        c80_header = "fpt_feedback_z_ens1040_c80_N100"
        c88_header = "fpt_feedback_z_ens1040_c88_N100"
        c95_header = "fpt_feedback_z_ens1040_c95_N100"
        fp_times_xyz_c80, params_a = read_fpt_and_params(dbdir, "%s_data.txt" % c80_header, "%s_params.csv" % c80_header)
        fp_times_xyz_c88, params_b = read_fpt_and_params(dbdir, "%s_data.txt" % c88_header, "%s_params.csv" % c88_header)
        fp_times_xyz_c95, params_c = read_fpt_and_params(dbdir, "%s_data.txt" % c95_header, "%s_params.csv" % c95_header)
        fpt_histogram(fp_times_xyz_c88, params_b, flag_ylog10=False, figname_mod="_xyz_feedbackz_N10k_c88_may25")
        plt.close('all')
        fpt_histogram(fp_times_xyz_c88, params_b, flag_ylog10=True, figname_mod="_xyz_feedbackz_N10k_c88_may25_logy")
        plt.close('all')
        multi_fpt = [fp_times_xyz_c80, fp_times_xyz_c88, fp_times_xyz_c95]
        labels = ("c=0.80 (Region I)", "c=0.88 (Region IV)", "c=0.95 (Region III)")
        fpt_histogram_multi(multi_fpt, labels, flag_show=True, flag_ylog10=False, flag_norm=flag_norm, fs=FS, ec=EC, lw=LW, figsize=FIGSIZE)
        fpt_histogram_multi(multi_fpt, labels, flag_show=True, flag_ylog10=True, flag_norm=flag_norm, fs=FS, ec=EC, lw=LW, figsize=FIGSIZE)
        fpt_histogram_multi(multi_fpt, labels, flag_show=True, flag_ylog10=True, flag_norm=False, fs=FS, ec=EC, lw=LW, figsize=FIGSIZE, flag_disjoint=True)

    if run_means_read_and_plot:
        datafile = OUTPUT_DIR + sep + "fpt_stats_collected_mean_sd_varying_N.txt"
        paramfile = OUTPUT_DIR + sep + "fpt_stats_collected_mean_sd_varying_N_params.csv"
        samplesize=48
        mean_fpt_varying, sd_fpt_varying, param_to_vary, param_set, params = \
            read_varying_mean_sd_fpt_and_params(datafile, paramfile)
        plt_axis = plot_mean_fpt_varying(mean_fpt_varying, sd_fpt_varying, param_to_vary, param_set, params, samplesize,
                                         SEM_flag=True, show_flag=True, figname_mod="_%s_n%d" % (param_to_vary, samplesize))
        """
        mu = params.mu
        mixed_fp_zinf_at_N = [0.0]*len(param_set)
        for idx, N in enumerate(param_set):
            params_at_N = params.mod_copy( {'N': N} )
            fps = get_physical_and_stable_fp(params_at_N)
            assert len(fps) == 1
            mixed_fp_zinf_at_N[idx] = fps[0][2]
        plt_axis.plot(param_set, [1/(mu*n) for n in param_set], '-o', label="(mu*N)^-1")
        plt_axis.plot(param_set, [1/(mu*zinf) for zinf in mixed_fp_zinf_at_N], '-o', label="(mu*z_inf)^-1")
        plt_axis.set_yscale("log", nonposx='clip')
        plt_axis.set_xscale("log", nonposx='clip')
        plt_axis.legend()
        plt.savefig(OUTPUT_DIR + sep + "theorycompare_loglog" + '.png', bbox_inches='tight')
        plt.show()
        """

    if run_means_collect_and_plot:
        dbdir = OUTPUT_DIR + sep + "tocollect" + sep + "runset_june17_FPT_cvary_44_ens240"
        datafile, paramfile = collect_fpt_mean_stats_and_params(dbdir)
        samplesize=240
        mean_fpt_varying, sd_fpt_varying, param_to_vary, param_set, params = \
            read_varying_mean_sd_fpt_and_params(datafile, paramfile)
        plot_mean_fpt_varying(mean_fpt_varying, sd_fpt_varying, param_to_vary, param_set, params, samplesize,
                              SEM_flag=True, show_flag=True, figname_mod="_%s_n%d" % (param_to_vary, samplesize))

