using Plots
using QuantEcon
using StatsBase

plotlyjs()

prettyprint(arr) = Base.showarray(STDOUT, arr, false, header=false)

P = zeros(6, 6)
P[1, 1] = 1
P[2, 5] = 1
P[3, 3], P[3, 4], P[3, 5] = 1/3, 1/3, 1/3
P[4, 1], P[4, 6] = 1/2, 1/2
P[5, 2], P[5, 5] = 1/2, 1/2
P[6, 1], P[6, 4] = 1/2, 1/2
P

mc1 = MarkovChain(P)

is_irreducible(mc1)

length(communication_classes(mc1))

communication_classes(mc1)

recurrent_classes(mc1)

recurrent_states = vcat(recurrent_classes(mc1)...)

transient_states = setdiff(collect(1:n_states(mc1)), recurrent_states)

permutation = vcat(recurrent_states, transient_states)
mc1.p[permutation, permutation]

is_aperiodic(mc1)

for recurrent_class in recurrent_classes(mc1)
    sub_matrix = mc1.p[recurrent_class, recurrent_class]
    d = period(MarkovChain(sub_matrix))
    println("Period of the sub-chain")
    prettyprint(sub_matrix)
    println("\n = $d")
end

stationary_distributions(mc1)

stationary_distributions(mc1)[1]'* mc1.p

stationary_distributions(mc1)[2]'* mc1.p

"""
Plot the given distribution.
"""
function draw_histogram(distribution;
        title="", xlabel="", ylabel="", ylim=(0, 1),
        show_legend=false, show_grid=false)
    
    n = length(distribution)
    p = bar(collect(1:n),
            distribution,
            xlim=(0.5, n+0.5),
            ylim=ylim,
            title=title,
            xlabel=xlabel,
            ylabel=ylabel,
            xticks=1:n,
            legend=show_legend,
            grid=show_grid)
    
    return p
end

titles = ["$recurrent_class"
          for recurrent_class in recurrent_classes(mc1)]

ps = []
for (title, dist) in zip(titles, stationary_distributions(mc1))
    push!(ps, draw_histogram(dist, title=title, xlabel="States"))
end

plot(ps..., layout=(1, 2))

simulate(mc1, 50, init=1)'

simulate(mc1, 50, init=2)'

simulate(mc1, 50)'

"""
Return the distribution of visits by a sample path of length t
of mc with an initial state init.
"""
function time_series_dist(mc, ts::Array{Int}; init=rand(1:n_states(mc)))
    t_max = maximum(ts)
    ts_size = length(ts)
    
    X = simulate(mc, t_max, init=init)
    dists = Array{Float64}(n_states(mc), ts_size)
    bins = 1:n_states(mc)+1
    for (i, t) in enumerate(ts)
        h = fit(Histogram, X[1:t], bins, closed=:left)
        dists[:, i] = h.weights / t
    end
    
    return dists
end

function time_series_dist(mc, t::Int; init=rand(1:n_states(mc)))
    return time_series_dist(mc, [t], init=init)[:, 1]
end

time_series_dist(mc1, 100, init=2)

time_series_dist(mc1, 10^4, init=2)

function plot_time_series_dists(mc, ts; init=rand(1:n_states(mc)), layout=(1,length(ts)))
    dists = time_series_dist(mc, ts, init=init)
    titles = ["t=$t" for t in ts]

    ps = []
    for (i, title) in enumerate(titles)
        p = draw_histogram(dists[:, i], title=title, xlabel="States")
        push!(ps, p)
    end
    plot(ps..., layout=layout)
end

init = 2
ts = [5, 10, 50, 100]
plot_time_series_dists(mc1, ts, init=init)

init = 3
ts = [5, 10, 50, 100]
plot_time_series_dists(mc1, ts, init=init)

init = 3
ts = [5, 10, 50, 100]

plot_time_series_dists(mc1, ts, init=init)

plot_time_series_dists(mc1, ts, init=init)

inits = [4, 5, 6]
t = 100

ps = []
for init in inits
    p = draw_histogram(
        time_series_dist(mc1, t, init=init),
        title="Initial state = $init",
        xlabel="States")
    push!(ps, p)
end

plot(ps..., layout=(1, 3))

function QuantEcon.simulate(mc, ts_length, num_reps::Int=10^4; init=1)
    X = Array{Int}(ts_length, num_reps)
    for i in 1:num_reps
        X[:, i] = simulate(mc, ts_length, init=init)
    end
    return X
end

"""
Return the distribution of visits at time T by num_reps times of simulation
of mc with an initial state init.
"""
function cross_sectional_dist(mc, ts::Array{Int}, num_reps=10^4; init=rand(1:n_states(mc)))
    t_max = maximum(ts)
    ts_size = length(ts)
    dists = Array{Float64}(n_states(mc), ts_size)
    X = simulate(mc, t_max+1, num_reps, init=init)
    bins = 1:n_states(mc)+1
    for (i, t) in enumerate(ts)
        h = fit(Histogram, X[t, :], bins, closed=:left)
        dists[:, i] = h.weights / num_reps
    end
    return dists
end

function cross_sectional_dist(mc, t::Int, num_reps=10^4; init=rand(1:n_states(mc)))
    return cross_sectional_dist(mc, [t], num_reps, init=init)[:, 1]
end

init = 2
t = 10
cross_sectional_dist(mc1, t, init=init)

t = 100
cross_sectional_dist(mc1, t, init=init)

function plot_cross_sectional_dists(mc, ts, num_reps=10^4; init=1)
    dists = cross_sectional_dist(mc, ts, num_reps, init=init)
    titles = map(t -> "t=$t", ts)

    ps = []
    for (i, title) in enumerate(titles)
        p = draw_histogram(dists[:, i], title=title, xlabel="States")
        push!(ps, p)
    end

    plot(ps..., layout=(1, length(ts)))
end

init = 2
ts = [2, 3, 5, 10]
plot_cross_sectional_dists(mc1, ts, init=init)

init = 3
t = 10
cross_sectional_dist(mc1, t, init=init)

t = 100
dist = cross_sectional_dist(mc1, t, init=init)

draw_histogram(dist,
               title="Cross sectional distribution at t=$t with init=$init",
               xlabel="States")

init = 3
ts = [2, 3, 5, 10]
plot_cross_sectional_dists(mc1, ts, init=init)

inits = [4, 5, 6]
t = 10

ps = []
for init in inits
    p = draw_histogram(cross_sectional_dist(mc1, t, init=init),
                   title="Initial state = $init",
                   xlabel="States")
    push!(ps, p)
end

plot(ps..., layout=(1, length(inits)))

ts = [10, 20, 30]
for t in ts
    P_t = mc1.p^t
    println("P^$t =")
    prettyprint(P_t)
    println()
end

Q = mc1.p[permutation, permutation]
println("Q =")
prettyprint(Q)
for t in ts
    Q_t = Q^t
    println("\nQ^$t =")
    prettyprint(Q_t)
end

P = zeros(10, 10)
P[1, 4] = 1
P[2, [1, 5]] = 1/2
P[3, 7] = 1
P[4, [2, 3, 8]] = 1/3
P[5, 4] = 1
P[6, 5] = 1
P[7, 4] = 1
P[8, [7, 9]] = 1/2
P[9, 10] = 1
P[10, 6] = 1
P

mc2 = MarkovChain(P)

is_irreducible(mc2)

is_aperiodic(mc2)

d = period(mc2)

cyclic_classes = [[1, 5, 7, 9], [4, 10], [2, 3, 6, 8]]

permutation = vcat(cyclic_classes...)
Q = mc2.p[permutation, permutation]

mc2 = MarkovChain(Q)

cyclic_classes = [[1, 2, 3, 4], [5, 6], [7, 8, 9, 10]]

P_blocks = []

for i in 1:d
    push!(P_blocks, mc2.p[cyclic_classes[(i-1)%d+1], :][:, cyclic_classes[i%d+1]])
    println("P_$i =")
    prettyprint(P_blocks[i])
    println()
end

P_power_d = mc2.p^d
prettyprint(P_power_d)

P_power_d_blocks = []
ordinals = ["1st", "2nd", "3rd"]

for (i, ordinal) in enumerate(ordinals)
    push!(P_power_d_blocks, P_power_d[cyclic_classes[i], :][:, cyclic_classes[i]])
    println("$ordinal diagonal block of P^d =")
    prettyprint(P_power_d_blocks[i])
    println()
end

products = []

for i in 1:d
    R = P_blocks[i]
    string = "P_$i"
    for j in 1:d-1
        R = R * P_blocks[(i+j-1)%d+1]
        string *= " P_$((i+j-1)%d+1)"
    end
    push!(products, R)
    println(string, " =")
    prettyprint(R)
    println()
end

for (matrix0, matrix1) in zip(P_power_d_blocks, products)
    println(matrix0 == matrix1)
end

length(stationary_distributions(mc2))

psi = stationary_distributions(mc2)[1]

draw_histogram(psi,
               title="Stationary distribution",
               xlabel="States",
               ylim=(0, 0.35))

psi_s = []

for i in 1:d
    push!(psi_s, stationary_distributions(MarkovChain(P_power_d_blocks[i]))[1])
    println("psi^$i =")
    println(psi_s[i])
end

ps = []

for i in 1:d
    psi_i_full_dim = zeros(n_states(mc2))
    psi_i_full_dim[cyclic_classes[i]] = psi_s[i]
    p = draw_histogram(psi_i_full_dim,
                       title="ψ^$i",
                       xlabel="States")
    push!(ps, p)
end

plot(ps..., layout=(1, 3))

for i in 1:d
    println("psi^$i P_$i =")
    println(psi_s[i]' * P_blocks[i])
end

# Right hand side of the above identity
rhs = zeros(n_states(mc2))

for i in 1:d
    rhs[cyclic_classes[i]] = psi_s[i]
end
rhs /= d
rhs

maximum(abs.(psi - rhs))

for i in 1:d+1
    println("P^$i =")
    prettyprint(mc2.p^i)
    println()
end

for i in [k*d for k in [2, 4, 6]]
    println("P^$i =")
    prettyprint(mc2.p^i)
    println()
end

for i in 10*d+1:11*d
    println("P^$i =")
    prettyprint(mc2.p^i)
    println()
end

init = 1
dist = time_series_dist(mc2, 10^4, init=init)

draw_histogram(dist,
               title="Time series distribution with init=$init",
               xlabel="States", ylim=(0, 0.35))

psi

bar(psi, legend=false, xlabel="States", title="ψ")

init = 1
k = 10
ts = [k*d + i + 1 for i in 1:2*d]
num_reps = 10^2
dists = cross_sectional_dist(mc2, ts, num_reps, init=init)

ps = []
for (i, t) in enumerate(ts)
    p = draw_histogram(dists[:, i],
                       title="t = $t")
    push!(ps, p)
end
plot(ps..., layout=(2, d))

function P_epsilon(eps, p=0.5)
    P = [1-(p+eps) p         eps;
         p         1-(p+eps) eps;
         eps       eps       1-2*eps]
    return P
end

P_epsilon(0)

recurrent_classes(MarkovChain(P_epsilon(0)))

P_epsilon(0.001)

recurrent_classes(MarkovChain(P_epsilon(0.001)))

epsilons = [10.0^(-i) for i in 11:17]
for eps in epsilons
    println("epsilon = $eps")
    
    w, v = eig(P_epsilon(eps)')
    i = indmax(w)
    println(v[:, i]/sum(v[:, i]))
end

epsilons = [10.0^(-i) for i in 11:17]
for eps in epsilons
    println("epsilon = $eps")
    
    w, v = eigs(P_epsilon(eps)', nev=2)
    i = indmax(w)
    println(v[:, i]/sum(v[:, i]))
end

epsilons = [10.0^(-i) for i in 12:17]
push!(epsilons, 1e-100)
for eps in epsilons
    println("epsilon = $eps")
    println(stationary_distributions(MarkovChain(P_epsilon(eps)))[1])
end