/* ara.prg: approximates the policy function for consumption with both a and k
                        as state variables, using an orthogonal collocation scheme. 
			a is an ar(1) process */

/* initialize program */

new;
library optmum,pgraph,nlsys;
graphset;
nlset;
_pdate = "";
_psymsiz =  3;
_paxht   = .2;
_pnumht  = .15;
_ptitlht = .3; 
_paxes   = { 1, 1 }; 

/* set dimension of model  */

nvars = 5;
nshocks = 1;
eulers = 1;						@ indicates equation number(s) of euler equations @
neulers = rows(eulers);					@ counts # of euler equations in the model  @
nobservables = 2;

@ establish parameters @

alp = 0.33;
bet = 0.96;
rho = 0.8;
sige = 0.0067;
p = alp|bet|rho|sige;

/* load log-linear approximation routines */

@ Note: p, xbar and tme must be global before loglin.prc can be compiled in stogrow.src.  @

xbar = zeros(nvars,1);
tme = 0;

#include stogrowm_nonlin_exact.src

/* given drawing pdraw, obtain non-linear model solution */

pdraw = p;

ss = steady(pdraw);
kss = ss[4];
css = ss[2];
ass = ss[5];

/* establish ranges for states */

omegaz = (1/(1-rho))*3*sige;
zl = -omegaz;
zh = omegaz;

omegak = 0.25*kss;
kl = kss-omegak;
kh = kss+omegak;

/* locate zeros of polynomials */
ordz = 4; 
ordk = 5;

nzersk=ordk;
nzersz=ordz;
zersk=zeros(ordk,1);
zersz=zeros(ordz,1);

iter=1;do while iter<=nzersk;
zersk[iter]=-1.0*cos((2*iter-1)*PI/(2*ordk));
iter=iter+1;endo;

iter=1;do while iter<=nzersz;
zersz[iter]=-1.0*cos((2*iter-1)*PI/(2*ordz));
iter=iter+1;endo;

//read points for gauss-hermite quadrature.
num_gh_pts=7;
ghermite=zeros(num_gh_pts,2);
ghermite =
-2.65196135683523381842974231403786689E+0000~9.71781245099519251466613223300328173E-0004|
-1.67355162876747143307909482246031985E+0000~5.45155828191270785954003486040164717E-0002|
-8.16287882858964697341264127317117527E-0001~4.25607252610127773095882730558514595E-0001|
-1.05978634916828728441807676992045652E-0016~8.10264617556807453802036889101145789E-0001|
8.16287882858964586318961664801463485E-0001~4.25607252610127773095882730558514595E-0001|
1.67355162876747165512369974749162793E+0000~5.45155828191271202287637720473867375E-0002|
2.65196135683523381842974231403786689E+0000~9.71781245099518926205961477648997970E-0004
;

/* obtain linear solution to the model, construct starting values */

{f,q} = modelsol(p);

sigk = f[2,4];
sigz = f[2,5];

startval = zeros(ordz*ordk,1);
startval[1] = css;
startval[2] = omegak*sigk*css/kss;
startval[ordk+1] = omegaz*sige*css;
startval[ordk+2] = 0.5*(omegak*sigk*css/kss)*omegaz*sigz*css;

/* construct policy function */

gamopt = polfcn(ss,f,startval);	

/* reset grids for states for graphing */

zsteps = 41;
zstep = (zh-zl)/(zsteps-1);
zrange = seqa(zl,zstep,zsteps);
znorm = zrange/omegaz;
nz = rows(zrange);

ksteps = 101;
kstep = (kh-kl)/(ksteps-1);
krange = seqa(kl,kstep,ksteps);
knorm = (krange-kss)/omegak;
nk = rows(krange);

optcs = zeros(ksteps,zsteps);
excs = zeros(ksteps,zsteps);
i=1; do while i<=ksteps;
    j=1; do while j<=zsteps;
        optcs[i,j] = ceeofzk(znorm[j],knorm[i],gamopt);
        excs[i,j] = (1-alp*bet)*(exp(zrange[j])*(krange[i].^alp));
    j=j+1; endo;
i=i+1; endo;

"measure accuracy:";?;

err = maxc(maxc(abs(optcs-excs)));

err/css;


/* collect procedures here */

proc ceeofzk(z,k,gam);
    /* calculates c as a function of a. spits out c, and value of polynomials over the range of k */

    local  nteesz,nteesk,teesz,teesk,j,tmat,gammat,i,cee;

    nteesz = ordz;
    nteesk = ordk; 
    teesz = zeros(nteesz,1);
    teesk = zeros(nteesk,1);
    teesz[1] = 1;
    teesz[2] = z;
    teesk[1] = 1;
    teesk[2] = k;
    j=3; do while j<=nteesz;
        teesz[j] = 2*z*teesz[j-1]-teesz[j-2];
    j=j+1; endo;    
    j=3; do while j<=nteesk;
        teesk[j] = 2*k*teesk[j-1]-teesk[j-2];
    j=j+1; endo;    

    tmat = teesk*teesz'; 
    gammat = zeros(nteesk,nteesz);
    i=1; do while i<=nteesz;
        gammat[.,i] = gam[(i-1)*nteesk+1:i*nteesk]; 
    i=i+1; endo;
    
    cee = sumc(sumc(tmat.*gammat));

    retp(cee);
endp;

proc feval(gam);

    local nrowsz,nrowsk,ncees,cee,f,i,j,rhs,
    expec,ztilde,ktilde,z_t,k_t,c_t,z_tp1,k_tp1,c_tp1,
    kk;

nrowsz = rows(zersz);
nrowsk = rows(zersk); 
ncees = nrowsz*nrowsk;
cee = zeros(ncees,1);
f = zeros(ncees,1);

i=1; do while i<=nrowsz;
    j=1; do while j<=nrowsk;
    cee[(i-1)*nrowsk+j] = ceeofzk(zersz[i],zersk[j],gam);
	z_t=omegaz*zersz[i];
	k_t=kss+omegak*zersk[j] ;
	c_t=cee[(i-1)*nrowsk+j];
	k_tp1=exp(z_t)*(k_t^alp)-c_t;
	ktilde=(k_tp1-kss)/omegak ;
	
	expec=zeros(num_gh_pts,1);
	kk=1;do while kk<=num_gh_pts;
		z_tp1=rho*z_t+sqrt(2)*sige*ghermite[kk,1];
		ztilde=z_tp1/omegaz;
		c_tp1=ceeofzk(ztilde,ktilde,gam);
		expec[kk]=(exp(z_tp1)*(k_tp1^(alp-1))/c_tp1)*ghermite[kk,2];
	kk=kk+1;endo;

	rhs=bet*alp*sumc(expec)/sqrt(PI);
	
    rhs = maxc( rhs|((exp(z_t)*(k_t^alp))^-1) );
    f[(i-1)*nrowsk+j] = cee[(i-1)*nrowsk+j]^-1 - rhs;

    j=j+1; endo;
i=i+1; endo;
retp(f);

endp;



proc polfcn(ss,f,startval);
	/* produces policy function, taking linearized model as input to construct ranges over a, k and starting values for gamma */

	local gamopt,fopt,gopt,retcode,ceeopt,i,j;

	/* the following are used as inputs to construct starting values */

	{ gamopt,fopt,gopt,retcode } = nlsys(&feval,startval);


	retp(gamopt);

endp;