from models.LIFACnoise import LifacNoiseModel
from stimuli.SinusoidalStepStimulus import SinusoidalStepStimulus
from CellData import CellData, icelldata_of_dir
from Baseline import get_baseline_class
from FiCurve import get_fi_curve_class
from AdaptionCurrent import Adaption
import numpy as np
from warnings import warn
from scipy.optimize import minimize
import time
import os


SAVE_PATH_PREFIX = ""


def main():
    run_with_real_data()


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 = [40, 60, 80]
    noise_strength_list = [0.03]  # [0.02, 0.06]
    dend_tau_list = [0.001, 0.002]

    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:
                    yield {"mem_tau": mem_tau, "input_scaling": input_scaling,
                           "noise_strength": noise_strength, "dend_tau": dend_tau}


def run_with_real_data():
    for cell_data in icelldata_of_dir("./data/"):

        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

        results_path = "results/" + 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)

            print(fmin)
            print(parameters)
            end_time = time.time()
            parameter_set_path = results_path + "start_par_set_{}_fmin_{:.2f}".format(start_par_count, fmin["fun"]) + "/"
            if not os.path.exists(parameter_set_path):
                os.makedirs(parameter_set_path)
            with open(parameter_set_path + "parameters_info.txt".format(start_par_count), "w") as file:
                file.writelines(["start_parameters:\t" + str(start_parameters),
                                 "\nfinal_parameters:\t" + str(parameters),
                                 "\nfinal_fmin:\t" + str(fmin)])

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

    # from Sounds import play_finished_sound
    # play_finished_sound()
    pass


def print_comparision_cell_model(cell_data, data_baseline, data_fi_curve, parameters, plot=False, save_path=None):
    model = LifacNoiseModel(parameters)
    eod_frequency = cell_data.get_eod_frequency()

    model_baseline = get_baseline_class(model, eod_frequency)
    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()

    model_ficurve = get_fi_curve_class(model, cell_data.get_fi_contrasts(), eod_frequency)
    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()

    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_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
    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("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_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)


class Fitter:

    def __init__(self, params=None):
        if params is None:
            self.base_model = LifacNoiseModel({"step_size": 0.00005})
        else:
            self.base_model = LifacNoiseModel(params)
            if "step_size" not in params:
                self.base_model.set_variable("step_size", 0.00005)

        #
        self.fi_contrasts = []
        self.eod_freq = 0

        self.sc_max_lag = 1

        # values to be replicated:
        self.baseline_freq = 0
        self.vector_strength = -1
        self.serial_correlation = []
        self.coefficient_of_variation = 0

        self.f_inf_values = []
        self.f_inf_slope = 0

        self.f_zero_values = []
        # self.f_zero_slopes = []
        self.f_zero_slope_at_straight = 0
        self.f_zero_straight_contrast = 0
        self.f_zero_fit = []

        self.tau_a = 0
        self.delta_a = 0

        # counts how often the cost_function was called
        self.counter = 0

    def fit_model_to_data(self, data: CellData, start_parameters=None):
        self.eod_freq = data.get_eod_frequency()

        data_baseline = get_baseline_class(data)
        self.baseline_freq = data_baseline.get_baseline_frequency()
        self.vector_strength = data_baseline.get_vector_strength()
        self.serial_correlation = data_baseline.get_serial_correlation(self.sc_max_lag)
        self.coefficient_of_variation = data_baseline.get_coefficient_of_variation()

        fi_curve = get_fi_curve_class(data, data.get_fi_contrasts())
        self.fi_contrasts = fi_curve.stimulus_values
        self.f_inf_values = fi_curve.f_inf_frequencies
        self.f_inf_slope = fi_curve.get_f_inf_slope()

        self.f_zero_values = fi_curve.f_zero_frequencies
        self.f_zero_fit = fi_curve.f_zero_fit
        # self.f_zero_slopes = [fi_curve.get_f_zero_fit_slope_at_stimulus_value(c) for c in self.fi_contrasts]
        self.f_zero_slope_at_straight = fi_curve.get_f_zero_fit_slope_at_straight()
        self.f_zero_straight_contrast = self.f_zero_fit[3]

        # around 1/3 of the value at straight
        # self.f_zero_slope = fi_curve.get_fi_curve_slope_at(fi_curve.get_f_zero_and_f_inf_intersection())

        # seems to work if divided by 1000...
        self.delta_a = (fi_curve.get_f_zero_fit_slope_at_straight() / self.f_inf_slope) / 1000

        adaption = Adaption(data, fi_curve)
        self.tau_a = adaption.get_tau_real()

        # print("delta_a: {:.3f}".format(self.delta_a), "tau_a: {:.3f}".format(self.tau_a))

        return self.fit_routine_5(data, start_parameters)

    def fit_routine_5(self, cell_data=None, start_parameters=None):
        self.counter = 0
        # fit only v_offset, mem_tau, input_scaling, dend_tau
        if start_parameters is None:
            x0 = np.array([0.02, 70, 0.001])
        else:
            x0 = np.array([start_parameters["mem_tau"], start_parameters["noise_strength"],
                           start_parameters["input_scaling"], self.tau_a, self.delta_a, start_parameters["dend_tau"]])
        initial_simplex = create_init_simples(x0, search_scale=2)

        # error_list = [error_bf, error_vs, error_sc, error_cv,
        #                       error_f_inf, error_f_inf_slope, error_f_zero, error_f_zero_slope]
        error_weights = (0, 1, 1, 1, 1, 1, 1, 1)
        fmin = minimize(fun=self.cost_function_all,
                        args=(error_weights,), x0=x0, method="Nelder-Mead",
                        options={"initial_simplex": initial_simplex, "xatol": 0.001, "maxfev": 400, "maxiter": 400})

        if cell_data is not None:
            print("##### After step 1:   (Everything)")
            # print_comparision_cell_model(cell_data, res_parameters_step2)

        return fmin, self.base_model.get_parameters()

    def fit_model(self, x0=None, initial_simplex=None, fit_adaption=False):
        self.counter = 0

        if fit_adaption:
            if x0 is None:
                x0 = np.array([0.02, 0.03, 70, self.tau_a, self.delta_a])
            if initial_simplex is None:
                initial_simplex = create_init_simples(x0)

            fmin = minimize(fun=self.cost_function_all, x0=x0,
                            method="Nelder-Mead", options={"initial_simplex": initial_simplex})
        else:
            if x0 is None:
                x0 = np.array([0.02, 0.03, 70])
            if initial_simplex is None:
                initial_simplex = create_init_simples(x0)

            fmin = minimize(fun=self.cost_function_with_fixed_adaption, x0=x0, args=(self.tau_a, self.delta_a),
                            method="Nelder-Mead", options={"initial_simplex": initial_simplex})

        return fmin, self.base_model.get_parameters()

    def cost_function_all(self, X, error_weights=None):
        self.base_model.set_variable("mem_tau", X[0])
        self.base_model.set_variable("noise_strength", X[1])
        self.base_model.set_variable("input_scaling", X[2])
        self.base_model.set_variable("tau_a", X[3])
        self.base_model.set_variable("delta_a", X[4])
        self.base_model.set_variable("dend_tau", X[5])

        base_stimulus = SinusoidalStepStimulus(self.eod_freq, 0)
        # find right v-offset
        test_model = self.base_model.get_model_copy()
        test_model.set_variable("noise_strength", 0)
        v_offset = test_model.find_v_offset(self.baseline_freq, base_stimulus)
        self.base_model.set_variable("v_offset", v_offset)

        # [error_bf, error_vs, error_sc, error_f_inf, error_f_inf_slope, error_f_zero, error_f_zero_slope]
        error_list = self.calculate_errors(error_weights)

        return sum(error_list)

    def cost_function_all_without_noise(self, X, error_weights=None):
        self.base_model.set_variable("mem_tau", X[0])
        self.base_model.set_variable("input_scaling", X[1])
        self.base_model.set_variable("tau_a", X[2])
        self.base_model.set_variable("delta_a", X[3])
        self.base_model.set_variable("dend_tau", X[4])
        self.base_model.set_variable("noise_strength", 0)

        base_stimulus = SinusoidalStepStimulus(self.eod_freq, 0)
        # find right v-offset
        test_model = self.base_model.get_model_copy()
        test_model.set_variable("noise_strength", 0)
        v_offset = test_model.find_v_offset(self.baseline_freq, base_stimulus)
        self.base_model.set_variable("v_offset", v_offset)

        # [error_bf, error_vs, error_sc, error_f_inf, error_f_inf_slope, error_f_zero, error_f_zero_slope]
        error_list = self.calculate_errors(error_weights)

        return sum(error_list)

    def cost_function_only_adaption(self, X, error_weights=None):
        self.base_model.set_variable("tau_a", X[0])
        self.base_model.set_variable("delta_a", X[1])
        self.base_model.set_variable("mem_tau", X[2])

        base_stimulus = SinusoidalStepStimulus(self.eod_freq, 0)
        # find right v-offset
        test_model = self.base_model.get_model_copy()
        test_model.set_variable("noise_strength", 0)
        v_offset = test_model.find_v_offset(self.baseline_freq, base_stimulus)
        self.base_model.set_variable("v_offset", v_offset)
        # [error_bf, error_vs, error_sc, error_f_inf, error_f_inf_slope, error_f_zero, error_f_zero_slope]
        error_list = self.calculate_errors(error_weights)

        return sum(error_list)

    def cost_function_with_fixed_adaption(self, X, tau_a, delta_a, error_weights=None):
        # set model parameters:
        model = self.base_model
        model.set_variable("mem_tau", X[0])
        model.set_variable("noise_strength", X[1])
        model.set_variable("input_scaling", X[2])
        model.set_variable("tau_a", tau_a)
        model.set_variable("delta_a", delta_a)

        base_stimulus = SinusoidalStepStimulus(self.eod_freq, 0)
        # find right v-offset
        test_model = model.get_model_copy()
        test_model.set_variable("noise_strength", 0)
        v_offset = test_model.find_v_offset(self.baseline_freq, base_stimulus)
        model.set_variable("v_offset", v_offset)

        error_list = self.calculate_errors(error_weights)

        return sum(error_list)

    def cost_function_with_fixed_adaption_with_dend_tau_no_noise(self, X, tau_a, delta_a, error_weights=None):
        # set model parameters:
        model = self.base_model
        model.set_variable("mem_tau", X[0])
        model.set_variable("input_scaling", X[1])
        model.set_variable("dend_tau", X[2])
        model.set_variable("tau_a", tau_a)
        model.set_variable("delta_a", delta_a)
        model.set_variable("noise_strength", 0)

        base_stimulus = SinusoidalStepStimulus(self.eod_freq, 0)
        # find right v-offset
        test_model = model.get_model_copy()
        test_model.set_variable("noise_strength", 0)
        v_offset = test_model.find_v_offset(self.baseline_freq, base_stimulus)
        model.set_variable("v_offset", v_offset)

        error_list = self.calculate_errors(error_weights)

        return sum(error_list)

    def cost_function_with_fixed_adaption_with_dend_tau(self, X, tau_a, delta_a, error_weights=None):
        # set model parameters:
        model = self.base_model
        model.set_variable("mem_tau", X[0])
        model.set_variable("noise_strength", X[1])
        model.set_variable("input_scaling", X[2])
        model.set_variable("dend_tau", X[3])
        model.set_variable("tau_a", tau_a)
        model.set_variable("delta_a", delta_a)

        base_stimulus = SinusoidalStepStimulus(self.eod_freq, 0)
        # find right v-offset
        test_model = model.get_model_copy()
        test_model.set_variable("noise_strength", 0)
        v_offset = test_model.find_v_offset(self.baseline_freq, base_stimulus)
        model.set_variable("v_offset", v_offset)

        error_list = self.calculate_errors(error_weights)

        return sum(error_list)

    def calculate_errors(self, error_weights=None):
        model_baseline = get_baseline_class(self.base_model, self.eod_freq)
        baseline_freq = model_baseline.get_baseline_frequency()
        vector_strength = model_baseline.get_vector_strength()
        serial_correlation = model_baseline.get_serial_correlation(self.sc_max_lag)
        coefficient_of_variation = model_baseline.get_coefficient_of_variation()

        fi_curve_model = get_fi_curve_class(self.base_model, self.fi_contrasts, self.eod_freq)
        f_zeros = fi_curve_model.get_f_zero_frequencies()
        f_infinities = fi_curve_model.get_f_inf_frequencies()
        f_infinities_slope = fi_curve_model.get_f_inf_slope()
        # f_zero_slopes = [fi_curve_model.get_f_zero_fit_slope_at_stimulus_value(x) for x in self.fi_contrasts]
        f_zero_slope_at_straight = fi_curve_model.get_f_zero_fit_slope_at_stimulus_value(self.f_zero_straight_contrast)

        # calculate errors with reference values
        error_bf = abs((baseline_freq - self.baseline_freq) / self.baseline_freq)
        error_vs = abs((vector_strength - self.vector_strength) / self.vector_strength)
        error_cv = abs((coefficient_of_variation - self.coefficient_of_variation) / self.coefficient_of_variation)

        error_sc = 0
        for i in range(self.sc_max_lag):
            error_sc = abs((serial_correlation[i] - self.serial_correlation[i]) / (self.serial_correlation[i]/10))
        error_sc = error_sc / self.sc_max_lag

        error_f_inf_slope = abs((f_infinities_slope - self.f_inf_slope) / (self.f_inf_slope/20))
        error_f_inf = calculate_list_error(f_infinities, self.f_inf_values)

        # error_f_zero_slopes = calculate_list_error(f_zero_slopes, self.f_zero_slopes)
        error_f_zero_slope_at_straight = abs(self.f_zero_slope_at_straight - f_zero_slope_at_straight) \
                                            / (self.f_zero_slope_at_straight / 20)
        error_f_zero = calculate_list_error(f_zeros, self.f_zero_values)

        error_list = [error_bf, error_vs, error_sc, error_cv,
                      error_f_inf, error_f_inf_slope, error_f_zero, error_f_zero_slope_at_straight]

        if error_weights is not None and len(error_weights) == len(error_list):
            for i in range(len(error_weights)):
                error_list[i] = error_list[i] * error_weights[i]

        if len(error_weights) != len(error_list):
            warn("Error weights had different length than errors and were ignored!")

        error = sum(error_list)
        self.counter += 1
        if self.counter % 200 == 0:  # and False:
            print("\nCost function run times: {:}\n".format(self.counter),
                  "Total weighted error: {:.4f}\n".format(error),
                  "Baseline frequency     - expected: {:.0f}, current: {:.0f}, error: {:.3f}\n".format(
                      self.baseline_freq, baseline_freq, error_bf),
                  "Vector strength        - expected: {:.2f}, current: {:.2f}, error: {:.3f}\n".format(
                      self.vector_strength, vector_strength, error_vs),
                  "Serial correlation     - expected: {:.2f}, current: {:.2f}, error: {:.3f}\n".format(
                      self.serial_correlation[0], serial_correlation[0], error_sc),
                  "Coefficient of variation - expected: {:.2f}, current: {:.2f}, error: {:.3f}\n".format(
                      self.coefficient_of_variation, coefficient_of_variation, error_cv),
                  "f-infinity slope   - expected: {:.0f}, current: {:.0f}, error: {:.3f}\n".format(
                      self.f_inf_slope, f_infinities_slope, error_f_inf_slope),
                  "f-infinity values:\nexpected:", np.around(self.f_inf_values), "\ncurrent: ", np.around(f_infinities),
                  "\nerror: {:.3f}\n".format(error_f_inf),
                  "f-zero slope   - expected: {:.0f}, current: {:.0f}, error: {:.3f}\n".format(
                      self.f_zero_slope_at_straight, f_zero_slope_at_straight, error_f_zero_slope_at_straight),
                  "f-zero values:\nexpected:", np.around(self.f_zero_values), "\ncurrent: ", np.around(f_zeros),
                  "\nerror: {:.3f}".format(error_f_zero))
        return error_list


def calculate_list_error(fit, reference):
    error = 0
    for i in range(len(reference)):
        error += abs_freq_error(fit[i] - reference[i])

    norm_error = error / len(reference)

    return norm_error


def abs_freq_error(diff, factor=10):
    return abs(diff) / factor


def create_init_simples(x0, search_scale=3.):
    dim = len(x0)
    simplex = [[x0[0]/search_scale], [x0[0]*search_scale]]
    for i in range(1, dim, 1):
        for vertex in simplex:
            vertex.append(x0[i]*search_scale)
        new_vertex = list(x0[:i])
        new_vertex.append(x0[i]/search_scale)
        simplex.append(new_vertex)

    return simplex


if __name__ == '__main__':
    main()