clear all
clc

%% DATA LOADING 

behavior = load('D:\Lorenzo\Nghia\data\NN11\210626_behavior_data\behavior.mat');

planes = 3;
fs = 10.42;


% Initialize the activity matrix
activity_F = [];
activity_Fneu = [];

% Loop over each plane
for plane = 1:planes

    Fneu_path = sprintf('D:\\Lorenzo\\Nghia\\data\\NN11\\210626_plane_%d\\Fneu.npy', plane);
    iscell_path = sprintf('D:\\Lorenzo\\Nghia\\data\\NN11\\210626_plane_%d\\iscell.mat', plane);
    F_path = sprintf('D:\\Lorenzo\\Nghia\\data\\NN11\\210626_plane_%d\\F.mat', plane);

    % Load the iscell file and take the first column.
    % It is assumed that iscell.mat contains a variable named 'iscell'.
    tmp = load(iscell_path).data;
    accepted_cells = tmp(:,1);
    
    % Load fluorescence data from F.mat and sum each row.
    tmp = double(load(F_path).data);
    F = tmp;
    Fneu = double(readNPY(Fneu_path));

    fluorescence = sum(F, 2);  % Summing across columns
    
    % Set accepted_cells to zero if fluorescence is zero
    accepted_cells(fluorescence == 0) = 0;
    
    % Convert accepted_cells to a logical vector (true if equal to 1).
    accepted_cells = (accepted_cells == 1);

    % Select only the accepted cells' activity data
    if plane == 1
        activity_F = F(accepted_cells, :);
        activity_Fneu = Fneu(accepted_cells, :);
    else

        to_delete_path = sprintf('D:\\Lorenzo\\Nghia\\data\\NN11\\210626_plane_%d\\to_delete_plane_%d.mat', plane, plane);

        to_delete = double(load(to_delete_path).data);

        activity_plane_F = F(accepted_cells, :);
        activity_plane_Fneu = Fneu(accepted_cells, :);

        % Keep only the rows of activity_plane for which to_delete equals 1.
        activity_plane_F = activity_plane_F(to_delete == 1, :);
        activity_plane_Fneu = activity_plane_Fneu(to_delete == 1, :);

        % Concatenate the activity data from all planes vertically.
        activity_F = [activity_F; activity_plane_F];
        activity_Fneu = [activity_Fneu; activity_plane_Fneu];

    end
end



clear F F_path activity_plane_F activity_plane_Fneu Fneu_path Fneu accepted_cells fluorescence iscell_path plane planes tmp to_delete to_delete_path

%% Getting CueIndexes and Baseline last frame

[cue1_idx, cue2_idx] = get_cueIndexes(behavior, 0, 0.5);

%% Correcting corrupted frames

activity_F(:,1) = [];
activity_Fneu(:,1) = []; %First frame orrible

cue1_idx = cue1_idx - 1;
cue2_idx = cue2_idx - 1; %Adjust indexes

bad_frames = [21333, 42666, 63999, 85332]; %Bad frames - 1
                                                                                   
for k = 1:4
    activity_F(:,bad_frames(k)) = (activity_F(:,bad_frames(k)+1) + activity_F(:,bad_frames(k)-1)) / 2;
    activity_Fneu(:,bad_frames(k)) = (activity_Fneu(:,bad_frames(k)+1) + activity_Fneu(:,bad_frames(k)-1)) / 2;
end

activity_F(:,1:49) = []; 
activity_Fneu(:,1:49) = []; %Discard also first 5 seconds of baseline for artifacts.

cue1_idx = cue1_idx - 49;
cue2_idx = cue2_idx - 49;

baseline_ends = min(cue1_idx(1,1),cue2_idx(1,1)) - 20;

%% Correct for neuropil

[activity_Fcorr, alpha_vec, corr_vec, energy_preserved, info] = ...
    correct_neuropil_dc_safe(activity_F, activity_Fneu, baseline_ends, ...
        'amax', 10, 'UseParallel', true);

%% Outlier detection \w Wavelet Decomposition 

%[Efrac, Eabs, Col_labels] = wavelet_component_energy_db3_parallel(activity_Fcorr, 16, 'energy', true);
load('C:\Users\Lorenzo\Documents\content\code\Fluorescence_Traces\data\Efrac.mat');
[thr_pct, keep_prop, keep_mask] = sweep_wavelet_energy_threshold(Efrac, 1:4, 5);
activity_Fsave = activity_Fcorr;
activity_Fcorr = activity_Fcorr(keep_mask(:,11), :);


%% Computing DF/F

[DF_F, F0safe, F0] = dff_percentile(activity_Fcorr, 10.42, 'q',20, 'win_s',60, 'sig_s',10);
figure; hist(log(mean(abs(DF_F), 2)), 1000);

%% DF_F STANDARDIZED
%DF_F_std = (DF_F - mean(DF_F,2)) ./ std(DF_F,0,2);

%% DF_F IN [0,1]

A_min = min(DF_F, [], 2);        % row-wise min (m x 1)
A_max = max(DF_F, [], 2);        % row-wise max (m x 1)

DF_F_mm = (DF_F - A_min) ./ (A_max - A_min);  % broadcasting row-wise

%% Framing ACTIVITY TENSOR AND CUES LABLES

N_neurons = size(DF_F_mm,1);
N_time    = 27;
nCues1    = size(cue1_idx,2);
nCues2    = size(cue2_idx,2);
N_trials  = nCues1 + nCues2;

activity_tensor = nan(N_neurons, N_time, N_trials);
labels = nan(1, N_trials);

j1 = 1; j2 = 1; i = 1;

while j1 <= nCues1 || j2 <= nCues2
    if j2 > nCues2 || (j1 <= nCues1 && cue1_idx(1,j1) < cue2_idx(1,j2))
        % take cue1
        idx = cue1_idx(1,j1):cue1_idx(2,j1);
        labels(i) = 1;
        j1 = j1+1;
    else
        % take cue2
        idx = cue2_idx(1,j2):cue2_idx(2,j2);
        labels(i) = 2;
        j2 = j2+1;
    end
    
    activity_tensor(:,:,i) = DF_F_mm(:, idx);
    i = i+1;
end

clear idx j1 j2 i k nCues1 nCues2 num_trials N_trials N_neurons N_time behavior
%% By Trials Normalization

n_trials = length(activity_tensor(1,1,:));
n_time = length(activity_tensor(1,:,1));


for i = 1:n_trials
        activity_tensor(:,:,i) = activity_tensor(:,:,i)/mean(activity_tensor(:,:,i),'all');
end