# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.3.3
#   kernelspec:
#     display_name: Julia 1.9.2
#     language: julia
#     name: julia-1.9
# ---

# ### Markov chain illustration
# - 2018-11-09 Jeff Fessler, University of Michigan
# - 2021-08-23 Julia 1.6.2
# - 2023-07-24 Julia 1.9.2

using LinearAlgebra
using Plots
using Interact
using Random: seed!
using Distributions: wsample
using Printf

P = [
    0.0 0.0 1.0 0.0;
    1.0 0.0 0.0 0.0;
    0.0 0.6 0.0 1.0;
    0.0 0.4 0.0 0.0;
]

function plot_circle!(x,y;r=10, ic=0)
    t = range(0, 2π, 101)
    plot!(x .+ r * cos.(t), y .+ r * sin.(t), label="")
    plot!(annotate=(x, y+0.8r, "$ic"))
end

xc = [-1, 1, -1, 1] * 20
yc = [1, 1, -1, -1] * 20

# +
# function for plotting the markov chain diagram

function plot_chain()

    plot()
    for ic=1:4
        plot_circle!(xc[ic], yc[ic], ic=ic)
    end

    for ii=1:4
        if P[ii,ii] != 0
            @show "Bug"
            #add_loop!()
        end
        for jj=1:4
            if ii != jj && P[ii,jj] != 0
            #    @show ii, jj, P[ii,jj]
                xi = xc[ii]
                xj = xc[jj]
                yi = yc[ii]
                yj = yc[jj]
                xd = xi - xj
                yd = yi - yj
                phi = atan(yd, xd) + π/2
                rd = sqrt(xd^2 + yd^2)
                frac = (rd-1*10)/rd
              #  @show rd, frac
                xs = frac * xj + (1-frac) * xi
                xe = frac * xi + (1-frac) * xj
                ys = frac * yj + (1-frac) * yi
                ye = frac * yi + (1-frac) * yj
                xm = (xi + xj) / 2
                ym = (yi + yj) / 2
                xo = 5 * cos(phi)
                yo = 5 * sin(phi)
                plot!([xs, xe], [ys, ye],
                    annotate=(xm+xo, ym+yo, "$(P[ii,jj])"), arrow=:arrow, label="")
            end
        end
    end
    plot!(aspect_ratio=1, axis=:on, title="$totalstep steps")
end

#plot_chain()
# -

# initial conditions for 100 points
xg = repeat(-4.5:4.5, 1, 10)
yg = xg'
xg = xg[:]
yg = yg[:]
node = 2*ones(Int, length(xg)) # all start in one node
totalstep = 0
plot_chain()
scatter!(xc[node] + xg[:], yc[node] + yg[:], label="")

# +
seed!(0)

function node_update!(node)
    global totalstep += 1
    for kk=1:length(node)
        node[kk] = wsample(1:4, P[:,node[kk]])
    end
end

# +
function run_and_plot(;niter=1)
    for iter = 1:niter
        node_update!(node)
        plot_chain()
        scatter!(xc[node] + xg[:], yc[node] + yg[:], label="")
    end
    for ii=1:4
        tmp = sum(node .== ii) / length(node)
        plot!(annotate=(xc[ii]+10, yc[ii]-10, "$tmp"))
    end
    plot!()
end

run_and_plot()
# -

# ### Clicker question:
# Which state has the lowest probability in equilibrium?
# - A 1
# - B 2
# - C 3
# - D 4
# - E None: they are all equally likely

@manipulate for niter = ["1", "5", "100"], go=button("go")
    run_and_plot(niter=parse(Int, niter))
end

function matprint(V) # function to nicely print the eigenvector matrix
    for i=1:size(V,1)
        for j=1:size(V,2)
            v = V[i,j]
            @printf("%5.2f", real(v))
            print(imag(v) < 0 ? " -" : " +")
            @printf("%5.2fı  ", abs(imag(v)))
            end
        println()
    end
end

(d, V) = eigen(P)

[d abs.(d)] # exactly one λ=1 and only one |λ| = 1

matprint(V)

v = Real.(V[:,4])
πss = v / sum(v)

# ### Clicker question
# (Do after covering irreducible matrix, aperiodic matrix, strongly connected, etc.)
#
# The transition matrix P in this example is
# (choose most specific correct answer):
# - A. Square
# - B. Nonnegative
# - C. Irreducible
# - D. Primitive
# - E. Positive

# +
# P^10