main.m

%% MAIN.M — Mind-Controlled Bionic Arm: EEG Classification Pipeline
% Master script for 8-class motor imagery EEG classification.
% Replaces the old YAMNet audio-based pipeline with EEG-native methods.
%
% Pipeline:
%   1. Load MILimbEEG dataset (60 subjects, 8 classes, 16 channels)
%   2. Preprocess EEG signals (bandpass, notch, CAR, artifact rejection)
%   3. Extract EEG features (band power, Hjorth, stats, PSD)
%   4. Train Classical ML models (SVM, RF, KNN) on features
%   5. Train EEGNet (EEG-native CNN) on raw preprocessed signals
%   6. Evaluate and compare all models
%
% Author: Dhruva Shaw et al.
% Project: Mind-Control Bionic Arm with Sense of Touch

clear; clc; close all;

%% ========================================================================
%  Configuration
%  ========================================================================

% Dataset path
dataDir = fullfile(pwd, 'datasets', 'datasourceDatasets', 'MILimbEEG', 'data');

% Sampling frequency
fs = 125;  % Hz (MILimbEEG dataset)

% Use motor-imagery trials (set to true for motor execution)
useMotorImagery = false;  % false = Imagery, which is what BCI needs

% Output directory for saved models
outputDir = fullfile(pwd, 'models');
if ~exist(outputDir, 'dir')
    mkdir(outputDir);
end

% Add source paths
addpath(genpath(fullfile(pwd, 'src')));
addpath(fullfile(pwd, 'Utility Functions'));

fprintf('==========================================================\n');
fprintf('  Mind-Controlled Bionic Arm — EEG Classification Pipeline\n');
fprintf('==========================================================\n');
fprintf('  Date:    %s\n', datestr(now));
fprintf('  MATLAB:  %s\n', version);
fprintf('  Dataset: MILimbEEG\n');
fprintf('  Mode:    %s\n', ternary(useMotorImagery, 'Motor Execution', 'Motor Imagery'));
fprintf('==========================================================\n\n');

%% ========================================================================
%  STEP 1: Load Dataset
%  ========================================================================
fprintf('\n%s\n', repmat('=', 1, 60));
fprintf('STEP 1: Loading MILimbEEG Dataset\n');
fprintf('%s\n\n', repmat('=', 1, 60));

tic;
[rawData, labels, subjectIds, taskNames] = load_milimbeeg(dataDir, useMotorImagery);
loadTime = toc;
fprintf('Dataset loaded in %.1f seconds.\n', loadTime);

%% ========================================================================
%  STEP 2: Preprocess EEG Signals
%  ========================================================================
fprintf('\n%s\n', repmat('=', 1, 60));
fprintf('STEP 2: Preprocessing EEG Signals\n');
fprintf('%s\n\n', repmat('=', 1, 60));

tic;
cleanData = eeg_preprocess(rawData, fs);
preprocTime = toc;
fprintf('Preprocessing completed in %.1f seconds.\n', preprocTime);

% Update labels and subjectIds to match valid trials
validIdx = ~cellfun(@isempty, cleanData);
fprintf('Valid trials after preprocessing: %d / %d\n', sum(validIdx), numel(cleanData));

%% ========================================================================
%  STEP 3: Extract Features (for Classical ML)
%  ========================================================================
fprintf('\n%s\n', repmat('=', 1, 60));
fprintf('STEP 3: Extracting EEG Features\n');
fprintf('%s\n\n', repmat('=', 1, 60));

tic;
[featureMatrix, featureNames] = extract_features(cleanData, fs);
featureTime = toc;
fprintf('Feature extraction completed in %.1f seconds.\n', featureTime);

%% ========================================================================
%  STEP 4: Train Classical ML Models
%  ========================================================================
fprintf('\n%s\n', repmat('=', 1, 60));
fprintf('STEP 4: Training Classical ML Models (SVM, RF, KNN)\n');
fprintf('%s\n\n', repmat('=', 1, 60));

tic;
classicalParams = struct();
classicalParams.numFolds = 5;
classicalParams.optimize = true;

[bestClassicalModel, classicalResults] = train_classical( ...
    featureMatrix, labels, classicalParams);
classicalTime = toc;
fprintf('Classical ML training completed in %.1f seconds.\n', classicalTime);

%% ========================================================================
%  STEP 5: Train EEGNet (Deep Learning)
%  ========================================================================
fprintf('\n%s\n', repmat('=', 1, 60));
fprintf('STEP 5: Training EEGNet (CPU Mode)\n');
fprintf('%s\n\n', repmat('=', 1, 60));

tic;
eegnetParams = struct();
eegnetParams.numFolds = 5;
eegnetParams.maxEpochs = 30;
eegnetParams.miniBatch = 64;
eegnetParams.learnRate = 1e-3;

[trainedEEGNet, eegnetResults] = train_eegnet( ...
    cleanData, labels, subjectIds, eegnetParams);
eegnetTime = toc;
fprintf('EEGNet training completed in %.1f seconds.\n', eegnetTime);

%% ========================================================================
%  STEP 6: Evaluate All Models
%  ========================================================================
fprintf('\n%s\n', repmat('=', 1, 60));
fprintf('STEP 6: Model Evaluation & Comparison\n');
fprintf('%s\n\n', repmat('=', 1, 60));

% Evaluate classical model
fprintf('\n--- %s Evaluation ---\n', classicalResults.bestModelName);
classicalEval = evaluate_model( ...
    labels(validIdx), ...
    classicalResults.predictions.(lower(classicalResults.bestModelName(1:3)))(validIdx), ...
    taskNames, ...
    classicalResults.bestModelName);

% Evaluate EEGNet
fprintf('\n--- EEGNet Evaluation ---\n');
validPredIdx = eegnetResults.predictions ~= '';
eegnetEval = evaluate_model( ...
    eegnetResults.trueLabels(validPredIdx), ...
    eegnetResults.predictions(validPredIdx), ...
    taskNames, ...
    'EEGNet');

%% ========================================================================
%  STEP 7: Final Comparison & Save
%  ========================================================================
fprintf('\n%s\n', repmat('=', 1, 60));
fprintf('FINAL COMPARISON\n');
fprintf('%s\n\n', repmat('=', 1, 60));

fprintf('  %-20s  %10s  %10s  %10s\n', 'Model', 'Accuracy', 'Macro F1', 'Kappa');
fprintf('  %s\n', repmat('-', 1, 56));
fprintf('  %-20s  %9.2f%%  %10.4f  %10.4f\n', ...
    classicalResults.bestModelName, ...
    classicalEval.accuracy * 100, classicalEval.macroF1, classicalEval.kappa);
fprintf('  %-20s  %9.2f%%  %10.4f  %10.4f\n', ...
    'EEGNet', ...
    eegnetEval.accuracy * 100, eegnetEval.macroF1, eegnetEval.kappa);

% Determine overall best
if eegnetEval.accuracy > classicalEval.accuracy
    bestOverall = 'EEGNet';
    bestNet = trainedEEGNet;
else
    bestOverall = classicalResults.bestModelName;
    bestNet = bestClassicalModel;
end
fprintf('\n  >> Best Model: %s <<\n', bestOverall);

% Save models and results
fprintf('\nSaving models and results...\n');
save(fullfile(outputDir, 'eegnet_model.mat'), 'trainedEEGNet', 'eegnetResults', 'eegnetEval');
save(fullfile(outputDir, 'classical_model.mat'), 'bestClassicalModel', 'classicalResults', 'classicalEval');
save(fullfile(outputDir, 'pipeline_results.mat'), ...
    'featureMatrix', 'featureNames', 'labels', 'subjectIds', 'taskNames', ...
    'classicalResults', 'eegnetResults');

fprintf('Models saved to: %s\n', outputDir);

% Timing summary
fprintf('\n=== Timing Summary ===\n');
fprintf('  Data Loading:     %.1f s\n', loadTime);
fprintf('  Preprocessing:    %.1f s\n', preprocTime);
fprintf('  Feature Extract:  %.1f s\n', featureTime);
fprintf('  Classical ML:     %.1f s\n', classicalTime);
fprintf('  EEGNet:           %.1f s\n', eegnetTime);
fprintf('  Total:            %.1f s (%.1f min)\n', ...
    loadTime + preprocTime + featureTime + classicalTime + eegnetTime, ...
    (loadTime + preprocTime + featureTime + classicalTime + eegnetTime) / 60);

fprintf('\n==========================================================\n');
fprintf('  Pipeline Complete! \n');
fprintf('==========================================================\n');


%% Helper function
function out = ternary(condition, trueVal, falseVal)
    if condition
        out = trueVal;
    else
        out = falseVal;
    end
end