%%% extact MUA - THINGS
% v.2-Gemini
% 2024 P. Papale fecit

clear all
addpath(genpath('\_code\code_utils_v2\'));

% constants
monkey = 'monkeyF';
datadir_gen = '\';
temp_folder = '\_temps\';
dates = {'20240112','20240115','20240116','20240118'};
instances = 1:2;
hubs = 1:2;
trial_length = .3;
pre_trial = .1;
post_trial = trial_length-pre_trial;
log_prefix = 'THINGS_';
plot_avg = 1;

%
downsample_factor = 30;
Fs = 30000;
Fbp=[500,5000];
N  = 2;    % filter order
Fn = Fs/2; % Nyquist frequency
Fl=200;
% FsD = Fs/downsample_factor;
% FnD = FsD/2; % Downsampled Nyquist frequency
[B, A] = butter(N, [min(Fbp)/Fn max(Fbp)/Fn]); % compute filter coefficients
[D, C] = butter(N,Fl/Fn,'low'); % compute filter coefficients

tb =-pre_trial*1000:1:(post_trial)*1000;
instance_chns = 1:256:256*4;
load('\\vs03\vs03-vandc-1\Latents\Passive_Fixation\monkeyN\_logs\1024chns_mapping_20220105.mat')
restarted = 0;
get_diode = 1;
for day = 1:length(dates)
    clear datadir all_folders blocks
    datadir = [datadir_gen,monkey,'\',dates{day},'\'];
    disp(['%%%%%% DAY: ',num2str(day),' ...']);
    
    all_folders = dir(datadir);
    all_folders = all_folders([all_folders(:).isdir]);
    all_folders = all_folders(~ismember({all_folders(:).name},{'.','..'}));
    blocks = length(all_folders);
    for b = 1:blocks
        block_number = str2num(all_folders(b).name(7:end));
        tic
        disp(['%%% Block: ',num2str(block_number),' ...']);
        
        % load and group the logs
        % MAT files are: [#trial #train_pic #test_pic #pic_rep #ncount #correct]
        % ALLMAT files will be: [#trial #train_pic #test_pic #pic_rep #ncount #block #day]
        clear MAT temp_MAT
        load([datadir_gen,monkey,'\_logs\',log_prefix,monkey,'_',dates{day},'_B',num2str(block_number),'.mat'],'MAT')
        temp_MAT = MAT(MAT(:,end)>0,1:5);
        temp_MAT = [temp_MAT ones([length(temp_MAT) 1])*block_number];
        temp_MAT = [temp_MAT ones([length(temp_MAT) 1])*day];
        save([datadir_gen,monkey,temp_folder,'TEMP_MAT_',dates{day},'_B',num2str(block_number),'.mat'],'temp_MAT');
        
        for instance_idx = instances
            clear instance_name
            instance_name = ['-instance',num2str(instance_idx),'_B',sprintf('%03.0f', block_number)];
            for hub_idx = hubs
                clear partition_idx partition_name
                partition_idx = hub_idx+2*(instance_idx-1);
                partition_name = ['partition',num2str(partition_idx),'_B',sprintf('%03.0f', block_number)];
                if ~isfile([datadir_gen,monkey,temp_folder,'TEMP_MUA_',dates{day},'_',partition_name,'.mat'])
                    if hub_idx == 1 || restarted
                        get_diode = 1;
                    else
                        get_diode = 0;
                    end
                    
                    % get the difference between the diode and the stimulus bit
                    if get_diode == 1
                        clear instance_NS6_filename RAW 
                        instance_NS6_filename = [datadir,'Block_',num2str(block_number),'\NSP',instance_name,'.ns6'];
                        RAW = openNSx(instance_NS6_filename);
                        
                        if RAW.RawData.PausedFile==1
                            RAW.Data = cell2mat(RAW.Data);
                        end

                        clear temp temp_trials idx_temp Trials
                        temp = smooth(double(RAW.Data(1,:)),100);
                        temp_trials = 1+find(diff(temp)>300);
                        idx_temp = [1; diff(temp_trials)>Fs/200];
                        Trials = temp_trials(idx_temp>0);
                        
                        assert(length(Trials)==length(MAT),'There are (%d) recorded trials while (%d) trails were logged',length(Trials),length(MAT));

                        clear diode dio_onset
                        diode = RAW.Data(15,:);
                        for trl=1:length(Trials)
                            clear diode_trl
                            temp_int = Trials(trl):Trials(trl)+Fs*trial_length-1;
                            diode_trl=diode(temp_int);
                            f=find(abs(diode_trl)>1*10^4,1,'first');
                            if isempty(f)
                                dio_onset(trl) = 0;
                            else
                                dio_onset(trl)= f;
                            end
                        end
                        restarted = 0;
                    end

                    % load the raw data
                    clear instance_NS6_filename RAW temp_block_mua
                    instance_NS6_filename = [datadir,'Block_',num2str(block_number),'\Hub',num2str(hub_idx),instance_name,'.ns6'];
                    RAW = openNSx(instance_NS6_filename);
                    
                    if RAW.RawData.PausedFile==1
                        RAW.Data = cell2mat(RAW.Data);
                    end
                    
                    temp_block_mua = nan([256 length(temp_MAT) length(tb)]);
                    for chn = 1:256
                        clear S dum1 dum2
                        S=double(squeeze(RAW.Data(chn,:)))';

                        % bandpass
                        dum1 = filtfilt(B, A, S); % apply filter to the data
                        % rectify
                        dum2 = abs(dum1);
                        % low-pass
                        clear muafilt
                        muafilt = filtfilt(D, C, dum2);
                        
                        % remove 50Hz line noise (+ harmonics)
                        buttLoop = muafilt;
                        for i = 1:5
                            d = designfilt('bandstopiir','FilterOrder',2, ...
                                'HalfPowerFrequency1',50*i-5,'HalfPowerFrequency2',50*i+5, ...
                                'DesignMethod','butter','SampleRate',Fs);
                            buttLoop = filtfilt(d,buttLoop);
                        end
                        
                        % remove 60Hz screen refresh rate (+ harmonics)
                        for i = 1:5
                            d = designfilt('bandstopiir','FilterOrder',2, ...
                                'HalfPowerFrequency1',60*i-5,'HalfPowerFrequency2',60*i+5, ...
                                'DesignMethod','butter','SampleRate',Fs);
                            buttLoop = filtfilt(d,buttLoop);
                        end
                        muafilt = buttLoop';
                        
                        % select only correct trials and downsample
                        clear  temp_chn
                        temp_chn = nan([length(temp_MAT) length(tb)]);
                        z = 1;
                        for trl=1:length(Trials)
                            if MAT(trl,end)>0
                                clear temp_int temp_data
                                temp_int = Trials(trl)+dio_onset(trl)-Fs*pre_trial:Trials(trl)+dio_onset(trl)+Fs*post_trial-1;
                                temp_chn(z,:) = downsample(muafilt(temp_int),downsample_factor);
                                z=z+1;
                            end
                        end
                        temp_block_mua(chn,:,:) = temp_chn;
                    end
                    if plot_avg
                        plot(squeeze(nanmean(temp_block_mua,[1,2])));
                    end
                    
                    clear mapping_inst
                    mapping_inst = mapping(instance_chns(instance_idx):instance_chns(instance_idx)+256-1);
                    mapping_inst = mapping_inst-((instance_idx-1)*256);
                    temp_block_mua = temp_block_mua(mapping_inst,:,:);
                    
                    % store temps
                    save([datadir_gen,monkey,temp_folder,'TEMP_MUA_',dates{day},'_',partition_name,'.mat'],'temp_block_mua','-v7.3');
                    pause(5)
                    disp(['% (Partiton: ',num2str(partition_name),' done!)']);
                else
                    restarted = 1;
                end
                disp(['%%% ... done in: ',num2str(round(toc)),'s'])            
            end
        end
    end
end