% routine to show ensemble LFP relation to generate fig6,8 and Sfig4, requires the results from
% Ensemble_detection routine, LFP_calculations
% --------------------------------------------
% parameters:
% K_opt: # of clusters
% distance_opt: distance to use for clustering

% important outputs: 
% idx_clus: clustering index for all the ensembles
% idx_clus_synch: clustering index for all the pairs with the significant
% zero lag interaction

clear all 
clc
close all

filepath = fileparts(pwd);
addpath(strcat(filepath,filesep,'Data',filesep));
addpath(strcat(filepath,filesep,'Results',filesep));
addpath(strcat(filepath,filesep,'Aux_functions',filesep));


%% load data
load('spectrogram_modulation.mat');
S_to_cluster = zeros(size(S_all,1),size(S_all,2)*size(S_all,3));
for i=1:size(S_all,1)
    tmp = squeeze(S_all(i,:,:));
    S_to_cluster(i,:) = tmp(:);
end

%% Finding the optimum parameters of clustering
center = 0;
%[K_opt,distance_opt] = kmeans_param_opt(S_to_cluster,center); % if u want
%to generate the Sfig4
K_opt = 4;
distance_opt = 'correlation';

%% Final Clustering and visualizing the results
[idx_clus,~,~] = kmeans(S_to_cluster,K_opt,'Replicates',10,'Start','plus','MaxIter',10000,'Distance',distance_opt,'EmptyAction','drop');

% load('Spectrograms_clustering_indices.mat'); % if u want the exact
% results in the paper load this

cmap = [-0.03 0.03]; % mapping for the colors
figure()
sgtitle('Average Spectrograms of Ensembles Inside Each Cluster','Fontsize',14)
for i=1:K_opt
    tmp_cluster_eucl = S_to_cluster(idx_clus == i,:);
    tmp_belong_clusters = reshape(tmp_cluster_eucl,[size(tmp_cluster_eucl,1),length(f),length(t)]);
    Avg_S_clusters = squeeze(mean(tmp_belong_clusters,1));
    subplot(2,ceil(K_opt/2),i);hndl = imagesc(Avg_S_clusters);caxis(cmap);
    xticks([]);yticks([]);title(strcat('Cluster: ',num2str(i)),'FontSize',12)
end

subplot(2,2,1);ylabel('Frequency (Hz)','FontSize',12);yticks(sort(36-[3,8,14,24]));yticklabels({'100','60','30','12'})
subplot(2,2,3);ylabel('Frequency (Hz)','FontSize',12);yticks(sort(36-[3,8,14,24]));yticklabels({'100','60','30','12'})
subplot(2,2,3);xlabel('Time (ms)','FontSize',12);xticks([1,21,31,41,61]);xticklabels({'-300','-100','0','100','300'})
subplot(2,2,4);xlabel('Time (ms)','FontSize',12);xticks([1,21,31,41,61]);xticklabels({'-300','-100','0','100','300'})

% save_dir = strcat(pwd,'/results/');
% save_file = strcat(save_dir,'Spectrograms_clustering_indices.mat');
% save(save_file,'idx_clus');

%% Significane assessment
p_val_uncorrected_all = zeros([K_opt,size(S_all,2),size(S_all,3)]);
h_uncorrected_all =  zeros([K_opt,size(S_all,2),size(S_all,3)]);
h_corrected_all = zeros([K_opt,size(S_all,2),size(S_all,3)]);
for i=1:K_opt
    S_temp = S_all(idx_clus == i,:,:);
    for j=1:size(S_temp,2)
        for k=1:size(S_temp,3)
            [p_val_uncorrected_all(i,j,k),h_uncorrected_all(i,j,k)] = signrank(squeeze(S_temp(:,j,k)));
        end
    end
    h_corrected_all(i,:,:) = stat_fdr_bh(p_val_uncorrected_all(i,:,:));
end

%% Visualizing significancy 
figure('units','normalized','outerposition',[0 0 1 1]);

for i=1:K_opt
    set(gcf,'papertype','a4');
    tmp_cluster_eucl = S_to_cluster(idx_clus == i,:);
    tmp_belong_clusters = reshape(tmp_cluster_eucl,[size(tmp_cluster_eucl,1),length(f),length(t)]);
    Avg_S_clusters = squeeze(mean(tmp_belong_clusters,1));
    h_prod = squeeze(h_corrected_all(i,:,:));
    h_prod(h_prod == 0) = nan;
    to_plot = Avg_S_clusters.*h_prod;
    subplot(2,ceil(K_opt/2),i);
    hndl = imagesc(to_plot);caxis(cmap);
    set(hndl,'AlphaData',~isnan(to_plot));xline(31,'r','linewidth',4);
    xticks([]);yticks([]);
    ylabel('Frequency (Hz)','FontSize',12);yticks(sort(36-[3,8,14,24]));yticklabels({'100','60','30','12'});ax = gca;ax.FontSize = 12; 
    xlabel('Time (ms)','FontSize',12);xticks([1,21,31,41,61]);xticklabels({'-300','-100','0','100','300'});ax = gca;ax.FontSize = 12;
    set(gca,'fontsize',12,'fontname','arial');
    colorbar('Ticks',[-0.03,0,0.03],'TickLabels',{'-0.03','0','0.03'})
end

subplot(2,2,1);ylabel('Frequency (Hz)','FontSize',30);yticks(sort(36-[3,8,14,24]));yticklabels({'100','60','30','12'});ax = gca;ax.FontSize = 24; 
subplot(2,2,3);ylabel('Frequency (Hz)','FontSize',30);yticks(sort(36-[3,8,14,24]));yticklabels({'100','60','30','12'});ax = gca;ax.FontSize = 24;
subplot(2,2,3);xlabel('Time (ms)','FontSize',30);xticks([1,21,31,41,61]);xticklabels({'-300','-100','0','100','300'});ax = gca;ax.FontSize = 24;
subplot(2,2,4);xlabel('Time (ms)','FontSize',30);xticks([1,21,31,41,61]);xticklabels({'-300','-100','0','100','300'});ax = gca;ax.FontSize = 24;
hp4 = get(subplot(2,2,4),'Position');
colorbar('Position', [hp4(1)+hp4(3)+0.03 hp4(2)   0.02  hp4(4)],'Ticks',[-0.03,0,0.03],'TickLabels',{'-0.03','0','0.03'})


%% Band assessment per cluster
bands = {[4,8],[8,12],[12,30],[30,70],[70,150]};
bands_name = {'theta','alpha','beta','low-gamma','high-gamma'};
n_bands = length(bands);
to_bar = zeros(K_opt,n_bands);
err_bar = to_bar;
f_idx = f(end:-1:1);
p_val_band = zeros(n_bands,K_opt);
for clst=1:K_opt
    for band=1:n_bands       
        idx_f = f_idx >=bands{band}(1) & f_idx <bands{band}(2);
        S_tmp = mean(mean(S_all(idx_clus == clst,idx_f,31:end),3,'omitnan'),2,'omitnan');
        to_bar(clst,band) = mean(S_tmp);
        err_bar(clst,band) = std(S_tmp)./sqrt(length(S_tmp));
        p_val_band(band,clst) = signrank(S_tmp);
    end
    p_val_band(:,clst) = bonf_holm(p_val_band(:,clst));
end
figure()
barsem(to_bar',err_bar');ylim([-0.03,0.03])
%% synch spectrogram clustering
load('spectrogram_modulation_synch.mat');

S_to_cluster_synch = zeros(size(S_all_synch,1),size(S_all_synch,2)*size(S_all_synch,3));
for i=1:size(S_all_synch,1)
    tmp = squeeze(S_all_synch(i,:,:));
    S_to_cluster_synch(i,:) = tmp(:);
end

eva = evalclusters(S_to_cluster_synch,'kmeans','silhouette','KList',[1:6],'Distance','correlation','ClusterPriors','equal');
K_opt = eva.OptimalK;distance_opt = 'correlation';
[idx_clus_synch,~,~] = kmeans(S_to_cluster_synch,K_opt,'Replicates',10,'Start','plus','MaxIter',10000,'Distance',distance_opt,'EmptyAction','drop');

%load('Spectrograms_clustering_indices_synch,mat'); % for the exact results
%of the paper 

p_val_uncorrected_all = zeros([K_opt,size(S_all_synch,2),size(S_all_synch,3)]);
h_uncorrected_all =  zeros([K_opt,size(S_all_synch,2),size(S_all_synch,3)]);
h_corrected_all = zeros([K_opt,size(S_all_synch,2),size(S_all_synch,3)]);
for i=1:K_opt
    S_temp = S_all_synch(idx_clus_synch == i,:,:);
    for j=1:size(S_temp,2)
        for k=1:size(S_temp,3)
            [p_val_uncorrected_all(i,j,k),h_uncorrected_all(i,j,k)] = signrank(squeeze(S_temp(:,j,k)));
        end
    end
    h_corrected_all(i,:,:) = stat_fdr_bh(p_val_uncorrected_all(i,:,:));
end


figure()
set(gcf,'papertype','a4');
for i=1:K_opt
    tmp_cluster = S_to_cluster_synch(idx_clus_synch == i,:);
    tmp_belong_clusters = reshape(tmp_cluster,[size(tmp_cluster,1),length(f),length(t)]);
    Avg_S_clusters = squeeze(mean(tmp_belong_clusters,1));
    h_prod = squeeze(h_corrected_all(i,:,:));
    h_prod(h_prod == 0) = nan;
    to_plot = Avg_S_clusters.*h_prod;
    subplot(2,ceil(K_opt/2),i);
    hndl = imagesc(to_plot);caxis([-1 1]*0.03);
    set(hndl,'AlphaData',~isnan(to_plot));xline(31,'r','linewidth',4);
    xticks([]);yticks([]);
    ylabel('Frequency (Hz)','FontSize',12);yticks(sort(36-[3,8,14,24]));yticklabels({'100','60','30','12'});ax = gca;ax.FontSize = 12; 
    xlabel('Time (ms)','FontSize',12);xticks([1,21,31,41,61]);xticklabels({'-300','-100','0','100','300'});ax = gca;ax.FontSize = 12;
    set(gca,'fontsize',12,'fontname','arial');
    colorbar('Ticks',[-0.03,0,0.03],'TickLabels',{'-0.03','0','0.03'})
end


%% reverse correlation scatter
load('reverse_correlation_modulation.mat');
load('reverse_correlation_modulation_sh.mat');
n_rats = length(reverse_corr_modulation);
bands_name = {'theta','low-gamma','high-gamma'};
n_bands = length(bands_name);
len_before = 0.3;
len_after = 0.4;
Fs = 1082.25;
t = -len_before:1/Fs:len_after;
t(1) = [];
bands_to_investigate = [1,4,5];

to_scatter = [];
to_scatter_sh = [];
for band=1:n_bands
    count = 1;
    for i=1:n_rats
        data_rat = reverse_corr_modulation{i};
        data_rat_sh = reverse_corr_modulation_sh{i};
        n_ensembles = length(data_rat);
        for j=1:n_ensembles
            data_ensemble = data_rat{j};
            data_ensemble_sh = data_rat_sh{j};
            data_band = data_ensemble{bands_to_investigate(band)};
            data_band_sh = data_ensemble_sh{bands_to_investigate(band)};
            to_scatter(count,band) = mean(data_band(t>0));
            to_scatter_sh(count,band) = mean(mean(data_band_sh(t>0,:)));
            count = count + 1;
        end
    end
end


figure('units','normalized','outerposition',[0 0 1 1]);
count = 1;
for band_1 = 1:n_bands
    for band_2 = band_1+1:n_bands
        subplot(1,3,count);
        tmp = [to_scatter(:,band_1),to_scatter(:,band_2)];
        top_right = round(sum(tmp(:,1)>0 & tmp(:,2)>0)./length(tmp)*100);
        bot_left = round(sum(tmp(:,1)<0 & tmp(:,2)<0)./length(tmp)*100);
        top_left = round(sum(tmp(:,1)<0 & tmp(:,2)>0)./length(tmp)*100);
        bot_right = round(sum(tmp(:,1)>0 & tmp(:,2)<0)./length(tmp)*100);
        scatter(to_scatter(:,band_1),to_scatter(:,band_2),'x')
        hold all
        scatter(to_scatter_sh(:,band_1),to_scatter_sh(:,band_2),'x')
        xlabel(bands_name{band_1});ylabel(bands_name{band_2})
        xline(0,'--r');yline(0,'--r');
        count = count + 1;
        axis square;xlim([-0.01,0.01]);ylim([-0.01,0.01]);
        text(min(xlim), min(ylim), strcat(num2str(bot_left),'%'))
        text(min(xlim), max(ylim), strcat(num2str(top_left),'%'))
        text(max(xlim), min(ylim), strcat(num2str(bot_right),'%'))
        text(max(xlim), max(ylim), strcat(num2str(top_right),'%'))
    end
end