text:
# This script finds a figure-8 solution to a
# three-body problem with arbitrary precision

println('\n', " "^4, "> Loading the packages...")

using LinearAlgebra # Norm
using Plots # Plotting

# Use the GR backend for plots
gr()

# Change the default font for plots
default(fontfamily = "Computer Modern", dpi = 300, legend = :topright)

# Define the paths to output directories
CURRENT_DIR = @__DIR__
ROOT_DIR = basename(CURRENT_DIR) == "scripts" ? dirname(CURRENT_DIR) : CURRENT_DIR
PLOTS_DIR = joinpath(ROOT_DIR, "plots")

# Prepare the output directory
mkpath(PLOTS_DIR)

# Define the mass of each body (sum is 1)
const m = 1 / 3

# Define the value of ϰ = -Gm (G = 1)
const ϰ = -m

"""
Integrate the three-body problem using the
4th-order Runge-Kutta's method, return the
values of position and velocity on each step
"""
function rk4(
    r₀₁::Vector{F},
    r₀₂::Vector{F},
    r₀₃::Vector{F},
    v₀₁::Vector{F},
    v₀₂::Vector{F},
    v₀₃::Vector{F},
    h::F,
    n::I,
)::Tuple{Matrix{F},Matrix{F},Matrix{F},Matrix{F},Matrix{F},Matrix{F}} where
{F<:AbstractFloat,I<:Unsigned}
    # Determine the length of the input vectors
    N = length(r₀₁)
    # Prepare output matrixes
    r₁ = Matrix{F}(undef, n + 1, N)
    r₂ = Matrix{F}(undef, n + 1, N)
    r₃ = Matrix{F}(undef, n + 1, N)
    v₁ = Matrix{F}(undef, n + 1, N)
    v₂ = Matrix{F}(undef, n + 1, N)
    v₃ = Matrix{F}(undef, n + 1, N)
    r₁[1, :] = copy(r₀₁)
    r₂[1, :] = copy(r₀₂)
    r₃[1, :] = copy(r₀₃)
    v₁[1, :] = copy(v₀₁)
    v₂[1, :] = copy(v₀₂)
    v₃[1, :] = copy(v₀₃)
    # Compute the increments to the position
    function k(i, v₁₊, v₂₊, v₃₊)
        k₁ = v₁[i-1, :] + v₁₊
        k₂ = v₂[i-1, :] + v₂₊
        k₃ = v₃[i-1, :] + v₃₊
        return k₁, k₂, k₃
    end
    # Compute the increments to the velocity
    function a(i, r₁₊, r₂₊, r₃₊)
        r₁⁺, r₂⁺, r₃⁺ = r₁[i-1, :] + r₁₊, r₂[i-1, :] + r₂₊, r₃[i-1, :] + r₃₊
        r₁₂, r₁₃, r₂₃ = r₁⁺ - r₂⁺, r₁⁺ - r₃⁺, r₂⁺ - r₃⁺
        r₂₁, r₃₁, r₃₂ = -r₁₂, -r₁₃, -r₂₃
        ρr₁₂ = ρr₂₁ = norm(r₁₂)^3
        ρr₁₃ = ρr₃₁ = norm(r₁₃)^3
        ρr₂₃ = ρr₃₂ = norm(r₂₃)^3
        a₁ = ϰ * r₁₂ / ρr₁₂ + ϰ * r₁₃ / ρr₁₃
        a₂ = ϰ * r₂₃ / ρr₂₃ + ϰ * r₂₁ / ρr₂₁
        a₃ = ϰ * r₃₁ / ρr₃₁ + ϰ * r₃₂ / ρr₃₂
        return a₁, a₂, a₃
    end
    # Compute the solution
    for i in 2:(n+1)
        k₁₁, k₁₂, k₁₃ = k(i, ntuple(_ -> repeat([0], N), 3)...)
        a₁₁, a₁₂, a₁₃ = a(i, ntuple(_ -> repeat([0], N), 3)...)
        k₂₁, k₂₂, k₂₃ = k(i, h * a₁₁ / 2, h * a₁₂ / 2, h * a₁₃ / 2)
        a₂₁, a₂₂, a₂₃ = a(i, h * k₁₁ / 2, h * k₁₂ / 2, h * k₁₃ / 2)
        k₃₁, k₃₂, k₃₃ = k(i, h * a₂₁ / 2, h * a₂₂ / 2, h * a₂₃ / 2)
        a₃₁, a₃₂, a₃₃ = a(i, h * k₂₁ / 2, h * k₂₂ / 2, h * k₂₃ / 2)
        k₄₁, k₄₂, k₄₃ = k(i, h * a₃₁, h * a₃₂, h * a₃₃)
        a₄₁, a₄₂, a₄₃ = a(i, h * k₃₁, h * k₃₂, h * k₃₃)
        r₁[i, :] = r₁[i-1, :] + h / 6 * (k₁₁ + 2 * k₂₁ + 2 * k₃₁ + k₄₁)
        r₂[i, :] = r₂[i-1, :] + h / 6 * (k₁₂ + 2 * k₂₂ + 2 * k₃₂ + k₄₂)
        r₃[i, :] = r₃[i-1, :] + h / 6 * (k₁₃ + 2 * k₂₃ + 2 * k₃₃ + k₄₃)
        v₁[i, :] = v₁[i-1, :] + h / 6 * (a₁₁ + 2 * a₂₁ + 2 * a₃₁ + a₄₁)
        v₂[i, :] = v₂[i-1, :] + h / 6 * (a₁₂ + 2 * a₂₂ + 2 * a₃₂ + a₄₂)
        v₃[i, :] = v₃[i-1, :] + h / 6 * (a₁₃ + 2 * a₂₃ + 2 * a₃₃ + a₄₃)
    end
    return r₁, r₂, r₃, v₁, v₂, v₃
end

println(" "^4, "> Integrating the three-body problem...")

# Set the precision (in bits)
setprecision(BigFloat, 128)

# Define the initial values of the position and velocity
v₀₂ = BigFloat.(["0.7494421910777922289898659", "1.1501789857502275024030202"])
v₀₁ = [-0.5 * v₀₂[1], -0.5 * v₀₂[2]]
v₀₃ = [-0.5 * v₀₂[1], -0.5 * v₀₂[2]]
r₀₃ = [-5 * m^2 / (1 + m * (norm(v₀₁)^2 + norm(v₀₂)^2 + norm(v₀₃)^2)), 0.0]
r₀₁ = [-r₀₃[1], 0.0]
r₀₂ = BigFloat.(["0.0", "0.0"])

# Define the time step and number of iterations
h = BigFloat("0.0001")
n = ceil(UInt, BigFloat("1.676118923759279") / h / 3)

# Compute the energy constant
E = m / 2 * (norm(v₀₁)^2 + norm(v₀₂)^2 + norm(v₀₃)^2) -
    m^2 * (1 / norm(r₀₁ - r₀₂) + 1 / norm(r₀₁ - r₀₃) + 1 / norm(r₀₂ - r₀₃))

println(" "^4, "  > E: ", E)

# Integrate the problem
r₁, r₂, r₃, _, _, _ = rk4(r₀₁, r₀₂, r₀₃, v₀₁, v₀₂, v₀₃, h, n)

println(" "^4, "> Plotting the figure...")

# Plot the solution
p = plot(r₁[:, 1], r₁[:, 2]; label = "1", legend_position = :topleft)
plot!(p, r₂[:, 1], r₂[:, 2]; label = "2")
plot!(p, r₃[:, 1], r₃[:, 2]; label = "3")

# Save the figure
savefig(p, joinpath(PLOTS_DIR, "solution.png"))
savefig(p, joinpath(PLOTS_DIR, "solution.pdf"))

println()


text:
@testset "Spaces" begin
  Qx, x = PolynomialRing(FlintQQ, "x")
  K, a = NumberField(x^2 - 2, "a1")
  OK = maximal_order(K)
  p = 3*OK
  Kt, t = K["t"]

  infp = infinite_places(K)

  E, b = NumberField(t^2 + 3)
  s = involution(E)

  F = GF(3)

  Hecke.change_base_ring(::Hecke.NfRel, ::Hecke.gfp_mat) = error("asd")
  @test_throws ErrorException hermitian_space(E, F[1 2; 2 1])

  Hecke.change_base_ring(::Hecke.NfRel, x::Hecke.gfp_mat) = x
  @test_throws ErrorException hermitian_space(E, F[1 2; 2 1])

  V = @inferred hermitian_space(E, 3)
  @test fixed_field(V) == base_field(E)
  @test sprint(show, V) isa String
  @test !is_quadratic(V)
  @test is_positive_definite(V)
  @test gram_matrix(V) == identity_matrix(E,3)
  @test V !== hermitian_space(E, 3, cached = false)
  v = matrix(E,1,3,[1,1,1])
  @test inner_product(V, v, v) == matrix(E,1,1,[3])

  @test_throws ArgumentError hermitian_space(E, E[1 b; b 1])
  @test_throws ArgumentError hermitian_space(E, FlintQQ[1 0; 1 1])

  V = @inferred hermitian_space(E, FlintQQ[1 2; 2 1])
  @test V === hermitian_space(E, FlintQQ[1 2; 2 1])
  @test V !== hermitian_space(E, FlintQQ[1 2; 2 1], cached = false)
  @test ishermitian(V)
  @test !is_definite(V)
  @test involution(V) == s
  @test det(V) == -discriminant(V)

  U = hermitian_space(E, E[a 1; 1 1])
  V = hermitian_space(E, E[0 1; 1 0])
  W = hermitian_space(E, E[0 1 0; 1 0 0; 0 0 1])
  @test !Hecke.is_locally_hyperbolic(U,p)
  @test !Hecke.is_isotropic(U,p)
  @test Hecke.is_locally_hyperbolic(V,p)
  @test Hecke.is_isotropic(V,p)
  @test all(v -> Hecke.is_isotropic(V, v), infp)
  @test !Hecke.is_locally_hyperbolic(W,p)
  @test Hecke.is_isotropic(W,p)
  @test_throws AssertionError Hecke.is_locally_hyperbolic(V, 2*OK)

  Qx, x = PolynomialRing(FlintQQ, "x")
  f = x^2-3
  K, a = NumberField(f, "a", cached = false)
  Kt, t = PolynomialRing(K, "t")
  g = t^2+(1)
  E, b = NumberField(g, "b", cached = false)
  D = matrix(E, 3, 3, [(a), 0, 0, 0, (2 * a), 0, 0, 0, (a + 5)])
  V = hermitian_space(E, D)
  DD = matrix(E, 2, 2, [a + 3, 0, 0, (a + 2)])
  W = hermitian_space(E, DD)
  @test @inferred Hecke.is_represented_by(W, V)
  @test @inferred Hecke.is_represented_by(V, V)
  DD = E(a) * identity_matrix(E, 5)
  W = hermitian_space(E, DD)
  @test @inferred Hecke.is_represented_by(V, W)

  # they are not isometric at some dyadic prime ideal:
  DD = identity_matrix(E, 4)
  W = hermitian_space(E, DD)
  @test !(@inferred Hecke.is_represented_by(W, V))

  K, a = rationals_as_number_field()
  OK = maximal_order(K)
  Kt, t = K["t"]

  E, b = NumberField(t^2+17)

  p = 2*OK
  q = 17*OK

  V = hermitian_space(E, E[102 b; -b 0])
  H = hermitian_space(E, E[0 1; 1 0])
  W = hermitian_space(E, E[1 1 2; 1 2 3; 2 3 1])

  @test_throws ErrorException hasse_invariant(V, p)
  @test_throws ErrorException witt_invariant(W, q)

  for r in [p,q]
    @test Hecke.is_locally_represented_by(V,H,r)
  end
  @test Hecke.is_represented_by(V,H)
  @test !Hecke.is_locally_represented_by(W,V,p)
  @test Hecke.is_locally_represented_by(V, W, p)
  @test !Hecke.is_represented_by(W,V)

  @test !is_isometric(V, W, ZZ(2))
  @test all(p -> is_isometric(V, H, ZZ(p)), PrimesSet(1, 20))

  Qx, x = QQ["x"]
  K, a = number_field(x^3+3, "a")
  Kt, t = K["t"]
  E, b = number_field(t^2-2, "b")

  U = hermitian_space(E, E[1 1;1 1])
  V1 = hermitian_space(E, E[102 b; -b 0])
  V2 = hermitian_space(E, E[1 0;0 1])
  H = hermitian_space(E, E[0 1; 1 0])
  W = hermitian_space(E, E[1 1 2; 1 2 3; 2 3 1])

  infp = infinite_places(K)

  @test_throws ArgumentError is_isometric(U, V)  # U is not regular
  @test is_isometric(V1, V1, infp[1])
  @test !is_isometric(V1, W, infp[2]) # infp[2] is complex but V and W don't have same rank
  @test count(v -> is_isometric(V1, H, v), infp) == 1 # V1 and H are not everywhere locally isometric
  @test !is_isometric(V1, H)
  @test is_isometric(V1, V2)

end

ext: jl


text:
@testset "struct_2_dict: without free parameters" begin
  pmodel = Model.m
  D = @struct_2_dict(pmodel)

  @test D["pmodel.a.fb.x"] ≈ 4.0
  @test D["pmodel.a.x"] ≈ 1.0
  @test D["pmodel.x"] ≈ 3.0
  @test D["pmodel.a.i"] ≈ 2
  @test D["pmodel.a.fb.i"] ≈ 5
end

@testset "Annotate dict" begin
  pmodel = Model.m
  D = @struct_2_dict(pmodel)

  @free D["pmodel.x"] Distributions.Normal
  @test D["pmodel.x"] isa FreeParameter

  @free D["pmodel.a.x"]
  @test D["pmodel.a.x"] isa FreeParameter
end

@testset "struct_2_dict complex: without free parameters" begin
  pmodel = DycomsModel.model
  D = @struct_2_dict(pmodel)

  @test D["pmodel.source.2.u_relaxation"]                  ≈ Float32[7.0, -5.5, 0.0]
  @test D["pmodel.boundarycondition.SHF"]                  ≈ 15.0
  @test D["pmodel.source.2.z_sponge"]                      ≈ 1500.0
  @test D["pmodel.radiation.F_0"]                          ≈ 70.0
  @test D["pmodel.source.2.γ"]                             ≈ 2.0
  @test D["pmodel.orientation"]                            == Main.DycomsModel.FlatOrientation()
  @test D["pmodel.source.4.v_geostrophic"]                 ≈ -5.5
  @test D["pmodel.ref_state.relativehumidity"]             ≈ 0.0
  @test D["pmodel.turbulence.C_smag"]                      ≈ 0.21
  @test D["pmodel.radiation.F_1"]                          ≈ 22.0
  @test D["pmodel.source.3.D"]                             ≈ 0.0
  @test D["pmodel.boundarycondition.C_drag"]               ≈ 0.0011
  @test D["pmodel.ref_state.temperatureprofile.Γ"]         ≈ 0.0097659705
  @test D["pmodel.radiation.α_z"]                          ≈ 1.0
  @test D["pmodel.radiation.D_subsidence"]                 ≈ 0.0
  @test D["pmodel.source.4.f_coriolis"]                    ≈ 0.000103
  @test D["pmodel.radiation.κ"]                            ≈ 85.0
  @test D["pmodel.source.2.z_max"]                         ≈ 2500.0
  @test D["pmodel.source.4.u_geostrophic"]                 ≈ 7.0
  @test D["pmodel.moisture.maxiter"]                       ≈ 5
  @test D["pmodel.radiation.ρ_i"]                          ≈ 1.13
  @test D["pmodel.source.1"]                               == Main.DycomsModel.Gravity()
  @test D["pmodel.radiation.z_i"]                          ≈ 840.0
  @test D["pmodel.source.2.α_max"]                         ≈ 1.0
  @test D["pmodel.boundarycondition.LHF"]                  ≈ 115.0
  @test D["pmodel.init_state"]                             == Main.DycomsModel.init_dycoms!
  @test D["pmodel.ref_state.temperatureprofile.T_min"]     ≈ 289.0
  @test D["pmodel.ref_state.temperatureprofile.T_surface"] ≈ 290.4
  @test D["pmodel.precipitation"]                          == Main.DycomsModel.NoPrecipitation()
end

@testset "Annotate dict" begin
  pmodel = DycomsModel.model
  D = @struct_2_dict(pmodel)

  @free D["pmodel.radiation.κ"] Distributions.Normal
  @test D["pmodel.radiation.κ"] isa FreeParameter

  @free D["pmodel.turbulence.C_smag"] Distributions.Normal
  @test D["pmodel.turbulence.C_smag"] isa FreeParameter
end


text:
using Mux, JSON

function todict(s::AbstractString)
    d = Dict{String,String}()
    for pair in split(s,"&")
        a = split(pair,"=")
        d[a[1]] = a[2]
    end
    d
end

function readresults(annualEmissions, results, firstyear, lastyear)
    years = iyear(firstyear):iyear(lastyear)
    res = results[years]
    temp = getfield.(res, :Temp_global)
    GAS3 = [:CO2, :CH4, :N2O]
    conc = getfield.(res, :Concentration)
    concentrations = Dict(g => get.(conc, g, 0.0) for g in GAS3)
    emissions = Dict(g => get(annualEmissions, g, 0.0)[years] for g in GAS3)
    return Dict(
            "temperature" => temp,
            "concentrations" => concentrations,
            "emissions" => emissions
    )
end

function runclimatemodel(req)
    cccdata = JSON.parse(transcode(String, req[:data]))
    firstyear = cccdata["firstyear"]
    lastyear = cccdata["lastyear"]
    # if cccdata["firstyear"] != 2010 || cccdata["lastyear"] != 2100
    #     throw("Received emission data for $(cccdata["firstyear"])-$(cccdata["lastyear"]), " *
    #         "but currently only 2010-2100 is supported.")
    # end
    rcp = "RCP3PD"
    annualEmissions = getscenario(rcp)
    annualEmissions[:CO2][iyear(firstyear):iyear(lastyear)] = 12/44*(cccdata["emissions"]["FossilCO2"] + cccdata["emissions"]["OtherCO2"])
    annualEmissions[:CH4][iyear(firstyear):iyear(lastyear)] = cccdata["emissions"]["CH4"] .+ 270    # natural background emissions: 270 MtCH4
    annualEmissions[:N2O][iyear(firstyear):iyear(lastyear)] = cccdata["emissions"]["N2O"] .+ 10.7   # natural background emissions: 10.7 MtN (TAR p253)
    results, p = solveclimate(annualEmissions, usecache=true, lambda=cccdata["climatesensitivity"]/3.7, rcp=rcp)
    printresults(firstyear:10:lastyear, results, p, annualEmissions, rcp)
    return Dict(
        :body => JSON.json(readresults(annualEmissions, results, firstyear, lastyear)),
        :headers => Dict("Content-Type" => "application/json")
    )
end

function printPOSTinfo(req)
    println("\n$(req[:method])\n")
    display(req[:headers])
    println("\n")
    d = JSON.parse(transcode(String, req[:data]))
    display(d)
    JSON.print(d, 4)
    "hello kitty! :)"
    #"
method: $(req[:method])
data: $(transcode(String, req[:data]))
headers: $(req[:headers])"
end

global webserver

function startserver()
    staticroute = route("static", files(@__DIR__), Mux.notfound())
    #POSTroute = page("/user/:user", req -> "
Hello, $(req[:params][:user])!
")
    #POSTroute = method("POST", respond("POST request received."))
    #POSTroute = branch(req -> req[:method] == "POST", req -> printPOSTinfo(req))
    POSTroute = branch(req -> req[:method] == "POST", req -> runclimatemodel(req))

    serverstarted = isdefined(ClimateCalculator, :webserver)
    defaults = Mux.stack(Mux.todict, Mux.splitquery, Mux.toresponse)
    # @app webserver = (defaults, Mux.basiccatch, POSTroute, staticroute, Mux.notfound())
    @app webserver = (defaults, POSTroute, staticroute, Mux.notfound())
    !serverstarted && serve(webserver)    # only run this once, modify the line with @app test = ... instead
    println("Open http://localhost:8000/static/UI_layout.html in your web browser.")
    # login to NAS as admin, do "sudo -i" to become root
    # cd /volume1/Sync/julia-1.3.1
    # nohup bin/julia -e 'using ClimateCalculator; @sync startserver(); wait(Condition())' > serverlog.txt 2>&1 &
    # the @sync prefix might not be needed anymore
    # nohup bin/julia -e 'using ClimateCalculator; @sync startserver(); wait(Condition())' > /dev/null 2>&1 &
    # ps ax | grep julia
    # kill procnumber
end

max_stars_repo_name: niclasmattsson/ClimateCalculator

max_issues_repo_name: niclasmattsson/ClimateCalculator

max_forks_repo_name: niclasmattsson/ClimateCalculator


text:
testfile = ARGS[1]

totab4(str) = " " ^ (div(length(str), 2) * 4)

function transform_line(line::AbstractString)
    m = match(r"^(\s*)describe\(\"(.*)\"\)\s+do(.*)\n", line)
    if m != nothing
        return "$(totab4(m.captures[1]))@testset \"$(m.captures[2])\" begin$(m.captures[3])\n"
    end

    m = match(r"^(\s*)test\(\"(.*)\"\)\s+do(.*)\n", line)
    if m != nothing
        return "$(totab4(m.captures[1]))@testset \"$(m.captures[2])\" begin$(m.captures[3])\n"
    end

    m = match(r"^(\s*)@repeat\s+test\(\"(.*)\"\)\s+do(.*)\n", line)
    if m != nothing
        return "$(totab4(m.captures[1]))@testset \"$(m.captures[2])\" for i in 1:NumReps $(m.captures[3])\n"
    end

    m = match(r"^(\s*)@check(\s+)(.+)\n", line)
    if m != nothing
        return "$(totab4(m.captures[1]))@test$(m.captures[2])$(m.captures[3])\n"
    end

    m = match(r"^(\s*)([^\s].*)\n", line)
    if m != nothing
        return "$(totab4(m.captures[1]))$(m.captures[2])\n"
    end

    line
end

newlines = open(testfile, "r") do fh
    map(transform_line, readlines(fh))
end

open(testfile * ".alt", "w") do fh
    println(fh, join(newlines, ""))
end


text:
import Dates
println("Last updated: ", Dates.now(), " (PT)")

# Load environment
import Pkg; Pkg.activate("../../../")

# Import Libraries
using Turing
import Turing.RandomMeasures.DirichletProcess
import Turing.RandomMeasures.ChineseRestaurantProcess
using Distributions
using PyPlot
import Random
using BenchmarkTools
import StatsBase.countmap

# rand(ChineseRestaurantProcess(DirichletProcess(1.0), [155, 10]))

function extract(chain, sym; burn=0)
    tail  = chain[sym].value.data[(burn + 1):end, :, :]
    return dropdims(tail, dims=3)
end

function counts_plot(x; density=false, color="C0", lw=4)
    cm_x = countmap(x)
    number = collect(keys(cm_x))
    count = let
        c = collect(values(cm_x))
        density ? c / sum(c) : c
    end
    plt.vlines(number, 0, count, color=color, lw=lw)
end

# FIXME: Not working???

# Define model
@model infiniteGMM(x) = begin
    nobs = length(x)

    # Hyper-parameters, i.e. concentration parameter and parameters of H.
    alpha = 1.0
    mu0 = 0.0
    sig0 = 1.0

    # Define random measure, e.g. Dirichlet process.
    rpm = DirichletProcess(alpha)

    # Define the base distribution, i.e. expected value of the Dirichlet process.
    H = Normal(mu0, sig0)

    # Latent assignment.
    z = tzeros(Int, nobs)

    # Locations of the infinitely many clusters.
    mu = tzeros(Float64, 0)

    for i in 1:nobs
        # Number of clusters.
        K = maximum(z)
        nk = Vector{Int}(map(k -> sum(z .== k), 1:K))

        # Draw the latent assignment.
        z[i] ~ ChineseRestaurantProcess(rpm, nk)

        # Create a new cluster?
        if z[i] > K
            push!(mu, 0.0)

            # Draw location of new cluster.
            mu[z[i]] ~ H
        end

        # Draw observation.
        x[i] ~ Normal(mu[z[i]], 1.0)
    end
end

# Generate data
Random.seed!(1)
data = vcat(randn(10), randn(10) .- 5, randn(10) .+ 10)
data .-= mean(data)
data /= std(data);

# plt.scatter(1:length(data), sort(data))
plt.hist(data, bins=15);

# Fit model
Random.seed!(2)
iterations = 1000
model_fun = infiniteGMM(data)
chain = sample(model_fun, SMC(), iterations);

samples = chain.value.data[:, 4:end, 1]
z = samples[:, end-30+1:end]
mu = samples[:, 1:end-30];

mu

nclus = [length(unique(z[i, :])) for i in 1:size(z, 1)]
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(nclus);
plt.subplot(1, 2, 2);
counts_plot(nclus, density=true);
plt.xlabel("Number of clusters")
plt.ylabel("Posterior Probability");

# counts_plot(z[end-2, :])

counts_plot(z[5,:]);

function rand_mu(mu, z, i)
    mu_vec =  filter(x -> !ismissing(x), mu[i, :])
    z = Int(rand(z[i, :]))
    mu_vec[z]
end
post_pred = [rand(Normal(rand_mu(mu, z, i), 1)) for i in 1:size(z, 1)]
plt.hist(post_pred);

mu

max_stars_repo_name: luiarthur/TuringBnpBenchmarks

max_issues_repo_name: luiarthur/TuringBnpBenchmarks

max_forks_repo_name: luiarthur/TuringBnpBenchmarks


text:
# # Coriolis Parameter
#
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/example_f.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/example_f.ipynb)
#
# [from wikipedia:](https://en.wikipedia.org/wiki/Coriolis_frequency)
#
# The Coriolis frequency `ƒ`, also called the Coriolis parameter or Coriolis
# coefficient, is equal to twice the rotation rate `Ω` of the Earth multiplied
# by the `sine` of the latitude `φ`.
#
# $$f=2\Omega \sin \varphi$$

using PhysicalOceanography
isapprox(f(45),1.03125867182e-4)


text:
using Random, Colors, Gtk, Cairo

#=========== Channel and flag (IPC) section ===================#

""" channel, communicates player's mouse choice of card or color to game logic """
const channel = Channel{Any}(100)

""" flush the channel from mouse choice to game logic """
flushchannel() = while !isempty(channel) take!(channel); end

""" challenge flag: true if challenge taken by player """
const challenge = [false]

#============ Game play section ==================#

""" The Uno card type. The first field is color, second field is number or command. """
UnoCard = Pair{String, String}
color(c::UnoCard) = first(c)
type(c::UnoCard) = last(c)

""" Each Uno player has a name, a score, may be a bot, and has a hand of UnoCards. """
mutable struct UnoCardGamePlayer
    name::String
    score::Int
    isabot::Bool
    hand::Vector{UnoCard}
end

"""
    mutable struct UnoCardGameState
Encapsulates a board state of the gane, including players, cards, color being played,
order of play, current player, and whether the card showing has had its command used
"""
mutable struct UnoCardGameState
    drawpile::Vector{UnoCard}
    discardpile::Vector{UnoCard}
    players::Vector{UnoCardGamePlayer}
    pnow::Int
    colornow::String
    lastcolor::String
    clockwise::Bool
    commandsvalid::Bool
end

""" classifications of colors and types for card faces """
const colors = ["Red", "Yellow", "Green", "Blue"]  # suit colors
const types = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "Skip", "Draw Two", "Reverse"]
const numtypes = types[begin:end-3]
const wildtypes = ["Wild", "Draw Four"]
const cmdtypes = vcat(types[end-2:end], wildtypes)
const alltypes = vcat(types, wildtypes)
const unopreferred = ["Skip", "Draw Two", "Reverse", "Draw Four"]
const ttypes = sort!(vcat(types, types))
const typeordering = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "Wild", "Skip",
    "Reverse", "Draw Two", "Draw Four"]
popfirst!(ttypes) # only 1 "0" card per color

""" The Uno card game deck, unshuffled. """
const originaldeck = [vec([UnoCard(c, v) for v in ttypes, c in colors]);
      fill(UnoCard("Wild", "Wild"), 4); fill(UnoCard("Wild", "Draw Four"), 4)]

""" Set the next player to play to game.pnow (clockwise or counterclockwise) """
function nextplayer!(game, idx)
   game.pnow = mod1(game.clockwise ? idx + 1 : idx - 1, length(game.players))
end
nextplayer!(game) = nextplayer!(game, game.pnow)

"""
    nextsaiduno(game)
Returns true if the next player to play has said Uno, which means they have only one card left.
If so, it is best for the current player if they play to make them draw or lose a turn.
"""
function nextsaiduno(game)
    idx = game.pnow
    nextp = mod1(game.pnow + (game.clockwise ? 1 : -1), length(game.players))
    return length(game.players[nextp].hand) == 1
end

cardscore(c) = (t = type(c); t == "Draw Four" ? 50 : t in cmdtypes ? 20 : parse(Int, t))
handscore(hand) = isempty(hand) ? 0 : sum(cardscore, hand)

"""
    UnoCardGameState(playernames = ["Player", "Bot1", "Bot2", "Bot3"])
Construct and initialize Uno game. Includes dealing hands and picking who is to start.
"""
function UnoCardGameState(playernames = ["Player", "Bot1", "Bot2", "Bot3"])
    deck = shuffle(originaldeck)
    discardpile, drawpile = UnoCard[], UnoCard[]
    while true  # cannot start with a Draw Four on discard pile top
        discardpile, drawpile = [deck[29]], deck[30:end]
        last(last(discardpile)) != "Draw Four" && break
        deck[29:end] .= shuffle(deck[29:end])
    end
    hands = [deck[i:i+6] for i in 1:7:27]
    game = UnoCardGameState(drawpile, discardpile, [UnoCardGamePlayer(playernames[i], 0,
        startswith(playernames[i], "Bot"), hands[i])
        for i in 1:length(playernames)], 1, "Wild", "Wild", true, true)
    dealer = rand(1:length(playernames))
    logline("Player $(playernames[dealer]) is dealer.")
    # handle an initial Reverse card
    if type(last(discardpile)) == "Reverse"
            game.clockwise = false
            logline("First card Reverse, so starting counterclockwise with dealer.")
            game.commandsvalid = false
            game.pnow = dealer
    else
        nextplayer!(game, dealer)
    end
    logline("Player $(playernames[game.pnow]) starts play.")
    if color(last(discardpile)) == "Wild"
        choosecolor!(game)
        game.commandsvalid = false
    else
        game.colornow = color(last(discardpile))
    end
    return game
end

function nextgame(game::UnoCardGameState)
    newgame = UnoCardGameState()
    for i in eachindex(newgame.players)
        newgame.players[i].score = game.players[i].score
    end
    return newgame
end

cardvalue(c::UnoCard) = something(findfirst(x -> x == type(c), typeordering), 0)
countcolor(cards, clr) = count(c -> color(c) == clr, cards)
colorcounts(cards) = sort!([(countcolor(cards, clr), clr) for clr in colors])

""" Preferred color is the one that is most counted in the hand. """
preferredcolor(hand) = last(last(colorcounts(hand)))

"""
    playableindices(game)
Return a vector of indices of cards in hand that are legal to discard
"""
function playableindices(game)
    hand = game.players[game.pnow].hand
    mcolor, mtype = game.colornow, type(game.discardpile[end])
    return [i for (i, c) in enumerate(hand) if
        color(c) == mcolor || type(c) == mtype || color(c) == "Wild"]
end

""" Current player to draw n cards from the draw pile. """
function drawcardsfromdeck!(game, n=1)
    n in [2, 4] && (game.commandsvalid = false)
    if n == 4  # draw four
        # bot will challenge half the time, player must challenge in 5 seconds.
        if game.players[game.pnow].isabot && rand() < 0.5  ||
            (!game.players[game.pnow].isabot && challenge[begin])
            challenge[begin] = false
            logline("$(game.players[game.pnow].name) challenged Draw Four!")
            challenger, savecolor = game.pnow, game.colornow
            nextplayer!(game); nextplayer!(game); nextplayer!(game); # prior player
            game.colornow = game.lastcolor
            hand = game.players[game.pnow].hand
            if any(c -> color(c) == game.colornow, hand)
                logline("Challenge sustained! Challenged Draw Four player draws 4.")
                drawcardsfromdeck!(game, 2); drawcardsfromdeck!(game, 2)
                game.pnow, game.colornow = challenger, savecolor
                return
            else
                logline("Challenge fails. Challenging player now draws 6.")
                n = 6
            end
            game.pnow, game.colornow = challenger, savecolor
        end
        challenge[begin] = false
    end
    logline("Player $(game.players[game.pnow].name) draws $n card$(n == 1 ? "" : "s").")
    for _ in 1:n
        push!(game.players[game.pnow].hand, pop!(game.drawpile))
        if isempty(game.drawpile)  # out of draw pile
            game.drawpile = shuffle(game.discardpile[begin:end-1])
            game.discardpile = [game.discardpile[end]]
        end
    end
end

"""
    discard!(game, idx = -1)
Current player to discard card at index idx in hand (last card in hand as default).
Handle wild card discard by having current player choose the new game.colornow.
"""
function discard!(game, idx = -1)
    hand = game.players[game.pnow].hand
    if idx != -1
        hand[idx], hand[end] = hand[end], hand[idx]
    end
    push!(game.discardpile, pop!(hand))
    lastdiscard = last(game.discardpile)
    logline("Player $(game.players[game.pnow].name) discarded $lastdiscard")
    game.lastcolor = game.colornow
    if color(lastdiscard) == "Wild"  # wild card discard, so choose a color to be colornow
        choosecolor!(game)
        logline("New color chosen: $(game.colornow)")
    else
        game.colornow = color(lastdiscard)
    end
    game.commandsvalid = true
end

"""
    turn!(game)
Execute a single turn of the game.  Command cards are followed only the first turn after played.
"""
function turn!(game)
    name, hand = game.players[game.pnow].name, game.players[game.pnow].hand
    lastdiscard, indices = game.discardpile[end], playableindices(game)
    mcolor, mtype = game.colornow, type(lastdiscard)
    isempty(hand) && error("Empty hand held by $name")
    if mtype in cmdtypes && game.commandsvalid && mtype != "Wild"
        game.commandsvalid = false
        if mtype == "Draw Four"
            drawcardsfromdeck!(game, 4)
        elseif mtype == "Draw Two"
            drawcardsfromdeck!(game, 2)
        elseif mtype == "Skip"    # skip, no uno check
            logline("$name skips a turn.")
        elseif mtype == "Reverse"
            game.clockwise = !game.clockwise
            logline("Reverse: now going $(game.clockwise ? "clockwise." : "counter-clockwise.")")
            nextplayer!(game)
        end
        nextplayer!(game)
        return
    else  # num card, or command card is already used
        if isempty(indices)
            drawcardsfromdeck!(game, 1)  # draw, then discard if drawn card is a match
            indices  = playableindices(game)
            if !isempty(indices)
                logline("Drawn card was discardable.")
                discard!(game, first(indices))
            end
        elseif !startswith(name, "Bot")  # not bot, so player moves
            logline("Click on a card to play.")
            flushchannel()
            while true
                item = take!(channel)
                if item isa Int && item in indices
                    discard!(game, item)
                    break
                end
                logline("That card is not playable.")
            end
        elseif nextsaiduno(game)  # bot might need to stop next player win
            sort!(hand, lt = (x, y) -> cardvalue(x) < cardvalue(y))
            indices = playableindices(game)
            discard!(game, last(indices))
        else # bot play any playable in hand
            discard!(game, rand(indices))
        end
    end
    length(hand) == 1 && logline("$name says UNO!")
    nextplayer!(game)
end

"""
    choosecolor!(game)
Choose a new game.colornow, automatically if a bot, via player choice if not a bot.
"""
function choosecolor!(game)
    logline("Player $(game.players[game.pnow].name) choosing color")
    hand = game.players[game.pnow].hand
    isempty(hand) && return rand(colors)
    if game.players[game.pnow].isabot
        game.colornow = preferredcolor(hand)
    else
        flushchannel()
        while true
            item = take!(channel)
            if item isa String && item in colors
                game.colornow = item
                break
            end
        end
    end
    logline("Current color is now $(game.colornow).")
end

#================  required documentation section ======================#

const unodocshtml = """
Official Rules For Uno Card Game
The aim of the game is to be the first player to score 500 points, achieved (usually over several rounds of play) by being the first to play all of one's own cards and scoring points for the cards still held by the other players.
The deck consists of 108 cards: four each of "Wild" and "Wild Draw Four", and 25 each of four colors (red, yellow, green, blue). Each color consists of one zero, two each of 1 through 9, and two each of "Skip", "Draw Two", and "Reverse". These last three types are known as "action cards".
To start a hand, seven cards are dealt to each player, and the top card of the remaining deck is flipped over and set aside to begin the discard pile. The player to the dealer's left plays first unless the first card on the discard pile is an action or Wild card (see below). On a player's turn, they must do one of the following:
*    play one card matching the discard in color, number, or symbol
*    play a Wild card, or a playable Wild Draw Four card (see restriction below)
*    draw the top card from the deck, then play it if possible
Cards are played by laying them face-up on top of the discard pile. Play proceeds clockwise around the table.
Action or Wild cards have the following effects:
===============================================================================================================================================================
Card            Effect when played from hand                                                                                Effect as first discard
---------------------------------------------------------------------------------------------------------------------------------------------------------------
Skip            Next player in sequence misses a turn                                                                                Player to dealer's left misses a turn
Reverse         Order of play switches directions (clockwise to counterclockwise, or vice versa)            Dealer plays first; play proceeds counterclockwise
Draw Two        Next player in sequence draws two cards and misses a turn                                             Player to dealer's left draws two cards and misses a turn
Wild            Player declares the next color to be matched ; current color may be chosen                       Player to dealer's left declares the first color to be matched and plays a card in it
Wild Draw Four  Player declares the next color to be matched; next player in sequence draws four   Return card to the deck, shuffle, flip top card to start discard pile
A player who draws from the deck must either play or keep that card and may play no other card from their hand on that turn.
A player may play a Wild card at any time, even if that player has other playable cards.
A player may play a Wild Draw Four card only if that player has no cards matching the current color. The player may have cards of a different color matching the current number or symbol or a Wild card and still play the Wild Draw Four card.[5] A player who plays a Wild Draw Four may be challenged by the next player in sequence (see Penalties) to prove that their hand meets this condition.
If the entire deck is used during play, the top discard is set aside and the rest of the pile is shuffled to create a new deck. Play then proceeds normally.
It is illegal to trade cards of any sort with another player.
A player who plays their next-to-last-card must call "uno" as a warning to the other players.[6]
The first player to get rid of their last card ("going out") wins the hand and scores points for the cards held by the other players. Number cards count their face value, all action cards count 20, and Wild and Wild Draw Four cards count 50. If a Draw Two or Wild Draw Four card is played to go out, the next player in the sequence must draw the appropriate number of cards before the score is tallied.
The first player to score 500 points wins the game.
Penalties
=========
If a player does not call "uno" after laying down their next-to-last card and is caught before the next player in sequence takes a turn (i.e., plays a card from their hand, draws from the deck, or touches the discard pile), they must draw two cards as a penalty. If the player is not caught in time (subject to interpretation) or remembers to call "uno" before being caught, they suffer no penalty.
If a player plays a Wild Draw Four card, the following player can challenge its use. The player who used the Wild Draw Four must privately show their hand to the challenging player, in order to demonstrate that they had no matching colored cards. If the challenge is correct, then the challenged player draws four cards instead. If the challenge is wrong, then the challenger must draw six cards; the four cards they were already required to draw plus two more cards.
"""

#============ GUI interface section =======================#

const logwindow = GtkScrolledWindow()
const logtxt = GtkTextBuffer()
logtxt.text[String] = "Started a game of Uno."
const tview = GtkTextView(logtxt)
push!(logwindow, tview)

""" Lines are logged by extending logtxt at its start. """
function logline(txt)
    set_gtk_property!(logtxt, :text, txt * "\n" * get_gtk_property(logtxt, :text, String))
end

const cairocolor = Dict("Red" => colorant"red", "Yellow" => colorant"gold",
    "Green" => colorant"green", "Blue" => colorant"blue", "Wild" => colorant"black")

""" CSS style button for bold colored text """
function colorbutton(txt::String, clr::String)
   button = GtkButton(txt)
   sc = Gtk.GAccessor.style_context(button)
   pr = Gtk.CssProviderLeaf(data="button {color:$(clr); font-weight: bolder}")
   push!(sc, Gtk.StyleProvider(pr), 600)
   return button
end

""" Draw a UnoCard as a rectangle with rounded corners. """
function cairocard(ctx, card, x0, y0, width, height, bcolor=colorant"white")
    fcolor = cairocolor[color(card)]
    set_source(ctx, fcolor)
    radius = (width + height) / 4
    set_line_width(ctx, radius / 5)
    x1 = x0 + width
    y1 = y0 + height
    if width / 2 < radius
        if height / 2 < radius
            move_to(ctx, x0, (y0 + y1) / 2)
            curve_to(ctx, x0 ,y0, x0, y0, (x0 + x1) / 2, y0)
            curve_to(ctx, x1, y0, x1, y0, x1, (y0 + y1) / 2)
            curve_to(ctx, x1, y1, x1, y1, (x1 + x0) / 2, y1)
            curve_to(ctx, x0, y1, x0, y1, x0, (y0 + y1) / 2)
        else
            move_to(ctx, x0, y0 + radius)
            curve_to(ctx, x0 ,y0, x0, y0, (x0 + x1) / 2, y0)
            curve_to(ctx, x1, y0, x1, y0, x1, y0 + radius)
            line_to(ctx, x1 , y1 - radius)
            curve_to(ctx, x1, y1, x1, y1, (x1 + x0) / 2, y1)
            curve_to(ctx, x0, y1, x0, y1, x0, y1 - radius)
        end
    else
        if rect_height / 2 < radius
            move_to(ctx, x0, (y0 + y1)  /2)
            curve_to(ctx, x0 , y0, x0 , y0, x0 + radius, y0)
            line_to(ctx, x1 - radius, y0)
            curve_to(ctx, x1, y0, x1, y0, x1, (y0 + y1) / 2)
            curve_to(ctx, x1, y1, x1, y1, x1 - radius, y1)
            line_to(ctx, x0 + radius, y1)
            curve_to(ctx, x0, y1, x0, y1, x0, (y0 + y1) / 2)
        else
            move_to(ctx, x0, y0 + radius)
            curve_to(ctx, x0 , y0, x0 , y0, x0 + radius, y0)
            line_to(ctx, x1 - radius, y0)
            curve_to(ctx, x1, y0, x1, y0, x1, y0 + radius)
            line_to(ctx, x1 , y1 - radius)
            curve_to(ctx, x1, y1, x1, y1, x1 - radius, y1)
            line_to(ctx, x0 + radius, y1)
            curve_to(ctx, x0, y1, x0, y1, x0, y1- radius)
        end
    end
    close_path(ctx)
    set_source(ctx, bcolor)
    fill_preserve(ctx)
    set_source(ctx, fcolor)
    stroke(ctx)
    move_to(ctx, x0 + width / 3, y0 + height / 3)
    txt = uppercase(type(card))
    if first(txt) in ['R', 'S', 'W']
        txt = string(first(txt))
    elseif first(txt) == 'D'
        txt = "D" * (txt[end] == 'O' ? "2" : "4")
    end
    show_text(ctx, txt)
    stroke(ctx)
end

""" Face down Uno cards are displayed as blank black rectangles with rounded corners. """
function cairodrawfacedowncard(ctx, x0, y0, width, height, bcolor=colorant"darkgray")
    cairocard(ctx, UnoCard("Wild", " "), x0, y0, width, height, bcolor)
end

"""
    UnoCardGameApp(w = 800, hcan = 600, hlog = 100)
Uno card game Gtk app. Draws game on a canvas, logs play on box below canvas.
"""
function UnoCardGameApp(w = 1120, wcan = 810, h = 700)
    win = GtkWindow("Uno Card Game", w, h) |> (GtkFrame() |> (hbox = GtkBox(:h)))
    vbox = GtkBox(:v)
    can = GtkCanvas(wcan, h)
    push!(hbox, can)
    push!(vbox, logwindow)  # from log section
    set_gtk_property!(logwindow, :expand, true)
    b, g = colorbutton("Blue", "blue"), colorbutton("Green", "green")
    r, y = colorbutton("Red", "red"), colorbutton("Yellow", "gold")
    chal = GtkButton("Challenge")
    signal_connect(w -> push!(channel, "Blue"), b, "clicked")
    signal_connect(w -> push!(channel, "Green"), g, "clicked")
    signal_connect(w -> push!(channel, "Red"), r, "clicked")
    signal_connect(w -> push!(channel, "Yellow"), y, "clicked")
    signal_connect(w -> (challenge[begin] = true), chal, "clicked")
    push!(vbox, b, g, r, y, chal)
    push!(hbox, vbox)
    fontpointsize = w / 50
    cardpositions = Dict{Int, Vector{Int}}()

    # announce the rules and penalties per task description
    info_dialog(unodocshtml, win)

    # create a game instance to start play
    game = UnoCardGameState()

    """ Draw the game board on the canvas including player's hand """
    @guarded Gtk.draw(can) do widget
        ctx = Gtk.getgc(can)
        height, width = Gtk.height(ctx), Gtk.width(ctx)
        select_font_face(ctx, "Courier", Cairo.FONT_SLANT_NORMAL, Cairo.FONT_WEIGHT_BOLD)
        set_font_size(ctx, fontpointsize)
        boardcolor = colorant"lightyellow"
        set_source(ctx, boardcolor)
        rectangle(ctx, 0, 0, width, height)
        fill(ctx)
        color = colorant"navy"
        set_source(ctx, color)
        move_to(ctx, 360, 420)
        show_text(ctx, game.players[1].name)
        stroke(ctx)
        move_to(ctx, 60, 300)
        show_text(ctx, game.players[2].name)
        stroke(ctx)
        move_to(ctx, 370, 60)
        show_text(ctx, game.players[3].name)
        stroke(ctx)
        move_to(ctx, 680, 300)
        show_text(ctx, game.players[4].name)
        stroke(ctx)
        cairocard(ctx, last(game.discardpile), 350, 240, 40, 80)
        cairodrawfacedowncard(ctx, 410, 240, 40, 80)
        hand = first(game.players).hand
        isempty(hand) && return
        nrow = (length(hand) + 15) ÷ 15
        for row in 1:nrow
            cards = hand[(row - 1) * 15 + 1 : min(length(hand), row * 15 - 1)]
            startx, starty = 40 + (15 - length(cards)) * 20, 500 + 85 * (row - 1)
            for (i, card) in enumerate(cards)
                idx, x0 = (row - 1) * 15 + i, startx + 50 * (i - 1)
                cardpositions[idx] = [x0, starty, x0 + 40, starty + 80]
                cairocard(ctx, card, x0, starty, 40, 80)
            end
        end
    end

    """ Gtk mouse callback: translates valid mouse clicks to a channel item """
    signal_connect(can, "button-press-event") do b, evt
        for p in cardpositions
            x0, y0, x1, y1 = last(p)
            if x0 < evt.x < x1 && y0 < evt.y < y1
                push!(channel, first(p))
                return
            end
        end
    end

    # do the turns of play in the game, keeping score totals from the rounds
    for n in 1:1000
        draw(can)
        Gtk.showall(win)
        while !any(i -> isempty(game.players[i].hand), 1:4)
            turn!(game)
            if startswith(game.players[game.pnow].name, "Play") &&
                type(game.discardpile[end]) == "Draw Four" && game.commandsvalid
                draw(can)
                Gtk.showall(win)
                info_dialog("Click Challenge within 5 seconds to challenge Draw Four")
                sleep(3)
            end
            sleep(2)
            draw(can)
            Gtk.showall(win)
        end
        winner = findfirst(i -> isempty(game.players[i].hand), 1:length(game.players))
        if type(game.discardpile[end]) == "Draw Two"
            nextplayer!(game)  # next player might have to draw before scoring done
            drawcardsfromdeck!(game, 2)
        elseif type(game.discardpile[end]) == "Draw Four"
            nextplayer!(game)
            drawcardsfromdeck!(game, 2)
            drawcardsfromdeck!(game, 2)   # D2 twice because not to be contested as a D4
        end
        roundpoints = sum(x -> handscore(x.hand), game.players)
        game.players[winner].score += roundpoints

        logline("Player $(game.players[winner].name) wins round $(n)!")
        info_dialog("The winner of round $n is $(game.players[winner].name).\n" *
            "Winner gains $roundpoints points.", win)
        logline("Scores: $([x.score for x in game.players])")
        if any(x -> x.score >= 500, game.players)
            s = "Game over. Scores:\n\n"
            for p in game.players
                s *= "   $(p.name):  $(p.score)  $(p.score >= 500 ? "WINNER!" : "")\n"
            end
            info_dialog(s)
            break
        end
        logline("-------------------------\nNew round!\n-----------------------")
        game = nextgame(game)
        draw(can)
        Gtk.showall(win)
    end
end

UnoCardGameApp()


text:
using LinearAlgebra

include("./grad_projection.jl")

mutable struct AugmentedLagrangian
    f::Function
    c::Function
    λ::Vector{Float64}
    μ::Float64
end

function (aug::AugmentedLagrangian)(x, ::Val{false})
    fx = aug.f(x, Val(false))
    cx = aug.c(x, Val(false))
    return fx - (dot(cx, aug.λ) - aug.μ * dot(cx, cx) / 2)
end

function (aug::AugmentedLagrangian)(x, ::Val{true})
    fx, gx = aug.f(x, Val(true))
    cx, Jx = aug.c(x, Val(true))

    ret = fx - (dot(cx, aug.λ) - aug.μ * dot(cx, cx) / 2)
    ret′ = mul!(gx, Jx', aug.λ, -1, 1)
    ret′ = mul!(ret′, Jx', cx, aug.μ, 1)

    return ret, ret′
end

function descent(f, x, rtol)
    fx, gx = f(x, Val(true))

    for i in 1:100
        if norm(gx) <= rtol
            return x, (x - (x + gx))
        end
        x .-= 0.01 .* gx
        fx, gx = f(x, Val(true))
    end
    return x, (x - (x + gx))
end

function _augmented_lagrangian!(
    aug, c, xk, ωk, ηk;
    lbound, ubound,
    maxiter, rtol, ctol, sr1tol, μpenalty, ηpenalty, ηupdate, monitor
)
    for k in 1:maxiter
        xk, r = descent(aug, xk, rtol)
        c_val = c(xk, Val(false))
        monitor(xk, norm(r))
        if norm(c_val) <= ηk
            if norm(c_val) <= ctol && norm(r) <= rtol
                return xk, true
            end
            aug.λ .-= aug.μ .* c_val
            ηk /= aug.μ ^ ηupdate
            ωk /= aug.μ
        else
            aug.μ *= μpenalty
            ηk = 1 / aug.μ ^ ηpenalty
            ωk = 1 / aug.μ
        end
    end
    return xk, false
end

function augmented_lagrangian(
    f, c, x0, λ0, μ0;
    lbound=fill(-Inf, size(x0)), ubound=fill(Inf, size(x0)),
    μpenalty=100, ηpenalty=0.1, ηupdate=0.9,
    ctol=1e-8, rtol=1e-8, sr1tol=1e-8,
    maxiter=100, monitor=(x, rnorm) -> nothing
)
    aug = AugmentedLagrangian(f, c, λ0, μ0)
    ωk, ηk = 1 / aug.μ, 1 / aug.μ ^ ηpenalty
    return _augmented_lagrangian!(
        aug, c, x0, ωk, ηk;
        lbound=lbound, ubound=ubound, μpenalty=μpenalty, ηpenalty=ηpenalty,
        ηupdate=ηupdate, ctol=ctol, rtol=rtol, sr1tol=sr1tol, maxiter=maxiter,
        monitor=monitor
    )
end

max_stars_repo_name: dbindel/btg

max_issues_repo_name: dbindel/btg

max_forks_repo_name: dbindel/btg


text:
# Install scikit-learn if not installed
ENV["PYTHON"]=""
using Pkg
Pkg.build("PyCall")

import PyCall
PyCall.pyimport_conda("sklearn", "scikit-learn")

using WaveletInformation
using Test


text:
using MortalityTables
using Test

@testset "basic `Ultimate` MortalityTable" begin
    tbl = MortalityTable([i/100 for i in 1:100],0)
    @test q(tbl,0,1) ≈ 0.01
    @test q(tbl,0,2) ≈ 1.0 - (1.0 - 0.01) * (1.0 - 0.02)
    @test q(tbl,0) ≈ 0.01

    @test p(tbl,0) ≈ 0.99
    @test p(tbl,0,1) ≈ 0.99
    @test p(tbl,0,2) ≈ 1.0 - q(tbl,0,2)
    @test q(tbl,0,1,1) ≈ 0.02
    @test p(tbl,0,1,1) ≈ 0.98
    @test_throws BoundsError q(tbl,200)
end

@testset "basic `Ultimate` MortalityTable with non-zero start age" begin
    tbl = MortalityTable([i/100 for i in 5:100],5)
    @test ismissing(q(tbl,0))
    @test q(tbl,5) ≈ 0.05
    @test q(tbl,5,1) ≈ 0.05
    @test q(tbl,5,2) ≈ 1.0 - (1.0 - 0.05) * (1.0 - 0.06)

    @test p(tbl,5) ≈ 0.95
    @test p(tbl,5,1) ≈ 0.95
    @test p(tbl,5,2) ≈ 1.0 - q(tbl,5,2)

    @test q(tbl,5,1,1) ≈ 0.06
    @test p(tbl,5,1,1) ≈ 0.94
end

@testset "test table names" begin
    tbl = MortalityTable([i/100 for i in 1:100],0,TableMetaData(name="test"))
    @test q(tbl,0,1) ≈ 0.01
    @test tbl.d.name == "test"
end

@testset "SOA tables" begin

    tables = MortalityTables.tables()
    cso2001 = tables["2001 CSO Super Preferred Select and Ultimate - Male Nonsmoker, ANB"]
    vbt2001 = tables["2001 VBT Residual Standard Select and Ultimate - Male Nonsmoker, ANB"]
    cso1980 = tables["1980 CSO Basic Table – Male, ANB"]

    @test q(cso1980,35,0,1) ≈ .00118
    @test qx(cso1980,35,1) ≈ .00118
    @test qx(cso1980,35,61) ≈ .27302
    @test qx(cso1980,95) ≈ .27302
    @test ismissing(qx(cso1980,35,95))
    @test ismissing(qx(cso1980,101))

    @test qx(cso2001,35,1) ≈ .00037
    @test qx(cso2001,35,61) ≈ .26719
    @test qx(cso2001,16) ≈ .00041
    @test qx(cso2001,95) ≈ .26719
    @test ismissing(qx(cso2001,15))
    @test_throws BoundsError qx(cso2001,150)
    @test ismissing(qx(cso2001,35,95))

    @test qx(vbt2001,35,1) ≈ .00036
    @test qx(vbt2001,35,61) ≈ .24298
    @test qx(vbt2001,95) ≈ .24298
    @test ismissing(qx(vbt2001,35,95))
    @test_throws BoundsError qx(vbt2001,150)

    @test cso2001[29,1:3] == [0.00029, 0.00035, 0.0004]

end