% Created by Terence Sanger 2019
% "A cerebellar computational mechanism for delay conditioning at precise time intervals"

rng('shuffle');  %random seed for random number generator

Ne=2000;     %total number of excitatory granule cells          

%fast spiking type Izhikevich parameters with random jitter
re=rand(Ne,1);        
a=0.016 + 0.008*re;    
b=0.225-0.05*re;     
c=-65+15*re.^2;      
d=8-6*re.^2;      

v=-65*ones(Ne,1);     % Initial values of membrane potential v
u=b.*v;               % Initial values of repolarization variable u
firings=[];           % resulting spike occurrences

%stimulation parameters
nt = 500;           %total time points (msec)
CSon = 200;         %conditioned stimulus
CSoff = 300;
CS = zeros(1,nt);
CS(CSon:CSoff) = 1; 

%create three different Unconditioned stimulus patterns
USX = {};
USon = 270;         %unconditioned stimulus at 270msec
USoff = 280;
USon2 = 230;        %second US
USoff2 = 240;

US = zeros(1,nt);
US(USon:USoff) = 1;    %single eyeblink
USX{1} = US;

US(USon2:USoff2) = 1;    %double eyeblink
USX{2} = US;

US(CSon:CSoff) = 0.5-0.5*cos(linspace(0,6*pi, 1+CSoff-CSon));   %sinusoid
USX{3} = US;

%spike input on mossy fibers
nMossy = 100;                   %number of incoming mossy fibers
MF = rand(nMossy,nt) > 0.8;     %random spike input signal, average rate 200Hz
epsc = exp(-linspace(0,5,100));  %time profile for synaptic conductance in granule cells
MF = filter(epsc,1,MF,[],2);    %incorporate the conductance
MF = MF .* (ones(nMossy,1)*CS); %signal is on only during the CS presentation
mMap = zeros(Ne,nMossy);     %random mapping from mossy fibers to granule cells
for i=1:Ne
    mMap(i,1+floor(nMossy*rand(1,4))) = 0.25;  %each granule cell contacts 4 random mossy fibers
end
epscV = 10 *(0.95 + 0.1 * rand(size(v)));   %epsc amplitude variation for each granule cell

%initialization for learning
PCweightsX = {};  %hold weights for each type of US
for PCn = 1:3,
    PCweights = ones(size(v));  %weights start at full conductance (1.0)
    PCweightsX{PCn} = PCweights;
end
LTDrate = 0.03;             %LTD caused by complex spikes
LTPrate = 0.0001;           %LTP caused by granule cell synaptic events
nTrials = 50;               %training repetitions
allSpMat = {};              %save the spike time matrices for averaging

%%
for r = 1:nTrials     %each example

    %find the full temporal spike pattern for this trial
    for t=1:nt          % in msec 
       I= epscV.*(mMap*MF(:,t));     %input current to granule cells from mossy fibers
       fired=find(v>=30); % indices of cells that exceed spiking threshold

       if ~isempty(fired)       %save the cells that are active, and update the Izhikevich parameters
          firings=[firings; t+0*fired, fired];     
          v(fired)=c(fired);  
          u(fired)=u(fired)+d(fired);
       end;
       v=v+0.5*(0.04*v.^2+5*v+140-u+I);     %integrate the Izhikevich equation in 2 steps 
       v=v+0.5*(0.04*v.^2+5*v+140-u+I);
       u=u+a.*(b.*v-u);   

    end

    %convert the spike timings to a raster
    spMat = zeros(Ne,nt);
    for i=1:length(firings),
        spMat(firings(i,2),firings(i,1)) = 1;
    end

    allSpMat{r} = spMat;        %save the raster

    %Purkinje cell learning
    for PCn = 1:length(PCweightsX),     %do for each type of US in parallel
        PCweights = PCweightsX{PCn};
        US = USX{PCn};                  %the US for this type
        for t=1:nt
            if PCn==3,                  %smooth function is learned using Widrow-Hoff LMS
                PCoutput = PCweights' * spMat(:,t)/50;
                PCweights = PCweights + LTDrate * (1-US(t) - PCoutput) * spMat(:,t);  %LMS
            else
                PCweights = PCweights + spMat(:,t) * ((1-US(t))*LTPrate - US(t)*LTDrate);   %learning rule
                    %LTD occurs during the US only.  Otherwise, LTP for all granule cells with spikes
            end
            PCweights(PCweights>1) = 1;     %maintain bounds for synaptic weights
            PCweights(PCweights<0) = 0;
        end
        PCweightsX{PCn} = PCweights;        %save
    end
end %nreps


%%  Output figure

figure('Position',[16 714 1106 226]);
for PCn = 1:length(PCweightsX)      %for each type of US
    PCweights = PCweightsX{PCn};    %learned weights for this case
    US = USX{PCn};                  %unconditioned stimulus for this case
    
    PCspikeT = [];
    for i=1:nTrials        %find histogram of trained PC firing for all trials
        spMat = allSpMat{i};
        spMatNorm = sum(spMat);      %normalize
        spMatNorm(spMatNorm==0) = 1; %prevent divide by zero
        spMatN = spMat ./ (ones(Ne,1) * sqrt(spMatNorm));  %sum squared will be 1

        PCepsc = PCweights' * spMatN;  %calculate Purkinje cell depolarization current
        epscAvg = sum(PCepsc .* CS) / sum(CS);  %average epsc during CS

        PCspike = rand(size(PCepsc)) < PCepsc/(epscAvg*2);  %Poisson spiking
        PCspikeT = [PCspikeT find(PCspike)];                %total spike times for all trials
    end
    
    subplot(1, length(PCweightsX), PCn);
    histogram(PCspikeT,CSon-50:CSoff+50, 'FaceColor', [.5 .5 .5]);
    xlabel('time (msec)', 'fontsize', 14, 'fontweight', 'bold');
    ylabel('PC spike count', 'fontsize', 14, 'fontweight', 'bold');
    hold();
    plot((CSon-50:CSoff+50),CS(CSon-50:CSoff+50)*5+62,'k', 'linewidth',2);
    plot((CSon-50:CSoff+50),US(CSon-50:CSoff+50)*5+55,'k--', 'linewidth',2);
    text(CSon-50,65,'CS', 'fontsize', 12, 'fontweight', 'bold');
    text(CSon-50,58,'US', 'fontsize', 12, 'fontweight', 'bold');

end     %repeat for next US pattern

