using InstantiateFromURL
# optionally add arguments to force installation: instantiate = true, precompile = true
github_project("QuantEcon/quantecon-notebooks-julia", version = "0.8.0")

using LinearAlgebra, Statistics

using QuantEcon, Plots, LinearAlgebra, Parameters
gr(fmt = :png);

abstract type AbstractStochProcess end

struct ContStochProcess{TF <: AbstractFloat} <: AbstractStochProcess
    A::Matrix{TF}
    C::Matrix{TF}
end


struct DiscreteStochProcess{TF <: AbstractFloat} <: AbstractStochProcess
    P::Matrix{TF}
    x_vals::Matrix{TF}
end

struct Economy{TF <: AbstractFloat, SP <: AbstractStochProcess}
    β::TF
    Sg::Matrix{TF}
    Sd::Matrix{TF}
    Sb::Matrix{TF}
    Ss::Matrix{TF}
    proc::SP
end

function compute_exog_sequences(econ, x)
    # compute exogenous variable sequences
    Sg, Sd, Sb, Ss = econ.Sg, econ.Sd, econ.Sb, econ.Ss
    g, d, b, s = [dropdims(S * x, dims = 1) for S in (Sg, Sd, Sb, Ss)]

    #= solve for Lagrange multiplier in the govt budget constraint
    In fact we solve for ν = λ / (1 + 2*λ).  Here ν is the
    solution to a quadratic equation a(ν^2 - ν) + b = 0 where
    a and b are expected discounted sums of quadratic forms of the state. =#
    Sm = Sb - Sd - Ss

    return g, d, b, s, Sm
end


function compute_allocation(econ, Sm, ν, x, b)
    Sg, Sd, Sb, Ss = econ.Sg, econ.Sd, econ.Sb, econ.Ss

    # solve for the allocation given ν and x
    Sc = 0.5 .* (Sb + Sd - Sg - ν .* Sm)
    Sl = 0.5 .* (Sb - Sd + Sg - ν .* Sm)
    c = dropdims(Sc * x, dims = 1)
    l = dropdims(Sl * x, dims = 1)
    p = dropdims((Sb - Sc) * x, dims = 1)  # Price without normalization
    τ = 1 .- l ./ (b .- c)
    rvn = l .* τ

    return Sc, Sl, c, l, p, τ, rvn
end


function compute_ν(a0, b0)
    disc = a0^2 - 4a0 * b0

    if disc ≥ 0
        ν = 0.5 *(a0 - sqrt(disc)) / a0
    else
        println("There is no Ramsey equilibrium for these parameters.")
        error("Government spending (economy.g) too low")
    end

    # Test that the Lagrange multiplier has the right sign
    if ν * (0.5 - ν) < 0
        print("Negative multiplier on the government budget constraint.")
        error("Government spending (economy.g) too low")
    end

    return ν
end


function compute_Π(B, R, rvn, g, ξ)
    π = B[2:end] - R[1:end-1] .* B[1:end-1] - rvn[1:end-1] + g[1:end-1]
    Π = cumsum(π .* ξ)
    return π, Π
end


function compute_paths(econ::Economy{<:AbstractFloat, <:DiscreteStochProcess}, T)
    # simplify notation
    @unpack β, Sg, Sd, Sb, Ss = econ
    @unpack P, x_vals = econ.proc

    mc = MarkovChain(P)
    state = simulate(mc, T, init=1)
    x = x_vals[:, state]

    # Compute exogenous sequence
    g, d, b, s, Sm = compute_exog_sequences(econ, x)

    # compute a0, b0
    ns = size(P, 1)
    F = I - β.*P
    a0 = (F \ ((Sm * x_vals)'.^2))[1] ./ 2
    H = ((Sb - Sd + Sg) * x_vals) .* ((Sg - Ss)*x_vals)
    b0 = (F \ H')[1] ./ 2

    # compute lagrange multiplier
    ν = compute_ν(a0, b0)

    # Solve for the allocation given ν and x
    Sc, Sl, c, l, p, τ, rvn = compute_allocation(econ, Sm, ν, x, b)

    # compute remaining variables
    H = ((Sb - Sc) * x_vals) .* ((Sl - Sg) * x_vals) - (Sl * x_vals).^2
    temp = dropdims(F * H', dims = 2)
    B = temp[state] ./ p
    H = dropdims(P[state, :] * ((Sb - Sc) * x_vals)', dims = 2)
    R = p ./ (β .* H)
    temp = dropdims(P[state, :] *((Sb - Sc) * x_vals)', dims = 2)
    ξ = p[2:end] ./ temp[1:end-1]

    # compute π
    π, Π = compute_Π(B, R, rvn, g, ξ)

    return (g = g, d = d, b = b, s = s, c = c,
            l = l, p = p, τ = τ, rvn = rvn, B = B,
            R = R, π = π, Π = Π, ξ = ξ)
end

function compute_paths(econ::Economy{<:AbstractFloat, <:ContStochProcess}, T)
    # simplify notation
    @unpack β, Sg, Sd, Sb, Ss = econ
    @unpack A, C = econ.proc

    # generate an initial condition x0 satisfying x0 = A x0
    nx, nx = size(A)
    x0 = nullspace(I - A)
    x0 = x0[end] < 0 ? -x0 : x0
    x0 = x0 ./ x0[end]
    x0 = dropdims(x0, dims = 2)

    # generate a time series x of length T starting from x0
    nx, nw = size(C)
    x = zeros(nx, T)
    w = randn(nw, T)
    x[:, 1] = x0
    for t in 2:T
        x[:, t] = A *x[:, t-1] + C * w[:, t]
    end

    # compute exogenous sequence
    g, d, b, s, Sm = compute_exog_sequences(econ, x)

    # compute a0 and b0
    H = Sm'Sm
    a0 = 0.5 * var_quadratic_sum(A, C, H, β, x0)
    H = (Sb - Sd + Sg)'*(Sg + Ss)
    b0 = 0.5 * var_quadratic_sum(A, C, H, β, x0)

    # compute lagrange multiplier
    ν = compute_ν(a0, b0)

    # solve for the allocation given ν and x
    Sc, Sl, c, l, p, τ, rvn = compute_allocation(econ, Sm, ν, x, b)

    # compute remaining variables
    H = Sl'Sl - (Sb - Sc)' *(Sl - Sg)
    L = zeros(T)
    for t in eachindex(L)
        L[t] = var_quadratic_sum(A, C, H, β, x[:, t])
    end
    B = L ./ p
    Rinv = dropdims(β .* (Sb- Sc)*A*x, dims = 1) ./ p
    R = 1 ./ Rinv
    AF1 = (Sb - Sc) * x[:, 2:end]
    AF2 = (Sb - Sc) * A * x[:, 1:end-1]
    ξ =  AF1 ./ AF2
    ξ = dropdims(ξ, dims = 1)

    # compute π
    π, Π = compute_Π(B, R, rvn, g, ξ)

    return (g = g, d = d, b = b, s = s,
            c = c, l = l, p = p, τ = τ,
            rvn = rvn, B = B, R = R,
            π = π, Π = Π, ξ = ξ)
end

function gen_fig_1(path)
    T = length(path.c)

    plt_1 = plot(path.rvn, lw=2, label = "tau_t l_t")
    plot!(plt_1, path.g, lw=2, label= "g_t")
    plot!(plt_1, path.c, lw=2, label= "c_t")
    plot!(xlabel="Time", grid=true)

    plt_2 = plot(path.rvn, lw=2, label="tau_t l_t")
    plot!(plt_2, path.g, lw=2, label="g_t")
    plot!(plt_2, path.B[2:end], lw=2, label="B_(t+1)")
    plot!(xlabel="Time", grid=true)

    plt_3 = plot(path.R, lw=2, label="R_(t-1)")
    plot!(plt_3, xlabel="Time", grid=true)

    plt_4 = plot(path.rvn, lw=2, label="tau_t l_t")
    plot!(plt_4, path.g, lw=2, label="g_t")
    plot!(plt_4, path.π, lw=2, label="pi_t")
    plot!(plt_4, xlabel="Time", grid=true)

    plot(plt_1, plt_2, plt_3, plt_4, layout=(2,2), size = (800,600))
end

function gen_fig_2(path)

    T = length(path.c)

    paths = [path.ξ, path.Π]
    labels = ["xi_t", "Pi_t"]
    plt_1 = plot()
    plt_2 = plot()
    plots = [plt_1, plt_2]

    for (plot, path, label) in zip(plots, paths, labels)
        plot!(plot, 2:T, path, lw=2, label=label, xlabel="Time", grid=true)
    end
    plot(plt_1, plt_2, layout=(2,1), size = (600,500))
end

# for reproducible results
using Random
Random.seed!(42)

# parameters
β = 1 / 1.05
ρ, mg = .7, .35
A = [ρ mg*(1 - ρ); 0.0 1.0]
C = [sqrt(1 - ρ^2) * mg / 10 0.0; 0 0]
Sg = [1.0 0.0]
Sd = [0.0 0.0]
Sb = [0 2.135]
Ss = [0.0 0.0]
proc = ContStochProcess(A, C)

econ = Economy(β, Sg, Sd, Sb, Ss, proc)
T = 50
path = compute_paths(econ, T)

gen_fig_1(path)

gen_fig_2(path)

# Parameters
β = 1 / 1.05
P = [0.8 0.2 0.0
    0.0 0.5 0.5
    0.0 0.0 1.0]

# Possible states of the world
# Each column is a state of the world. The rows are [g d b s 1]
x_vals = [0.5 0.5 0.25;
        0.0 0.0  0.0;
        2.2 2.2  2.2;
        0.0 0.0  0.0;
        1.0 1.0  1.0]
Sg = [1.0 0.0 0.0 0.0 0.0]
Sd = [0.0 1.0 0.0 0.0 0.0]
Sb = [0.0 0.0 1.0 0.0 0.0]
Ss = [0.0 0.0 0.0 1.0 0.0]
proc = DiscreteStochProcess(P, x_vals)

econ = Economy(β, Sg, Sd, Sb, Ss, proc)
T = 15
path = compute_paths(econ, T)

gen_fig_1(path)

gen_fig_2(path)

# parameters
β = 1 / 1.05
ρ, mg = .95, .35
A = [0. 0. 0. ρ  mg*(1-ρ);
     1. 0. 0. 0.       0.;
     0. 1. 0. 0.       0.;
     0. 0. 1. 0.       0.;
     0. 0. 0. 0.       1.]
C = zeros(5, 5)
C[1, 1] = sqrt(1 - ρ^2) * mg / 8
Sg = [1. 0. 0. 0. 0.]
Sd = [0. 0. 0. 0. 0.]
Sb = [0. 0. 0. 0. 2.135]
Ss = [0. 0. 0. 0. 0.]
proc = ContStochProcess(A, C)
econ = Economy(β, Sg, Sd, Sb, Ss, proc)

T = 50
path = compute_paths(econ, T)

gen_fig_1(path)

gen_fig_2(path)