
using Atom
using Distributions
using Clustering
using Optim
using Distances
using StatsBase
using DataFrames
using CSV
using Random
using LinearAlgebra

#--------------------------------------------------------------------------
# Include procedures for model solution
#--------------------------------------------------------------------------
readDirSol = "$(pwd())/FunctionsSolution/"
include(readDirSol * "cga_grid_h.jl")
include(readDirSol * "cga_grid_km.jl")
include(readDirSol * "cga_grid_kms.jl")
include(readDirSol * "eval_hh_33_revised.jl")
include(readDirSol * "get_coefs_plc.jl")
include(readDirSol * "get_GH_nodes.jl")
include(readDirSol * "get_order1_sgu.jl")
include(readDirSol * "ghpoints.jl")
include(readDirSol * "gx_hx.jl")
include(readDirSol * "my_levenberg_marquardt2.jl")
include(readDirSol * "options.jl")
include(readDirSol * "PLC_solve.jl")
include(readDirSol * "PLC_main.jl")
#include(readDirSol * "simulate_plc.jl")
include(readDirSol * "solve_sgu.jl")
#include(readDirSol * "system_piecewise_linear_jacob_33_revised.jl")
include(readDirSol * "tsga_grid.jl")
include(readDirSol * "Monomials_2_N4.jl")

#--------------------------------------------------------------------------
# Include procedures for model estimation
#--------------------------------------------------------------------------
readDir = "FunctionsEstimation/"
include(readDir * "fCreateParamDict.jl")
include(readDir * "fStateSpaceForm.jl")
include(readDir * "fAddSteadyStateParamDict.jl")
include(readDir * "fCanonicalForm.jl")
include(readDir * "fStateSpaceSimulationaug.jl")

#--------------------------------------------------------------------------
# Inputs
#--------------------------------------------------------------------------
nLoadParam = "COPFexactZLB"
nParamDGP = "DGP1" # Name of file with in \PriorSpec with parameter values
nSim     = 1100; # Number of simulated observations
nBurnSim = 100; # Number of initial simulations to discard
#--------------------------------------------------------------------------
# Set solution type and Grid
#--------------------------------------------------------------------------
nSolutionRCLinear = 1
nSmolyakGrid = 1 # = 1 use Smolyak Grid. Note: need to set in options.jl file options_int_out["ext"] = 1 and options_int_out["k"] = 0

options_int_use,options_float_use = @eval options()

if nSmolyakGrid == 1
    options_int_use["k"]   = 0 # Number of non-linear solve steps
    options_int_use["ext"] = 1 # 1 = load cluster from input

    include("FunctionsSolutionSmolyak/Smolyak_Elem_isotrop.jl")
    include("FunctionsSolutionSmolyak/Smolyak_Grid.jl")
    include("FunctionsSolutionSmolyak/fGenSmolyakGrid.jl")
    include("FunctionsSolutionSmolyak/fScaleSmolyakGrid2.jl")

    d          = 6
    mu         = 2
    sd         = 1.282
    wSmolyak   = 1
    #R_quantile = quantile(mData[nStartDate:nEndDate,5]./400 .+ 1,0.90) # R data quantile
    R_quantile = 1.01

    vSmolyakGridSpec = [d, mu, sd, wSmolyak, R_quantile]
    fixed_grid = fGenSmolyakGrid(d,mu,sd)

else
    fixed_grid = zeros(1,1)
    vSmolyakGridSpec = ones(1)
end

if nSolutionRCLinear == 1
    include(readDirSol * "simulate_plc_lin.jl")
    include(readDirSol * "system_piecewise_linear_jacob_33_lin.jl")
else
    include(readDirSol * "simulate_plc_nlin.jl")
    include(readDirSol * "system_piecewise_linear_jacob_33_nlin.jl")
end



#--------------------------------------------------------------------------
# Script Starts
#--------------------------------------------------------------------------
nPath = pwd()
rnd_off = 12
Random.seed!(rnd_off+1)
savedir      = "$(pwd())/Data/"

#--------------------------------------------------------------------------
# Load parameters
#--------------------------------------------------------------------------
readDir     = "$(pwd())/Data/"
vParameters = CSV.read(readDir * nLoadParam * "_" * nParamDGP * "_Param.csv" ,
              datarow=1 , limit=1)
vParameters = convert(Array,vParameters)'

#--------------------------------------------------------------------------
# Solve model and create state-space representation
#--------------------------------------------------------------------------
sParam            = fCreateParamDict(vParameters)
sParam_StateSpace = fStateSpaceForm(sParam,options_int_use,options_float_use,fixed_grid,vSmolyakGridSpec,nSolutionRCLinear)

#--------------------------------------------------------------------------
# Simulate Model
#--------------------------------------------------------------------------
vInitS     = sParam_StateSpace["vInitS"];
pSigmaTE   = sParam_StateSpace["pSigmaTE"];
pA         = sParam_StateSpace["pA"];
pA0        = sParam_StateSpace["pA0"];
D_EpsToEta = sParam_StateSpace["D_EpsToEta"]

pSigmaTE  = sParam_StateSpace["pSigmaTE"];
pSigmaEta = sParam_StateSpace["pSigmaTE"];

eta1BarCoeff = sParam_StateSpace["eta1BarCoeff"];
pPhi1_0      = sParam_StateSpace["Phi01"];
Phi01        = pPhi1_0
pPhi1_1      = sParam_StateSpace["Phi11"];
Phi11        = pPhi1_1
pPhi1_Eta    = sParam_StateSpace["PhiEta1"];
PhiEta1      = pPhi1_Eta
pPhi2_0      = sParam_StateSpace["Phi02"];
Phi02        = pPhi2_0
pPhi2_1      = sParam_StateSpace["Phi12"];
Phi12        = pPhi2_1
pPhi2_Eta    = sParam_StateSpace["PhiEta2"];
PhiEta2      = pPhi2_Eta

# Draw all the shocks
mShockEps = rand(MvNormal(zeros(size(pSigmaTE,1)),pSigmaTE), nSim)';
mShockEta = (mShockEps*D_EpsToEta');

# Create matrices
nNumStates = size(pA,2)
nNumObs    = size(pA,1)
mState     = zeros(nSim,nNumStates)
mSimObs    = zeros(nSim,nNumObs)

for iTime = 1:nSim
    if iTime == 1
        (mState[iTime,:], mSimObs[iTime,:]) = fStateSpaceSimulationaug(eta1BarCoeff,Phi01,Phi11,PhiEta1,Phi02,Phi12,PhiEta2,pA,pA0,vInitS,mShockEta[iTime,:])
    else
        (mState[iTime,:], mSimObs[iTime,:]) = fStateSpaceSimulationaug(eta1BarCoeff,Phi01,Phi11,PhiEta1,Phi02,Phi12,PhiEta2,pA,pA0,mState[iTime-1,:],mShockEta[iTime,:])
    end
end

# Discard first nBurnSim
mState  = mState[nBurnSim + 1:end, :]
mSimObs = mSimObs[nBurnSim + 1:end, :]

# Order observation like input data
mSimObs = [ collect(1:(nSim-nBurnSim)) collect(1:(nSim-nBurnSim)) mSimObs[:,1] mSimObs[:,3] mSimObs[:,4] mSimObs[:,2] mSimObs[:,2] mSimObs[:,2]]
vName   = ["date"; "date"; "ygrowth"; "infl"; "r"; "logcy4"; "logcy4"; "logcy4"]

# Save Output
savename     = nLoadParam * "_" * nParamDGP * "_Data.csv"
CSV.write(savedir * savename,  DataFrame(mSimObs), header=vName)

mState = [ collect(1:(nSim-nBurnSim)) collect(1:(nSim-nBurnSim)) mState]
vName   = ["date"; "date"; "er"; "g"; "z"; "d"; "y"; "pi"; "c"; "R"; "ylag"]

# Save Output
savename     = nLoadParam * "_" * nParamDGP * "_States.csv"
CSV.write(savedir * savename,  DataFrame(mState), header=vName)
