% This file is licensed under the Allen Institute Terms of Use, which can be found at:
%http://www.alleninstitute.org/Media/policies/terms_of_use_content.html

%------------------------------------------------------------------------------------

function [cellpop connect t]=dipde_multi_pop2(cellpop,connect,dt,tf,haveleak);

%Function to propagate numerical solution to probability distribution of
%neuron voltage with a current or conductance-based synaptic weight distribution
%given by a 3-row array qdw and linear leak L(v)=lv. Please refer to the
%section METHODS: NUMERICAL SOLUTIONS in the paper. This function is the
%SAME AS dipde_multi_pop_fig1, EXCEPT THAT IT DOES NOT
%compute the 95% confidence intervals for the output firing rate.
%THIS IS IMPLEMENTED ONLY FOR LIF NEURONS.
%In addition, the function implements short-term synaptic depression
%with the full joint probability p(v,u,t) expressed as a product of
%marginals over v and u separately so that p(v,u,t) = p(v,t)*p(u,t). The synaptic
%utility (u) is modeled as a separate ODE using the Tsodyks-Markram mechanism. For
%further details, please see METHODS: Implementation of Short-term Synaptic
%Depression in the paper.
%
%Function Inputs :
%-----------------
%cellpop : structure containing source and target populations
%connect : structure containing object connecting source and target
%dt : time-step for time-discretization
%tf : simulation time
%haveleak : (= 1) specifies if leak is implemented
%
%Function Outputs:
%-----------------
%cellpop: structure containing evolved populations
%connect: connection object after simulation (will change if synaptic utility present)
%t : time-discretization vector used in simulation (same as initial)
%
%NOTE: To avoid discretizing the spatial derivative, this function uses
%geometric binning to calculate the analytical solution of the transport
%equation dp/dt = d(lvp)/dx : this is hard-coded only for linear leak=lv, so will
%need to be changed if the form of the leak is different.
%
%The probability distribution is first evolved with the leak term and then
%the synaptic input. The geometric binning scheme allows the implementing the
%effect of the leak term as a single-index shift towards the resting
%potential. The effect of synaptic input is then implemented by generating a
%transition matrix as explained in METHODS: NUMERICAL SOLUTIONS in the
%paper.

if nargin<5
    haveleak=1;
end

%Number of cell populations and connections
num_pop=length(cellpop);
num_con=length(connect);

%interpolation method for initial probability distribution
interpmethod='nearest';

%add input connection info to cellpop
for ii=1:num_pop
    cellpop(ii).inputs=[];
    cellpop(ii).cons=[];
end;
for ii=1:num_con
    cellpop(connect(ii).target).inputs=[cellpop(connect(ii).target).inputs connect(ii).source];
    cellpop(connect(ii).target).cons=[cellpop(connect(ii).target).cons ii];
    %normalize probabilities, if not already done
    qdw=connect(ii).qdw;
    qdw(3,:)=qdw(3,:)/sum(qdw(3,:));
    connect(ii).qdw=qdw;
end;

dt_new=dt;

%start with dt, check if it needs to be made smaller to ensure that the
%zero bin is smaller than than the smallest synaptic weight
for ii=1:num_pop
    expval=exp(-cellpop(ii).lk*dt);
    numbins=ceil(log(1-expval)/(-cellpop(ii).lk*dt));        %to ensure that left-most (0) bin is as small as/smaller than right-most bin
    %check that bin size is smaller than the smallest possible synaptic dV from
    %rest
    dv_max=cellpop(ii).th*(1-expval);
    min_syn=dv_max;
    for jj=1:length(cellpop(ii).cons)
        qdw=connect(cellpop(ii).cons(jj)).qdw;
        keep_cols=find(abs(qdw(1,:))>0);
        min_syn=min(min_syn,abs(qdw(1,keep_cols).*qdw(2,keep_cols)));
    end;
    if dv_max>min_syn
        dv_max=min_syn;
        disp(sprintf('Max bin size is too large compared to synaptic weight. Changing max bin size to %g',dv_max));
        %recalculate dt to ensure new max bin size
        dt_new=log((1-dv_max/cellpop(ii).th))/(-cellpop(ii).lk);
    end;
end;
if dt_new~=dt
    %reinterpolate external firing rate vectors
    for ii=1:num_pop
        if (cellpop(ii).flag_ext==1)
            cellpop(ii).external_fr=interp1(0:dt:tf,cellpop(ii).external_fr,0:dt_new:tf);
        end
    end
    dt=dt_new;
end;
t=0:dt:tf;
nt=length(t);
disp(sprintf('Time step = %g',dt));

%generate exponential voltage bins
for ii=1:num_pop
    cellpop(ii).v=generate_geom_bins(1/cellpop(ii).lk,dt,cellpop(ii).th);
    cellpop(ii).nv=length(cellpop(ii).v);
    %check to see if any reason to look at negative voltages
    vminus=[];
    vr_min=cellpop(ii).v0(1);
    for jj=1:length(cellpop(ii).cons)
        qdw=connect(cellpop(ii).cons(jj)).qdw;
        if connect(cellpop(ii).cons(jj)).flag_cond_syn==1      %conductance-based synapses
            if sum(qdw(2,:)<0)>0                               %synapse has a reversal potential <0
                vr_min=min([qdw(2,:) vr_min]);
            end;
        else                                                   %current-based synapses
            if sum(qdw(1,:)<0)>0
                vr_min=min([-cellpop(ii).th vr_min]);
            end
        end
    end
    if vr_min<0
        expval=exp(-cellpop(ii).lk*dt);
        numbins=ceil(log(1-expval)/(-cellpop(ii).lk*dt));
        vminus=[cumprod([vr_min expval*ones(1,numbins)])];
        cellpop(ii).v=[vminus cellpop(ii).v(2:end)];
        cellpop(ii).zero_index=length(vminus);            %index of bin whose left edge is at zero
        cellpop(ii).nv=length(v);
    else
        cellpop(ii).zero_index=1;
    end;
    disp(sprintf('Population: %g\t, number of voltage bins: %d\n',ii,cellpop(ii).nv));
end;

%generate synaptic utility vectors for each population
for ii=1:num_pop
    if cellpop(ii).u_tau>0                                            %time constant for synaptic utility recovery
        u_vec=fliplr(1-(generate_geom_bins(cellpop(ii).u_tau,dt,1))); %geometric bins with u_tau
        u0_interp=zeros(size(u_vec));
        if (sum(cellpop(ii).u_prob>0)==1)           %delta function, presumably, so assign to a single bin
            [val ind]=min(abs(u_vec-cellpop(ii).u_vec(cellpop(ii).u_prob>0)));
            u0_interp(ind)=1;
        else                                        %interpolate u_prob from uniform to geometric u-bins 
            u0_interp=interp1(cellpop(ii).u_vec,cellpop(ii).u_prob,u_vec,interpmethod);
        end
        cellpop(ii).u_vec=u_vec;
        cellpop(ii).u_prob=u0_interp;
        %u_pts are the set of points to which to interpolate the utility values (to speed up computation)
        if (cellpop(ii).u_pts(1)<0)
            cellpop(ii).u_pts=cellpop(ii).u_vec;
        end
        %matrix to collapse utility vector from geometric bins to u_pts
        collapse_mat=(zeros(length(cellpop(ii).u_pts),length(cellpop(ii).u_vec)));
        for jj=1:length(cellpop(ii).u_vec)
            diffs=cellpop(ii).u_pts-cellpop(ii).u_vec(jj);
            tempdiffs=diffs;
            tempdiffs(diffs>0)=-Inf;
            [lowerbnd lowerind]=max(tempdiffs);
            tempdiffs=diffs;
            tempdiffs(diffs<0)=Inf;
            [upperbnd upperind]=min(tempdiffs);
            denom=cellpop(ii).u_pts(upperind)-cellpop(ii).u_pts(lowerind);
            if (denom>0)
                collapse_mat(lowerind,jj)=upperbnd/denom;
                collapse_mat(upperind,jj)=-lowerbnd/denom;
            else
                collapse_mat(lowerind,jj)=1;
            end;
        end;
        cellpop(ii).collapse_mat=collapse_mat;
    else
        cellpop(ii).u_vec=1;
        cellpop(ii).u_prob=1;
        cellpop(ii).collapse_mat=1;
        cellpop(ii).util_change=1;
    end;
    %Make change in utility for non-trivial u_frac (fraction of synaptic utility available after a spike)
    if cellpop(ii).u_frac<1
        cellpop(ii).util_change=generate_u_transition(cellpop(ii).u_vec,cellpop(ii).u_frac);
    end
    %average utility
    cellpop(ii).u_avg=zeros(size(nt));
    cellpop(ii).u_avg(1)=sum(cellpop(ii).u_prob.*cellpop(ii).u_vec);
end;

%generate leak matrix for each population
for ii=1:num_pop
    nv=cellpop(ii).nv;
    z=cellpop(ii).zero_index;
    diagidx=[sub2ind([nv nv],z,z) sub2ind([nv nv],z+1:nv-1,z:nv-2)];  %indices of the diagonal
    if z>1                                                            %dealing with negative voltages
        diagidx=[diagidx sub2ind([nv nv],1:z-1,2:z)];
    end;
    cellpop(ii).leak_mat=sparse(zeros(nv));
    cellpop(ii).leak_mat(diagidx)=1;
end;


%generate separate synaptic-event transition matrices for each connection
for ii=1:num_con
    v=cellpop(connect(ii).target).v;
    u_pts=cellpop(connect(ii).source).u_pts;
    connect(ii).syn_mat=cell(1,length(u_pts));
    for jj=1:length(u_pts)
        qdw=connect(ii).qdw;
        qdw(1,:)=qdw(1,:)*u_pts(jj);
        syn_mat=generate_synaptic_matrix2_testsyn2(qdw,v,1,1,connect(ii).flag_cond_syn,0,connect(ii).voltage_frac);
        connect(ii).syn_mat{jj}=syn_mat;
    end;
end;

%set up solution matrices for each population
for ii=1:num_pop
    p0=cellpop(ii).p0;
    v0=cellpop(ii).v0;
    p0(v0==cellpop(ii).th)=0;
    v=cellpop(ii).v;
    p0_interp=zeros(size(v));
    if (sum(p0>0)==1)                   %delta function, presumably, so assign to a single bin
        [val ind]=min(abs(v-v0(p0>0)));
        p0_interp(ind)=1;
    else
        p0_interp=interp1(v0,p0,v,interpmethod);
    end;
    dv=v(2:end)-v(1:end-1);
    cellpop(ii).fout=zeros(size(t));          %vector to record firing rate
    cellpop(ii).prob=zeros(nt,length(v));     %matrix to store probabilities
    cellpop(ii).prob(1,1:end-1)=p0_interp(1:end-1).*dv;
    cellpop(ii).prob(1,end)=0;
    cellpop(ii).prob(1,:)=cellpop(ii).prob(1,:)/sum(cellpop(ii).prob(1,:));
    cellpop(ii).prob_unif=zeros(nt,length(v0));
    cellpop(ii).prob_unif(1,:)=p0;
end;

%march through time steps
for ii=2:nt %time counter
    for jj=1:num_pop
        if (cellpop(jj).flag_ext==0)        %not external population
            probnew=cellpop(jj).prob(ii-1,:);
            
            %apply leak
            if haveleak==1
                probnew=probnew*cellpop(jj).leak_mat;
            end;
            %apply synaptic input
            inputs=cellpop(jj).inputs;
            cons=cellpop(jj).cons;
            for kk=1:length(cons)
                %get utility values for input population
                input_u=cellpop(inputs(kk)).collapse_mat*cellpop(inputs(kk)).u_prob';
                if cellpop(inputs(kk)).flag_ext==1  %external firing
                    fr=cellpop(inputs(kk)).external_fr(ii)*dt; %no of inputs in time dt
                else
                    fr=cellpop(inputs(kk)).fout(ii-1)*connect(cons(kk)).nsyn;
                end;
                
                if connect(cons(kk)).flag_poisson==1
                    %calculate number of synaptic events needed to account for
                    %>99.999% of Poisson probability
                    syn_probs=poisson_prob(fr,0.99999);
                else
                    syn_probs=zeros(1,ceil(fr)+1);
                    syn_probs(end)=fr-floor(fr);
                    syn_probs(end-1)=ceil(fr)-fr;
                end;
                
                %if needed, generate more synaptic matrices depending on
                %input firing rate fr
                [r c]=size(connect(cons(kk)).syn_mat);
                for rowcount=1:length(syn_probs)-r-1;   
                    v=cellpop(connect(cons(kk)).target).v;
                    u_pts=cellpop(connect(cons(kk)).source).u_pts;
                    temprow=cell(1,length(u_pts));
                    for mm=1:length(u_pts)
                        qdw=connect(cons(kk)).qdw;
                        qdw(1,:)=qdw(1,:)*u_pts(mm);
                        syn_mat=generate_synaptic_matrix2_testsyn2(qdw,v,1,r+rowcount,connect(cons(kk)).flag_cond_syn,0,connect(cons(kk)).voltage_frac);
                        temprow{mm}=syn_mat;
                    end;
                    connect(cons(kk)).syn_mat=[connect(cons(kk)).syn_mat; temprow];
                end;
                
                %Full transition matrix for synaptic events 
                full_synaptic_mat=sparse(zeros(cellpop(jj).nv));
                
                %run through each utility value
                for util=1:length(input_u)
                    if (input_u(util)>0)
                        full_synaptic_mat=full_synaptic_mat+eye(cellpop(jj).nv)*input_u(util)*syn_probs(1);   %probability of zero events
                        for num_syn=2:length(syn_probs)
                            full_synaptic_mat=full_synaptic_mat+connect(cons(kk)).syn_mat{num_syn-1,util}*input_u(util)*syn_probs(num_syn);
                        end;
                    end;
                end;
                
                %restore truncated probability to diagonal of
                %full_synaptic_mat
                diagIdx = 1:cellpop(jj).nv+1:cellpop(jj).nv^2;
                missing = 1-sum(full_synaptic_mat(:,1:end-1)');
                full_synaptic_mat(diagIdx)=full_synaptic_mat(diagIdx)+missing;
                probnew=probnew*full_synaptic_mat;
                cellpop(jj).fout(ii)=cellpop(jj).fout(ii)+sum(probnew(:,end));
                probnew(:,end)=0;
            end;
            cellpop(jj).prob(ii,:)=probnew;
        end;
    end;
    
    %update utilities
    for jj=1:num_pop
        %recovery from depression
        if (cellpop(jj).u_frac<1)
            if (cellpop(jj).external_fr(ii)>0)
                %find number of synapses per output neuron from the
                %input population
                nsyn=connect(connect(:).source==jj).nsyn;
                cellpop(jj).fout(ii)=cellpop(jj).external_fr(ii)*dt/nsyn;
            end;
            if cellpop(jj).fout(ii)>0   %spiking occurred, update utility vector
                full_u_transition_mat=(sparse(eye(length(cellpop(jj).u_vec)))*(1-cellpop(jj).fout(ii)))+(cellpop(jj).util_change*cellpop(jj).fout(ii));
                cellpop(jj).u_prob=cellpop(jj).u_prob*full_u_transition_mat;
            end
            if (haveleak==1)
                cellpop(jj).u_prob(end)=sum(cellpop(jj).u_prob(end-1:end));
                cellpop(jj).u_prob(2:end-1)=cellpop(jj).u_prob(1:end-2);
                cellpop(jj).u_prob(1)=0;
            end
            cellpop(jj).u_avg(ii)=sum(cellpop(jj).u_prob.*cellpop(jj).u_vec);
        end
    end
end


%normalize and generate uniform probability vectors
for ii=1:num_pop
    if (cellpop(ii).flag_ext==0)        %not external population
        dv=cellpop(ii).v(2:end)-cellpop(ii).v(1:end-1);
        scale=(length(cellpop(ii).v0)-1)/(cellpop(ii).v0(end)-cellpop(ii).v0(1));
        for jj=1:nt
            temp=rebin_uniform_vector(cellpop(ii).v,cellpop(ii).v0,cellpop(ii).prob(jj,:));
            cellpop(ii).prob_unif(jj,:)=temp*scale;
        end;
        cellpop(ii).prob=cellpop(ii).prob*diag([1./dv 1]);
        cellpop(ii).prob(:,end)=cellpop(ii).fout;
        cellpop(ii).fout=cellpop(ii).fout/dt;
    end
end
end

function outvec=generate_geom_bins(tau,dt,th)   %function to generate geometric v-bins for LIF
expval2=exp(-dt/tau);
numbins2=ceil(log(1-expval2)/(-dt/tau));        %to ensure that left-most (0) bin is as small as/smaller than right-most bin
outvec=[0 fliplr(cumprod([th expval2*ones(1,numbins2)]))];
end

function syn_probs=poisson_prob(f,problimit)    %function to generate Poisson probabilities 
probtot=exp(-f);
syn_probs=probtot;
syn=0;
while (probtot<problimit)
    syn=syn+1;
    val2=f^syn*exp(-f)/factorial(syn);
    syn_probs=[syn_probs val2];
    probtot=probtot+val2;
end;
end
