% MODEL: Compute the model derivatives using symbolic differentiation
%
% usage
%
% [fyn, fxn, fypn, fxpn, eta, param] = model(param)
%
% where
%
% param = object containing all parameter values
 

function [fyn fxn  fypn fxpn eta param] = model(param)

neps = 8;

%Compute Model SS
[ss param] = model_ss(param);

%Put parameters as variables in workspace
assign_vars(param);

%Declare symbols and X and Y variables
syms V   K   KT   TK   TK_l   TN   TN_l  X_l  D_l  I_l  G ET A MU 
syms V_p K_p KT_p TK_p TK_lp  TN_p TN_lp X_lp D_lp I_lp G_p ET_p A_p MU_p
syms XX   N   C   D   I   U    W   R   Y   L1   L2   L3   L4    L5   TAU   GT     EPST
syms XX_p N_p C_p D_p I_p U_p  W_p R_p Y_p L1_p L2_p L3_p L4_p  L5_p TAU_p GT_p   EPST_p


ce_X =  [V   K   KT   TN   TN_l   TK   TK_l  X_l  D_l  I_l  A   MU    G   EPST];
ce_XP = [V_p K_p KT_p TN_p TN_lp  TK_p TK_lp X_lp D_lp I_lp A_p MU_p  G_p EPST_p];
ce_Y =  [XX   N   C   D   I   U    W   R   Y   L1   L2   L3   L4   L5   TAU GT];
ce_YP = [XX_p N_p C_p D_p I_p U_p  W_p R_p Y_p L1_p L2_p L3_p L4_p L5_p TAU_p GT_p];


%Add in TAX news
news_per = 6;
[nvar, nvar_p, t_news] = news_process(news_per, 'EPST_');
assign_vars(nvar);
assign_vars(nvar_p);
ce_X = [ce_X, transpose(nvar)];
ce_XP = [ce_XP, transpose(nvar_p)];


%Add in Tech news
news_per = 6;
[nvar, nvar_p, a_news] = news_process(news_per, 'EPSA_');
assign_vars(nvar);
assign_vars(nvar_p);
ce_X = [ce_X, transpose(nvar)];
ce_XP = [ce_XP, transpose(nvar_p)];

%Add in Pref news
news_per = 6;
[nvar, nvar_p, m_news] = news_process(news_per, 'EPSM_');
assign_vars(nvar);
assign_vars(nvar_p);
ce_X = [ce_X, transpose(nvar)];
ce_XP = [ce_XP, transpose(nvar_p)];

%Add in GOV't news
news_per = 6;
[nvar, nvar_p, g_news] = news_process(news_per, 'EPSG_');
assign_vars(nvar);
assign_vars(nvar_p);
ce_X = [ce_X, transpose(nvar)];
ce_XP = [ce_XP, transpose(nvar_p)];

%Some expressions
g1 = gam_z^-sig;
UX = (XX-b*X_l/gam_z)^-sig;
UX_p = (XX_p*gam_z-b*X_lp)^-sig;

%2nd Onder Functional Forms
psi = @(x) psi_p*(x-1) + psi_pp/2*(x-1)^2;
psi_1 = @(x) psi_p + psi_pp*(x-1);

phik = @(x)   phi_kk/2*(x-gam_z)^2;
phik_1 = @(x) phi_kk*(x-gam_z);

phiv = @(x) phi_vv/2*(x-gam_z)^2;
phiv_1 = @(x) phi_vv*(x-gam_z);


%Consumer FOC's
eq1 = UX - L1 - b*bet*UX_p;
eq2 = -MU*N^kappa + L4*(1-TN)*W;
eq3 = -L2 + bet*L1_p*g1*(1-nu)*(C_p/V_p)^nu + bet*L2_p*g1*(1-delta_v);
eq4 = -L4 + L2*(1-phiv(gam_z*D/D_l)-gam_z*(D/D_l)*phiv_1(gam_z*D/D_l)) + bet*L2_p*g1*(gam_z*D_p/D_lp)^2*phiv_1(gam_z*D_p/D_lp); 
eq6 = -L4 + L3*(1-phik(gam_z*I/I_l)-gam_z*(I/I_l)*phik_1(gam_z*I/I_l)) + bet*L3_p*g1*(gam_z*I_p/I_lp)^2*phik_1(gam_z*I_p/I_lp) +L5;
eq5 = -L3 + bet*L3_p*g1*(1-delta_k-psi(U_p)) + bet*L4_p*g1*(1-TK_p)*R_p*U_p;
eq7 = -L3*psi_1(U)+L4*(1-TK)*R;
eq8 = -L5 + bet*L5_p*g1*(1-delta_tau) + bet*L4_p*g1*TK_p*delta_tau;
eq9 = -L1*nu*(C/V)^(nu-1) + L4;

%Consumer Constraints
eq10 = XX - C^nu*V^(1-nu);
eq11 = V_p*gam_z - (1-delta_v       )*V - D*(1-phiv(gam_z*D/D_lp));
eq12 = K_p*gam_z - (1-delta_k-psi(U))*K - I*(1-phik(gam_z*I/I_lp));
eq13 = KT_p*gam_z -(1-delta_tau     )*KT - I;

%Firm FOC's
eq14 = W - (1-alph)*A*(U*K)^alph*N^-alph;
eq15 = R - alph*A*(U*K)^(alph-1)*N^(1-alph);

%RC
eq16 = C+I+D+G-Y;
eq17 = Y - A*(U*K)^alph*N^(1-alph);

%Exogenous Processes
eq18 = TN_p - (1-rho1n - rho2n)*tau_n - rho1n*TN - rho2n*TN_l - EPST_6_0p ; 
eq19 = TK_p - (1-rho1k - rho2k)*tau_k - rho1k*TK - rho2k*TK_l - EPST_6_0p; 
eq20 = log(G_p/(ybar*gbar)) - rhog*log(G/(ybar*gbar)) - EPSG_6_0p;
eq21 = GT - G - pi_g*log(TAU/tbar);
eq22 = log(A_p/A_bar) - rhoa*log(A/A_bar) - EPSA_6_0p;
eq23 = log(MU_p/mu_bar) - rhomu*log(MU/mu_bar) - EPSM_6_0p;
eq24 = TAU - TN*W*N - TK*(R*U*K - delta_tau*KT);

%Linking EQ's
eq25 = D_lp - D;
eq26 = I_lp - I;
eq27 = X_lp - XX;
eq28 = TN_lp - TN;
eq29 = TK_lp - TK;
eq30 = EPST_p;

%All FOC's
f = [eq1 eq2 eq3 eq4 eq5 eq6 eq7 eq8 eq9 eq10 eq11 eq12 eq13 eq14 eq15 eq16 eq17 eq18 eq19 eq20 eq21 eq22 eq23 eq24 eq25 eq26 eq27 eq28 eq29 eq30 transpose(t_news) transpose(g_news) transpose(a_news) transpose(m_news)];

%For log-linearization
xlog = [1:3, 8:13];
ylog = 1:length(ce_Y);
logvar = [ce_X(xlog) ce_Y(ylog) ce_XP(xlog) ce_YP(ylog)];

f = subs(f, logvar, exp(logvar));

%differentiate
fx =  jacobian(f, ce_X);
fy =  jacobian(f, ce_Y);
fxp = jacobian(f, ce_XP);
fyp = jacobian(f, ce_YP);

%Plug back into levels
fx =  subs(fx,  logvar, log(logvar));
fy =  subs(fy,  logvar, log(logvar));
fxp = subs(fxp, logvar, log(logvar));
fyp = subs(fyp, logvar, log(logvar));

%Numerical
fxn =  double(subs(fx,  [ce_Y ce_X ce_YP ce_XP], [ss, ss]));
fyn =  double(subs(fy,  [ce_Y ce_X ce_YP ce_XP], [ss, ss]));
fxpn = double(subs(fxp, [ce_Y ce_X ce_YP ce_XP], [ss, ss]));
fypn = double(subs(fyp, [ce_Y ce_X ce_YP ce_XP], [ss, ss]));


%Shocks
eta = zeros(length(ce_X),neps);
eta(4,1) = sigT;
eta(6,1) = sigT;
eta(14,1) = sigT;

eta(11,2) = sigA;
eta(12,3) = sigMU;
eta(13,4) = sigG;

eta(21,5) = sigT_ant;
eta(28,6) = sigA_ant;
eta(35,7) = sigMU_ant;
eta(42,8) = sigG_ant;


%**********************************************************************************
% MODEL_SS - Return the steady state of the model
%
% usage:
% 
% [ss, param] =model_ss(param)
%**********************************************************************************

function [ss, param] = model_ss(param)

%Put parameters as variables in workspace
assign_vars(param);

%Compute SS
A = A_bar;
n = n_bar;
g1 = gam_z^(-sig);
psi_p = (1 - bet*g1*(1-delta_k))/(bet*g1);
r = (1 - (bet*g1*tau_k*delta_tau)/(1-bet*g1*(1-delta_tau)))*psi_p/(1-tau_k);
k = n_bar*(r/(alph*A))^inv(alph-1);
i = k*(gam_z-1+delta_k);
y = A*k^alph*n^(1-alph);
g = gbar*y;
w = (1-alph)*A*(k/n)^alph;
d = dshare*(y-i-gbar*y);
c = (1-dshare)*(y-i-gbar*y);
v = d/(gam_z -1 + delta_v);
nu = bet*g1/(bet*g1 + inv(c/v)*(1-bet*g1*(1-delta_v)));

ktau = i/(gam_z-1+delta_tau);
x = c^nu*v^(1-nu);
l1 = (x-b*x/gam_z)^(-sig) - b*bet*(x*gam_z-b*x)^(-sig);

mu = l1*nu*(c/v)^(nu-1)*(1-tau_n)*w/(n^kappa);
tau = tau_n*w*n + tau_k*r*k - delta_tau*tau_k*ktau;

%lagrange multipliers
l2 = l1*nu*(c/v)^(nu-1);
l3 = l2*(1-tau_k)*r/psi_p;
l4 = l2;
l5 = l4-l3;

%cap util
u = 1;

%Parameters Determined in SS
param.psi_p = psi_p;
param.psi_pp = psi_p*psi_pp_p;
param.mu_bar = mu;
param.nu = nu;
param.ybar = y;
param.tbar = tau;

%Steady States of Shocks

EPST =0;

EPST_6_0 = 0;
EPST_6_1 = 0;
EPST_6_2 = 0;
EPST_6_3 = 0;
EPST_6_4 = 0;
EPST_6_5 = 0;
EPST_6_6 = 0;

EPSG_6_0 = 0;
EPSG_6_1 = 0;
EPSG_6_2 = 0;
EPSG_6_3 = 0;
EPSG_6_4 = 0;
EPSG_6_5 = 0;
EPSG_6_6 = 0;

EPSA_6_0 = 0;
EPSA_6_1 = 0;
EPSA_6_2 = 0;
EPSA_6_3 = 0;
EPSA_6_4 = 0;
EPSA_6_5 = 0;
EPSA_6_6 = 0;


EPSM_6_0 = 0;
EPSM_6_1 = 0;
EPSM_6_2 = 0;
EPSM_6_3 = 0;
EPSM_6_4 = 0;
EPSM_6_5 = 0;
EPSM_6_6 = 0;


%SS values for ce_Y
y_ss = [x,n,c,d,i,u,w,r,y,l1,l2,l3,l4,l5,tau,g];


%SS values fors ce_X
x_ss = [v,k,ktau,tau_n,tau_n,tau_k,tau_k,x,d,i,A,mu,g,EPST,EPST_6_0,EPST_6_1,EPST_6_2,EPST_6_3,EPST_6_4,EPST_6_5,EPST_6_6,EPSA_6_0,EPSA_6_1,EPSA_6_2,EPSA_6_3,EPSA_6_4,EPSA_6_5,EPSA_6_6,EPSM_6_0,EPSM_6_1,EPSM_6_2,EPSM_6_3,EPSM_6_4,EPSM_6_5,EPSM_6_6,EPSG_6_0,EPSG_6_1,EPSG_6_2,EPSG_6_3,EPSG_6_4,EPSG_6_5,EPSG_6_6];

%Combine
ss = [y_ss x_ss];



%**********************************************************************************
% NEWS_PROCESS: Create the AR represnetation of a disturbance announced
% advance.  Automatically creates and names relevant symbolic variables.
%
% usage
% 
% [nvar nvar_p x_process M] = news_process(news_per, vname)
%
% news_per = periods in advance
% vname = the prefix of the shock's name (to have _1, _2, etc appended)
%
% nvar      = vector of time t states created
% nvar_p    = vector of time t+1 states created
% x_process = AR relation between nvar and nvar_p
%**********************************************************************************
function [nvar nvar_p x_process] = news_process(news_per, vname)


M = zeros(1+news_per);
nvar = [];
nvar_p = [];
for jj = 1:news_per+1
        vrbl = [vname, num2str(news_per), '_', num2str(jj-1)];
        vrbl_p = [vrbl,'p'];
        
        %Put into vector
        nvar = [nvar; sym(vrbl)];
        nvar_p = [nvar_p; sym(vrbl_p)];
end
M(1:news_per,2:news_per+1) = eye(news_per);       
x_process = nvar_p - M*nvar;


%**********************************************************************************
% ASSIGN_VARS: assign variables in the current workspace for the elements of a data
% structure or symbolic variables contained in a symbolic vector.
%
% usage
%
% assign_vars(obj)
%
% where
% 
% obj = a symbolic vector (e.g. obj = [sym('a'), sym('b')^2];) or data structure (e.g. obj.a =1; obj.b = 2^2;); 
%**********************************************************************************

function out = assign_vars(F)

if isa(F, 'sym')
    try
        list_l = [findsym(F), '  '];
    catch
        s = lasterror;
        if strcmp(s.identifier, 'MATLAB:UndefinedFunction');
            out = F;
            return
        else
            rethrow(s);
        end
    end

    idx = [-1,findstr(list_l, ','), length(list_l)-1];
    for j = 1:length(idx)-1
        vname = list_l(idx(j)+2:idx(j+1)-1);
        if ~isempty(vname)
            assignin('caller',vname,sym(vname));
        end
    end


elseif isa(F,'struct')

    %Creat variables for each parameter
    param_list = fieldnames(F);
    param_val = struct2array(F);
    for j = 1:length(param_list)
        assignin('caller',param_list{j},param_val(j));
    end
end

