from models.LIFACnoise import LifacNoiseModel
from CellData import icelldata_of_dir, CellData
from Baseline import get_baseline_class
from FiCurve import get_fi_curve_class
from Fitter import Fitter
from ModelFit import ModelFit

import time
import os
import copy
import argparse
import numpy as np

import multiprocessing as mp


SAVE_PATH_PREFIX = ""
FIT_ROUTINE = ""


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--cell", help="folder (with .dat files) containing the cell data")
    args = parser.parse_args()
    if args.cell is not None:
        cell_data = CellData(args.cell)
        start_parameters = [p for p in iget_start_parameters()]
        fit_cell_parallel(cell_data, start_parameters)
        quit()

    # test_single_cell("invivo_data/2012-01-17-ap/")
    #
    # quit()

    start_parameters = [p for p in iget_start_parameters()]
    # start_data = 8
    # count = 0
    # for cell_data in icelldata_of_dir("./invivo_data/"):
    #     count += 1
    #     if count < start_data:
    #         continue
    #     fit_cell_parallel(cell_data, start_parameters)

    cell_data = CellData("invivo_data/2012-04-20-ab-invivo-1/")
    fit_cell_parallel(cell_data, start_parameters)


def test_single_cell(path):
    cell_data = CellData(path)
    start_parameters = [p for p in iget_start_parameters()]

    for i, p in enumerate(start_parameters):
        fitter = Fitter()
        fitter.set_data_reference_values(cell_data)
        fmin, res_par = fitter.fit_routine_1(p)

        cell_path = os.path.basename(cell_data.get_data_path())

        error = fitter.calculate_errors(model=LifacNoiseModel(res_par))
        save_path = "results/invivo_results/" + cell_path + "/start_parameter_{:}_err_{:.2f}/".format(i, sum(error))
        save_fitting_run_info(cell_data, res_par, p, plot=True, save_path=save_path)
        print("Done with start parameters {}".format(str(i)))


def fit_cell_base(parameters):
    # parameter = (cell_data, start_parameter_index, start_parameter, results_base_folder)
    time1 = time.time()
    fitter = Fitter()
    fitter.set_data_reference_values(parameters[0])
    fmin, res_par = fitter.fit_routine_1(parameters[2])

    cell_data = parameters[0]
    cell_path = os.path.split(cell_data.get_data_path())[-1]

    error = fitter.calculate_errors(model=LifacNoiseModel(res_par))
    save_path = parameters[3] + "/" + cell_path + "/start_parameter_{:}_err_{:.2f}/".format(parameters[1], sum(error))
    save_fitting_run_info(parameters[0], res_par, parameters[2], plot=True, save_path=save_path)
    time2 = time.time()

    del fitter

    print("Time taken for " + cell_path +
          "\n and start parameters ({:}): {:.2f}s thread time".format(parameters[1]+1, time2 - time1) +
          "\n error: {:.2f}".format(sum(error)))


def fit_all_cells_parallel_sync(cells, start_parameters, thread_pool, results_base_folder):
    parameter = []
    for cell in cells:
        for i, s_pars in enumerate(start_parameters):
            parameter.append((cell, i, s_pars, results_base_folder))

    time1 = time.time()
    thread_pool.map(fit_cell_base, parameter)
    time2 = time.time()
    print("Time taken for all ({:}): {:.2f}s".format(len(parameter)*len(cells), time2 - time1))


def fit_cell_parallel(cell_data, start_parameters):
    cell_path = os.path.basename(cell_data.get_data_path())
    save_directory = "./results/invivo_results/"
    save_path_cell = os.path.join(save_directory, cell_data.get_cell_name())
    print(cell_path)
    core_count = mp.cpu_count()
    pool = mp.Pool(core_count - 1)

    parameters = []
    for i, p in enumerate(start_parameters):
        parameters.append((cell_data, i, p, save_directory))

    time1 = time.time()
    pool.map(fit_cell_base, parameters)
    time2 = time.time()
    print("Time taken for all start parameters ({:}): {:.2f}s".format(len(start_parameters), time2-time1))
    del pool
    del cell_data

    best_fit = None
    min_err = np.inf
    for fit in os.listdir(save_path_cell):
        cur_fit = ModelFit(os.path.join(save_path_cell, fit))
        if cur_fit.comparable_error() < min_err:
            min_err = cur_fit.comparable_error()
            best_fit = cur_fit

    best_fit.generate_master_plot("./results/invivo_best/singles/")



def test_fit_routines():
    fitter = Fitter()
    names = ("routine_1", "routine_2")
    global FIT_ROUTINE
    for i, routine in enumerate([fitter.fit_routine_1, fitter.fit_routine_2]):
        FIT_ROUTINE = names[i]
        run_with_real_data(fitter, routine)

    best = []
    cells = sorted(os.listdir("test_routines/" + names[0] + "/"))
    for name in names:

        save_path = "test_routines/" + name + "/"
        cell_best = []
        for directory in sorted(os.listdir(save_path)):
            path = os.path.join(save_path, directory)
            if os.path.isdir(path):
                cell_best.append(find_best_run(path))
        best.append(cell_best)

    with open("test_routines/comparision.csv", "w") as res_file:
        res_file.write("routine")
        for cell in cells:
            res_file.write("," + cell)

        for i, routine_results in enumerate(best):
            res_file.write(names[i])
            for cell_best in routine_results:
                res_file.write("," + str(cell_best))


def find_best_run(cell_path):
    values = []
    for directory in sorted(os.listdir(cell_path)):
        start_par_path = os.path.join(cell_path, directory)
        if os.path.isdir(start_par_path):
            values.append(float(start_par_path.split("_")[-1]))

    return min(values)


def iget_start_parameters():
    # mem_tau, input_scaling, noise_strength, dend_tau,
    # expand by tau_a, delta_a ?

    mem_tau_list = [0.01]
    input_scaling_list = [60]
    noise_strength_list = [0.03]  # [0.02, 0.06]
    dend_tau_list = [0.001, 0.002]
    delta_a_list = [0.035, 0.065]
    tau_a_list = [0.1, 0.4]
    ref_time_list = [0.00065]

    for mem_tau in mem_tau_list:
        for input_scaling in input_scaling_list:
            for noise_strength in noise_strength_list:
                for dend_tau in dend_tau_list:
                    for delta_a in delta_a_list:
                        for tau_a in tau_a_list:
                            for ref_time in ref_time_list:
                                yield {"mem_tau": mem_tau, "input_scaling": input_scaling,
                                       "noise_strength": noise_strength, "dend_tau": dend_tau,
                                       "delta_a": delta_a, "tau_a": tau_a, "refractory_period": ref_time}


def run_with_real_data(fitter, fit_routine_func, parallel=False):
    count = 0
    for cell_data in icelldata_of_dir("./data/"):
        count += 1
        if count < 7:
            pass
            #continue

        print("cell:", cell_data.get_data_path())
        trace = cell_data.get_base_traces(trace_type=cell_data.V1)
        if len(trace) == 0:
            print("NO V1 TRACE FOUND")
            continue

        global FIT_ROUTINE
        # results_path = "results/" + os.path.split(cell_data.get_data_path())[-1] + "/"
        results_path = "test_routines/" + FIT_ROUTINE + "/" + os.path.split(cell_data.get_data_path())[-1] + "/"
        print("results at:", results_path)

        if not os.path.exists(results_path):
            os.makedirs(results_path)

        # plot cell images:
        cell_save_path = results_path + "cell/"
        if not os.path.exists(cell_save_path):
            os.makedirs(cell_save_path)
        data_baseline = get_baseline_class(cell_data)
        data_baseline.plot_baseline(cell_save_path)
        data_baseline.plot_interspike_interval_histogram(cell_save_path)
        data_baseline.plot_serial_correlation(6, cell_save_path)

        data_fi_curve = get_fi_curve_class(cell_data, cell_data.get_fi_contrasts())
        data_fi_curve.plot_fi_curve(cell_save_path)

        start_par_count = 0
        for start_parameters in iget_start_parameters():
            start_par_count += 1
            print("START PARAMETERS:", start_par_count)

            start_time = time.time()
            # fitter = Fitter()
            fmin, parameters = fitter.fit_model_to_data(cell_data, start_parameters, fit_routine_func)

            print(fmin)
            print(parameters)
            end_time = time.time()
            parameter_set_path = results_path + "start_par_set_{}_fmin_{:.2f}".format(start_par_count, fmin["fun"]) + "/"

            print('Fitting of cell took function took {:.3f} s'.format((end_time - start_time)))
            # print(results_path)
            save_fitting_run_info(cell_data, parameters, start_parameters,
                                  plot=True, save_path=parameter_set_path)

    # from Sounds import play_finished_sound
    # play_finished_sound()
    pass


def save_fitting_run_info(cell_data, parameters, start_parameters, plot=False, save_path=None):
    if save_path is not None:
        if not os.path.exists(save_path):
            os.makedirs(save_path)

    if save_path is None:
        return

    with open(save_path + "parameters_info.txt", "w") as file:
        file.writelines(["start_parameters:\t" + str(start_parameters),
                         "\nfinal_parameters:\t" + str(parameters)])

    model = LifacNoiseModel(parameters)
    eod_frequency = cell_data.get_eod_frequency()

    data_baseline = get_baseline_class(cell_data)
    c_bf = data_baseline.get_baseline_frequency()
    c_vs = data_baseline.get_vector_strength()
    c_sc = data_baseline.get_serial_correlation(1)
    c_cv = data_baseline.get_coefficient_of_variation()
    c_burst = data_baseline.get_burstiness()

    data_fi_curve = get_fi_curve_class(cell_data, cell_data.get_fi_contrasts())
    c_f_inf_slope = data_fi_curve.get_f_inf_slope()
    c_f_inf_values = data_fi_curve.f_inf_frequencies

    c_f_zero_slope = data_fi_curve.get_f_zero_fit_slope_at_straight()
    c_f_zero_values = data_fi_curve.f_zero_frequencies

    if c_f_inf_slope < 0:
        contrasts = np.array(cell_data.get_fi_contrasts()) * -1
    else:
        contrasts = np.array(cell_data.get_fi_contrasts())
    model_baseline = get_baseline_class(model, eod_frequency, trials=15)
    m_bf = model_baseline.get_baseline_frequency()
    m_vs = model_baseline.get_vector_strength()
    m_sc = model_baseline.get_serial_correlation(1)
    m_cv = model_baseline.get_coefficient_of_variation()
    m_burst = model_baseline.get_burstiness()

    model_ficurve = get_fi_curve_class(model, cell_data.get_fi_contrasts(), eod_frequency, trials=60)
    m_f_infinities = model_ficurve.get_f_inf_frequencies()
    m_f_zeros = model_ficurve.get_f_zero_frequencies()
    m_f_infinities_slope = model_ficurve.get_f_inf_slope()
    m_f_zero_slope = model_ficurve.get_f_zero_fit_slope_at_straight()

    np.save(os.path.join(save_path, "model_fi_inf_values.npy"), np.array(m_f_infinities))
    np.save(os.path.join(save_path, "cell_fi_inf_values.npy"), np.array(c_f_inf_values))
    np.save(os.path.join(save_path, "model_fi_zero_values.npy"), np.array(m_f_zeros))
    np.save(os.path.join(save_path, "cell_fi_zero_values.npy"), np.array(c_f_zero_values))

    with open(os.path.join(save_path, "cell_data_path.txt"), "w") as f:
        path = cell_data.get_data_path() + "\n"
        f.write(path)

    if c_f_inf_slope < 0:
        model_ficurve.stimulus_values = contrasts * -1

    # print("EOD-frequency: {:.2f}".format(cell_data.get_eod_frequency()))
    # print("bf: cell - {:.2f} vs model {:.2f}".format(c_bf, m_bf))
    # print("vs: cell - {:.2f} vs model {:.2f}".format(c_vs, m_vs))
    # print("sc: cell - {:.2f} vs model {:.2f}".format(c_sc[0], m_sc[0]))
    # print("cv: cell - {:.2f} vs model {:.2f}".format(c_cv, m_cv))
    # print("f_inf_slope: cell - {:.2f} vs model {:.2f}".format(c_f_inf_slope, m_f_infinities_slope))
    # print("f infinity values:\n cell  -", c_f_inf_values, "\n model -", m_f_infinities)
    #
    # print("f_zero_slope: cell - {:.2f} vs model {:.2f}".format(c_f_zero_slope, m_f_zero_slope))
    # print("f zero values:\n cell  -", c_f_zero_values, "\n model -", m_f_zeros)
    if save_path is not None:
        with open(save_path + "value_comparision.tsv", 'w') as value_file:
            value_file.write("Variable\tCell\tModel\n")
            value_file.write("baseline_frequency\t{:.2f}\t{:.2f}\n".format(c_bf, m_bf))
            value_file.write("vector_strength\t{:.2f}\t{:.2f}\n".format(c_vs, m_vs))
            value_file.write("serial_correlation\t{:.2f}\t{:.2f}\n".format(c_sc[0], m_sc[0]))
            value_file.write("Burstiness\t{:.2f}\t{:.2f}\n".format(c_burst, m_burst))
            value_file.write("coefficient_of_variation\t{:.2f}\t{:.2f}\n".format(c_cv, m_cv))
            value_file.write("f_inf_slope\t{:.2f}\t{:.2f}\n".format(c_f_inf_slope, m_f_infinities_slope))
            value_file.write("f_zero_slope\t{:.2f}\t{:.2f}\n".format(c_f_zero_slope, m_f_zero_slope))

    if plot:
        # plot model images
        model_baseline.plot_baseline(save_path)
        model_baseline.plot_interspike_interval_histogram(save_path)
        model_baseline.plot_isi_histogram_comparision(data_baseline.get_interspike_intervals(),
                                                      model_baseline.get_interspike_intervals(), save_path)
        model_baseline.plot_serial_correlation(6, save_path)

        model_ficurve.plot_fi_curve(save_path)
        model_ficurve.plot_fi_curve_comparision(data_fi_curve, model_ficurve, save_path)


def test_effect_of_refractory_period():
    ref_periods = [0.0006, 0.001, 0.0015]
    counter = 0
    core_count = mp.cpu_count()

    for cell in icelldata_of_dir("./data/"):
        pool = mp.Pool(core_count - 1)
        counter += 1
        if counter < 10:
            continue
        elif counter >= 14:
            return
        start_parameters_base = [p for p in iget_start_parameters()]
        for ref_period in ref_periods:
            print(cell.get_data_path())
            print("ref period: {:.4f}".format(ref_period))
            results_base_folder = "./test_routines/ref_period_{:.4f}/".format(ref_period)
            all_start_parameters = copy.deepcopy(start_parameters_base)

            for par_set in all_start_parameters:
                par_set["refractory_period"] = ref_period
            fit_all_cells_parallel_sync([cell], all_start_parameters, pool, results_base_folder)

        del cell
        del pool


if __name__ == '__main__':
    main()