import numpy
import pandas as pd

from tools.helper import get_filters
from tools.definitions import OUTPUT_FOLDER, HELPER_TABLE_FOLDER, STIMULATION_METADATA_FILTERED

print "# Analyse persistent activity"
print "Test whether in the UP protocols frequency was increased/in DOWN protocols frequency was decreased"

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

path_to_filtered_stimulus_file = OUTPUT_FOLDER+HELPER_TABLE_FOLDER+STIMULATION_METADATA_FILTERED

stimulations = pd.read_csv(path_to_filtered_stimulus_file, index_col="stimulation_id")
stimulations["was_successful"] = numpy.nan

# get protocols
type_number_combinations = stimulations[["protocol_type", "pulse_number"]].drop_duplicates().values


def analyse_up_protocols(stimulation):
    after_pulse_firing_columns = [after_pulse_mean_string.format(number) for number in
                                  range(1, stimulation["pulse_number"] + 1)]
    after_pulse_firing = stimulation[after_pulse_firing_columns].values.astype(numpy.float)
    after_pulse_firing = map(lambda firing: numpy.Infinity if numpy.isnan(firing) else firing, after_pulse_firing)

    isi_decreases = numpy.greater_equal(after_pulse_firing[:-1], after_pulse_firing[1:])

    isi_is_decreasing = numpy.all(isi_decreases)
    isi_is_finite_after_the_last_pulse = after_pulse_firing[-1] <= 10
    return isi_is_decreasing and isi_is_finite_after_the_last_pulse


def analyse_down_protocols(stimulation):
    during_pulse_firing_columns = [pulse_mean_string.format(number) for number in
                                   range(1, stimulation["pulse_number"] + 1)]
    during_pulse_firing = stimulation[during_pulse_firing_columns].values.astype(numpy.float)
    during_pulse_firing = map(lambda isi: numpy.Infinity if numpy.isnan(isi) else isi,
                              during_pulse_firing)
    after_protocol_firing = stimulation["isi_after"]

    firing = list(during_pulse_firing) + [after_protocol_firing]

    isi_is_increasing = True
    for isi_pre, isi_after in zip(firing[:-1], firing[1:]):
        isi_pre = numpy.Infinity if numpy.isnan(isi_pre) else isi_pre
        isi_after = numpy.Infinity if numpy.isnan(isi_after) else isi_after
        isi_increases = (isi_pre < isi_after) or (isi_pre == numpy.Infinity and isi_pre == numpy.Infinity)
        isi_is_increasing = isi_is_increasing and isi_increases

    return isi_is_increasing


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)
    path_to_firing_rates_file = OUTPUT_FOLDER+HELPER_TABLE_FOLDER+"{}.csv".format(protocol_id)
    firing_rates = pd.read_csv(path_to_firing_rates_file, index_col="stimulation_id")

    if protocol_type == "UP":
        firing_rates["was_successful"] = firing_rates.apply(analyse_up_protocols,
                                                            axis=1)
    elif protocol_type == "DOWN":
        firing_rates["was_successful"] = firing_rates.apply(analyse_down_protocols,
                                                            axis=1)
    else:
        RuntimeWarning("Unknown protocol {}".format(protocol_type))

    stimulations.loc[(stimulations["protocol_type"] == protocol_type) & (stimulations["pulse_number"] == number) & (
        get_filters(stimulations)),
                     "was_successful"] = firing_rates["was_successful"].values

stimulations.to_csv(path_or_buf=path_to_filtered_stimulus_file)