import sys
import os
import shutil
import wfdb
import numpy as np
from scipy.io import savemat
import tkinter as tk
from tkinter import filedialog

# =============================================================================
# README: WFDB to MAT Conversion Script (GRABMyoFlow dynamic Dataset)
# =============================================================================
# This script processes raw WFDB data files (.dat, .hea) from the GRABMyoFlow
# dynamic dataset (30 gestures, 2 trials each) and converts them into MATLAB (.mat) files.
#
# UPDATED LOGIC:
# - Removes channels 1, 8, 9, 16 (corresponding to U1, U2, U3, U4).
# - Keeps only the 12 Wrist channels.
# =============================================================================


# =============================================================================
# STEP 0: CONSTANTS AND INITIAL SETUP
# =============================================================================

root = tk.Tk()
root.withdraw()

# --- Dataset Constants ---
NSESSION = 3
NGEESTURE = 30 
NTRIALS = 2    

# --- Dataset Mask (16 Channels) ---
# Keeps 12 Wrist channels (Skipping indices 1, 7, 8, 15)
EXTENSION_INDICES_KEEP = list(range(1, 7)) + list(range(9, 15))


# =============================================================================
# STEP 1: USER CHOICE AND PATH SELECTION
# =============================================================================

print("Please select the base folder ...\GRABMyoFlow_1.0\dynamic")
BASE_INPUT = filedialog.askdirectory(title="Select base folder containing session data")

if not BASE_INPUT:
    print("No input folder selected. Exiting Script!")
    sys.exit()

# --- Output Path Determination ---
parent_dir = os.path.dirname(BASE_INPUT.rstrip("\\/"))
OUTPUT_ROOT = os.path.join(parent_dir, f"dynamic_MATLAB")

print(f"\nSelected input: {BASE_INPUT}")
print(f"Output folder: {OUTPUT_ROOT}")


# =============================================================================
# STEP 2: OUTPUT FOLDER HANDLING (Overwrite check)
# =============================================================================
if not os.path.exists(OUTPUT_ROOT):
    os.makedirs(OUTPUT_ROOT)
    print(f"Created output folder: {OUTPUT_ROOT}")
else:
    while True:
        print(f"Found existing folder in: {OUTPUT_ROOT}")
        cont = input("Overwrite it (Y/N)? ").upper()
        if cont in ('Y', 'N'):
            if cont == 'Y':
                print("Overwriting...")
                shutil.rmtree(OUTPUT_ROOT)
                os.makedirs(OUTPUT_ROOT)
                break
            else:
                print("Exiting Script!")
                sys.exit()

session_paths = [os.path.join(BASE_INPUT, f"session{i}") for i in (1, 2, 3)]
session1_path = session_paths[0]

if not os.path.exists(session1_path):
    print(f"[ERROR] Session folder not found: {session1_path}")
    sys.exit()

nsub = len([d for d in os.listdir(session1_path) if os.path.isdir(os.path.join(session1_path, d))])


# =============================================================================
# STEP 3: MAIN CONVERSION LOOP
# =============================================================================
count = 0

for isession in range(1, NSESSION + 1):
    converted_folder = f"session{isession}"
    output_session_dir = os.path.join(OUTPUT_ROOT, converted_folder)
    os.makedirs(output_session_dir, exist_ok=True) 
    
    session_dir = os.path.join(BASE_INPUT, f"session{isession}")
    if not os.path.exists(session_dir):
        print(f"[WARN] Input session folder not found: {session_dir}")
        continue

    for isub in range(1, nsub + 1):
        foldername = f"session{isession}_participant{isub}"
        participant_dir = os.path.join(session_dir, foldername)

        if not os.path.exists(participant_dir):
            print(f"[WARN] Participant folder missing: {participant_dir}")
            continue

        matrices_wrist = np.empty((NTRIALS, NGEESTURE), dtype=object)

        for igesture in range(1, NGEESTURE + 1):
            for itrial in range(1, NTRIALS + 1):
                filename = f"session{isession}_participant{isub}_gesture{igesture}_trial{itrial}"
                filepath = os.path.join(participant_dir, filename)

                if not os.path.exists(filepath + ".dat"):
                    print(f"[WARN] Missing file: {filepath}.dat")
                    continue
                
                try:
                    record = wfdb.rdrecord(filepath)
                    data_emg = record.p_signal 

                    # Split 16ch (Wrist only) to 12ch.                        
                    if data_emg.shape[1] >= 16:
                        data_wrist = data_emg[:, EXTENSION_INDICES_KEEP]
                    else:
                        print(f"{filename}: Unexpected channel count ({data_emg.shape[1]}ch). Using all available channels.")
                        data_wrist = data_emg

                    matrices_wrist[itrial - 1, igesture - 1] = data_wrist

                except Exception as e:
                    print(f"[ERROR] Reading {filename}: {e}")

        count += 1
        print(f"[OK] Converted participant {isub} of session {isession} ({count} total)")

        # =============================================================================
        # STEP 4: SAVE OUTPUT
        # =============================================================================
        mat_filename = f"{foldername}.mat"
        try:
            savemat(os.path.join(output_session_dir, mat_filename),
                    {"DATA_WRIST": matrices_wrist})
        except Exception as e:
            print(f"[FATAL ERROR] Could not save {mat_filename}: {e}")

print("\nEnd of Script.")