from CellData import CellData
import numpy as np
from scipy.optimize import curve_fit
from scipy.stats import linregress
import matplotlib.pyplot as plt
from warnings import warn
import functions as fu
import helperFunctions as hF


class FICurve:

    def __init__(self, cell_data: CellData, contrast: bool = True):
        self.cell_data = cell_data
        self.using_contrast = contrast

        if contrast:
            self.stimulus_value = cell_data.get_fi_contrasts()
        else:
            self.stimulus_value = cell_data.get_fi_intensities()

        self.f_zeros = []
        self.f_infinities = []
        self.f_baselines = []

        # f_max, f_min, k, x_zero
        self.boltzmann_fit_vars = []
        # increase, offset
        self.f_infinity_fit = []

        self.all_calculate_frequency_points()
        self.f_infinity_fit = hF.fit_clipped_line(self.stimulus_value, self.f_infinities)
        self.boltzmann_fit_vars = hF.fit_boltzmann(self.stimulus_value, self.f_zeros)

    def all_calculate_frequency_points(self):
        mean_frequencies = self.cell_data.get_mean_isi_frequencies()
        time_axes = self.cell_data.get_time_axes_mean_frequencies()
        stimulus_start = self.cell_data.get_stimulus_start()
        stimulus_duration = self.cell_data.get_stimulus_duration()
        sampling_interval = self.cell_data.get_sampling_interval()

        if len(mean_frequencies) == 0:
            warn("FICurve:all_calculate_frequency_points(): mean_frequencies is empty.\n"
                 "Was all_calculate_mean_isi_frequencies already called?")

        for i in range(len(mean_frequencies)):

            if time_axes[i][0] > self.cell_data.get_stimulus_start():
                warn("TODO: Deal with to strongly cut frequency traces in cell data! ")
                self.f_zeros.append(-1)
                self.f_baselines.append(-1)
                self.f_infinities.append(-1)
                continue

            f_zero = hF.detect_f_zero_in_frequency_trace(time_axes[i], mean_frequencies[i],
                                                         stimulus_start, sampling_interval)
            self.f_zeros.append(f_zero)
            f_baseline = hF.detect_f_baseline_in_freq_trace(time_axes[i], mean_frequencies[i],
                                                            stimulus_start, sampling_interval)
            self.f_baselines.append(f_baseline)
            f_infinity = hF.detect_f_infinity_in_freq_trace(time_axes[i], mean_frequencies[i],
                                                            stimulus_start, stimulus_duration, sampling_interval)
            self.f_infinities.append(f_infinity)

    # def __calculate_f_baseline__(self, time, frequency, buffer=0.025):
    #
    #     stim_start = self.cell_data.get_stimulus_start() - time[0]
    #     sampling_interval = self.cell_data.get_sampling_interval()
    #     if stim_start < 0.1:
    #         warn("FICurve:__calculate_f_baseline__(): Quite short delay at the start.")
    #
    #     start_idx = 0
    #     end_idx = int((stim_start-buffer)/sampling_interval)
    #     f_baseline = np.mean(frequency[start_idx:end_idx])
    #
    #     return f_baseline
    #
    # def __calculate_f_zero__(self, time, frequency, peak_buffer_percent=0.05, buffer=0.025):
    #
    #     stimulus_start = self.cell_data.get_stimulus_start() - time[0]  # time start is generally != 0 and != delay
    #     sampling_interval = self.cell_data.get_sampling_interval()
    #
    #     freq_before = frequency[0:int((stimulus_start - buffer) / sampling_interval)]
    #     min_before = min(freq_before)
    #     max_before = max(freq_before)
    #     mean_before = np.mean(freq_before)
    #
    #     # time where the f-zero is searched in
    #     start_idx = int((stimulus_start-0.1*buffer) / sampling_interval)
    #     end_idx = int((stimulus_start + buffer) / sampling_interval)
    #
    #     min_during_start_of_stim = min(frequency[start_idx:end_idx])
    #     max_during_start_of_stim = max(frequency[start_idx:end_idx])
    #
    #     if abs(mean_before-min_during_start_of_stim) > abs(max_during_start_of_stim-mean_before):
    #         f_zero = min_during_start_of_stim
    #     else:
    #         f_zero = max_during_start_of_stim
    #
    #     peak_buffer = (max_before - min_before) * peak_buffer_percent
    #     if min_before - peak_buffer <= f_zero <= max_before + peak_buffer:
    #         end_idx = start_idx + int((end_idx-start_idx)/2)
    #         f_zero = np.mean(frequency[start_idx:end_idx])
    #
    #     return f_zero
    #
    #     # start_idx = int(stimulus_start / sampling_interval)
    #     # end_idx = int((stimulus_start + buffer*2) / sampling_interval)
    #     #
    #     # freq_before = frequency[start_idx-(int(length_of_mean/sampling_interval)):start_idx]
    #     # fb_mean = np.mean(freq_before)
    #     # fb_std = np.std(freq_before)
    #     #
    #     # peak_frequency = fb_mean
    #     # count = 0
    #     # for i in range(start_idx + 1, end_idx):
    #     #     if fb_mean-3*fb_std <= frequency[i] <= fb_mean+3*fb_std:
    #     #         continue
    #     #
    #     #     if abs(frequency[i] - fb_mean) > abs(peak_frequency - fb_mean):
    #     #         peak_frequency = frequency[i]
    #     #         count += 1
    #
    #     # return peak_frequency
    #
    # def __calculate_f_infinity__(self, time, frequency, length=0.1, buffer=0.025):
    #     stimulus_end_time = self.cell_data.get_stimulus_start() + self.cell_data.get_stimulus_duration() - time[0]
    #
    #     start_idx = int((stimulus_end_time - length - buffer) / self.cell_data.get_sampling_interval())
    #     end_idx = int((stimulus_end_time - buffer) / self.cell_data.get_sampling_interval())
    #
    #     # TODO add way to plot detected f_zero, f_inf, f_base. With detection of remaining slope?
    #     # x = np.arange(start_idx, end_idx, 1)  # time[start_idx:end_idx]
    #     # slope, intercept, r_value, p_value, std_err = linregress(x, frequency[start_idx:end_idx])
    #     # if p_value < 0.0001:
    #     #     plt.title("significant slope: {:.2f}, p: {:.5f}, r: {:.5f}".format(slope, p_value, r_value))
    #     #     plt.plot(x, [i*slope + intercept for i in x], color="black")
    #     #
    #     #
    #     # plt.plot((start_idx, end_idx), (np.mean(frequency[start_idx:end_idx]), np.mean(frequency[start_idx:end_idx])), label="f_inf")
    #     # plt.legend()
    #     # plt.show()
    #     # plt.close()
    #
    #     return np.mean(frequency[start_idx:end_idx])

    def get_f_zero_inverse_at_frequency(self, frequency):
        b_vars = self.boltzmann_fit_vars
        return fu.inverse_full_boltzmann(frequency, b_vars[0], b_vars[1], b_vars[2], b_vars[3])

    def get_f_infinity_frequency_at_stimulus_value(self, stimulus_value):
        infty_vars = self.f_infinity_fit
        return fu.clipped_line(stimulus_value, infty_vars[0], infty_vars[1])

    def get_f_infinity_slope(self):
        return self.f_infinity_fit[0]

    def get_fi_curve_slope_at(self, stimulus_value):
        fit_vars = self.boltzmann_fit_vars
        return fu.derivative_full_boltzmann(stimulus_value, fit_vars[0], fit_vars[1], fit_vars[2], fit_vars[3])

    def get_fi_curve_slope_of_straight(self):
        fit_vars = self.boltzmann_fit_vars
        return fu.full_boltzmann_straight_slope(fit_vars[0], fit_vars[1], fit_vars[2], fit_vars[3])

    def get_f_zero_and_f_inf_intersection(self):
        x_values = np.arange(min(self.stimulus_value), max(self.stimulus_value), 0.0001)
        fit_vars = self.boltzmann_fit_vars
        f_zero = fu.full_boltzmann(x_values, fit_vars[0], fit_vars[1], fit_vars[2], fit_vars[3])
        f_inf = fu.clipped_line(x_values, self.f_infinity_fit[0], self.f_infinity_fit[1])

        intersection_indicies = np.argwhere(np.diff(np.sign(f_zero - f_inf))).flatten()
        # print("fi-curve calc intersection:", intersection_indicies, x_values[intersection_indicies])
        if len(intersection_indicies) > 1:
            f_baseline = np.median(self.f_baselines)
            best_dist = np.inf
            best_idx = -1
            for idx in intersection_indicies:
                dist = abs(fu.clipped_line(x_values[idx], self.f_infinity_fit[0], self.f_infinity_fit[1]) - f_baseline)
                if dist < best_dist:
                    best_dist = dist
                    best_idx = idx

            return x_values[best_idx]

        elif len(intersection_indicies) == 0:
            raise ValueError("No intersection found!")
        else:
            return x_values[intersection_indicies[0]]

    def get_fi_curve_slope_at_f_zero_intersection(self):
        x = self.get_f_zero_and_f_inf_intersection()
        fit_vars = self.boltzmann_fit_vars
        return fu.derivative_full_boltzmann(x, fit_vars[0], fit_vars[1], fit_vars[2], fit_vars[3])

    def plot_fi_curve(self, savepath: str = None, comp_f_baselines=None, comp_f_zeros=None, comp_f_infs=None):
        min_x = min(self.stimulus_value)
        max_x = max(self.stimulus_value)
        step = (max_x - min_x) / 5000
        x_values = np.arange(min_x, max_x, step)

        plt.plot(self.stimulus_value, self.f_baselines, color='blue', label='f_base')
        if comp_f_baselines is not None:
            plt.plot(self.stimulus_value, comp_f_baselines, 'o', color='skyblue', label='comp_values base')

        plt.plot(self.stimulus_value, self.f_infinities, 'o', color='green', label='f_inf')
        plt.plot(x_values, [fu.clipped_line(x, self.f_infinity_fit[0], self.f_infinity_fit[1]) for x in x_values],
                 color='darkgreen', label='f_inf_fit')
        if comp_f_infs is not None:
            plt.plot(self.stimulus_value, comp_f_infs, 'o', color='lime', label='comp values f_inf')


        plt.plot(self.stimulus_value, self.f_zeros, 'o', color='orange', label='f_zero')
        popt = self.boltzmann_fit_vars
        plt.plot(x_values, [fu.full_boltzmann(x, popt[0], popt[1], popt[2], popt[3]) for x in x_values],
                 color='red', label='f_0_fit')
        if comp_f_zeros is not None:
            plt.plot(self.stimulus_value, comp_f_zeros, 'o', color='wheat', label='comp_values f_zero')

        plt.legend()
        plt.ylabel("Frequency [Hz]")
        if self.using_contrast:
            plt.xlabel("Stimulus contrast")
        else:
            plt.xlabel("Stimulus intensity [mv]")
        if savepath is None:
            plt.show()
        else:
            print("save")
            plt.savefig(savepath + "fi_curve.png")
        plt.close()

    def plot_f_point_detections(self):
        mean_frequencies = np.array(self.cell_data.get_mean_isi_frequencies())
        time_axes = self.cell_data.get_time_axes_mean_frequencies()

        for i in range(len(mean_frequencies)):
            fig, axes = plt.subplots(1, 1, sharex="all")
            axes.plot(time_axes[i], mean_frequencies[i], label="voltage")
            axes.plot((time_axes[i][0], time_axes[i][-1]), (self.f_zeros[i], self.f_zeros[i]), label="f_zero")
            axes.plot((time_axes[i][0], time_axes[i][-1]), (self.f_infinities[i], self.f_infinities[i]), '--', label="f_inf")
            axes.plot((time_axes[i][0], time_axes[i][-1]), (self.f_baselines[i], self.f_baselines[i]), label="f_base")
            axes.set_title(str(self.stimulus_value[i]))
            plt.legend()

        plt.show()