"""
    NoncentralF(ν1, ν2, λ)
"""
struct NoncentralF{T<:Real} <: ContinuousUnivariateDistribution
    ν1::T
    ν2::T
    λ::T

    function NoncentralF{T}(ν1::T, ν2::T, λ::T) where T
        @check_args(NoncentralF, ν1 > zero(T) && ν2 > zero(T))
        @check_args(NoncentralF, λ >= zero(T))
        new{T}(ν1, ν2, λ)
    end
end

NoncentralF(ν1::T, ν2::T, λ::T) where {T<:Real} = NoncentralF{T}(ν1, ν2, λ)
NoncentralF(ν1::Real, ν2::Real, λ::Real) = NoncentralF(promote(ν1, ν2, λ)...)
NoncentralF(ν1::Integer, ν2::Integer, λ::Integer) = NoncentralF(Float64(ν1), Float64(ν2), Float64(λ))

@distr_support NoncentralF 0.0 Inf

#### Conversions

function convert(::Type{NoncentralF{T}}, ν1::S, ν2::S, λ::S) where {T <: Real, S <: Real}
    NoncentralF(T(ν1), T(ν2), T(λ))
end
function convert(::Type{NoncentralF{T}}, d::NoncentralF{S}) where {T <: Real, S <: Real}
    NoncentralF(T(d.ν1), T(d.ν2), T(d.λ))
end

### Parameters

params(d::NoncentralF) = (d.ν1, d.ν2, d.λ)
@inline partype(d::NoncentralF{T}) where {T<:Real} = T

### Statistics

function mean(d::NoncentralF{T}) where T<:Real
    d.ν2 > 2 ? d.ν2 / (d.ν2 - 2) * (d.ν1 + d.λ) / d.ν1 : T(NaN)
end

var(d::NoncentralF{T}) where {T<:Real} = d.ν2 > 4 ? 2d.ν2^2 *
               ((d.ν1 + d.λ)^2 + (d.ν2 - 2)*(d.ν1 + 2d.λ)) /
               (d.ν1 * (d.ν2 - 2)^2 * (d.ν2 - 4)) : T(NaN)

### Evaluation & Sampling

@_delegate_statsfuns NoncentralF nfdist ν1 ν2 λ

function rand(rng::AbstractRNG, d::NoncentralF)
    r1 = rand(rng, NoncentralChisq(d.ν1,d.λ)) / d.ν1
    r2 = rand(rng, Chisq(d.ν2)) / d.ν2
    r1 / r2
end

# TODO: remove RFunctions dependency once NoncentralChisq has its removed
@rand_rdist(NoncentralF)
function rand(d::NoncentralF)
    r1 = rand(NoncentralChisq(d.ν1,d.λ)) / d.ν1
    r2 = rand(Chisq(d.ν2)) / d.ν2
    r1 / r2
end


using Revise
using ADCME
using ADCMEKit
using NNFEM
using PyPlot
using ProgressMeter
using Printf
close("all")
include("load_domain_function.jl")

NT = 100
Δt = 20/NT

#################### solve static elastic response  ####################
domain = load_crack_domain()
H = domain.elements[1].mat[1].H
EBC_func = nothing
FBC_func = nothing
body_func = nothing
globaldata = GlobalData(missing, missing, missing, missing, domain.neqs, EBC_func, FBC_func, body_func)
Fext = compute_external_force(globaldata, domain)
d = LinearStaticSolver(globaldata, domain, domain.state, H, Fext)
sess = Session(); init(sess)
domain.state = run(sess, d)

figure()
subplot(211)
visualize_scalar_on_scoped_body(domain.state[1:domain.nnodes], zeros(size(domain.state)...), domain)
autoscale(enable=true, axis="both", tight=true)
xticks([])
xlabel("")
title(L"$u_x$")
subplot(212)
visualize_scalar_on_scoped_body(domain.state[domain.nnodes+1:end], zeros(size(domain.state)...), domain)
autoscale(enable=true, axis="both", tight=true)
title(L"$u_y$")
savefig("figures/dippingfault_disp_elasticity.png")

#################### solve viscoelastic relaxation  ####################
strain = getStrain(domain)
stress = getStress(domain)
state = domain.state

domain = load_crack_domain(option="mixed")
ts = GeneralizedAlphaSolverTime(Δt, NT)
ubd, abd = compute_boundary_info(domain, globaldata, ts)
Fext = compute_external_force(domain, globaldata, ts)

d0 = state
v0 = zeros(2domain.nnodes)
a0  = zeros(2domain.nnodes)
# σ0 = zeros(getNGauss(domain),3)
# ϵ0 = zeros(getNGauss(domain),3)
σ0 = stress
ϵ0 = strain

μ = zeros(getNGauss(domain))
λ = zeros(getNGauss(domain))
η = zeros(getNGauss(domain))
k = 0
for i = 1:domain.neles
  e = domain.elements[i]
  for mat in e.mat
    global k += 1
    μ[k] = mat.μ
    λ[k] = mat.λ
    η[k] = mat.η
  end
end

d, v, a, σ, ϵ = ViscoelasticitySolver(
  globaldata, domain, d0, v0, a0, σ0, ϵ0, Δt, NT, μ, λ, η, Fext, ubd, abd
)

sess = Session(); init(sess)

d_, σ_ = run(sess, [d, σ])
domain.history["state"] = []
domain.history["stress"] = []
for i = 1:NT+1
  push!(domain.history["state"], d_[i,:])
  push!(domain.history["stress"], σ_[i,:,:])
end

#################### visulalize results ####################

## get surface slip
slip_idx = findall(domain.nodes[:,2] .≈ 0.0)
coords = domain.nodes[slip_idx,1]
ii = sortperm(coords)
coords = coords[ii]
x_ = d_[:, slip_idx[ii]]
y_ = d_[:, slip_idx[ii] .+ domain.nnodes]

x0 = d0[slip_idx[ii]]
y0 = d0[slip_idx[ii] .+ domain.nnodes]

figure(figsize=(6,3))
subplot(211)
plot(coords, (x_' .- x0)[:,1:20:end])
title(L"u_x - u_{x0}")
autoscale(enable=true, axis="x", tight=true)
xticks([])
subplot(212)
plot(coords, x_[1:20:end, :]')
title(L"u_x")
autoscale(enable=true, axis="x", tight=true)
tight_layout()
xlabel("x")
savefig("figures/dippingfault_ux_viscosity.png")

figure(figsize=(6,3))
subplot(211)
plot(coords, -(y_' .- y0)[:,1:20:end])
title(L"u_y - u_{y0}")
autoscale(enable=true, axis="x", tight=true)
xticks([])
subplot(212)
plot(coords, -y_[1:20:end, :]')
title(L"u_y")
autoscale(enable=true, axis="x", tight=true)
tight_layout()
xlabel("x")
savefig("figures/dippingfault_uy_viscosity.png")

figure(figsize=(5,3))
subplot(211)
pl1, = plot(coords, x0, "o-", markersize = 3)
autoscale(enable=true, axis="x", tight=true)
t = title("time = 0")
xticks([])
ylabel(L"u_x")
subplot(212)
pl2, = plot(coords, y0, "o-", markersize = 3)
autoscale(enable=true, axis="x", tight=true)
xlabel("x")
ylabel(L"u_y")
tight_layout()
function update(i)
  pl1.set_data(coords, x_[i, :])
  pl2.set_data(coords, -y_[i, :])
  t.set_text("time = $(round(i*Δt, digits=1))")
end
p = animate(update, collect(1:2:NT))
saveanim(p, "figures/dippingfault_viscosity.gif")

## get surface stress
idx = []
k = 0
pts = getGaussPoints(domain)
xcoord = Float64[]
for e in domain.elements
    global k
    if sum(map(x->x≈0.0, e.coords[:,2]))≠2
        k += length(e.weights)
        continue
    end
    l = argmin(pts[k+1:k+length(e.weights), 2])
    push!(idx, k + l)
    push!(xcoord, pts[k+l, 1])
    k += length(e.weights)
end
ii = sortperm(xcoord)
xcoord = xcoord[ii]
idx = idx[ii]

σ_surface = σ_[:, idx, :]
s = zeros(size(σ_surface, 1), size(σ_surface, 2))
for i = 1:size(σ_surface, 1)
    for j = 1:size(σ_surface, 2)
        s[i,j] = NNFEM.postprocess_stress(σ_surface[i,j,:], "vonMises")
    end
end

figure(figsize=(8,3))
plot(s[1:20:end, :]')
title("von Mises stress")
savefig("figures/dippingfault_vonMises_stress.png")

figure()
visualize_boundary(domain)
# visualize_mesh(domain)
visualize_scalar_on_scoped_body(η[1:9:end], zeros(domain.nnodes*2), domain)
gca().invert_yaxis()
# autoscale(enable=true, axis="both", tight=true)
scatter(domain.nodes[slip_idx,1], domain.nodes[slip_idx,2], color="m", marker=".", label="Tractrion Free")
legend()
title(L"Viscosity ($\eta$) model")
savefig("figures/dippingfault_mesh.png")

# figure()
# p = visualize_scalar_on_scoped_body(σ_[:, 1:domain.nnodes, 2], d_, domain, scale_factor=1.0)
# # p = visualize_von_mises_stress_on_scoped_body(d_, domain, scale_factor=10.0)
# # p = visualize_total_deformation_on_scoped_body(d_, domain, scale_factor=10.0)
# saveanim(p, "ad_solver.gif")
# close("all")


#function DemoLinRegCUDA4()

    # note that this linear regression code is not the fastest way to do this since it is wasteful in terms of storage. We would need to code an operation Ax+beta*b where A,x,b are CudaArrays and beta is a Float64. This would mean we can then avoid the separate A*x (which stores this on the graph) and then Ax-y. Hence we can cut the storage significantly this way.

    #include("setup.jl")

    #AtransB(A,B)=A'*B

    StartCode()

    A=ADvariable()
#B=ADvariable()
#    loss=sum(AtransB(A,A))
#    loss=meanSquare(A')
#     loss=sum(A'-B')
    # loss=sumSquare(diagm(B)*A*diagm(B))
#    loss=sumSquare(diagm(A.*A))
    loss=sumSquare(abs(A))

    net=EndCode() # defines the graph

    # instantiate parameter nodes and inputs:
    #value=Array(Any,NodeCounter()) # function values on the nodes

    net.value[A]=randn(2,2)
    #net.value[B]=randn(2,1)

    net=compile(net;debug=true) # compile and preallocate memory

    @gpu CUDArt.init([0])
@gpu net=convert(net,"GPU")

gradcheck(net;showgrad=true) # use a small number of datapoints and small network to check the gradient, otherwise this will be very slow
 #   @cpu gradcheck(net;showgrad=true) # use a small number of datapoints and small network to check the gradient, otherwise this will be very slow
#@gpu gradcheckGPU(net;showgrad=false)

#end


#
using DiffMesh, DiffEqFlux, LinearAlgebra, Plots
using GalacticOptim, ParameterSchedulers
using Test
#----------------------------------------------------------------------------#
# model

nprt = 4
ndim = 4

pou = nnPOU(nprt;depth=4,act=tanh)
#pou = femPOU(nprt,x -> @. sin(pi/2*x))
bsp = nnBanachSpace(nprt,ndim;depth=4,act=cos)

net  = MeshNet(pou,bsp)
p,re = Flux.destructure(net)
#--------------------------------------#
function utrue(x)
    u  = similar(x)
#   u .= @. 2 + sin(3.7*pi*x + pi*0.3)
    u .= @. sin(x^2 + 3*x)^2 + exp(x)
#   u .= @. 3x^2
#   u .= @. 2 + DiffMesh.TRI(x,1)
    u .= @. 2 + DiffMesh.TRI2(x,1)
#   u .= @. 2 + abs(x - 0.5)
#   u .= @. 2 + sin(10*pi*x)
    return u
end

# evaluation points
xp = Array(range(-1f0,1f0,length=1000))'
up = utrue(xp)

# data poins (for least sq)
xd = xp
#xd = Array(range(-1f0,1f0,length=50))'
ud = utrue(xd)

it = 0 # iter
#----------------------------------------------------------------------------#
function loss(p,xp)
    phi = re(p)

    fp = phi(xp)
    ep = up .- fp
    loss = Flux.Losses.mse(fp,up)
    return loss
end
#--------------------------------------#
function makeplot(phi)

    fp = phi(xp)

    plt = plot()
    plt = plot!(title="MeshNet, nprt=$nprt, ndim=$ndim"
               ,xlabel="x",ylabel="y")
    plt = plot!(ylims=(0,Inf))

    plt = plot!(xp',up',width=3,label="data")
    plt = plot!(xp',fp',width=3,label="prediction")

    # plot partition functions
    if(phi isa MeshNet)
        par = phi.pou(xp)
        plt = plot!(xp',par',width=3
                    ,label=:none)#,color=:black)
    end

    display(plt)
    return
end
#--------------------------------------#
function cb1(p,l)

    phi = re(p)

    if(phi.bsp isa GL)
        phi.C .= utrue(phi.bsp.xm)
    else
        leastSq!(phi,xd,ud)
    end

    fp  = phi(xp)
    ep  = up .- fp
    er  = norm(ep,Inf)

    makeplot(phi)
    println("loss: $l, pointwise er: $er")
    (er < 1e-6) && return true
    return false
end
#--------------------------------------#
function cb2(p,l)

    phi = re(p)

    fp  = phi(xp)
    ep  = up .- fp
    er  = norm(ep,Inf)

    makeplot(phi)
    println("loss: $l, pointwise er: $er")
    (er < 1e-6) && return true
    return false
end
#----------------------------------------------------------------------------#
# training

func = OptimizationFunction(loss,GalacticOptim.AutoZygote())
prob = OptimizationProblem(func,p,xp)

sch  = ParameterSchedulers.SinExp(λ0=1e-0,λ1=1e-3,γ=5e-1,period=50)
opt1 = ParameterSchedulers.Scheduler(sch,ADAM())
opt1 = ADAM(1e-2)
opt2 = BFGS()

res  = GalacticOptim.solve(prob,opt1,cb=cb1,maxiters=200)
prob = remake(prob,u0=res.minimizer)
res  = GalacticOptim.solve(prob,opt2,cb=cb1,maxiters=200)
prob = remake(prob,u0=res.minimizer)
res  = GalacticOptim.solve(prob,opt2,cb=cb2,maxiters=100)

mzr = res.minimizer
phi = re(mzr)
#--------------------------------------#
return


#=
This file is part of the replication code for: Hasenzagl, T., Pellegrino, F., Reichlin, L., & Ricco, G. (2020). A Model of the Fed's View on Inflation.
Please cite the paper if you are using any part of the code for academic work (including, but not limited to, conference and peer-reviewed papers).
=#

function get_par_unb(θ_bound::Array{Float64,1}, MIN::Array{Float64,1}, MAX::Array{Float64,1}, opt_transf::Array{Int64, 1})

# ----------------------------------------------------------------------------------------------------------------------
# Change domain: from bounded to unbounded domain
# - Transformation: 1 natural logarithm, 2 no transformations, 3 generalized logit
# ----------------------------------------------------------------------------------------------------------------------

     i1 = opt_transf .== 1;
     i2 = opt_transf .== 2;
     i3 = opt_transf .== 3;

     θ_unb     = zeros(size(θ_bound));
     θ_unb[i1] = log.(θ_bound[i1] - MIN[i1]);
     θ_unb[i2] = θ_bound[i2];
     θ_unb[i3] = log.((θ_bound[i3] - MIN[i3]) ./ (MAX[i3] - θ_bound[i3]));

     return θ_unb;
end


##########################
# FunctionCell Interface #
##########################
"""
	tree(nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0  [;kwargs])

Generates a function cell that when piped data, trains a decision tree classifier using the DecisionTree.jl interface.

# Arguments (most from `DecisionTree`)
  * `nsubfeatures` number of features to consider for node thresholding (default 0 e.g. all)
  * `maxdepth` maximum tree depth (default -1 e.g. grow fully)
  * `prune_purity` fraction in [0.0,1.0] of combined purity for leaves to be merged (-1.0 indicates no pruning)

# Keyword arguments (from `DecisionTree`)
  * `rng::AbstractRNG=Base.GLOBAL_RNG` is the random number generator

Read the `DecisionTree.jl` documentation for more information.
"""
tree(nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	FunctionCell(tree, (nsubfeatures, maxdepth, prune_purity), ModelProperties(),
	      kwtitle("Decision Tree (classifier)", kwargs);  kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	tree(x, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0  [;kwargs])

Generates a trained function cell by training a decision tree classifier using the DecisionTree.jl interface.
"""
# Training
tree(x::T where T<:CellDataL, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	tree((getx!(x), gety(x)), nsubfeatures, maxdepth, prune_purity; kwargs...)

tree(x::Tuple{T,S} where T<:AbstractVector where S<:AbstractVector, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	tree((mat(x[1], LearnBase.ObsDim.Constant{2}()), x[2]), nsubfeatures, maxdepth, prune_purity; kwargs...)

tree(x::Tuple{T,S} where T<:AbstractMatrix where S<:AbstractVector, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) = begin

	@assert nobs(x[1]) == nobs(x[2]) "[tree] Expected $(nobs(x[1])) labels/values, got $(nobs(x[2]))."

	# Transform labels first
	enc = labelencn(x[2])
	yenc = label2ind.(x[2],enc)

	# Train model
	treedata = DecisionTree.build_tree(yenc, getobs(x[1])', nsubfeatures, maxdepth; kwargs...)

	# Prune if necessary
	if prune_purity >=0
		treedata = DecisionTree.prune_tree(treedata, prune_purity)
	end

	# Build model properties
	modelprops = ModelProperties(nvars(x[1]), length(enc.label), enc)

	FunctionCell(tree, Model(treedata, modelprops), kwtitle("Decision Tree (classifier)", kwargs))
end

# Execution
tree(x::T where T<:CellData, model::Model{<:DecisionTree.LeafOrNode}) =
	datacell(tree(getx!(x), model), gety(x))
tree(x::T where T<:AbstractVector, model::Model{<:DecisionTree.LeafOrNode}) =
	tree(mat(x, LearnBase.ObsDim.Constant{2}()), model)
tree(x::T where T<:AbstractMatrix, model::Model{<:DecisionTree.LeafOrNode}) =
	# Always return probabilities
	DecisionTree.apply_tree_proba(model.data, getobs(x)', collect(1:model.properties.odim))'::Matrix{Float64} # labels are integer-coded, just use [1,2, ..., C]

##########################
# FunctionCell Interface #
##########################
"""
	treer(maxlabels::Int=5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0  [;kwargs])

Generates a function cell that when piped data, trains a decision tree regressor using the DecisionTree.jl interface.

# Arguments (most from `DecisionTree`)
  * `maxlabels` number of samples used in each leaf to construct estimates (default 5)
  * `nsubfeatures` number of features to consider for node thresholdingi (default 0 e.g. all)
  * `maxdepth` maximum tree depth (default -1 e.g. grow fully)
  * `prune_purity` fraction in [0.0,1.0] of combined purity for leaves to be merged (-1.0 indicates no pruning)

# Keyword arguments (from `DecisionTree`)
  * `rng::AbstractRNG=Base.GLOBAL_RNG` is the random number generator

!!! note
	The effects of `prune_purity` in the case of regression are not quite clear. Should work well for very
	low values e.g. leaves containing many distinct values are combined.

Read the `DecisionTree.jl` documentation for more information.
"""
treer(maxlabels::Int = 5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	FunctionCell(treer, (maxlabels, nsubfeatures, maxdepth, prune_purity), ModelProperties(),
	      kwtitle("Decision Tree (regressor)", kwargs);  kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	treer(x, maxlabels::Int=5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0  [;kwargs])

Generates a trained function cell by training a decision tree regressor using the DecisionTree.jl interface.
"""
# Training
treer(x::T where T<:CellDataL, maxlabels::Int=5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	treer((getx!(x), gety(x)), maxlabels, nsubfeatures, maxdepth, prune_purity; kwargs...)

treer(x::Tuple{T,S} where T<:AbstractVector where S<:AbstractVector, maxlabels::Int=5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) =
	treer((mat(x[1], LearnBase.ObsDim.Constant{2}()), x[2]), maxlabels, nsubfeatures, maxdepth, prune_purity; kwargs...)

treer(x::Tuple{T,S} where T<:AbstractMatrix where S<:AbstractVector, maxlabels::Int=5, nsubfeatures::Int=0, maxdepth::Int=-1, prune_purity::Real=-1.0; kwargs...) = begin

	@assert nobs(x[1]) == nobs(x[2]) "[treer] Expected $(nobs(x[1])) labels/values, got $(nobs(x[2]))."

	# Convert targets to vector of Float first
	yenc = Vector{typeof(1.0)}(vec(x[2]))

	# Train model
	treedata = DecisionTree.build_tree(yenc, getobs(x[1])', maxlabels, nsubfeatures, maxdepth; kwargs...)

	# Prune if necessary
	if prune_purity >=0
		treedata = DecisionTree.prune_tree(treedata, prune_purity)
	end

	# Build model properties
	modelprops = ModelProperties(nvars(x[1]),1)

	FunctionCell(treer, Model(treedata, modelprops), kwtitle("Decision Tree (regressor)", kwargs))
end

# Execution
treer(x::T where T<:CellData, model::Model{<:DecisionTree.LeafOrNode}) =
	datacell(treer(getx!(x), model), gety(x))
treer(x::T where T<:AbstractVector, model::Model{<:DecisionTree.LeafOrNode}) =
	treer(mat(x, LearnBase.ObsDim.Constant{2}()), model)
treer(x::T where T<:AbstractMatrix, model::Model{<:DecisionTree.LeafOrNode})::Matrix{Float64} =
	Matrix(DecisionTree.apply_tree(model.data, getobs(x)')')

##########################
# FunctionCell Interface #
##########################
"""
	randomforest(ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1 [;kwargs])

Generates a function cell that when piped data, trains a random forest classifier using the DecisionTree.jl interface.
The trees are not pruned.

# Arguments (from `DecisionTree`)
  * `ntrees` number of trees (default 50)
  * `nsubfeatures` number of features to consider for node thresholding (default 0 e.g. all)
  * `partialsampling` is the random fraction in [0.0,1.0] of training samples to be used for each tree (default 0.7)
  * `maxdepth` maximum tree depth (default -1 e.g. grow fully)

# Keyword arguments (from `DecisionTree`)
  * `rng::AbstractRNG=Base.GLOBAL_RNG` is the random number generator

Read the `DecisionTree.jl` documentation for more information.
"""
randomforest(ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1; kwargs...) =
	FunctionCell(randomforest, (ntrees, nsubfeatures, partialsampling, maxdepth), ModelProperties(),
	      	     kwtitle("Random Forest (classifier, $ntrees trees)", kwargs);  kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	randomforest(x, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1 [;kwargs])

Generates a trained function cell by training a random forest classifier using the DecisionTree.jl interface.
"""
# Training
randomforest(x::T where T<:CellDataL, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1; kwargs...) =
	randomforest((getx!(x), gety(x)), ntrees, nsubfeatures, partialsampling, maxdepth; kwargs...)

randomforest(x::Tuple{T,S} where T<:AbstractVector where S<:AbstractVector, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1; kwargs...) =
	randomforest((mat(x[1], LearnBase.ObsDim.Constant{2}()), x[2]), ntrees, nsubfeatures, partialsampling, maxdepth; kwargs...)

randomforest(x::Tuple{T,S} where T<:AbstractMatrix where S<:AbstractVector, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1; kwargs...) =
begin
	@assert nobs(x[1]) == nobs(x[2]) "[randomforest] Expected $(nobs(x[1])) labels/values, got $(nobs(x[2]))."

	# Transform labels first
	enc = labelencn(x[2])
	yenc = label2ind.(x[2],enc)

	# Train model
	rforestdata = DecisionTree.build_forest(yenc, getobs(x[1])', nsubfeatures, ntrees, partialsampling, maxdepth; kwargs...)

	# Build model properties
	modelprops = ModelProperties(nvars(x[1]), length(enc.label), enc)

	FunctionCell(randomforest, Model(rforestdata, modelprops), kwtitle("Random Forest (classifier, $ntrees trees)", kwargs))
end

# Execution
randomforest(x::T where T<:CellData, model::Model{<:DecisionTree.Ensemble}) =
	datacell(randomforest(getx!(x), model), gety(x))
randomforest(x::T where T<:AbstractVector, model::Model{<:DecisionTree.Ensemble}) =
	randomforest(mat(x, LearnBase.ObsDim.Constant{2}()), model)
randomforest(x::T where T<:AbstractMatrix, model::Model{<:DecisionTree.Ensemble})::Matrix{Float64} =
	DecisionTree.apply_forest_proba(model.data, getobs(x)', collect(1:model.properties.odim))' # labels are integer-coded, just use [1,2, ..., C]

##########################
# FunctionCell Interface #
##########################
"""
	randomforestr(ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5 [;kwargs])

Generates a function cell that when piped data, trains a random forest regressor using the DecisionTree.jl interface.
The trees are not pruned.

# Arguments (from `DecisionTree`)
  * `ntrees` number of trees (default 50)
  * `nsubfeatures` number of features to consider for node thresholding (default 0 e.g. all)
  * `partialsampling` is the random fraction in [0.0,1.0] of training samples to be used for each tree (default 0.7)
  * `maxdepth` maximum tree depth (default -1 e.g. grow fully)
  * `maxlabels` number of samples used in each leaf to construct estimates (default 5)

# Keyword arguments (from `DecisionTree`)
  * `rng::AbstractRNG=Base.GLOBAL_RNG` is the random number generator

Read the `DecisionTree.jl` documentation for more information.
"""
randomforestr(ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5; kwargs...) =
	FunctionCell(randomforestr, (ntrees, nsubfeatures, partialsampling, maxdepth, maxlabels), ModelProperties(),
	             kwtitle("Random Forest (regressor, $ntrees trees)", kwargs);  kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	randomforestr(x, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5 [;kwargs])

Generates a trained function cell by training a random forest regressor using the DecisionTree.jl interface.
"""
# Training
randomforestr(x::T where T<:CellDataL, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5; kwargs...) =
	randomforestr((getx!(x), gety(x)), ntrees, nsubfeatures, partialsampling, maxdepth, maxlabels; kwargs...)

randomforestr(x::Tuple{T,S} where T<:AbstractVector where S<:AbstractVector, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5; kwargs...) =
	randomforestr((mat(x[1], LearnBase.ObsDim.Constant{2}()), x[2]), ntrees, nsubfeatures, partialsampling, maxdepth, maxlabels; kwargs...)

randomforestr(x::Tuple{T,S} where T<:AbstractMatrix where S<:AbstractVector, ntrees::Int=50, nsubfeatures::Int=0, partialsampling::Float64=0.7, maxdepth::Int=-1, maxlabels::Int=5; kwargs...) =
begin
	@assert nobs(x[1]) == nobs(x[2]) "[randomforestr] Expected $(nobs(x[1])) labels/values, got $(nobs(x[2]))."

	# Convert targets to vector of Float first
	yenc = Vector{typeof(1.0)}(vec(x[2]))

	# Train model
	rforestdata = DecisionTree.build_forest(yenc, getobs(x[1])', nsubfeatures, ntrees, maxlabels, partialsampling, maxdepth; kwargs...)

	# Build model properties
	modelprops = ModelProperties(nvars(x[1]),1)

	FunctionCell(randomforestr, Model(rforestdata, modelprops), kwtitle("Random Forest (regressor, $ntrees trees)", kwargs))
end

# Execution
randomforestr(x::T where T<:CellData, model::Model{<:DecisionTree.Ensemble}) =
	datacell(randomforestr(getx!(x), model), gety(x))
randomforestr(x::T where T<:AbstractVector, model::Model{<:DecisionTree.Ensemble}) =
	randomforestr(mat(x, LearnBase.ObsDim.Constant{2}()), model)
randomforestr(x::T where T<:AbstractMatrix, model::Model{<:DecisionTree.Ensemble}) =
	Matrix(DecisionTree.apply_forest(model.data, getobs(x)')')

##########################
# FunctionCell Interface #
##########################
"""
	adaboostump(niterations::Int=10 [;kwargs])

Generates a function cell that when piped data, trains an Adaboosted-stump classifier using the DecisionTree.jl interface.

# Arguments (most from `DecisionTree`)
  * `niterations` number of iterations e.g ensemble size (default 10)

# Keyword arguments (from `DecisionTree`)
  * `rng::AbstractRNG=Base.GLOBAL_RNG` is the random number generator

Read the `DecisionTree.jl` documentation for more information.
"""
adaboostump(niterations::Int=10; kwargs...) =
	FunctionCell(adaboostump, (niterations,), ModelProperties(),
	      kwtitle("Adaboost Stumps (classifier, $niterations stumps)", kwargs);  kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	adaboostump(x, niterations::Int=10 [;kwargs])

Generates a trained function cell by training an Adaboosted-stump classifier using the DecisionTree.jl interface.
"""
# Training
adaboostump(x::T where T<:CellDataL, niterations::Int=10; kwargs...) = adaboostump((getx!(x), gety(x)), niterations; kwargs...)

adaboostump(x::Tuple{T,S} where T<:AbstractVector where S<:AbstractVector, niterations::Int=10; kwargs...) =
	adaboostump((mat(x[1], LearnBase.ObsDim.Constant{2}()), x[2]), niterations; kwargs...)

adaboostump(x::Tuple{T,S} where T<:AbstractMatrix where S<:AbstractVector, niterations::Int=10; kwargs...) = begin

	@assert nobs(x[1]) == nobs(x[2]) "[adaboostump] Expected $(nobs(x[1])) labels/values, got $(nobs(x[2]))."

	# Transform labels first
	enc = labelencn(x[2])
	yenc = label2ind.(x[2],enc)

	# Train model
	stumpdata = DecisionTree.build_adaboost_stumps(yenc, getobs(x[1])', niterations; kwargs...)

	# Build model properties
	modelprops = ModelProperties(nvars(x[1]), length(enc.label), enc)

	FunctionCell(adaboostump, Model(stumpdata, modelprops), kwtitle("Adaboost Stumps (classifier, $niterations stumps)", kwargs));
end

# Execution
adaboostump(x::T where T<:CellData, model::Model{<:Tuple{DecisionTree.Ensemble,Vector}}) =
	datacell(adaboostump(getx!(x), model), gety(x))
adaboostump(x::T where T<:AbstractVector, model::Model{<:Tuple{DecisionTree.Ensemble,Vector}}) =
	adaboostump(mat(x, LearnBase.ObsDim.Constant{2}()), model)
adaboostump(x::T where T<:AbstractMatrix, model::Model{<:Tuple{DecisionTree.Ensemble,Vector}}) =
	DecisionTree.apply_adaboost_stumps_proba(model.data[1], model.data[2],
					  getobs(x)', collect(1:model.properties.odim))'::Matrix{Float64}


# Tracing charged particle in the Harris current sheet.
# This example shows how to trace protons in a stationary magnetic field that
# corresponds to the 1D Harris current sheet.
# Reference: https://en.wikipedia.org/wiki/Current_sheet
#
# Hongyang Zhou, hyzhou@umich.edu

using TestParticle
using TestParticle: getB_CS_harris
using OrdinaryDiffEq
using StaticArrays
using PyPlot
using Statistics: mean

## Obtain field

# Harris current sheet parameters in SI units
const B₀, L = 20e-9, 0.4TestParticle.Rₑ

function getB(xu)
   SVector{3}(getB_CS_harris(xu[1:3], B₀, L))
end

function getE(xu)
   SA[0.0, 0.0, 0.0]
end

## Initialize particles

m = TestParticle.mᵢ
q = TestParticle.qᵢ
c = TestParticle.c
Rₑ = TestParticle.Rₑ

Ek = 5e7 # [eV]

# initial velocity, [m/s]
v₀ = [c*√(1-1/(1+Ek*q/(m*c^2))^2), 0.0, 0.0]
# initial position, [m]
r₀ = [-5.0Rₑ, 0.0, 0.0]
stateinit = [r₀..., v₀...]

param = prepare(getE, getB)
tspan = (0.0, 10.0)

prob = ODEProblem(trace_analytic!, stateinit, tspan, param)

sol = solve(prob, Tsit5(); save_idxs=[1,2,3])

## Visualization

using3D()
fig = plt.figure(figsize=(10,6))
ax = fig.gca(projection="3d")

n = 100 # number of timepoints
ts = range(tspan..., length=n)
ax.plot(sol(ts, idxs=1)./Rₑ,sol(ts, idxs=2)./Rₑ, sol(ts, idxs=3)./Rₑ,
   label="50 MeV proton, B0 = 20 nT")

ax.legend()
title("Particle trajectory near the Harris current sheet")
xlabel("x [Re]")
ylabel("y [Re]")
zlabel("z [Re]")

x = range(-8, 8, length=20)

X = collect(Float32, x)
Y = zeros(Float32, size(X)...)
Z = zeros(Float32, size(X)...)

B = zeros(Float32, 3, size(X)...)

i = 1
for (x,y) in zip(X,Y)
   B[1+3*(i-1):3*i] = getB_CS_harris([x,0.0,0.0], 4e-2, 1.0)
   global i += 1
end

for s = 1:3
   ax.quiver(X, Y, Z, vec(B[1,:,:]), vec(B[2,:,:]), vec(B[3,:,:]),
      color="k", alpha=0.6)
   @. Y -= 15
end

ax.set_box_aspect([2,5,1])
TestParticle.set_axes_equal(ax)

#savefig("test.png", bbox_inches="tight")


"""
    UniGrid(m::Int)

where `m` is the number of points on each dimension
Adaptive uniform grid based on [1]

[1] Moreno-Muñoz, P., Artés-Rodríguez, A. & Álvarez, M. A. Continual Multi-task Gaussian Processes. (2019).
"""
struct UniGrid <: OnIPSA
    m::Int # Number of points per dimension
end

Base.show(io::IO, Z::UniGrid) = print(io, "Uniform grid with side length $(Z.m).")

function initZ(
    ::AbstractRNG,
    alg::UniGrid,
    X::Union{AbstractVector{<:Real},AbstractVector{<:AbstractVector{<:Real}}};
    kwargs...,
)
    ndim = length(first(X)) # Take the dimensionality
    bounds = [extrema(x -> getindex(x, i), X) for i in 1:ndim]
    Z = map(bounds) do lims
        LinRange(lims..., alg.m)
    end
    return Z
end

function updateZ!(
    ::AbstractRNG, Z::AbstractVector, alg::UniGrid, X::AbstractVector; kwargs...
)
    ndim = length(Z)
    new_bounds = [extrema(x -> getindex(x, i), X) for i in 1:ndim]
    map!(Z, Z, new_bounds) do Z_d, new_b
        x_start = min(Z_d.start, new_b[1]) # Find the new limits
        x_stop = max(Z_d.stop, new_b[2])
        return LinRange(x_start, x_stop, alg.m) # readapt bounds
    end
    return Z
end


module Simpson
    export simpson

    function simpson(n)
        # parameters
        f(x) = (3/2)*sin(x)^3
        a = 0
        b = π
        # calculation
        h = (b-a)/n
        sum = f(a)+f(b)
        for i in 1:2:n
            ai = muladd(h,i,a)
            sum += 4f(ai)
            sum += 2f(ai+h)
        end
        h*sum/3
    end

end


module Day11

using SparseArrays
using ..Utils

neighbors_matrix(n, m) = begin
    bidiag(k) = spdiagm(-1 => ones(Bool, k-1), 1 => ones(Bool, k-1))
    tridiag(k) = spdiagm(-1 => ones(Bool, k-1), 0 => ones(Bool, k), 1 => ones(Bool, k-1))

    kron(bidiag(m), tridiag(n)) .| kron(tridiag(m), bidiag(n))
end

function step!(octopuses::Matrix{Int8})
    n, m = size(octopuses)
    octopuses = vec(octopuses)
    neighbors = neighbors_matrix(n, m)

    # Increase each by 1
    octopuses .+= 1

    # Any which are above 9
    already_flashed = spzeros(Bool, length(octopuses))

    flashing = sparse(octopuses .> 9)
    already_flashed .|= flashing
    while any(flashing)
        octopuses .+= neighbors * flashing
        flashing = sparse(octopuses .> 9) .⊻ already_flashed
        already_flashed .|= flashing
    end
    octopuses[already_flashed] .= 0
    return sum(already_flashed)
end

function part1(octopuses)
    sum = 0
    for i in 1:100
        sum += step!(octopuses)
    end
    return sum
end

function part2(octopuses)
    i = 0
    while true
        i += 1
        count = step!(octopuses)
        if count == length(octopuses)
            return i
        end
    end
end

function run()
    test = """5483143223
2745854711
5264556173
6141336146
6357385478
4167524645
2176841721
6882881134
4846848554
5283751526"""

    test_octopuses = parse_array(test)
    @assert part1(test_octopuses) == 1656
    @assert part2(parse_array(test)) == 195

    input = """4764745784
4643457176
8322628477
7617152546
6137518165
1556723176
2187861886
2553422625
4817584638
3754285662"""

    octopuses = parse_array(input)
    @show part1(octopuses)
    octopuses = parse_array(input)
    @show part2(octopuses)
end

end


using LabelledArrays

export rossler_rossler

"""
    eom_rossler_rossler(u, p, t) → Function

Equations of motions for a 6D dynamical system consisting
of two 3D Rössler attractors.
"""
function eom_rossler_rossler(u, p, t)
    a, b, c, ϵ₁, ϵ₂, ω₁, ω₂ = ([p[i] for i = 1:7]...,)
    x₁, y₁, z₁, x₂, y₂, z₂ = ([u[i] for i = 1:6]...,)
    # First Rössler system
    dx₁ = ω₁*(-y₁) - z₁ + ϵ₂*(x₁ - x₂)
    dy₁ = ω₁*x₁ + a*y₁
    dz₁ = b + z₁*(x₁ - c)

    # Second Rössler system
    dx₂ = ω₂*(-y₂) - z₂ + ϵ₁*(x₂ - x₁)
    dy₂ = ω₂*x₂ + a*y₂
    dz₂ = b + z₂*(x₂ - c)
    return SVector{6}(dx₁, dy₁, dz₁, dx₂, dy₂, dz₂)
end

"""
	rossler_rossler(u₀, a, b, c, ϵ₁, ϵ₂, ω₁, ω₂) → ContinuousDynamicalSystem

A coupled 6D Rossler system from Krakovská et al. (2018).
The system consists of two separate subsystems, each being a 3D Rossler
attractor. The subsystems are bidirectionally coupled, influencing each other
through variables `x₁` and `x₂`.

The parameters and default values (if calling the keyword version of `rrbi`)
are as follows:

```
a = 0.1      # parameter in the Rössler attractor
b = 0.1      # parameter in the Rössler attractor
c = 14.0     # parameter in the Rössler attractor
ϵ₁ = 0.0     # influence of subsystem 1 on subsystem 2
ϵ₂ = 0.0     # influence of subsystem 2 on subsystem 1
ω₁ = 1+0.015 # the frequency of the first system
ω₂ = 1-0.015 # the frequency of the second system
```

The dynamics is generated by the vector field:

```math
\\begin{aligned}
\\dot x₁ = ω₁ (-y₁) - z₁ + ϵ₂(x₁ - x₂)
\\dot y₁ = ω₁ x₁ + ay₁
\\dot z₁ = b + z₁(x₁ - c)

\\dot x₂ = ω₂(-y₂) - z₂ + ϵ₁(x₂ - x₁)
\\dot y₂ = ω₂x₂ + ay₂
\\dot z₂ = b + z₂(x₂ - c)
\\end{aligned}
```

with the coupling constant ``c \\geq 0``.

# References
Krakovská, A., Jakubík, J., Chvosteková, M., Coufal, D., Jajcay, N., & Paluš, M. (2018). Comparison of six methods for the detection of causality in a bivariate time series. Physical Review E, 97(4), 042207.
"""
function rossler_rossler(u₀, a, b, c, ϵ₁, ϵ₂, ω₁, ω₂)
    p = @LArray [a, b, c, ϵ₁, ϵ₂, ω₁, ω₂] (:a, :b, :c, :ϵ₁, :ϵ₂, :ω₁, :ω₂)
    return ContinuousDynamicalSystem(eom_rossler_rossler, u₀, p)
end

"""
	rossler_rossler(;u₀ = rand(6), a = 0.1, b = 0.1, c = 14.0, ϵ₁ = 0.0,
		ϵ₂ = 0.0, ω₁ = 1 + 0.015, ω₂ = 1 - 0.015) → ContinuousDynamicalSystem

A coupled 6D Rossler system from Krakovská et al. (2018).
The system consists of two separate subsystems, each being a 3D Rossler
attractor. The subsystems are bidirectionally coupled, influencing each other
through variables `x₁` and `x₂`.

The parameters and default values (if calling the keyword version of `rrbi`)
are as follows:

```
a = 0.1      # parameter in the Rössler attractor
b = 0.1      # parameter in the Rössler attractor
c = 14.0     # parameter in the Rössler attractor
ϵ₁ = 0.0     # influence of subsystem 1 on subsystem 2
ϵ₂ = 0.0     # influence of subsystem 2 on subsystem 1
ω₁ = 1+0.015 # the frequency of the first system
ω₂ = 1-0.015 # the frequency of the second system
```

The dynamics is generated by the vector field:

```math
\\begin{aligned}
dx₁ = ω₁*(-y₁) - z₁ + ϵ₂*(x₁ - x₂)
dy₁ = ω₁*x₁ + a*y₁
dz₁ = b + z₁*(x₁ - c)

dx₂ = ω₂*(-y₂) - z₂ + ϵ₁*(x₂ - x₁)
dy₂ = ω₂*x₂ + a*y₂
dz₂ = b + z₂*(x₂ - c)
\\end{aligned}
```

with the coupling constant ``c \\geq 0``.

# References
Krakovská, A., Jakubík, J., Chvosteková, M., Coufal, D., Jajcay, N., & Paluš, M. (2018). Comparison of six methods for the detection of causality in a bivariate time series. Physical Review E, 97(4), 042207.
"""
rossler_rossler(;u₀ = rand(6), a = 0.1, b = 0.1, c = 14.0,
            ϵ₁ = 0.0, ϵ₂ = 0.0, ω₁ = 1 + 0.015, ω₂ = 1 - 0.015) =
    rossler_rossler(u₀, a, b, c, ϵ₁, ϵ₂, ω₁, ω₂)


fish = parse.(Int, split(readline("day6/fish.txt"), ','))

function forecast_fish(init, days)
    curr = merge(
        Dict((i, 0) for i in 0:8),
        Dict((i, count(==(i), init)) for i in unique(init))
    )
    for _ in 1:days
        prev = copy(curr)
        for i in 8:-1:1
            curr[i - 1] = prev[i]
        end
        curr[6] += prev[0]
        curr[8] = prev[0]
    end
    curr
end

p1 = sum(values(forecast_fish(fish, 80)))

@show p1

p2 = sum(values(forecast_fish(fish, 256)))

@show p2


# coding: utf-8

# In[1]:

#shaker sort, or cocktail sort, or cocktail shaker sort
#in this script, we will call it cocktail shaker sort
#with so many names, it is merely a variation of bubble sort
#go the link below to check out the details of bubble sort
# https://github.com/je-suis-tm/search-and-sort/blob/master/bubble,%20selection%20and%20insertion%20sort.py
#bubble sort repeatedly takes comparison from one side to the other
#usually left to right, takes n*(n-1) steps, assuming n is the size
#cocktail shaker sort has a similar mechanism
#except it compares its elements from left to right
#then right to left, later left to right
#the process looks like a cocktail shaker (not really)
#that is how the sort gets its name
#in theory, for a partially sorted list
#cocktail shaker sort takes less time than a bubble sort

# In[2]:

function cocktail_shaker_sort(arr)

    swap=true

    #we use left and right to define each round's start and end point
    left=1
    right=length(arr)

    while (left<right) || !(swap)

        swap=false

        #as usual, we start from left to right
        for i in left+1:right

            if arr[i]<arr[i-1]

                arr[i],arr[i-1]=arr[i-1],arr[i]
                swap=true

            end

        end

        right-=1

        #swap is the key to increase the efficiency
        #once there is no swap happened in the loop
        #we can conclude the target list is fully sorted
        if !(swap)
            return arr
        end

        swap=false

        #then right to left
        for i in right:-1:left+1

            if arr[i]<arr[i-1]

                arr[i],arr[i-1]=arr[i-1],arr[i]
                swap=true

            end

        end

        left+=1

    end

    return arr

end

# In[3]:

for _ in 1:100

    test_arr=rand(1000)

    if !(cocktail_shaker_sort(test_arr)==sort(test_arr))
        printstyled("Erreur",color=:red)
    end

end


module Utils

using GraphUtils

export
    interpolate,
    linear_interp,
    prune_sorted_array

interpolate(a,b,t) = (1 - t)*a + t*b

function linear_interp(v,t)
    idx = find_index_in_sorted_array(v,t)-1
    idx = max(1, min(length(v)-1, idx))
    if isapprox(v[idx+1] - v[idx],0.0)
        return idx,0.0
    end
    Δt = (t - v[idx]) / (v[idx+1] - v[idx])
    return idx, Δt
end

"""
    `prune_array(arr,v,t1,t2)`

    prunes `arr` on either side based on where t1 and t2 fall in v.
"""
function prune_sorted_array(v,idx1,idx2,Δt1,Δt2)
    @assert [idx1,Δt1] <= [idx2,Δt2] "[idx1,Δt1] > [idx2,Δt2]"
    vp = [interpolate(v[idx1],v[idx1+1],Δt1)]
    for idx in min(idx1+1,length(v)):min(idx2,length(v))
        push!(vp,v[idx])
    end
    push!(vp,interpolate(v[idx2],v[idx2+1],Δt2))
    return vp
end

end


function readinput(filename)
  input = readlines("input")
  midpoint = findfirst(isequal(""), input)
  p1 = map(n -> parse(Int, n), input[2 : midpoint - 1])
  p2 = map(n -> parse(Int, n), input[midpoint + 2 : end])
  [p1, p2]
end

function combat(decks)
  queues = [copy(deck) for deck=decks]
  while !any(isempty.(queues))
    cards = [popfirst!(q) for q=queues]
    winner = queues[argmax(cards)]
    for card in sort(cards, rev=true)
      push!(winner, card)
    end
  end
  winner = argmax(length.(queues))
  winner, queues[winner]
end

function recursivecombat(decks, game = 1)
  queues = [copy(deck) for deck=decks]
  rounds = Set()
  round = 1
  while !any(isempty.(queues))
    round += 1
    if queues in rounds
      return 1, queues[1]
    end
    push!(rounds, deepcopy(queues))
    cards = [popfirst!(q) for q=queues]
    if all(cards .<= length.(queues))
      recdecks = [ queues[i][1:cards[i]] for i=1:length(queues)]
      windex, _ = recursivecombat(recdecks, game + 1)
    else
      windex = argmax(cards)
    end
    push!(queues[windex], cards[windex])

    for loser in cards[1:end .!=windex]
      push!(queues[windex], loser)
    end
  end
  winner = argmax(length.(queues))
  winner, queues[winner]
end

function deckvalue(deck)
  decklength = length(deck)
  sum([ (decklength - i + 1) * deck[i] for i=1:decklength ])
end

decks = readinput("input")

v1, d1 = combat(decks)
v1value = deckvalue(d1)
println("Part 1: $v1value")

v2, d2 = recursivecombat(decks)
v2value = deckvalue(d2)
println("Part 2: $v2value")


using StochasticDiffEq, DiffEqDevTools, Test

linear = (u,p,t) -> (p*u)
g = (u,p,t) -> zero(u)
linear_analytic = (u0,p,t,W) -> u0*exp(p*t)
prob = SDEProblem(SDEFunction(linear,g,analytic=linear_analytic),g,
                  1/2,(0.0,1.0),1.01)

dts = (1/2) .^ (7:-1:4) #14->7 good plot
sol = solve(prob,SKenCarp())
sim2 = test_convergence(dts,prob,SKenCarp(),trajectories=20)
@test abs(sim2.𝒪est[:l∞]-3) <.1 #High tolerance since low dts for testing!

linear = (du,u,p,t) -> (du.=p.*u)
g = (du,u,p,t) -> (du.=0)
linear_analytic = (u0,p,t,W) -> u0*exp.(p.*t)
prob = SDEProblem(SDEFunction(linear,g,analytic=linear_analytic),g,
                  rand(4,2),(0.0,1.0),1.01)

dts = (1/2) .^ (7:-1:4) #14->7 good plot
sol = solve(prob,SKenCarp())
sim2 = test_convergence(dts,prob,SKenCarp(),trajectories=20)
@test abs(sim2.𝒪est[:l∞]-3) <.1 #High tolerance since low dts for testing!


using Oceananigans.Solvers: solve_for_pressure!, solve_poisson_equation!, set_source_term!
using Oceananigans.Solvers: poisson_eigenvalues
using Oceananigans.Models.HydrostaticFreeSurfaceModels: _compute_w_from_continuity!
using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions

function poisson_solver_instantiates(arch, grid, planner_flag)
    solver = FFTBasedPoissonSolver(arch, grid, planner_flag)
    return true  # Just making sure the FFTBasedPoissonSolver does not error/crash.
end

function random_divergent_source_term(arch, grid)
    default_bcs = FieldBoundaryConditions()
    u_bcs = regularize_field_boundary_conditions(default_bcs, grid, :u)
    v_bcs = regularize_field_boundary_conditions(default_bcs, grid, :v)
    w_bcs = regularize_field_boundary_conditions(default_bcs, grid, :w)

    Ru = CenterField(arch, grid, u_bcs)
    Rv = CenterField(arch, grid, v_bcs)
    Rw = CenterField(arch, grid, w_bcs)
    U = (u=Ru, v=Rv, w=Rw)

    Nx, Ny, Nz = size(grid)
    set!(Ru, rand(Nx, Ny, Nz))
    set!(Rv, rand(Nx, Ny, Nz))
    set!(Rw, rand(Nx, Ny, Nz))

    fill_halo_regions!(Ru, arch)
    fill_halo_regions!(Rv, arch)
    fill_halo_regions!(Rw, arch)

    # Compute the right hand side R = ∇⋅U
    ArrayType = array_type(arch)
    R = zeros(Nx, Ny, Nz) |> ArrayType
    event = launch!(arch, grid, :xyz, divergence!, grid, U.u.data, U.v.data, U.w.data, R,
                    dependencies=Event(device(arch)))
    wait(device(arch), event)

    return R, U
end

function random_divergence_free_source_term(arch, grid)
    default_bcs = FieldBoundaryConditions()
    u_bcs = regularize_field_boundary_conditions(default_bcs, grid, :u)
    v_bcs = regularize_field_boundary_conditions(default_bcs, grid, :v)
    w_bcs = regularize_field_boundary_conditions(default_bcs, grid, :w)

    # Random right hand side
    Ru = CenterField(arch, grid, u_bcs)
    Rv = CenterField(arch, grid, v_bcs)
    Rw = CenterField(arch, grid, w_bcs)
    U = (u=Ru, v=Rv, w=Rw)

    Nx, Ny, Nz = size(grid)
    set!(Ru, rand(Nx, Ny, Nz))
    set!(Rv, rand(Nx, Ny, Nz))
    set!(Rw, zeros(Nx, Ny, Nz))

    fill_halo_regions!(Ru, arch, nothing, nothing)
    fill_halo_regions!(Rv, arch, nothing, nothing)
    fill_halo_regions!(Rw, arch, nothing, nothing)

    event = launch!(arch, grid, :xy, _compute_w_from_continuity!, U, grid,
                    dependencies=Event(device(arch)))
    wait(device(arch), event)

    fill_halo_regions!(Rw, arch, nothing, nothing)

    # Compute the right hand side R = ∇⋅U
    ArrayType = array_type(arch)
    R = zeros(Nx, Ny, Nz) |> ArrayType
    event = launch!(arch, grid, :xyz, divergence!, grid, Ru.data, Rv.data, Rw.data, R,
                    dependencies=Event(device(arch)))
    wait(device(arch), event)

    return R
end

#####
##### Regular rectilinear grid Poisson solver
#####

function divergence_free_poisson_solution(arch, grid, planner_flag=FFTW.MEASURE)
    ArrayType = array_type(arch)
    FT = eltype(grid)

    solver = FFTBasedPoissonSolver(arch, grid, planner_flag)
    R, U = random_divergent_source_term(arch, grid)

    p_bcs = FieldBoundaryConditions(grid, (Center, Center, Center))
    ϕ   = CenterField(arch, grid, p_bcs)  # "kinematic pressure"
    ∇²ϕ = CenterField(arch, grid, p_bcs)

    # Using Δt = 1 but it doesn't matter since velocities = 0.
    solve_for_pressure!(ϕ.data, solver, arch, grid, 1, datatuple(U))

    compute_∇²!(∇²ϕ, ϕ, arch, grid)

    return CUDA.@allowscalar interior(∇²ϕ) ≈ R
end

#####
##### Test that Poisson solver error converges as error ~ N⁻²
#####

ψ(::Type{Bounded}, n, x) = cos(n*x/2)
ψ(::Type{Periodic}, n, x) = cos(n*x)

k²(::Type{Bounded}, n) = (n/2)^2
k²(::Type{Periodic}, n) = n^2

function analytical_poisson_solver_test(arch, N, topo; FT=Float64, mode=1)
    grid = RegularRectilinearGrid(FT, topology=topo, size=(N, N, N), x=(0, 2π), y=(0, 2π), z=(0, 2π))
    solver = FFTBasedPoissonSolver(arch, grid)

    xC, yC, zC = nodes((Center, Center, Center), grid, reshape=true)

    TX, TY, TZ = topology(grid)
    Ψ(x, y, z) = ψ(TX, mode, x) * ψ(TY, mode, y) * ψ(TZ, mode, z)
    f(x, y, z) = -(k²(TX, mode) + k²(TY, mode) + k²(TZ, mode)) * Ψ(x, y, z)

    solver.storage .= convert(array_type(arch), f.(xC, yC, zC))

    solve_poisson_equation!(solver)

    ϕ = real(Array(solver.storage))

    L¹_error = mean(abs, ϕ - Ψ.(xC, yC, zC))

    return L¹_error
end

function poisson_solver_convergence(arch, topo, N¹, N²; FT=Float64, mode=1)
    error¹ = analytical_poisson_solver_test(arch, N¹, topo; FT, mode)
    error² = analytical_poisson_solver_test(arch, N², topo; FT, mode)

    rate = log(error¹ / error²) / log(N² / N¹)

    TX, TY, TZ = topo
    @info "Convergence of L¹-normed error, $(typeof(arch)), $FT, ($(N¹)³ -> $(N²)³), topology=($TX, $TY, $TZ): $rate"

    return isapprox(rate, 2, rtol=5e-3)
end

#####
##### Vertically stretched Poisson solver
#####

get_grid_size(TX, TY, TZ, Nx, Ny, Nz) = (Nx, Ny, Nz)
get_grid_size(::Type{Flat}, TY, TZ, Nx, Ny, Nz) = (Ny, Nz)
get_grid_size(TX, ::Type{Flat}, TZ, Nx, Ny, Nz) = (Nx, Nz)

get_xy_interval_kwargs(TX, TY, TZ) = (x=(0, 1), y=(0, 1))
get_xy_interval_kwargs(TX, ::Type{Flat}, TZ) = (x=(0, 1),)
get_xy_interval_kwargs(::Type{Flat}, TY, TZ) = (y=(0, 1),)

function vertically_stretched_poisson_solver_correct_answer(FT, arch, topo, Nx, Ny, zF)
    Nz = length(zF) - 1
    sz = get_grid_size(topo..., Nx, Ny, Nz)
    xy_intervals = get_xy_interval_kwargs(topo...)
    vs_grid = VerticallyStretchedRectilinearGrid(FT; architecture=arch, topology=topo, size=sz, z_faces=zF, xy_intervals...)
    solver = FourierTridiagonalPoissonSolver(arch, vs_grid)

    p_bcs = FieldBoundaryConditions(vs_grid, (Center, Center, Center))
    ϕ   = CenterField(arch, vs_grid, p_bcs)  # "kinematic pressure"
    ∇²ϕ = CenterField(arch, vs_grid, p_bcs)

    R = random_divergence_free_source_term(arch, vs_grid)

    set_source_term!(solver, R)
    solve_poisson_equation!(solver)

    # interior(ϕ) = solution(solver) or solution!(interior(ϕ), solver)
    CUDA.@allowscalar interior(ϕ) .= real.(solver.storage)
    compute_∇²!(∇²ϕ, ϕ, arch, vs_grid)

    return CUDA.@allowscalar interior(∇²ϕ) ≈ R
end

#####
##### Run pressure solver tests
#####

PB = (Periodic, Bounded)
topos = collect(Iterators.product(PB, PB, PB))[:]

two_dimensional_topologies = [(Flat,     Bounded,  Bounded),
                              (Bounded,  Flat,     Bounded),
                              (Bounded,  Bounded,  Flat),
                              (Flat,     Periodic, Bounded),
                              (Periodic, Flat,     Bounded),
                              (Periodic, Bounded,  Flat)]

@testset "Poisson solvers" begin
    @info "Testing Poisson solvers..."

    for arch in archs
        @testset "Poisson solver instantiation [$(typeof(arch))]" begin
            @info "  Testing Poisson solver instantiation [$(typeof(arch))]..."
            for FT in float_types

                grids_3d = [RegularRectilinearGrid(FT, size=(2, 2, 2), extent=(1, 1, 1)),
                            RegularRectilinearGrid(FT, size=(1, 2, 2), extent=(1, 1, 1)),
                            RegularRectilinearGrid(FT, size=(2, 1, 2), extent=(1, 1, 1)),
                            RegularRectilinearGrid(FT, size=(2, 2, 1), extent=(1, 1, 1))]

                grids_2d = [RegularRectilinearGrid(FT, size=(2, 2), extent=(1, 1), topology=topo)
                            for topo in two_dimensional_topologies]

                grids = []
                push!(grids, grids_3d..., grids_2d...)

                for grid in grids
                    @test poisson_solver_instantiates(arch, grid, FFTW.ESTIMATE)
                    @test poisson_solver_instantiates(arch, grid, FFTW.MEASURE)
                end
            end
        end

        @testset "Divergence-free solution [$(typeof(arch))]" begin
            @info "  Testing divergence-free solution [$(typeof(arch))]..."

            for topo in topos
                for N in [7, 16]

                    grids_3d = [RegularRectilinearGrid(topology=topo, size=(N, N, N), extent=(1, 1, 1)),
                                RegularRectilinearGrid(topology=topo, size=(1, N, N), extent=(1, 1, 1)),
                                RegularRectilinearGrid(topology=topo, size=(N, 1, N), extent=(1, 1, 1)),
                                RegularRectilinearGrid(topology=topo, size=(N, N, 1), extent=(1, 1, 1))]

                    grids_2d = [RegularRectilinearGrid(size=(N, N), extent=(1, 1), topology=topo)
                                for topo in two_dimensional_topologies]

                    grids = []
                    push!(grids, grids_3d..., grids_2d...)

                    for grid in grids
                        N == 7 && @info "    Testing $(topology(grid)) topology on square grids [$(typeof(arch))]..."
                        @test divergence_free_poisson_solution(arch, grid)
                    end
                end
            end

            Ns = [11, 16]
            for topo in topos
                @info "    Testing $topo topology on rectangular grids with even and prime sizes [$(typeof(arch))]..."
                for Nx in Ns, Ny in Ns, Nz in Ns
                    grid = RegularRectilinearGrid(topology=topo, size=(Nx, Ny, Nz), extent=(1, 1, 1))
                    @test divergence_free_poisson_solution(arch, grid)
                end
            end

            # Do a couple at Float32 (since its too expensive to repeat all tests...)
            Float32_grids = [RegularRectilinearGrid(Float32, topology=(Periodic, Bounded, Bounded), size=(16, 16, 16), extent=(1, 1, 1)),
                             RegularRectilinearGrid(Float32, topology=(Bounded, Bounded, Periodic), size=(7, 11, 13), extent=(1, 1, 1))]

            for grid in Float32_grids
                @test divergence_free_poisson_solution(arch, grid)
            end
        end

        @testset "Convergence to analytic solution [$(typeof(arch))]" begin
            @info "  Testing convergence to analytic solution [$(typeof(arch))]..."
            for topo in topos
                @test poisson_solver_convergence(arch, topo, 2^6, 2^7)
                @test poisson_solver_convergence(arch, topo, 67, 131, mode=2)
            end
        end
    end

    # Vertically stretched topologies to test.
    vs_topos = [
        (Periodic, Periodic, Bounded),
        (Periodic, Bounded,  Bounded),
        (Bounded,  Periodic, Bounded),
        (Bounded,  Bounded,  Bounded),
        (Flat,     Bounded,  Bounded),
        (Flat,     Periodic, Bounded),
        (Bounded,  Flat,     Bounded),
        (Periodic, Flat,     Bounded)
    ]

    for arch in archs, topo in vs_topos
        @testset "Vertically stretched Poisson solver [FACR, $(typeof(arch)), $topo]" begin
            @info "  Testing vertically stretched Poisson solver [FACR, $(typeof(arch)), $topo]..."

            @test vertically_stretched_poisson_solver_correct_answer(Float64, arch, topo, 8, 8, 1:8)
            @test vertically_stretched_poisson_solver_correct_answer(Float64, arch, topo, 7, 7, 1:7)
            @test vertically_stretched_poisson_solver_correct_answer(Float32, arch, topo, 8, 8, 1:8)

            zF_even = [1, 2, 4, 7, 11, 16, 22, 29, 37]      # Nz = 8
            zF_odd  = [1, 2, 4, 7, 11, 16, 22, 29, 37, 51]  # Nz = 9

            for zF in [zF_even, zF_odd]
                @test vertically_stretched_poisson_solver_correct_answer(Float64, arch, topo, 8,  8, zF)
                @test vertically_stretched_poisson_solver_correct_answer(Float64, arch, topo, 16, 8, zF)
                @test vertically_stretched_poisson_solver_correct_answer(Float64, arch, topo, 8, 16, zF)
                @test vertically_stretched_poisson_solver_correct_answer(Float64, arch, topo, 8, 11, zF)
                @test vertically_stretched_poisson_solver_correct_answer(Float64, arch, topo, 5,  8, zF)
                @test vertically_stretched_poisson_solver_correct_answer(Float64, arch, topo, 7, 13, zF)
            end
        end
    end
end


struct MatrixAngularCentralGaussian{N,T,n,k} <: ParameterizedMeasure{N}
    par::NamedTuple{N,T}
    manifold::Stiefel{n,k,ℝ}
end

function MatrixAngularCentralGaussian(n, k; kwargs...)
    par = NamedTuple(kwargs)
    return MatrixAngularCentralGaussian(par, Stiefel(n, k))
end

Manifolds.base_manifold(d::MatrixAngularCentralGaussian) = getfield(d, :manifold)

function MeasureTheory.basemeasure(μ::MatrixAngularCentralGaussian{(:Σ⁻¹,)})
    return normalize(Hausdorff(base_manifold(μ), μ.Σ⁻¹))
end
function MeasureTheory.basemeasure(μ::MatrixAngularCentralGaussian{(:L,)})
    return normalize(Hausdorff(base_manifold(μ), μ.L))
end

# Chikuse, 2003 Eq. 2.4.3
function MeasureTheory.logdensity(
    d::MatrixAngularCentralGaussian{(:Σ⁻¹,)}, x::AbstractMatrix
)
    n, k = representation_size(base_manifold(d))
    Σ⁻¹ = d.Σ⁻¹
    return -n//2 * log(dot(x, Σ⁻¹, x)) + k//2 * logdet(Σ⁻¹)
end
function MeasureTheory.logdensity(d::MatrixAngularCentralGaussian{(:L,)}, x::AbstractMatrix)
    n, k = representation_size(base_manifold(d))
    L = d.L
    z = LowerTriangular(L) \ x
    return -n//2 * log(dot(z, z)) - k * logdet(L)
end

function Base.rand(rng::AbstractRNG, ::Type, d::MatrixAngularCentralGaussian{(:L,)})
    L = d.L
    z = randn!(rng, similar(d.L, representation_size(base_manifold(d))))
    y = lmul!(LowerTriangular(L), z)
    return y * (Symmetric(y'y))^(-1//2)
end


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ a31e1a3a-f637-11ea-1e0f-2d61b88049db
using Suppressor, InteractiveUtils

# ╔═╡ 80b2135e-f1ed-11ea-17ee-656b8062026a
md"""
# What are macros?

**Macros** change the syntax of parts of a program.

While functions work on values, a macro works on a **piece of Julia code**, which is represented as a data structure in Julia itself, and **rewrites it to a different piece of code**! This can be a very powerful thing to do.
"""

# ╔═╡ 7acec72a-f1fa-11ea-2753-efcf2cee8335
md"""As a simple example, let's look at the macro called `@elapsed`, which measures how much time it takes to run a piece of code.

It does this by taking the input expression and adding extra code around it.

First let's see it in action! We'll use a function that takes some time to run: the function `peakflops()` estimates how fast your computer runs. (It reports an estimate of the number of flops, i.e. floating-point operations per second, when doing matrix multiplications.)
"""

# ╔═╡ 0f7c908c-f1fb-11ea-18fb-c7e70278ec46
peakflops()

# ╔═╡ f43139a2-f1ed-11ea-0cdf-678ffac4374a
@elapsed peakflops()

# ╔═╡ 67b3fb78-f1ee-11ea-2d0a-a92e81580917
md"""
**NOTE:** a macro is always called using the `@` prefix character; this is a visual cue that the following piece of code is *not what is actually executed*, but rather is *transformed in some way*.
"""

# ╔═╡ 5677f2d6-f1ee-11ea-3f9b-b59c8a8d07ff
md"""
You can look at the transformed expression by using `@macroexpand`. Since this name begins with `@`, we can see that this is actually itself a macro!
"""

# ╔═╡ 4b2cbfc4-f1ee-11ea-17d6-a72fd22e6d89
@macroexpand @elapsed peakflops()

# ╔═╡ c4e4e402-f1fa-11ea-3e5e-4344dcc7e17b
md"Let's remove the line number information:"

# ╔═╡ 9639e822-f1ee-11ea-19b6-5f377bb4b809
Base.remove_linenums!( @macroexpand @elapsed peakflops() )

# ╔═╡ ea65d68c-f1fa-11ea-1cde-ebf4b7958b8a
md"We see that the macro has added more code around the original `peakflops()` call.
Julia *replaces* the original code by the new code; this new code is what will actually be used (compiled)."

# ╔═╡ a5816132-f1ee-11ea-268b-05086212324b
md"""
Here, the `remove_linenums!` function actually receives a Julia **expression** as its argument.
"""

# ╔═╡ e3bd8fbc-f1fa-11ea-331f-bbe8f038f46a
md"## Expressions in Julia"

# ╔═╡ cb5e6a8a-f1fb-11ea-2f15-e7486afdfa4a
md"One feature that makes Julia powerful is its ability (influenced by the Lisp language) to represent Julia code as Julia objects, which can be manipulated using the usual Julia features.
"

# ╔═╡ faa18946-f1fb-11ea-0f5b-290683f29ebf
md"Let's look at some simple expressions.

If we type an expression, Julia will **evaluate** it and return its value:
"

# ╔═╡ c5465fe0-f1ee-11ea-2f61-3172509f693d
x = 1 + 2

# ╔═╡ 6bb8d276-f1ff-11ea-2ada-1916761bc6a8
md"The variable `x` now contains the value `3`, the result of evaluating the expression `1 + 2`."

# ╔═╡ 10d909b4-f1fc-11ea-080a-f9f87b3512c5
md"We can prevent Julia from evaluating an expression by **quoting** it, which we can do in one of two ways:"

# ╔═╡ c88dfe88-f1ee-11ea-0327-8b18a25c511c
expr = :(1 + 2)

# ╔═╡ 80ec2c30-f1ff-11ea-30ea-ef5dd01e26ae
expr2 = quote
			1 + 2
		end

# ╔═╡ 60778b00-f1ff-11ea-011f-231f3ab18de1
md"This *leaves the expression alone, without evaluating it*, giving as result the unevaluated expression. This is a Julia object of type `Expr` that we have not seen before:"

# ╔═╡ bdb08e52-f1ff-11ea-1881-4bb2bb3f5c77
typeof(expr)

# ╔═╡ 6a0fdfc8-f637-11ea-04d8-a55a4c42df3e
md"As an example, we can create the following expression:"

# ╔═╡ cbce40ee-f1ee-11ea-2e05-1d523acbf8b9
expr3 = :(y + 1)

# ╔═╡ 7f1cad88-f637-11ea-1cf2-4d5b9a4fb50f
md"If we evaluate this expression, it gives an error, since there is no variable `y` in the current notebook:"

# ╔═╡ 906e3234-f637-11ea-1f77-679dc897ec69
y + 1

# ╔═╡ 940f888c-f637-11ea-20a0-1ddda8fb7908
md"Since expressions are Julia objects, we can examine them using Julia's usual mechanisms. For example, the `dump` function allows us to inspect the fields inside an object:"

# ╔═╡ 8dab58d4-f639-11ea-1b0b-536927639c6c
md"We can extract the fields inside the object:"

# ╔═╡ e46e6c96-f1ee-11ea-3e7f-4f613282ccb1
expr3.head

# ╔═╡ e631bede-f1ee-11ea-0fab-7516ab535cf5
expr3.args

# ╔═╡ 95cf0d76-f639-11ea-34dc-f32852ac32b8
md"and even modify them:"

# ╔═╡ 9a587ba4-f639-11ea-2bef-37fed4d45a7f
expr3.args[2]

# ╔═╡ 9cd4334e-f639-11ea-26c6-e9782451ef02
expr3.args[2] = :x

# ╔═╡ a037391e-f639-11ea-2723-455ccff4b024
md"(`:x` and `:y` here are `Symbol`s, a special type of string.)"

# ╔═╡ b6e67c1a-f639-11ea-3abe-f75671ed4ca2
expr3

# ╔═╡ cb0cb434-f639-11ea-1ea2-93009c372885
md"We see that we have been able to *modify the `Expr` object, which is a piece of Julia code, using Julia itself! This is the simplest example of **metaprogramming**, i.e. writing a program that modifies another *program* (instead of modifying data).

This is what macros do internally: transform Julia expressions into something else.
"

# ╔═╡ 0ff6e154-f1ef-11ea-2ee8-61eb98523642
md"""
## Conclusion

When we see something like `@foo`, we recognise that it's a **syntactic transformation** and may modify the behavior of the language in some way!

Being able to represent Julia code within Julia is very useful, and hence we frequently see macros and expression manipulation being used in Julia code.
"""

# ╔═╡ aa3a731c-f637-11ea-3afc-9144bdeb7d6f
md"## Function library"

# ╔═╡ b440b66c-f637-11ea-393d-9d444aa58501
function with_terminal(f)
	local spam_out, spam_err
	@color_output false begin
		spam_out = @capture_out begin
			spam_err = @capture_err begin
				f()
			end
		end
	end
	spam_out, spam_err

	HTML("""

$(Markdown.htmlesc(spam_out))

	""")
end

# ╔═╡ 77437eaa-f639-11ea-0c4d-3756cbd589da
with_terminal() do
	dump(expr3)
end

# ╔═╡ Cell order:
# ╟─80b2135e-f1ed-11ea-17ee-656b8062026a
# ╟─7acec72a-f1fa-11ea-2753-efcf2cee8335
# ╠═0f7c908c-f1fb-11ea-18fb-c7e70278ec46
# ╠═f43139a2-f1ed-11ea-0cdf-678ffac4374a
# ╟─67b3fb78-f1ee-11ea-2d0a-a92e81580917
# ╟─5677f2d6-f1ee-11ea-3f9b-b59c8a8d07ff
# ╠═4b2cbfc4-f1ee-11ea-17d6-a72fd22e6d89
# ╟─c4e4e402-f1fa-11ea-3e5e-4344dcc7e17b
# ╠═9639e822-f1ee-11ea-19b6-5f377bb4b809
# ╟─ea65d68c-f1fa-11ea-1cde-ebf4b7958b8a
# ╟─a5816132-f1ee-11ea-268b-05086212324b
# ╟─e3bd8fbc-f1fa-11ea-331f-bbe8f038f46a
# ╟─cb5e6a8a-f1fb-11ea-2f15-e7486afdfa4a
# ╟─faa18946-f1fb-11ea-0f5b-290683f29ebf
# ╠═c5465fe0-f1ee-11ea-2f61-3172509f693d
# ╟─6bb8d276-f1ff-11ea-2ada-1916761bc6a8
# ╟─10d909b4-f1fc-11ea-080a-f9f87b3512c5
# ╠═c88dfe88-f1ee-11ea-0327-8b18a25c511c
# ╠═80ec2c30-f1ff-11ea-30ea-ef5dd01e26ae
# ╟─60778b00-f1ff-11ea-011f-231f3ab18de1
# ╠═bdb08e52-f1ff-11ea-1881-4bb2bb3f5c77
# ╟─6a0fdfc8-f637-11ea-04d8-a55a4c42df3e
# ╠═cbce40ee-f1ee-11ea-2e05-1d523acbf8b9
# ╟─7f1cad88-f637-11ea-1cf2-4d5b9a4fb50f
# ╠═906e3234-f637-11ea-1f77-679dc897ec69
# ╟─940f888c-f637-11ea-20a0-1ddda8fb7908
# ╠═77437eaa-f639-11ea-0c4d-3756cbd589da
# ╠═8dab58d4-f639-11ea-1b0b-536927639c6c
# ╠═e46e6c96-f1ee-11ea-3e7f-4f613282ccb1
# ╠═e631bede-f1ee-11ea-0fab-7516ab535cf5
# ╠═95cf0d76-f639-11ea-34dc-f32852ac32b8
# ╠═9a587ba4-f639-11ea-2bef-37fed4d45a7f
# ╠═9cd4334e-f639-11ea-26c6-e9782451ef02
# ╠═a037391e-f639-11ea-2723-455ccff4b024
# ╠═b6e67c1a-f639-11ea-3abe-f75671ed4ca2
# ╟─cb0cb434-f639-11ea-1ea2-93009c372885
# ╟─0ff6e154-f1ef-11ea-2ee8-61eb98523642
# ╟─aa3a731c-f637-11ea-3afc-9144bdeb7d6f
# ╠═a31e1a3a-f637-11ea-1e0f-2d61b88049db
# ╠═b440b66c-f637-11ea-393d-9d444aa58501


using Catalyst, Test
using LinearAlgebra

include("test_networks.jl")

rn = @reaction_network begin
    (1, 2), A + B <--> C
    (3, 2), D <--> E
    (0.1, 0.2), E <--> F
    6, F --> G
    7, H --> G
    5, 0 --> K
    3, A + B --> Z + C
end

S = netstoichmat(rn)
C = conservationlaws(S)
@test size(C,1) == 3
b = [ 0, 0, 0, 1, 1, 1, 1, 1, 0, 0 ]
@test any(b == C[i,:] for i in 1:size(C,1))

# For the A + B <--> C subsystem one of these must occur
# as a conservation law
D = [ 1 -1 0 0 0 0 0 0 0 0;
      -1 1 0 0 0 0 0 0 0 0
      0  1 1 0 0 0 0 0 0 0 ]
@test any(D[j,:] == C[i,:] for i in 1:size(C,1), j in 1:size(D,1))

C = conservationlaws(rn)
@test size(C,1) == 3
@test Catalyst.get_networkproperties(rn).nullity == 3
@test any(b == C[i,:] for i in 1:size(C,1))
@test any(D[j,:] == C[i,:] for i in 1:size(C,1), j in 1:size(D,1))

Cs_standard = map(conservationlaws, reaction_networks_standard)
@test all(size(C, 1) == 0 for C in Cs_standard)

Cs_hill = map(conservationlaws, reaction_networks_hill)
@test all(size(C, 1) == 0 for C in Cs_hill)

function consequiv(A, B)
    rank([A; B]) == rank(A) == rank(B)
end
Cs_constraint = map(conservationlaws, reaction_networks_constraint)
@test all(consequiv.(Matrix{Int}.(Cs_constraint), reaction_network_constraints))


include("utils.jl")

@testset "Rotations" begin
    M = Manifolds.Rotations(2)
    @test repr(M) == "Rotations(2)"
    @test representation_size(M) == (2, 2)
    @test injectivity_radius(M) == π * sqrt(2.0)
    @test injectivity_radius(M, [1.0 0.0; 0.0 1.0]) == π * sqrt(2.0)
    @test injectivity_radius(M, ExponentialRetraction()) == π * sqrt(2.0)
    @test injectivity_radius(M, [1.0 0.0; 0.0 1.0], ExponentialRetraction()) ==
          π * sqrt(2.0)
    @test injectivity_radius(M, PolarRetraction()) ≈ π / sqrt(2)
    @test injectivity_radius(M, [1.0 0.0; 0.0 1.0], PolarRetraction()) ≈ π / sqrt(2)
    types = [Matrix{Float64}]
    TEST_FLOAT32 && push!(types, Matrix{Float32})
    TEST_STATIC_SIZED && push!(types, MMatrix{2,2,Float64,4})
    retraction_methods = [Manifolds.PolarRetraction(), Manifolds.QRRetraction()]

    inverse_retraction_methods =
        [Manifolds.PolarInverseRetraction(), Manifolds.QRInverseRetraction()]

    basis_types = (DefaultOrthonormalBasis(), ProjectedOrthonormalBasis(:svd))

    @testset "vee/hat" begin
        M = Manifolds.Rotations(2)
        v = [1.23]
        x = Matrix{Float64}(I, 2, 2)
        V = Manifolds.hat(M, x, v)
        @test isa(V, AbstractMatrix)
        @test norm(M, x, V) / sqrt(2) ≈ norm(v)
        @test Manifolds.vee(M, x, V) == v

        V = project(M, x, randn(2, 2))
        v = Manifolds.vee(M, x, V)
        @test isa(v, AbstractVector)
        @test Manifolds.hat(M, x, v) == V
    end

    for T in types
        angles = (0.0, π / 2, 2π / 3, π / 4)
        pts = [convert(T, [cos(ϕ) -sin(ϕ); sin(ϕ) cos(ϕ)]) for ϕ in angles]
        test_manifold(
            M,
            pts;
            test_reverse_diff=false,
            test_injectivity_radius=false,
            test_project_tangent=true,
            test_musical_isomorphisms=true,
            retraction_methods=retraction_methods,
            inverse_retraction_methods=inverse_retraction_methods,
            point_distributions=[Manifolds.normal_rotation_distribution(M, pts[1], 1.0)],
            tvector_distributions=[Manifolds.normal_tvector_distribution(M, pts[1], 1.0)],
            basis_types_to_from=basis_types,
        )

        @testset "log edge cases" begin
            v = Manifolds.hat(M, pts[1], [Float64(π)])
            x = exp(M, pts[1], v)
            @test isapprox(x, exp(M, pts[1], log(M, pts[1], x)))
        end

        v = log(M, pts[1], pts[2])
        @test norm(M, pts[1], v) ≈ (angles[2] - angles[1]) * sqrt(2)

        v14_polar = inverse_retract(M, pts[1], pts[4], Manifolds.PolarInverseRetraction())
        p4_polar = retract(M, pts[1], v14_polar, Manifolds.PolarRetraction())
        @test isapprox(M, pts[4], p4_polar)

        v14_qr = inverse_retract(M, pts[1], pts[4], Manifolds.QRInverseRetraction())
        p4_qr = retract(M, pts[1], v14_qr, Manifolds.QRRetraction())
        @test isapprox(M, pts[4], p4_qr)
    end

    @testset "Distribution tests" begin
        usd_mmatrix =
            Manifolds.normal_rotation_distribution(M, (@MMatrix [1.0 0.0; 0.0 1.0]), 1.0)
        @test isa(rand(usd_mmatrix), MMatrix)

        usd1_mmatrix =
            Manifolds.normal_rotation_distribution(Rotations(1), (@MMatrix [1.0]), 1.0)
        @test isa(rand(usd1_mmatrix), MMatrix)
        @test rand(usd1_mmatrix) == @MMatrix [1.0]

        gtsd_mvector =
            Manifolds.normal_tvector_distribution(M, (@MMatrix [1.0 0.0; 0.0 1.0]), 1.0)
        @test isa(rand(gtsd_mvector), MMatrix)
    end

    Random.seed!(42)
    for n in (3, 4, 5)
        @testset "Rotations: SO($n)" begin
            SOn = Manifolds.Rotations(n)
            ptd = Manifolds.normal_rotation_distribution(SOn, Matrix(1.0I, n, n), 1.0)
            tvd = Manifolds.normal_tvector_distribution(SOn, Matrix(1.0I, n, n), 1.0)
            pts = [rand(ptd) for _ in 1:3]
            test_manifold(
                SOn,
                pts;
                test_forward_diff=n == 3,
                test_reverse_diff=false,
                test_injectivity_radius=false,
                test_musical_isomorphisms=true,
                test_mutating_rand=true,
                retraction_methods=retraction_methods,
                inverse_retraction_methods=inverse_retraction_methods,
                point_distributions=[ptd],
                tvector_distributions=[tvd],
                basis_types_to_from=basis_types,
                exp_log_atol_multiplier=20,
                retraction_atol_multiplier=12,
            )

            @testset "vee/hat" begin
                x = Matrix(1.0I, n, n)
                v = randn(manifold_dimension(SOn))
                V = Manifolds.hat(SOn, x, v)
                @test isa(V, AbstractMatrix)
                @test norm(SOn, x, V) / sqrt(2) ≈ norm(v)
                @test Manifolds.vee(SOn, x, V) == v

                V = project(SOn, x, randn(n, n))
                v = Manifolds.vee(SOn, x, V)
                @test isa(v, AbstractVector)
                @test Manifolds.hat(SOn, x, v) ≈ V
            end

            if n == 4
                @testset "exp/log edge cases" begin
                    vs = [
                        [0, 0, π, 0, 0, π],  # θ = (π, π)
                        [0, 0, π, 0, 0, 0],  # θ = (π, 0)
                        [0, 0, π / 2, 0, 0, π],  # θ = (π, π/2)
                        [0, 0, π, 0, 0, 0] ./ 2,  # θ = (π/2, 0)
                        [0, 0, π, 0, 0, π] ./ 2,  # θ = (π/2, π/2)
                        [0, 0, 0, 0, 0, 0],  # θ = (0, 0)
                        [0, 0, 1, 0, 0, 1] .* 1e-100, # α = β ≈ 0
                        [0, 0, 1, 0, 0, 1] .* 1e-6, # α = β ⩰ 0
                        [0, 0, 10, 0, 0, 1] .* 1e-6, # α ⪆ β ⩰ 0
                        [0, 0, π / 4, 0, 0, π / 4 - 1e-6], # α ⪆ β > 0
                    ]
                    for v in vs
                        @testset "rotation vector $v" begin
                            V = Manifolds.hat(SOn, Matrix(1.0I, n, n), v)
                            x = exp(V)
                            @test x ≈ exp(SOn, one(x), V)
                            @test ForwardDiff.derivative(t -> exp(SOn, one(x), t * V), 0) ≈
                                  V
                            x2 = exp(log(SOn, one(x), x))
                            @test isapprox(x, x2; atol=1e-6)
                        end
                    end
                end
            end

            v = Matrix(
                Manifolds.hat(SOn, pts[1], π * normalize(randn(manifold_dimension(SOn)))),
            )
            x = exp(SOn, pts[1], v)
            v2 = log(SOn, pts[1], x)
            @test x ≈ exp(SOn, pts[1], v2)
        end
    end
    @testset "Test Manifold Point and Tangent Vector checks" begin
        M = Manifolds.Rotations(2)
        for x in [1, [2.0 0.0; 0.0 1.0], [1.0 0.5; 0.0 1.0]]
            @test_throws DomainError is_manifold_point(M, x, true)
            @test !is_manifold_point(M, x)
        end
        x = one(zeros(2, 2))
        @test is_manifold_point(M, x)
        @test is_manifold_point(M, x, true)
        for v in [1, [0.0 1.0; 0.0 0.0]]
            @test_throws DomainError is_tangent_vector(M, x, v, true)
            @test !is_tangent_vector(M, x, v)
        end
        v = [0.0 1.0; -1.0 0.0]
        @test is_tangent_vector(M, x, v)
        @test is_tangent_vector(M, x, v, true)
    end
    @testset "Project point" begin
        M = Manifolds.Rotations(2)
        x = Matrix{Float64}(I, 2, 2)
        x1 = project(M, x)
        @test is_manifold_point(M, x1, true)

        M = Manifolds.Rotations(3)
        x = collect(reshape(1.0:9.0, (3, 3)))
        x2 = project(M, x)
        @test is_manifold_point(M, x2, true)

        rng = MersenneTwister(44)
        x3 = project(M, randn(rng, 3, 3))
        @test is_manifold_point(M, x3, true)
    end
    @testset "Edge cases of Rotations" begin
        @test_throws OutOfInjectivityRadiusError inverse_retract(
            Rotations(2),
            [1.0 0.0; 0.0 1.0],
            [1.0 0.0; 0.0 -1.0],
            PolarInverseRetraction(),
        )
    end
end


include("../examples/nonconvex_smooth.jl")
using PyPlot
using Printf
function plot_loss_landscape(s=4.0)
	npts = 10000
	w = rand(npts)
	lw = rand(npts)
	fig, ax = subplots()
	ax.xaxis.set_tick_params(labelsize=30)
	ax.yaxis.set_tick_params(labelsize=30)
	for (i, wi) = enumerate(w)
       lw[i] = loss(wi,s)
	end
	ax.plot(w, lw,".")
end
function plot_bifurcation()
	npts = 50
	wstar = rand(npts)
	nSteps = 50000
	η = LinRange(0.0001, 0.1, 20)
	fig, ax = subplots()
	ax.xaxis.set_tick_params(labelsize=30)
	ax.yaxis.set_tick_params(labelsize=30)

	for (i, ηi) = enumerate(η)
		@show i, ηi
		wstar .= rand(npts)
		for j = 1:npts
			wstar[j] = run(wstar[j], ηi, nSteps)
		end
		ax.plot(ηi*ones(npts), wstar, "bo", ms=3.0)
	end
	ax.grid(true)
	ax.set_xlabel(L"$ \eta $", fontsize=30)
	ax.set_ylabel(L"$ w^* $", fontsize=30)
end


using Statistics: median

"""
    KNNDetector(k=5,
                metric=Euclidean,
                algorithm=:kdtree,
                leafsize=10,
                reorder=true,
                reduction=:maximum)

Calculate the anomaly score of an instance based on the distance to its k-nearest neighbors.

Parameters
----------
$K_PARAM

$KNN_PARAMS

    reduction::Symbol
One of `(:maximum, :median, :mean)`. (`reduction=:maximum`) was proposed by [1]. Angiulli et al. [2] proposed sum to
reduce the distances, but mean has been implemented for numerical stability.

Examples
--------
$(SCORE_UNSUPERVISED("KNNDetector"))

References
----------
[1] Ramaswamy, Sridhar; Rastogi, Rajeev; Shim, Kyuseok (2000): Efficient Algorithms for Mining Outliers from Large Data
Sets.

[2] Angiulli, Fabrizio; Pizzuti, Clara (2002): Fast Outlier Detection in High Dimensional Spaces.
"""
OD.@detector mutable struct KNNDetector <: UnsupervisedDetector
    k::Integer = 5::(_ > 0)
    metric::DI.Metric = DI.Euclidean()
    algorithm::Symbol = :kdtree::(_ in (:kdtree, :balltree))
    static::Union{Bool,Symbol} = :auto::(_ in (true, false, :auto))
    leafsize::Integer = 10::(_ ≥ 0)
    reorder::Bool = true
    parallel::Bool = false
    reduction::Symbol = :maximum::(_ in (:maximum, :median, :mean))
end

struct KNNModel <: DetectorModel
    tree::NN.NNTree
end

function OD.fit(detector::KNNDetector, X::Data; verbosity)::Fit
    X = prepare_data(X, detector.static)

    # create the specified tree
    tree = @tree detector X

    # use tree to calculate distances
    _, dists = knn_others(tree, X, detector.k)

    # reduce distances to outlier score
    scores = _knn(dists, detector.reduction)
    return KNNModel(tree), scores
end

function OD.transform(detector::KNNDetector, model::KNNModel, X::Data)::Scores
    X = prepare_data(X, detector.static)

    if detector.parallel
        idxs, dists = knn_parallel(model.tree, X, detector.k)
        return _knn(dists, detector.reduction)
    else
        idxs, dists = NN.knn(model.tree, X, detector.k)
        return _knn(dists, detector.reduction)
    end
end

@inline function _knn(distances::AbstractVector{<:AbstractVector}, reduction::Symbol)::Scores
    # Helper function to reduce `k` distances to a single distance.
    if reduction == :maximum
        return maximum.(distances)
    elseif reduction == :median
        return median.(distances)
    elseif reduction == :mean
        return mean.(distances)
    end
end


ENV["NN_DEBUG"] = true # or change DEBUG = true in NearestNeighbors.jl file

using NearestNeighbors

const data = [[2, 3] [5, 4] [9, 6] [4, 7] [8, 1] [7, 2]]

NearestNeighbors.reset_stats()
const kdtree = KDTree(Float64.(data))
const indexpoint = [9,2]
idx, dist = knn(kdtree, indexpoint, 1)
println("Wikipedia example: The nearest neighbor to $indexpoint is ",
    "$(data[1:2, idx[1]]) at distance $(dist).")
NearestNeighbors.print_stats()

NearestNeighbors.reset_stats()
const cubedis = rand(3, 1000)
const kdcubetree = KDTree(cubedis)
const rand3point = rand(3, 1)
idx, dist = knn(kdcubetree, rand3point, 1)
println("\n\n1000 cube points: The point $rand3point is closest to the point $(cubedis[1:3, idx[1]])",
    " at distance $(dist[1]).")
NearestNeighbors.print_stats()

NearestNeighbors.reset_stats()
const cubedis2 = rand(3, 500000)
const kdcubetree2 = KDTree(cubedis2)
const rand3point2 = rand(3, 1)
idx, dist = knn(kdcubetree2, rand3point2, 1)
println("\n\nExtra: The point $rand3point2 is closest to the point $(cubedis2[1:3, idx[1]]) out of $(size(cubedis2)[2]) points,",
    " at distance $(dist[1]).")
NearestNeighbors.print_stats()


#############################################################################
# real_imag.jl
# Handles real and imaginary part of the variables, constants
# and expressions.
#############################################################################

import Base.real, Base.imag
export real, imag
export sign, monotonicity, curvature, evaluate

### Real
struct RealAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int, Int}

    function RealAtom(x::AbstractExpr)
        children = (x,)
        return new(:real, hash(children), children, x.size)
    end
end

function sign(x::RealAtom)
    if sign(x.children[1]) == ComplexSign()
        return NoSign()
    else
        return sign(x.children[1])
    end
end

function monotonicity(x::RealAtom)
    return (Nondecreasing(),)
end

function curvature(x::RealAtom)
    return ConstVexity()
end

function evaluate(x::RealAtom)
    return real.(evaluate(x.children[1]))
end

function conic_form!(x::RealAtom, unique_conic_forms::UniqueConicForms=UniqueConicForms())
    if !has_conic_form(unique_conic_forms, x)
        new_objective = ConicObj()
        objective = conic_form!(x.children[1], unique_conic_forms)

        for var in keys(objective)
            re = real.(objective[var][1])
            im = real.(objective[var][2])
            new_objective[var] = (re,im)
        end

        cache_conic_form!(unique_conic_forms, x, new_objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

real(x::AbstractExpr) = RealAtom(x)
real(x::Value) = RealAtom(Constant(x))

### Imaginary
struct ImaginaryAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int, Int}

    function ImaginaryAtom(x::AbstractExpr)
        children = (x,)
        return new(:imag, hash(children), children, x.size)
    end
end

function sign(x::ImaginaryAtom)
    sign(x.children[1]) == ComplexSign()
    return NoSign()
end

function monotonicity(x::ImaginaryAtom)
    return (Nondecreasing(),)
end

function curvature(x::ImaginaryAtom)
    return ConstVexity()
end

function evaluate(x::ImaginaryAtom)
    return imag.(evaluate(x.children[1]))
end

function conic_form!(x::ImaginaryAtom, unique_conic_forms::UniqueConicForms=UniqueConicForms())
    if !has_conic_form(unique_conic_forms, x)
        new_objective = ConicObj()
        objective = conic_form!(x.children[1], unique_conic_forms)

        for var in keys(objective)
            re = imag.(objective[var][1])
            im = imag.(objective[var][2])
            new_objective[var] = (re,im)
        end
        cache_conic_form!(unique_conic_forms, x, new_objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

imag(x::AbstractExpr) = ImaginaryAtom(x)
imag(x::Value) = ImaginaryAtom(Constant(x))


using Homotopies
using Base.Test
import FixedPolynomials
const FP = FixedPolynomials
import DynamicPolynomials: @polyvar

include("straightline_test.jl")
include("geodesic_test.jl")
include("misc_test.jl")
include("condition_test.jl")


using BenchmarkTools

include("../../src/part/BarycentricInterpolation.jl")

function computeBaryWeights_TEST()
    x = [0.0, 1.0, 0.8, 0.0]
    y = [0.0, 0.0, 0.8, 0.4]
    p = [0.8, 0.2]

    w = computeBaryWeights(x, y, p)

    @btime computeBaryWeights($x,$y,$p)

    z = [0.0, 0.5, 1.0, 1.5]
    zi = dot(z, w)
end


include("../src/TrueSkill.jl")
using .TrueSkill
global const ttt = TrueSkill
using Test
using CSV
using DataFrames

@testset "Examples" begin

@testset "same_strength" begin
    events = [ [["aj"],["bj"]],[["bj"],["cj"]], [["cj"],["aj"]] ,[["aj"],["bj"]],[["bj"],["cj"]], [["cj"],["aj"]]]
    results = [[1.,0.],[1.,0.],[1.,0.],[1.,0.],[1.,0.],[1.,0.]]
    priors = Dict{String,ttt.Player}()
    for k in ["aj", "bj", "cj"]
        priors[k] = ttt.Player(ttt.Gaussian(0., 3.), 0.5, 0.0)
    end
    h = ttt.History(events, results, [1,2,3,4,5,6], priors)
    lc_ts = ttt.learning_curves(h)
    ttt.convergence(h)
    lc_ttt = ttt.learning_curves(h)

    df = DataFrame(mu_a_ts = [ N.mu for (k,N) in lc_ts["aj"]]
                    ,sigma_a_ts = [ N.sigma for (k,N) in lc_ts["aj"]]
                    ,mu_a_ttt = [ N.mu for (k,N) in lc_ttt["aj"]]
                    ,sigma_a_ttt = [ N.sigma for (k,N) in  lc_ttt["aj"]]
                    ,time_a = [ k for (k,N) in lc_ts["aj"]]
                    ,mu_b_ts = [ N.mu for (k,N) in lc_ts["bj"]]
                    ,sigma_b_ts = [ N.sigma for (k,N) in lc_ts["bj"]]
                    ,mu_b_ttt = [ N.mu for (k,N) in lc_ttt["bj"]]
                    ,sigma_b_ttt = [ N.sigma for (k,N) in lc_ttt["bj"]]
                    ,time_b = [ k for (k,N) in lc_ts["bj"]]
                    ,mu_c_ts = [ N.mu for (k,N) in lc_ts["cj"]]
                    ,sigma_c_ts = [ N.sigma for (k,N) in lc_ts["cj"]]
                    ,mu_c_ttt = [ N.mu for (k,N) in lc_ttt["cj"]]
                    ,sigma_c_ttt = [ N.sigma for (k,N) in lc_ttt["cj"]]
                    ,time_c = [ k for (k,N) in lc_ts["cj"]]
                    )

    CSV.write("output/same_strength.csv", df; header=true)
    @test true
end
@testset "same_strength_two_groups" begin
    predicciones_mle = Float64[]
    predicciones_aiaj = Float64[]
    predicciones_bibj = Float64[]
    predicciones_aibi = Float64[]
    predicciones_bici = Float64[]
    predicciones_ajbj = Float64[]
    predicciones_bjcj = Float64[]
    composition = [ [["aj"],["bj"]], [["bj"],["cj"]], [["cj"],["aj"]],[["aj"],["bj"]], [["bj"],["cj"]], [["cj"],["aj"]],
                [["ai"],["bi"]], [["bi"],["ci"]], [["ci"],["ai"]], [["ai"],["bi"]], [["bi"],["ci"]], [["ci"],["ai"]],
                [["aj"],["ai"]] ]
    times = [1,1,1,1,1,1,1,1,1,1,1,1,1]
    Nbeta = ttt.Gaussian(0.0,0.5)

    for i in 1:150
        push!(composition, [["ai"],["aj"]])
        push!(times, 1)
        priors = Dict{String,ttt.Player}()
        for k in ["aj", "bj", "cj", "ai", "bi", "ci"]
            priors[k] = ttt.Player(ttt.Gaussian(0., 5.0), 0.5, 0.0)
        end

        h = ttt.History(composition=composition, times=times, priors=priors, iter=40)
        ttt.convergence(h)

        lc = ttt.learning_curves(h)

        push!(predicciones_mle, i/(1.0+i))

        push!(predicciones_aiaj, 1-ttt.cdf(lc["ai"][1][2]+Nbeta - lc["aj"][1][2]+Nbeta, 0.))
        push!(predicciones_bibj, 1-ttt.cdf(lc["bi"][1][2]+Nbeta - lc["bj"][1][2]+Nbeta, 0.))
        push!(predicciones_aibi, 1-ttt.cdf(lc["ai"][1][2]+Nbeta - lc["bi"][1][2]+Nbeta, 0.))
        push!(predicciones_bici, 1-ttt.cdf(lc["bi"][1][2]+Nbeta - lc["ci"][1][2]+Nbeta, 0.))
        push!(predicciones_ajbj, 1-ttt.cdf(lc["aj"][1][2]+Nbeta - lc["bj"][1][2]+Nbeta, 0.))
        push!(predicciones_bjcj, 1-ttt.cdf(lc["bj"][1][2]+Nbeta - lc["cj"][1][2]+Nbeta, 0.))
    end

    # S - Shape
    using Plots
    plot(predicciones_mle,predicciones_mle,legend=(-10.0, 10.0))
    plot!(predicciones_mle,predicciones_aiaj)
    #plot!(predicciones_mle,predicciones_bibj)

    df = DataFrame(mle = predicciones_mle
                    ,aiaj = predicciones_aiaj
                    ,bibj = predicciones_bibj
                    ,aibi = predicciones_aibi
                    ,bici = predicciones_bici
                    )

    CSV.write("output/same_strength_two_groups.csv", df; header=true)
    @test true
end
@testset "smoothing" begin
    events = [ [["a"],["b"]], [["a"],["b"]]]
    results = [[1.,0.],[0.,1.]]
    times = [1,2]
    priors = Dict{String,ttt.Player}()
    Nbeta = ttt.Gaussian(0.,0.5)
    for k in ["a", "b"]
        priors[k] = ttt.Player(ttt.Gaussian(0., 3.), 0.5, 0.0 )
    end
    h = ttt.History(events, results, times, priors)
    fp_a_1 = Vector{ttt.Gaussian}()
    bp_a_1 = Vector{ttt.Gaussian}()
    lh_a_1 = Vector{ttt.Gaussian}()
    wp_a_1 = Vector{ttt.Gaussian}()

    fp_a_2 = Vector{ttt.Gaussian}()
    bp_a_2 = Vector{ttt.Gaussian}()
    lh_a_2 = Vector{ttt.Gaussian}()
    wp_a_2 = Vector{ttt.Gaussian}()

    lh_b_1 = Vector{ttt.Gaussian}()
    lh_b_2 = Vector{ttt.Gaussian}()

    wp_b_1 = Vector{ttt.Gaussian}()
    wp_b_2 = Vector{ttt.Gaussian}()

    p_a_1 =  Vector{ttt.Gaussian}()
    p_a_2 =  Vector{ttt.Gaussian}()
    p_b_1 =  Vector{ttt.Gaussian}()
    p_b_2 =  Vector{ttt.Gaussian}()

    e_1 = Float64[]
    e_2 = Float64[]
    d_div_1 = Vector{ttt.Gaussian}()
    d_div_2 = Vector{ttt.Gaussian}()

    push!(fp_a_1, h.batches[1].skills["a"].forward)
    push!(bp_a_1, h.batches[1].skills["a"].backward)
    push!(lh_a_1, h.batches[1].skills["a"].likelihood)
    push!(wp_a_1, fp_a_1[end]*bp_a_1[end])
    push!(wp_b_1, h.batches[1].skills["b"].forward*h.batches[1].skills["b"].backward)
    push!(fp_a_2, h.batches[2].skills["a"].forward)
    push!(bp_a_2, h.batches[2].skills["a"].backward)
    push!(lh_a_2, h.batches[2].skills["a"].likelihood)
    push!(wp_a_2, fp_a_2[end]*bp_a_2[end])
    push!(wp_b_2, h.batches[2].skills["b"].forward*h.batches[2].skills["b"].backward)
    push!(e_1, h.batches[1].events[1].evidence)
    push!(e_2, h.batches[2].events[1].evidence)
    push!(lh_b_1, h.batches[1].skills["b"].likelihood)
    push!(lh_b_2, h.batches[2].skills["b"].likelihood)

    push!(p_a_1, wp_a_1[end]*lh_a_1[end])
    push!(p_a_2, wp_a_2[end]*lh_a_2[end])
    push!(p_b_1, wp_b_1[end]*lh_b_1[end])
    push!(p_b_2, wp_b_2[end]*lh_b_2[end])

    d_1 = wp_a_1[end]+Nbeta - wp_b_1[end]+Nbeta
    push!(d_div_1 , ttt.approx(d_1,0.,false)/d_1)
    d_2 = wp_a_1[end]+Nbeta - wp_b_1[end]+Nbeta
    push!(d_div_2 , ttt.approx(d_2,0.,false)/d_2)

    for _ in 1:10
        ttt.iteration(h)
        push!(fp_a_1, h.batches[1].skills["a"].forward)
        push!(bp_a_1, h.batches[1].skills["a"].backward)
        push!(lh_a_1, h.batches[1].skills["a"].likelihood)
        push!(wp_a_1, fp_a_1[end]*bp_a_1[end])
        push!(wp_b_1, h.batches[1].skills["b"].forward*h.batches[1].skills["b"].backward)
        push!(fp_a_2, h.batches[2].skills["a"].forward)
        push!(bp_a_2, h.batches[2].skills["a"].backward)
        push!(lh_a_2, h.batches[2].skills["a"].likelihood)
        push!(wp_a_2, fp_a_2[end]*bp_a_2[end])
        push!(wp_b_2, h.batches[2].skills["b"].forward*h.batches[2].skills["b"].backward)
        push!(e_1, h.batches[1].events[1].evidence)
        push!(e_2, h.batches[2].events[1].evidence)
        push!(lh_b_1, h.batches[1].skills["b"].likelihood)
        push!(lh_b_2, h.batches[2].skills["b"].likelihood)

        d_1 = wp_a_1[end]+Nbeta - wp_b_1[end]+Nbeta
        push!(d_div_1 , ttt.approx(d_1,0.,false)/d_1)
        d_2 = wp_a_1[end]+Nbeta - wp_b_1[end]+Nbeta
        push!(d_div_2 , ttt.approx(d_2,0.,false)/d_2)
        push!(p_a_1, wp_a_1[end]*lh_a_1[end])
        push!(p_a_2, wp_a_2[end]*lh_a_2[end])
        push!(p_b_1, wp_b_1[end]*lh_b_1[end])
        push!(p_b_2, wp_b_2[end]*lh_b_2[end])
    end

    df = DataFrame(fp_a_1_mu = [N.mu for N in fp_a_1]
                    ,fp_a_1_sigma = [N.sigma for N in fp_a_1]
                    ,bp_a_1_mu = [N.mu for N in bp_a_1]
                    ,bp_a_1_sigma = [N.sigma for N in bp_a_1]
                    ,lh_a_1_mu = [N.mu for N in lh_a_1]
                    ,lh_a_1_sigma = [N.sigma for N in lh_a_1]
                    ,wp_a_1_mu = [N.mu for N in wp_a_1]
                    ,wp_a_1_sigma = [N.sigma for N in wp_a_1]
                    ,p_a_1_mu = [N.mu for N in p_a_1]
                    ,p_a_1_sigma = [N.sigma for N in p_a_1]
                    ,p_a_2_mu = [N.mu for N in p_a_2]
                    ,p_a_2_sigma = [N.sigma for N in p_a_2]
                    ,p_b_1_mu = [N.mu for N in p_b_1]
                    ,p_b_1_sigma = [N.sigma for N in p_b_1]
                    ,p_b_2_mu = [N.mu for N in p_b_2]
                    ,p_b_2_sigma = [N.sigma for N in p_b_2]
                    ,wp_b_1_mu = [N.mu for N in wp_b_1]
                    ,wp_b_1_sigma = [N.sigma for N in wp_b_1]
                    ,fp_a_2_mu = [N.mu for N in fp_a_2]
                    ,fp_a_2_sigma = [N.sigma for N in fp_a_2]
                    ,bp_a_2_mu = [N.mu for N in bp_a_2]
                    ,bp_a_2_sigma = [N.sigma for N in bp_a_2]
                    ,lh_a_2_mu = [N.mu for N in lh_a_2]
                    ,lh_a_2_sigma = [N.sigma for N in lh_a_2]
                    ,wp_a_2_mu = [N.mu for N in wp_a_2]
                    ,wp_a_2_sigma = [N.sigma for N in wp_a_2]
                    ,wp_b_2_mu = [N.mu for N in wp_b_2]
                    ,wp_b_2_sigma = [N.sigma for N in wp_b_2]
                    ,lh_b_1_mu = [N.mu for N in lh_b_1]
                    ,lh_b_2_mu = [N.mu for N in lh_b_2]
                    ,lh_b_1_sigma = [N.sigma for N in lh_b_1]
                    ,lh_b_2_sigma = [N.sigma for N in lh_b_2]
                    ,e_1 = e_1
                    ,e_2 = e_2
                    ,d_div_1_mu = [N.mu for N in d_div_1]
                    ,d_div_1_sigma = [N.sigma for N in d_div_1]
                    ,d_div_2_mu = [N.mu for N in d_div_2]
                    ,d_div_2_sigma = [N.sigma for N in d_div_2]
                    )
    CSV.write("output/smoothing.csv", df; header=true)
    @test true
end
@testset "ttt_vs_ts" begin

    events = Array{Array{Array{String,1},1},1}()
    results =  Array{Array{Float64,1},1}()
    times = Int64[]
    ts_log_evidence = Float64[]
    ttt_log_evidence = Float64[]

    ts_last = Float64[]; ts_midle = Float64[]; ts_second = Float64[]
    ttt_last = Float64[]; ttt_midle = Float64[]; ttt_second = Float64[]
    for i in 0:127
        push!(events ,  [["a"],["b"]], [["a"],["b"]])
        push!(results ,  [1.,0.],[0.,1.] )
        push!(times,  i*2+1, i*2+2 )
        priors = Dict{String,ttt.Player}()
        for k in ["a", "b"]
            priors[k] = ttt.Player(ttt.Gaussian(0., 3.0), 0.5, 0.0)
        end
        h = ttt.History(events, results, times, priors)
        push!(ts_log_evidence , ttt.log_evidence(h))
        push!(ts_last, h.batches[end].events[1].evidence); push!(ts_midle, h.batches[i*2+1].events[1].evidence);
        push!(ts_second, h.batches[2].events[1].evidence);
        ttt.convergence(h)
        push!(ttt_last, h.batches[end].events[1].evidence); push!(ttt_midle, h.batches[i*2+1].events[1].evidence);
        push!(ttt_log_evidence, ttt.log_evidence(h))
        push!(ttt_second, h.batches[2].events[1].evidence)
    end
    df = DataFrame(ts_log_evidence = ts_log_evidence
                    ,ttt_log_evidence = ttt_log_evidence
                    ,ts_last = ts_last
                    ,ts_midle = ts_midle
                    ,ts_second = ts_second
                    ,ttt_second = ttt_second
                    ,ttt_last = ttt_last
                    ,ttt_midle = ttt_midle )

    CSV.write("output/ttt_vs_ts.csv", df; header=true)
    @test true
end
@testset "Best gamma" begin

    function skill(exp::Int64, alpha::Float64=0.133)
        return exp^alpha - 1
    end
    mean_agent = [skill(i) for i in 1:1000]
    beta = 0.5

    using Random
    Random.seed!(1)

    mean_target = [(Random.randn(1)[1]*beta + skill(i)) for i in 1:1000]
    perf_target = [(Random.randn(1)[1]*beta + mean_target[i]) for i in 1:1000]
    perf_agent = [(Random.randn(1)[1]*beta + mean_agent[i]) for i in 1:1000]
    events = [ [["a"], [string(i)] ] for i in 1:1000]
    results = [ perf_agent[i] > perf_target[i] ? [1.,0.] : [0.,1.] for i in 1:1000 ]
    batches= [i for i in 1:1000 ]

    selected_gammas = [0.005,0.01,0.015,0.02,0.025]

    gammas = [gamma for gamma in 0.001:0.001:0.04]
    evidencias_ts = Float64[]
    evidencias_ttt = Float64[]
    learning_curves = Vector{Vector{Float64}}()
    for gamma in gammas#gamma=0.015

        priors = Dict{String,ttt.Player}()
        priors["a"] = ttt.Player(ttt.Gaussian(0., 3.0), beta, gamma)
        for k in 1:1000
            priors[string(k)] = ttt.Player(ttt.Gaussian(mean_target[k], 0.5), beta, 0.0)
        end
        h = ttt.History(events, results, batches, priors)
        push!(evidencias_ts, ttt.log_evidence(h))
        ttt.convergence(h)
        push!(evidencias_ttt, ttt.log_evidence(h))
        if gamma in selected_gammas
            push!(learning_curves, [r.mu for (t, r) in ttt.learning_curves(h)["a"]] )
        end
    end
    #using Plots
    #plot([g for g in gammas ],evidencias_ts)
    #plot!([g for g in gammas ] ,evidencias_ttt)

    #plot(mean_agent )
    #plot!(learning_curves[3])

    @test 0.015 == gammas[argmax(evidencias_ttt)]

    df = DataFrame(gammas = gammas
                    ,evidencias_ts = evidencias_ts
                    ,evidencias_ttt = evidencias_ttt)
    CSV.write("output/best_gamma-evidences.csv", df; header=true)

    df = DataFrame(mean_agent = mean_agent
                    ,lc_005 = learning_curves[1]
                    ,lc_01  = learning_curves[2]
                    ,lc_015 = learning_curves[3]
                    ,lc_020 = learning_curves[4]
                    ,lc_025 = learning_curves[5])
    CSV.write("output/best_gamma-learning_curves.csv", df; header=true)

    #plot(mean_agent)
    #plot!(learning_curves[3])
end
@testset "Chain of betas" begin
    using Random
    Random.seed!(1)
    beta = 1.0
    agents = [string(i) for i in -5:5]
    true_skills = [i for i in -5.0:5.0]
    mu_prior = [ abs(i)==5.0 ? i : 0.0 for i in -5.0:5.0]
    sigma_prior = [ abs(i)==5.0 ? 1e-7 : 6.0 for i in -5.0:5.0]
    priors = Dict{String,ttt.Player}()
    for i in -5:5
        priors[string(i)] = ttt.Player(ttt.Gaussian(mu_prior[i+6], sigma_prior[i+6]), beta, 0.0)
    end

    events = [ [[string(a-1)],[string(a)]]  for e in 1:100 for a in -4:5]
    results = [ (Random.randn(1)[1]*beta + a-1) > (Random.randn(1)[1]*beta + a) ? [0.,1.] : [1.,0.] for i in 1:100 for a in -4.0:5.0]
    times = [ e  for e in 1:100 for a in -4:5]
    h = ttt.History(events , results, times, priors, iter=100)
    ttt.convergence(h)
    diffs = [(ttt.posterior(h.batches[1],string(a))-ttt.posterior(h.batches[1],string(a-1))).mu for a in -3:4]
    probs = [ 1-ttt.cdf(ttt.Gaussian(d,sqrt(2)),0.0)  for d in diffs]
    for p in probs
        @test abs(1-ttt.cdf(ttt.Gaussian(1.0,sqrt(2)),0.0)-p) < 0.04
    end
end
@testset "Online predictions" begin
    priors = Dict{String,ttt.Player}()
    priors["a"] = ttt.Player(ttt.Gaussian(0.0, 3.0), 1.0, 0.0)
    priors["b"] = ttt.Player(ttt.Gaussian(0.0, 3.0), 1.0, 0.0)
    priors["c"] = ttt.Player(ttt.Gaussian(2.0, 0.5), 1.0, 0.0)
#
    composition = [[["a"],["b"]], [["a"],["c"]], [["a"],["b"]]]
    results = [ [0.,1.], [1.,0.], [1.,0.]]
    times = [1,2,3]
    h = ttt.History(composition , results, times, priors, iter=100, gamma=0.0,online=true)

    h.batches[3].skills["b"].forward
    h.batches[3].skills["b"].online
end
end