import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import pandas as pd
import tools.pyrelacs as pr
from tools.definitions import OUTPUT_FOLDER, HELPER_TABLE_FOLDER, PLOTS_FOLDER, STIMULATION_METADATA_RAW
from tools.helper import get_times, get_traces

path_to_raw_stimulus_file = OUTPUT_FOLDER+HELPER_TABLE_FOLDER+STIMULATION_METADATA_RAW
stimulations = pd.read_csv(path_to_raw_stimulus_file, index_col="stimulation_id")

print "# Plotting"
path_to_plot_folder = OUTPUT_FOLDER+PLOTS_FOLDER

for id, stimulation in stimulations.iterrows():
    print "Loading stimulation {}".format(id)
    times = get_times(stimulation)
    trace_infos = pr.get_available_traces(stimulation["recording"])
    traces = get_traces(stimulation)

    labels = ["{} [{}]".format(trace_info["identifier"], trace_info["unit"]) for trace_info in trace_infos]
    print "Plotting stimulation {}".format(id)
    fig = plt.figure(figsize=(9, 16))
    gs = gridspec.GridSpec(len(labels) + 1, 1)
    axes = [fig.add_subplot(gs[trace_idx, :]) for trace_idx in
            range(1, len(labels) + 1)]
    for ax, label, trace in zip(axes, labels, traces):
        ax.plot(times, trace)
        ax.axvline(0,  color='k', ls='--', lw=2)
        ax.axvline(stimulation["stimulus_length"], color='k', ls='--', lw=2)
        ax.set_ylabel(label)
    fig.tight_layout()
    fig.savefig("{}/{}.png".format(path_to_plot_folder, id))
    plt.close(fig)
    print ""