% =========================================================================
% Feature Extraction Script for GRABMyoFlow (Wrist Data)
% =========================================================================
% This script processes the .mat files generated by the Python converter.
% It extracts features for the WRIST electrodes (using the first 6 channels).
%
% INPUT:
% - Folder generated by Python script (e.g., 'static_extension_MATLAB')
% - Structure: Root/SessionX/sessionX_participantY.mat
%
% OPTIONS:
% 1. Whole Set (W): IDs 1-63
% 2. Extension Set (E): IDs 44-63
% 3. Original Set (O): IDs 1-43
%
% OUTPUT:
% - Folder: 'Feature_Extracted_BM_[POSTFIX]'
% - Files: 'Wrist_Session1_[POSTFIX].mat', etc.
% =========================================================================

clear; clc;

% =========================================================================
% STEP 1: USER CONFIGURATION
% =========================================================================

disp('---------------------------------------------------------');
disp('   GRABMyoFlow Wrist Feature Extraction');
disp('---------------------------------------------------------');
disp('Which dataset do you want to process?');
disp(' [W] Whole Dataset     (IDs 1-63)');
disp(' [E] Extension Only    (IDs 44-63)');
disp(' [O] Original Only     (IDs 1-43)');
disp('---------------------------------------------------------');

choice = input('Enter selection (W/E/O): ', 's');

if strcmpi(choice, 'W')
    sub_list = 1:63;
    dataset_postfix = '_whole_set';
    default_title = 'Select the "static_whole_MATLAB" folder';
elseif strcmpi(choice, 'E')
    sub_list = 44:63;
    dataset_postfix = '_extension_set';
    default_title = 'Select the "static_extension_MATLAB" or the "static_whole_MATLAB" folder';
elseif strcmpi(choice, 'O')
    sub_list = 1:43;
    dataset_postfix = '_original_set';
    default_title = 'Select the "static_whole_MATLAB" or the original "Output_BM"(from grabmyo) folder';
else
    error('Invalid selection. Please run script again and type W, E, or O.');
end

% Select the input folder
fprintf('\nPlease select the folder containing the converted MAT files...\n');
base_input_folder = uigetdir(pwd, default_title);

if base_input_folder == 0
    error('No folder selected.');
end

% Constants
fs = 2048;              % Sampling frequency
NGESTURE = 16;          % Total number of gestures (Script loads 1:NGESTURE+1 for rest)
NTRIALS = 7;            % Total number of trials
NSESSION = 3;           % Number of sessions

% Define Output Folder with Postfix
output_dir = [pwd filesep 'Feature_Extracted_BM' dataset_postfix];
if ~exist(output_dir, 'dir')
    mkdir(output_dir);
    fprintf('Created output folder: %s\n', output_dir);
else
    fprintf('Output folder exists (%s). Overwriting content...\n', output_dir);
end

% =========================================================================
% STEP 2: LOAD AND FLATTEN DATA
% =========================================================================

for isession = 1:NSESSION
    CompleteSet = []; 
    current_session_folder = [base_input_folder filesep 'Session' num2str(isession)];

    if ~exist(current_session_folder, 'dir')
        warning(['Session folder not found: ' current_session_folder]);
        continue;
    end    

    fprintf('\nProcessing Session %d (%s)...\n', isession, dataset_postfix);   
    valid_sub_count = 0;
    
    for isub = sub_list
        fileName = ['session' num2str(isession) '_participant' num2str(isub) '.mat'];
        filePath = [current_session_folder filesep fileName];
        
        if ~exist(filePath, 'file')
            warning(['File missing: ' fileName]);
            continue;
        end
        
        temp_load = load(filePath, 'DATA_WRIST');
        datafile = temp_load.DATA_WRIST;
        a2 = []; 
        
        for igesture = 1:(NGESTURE + 1)
            a1 = []; 
            for itrial = 1:NTRIALS
                if isempty(datafile{itrial, igesture})
                    warning(['Empty data for Sub ' num2str(isub) ' Ges ' num2str(igesture)]);
                    continue;
                end
                a1 = [a1; datafile{itrial, igesture}];
            end
            a2 = [a2, {a1}];
        end
        
        CompleteSet = [CompleteSet; a2];
        valid_sub_count = valid_sub_count + 1;
        
        if mod(valid_sub_count, 10) == 0
            fprintf('Loaded %d subjects...\n', valid_sub_count);
        end
    end
    
    if isempty(CompleteSet)
        warning(['No data loaded for Session ' num2str(isession) '. Skipping.']);
        continue;
    end

    % =========================================================================
    % STEP 3: SEGMENTATION AND FEATURE EXTRACTION
    % =========================================================================
    
    fprintf('Extracting Features for Session %d...\n', isession);
    
    FeatSet = {};
    num_rows_processed = 0;
    
    for idx = 1:size(CompleteSet, 1)
        
        for igesture = 1:NGESTURE
            OneSet = CompleteSet{idx, igesture}'; 
            post_process = zeros(6, size(OneSet, 2));
            temp_6ch = OneSet(1:6, :); 
            common_avg = mean(temp_6ch, 1);
            for ich = 1:6
                post_process(ich, :) = OneSet(ich, :) - common_avg;
            end
            
            % --- SEGMENTATION ---
            % 0.2s window, 0.15s slide
            segData = segmentEMG(post_process', 0.2, 0.15, NTRIALS*5, fs, 1);
            
            % --- FEATURE EXTRACTION ---
            feat = featiDFTl(2048, 6, segData);
            FeatSet(idx, igesture) = {feat};
        end
        
        num_rows_processed = num_rows_processed + 1;
    end
    
    % Save Result with Postfix
    saveName = [output_dir filesep 'Wrist_Session' num2str(isession) dataset_postfix '.mat'];
    disp(['Saving: ' saveName]);
    save(saveName, 'FeatSet');
    
end

disp('Feature Extraction Complete.');