import matplotlib.pyplot as plt
import numpy as np
import os

import helperFunctions as hF

from CellData import CellData
from Baseline import BaselineCellData
from FiCurve import FICurveCellData

MODEL_COLOR = "orange"
DATA_COLOR = "blue"

DATA_SAVE_PATH = "data/figure_data/"


def main():
    # data_isi_histogram()
    data_mean_freq_step_stimulus_examples()
    # data_mean_freq_step_stimulus_with_detections()
    # data_fi_curve()
    pass


def data_fi_curve():
    cell = "data/invivo/2013-04-17-ac-invivo-1/"
    cell_data = CellData(cell)
    fi = FICurveCellData(cell_data, cell_data.get_fi_contrasts())
    fi.plot_fi_curve()


def data_mean_freq_step_stimulus_with_detections():
    cell = "data/invivo/2013-04-17-ac-invivo-1/"
    cell_data = CellData(cell)
    fi = FICurveCellData(cell_data, cell_data.get_fi_contrasts())

    time_traces, freq_traces = fi.get_time_and_freq_traces()
    mean_times, mean_freqs = fi.get_mean_time_and_freq_traces()
    idx = -1
    time = np.array(mean_times[idx])
    freq = np.array(mean_freqs[idx])
    f_inf = fi.f_inf_frequencies[idx]
    f_zero = fi.f_zero_frequencies[idx]

    plt.plot(time, freq, color=DATA_COLOR)
    plt.plot(time[freq == f_zero][0], f_zero, "o", color="black")
    f_inf_time = time[(0.2 < time) & (time < 0.4)]
    plt.plot(f_inf_time, [f_inf for _ in f_inf_time], color="black")
    plt.xlim((-0.1, 0.6))
    plt.show()


def data_mean_freq_step_stimulus_examples():
    # todo smooth! add f_0, f_inf, f_base to it?
    cell = "data/invivo/2013-04-17-ac-invivo-1/"
    cell_data = CellData(cell)
    fi = FICurveCellData(cell_data, cell_data.get_fi_contrasts())

    time_traces, freq_traces = fi.get_time_and_freq_traces()
    mean_times, mean_freqs = fi.get_mean_time_and_freq_traces()

    used_idicies = (0, 7, -1)
    fig, axes = plt.subplots(len(used_idicies), figsize=(8, 12), sharex=True, sharey=True)
    for ax_idx, idx in enumerate(used_idicies):
        sv = fi.stimulus_values[idx]

        # for j in range(len(time_traces[i])):
        #     axes[i].plot(time_traces[i][j], freq_traces[i][j], color="gray", alpha=0.5)

        axes[ax_idx].plot(mean_times[idx], mean_freqs[idx], color=DATA_COLOR)
        # plt.plot(mean_times[i], mean_freqs[i], color="black")

        axes[ax_idx].set_ylabel("Frequency [Hz]")
        axes[ax_idx].set_xlim((-0.2, 0.6))
        axes[ax_idx].set_title("Contrast {:.2f} ({:} trials)".format(sv, len(time_traces[idx])))
    axes[ax_idx].set_xlabel("Time [s]")
    plt.show()


def data_isi_histogram(recalculate=True):
    # if isis loadable - load
    name = "isi_cell_data.npy"
    path = os.path.join(DATA_SAVE_PATH, name)
    if os.path.exists(path) and not recalculate:
        isis = np.load(path)
        print("loaded")
    else:
        # if not get them from the cell
        cell = "data/invivo/2013-04-17-ac-invivo-1/"  # not bursty
        # cell = "data/invivo/2014-12-03-ad-invivo-1/"  # half bursty
        # cell = "data/invivo/2015-01-20-ad-invivo-1/"  # does triple peaks...
        # cell = "data/invivo/2018-05-08-ae-invivo-1/"  # a bit bursty
        # cell = "data/invivo/2013-04-10-af-invivo-1/"  # a bit bursty
        cell_data = CellData(cell)
        base = BaselineCellData(cell_data)
        isis = np.array(base.get_interspike_intervals())
        # base.plot_baseline(position=0,time_length=10)

        # save isis
        np.save(path, isis)
    isis = isis * 1000
    # plot histogram
    bins = np.arange(0, 30.1, 0.1)
    plt.hist(isis, bins=bins, color=DATA_COLOR)
    plt.xlabel("Inter spike intervals [ms]")
    plt.ylabel("Count")
    plt.tight_layout()

    plt.show()


if __name__ == '__main__':
    main()