%% paths 

addpath /home/lau/matlab/fieldtrip/
addpath /home/lau/matlab/fieldtrip_private_functions/
ft_defaults

subject_path = '/home/lau/analyses/sef_data_second_benchmarking/to_be_uploaded_to_donders/data/MEG/subject_0001/recording_1/';
save_path = '/home/lau/analyses/sef_data_second_benchmarking/test_scripts_for_upload/';

recordings = {
    {'E1-raw' 'E1-raw-1'}
    {'E2-raw' 'E2-raw-1'}
    };

position_names = {'A1' 'A3' 'A6' 'C1' 'C3' 'C6' 'B1' 'B2' 'B3' 'B4'};

%% preprocess data

n_recordings = length(recordings);
preprocessed_recordings = cell(1, n_recordings);
variable_as_string = get_variable_name_as_string(preprocessed_recordings);
overwrite = true;
output_name = [save_path 'data_low_tc/' variable_as_string '.mat'];
lowpass_frequency = 300;
dft_frequency = 50;

if ~exist(output_name, 'file') || overwrite

    for recording_index = 1:n_recordings

        recording = recordings{recording_index};
        n_split_recordings = length(recording);
        split_recordings = cell(1, n_split_recordings);

        for split_recording_index = 1:n_split_recordings

            data_path = [subject_path '1_raw/' recording{split_recording_index} '.fif'];

            cfg = [];
            cfg.dataset = data_path;
            cfg.trialdef.eventtype = 'STI102';
            cfg.trialdef.eventvalue = 128;
            cfg.trialdef.prestim = 0.1;
            cfg.trialdef.poststim = 0.2;

            cfg = ft_definetrial(cfg);

            cfg.baseline = [-0.1 0];
            cfg.demean = 'yes';
            cfg.lpfilter = 'yes';
            cfg.dftfilter = 'yes';
            cfg.lpfreq = lowpass_frequency;
            cfg.dftfreq = dft_frequency:dft_frequency:lowpass_frequency;
            cfg.coilaccuracy = 2;

            split_recordings{split_recording_index} = ft_preprocessing(cfg);

        end

        cfg = [];

        preprocessed_recordings{recording_index} = ft_appenddata(cfg, split_recordings{1:end});

    end
    
    disp(['Saving: ' output_name])
    save(output_name, variable_as_string, '-v7.3');
    
end

%% load preprocessed recordings

load([save_path, 'data_low_tc/preprocessed_recordings'])

%% eog clean data threshold (250 microvolt)

n_recordings = length(preprocessed_recordings);
eog_preprocessed_recordings = cell(1, n_recordings);
variable_as_string = get_variable_name_as_string(eog_preprocessed_recordings);
output_name = [save_path 'data_low_tc/' variable_as_string '.mat'];
overwrite = true;
eog_threshold = 250e-6; % V

if ~exist(output_name, 'file') || overwrite

    for recording_index = 1:n_recordings

        preprocessed_recording = preprocessed_recordings{recording_index};
        n_trials = length(preprocessed_recording.trial);

        cfg = [];
        cfg.channel = 'EOG';

        eog_data = ft_selectdata(cfg, preprocessed_recording);
        trial_indices = ones(1, n_trials);

        for trial_index = 1:n_trials

            current_trial = eog_data.trial{trial_index};
            max_eog_value = max(max(current_trial));

            if max_eog_value > eog_threshold

                trial_indices(trial_index) = 0;
                disp(['Trial no. ' num2str(trial_index) ' is rejected due to EOG value of: ' num2str(max_eog_value) ...
                    ' V, which is higher than the allowed threshold: ' num2str(eog_threshold) ' V'])

            end

        end

        trial_indices = logical(trial_indices);

        cfg = [];
        cfg.channel = 'MEG';
        cfg.trials = trial_indices;

        eog_preprocessed_recordings{recording_index} = ft_selectdata(cfg, preprocessed_recording);


    end
    
    disp(['Saving: ' output_name])
    save(output_name, variable_as_string, '-v7.3');
    
end

%% load eog preprocessed recordings

load([save_path, 'data_low_tc/eog_preprocessed_recordings'])
        
%% clean data manually

n_recordings = length(eog_preprocessed_recordings);
cleaned_recordings = cell(1, n_recordings);
variable_as_string = get_variable_name_as_string(cleaned_recordings);
output_name = [save_path 'data_low_tc/' variable_as_string '.mat'];
overwrite = true;

if ~exist(output_name, 'file') || overwrite

    for recording_index = 1:n_recordings


        cfg = [];
        cfg.method = 'summary';
        cfg.keepchannel = 'yes';
        cfg.channel = 'MEGMAG';
        cfg.layout = 'neuromag306all.lay';
        
        cleaned_mag = ft_rejectvisual(cfg, eog_preprocessed_recordings{recording_index});
        
        cfg.channel = 'MEGGRAD';
        
        cleaned_recordings{recording_index} = ft_rejectvisual(cfg, cleaned_mag);

    end
    
    disp(['Saving: ' output_name])
    save(output_name, variable_as_string, '-v7.3');
    
end

%% load cleaned recordings

load([save_path, 'data_low_tc/cleaned_recordings'])

%% timelock data

n_recordings = length(cleaned_recordings);
timelockeds = cell(1, n_recordings);
variable_as_string = get_variable_name_as_string(timelockeds);
output_name = [save_path 'data_low_tc/' variable_as_string '.mat'];
overwrite = true;

if ~exist(output_name, 'file') || overwrite

    for recording_index = 1:n_recordings

        cleaned_recording = cleaned_recordings{recording_index};

        cfg = [];
        cfg.covariance = 'yes';
        cfg.covariancewindow = 'prestim';
 
        timelockeds{recording_index} = ft_timelockanalysis(cfg, cleaned_recording);
        
    end
    
    disp(['Saving: ' output_name])
    save(output_name, variable_as_string, '-v7.3');
    
end

%% timelock data magnetometer

n_recordings = length(cleaned_recordings);
timelockeds_mag = cell(1, n_recordings);
variable_as_string = get_variable_name_as_string(timelockeds_mag);
output_name = [save_path 'data_low_tc/' variable_as_string '.mat'];
overwrite = true;
lowpass_frequency = 300;
dft_frequency = 50;

if ~exist(output_name, 'file') || overwrite

    for recording_index = 1:n_recordings

        cleaned_recording = cleaned_recordings{recording_index};

        cfg = [];
        cfg.covariance = 'yes';
        cfg.covariancewindow = 'prestim';
        cfg.channel = 'megmag';
 
        timelockeds_mag{recording_index} = ft_timelockanalysis(cfg, cleaned_recording);
        
    end
    
    disp(['Saving: ' output_name])
    save(output_name, variable_as_string, '-v7.3');
    
end


%% get headmodels

load([save_path, 'mris/mri_segmented'])

cfg = [];
cfg.method = 'projectmesh';
cfg.tissue = 'brain';
cfg.numvertices = 3000;

mesh_brain = ft_prepare_mesh(cfg, mri_segmented);

headmodel_names = {'singlesphere' 'singleshell'};
n_headmodels = length(headmodel_names);

headmodels = cell(1, n_headmodels);

for headmodel_index = 1:n_headmodels
    
    cfg = [];
    cfg.method = headmodel_names{headmodel_index};
    
    headmodel = ft_prepare_headmodel(cfg, mesh_brain);
    headmodel = ft_convert_units(headmodel, 'm');
    
    headmodels{headmodel_index} = headmodel;
    
end

save([save_path 'headmodels/headmodels'], 'headmodels');

%% load headmodels and averages

load([save_path 'headmodels/headmodels'])
load([save_path 'data_low_tc/timelockeds'])

%% fit dipoles

latency = [0.015 0.025];

headmodel_names = {'singlesphere' 'singleshell'};
n_headmodels = length(headmodel_names);

sensor_groups = {'megmag' 'meggrad' 'all'};
n_sensor_groups = length(sensor_groups);

recordings = {1 2};
n_recordings = length(recordings);

reduce_ranks = {2 'no'};
n_ranks = length(reduce_ranks);

dipoles = cell(2, n_headmodels * n_sensor_groups * n_recordings * n_ranks);

variable_as_string = get_variable_name_as_string(dipoles);
output_name = [save_path 'data_low_tc/' variable_as_string '.mat'];
overwrite = true;

if ~exist(output_name, 'file') || overwrite

    dipole_index = 1;

    for recording_index = 1:n_recordings

        for headmodel_index = 1:n_headmodels

            headmodel = headmodels{headmodel_index};
            headmodel = ft_convert_units(headmodel, 'm');

            for sensor_group_index = 1:n_sensor_groups

                cfg = [];
                cfg.channel = sensor_groups{sensor_group_index};

                timelockeds_subset = ft_selectdata(cfg, timelockeds{recording_index});
                timelockeds_subset.grad = ft_convert_units(timelockeds_subset.grad, 'm');

                for rank_index = 1:n_ranks

                    reduce_rank = reduce_ranks{rank_index};
                    
                    if strcmp(headmodel.type, 'singlesphere')
                    
                        cfg = [];
                        cfg.headmodel = headmodel;
                        cfg.senstype = 'meg';

                        grid = ft_prepare_leadfield(cfg, timelockeds_subset);
                        
                    end

                    cfg = [];
                    cfg.latency = latency;
                    cfg.numdipoles = 1;
                    cfg.gridsearch = 'yes';
                    
                    if strcmp(headmodel.type, 'singlesphere')
                    
                        cfg.grid = grid;
                        
                    else
                        
                        cfg.grid.unit = 'm';
                        cfg.grid.resolution = 0.01;
                        
                    end
                        
                    cfg.headmodel = headmodel;
                    cfg.dipfit.metric = 'rv';
                    cfg.dipfit.noisecov = timelockeds_subset.cov;
                    cfg.model = 'moving';
                    cfg.senstype = 'meg';
                    cfg.reducerank = reduce_rank;
                    cfg.meanpos = 'yes';
                    cfg.sequentialguess = 'no';

                    dipoles{1, dipole_index} = ft_dipolefitting(cfg, timelockeds_subset);

                    % get name
                    if recording_index == 1

                        recording_name = 'before';

                    else

                        recording_name = 'after';

                    end


                    headmodel_name = headmodel.type;

                    sensor_group_name = sensor_groups{sensor_group_index};

                    if reduce_rank == 2

                        rank_name = 'reduced_rank';

                    else

                        rank_name = 'full_rank';

                    end

                    dipoles{2, dipole_index} = ['dipoles_' recording_name '_' headmodel_name '_' sensor_group_name '_' rank_name];
                    
                    dipole_index = dipole_index + 1;

                end

            end

        end

    end
                
    disp(['Saving: ' output_name])
    save(output_name, variable_as_string, '-v7.3');             

end



%% get headmodel with fine scalp (10000 vertices)

load([save_path, 'mris/mri_segmented'])

cfg = [];
cfg.method = 'projectmesh';
cfg.tissue = 'brain';
cfg.numvertices = 3000;

mesh_brain = ft_prepare_mesh(cfg, mri_segmented);

cfg.tissue = 'skull';
cfg.numvertices = 2000;

mesh_skull = ft_prepare_mesh(cfg, mri_segmented);

cfg.tissue = 'scalp';
cfg.numvertices = 10000;

mesh_scalp = ft_prepare_mesh(cfg, mri_segmented);

cfg = [];
cfg.method = 'bemcp';
cfg.conductivity = [1 1/20 1] .* 0.33;

mesh_eeg = [mesh_brain mesh_skull mesh_scalp];

headmodel_eeg = ft_prepare_headmodel(cfg, mesh_eeg);

save([save_path 'headmodels/headmodel_eeg'], 'headmodel_eeg');


%% create artificial magnetometers

load([save_path 'headmodels/headmodel_eeg'])
load([save_path 'data_low_tc/timelockeds'])

headmodel_eeg = ft_convert_units(headmodel_eeg, 'm');
skin_vertices = headmodel_eeg.bnd(1, 3).pos;

avg_artificial = timelockeds{1};
avg_artificial.grad = ft_convert_units(avg_artificial.grad, 'm');

avg_artificial.grad.coilpos = skin_vertices;
avg_artificial.grad.coilori = normals(headmodel_eeg.bnd(1, 3).pos, ...
    headmodel_eeg.bnd(1, 3).tri);

n_sensors = length(skin_vertices);
chan_labels = cell(n_sensors, 1);
chan_type = cell(n_sensors, 1);
chan_unit = cell(n_sensors, 1);

for sensor_index = 1:n_sensors
    
    chan_labels{sensor_index} = strjoin({'MEG', num2str(sensor_index)}, '');
    chan_type{sensor_index} = 'megmag';
    chan_unit{sensor_index} = 'T';
    
end

avg_artificial.grad.label = chan_labels;
avg_artificial.grad.chantype = chan_type;
avg_artificial.grad.chanunit = chan_unit;

avg_artificial.grad.tra = eye(n_sensors); %% identity matrix

avg_artificial.grad.chanori = avg_artificial.grad.coilori;
avg_artificial.grad.chanpos = skin_vertices;

save([save_path 'data_low_tc/avg_artificial'], 'avg_artificial');


%% create forward models (load)

load([save_path 'headmodels/headmodels'])
load([save_path 'data_low_tc/dipoles'])
load([save_path 'data_low_tc/avg_artificial'])

%% create forward models 

n_dipoles = length(dipoles);
forwards = cell(1, n_dipoles);
variable_as_string = get_variable_name_as_string(forwards);
output_name = [save_path 'forwards/' variable_as_string '.mat'];
overwrite = false;

if ~exist(output_name, 'file') || overwrite

    for dipole_index = 1:n_dipoles

        dipole = dipoles{1, dipole_index};
        rv = [dipole.dip.rv];
        min_rv = min(rv);
        min_index = find(rv == min_rv);
        time_min = dipole.time(min_index);
        
        % make forward model based on that dipole (scalar strength)
        
        if isfield(dipole.cfg.grid, 'pos')
            
            headmodel_index = 1;
            
        else
            
            headmodel_index = 2;
            
        end

        headmodel = headmodels{headmodel_index};

        cfg = [];
        cfg.channel = 'all';
        cfg.grid.pos = dipole.dip(min_index).pos;
        cfg.grid.unit = dipole.dip(min_index).unit;
        cfg.headmodel = headmodel;
        cfg.grad = avg_artificial.grad;
        cfg.senstype = 'MEG';

        forward = ft_prepare_leadfield(cfg);

        forwards{dipole_index} = forward;
    
    end
    
    disp(['Saving: ' output_name])
    save(output_name, variable_as_string, '-v7.3');   

end


%% make forward model based on three orientations

load([save_path 'headmodels/headmodels'])

grad_X = avg_artificial.grad;
grad_X.coilori(:, :) = 0; grad_X.coilori(:, 1) = 1;

grad_Y = avg_artificial.grad;
grad_Y.coilori(:, :) = 0; grad_Y.coilori(:, 2) = 1;

grad_Z = avg_artificial.grad;
grad_Z.coilori(:, :) = 0; grad_Z.coilori(:, 3) = 1;

forwards_XYZ = cell(3, n_dipoles);
n_dipoles = length(dipoles);
variable_as_string = get_variable_name_as_string(forwards_XYZ);
output_name = [save_path 'forwards/' variable_as_string '.mat'];
overwrite = false;

if ~exist(output_name, 'file') || overwrite

    for dipole_index = 1:n_dipoles

        dipole = dipoles{1, dipole_index};
        rv = [dipole.dip.rv];
        min_rv = min(rv);
        min_index = find(rv == min_rv);
        time_min = dipole.time(min_index);

        if isfield(dipole.cfg.grid, 'pos')

                headmodel_index = 1;

        else

            headmodel_index = 2;

        end

        headmodel = headmodels{headmodel_index};

        cfg = [];
        cfg.channel = 'all';
        cfg.grid.pos = dipole.dip(min_index).pos;
        cfg.grid.unit = dipole.dip(min_index).unit;
        cfg.headmodel = ft_convert_units(headmodel, 'm');
        cfg.senstype = 'MEG';

        cfg.grad = grad_X;
        forward_X = ft_prepare_leadfield(cfg);
        forwards_XYZ{1, dipole_index} = forward_X;

        cfg.grad = grad_Y;
        forward_Y = ft_prepare_leadfield(cfg);
        forwards_XYZ{2, dipole_index} = forward_Y;

        cfg.grad = grad_Z;
        forward_Z = ft_prepare_leadfield(cfg);
        forwards_XYZ{3, dipole_index} = forward_Z;
    end
    
    disp(['Saving: ' output_name])
    save(output_name, variable_as_string, '-v7.3');   
    
end

%% predicted field vector

n_dipoles = length(dipoles);
predicted_fields = cell(1, n_dipoles);
variable_as_string = get_variable_name_as_string(predicted_fields);
output_name = [save_path 'forwards/' variable_as_string '.mat'];
overwrite = true;

if ~exist(output_name, 'file') || overwrite

    for dipole_index = 1:length(predicted_fields)

        dipole = dipoles{1, dipole_index};
        predicted_field = cell(1, length(dipole.time));

        for time = dipole.time

            dipole_time_index = find(dipole.time == time);
            forward_X = forwards_XYZ{1, dipole_index};
            forward_Y = forwards_XYZ{2, dipole_index};
            forward_Z = forwards_XYZ{3, dipole_index};
            predicted_field{dipole_time_index} = ...
                [forward_X.leadfield{1} * dipole.dip(dipole_time_index).mom ...
                forward_Y.leadfield{1} * dipole.dip(dipole_time_index).mom ...
                forward_Z.leadfield{1} * dipole.dip(dipole_time_index).mom];

        end

        predicted_fields{dipole_index} = predicted_field;

    end
    
    disp(['Saving: ' output_name])
    save(output_name, variable_as_string, '-v7.3'); 
    
end

%% forward models helmet

load([save_path 'headmodels/headmodels'])
load([save_path 'data_low_tc/dipoles'])

recordings = {
    {'E1-raw' 'E1-raw-1'}
    {'E2-raw' 'E2-raw-1'}
    };

sens_elekta = ft_read_sens([subject_path '1_raw/' recordings{1}{1} '.fif'], 'senstype', 'meg');
% sens_elekta = ft_convert_units(sens_elekta, 'm');

n_dipoles = length(dipoles);
forwards_helmet = cell(1, n_dipoles);
variable_as_string = get_variable_name_as_string(forwards_helmet);
output_name = [save_path 'forwards/' variable_as_string '.mat'];
overwrite = true;

if ~exist(output_name, 'file') || overwrite

    for dipole_index = 1:n_dipoles
       
        dipole = dipoles{1, dipole_index};
        rv = [dipole.dip.rv];
        min_rv = min(rv);
        min_index = find(rv == min_rv);
        time_min = dipole.time(min_index);

        % make forward model based on that dipole (scalar strength)
        if isfield(dipole.cfg.grid, 'pos')

            headmodel_index = 1;

        else

            headmodel_index = 2;

        end

        headmodel = headmodels{headmodel_index};

        cfg = [];
        cfg.channel = 'all';
        cfg.grid.pos = dipole.dip(min_index).pos;
        cfg.grid.unit = dipole.dip(min_index).unit;
        cfg.headmodel = headmodel;
        cfg.grad = sens_elekta;
        cfg.senstype = 'MEG';

        forward_low_tc = ft_prepare_leadfield(cfg);
        
        forwards_helmet{dipole_index} = forward_low_tc;
        
    end

    disp(['Saving: ' output_name])
    save(output_name, variable_as_string, '-v7.3');     

end

%% get bootstrapped low-tc values

load([save_path, 'data_low_tc/cleaned_recordings'])

n_recordings = length(recordings);
bootstrapped_std = cell(1, n_recordings);
variable_as_string = get_variable_name_as_string(bootstrapped_std);
output_name = [save_path 'data_low_tc/' variable_as_string '.mat'];
overwrite = true;
channels_to_plot{1} = {'MEG1311' 'MEG1311' 'MEG1311' 'MEG2241' 'MEG2211' 'MEG2241' 'MEG2211' 'MEG2211' 'MEG1131' 'MEG1131'}; %% chosen manually

if ~exist(output_name, 'file') || overwrite

    for recording_index = 1:n_recordings

        data_clean = cleaned_recordings{recording_index};
        data_clean.grad = ft_convert_units(data_clean.grad, 'm');
        
        cfg = [];
        cfg.lpfilter = 'yes';
        cfg.lpfreq = 300;
        cfg.dftfilter = 'yes';
        cfg.dftfreq = 50:50:300;
        
        data_clean = ft_preprocessing(cfg, data_clean);
        
        n_trials = length(data_clean.trial);
        n_samples = length(data_clean.trial{1});
        epochs = zeros(n_trials, n_samples);

        all_channels = unique([channels_to_plot{1}]);
        std_dev_cell = cell(2, length(all_channels));

        for channel = 1:length(all_channels)

                boot_index = strcmp(data_clean.label, all_channels(channel));

                for n_trial = 1:n_trials

                    epochs(n_trial, :) = data_clean.trial{n_trial}(boot_index, :);

                end

                [evoked_boot, std_dev, ~, ~] = bootstrap_epochs({epochs}, 1000);
                std_dev_cell{1, channel} = all_channels{channel};
                std_dev_cell{2, channel} = std_dev;

        end

        bootstrapped_std{recording_index} = std_dev_cell;

    end
    
        disp(['Saving: ' output_name])
        save(output_name, variable_as_string, '-v7.3');     
        
end

%% get matlab files based on fif files

for position_name_index = 1:length(position_names)
    
    filepath = [subject_path '4_cleaned_epochs/' position_names{position_name_index} '_manually_cleaned-epo'];
    
    data = fiff_read_epochs([filepath '.fif']);
    high_tc_data = data.data(:, 1, :);
    save([filepath '.mat'], 'high_tc_data')
    
end

%% bootstrap hightc data

epoch = cell(1,10);

for position_name_index = 1:length(position_names)
    
    filepath = [subject_path '4_cleaned_epochs/' position_names{position_name_index} '_manually_cleaned-epo.mat'];
    S = load(filepath);
    epoch{position_name_index} = S.high_tc_data; %% from load(fname)
end
[evoked_boot, std_dev, evoked_orig, epoch] = bootstrap_epochs(epoch, 1000);
save([save_path 'data_high_tc/std_dev'], 'std_dev')