using JSON, Colors, Memoize

# Grabbing the open league rankings from PvPoke. These are common enough that
# automatically downloading these and making them constant makes sense.
const gamemaster =     JSON.parsefile(download(
    "https://raw.githubusercontent.com/pvpoke/pvpoke/master/src/data/" *
    "gamemaster.json"))

"""
    get_cp_limit(league)

Given a league ("great", "ultra", "master") return its cp limit
(1500, 2500, 10000). Masters has no cp limit but is signified in PvPoke with the
arbitrarily large 10000.

# Examples
```jldoctest
julia> get_cp_limit("great")
1500
"""
get_cp_limit(league::String) = league == "master" ? 10_000 :
    league == "ultra" ? 2_500 : 1_500

"""
    get_rankings(cup; league = "great")

Download the PvPoke rankings for a particular meta, optionally within a league
that is not great league. This function is memoized to avoid downloading the
same file multiple times.
"""
@memoize function get_rankings(cup::String; league = "great")
    return JSON.parsefile(download(
        "https://raw.githubusercontent.com/pvpoke/pvpoke/master/src/" *
        "data/rankings/$(cup)/overall/rankings-$(get_cp_limit(league)).json"))
end

"""
    get_overrides(cup; league = "great")

Download the PvPoke overrides used to generate rankings for a particular meta,
optionally within a league that is not great league. This function is memoized
to avoid downloading the same file multiple times.
"""
@memoize function get_overrides(cup::String; league = "great")
    cup == "great" && return get_overrides("overall", league = "great")
    cup == "ultra" && return get_overrides("overall", league = "ultra")
    cup == "master" && return get_overrides("overall", league = "master")
    return JSON.parsefile(download(
        "https://raw.githubusercontent.com/pvpoke/pvpoke/master/src/data/" *
        "overrides/$(cup)/overall/$(get_cp_limit(league)).json"))
end

@memoize get_silph_usage(cup::String) =
    JSON.parsefile(download("https://silph.gg/api/cup/" * cup * "/stats/.json"))

"""
    get_gamemaster_mon_id(name)

Given a mon's PvPoke id, get its location within the gamemaster JSON file.
"""
function get_gamemaster_mon_id(name::String)
    for i = 1:length(gamemaster["pokemon"])
        gamemaster["pokemon"][i]["speciesId"] == name && return i
    end
end

"""
    get_gamemaster_move_id(name)

Given a move's PvPoke id, get its location within the gamemaster JSON file.
"""
function get_gamemaster_move_id(name::String)
    for i = 1:length(gamemaster["moves"])
        gamemaster["moves"][i]["moveId"] == name && return i
    end
end

"""
    get_rankings_mon_id(name; league = "great", cup = "all")

Given a mon's PvPoke id, find its location within a particular meta's rankings
"""
function get_rankings_mon_id(name::String;
    league::String = "great", cup::String = "all")
    rankings = get_rankings(cup, league = league)
    for i = 1:length(rankings)
        rankings[i]["speciesId"] == name && return i
    end
    return 0
end

const typings = ("normal", "fighting", "flying", "poison", "ground",
    "rock", "bug", "ghost", "steel", "fire", "water", "grass",
    "electric", "psychic", "ice", "dragon", "dark", "fairy", "none")

# CP multipliers from PvPoke
const cpm = Dict(
    1.0 => 0.0939999967813492, 1.5 => 0.135137432089339,
    2.0 => 0.166397869586945, 2.5 => 0.192650913155325,
    3.0 => 0.215732470154762, 3.5 => 0.236572651424822,
    4.0 => 0.255720049142838, 4.5 => 0.273530372106572,
    5.0 => 0.290249884128571, 5.5 => 0.306057381389863,
    6.0 => 0.321087598800659, 6.5 => 0.335445031996451,
    7.0 => 0.349212676286697, 7.5 => 0.362457736609939,
    8.0 => 0.375235587358475, 8.5 => 0.387592407713878,
    9.0 => 0.399567276239395, 9.5 => 0.4111935532161,
    10.0 => 0.422500014305115, 10.5 => 0.432926420512509,
    11.0 => 0.443107545375824, 11.5 => 0.453059948165049,
    12.0 => 0.46279838681221, 12.5 => 0.472336085311278,
    13.0 => 0.481684952974319, 13.5 => 0.490855807179549,
    14.0 => 0.499858438968658, 14.5 => 0.5087017489616,
    15.0 => 0.517393946647644, 15.5 => 0.525942516110322,
    16.0 => 0.534354329109192, 16.5 => 0.542635753803599,
    17.0 => 0.550792694091797, 17.5 => 0.558830584490385,
    18.0 => 0.566754519939423, 18.5 => 0.57456912814537,
    19.0 => 0.582278907299042, 19.5 => 0.589887907888945,
    20.0 => 0.597400009632111, 20.5 => 0.604823648665171,
    21.0 => 0.61215728521347, 21.5 => 0.619404107958234,
    22.0 => 0.626567125320435, 22.5 => 0.633649178748576,
    23.0 => 0.6406529545784, 23.5 => 0.647580971386554,
    24.0 => 0.654435634613037, 24.5 => 0.661219265805859,
    25.0 => 0.667934000492095, 25.5 => 0.674581885647492,
    26.0 => 0.681164920330048, 26.5 => 0.687684901255373,
    27.0 => 0.694143652915955, 27.5 => 0.700542901033063,
    28.0 => 0.706884205341339, 28.5 => 0.713169074873823,
    29.0 => 0.719399094581604, 29.5 => 0.725575586915154,
    30.0 => 0.731700003147125, 30.5 => 0.734741038550429,
    31.0 => 0.737769484519958, 31.5 => 0.740785579737136,
    32.0 => 0.743789434432983, 32.5 => 0.746781197247765,
    33.0 => 0.749761044979095, 33.5 => 0.752729099732281,
    34.0 => 0.75568550825119, 34.5 => 0.758630370209851,
    35.0 => 0.761563837528229, 35.5 => 0.76448604959218,
    36.0 => 0.767397165298462, 36.5 => 0.770297293677362,
    37.0 => 0.773186504840851, 37.5 => 0.776064947064992,
    38.0 => 0.778932750225067, 38.5 => 0.781790050767666,
    39.0 => 0.784636974334717, 39.5 => 0.787473608513275,
    40.0 => 0.790300011634827, 40.5 => 0.792803968023538,
    41.0 => 0.795300006866455, 41.5 => 0.797803898371622,
    42.0 => 0.800300002098083, 42.5 => 0.802803871877596,
    43.0 => 0.805299997329711, 43.5 => 0.807803850847053,
    44.0 => 0.81029999256134, 44.5 => 0.812803835179168,
    45.0 => 0.815299987792968, 45.5 => 0.817803806620319,
    46.0 => 0.820299983024597, 46.5 => 0.822803778631297,
    47.0 => 0.825299978256225, 47.5 => 0.827803750922782,
    48.0 => 0.830299973487854, 48.5 => 0.832803753381377,
    49.0 => 0.835300028324127, 49.5 => 0.837803755931569,
    50.0 => 0.840300023555755
)

# Color scheme to match PvPoke
const colors = [RGBA(153/255, 159/255, 161/255, 1.0),
                RGBA(213/255,  63/255,  91/255, 1.0),
                RGBA(148/255, 171/255, 225/255, 1.0),
                RGBA(193/255,  98/255, 212/255, 1.0),
                RGBA(212/255, 141/255,  91/255, 1.0),
                RGBA(208/255, 196/255, 142/255, 1.0),
                RGBA(158/255, 195/255,  49/255, 1.0),
                RGBA( 89/255, 107/255, 181/255, 1.0),
                RGBA( 82/255, 142/255, 160/255, 1.0),
                RGBA(254/255, 163/255,  84/255, 1.0),
                RGBA( 86/255, 158/255, 222/255, 1.0),
                RGBA( 94/255, 189/255,  91/255, 1.0),
                RGBA(246/255, 215/255,  75/255, 1.0),
                RGBA(245/255, 126/255, 121/255, 1.0),
                RGBA(120/255, 212/255, 192/255, 1.0),
                RGBA( 14/255, 104/255, 184/255, 1.0),
                RGBA( 86/255,  86/255,  99/255, 1.0),
                RGBA(240/255, 152/255, 228/255, 1.0)]

# Purple-ish shield color to match game
const shieldColor = RGBA(235/255,13/255,199/255, 1.0)

"""
    silph_to_pvpoke(name)

Given the name of a pokemon from the Silph API, return the pvpoke id

# Examples
```jldoctest
julia> silph_to_pvpoke("Mr. Mime")
mr_mime
"""
function silph_to_pvpoke(name::String)
    name = lowercase(name)
    name = replace(name, "-" => "_")
    name = replace(name, "_normal" => "")
    name = replace(name, "mr. " => "mr_")
    name = replace(name, "alola" => "alolan")
    name = replace(name, "galar" => "galarian")
    name = replace(name, "sunshine" => "sunny")
    name = replace(name, "basculin_red" => "basculin")
    name = replace(name, "basculin_blue" => "basculin")
    name = replace(name, "darmanitan" => "darmanitan_standard")
    name = replace(name, "gastrodon_blue" => "gastrodon_west_sea")
    name = replace(name, "gastrodon_pink" => "gastrodon_east_sea")
    name = replace(name, "sirfetch'd" => "sirfetchd")
    return name
end


# common args for all hdmr tests
testargs = Dict(:verbose=>false,:showwarnings=>true)

@testset "hdmr" begin

hdmrmodels = [
    @model(h ~ x + z + cat + y, c ~ x + z + cat + y)
    @model(h ~ x + z + cat + y, c ~ z + cat + y)
    @model(h ~ x + z + cat, c ~ x + z + cat + y)
    @model(h ~ x + z + cat, c ~ z + cat + y)
    ]

@testset "$M" for M in (InclusionRepetition,Hurdle)

includel = (M == InclusionRepetition)

@testset "$(replace(string(f),"\n" => " "))" for f in hdmrmodels

# # parse and merge rhs terms
trmszero = HurdleDMR.getrhsterms(f, :h)
trmspos = HurdleDMR.getrhsterms(f, :c)
trms, inzero, inpos = HurdleDMR.mergerhsterms(trmszero,trmspos)

# create model matrix
covarsb, countsb, as = modelcols(trms, covarsdf, counts; model=M)

# inzero and inpos may be different in mm with factor variables
inzero, inpos = HurdleDMR.mapins(inzero, inpos, as)
pzero = length(inzero)
ppos = length(inpos)

# resolve projdirs
projdir = HurdleDMR.ixprojdir(as, :y)
dirpos = something(findfirst(isequal(projdir), inpos), 0)
dirzero = something(findfirst(isequal(projdir), inzero), 0)

# hurdle dmr parallel local cluster
hdmrcoefs = hdmr(covars, counts, M; inzero=inzero, inpos=inpos, parallel=true, testargs...)
coefsHppos, coefsHpzero = coef(hdmrcoefs)
@test size(coefsHppos) == (ppos+1, d)
@test size(coefsHpzero) == (pzero+1, d)
@test_throws MethodError coef(hdmrcoefs, AllSeg())

# test Int matrix for counts
hdmrcoefsint = hdmr(covars, countsint, M; inzero=inzero, inpos=inpos, parallel=true, testargs...)
coefsHpposint, coefsHpzeroint = coef(hdmrcoefsint)
@test coefsHpposint == coefsHppos
@test coefsHpzeroint == coefsHpzero

# test standardize=false
hdmrcoefsunstd = hdmr(covars, counts, M; standardize=false, inzero=inzero, inpos=inpos, parallel=true, testargs...)
coefsHpposunstd, coefsHpzerounstd = coef(hdmrcoefsunstd)
@test coefsHpposunstd != coefsHppos
@test coefsHpzerounstd != coefsHpzero
hdmrcoefsunstd2 = hdmr(covars, counts, M; standardize=false, inzero=inzero, inpos=inpos, local_cluster=false, parallel=true, testargs...)
coefsHpposunstd2, coefsHpzerounstd2 = coef(hdmrcoefsunstd2)
@test coefsHpposunstd ≈ coefsHpposunstd2
@test coefsHpzerounstd ≈ coefsHpzerounstd2

hdmrcoefsb = fit(HDMRCoefs{M}, covars, counts; inzero=inzero, inpos=inpos, parallel=true, testargs...)
@test coef(hdmrcoefsb)[1] == coefsHppos
@test coef(hdmrcoefsb)[2] == coefsHpzero
@test n == nobs(hdmrcoefs)
@test d == ncategories(hdmrcoefs)
@test ppos == ncovarspos(hdmrcoefs)
@test pzero == ncovarszero(hdmrcoefs)

# select=MinBIC()
hdmrcoefsb = fit(HDMRCoefs{M}, covars, counts; inzero=inzero, inpos=inpos, select=MinBIC(), testargs...)
@test coef(hdmrcoefsb)[1] != coefsHppos
# @test coef(hdmrcoefsb)[2] != coefsHpzero
hdmrb = fit(HDMR{M}, covars, counts; inzero=inzero, inpos=inpos, select=MinBIC(), testargs...)
@test coef(hdmrb)[1] == coef(hdmrcoefsb)[1]
@test coef(hdmrb)[2] == coef(hdmrcoefsb)[2]

# hurdle dmr parallel remote cluster
hdmrcoefs2 = hdmr(covars, counts, M; inzero=inzero, inpos=inpos, parallel=true, local_cluster=false, testargs...)
coefsHppos2, coefsHpzero2 = coef(hdmrcoefs2)
@test coefsHppos ≈ coefsHppos2
@test coefsHpzero ≈ coefsHpzero2
hdmrcoefs2 = fit(HDMRCoefs{M}, covars, counts; inzero=inzero, inpos=inpos, parallel=true, local_cluster=false, testargs...)
@test coef(hdmrcoefs2)[1] ≈ coefsHppos2
@test coef(hdmrcoefs2)[2] ≈ coefsHpzero2
@test_throws ErrorException predict(hdmrcoefs2,newcovars)

hdmrpaths3 = fit(HDMRPaths{M}, covars, counts; inzero=inzero, inpos=inpos, parallel=true, testargs...)
coefsHppos3, coefsHpzero3 = coef(hdmrpaths3)
@test coefsHppos3 ≈ coefsHppos
@test coefsHpzero3 ≈ coefsHpzero
@test n == nobs(hdmrpaths3)
@test d == ncategories(hdmrpaths3)
@test ppos == ncovarspos(hdmrpaths3)
@test pzero == ncovarszero(hdmrpaths3)
@test size(hdmrpaths3.nlpaths,1) == d
η = predict(hdmrpaths3,newcovars)
@test sum(η, dims=2) ≈ ones(size(η, 1))
coefsallpos, coefsallzero = coef(hdmrpaths3, AllSeg())
@test size(coefsallpos,1) > 1
@test size(coefsallpos,2) == ppos+1
@test size(coefsallpos,3) == d
@test size(coefsallzero,1) > 1
@test size(coefsallzero,2) == pzero+1
@test size(coefsallzero,3) == d

# hurdle dmr serial
hdmrcoefs3 = hdmr(covars, counts, M; inzero=inzero, inpos=inpos, parallel=false, testargs...)
coefsHspos, coefsHszero = coef(hdmrcoefs3)
@test coefsHppos ≈ coefsHspos
@test coefsHpzero ≈ coefsHszero
hdmrcoefs3 = fit(HDMRCoefs{M}, covars, counts; inzero=inzero, inpos=inpos, parallel=false, testargs...)
@test coef(hdmrcoefs3)[1] ≈ coefsHspos
@test coef(hdmrcoefs3)[2] ≈ coefsHszero

# serial paths
hdmrcoefs4 = hdmrpaths(covars, counts, M; inzero=inzero, inpos=inpos, parallel=false, testargs...)
coefsHspos4, coefsHszero4 = coef(hdmrcoefs4)
@test coefsHspos == coefsHspos4
@test coefsHszero == coefsHszero4

# using a dataframe and formula
hdmrcoefsdf = fit(HDMRCoefs{M}, f, covarsdf, counts; inzero=inzero, inpos=inpos, parallel=true, testargs...)
@test coef(hdmrcoefsdf)[1] == coefsHppos
@test coef(hdmrcoefsdf)[2] ≈ coefsHpzero
@test_throws ErrorException predict(hdmrcoefsdf,newcovars)

hdmrpathsdf = fit(HDMRPaths{M}, f, covarsdf, counts; inzero=inzero, inpos=inpos, parallel=true, testargs...)
@test coef(hdmrpathsdf)[1] == coefsHppos3
@test coef(hdmrpathsdf)[2] ≈ coefsHpzero3
@test predict(hdmrpathsdf,newcovars) ≈ η

if M == InclusionRepetition
    srprojcounts = counts - posindic(counts)
else
    srprojcounts = counts
end

zHpos = srproj(coefsHppos, srprojcounts)
@test size(zHpos) == (n,ppos+1)

zHzero = srproj(coefsHpzero, posindic(counts))
@test size(zHzero) == (n,pzero+1)

Z1 = srproj(coefsHppos, coefsHpzero, counts, dirpos, dirzero;
    includel=includel, intercept=true)

if dirpos > 0
    z1pos = srproj(coefsHppos, srprojcounts, dirpos)
    @test z1pos ≈ zHpos[:,[dirpos,ppos+1]]
    @test hascol(Z1,z1pos[:,1])
    @test !includel || hascol(Z1,z1pos[:,2])
end

if dirzero > 0
    z1zero = srproj(coefsHpzero, posindic(counts), dirzero)
    @test z1zero ≈ zHzero[:,[dirzero,pzero+1]]
    @test hascol(Z1,z1zero[:,1])
    @test !includel || hascol(Z1,z1zero[:,2])
end

Z1b = srproj(hdmrcoefs, counts, dirpos, dirzero; intercept=true)
@test Z1 == Z1b

@test_throws ErrorException srproj(coefsHppos, coefsHpzero, counts, 0, 0; intercept=true)

X1, X1_nocounts, inz = srprojX(coefsHppos,coefsHpzero,counts,covars,projdir;
    inzero=inzero, inpos=inpos, includel=includel, includem=true)
ix = filter!(x->x!=projdir,collect(1:p))
@test X1_nocounts ≈ [ones(n) covars[:,ix]] rtol=1e-8
@test X1 ≈ [X1_nocounts Z1] rtol=1e-8
X1b, X1_nocountsb, inzb = srprojX(hdmrcoefs,counts,covars,projdir;
    inzero=inzero, inpos=inpos, includem=true)
@test X1 ≈ X1b rtol=1e-8
@test X1_nocountsb ≈ X1_nocountsb rtol=1e-8
@test inz == inzb

X2, X2_nocounts, inz = srprojX(coefsHppos,coefsHpzero,counts,covars,projdir;
    inzero=inzero, inpos=inpos, includel=includel, includem=false)
@test X2_nocounts ≈ X1_nocounts rtol=1e-8
if includel && dirzero > 0
    @test X2 ≈ [X1[:,1:end-2] X1[:,end]] rtol=1e-8
else
    @test X2 ≈ X1[:,1:end-1] rtol=1e-8
end
X2b, X2_nocountsb, inzb = srprojX(hdmrcoefs,counts,covars,projdir;
    inzero=inzero, inpos=inpos, includem=false)
@test X2 ≈ X2b rtol=1e-8
@test X2_nocountsb ≈ X2_nocountsb rtol=1e-8
@test inz == inzb

X3, X3_nocounts, inz3 = @test_logs (:info, r"rank\(X\)") srprojX(coefsHppos,coefsHpzero,zero(counts),covars,projdir;
    inzero=inzero, inpos=inpos, includel=includel, includem=true)
@test X3_nocounts == [ones(n) covars[:,setdiff(1:p,[projdir])]]
@test inz3 == [2]

X3, X3_nocounts, inz3b = @test_logs (:info, "includezpos == false. dropping zpos.") srprojX(coefsHppos,coefsHpzero,zero(counts),covars,projdir;
    inzero=inzero, inpos=inpos, includel=includel, includem=true, inz=inz3)
@test X3_nocounts == [ones(n) covars[:,setdiff(1:p,[projdir])]]
@test inz3 == inz3b

# HIR
hir = fit(CIR{HDMR{M}, LinearModel},covars,counts,projdir;
    inzero=inzero, inpos=inpos, nocounts=true, testargs...)
@test coefbwd(hir)[1] ≈ coef(hdmrcoefs)[1]
@test coefbwd(hir)[2] ≈ coef(hdmrcoefs)[2]

hirglm = fit(CIR{HDMR{M}, GeneralizedLinearModel},covars,counts,projdir,Gamma();
    inzero=inzero, inpos=inpos, nocounts=true, testargs...)
@test coefbwd(hirglm)[1] ≈ coef(hdmrcoefs)[1]
@test coefbwd(hirglm)[2] ≈ coef(hdmrcoefs)[2]
@test !(coeffwd(hirglm)[1] ≈ coeffwd(hir)[1])
@test !(coeffwd(hirglm)[2] ≈ coeffwd(hir)[2])

hirdf = fit(CIR{HDMR{M}, LinearModel},f,covarsdf,counts,:y;
    inzero=inzero, inpos=inpos, nocounts=true, testargs...)
@test coefbwd(hirdf)[1] ≈ coef(hdmrcoefs)[1]
@test coefbwd(hirdf)[2] ≈ coef(hdmrcoefs)[2]
@test coeffwd(hirdf) ≈ coeffwd(hir)
hirglmdf = fit(CIR{HDMR{M}, GeneralizedLinearModel},f,covarsdf,counts,:y,Gamma();
    inzero=inzero, inpos=inpos, nocounts=true, testargs...)
@test coefbwd(hirglmdf)[1] ≈ coef(hdmrcoefs)[1]
@test coefbwd(hirglmdf)[2] ≈ coef(hdmrcoefs)[2]
@test !(coeffwd(hirglmdf)[2] ≈ coeffwd(hirdf)[2])
@test !(coeffwd(hirglmdf)[1] ≈ coeffwd(hirdf)[1])

# select=MinBIC()
hirdfb = fit(CIR{HDMR{M}, LinearModel},f,covarsdf,counts,:y;
    inzero=inzero, inpos=inpos, nocounts=true, select=MinBIC(), testargs...)
@test coefbwd(hirdf)[1] != coefbwd(hirdfb)[1]
# @test coefbwd(hirdf)[2] != coefbwd(hirdfb)[2]
@test coeffwd(hirdf)[1] != coeffwd(hirdfb)[1]
@test coeffwd(hirdf)[2] != coeffwd(hirdfb)[2]

zlm = lm(hcat(ones(n,1),Z1,covars[:,1:4]),covars[:,projdir])
@test r2(zlm) ≈ r2(hir)
@test adjr2(zlm) ≈ adjr2(hir)
@test predict(zlm,hcat(ones(10,1),Z1[1:10,:],covars[1:10,1:4])) ≈ predict(hir,covars[1:10,:],counts[1:10,:])

zlmnocounts = lm(hcat(ones(n,1),covars[:,1:4]),covars[:,projdir])
@test r2(zlmnocounts) ≈ r2(hir; nocounts=true)
@test adjr2(zlmnocounts) ≈ adjr2(hir; nocounts=true)
@test predict(zlmnocounts,hcat(ones(10,1),covars[1:10,1:4])) ≈ predict(hir,covars[1:10,:],counts[1:10,:]; nocounts=true)

end

@testset "HDMR{$M} degenerate cases" begin

@info("Testing hdmr degenerate cases. The 2 following warnings by workers are expected ...")

# column j is never zero, so hj=1 for all observations
zcounts = deepcopy(counts)
zcounts[:,2] = zeros(n)
zcounts[:,3] = ones(n)

# make sure we are not adding all zero obseravtions
m = sum(zcounts, dims=2)
@test sum(m .== 0) == 0

# hurdle dmr parallel local cluster
hdmrcoefs = fit(HDMR{M}, covars, zcounts; parallel=true, testargs...)
coefsHppos, coefsHpzero = coef(hdmrcoefs)
@test size(coefsHppos) == (p+1, d)
@test size(coefsHpzero) == (p+1, d)
@test coefsHppos[:,2] == zeros(p+1)
@test coefsHpzero[:,2] == zeros(p+1)
@test coefsHppos[:,3] == zeros(p+1)
@test coefsHpzero[:,3] == zeros(p+1)
@test HurdleDMR.includelinX(hdmrcoefs) == includel

if M == Hurdle
    η3 = 0.36
    η4 = 0.6
else
    η3 = 0.32
    η4 = 0.64
end

# hurdle dmr parallel remote cluster
hdmrcoefs2 = fit(HDMRPaths{M},covars, zcounts; parallel=true, local_cluster=false, testargs...)
coefsHppos2, coefsHpzero2 = coef(hdmrcoefs)
@test coefsHppos ≈ coefsHppos2
@test coefsHpzero ≈ coefsHpzero2
η = predict(hdmrcoefs2,newcovars)
@test sum(η, dims=2) ≈ ones(size(newcovars, 1))
@test η[:,2] == zeros(size(newcovars,1))
@test η[:,3] ≈ ones(size(newcovars,1))*η3 rtol=0.05
@test η[:,4] ≈ ones(size(newcovars,1))*η4 rtol=0.06
@test HurdleDMR.includelinX(hdmrcoefs2) == includel

# hurdle dmr serial paths
rx = Regex("fit\\($M...\\) failed for countsj")
hdmrcoefs3 = @test_logs (:warn, rx) fit(HDMRPaths{M},covars, zcounts; parallel=false, testargs...)
coefsHppos3, coefsHpzero3 = coef(hdmrcoefs3)
@test coefsHppos ≈ coefsHppos3
@test coefsHpzero ≈ coefsHpzero3
η = predict(hdmrcoefs3,newcovars)
@test sum(η, dims=2) ≈ ones(size(newcovars, 1))
@test η[:,2] == zeros(size(newcovars,1))
@test η[:,3] ≈ ones(size(newcovars,1))*η3 rtol=0.05
@test η[:,4] ≈ ones(size(newcovars,1))*η4 rtol=0.06
@test HurdleDMR.includelinX(hdmrcoefs3) == includel

# hurdle dmr serial coefs
hdmrcoefs4 = @test_logs (:warn, r"failed on count dimension 2") fit(HDMR{M},covars, zcounts; parallel=false, testargs...)
coefsHppos4, coefsHpzero4 = coef(hdmrcoefs4)
@test size(coefsHppos4) == (p+1, d)
@test size(coefsHpzero4) == (p+1, d)
@test coefsHppos4[:,2] == zeros(p+1)
@test coefsHpzero4[:,2] == zeros(p+1)
@test coefsHppos4[:,3] == zeros(p+1)
@test coefsHpzero4[:,3] == zeros(p+1)
@test HurdleDMR.includelinX(hdmrcoefs4) == includel

if M == InclusionRepetition
    @info("Testing hdmr degenerate cases. The 12 following warnings by workers are expected ...")
end

zcounts = Matrix(deepcopy(counts))
Random.seed!(13)
for I = eachindex(zcounts)
    if iszero(zcounts[I])
        zcounts[I] = rand(1:10)
    end
end

# make sure we are not adding all zero obseravtions
m = sum(zcounts, dims=2)
@test sum(m .== 0) == 0
if M == InclusionRepetition
    l = sum(posindic(zcounts), dims=2)
    @test sum((m .- l) .== 0) == 0
    @test sum((d .- l) .== 0) == n
end

# hurdle dmr parallel local cluster
hdmrcoefs = fit(HDMR{M}, covars, zcounts; parallel=true, testargs...)
coefsHppos, coefsHpzero = coef(hdmrcoefs)
@test size(coefsHppos) == (p+1, d)
@test size(coefsHpzero) == (p+1, d)
if M == Hurdle
	@test coefsHppos[:,2] != zeros(p+1)
	@test coefsHppos[:,3] != zeros(p+1)
else
	@test coefsHpzero[:,2] == zeros(p+1)
	@test coefsHpzero[:,3] == zeros(p+1)
end
@test HurdleDMR.includelinX(hdmrcoefs) == includel

# hurdle dmr parallel remote cluster
hdmrcoefs2 = fit(HDMRPaths{M}, covars, zcounts; parallel=true, local_cluster=false, testargs...)
coefsHppos2, coefsHpzero2 = coef(hdmrcoefs2)
@test coefsHppos ≈ coefsHppos2
@test coefsHpzero ≈ coefsHpzero2
@test HurdleDMR.includelinX(hdmrcoefs2) == includel

# just checking the fit(HDMRPaths...) ignore local_cluster
hdmrcoefs3 = fit(HDMRPaths{M},covars, zcounts; parallel=true, local_cluster=true, testargs...)
coefsHppos3, coefsHpzero3 = coef(hdmrcoefs3)
@test coefsHppos2 == coefsHppos3
@test coefsHpzero2 == coefsHpzero3
@test HurdleDMR.includelinX(hdmrcoefs3) == includel

# hurdle dmr parallel remote cluster
if M == Hurdle
    hdmrcoefs4 = fit(HDMRPaths{M},covars, zcounts; parallel=false, testargs...)
else
    rx = Regex("fit\\($M...\\) failed for countsj")
    warnings = [(:warn, rx) for i=1:4]
    hdmrcoefs4 = @test_logs(warnings..., fit(HDMRPaths{M},covars, zcounts; parallel=false, testargs...))
end
coefsHppos4, coefsHpzero4 = coef(hdmrcoefs4)
@test coefsHppos ≈ coefsHppos4
@test coefsHpzero ≈ coefsHpzero4
@test HurdleDMR.includelinX(hdmrcoefs4) == includel

end

end

end


module ArbitrarySampler
export get_sample, @spec, generate
using MLStyle
import InteractiveUtils: subtypes
# for typed test generator

struct Iterator{T}
    f :: Function
end

function next_item!(iter::Iterator{T}):: Union{Nothing, Some{T}} where T
    e  = iter.f() :: Union{Nothing, Some{<:T}}
    e === nothing && return nothing
    e.value isa T && return e
    Some{T}(e.value)
end

function take_all!(iter::Iterator{T})::Vector{T} where T
    xs = T[]
    while (a = next_item!(iter); a !== nothing)
        push!(xs, a.value)
    end
    return xs
end

function take_n!(iter::Iterator{T}, n :: Int)::Vector{T} where T
    xs = T[]
    while (a = next_item!(iter); a !== nothing && length(xs) < n)
        push!(xs, a.value)
    end
    return xs
end

abstract type SeqElt{T}  end
abstract type Spec{T} end

→(A, B) = Function

@data SeqElt{T} begin
    One(Spec{T})
    Ellipsis(Spec{Vector{T}})
end

@data Spec{T} begin
    S_of{T}           :: (() → T)           => Spec{T}
    S_modif{T}        :: ((T → T), Spec{T}) => Spec{T}
    S_rep_of{T}       :: Spec{T}            => Spec{Iterator{T}}
    S_app{F}          :: (f, args)          => Spec{F}
    S_tp{Tp <: Tuple} :: tup                => Spec{Tp} where
      { Tp = Tuple{map(get_spec_type, tup)...} }

    S_guard{T}  :: (filter::(T → Bool), spec::Spec{T}) => Spec{T}
    S_map{A, B} :: (f::(A → B), spec::Spec{A})         => Spec{B}
    S_or{A, B}  :: (Spec{A}, Spec{B})                  => Spec{Union{A, B}}

    S_vec_of{T} :: (xs::Vector{SeqElt})                => Spec{Vector{T}} where
      { T =  isempty(xs) ? Any : reduce(typejoin, map(get_elt_type, xs)) }
end

S_tp(hd, args...) = S_tp((hd, args...))
get_spec_type(::Spec{T}) where T = T
get_elt_type(::SeqElt{T}) where T = T

function int_mod(n::Int, spec::Spec{Int})
    function modify(orig::Int)::Int
        orig % n
    end
    S_modif(modify, spec)
end

@nospecialize
function _take_method_arg_type(t)
    while t isa UnionAll
        t = t.body
    end
    t = t.parameters[2].parameters[1]
    while t isa TypeVar
        t = t.ub
    end
    return t
end
@specialize

get_subtypes(t) = get_subtypes(Val(Base.get_world_counter()), t)
@generated function get_subtypes(::Val{WorldAge}, ::Type{T}) where {WorldAge, T}
    subtypes(T)
end

get_registered_types() = get_registered_types(Val(Base.get_world_counter()))
@generated function get_registered_types(::Val{WorldAge}) where WorldAge
    meths = methods(get_sample)
    Any[_take_method_arg_type(m.sig) for m in meths.ms]
end

get_registered_datatypes() = get_registered_datatypes(Val(Base.get_world_counter()))
@generated function get_registered_datatypes(::Val{WorldAge}) where WorldAge
    regs = get_registered_types(Val(WorldAge))
    DataType[t for t in regs if t isa DataType]
end

function _get_all_concrete!(sup, cts::Set)
    if isabstracttype(sup)
        if Core.Compiler.unwrap_unionall(sup).name.module === Core
            push!(cts, sup)
        end
        for sub in get_subtypes(sup)
            _get_all_concrete!(sub, cts)
        end
    else
        if Core.Compiler.unwrap_unionall(sup).name.module === Core
            push!(cts, sup)
        end
    end
end
const Cores = Set{Any}([BigFloat, BigInt])
_get_all_concrete!(Real, Cores)

rand′(x) = rand(x)
rand′(::Type{BigInt}) = BigInt(rand(Int))
rand′(::Type{<:Irrational}) = pi
rand′(::Type{<:Irrational}) = pi
rand′(::Type{Complex}) =
    let T = rand([Int, Float64])
        rand(Complex{T})
    end
rand′(::Type{Rational}) =
    let T = rand([Integer])
        get_sample(T) // (abs(get_sample(T)) + 1)
    end

function _get_sample_generic(x)
    regs = get_registered_types()
    if isabstracttype(x)
        possible = [t for t in get_subtypes(x) if t in Cores || t in regs]
        isempty(possible) && return nothing
        return get_sample(rand(possible))
    else
        (x in Cores || x in regs) && return rand′(x)
        return nothing
    end
end

get_sample(::Type{Complex}) = rand′(Complex)
function get_sample(::Type{T}) where T <: Real
    ret = _get_sample_generic(T)
    ret === nothing && error("cannot find random generator for $T.")
    return ret
end
get_sample(::Type{String}) =
    String([rand('0':'z')  for i in 1:get_sample(Int) % 100])
get_sample(::Type{Symbol}) =
    Symbol(String([rand('0':'z')  for i in 1:get_sample(Int) % 7]))

function get_sample(::Type{DataType})
    dts = get_registered_datatypes()
    dts[rand(1:length(dts))]
end

function get_sample(::Type{Any})
    get_sample(rand(get_registered_types()))
end

function generate(spec::Spec{T})::T where T
    @switch spec begin
        @case S_of(gen)
            return gen()

        @case S_or(a, b)
            return rand(Bool) === true ?
                generate(a) :
                generate(b)

        @case S_modif(modif, spec)
            return modif(generate(spec))

        @case S_rep_of{A}(unit) where A
            # infinite generator
            return Iterator{A}() do
                return Some(generate(unit))
            end

        @case S_app(f, args)
            return f(generate(args)...)

        @case S_tp(tup)
            tp = Any[]
            for e in tup
                push!(tp, generate(e))
            end
            return Tuple(tp)

        @case S_guard(filt, spec)
            a = nothing
            while (a = generate(spec); !filt(a))
            end
            return a

        @case S_map(f, spec)
            a = generate(spec)
            return f(a)

        @case S_vec_of{G}(elts) where G
            ret = G[]
            for each in elts
                @switch each begin
                    @case One(spec)
                        push!(ret, generate(spec))
                        continue
                    @case Ellipsis(spec)
                        append!(ret, generate(spec))
                end
            end
            return ret
    end
end

const any_spec = S_of{Any}() do
    get_sample(Any)
end

const sym_spec = S_of{Symbol}() do
    get_sample(Symbol)
end

function gen_spec(s, __module__::Module)
    S_of{typeof(s)}() do
        s
    end
end

function gen_spec(s::QuoteNode, __module__::Module)
    s.value isa Symbol && return let s=s.value
            S_of{Symbol}() do
                s
            end
    end

    gen_spec(s.value)
end

function gen_spec(s::Symbol, __module__::Module)
    if s === :_
        return any_spec
    end
    error("Only '_' is a supported Symbol for spec syntax so far.")
end

function qt2ex(ex::Any)
    if ex isa Expr
        Meta.isexpr(ex, :$) && return ex.args[1]
        ret = Expr(:call)
        push!(ret.args, Expr, QuoteNode(ex.head))
        for each in ex.args
            if each isa LineNumberNode
                continue
            end
            push!(ret.args, qt2ex(each))
        end
        ret
    elseif ex isa Symbol
        QuoteNode(ex)
    else
        ex
    end
end

function gen_spec(ex::Expr, __module__::Module)
    rec(ex) = gen_spec(ex, __module__)
    @switch ex begin
        @case Expr(:quote, a)
            return rec(qt2ex(a))
        @case Expr(:(::), val, ty)
            T = __module__.eval(ty)
            val = T(val)
            return S_of{T}() do
                val
            end
        @case Expr(:(::), ty)
            T = __module__.eval(ty)
            return S_of{T}() do
                get_sample(T)
            end
        @case Expr(:$, a)
            spec = __module__.eval(a)
            spec isa Spec || error("invalid spec")
            return spec
        @case Expr(:||, a, b)
            return S_or(rec(a), rec(b))
        @case Expr(:tuple, args...)
            return S_tp(Tuple(rec(arg) for arg in args))
        @case Expr(:vect, args...)
            specs = SeqElt[]
            for arg in args
                @switch arg begin
                    @case :($a...)
                        push!(specs, Ellipsis(rec(a)))
                        continue
                    @case _
                        push!(specs, One(rec(arg)))
                end
            end
            return S_vec_of(specs)
        @case Expr(:comparison, args...)
            isodd(length(args)) || error("invalid spec: $ex")
            args[2] === :isa || error("invalid spec: $ex, require 'a isa b isa c'")
            return foldl(3:2:length(args), init=rec(args[1])) do last, i
                    arg = args[i]
                    args[i-1] === :isa || error("invalid spec: $ex, require 'a isa b isa c'")
                    @match arg begin
                        :($f.?) => S_guard(__module__.eval(f), last)
                        _ => S_modif(__module__.eval(arg), last)
                    end
            end
        @case :($a isa $f)
            return @match f begin
                :($f.?) => S_guard(__module__.eval(f), rec(a))
                _ => S_modif(__module__.eval(f), rec(a))
            end

        @case :($a.?($f))
            return S_guard(__module__.eval(f), rec(a))

        @case :($a{$min_v, $max_v}) || :($a{$min_v}) && let max_v = min_v end
            min_v = __module__.eval(min_v)::Int
            max_v = __module__.eval(max_v)::Int
            spec = rec(a)
            T = get_spec_type(spec)
            rep_spec = S_rep_of(spec)
            if min_v === max_v
                return  S_map{Iterator{T}, Vector{T}}(rep_spec) do iter
                    take_n!(iter, min_v)
                end
            else
                return  S_map{Iterator{T}, Vector{T}}(rep_spec) do iter
                    take_n!(iter, rand(min_v:max_v))
                end
            end
        @case Expr(:call, f, args...)
            f = __module__.eval(f)
            return S_app{f}(f, rec(Expr(:vect, args...)))
        @case _
            error("not supported syntax: $ex")
    end
end

macro spec(ex)
    gen_spec(ex, __module__)
end

# mod′(n) = x -> x % n
# len_lt(n) = x -> length(x) < n && length(x) !== 0
# not_num(x) = !(x isa Number)
# struct S{T}
#     a :: T
#     v :: Int
# end

# spec0 = @spec (::Int) isa iseven.?
# spec1 = @spec [_, (2||3, _, 0)]
# spec2 = @spec S(::Symbol, ::Int isa mod′(12))
# spec3 = @spec 1 || ::DataType
# spec4 = @spec [1, (2, ::String isa len_lt(6).?), 3]
# spec5 = @spec (_, S(::Bool, 2), ::Int isa mod′(3))

# for (id, spec) in [0=>spec0, 1=>spec1, 2=>spec2, 3=>spec3, 4=>spec4, 5=>spec5]
#     for nth in 1:2
#         sample = generate(spec)
#         println("spec$id, ($nth)th sample: $sample")
#     end
# end
end


using DiffEqBase: DEProblem, remake
using Distributions: Normal, cquantile

dimension(prob::DEProblem) = length(prob.u0)

PhaseRelaxer(prob::LEProblem,
             ::LEProblem,
             ::LESolution,
             ::Terminator,
             integrator_options) =
    PhaseRelaxer(get_integrator(prob.phase_prob; integrator_options...),
                 prob.t_tran)

function step!(relaxer::PhaseRelaxer)
    step!(relaxer.integrator, relaxer.t_tran)
    relaxer.i += 1
end
Base.length(relaxer::PhaseRelaxer) = 1  # FIXME

function get_tangent_prob(prob;
                          x0 = prob.phase_prob.u0,
                          kwargs...)
    u0 = copy(prob.tangent_prob.u0)
    u0[:, 1] = x0
    return remake(prob.tangent_prob; u0=u0, kwargs...)
end

current_state(relaxer::Union{PhaseRelaxer, AbstractRenormalizer}) =
    current_state(relaxer.integrator)

function get_tangent_integrator(prob::LEProblem, relaxer;
                                integrator_options...)
    # Carry on simulated time to tangent integrator [*]:
    tspan = if relaxer.integrator isa ODEIntegrator
        (relaxer.integrator.t, Inf)
    else
        # TODO: Support this in discrete problem
        prob.phase_prob.tspan
    end
    x0 = current_state(relaxer)
    tangent_prob = get_tangent_prob(prob; x0=x0, tspan=tspan)
    return get_integrator(tangent_prob; integrator_options...)
end
# [*] Another way to propagate simulated time to tangent integrator is
# to use set_t!() after init().  However, init() calls phase_prob.f
# with phase_prob.tspan[1] which is less than relaxer.integrator.t.
# It causes problems in some rare cases with stateful problem
# definitions.
# See: [[./clv/clv_problem.jl::get_integrator]]

TangentRenormalizer(relaxer::PhaseRelaxer, prob::LEProblem, sol::LESolution,
                    tmnr::Terminator, integrator_options) =
    TangentRenormalizer(get_tangent_integrator(prob, relaxer;
                                               integrator_options...),
                        prob.t_attr, prob.t_renorm, sol, tmnr)

MLERenormalizer(relaxer::PhaseRelaxer, prob::LEProblem, sol::LESolution,
                tmnr::Terminator, integrator_options) =
    MLERenormalizer(get_tangent_integrator(prob, relaxer;
                                           integrator_options...),
                    prob.t_attr, prob.t_renorm, sol, tmnr)

Base.length(stage::AbstractRenormalizer) =
    ceil(Int, stage.t_attr / stage.t_renorm)

@inline function keepgoing!(stage::AbstractRenormalizer)
    u0 = current_state(stage)
    u0[:, 1] = stage.phase_state
    u0[:, 2:end] = stage.tangent_state
    set_u!(stage.integrator, u0)
    step!(stage.integrator, stage.t_renorm, true)
    assert_success(stage.integrator)
    # TODO: maybe use step!(solver.integrator, dt) aka "inexact step"
    u0 = current_state(stage)
    stage.phase_state[:] = @view u0[:, 1]
    stage.tangent_state[:] = @view u0[:, 2:end]
end

"""
S_n = ((n-1)/n) S_{n-1} + r_n / n
"""
@inline function lyap_add_R!(n, lyap, R)
    for i = 1:length(lyap)
        lyap[i] = lyap_add_r(n, lyap[i], R[i, i])
    end
end

@inline function lyap_add_r(n, lyap, r)
    return (n - 1) / n * lyap + log(r) / n
end

@inline function eye!(A)
    A .= 0
    for i in 1:min(size(A)...)
        @inbounds A[i, i] = 1
    end
    A
end

""" A = A * diag(sgn) """
@inline function A_mul_sign!(A, sgn)
    for i = 1:size(A)[2]
        if !sgn[i]
            A[:, i] *= -1
        end
    end
    A
end

""" A = diag(sgn) * A """
@inline function sign_mul_A!(sgn, A)
    for i = 1:size(A)[1]
        if !sgn[i]
            A[i, :] *= -1
        end
    end
    A
end

""" sgn = sign(diag(A))  """
@inline function sign_diag!(sgn, A)
    for i = 1:size(A)[1]
        sgn[i] = A[i, i] >= 0
    end
    sgn
end

"""
    step!(stage::AbstractRenormalizer)

Evolve the dynamics and then do an orthonormalization.
"""
function step!(stage::AbstractRenormalizer)
    keepgoing!(stage)
    post_evolve!(stage)
    record!(stage, Val{:OS})
    record!(stage, Val{:FTLE})
    stage.is_finished = should_terminate!(stage.tmnr, stage.sol)
end

is_finished(stage::AbstractRenormalizer) =
    stage.is_finished || (stage.i >= length(stage))

function post_evolve!(stage::TangentRenormalizer)
    dim_lyap = length(stage.inst_exponents)
    P = stage.tangent_state
    F = qrfact!(P)
    Q = stage.Q
    A_mul_B!(F[:Q], eye!(Q))  # Q = Matrix(F[:Q])[...]; but lesser allocation
    R = stage.R = F[:R]

    sign_R = stage.sign_R
    sign_diag!(sign_R, R)        # sign_R = diagm(sign(diag(R)))
    A_mul_sign!(Q, sign_R)       # Q = Q * sign_R
    sign_mul_A!(sign_R, R)       # R = signR * R

    stage.i += 1
    stage.tangent_state, stage.Q = Q, stage.tangent_state
    # At this point:
    # stage.tangent_state === Q == 𝑮ₙ₊ₖ
    # stage.Q is a mess, as qrfact! use it as a "buffer"
end

function post_evolve!(stage::MLERenormalizer)
    v = stage.tangent_state[:, 1]
    r = norm(v)
    n = (stage.i += 1)
    stage.exponent = lyap_add_r(n, stage.exponent, r)
    stage.inst_exponent = log(r) / t_chunk(stage)
    # TODO: don't re-calculate log(r)
    stage.tangent_state[:, 1] .= v ./ r
end

t_chunk(stage) = stage.t_renorm

function record!(stage::TangentRenormalizer{<: LESolRecOS}, ::Type{Val{:OS}})
    dim_lyap = length(stage.inst_exponents)
    R = stage.R
    dt = t_chunk(stage)
    @assert size(R) == (dim_lyap, dim_lyap)
    for i in 1:dim_lyap
        @inbounds stage.inst_exponents[i] = log(R[i, i]) / dt
    end
    OnlineStats.fit!(stage.sol.series, stage.inst_exponents)
end

function record!(stage::MLERenormalizer{<: LESolRecOS}, ::Type{Val{:OS}})
    # FIXME: Assuming stage.sol.main is VecMean and alike
    OnlineStats.fit!(stage.sol.series, [stage.inst_exponent])
end

function record!(stage::AbstractRenormalizer{<: LESolRecFTLE},
                 ::Type{Val{:FTLE}})
    stage.sol.ftle_history[stage.i] .= ftle(stage)
    stage.sol.num_orth = stage.i
end

current_result(stage::TangentRenormalizer) = stage.inst_exponents
current_result(stage::MLERenormalizer) = stage.inst_exponent

"""
    lyapunov_exponents(solver)

Get the result of Lyapunov exponents calculation stored in `solver`.
"""
lyapunov_exponents(solver::Union{LESolverRecOS,
                                 AbstractRenormalizer{<: LESolRecOS}}) =
    lyapunov_exponents(solver.sol)

lyapunov_exponents(sol::LESolRecOS) = mean(sol.main_stat)

@inline function lyapunov_exponents(stage::MLERenormalizer)
    [stage.exponent / t_chunk(stage)]
end

"""Get finite-time Lyapunov exponents (FTLE)"""
ftle(stage::TangentRenormalizer) = stage.inst_exponents
ftle(stage::MLERenormalizer) = [stage.inst_exponent]

ftle_history(solver::Union{LESolverRecFTLE,
                           AbstractRenormalizer{<: LESolRecFTLE}},
             args...) =
    ftle_history(solver.sol, args...)

ftle_history(sol::LESolRecFTLE) = @view sol.ftle_history[1:sol.num_orth]
ftle_history(sol::LESolRecFTLE, i) =
    [x[i] for x in sol.ftle_history[1:sol.num_orth]]

function ftle_history(::Type{A}, sol::LESolRecFTLE) where {A <: AbstractMatrix}
    dim_lyap = length(sol.ftle_history[1])
    num_orth = sol.num_orth
    arr = A(dim_lyap, num_orth)
    for (i, ftle) in enumerate(ftle_history(sol))
        arr[:, i] .= ftle
    end
    return arr
end

const Vecs = AbstractVector{<: AbstractVector}

exponents_history(sol) = exponents_history(ftle_history(sol))

function exponents_history(ftle_history::Vecs)
    t_attr = length(ftle_history)
    dim_lyap = length(ftle_history[1])

    m = VecMean(dim_lyap)
    s = OnlineStats.Series(m)
    le_hist = similar(ftle_history[1], (dim_lyap, t_attr))
    for i in 1:t_attr
        OnlineStats.fit!(s, ftle_history[i])
        le_hist[:, i] .= mean(m)
    end
    return le_hist
end

exponents_stat_history(sol::Union{LESolRecFTLE,
                                  AbstractRenormalizer{<: LESolRecFTLE},
                                  LESolverRecFTLE}) =
    exponents_stat_history(ftle_history(sol))

function exponents_stat_history(ftle_history::Vecs, coverageprob = 0.95)
    t_attr = length(ftle_history)
    dim_lyap = length(ftle_history[1])
    c = cquantile(Normal(), (1 - coverageprob) / 2)
    # Usually t_attr is big so using TDist does not change much here.

    v = VecVariance(dim_lyap)
    s = OnlineStats.Series(v)
    le_hist = similar(ftle_history[1], (dim_lyap, t_attr))
    ci_hist = similar(le_hist)

    OnlineStats.fit!(s, ftle_history[1])
    le_hist[:, 1] .= mean(v)
    ci_hist[:, 1] .= NaN
    for i in 2:t_attr
        OnlineStats.fit!(s, ftle_history[i])
        # c = cquantile(TDist(i - 1), (1 - coverageprob) / 2)
        le_hist[:, i] .= mean(v)
        ci_hist[:, i] .= std(v) / sqrt(i) * c
    end
    return le_hist, ci_hist
end
# MAYBE: use OnlineStats.Bootstrap
# http://joshday.github.io/OnlineStats.jl/latest/api.html#OnlineStats.Bootstrap


@testset "king" begin
    b = Bobby.set_board()

    kv = Bobby.gen_king_valid(Bobby.INT2UINT[61])
    @test length(kv) == 5
    kvw = [0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,
           0,0,0,1,1,1,0,0,
           0,0,0,1,0,1,0,0]
    @test all(Int.(Bobby.cvt_to_bitarray(kv)) .== kvw)

    akv = Bobby.gen_all_king_valid_moves()
    @test all(Int.(Bobby.cvt_to_bitarray(akv[Bobby.INT2UINT[61]])) .== kvw)

    @test test_fen("k7/8/8/8/8/8/8/4K3 w - - 0 1", 1, [5])
    @test test_fen("k7/8/8/8/8/8/8/4K2R w - - 0 1", 1, [14])
    @test test_fen("k7/8/8/8/8/8/8/4K2R w K - 0 1", 1, [15])
    @test test_fen("1k6/8/8/8/8/8/8/R3K3 w K - 0 1", 1, [15])
    @test test_fen("1k6/8/8/8/8/8/8/R3K3 w Q - 0 1", 1, [16])
    @test test_fen("1k6/8/8/8/8/3r4/8/R3K3 w Q - 0 1", 1, [13])
    @test test_fen("1k6/8/8/8/8/2r5/8/R3K3 w Q - 0 1", 1, [15])
    @test test_fen("1k6/8/8/8/8/5r2/8/4K2R w Q - 0 1", 1, [12])
    @test test_fen("r3k3/8/8/8/8/8/8/6K1 b - - 0 1", 1, [15])
    @test test_fen("r3k3/8/8/8/8/8/8/6K1 b q - 0 1", 1, [16])
    @test test_fen("r3k3/8/8/8/8/8/8/1R4K1 b q - 0 1", 1, [16])
    @test test_fen("r3k3/8/8/8/8/8/8/2R3K1 b q - 0 1", 1, [15])
    @test test_fen("r3k3/8/8/8/8/8/8/3R2K1 b q - 0 1", 1, [13])
    @test test_fen("r3k3/8/8/8/8/8/8/4R1K1 b q - 0 1", 1, [4])
    @test test_fen("k3r3/8/8/8/8/8/8/4K2R w K - 0 1", 1, [4])
    @test test_fen("k7/8/8/8/6p1/6Pp/6PP/r6K w - - 0 1", 1, [0])
end


#
# This file is part of module
#   ModiaMedia (ModiaMedia/src/ModiaMedia.jl)
#

### Data structures that are additionally available for media <: PartialPureSubstance -----------------------

const X_dummy = fill(0.0,0)

"""
    state = setState_pT(medium, p,T)

Generate a state object for medium `medium::PureSubstance` for
pressure `p` [Pa] and temperature `T` [K].
"""
setState_pT( m::PureSubstance,p,T) = setState_pTX(m,p,T,X_dummy)

"""
    setState_pT!(state, p,T)

Update the `state::ThermodynamicState` of a `PureSubstance` medium with
pressure `p` [Pa] and temperature `T` [K].
"""
setState_pT!(state::ThermodynamicState,p,T) = setState_pTX!(state,p,T,X_dummy)

"""
    state = setState_ph(medium, p,h)

Generate a state object for medium `medium::PureSubstance` for
pressure `p` [Pa] and specific enthalpy `h` [J/kg].
"""
setState_ph(m::PureSubstance,p,h) = setState_phX(m,p,h,X_dummy)

"""
    setState_ph!(state, p,h)

Update the `state::ThermodynamicState` of a `PureSubstance` medium with
pressure `p` [Pa] and specific enthalpy `h` [J/kg]].
"""
setState_ph!(state::ThermodynamicState,p,h) = setState_phX!(state,p,h,X_dummy)

"""
    state = setState_ps(medium, p,s)

Generate a state object for medium `medium::PureSubstance` for
pressure `p` [Pa] and specific entropy `s` [J/(kg*K)].
"""
setState_ps( m::PureSubstance,p,s) = setState_psX(m,p,s,X_dummy)

"""
    setState_ps!(state, p,s)

Update the `state::ThermodynamicState` of a `PureSubstance` medium with
pressure `p` [Pa] and specific entropy `s` [J/(kg*K)].
"""
setState_ps!(state::ThermodynamicState,p,s) = setState_psX!(state,p,s,X_dummy)

"""
    state = setState_dT(medium, d,T)

Generate a state object for medium `medium::PureSubstance` for
density `d` [kg/m^3] and temperature `T` [K].
"""
setState_dT( m::PureSubstance,d,T) = setState_dTX(m,d,T,X_dummy)

"""
    setState_dT!(state, d,T)

Update the `state::ThermodynamicState` of a `PureSubstance` medium with
density `d` [kg/m^3] and temperature `T` [K].
"""
setState_dT!(state::ThermodynamicState,d,T) = setState_dTX!(state,d,T,X_dummy)

### Medium functions that are additionally available for media <: PartialPureSubstance -----------------------

"""
    density_ph(medium,p,h)

Return density in [kg/m^3] for `medium::PureSubstance` from
pressure `p` [Pa] and specific enthalpy `h` [J/kg].
"""
density_ph(m::PureSubstance, p,h) = density(setState_ph(m,p,h))

"""
    temperature_ph(medium,p,h)

Return temperature in [K] for `medium::PureSubstance` from
pressure `p` [Pa] and specific enthalpy `h` [J/kg].
"""
temperature_ph(m::PureSubstance, p,h) = temperature(setState_ph(m,p,h))

"""
    pressure_dT(medium,d,T)

Return pressure in [Pa] for `medium::PureSubstance` from density `d` in [kg/m^3]
and temperature `T` [K].
"""
pressure_dT(m::PureSubstance, d,T) = pressure(setState_dT(m,d,T))

"""
    specificEnthalpy_dT(medium,d,T)

Return specific enthalpy in [J/kg] for `medium::PureSubstance` from
density `d` [kg/m^3] and and temperature `T` [K].
"""
specificEnthalpy_dT(m::PureSubstance, d,T) = specificEnthalpy(setState_dT(m,d,T))

### Derivatives needed for mass and energy balance ------------------------------------------------------

density_ph_der_1(m::PureSubstance, p, h) = 0.0
density_ph_der_2(m::PureSubstance, p, h) = undefinedFunction("density_ph_der_2", m)
density_ph_der_3(m::PureSubstance, p, h) = undefinedFunction("density_ph_der_3", m)


module Zm_par

using Printf, JuLIP, Distributions, LinearAlgebra, ASE

export Stationary, VelocityVerlet, Zmethod, MaxwellBoltzmann_scale, MaxwellBoltzmann, kB, fs

kB = 8.617330337217213e-05 #units taken from Python ASE
fs = 0.09822694788464063

function MaxwellBoltzmann(at, temp)
    d = Normal()
    M = reshape(collect(rand(d, 3*length(at.M))), (length(at.M),3))
    M2 = M .* sqrt.(at.M .* temp)
    set_momenta!(at, collect(transpose(M2)))

    return at
end

function Stationary(at)
    p0 = sum(at.P)
    mtot = sum(at.M)
    v0 = p0 / mtot

    momenta = [at.M[i] * v0 for i in 1:length(at.M)]
    set_momenta!(at, collect(at.P - momenta))

    return at
end

function VelocityVerlet(IP, at, dt)
    V = at.P ./ at.M
    A = forces(IP, at) ./ at.M

    set_positions!(at, at.X + (V .* dt) + (.5 * A * dt^2))

    nA = forces(IP, at) ./ at.M
    nV = V + (.5 * (A + nA) * dt)
    set_momenta!(at, collect(nV .* at.M))

    return at
end

function MaxwellBoltzmann_scale(at, temp)
    d = Normal()
    M = reshape(collect(rand(d, 3*length(at.M))), (length(at.M),3))

    N = []

    for j in 1:1000
        d = Normal()
        M = reshape(collect(rand(d, 3*length(at.M))), (length(at.M),3))
        n = 0
        for i in 1:1:length(M[:,1])
            n += norm(M[i,:])
        end
        push!(N,n)
    end

    d = Normal()
    M = reshape(collect(rand(d, 3*length(at.M))), (length(at.M),3))
    Mnm = mean(N)

    for i in 1:10000
        n = 0
        for i in 1:length(M[:,1])
            n += norm(M[i,:])
        end
        if n < Mnm
            M = M*1.0001
        else
            M = M*0.9999
        end
    end

    M2 = M .* sqrt.(at.M .* temp)
    set_momenta!(at, transpose(M2))

    return at #, Ml, Mnm, Sl, M
end

function Zmethod(IP, at, nsteps, R, dt, A, N, save_config)
    E0 = energy(IP, at)

    m = at.M

    E_tot = zeros(nsteps*R)
    E_pot = zeros(nsteps*R)
    E_kin = zeros(nsteps*R)
    P = zeros(nsteps*R)
    T = zeros(nsteps*R)

    al = []

    for j in 0:R-1
        for i in 1:nsteps
            k = (j*nsteps)+i
            at = VelocityVerlet(IP, at, dt * fs)
            Ek = ((0.5 * sum(at.M) * norm(at.P ./ at.M)^2)/length(at.M)) / length(at.M)
            Ep = (energy(IP, at) - E0) / length(at.M)
            E_tot[k] = Ek + Ep
            E_pot[k] = Ep
            E_kin[k] = Ek
            T[k] = Ek / (1.5 * kB)
            P[k] = -tr(stress(IP, at))/3.0

            if k % 100 == 0
                Temp =  (Ek / (1.5 * kB))
                println("iteration: ($k), temp: ($Temp)")
            end

            if i % save_config == 0
                push!(al, deepcopy(at))
            end
        end

        v = at.P ./ m
        C = A/norm(v)

        set_momenta!(at, collect((v + C*v) .* m))
    end

    return E_tot, E_pot, E_kin, P, T, al
end

end


using Images

import Pkg
Pkg.activate("../../WCvsLHE")
using WCvsLHE

normalize(x) = typeof(x)( (x-minimum(x).*ones(size(x)))*0.7./(maximum(x)-minimum(x)) + 0.15*(ones(size(x))) )

function plot_comparison_middle(A, B,C;args...)
    mA = floor(size(A,1)/2) |> Int64
    mB = floor(size(B,1)/2) |> Int64
    mC = floor(size(C,1)/2)|> Int64
    Plots.plot([Float64.(A[mA,:]),Float64.(B[mB,:]), Float64.(C[mC,:])]; args...)
end

function test(name, params; lhe3d = false, dest = name)
	println(string("Processing: ", name))
	img = Gray.(load(string("originals/",name,".png")))
	img = Gray.(normalize(Float64.(img)))
	print(string("\tContrast enhancement: WC2D "))
	wc2d = wc(img, params[:WC2D], algo_type = :planar) |> normalize
	print(string("LHE2D "))
	lhe2d = lhe(img, params[:LHE2D], algo_type = :planar) |> normalize
	print(string("WC3D "))
	wc3d = wc(img, params[:WC3D], algo_type = :cortical) |> normalize

	if lhe3d
		print(string("LHE3D "))
		lhe3d = lhe(img, params[:LHE3D], algo_type = :cortical) |> normalize
		save(string("results/",dest,"_lhe3d.png"), map(clamp01nan,lhe3d))
	end

	println("\n\tSaving.")

	save(string("results/",dest,".png"), img)
	save(string("results/",dest,"_wc2d.png"), map(clamp01nan, wc2d))
	save(string("results/",dest,"_lhe2d.png"), map(clamp01nan, lhe2d))
	save(string("results/",dest,"_wc3d.png"), map(clamp01nan, wc3d))
end

# Illusion parameters
illusions = Dict(
"white" => Dict(:WC2D => Params(10,20,.7,1.4),
				:LHE2D => Params(10,50,.7,1),
				:WC3D => Params(20,30,.7,1.4),
				:LHE3D => Params(2, 50,.7,1))

, "brightness" => Dict(:WC2D => Params(2,10,.7,1.4),
				:LHE2D => Params(2,10,.7,1),
				:WC3D => Params(2,10,.7,1.4),
				:LHE3D => Params(2, 10,.7,1))

, "checkerboard" => Dict(:WC2D => Params(10,70,.7,1.4),
				:LHE2D => Params(10,70,.7,1),
				:WC3D => Params(10,70,.7,1.4),
				:LHE3D => Params(10,70,.7,1))

, "chevreul" => Dict(:WC2D => Params(2, 5,.7,1),
				:LHE2D => Params(2,10,.7,1),
				:WC3D => Params(2, 40,.5,1),
				:LHE3D => Params(5,7,.7,1))

, "chevreulCanc" => Dict(:WC2D => Params(5, 3,.9,1),
				:LHE2D => Params(5, 3,.9,1),
				:WC3D => Params(2, 20,.5,1.4),
				:LHE3D => Params(5, 3,.9,1))

, "dungeon" => Dict(:WC2D => Params(6,10,.7,1.4),
				:LHE2D => Params(5,40,.7,1),
				:WC3D => Params(2,50,.7,1.4),
				:LHE3D => Params(5,50,.7,1))

, "gratings" => Dict(:WC2D => Params(2,6,.7,1),
				:LHE2D => Params(2,6,.7,1),
				:WC3D => Params(2,6,.7,1),
				:LHE3D => Params(2, 6,.7,1))

, "hong_shevell" => Dict(:WC2D => Params(5,20,.7,1),
				:LHE2D => Params(5,.5,.7,1),
				:WC3D => Params(10,30,.7,1),
				:LHE3D => Params(10,30,.7,1))

, "luminance" => Dict(:WC2D => Params(2,6,.7,1),
				:LHE2D => Params(2,6,.7,1),
				:WC3D => Params(2,6,.7,1),
				:LHE3D => Params(2,6,.7,1))

, "poggendorff" => Dict(:WC2D => Params(3,10,.5,1),
				:LHE2D => Params(3,10,.5,1),
				:WC3D => Params(3,10,.5,1),
				:LHE3D => Params(3,10,.5,1))

, "tilt" => Dict(:WC2D => Params(15,20,.7,1),
				:LHE2D => Params(15,20,.7,1),
				:WC3D => Params(15,20,.7,1),
				:LHE3D => Params(15,20,.7,1))
)

function batch_test(dict; args...)
	for k in keys(dict)
		if k ∉ ["hong_shevell", "tilt", "dungeon"]
			test(k, dict[k]; args...)
		end
	end

	test("hs1", dict["hong_shevell"]; args...)
	test("hs2", dict["hong_shevell"]; args...)

	test("dungeon1", dict["dungeon"]; args...)
	test("dungeon2", dict["dungeon"]; args...)

	test("tilt1", dict["tilt"]; args...)
	test("tilt2", dict["tilt"]; args...)
end

mkpath("results")
batch_test(illusions, lhe3d = true)


import Distributions.Bernoulli
using JLD
using ArgParse
using Random

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table s begin
        "q"
            help = "number of levels"
            arg_type = Int
            required = true
        "ρ"
            help = "oversampling radius"
            arg_type = Float64
            required = true
        "dense"
            help = "compute reference solution using dense linear algebra"
            arg_type = Bool
            required = true
    end

    return parse_args(s)
end

parsed_args = parse_commandline()
println("Parsed args:")
for (arg,val) in parsed_args
    println("  $arg  =>  $val")
end

for key in keys(parsed_args)
  symb = Symbol(key)
  data = parsed_args[key]
  @eval $(symb) = $(data)
end

let
  Random.seed!(123)

  include("BEMUtils.jl")

  λ = 1.0

  @show Threads.nthreads()

  hwv = haarCube(q)

  n_test = 100
  test_points = points([1.0, 1.0, 1.0] .- [1.0, 0.1, 0.1] .* rand(3, n_test))

  n_charges = 1000
  charges_weights = [sign.(rand(Bernoulli(0.8), n_charges) .- 0.5);
                     sign.(rand(Bernoulli(0.3), n_charges) .- 0.5)]
  charges_points = points([(rand(3, n_charges) .* [1.0, 1.0, 0.1] .+ [0.0, 0.0, 1.0])';
                           (rand(3, n_charges) .* [1.0, 0.1, 1.0] .+ [0.0, 1.0, 0.0])']')

  #Either loads The exact solution, or computes it if it is not available
  if isfile("./out/jld/BEM_predict_q_$(q)_rho_$(Int(ρ * 10)).jld")
    # Loads the ordering
    println("Loading ordering from file")
    ld = load("./out/jld/BEM_sort_q_$(q).jld")
    P = ld["P"]
    revP = ld["revP"]
    distances = ld["distances"]
    hwv = hwv[P]

    println("loading boundary data from file")
    ld = load("./out/jld/BEM_boundaryData_q_$(q).jld")
    boundaryData = ld["boundaryData"]

    println("loading prediction from file")
    ld = load("./out/jld/BEM_predict_q_$(q)_rho_$(Int(ρ * 10)).jld")
    u_predict = ld["u_predict"]
    time_predict = ld["time_predict"]
    length_rowval = ld["length_rowval"]

  else
    println("Computing ordering from scratch")
    colptr, rowval, P, revP, distances = sortSparseRev(hwv, ρ)
    save("./out/jld/BEM_sort_q_$(q).jld",
    "P", P,
    "revP", revP,
    "distances", distances)

    # Change distance to ensure that no grouping into skeletons takes place.

    distances .+= 1e-12 * (1 : length(distances))
    length_rowval = length(rowval)
    skeletons = construct_skeletons(colptr, rowval, distances, λ)
    hwv = hwv[P]

    println("computing boundary data from scratch")
    boundaryData = assembleMatrix(hwv, charges_points) * charges_weights
    save("./out/jld/BEM_boundaryData_q_$(q).jld", "boundaryData", boundaryData)

    println("computing prediction from scratch")
    time_predict = @elapsed u_predict = predict(skeletons, boundaryData, hwv, test_points)[1]
    save("./out/jld/BEM_predict_q_$(q)_rho_$(Int(ρ * 10)).jld", "u_predict", u_predict, "time_predict", time_predict, "length_rowval", length_rowval)
  end

  u_true = assembleMatrix(test_points, charges_points) * charges_weights

  # >>>>>>>>>>>>>>>>>>>>>>>> dense part >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
  #Either loads The exact solution, or computes it if it is not available
  if dense
    if isfile("./out/jld/BEM_exact_q_$(q).jld")
      println("Loading full matrix from file")
      ld = load("./out/jld/BEM_exact_q_$(q).jld")
      Θ = ld["Θ"]
      time_assemble_dense = ld["timing_assemble_dense"]
    else
      println("Computing full matrix from scratch")
      time_assemble_dense = @elapsed Θ = assembleMatrix(hwv)
      save("./out/jld/BEM_exact_q_$(q).jld", "Θ", Θ, "timing_assemble_dense", time_assemble_dense)
    end

    time_compute_dense = @elapsed u_BEM_exact = assembleMatrix(test_points, hwv) * (Θ \ boundaryData)
    time_dense = time_compute_dense + time_assemble_dense
    @show mean((u_true - u_BEM_exact).^2)
    @show mean((u_predict - u_BEM_exact).^2)
    @show time_dense
  end
  # <<<<<<<<<<<<<<<<<<<<<<<< dense part <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

  @show mean((u_true - u_predict).^2)
  @show time_predict

  @show 2 * length_rowval / length(hwv) / length(hwv)
end


module PDBenchmark

# Write your package code here.
using PeriDyn
using PDMesh

include("./tests/test.jl")
include("./tests/standardtests.jl")
include("./geom/geom.jl")

end


# trying to see range bearing AprilTag detections

using Images, ImageDraw, ImageView, FileIO
using YAML, JLD, HDF5
using Colors

using AprilTags

include(joinpath(dirname(@__FILE__),"racecarUtils.jl"))

cfg = loadConfig()
# some default, probably not right
# cw, ch = 330.4173, 196.32587
# focal = 340.97913
# tagsize = 0.172

cw, ch = cfg[:intrinsics][:cx], cfg[:intrinsics][:cy]
focal = cfg[:intrinsics][:fx]
tagsize = 0.172

# datafolder = ENV["HOME"]*"/data/racecar/smallloop/"
datafolder = ENV["HOME"]*"/data/racecar/thursday/"
@load datafolder*"tag_det_per_pose.jld" tag_bag
# tag_bag[0]

imgfolder = "newCalibrationTest3"
# camcount = readdlm(datafolder*"$(imgfolder)/cam-count.csv",',')
camlookup = Dict{Int, String}()
for i in 1:length(tag_bag)
  camlookup[i-1] = "camera_image$(5*(i-1)).jpeg"
  # camlookup[camcount[i,1]] = strip(camcount[i,2])
end

IMGS = []
psid = 1
for psid in 0:(length(tag_bag)-1)
  img = load(datafolder*"$(imgfolder)/$(camlookup[psid])")
  push!(IMGS, deepcopy(img))
  for (tid, vals) in tag_bag[psid]
    drawTagLine!(IMGS[psid+1], vals)
  end
end

IMGS[1]
IMGS[20]
IMGS[30]
IMGS[40]
IMGS[50]
IMGS[59]
IMGS[60]
IMGS[61]
IMGS[62]
IMGS[63]
IMGS[64]
IMGS[70]
IMGS[71]
IMGS[72]
IMGS[73]
IMGS[74]
IMGS[75]
IMGS[76]
IMGS[77]
IMGS[78]
IMGS[79]
IMGS[87]
IMGS[88]
IMGS[90]
IMGS[99]
IMGS[100]
IMGS[101]
IMGS[102]
IMGS[109]
IMGS[110]
IMGS[119]
IMGS[120]
IMGS[121]
IMGS[122]
IMGS[123]
IMGS[124]
IMGS[125]
IMGS[126]
IMGS[127]
IMGS[128]
IMGS[129]
IMGS[130]
IMGS[131]
IMGS[132]
IMGS[139]
IMGS[140]
IMGS[149]
IMGS[150]
IMGS[159]
IMGS[160]
IMGS[163]

tag_bag[60][18][:bearing]
tag_bag[61][18][:bearing]
tag_bag[62][18][:bearing]
tag_bag[63][18][:bearing]

tag_bag[77]

tag_bag[86]
tag_bag[87]
tag_bag[88]

println("$(tag_bag[98][0][:bearing]), $(tag_bag[98][0][:range])")
println("$(tag_bag[99][0][:bearing]), $(tag_bag[99][0][:range])")
println("$(tag_bag[100][0][:bearing]), $(tag_bag[100][0][:range])")
println("$(tag_bag[101][0][:bearing]), $(tag_bag[101][0][:range])")
println("$(tag_bag[102][0][:bearing]), $(tag_bag[102][0][:range])")

println("$(tag_bag[98][1][:bearing]), $(tag_bag[98][1][:range])")
println("$(tag_bag[99][1][:bearing]), $(tag_bag[99][1][:range])")
println("$(tag_bag[100][1][:bearing]), $(tag_bag[100][1][:range])")
println("$(tag_bag[101][1][:bearing]), $(tag_bag[101][1][:range])")
println("$(tag_bag[102][1][:bearing]), $(tag_bag[102][1][:range])")

tag_bag[109]
tag_bag[110]

tag_bag[123]

IMGS[125]
tag_bag[124]

tag_bag[125]

# pose :x1
psid = 1

pose_tags = tag_bag[psid]
img = load(datafolder*"images/$(camlookup[psid])")
push!(IMGS, deepcopy(img))

drawTagLine!(IMGS[2], pose_tags[0])
drawTagLine!(IMGS[2], pose_tags[1])
drawTagLine!(IMGS[2], pose_tags[16])
drawTagLine!(IMGS[2], pose_tags[584])

IMGS[2]

# pose :x10
psid = 10

pose_tags = tag_bag[psid]
img = load(datafolder*"images/$(camlookup[psid])")
push!(IMGS, deepcopy(img))

drawTagLine!(IMGS[end], pose_tags[0])
drawTagLine!(IMGS[end], pose_tags[1])
drawTagLine!(IMGS[end], pose_tags[16])

IMGS[end]

# pose :x25
psid = 25

pose_tags = tag_bag[psid]
img = load(datafolder*"images/$(camlookup[psid])")
push!(IMGS, deepcopy(img))

drawTagLine!(IMGS[end], pose_tags[8])
drawTagLine!(IMGS[end], pose_tags[9])
drawTagLine!(IMGS[end], pose_tags[10])
drawTagLine!(IMGS[end], pose_tags[16])

IMGS[end]

# pose :x25
psid = 26

pose_tags = tag_bag[psid]
img = load(datafolder*"images/$(camlookup[psid])")
push!(IMGS, deepcopy(img))

drawTagLine!(IMGS[end], pose_tags[9])
drawTagLine!(IMGS[end], pose_tags[10])
drawTagLine!(IMGS[end], pose_tags[16])

IMGS[end]

using TransformUtils

# redetect tags

detector = AprilTagDetector()

tags = detector(img)

pose = homography_to_pose(tags[5].H, -focal, focal, cw, ch)

cV = pose[1:3,4]

bV = cfg[:bRc]*cV

bearing = atan2(-bV[1], bV[3])
# tag_det[id][:range] = [norm(tag_det[id][:pos][[1;3]]);]

foreach(tag->drawTagBox!(img,tag, width = 5, drawReticle = false), tags)

img

#


module models

using ..gpflow
import ..gpflow: instantiate!, predict_f, predict_f_samples
export
    GPR,
    SGPR,
    VGP,
    SVGP,
    GPMC,
    SGPMC,
    instantiate!,
    predict_f,
    predict_f_samples

abstract type GPModel <: Model end

function predict_f(m::GPModel, Xnew)
    return m.o.predict_f(Xnew)
end
function predict_f_samples(m::GPModel, Xnew, num_samples)
    return m.o.predict_f_samples(Xnew, num_samples)
end

"""
    Gaussian Process Regression.

    This is a vanilla implementation of GP regression with a Gaussian
    likelihood. In this case inference is exact, but costs O(N^3). This means
    that we can compute the predictive distributions (predict_f, predict_y) in
    closed-form, as well as the marginal likelihood, which we use to estimate
    (optimize) the kernel parameters.

    Multiple columns of Y are treated independently, using the same kernel.

    The log likelihood of this model is sometimes referred to as the
    'marginal log likelihood', and is given by

    .. math::

       \\log p(\\mathbf y | \\mathbf f) = \\mathcal N(\\mathbf y | 0, \\mathbf K + \\sigma_n \\mathbf I)
"""
mutable struct GPR{T1,T2} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    mean_function::Union{AbstractMeanFunction, Nothing}
    name::Union{String, Nothing}
    likelihood::Union{AbstractLikelihood, Nothing}
    o::Union{PyObject, Nothing}
end

function GPR(
    X,
    Y,
    kern::Kernel;
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    name::Union{String, Nothing}=nothing
)
    out = GPR(X, Y, kern, mean_function, name, nothing, nothing)
    instantiate!(out)
    return out
end

function instantiate!(o::Union{GPR, Nothing})
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.GPR(o.X, o.Y, kern_, mean_function_, o.name)
    #= TODO: Shoudld we manually initialize likelihood?
    o.likelihood = likelihoods.Gaussian()
    o.likelihood.o = o.o.likelihood =#
    return o.o
end

# TODO: SGPRUpperMixin required?
"""
    Sparse Variational GP regression. The key reference is

    ::

      @inproceedings{titsias2009variational,
        title={Variational learning of inducing variables in
               sparse Gaussian processes},
        author={Titsias, Michalis K},
        booktitle={International Conference on
                   Artificial Intelligence and Statistics},
        pages={567--574},
        year={2009}
      }

"""
mutable struct SGPR{T1,T2,T3} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    feat
    mean_function::Union{AbstractMeanFunction, Nothing}
    Z::T3
    name::Union{String, Nothing}
    o::Union{PyObject, Nothing}
end

function SGPR(
    X,
    Y,
    kern::Union{Kernel, Nothing};
    feat=nothing,
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    Z=nothing,
    name::Union{String, Nothing}=nothing
)
    out = SGPR(X, Y, kern, feat, mean_function, Z, name, nothing)
    instantiate!(out)
    return out
end

function instantiate!(o::Union{SGPR, Nothing})
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.SGPR(
        o.X, o.Y, kern=kern_;
        mean_function=mean_function_,
        name=o.name
    )
    return o.o
end

"""
    This method approximates the Gaussian process posterior using a multivariate Gaussian.

    The idea is that the posterior over the function-value vector F is
    approximated by a Gaussian, and the KL divergence is minimised between
    the approximation and the posterior.

    This implementation is equivalent to svgp with X=Z, but is more efficient.
    The whitened representation is used to aid optimization.

    The posterior approximation is

    .. math::

       q(\\mathbf f) = N(\\mathbf f \\,|\\, \\boldsymbol \\mu, \\boldsymbol \\Sigma)

"""
mutable struct VGP{T1,T2} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    likelihood::Union{AbstractLikelihood, Nothing}
    mean_function::Union{AbstractMeanFunction, Nothing}
    num_latent::Union{Int, Nothing}
    o::Union{PyObject, Nothing}
end

function VGP(
    X,
    Y,
    kern::Union{Kernel, Nothing},
    likelihood::Union{AbstractLikelihood, Nothing};
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    num_latent::Union{Int, Nothing}=nothing
)
    out = VGP(X, Y, kern, likelihood, mean_function, num_latent, nothing)
    instantiate!(out)
    return out
end

function instantiate!(o::VGP)
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    likelihood_ = instantiate!(o.likelihood)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.VGP(
        o.X, o.Y, kern=kern_, likelihood=likelihood_;
        mean_function=mean_function_,
        num_latent=o.num_latent
    )
    return o.o
end

"""
    This is the Sparse Variational GP (SVGP). The key reference is

    ::

    @inproceedings{hensman2014scalable,
        title={Scalable Variational Gaussian Process Classification},
        author={Hensman, James and Matthews,
                Alexander G. de G. and Ghahramani, Zoubin},
        booktitle={Proceedings of AISTATS},
        year={2015}
    }

"""
mutable struct SVGP{T1,T2,T3,T4,T5,T6} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    likelihood::Union{AbstractLikelihood, Nothing}
    feat::T3
    mean_function::Union{AbstractMeanFunction, Nothing}
    num_latent::Union{Int, Nothing}
    q_diag::Bool
    whiten::Bool
    minibatch_size::Union{Int, Nothing}
    Z::T4
    num_data::Union{Int, Nothing}
    q_mu::T5
    q_sqrt::T6
    o::Union{PyObject, Nothing}
end

function SVGP(
    X,
    Y,
    kern::Union{Kernel, Nothing},
    likelihood::Union{AbstractLikelihood, Nothing};
    feat=nothing,
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    num_latent::Union{Int, Nothing}=nothing,
    q_diag::Bool=false,
    whiten::Bool=true,
    minibatch_size::Union{Int, Nothing}=nothing,
    Z=nothing,
    num_data=nothing,
    q_mu=nothing,
    q_sqrt=nothing
)

    out = SVGP(
        X,
        Y,
        kern,
        likelihood,
        feat,
        mean_function,
        num_latent,
        q_diag,
        whiten,
        minibatch_size,
        Z,
        num_data,
        q_mu,
        q_sqrt,
        nothing
    )
    instantiate!(out)
    return out
end

function instantiate!(o::SVGP)
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    likelihood_ = instantiate!(o.likelihood)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.SGVP(
        o.X,
        o.Y,
        kern=kern_,
        likelihood=likelihood_;
        feat=o.feat,
        mean_function=mean_function_,
        num_latent=o.num_latent,
        q_diag=o.q_diag,
        whiten=o.whiten,
        minibatch_size=o.minibatch_size,
        Z=o.Z,
        num_data=o.num_data,
        q_mu=o.q_mu,
        q_sqrt=o.q_sqrt,
    )
    return o.o
end

mutable struct GPMC{T1,T2} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    likelihood::Union{AbstractLikelihood, Nothing}
    mean_function::Union{AbstractMeanFunction, Nothing}
    num_latent::Union{Int, Nothing}
    o::Union{PyObject, Nothing}
end

function GPMC(
    X,
    Y,
    kern::Union{Kernel, Nothing},
    likelihood::Union{AbstractLikelihood, Nothing};
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    num_latent::Union{Int, Nothing}=nothing,
)
    out = GPMC(
        X,
        Y,
        kern,
        likelihood,
        mean_function,
        num_latent,
        nothing
    )
    instantiate!(out)
    return out
end

function instantiate!(o::GPMC)
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    likelihood_ = instantiate!(o.likelihood)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.GPMC(
        o.X,
        o.Y,
        kern=kern_,
        likelihood=likelihood_;
        mean_function=mean_function_,
        num_latent=o.num_latent
    )
    return o.o
end

"""
    This is the Sparse Variational GP using MCMC (SGPMC). The key reference is

    ::

      @inproceedings{hensman2015mcmc,
        title={MCMC for Variatinoally Sparse Gaussian Processes},
        author={Hensman, James and Matthews, Alexander G. de G.
                and Filippone, Maurizio and Ghahramani, Zoubin},
        booktitle={Proceedings of NIPS},
        year={2015}
      }

    The latent function values are represented by centered
    (whitened) variables, so

    .. math::
       :nowrap:

       \\begin{align}
       \\mathbf v & \\sim N(0, \\mathbf I) \\\\
       \\mathbf u &= \\mathbf L\\mathbf v
       \\end{align}

    with

    .. math::
        \\mathbf L \\mathbf L^\\top = \\mathbf K

"""
mutable struct SGPMC{T1,T2,T3,T4} <: GPModel
    X::T1
    Y::T2
    kern::Union{Kernel, Nothing}
    likelihood::Union{AbstractLikelihood, Nothing}
    feat::T3
    mean_function::Union{AbstractMeanFunction, Nothing}
    num_latent::Union{Int, Nothing}
    Z::T4
    o::Union{PyObject, Nothing}
end

function SGPMC(
    X,
    Y,
    kern::Union{Kernel, Nothing},
    likelihood::Union{AbstractLikelihood, Nothing};
    feat=nothing,
    mean_function::Union{AbstractMeanFunction, Nothing}=nothing,
    num_latent::Union{Int, Nothing}=nothing,
    Z=nothing
    )
    out = SGPMC(
        X,
        Y,
        kern,
        likelihood,
        feat,
        mean_function,
        num_latent,
        Z,
        nothing
    )
    instantiate!(out)
    return out
end

function instantiate!(o::SGPMC)
    if o === nothing return nothing end
    if typeof(o.o)<:PyObject return o.o end
    kern_ = instantiate!(o.kern)
    likelihood_ = instantiate!(o.likelihood)
    mean_function_ = instantiate!(o.mean_function)
    o.o = py_gpflow.models.SGPMC(
        o.X,
        o.Y,
        kern=kern_,
        likelihood=likelihood_;
        fwat=o.feat,
        mean_function=mean_function_,
        num_latent=o.num_latent,
        Z=o.Z
    )
    return o.o
end

end # module


# This module extends the julia markdown parser to improve compatibility with Jupyter, Pandoc etc.
module WeaveMarkdown

using ..Weave: isnothing, take2string!
using Markdown
import Markdown: @trigger, @breaking, Code, MD, withstream, startswith, LaTeX

# HACK: that this definition is type-piracy. It is required since `Markdown`
# does not have a built in system for contextual rendering by users. `io` here
# should always be either `IOBuffer` or `IOContext` since it is reached via
# `sprint` in all user-facing code paths in `Markdown`.
function Markdown.latex(io::Union{IOBuffer,IOContext}, tex::Markdown.LaTeX)
    math_envs = ["align", "equation", "eqnarray"]
    use_dollars =
        !any([occursin("\\begin{$me", tex.formula) for me in math_envs])
    use_dollars && write(io, "\\[")
    write(io, string("\n", tex.formula, "\n"))
    use_dollars && write(io, "\\]\n")
end

mutable struct Comment
    text::String
end

@breaking true -> function dollarmath(stream::IO, block::MD)
    withstream(stream) do
        str = Markdown.startswith(stream, r"^\$\$$"m)
        isempty(str) && return false
        trailing = strip(readline(stream))
        buffer = IOBuffer()
        while !eof(stream)
            line_start = position(stream)
            estr = Markdown.startswith(stream, r"^\$\$$"m)
            if !isempty(estr)
                estr = Markdown.startswith(stream, r"^\$\$$"m)
                if isempty(estr)
                    push!(block, LaTeX(take2string!(buffer) |> chomp))
                end
                return true
            else
                seek(stream, line_start)
            end
            write(buffer, readline(stream, keep = true))
        end
        return false
    end
end

@breaking true -> function topcomment(stream::IO, block::MD)
    buffer = IOBuffer()
    withstream(stream) do
        str = Markdown.startswith(stream, r"^<!--")
        isempty(str) && return false
        while !eof(stream)
            line = readline(stream, keep = true)
            write(buffer, line)
            if occursin(r"-->$", line)
                s = replace(take2string!(buffer) |> chomp, r"-->$" => "")
                push!(block, Comment(s))
                return true
            end
        end
        return false
    end
end

@trigger '<' -> function comment(stream::IO, md::MD)
    withstream(stream) do
        Markdown.startswith(stream, "<!--") || return
        text = Markdown.readuntil(stream, "-->")
        isnothing(text) && return
        return Comment(text)
    end
end

# Create own flavor and copy all the features from julia flavor
Markdown.@flavor weavemd [dollarmath, comment, topcomment]
weavemd.breaking = [weavemd.breaking; Markdown.julia.breaking]
weavemd.regular = [weavemd.regular; Markdown.julia.regular]
for key in keys(Markdown.julia.inner)
    if haskey(weavemd.inner, key)
        weavemd.inner[key] = [weavemd.inner[key]; Markdown.julia.inner[key]]
    else
        weavemd.inner[key] = Markdown.julia.inner[key]
    end
end

include("html.jl")
include("latex.jl")

end  # module


#===============================================

    TETRIS FROM TETRIS
    @author: Gabriel M. Trettel

===============================================#

module Tetris

export start_game,
       Tetromino,
       Board,
       GameStatus,
       move_tetromino!,
       remove_tetromino_from_board,
       collides,
       cpy_piece_to_board,
       rotate!,
       fall_dist

const I = ones(Int64, 4, 1)|>rotr90

const J = [2 0 0;
           2 2 2]

const L = [0 0 3;
           3 3 3]

const O = [4 4;
           4 4]

const S = [0 5 5;
           5 5 0]

const T = [0 6 0;
           6 6 6]

const Z = [7 7 0;
           0 7 7]

const tetrominoes = [I, J, L, O, S, T, Z]
const Board = Matrix{Int64}

who_am_I(a::Array) = Dict(I=>:I, J=>:J, L=>:L, O=>:O, S=>:S, T=>:T, Z=>:Z)[a]

calc_iterable(x_len::Integer, y_len::Integer) = Tuple.(Iterators.product(1:y_len, 1:x_len))

mutable struct Tetromino
    type     :: Array
    x        :: Integer
    y        :: Integer
    is_alive :: Bool
    x_len    :: Integer # Column
    y_len    :: Integer # Row
    iter     :: Array
    name     :: Symbol
    function Tetromino(t::Array)
        x = rand(1:6)
        y = 1

        x_len = length(t[1,:]) # Column
        y_len = length(t[:,1]) # Row

        iter = calc_iterable(x_len, y_len)

        name = who_am_I(t)

        new(t, x, y, true, x_len, y_len, iter, name)
    end

    function Tetromino(type, x, y, is_alive, x_len, y_len, iter)
        new(type, x, y, is_alive, x_len, y_len, iter)
    end

    function Tetromino(t::Integer)
        Tetromino(tetrominoes[t])
    end
end

mutable struct GameStatus
    score  :: Real
    level  :: Real
    record :: Real
    board  :: Union{Board, Nothing}
    curr_tetromino :: Union{Tetromino, Nothing}
    next_tetromino :: Union{Tetromino, Nothing}
    status ::Bool
    function GameStatus(s::Real, l::Real, r::Real)
        new(s, l, r, nothing, nothing, nothing, true)
    end
    function GameStatus(s::Real, l::Real, r::Real, b::Board, ct::Tetromino, nt::Tetromino)
        new(s,l,r,b,ct,nt,true)
    end
end

function rotate_tetromino!(t::Tetromino)
    t.type = rotr90(t.type)
    aux = t.x_len
    t.x_len = t.y_len
    t.y_len = aux
    t.iter = calc_iterable(t.x_len, t.y_len)

    y_offset = 0
    x_offset = 0

    if t.name == :I
        x_offset = t.x_len == 1 ? 1 : -1
        y_offset = t.x_len == 1 ? -2 : 2
    elseif t.name != :O
        rotation_offset = abs(t.x_len - t.y_len)
        x_offset = t.x_len > t.y_len ? -rotation_offset : rotation_offset
        y_offset = t.x_len > t.y_len ? rotation_offset : -rotation_offset
    end
    t.y += t.y+y_offset < 0 ? 0 : y_offset
    t.x += t.x+x_offset < 0 ? 0 : x_offset
end

function pick_a_piece(lastpiece::Integer, nextpiece::Integer)
    lastpiece = nextpiece
    nextpiece = rand(0:length(tetrominoes))

    while nextpiece == 0 || nextpiece == lastpiece
        nextpiece = rand(0:length(tetrominoes))
    end

    # lastpiece = 1
    curr_t = Tetromino(tetrominoes[lastpiece])

    return curr_t, lastpiece, nextpiece
end

function is_a_valid_game(board::Board)
    all(x->x==0, board[4,:])
end

function parse_input(input::String)
    lwc_input = lowercase(input)
    if "a" == lwc_input    || "Ϩ" == input
        return :left
    elseif "s" == lwc_input || "ϫ" == input
        return :down
    elseif "d" == lwc_input || "ϩ" == input
        return :right
    elseif "w" == lwc_input || "Ϫ" == input
        return :rotate
    elseif "l" == lwc_input
        return :lay
    elseif "q" == lwc_input
        exit()
    end

    return :none
end

function clean_channel(input_channel::Channel)
    while isready(input_channel)
        take!(input_channel)
    end
end

function cpy_piece_to_board(board::Board, tetromino::Tetromino)
    for (i, j) in tetromino.iter
        if tetromino.type[i,j] != 0
            board[tetromino.y+i-1, tetromino.x+j-1] = tetromino.type[i, j]
        end
    end
end

function collides(board::Board, t::Tetromino)
    for (i,j) in t.iter
        if t.type[i,j] != 0 && board[t.y+i-1, t.x+j-1] != 0
             return true
         end
     end
     return false
end

function Base.copy(ta::Tetromino, tb::Tetromino)
    for f in fieldnames(Tetromino)
        setproperty!(ta, f, getproperty(tb, f))
    end
end

function rotate!(board::Board, tetromino::Tetromino)
    # Cant rotate because there is no vertical space
    # Using t.x_len, which will be the y_len after rotating
    if tetromino.name != :I
        (tetromino.y + tetromino.x_len -2) > 24 && return # for y positioning adjust
    else
        tetromino.y + tetromino.x_len > 24 && return
    end

    new_t = deepcopy(tetromino)

    rotate_tetromino!(new_t)

    if (new_t.x + new_t.x_len) >= 10
        new_t.x = 11 - new_t.x_len
    end
    new_t.x = max(1, new_t.x)
    new_t.y = max(1, new_t.y)

    collides(board, new_t) && return
    copy(tetromino, new_t)
    return
end

function fall_dist(board::Board, tetromino::Tetromino)
    for l in 4:24
        for (i,j) in tetromino.iter
            # check if there is a collision
            if l+i-1 > 24 || tetromino.type[i,j] != 0 && board[l+i-1, tetromino.x+j-1] != 0
                return l-1 - tetromino.y
            end
        end
    end
    return 0
end

function remove_tetromino_from_board(b::Board, t::Tetromino)
    for (i,j) in t.iter
        if t.type[i,j] != 0
            b[t.y+i-1, t.x+j-1] = 0
        end
    end
end

function move_tetromino!(board::Board, tetromino::Tetromino, move_to::Symbol)
    if move_to == :none return end
    if !tetromino.is_alive return end

    if tetromino.y + tetromino.y_len > 24
        return tetromino.is_alive = false
    end

    # Delete's the old Tetromino positioning
    remove_tetromino_from_board(board::Board, tetromino::Tetromino)

    newx = tetromino.x
    newy = tetromino.y

    # Assumes the next move is possible
    if move_to == :left && tetromino.x > 1
        newx -= 1
    elseif move_to == :right && (tetromino.x < (11 - tetromino.x_len))
        newx += 1
    elseif move_to == :down && tetromino.y < (25 - tetromino.y_len)
        newy += 1
    elseif move_to == :lay && tetromino.y < (25 - tetromino.y_len)
        newy += fall_dist(board, tetromino)
    elseif move_to == :rotate
        rotate!(board, tetromino)
        cpy_piece_to_board(board, tetromino)
        return
    end

    # Check if it is possible
    for (i,j) in tetromino.iter
        # check if there is a collision
        if tetromino.type[i,j] != 0 && board[newy+i-1, newx+j-1] != 0
            # Collision in x-axis don't break the game but tetromino can't move sideways
            if newx != tetromino.x
                newx = tetromino.x
            else
                # Collision in y-axis break the game
                cpy_piece_to_board(board, tetromino)
                tetromino.is_alive = false
                return
            end
        end
    end

    tetromino.x = newx
    tetromino.y = newy

    # Persists the moove
    cpy_piece_to_board(board, tetromino)
    return
end

function shift_1_down(b::Board, row::Integer)
    for r in row:-1:4
        b[r,:] = b[r-1,:]
    end
end

function cleaned_points(b::Board)
    points = 0
    row = 24
    while row > 4 && !all(x->x==0, b[row,:])
        if all(x->x!=0, b[row,:])
            shift_1_down(b, row)
            points += 1
        else
            row -= 1
        end
    end
    return points
end

function gravity(input_channel::Channel, gs::GameStatus)
    @async begin
        while true
            put!(input_channel, :down)
            sleep(0.4 / (1 + (0.1*gs.level)))
        end
    end
end

function start_game(input_channel::Channel)
    output_channel = Channel(1000)
    @async loop(input_channel, output_channel)
    return output_channel
end

function loop(input_channel::Channel, output_channel::Channel)

    board = zeros(Int, 24,10)
    last_piece = 1
    next_piece = 2
    tetromino = nothing
    move_to = :down
    down_move = :none
    quit_game = :continue

    down_channel = Channel(1)
    need_new_piece = true

    gs = GameStatus(0, 0, 0)
    t = gravity(down_channel, gs)

    while true
        if tetromino == nothing || !tetromino.is_alive
            clean_channel(input_channel)
            gs.score += cleaned_points(board)
            if (gs.score > 0 && gs.score % 10 == 0) gs.level += 1 end

            if !is_a_valid_game(board) return end
            tetromino,last_piece,next_piece = pick_a_piece(last_piece, next_piece)
            cpy_piece_to_board(board, tetromino)
        end

        if isready(input_channel)
            key_event = take!(input_channel)

            if "q" == key_event gs.level += 20 end
            move_to = parse_input(key_event)

            move_tetromino!(board, tetromino, move_to)
        end
        if isready(down_channel)
            down_move = take!(down_channel)
            move_tetromino!(board, tetromino, down_move)
        end

        if move_to != :none || down_move != :none
            gs.board = board[4:end,:]
            gs.curr_tetromino = tetromino
            gs.next_tetromino = Tetromino(next_piece)
            put!(output_channel, gs)

            move_to = :none
            down_move = :none
        end
        sleep(0.0001)
    end
    gs.status = false
    put!(output_channel, gs)
end
end #module


module ImageComparisons
    using Winston, Compat, Colors
    using Printf, Random
    include("examples.jl")
    include("issues.jl")
    include("plot.jl")
end

function read_png_data(fn::String)
    surface = Cairo.read_from_png(fn)
    w = Cairo.width(surface)
    h = Cairo.height(surface)
    a = Array{UInt32}(undef, convert(Int, w), convert(Int, h))

    p = ccall((:cairo_image_surface_get_data,Cairo._jl_libcairo),
              Ptr{UInt32}, (Ptr{Cvoid},), surface.ptr)

    for i = 1:length(a)
        a[i] = unsafe_load(p, i)
    end
    return a
end

function img_dist(img1::Matrix{UInt32}, img2::Matrix{UInt32})
    @assert size(img1) == size(img2)
    s = 0.
    for i = 1:length(img1)
        a = convert(RGB, reinterpret(RGB24, img1[i]))
        b = convert(RGB, reinterpret(RGB24, img2[i]))
        s += colordiff(a, b)
    end
    s
end

@testset "Image comparison" begin
    root = dirname(task_local_storage()[:SOURCE_PATH])
    dir1, dir2 = joinpath(root, "_baseline"), joinpath(root, "_current")
    isdir(dir1) || mkdir(dir1)
    isdir(dir2) || mkdir(dir2)

    for name in sort!(names(ImageComparisons))
        name == :ImageComparisons && continue

        @testset "$name" begin
            println("Testing $name")
            figure()
            func = eval(:(ImageComparisons.$name))
            global p = func()

            fn1 = joinpath(dir1, "$name.png")
            fn2 = joinpath(dir2, "$name.png")
            isfile(fn1) || savefig(p, fn1)

            savefig(p, fn2)

            img1 = read_png_data(fn1)
            img2 = read_png_data(fn2)
            @test size(img1) == size(img2)

            @test img_dist(img1, img2) < 0.1
        end
    end
end


"""
printFreq: defaults to 0 => prints every 10%. turn off printing by
           setting to -1.
"""
function cytof5_fit(init::State, c::Constants, d::Data;
                    nmcmc::Int=1000, nburn::Int=1000,
                    monitors=[[:Z, :lam, :W, :v,
                               :sig2, :delta, :alpha, :v,
                               :eta, :eps]],
                    fix::Vector{Symbol}=Symbol[],
                    thins::Vector{Int}=[1],
                    thin_dden::Int=1,
                    printFreq::Int=0, flushOutput::Bool=false,
                    computeDIC::Bool=false, computeLPML::Bool=false,
                    computedden::Bool=false,
                    sb_ibp::Bool=false,
                    use_repulsive::Bool=false, joint_update_Z::Bool=false,
                    verbose::Int=1)

  if verbose >= 1
    fixed_vars_str = join(fix, ", ")
    if fixed_vars_str == ""
      fixed_vars_str = "nothing"
    end
    println("fixing: $fixed_vars_str")
    println("Use stick-breaking IBP: $(sb_ibp)")
  end

  @assert printFreq >= -1
  if printFreq == 0
    numPrints = 10
    printFreq = Int(ceil((nburn + nmcmc) / numPrints))
  end

  y_tuner = begin
    dict = Dict{Tuple{Int, Int, Int}, MCMC.TuningParam{Float64}}()
    for i in 1:d.I
      for n in 1:d.N[i]
        for j in 1:d.J
          if d.m[i][n, j] == 1
            dict[i, n, j] = MCMC.TuningParam(1.0)
          end
        end
      end
    end
    dict
  end

  tuners = Tuners(y_imputed=y_tuner, # yinj, for inj s.t. yinj is missing
                  Z=MCMC.TuningParam(MCMC.sigmoid(c.probFlip_Z, a=0.0, b=1.0)),
                  v=[MCMC.TuningParam(1.0) for k in 1:c.K])

  # Loglike
  loglike = Float64[]

  function printMsg(iter::Int, msg::String)
    if printFreq > 0 && iter % printFreq == 0
      print(msg)
    end
  end

  # Instantiate (but not initialize) CPO stream
  if computeLPML
    cpoStream = MCMC.CPOstream{Float64}()
  end

  # DIC
  if computeDIC
    local tmp = DICparam(p=deepcopy(d.y),
                         mu=deepcopy(d.y),
                         sig=[zeros(Float64, d.N[i]) for i in 1:d.I],
                         y=deepcopy(d.y))
    dicStream = MCMC.DICstream{DICparam}(tmp)

    function updateParams(d::MCMC.DICstream{DICparam}, param::DICparam)
      d.paramSum.p += param.p
      d.paramSum.mu += param.mu
      d.paramSum.sig += param.sig
      d.paramSum.y += param.y
      return
    end

    function paramMeanCompute(d::MCMC.DICstream{DICparam})::DICparam
      return DICparam(d.paramSum.p ./ d.counter,
                      d.paramSum.mu ./ d.counter,
                      d.paramSum.sig ./ d.counter,
                      d.paramSum.y ./ d.counter)
    end

    function loglikeDIC(param::DICparam)::Float64
      ll = 0.0

      for i in 1:d.I
        for j in 1:d.J
          for n in 1:d.N[i]
            y_inj_is_missing = (d.m[i][n, j] == 1)

            # NOTE: Refer to `../compute_loglike.jl` for reasoning.
            if y_inj_is_missing

              # Compute p(m_inj | y_inj, theta) term.
              ll += log(param.p[i][n, j])
            end

            # Compute p(y_inj | theta) term.
            ll += logpdf(Normal(param.mu[i][n, j], param.sig[i][n]),
                         param.y[i][n, j])
          end
        end
      end

      return ll
    end

    # FIXME: Doesn't work for c.noisyDist not Normal
    function convertStateToDicParam(s::State)::DICparam
      p = [[prob_miss(s.y_imputed[i][n, j], c.beta[:, i])
            for n in 1:d.N[i], j in 1:d.J] for i in 1:d.I]

      mu = [[s.lam[i][n] > 0 ? mus(i, n, j, s, c, d) : 0.0
             for n in 1:d.N[i], j in 1:d.J] for i in 1:d.I]

      sig = [[s.lam[i][n] > 0 ? sqrt(s.sig2[i]) : std(c.noisyDist)
              for n in 1:d.N[i]] for i in 1:d.I]

      y = deepcopy(s.y_imputed)

      return DICparam(p, mu, sig, y)
    end
  end

  dden = Matrix{Vector{Float64}}[]

  function update!(s::State, iter::Int, out)
    update_state!(s, c, d, tuners, loglike, fix, use_repulsive,
                  joint_update_Z, sb_ibp)

    if computedden && iter > nburn && (iter - nburn) % thin_dden == 0
      append!(dden,
              [[datadensity(i, j, s, c, d) for i in 1:d.I, j in 1:d.J]])
    end

    if computeLPML && iter > nburn
      # Inverse likelihood for each data point
      like = [[compute_like(i, n, s, c, d) for n in 1:d.N[i]] for i in 1:d.I]

      # Update (or initialize) CPO
      MCMC.updateCPO(cpoStream, vcat(like...))

      # Add to printMsg
      printMsg(iter, " -- LPML: $(MCMC.computeLPML(cpoStream))")
    end

    if computeDIC && iter > nburn
      # Update DIC
      MCMC.updateDIC(dicStream, s, updateParams,
                     loglikeDIC, convertStateToDicParam)

      # Add to printMsg
      printMsg(iter, " -- DIC: $(MCMC.computeDIC(dicStream, loglikeDIC,
                                                 paramMeanCompute))")
    end

    printMsg(iter, "\n")
    if flushOutput
      flush(stdout)
    end
  end

  if isinf(compute_loglike(init, c, d, normalize=true))
    println("Warning: Initial state yields likelihood of zero.")
    msg = """
    It is likely the case that the initialization of missing values
    is not consistent with the provided missing mechanism. The MCMC
    will almost certainly reject the initial values and sample new
    ones in its place.
    """
    println(join(split(msg), " "))
  else
    println("")
  end

  out, lastState = MCMC.gibbs(init, update!, monitors=monitors,
                              thins=thins, nmcmc=nmcmc, nburn=nburn,
                              printFreq=printFreq,
                              loglike=loglike, flushOutput=flushOutput,
                              printlnAfterMsg=!(computeDIC || computeLPML))

  mega_out = out, lastState, loglike

  if computeDIC || computeLPML
    LPML = computeLPML ? MCMC.computeLPML(cpoStream) : NaN
    Dmean, pD = computeDIC ? MCMC.computeDIC(dicStream,
                                             loglikeDIC,
                                             paramMeanCompute,
                                             return_Dmean_pD=true) : (NaN, NaN)
    metrics = Dict(:LPML => LPML,
                   :DIC => Dmean + pD,
                   :Dmean => Dmean,
                   :pD => pD)
    println()
    println("metrics:")
    for (k, v) in metrics
      println("$k => $v")
    end

    mega_out = mega_out ..., metrics
  end

  if computedden
    mega_out = mega_out ..., dden
  end

  return mega_out
end