import os
import sys
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from run import run

from matplotlib.colors import LinearSegmentedColormap


def run_experiment2(data_path):
    '''
    Function that conducts a series of experiments with varying:
     - number of segments used for training and testing (N)
     - segment length (T)

    Parameters:
    data_path (str): path to the preprocessed dataset.

    Saves:
    - CSV files for accuracy, training time, and evaluation time grids, and 
    - a heatmap plot of model accuracy in the 'exp2_results' directory.
    '''

    segments_range = [500,350,200,100,50]
    seconds_range = [1,2,5,7,10]
    leads = [0,1]

    accuracy_grid = np.zeros((len(segments_range), len(seconds_range)))
    train_time_grid = np.zeros_like(accuracy_grid)
    eval_time_grid = np.zeros_like(accuracy_grid)


    for i, num_segments in enumerate(segments_range):
        for j, num_seconds in enumerate(seconds_range):
            print(f'Training for {num_segments} segments, {num_seconds} s long')
            accuracy, total_train_time, total_eval_time = run(NUM_SEGMENTS=num_segments, 
                                                              NUM_SECONDS=num_seconds, 
                                                              NUM_BATCH=16, 
                                                              LEADS=leads, 
                                                              NUM_EPOCHS=200, 
                                                              DATA_PATH = data_path, 
                                                              FS=128)
            accuracy_grid[i, j] = accuracy
            train_time_grid[i,j] = total_train_time
            eval_time_grid[i,j] = total_eval_time


    os.makedirs('exp2_results', exist_ok=True)
    np.savetxt('exp2_results/accuracy_grid.csv', accuracy_grid, delimiter=',', fmt='%.2f')
    np.savetxt('exp2_results/train_time_grid.csv', train_time_grid, delimiter=',', fmt='%.2f')
    np.savetxt('exp2_results/eval_time_grid.csv', eval_time_grid, delimiter=',', fmt='%.2f')

    color_list = ["#FCDDD9", "#FABEB7", "#FAA99F", "#FA968A", "#F87F71",
                "#DC7366", "#B15F56", "#884C45", "#593633", "#2E2322"]
    custom_colormap = LinearSegmentedColormap.from_list("custom_salmon", color_list)

    yticklabels = [int(w/2) for w in segments_range]
    plt.figure(figsize=(10, 8), dpi=100)
    sns.heatmap(accuracy_grid, annot=True, fmt=".2f", xticklabels=seconds_range, yticklabels=yticklabels, 
                cmap=custom_colormap, cbar_kws={'label': 'Accuracy %'})
    plt.xlabel('NUMBER OF SECONDS (T)')
    plt.ylabel('TRAIN SAMPLES PER SUBJECT (N)')
    plt.savefig('exp2_results/fig_exp2.png')
    plt.close()

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python run_exp_2.py <data_path>")
        sys.exit(1)

    data_path = sys.argv[1]

    run_experiment2(data_path=data_path)