import pandas as pd


from tools.definitions import STIMULATION_METADATA_RAW, STIMULATION_METADATA_EXTENDED, OUTPUT_FOLDER, HELPER_TABLE_FOLDER



def get_protocol_type(stimulation):
    return "UP" if stimulation["pulse_offset"] > 0 else "DOWN"


def get_pulse_number(stimulation):
    length = stimulation["stimulus_length"]
    period = stimulation["pulse_period"]
    return int(length / period)


def is_control(stim):
    return stim["J"] == 0

path_to_raw_stimulus_file = OUTPUT_FOLDER + HELPER_TABLE_FOLDER + STIMULATION_METADATA_RAW

path_to_extended_stimulus_file = OUTPUT_FOLDER + HELPER_TABLE_FOLDER + STIMULATION_METADATA_EXTENDED

print "# Assign protocol types"
print "Go through stimulations in {} and determine pulse number and amplitude".format(path_to_raw_stimulus_file)

stimulations = pd.read_csv(path_to_raw_stimulus_file, index_col="stimulation_id")
stimulations["protocol_type"] = stimulations.apply(lambda stimulation: get_protocol_type(
    stimulation), axis=1)
stimulations["pulse_number"] = stimulations.apply(lambda stimulation: get_pulse_number(stimulation), axis=1)

stimulations["is_control"] = stimulations.apply(lambda stim: is_control(stim), axis=1)

stimulations.to_csv(path_or_buf=path_to_extended_stimulus_file)
print "Stimulations with pulse information saved to {}".format(path_to_extended_stimulus_file)
print ""