% 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 [v t fout ren_prob full_ren_prob raw_prob dt]=fp_lv_leak_curr_ftcs_new(th,lk,dt,tf,v0,p0,mu,sig2)%,wrap_around_frac);
%This function generates the forward-time centered-space (FTCS) solution to Fokker-Planck equation
%with additional empty points to account for no flux through left boundary.
%The solution is generated with two separate 'transition' matrices for the diffusion 
%and drift respectively. This renders the solution stable even with zero
%diffusion
%
%Inputs: 
%th             Threshold (V)
%lk             Leak (1/membrane time-constant)(1/s)
%dt             Time step (s)
%tf             Final time (s)
%v0             Uniform voltage discretization vector (vector, s)
%p0             Initial probability distribution
%mu             Drift = Mean syn wt*firing rate*tau
%sig2           Diffusion = Mean(syn wt^2)*firing rate
%
%Outputs:
%v              Same as v0
%t              Time discretization vector
%fout           Output firing rate
%ren_prob       Normalized probability distribution matrix (between resting potential and threshold) for heat map 
%full_ren_prob  Normalized probability distribution matrix over entire range of appended voltage discretization vector
%raw_prob       Non-normalized probability distribution matrix (between resting potential and threshold)
%dt             time-step used by the FTCS scheme


if nargin<1                 %parameters
    th=20*10^-3;            %V Threshold
end
if nargin<2
    lk=1/(20*10^-3);        %s^-1 leak constant 1/RC
end
if nargin<4
    dt=0.1*10^-3;           %s time mesh
    tf=50*10^-3;            %s final time
end
if nargin<6
    v0=(0:0.5:20)*10^-3;    %mV defines the grid for the membrane potential
    p0=zeros(size(v0));
    p0(1)=1;                %start from rest
end
if nargin<8                 %drift and diffusion
    mu = (2*10^-3)*200/lk;  %Drift for mean syn wt = 2 mV, fr=200 Hz, and tau = 20 ms
    sig2 = (2*10^-3)*(2*10^-3)*200; %Diffusion for delta fn of synaptic weight at 2 mV and fr=200 Hz
end;

%Shift initial v0 by mu, so that zero is at the leak+drift position
v0=v0-mu;
dv=v0(2)-v0(1);
if (v0(end)<0)
    v0=[v0 v0(end)+dv:dv:0];
end;
th=th-mu;

%if required, change time-step (WHY?)
sig2=sig2/2;
dt=min(dt,dv*dv/(2*sig2));            
t=0:dt:tf;
nt=length(t);

%sufficiently extend the uniform voltage discretization vector
vext = 100*10^-3;
vliml = min(v0)-vext;
vlimr = max(v0)+vext;
vleftext = vliml:dv:v0(1)-dv/2;
vrightext = v0(end)+dv:dv:vlimr;
v=[vleftext v0 vrightext];
nv=length(v);

%find threshold index
[C I]=min(abs(v-th));
thresh_index=I;

%find index to reinject probability (vr)
[C I]=min(abs(v+mu));
rest_index=I;


p0(end)=0;                             %no starting probability at threshold
p0 = [zeros(1,rest_index-1) p0];
p0 = [p0 zeros(1,nv-length(p0))];
prob=zeros(nt,nv);                     %matrix to store probabilities
dv=v(2:end)-v(1:end-1);                %vector of bin widths, to be used for normalization
prob(1,:)=p0;
prob(1,:)=prob(1,:)/sum(prob(1,:));


%Generate RHS matrices for FTCS diffusion
dv=dv(1);
dv2=dv^2;

RHSmat=sparse(diag((1-2*sig2*dt/dv2)*ones(size(v)),0));
RHSmat=RHSmat+sparse(diag((sig2*dt/dv2)*ones(nv-1,1),1));
RHSmat=RHSmat+sparse(diag((sig2*dt/dv2)*ones(nv-1,1),-1));
RHSmat(1,1)=1-sig2*dt/dv2;

fout=zeros(size(t));                 %vector to record firing rate

%precalculate exponential shifts, derived from characteristic equation from
%the transport equation portion: dp/dt = d(Lvp)/dx
expval=exp(lk*dt);

%generate transition matrix for drift (RIGHT?)
full_fp_mat = generate_fp_mat(v,expval);

%march through time steps
probnew=prob(1,:);
for ii=2:nt 
   
    %diffusion term
    if (sig2>0)
        diff_term=RHSmat*probnew';
        probnew=diff_term';
    end;
    
    %drift term
    probnew = probnew*full_fp_mat;
    
    %flux out through threshold reinjected at zero
    fout(ii)=sum(probnew(thresh_index:end));
    probnew(thresh_index:end)=0;
    probnew(rest_index)=probnew(rest_index)+fout(ii);
    prob(ii,:)=probnew;
end;

%renormalize probabilities
raw_prob = prob(:,rest_index:thresh_index);
scale=(v(thresh_index)-v(rest_index))/(thresh_index-rest_index+1);
ren_prob=prob/scale;
v=v+mu;
full_ren_prob= ren_prob;
ren_prob=ren_prob(:,rest_index:thresh_index);
v=v(rest_index:thresh_index);


