#=
This file is part of the replication code for: Hasenzagl, T., Pellegrino, F., Reichlin, L., & Ricco, G. (2020). A Model of the Fed's View on Inflation.
Please cite the paper if you are using any part of the code for academic work (including, but not limited to, conference and peer-reviewed papers).
=#

__precompile__()

module MetropolisWithinGibbs
# ----------------------------------------------------------------------------------------------------------------------
# This module define the functions used to estimate state-space models with a Metropolis-Within-Gibbs algorithm.
#
# The Metropolis step of the Metropolis-Within-Gibbs follows closely the algorithm used in An and Schorfheide (2007).
# The major difference between the two of them is that An and Schorfheide (2007) draw the parameters directly with
# the Metropolis algorithm. Here, the supports of most parameters are constrained in some dimension.
# Therefore, the parameters are transformed, so that their support is unconstrained, and then the Metropolis algorithm
# is executed. The sum of the natural logarithm of the Jacobian is used to evaluate correctly the prior densities.
# DSGE software (e.g. YADA) uses similar modifications to deal with constrained parameters.
# ----------------------------------------------------------------------------------------------------------------------

	using DataFrames, Distributions, Distributed, LinearAlgebra;
	local_path = dirname(@__FILE__);


	# -----------------------------------------------------------------------------------------------------------------
	# Types
	# -----------------------------------------------------------------------------------------------------------------

	mutable struct ParSsm{X <: Float64}
		y::Array{Union{X, Missing}, 2}
		d::Union{Array{X, 1}}
		Z::Union{Array{X, 1}, Array{X, 2}}
		R::Union{Array{X, 1}, Array{X, 2}}
		c::Union{Array{X, 1}}
		T::Union{Array{X, 1}, Array{X, 2}}
		Q::Union{Array{X, 1}, Array{X, 2}}
		α¹::Union{Array{X, 1}}
		P¹::Union{Array{X, 1}, Array{X, 2}}
		P̄¹::Union{Array{X, 1}, Array{X, 2}}
		λ::Array{X, 1}
		ρ::Array{X, 1}
		logprior::X
		loglik::X
		logposterior::X
	end

	struct SizeParSsm{X <: Int64}
		d::X
		Z::X
		R::X
		c::X
		T::X
		Q::X
		λ::X
		ρ::X
		θ::X
	end

	struct BoolParSsm{X1 <: BitArray{1}, X2 <: BitArray{2}}
		d::X1
		Z::X2
		R::X2
		c::X1
		T::X2
		Q::X2
		λ::X1
		ρ::X1
	end

	struct PriorOpt{X <: Float64} # the logpdf for λ and ρ are constants
		N::Distributions.Normal{X}
		IG::Distributions.InverseGamma{X} #TruncatedNormal{X}
		λ::X
		ρ::X
	end

	import Base.copy
	Base.copy(x::T) where T = T([getfield(x, k) for k ∈ fieldnames(T)]...)

	#=
	Base.copy(m::ParSsm) 	= ParSsm([ copy(getfield(m, k)) for k = 1:length(fieldnames(m)) ]...);
	Base.copy(m::SizeParSsm) = SizeParSsm([ copy(getfield(m, k)) for k = 1:length(fieldnames(m)) ]...);
	Base.copy(m::BoolParSsm) = BoolParSsm([ copy(getfield(m, k)) for k = 1:length(fieldnames(m)) ]...);
	Base.copy(m::PriorOpt) 	= PriorOpt([ copy(getfield(m, k)) for k = 1:length(fieldnames(m)) ]...);
	=#

	# -----------------------------------------------------------------------------------------------------------------
	# Subroutines
	# -----------------------------------------------------------------------------------------------------------------

	# Estimation
	include("$local_path/subroutines/estimation/kalman_diffuse!.jl");
	include("$local_path/subroutines/estimation/mwg_main.jl");
	include("$local_path/subroutines/estimation/mwg_run.jl");

	# Julia's subroutines
	include("$local_path/subroutines/extra/ex_blkdiag.jl");
	include("$local_path/subroutines/extra/ex_inv.jl");
	include("$local_path/subroutines/extra/ex_ismember.jl");

	# Get
	include("$local_path/subroutines/get/get_logjacobian.jl");
	include("$local_path/subroutines/get/get_par_bound.jl");
	include("$local_path/subroutines/get/get_par_unb.jl");
	include("$local_path/subroutines/get/get_progress.jl");
	include("$local_path/subroutines/get/get_random_disturbance.jl");
	include("$local_path/subroutines/get/get_mwg_jump.jl");

	# Set
	include("$local_path/subroutines/set/set_par.jl");


	# -----------------------------------------------------------------------------------------------------------------
	# Export functions
	# -----------------------------------------------------------------------------------------------------------------

	export kalman_diffuse!, mwg_main, ex_blkdiag, ex_inv, ex_ismember, set_par_fast!, ParSsm, SizeParSsm, BoolParSsm, PriorOpt;
end
