clear all;
% close all;

%1D FSP, Brian munsky, qbss 2019
%Naomi Ziv

%simple birth/death model

%loop on size of time (only done after working through some examples)
N=200;
%50 is small, error will be big
%N=50;
%K is the birth rate
K=30*ones(1,N+1);
%G is the death rate
G=0.02*(0:N);

%intialise A matix
A=zeros(N+1);
%intialise B, the last column
B=zeros(1,N+1);


DIV = zeros(N+1);


alpha = 0.9;  % symmetry of division.
alphaB = 0.6;  % symmetry of division.

divprop = @(n)100*(n>=180);

for n=0:N
    %fill diagonal, you leave state with birth and death
    A(n+1,n+1)=-K(n+1)-G(n+1);
    
    if n<N %not last step
        %off diagonal, birth leads to +1 state
        A(n+2,n+1)=K(n+1);
    else
        %B is beyond approximation
        B(1,n+1)=K(n+1);
    end
    
    if n>0 %not first step
        %off diagonal, death leads to -1 state
        A(n,n+1)=G(n+1);
    end
    
    DIV(1:n+1,n+1) = 1/2*binopdf([0:n],n,alpha)+1/2*binopdf([0:n],n,1-alpha);
%     div(1:n+1,n+1) = divprop(n)*(1/2*binopdf([0:n],n,alpha)+1/2*binopdf([0:n],n,1-alpha));
    div(1:n+1,n+1) = divprop(n)*(1/4*binopdf([0:n],n,alpha)+1/4*binopdf([0:n],n,1-alpha))+...
        1/4*binopdf([0:n],n,alphaB)+1/4*binopdf([0:n],n,1-alphaB);
%     div(1:n+1,n+1) = divprop(n)*(1/3*binopdf([0:n],n,alpha)+2/3*binopdf([0:n],n,1-alpha));
    
end
DIV(end+1,end+1) = 1;

div = div - diag(sum(div));
div(end+1,end+1) = 0;
pause
%%
%N is the size of the approximation
%100 is pretty big, error will be small
%loop on col/row of A
%make big matrix
A=[A,zeros(N+1,1);...
    B,0];



%%
%%%%%%%%%%%%%%%%%%%
%intial condidtions:
X0=3;
%X0=0;


%first example for P0, gives poission for X0=0 not X0=3, at t=5
P0=zeros(N+2,1);
P0(X0+1)=1;


tarr=[0:.1:3*30];
Pt = P0;
for j=2:length(tarr)
    
    %time when looping on time
    dt=tarr(j)-tarr(j-1);
    %%%%%%%%%%%%%%%%%%%
    
    %solve
    Pt=expm((A+div)*dt)*Pt;
    
    mu(2*j-1)=[0:N]*Pt(1:N+1);
    tarr_plot((2*j-1)) = tarr(j);
    
%     if mod(tarr(j),30)==0
%         Pt = DIV*Pt;
%     end
    mu(2*j)=[0:N]*Pt(1:N+1);
    tarr_plot((2*j)) = tarr(j);
    
    %ploting probabilities (used for specfic examples but also nice across
    %iterations if you look at all distributions on same plot)
    set(gca,'ylim',[0 0.1],'xlim',[0 80]);
    plot([0:N],Pt(1:end-1));
    title(['t = ',num2str(tarr_plot((2*j)))]);
    drawnow
    pause(0.1);
    
    %save error for each iteration on time
%     C(j)=Pt(end);
end


%this plot is for example of looping on time durations
%sould be cumalative distribution (to really see, have N=50)
%it is distribution of "first passage times"
%plot(tarr,C)
figure(2);

plot(tarr_plot,mu)
% plot(log(tarr),C)


%these calculations are for demonstration during specfic examples

%calculate moments
%mean
mu=[0:N]*Pt(1:N+1)
%variance
mux2=[0:N].^2*Pt(1:N+1);
s2=mux2-mu^2

%calculate analytically (5, 0.1 are parameters of model)
% Z=(5/0.1)*(1-exp(-0.1*t))
%calculate analytically when intial conditions are distribution
% Z2=(5/0.1)*(1-exp(-0.1*t))+exp(-0.1*t)*X0

%compare to theoretical distribution (you should have the plot of
%probabilties up)
% Q=poisspdf([0:N],Z);
% hold on;
% plot([0:N],Q,'--')

%error, prob leaving approximation
Pt(end)