% This version of the file includes an additional state variable which is
% the multiplier on the lower bound on reserves

function td = compute_td(b0,T,P,shock,previousrun)

% aux
P.b0 = b0;
reltol = read_opt(P,'reltol',1e-5);
abstol = read_opt(P,'abstol',1e-5); % set to 1e-9
NMax   = read_opt(P,'NMax',1200);
options = bvpset('RelTol',reltol,'AbsTol',abstol,'NMax',NMax);


%%% Set initial guess
% 1. ad-hoc guess
t_guess   = linspace(0,T,30); %[linspace(0,T/2,300) linspace(0,T/2+1e-3,300)];
bT_guess  = read_opt(P,'bT_guess',P.b0); a_guess = P.h_inv(bT_guess);
lam_guess = P.Vp(a_guess) / P.hp(a_guess);
y_guess   = [bT_guess, a_guess, 1/(P.rho+P.delta),1/(P.rho+P.delta), lam_guess, 0, 0, P.rho];
P.n_y     = length(y_guess); 

% 2. load guess from previous run
from_prev_run = 0;
if exist('previousrun','var')
    if isstruct(previousrun)
        % use guess from previous run
        t_guess = previousrun.t;
        y_guess = previousrun.y_all(:,end)';
        from_prev_run = 1;
    end
end

% 3. Set the functions used for bvp4c
myodefun = @(t,y) odefun(t,y,P,shock);
mybcfun  = @(y0,yT) bcfun(y0,yT,P,shock);

% 4. Special cases:
% 4a. With reserve target, add additional parameter to solve for
if strcmp(P.mode,'opt_res_acc')
    t_guess  = sort([t_guess P.res_goal_T P.res_goal_T]);
    myodefun = @(t,y,region) odefun(t,y,P,shock);
end
% 4b. With uncertainty over shocks, need to double the state space
if strcmp(P.mode,'opt_unc')
    
    if from_prev_run == 0
        t_guess  = sort([t_guess P.unc_T P.unc_T]);
        y_guess  = [y_guess, y_guess];
    end
    assert(sum(t_guess == P.unc_T) == 2);
    
    % prepare second P with alternative shock
    P2 = P; P2.D = P.unc_D2; shock2 = P.unc_shock2;
    if isfield(P,'unc_res_lbar_rstardot2')
        P2.res_lbar_rstardot = P.unc_res_lbar_rstardot2;
    end
    myodefun = @(t,y,region) [odefun(t,y(1:P.n_y),P,shock); odefun(t,y(P.n_y+1:end),P2,shock2)];
    mybcfun  = @(y0,yT) bcfun(y0,yT,P,shock);
    
    % test:
    mybcfun([y_guess',y_guess'],[y_guess',y_guess']);
end
% 4c. With limited commitment, need to add dates where CB can reoptimize
if strcmp(P.mode,'opt_LC')
    t_guess  = sort([t_guess P.LC_dates(:)' P.LC_dates(:)']);
    myodefun = @(t,y,region) odefun(t,y,P,shock);
end


solinit = bvpinit(t_guess,y_guess);

tic;
sol = bvp4c(myodefun,...
            mybcfun,...
            solinit,...
            options);
toc;

% read out the rest
td.y_all = sol.y;
td.t     = sol.x;i=1;
td.b     = sol.y(i,:);i=i+1;
td.a     = sol.y(i,:);i=i+1;
td.qstar = sol.y(i,:);i=i+1;
td.q     = sol.y(i,:);i=i+1;
td.lam   = sol.y(i,:);i=i+1;
td.mu    = sol.y(i,:); i=i+1;
td.chi   = sol.y(i,:);i=i+1;
td.rr    = sol.y(i,:);i=i+1;
if strcmp(P.mode, 'opt_unc')
    td.b2   = sol.y(i,:);i=i+1;
    td.a2   = sol.y(i,:);i=i+1;
    td.qstar2 = sol.y(i,:);i=i+1;
    td.q2   = sol.y(i,:);i=i+1;
    td.lam2 = sol.y(i,:);i=i+1;
    td.mu2  = sol.y(i,:); i=i+1;
    td.chi2 = sol.y(i,:);i=i+1;
    td.rr2  = sol.y(i,:);i=i+1;
end

% r, rstar, bI, bIH
for i = 1:length(td.t)
    t           = td.t(i);
    [r,rstar,rho] = compute_aux(t,td.y_all(:,i),P,shock);
    td.r(i)     = r;
    td.rstar(i) = rstar;
    td.rho(i)   = rho;
    
    if P.res_lbar > -Inf
        td.r(i) = td.rr(i);
        r = td.r(i);
    end
    
    D   = P.D.fn(r,rstar,td.q(i),td.qstar(i),t);
    Dr  = P.D.r(r,rstar,td.q(i),td.qstar(i),t);
    DHr = P.D.Hr(r,rstar,td.q(i),td.qstar(i),t);
    td.zeta(i) = ( (D + Dr*(r-rstar))*td.lam(i) - td.mu(i))/(Dr+DHr);
    td.bI(i)    = P.D.plots(r,rstar,td.q(i),td.qstar(i),t);
    td.bIH(i)   = P.D.Hfn(r,rstar,td.q(i),td.qstar(i),t);
    
    if strcmp(P.mode, 'opt_unc')
        [r,rstar,rho] = compute_aux(t,td.y_all(9:end,i),P,shock2);
        td.rstar2(i) = rstar;
    end
    
end

% solve for more variables
switch P.addl_vars
    case 'StoneGeary'
        td = compute_addl_vars_StoneGeary(P,td);
    case 'homebias'
        td = compute_addl_vars_homebias(P,td);
end


end



%%% ODEFUN %%%
function ydot = odefun(t,y,P,shock)
% this function computes the derivative of the ode system
i=1;
b=y(i);i=i+1;
a=y(i);i=i+1;
qstar=y(i);i=i+1;
q=y(i);i=i+1;
lam=y(i);i=i+1;
mu=y(i);i=i+1;
chi=y(i);i=i+1;
rr=y(i);i=i+1;

% compute effective sigma (this is needed in the Euler equation)
%  (a corrected version of sigma in case sigma ? 1 since then the RER
%  matters)
elast_p   = P.pp_fn(exp(a))*exp(a)/P.p_fn(exp(a)); % elasticity of NT price w.r.t. Ricardian dollar spending
sigma_eff = P.sigma + (1-P.alpha)*(1-P.sigma)*elast_p;

if P.res_lbar == -Inf
    % compute r (could be optimal, could be equilibrium, could be ad-hoc)
    [r,rstar,rho] = compute_aux(t,y,P,shock);

    % compute bI
    bI = P.D.fn(r,rstar,q,qstar,t);

    %%% ODEs
    % eqm objects:
    bdot     = rstar*b - P.h(a) - bI*(r-rstar);
    adot     = (r - P.rho)/sigma_eff;
    qstardot = (rstar+P.delta)*qstar - 1;
    qdot     = (r+P.delta)*q - 1;
    rrdot    = 0;

    % multipliers:
    if strcmp(P.mode,'opt') || strcmp(P.mode,'opt_iniER') || strcmp(P.mode,'tau0') ...
            || strcmp(P.mode,'opt_res_acc') || strcmp(P.mode,'opt_unc') || strcmp(P.mode,'opt_LC')

        gamma  = lam*(r-rstar);

        lamdot = lam*(rho - rstar);
        mudot  = P.rho*mu - P.Vp(a) + lam*P.hp(a);
        chidot = P.rho*chi + gamma * P.D.q(r,rstar,q,qstar,t) - r*chi;
    elseif strcmp(P.mode,'eqm')
        mudot  = 0;
        chidot = 0;
        lamdot = 0;
    end

else % lower bound active
    
    % compute rstar
    rstar = P.rho;
    fields = fieldnames(shock);
    for i = 1:length(fields)
        cur_field = fields{i};
        if strcmp(cur_field,'rstar')
            rstar = shock.rstar(t);
        else
            eval([cur_field ' = ' num2str(shock.(cur_field)(t)) ';']);
        end
    end
    rstardot = P.res_lbar_rstardot(t);
    qstardot = (rstar+P.delta)*qstar - 1;
    qdot     = (rr+P.delta)*q - 1;
    
    % auxiliary terms
    D  = P.D.fn(rr,rstar,q,qstar,t); DH = P.D.Hfn(rr,rstar,q,qstar,t);
    Dr = P.D.r(rr,rstar,q,qstar,t); DHr = P.D.Hr(rr,rstar,q,qstar,t);
    Drstar = P.D.rstar(rr,rstar,q,qstar,t); DHrstar = P.D.Hrstar(rr,rstar,q,qstar,t);
    Dq = P.D.q(rr,rstar,q,qstar,t); DHq = P.D.Hq(rr,rstar,q,qstar,t);
    Dqstar = P.D.qstar(rr,rstar,q,qstar,t); DHqstar = P.D.Hqstar(rr,rstar,q,qstar,t);
    Dt = P.D.t(rr,rstar,q,qstar,t); DHt = P.D.Ht(rr,rstar,q,qstar,t);
    Drt = 0; % load as function
    
    Dt_total  = Dt + Dq*qdot + Dqstar*qstardot + Drstar*rstardot;
    DHt_total = DHt + DHq*qdot + DHqstar*qstardot + DHrstar*rstardot;
    
    % compute reserves
    res = b + D + DH;
    
    % chi not needed (shouldnt have been there from get go)
    chidot = 0; % not needed!
    

    
    if res > P.res_lbar
        zeta = 0;
    else
        zeta = ( (D + Dr*(rr-rstar))*lam - mu)/(Dr+DHr);
        % rule out negative ones
        zeta = max(zeta,0);
    end
    
    % write derivatives
    bdot     = rstar*b - P.h(a) - D*(rr-rstar);
    adot     = (rr - P.rho)/sigma_eff;
    lamdot   = (P.rho-rstar)*lam - zeta;
    mudot    = P.rho*mu - P.Vp(a) + lam*P.hp(a);
    
    if res > P.res_lbar || zeta == 0
        rrdot = (mudot/lam - Dt_total - Drt*(rr-rstar) + Dr*rstardot - (D + Dr*(rr-rstar))*lamdot/lam )/(2*Dr);
    else
        all_deriv = Dt_total + DHt_total;
        rrdot = - (bdot + all_deriv)/(Dr + DHr);
    end
    
end

% write derivatives
ydot = [bdot,adot,qstardot,qdot,lamdot,mudot,chidot,rrdot]';

assert(sum(isnan(ydot))==0);assert(sum(~isreal(ydot))==0);
end



%%% COMPUTE_AUX %%%
function [r,rstar,rho] = compute_aux(t,y,P,shock)
% this function computes auxiliary objects based on state variables

i=1;
b=y(i);i=i+1;
a=y(i);i=i+1;
qstar=y(i);i=i+1;
q=y(i);i=i+1;
lam=y(i);i=i+1;
mu=y(i);i=i+1;
chi=y(i);i=i+1;
rr=y(i);i=i+1;

% set default parameters
rho = P.rho;
rstar = P.rho;
reserves = 0;

fields = fieldnames(shock);
for i = 1:length(fields)
    cur_field = fields{i};
    if strcmp(cur_field,'rstar')
        rstar = shock.rstar(t);
    else
        eval([cur_field ' = ' num2str(shock.(cur_field)(t)) ';']);
    end
end

% compute optimal r
if strcmp(P.mode,'opt') || strcmp(P.mode,'opt_iniER') || strcmp(P.mode,'opt_res_acc') || strcmp(P.mode,'opt_unc') || strcmp(P.mode,'opt_LC')
    
    x = 1/lam * ( mu + chi*q );
    if P.D.explicit1 == true
        r = P.D.explicit1_fn(x,rstar,q,qstar,t);
    else
        myfun = @(r) P.D.fn(r,rstar,q,qstar,t) + (r-rstar)*P.D.r(r,rstar,q,qstar,t) - x;
        r = fzero(myfun,rstar);
    end
elseif strcmp(P.mode,'eqm')
    x = reserves-b;
    if P.D.explicit3 == true
        r = P.D.explicit3_fn(x,rstar,q,qstar,t);
    else
        myfun = @(r) P.D.fn(r,rstar,q,qstar,t) + P.D.Hfn(r,rstar,q,qstar,t) - x;
        r = fzero(myfun,rstar);
    end
elseif strcmp(P.mode,'tau0')
    r = rstar;
end

% possible overwrite it if r is in shock (to get 'ad_hoc' policies like tau
% = 0 or even more complicated)
fields = fieldnames(shock);
for i = 1:length(fields)
    cur_field = fields{i};
    if strcmp(cur_field,'rstar')
        rstar = shock.rstar(t);
    else
        eval([cur_field ' = ' num2str(shock.(cur_field)(t)) ';']);
    end
end

assert(sum(isnan([r, rstar, rho]))==0);assert(sum(~isreal([r, rstar, rho]))==0);
end



%%% BCFUN %%%
function res = bcfun(yL,yR,P,shock)

% this function imposes boundary conditions

% aux: compute initial rstar
t0 = 0;
rstar = P.rho;
fields = fieldnames(shock);
for i = 1:length(fields)
    cur_field = fields{i};
    if strcmp(cur_field,'rstar')
        rstar = shock.rstar(t0);
    else
        eval([cur_field ' = ' num2str(shock.(cur_field)(t0)) ';']);
    end
end

% read out values
b_T=yR(1,end);a_T=yR(2,end);lam_T=yR(5,end);mu_T=yR(6,end);

% common boundary conditions
k=1;
res(k) = yL(1,1) - P.b0;k=k+1; % b(0) given
res(k) = P.rho*b_T - P.h(a_T);k=k+1; % bT aT compatibility
res(k) = (P.rho+P.delta)*yR(3,end) - 1;k=k+1; % terminal condition for qstar
res(k) = (P.rho+P.delta)*yR(4,end) - 1;k=k+1; % terminal condition for q
res(k) = yL(6,1) - 0;k=k+1;     % mu_0 initial condition
res(k) = yR(7,end) - 0;k=k+1;   % chi terminal condition
res(k) = yL(8,1) - rstar;k=k+1; % rr initial condition
    
% specific boundary conditions
if strcmp(P.mode,'opt') || strcmp(P.mode,'opt_iniER') || strcmp(P.mode,'tau0') ...
        || strcmp(P.mode,'opt_res_acc') || strcmp(P.mode,'opt_unc') || strcmp(P.mode,'opt_LC')
    
    % add optimality at infinity
    res(k) = P.rho*mu_T - P.Vp(a_T) + lam_T*P.hp(a_T);k=k+1;
    
    if strcmp(P.mode,'opt_iniER')
        % assume before shock: b0 was nfa, bI was zero, h(a)=rho*b0
        % this frees up mu_0 !
        a0 = yL(2,1);
        
        % before shock:
        a_before   = P.h_inv(P.rho*P.b0);
        RER_before = P.RER_fn(a_before); %(P.alpha + (1-P.alpha)*exp(a_before))^(-(1-P.alpha));
        RER_target = RER_before*(1+P.iniER);
        a_target   = P.RER_inv_fn(RER_target); %log( (RER_target^(-1/(1-P.alpha)) - P.alpha)/(1-P.alpha) );
        res(5) = a_target - a0;
    end
    
    if strcmp(P.mode,'opt_res_acc') 
        % this works with 2 regions. Therefore, need an additional 8 constraints! 
        for i = [1,2,3,4,6,7,8] % all except lambda
            res(i+P.n_y) = yR(i,1) - yL(i,2); % continuity
        end
        
        % load bI and bIH
        [r,rstar,~] = compute_aux(P.res_goal_T,yL(:,end),P,shock);
        b_I =  P.D.fn(r,rstar,yL(4,end),yL(3,end),P.res_goal_T);
        b_IH = P.D.Hfn(r,rstar,yL(4,end),yL(3,end),P.res_goal_T);
        
        % add condition instead of lambda: b needs to hit target!
        res(5+P.n_y) = yR(1,1)+b_I+b_IH - P.res_goal - P.b0;
    end
    
    if strcmp(P.mode,'opt_LC')
        % for every date in P.LC_dates, need an additional 8 constraints
        % idea: assume everything is continuous except that mu = 0 after
        % reoptimization.
        n_regions = length(P.LC_dates)+1;
        for region = 1:(n_regions-1)
            indices = [1:5, 7:8];
            for l = indices
                res(k) = yR(l,region) - yL(l,region+1); k=k+1;
            end
            res(k) = yL(6,region+1) - 0;k=k+1; % mu_T = 0 at times of reoptimization
        end
    end
    
    if strcmp(P.mode,'opt_unc') % case with uncertainty resolved at t=T
        
        % check sizes of inputs
        sL = size(yL); sR = size(yR);
        assert(sL(2)==2); assert(sR(2)==2);
        
        % 8 'common' conditions are already set for indices 1:8.
        % Set the same 8 again for indices 9:16
        res(k) = yL(P.n_y+1,1) - P.b0;k=k+1; % b(0) given
        res(k) = P.rho*yR(P.n_y+1,end) - P.h(yR(P.n_y+2,end));k=k+1; % bT aT compatibility
        res(k) = (P.rho+P.delta)*yR(P.n_y+3,end) - 1;k=k+1; % terminal condition for qstar
        res(k) = (P.rho+P.delta)*yR(P.n_y+4,end) - 1;k=k+1; % terminal condition for q
        res(k) = yL(P.n_y+6,1) - 0;k=k+1; % mu_0 initial condition
        res(k) = yR(P.n_y+7,end) - 0;k=k+1; % chi terminal condition
        res(k) = P.rho*yR(P.n_y+6,end) - P.Vp(yR(P.n_y+2,end)) + yR(P.n_y+5,end)*P.hp(yR(P.n_y+2,end));k=k+1;
        res(k) = yL(P.n_y+8,1) - rstar;k=k+1; % rr initial condition


        % This works with 2 regions and the doubled state space. 
        % Therefore, need an additional 16 constraints.

        % 4 equalities before the shock across the two sets of indices:
        %  qstar, q, a, lam
        set = [2,3,4,5];%[1,2,3,4,5,7];
        for i = set
            res(k) = yR(i,1) - yR(P.n_y+i,1);k=k+1;
        end

        % 4 variables are continuous "in expectation"
        % lambda, q, qstar + marginal utility
        set = [3,4,5];
        for i = set
            res(k) = yR(i,1) - P.unc_p1 * yL(i,2) - P.unc_p2 * yL(P.n_y+i,2);k=k+1;
        end
        res(k) = P.unc_MU(yR(2,1)) - P.unc_p1 * P.unc_MU(yL(2,2)) - P.unc_p2 * P.unc_MU(yL(P.n_y+2,2)); k=k+1; % Euler

        % 6 conditions since b, chi, rr are continuous
        indices = [1,7,8];
        for l = indices
            res(k) = yR(l,1) - yL(l,2); k=k+1;
            res(k) = yR(P.n_y+l,1) - yL(P.n_y+l,2); k=k+1;
        end

        % 2 conditions on mu coming from Euler
        res(k) = yR(6,1) - exp(yL(2,2) - yR(2,1)) * yL(6,2);k=k+1;
        res(k) = yL(6,2) - exp(yL(P.n_y+2,2) - yL(2,2)) * yL(P.n_y+6,2);k=k+1;
        
    end
    
elseif strcmp(P.mode,'eqm')
    % set lam to 1 for simplicity
    res(k) = lam_T - 1;
end

assert(sum(isnan(res))==0);assert(sum(~isreal(res))==0);

end







