# This function solves the ACS model using the PLC method 3.3
# First Version: 2018-03-27
# Current Version: 2018-11-01
# Version Notes:

# Inputs:
#=
-params_in: Model parameters, in String->Float64 dictionary
-shocks_use: Series of 4 shocks, in String->Vector dictionary
-theta_0: Initial guess of solution, as Vector
-options_int_in: Function options with integer values, in String->Int64 dictionary
-options_float_in: Function options with non-integer values, in String->Float64 dictionary
=#

# Outputs:
#=
-theta_use: Solution parameters
-theta_keep: Matrix with the thetas for every iteration
-time_keep: Matrix of solution times
-flag: Error flags, 1= clustering error, 2=non-converged solution
-out_x_all: Matrix with the sequence of thetas inside the otpimization routine
-out_res_all: Matrix with the sequence of residual inside the optimization routine
-out_conv_all: Matrix with the sequence of convergence criteria inside the optimization routine
-out_L: Matrix with the thetas for the L iterations
-res_out_1: Value of the residual vector in the final iteration
=#

function PLC_solve(params_in::Dict{String}{Float64},shocks_use::Dict{String}{Array{Float64,1}},theta_0::Array{Float64,1},options_int_in::Dict{String}{Int64},options_float_in::Dict{String}{Float64},fixed_grid::Array{Float64,2},vSmolyakGridSpec::Array{Float64,1})

#--------------------------------------------------------------------------
# FIXED PARAMETERS
#--------------------------------------------------------------------------

sig_z  = params_in["sig_z"];
sig_g  = params_in["sig_g"];
sig_r  = params_in["sig_r"];
sig_d  = params_in["sig_d"];
y_ss   = params_in["y_ss"];
R_ss   = params_in["R_ss"];

#--------------------------------------------------------------------------
# OPTIONS
#--------------------------------------------------------------------------

L = options_int_in["L"]
M = options_int_in["M"]
skip1 = options_int_in["skip"]
k = options_int_in["k"]
cluster_d = options_int_in["cluster_d"]
trans = options_int_in["trans"]
tolx_in = options_float_in["tolx"]
tolf_in = options_float_in["tolf"]
tolg_in = options_float_in["tolg"]
maxIt_in = options_int_in["maxIt"]
ext     = options_int_in["ext"]

#--------------------------------------------------
# Gauss-Hermite Nodes and Weights
#--------------------------------------------------
if options_int_in["int_d"]==1
EG_nodes,ER_nodes,EZ_nodes,ED_nodes, GH_weight = get_GH_nodes(options_int_in["Q"],sig_r,sig_z,sig_g,sig_d);
elseif options_int_in["int_d"]==2
EG_nodes,ER_nodes,EZ_nodes,ED_nodes, GH_weight = get_Monomial_nodes("M1",params_in,4);
elseif options_int_in["int_d"]==3
vcv=diagm(0=>[sig_r^2,sig_z^2,sig_g^2,sig_d^2])
n_nodes,nodes, weight = Monomials_2_N4(vcv);
ER_nodes=nodes[:,1]
EZ_nodes=nodes[:,2]
EG_nodes=nodes[:,3]
ED_nodes=nodes[:,4]
GH_weight=vec(weight)
end

#--------------------------------------------------------------------------
#                       Check Grid
#--------------------------------------------------------------------------

if ext==1 && k>0
    error("k>0 with external grid")
end

if ext==1
    M=size(fixed_grid,1)
end

#--------------------------------------------------------------------------
#                       INITIAL CONDITIONS
#--------------------------------------------------------------------------

# Set Initial Values for Simulation

init=[R_ss,y_ss,0,0,0];

theta_use=theta_0;

time_keep=zeros(L+k+1,4);
time_one=0.0;
time_clus=0.0;
flag=zeros(L+k+2)
Mean_res=zeros(3);
VCV=Matrix{Float64}(I,3,3);
inv_VCV=Matrix{Float64}(I,3,3);
err_weight=Matrix{Float64}(I,3,3);
CGA_GRID=zeros(M,6)
theta_keep=zeros(24,L+k+2);
theta_keep[:,1]=theta_0
ZLB_hat=log(1/R_ss)
out_L=zeros(L,24)
out_x_all=zeros((maxIt_in+1)*(k+1),24)
res_out_1=Array{Float64}(undef,1)
out_res_all=zeros((maxIt_in+1)*(k+1))
out_conv_all=zeros((maxIt_in+1)*(k+1),3)
iter_c=0
cweights=ones(M)
GRID=Array{Float64,2}

save_GMM=zeros(L+k+1,3,3)

c_ss=params_in["c_ss"]
tau=params_in["tau"]

#err_weight[1,1]=1/c_ss^-tau
#err_weight[2,2]=1/c_ss^-tau
#err_weight[3,3]=1/c_ss^-tau

err_weight=Matrix{Float64}(I,3,3);
err_weight[3,3]=1.0;
err_weight=err_weight./sum(err_weight)*3



# Start the iteration

save_R_GH=zeros(k+1,M*3)
save_R_MONO=zeros(k+1,M*3)
save_distJ=zeros(k+1,M*3)
save_sizeJ_GH=zeros(k+1,M*3)
save_GH_ZLB=zeros(k+1,M)
save_MONO_ZLB=zeros(k+1,M)

for iter=1:(L+k+1)

    #println(string("*** Solving PLC Iteration = ", iter ," ***"));

    #--------------------------------------------------------------------------
    #               Simulate and Cluster
    #--------------------------------------------------------------------------

    ini_sim=time_ns()

    if ext==1
        GRID, cweights = fScaleSmolyakGrid(fixed_grid,vSmolyakGridSpec,params_in)
        #GRID=fixed_grid
        #cweights=fixed_w
    else

    simout::Dict{String}{Array{Float64,1}} = simulate_plc(params_in::Dict{String}{Float64},theta_use::Array{Float64,1},shocks_use::Dict{String}{Array{Float64,1}},init,options_int_in["Tsim"]*skip1+1,options_int_in["Tdrop"]);
    #simout = simulate_plc(params_in::Dict{String}{Float64},theta_use::Array{Float64,1},shocks_use::Dict{String}{Array{Float64,1}},init,options_int_in["Tsim"]*skip1+1,options_int_in["Tdrop"]);

    #--------------------------------------------------------------------------
    # CREATE SOLUTION GRID
    #--------------------------------------------------------------------------

    simtemp=zeros(size(simout["Rhat"][1:end-1],1),6);

    simtemp[:,1]     = simout["Rhat"][1:end-1];    # Rhat_lag;
    simtemp[:,2]     = simout["yhat"][1:end-1];    # yhat_lag;
    simtemp[:,3]     = simout["er"][2:end];        # er
    simtemp[:,4]     = simout["g"][2:end];         # g
    simtemp[:,5]     = simout["z"][2:end];         # z
    simtemp[:,6]     = simout["d"][2:end];         # d

    if skip1>1
        draw=collect(1:size(simout["Rhat"][1:end-1],1))
        simtemp=simtemp[(draw.%skip1).==1,:]
    end

    time_keep[iter,1]=signed(time_ns()-ini_sim)/1000000000;

    ini_clus=time_ns();
    try

        if cluster_d==1
            GRID,temp2,cweights = cga_grid_kms(simtemp,M,trans);
        elseif cluster_d==2
            GRID,temp2,cweights = cga_grid_km(simtemp,M,trans);
        elseif cluster_d==3
            GRID = tsga_grid(simtemp,M);
        elseif cluster_d==4
            GRID,temp2,cweights = cga_grid_h(simtemp,M,trans);
        end

    catch
       flag[iter:end].=1
        break
    end

    time_keep[iter,2]=signed(time_ns()-ini_clus)/1000000000;


    if options_int_in["ind_cweight"] ==0
        cweights=ones(M)
    end

    end

    ini_b=time_ns();

    RHATLAGGRID = GRID[:,1];
    YHATLAGGRID = GRID[:,2];
    EGRID       = GRID[:,3];
    GGRID       = GRID[:,4];
    ZGRID       = GRID[:,5];
    DGRID       = GRID[:,6];

    #--------------------------------------------------------------------------
    # Minimization
    #--------------------------------------------------------------------------
    #println("PLC Algorithm. Solving for the Policy Function. Please wait...");

    e_grid=vec(EGRID)
    g_grid=vec(GGRID)
    d_grid=vec(DGRID)
    z_grid=vec(ZGRID)
    Rhat_grid=vec(RHATLAGGRID)
    yhat_grid=vec(YHATLAGGRID)

    #GH_weight_1=vec(GH_weight)

    time_keep[iter,3]=signed(time_ns()-ini_b)/1000000000;;

    ini_solve=time_ns()
    function model(p::Array{Float64,1})

        R_temp1::Array{Float64,1},jac_temp1::Array{Float64,2}=system_piecewise_linear_jacob_33_revised(p,params_in,e_grid,g_grid,z_grid,d_grid,Rhat_grid,yhat_grid,EG_nodes,ER_nodes,EZ_nodes,ED_nodes,GH_weight);
        R_temp=R_temp1
        jac_temp=jac_temp1

        #Weighting the residuals by cluster weights (cweights) and GMM weights (err_weights)
        #If this options are not active then this should not affect the residuals.
        for i=1:M
            R_temp1[i] = R_temp[i]*(cweights[i].^0.5)
            R_temp1[i+M] = R_temp[i+M]*(cweights[i].^0.5)
#            R_temp1[i+2*M] = R_temp[i+2*M]*(cweights[i].^0.5)
            R_temp[i] = R_temp1[i]*err_weight[1,1]+R_temp1[i+M]*err_weight[2,1] #+R_temp1[i+2*M]*err_weight[3,1]
            R_temp[i+M] = R_temp1[i]*err_weight[1,2]+R_temp1[i+M]*err_weight[2,2] #+R_temp1[i+2*M]*err_weight[3,2]
#            R_temp[i+2*M] = R_temp1[i]*err_weight[1,3]+R_temp1[i+M]*err_weight[2,3]+R_temp1[i+2*M]*err_weight[3,3]
            for j=1:24
                jac_temp1[i,j]     = jac_temp1[i,j]*(cweights[i].^0.5)
                jac_temp1[i+M,j]   = jac_temp1[i+M,j]*(cweights[i].^0.5)
                #jac_temp1[i+2*M,j] = jac_temp1[i+2*M,j]*(cweights[i].^0.5)
#                jac_temp1[i,j]     = jac_temp1[i,j]*(cweights[i])
#                jac_temp1[i+M,j]   = jac_temp1[i+M,j]*(cweights[i])
#                jac_temp1[i+2*M,j] = jac_temp1[i+2*M,j]*(cweights[i])

                jac_temp[i,j]      = jac_temp1[i,j]*err_weight[1,1]+jac_temp1[i+M,j]*err_weight[2,1] #+jac_temp1[i+2*M,j]*err_weight[3,1]
                jac_temp[i+M,j]    = jac_temp1[i,j]*err_weight[1,2]+jac_temp1[i+M,j]*err_weight[2,2] #+jac_temp1[i+2*M,j]*err_weight[3,2]
#                jac_temp[i+2*M,j]  = jac_temp1[i,j]*err_weight[1,3]+jac_temp1[i+M,j]*err_weight[2,3]+jac_temp1[i+2*M,j]*err_weight[3,3]
            end
        end
        return R_temp,jac_temp
    end

    if iter<=L
        theta_out_1,res_out_1,iter_c,conv,out_x,out_res,out_conv=my_levenberg_marquardt2(model,theta_use,tolG=tolg_in,tolX=tolx_in,tolF=tolf_in,maxIter=1,show_trace=true,lambda=10)
        out_L[iter,:]=theta_out_1
        theta_use = theta_out_1
        time_keep[iter,4]=toc()
        theta_keep[:,iter+1]=theta_out_1
    else

        theta_out_1,res_out_1,iter_c,conv,out_x,out_res,out_conv=my_levenberg_marquardt2(model,theta_use,tolG=tolg_in,tolX=tolx_in,tolF=tolf_in,maxIter=maxIt_in,show_trace=true,lambda=1)
        out_x_all[(iter-L-1)*(maxIt_in+1)+1:(iter-L)*(maxIt_in+1),:]=out_x
        out_res_all[(iter-L-1)*(maxIt_in+1)+1:(iter-L)*(maxIt_in+1)]=out_res
        out_conv_all[(iter-L-1)*(maxIt_in+1)+1:(iter-L)*(maxIt_in+1),:]=out_conv

        time_keep[iter,4]=signed(time_ns()-ini_solve)/1000000000;

        if ~conv
            flag[iter]=2
        end

        # Update initial guess
        theta_use = theta_out_1
        theta_keep[:,iter+1]=theta_out_1

        # Update weights

        if options_int_in["ind_gmm"] ==1

            Mean_res=[sum(res_out_1[1:M].*cweights)/sum(cweights), sum(res_out_1[M+1:2*M].*cweights)/sum(cweights), sum(res_out_1[2*M+1:end].*cweights)/sum(cweights)]
            VCV[1,1]=sum(((res_out_1[1:M].-Mean_res[1]).*(res_out_1[1:M].-Mean_res[1])).*cweights)/sum(cweights)
            VCV[2,2]=sum(((res_out_1[M+1:2*M].-Mean_res[2]).*(res_out_1[M+1:2*M].-Mean_res[2])).*cweights)/sum(cweights)
            VCV[3,3]=sum(((res_out_1[2*M+1:end].-Mean_res[3]).*(res_out_1[2*M+1:end].-Mean_res[3])).*cweights)/sum(cweights)
            VCV[1,2]=sum(((res_out_1[1:M].-Mean_res[1]).*(res_out_1[M+1:2*M].-Mean_res[2])).*cweights)/sum(cweights)
            VCV[2,1]=VCV[1,2]
            VCV[1,3]=sum(((res_out_1[1:M].-Mean_res[1]).*(res_out_1[2*M+1:end].-Mean_res[3])).*cweights)/sum(cweights)
            VCV[3,1]=VCV[1,3]
            VCV[2,3]=sum(((res_out_1[M+1:2*M].-Mean_res[2]).*(res_out_1[2*M+1:end].-Mean_res[3])).*cweights)/sum(cweights)
            VCV[3,2]=VCV[2,3]
            inv_VCV=VCV^-1
            inv_VCV[1,2]=inv_VCV[2,1]
            inv_VCV[1,3]=inv_VCV[3,1]
            inv_VCV[2,3]=inv_VCV[3,2]
            err_weight=(cholesky(inv_VCV)).U'
            err_weight=err_weight/sum(err_weight)*3


        end

        save_GMM[iter,:,:]=inv_VCV

    end
end

return theta_use,theta_keep,time_keep,flag, out_x_all, out_res_all, out_conv_all,out_L,res_out_1,save_GMM,save_R_GH,save_R_MONO,save_distJ,save_sizeJ_GH,save_GH_ZLB,save_MONO_ZLB

end
