% Code for NREM to REM sleep transition project
%
% The file NREM_to_REM_transitions.mat contains the variables necessary for
% reproducing the results, the code in this file 
% (NREM_to_REM_transitions_Script.m) computes those results, and the file
% NREM_to_REM_transitions_results.mat contains the result variables.
% Variables in NREM_to_REM_transitions.mat:
% Channels: 718x2, one row per channel, the first column indicates the
% patient, and the second column the brain region from which the channel is
% recording.
% Features: 718x1500x24, for each channel (718 channels) this variable 
% contains the time-course (up to 1500 epochs) of the normalized feature 
% values (24 features). The features are features 1-3: first to third 
% log-cumulants of the multifractal spectrum, features 4-10: logarithm of 
% the mean value of the magnitude of the coefficients of the wavelet 
% decomposition of the signal (from lower to higher scale, i.e. higher to 
% lower frequency; 64-32, 32-16, 16-8, 8-4, 4-2, 2-1, 0.5-1 Hz), features 
% 11-17: approximation (sum of squared-magnitude of the coefficients of the
% wavelet decomposition) of the energy in each scale relative to the total 
% energy in the epoch, features 18-24: logarithm of the variance of the
% magnitude of the coefficients in each scale. The fill value for shorter
% nights and missing values is NaN.
% Regions: 28x1 cell array, names of the analyzed brain regions. The 
% 'undetermined' region label corresponds to channels just outside the 
% brain or limbic regions, or white matter far away from any cortical 
% region.
% ScalpStages: 29x1500, expert sleep scoring for the whole night (up to
% 1500 epochs) for the 29 patients included in the study. Stage W is 0, N1
% is 1, N2 is 2, N3 is 3, R is 5, and disconnected (W) is 9, fill value for
% shorter nights is NaN.
% ScalpTransitions: 123x6, list of transitions and expert marks. Each row
% corresponds to an analyzed transition. First column is the patient,
% second column the sleep cycle, third column the index of the epoch
% preceding the transition, fourth column time of the Onset_R marker (first
% rapid eye-movement) realtive to the transition time, fifth column 
% relative time of the Onset_S marker (first saw-tooth wave), sixth column
% relative time of the Onset_M marker (first mscle atonia). 
%
% The output variables are described at the end of this file.
%
% 2023-04-17, Nicolás von Ellenrieder
 
load NREM_to_REM_sleep_transitions ScalpStages Channels;

% Create a stage variable indicating ignore (0), NREM (1), or REM (2) stage
% for each channel, and a level variable indicating how far from a stage
% transtion each epoch is.
ScalpStages(ScalpStages==9)=0;
for ii=1:29
    ScalpStages(ii,1:find(ScalpStages(ii,:)==2,1)-1)=-1;
    ScalpStages(ii,find(ScalpStages(ii,:)==5,1,'last')+1:end)=-1;
end
ScalpStages(ScalpStages==0)=1;
ScalpStages(ScalpStages==-1)=0;
ScalpStages(ScalpStages==1)=1;
ScalpStages(ScalpStages==2)=1;
ScalpStages(ScalpStages==3)=1;
ScalpStages(ScalpStages==5)=2;
stage=ScalpStages(Channels(:,1),:);
N=size(stage,2);
level=stage+NaN;
for ii=1:size(stage,1)
    se=ones(N,2);
    for jj=1:N-1
        if stage(ii,jj+1)==stage(ii,jj), se(jj+1,1)=se(jj,1)+1; end
        if stage(ii,N-jj)==stage(ii,N-jj+1), se(N-jj,2)=se(N-jj+1,2)+1; end
    end
    level(ii,:)=min(se,[],2)';
end  
clearvars ii se jj ScalpStage;

% Automatic classification into REM or not REM

load NREM_to_REM_sleep_transitions Features;
% Parameters for the definition of "stable" runs
Lmax=10; % Maximum number of minutes away from transition
Lmin=2; % Minimum number of minutes away from a sleep state change
Lnum=30; % Minimum number of epochs in the stable period

Lmax=Lmax*2; % in epochs
Lmin=Lmin*2; % in epochs
P=stage+NaN;
FeatureImportance=zeros(24,718);
% For each channel
for nc=1:718
    F=squeeze(Features(nc,:,:));
    ix=isnan(sum(F,2));
    F(ix,:)=[];
    lev=level(nc,~ix);
    col=stage(nc,~ix);
    L=Lmax;
    while L>=Lmin
        id=lev>=L;
        n=sum(col(id)==1);
        if n<Lnum, L=L-1; else, break; end
    end
    L=max(L,Lmin); % To correct for extreme case
    pon=F(lev>=L&col==1,:); % Stable not REM
    L=Lmax;
    while L>=Lmin
        id=lev>=L;
        n=sum(col(id)==2);
        if n<Lnum, L=L-1; else, break; end
    end
    L=max(L,Lmin); % To correct for extreme case
    por=F(lev>=L&col==2,:); % Stable REM
    if ~isempty(pon)&&~isempty(por)
        % Train channel classifier with stale not REM and stable REM epochs
        M=fitctree([pon;por],[0*pon(:,1);1+0*por(:,1)],'MaxNumSplits',50);
        % Classify all epochs of channel
        [~,p]=predict(M,F);
        FeatureImportance(:,nc)=predictorImportance(M)'; 
        P(nc,~ix)=p(:,1)';
    end
end
clearvars Features Lmax Lmin Lnum nc F ix lev col L id n pon por M level N p ScalpStages stage;

% Focus on epochs around transitions
load NREM_to_REM_sleep_transitions ScalpTransitions Regions;
IntracranialTransitions=zeros(3000,4);
tran=zeros(3000,110); 
cnt=0;
for ii=1:size(ScalpTransitions,1)
    ix=ScalpTransitions(ii,1)==Channels(:,1);
    nt=nnz(ix);    
    tran(cnt+(1:nt),:)=P(ix,ScalpTransitions(ii,3)-59:ScalpTransitions(ii,3)+50);
    IntracranialTransitions(cnt+(1:nt),1)=ii;
    IntracranialTransitions(cnt+(1:nt),2)=find(ix);    
    cnt=cnt+nt;
end
tran=tran(1:cnt,:);
IntracranialTransitions=IntracranialTransitions(1:cnt,:);
clearvars cnt N P ii tr ix iy nt;

% Compute clarity and delay
x=[ones(20,1); -ones(10,1)];
M=zeros(size(tran,1),2);
for ii=1:size(tran,1)
    t=tran(ii,:);
    b=find(t(41:81)<.5,1)+40; 
    if isempty(b)
        M(ii,:)=NaN;
    else
        last=min(b+9,110);
        u=corrcoef(t(b-20:last),x(1:length(b-20:last)));
        M(ii,1)=u(1,2);
        M(ii,2)=b;
    end
end
M(:,2)=(M(:,2)-61)*30; % Convert epochs to seconds 
IntracranialTransitions(:,3:4)=M;
clearvars x ii t b last u M tran;

% Remove cases in which no stage R was found intracranially
id=isnan(IntracranialTransitions(:,3));
IntracranialTransitions(id,:)=[];
clearvars id;

% Clarity
GlobalClarity=zeros(size(ScalpTransitions,1),1); 
Nr=size(Regions,1);
RegionClarity=zeros(Nr,1);
for ii=1:size(ScalpTransitions,1)
    GlobalClarity(ii,1)=mean(IntracranialTransitions(IntracranialTransitions(:,1)==ii,3));
end
for ii=1:Nr
    RegionClarity(ii,1)=mean(IntracranialTransitions(Channels(IntracranialTransitions(:,2),2)==ii,3));
end
clearvars ii; 

% Model selection for local delays
y=IntracranialTransitions(:,4);
n=size(y,1);
k=2; 
RSS=sum((y-mean(y)).^2);
dAICc0=n*log(RSS)+2*k+(2*k^2+2*k)/(n-k-1); % No model
A(:,1)=IntracranialTransitions(:,2); % Channel
A(:,5)=IntracranialTransitions(:,1); % Transition
A(:,2)=Channels(IntracranialTransitions(:,2),2); % Brain region
A(:,3)=ScalpTransitions(A(:,5),2); A(A(:,3)>5,3)=5; % Sleep cycle
A(:,4)=ScalpTransitions(A(:,5),1); % Patient
va={1,[1 3],2,3,4,[1 5],[2 4],[2 3],[3 4],[2 3 4],5,[2 5]};
k=zeros(length(va),1); RSS=k; dAICc=k;
for ii=1:length(va)
    B=[];
    for jj=1:length(va{ii})
        for kk=1:max(A(:,va{ii}(jj)))
            b=zeros(n,1);
            b(A(:,va{ii}(jj))==kk)=1;
            B=cat(2,B,b);
        end
    end
    c=B\y;
    k(ii)=size(B,2)+1;
    RSS(ii)=sum((y-B*c).^2);
    dAICc(ii)=n*log(RSS(ii))+2*k(ii)+(2*k(ii)^2+2*k(ii))/(n-k(ii)-1);
end
dAICc=dAICc-dAICc0;
RegionDelay=c(1:Nr)-mean(c(1:Nr));
GlobalDelay=c(Nr+1:end)+mean(c(1:Nr));
clearvars A va k RSS dAICc0 ii B jj kk b c y Nr n Channels Regions ScalpTransitions;

% Results are in the following variables:
% dAICc: Results of the model selection, AIC values (corrected for small
% sample size, i.e. AICc), relative to the no-model case (dAICc=0). The
% values correspond to models with one coefficient per (1) channel, (2)
% channel and cycle, (3) brain region, (4) cycle, (5) patient, (6) Channel+
% transition, (7) region+patient, (8) region+cycle, (9) patient+cycle, 
% (10) region+cycle+patient, (11) transition, (12) transition+region.
% FeatureImportance: 718x24, importance of each feature for the
% classification in each channel.
% GlobalClarity: Average clarity of all intracranial transitions assoicated
% to each scalp transition.
% GlobalDelay: Global intracranial delay associated to each scalp
% transition, as porvided by the model with transition+brain region, i.e.
% the best model.
% IntracranialTransitions: 2837x4, one row per local intracranial
% transition. First column associated ScalpTransition, second column
% channel, third column clarity, fourth column time with respect to scalp
% transition.
% RegionClarity: 28x1, average clarity of all local transitions per region.
% RegionDelay: 28x1, time of the transitions in each brain region with 
% respect to the scalp transition. Corresponds to the values of the fitted
% best model, with coefficients for each region and each global transition.
