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

function h_j(j, nk, s1, s2, θ, δ, ρ)
    # Find out who's h we are evaluating
    if j == 1
        sj = s1
        sk = s2
    else
        sj = s2
        sk = s1
    end

    # Coefficients on the quadratic a x^2 + b x + c = 0
    a = 1.0
    b = ((ρ + 1 / ρ) * nk - sj - sk)
    c = (nk * nk - (sj * nk) / ρ - sk * ρ * nk)

    # Positive solution of quadratic form
    root = (-b + sqrt(b * b - 4 * a * c)) / (2 * a)

    return root
end

DLL(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ) =
    (n1 ≤ s1_ρ) && (n2 ≤ s2_ρ)

DHH(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ) =
    (n1 ≥ h_j(1, n2, s1, s2, θ, δ, ρ)) && (n2 ≥ h_j(2, n1, s1, s2, θ, δ, ρ))

DHL(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ) =
    (n1 ≥ s1_ρ) && (n2 ≤ h_j(2, n1, s1, s2, θ, δ, ρ))

DLH(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ) =
    (n1 ≤ h_j(1, n2, s1, s2, θ, δ, ρ)) && (n2 ≥ s2_ρ)

function one_step(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
    # Depending on where we are, evaluate the right branch
    if DLL(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
        n1_tp1 = δ * (θ * s1_ρ + (1 - θ) * n1)
        n2_tp1 = δ * (θ * s2_ρ + (1 - θ) * n2)
    elseif DHH(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
        n1_tp1 = δ * n1
        n2_tp1 = δ * n2
    elseif DHL(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
        n1_tp1 = δ * n1
        n2_tp1 = δ * (θ * h_j(2, n1, s1, s2, θ, δ, ρ) + (1 - θ) * n2)
    elseif DLH(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
        n1_tp1 = δ * (θ * h_j(1, n2, s1, s2, θ, δ, ρ) + (1 - θ) * n1)
        n2_tp1 = δ * n2
    end

    return n1_tp1, n2_tp1
end

new_n1n2(n1_0, n2_0, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ) =
    one_step(n1_0, n2_0, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)

function pers_till_sync(n1_0, n2_0, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ,
                        maxiter, npers)

    # Initialize the status of synchronization
    synchronized = false
    pers_2_sync = maxiter
    iters = 0

    nsync = 0

    while (~synchronized) && (iters < maxiter)
        # Increment the number of iterations and get next values
        iters += 1

        n1_t, n2_t = new_n1n2(n1_0, n2_0, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)

        # Check whether same in this period
        if abs(n1_t - n2_t) < 1e-8
            nsync += 1
        # If not, then reset the nsync counter
        else
            nsync = 0
        end

        # If we have been in sync for npers then stop and countries
        # became synchronized nsync periods ago
        if nsync > npers
            synchronized = true
            pers_2_sync = iters - nsync
        end
        n1_0, n2_0 = n1_t, n2_t
    end
    return synchronized, pers_2_sync
end

function create_attraction_basis(s1_ρ, s2_ρ, s1, s2, θ, δ, ρ,
                                 maxiter, npers, npts)
    # Create unit range with npts
    synchronized, pers_2_sync = false, 0
    unit_range = range(0.0,  1.0, length = npts)

    # Allocate space to store time to sync
    time_2_sync = zeros(npts, npts)
    # Iterate over initial conditions
    for (i, n1_0) in enumerate(unit_range)
        for (j, n2_0) in enumerate(unit_range)
            synchronized, pers_2_sync = pers_till_sync(n1_0, n2_0, s1_ρ, s2_ρ,
                                                       s1, s2, θ, δ, ρ,
                                                       maxiter, npers)
            time_2_sync[i, j] = pers_2_sync
        end
    end
    return time_2_sync
end

# model
function MSGSync(s1 = 0.5, θ = 2.5, δ = 0.7, ρ = 0.2)
    # Store other cutoffs and parameters we use
    s2 = 1 - s1
    s1_ρ = min((s1 - ρ * s2) / (1 - ρ), 1)
    s2_ρ = 1 - s1_ρ

    return (s1 = s1, s2 = s2, s1_ρ = s1_ρ, s2_ρ = s2_ρ, θ = θ, δ = δ, ρ = ρ)
end

function simulate_n(model, n1_0, n2_0, T)
    # Unpack parameters
    @unpack s1, s2, θ, δ, ρ, s1_ρ, s2_ρ = model

    # Allocate space
    n1 = zeros(T)
    n2 = zeros(T)

    # Simulate for T periods
    for t in 1:T
        # Get next values
        n1[t], n2[t] = n1_0, n2_0
        n1_0, n2_0 = new_n1n2(n1_0, n2_0, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
    end

    return n1, n2
end

function pers_till_sync(model, n1_0, n2_0,
                        maxiter = 500, npers = 3)
    # Unpack parameters
    @unpack s1, s2, θ, δ, ρ, s1_ρ, s2_ρ = model

    return pers_till_sync(n1_0, n2_0, s1_ρ, s2_ρ, s1, s2,
                        θ, δ, ρ, maxiter, npers)
end

function create_attraction_basis(model;
                                 maxiter = 250,
                                 npers = 3,
                                 npts = 50)
    # Unpack parameters
    @unpack s1, s2, θ, δ, ρ, s1_ρ, s2_ρ = model

    ab = create_attraction_basis(s1_ρ, s2_ρ, s1, s2, θ, δ,
                                 ρ, maxiter, npers, npts)
    return ab
end

function plot_timeseries(n1_0, n2_0, s1 = 0.5, θ = 2.5, δ = 0.7, ρ = 0.2)
    model = MSGSync(s1, θ, δ, ρ)
    n1, n2 = simulate_n(model, n1_0, n2_0, 25)
    return [n1 n2]
end

# Create figures
data_ns = plot_timeseries(0.15, 0.35)
data_s = plot_timeseries(0.4, 0.3)

plot(data_ns, title = "Not Synchronized", legend = false)

plot(data_s, title = "Synchronized", legend = false)

function plot_attraction_basis(s1 = 0.5, θ = 2.5, δ = 0.7, ρ = 0.2; npts = 250)
    # Create attraction basis
    unitrange = range(0,  1, length = npts)
    model = MSGSync(s1, θ, δ, ρ)
    ab = create_attraction_basis(model,npts=npts)
    plt = Plots.heatmap(ab, legend = false)
end

params = [[0.5, 2.5, 0.7, 0.2],
          [0.5, 2.5, 0.7, 0.4],
          [0.5, 2.5, 0.7, 0.6],
          [0.5, 2.5, 0.7, 0.8]]

plots = (plot_attraction_basis(p...) for p in params)
plot(plots..., size = (1000, 1000))