import numpy as np
import os
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

from ModelFit import get_best_fit
import functions as fu


def main():
    folder = "results/final_1/"

    param_values = get_parameter_distributions(folder)
    plot_distributions(param_values)
    pass


def plot_distributions(param_values):
    bin_number = 30
    fig, axes = plt.subplots(len(param_values.keys()), 2)

    for i, key in enumerate(sorted(param_values.keys())):

        # normal hist:
        values = param_values[key]
        normal, n_bins, patches = axes[i, 0].hist(values, bins=bin_number, density=True)
        axes[i, 0].set_title(key)

        # fit gauss:
        bin_width = np.mean(np.diff(n_bins))
        middle_of_bins = n_bins + bin_width / 2
        axes[i, 0].plot(middle_of_bins[:-1], normal, 'o')
        try:
            n_gauss_pars = fit_gauss(middle_of_bins[:-1], normal)
            x = np.arange(min(param_values[key]), max(param_values[key]),
                          (max(param_values[key]) - min(param_values[key])) / 100)
            axes[i, 0].plot(x, fu.gauss(x, n_gauss_pars[0], n_gauss_pars[1], n_gauss_pars[2]))
        except RuntimeError as e:
            pass

        # log transformed:
        if key != "v_offset":
            log_values = np.log(np.array(param_values[key]))
            log_trans, l_bins, patches = axes[i, 1].hist(log_values, bins=bin_number, density=True)
            bin_width = np.mean(np.diff(l_bins))
            middle_of_bins = l_bins + bin_width / 2
            axes[i, 1].plot(middle_of_bins[:-1], log_trans, 'o')
            try:
                l_gauss_pars = fit_gauss(middle_of_bins[:-1], log_trans)
                x = np.arange(min(log_values), max(log_values),
                              (max(log_values) - min(log_values)) / 100)
                axes[i, 1].plot(x, fu.gauss(x, l_gauss_pars[0], l_gauss_pars[1], l_gauss_pars[2]))
            except RuntimeError as e:
                pass







    plt.tight_layout()
    plt.show()
    plt.close()


def fit_gauss(x_values, y_values):
    mean_v = np.mean(x_values)
    std_v = np.std(x_values)
    amp = max(y_values)
    popt, pcov = curve_fit(fu.gauss, x_values, y_values, p0=(amp, mean_v, std_v))

    return popt

def get_parameter_distributions(folder, param_keys=None):
    if param_keys is None:
        param_keys = ["v_offset", 'input_scaling', 'dend_tau', 'tau_a', 'delta_a',
                      'refractory_period', 'noise_strength', 'mem_tau']
    parameter_values = {}

    for key in param_keys:
        parameter_values[key] = []

    for cell in sorted(os.listdir(folder)):
        fit = get_best_fit(folder + cell)

        final_params = fit.get_final_parameters()

        for key in param_keys:
            parameter_values[key].append(final_params[key])

    return parameter_values








if __name__ == '__main__':
    main()