import warnings

import pandas as pd

from tools.helper import get_trace, get_times, get_pulse_times, get_filters
from tools.definitions import OUTPUT_FOLDER, HELPER_TABLE_FOLDER, STIMULATION_METADATA_FILTERED
from tools.voltage_trace_analyzer import VoltageTraceAnalyzer

path_to_filtered_stimulus_file = OUTPUT_FOLDER+HELPER_TABLE_FOLDER+STIMULATION_METADATA_FILTERED

print "# Analyse firing during and after pulses"
print ""

spiking_threshold = -10
minimum_width_of_a_spike_in_time_points = 100

pulse_mean_string = "p_{:d}_ISI_mean"
after_pulse_mean_string = "ap_{:d}_ISI_mean"

stimulations = pd.read_csv(path_to_filtered_stimulus_file, index_col="stimulation_id")
stimulations = stimulations[get_filters(stimulations)]
# get protocols
type_number_combinations = stimulations[["protocol_type", "pulse_number"]].drop_duplicates().values


def get_firing_rates(stimulation):
    times = get_times(stimulation)
    voltage_trace = get_trace(stimulation, "V-1")
    v_analyser = VoltageTraceAnalyzer(times, voltage_trace, spiking_threshold, minimum_width_of_a_spike_in_time_points)

    transient_period = 2

    pulses = get_pulse_times(stimulation)

    # Suppress RuntimeWarnings generated by empty ISI arrays
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)

        for idx, pulse in enumerate(pulses):
            during_pulse_column = pulse_mean_string.format(idx + 1)
            after_pulse_column = after_pulse_mean_string.format(idx + 1)
            stimulation[during_pulse_column] = v_analyser.get_interspike_interval(pulse["start"]+transient_period, pulse["pulse_end"]).mean()

            stimulation[after_pulse_column] = v_analyser.get_interspike_interval(pulse["pulse_end"]+transient_period, pulse["end"]).mean()


        stimulation["isi_before"] = v_analyser.get_interspike_interval(-stimulation["before_protocol"]+transient_period, 0).mean()
        stimulation["isi_after"] = v_analyser.get_interspike_interval(stimulation["stimulus_length"]+transient_period, stimulation[
            "stimulus_length"] + stimulation["after_protocol"]).mean()
    return stimulation


for type_number in type_number_combinations:
    protocol_type = type_number[0]
    number = type_number[1]

    # Determine the protocol
    protocol_id = "{}-{:d}".format(protocol_type, number)
    stimulation_of_current_protocol = stimulations[(stimulations["protocol_type"] == protocol_type) & (stimulations[
                                                                                                           "pulse_number"] == number)]
    print "{}: {} stimulations".format(protocol_id, len(stimulation_of_current_protocol))

    # Analyse firing rates
    stimulation_with_firing_info = stimulation_of_current_protocol.apply(get_firing_rates, axis=1)

    # Save firing rates
    firing_info_columns = ["protocol_type", "pulse_number", "isi_before", "isi_after"] + \
                          [pulse_mean_string.format(pulse_idx) for pulse_idx in range(1, number + 1)] + \
                          [after_pulse_mean_string.format(pulse_idx) for pulse_idx in range(1, number + 1)]
    path_to_firing_rate_file = OUTPUT_FOLDER+HELPER_TABLE_FOLDER+"{:s}.csv".format(protocol_id)
    stimulation_with_firing_info[firing_info_columns].to_csv(path_to_firing_rate_file)
    print "Firing rates saved in {:s}".format(path_to_firing_rate_file)
    print ""