clear all; close all; clc
%Constants and instantiating
ger_inf=readtable("germany_infected.csv");ger_rec=readtable("germany_removed.csv");
ger_inf=ger_inf{:,:}; ger_rec=ger_rec{:,:}; %Turn our tables into vectors
ger_pop=83000000;

ger_sus=ones(height(ger_inf),1)*ger_pop;
ger_sus=ger_sus-ger_inf-ger_rec;
phi=0.995; alpha=2; %Given from problem statement
a=2; beta = 3; b = 3;

%Calculate our deltas using our susceptible and infected individuals
%Where diff(ger_sus) is the difference each day
deltaS = [ger_pop-ger_sus(1,1);abs(diff(ger_sus))];
deltaI = -ger_inf(1,1);
delta_S = ger_pop - ger_sus(1,1)-ger_inf(1,1);
delta_for_beta_1 = -deltaI + sum(diff(ger_inf)); delta_for_beta_2 = delta_S - sum(diff(ger_sus));
beta_beta = delta_for_beta_1 + delta_for_beta_2;

g = @(lambda,Y,kappa) sum(gammaln(kappa+Y)-gammaln(Y+1)-gammaln(kappa)+kappa*log(1-phi)+Y*log(phi)); %Draw for t

%Tuning parameters
breakpoint = 3;
N = 15000;
M = 2; sigma = 0.005;
%Initialization
brk_pts = round(linspace(1,length(ger_inf),breakpoint+2));
lambdas = ones(1,breakpoint+1); %Starting guess for lambda
p_ir_vec = zeros(N,1);
params = zeros(breakpoint*2+2,1);
lambda_vec = zeros(N,breakpoint+1); %Lambda needs to be n+1
suggested_breakpoint = zeros(N,breakpoint);

for i = 1:N
    dummy = 0;
    kappas = [];
    c = 1;
    for j = 1:length(params)

        epsilon_1 = normrnd(0,1);
        epsilon_M = randi([-M,M]);
        %Starting with lambda, we have that lambdastar = lambda +
        %sigma*epsilon, so we need to update our alpha using kappa with
        %lambdastar, we start with just calculating lambda as:
        if j <= breakpoint+1


            fl_hol = 1-exp(-lambdas(1,j)*ger_inf(brk_pts(1,j):brk_pts(1,j+1),:)/ger_pop);
            kappa1 = (1/phi-1)*ger_sus(brk_pts(1,j):brk_pts(1,j+1),:).*fl_hol;
            lambda_star = lambdas(1,j) + sigma*epsilon_1;

                if lambda_star < 0 %Don't bother with negative lambda_stars
                    lambda_star = lambdas(1,j);
                end


            fl_hol = 1-exp(-lambda_star*ger_inf(brk_pts(1,j):brk_pts(1,j+1),:)/ger_pop);
            kappa2 = (1/phi-1).*ger_sus(brk_pts(1,j):brk_pts(1,j+1),:).*fl_hol;

            U = rand(1); %draw from uniform distribution
            %disp(j);
            delta_dummy = deltaS(brk_pts(1,j):brk_pts(1,j+1),:);
            alpha_prob = min(1,exp(dummy+(prob(lambda_star,delta_dummy,kappa2))-(dummy+(prob(lambdas(1,j),delta_dummy,kappa1)))));

            if U <= alpha_prob
                %Accept lambdastar if U <= alpha and update our dummy
                lambdas(1,j) = lambda_star;
                dummy = dummy + prob(lambda_star,delta_dummy,kappa2);
                lambda_vec(i,j) = lambda_star;
            else
               %If we reject lambda_star we update our dummy variable with our
               %original kappa1
               lambda_vec(i,j) = lambdas(1,j);
               dummy = dummy + prob(lambdas(1,j),delta_dummy,kappa1);
            end



        elseif breakpoint+1 < j && j <= 2*breakpoint+1
            %Looking at t now, t_star is given as t_star = t_i + epsilon_M
            %where epsilon_M is uniformly distributed on +-M.
            U = rand(1);
            t_star = brk_pts(1,c+1) + epsilon_M;

            if brk_pts(1,c)>t_star || brk_pts(1,c+2) < t_star
                t_star = brk_pts(1,c+1); %If we're between two breakpts we set t_star to simply be the mid-point of the breakpts
            end

            fl_hol=1-exp(-lambdas(1,c)*ger_inf(brk_pts(1,c):brk_pts(1,c+1),:)/ger_pop);
            kappa1=(1/phi-1)*ger_sus(brk_pts(1,c):brk_pts(1,c+1),:).*fl_hol;

            %Generate Kappa2 but now using our suggested t_star as new
            %breakpoint
            fl_hol=1-exp(-lambdas(1,c)*ger_inf(brk_pts(1,c):t_star,:)/ger_pop);
            kappa2=(1/phi-1)*ger_sus(brk_pts(1,c):t_star,:).*fl_hol;


            deltas_temp1=deltaS(brk_pts(1,c):brk_pts(1,c+1),:);
            deltas_temp2=deltaS(brk_pts(1,c):t_star,:);


            prob_1=inverseProb(deltas_temp2,kappa2);
            prob_2=inverseProb(deltas_temp1,kappa1);


            %Repeat after drawing from the first section up to t_star
            fl_hol=1-exp(-lambdas(1,c+1)*ger_inf(brk_pts(1,c+1):brk_pts(1,c+2),:)/ger_pop);
            kappa1=(1/phi-1)*ger_sus(brk_pts(1,c+1):brk_pts(1,c+2),:).*fl_hol;


            fl_hol=1-exp(-lambdas(1,c+1)*ger_inf(t_star:brk_pts(1,c+2),:)/ger_pop);
            kappa2=(1/phi-1)*ger_sus(t_star:brk_pts(1,c+2),:).*fl_hol;


            deltas_temp1=deltaS(brk_pts(1,c+1):brk_pts(1,c+2),:);
            deltas_temp2=deltaS(t_star:brk_pts(1,c+2),:);

            prob_1=prob_1+inverseProb(deltas_temp2,kappa2);
            prob_2=prob_2+inverseProb(deltas_temp1,kappa1);

            alpha_prob = min(1,exp(prob_1-prob_2));

            if U <= alpha_prob
                %Accepted breakpoint
                brk_pts(1,c+1) = t_star;

            else
            brk_pts(1,c+1) = brk_pts(1,c+1);
            suggested_breakpoint(i,c) = t_star;
            c = c+1;
            end

        elseif j > 2*breakpoint+1 && j <=2*breakpoint+2
            p_ir = betarnd(beta_beta-a,sum(ger_inf)-beta_beta+b);
            p_ir_vec(i,1)=p_ir;

        end

        delta_I = -ger_inf(j+1,1)+ger_inf(j,1);
        delta_S = -ger_sus(j+1,1)+ger_sus(j,1);
    end
end



 %Plots

figure(1)
plot(suggested_breakpoint(:,1))
hold on
plot(suggested_breakpoint(:,2))
hold on
plot(suggested_breakpoint(:,3))
legend('t_1','t_2','t_3')
xlabel("Iterations")
ylabel("Proposed breakpoint")
title("German data, 3 Breakpoints")



figure(2)
plot(lambda_vec(:,1))
hold on
plot(lambda_vec(:,2))
hold on
plot(lambda_vec(:,3))
hold on
plot(lambda_vec(:,4))
legend('lambda_1','lambda_2','lambda_3','lambda_4')
xlabel("Iterations")
ylabel("Proposed Lambdas")
title("German data, 4 lambdas")



%p_ir = betarnd(beta_beta - a, sum(ger_inf)-beta_beta+b);
%p_ir_vec = [p_ir_vec, p_ir];

function f = prob(lambda,Y,kappa)
    alpha = 2;
    beta = 3;
    phi = 0.995;
    %Density function
    f = alpha*log(beta)-gammaln(alpha)+(alpha-1)*log(lambda)-beta*lambda+sum(gammaln(kappa+Y)-gammaln(Y+1)-gammaln(kappa)+kappa*log(1-phi)+Y*log(phi));
end

function g = inverseProb(Y,kappa)

    phi = 0.995;
    %Transform function when sampling for t
    g = sum(gammaln(kappa+Y)-gammaln(Y+1)-gammaln(kappa)+kappa*log(1-phi)+Y*log(phi));
end