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 Distributions, Parameters, Plots, QuantEcon
gr(fmt = :png);

function AMF_LSS_VAR(A, B, D, F = nothing; ν = nothing)

    if B isa AbstractVector
        B = reshape(B, length(B), 1)
    end
    # unpack required elements
    nx, nk = size(B)

    # checking the dimension of D (extended from the scalar case)
    if ndims(D) > 1
        nm = size(D, 1)
        if D isa Union{Adjoint, Transpose}
            D = convert(Matrix, D)
        end
    else
        nm = 1
        D = reshape(D, 1, length(D))
    end

    # set F
    if isnothing(F)
        F = zeros(nk, 1)
    elseif ndims(F) == 1
        F = reshape(F, length(F), 1)
    end

    # set ν
    if isnothing(ν)
        ν = zeros(nm, 1)
    elseif ndims(ν) == 1
        ν = reshape(ν, length(ν), 1)
    else
        throw(ArgumentError("ν must be column vector!"))
    end

    if size(ν, 1) != size(D, 1)
        error("The size of ν is inconsistent with D!")
    end

    # construct BIG state space representation
    lss = construct_ss(A, B, D, F, ν, nx, nk, nm)

    return (A = A, B = B, D = D, F = F, ν = ν, nx = nx, nk = nk, nm = nm, lss = lss)
end

AMF_LSS_VAR(A, B, D) =
    AMF_LSS_VAR(A, B, D, nothing, ν=nothing)
AMF_LSS_VAR(A, B, D, F, ν) =
    AMF_LSS_VAR(A, B, D, [F], ν=[ν])

function construct_ss(A, B, D, F,
                    ν, nx, nk, nm)

    H, g = additive_decomp(A, B, D, F, nx)

    # auxiliary blocks with 0's and 1's to fill out the lss matrices
    nx0c = zeros(nx, 1)
    nx0r = zeros(1, nx)
    nx1 = ones(1, nx)
    nk0 = zeros(1, nk)
    ny0c = zeros(nm, 1)
    ny0r = zeros(1, nm)
    ny1m = I + zeros(nm, nm)
    ny0m = zeros(nm, nm)
    nyx0m = similar(D)

    # build A matrix for LSS
    # order of states is: [1, t, xt, yt, mt]
    A1 = hcat(1, 0, nx0r, ny0r, ny0r)          # transition for 1
    A2 = hcat(1, 1, nx0r, ny0r, ny0r)          # transition for t
    A3 = hcat(nx0c, nx0c, A, nyx0m', nyx0m')   # transition for x_{t+1}
    A4 = hcat(ν, ny0c, D, ny1m, ny0m)          # transition for y_{t+1}
    A5 = hcat(ny0c, ny0c, nyx0m, ny0m, ny1m)   # transition for m_{t+1}
    Abar = vcat(A1, A2, A3, A4, A5)

    # build B matrix for LSS
    Bbar = vcat(nk0, nk0, B, F, H)

    # build G matrix for LSS
    # order of observation is: [xt, yt, mt, st, tt]
    G1 = hcat(nx0c, nx0c, I, nyx0m', nyx0m')          # selector for x_{t}
    G2 = hcat(ny0c, ny0c, nyx0m, ny1m, ny0m)          # selector for y_{t}
    G3 = hcat(ny0c, ny0c, nyx0m, ny0m, ny1m)          # selector for martingale
    G4 = hcat(ny0c, ny0c, -g, ny0m, ny0m)             # selector for stationary
    G5 = hcat(ny0c, ν, nyx0m, ny0m, ny0m)             # selector for trend
    Gbar = vcat(G1, G2, G3, G4, G5)

    # build LSS type
    x0 = hcat(1, 0, nx0r, ny0r, ny0r)
    S0 = zeros(length(x0), length(x0))
    lss = LSS(Abar, Bbar, Gbar, zeros(nx+4nm, 1), x0, S0)

    return lss
end

function additive_decomp(A, B, D, F, nx)
    A_res = \(I - A, I)
    g = D * A_res
    H = F .+ D * A_res * B

    return H, g
end

function multiplicative_decomp(A, B, D, F, ν, nx)
    H, g = additive_decomp(A, B, D, F, nx)
    ν_tilde = ν .+ 0.5 * diag(H * H')

    return H, g, ν_tilde
end

function loglikelihood_path(amf, x, y)
    @unpack A, B, D, F = amf
    k, T = size(y)
    FF = F * F'
    FFinv = inv(FF)
    temp = y[:, 2:end]-y[:, 1:end-1] - D*x[:, 1:end-1]
    obs =  temp .* FFinv .* temp
    obssum = cumsum(obs)
    scalar = (logdet(FF) + k * log(2π)) * (1:T)

    return -(obssum + scalar) / 2
end

function loglikelihood(amf, x, y)
    llh = loglikelihood_path(amf, x, y)

    return llh[end]
end

function plot_additive(amf, T; npaths = 25, show_trend = true)

    # pull out right sizes so we know how to increment
    @unpack nx, nk, nm = amf

    # allocate space (nm is the number of additive functionals - we want npaths for each)
    mpath = zeros(nm*npaths, T)
    mbounds = zeros(2nm, T)
    spath = zeros(nm*npaths, T)
    sbounds = zeros(2nm, T)
    tpath = zeros(nm*npaths, T)
    ypath = zeros(nm*npaths, T)

    # simulate for as long as we wanted
    moment_generator = moment_sequence(amf.lss)

    # pull out population moments
    for (t, x) in enumerate(moment_generator)
        ymeans = x[2]
        yvar = x[4]

        # lower and upper bounds - for each additive functional
        for ii in 1:nm
            li, ui = 2(ii - 1) + 1, 2ii
            if sqrt(yvar[nx + nm + ii, nx + nm + ii]) != 0.0
                madd_dist = Normal(ymeans[nx + nm + ii], sqrt(yvar[nx + nm + ii, nx + nm + ii]))
                mbounds[li, t] = quantile(madd_dist, 0.01)
                mbounds[ui, t] = quantile(madd_dist, 0.99)
            elseif sqrt(yvar[nx + nm + ii, nx + nm + ii]) == 0.0
                mbounds[li, t] = ymeans[nx + nm + ii]
                mbounds[ui, t] = ymeans[nx + nm + ii]
            else
                error("standard error is negative")
            end

            if sqrt(yvar[nx + 2nm + ii, nx + 2nm + ii]) != 0.0
                sadd_dist = Normal(ymeans[nx + 2nm + ii], sqrt(yvar[nx + 2nm + ii, nx + 2nm + ii]))
                sbounds[li, t] = quantile(sadd_dist, 0.01)
                sbounds[ui, t] = quantile(sadd_dist, 0.99)
            elseif sqrt(yvar[nx + 2nm + ii, nx + 2nm + ii]) == 0.0
                sbounds[li, t] = ymeans[nx + 2nm + ii]
                sbounds[ui, t] = ymeans[nx + 2nm + ii]
            else
                error("standard error is negative")
            end
        end
        t == T && break
    end

    # pull out paths
    for n in 1:npaths
        x, y = simulate(amf.lss,T)
        for ii in 0:nm - 1
            ypath[npaths * ii + n, :] = y[nx + ii + 1, :]
            mpath[npaths * ii + n, :] = y[nx + nm + ii + 1, :]
            spath[npaths * ii + n, :] = y[nx + 2nm + ii + 1, :]
            tpath[npaths * ii + n, :] = y[nx + 3nm + ii + 1, :]
        end
    end

    add_figs = []

    for ii in 0:nm-1
        li, ui = npaths*(ii), npaths*(ii + 1)
        LI, UI = 2ii, 2(ii + 1)
        push!(add_figs,
            plot_given_paths(T, ypath[li + 1:ui, :], mpath[li + 1:ui, :], spath[li + 1:ui, :],
                            tpath[li + 1:ui, :], mbounds[LI + 1:UI, :], sbounds[LI + 1:UI, :],
                            show_trend = show_trend))
    end
    return add_figs
end

function plot_multiplicative(amf, T, npaths = 25, show_trend = true)
    # pull out right sizes so we know how to increment
    @unpack nx, nk, nm = amf
    # matrices for the multiplicative decomposition
    H, g, ν_tilde = multiplicative_decomp(A, B, D, F, ν, nx)

    # allocate space (nm is the number of functionals - we want npaths for each)
    mpath_mult = zeros(nm * npaths, T)
    mbounds_mult = zeros(2nm, T)
    spath_mult = zeros(nm * npaths, T)
    sbounds_mult = zeros(2nm, T)
    tpath_mult = zeros(nm * npaths, T)
    ypath_mult = zeros(nm * npaths, T)

    # simulate for as long as we wanted
    moment_generator = moment_sequence(amf.lss)

    # pull out population moments
    for (t, x) in enumerate(moment_generator)
        ymeans = x[2]
        yvar = x[4]

        # lower and upper bounds - for each multiplicative functional
        for ii in 1:nm
            li, ui = 2(ii - 1)+1, 2ii
            if yvar[nx + nm + ii, nx + nm + ii] != 0.0
                Mdist = LogNormal(ymeans[nx + nm + ii]- 0.5t * diag(H * H')[ii],
                                sqrt(yvar[nx + nm + ii, nx + nm + ii]))
                mbounds_mult[li, t] = quantile(Mdist, 0.01)
                mbounds_mult[ui, t] = quantile(Mdist, 0.99)
            elseif yvar[nx + nm + ii, nx + nm + ii] == 0.0
                mbounds_mult[li, t] = exp.(ymeans[nx + nm + ii] - 0.5t * diag(H * H')[ii])
                mbounds_mult[ui, t] = exp.(ymeans[nx + nm + ii] - 0.5t * diag(H * H')[ii])
            else
                error("standard error is negative")
            end
            if yvar[nx + 2nm + ii, nx + 2nm + ii] != 0.0
                Sdist = LogNormal(-ymeans[nx + 2nm + ii],
                                sqrt(yvar[nx + 2nm + ii, nx + 2nm + ii]))
                sbounds_mult[li, t] = quantile(Sdist, 0.01)
                sbounds_mult[ui, t] = quantile(Sdist, 0.99)
            elseif yvar[nx + 2nm + ii, nx + 2nm + ii] == 0.0
                sbounds_mult[li, t] = exp.(-ymeans[nx + 2nm + ii])
                sbounds_mult[ui, t] = exp.(-ymeans[nx + 2nm + ii])
            else
                error("standard error is negative")
            end
        end
        t == T && break
    end

    # pull out paths
    for n in 1:npaths
        x, y = simulate(amf.lss,T)
        for ii in 0:nm-1
            ypath_mult[npaths * ii + n, :] = exp.(y[nx+ii+1, :])
            mpath_mult[npaths * ii + n, :] =
                exp.(y[nx+nm + ii+1, :] - collect(1:T)*0.5*diag(H * H')[ii+1])
            spath_mult[npaths * ii + n, :] = 1 ./exp.(-y[nx+2*nm + ii+1, :])
            tpath_mult[npaths * ii + n, :] =
                exp.(y[nx + 3nm + ii+1, :] + (1:T) * 0.5 * diag(H * H')[ii + 1])
        end
    end

    mult_figs = []

    for ii in 0:nm-1
        li, ui = npaths * ii, npaths * (ii + 1)
        LI, UI = 2ii, 2(ii + 1)
        push!(mult_figs,
            plot_given_paths(T, ypath_mult[li+1:ui, :], mpath_mult[li+1:ui, :],
                            spath_mult[li+1:ui, :], tpath_mult[li+1:ui, :],
                            mbounds_mult[LI+1:UI, :], sbounds_mult[LI+1:UI, :],
                            horline = 1.0, show_trend=show_trend))
    end

    return mult_figs
end

function plot_martingales(amf, T, npaths = 25)

    # pull out right sizes so we know how to increment
    @unpack A, B, D, F, ν, nx, nk, nm = amf
    # matrices for the multiplicative decomposition
    H, g, ν_tilde = multiplicative_decomp(A, B, D, F, ν, nx)

    # allocate space (nm is the number of functionals - we want npaths for each)
    mpath_mult = zeros(nm * npaths, T)
    mbounds_mult = zeros(2nm, T)

    # simulate for as long as we wanted
    moment_generator = moment_sequence(amf.lss)
    # pull out population moments
    for (t, x) in enumerate(moment_generator)
        ymeans = x[2]
        yvar = x[4]

        # lower and upper bounds - for each functional
        for ii in 1:nm
            li, ui = 2(ii - 1) + 1, 2ii
            if yvar[nx + nm + ii, nx + nm + ii] != 0.0
                Mdist = LogNormal(ymeans[nx + nm + ii] - 0.5^2 * t * diag(H * H')[ii],
                            sqrt(yvar[nx + nm + ii, nx + nm + ii]))
                mbounds_mult[li, t] = quantile(Mdist, 0.01)
                mbounds_mult[ui, t] = quantile(Mdist, 0.99)
            elseif yvar[nx + nm + ii, nx + nm + ii] == 0.0
                mbounds_mult[li, t] = ymeans[nx + nm + ii] - 0.5^2 * t * diag(H * H')[ii]
                mbounds_mult[ui, t] = ymeans[nx + nm + ii]- 0.5t * diag(H * H')[ii]
            else
                error("standard error is negative")
            end
        end
        t == T && break
    end

    # pull out paths
    for n in 1:npaths
        x, y = simulate(amf.lss, T)
        for ii in 0:nm-1
            mpath_mult[npaths * ii + n, :] =
                exp.(y[nx+nm + ii+1, :] - (1:T) * 0.5 * diag(H * H')[ii+1])
        end
    end

    mart_figs = []

    for ii in 0:nm-1
        li, ui = npaths*(ii), npaths*(ii + 1)
        LI, UI = 2ii, 2(ii + 1)
        push!(mart_figs,
            plot_martingale_paths(T, mpath_mult[li + 1:ui, :],
                                                    mbounds_mult[LI + 1:UI, :], horline = 1))
        plot!(mart_figs[ii + 1], title = "Martingale components for many paths of y_(ii + 1)")
    end

    return mart_figs
end

function plot_given_paths(T, ypath, mpath, spath, tpath, mbounds, sbounds;
                          horline = 0.0, show_trend = true)

    # allocate space
    trange = 1:T

    # allocate transpose
    mpathᵀ = Matrix(mpath')

    # create figure
    plots=plot(layout = (2, 2), size = (800, 800))

    # plot all paths together

    plot!(plots[1], trange, ypath[1, :], label = "y_t", color = :black)
    plot!(plots[1], trange, mpath[1, :], label = "m_t", color = :magenta)
    plot!(plots[1], trange, spath[1, :], label = "s_t", color = :green)
    if show_trend
        plot!(plots[1], trange, tpath[1, :], label = "t_t", color = :red)
    end
    plot!(plots[1], seriestype = :hline, [horline], color = :black, linestyle=:dash, label = "")
    plot!(plots[1], title = "One Path of All Variables", legend=:topleft)

    # plot martingale component
    plot!(plots[2], trange, mpath[1, :], color = :magenta, label = "")
    plot!(plots[2], trange, mpathᵀ, alpha = 0.45, color = :magenta, label = "")
    ub = mbounds[2, :]
    lb = mbounds[1, :]
    plot!(plots[2], ub, fillrange = [lb, ub], alpha = 0.25, color = :magenta, label = "")
    plot!(plots[2], seriestype = :hline, [horline], color = :black, linestyle =:dash, label = "")
    plot!(plots[2], title = "Martingale Components for Many Paths")

    # plot stationary component
    plot!(plots[3], spath[1, :], color = :green, label = "")
    plot!(plots[3], Matrix(spath'), alpha = 0.25, color = :green, label = "")
    ub = sbounds[2, :]
    lb = sbounds[1, :]
    plot!(plots[3], ub, fillrange = [lb, ub], alpha = 0.25, color = :green, label = "")
    plot!(plots[3], seriestype = :hline, [horline], color = :black, linestyle=:dash, label = "")
    plot!(plots[3], title = "Stationary Components for Many Paths")

    # plot trend component
    if show_trend == true
        plot!(plots[4], Matrix(tpath'), color = :red, label = "")
    end
    plot!(plots[4], seriestype = :hline, [horline], color = :black, linestyle =:dash, label = "")
    plot!(plots[4], title = "Trend Components for Many Paths")

    return plots
end

function plot_martingale_paths(T, mpath, mbounds;
                               horline = 1, show_trend = false)
    # allocate space
    trange = 1:T

    # create the plot
    plt = plot()

    # plot martingale component
    ub = mbounds[2, :]
    lb = mbounds[1, :]
    plot!(plt, lb, fillrange = [lb, ub], alpha = 0.25, color = :magenta, label = "")
    plot!(plt, seriestype = :hline, [horline], color = :black, linestyle =:dash, label = "")
    plot!(plt, trange, Matrix(mpath'), linewidth=0.25, color = :black, label = "")

    return plt
end

ϕ_1, ϕ_2, ϕ_3, ϕ_4 = 0.5, -0.2, 0, 0.5
σ = 0.01
ν = 0.01 # growth rate

# A matrix should be n x n
A = [ϕ_1 ϕ_2 ϕ_3 ϕ_4;
       1   0   0   0;
       0   1   0   0;
       0   0   1   0]

# B matrix should be n x k
B = [σ, 0, 0, 0]

D = [1 0 0 0] * A
F = [1, 0, 0, 0] ⋅ vec(B)

amf = AMF_LSS_VAR(A, B, D, F, ν)

T = 150
x, y = simulate(amf.lss, T)

plots = plot(layout = (2, 1))

plot!(plots[1], 1:T, y[amf.nx + 1, :], color = :black, lw = 2, label = "")
plot!(plots[1], title =  "A particular path of y_t")
plot!(plots[2], 1:T, y[1, :], color = :green, lw = 2, label = "")
plot!(plots[2], seriestype = :hline, [0], color = :black, lw = 2, linestyle=:dashdot,
      label = "")
plot!(plots[2], title = "Associated path of x_t")

plot(plots)

plt = plot_additive(amf, T)
plt[1]

plt = plot_multiplicative(amf, T)
plt[1]

plt = plot_martingales(amf, 12000)
plt[1]