#!/bin/python
# author: phga
# date: 2021-04-08
# desc: Script that turns flexvolt exports/recordings into 2x2 plots
import os
import seaborn as sns
import matplotlib.pyplot as mp
from pandas import read_csv

sns.set_theme(style="white", color_codes=True)
sns.set_palette("colorblind")

plot_dir = "../flexvolt_plots"
typing_tests = []

# Collect all tests
for root, sub, files in os.walk("../flexvolt"):
    for f in files:
        if f.endswith("txt"):
            typing_tests.append(os.path.join(root, f))

for f in typing_tests:
    # Get participants name and datetime for test
    subject = f.split("/")[-2]
    date_time = f.split("flexvolt-recorded-data--")[-1].strip(".txt")
    # Skip first row of meta_data; Only use Processed Data
    flex_data = read_csv(f, skiprows=1, usecols=[0,5,6,7,8])
    # Get min values to normalize data
    start_time = flex_data["Time (ms)"][0]
    min_c1 = flex_data["Chan 1Processed"].min()
    min_c2 = flex_data["Chan 2Processed"].min()
    min_c3 = flex_data["Chan 3Processed"].min()
    min_c4 = flex_data["Chan 4Processed"].min()
    # Subtract min values and start time from respective columns
    flex_data = flex_data.sub([start_time, min_c1, min_c2, min_c3, min_c4])
    # Convert ms to s
    flex_data["Time (ms)"] *= 10**-3
    # Create 2x2 plot
    plot, axi = mp.subplots(2, 2)
    # Settings for whole image
    plot.suptitle(f'{subject}: {date_time}')
    mp.subplots_adjust(hspace=0.5, wspace=0.5)
    # Subplots
    sns.lineplot(x="Time (ms)", y="Chan 2Processed", data=flex_data, ax=axi[0][0])
    sns.lineplot(x="Time (ms)", y="Chan 1Processed", data=flex_data, ax=axi[0][1])
    sns.lineplot(x="Time (ms)", y="Chan 4Processed", data=flex_data, ax=axi[1][0])
    sns.lineplot(x="Time (ms)", y="Chan 3Processed", data=flex_data, ax=axi[1][1])
    # title for each plot and counter var
    plot_title = ["Extensor L", "Extensor R", "Flexor L", "Flexor R"]
    i = 0
    # Label axis
    for rows in axi:
        for plot in rows:
            plot.set_ylim(-5, 300)
            plot.set_title(plot_title[i], fontweight="bold")
            plot.set_ylabel("Muscle Signal [µV]", fontsize="x-small")
            plot.set_xlabel("Time [s]", fontsize="x-small")
            i += 1

    try:
        os.mkdir(f"{plot_dir}/{subject}")
    except OSError as error:
        pass
    mp.savefig(f"{plot_dir}/{subject}/{date_time}.png")