text:
abstract type AbstractTimeSeries{N,T} end

struct FieldTimeSeries{N,T} <: AbstractTimeSeries{N,T}
    ts     :: T
    field  :: String
end

# slightly adapted from openPMD-viewer
mutable struct OpenPMDTimeSeries{S}
    iterations        :: Array{Int64}
    files             :: Array{String}
    current_i         :: Int64
    current_iteration :: Int64
    current_t         :: Float64
    params            :: S
end

function OpenPMDTimeSeries(foldername::String, prefix::String)
    iterations, files = try
        list_h5_files(foldername, prefix=prefix)
    catch e
        if isa(e, SystemError) && e.errnum == 2 # "No such file or directory"
            throw(ErrorException("No files found."))
        else
            throw(e)
        end
    end

    if length(iterations) == 0
        throw(ErrorException("No files found."))
    end

    # open first file to extract parameters
    fid    = h5open(files[1], "r")
    it     = iterations[1]
    grp, t = read_openpmd_file(fid, it)

    # read group attributes
    params = read_group_attributes(grp)

    OpenPMDTimeSeries{typeof(params)}(iterations, files, 1, it, t, params)
end

"""
    OpenPMDTimeSeries(foldername::String; prefix::String="")

Initialize an openPMD time series, ie, scan the directory and extract the
openPMD files starting with the given prefix (can be left empty)

# Example
```
julia> ts = OpenPMDTimeSeries("./data"; prefix="wave_")
```
"""
OpenPMDTimeSeries(foldername::String; prefix::String="") =
    OpenPMDTimeSeries(foldername, prefix)

FieldTimeSeries{N}(ts::OpenPMDTimeSeries, field::String) where{N} =
    FieldTimeSeries{N,typeof(ts)}(ts, field)

"""
    FieldTimeSeries(foldername::String; prefix::String, field::String)

Initialize a (openPMD) time series for the given `N`-dimensional `field`, ie,
scan the directory and extract the openPMD files starting with the given
`prefix`. The data corresponding to the given `field` will then be extracted
from the corresponding hdf5 file when requested.

# Example
```
julia> xi_ts = FieldTimeSeries("./", prefix="gauge_", field="xi")

julia> xi_ts[1:2,1,2:4,4]
2×3 Array{Float64,2}:
 0.0466351  0.0466351  0.0466351
 0.0738048  0.0738148  0.0738312
```
"""
function FieldTimeSeries(foldername::String; prefix::String, field::String)
    ts = OpenPMDTimeSeries(foldername, prefix)
    f, chart = get_field(ts, it=1, field=field)
    N = ndims(chart)
    FieldTimeSeries{N}(ts, field)
end

function Base.size(ff::FieldTimeSeries)
    Nt  = length(ff.ts.iterations)
    it0 = ff.ts.iterations[1]
    f0, = get_field(ff.ts, it=it0, field=ff.field)
    size_ = size(f0)
    Nt, size_...
end

@inline Base.firstindex(ff::FieldTimeSeries, d::Int) = 1
@inline Base.lastindex(ff::FieldTimeSeries, d::Int) = size(ff)[d]

function Base.getindex(ff::FieldTimeSeries, a::Int, idx::Vararg)
    it = ff.ts.iterations[a]
    f, = get_field(ff.ts, it=it, field=ff.field)
    f[idx...]
end

function Base.getindex(ff::FieldTimeSeries, aa, idx::Vararg)
    it  = ff.ts.iterations[aa[1]]
    f0, = get_field(ff.ts, it=it, field=ff.field)

    Na    = length(aa)
    size_ = size(f0[idx...])
    f     = zeros(Na, size_...)

    slicer = [Colon() for _ in 1:ndims(f0)]
    for (i,a) in enumerate(aa)
        it  = ff.ts.iterations[a]
        f0, = get_field(ff.ts, it=it, field=ff.field)
        f[i,slicer...] .= f0[idx...]
    end
    f
end

function Base.getindex(ff::FieldTimeSeries, ::Colon, idx::Vararg)
    Na  = length(ff.ts.iterations)
    getindex(ff, 1:Na, idx...)
end

"""
    get_coords(ff::FieldTimeSeries, a::Int, idx::Vararg)

# Example
```
julia> t, u, x, y = Jecco.get_coords(xi_ts,1:2,1,2:4,4)
([0.0, 0.10214880285375176], NaN, -4.6875:0.3125:-4.0625, -4.0625)
```
"""
function get_coords(ff::FieldTimeSeries, a::Int, idx::Vararg)
    it = ff.ts.iterations[a]
    f, chart = get_field(ff.ts, it=it, field=ff.field)
    t = ff.ts.current_t
    t, chart[idx...]...
end

"""
    get_coords(ff::FieldTimeSeries, aa, idx::Vararg)

# Example
```
julia> t, u, x, y = Jecco.get_coords(xi_ts,1,1,2,4)
(0.0, NaN, -4.6875, -4.0625)
```
"""
function get_coords(ff::FieldTimeSeries, aa, idx::Vararg)
    it = ff.ts.iterations[aa[1]]
    f0, chart = get_field(ff.ts, it=it, field=ff.field)
    t0 = ff.ts.current_t

    Na = length(aa)
    t  = Vector{typeof(t0)}(undef, Na)
    for (i,a) in enumerate(aa)
        it   = ff.ts.iterations[a]
        get_field(ff.ts, it=it, field=ff.field)
        t0   = ff.ts.current_t
        t[i] = t0
    end
    t, chart[idx...]...
end

function Jecco.get_coords(ff::FieldTimeSeries, ::Colon, idx::Vararg)
    Na = length(ff.ts.iterations)
    get_coords(ff, 1:Na, idx...)
end

function list_h5_files(foldername::String; prefix::String="")
    path     = abspath(foldername)
    allfiles = readdir(path)

    Ns = length(prefix)

    its_names = Tuple[]
    # append only the files whose names start with the given prefix
    for file in allfiles
        try
            if (file[1:Ns] == prefix && (file[end-2:end] == ".h5" ||
                                         file[end-4:end] == ".hdf5"))
                fullname = joinpath(path, file)
                # extract all iterations from the file
                fid      = h5open(fullname, "r")
                its      = names(fid["/data"])
                close(fid)
                # for each iteration add to list of tuples with iteration and
                # name
                for it in its
                    push!(its_names, (parse(Int64, it), fullname))
                end
            end
        catch ex
            if isa(ex, BoundsError)
                # probably triggered by string comparison; do nothing
            else
                throw(ex)
            end
        end
    end

    # sort according to iteration
    sort!(its_names)
    # and extract the list of filenames and iterations
    filenames = [name for (it, name) in its_names]
    its       = [it for (it, name) in its_names]

    (its, filenames)
end

"""
    get_field(ts::OpenPMDTimeSeries; it::Int, field::String)

Given a time series, extract the requested field (and corresponding chart) from
an HDF5 file in the openPMD format. As side-effects, ```ts.current_i```,
```ts.current_iteration``` and ```ts.current_t``` are correspondingly modified.

# Example
```
julia> psi, chart=get_field(ts, it=20, field="psi");

julia> ts.current_t
20.0

julia> ts.current_iteration
20

julia> ts.current_i
21

```
"""
function get_field(ts::OpenPMDTimeSeries; it::Int, field::String, verbose::Bool=false)
    # index that corresponds to the closest iteration requested
    ts.current_i = argmin(abs.(it .- ts.iterations))
    # the closest iteration found (it need not be the requested one)
    ts.current_iteration = ts.iterations[ts.current_i]
    # and corresponding file
    filename = ts.files[ts.current_i]

    if verbose
        println("Reading file $filename")
    end
    # open file
    fid = h5open(filename, "r")

    # read in openPMD structure
    grp, ts.current_t = read_openpmd_file(fid, ts.current_iteration)

    # read actual data
    data, chart = read_dataset(grp, field)

    # close file
    close(fid)

    data, chart
end

function read_openpmd_file(fid::HDF5.File, it::Integer)
    basePath   = read(attrs(fid)["basePath"])
    basePath   = replace(basePath, "%T" => it)
    meshesPath = read(attrs(fid)["meshesPath"])

    # pointer to base group (ie, with the information for the requested time
    # level) within the given hdf5 file
    grp_base = fid[basePath]

    time = read(attrs(grp_base)["time"])

    # pointer to mesh group (with the actual chart function data)
    grp_mesh = grp_base[meshesPath]

    grp_mesh, time
end

function read_group_attributes(grp::HDF5.Group)
    grp_attrs = attrs(grp)
    keys      = names(grp_attrs)
    vals      = read.(Ref(grp_attrs), keys)
    Dict(keys[i] => vals[i] for i in 1:length(keys))
end

function read_dataset(grp::HDF5.Group, var::String)
    dset       = grp[var]
    dset_attrs = attrs(dset)

    func  = read(dset)
    nodes = size(func)
    T     = eltype(func)
    dim_  = length(nodes)

    if dim_ == 1
        names         = read(dset_attrs["axisLabels"])
        mins          = read(dset_attrs["gridGlobalOffset"])
        maxs          = read(dset_attrs["gridMax"])
        gridtypes     = read(dset_attrs["gridType"])
    else
        # remember to flip the order since HDF5 uses row-major order to store
        # arrays, as opposed to Julia's column-major order
        names         = read(dset_attrs["axisLabels"])[end:-1:1]
        mins          = read(dset_attrs["gridGlobalOffset"])[end:-1:1]
        maxs          = read(dset_attrs["gridMax"])[end:-1:1]
        gridtypes     = read(dset_attrs["gridType"])[end:-1:1]
    end

    chart = Chart(gridtypes, names, mins, maxs, nodes)

    func, chart
end


text:
#fist we check if all package are installed:
#include("installer.jl")

#include file
include("../ALAS/OutputLSmod.jl")
include("../ALAS/OutputALASmod.jl")
include("../ALAS/OutputRelaxationmod.jl")

include("../ALAS/Thetafunc.jl")
include("../ALAS/Relaxation.jl")
include("../ALAS/Penalty.jl")
include("../ALAS/ParamSetmod.jl")
include("../ALAS/ActifMPCCmod.jl")
include("../ALAS/DDirection.jl")
include("../ALAS/LineSearch.jl")
include("../ALAS/AlgoSetmod.jl")
include("../ALAS/UnconstrainedMPCCActif.jl")
include("../ALAS/MPCCmod.jl")
include("../ALAS/ALASMPCCmod.jl")
include("../ALAS/MPCCsolve.jl")

ext: jl


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

"""
    struct ModeAsDefined <: AbstractModeEstimator

Get the mode as defined by the density, resp. the underlying distribution
(if available), via `StatsBase.mode`.

Constructors:

* ```$(FUNCTIONNAME)()```
"""
struct ModeAsDefined <: AbstractModeEstimator end
export ModeAsDefined

function bat_findmode_impl(rng::AbstractRNG, target::AnySampleable, algorithm::ModeAsDefined)
    (result = StatsBase.mode(target),)
end

function bat_findmode_impl(rng::AbstractRNG, target::Distribution, algorithm::ModeAsDefined)
    (result = varshape(target)(StatsBase.mode(unshaped(target))),)
end

function bat_findmode_impl(rng::AbstractRNG, target::DistributionDensity, algorithm::ModeAsDefined)
    bat_findmode_impl(rng, parent(target), algorithm)
end

"""
    MaxDensitySampleSearch <: AbstractModeEstimator

Constructors:

    MaxDensitySampleSearch()

Estimate the mode as the variate with the highest posterior density value
within a given set of samples.
"""
struct MaxDensitySampleSearch <: AbstractModeEstimator end
export MaxDensitySampleSearch

function bat_findmode_impl(rng::AbstractRNG, target::DensitySampleVector, algorithm::MaxDensitySampleSearch)
    v, i = _get_mode(target)
    (result = v, mode_idx = i)
end


text:
using Pitot
using Test

@testset "Pitot.jl" begin

    P = Pitot()
    @test 1.0 ≈ P.Cd
    @test velocity(P, 50.0, 1.0) ≈ 10.0
    @test dynamicpress(P, 50.0) ≈ 50.0
    @test pressdiff(P, 10.0, 1.0) ≈ 50.0

    P = Pitot(Cd=0.9)
    @test 0.9 == P.Cd
    @test velocity(P, 50.0, 1.0) ≈ 9.0
    @test dynamicpress(P, 50.0) ≈ 40.5
    @test pressdiff(P, 9.0, 1.0) ≈ 50.0

end


text:
prog = "1 + 2"
ex = Meta.parse(prog)
println(typeof(ex))
dump(ex)

a = :(3 + 4)
println(typeof(a))
b = quote
    5 + 6
end
println(typeof(b))
println(eval(a))
println(eval(b))

macro containervariable(container, element)
    return esc(:($(Symbol(container,element)) = $container[$element]))
end

println(@containervariable "abcdef" 3)
println(@macroexpand @containervariable "abcdef" 3)

# Generated functions
# The input is a type, and the return value is a expression to be evaluated.
@generated function square(x)
    println(x)
    :(x * x)
end

x = square(2);
println(x)
y = square("spam");
println(y)


text:
"""
@default type Foo
    x::Int64=6
    y::Float64=7
end

yields...

type Foo
    x::Int64
    y::Float64
end

function Foo(;x=6, y=7)
    Foo(x,y)
end
"""
macro default(expr)
    if expr.head != symbol("type")
        error("@default only works on composite types")
    else

        field_name(a::Symbol) = a
        field_name(a::Expr) = a.args[1]
        type_symbol(a::Symbol) = a
        type_symbol(a::Expr) = a.args[1]

        type_name = type_symbol(expr.args[2])

        for i = 3:length(expr.args)
            block = expr.args[i]
            if block.head == :(block)
                field_names = Any[]
                defaults = Any[]
                fields = Any[]
                for arg in block.args
                    if arg.head == :(=)
                        val = arg.args[2]
                        field = field_name(arg.args[1])
                        push!(field_names, field)
                        push!(defaults, :($field = $val))
                        push!(fields, arg.args[1])
                    end
                end
                # replace block with just fields names and type annotations
                block.args = fields
                for default in defaults
                    default.head = :kw
                end
                kwarg_func = :(function $type_name(;$(defaults...))
                                    $type_name($(field_names...))
                               end)
                default_type = :(begin
                                    $expr
                                    $kwarg_func
                                end)
                return esc(default_type)
            end
        end
    end
end

max_stars_repo_name: sjkelly/JuliaPlayground

max_issues_repo_name: sjkelly/JuliaPlayground

max_forks_repo_name: sjkelly/JuliaPlayground


text: function f(x::yo) end

max_stars_repo_name: gabriel-doriath-dohler/pjulia

max_issues_repo_name: gabriel-doriath-dohler/pjulia

max_forks_repo_name: gabriel-doriath-dohler/pjulia


text:
function spgreqn(x::Vector{Float64}, p::Vector{Float64}, TR::Float64)
  y = similar(x)
  E = exp(-TR*p[2])  # hoisted expensive, invariant result
  @simd for k in 1:length(x)
     @inbounds y[k] = p[1]*sin(x[k])*(1.0 - E) / (1.0 - E*cos(x[k]))
  end
  y
end

function toftskety(t::Vector{Float64}, p::Vector{Float64}, Cp::Vector{Float64})
  # Standard Tofts-Kety Model.  Works when t_dce = t_aif only, so you should
  # resample the AIF to match the DCE dynamic spacing before using.
  Ct = zeros(t)
  kep = p[2]  # kep = Ktrans / ve
  for k in 1:length(Ct)
    #Ct[k] = p[1]*trapz(exp(-p[1]*(t[k] - t[1:k])/p[2]) .* Cp[1:k], t[1:k])
    # inlined trapz here ... original equations are up here ^
    tk = t[k]
    @simd for j = 1:k
      @inbounds y = exp(kep*(t[j] - tk)) * Cp[j]
      @inbounds Ct[k] += ifelse((j == 1 || j == k) && k > 1, 0.5*y, y)
    end
  end
  dtp = (t[2] - t[1]) * p[1]
  @simd for k = 1:length(Ct)
    @inbounds Ct[k] *= dtp
  end
  Ct
end

function extendedtoftskety(t::Vector{Float64}, p::Vector{Float64}, Cp::Vector{Float64})
  # Extended Tofts-Kety Model. Works when t_dce = t_aif only, so you should
  # resample the AIF to match the DCE dynamic spacing before using.
  Ct = toftskety(t, p, Cp)
  @simd for k in 1:length(Ct)
    @inbounds Ct[k] += p[3]*Cp[k]
  end
  Ct
end

function onecompartment(t::Vector{Float64}, p::Vector{Float64}, Cp::Vector{Float64})
  # plasma only model
  Ct = similar(Cp)
  @simd for k in 1:length(Ct)
    @inbounds Ct[k] = p[1]*Cp[k]
  end
  Ct
end


text:
"""
```julia
h = SphericalAltitude(Body::T, PosRsph::Vector{Float64}) where {T <: abstractCelestialBody}
```
Computes the altitude above a spherical central body with radius Re.
"""
function SphericalAltitude(Body::T, PosRsph::Vector{Float64}) where {T <: abstractCelestialBody}
  return PosRsph[1] - mean_radius(Body)
end


text:
using Test
import AbstractAlgebra
using Groups

import KnuthBendix: Word

using GroupsCore
include(joinpath(pathof(GroupsCore), "..", "..", "test", "conformance_test.jl"))

@testset "Groups" begin

    @testset "wlmetric_ball" begin
        M = AbstractAlgebra.MatrixAlgebra(AbstractAlgebra.zz, 3)
        w = one(M); w[1,2] = 1;
        r = one(M); r[2,3] = -3;
        s = one(M); s[1,3] = 2; s[3,2] = -1;

        S = [w,r,s]; S = unique([S; inv.(S)]);
        _, sizes = Groups.wlmetric_ball(S, radius=4);
        @test sizes == [7, 33, 141, 561]
        _, sizes = Groups.wlmetric_ball_serial(S, radius=4);
        @test sizes == [7, 33, 141, 561]
    end

    include("free_groups.jl")
    include("fp_groups.jl")

    include("AutFn.jl")
    include("AutSigma_41.jl")
    include("AutSigma3.jl")

    # if !haskey(ENV, "CI")
    #    include("benchmarks.jl")
    # end
end


text:
include("go_compile.jl")
include("rr.jl")


text:
using Wakame
using Test

@testset "Wakame.jl" begin
    # Write your tests here.
end


text:
# check polylog
#    compare Series~2 and Series~3 near s=n > 0
#
using Polylogarithms
using BenchmarkTools
using DataFrames, CSV
using PyPlot
using Printf
include("utilities.jl")

series_1 = Polylogarithms.polylog_series_1
series_2 = Polylogarithms.polylog_series_2
series_3 = Polylogarithms.polylog_series_3
L = Symbol("Li_s(z)")

# from https://github.com/JuliaLang/julia/issues/18328
function parse(::Type{Complex{T}}, s::AbstractString) where {T<:Real}
    s = replace(s, "*^" => "e") # Mathematica seems to export in forms like "3.061616997868383*^-18"
    real_r = r"([+-]?\d*\.\d*(e[+-]\d+)?)([\s+-]|$)"
    imag_r = r"([+-]?\d*\.\d*(e[+-]\d+)?)\s*\*\s*I"

    z = zero(Complex{T})

    m = match( imag_r, s )
    if m != nothing
        z += im*parse(T, m[1])
    end

    m = match( real_r, s )
    if m != nothing
        z += parse(T, m[1])
    end
    return z
end

# input data from Mathematica and reparse into complex numbers
data1 = CSV.read("../data/polylog_test_data_b_1.csv"; delim=",", type=String)
#    has trouble reading in numbers like "2." so read all into strings, and parse
data1[!,:n] = Int.(parse.(Float64, data1[!,:n] ))
data1[!,:s] = parse.(Complex{Float64}, data1[!,:s] )
data1[!,:r] = parse.(Float64, data1[!,:r] )
data1[!,:theta] = parse.(Float64, data1[!,:theta] )
data1[!,:z] = parse.(Complex{Float64}, data1[!,:z] )
data1[!,L] = parse.(Complex{Float64}, data1[!,L] )
# check that complex numbers are being read in correctly
z = data1[!,:z]
z2 = data1[!,:r] .* exp.( im* data1[!,:theta] )
error = data1[!,:z] .- z2
maximum_abs_error = maximum(abs.(error))
# [data1[1:10,:z] parse.(Complex{Float64}, data1[1:10,:z] ) z2[1:10] ]
# [data1[1:10,:z]  z2[1:10] ]
println("max error in parsing is ", maximum_abs_error)
θ = unique(data1[!,:theta])
r = unique(data1[!,:r])

# figure(1)
# plot(real.(z), imag.(z), "o")
# axis("equal")

m = size(data1,1)
Li = data1[!,L]
s = data1[!,:s]

S1 = zeros(Complex{Float64}, m)
n1 = zeros(Int64, m)
error1 = zeros(Float64, m)
rel_error1 = zeros(Float64, m)

S2 = zeros(Complex{Float64}, m)
n2 = zeros(Int64, m)
error2 = zeros(Float64, m)
rel_error2 = zeros(Float64, m)

S3 = zeros(Complex{Float64}, m)
n3 = zeros(Int64, m)
error3 = zeros(Float64, m)
rel_error3 = zeros(Float64, m)

S5 = zeros(Complex{Float64}, m)
n5 = zeros(Int64, m)
error5 = zeros(Float64, m)
rel_error5 = zeros(Float64, m)

benchmark = series_3(data1[1,:n], z[1]; n_terms=3)[1]
bench_error = abs( benchmark - Li[1] )
bench_rel_error = bench_error / abs( Li[1] )
for i=1:m
    print(".")
    # result1 = series_1(s[i], z[i])
    # S1[i] = result1[1]
    # n1[i] = result1[2]
    # error1[i] = abs( S1[i] - Li[i]  )
    # rel_error1[i] = error1[i]/abs( Li[i] )

    result2 = series_2(s[i], z[i])
    S2[i] = result2[1]
    n2[i] = result2[2]
    error2[i] = abs( S2[i] - Li[i]  )
    rel_error2[i] = error2[i]/abs( Li[i] )

    result3 = series_3(s[i], z[i]; n_terms=3)
    S3[i] = result3[1]
    n3[i] = result3[2]
    error3[i] = abs( S3[i] - Li[i]  )
    rel_error3[i] = error3[i]/abs( Li[i] )

    result5 = series_3(s[i], z[i]; n_terms=5)
    S5[i] = result5[1]
    n5[i] = result5[2]
    error5[i] = abs( S5[i] - Li[i]  )
    rel_error5[i] = error5[i]/abs( Li[i] )
end
println("")
# println("max abs. error1 = $(maximum( abs.(error1) ))")
println("max abs. error2 = $(maximum( abs.(error2) ))")
println("max abs. error3 = $(maximum( abs.(error3) ))")
println("max abs. error5 = $(maximum( abs.(error5) ))")

fig = figure("polylog_bench_b", figsize=(6,4))
clf()
d = 0.00

# loglog( data1[!,:r] .- d, rel_error1, "o")
loglog( data1[!,:r] .+ d, rel_error2, "+"; label="Series 2" )
plot( data1[!,:r] .+ d, rel_error5, "rd"; label="Series 3 (5 terms)" )
plot( data1[!,:r] .+ d, rel_error3, "gx"; label="Series 3 (3 terms)" )
# plot( [minimum(r), maximum(r)], [1,1].*bench_rel_error, "--")
legend()
xlabel("|τ|")
ylabel("relative absolute error")
plot([minimum(r),maximum(r)], [1,1]*Polylogarithms.default_accuracy)
# xlim([0, 1.15])
# ylim([1.0-e15, 1.0e-10])
savefig("plots/polylog_bench_b.svg")
savefig("plots/polylog_bench_b.pdf")

# subplot(222)
# # loglog( data1[!,:r] .- d, rel_error1, "o")
# semilogy( data1[!,:r] .+ d, rel_error2, "+")
# plot( data1[!,:r] .+ d, rel_error3, "d")
# plot( [minimum(r), maximum(r)], [1,1].*bench_rel_error, "--")
# xlabel("r")
# ylabel("relative absolute error")
# plot([minimum(r),maximum(r)], [1,1]*Polylogarithms.default_accuracy)
# # xlim([0, 1.15])
# # ylim([1.0-e15, 1.0e-10])

# subplot(223)
# # loglog( data1[!,:r] .- d, n1, "o")
# loglog( data1[!,:r] .+ d, n2, "+")
# loglog( data1[!,:r] .+ d, n3, "d")
# # for i=1:length(θ)
# #     k = findall(data1[!,:theta] .== θ[i])
# #     plot(data1[k,:r] .+ d, n2[k], "g:")
# #     text(data1[k[end],:r]+2*d, n2[k[end]], "$(θ[i]/π)"; verticalalignment="center")
# # end
# # text(1.0, 72, "θ/π"; verticalalignment="center")
# xlabel("r")
# ylabel("number of terms")
# plot([minimum(r),maximum(r)], [1,1]*Polylogarithms.default_max_iterations)
# # xlim([0, 1.15])

# subplot(222)
# loglog( data1[!,:theta] ./ π .- d, rel_error1, "o")
# loglog( data1[!,:theta] ./ π .+ d, rel_error2, "+")
# xlabel("θ/π")
# ylabel("relative absolute error")
# plot([0,1], [1,1]*Polylogarithms.default_accuracy)

# subplot(224)
# loglog( data1[!,:theta] ./ π .- d, n1, "o")
# loglog( data1[!,:theta] ./ π .+ d, n2, "+")
# xlabel("θ/π")
# ylabel("number of terms")
# plot([0,1], [1,1]*Polylogarithms.default_max_iterations)


text:
import Documenter, DocumenterCitations, Literate
import ClimaCore, ClimaCoreVTK, ClimaCoreMakie, ClimaCorePlots

if !@isdefined(TUTORIALS)
    TUTORIALS = ["introduction"]
end

rm(joinpath(@__DIR__, "src", "tutorials"), force = true, recursive = true)

function preprocess_markdown(input)
    line1, rest = split(input, '\n', limit = 2)
    string(
        line1,
        "\n# *This tutorial is available as a [Jupyter notebook](@__NAME__.ipynb).*\n",
        rest,
    )
end
for tutorial in TUTORIALS
    Literate.markdown(
        joinpath(@__DIR__, "tutorials", tutorial * ".jl"),
        joinpath(@__DIR__, "src", "tutorials");
        preprocess = preprocess_markdown,
    )
    Literate.notebook(
        joinpath(@__DIR__, "tutorials", tutorial * ".jl"),
        joinpath(@__DIR__, "src", "tutorials");
        execute = false,
    )
end

withenv("GKSwstype" => "nul") do

    bib =
        DocumenterCitations.CitationBibliography(joinpath(@__DIR__, "refs.bib"))

    format = Documenter.HTML(
        prettyurls = !isempty(get(ENV, "CI", "")),
        collapselevel = 1,
    )

    Documenter.makedocs(
        bib,
        sitename = "ClimaCore.jl",
        strict = [:example_block],
        format = format,
        checkdocs = :exports,
        clean = true,
        doctest = true,
        modules = [ClimaCore, ClimaCoreVTK, ClimaCorePlots, ClimaCoreMakie],
        pages = Any[
            "Home" => "index.md",
            "API" => "api.md",
            "Operators" => "operators.md",
            "Tutorials" => [
                joinpath("tutorials", tutorial * ".md") for
                tutorial in TUTORIALS
            ],
            "Libraries" => [joinpath("lib", "ClimaCoreVTK.md")],
            "references.md",
        ],
    )
end

Documenter.deploydocs(
    repo = "github.com/CliMA/ClimaCore.jl.git",
    target = "build",
    push_preview = true,
    devbranch = "main",
    forcepush = true,
)


text:
using Compose

function sierpinski(n)
    if n == 0
        compose(context(), polygon([(1,1), (0,1), (1/2, 0)]))
    else
        t = sierpinski(n - 1)
        compose(context(),
                (context(1/4,   0, 1/2, 1/2), t),
                (context(  0, 1/2, 1/2, 1/2), t),
                (context(1/2, 1/2, 1/2, 1/2), t))
    end
end

img = SVG("sierpinski.svg", 4inch, 4(sqrt(3)/2)inch)
draw(img, compose(sierpinski(8), linewidth(0.1mm), fill(nothing), stroke("black")))


text:
########################################################################
# Synchronous Vertical Federated Logistic Regression
########################################################################

function vertical_lr_train(server::Server, clients::Vector{Client})
    # number of epoches
    num_epoches = server.num_epoches
    # number of training data
    num_train_data = length(server.Ytrain)
    # batch size
    batch_size = server.batch_size
    # number of batches
    num_batches = div(num_train_data, batch_size)
    # round
    round = 1
    # start training
    @inbounds for epoch = 1:num_epoches
        # generate mini-batches
        batches = generate_batches(num_train_data, num_batches)
        @inbounds for i = 1:num_batches
            batch = batches[i]
            # server updates batch information
            update_batch(server, batch)
            Threads.@threads for c in clients
                # client updates batch information
                update_batch(c, batch)
                # client compute and upload embeddings
                send_embedding(c, server)
            end
            # server compute the loss and the gradient
            batch_loss = compute_mini_batch_gradient(server)
            if i % 1 == 0
                @printf "Epoch %d, Batch %d, Loss %.2f\n" epoch i batch_loss
            end
            # server and clients update model
            update_model(server)
            Threads.@threads for c in clients
                update_model(c)
            end
            # update round
            round += 1
        end
    end
end

function evaluation(server::Server, clients::Vector{Client})
    # test and train accuracy
    for c in clients
        # client compute and upload training embeddings
        send_embedding(c, server, tag = "training")
        # client compute and upload test embeddings
        send_embedding(c, server, tag = "test")
    end
    train_loss, train_acc, test_loss, test_acc = eval(server)
    @printf "Train Loss %.2f, Train Accuracy %.2f, Test Loss %.2f, Test Accuracy %.2f\n" train_loss train_acc test_loss test_acc
end

########################################################################
# Asynchronous Vertical Federated Logistic Regression
########################################################################

function vertical_lr_train(server::AsynServer, clients::Vector{AsynClient}, time_limit::Float64)
    tag = true
    # set time limit
    @async begin
        sleep(time_limit)
        tag = false
    end
    # start training
    Threads.@threads for c in clients
        while tag
            # client compute and send embedding to server
            batch = send_embedding(c, server)
            # server compute and send back the gradient
            send_gradient(server, c.id, batch)
            # client update model
            update_model(c, batch)
            # time break
            sleep(c.ts)
        end
    end
    @printf "Finish training after %.2f seconds\n" time_limit
    # print number of communication rounds
    for c in clients
        @printf "Client %i communicate %i times with server \n" c.id c.num_commu
    end
end

function evaluation(server::AsynServer, clients::Vector{AsynClient})
    # test and train accuracy
    for c in clients
        # client compute and upload training embeddings
        send_embedding(c, server, tag = "training")
        # client compute and upload test embeddings
        send_embedding(c, server, tag = "test")
    end
    train_loss, train_acc, test_loss, test_acc = eval(server)
    @printf "Train Loss %.2f, Train Accuracy %.2f, Test Loss %.2f, Test Accuracy %.2f\n" train_loss train_acc test_loss test_acc
end


text:
using KernelGradients
using KernelAbstractions
using Test
using Enzyme

include("testsuite.jl")

struct CPUDeviceArray{T,N,A} end # Fake and unused
GradientsTestsuite.testsuite(CPU, "CPU", Base, Array, CPUDeviceArray)


text:
# plot first 5000 time points (before stabilization)
include("../Model_scripts/Functions.jl");
# plot the number of infected & uninfected hosts
plot_population_numbers(
    hcat(labels...),
    "c=0.5\n \nNumber of infected and uninfected hosts",
    100,
    5000,
    "$figure_directory/supplemental/supp_2a.png"
)
# plot the average weighted R0
plot_reproductive_rate(5, 5000, "$figure_directory/supplemental/supp_2c.png")
# plot the average V0
plot_vertical_reproductive_ratio(1, 5000, "$figure_directory/supplemental/supp_2e.png")
# plot the average weighted horizontal transmission
plot_horizontal_and_virulence(1, 5000, "$figure_directory/supplemental/supp_2g.png")
# plot the evenness through time
plot_evenness(1, 5000, "$figure_directory/supplemental/supp_2i.png")


text:
# DenseSampleSpace
s = samplespace(zeros(4,2), 100)
@test size(ShapleyValues.data(s)) == (4,0)
ShapleyValues.addsample!(s, ones(4))
@test ShapleyValues.data(s) == ones(4,1)
ShapleyValues.reset!(s)
@test size(ShapleyValues.data(s)) == (4,0)
ShapleyValues.addsample!(s, ones(4))
ShapleyValues.addsample!(s, ones(4))
@test ShapleyValues.data(s) == ones(4,2)

# SparseSampleSpace
s = samplespace(sparse(zeros(4,2)), 100)
@test size(ShapleyValues.data(s)) == (4,0)
ShapleyValues.addsample!(s, ones(4))
@test ShapleyValues.data(s) == ones(4,1)
ShapleyValues.reset!(s)
@test size(ShapleyValues.data(s)) == (4,0)
ShapleyValues.addsample!(s, ones(4))
ShapleyValues.addsample!(s, ones(4))
@test ShapleyValues.data(s) == ones(4,2)


text:
"""
    impsampling(lltarget, proposal, nsamples, nattemps)

Tries to produce samples using an importance sampling procedure with a proposal
in the Exponential Family (see module `ExpFamily`). The function lltarget takes
an array of points of dimensions `p x N` with `N` the number of samples to be
generated.
"""
function impsampling(lltarget::Function, proposal::EFamily, nsamples::Int;
                     nattempts=5)::Tuple{Matrix{Float},Vector{Float}}
    attempt   = 1
    notvalid  = true
    samples,w = [],[]
    while notvalid && attempt <= nattempts
        samples       = rand( proposal, nsamples )
        logliksamples = uloglik( proposal, samples)

        # IS weights
        lW  = lltarget( samples ) - logliksamples
        lW -= maximum(lW) # avoid underflow
        w   = exp.(lW)
        w  /= sum(w)

        notvalid = any(isnan.(w))
        attempt += 1
    end
    if attempt == nattempts+1
        throw("Could not generate a sample in $nattempts attempts")
    end
    (samples, w)
end


text:
function read_solution(filepath)
    sol = readlines(filepath)
    n_nodes = sol[1]
    tour = parse.(Int, split(join(sol[2:end]))) .+ 1
    return tour
end

function tour_length(tour, M)
    n_nodes = length(tour)
    len = 0
    for i in 1:n_nodes
        j = i + 1
        if i == n_nodes
            j = 1
        end

        len += M[tour[i], tour[j]]
    end
    return len
end

function cleanup(name)
    exts = ["mas", "pul", "sav", "sol", "tsp", "res"]
    for ext in exts
        file =  "$(name).$(ext)"
        rm(file, force=true)
        file =  "O$(name).$(ext)"
        rm(file, force=true)
    end
end


text:
pbex1 = ex1()

x0 = 1.5 * ones(2)
stop_1 = MPCCStopping(pbex1, MPCCAtX(x0, zeros(1)), optimality_check = WStat)

mpcc_solve(stop_1, verbose = false)

@show stop_1.current_state.x, status(stop_1)

pbex3 = ex3()

x0 = 1.5 * ones(3)
#stop_3 = MPCCStopping(pbex3, WStat, MPCCAtX(x0, zeros(2)), atol = 1e-3)
stop_3 = MPCCStopping(pbex3, MPCCAtX(x0, zeros(2)), atol = 1e-3, optimality_check = (x,y) -> WStat(x,y, actif = 1e-3))

mpcc_solve(stop_3)

@show stop_3.current_state.x, status(stop_3)


text:
using Distributed: @everywhere, @distributed, addprocs; addprocs(3)
@everywhere using Distributed: @spawnat
@everywhere include("GSTM_v2.jl")
using Statistics: norm

const hm_vectors   = 4
const vector_size  = 13
const storage_size = 25

const is_layers = [30]
const gs_layers = [35]
const go_layers = [35]

const max_t     = 100
const hm_data   = 20

const lr        = .0005
const hm_epochs = 100

make_model(hm_vectors, vector_size, storage_size, is_layers, gs_layers, go_layers) =
begin
    enc, enc_states = ENCODER(hm_vectors, vector_size, storage_size, is_layers, gs_layers, go_layers)
    dec, dec_states = DECODER(hm_vectors, vector_size, storage_size, is_layers, gs_layers, go_layers)
    model = [enc,dec]
    state = [enc_states,dec_states]
[model, state]
end

(encoder, decoder), (enc_zerostate, dec_zerostate) = make_model(hm_vectors, vector_size, storage_size, is_layers, gs_layers, go_layers)

make_data(hm_data; max_timesteps=50) =
begin
    data = []
    for i in 1:hm_data
        push!(data,
            [
                [[randn(1, vector_size) for iii in 1:hm_vectors] for ii in 1:rand(max_timesteps/2:max_timesteps)],
                [[randn(1, vector_size) for iii in 1:hm_vectors] for ii in 1:rand(max_timesteps/2:max_timesteps)]
            ]
        )
    end
data
end

make_accugrads(hm_vectors, vector_size, storage_size, is_layers, gs_layers, go_layers) =
begin
    accu_grads = []
    for model in [encoder, decoder]
        for mfield in fieldnames(typeof(model))
            net = getfield(model, mfield)
            for nfield in fieldnames(typeof(net))
                layer = getfield(net, nfield)
                for lfield in fieldnames(typeof(layer))
                    field_size = size(getfield(layer, lfield))
                    push!(accu_grads, zeros(field_size))
                end
            end
        end
    end
accu_grads
end

shuffle(arr_in) =
begin
    array_copy = copy(arr_in)
    array_new  = []
    while length(array_copy) > 0
        index = rand(1:length(array_copy))
        e = array_copy[index]
        deleteat!(array_copy, index)
        push!(array_new, e)
    end
array_new
end

accu_grads = make_accugrads(hm_vectors, vector_size, storage_size, is_layers, gs_layers, go_layers)

train(data, (encoder, decoder), enc_zerostate, lr, ep; accu_grads=nothing) =
begin
    @show lr

    losses = []

    for epoch in 1:ep

        print("Epoch ", epoch, ": ")

        loss = 0.0

        # for (g,l) in
        #
        #     (@distributed (vcat) for (x,y) in shuffle(data)
        #
        #         d = @diff sequence_loss(
        #             propogate(encoder, decoder, enc_zerostate, x, length(y)),
        #             y
        #         )
        #
        #         grads(d, encoder, decoder), value(d)
        #
        #     end)
        #
        #     upd!(encoder, decoder, g, lr)
        #     #upd_rms!(encoder, decoder, g, lr, accu_grads, alpha=.9)
        #     loss += sum(l)
        #
        #     for e in g
        #         @show norm(e)
        #     end
        #
        # end

        results =
            @distributed (vcat) for (x,y) in shuffle(data)

                d = @diff sequence_loss(
                        propogate(encoder, decoder, enc_zerostate, x, length(y)),
                        y
                )

                @spawnat 1 print("/")

                grads(d, encoder, decoder), value(d)

            end ; print("\n")

        gs = [zeros(size(e)) for e in results[1][1]]

        for (g,l) in results

            loss += l
            gs   += g

        end

        upd!(encoder, decoder, gs, lr)

        # for (name,g) in zip(names, gs)
        #
        #     # g = trunc(Int, norm(g))
        #     # @show name..., norm(g)
        #     # @info(name, norm(g))
        #c
        # end

        # show_details(encoder, decoder)

        @show loss ; push!(losses, loss)

    end
[encoder, decoder, losses]
end

names = []
for model in [encoder, decoder]
    m = typeof(model)
    for mfield in fieldnames(typeof(model))
        net = getfield(model, mfield)
        for nfield in fieldnames(typeof(net))
            layer = getfield(net, nfield)
            for lfield in fieldnames(typeof(layer))
                push!(names, [m, mfield, nfield, lfield])
            end
        end
    end
end

show_details(enc, dec) =
begin
    for model in [enc, dec]
        m = typeof(model)
        for mfield in fieldnames(typeof(model))
            net = getfield(model, mfield)
            for nfield in fieldnames(typeof(net))
                layer = getfield(net, nfield)
                for lfield in fieldnames(typeof(layer))
                    w_norm = trunc(Int, norm(getfield(layer, lfield)))
                    @info((m, mfield, nfield, lfield), w_norm)
                end
            end
        end
    end
end

@time (enc, dec, loss) =

    train(make_data(hm_data,
           max_timesteps=max_t),
         (encoder, decoder),
          enc_zerostate,
          lr,
          hm_epochs,
          accu_grads=accu_grads)

# using PyPlot: plot
# plot(loss, collect(1:hm_epochs), color="red", linewidth=2.0, linestyle="--")

max_stars_repo_name: developersfox/GSTM

max_issues_repo_name: developersfox/GSTM

max_forks_repo_name: developersfox/GSTM


text:
export MonteCarloLearner,
    FIRST_VISIT,
    EVERY_VISIT,
    NO_SAMPLING,
    ORDINARY_IMPORTANCE_SAMPLING,
    WEIGHTED_IMPORTANCE_SAMPLING

using StatsBase: countmap

abstract type AbstractVisitType end
struct FirstVisit <: AbstractVisitType end
struct EveryVisit <: AbstractVisitType end
const FIRST_VISIT = FirstVisit()
const EVERY_VISIT = EveryVisit()

abstract type AbstractSamplingStyle end
struct NoSampling <: AbstractSamplingStyle end
struct OrdinaryImportanceSampling <: AbstractSamplingStyle end
struct WeightedImportanceSampling <: AbstractSamplingStyle end
const NO_SAMPLING = NoSampling()
const ORDINARY_IMPORTANCE_SAMPLING = OrdinaryImportanceSampling()
const WEIGHTED_IMPORTANCE_SAMPLING = WeightedImportanceSampling()

"""
    MonteCarloLearner(;kwargs...)

Use monte carlo method to estimate state value or state-action value.

# Fields
- `approximator`::[`TabularApproximator`](@ref), can be either
  `TabularVApproximator` or `TabularQApproximator`.
- `γ=1.0`, discount rate.
- `kind=FIRST_VISIT`. Optional values are `FIRST_VISIT` or `EVERY_VISIT`.
- `sampling=NO_SAMPLING`. Optional values are `NO_SAMPLING`,
  `WEIGHTED_IMPORTANCE_SAMPLING` or `ORDINARY_IMPORTANCE_SAMPLING`.
"""
Base.@kwdef struct MonteCarloLearner{A,K,S} <: AbstractLearner
    approximator::A
    γ::Float64 = 1.0
    kind::K = FIRST_VISIT
    sampling::S = NO_SAMPLING
end

(learner::MonteCarloLearner)(env::AbstractEnv) = learner(state(env))
(learner::MonteCarloLearner)(s) = learner.approximator(s)
(learner::MonteCarloLearner)(s, a) = learner.approximator(s, a)

function RLBase.update!(::VBasedPolicy{<:MonteCarloLearner}, ::AbstractTrajectory) end

"Only update at the end of an episode"
function RLBase.update!(
    p::VBasedPolicy{<:MonteCarloLearner},
    t::AbstractTrajectory,
    ::AbstractEnv,
    ::PostEpisodeStage,
)
    update!(p.learner, t)
end

function RLBase.update!(
    L::MonteCarloLearner,
    t::AbstractTrajectory,
    e::AbstractEnv,
    s::PreActStage,
) end

function RLBase.update!(
    L::MonteCarloLearner,
    t::AbstractTrajectory,
    e::AbstractEnv,
    s::PostEpisodeStage,
)
    update!(L, t)
end

"Empty the trajectory at the end of an episode"
function RLBase.update!(
    t::AbstractTrajectory,
    ::Union{
        VBasedPolicy{<:MonteCarloLearner},
        QBasedPolicy{<:MonteCarloLearner},
        NamedPolicy{<:VBasedPolicy{<:MonteCarloLearner}},
    },
    ::AbstractEnv,
    ::PreEpisodeStage,
)
    empty!(t)
end

function RLBase.update!(L::MonteCarloLearner, t::AbstractTrajectory)
    _update!(L.kind, L.approximator, L.sampling, L, t)
end

function _update!(
    ::FirstVisit,
    ::Union{TabularVApproximator,LinearVApproximator},
    ::NoSampling,
    L::MonteCarloLearner,
    t::AbstractTrajectory,
)
    S, R = t[:state], t[:reward]
    V, G, γ = L.approximator, 0.0, L.γ
    state_counts = countmap(@view(S[1:end-1]))
    for i in length(R):-1:1
        s, r = S[i], R[i]
        G = γ * G + r
        if state_counts[S[i]] == 1
            update!(V, s => V(s) - G)  # first visit
        else
            state_counts[s] -= 1
        end
    end
end

function _update!(
    ::EveryVisit,
    ::Union{TabularVApproximator,LinearVApproximator},
    ::NoSampling,
    L::MonteCarloLearner,
    t::AbstractTrajectory,
)
    S, R = t[:state], t[:reward]
    V, G, γ = L.approximator, 0.0, L.γ
    for i in length(R):-1:1
        s, r = S[i], R[i]
        G = γ * G + r
        update!(V, s => V(s) - G)
    end
end

function _update!(
    ::EveryVisit,
    ::TabularQApproximator,
    ::NoSampling,
    L::MonteCarloLearner,
    t::AbstractTrajectory,
)
    S, A, R = t[:state], t[:action], t[:reward]
    γ, Q, G = L.γ, L.approximator, 0.0
    for i in length(R):-1:1
        s, a, r = S[i], A[i], R[i]
        G = γ * G + R[i]
        update!(Q, (s, a) => (Q(s, a) - G))
    end
end

function _update!(
    ::FirstVisit,
    ::TabularQApproximator,
    ::NoSampling,
    L::MonteCarloLearner,
    t::AbstractTrajectory,
)
    S, A, R = t[:state], t[:action], t[:reward]
    γ, Q, G = L.γ, L.approximator, 0.0
    seen_pairs = countmap(zip(@view(S[1:end-1]), @view(A[1:end-1])))

    for i in length(R):-1:1
        s, a, r = S[i], A[i], R[i]
        pair = (s, a)
        G = γ * G + r
        if seen_pairs[pair] == 1  # first visit
            update!(Q, (s, a) => (Q(s, a) - G))
        else
            seen_pairs[pair] -= 1
        end
    end
end

function _update!(
    ::FirstVisit,
    ::Tuple{
        <:Union{TabularVApproximator,LinearVApproximator},
        <:Union{TabularVApproximator,LinearVApproximator},
    },
    ::OrdinaryImportanceSampling,
    L::MonteCarloLearner,
    t::AbstractTrajectory,
)
    S, R, W = t[:state], t[:reward], t[:weight]
    (V, G), g, γ, ρ = L.approximator, 0.0, L.γ, 1.0
    seen_states = countmap(@view(S[1:end-1]))

    # @info "debug" S R W seen_states G V t[:action]
    for i in length(R):-1:1
        s, r = S[i], R[i]
        g = γ * g + r
        ρ *= W[i]
        if seen_states[s] == 1  # first visit
            update!(G, s => G(s) - ρ * g)
            update!(V, s => V(s) - G(s))
        else
            seen_states[s] -= 1
        end
    end
end

function _update!(
    ::FirstVisit,
    ::Tuple,
    ::WeightedImportanceSampling,
    L::MonteCarloLearner,
    t::AbstractTrajectory,
)
    S, R, W = t[:state], t[:reward], t[:weight]
    (V, G, Ρ), g, γ, ρ = L.approximator, 0.0, L.γ, 1.0
    seen_states = countmap(@view(S[1:end-1]))

    for i in length(R):-1:1
        s, r = S[i], R[i]
        g = γ * g + r
        ρ *= W[i]
        if seen_states[s] == 1  # first visit
            update!(G, s => G(s) - ρ * g)
            update!(Ρ, s => Ρ(s) - ρ)
            val = Ρ(s) == 0 ? 0 : G(s) / Ρ(s)
            update!(V, s => V(s) - val)
        else
            seen_states[s] -= 1
        end
    end
end


text:
# Library routines

function agexp(r, x, y)
    ccall((:agexp, libSPTK), Cdouble, (Cdouble, Cdouble, Cdouble), r, x, y)
end

gexp(r, x) = ccall((:gexp, libSPTK), Cdouble, (Cdouble, Cdouble), r, x)
glog(r, x) = ccall((:glog, libSPTK), Cdouble, (Cdouble, Cdouble), r, x)

mseq() = ccall((:mseq, libSPTK), Cint, ())

# solve (T + H)a = b
# NOTE: the order of arguments is different in theq in SPTK
# I would put the solution vector `a` as the first argument instead of third.
function theq!(a::StridedVector{Cdouble}, t::StridedVector{Cdouble},
               h::StridedVector{Cdouble}, b::StridedVector{Cdouble};
               min_det::Float64=1.0e-6)
    n = length(a)
    if length(b) != n || length(t) != n || length(h) != 2n-1
        throw(DimensionMismatch("inconsistent vector length"))
    end
    if min_det < 0.0
        throw(ArgumentError("min_det must be positive: min_det = $min_det"))
    end

    ret = ccall((:theq, libSPTK), Cint,
                (Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Cint,
                 Cdouble), t, h, a, b, n, min_det)
    if ret != 0
        error("failed to solve Toeplitz plus Hankel matrix system (T + H)a = b")
    end

    a
end

function theq(t::StridedVector{Cdouble}, h::StridedVector{Cdouble},
              b::StridedVector{Cdouble};
              kargs...)
    a = similar(t)
    theq!(a, t, h, b, kargs...)
end

# solve Ta = b
# NOTE: the order of arguments is different in toeplitz in SPTK
# I would put the solution vector `a` as the first argument instead of second.
function toeplitz!(a::StridedVector{Cdouble}, t::StridedVector{Cdouble},
                   b::StridedVector{Cdouble};
                   eps::Float64=1.0e-6)
    if length(a) != length(t) || length(a) != length(b)
        throw(DimensionMismatch("inconsistent vector length"))
    end
    if eps < 0.0
        throw(ArgumentError("eps must be positive: eps = $eps"))
    end

    ret = ccall((:toeplitz, libSPTK), Cint,
                (Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Cint, Cdouble),
                t, a, b, length(t), eps)
    if ret != 0
        error("failed to solve a symmetric Toeplitz set of linear equations")
    end

    a
end

function toeplitz(t::StridedVector{Cdouble}, b::StridedVector{Cdouble};
                  kargs...)
    a = similar(t)
    toeplitz!(a, t, b, kargs...)
end


text:
using LightXML
import LightXML: child_elements

include("schinke_stemmer.jl")

child_elements(x::Union{XMLElement, XMLNode}, name::AbstractString) =
    (XMLElement(n) for n in child_nodes(x) if is_elementnode(n) && LightXML.name(n) == name)

xdoc = parse_file("epigrammata.tei")
xroot = root(xdoc)
xbody = find_element(find_element(xroot, "text"), "body")
books = child_elements(xbody, "div1")

open("epigrams.txt", "w") do ef
    for book in books
        epigrams = child_elements(book, "div2")
        for epigram in epigrams
            isnothing(match(r"^[0-9]+$", attribute(epigram, "n"))) && continue
            lines = child_nodes(find_element(epigram, "p"))
            for node in lines
                !is_textnode(node) && continue
                line = content(node)
                words = (schinke_stemmer(m.match) for m in eachmatch(r"\w+", line))
                join(ef, words, ' ')
                write(ef, ' ')
            end
            write(ef, '\n')
        end
    end
end

max_stars_repo_name: phipsgabler/turing-lda

max_issues_repo_name: phipsgabler/turing-lda

max_forks_repo_name: phipsgabler/turing-lda


text:
function call(::Types.Create, args::Vector, storage, conn::IO)
    ip_addr = string(first(getpeername(conn)))
    name = args[1]
    password = "no password"

    if (length(args) == 2)
        password = String(Bcrypt.GenerateFromPassword(Array{UInt8,1}(args[2]), 0))
    end

    protected = password == "no password" ? 0 : 1
    user_id = storage.active_clients[ip_addr].id

    response = Requests.insert_channel(
        storage.sql_conn, name,
        protected, password,
        user_id)

    if (response)
        write(conn, "Successfully created the channel $name\n")
    else
        write(conn, "A channel with the name $name already exists\n")
    end
end


text:
# Load the CytofResearch environment
import Pkg; Pkg.activate("../../")

println("Loading libraries. Expect this to take a minute or two ...")
using CytofResearch
using Random
using BSON
println("Finished loading libraries ...")

flipbits(x; prob) = [prob > rand() ? 1 - xi : xi for xi in x]

function createSimData(;seed::Int, nfac::Int,
                       K=5, L=Dict(0=>3, 1=>3), J=20, eps=zeros(3))

  N = [8, 1, 2] * nfac
  I = length(N)

  # Set rng seed for reproducibility.
  Random.seed!(seed)

  # Initialize true Z matrix.
  Z = zeros(Int, J, K)

  # Regenerate Z unitl it is valid.
  while !CytofResearch.Model.isValidZ(Z)
    Z = CytofResearch.Model.genZ(J, K, .5)
    Z[:, 2] .= flipbits(Z[:, 1], prob=.1)
  end

  # Create true mu*
  mus = Dict(0 => -[1.0, 2.3, 3.5],
             1 => +[1.0, 2.0, 3.0])

  # Create true a_W
  a_W = rand(K) * 10

  # Create true a_eta
  a_eta = Dict(z => rand(L[z]) * 10 for z in 0:1)

  # Make simulated data
  simdat = CytofResearch.Model.genData(J=J, N=N, K=K, L=L, Z=Z,
                                beta=[-9.2, -2.3],
                                sig2=[0.2, 0.1, 0.3],
                                mus=mus,
                                a_W=a_W,
                                a_eta=a_eta,
                                sortLambda=false, propMissingScale=0.7,
                                eps=eps)
  return simdat
end

### Main ###
settings = [Dict(:nfac => 500, :K => 5, :seed => 90),
            Dict(:nfac => 5000, :K => 10, :seed => 1)]

for setting in settings
  nfac = setting[:nfac]
  seed = setting[:seed]
  K = setting[:K]

  println("Creating simulated data with: $(setting)")
  simdat = createSimData(seed=seed, nfac=nfac,
                         K=K, L=Dict(0=>3, 1=>3), J=20, eps=fill(.005, 3))

  # Save simulated data
  path_to_simdat = "../data/simdat-nfac$(nfac).bson"
  BSON.bson(path_to_simdat, simdat)
end

# NOTE: The ".bson" file is cross-platform, but requires the same BSON version.
#       To load the simulated data:
# simdat_small = BSON.load("../data/simdat-nfac500.bson")
# simdat_large = BSON.load("../data/simdat-nfac5000.bson")

max_stars_repo_name: luiarthur/CytofResearch

max_issues_repo_name: luiarthur/CytofResearch

max_forks_repo_name: luiarthur/CytofResearch


text:
using LinearAlgebra
import Base: eltype

struct Basis{T,U}
    grid::Grid{T,U}
    L′::Array{T,3}
end

function Basis(r::AbstractVector, n::Integer, args...)
    grid = Grid(r, n, args...)
    L′ = lagrangeder(grid)
    Basis(grid, L′)
end

function eval_element!(xⁱ, Nⁱ, x, χ, mrange)
    n = length(xⁱ)
    for (mi,m) in enumerate(mrange)
        Lₘ = lagrange(xⁱ, m, x)
        χ[:,mi] = Nⁱ[m]*Lₘ
    end
end

function evaluate!(basis::Basis, x::AbstractVector, χ::AbstractMatrix)
    g = basis.grid
    n = order(g)

    (last(x) < first(basis.grid.X) || first(x) > last(basis.grid.X)) && return χ

    sel = 1:1
    el1shift = n - length(bases(g, 1))
    for i in elems(g)
        sel = find_interval(g.X, x, i, sel)
        last(sel) < first(sel) && continue # This element is not covered by x
        b = bases(g, i)
        eval_element!(g.X[i,:], g.N[i,:], x[sel],
                      view(χ, sel,
                           eachindex(b) .+ (i-1)*(n-1) .- (i > 1 ? el1shift : 0)),
                      b)
    end
    χ
end

function block_structure(basis::Basis)
    g = basis.grid
    n = order(g)
    elrange = elems(g)
    if length(elrange) > 1
        vcat([n-1-(g.bl == :dirichlet0),1],
             vcat([[n-2,1] for i=2:length(elrange)-1]...),
             n-1-(g.br == :dirichlet0))
    else
        [n]
    end
end

function (basis::Basis)(x::AbstractVector)
    χ = spzeros(length(x),basecount(basis.grid))
    evaluate!(basis, x, χ)
end

DiagonalBlockDiagonal(A::AbstractMatrix, (rows,cols)::Tuple) =
    BandedBlockBandedMatrix(A, (rows,cols), (0,0), (0,0))

DiagonalBlockDiagonal(A::AbstractMatrix, rows) =
    DiagonalBlockDiagonal(A, (rows,rows))

"""
    basis(f)

Generate the scalar operator corresponding to `f(x)` on the FEDVR
`basis`.
"""
(basis::Basis)(f::Function) = basis(Diagonal(f.(locs(basis.grid))))

"""
    basis(D)

Generate the diagonal matrix with the appropriate block structure for
the FEDVR `basis`.
"""
function (basis::Basis)(D::Diagonal)
    n = basecount(basis)
    @assert size(D) == (n,n)
    DiagonalBlockDiagonal(D, block_structure(basis))
end

"""
    basis(I)

Return the identity matrix of `basis`.
"""
(basis::Basis)(::UniformScaling) = Diagonal(ones(basecount(basis.grid)))

locs(basis::Basis) = locs(basis.grid)
order(basis::Basis) = order(basis.grid)

eltype(basis::Basis) = eltype(basis.grid)

show(io::IO, basis::Basis) = write(io, "FEDVR Basis($(basis.grid))")

@recipe function plot(basis::Basis, n=1001)
    x = range(basis.grid.X[1,1], stop=basis.grid.X[end,end], length=n)
    xticks --> [basis.grid.X[:,1]..., basis.grid.X[end,end]]
    xlabel --> "x"
    legend --> false
    x,basis(x)
end

export evaluate


text:
"""
    cluster(inputs, radius::Float64; distfunc=euclidean, center=mean, verbose::Int64=1,
                cycle_lim::Int64=30, triangle=false)

A wrapper that decides whether or not to use triangle inequality.

"""
function cluster(inputs, radius::Float64; distfunc=euclidean, center=mean, verbose::Int64=1,
                cycle_lim::Int64=30, triangle=false)
    if triangle
        #use triangle, since euclid is a metric or approximation is okay
        return C2MC_dp_centers(inputs, radius; distfunc=distfunc, center=center, verbose=verbose,
                        cycle_lim=cycle_lim)
    else
        if verbose > 1
            verb = true
        else
            verb = false
        end
        #distfunc is non-euclidean and approximations are not acceptable
        return dp_centers(inputs, radius, distfunc=distfunc, center=center, verbose=verb,
                        cycle_lims=cycle_lim)
    end
end

"""
    get_initial_centroids(inputs, radius::Float64, input_length::Int64, distfunc, center, verbose::Int64)

Performs the first cycle of clustering in preparation for the creation of metacentroids
"""
function get_initial_centroids(inputs, radius::Float64, input_length::Int64, distfunc, center, verbose::Int64)
    num_C = 1
    C_arr = [inputs[1]]
    C_matrix = [falses(1, input_length)]
    for i in 1:length(inputs)
        curr_input = inputs[i]
        dists = [distfunc(C_arr[j], curr_input) for j in 1:num_C]
        if minimum(dists) > radius
            # Make a new cluster
            for k in 1:num_C #eliminate the input from its old cluster
                C_matrix[k][i] = false
            end
            num_C += 1 #increment clusters
            push!(C_matrix, falses(1, input_length)) #add to matrix
            C_matrix[num_C][i] = true #mark which input is in the new cluster
            push!(C_arr, curr_input)
        else
            # Assign point to best cluster
            min_index = indmin(dists)
            if !(C_matrix[min_index][i])
                for k in 1:num_C #eliminate the input from its old cluster
                    C_matrix[k][i] = false
                end
                C_matrix[min_index][i] = true #mark which input is in the new cluster
            end
        end
    end
    C_arr = [center(inputs[find(C_matrix[i])]) for i in 1:num_C]
    if verbose == 2
        println("cycle 1. n_clusters=$(num_C)")
    end
    return C_arr, C_matrix
end

"""
    get_initial_metacentroids(inputs, C_arr, MC_radius::Float64, distfunc, center, verbose::Int64)

Create metacentroids for the triangle inequality
"""
function get_initial_metacentroids(inputs, C_arr, MC_radius::Float64, distfunc, center, verbose::Int64)
    num_MC = 1
    changed = true
    num_C = length(C_arr)
    MC_arr = [C_arr[1]]
    MC_matrix = [falses(1, num_C)]
    C2MC_dist_matrix = Array{Array{Float64, 1}, 1}()
    changed = true
    for i in 1:num_C
        curr_C = C_arr[i]
        dists = [distfunc(MC_arr[j], curr_C) for j in 1:num_MC]
        if minimum(dists) > MC_radius
            for k in 1:num_MC #eliminate the input from its old cluster
                MC_matrix[k][i] = false
            end
            num_MC += 1 #increment clusters
            push!(MC_matrix, falses(1, num_C)) #add to matrix
            MC_matrix[num_MC][i] = true #mark which input is in the new cluster
            push!(MC_arr, curr_C)
        else
            # Assign point to best cluster
            min_index = indmin(dists)
            if !(MC_matrix[min_index][i])
                for k in 1:num_MC #eliminate the input from its old cluster
                    MC_matrix[k][i] = false
                end
                MC_matrix[min_index][i] = true #mark which input is in the new cluster
            end
        end
    end
    MC_arr = [center(C_arr[find(MC_matrix[MC_index])]) for MC_index in 1:num_MC]
    for MC_index in 1:num_MC #get distances between centroids and metacentroids
        C_indices = find(MC_matrix[MC_index])
        push!(C2MC_dist_matrix, [distfunc(MC_arr[MC_index], C_arr[index]) for index
                               in C_indices])
    end
    return MC_arr, C2MC_dist_matrix, MC_matrix
end

"""
    get_possible_centroids(MC_matrix, C_arr, R2MC_dist_arr::Array{Float64, 1},
                            C2MC_dist_matrix::Array{Array{Float64, 1}, 1}, radius::Float64, num_MC::Int64,
                            distfunc, MC_radius::Float64, verbose::Int64)

Uses the triangle inequality in order to determine which centroids a read may cluster to. This is
done by finding the difference between the read to metacentroid(R2MC_dist_arr) as well as centroid
to metacentroid distance(C2MC_dist_matrix), since the triangle inequality mandates that the
distance from read to centroid must be at least the difference between the other two sides

Returns (possible_centroids_indices) which contains the indices of the centroids a read may belong
to. The indices correspond to the centroids' locations in the C_arr.
"""

function get_possible_centroids(MC_matrix, C_arr, R2MC_dist_arr::Array{Float64, 1},
                            C2MC_dist_matrix::Array{Array{Float64, 1}, 1}, radius::Float64, num_MC::Int64,
                            distfunc, MC_radius::Float64, verbose::Int64)
    possible_centroids_indices = Array{Int32, 1}()
    total_radius = MC_radius + radius
    possible_metacentroids_indices = Array{Int32, 1}()
    for MC_index in 1:num_MC
        R2MC = R2MC_dist_arr[MC_index]
        if R2MC <= total_radius  #only reads at a certain dist from an MC qualifies
            push!(possible_metacentroids_indices, MC_index)
        end
    end

    for MC in possible_metacentroids_indices #go through each predetermined MC
        R2MC_arr = R2MC_dist_arr[MC]
        C_indices = find(MC_matrix[MC]) #indices of centroids belonging to this MC
        C2MC_distances = C2MC_dist_matrix[MC] #distances between each centroid of this MC
        num_possible_C = length(C2MC_distances)
        arr = abs.(C2MC_distances - R2MC_arr) #minimum read to centroid distance
        for C_index in 1:num_possible_C # go through each C in this MC
            if(arr[C_index] <= radius) #the minimum read to centroid must be less than radius
                push!(possible_centroids_indices, C_indices[C_index])
            end
        end
    end
    return possible_centroids_indices
end

"""
    C2MC_dp_centers(inputs, radius::Float64; distfunc=euclidean, center=mean, verbose::Int64=1,
                        cycle_lim::Int64=30)

Clusters the input kmer arrays using the triangle inequality for given radius
"""
function C2MC_dp_centers(inputs, radius::Float64; distfunc=euclidean, center=mean, verbose::Int64=1,
                        cycle_lim::Int64=30)
    input_length = length(inputs)

    #1. Get initial centroids. Looks correct(ran comparison with old_dp)
    C_arr, C_matrix = get_initial_centroids(inputs, radius, input_length, distfunc,
                                            center, verbose)
    MC_radius =  .5 * radius #how to determine optimum radius size for MCs?
    num_C = length(C_arr)

    #2. Create initial metacentroids as well as C2MC distances
    MC_arr, C2MC_dist_matrix, MC_matrix = get_initial_metacentroids(inputs,C_arr, MC_radius, distfunc,
                                                                    center, verbose)
    num_MC = length(MC_arr)
    #3. Get a matrix of the distances between the reads and metacentroids
    R2MC_dist_matrix = [distfunc(MC_arr[MC_index], inputs[input_index]) for input_index in
                        1:input_length, MC_index in 1:num_MC]

    total_proportions = 0
    cycles = 1 #initial centroid cycle already complete
    changed = true

    while changed && cycles < cycle_lim
        changed = false
        for i in 1:input_length
            curr_input = inputs[i]
            R2MC_dist_arr = R2MC_dist_matrix[i,:]
            #5. Triangle to get possible C's
            possible_centroids_arr = get_possible_centroids(MC_matrix, C_arr, R2MC_dist_arr,
                                                C2MC_dist_matrix, radius, num_MC, distfunc, MC_radius, verbose)
            dists = [distfunc(C_arr[j], curr_input) for j in possible_centroids_arr]

            new_proportion = length(possible_centroids_arr) / num_C
            total_proportions+=new_proportion
            #6. Insert read into C
            #the read neads a new centroid
            if length(dists) == 0 || minimum(dists) > radius
                # Make a new cluster
                for k in 1:num_C #eliminate the input from its old cluster
                    C_matrix[k][i] = false
                end
                num_C += 1 #increment clusters
                push!(C_matrix, falses(1, input_length)) #add to matrix
                C_matrix[num_C][i] = true #mark which input is in the new cluster
                push!(C_arr, curr_input)
                changed = true

                MC_distances = [distfunc(MC_arr[MC_index], curr_input) for MC_index in 1:num_MC]
                min_MC_index = indmin(MC_distances)
                minimum_distance = minimum(MC_distances)
                #the new centroid can be added to an existing metacentroid
                if minimum_distance <= MC_radius + radius
                    #add C to MC_matrix
                    for i in 1:num_MC
                        MC_matrix[i] = hcat(MC_matrix[i], false)
                    end
                    MC_matrix[min_MC_index][num_C] = true
                    push!(C2MC_dist_matrix[min_MC_index], minimum_distance)
                else #the new centroid requires a new metacentroid
                    for i in 1:num_MC
                        MC_matrix[i] = hcat(MC_matrix[i], false)
                    end
                    num_MC+=1

                    push!(MC_matrix, falses(1, num_C))
                    MC_matrix[num_MC][num_C] = true
                    push!(MC_arr, curr_input)
                    push!(C2MC_dist_matrix, [0.0])
                    new_dist_arr = [distfunc(inputs[j], curr_input) for j in 1:input_length]
                    R2MC_dist_matrix = hcat(R2MC_dist_matrix, new_dist_arr)
                end
            else
                # Assign point to best cluster
                min_index = possible_centroids_arr[indmin(dists)]
                if !(C_matrix[min_index][i]) #do a check to see if already in this cluster
                    for k in 1:num_C #eliminate the input from its old cluster
                        C_matrix[k][i] = false
                    end
                    C_matrix[min_index][i] = true #mark which cluster the input was moved to
                    changed= true
                end
            end
        end

        #7. Update C's
        C_arr = [center(inputs[find(C_matrix[i])]) for i in 1:num_C]

        #8. Update C2MC. Exactly the same as getting initials.
        cycles+=1

        empty!(C2MC_dist_matrix)
        for MC_index in 1:num_MC #get distances between centroids and metacentroids
            C_indices = find(MC_matrix[MC_index])
            push!(C2MC_dist_matrix, [distfunc(MC_arr[MC_index], C_arr[index]) for index
                                   in C_indices])
        end
        if verbose == 2
            println("cycle $(cycles). n_clusters=$(num_C)")
        end
    end

    total_proportions = total_proportions / (input_length * (cycles - 1))
    if verbose > 0
        println("A clustering call has converged!")
        println("The average proportion of centroids passing triangle inequality
                is $(total_proportions)")
        println("There are $(num_C) centroids and $(num_MC) metacentroids")
    end

    cluster_indices = [find(C_matrix[i]) for i in 1:num_C]
    # Cleaning out the empty clusters. Maybe only run this if there
    # are length 0 elements in cluster_indices.
    sizes = zeros(Int, num_C)
    for i in 1:num_C
        sizes[i] = length(cluster_indices[i])
    end

    clean_cluster_indices = Vector{Int}[]
    clean_μs = []  # TODO: set  type. But mean may not be same type as inputs!!!
    clean_sizes = Int[]
    for i in 1:num_C
        if length(cluster_indices[i]) > 0
            push!(clean_cluster_indices, cluster_indices[i])
            push!(clean_μs, C_arr[i])
            push!(clean_sizes, sizes[i])
        end
    end

    n_keep = length(clean_μs)
    centroid_inds = [-1 for i in 1:n_keep]
    for j in 1:n_keep
    dists = [distfunc(clean_μs[j], inputs[clean_cluster_indices[j][i]])
    for i in 1:length(clean_cluster_indices[j])]
        centroid_inds[j] = clean_cluster_indices[j][indmin(dists)]
    end
    return clean_μs, clean_sizes, clean_cluster_indices, centroid_inds
end

"""
    dp_means(input_vectors::Array{Vector{Int}, 1}, radius::Float64; verbose = false)

Cluster `input_vectors` using euclidean distance metric and arithmetic mean, where a
vector with distance greater than `radius` from the nearest cluster forms a new cluster.

Returns (centroid_vectors, cluster_sizes, cluster_indices, centroid_indices) where cluster_indices
is an array of arrays of indices of `input_vectors` grouped by cluster, and centroid_indices
is an array of values where each value indexes the vector in `input_vectors` closest to
the respective centroid.

Example:

```
inputs = [[5, 0, 0, 1, 1, 5],
          [5, 1, 0, 0, 1, 5],
          [5, 0, 1, 0, 1, 5],
          [0, 4, 6, 2, 0, 0],
          [0, 4, 6, 1, 1, 0],
          [0, 4, 6, 1, 0, 1]]
radius = 3.0
μs, sizes, indices, centroids = dp_means(inputs, radius)
μs == [[5, 1/3, 1/3, 1/3, 1, 5],
             [0, 4, 6, 4/3, 1/3, 1/3]]
sizes == [3, 3]
indices == [[1, 2, 3], [4, 5, 6]]
centroids == [1, 4]
```
"""
function dp_means(input_vectors::Vector{Vector{Int}}, radius::Float64; verbose = false)
    # TODO:
    # 0) Return exemplar indices for each cluster, where the vector
    #    closest to the cluster mean is selected. [DONE]
    # 0.1) Check if exemplars close to cluster means tend to have better accuracies.
    # 0.2) Return cluster assignments more conveniently. [DONE]
    # 1) Experiment with removing singleton clusters after convergence
    #    and starting the iteration again. [maybe nah]
    n_clusters = 1
    oldZs = [-1 for i in 1:length(input_vectors)]
    μs = [input_vectors[1]]
    newZs = [1 for i in 1:length(input_vectors)]
    sizes = zeros(n_clusters)
    cycles = 0
    while oldZs != newZs
        oldZs = copy(newZs)
        for i in 1:length(input_vectors)
            dists = [euclidean(μs[j], input_vectors[i]) for j in 1:n_clusters]
            if minimum(dists) > radius
                # Make a new cluster
                n_clusters += 1
                newZs[i] = n_clusters  #  Note we've just incremented n_clusters!
                push!(μs, input_vectors[i])
            else
                # Assign point to best cluster
                newZs[i] = indmin(dists)
            end
        end
        # Reset all means
        μs = μs * 0.0
        sizes = zeros(n_clusters)
        for i in 1:length(input_vectors)
            μs[newZs[i]] += input_vectors[i]
            sizes[newZs[i]] += 1
        end
        for j in 1:n_clusters
            μs[j] = μs[j] / sizes[j]
        end
        cycles += 1
        if verbose
            println("cycle $(cycles). n_clusters=$(n_clusters)")
        end
    end
    if verbose
        println("Converged!")
    end
    cluster_indices = [[] for i in 1:n_clusters]
    for i in 1:length(input_vectors)
        push!(cluster_indices[newZs[i]], i)
    end

    # Cleaning out the empty clusters. Maybe only run this if there
    # are length 0 elements in cluster_indices.
    clean_cluster_indices = Vector{Int}[]
    clean_μs = Vector{Float64}[]
    clean_sizes = Int[]
    for i in 1:length(cluster_indices)
        if length(cluster_indices[i]) > 0
            push!(clean_cluster_indices, cluster_indices[i])
            push!(clean_μs, μs[i])
            push!(clean_sizes, sizes[i])
        end
    end

    n_keep = length(clean_μs)
    if verbose
        println("keeping $(n_keep) after removing empty clusters.")
    end
    centroid_indices = [-1 for i in 1:n_keep]
    for j in 1:n_keep
        dists = [euclidean(clean_μs[j], input_vectors[clean_cluster_indices[j][i]])
                 for i in 1:length(clean_cluster_indices[j])]
        centroid_indices[j] = clean_cluster_indices[j][indmin(dists)]
    end
    return clean_μs, clean_sizes, clean_cluster_indices, centroid_indices
end

"""
    dp_centers(inputs, radius::Float64; distfunc = euclidean, center = mean, verbose = false, cycle_lim = 30)

Cluster `input_vectors` using given distance and mean calculations, where a
vector with distance greater than `radius` from the nearest cluster forms a new cluster.
Runs a maximum of `cycle_lim` iterations.

Returns (centroid_vectors, cluster_sizes, cluster_indices, centroid_indices) where cluster_indices
is an array of arrays of indices of `input_vectors` grouped by cluster, and centroid_indices
is an array of values where each value indexes the vector in `input_vectors` closest to
the respective centroid.

This is a generalization of `dp_means`.
"""
function dp_centers(inputs, radius::Float64;
                    distfunc = euclidean,
                    center = mean,
                    verbose = false,
                    cycle_lims = 30)

    # ToDo:
    # 0) Return exemplar indices for each cluster, where the vector
    #    closest to the cluster mean is selected. [DONE]
    # 0.1) Check if exemplars close to cluster means tend to have better accuracies. [YES THEY DOOOOOOO!]
    # 0.2) Return cluster assignments more conveniently. [DONE]
    # 1) Experiment with removing singleton clusters after convergence
    #    and starting the iteration again.
    n_clusters = 1
    oldZs = [-1 for i in 1:length(inputs)]
    μs = [inputs[1]]
    newZs = [1 for i in 1:length(inputs)]
    sizes = zeros(Int, n_clusters)
    cluster_indices = [[] for i in 1:n_clusters]
    cycles = 0
    while oldZs != newZs && cycles < cycle_lims
        oldZs = copy(newZs)
        for i in 1:length(inputs)
            dists = [distfunc(μs[j], inputs[i]) for j in 1:n_clusters]
            if minimum(dists) > radius
                # Make a new cluster
                n_clusters += 1
                # Note we've just incremented n_clusters!
                newZs[i] = n_clusters
                push!(μs, inputs[i])
            else
                # Assign point to best cluster
                newZs[i] = indmin(dists)
            end
        end

        # Reset all means
        cluster_indices = [[] for i in 1:n_clusters]
        for i in 1:length(inputs)
            push!(cluster_indices[newZs[i]], i)
        end

        μs = [center(inputs[cluster_indices[i]]) for i in 1:n_clusters]
        cycles += 1
        if verbose
            println("cycle $(cycles). n_clusters=$(n_clusters)")
        end
    end
    if verbose
        println("Converged!")
    end
    # Cleaning out the empty clusters. Maybe only run this if there
    # are length 0 elements in cluster_indices.
    sizes = zeros(Int, n_clusters)
    for i in 1:length(inputs)
        sizes[newZs[i]] += 1
    end
    clean_cluster_indices = Vector{Int}[]
    clean_μs = []  # TODO: set  type. But mean may not be same type as inputs!!!
    clean_sizes = Int[]
    for i in 1:length(cluster_indices)
        if length(cluster_indices[i]) > 0
            push!(clean_cluster_indices, cluster_indices[i])
            push!(clean_μs, μs[i])
            push!(clean_sizes, sizes[i])
        end
    end

    n_keep = length(clean_μs)
    if verbose
        println("n_keep=$(n_keep) after removing empty clusters.")
    end

    centroid_inds = [-1 for i in 1:n_keep]
    for j in 1:n_keep
        dists = [distfunc(clean_μs[j], inputs[clean_cluster_indices[j][i]])
                 for i in 1:length(clean_cluster_indices[j])]
        centroid_inds[j] = clean_cluster_indices[j][indmin(dists)]
    end
    return clean_μs, clean_sizes, clean_cluster_indices, centroid_inds
end

"""
    dp_centers(inputs, radii::Vector{Float64}; distfunc = euclidean, center = mean, verbose = false, cycle_lim = [30...])

Cluster `input_vectors` using given distance and mean calculations recursively; `radii` should
be an array of decreasing values, where each value is the radius of a successively finer clustering
operation. During each clustering, a vector with distance greater than `radius` from the nearest cluster forms a new cluster.
Runs a maximum of `cycle_lims` iterations during each cluster.

Returns (centroid_vectors, cluster_sizes, cluster_indices, centroid_indices) where cluster_indices
is an array of arrays of indices of `input_vectors` grouped by cluster, and centroid_indices
is an array of values where each value indexes the vector in `input_vectors` closest to
the respective centroid.

This is a recursive implementation of `dp_centers(inputs, radius)` for faster clustering.
"""
function dp_centers(inputs, radii::Array{Float64};
                    distfunc = euclidean,
                    center = mean,
                    verbose = false,
                    cycle_lims = 20)#fill(30, length(radii)))
    cycle_lims = 20; #remove this later
    μs, sizes, cluster_inds, centroid_inds = dp_centers(inputs, radii[1];
                                                        distfunc=distfunc,
                                                        center=center,
                                                        verbose=verbose,
                                                        cycle_lims=cycle_lims)#cycle_lims[1])
    if length(radii) > 1
        newClustsZ = [collect(dp_centers(inputs[cluster_inds[i]], radii[2:end],
                                                         distfunc=distfunc,
                                                         center=center,
                                                         verbose=verbose,
                                                         cycle_lims=cycle_lims))#cycle_lims[2:end]))
                for i in 1:length(cluster_inds)]

        # reformat to be consistent with nonrecursive dp_centers output.
        # "transpose" arrays, ie. [[[a, a, a], [b, b, b]], [[A, A, A], [B, B, B]]] to
        # [[[a, a, a], [A, A, A]], [[b, b, b], [B, B, B]]]
        newClusts = [[newClustsZ[i][j] for i in 1:length(newClustsZ)] for j in 1:4]

        # new, finer sizes, means, and centroid inds are calculated
        # but here we get the new cluster_inds and centroid_inds in terms of the newly computed clusters
        for i in 1:length(cluster_inds)
            for j in 1:length(newClusts[3][i])
                newClusts[3][i][j] = cluster_inds[i][newClusts[3][i][j]]
            end
            # merge finer clusters into bigger one to get relative indices of new centroids
            # FIX: currently relies on dp_centers returning sorted lists of indices
            newClusts[4][i] = sort(vcat(newClusts[3][i]...))[newClusts[4][i]]
        end

        # merge sub-clusters into a single array
        for i in 1:4
            newClusts[i] = vcat(newClusts[i]...)
        end
        return newClusts
    else
        return μs, sizes, cluster_inds, centroid_inds
    end
end


text:
#####

function prepcfull(P::Array{Vector{RT},D})::SharedArray{Float64,1} where {RT,D}
    N = sum( prod( length(P[i][d]) for d = 1:length(P[i])) for i = 1:length(P) )

    return SharedArray{Float64}(N)
end

# common kernel.

function fitpatchRKHS!( c_full::SharedArray{T,1},
                        X_full_nD::Array{Vector{T},D},
                        y_full::SharedArray{T,D},
                        P,
                        θ::KT,
                        σ²) where {T,D,KT}

    #
    workerfunc! = ii->fitpatchRKHSworker!(c_full,
                            X_full_nD,
                            y_full,
                            P,
                            θ,
                            σ²,
                            ii)

    pmap(workerfunc!, collect(1:length(P)))

    return nothing
end

function fitpatchRKHSworker!( c_full::SharedArray{T,1},
                        X_full_nD::Array{Vector{T},D},
                        y_full::SharedArray{T,D},
                        P,
                        θ::KT,
                        σ²,
                        𝑖::Int) where {T,D, KT}

    # quantities that won't change.
    X = vec(X_full_nD[P[𝑖]...])
    #x_ranges = P[𝑖]
    #X_nD = Utilities.ranges2collection(x_ranges, Val(D))
    #X = vec(X_nD)

    N = length(X)
    w_X = Vector{T}(undef, N)

    # kernel matrix.
    K = Matrix{T}(undef, N, N)
    constructkernelmatrix!(K,
                                w_X,
                                X,
                                θ)

    #
    for i = 1:N
        K[i,i] += σ²[1]
    end

    # compute st, fin.
    st = 1
    if 𝑖 > 1
        st = 1 + sum( prod( length(P[i][d]) for d = 1:length(P[i])) for i = 1:𝑖-1 )
    end
    fin = st + N -1

    c_full[st:fin] = K\vec(y_full[P[𝑖]...])

    return nothing
end

function querypatchRKHS!(   fq::SharedArray{T,D},
                            c_full::SharedArray{T,1},
                            X_full_nD::Array{Vector{T},D},
                            P,
                            θ::KT,
                            Xq_full_nD::Array{Vector{T},D},
                            J) where {T,D,KT}

    #
    workerfunc! = ii->querypatchRKHSworker!(fq,
                            c_full,
                            X_full_nD,
                            P,
                            θ,
                            Xq_full_nD,
                            J,
                            ii)

    pmap(workerfunc!, collect(1:length(P)))

    return nothing
end

"""
P contain the ranges for the observed image.
J contain the ranges for the query image.
"""
function querypatchRKHSworker!( fq::SharedArray{T,D},
                        c_full::SharedArray{T,1},
                        X_full_nD::Array{Vector{T},D},
                        P,
                        θ::KT,
                        Xq_full_nD::Array{Vector{T},D},
                        J,
                        𝑖::Int) where {T,D, KT}

    # quantities that won't change.
    X = vec(X_full_nD[P[𝑖]...])
    Xq_nD = Xq_full_nD[J[𝑖]...]

    N = length(X)
    w_X::Vector{T} = θ.warpfunc[1].(X)

    # compute st, fin.
    st = 1
    if 𝑖 > 1
        st = 1 + sum( prod( length(P[i][d]) for d = 1:length(P[i])) for i = 1:𝑖-1 )
    end
    fin = st + N -1
    c = c_full[st:fin]

    fq_a = xx->evalquery(xx, c, X, w_X, θ)

    fq[J[𝑖]...] = fq_a.(Xq_nD)

    return nothing
end

# function getXcoords(P, X::Vector{Vector{T}}) where T
#     #
#
#     # interior patches.
#     X_set = Vector{Vector{LinRange{T,Int}}}(undef, size(P))
#     for j = 1:size(P,2)
#         for i = 1:size(P,1)
#             out[i,j] = X[P[i,j]...]
#
#     return
# end

function ind2sub(a, i)
    i2s = CartesianIndices(a)
    return i2s[i]
end

function extractnonoverlapboundaries(  P::Array{Vector{UnitRange{Int}},D},
                                        overlap::Int,
                                        y_sz,
                                        dummy_val::T) where {T,D}
    #
    U = Array{Vector{Vector{T}},D}(undef, size(P))

    for n = 1:length(U)

        U[n] = Vector{Vector{T}}(undef,D)
        for d = 1:D
            st = P[n][d][1]
            fin = P[n][d][end]

            u_st = st + overlap/2 - 0.5
            if st == 1
                u_st = 1
                #u_st = -Inf
            end

            u_fin = fin - overlap/2 + 0.5
            if fin == y_sz[d]
                u_fin = y_sz[d]
                #u_fin = Inf
            end

            U[n][d] = [u_st; u_fin]
        end
    end

    return U
end

# ϵ is numerical tolerance.
function isin(xq::Vector{T}, patch_ranges, x_ranges, y_sz; ϵ::T = eps(T)*2) where T
    D = length(xq)
    @assert D == length(patch_ranges) == length(x_ranges)

    for d = 1:D

        st = round(Int,max(patch_ranges[d][1],1))
        fin = round(Int,min(patch_ranges[d][end], y_sz[d]))

        xd_st = x_ranges[d][st]
        xd_fin = x_ranges[d][fin]
        if !(xd_st - ϵ < xq[d] < xd_fin + ϵ)
            return false
        end
    end

    return true
end

# to do: returns closest patch if extrapolation case.
# for now: interpolation only.
function findpatch(xq::Vector{T}, U, y_sz, x_ranges)::Int where T
    N_patches = length(P)
#println("hi")
    for i = 1:N_patches

        if isin(xq, U[i], x_ranges, y_sz)
            return i
        end
    end

    println("Warning: Extrapolation case detected. Returning last patch.")
    return N_patches
end

function prepXqset( Xq,
                    U::Array{Vector{Vector{T}},D},
                    y_sz,
                    x_ranges) where {T,D}

    Nq = length(Xq)

    Xq_set = Array{Vector{Vector{T}},D}(undef, size(U))
    𝓘q_set = Array{Vector{Int},D}(undef, size(U))
    for n = 1:length(Xq_set)
        Xq_set[n] = Vector{Vector{T}}(undef, 0)
        𝓘q_set[n] = Vector{Int}(undef,0)
    end

    for i = 1:Nq
        #n = 0
        #println("i = ", i, ", n = ", n)
        n = findpatch(Xq[i], U, y_sz, x_ranges)

        push!(Xq_set[n], Xq[i])
        push!(𝓘q_set[n], i)
    end

    return Xq_set, 𝓘q_set
end

function parseyq(   sol,
                    𝓘q_set::Array{Vector{Int},D},
                    Nq_array) where {T,D}

    yq = Array{Float64,D}(undef, Nq_array...)

    for i = 1:length(sol)
        𝑖 = sol[i][2]
        patch_query = sol[i][1]

        for j = 1:length(patch_query)
            yq[𝓘q_set[𝑖][j]] = patch_query[j]
        end
    end

    return yq
end

# function reassemblequery(   Xq_set::Array{Vector{Vector{T}},D},
#                             𝓘q_set::Array{Vector{Int},D},
#                             yq_vec) where {T,D}
#     #
#     for n = 1:length(𝓘q_set)
#         patch = 𝓘q_set[n]
#
#         for i = 1:length(patch)
#             k = 𝓘q_set[n][i]
#             yq[k] = Xq
#         end
#     end
#
#     return 0
# end

# speed up idea: make X_full_nD into a Matrix.
function querypatchRKHS2(   Nq_array,
                            c_full::SharedArray{T,1},
                            #X_full_nD::Array{Vector{T},D},
                            X_set,
                            P::Array{Vector{UnitRange{Int}},D},
                            Xq_set,
                            U,
                            θ::KT,
                            𝓘q_set) where {T,D,KT}

    #
    N_patches = length(P)
    @assert size(U) == size(P) == size(Xq_set) == size(𝓘q_set)

    # c_set = Vector{Vector{T}}(undef, N_patches)
    # for 𝑖 = 1:N_patches
    #     st = 1
    #     if 𝑖 > 1
    #         st = 1 + sum( prod( length(P[i][d]) for d = 1:length(P[i])) for i = 1:𝑖-1 )
    #     end
    #     N = prod( length(P[𝑖][d]) for d = 1:D )
    #     fin = st + N -1
    #
    #     c_set[𝑖] = c_full[st:fin]
    # end

    # workerfunc! = ii->querypatchRKHSworker2(
    #                         c_full,
    #                         #c_set,
    #                         #X_full_nD,
    #                         X_set,
    #                         P,
    #                         Xq_set,
    #                         U,
    #                         θ,
    #                         𝓘q_set,
    #                         ii)
    #
    # sol = pmap(workerfunc!, collect(1:length(P)))
    # yqx = parseyq(sol, 𝓘q_set, Nq_array)
    #
    # return yqx

    yq = SharedArray{Float64}(Nq_array...)
    fill!(yq,Inf)
    workerfunc! = ii->querypatchRKHSworker3!( yq,
                            c_full,
                            X_set,
                            P,
                            Xq_set,
                            U,
                            θ,
                            𝓘q_set,
                            ii)
    pmap(workerfunc!, collect(1:length(P)))

    yqx = convert(Array{T,D}, yq)
    return yqx
end

function querypatchRKHSworker3!( yq::SharedArray{T,D},
                        c_full::SharedArray{T,1},
                        X_set,
                        P,
                        Xq_set::Array{Vector{Vector{T}},D},
                        U,
                        θ::KT,
                        𝓘q_set,
                        𝑖::Int) where {T,D, KT}

    println("Start patch ", 𝑖)

    # quantities that won't change.
    X = X_set[𝑖]
    Xq = Xq_set[𝑖]

    N = length(X)

    w_X::Vector{T} = θ.warpfunc[1].(X)

    # compute st, fin.
    st = 1
    if 𝑖 > 1
        st = 1 + sum( prod( length(P[i][d]) for d = 1:length(P[i])) for i = 1:𝑖-1 )
    end
    fin = st + N -1

    c = c_full[st:fin]
    #c = c_set[𝑖]
    #@assert prod( length(P[𝑖][d]) for d = 1:D ) == N

    # fq_a = xx->evalquery(xx, c, X, w_X, θ)
    #
    # # 𝓘q = 𝓘q_set[𝑖]
    # # for j = 1:length(Xq)
    # #     fq[𝓘q[j]] = fq_a(Xq[j])
    # # end
    #
    # return fq_a.(Xq), 𝑖

    # Nq = length(Xq)
    # out = Vector{T}(undef,Nq)
    # for i = 1:Nq
    #     out[i] = evalquery(Xq[i], c, X, w_X, θ)
    # end
    #
    # return out, 𝑖

    Nq = length(Xq)
    for j = 1:Nq
        yq[𝓘q_set[𝑖][j]] = evalquery(Xq[j], c, X, w_X, θ)
    end
    println("Done patch ", 𝑖)

    return nothing
end

function querypatchRKHSworker2(
                        c_full::SharedArray{T,1},
                        #c_set,
                        #X_full_nD::Array{Vector{T},D},
                        X_set,
                        P,
                        Xq_set::Array{Vector{Vector{T}},D},
                        U,
                        θ::KT,
                        𝓘q_set,
                        𝑖::Int) where {T,D, KT}

    # quantities that won't change.

    #X = vec(X_full_nD[P[𝑖]...])
    X = X_set[𝑖]

    Xq = Xq_set[𝑖]

    N = length(X)

    w_X::Vector{T} = θ.warpfunc[1].(X)

    # compute st, fin.
    st = 1
    if 𝑖 > 1
        st = 1 + sum( prod( length(P[i][d]) for d = 1:length(P[i])) for i = 1:𝑖-1 )
    end
    fin = st + N -1

    c = c_full[st:fin]
    #c = c_set[𝑖]
    #@assert prod( length(P[𝑖][d]) for d = 1:D ) == N

    # fq_a = xx->evalquery(xx, c, X, w_X, θ)
    #
    # # 𝓘q = 𝓘q_set[𝑖]
    # # for j = 1:length(Xq)
    # #     fq[𝓘q[j]] = fq_a(Xq[j])
    # # end
    #
    # return fq_a.(Xq), 𝑖

    Nq = length(Xq)
    out = Vector{T}(undef,Nq)
    for i = 1:Nq
        out[i] = evalquery(Xq[i], c, X, w_X, θ)
    end

    return out, 𝑖
end

max_stars_repo_name: RoyCCWang/PatchMixtureKriging

max_issues_repo_name: RoyCCWang/PatchMixtureKriging


text:
"""
    add_dict_vals!(d)

Internal function to add the values of a dictionary as keys with the same values. So for instance
if the dictionary has a pair "s" => "square", it will add "square" => "square".
"""
function add_dict_vals!(d::Dict{String,String})
    for v ∈ values(d)
        get(d, v) do
            d[v] = v
        end
    end
    return nothing
end

const GLE_MARKERS = Dict{String,String}(
    "^"  => "triangle"   ,
    "w^" => "wtriangle"  ,
    "f^" => "ftriangle"  ,
    "s"  => "square"     ,
    "ws" => "wsquare"    ,
    "fs" => "fsquare"    ,
    "o"  => "circle"     ,
    "wo" => "wcircle"    ,
    "."  => "fcircle"    ,
    "fo" => "fcircle"    ,
    "d"  => "diamond"    ,
    "wd" => "wdiamond"   ,
    "fd" => "fdiamond"   ,
    "x"  => "cross"      ,
    "+"  => "plus"       ,
    "-"  => "minus"      ,
    "*"  => "asterisk"   ,
    "o." => "odot"       ,
    "o-" => "ominus"     ,
    "o+" => "oplus"      ,
    "ox" => "otimes"     ,
    "clubs"   => "club"  ,
    "hearts"  => "heart" ,
    "spades"  => "spade" ,
    "dagger"  => "dag"   ,
    "ddagger" => "ddag"  ,
    # ones without direct abbreviation / alias
    "star"    => "star"  ,
    "star2"   => "star2" ,
    "star3"   => "star3" ,
    "star4"   => "star4" ,
    "flower"  => "flower",
    "snake"   => "snake" ,
    "none"    => "none"  ,
    )
add_dict_vals!(GLE_MARKERS)

const GLE_LSTYLES = Dict{String,Int}(
    "-"    => 0,
    "--"   => 9,
    "-."   => 8,
    "none" => -1,
    )

const GLE_TEXSCALE = ("scale", "fixed", "none")

const GLE_FONTS = Dict{String,String}(
    "roman"                 => "rm"     ,
    "roman-bold"            => "rmb"    ,
    "roman-italic"          => "rmi"    ,
    "sans-serif"            => "ss"     ,
    "sans-serif-bold"       => "ssb"    ,
    "sans-serif-italic"     => "ssi"    ,
    "typewriter"            => "tt"     ,
    "typewriter-bold"       => "ttb"    ,
    "typewriter-italic"     => "tti"    ,
    "times"                 => "pstr"   ,
    "times-roman"           => "pstr"   ,
    "times-bold"            => "pstb"   ,
    "times-italic"          => "psti"   ,
    "times-bold-italic"     => "pstbi"  ,
    "courier"               => "psc"    ,
    "courier-bold"          => "pscb"   ,
    "courier-oblique"       => "psco"   ,
    "courier-bold-oblique"  => "pscbo"  ,
    "helvetica"                     => "psh"    ,
    "helvetica-bold"                => "pshb"   ,
    "helvetica-oblique"             => "psho"   ,
    "helvetica-bold-oblique"        => "pshbo"  ,
    "helvetica-narrow"              => "pshn"   ,
    "helvetica-narrow-bold"         => "pshnb"  ,
    "helvetica-narrow-oblique"      => "pshno"  ,
    "helvetica-narrow-bold-oblique" => "pshnbo" ,
    "avantgarde"                    => "psagb"  ,
    "avantgarde-book"               => "psagb"  ,
    "avantgarde-demi"               => "psagd"  ,
    "avantgarde-book-oblique"       => "psagbo" ,
    "avantgarde-demi-oblique"       => "psagdo" ,
    "bookman"               => "psbl"   ,
    "bookman-light"         => "psbl"   ,
    "bookman-demi"          => "psbd"   ,
    "bookman-light-italic"  => "psbli"  ,
    "bookman-demi-italic"   => "psbdi"  ,
    "newcentury"            => "psncsr" ,
    "newcentury-roman"      => "psncsr" ,
    "newcentury-bold"       => "psncsb" ,
    "newcentury-italic"     => "psncsi" ,
    "newcentury-bold-italic"=> "psncsbi",
    "palatino"              => "pspr"   ,
    "palatino-roman"        => "pspr"   ,
    "palantino-bold"        => "pspb"   ,
    "palatino-italic"       => "pspi"   ,
    "palatino-bold-italic"  => "pspbi"  ,
    "cmr"   => "texcmr"   ,
    "cmb"   => "texcmb"   ,
    "cmti"  => "texcmti"  ,
    "cmmi"  => "texcmmi"  ,
    "cmss"  => "texcmss"  ,
    "cmssb" => "texcmssb" ,
    "cmssi" => "texcmssi" ,
    "cmtt"  => "texcmtt"  ,
    "cmitt" => "texcmitt" ,
    "computer-modern-roman"             => "texcmr"  ,
    "computer-modern-bold"              => "texcmb"  ,
    "computer-modern-text-italic"       => "texcmti" ,
    "computer-modern-maths-italic"      => "texcmmi" ,
    "computer-modern-sans-serif"        => "texcmss" ,
    "computer-modern-sans-serif-bold"   => "texcmssb",
    "computer-modern-sans-serif-italic" => "texcmssi",
    "computer-modern-typewriter"        => "texcmtt" ,
    "computer-modern-italic-typewriter" => "texcmitt",
    # fonts with no description (yet)
    "plsr"   => "plsr",
    "pldr"   => "pldr",
    "pltr"   => "pltr",
    "plti"   => "plti",
    "plcr"   => "plcr",
    "plci"   => "plci",
    "plss"   => "plss",
    "plcs"   => "plcs",
    "plsa"   => "plsa",
    "plba"   => "plba",
    "plge"   => "plge",
    "plgg"   => "plgg",
    "plgi"   => "plgi",
    "plsg"   => "plsg",
    "pshc"   => "pshc",
    "pshcb"  => "pshcb",
    "pshcdo" => "pshcdo",
    "pssym"  => "pssym",
    "zapfchancery-medium-italic" => "pszd",
    "arial" => "arial8",
    "arial-bold" => "arial8b",
    "arial-bold-italic" => "arial8bi",
    "arial-italic" => "arial8i",
    "plcc" => "plcc",
    "plcg" => "plcg",
    "pscb" => "pscb",
    "pshcbo" => "pshcbo",
    "pszcmi" => "pszcmi",
    "texcmex" => "texcmex",
    "texcmsy" => "texcmsy",
    "texex" => "texex",
    "texmi" => "texmi",
    "texsy" => "texsy",
    "times" => "times8",
    "times-bold" => "times8b",
    "typerwiter-bold-italic" => "ttbi",
    )
add_dict_vals!(GLE_FONTS)

const GLE_POSITION = Dict{String,String}(
    "top-left"      => "tl",
    "bottom-left"   => "bl",
    "top-right"     => "tr",
    "bottom-right"  => "br",
    "top-center"    => "tc",
    "bottom-center" => "bc",
    "left-center"   => "lc",
    "right-center"  => "rc",
    "center"        => "cc",
    )
add_dict_vals!(GLE_POSITION)


text:
function deref(pointer::Int, program::Array{Int,1}, mode::Int, base::Int)::Int
    if mode == 0
        return program[pointer+1]
    elseif mode == 1
        return pointer
    elseif mode == 2
        return program[pointer+base+1]
    else
        throw(ErrorException("Invalid mode: " * mode))
    end
end

function write_pointer(pointer::Int, mode::Int, base::Int)::Int
    if mode == 0
        return pointer+1
    elseif mode == 2
        return pointer+base+1
    else
        throw(ErrorException("Invalid mode for writing: " * mode))
    end
end

function run_intcode(program::Vector{Int}, input::Channel, output::Channel)::Int
    program = vcat(program, zeros(Int, 1000))
    pointer = 1
    ret = nothing
    base = 0
    while pointer <= length(program)
        # println(program)
        instruction = reverse(digits(program[pointer], pad=5))
        # println(instruction)
        opcode = instruction[end] + instruction[end-1]*10
        modes = reverse(instruction[1:end-2])
        # println(join([opcode,": ",modes]))
        if opcode == 1 # add
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            pt3 = program[pointer+3]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            program[write_pointer(pt3, modes[3], base)] = val1 + val2
            pointer = pointer + 4
        elseif opcode == 2 # multiply
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            pt3 = program[pointer+3]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            program[write_pointer(pt3, modes[3], base)] = val1 * val2
            pointer = pointer + 4
        elseif opcode == 3 # store input
            pt1 = program[pointer+1]
            program[write_pointer(pt1, modes[1], base)] = take!(input)
            pointer = pointer + 2
        elseif opcode == 4 # output
            pt1 = program[pointer+1]
            val1 = deref(pt1, program, modes[1], base)
            pointer = pointer + 2
            ret = val1
            put!(output, val1)
        elseif opcode == 5 # jump if true
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            if val1 != 0
                pointer = val2 + 1
            else
                pointer = pointer + 3
            end
        elseif opcode == 6 # jump if false
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            if val1 == 0
                pointer = val2 + 1
            else
                pointer = pointer + 3
            end
        elseif opcode == 7 # less than
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            pt3 = program[pointer+3]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            program[write_pointer(pt3, modes[3], base)] = val1 < val2 ? 1 : 0
            pointer = pointer + 4
        elseif opcode == 8 # equal
            pt1 = program[pointer+1]
            pt2 = program[pointer+2]
            pt3 = program[pointer+3]
            val1 = deref(pt1, program, modes[1], base)
            val2 = deref(pt2, program, modes[2], base)
            program[write_pointer(pt3, modes[3], base)] = val1 == val2 ? 1 : 0
            pointer = pointer + 4
        elseif opcode == 9 # modify base
            pt1 = program[pointer+1]
            val1 = deref(pt1, program, modes[1], base)
            base = base + val1
            pointer = pointer + 2
        elseif opcode == 99
            close(output)
            return ret
        else
            throw(ErrorException("Invalid opcode: " * string(opcode)))
        end
    end
end

function draw_map(map, robot, direction, min=nothing, max=nothing)
    map = string.(map)
    map[robot[2], robot[1]] = if direction == 0
        "^"
    elseif direction == 90
        "<"
    elseif direction == 180
        "v"
    elseif direction == 270
        ">"
    end
    if min === nothing || max === nothing
        for row in eachrow(map[robot[2]-2:robot[2]+2,:])
            println(replace(replace(join(row[robot[1]-3:robot[1]+3]), "0" => "."), "1" => "#"))
        end
    else
        for row in eachrow(map[min[2]-1:max[2]+1,:])
            println(replace(replace(join(row[min[1]-1:max[1]+1]), "0" => " "), "1" => "#"))
        end
    end
end

function process_output(in::Channel, out::Channel)
    hull = zeros(Int, 1000, 1000)
    hull[500, 500] = 1
    robot = [500, 500]
    max = [500, 500]
    min = [500, 500]
    paint = turn = 0
    direction = 0 # up
    painted_coords = []
    counter = 0
    while true
        if counter > 10
            # throw(ErrorException)
        end
        counter = counter + 1
        put!(out, hull[robot[2], robot[1]])
        try
            paint = take!(in)
            turn = take!(in)
        catch e
            if isa(e, InvalidStateException)
                break
            else
                rethrow
            end
        end

        # println(counter, ": robot at ",robot," (", direction ,"): ",hull[robot[2], robot[1]]," -> ", paint)
        hull[robot[2], robot[1]] = paint
        push!(painted_coords, deepcopy(robot))

        if turn == 0
            direction = direction + 90
        elseif turn == 1
            direction = direction - 90
        else
            throw(ErrorException("Invalid turn: " * turn))
        end
        if direction >= 360
            direction = direction - 360
        elseif direction < 0
            direction = direction + 360
        end
        if direction == 0 # up
            robot[2] = robot[2] - 1
        elseif direction == 90 # left
            robot[1] = robot[1] - 1
        elseif direction == 180 # down
            robot[2] = robot[2] + 1
        elseif direction == 270 # right
            robot[1] = robot[1] + 1
        else
            throw(ErrorException("Invalid direction: " * direction))
        end

        if robot[1] < min[1]
            min[1] = robot[1]
        end
        if robot[1] > max[1]
            max[1] = robot[1]
        end
        if robot[2] < min[2]
            min[2] = robot[2]
        end
        if robot[2] > max[2]
            max[2] = robot[2]
        end
        # println("Turned ",turn == 0 ? "left" : "right", " (", turn, ") to face ",direction, ". Visited ", length(painted_coords), " points (", length(unique(painted_coords)), " unique)")
        # println(painted_coords)
        # draw_map(hull, robot, direction)
    end
    draw_map(hull, robot, direction, min, max)
    println("Visited ", length(painted_coords), " points (", length(unique(painted_coords)), " unique)")
    return length(unique(painted_coords))
end

program = open("paintrobot.csv") do f
    parse.(Int, split(readlines(f)[1], ","))
end

in = Channel(1)
out = Channel(2)

@async run_intcode(deepcopy(program), in, out)
println(process_output(out, in))