#####
##### First derivative operators
#####

for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :ᶠ)

    x_derivative = Symbol(:∂x, LX, LY, LZ)
    x_spacing    = Symbol(:Δx, LX, LY, LZ)
    x_difference = Symbol(:δx, LX, :ᵃ, :ᵃ)

    y_derivative = Symbol(:∂y, LX, LY, LZ)
    y_spacing    = Symbol(:Δy, LX, LY, LZ)
    y_difference = Symbol(:δy, :ᵃ, LY, :ᵃ)

    z_derivative = Symbol(:∂z, LX, LY, LZ)
    z_spacing    = Symbol(:Δz, LX, LY, LZ)
    z_difference = Symbol(:δz, :ᵃ, :ᵃ, LZ)

    @eval begin
        @inline $x_derivative(i, j, k, grid, c) = $x_difference(i, j, k, grid, c) / $x_spacing(i, j, k, grid)
        @inline $y_derivative(i, j, k, grid, c) = $y_difference(i, j, k, grid, c) / $y_spacing(i, j, k, grid)
        @inline $z_derivative(i, j, k, grid, c) = $z_difference(i, j, k, grid, c) / $z_spacing(i, j, k, grid)

        @inline $x_derivative(i, j, k, grid, f::Function, args...) = $x_difference(i, j, k, grid, f, args...) / $x_spacing(i, j, k, grid)
        @inline $y_derivative(i, j, k, grid, f::Function, args...) = $y_difference(i, j, k, grid, f, args...) / $y_spacing(i, j, k, grid)
        @inline $z_derivative(i, j, k, grid, f::Function, args...) = $z_difference(i, j, k, grid, f, args...) / $z_spacing(i, j, k, grid)
    end
end

#####
##### Second, Third, and Fourth derivatives
#####

@inline insert_symbol(dir, L, L1, L2) =
                      dir == :x ?
                      (L, L1, L2) :
                      dir == :y ?
                      (L1, L, L2) :
                      (L1, L2, L)

for dir in (:x, :y, :z), L1 in (:ᶜ, :ᶠ), L2 in (:ᶜ, :ᶠ)

     first_order_face = Symbol(:∂,  dir, insert_symbol(dir, :ᶠ, L1, L2)...)
    second_order_face = Symbol(:∂², dir, insert_symbol(dir, :ᶠ, L1, L2)...)
     third_order_face = Symbol(:∂³, dir, insert_symbol(dir, :ᶠ, L1, L2)...)
    fourth_order_face = Symbol(:∂⁴, dir, insert_symbol(dir, :ᶠ, L1, L2)...)

     first_order_center = Symbol(:∂,  dir, insert_symbol(dir, :ᶜ, L1, L2)...)
    second_order_center = Symbol(:∂², dir, insert_symbol(dir, :ᶜ, L1, L2)...)
     third_order_center = Symbol(:∂³, dir, insert_symbol(dir, :ᶜ, L1, L2)...)
    fourth_order_center = Symbol(:∂⁴, dir, insert_symbol(dir, :ᶜ, L1, L2)...)

    @eval begin
        $second_order_face(i, j, k, grid, c) =  $first_order_face(i, j, k, grid, $first_order_center,  c)
         $third_order_face(i, j, k, grid, c) =  $first_order_face(i, j, k, grid, $second_order_center, c)
        $fourth_order_face(i, j, k, grid, c) = $second_order_face(i, j, k, grid, $second_order_face,   c)

        $second_order_center(i, j, k, grid, c) =  $first_order_center(i, j, k, grid, $first_order_face,    c)
         $third_order_center(i, j, k, grid, c) =  $first_order_center(i, j, k, grid, $second_order_face,   c)
        $fourth_order_center(i, j, k, grid, c) = $second_order_center(i, j, k, grid, $second_order_center, c)

        $second_order_face(i, j, k, grid, f::Function, args...) =  $first_order_face(i, j, k, grid, $first_order_center,  f::Function, args...)
         $third_order_face(i, j, k, grid, f::Function, args...) =  $first_order_face(i, j, k, grid, $second_order_center, f::Function, args...)
        $fourth_order_face(i, j, k, grid, f::Function, args...) = $second_order_face(i, j, k, grid, $second_order_face,   f::Function, args...)

        $second_order_center(i, j, k, grid, f::Function, args...) =  $first_order_center(i, j, k, grid, $first_order_face,    f::Function, args...)
         $third_order_center(i, j, k, grid, f::Function, args...) =  $first_order_center(i, j, k, grid, $second_order_face,   f::Function, args...)
        $fourth_order_center(i, j, k, grid, f::Function, args...) = $second_order_center(i, j, k, grid, $second_order_center, f::Function, args...)
    end
end

#####
##### Operators of the form A*∂(q) where A is an area and q is some quantity.
#####

for dir in (:x, :y, :z), LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :ᶠ)

    operator   = Symbol(:A, dir, :_∂, dir, LX, LY, LZ)
    area       = Symbol(:A, dir, LX, LY, LZ)
    derivative = Symbol(:∂, dir, LX, LY, LZ)

    @eval begin
        $operator(i, j, k, grid, c) = $area(i, j, k, grid) * $derivative(i, j, k, grid, c)
    end
end


using PtFEM, Test

data = Dict(
  # Beam(nels, nn, ndim, nst, nip, finite_element(nod, nodof))
  :struc_el => Beam(2, 1, 4, 1, :x, Line(2, 2), false),
  :properties => [4.0e4; 2.0e4],
  :etype => [1, 1, 2, 2],
  :x_coords => [0.0, 2.5, 5.0, 8.0, 10.0],
  :loaded_nodes => [
  (2, [-20.0 0.0]),
  (3, [-6.0 -3.0]),
  (4, [-8.8 2.2]),
  (5, [-1.2 0.5333])],
  :fixed_freedoms => [
    (1, 1, 0.0),
    (1, 2, -0.001),
    (3, 1, -0.005),
    (4, 1, 0.0)],
  :penalty => 1e19
)

@time fem, dis_df, fm_df = p43(data)

@test round.(fem.displacements[:, 5], digits=7) == [0.0047132, 0.0023433]


function get_cartesian_mesh_size(dim, order, prob_size)
    dims = zeros(CeedInt, dim)
    # Use the approximate formula:
    #    prob_size ~ num_elem * order^dim
    num_elem = div(prob_size, order^dim)
    s = 0 # find s: num_elem/2 < 2^s <= num_elem
    while num_elem > 1
        num_elem = div(num_elem, 2)
        s += 1
    end
    r = s%dim
    for d = 1:dim
        sd = div(s, dim)
        if r > 0
            sd += 1
            r -= 1
        end
        dims[d] = 2^sd
    end
    dims
end

struct FormRestrictionMode{T} end
const RestrictionOnly = FormRestrictionMode{:restr}()
const StridedOnly = FormRestrictionMode{:restr_i}()
const RestrictionAndStrided = FormRestrictionMode{:both}()

function build_cartesian_restriction(
    c::Ceed,
    dim,
    nxyz,
    order,
    ncomp,
    num_qpts;
    mode::Mode=RestrictionOnly,
) where {Mode}
    p::CeedInt = order
    pp1::CeedInt = p + 1
    nnodes::CeedInt = pp1^dim # number of scal. nodes per element
    elem_qpts::CeedInt = num_qpts^dim # number of qpts per element

    nd = CeedInt.(p*nxyz .+ 1)
    num_elem::CeedInt = prod(nxyz)
    scalar_size::CeedInt = prod(nd)
    size::CeedInt = scalar_size*ncomp

    form_restr = (Mode() != StridedOnly)
    form_strided = (Mode() != RestrictionOnly)

    # elem:         0             1                 n-1
    #        |---*-...-*---|---*-...-*---|- ... -|--...--|
    # nnodes:   0   1    p-1  p  p+1       2*p             n*p

    if form_restr
        el_nodes = zeros(CeedInt, num_elem*nnodes)
        exyz = zeros(CeedInt, dim)
        @inbounds @simd for e = 0:(num_elem-1)
            re::CeedInt = e
            for d::CeedInt = 1:dim
                exyz[d] = re%nxyz[d]
                re ÷= nxyz[d]
            end
            for lnodes::CeedInt = 0:(nnodes-1)
                gnodes::CeedInt = 0
                gnodes_stride::CeedInt = 1
                rnodes::CeedInt = lnodes
                for d::CeedInt = 1:dim
                    q = rnodes÷pp1
                    r = rnodes - pp1*q
                    gnodes::CeedInt += (exyz[d]*p + r)*gnodes_stride
                    gnodes_stride::CeedInt *= nd[d]
                    rnodes = q
                end
                el_nodes[e*nnodes+lnodes+1] = gnodes
            end
        end
    end

    restr =
        form_restr ?
        create_elem_restriction(
            c,
            num_elem,
            nnodes,
            ncomp,
            scalar_size,
            ncomp*scalar_size,
            el_nodes,
        ) : nothing
    restr_i =
        form_strided ?
        create_elem_restriction_strided(
            c,
            num_elem,
            elem_qpts,
            ncomp,
            ncomp*elem_qpts*num_elem,
            STRIDES_BACKEND,
        ) : nothing

    return size, restr, restr_i
end

function set_cartesian_mesh_coords!(dim, nxyz, mesh_order, mesh_coords)
    p = mesh_order
    nd = p*nxyz .+ 1
    num_elem::CeedInt = prod(nxyz)
    scalar_size::CeedInt = prod(nd)

    # The H1 basis uses Lobatto quadrature points as nodes.
    nodes::Vector{CeedScalar} = lobatto_quadrature(p + 1) # nodes are in [-1,1]
    nodes = 0.5 .+ 0.5*nodes

    @witharray coords = mesh_coords begin
        @inbounds @simd for gsnodes = 0:scalar_size-1
            rnodes = gsnodes
            for d = 1:dim
                ndd = nd[d]
                q = rnodes÷ndd
                r = rnodes - ndd*q

                d1d = r
                coords[gsnodes+scalar_size*(d-1)+1] = ((d1d÷p) + nodes[d1d%p+1])/nxyz[d]
                rnodes = q
            end
        end
    end
end


@node MvNormalMeanCovariance Stochastic [ out, (μ, aliases = [ mean ]), (Σ, aliases = [ cov ]) ]

conjugate_type(::Type{ <: MvNormalMeanCovariance }, ::Type{ Val{ :out } }) = MvNormalMeanCovariance
conjugate_type(::Type{ <: MvNormalMeanCovariance }, ::Type{ Val{ :μ } })   = MvNormalMeanCovariance
conjugate_type(::Type{ <: MvNormalMeanCovariance }, ::Type{ Val{ :Σ } })   = InverseWishart

@average_energy MvNormalMeanCovariance (q_out::Any, q_μ::Any, q_Σ::Any) = begin
    (m_mean, v_mean) = mean(q_μ), cov(q_μ)
    (m_out, v_out)   = mean(q_out), cov(q_out)
    0.5 * (ndims(q_out) * log2π + logdet(mean(q_Σ)) + tr(cholinv(mean(q_Σ))*(v_out + v_mean + (m_out - m_mean)*(m_out - m_mean)')))
end

@average_energy MvNormalMeanCovariance (q_out_μ::Any, q_Σ::Any) = begin
    (m, V) = mean(q_out_μ), cov(q_out_μ)
    d = Int64(ndims(q_out_μ)/2)
    @views 0.5*(d*log2π + logdet(mean(q_Σ)) + tr(cholinv(mean(q_Σ))*( V[1:d,1:d] - V[1:d,d+1:end] - V[d+1:end,1:d] + V[d+1:end,d+1:end] + (m[1:d] - m[d+1:end])*(m[1:d] - m[d+1:end])' ) ))
end


using Knet
include("readdata.jl")

# Copy of VGG model A: http://arxiv.org/pdf/1409.1556v6.pdf

@knet function model04(x; T=51, D=4, W=3, o...)
    c1 = cbfp73(x;  o..., out=64, pwindow=T)
    c2 = cbfp73(c1; o..., out=128, pwindow=T)
    c3 = cbf73(c2; o..., out=256)
    c4 = cbfp73(c3; o..., out=256, pwindow=T)
    c5 = cbf73(c4; o..., out=512)
    c6 = cbfp73(c5; o..., out=512, pwindow=T)
    c7 = cbf73(c6; o..., out=512)
    c8 = cbfp73(c7; o..., out=512, pwindow=T)
    f1 = wbf73(c8; o..., out=4096)
    f2 = wbf73(f1; o..., out=4096)
    return wbf73(f2; o..., out=2, f=:soft)
end

function train04(; N=2^13, lr=0.001, adam=true, nbatch=128, o...)
    global f = compile(:model04; o...)
    setp(f, lr=lr, adam=adam)
    global d = GulayseData(nbatch)
    sloss = zloss = -1; n = nextn = 1
    for (x,ygold) in d
        ypred = forw(f, x)
        sl = softloss(ypred,ygold); sloss = (n==1 ? sl : 0.99 * sloss + 0.01 * sl)
        zl = zeroone(ypred,ygold);  zloss = (n==1 ? zl : 0.99 * zloss + 0.01 * zl)
        n==nextn && (println((n,sloss,1-zloss)); nextn*=2)
        back(f, ygold, softloss)
        update!(f)
        (n += 1) > N && break
    end
end

train04()


L = 1:1000; d = 1:1000;
L2d = Array{Float64}(undef, length(L), length(d));
c(L,d) = (10*L/(2*pi)*tanh(2*pi*d/L))^0.5;

Lmatrix = [fill(convert(Float64,i),(1,length(d))) for i in 1:length(d)]


using Random, Optim

Base.@kwdef struct Config
    fuel_dist::Float64 = 2.e6
    seed::Int = 1234
    rng::MersenneTwister = MersenneTwister(seed)
    optimizer::Optim.AbstractOptimizer = NelderMead()
    models::Int = 20
end

struct BookKeep
    mask::BitArray # keeps visited locations
    betas::Vector{Float64} # optimization weights
    pointwise::Vector{Float64}
    function BookKeep(dim, feature_count)
        new(trues(dim), Vector{Float64}(undef, feature_count), Vector{Float64}(undef, dim))
    end
end

# finds most desirable location (local maximum)
function masked_weighted_argmax(mask, scores)
    id = 1
    m = floatmax(Float64) * -1
    @inbounds for i ∈ eachindex(mask)
        if mask[i] && scores[i] > m
            id = i
            m = scores[i]
        end
    end
    return id
end

function find_betas!(coord, locinfo, config, bookkeep, feat, betas; output=false)
    beer_count, traveled = 0, 0.
    bookkeep.mask .= true

    calc_scores!(feat, betas)
    id = masked_weighted_argmax(bookkeep.mask, feat.scores)
    beer_count += size(locinfo.beers[id], 1)
    traveled += bookkeep.pointwise[id]
    bookkeep.mask[id] = false
    output && println("$(Int(traveled÷1000))km $(locinfo.locations[id, :name]) +$(size(locinfo.beers[id], 1))")

    while true
        calc_scores!(feat, betas, id)
        id_new = masked_weighted_argmax(bookkeep.mask, feat.scores)
        bookkeep.pointwise[id_new] + coord.pairwise[id, id_new] + traveled > config.fuel_dist && break
        beer_count += size(locinfo.beers[id], 1)
        traveled += coord.pairwise[id, id_new]
        id = id_new
        bookkeep.mask[id] = false
        output && println("$(Int(traveled÷1000))km $(locinfo.locations[id, :name]) +$(size(locinfo.beers[id], 1))")
    end

    return -beer_count # only beer count is optimized since task is ambiguous
end

function run_optimization(coord, locinfo, config, bookkeep, feat, output=true)
    minimum = 1.
    minimizer = []
    for _ ∈ 1:config.models # few models are run with different betas initializations
        rand!(config.rng, bookkeep.betas)
        res = optimize(betas -> find_betas!(coord, locinfo, config, bookkeep, feat, betas), bookkeep.betas, config.optimizer)
        if Optim.minimum(res) < minimum
            minimum = Optim.minimum(res)
            minimizer = Optim.minimizer(res)
        end
    end

    find_betas!(coord, locinfo, config, bookkeep, feat, minimizer; output=output)
    output && println("Beer count: $(Int(-minimum))")
    return (Int(-minimum))
end


# this potential is an α-attractor model with a quadratic potential
# this run script is set to run on a cluster with 24 total physical CPU cores
# the simulations are scanned over the model, for various values of α, N0, and Nf
# each simulation, if it successfully computes perturbations is saved in a data folder

const NUMCPU = 24
Nfs = [2,5,10]
runs = 100
alphas = [1//6, 1//60]
N0s = 75:25:400

datafolder = "/data/ANflation/"
subfolder = "batch1"

if !isdir(datafolder)
    println("Data folder doesn't exist!")
elseif !isdir(joinpath(datafolder,subfolder))
    println("Creating $(joinpath(datafolder,subfolder)) ...")
    mkdir(joinpath(datafolder,subfolder))
end

using Distributed
if length(workers()) == 1
    addprocs(NUMCPU - 1)
end

@everywhere using Inflation
@everywhere using SymPy

@everywhere begin
    using Random
    using Roots
    using ForwardDiff
    m = 1e-5
    β = 1/2
    a = m^2 * (1-sqrt(β))^2
    b = m^2 * (1+sqrt(β))^2
    function MarchenkoPasturPDF(x)
        return sqrt( (b-x)*(x-a)) / (2*pi*x*β*m^2)
    end
    support = (a,b)
    function rejection_sample(pdf,num,support=(-Inf,Inf))
        # for now assume pdf has exactly one critical point
        # will probably fail without finite support
        a = support[1]
        b = support[2]
        dpdf(x) = ForwardDiff.derivative(pdf,x)
        xmax = find_zero(dpdf,(a,b))
        pdf_max = pdf(xmax)
        samples = Array{Float64}(undef,num)
        i = 1
        while i <= num
            x = a+(b-a)*rand()
            s = rand()*pdf_max
            if s < pdf(x)
                samples[i] = x
                i += 1
            end
        end
        return samples
    end
end

for d in Nfs
    function G(Phi,p)
        d = length(Phi)
        α = p[1]
        Phi2 = sum([f^2 for f in Phi])
        g = Diagonal(ones(Sym,d) .* 6*α / (1 - Phi2)^2)
        return g
    end

    function V(Phi,p)
        m = p[2:end]
        α = p[1]
        v = 0
        for i in 1:d
            v += α/2.0*(m[i]*Phi[i]^2)
        end
        return v
    end

    runs_to_run = []
    for N0 in N0s, alpha in alphas, run in 1:runs
        fname = "N$(N0)_a$(Int(1/alpha))_Nf$(d)_r$run.jld2"
        fullfname = joinpath(datafolder,subfolder,fname)
        if !isfile(fullfname)
            push!(runs_to_run,(N0,alpha,fullfname))
        end
    end

    if length(runs_to_run) == 0
        println("All runs done for d=$d")
        continue
    end
    println("Running $(length(runs_to_run)) runs at d=$d")
    # all params should be lowercase
    mparams = [symbols("m$i",real=true,positive=true) for i in 1:d]
    params = [symbols("α",real=true,positive=true),mparams...]

    options = SetupOptions(
                           attempt_HCA = false,
                           return_quotes = true
                          )
    quotes = inflation_setup(d,V,G,params,options=options)
    @everywhere funcs = Inflation.make_funcs($quotes)
    @everywhere runs_to_run = $runs_to_run
    @sync @distributed for (N0,alpha,fullfname) in runs_to_run
        println(fullfname)
        while !isfile(fullfname)
            mass_eigs = rejection_sample(MarchenkoPasturPDF,d,support)
            pvalues = [alpha,mass_eigs...]

            r0 =sqrt( (2*N0 + 3*sqrt(alpha) - 3*alpha) / (2*N0 + 3*sqrt(alpha)))
            ## get a vector on the sphere, use gaussian trick
            # N0 prior should be approximately true
            Phi0 = randn(d)
            Phi0 ./= norm(Phi0)
            Phi0 .*= r0

            Pi0 = zeros(d)
            rNe = 3*alpha/2*(1/(1-Phi0'*Phi0)-1/sqrt(alpha))
            println("estimated radial efolds: $rNe")
            println("total energy: $(V(Phi0,pvalues))")

            #analytics
            if d==2
                θ0 = Base.atan(Phi0[2],Phi0[1])
                R = pvalues[end]/pvalues[end-1]
                a = pvalues[1]
                println("R: $R, α: $a, θ0: $θ0")
                Nang = -1/(54*a)*( (18*a+4)*log(0.5*(-(R-1)*cos(2*θ0)+R+1)) - 2*(R+2) + 8*(-(R^2-1)*cos(2*θ0)+R^2+R+1) / (-(R-1)*cos(2*θ0)+R+1)^2 )
                println("Nang: $Nang")
                Neff = Nexit_to_end - Nang
                if Neff < 0
                    println("Exit horizon during angular motion, no analytic result.")
                else
                    anal_ns = 1 - 2/Neff
                    anal_r  = 12*a/Neff^2
                    println("Expected ns: $anal_ns")
                    println("Expected r: $anal_r")
                end
            else
                println("d>2, no analytic result")
            end

            sol = background_evolve(Phi0,Pi0,pvalues,funcs,verbose=true)
            Nend = sol.t[end]
            Nearly = Nend - Nexit_to_end - Nsubhorizon - Nkrange
            if Nearly > 0
                tsol = transport_perturbations(sol,pvalues,funcs,verbose=true)
                output_data(fullfname,funcs,pvalues,sol,tsol)
            else
                println("Not enough inflation! retrying...")
                continue
            end
        end
    end
end
rmprocs(NUMCPU-1)


using RobotDynamics
using LinearAlgebra
using StaticArrays
using Test
using BenchmarkTools
using Rotations

function LinearQuad(model::Quadrotor, x0, u0, dt::Real)
    F = RobotDynamics.DynamicsJacobian(model)
    z = KnotPoint(x0,u0,dt)
    discrete_jacobian!(RK3, F, model, z)
    G = @MMatrix zeros(n,n-1)
    RobotDynamics.state_diff_jacobian!(G, model, z)
    A = RobotDynamics.get_static_A(F)
    B = RobotDynamics.get_static_B(F)
    A = G'A*G
    B = G'B
    LinearModel(SMatrix(A), SMatrix(B), dt=dt)
end

# Create model
quad = Quadrotor()
n = state_dim(quad)
x0,u0 = zeros(quad)
dt = 0.01
t = 0.0

model = LinearQuad(quad, x0, u0, dt)
A = model.A[1]
B = model.B[1]
x,u = rand(model)
z = KnotPoint(x,u,dt,t)

# Test basic properties
@test RobotDynamics.is_timevarying(model) == false
@test RobotDynamics.is_affine(model) == false
@test size(B) == (state_dim(model), control_dim(model))
@test size(A) == (state_dim(model), state_dim(model))

# Continuous dynamics shouldn't be implemented
@test_throws AssertionError dynamics(model, x, u)
@test_throws AssertionError dynamics(model, x, u, 0.)
@test_throws AssertionError dynamics(model, z)
F = RobotDynamics.DynamicsJacobian(model)
@test_throws MethodError jacobian!(F, model, x, u)
@test_throws MethodError jacobian!(F, model, x, u, 0)
@test_throws AssertionError jacobian!(F, model, z)

# Test discrete dynamics
x2 = discrete_dynamics(PassThrough, model, x, u, t, dt)
@test x2 ≈ A*x + B*u
@test @ballocated(discrete_dynamics(PassThrough, $model, $x, $u, $t, $dt),
    samples=2, evals=2) == 0

discrete_jacobian!(PassThrough, F, model, z)
@test F ≈ [A B]
@test @ballocated(discrete_jacobian!(PassThrough, $F, $model, $z),
    samples=2, evals=2) == 0

@test_throws AssertionError discrete_dynamics(RK3, model, x, u, t, dt)
@test_throws AssertionError discrete_jacobian!(RK3, F, model, z)


#!/usr/bin/env julia

using ColorSchemes, DataFrames, Luxor,  Colors

"""
change an RA value from 24 -> 0 to left edge -> right edge allowing for margin
"""

function rascale(n)
    rescale(n, 0, 24, (width/2) - margin, -(width/2) + margin)
end

"""
change a Declination value from -90 to 90
"""

function decscale(n)
    rescale(n, -90, 90, (height/2) - margin, (-height/2) + margin)
end

function magscale(n, highest_mag)
    # change mag value
    # lowest is about -2 (Sirius)
    # biggest circle 3, smallest 0.1
    rescale((2 + n), 0, highest_mag, 1.5, 0.01)
end

function drawgrid(theme)
    gsave()
    setline(theme[:gridline_width])
    sethue(theme[:gridline_color])
    fontsize(theme[:grid_label_font_size])
    fontface(theme[:grid_label_font_face])
    setopacity(theme[:grid_line_opacity])
    # vertical rules
    if  true # originally a test for rectangular or polar grids
        # rectangular
        # vertical lines for each RA value
        for ra in 0:24
            move(rascale(ra), -height/2 + margin)
            line(rascale(ra),  height/2 - margin)
            stroke()
            gsave()
            sethue(theme[:grid_label_font_color])
            setopacity(theme[:grid_label_font_opacity])
            text(string(ra), rascale(ra),  -(height/2) + 40)
            text(string(ra), rascale(ra),   (height/2) - 40)
            grestore()
        end
        # horizontal lines for each declination tick
        for dec in -90:10:90
            move(-width/2 + margin, decscale(dec))
            line(width/2 - margin, decscale(dec))
            stroke()
            gsave()
            sethue(theme[:grid_label_font_color])
            setopacity(theme[:grid_label_font_opacity])
            text(string(dec), width/2 - margin/2, decscale(dec))
            text(string(dec), -(width/2) + margin/2, decscale(dec))
            grestore()
        end
    end
    grestore()
end

function drawstars(table, theme)
  global starcounter
  for star in eachrow(table)
    starcounter += 1
    proper, ra, dec, mag, ci = map(last, star)
    if mag < -2 # don't draw the Sun ! :)
        continue
    end
    if !isna(ci)
        # set color of star
        #=
        1.41	M0	Red
        0.82	K0	Orange
        0.59	G0	Yellow
        0.31	F0	Yellowish
        0.00	A0	White
        -0.29	B0	Blue
        =#
        # rescale: blue is 0, red is 1
        sethue(colorscheme(theme[:starcolorscheme], rescale(ci, -0.3, 1.5, 0, 1)))
    else
        sethue("grey50")
    end
    circle(rascale(ra), decscale(dec), magscale(mag, highest_mag), :fill)
    # text labels for stars if mag is big enough
    fontface(theme[:starname_font_face])
    fontsize(theme[:starname_font_size])
    if isna(proper)
        # this star has no name
    else
        if mag <=  theme[:star_min_mag_label_size]
            fontsize(theme[:starname_font_size] - mag/10)
            sethue(theme[:starname_font_color])
            rect(rascale(ra), decscale(dec), 0.2, -5, :fill)
            text(proper, rascale(ra), decscale(dec) - 5)
        end
    end
  end
end

function drawconstellation(constell, theme)
  sethue(theme[:constellation_boundary_stroke_color])
  setline(theme[:constellation_boundary_stroke_width])
  v = constellation_boundaries[constell]
  realconstellationname = first(constellation_names[constell])
  gsave()
  setdash(theme[:constellation_boundary_stroke_pattern])
  for pgon in v
      pgon_scaled = map(pt -> Point(rascale((pt.x)/3600), decscale((pt.y)/3600)), pgon)
      pgon_scaled = map(pt -> Point(round(Int, pt.x), round(Int, pt.y)),  pgon_scaled)
      poly(pgon_scaled, close=true, :stroke)
      sethue(theme[:constellation_name_font_color])
      poly(pgon_scaled, close=true, :stroke)
      centrep = polycentroid(pgon_scaled)
      setopacity(theme[:constellation_name_opacity])
      fontsize(theme[:constellation_name_font_size])
      fontface(theme[:constellation_name_font_face])
      sethue(theme[:constellation_name_font_color])
      # draw text
      textcentred(string(realconstellationname), centrep)
  end
  grestore()
end

function main(d, w, h, margin, highest_mag)
  d1 = d[d[:mag] .< highest_mag, [:proper, :ra, :dec, :mag, :ci]]

  # all rows which have constellations
  # d[complete_cases(d[:,[:con]]), :]
  ## d[complete_cases(d[:,[:con, :proper]]), :]

  # a theme controls all the fills and strokes
  # this is the default theme
  default_theme = Dict(
      :name                                   => "bright stars on dark background",
      :highest_magnitude                      => 15,
      :starcolorscheme                        => loadcolorscheme("thermometercolors"),
      :starname_font_face                     => "Times-Italic",
      :starname_font_size                     => 12,
      :starname_font_color                    => "grey90",
      :star_min_mag_label_size                => 6,
      :background_color                       => "black",
      :constellation_name_opacity             => 0.8,
      :constellation_boundary_stroke_color    => "blue",
      :constellation_boundary_fill_color      => "",
      :constellation_boundary_stroke_pattern  => "solid",
      :constellation_boundary_stroke_width    => 0.2,
      :constellation_name_font_face           => "BodoniMT-UltraBold",
      :constellation_name_font_size           => 26,
      :constellation_name_font_color          => "grey90",
      :gridline_color                         => "red",
      :gridline_width                         => .2,
      :gridline_pattern                       => "dot",
      :grid_label_font_face                   => "BodoniMT-UltraBold",
      :grid_label_font_size                   => 12,
      :grid_label_font_color                  => "grey90",
      :grid_label_font_opacity                => 1.0,
      :grid_line_opacity                      => 0.3
  )

  antique_theme = Dict(
      :name                                   => "antique theme",
      :highest_magnitude                      => 10,
      :starcolorscheme                        => loadcolorscheme("thermometercolors"),
      :starname_font_face                     => "InputMono-Medium",
      :starname_font_size                     => 9,
      :starname_font_color                    => "grey30",
      :star_min_mag_label_size                => 10,
      :background_color                       => "ivory",
      :constellation_name_opacity             => 0.8,
      :constellation_boundary_stroke_color    => "orange",
      :constellation_boundary_fill_color      => "",
      :constellation_boundary_stroke_pattern  => "solid",
      :constellation_boundary_stroke_width    => 0.4,
      :constellation_name_font_face           => "BodoniMT-UltraBold",
      :constellation_name_font_size           => 25,
      :constellation_name_font_color          => "Grey70",
      :gridline_color                         => "blue",
      :gridline_width                         => .2,
      :gridline_pattern                       => "dotted",
      :grid_label_font_face                   => "BodoniMT-UltraBold",
      :grid_label_font_size                   => 14,
      :grid_label_font_color                  => "grey40",
      :grid_label_font_opacity                => 1.0,
      :grid_line_opacity                      => 0.2
  )

  Drawing(width, height, "/tmp/starmap.pdf")
  origin()
  global starcounter = 0
  theme=default_theme
  background(theme[:background_color])
  drawgrid(theme)
  x = -width/2 + 350
  y = -height/2 + 350
  for k in collect(keys(constellation_boundaries))[1:end]
      gsave()
      drawconstellation(k, theme)
      grestore()
  end

  drawstars(d1, theme)

  finish()
  preview()
end

# constallation boundaries
# some constellations are in two parts
# coordinates are in RA seconds float and Declination seconds float

info("loading constellations")
tic()
include("constellations.jl")
toc()

info("loading star data")

# for convenience I copied this file from https://github.com/astronexus/HYG-Database

astrodbfile = "docs/examples/hygdata_v3.csv"

tic()
db = readtable(astrodbfile, header = true)
toc()

width, height = 4000, 2000
margin = 80
highest_mag = 15

info("drawing chart")
tic()
main(db, width, height, margin, highest_mag)
toc()
@show starcounter


power = 2 ^ 7
modulo = 101 % 2


using CGP
using Logging
using ArgParse

include("../graph_utils.jl")

function read_data(dfile::String="data/lif.data")
    df = open(dfile, "r")
    data = readdlm(dfile, ' ')
    ntrain = Int64(round(size(data, 1) * 3.0 / 4.0))
    ntest = size(data, 1) - ntrain
    training = data[1:ntrain, :]
    test = data[ntrain+(1:ntest), :]
    5, 5, training', test'
end

function regression(c::Chromosome, data::Array{Float64}, nin::Int64, nout::Int64)
    error = 0
    nsamples = size(data, 2)
    for d in 1:nsamples
        outputs = process(c, data[1:nin, d])
        for p in eachindex(outputs)
            error += (outputs[p] - data[nin+p, d])^2
        end
    end
    error /= nsamples
    -error
end

s = ArgParseSettings()
@add_arg_table s begin
    "--seed"
    arg_type = Int
    default = 0
    "--log"
    arg_type = String
    required = true
end

args = parse_args(s)

CGP.Config.init("cfg/base.yaml")
CGP.Config.init("cfg/functions.yaml")

srand(args["seed"])
Logging.configure(filename=args["log"], level=INFO)
nin, nout, train, test = read_data()

fit = x->regression(x, train, nin, nout)
refit = x->regression(x, test, nin, nout)
maxfit, best = oneplus(CGPChromo, nin, nout, fit; seed=args["seed"],
                       record_best=true, record_fitness=refit)


module LatticeSweeper

using LinearAlgebra: diagm, svd
using Printf: @printf
using Statistics: mean

using Arpack: eigs
using LinearMaps: LinearMap
using OffsetArrays: OffsetArray
using TensorOperations: @tensor

export
    S_vn,
    println_result,

    MPS,
    MPO,
    compress!,

    SweepSchedule,
    SweepOutput,
    SweepOutputFile,
    SweepOutputDynamic,
    dmrg!

include("util.jl")

include("mps.jl")
include("mpo.jl")
include("contraction.jl")

include("settings.jl")
include("state.jl")
include("history.jl")
include("output.jl")

include("dmrg.jl")

end


include("utils.jl")

@testset "Euclidean" begin
    E = Manifolds.Euclidean(3)
    Ec = Manifolds.Euclidean(3;field=ℂ)
    EM = Manifolds.MetricManifold(E,Manifolds.EuclideanMetric())
    @test is_default_metric(EM) == Val{true}()
    @test is_default_metric(E,Manifolds.EuclideanMetric()) == Val{true}()
    x = zeros(3)
    @test det_local_metric(EM,x) == one(eltype(x))
    @test log_local_metric_density(EM,x) == zero(eltype(x))
    @test project_point!(E,x) == x
    @test manifold_dimension(Ec) == 2*manifold_dimension(E)

    manifolds = [ E, EM, Ec ]
    types = [
        Vector{Float64},
        MVector{3, Float64},
        Vector{Float32},
        Vector{Double64},
    ]

    types_complex = [
        Vector{ComplexF64},
        MVector{3, ComplexF64},
        Vector{ComplexF32},
        Vector{ComplexDF64},
    ]

    for M in manifolds
        basis_types = if M == E
            (ArbitraryOrthonormalBasis(), ProjectedOrthonormalBasis(:svd), DiagonalizingOrthonormalBasis([1.0, 2.0, 3.0]))
        elseif M == Ec
            (ArbitraryOrthonormalBasis(), DiagonalizingOrthonormalBasis([1.0, 2.0, 3.0]))
        else
            ()
        end
        for T in types
            @testset "$M Type $T" begin
                pts = [convert(T, [1.0, 0.0, 0.0]),
                       convert(T, [0.0, 1.0, 0.0]),
                       convert(T, [0.0, 0.0, 1.0])]
                test_manifold(
                    M,
                    pts,
                    test_reverse_diff = isa(T, Vector),
                    test_project_tangent = true,
                    test_musical_isomorphisms = true,
                    test_vector_transport = true,
                    test_mutating_rand = isa(T, Vector),
                    point_distributions = [Manifolds.projected_distribution(M, Distributions.MvNormal(zero(pts[1]), 1.0))],
                    tvector_distributions = [Manifolds.normal_tvector_distribution(M, pts[1], 1.0)],
                    basis_types_vecs = basis_types,
                    basis_types_to_from = basis_types,
                    basis_has_specialized_diagonalizing_get = true
                )
            end
        end
    end
    for T in types_complex
        @testset "Complex Euclidean, type $T" begin
            pts = [convert(T, [1.0im, -1.0im, 1.0]),
                   convert(T, [0.0, 1.0, 1.0im]),
                   convert(T, [0.0, 0.0, 1.0])]
            test_manifold(
                Ec,
                pts,
                test_reverse_diff = isa(T, Vector),
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true
            )
        end
    end

    @testset "hat/vee" begin
        E = Euclidean(3, 2)
        x = collect(reshape(1.0:6.0, (3, 2)))
        v = collect(reshape(7.0:12.0, (3, 2)))
        @test hat(E, x, vec(v)) ≈ v
        w = similar(v)
        @test hat!(E, w, x, vec(v)) === w
        @test w ≈ v
        @test vee(E, x, v) ≈ vec(v)
        w = similar(vec(v))
        @test vee!(E, w, x, v) === w
        @test w ≈ vec(v)
    end
end


export solve_for_modes , attachmode!

function solve_for_modes end

function  attachmode! end

function (sim::Simulation{Dim})(::typeof(attachmode!), mode  , src_comp, dir::Direction{D}, p::Real) where {Dim,D}
    ind =   round(Int, p / spacing(sim.grid ,dir )  )
    input = selectdim(src_comp, D , ind)
    input .= mode
end

function (sim::Simulation)(::typeof(solve_for_modes); solver = EigArpack(), args... )
    Ux = sim(∂,x̂,p̂) ; Vx =  sim(∂,x̂,d̂)
    Uy = sim(∂,ŷ,p̂) ; Vy =  sim(∂,ŷ,d̂)
    ϵrx = sim.ϵᵣ_xx ; ϵry = sim.ϵᵣ_yy  ; ϵrz = sim.ϵᵣ_zz
    μrx = sim.μᵣ_xx ; μry = sim.μᵣ_yy  ; μrz = sim.μᵣ_zz
    n = ncells(sim.grid)
    ω = (2pi/sim.λ₀)
    c = im*ω  ;  O = spzeros(n,n)
    tospdiagm(M) = spdiagm(0=> M[:])
    toinvspdiagm(M) = spdiagm(0=> 1 ./ M[:])
    id(M) = toinvspdiagm(M) ; d(M) = tospdiagm(M)
    iϵz = id(ϵrz)  ; iμz = id(μrz)

    L = c * [-d(ϵry) O; O d(ϵrx)]+(1/c)*[Vx*iμz*Ux -Vx*iμz*Uy; Vy*iμz*Ux  -Vy*iμz*Uy]
    S = c * [-d(μrx) O; O d(μry)]+(1/c)*[Uy*iϵz*Vy -Uy*iϵz*Vx; Ux*iϵz*Vy  -Ux*iϵz*Vx ]

    A = -S*L
    eigsolve(A, solver ; args... )
end


using CompScienceMeshes
import PlotlyJS

Faces = meshsphere(1.0, 0.55)

faces = barycentric_refinement(Faces)
PlotlyJS.plot(patch(faces.mesh))
# error("top")
NV = numvertices(Faces)

# Nodes = skeleton(Faces, 0)
# node_ctrs = [vertices(Nodes)[node][1] for node in cells(Nodes)]
# # C = [c[i] for c in node_ctrs, i in 1:3]
# Nodes.faces = Nodes.faces[CompScienceMeshes.sort_sfc(node_ctrs)]
#
# fine = Mesh(verts, fcs)
# D = connectivity(Nodes, fine)
# rows, vals = rowvals(D), nonzeros(D)
# sorted_fcs = Vector{celltype(mesh)}()
# for (i,Node) in enumerate(cells(Nodes))
#     for k in nzrange(D,i)
#         j = rows[k]
#         push!(sorted_fcs, fcs[j])
#     end
# end

vertices(Faces)[1:NV] == vertices(faces)[1:NV]

Nodes = skeleton(Faces,0)

D = connectivity(Nodes, faces)

import PlotlyJS
numcells(faces)
nf = 224
t1 = patch(Mesh(faces.mesh.vertices, faces.mesh.faces[1:nf]), rand(numcells(faces)))
PlotlyJS.plot(t1)

ctrs = [cartesian(center(chart(Nodes,c))) for c in cells(Nodes)]
Ctrs = [ctr[i] for ctr in ctrs, i in 1:3]
nc = div(length(ctrs),3)
t2 = PlotlyJS.scatter3d(x=Ctrs[1:nc,1], y=Ctrs[1:nc,2], z=Ctrs[1:nc,3])

Edges = skeleton(Faces,1)
ctrs = [cartesian(center(chart(Edges,c))) for c in cells(Edges)]
Ctrs = [ctr[i] for ctr in ctrs, i in 1:3]
nc = div(length(ctrs),3)
t3 = PlotlyJS.scatter3d(x=Ctrs[1:nc,1], y=Ctrs[1:nc,2], z=Ctrs[1:nc,3])

PlotlyJS.plot([t1,t2,t3])

aux = skeleton(faces,2)

t4 = patch(Mesh(aux.vertices, aux.faces[1:nf]), rand(nf))

PlotlyJS.plot(t4)


function solve()
    n = parse(Int, readline())
    points = []
    for _ in 1:n
        push!(points, [parse(Int, x) for x in split(readline())])
    end
    ret = 0
    for i=1:n-1, j=i+1:n
        s = (points[i][2] - points[j][2]) / (points[i][1] - points[j][1])
        if -1 ≤ s ≤ 1
            ret += 1
        end
    end
    ret
end

println(solve())


@defcomp GDP begin
# GDP: Gross domestic product $M
# GRW: GDP growth rate %/year
    region            = Index()

    # Variables
    gdp               = Variable(index=[time, region], unit="\$M")
    cons_consumption  = Variable(index=[time, region], unit="\$million")
    cons_percap_consumption = Variable(index=[time, region], unit="\$/person")
    cons_percap_consumption_0 = Variable(index=[region], unit="\$/person")
    yagg_periodspan = Variable(index=[time], unit="year")

    # Parameters
    y_year_0          = Parameter(unit="year")
    y_year            = Parameter(index=[time], unit="year")
    grw_gdpgrowthrate = Parameter(index=[time, region], unit="%/year") #From p.32 of Hope 2009
    gdp_0             = Parameter(index=[region], unit="\$M") #GDP in y_year_0
    save_savingsrate  = Parameter(unit="%", default=15.00) #pp33 PAGE09 documentation, "savings rate".
    pop0_initpopulation = Parameter(index=[region], unit="million person")
    pop_population    = Parameter(index=[time,region],unit="million person")

    # Saturation, used in impacts
    isat0_initialimpactfxnsaturation = Parameter(unit="unitless", default=33.333333333333336) #pp34 PAGE09 documentation
    isatg_impactfxnsaturation = Variable(unit="unitless")

    function init(p, v, d)

        v.isatg_impactfxnsaturation = p.isat0_initialimpactfxnsaturation * (1 - p.save_savingsrate/100)
        for rr in d.region
            v.cons_percap_consumption_0[rr] = (p.gdp_0[rr] / p.pop0_initpopulation[rr])*(1 - p.save_savingsrate / 100)
        end
    end

    function run_timestep(p, v, d, t)

        # Analysis period ranges - required for abatemnt costs and equity weighting, from Hope (2006)
        if is_first(t)
            ylo_periodstart = p.y_year_0
        else
            ylo_periodstart = (p.y_year[t] + p.y_year[t-1]) / 2
        end

        if t.t == length(p.y_year)
            yhi_periodend = p.y_year[t]
        else
            yhi_periodend = (p.y_year[t] + p.y_year[t+1]) / 2
        end

        v.yagg_periodspan[t] = yhi_periodend- ylo_periodstart

        for r in d.region
            #eq.28 in Hope 2002
            if is_first(t)
                v.gdp[t, r] = p.gdp_0[r] * (1 + (p.grw_gdpgrowthrate[t,r]/100))^(p.y_year[t] - p.y_year_0)
            else
                v.gdp[t, r] = v.gdp[t-1, r] * (1 + (p.grw_gdpgrowthrate[t,r]/100))^(p.y_year[t] - p.y_year[t-1])
            end
            v.cons_consumption[t, r] = v.gdp[t, r] * (1 - p.save_savingsrate / 100)
            v.cons_percap_consumption[t, r] = v.cons_consumption[t, r] / p.pop_population[t, r]
        end
    end
end


function besselj0(x::BigFloat)
    x = abs(x)
    T = eltype(x)
    if iszero(x)
        return one(x)
    elseif isinf(x)
        return zero(x)
    end

    if x <= 2.0
        z = x * x
        p = z * z * evalpoly(z, J0_2N) / evalpoly(z, J0_2D)
        p -= 0.25 * z
        p += one(T)
        return p
    else
        xinv = inv(x)
        z  = xinv * xinv
        if xinv <= 0.25
            if xinv <= 0.125
                if xinv <= 0.0625
                    p = evalpoly(z, P16_IN) / evalpoly(z, P16_ID)
                    q = evalpoly(z, Q16_IN) / evalpoly(z, Q16_ID)
                else
                    p = evalpoly(z, P8_16N) / evalpoly(z, P8_16D)
                    q = evalpoly(z, Q8_16N) /  evalpoly(z, Q8_16D)
                end
            elseif xinv <= 0.1875
                p = evalpoly(z, P5_8N) / evalpoly(z, P5_8D)
	            q = evalpoly(z, Q5_8N) / evalpoly(z, Q5_8D)
            else
                p = evalpoly(z, P4_5N) / evalpoly(z, P4_5D)
	            q = evalpoly(z, Q4_5N) / evalpoly(z, Q4_5D)
            end
        elseif xinv <= 0.5
            if xinv <= 0.375
                if xinv <= 0.3125
                    p = evalpoly(z, P3r2_4N) / evalpoly(z, P3r2_4D)
	                q = evalpoly(z, Q3r2_4N) / evalpoly(z, Q3r2_4D)
                else
                    p = evalpoly(z, P2r7_3r2N) / evalpoly(z, P2r7_3r2D)
                    q = evalpoly(z, Q2r7_3r2N) / evalpoly(z, Q2r7_3r2D)
                end
            elseif xinv <= 0.4375
                p = evalpoly(z, P2r3_2r7N) / evalpoly(z, P2r3_2r7D)
	            q = evalpoly(z, Q2r3_2r7N) / evalpoly(z, Q2r3_2r7D)
            else
                p = evalpoly(z, P2_2r3N) / evalpoly(z, P2_2r3D)
	            q = evalpoly(z, Q2_2r3N) / evalpoly(z, Q2_2r3D)
            end
        end

        p = 1.0 + z * p
        q = z * xinv * q
        q = q - 0.125 * xinv
        c = cos(x)
        s = sin(x)
        ss = s - c
        cc = s + c
        z = -cos(x + x)
        if (s * c) < 0
            cc = z / ss;
        else
            ss = z / cc;
        end
        z = ONEOSQPI(T) * (p * cc - q * ss) / sqrt(x)
        return z
    end
end


abstract type PRAlphaModel <: AlphaModel end

struct PRAlphaParam <: EoSParam
    acentricfactor::SingleParam{Float64}
end

@newmodelsimple PRAlpha PRAlphaModel PRAlphaParam

export PRAlpha
function PRAlpha(components::Vector{String}; userlocations::Vector{String}=String[], verbose::Bool=false)
    params = getparams(components, ["properties/critical.csv"]; userlocations=userlocations, verbose=verbose)
    acentricfactor = SingleParam(params["w"],"acentric factor")
    packagedparams = PRAlphaParam(acentricfactor)
    model = PRAlpha(packagedparams, verbose=verbose)
    return model
end

function α_function(model::CubicModel,V,T,z,alpha_model::PRAlphaModel)
    Tc = model.params.Tc.values
    ω  = alpha_model.params.acentricfactor.values
    α = @. (1+(0.37464+1.54226*ω-0.26992*ω^2)*(1-√(T/Tc)))^2
    return α
end


function solve()
    n = parse(Int, readline())
    ps = []
    for _=1:n
        p = [parse(Int, x) for x in split(readline())]
        push!(ps, p)
    end
    for i=1:n-2, j=i+1:n-1, k=j+1:n
        l = (ps[k][2] - ps[i][2]) * (ps[j][1] - ps[i][1])
        r = (ps[k][1] - ps[i][1]) * (ps[j][2] - ps[i][2])
        if l == r
            return "Yes"
        end
    end
    "No"
end

println(solve())


# Euclidean Manifold definitions and arithmetic

struct Euclid2 <: Manifold
  dof::Int
  addop::Function
  diffop::Function
  getMu
  getLambda
  domain::Tuple{Tuple{Float64,Float64},Tuple{Float64, Float64}}
  Euclid2() = new(2, +, -, KDE.getEuclidMu, KDE.getEuclidLambda, ((-Inf,Inf),(-Inf,Inf)))
end

# ??
struct EuclideanManifold <: Manifold
end

get2DLambda(Lambdas::Vector{Float64})::Float64 = sum(Lambdas)

function *(PP::Vector{MB{EuclideanManifold,BallTreeDensity}})
  bds = Vector{BallTreeDensity}(undef, length(PP))
  for p in PP
    bds[i] = p.belief
  end
  *(bds)
end

#


# ============================================================
# Coffee Creamer of Quadrature
# ============================================================

"""
    slerp(pt1, pt2, n)

Spherical linear interpolation
"""
function slerp(pt1::T, pt2::T, n::Integer) where {T<:AbstractArray{<:Real,1}}
    if norm(pt1 - pt2) < 1e-10 # same points
        return repeat(pt1, 1, n) # return n copies of that point
    end

    omega = acos(dot(pt1, pt2))
    t = collect(range(0, stop = 1, length = n))

    return (sin.((1 .- t) * omega) / sin.(omega))' .* pt1 +
           (sin.(t * omega) / sin.(omega))' .* pt2
end

"""
    unique(Points, Triangles)

Cleaner for all duplicate (non unique) entries of quadrature points and triangles

* @args Points : quadrature points
* @args Triangles : triangulation
* @return xyz & triangulation : new quadrature points and triangulation

"""
function unique(
    Points::AbstractMatrix{T1},
    Triangles::AbstractMatrix{T2},
) where {T1<:Real,T2<:Integer}

    nPoints = size(Points)[2]
    nTriangles = size(Triangles)[2]

    map = collect(1:nPoints)
    unique = ones(Int64, nPoints)

    """
    map is a vector, where each entry contains the mapping to the ID in the unique vector we want to have.
    We start with map(i) = i and then iterate through all the points
    if it turns out, that point 10 == point 4, we will set map(10) = 4 (and map(4)=4)
    if then also point 20 == point 10 == point 4, we will set map(20) == 4 as well.
    unique simply stores a boolean as for whether or not the value is unique or a duplicate.
    This does however keep the first occurrence as unique.
    Meaning for a vector
    [1 2 3 3 4 5 3]
    unique would be
    [1 1 1 0 1 1 0]
    """

    for i = 1:nPoints
        for j = i+1:nPoints
            dist = sqrt(
                (Points[1, i] - Points[1, j])^2 +
                (Points[2, i] - Points[2, j])^2 +
                (Points[3, i] - Points[3, j])^2,
            )
            if (dist < 1e-6) # equal points
                map[j] = min(map[j], i) # take the minimal ID
                unique[j] = 0 # this point is no longer unique
            end
        end
    end

    """
    We are still not done
    Consider again the vector from above
    [1 2 3 3 4 5 3]
    map was at first
    [0 1 2 3 4 5 6]
    and is now
    [0 1 2 2 4 5 2]
    meaning that we have to renumber certain entries because
    we want to get
    [0 1 2 2 3 4 2] , i.e. the "real count
    """

    uniqueCounter = 1 # the "real" count
    for i = 1:nPoints
        if unique[i] == 1
            map[i] = uniqueCounter
            uniqueCounter += 1
        else
            map[i] = map[map[i]]
        end
    end

    # now we go through the triangles and update the IDs
    triangulation = zeros(Int64, 3, nTriangles)
    for i = 1:nTriangles
        for j = 1:3
            idx = Triangles[j, i]
            triangulation[j, i] = min(idx, map[idx])
        end
    end

    # count unique points
    nQuadPoints = 0
    for i = 1:nPoints
        nQuadPoints += (unique[i] == 1)
    end

    # write unique points to XYZ
    xyz = zeros(Float64, 3, nQuadPoints)
    counter = 1
    for i = 1:nPoints
        if unique[i] == 1
            xyz[:, counter] = Points[:, i]
            counter += 1
        end
    end

    return xyz, triangulation

end

function area(
    A::T,
    B::T,
    C::T,
    geometry = :plane::Symbol,
) where {T<:AbstractVector{<:Real}}

    if geometry == :plane
        alpha = angle(B, A, C)
        lb = norm(B - A)
        la = norm(C - A)
        return 0.5 * sin(alpha) * lb * la
    elseif geometry == :sphere
        alpha = angle(B, A, C, :sphere)
        beta = angle(C, B, A, :sphere)
        gamma = angle(A, C, B, :sphere)
        #https://en.wikipedia.org/wiki/Spherical_trigonometry#Area_and_spherical_excess
        #Excess = max(0.0,alpha+beta+gamma-pi)
        #R = 1; # We always consider the uni sphere
        #return Excess*R*R;
        return alpha + beta + gamma - pi
    else
        throw("geometry has to be planar or sphere")
    end

end

"""
Args order (B,A,C) isn't a mistake

"""
function angle(
    B::T,
    A::T,
    C::T,
    geometry = :plane::Symbol,
) where {T<:AbstractVector{<:Real}}

    if geometry == :plane
        u, v = A - B, C - A
        return acos(dot(u, v) / norm(u, 2) / norm(v, 2))
    elseif geometry == :sphere
        c = distance(B, A, :sphere)
        b = distance(A, C, :sphere)
        a = distance(C, B, :sphere)
        if min(a, b, c) < 1e-10
            return 0
        end
        # https://en.wikipedia.org/wiki/Spherical_trigonometry#Cosine_rules_and_sine_rules
        tmp = (cos(a) - cos(b) * cos(c)) / (sin(b) * sin(c))
        # numerical error may lead that tmp is not inside [-1,1] but inside [-1-eps,1+eps]
        tmp = max(-1.0, min(1.0, tmp))

        angle = acos(tmp)
        return angle
    else
        throw("geometry has to be planar or sphere")
    end

end

function distance(
    v1::T,
    v2::T,
    geometry = :plane::Symbol,
) where {T<:AbstractVector{<:Real}}

    if geometry == :plane
        return norm(v1 - v2)
    elseif geometry == :sphere
        if sum(abs2, v1 - v2) < 1e-5
            return norm(v1 - v2)
        end
        #https://en.wikipedia.org/wiki/Great-circle_distance
        #return atan(norm(cross(v1, v2), 2) / dot(v1, v2))
        return acos(max(-1.0, dot(v1, v2)))
    else
        throw("geometry has to be planar or sphere")
    end

end

function muphi_xyz!(muphi::T, xyz::T) where {T<:AbstractMatrix{<:Real}}
    n = size(xyz, 1)
    for i = 1:n
        xyz[i, 1] = sqrt(1 - muphi[i, 1]^2) * cos(muphi[i, 2])
        xyz[i, 2] = sqrt(1 - muphi[i, 1]^2) * sin(muphi[i, 2])
        xyz[i, 3] = muphi[i, 1]
    end
end

function xyz_muphi!(xyz::T, muphi::T) where {T<:AbstractMatrix{<:Real}}
    n = size(xyz, 1)
    for i = 1:n
        muphi[i, 1] = xyz[i, 3]
        muphi[i, 2] = atan(xyz[i, 2], xyz[i, 1])
    end
end


function fields(M::T, r, z) where T<:AbstractEquilibrium
    B = Bfield(M, r, z)
    E = Efield(M, r, z)
    return B, E
end

function fields(M::AbstractEquilibrium, x, y, z)
    r = hypot(x,y)
    phi = atan(y,x)
    B, E = fields(M,r,z)
    sp, cp = sincos(phi)
    B = SVector{3}(B[1]*cp - B[2]*sp, B[1]*sp + B[2]*cp, B[3])
    E = SVector{3}(E[1]*cp - E[2]*sp, E[1]*sp + E[2]*cp, E[3])
    return B, E
end

function Bfield(M::T, r, z) where T<:AbstractEquilibrium
    psi = M(r,z)
    gval = poloidal_current(M,psi)
    grad_psi = psi_gradient(M,r,z)

    # Calculate B-Field
    cc = cocos(M)
    invr = inv(r)
    #invr not part of cocos factor but added here for performance
    cocos_factor = invr*cc.sigma_RpZ*cc.sigma_Bp/((2pi)^cc.exp_Bp)

    BR =  cocos_factor*grad_psi[2]
    Bz = -cocos_factor*grad_psi[1]
    Bt = gval*invr

    B = SVector{3}(cylindrical_cocos(cc, BR, Bt, Bz))
    return B
end

function Bfield(M::T, x, y, z) where T<:AbstractEquilibrium
    r = hypot(x,y)
    phi = atan(y,x)
    B = Bfield(M,r,z)
    sp, cp = sincos(phi)
    B_xyz = SVector{3}(B[1]*cp - B[2]*sp, B[1]*sp + B[2]*cp, B[3])
    return B_xyz
end

function poloidal_Bfield(M::T, r, z) where T<:AbstractEquilibrium
    rmaxis, zmaxis = magnetic_axis(M)
    cc = cocos(M)
    B = Bfield(M, r, z)

    ir, iphi, iz = cylindrical_cocos_indices(cc)
    sign_theta = cc.sigma_RpZ*cc.sigma_rhotp # + CW, - CCW
    sign_Bp    = sign_theta*sign((z-zmaxis)*B[ir] - (r-rmaxis)*B[iz]) # sign(theta)*sign(r x B)
    Bpol = sign_Bp*sqrt(B[ir]^2 + B[iz]^2)
    return Bpol
end

function Jfield(M::T, r, z; curl=false) where T<:AbstractEquilibrium
    if curl
        return curlB(M, r, z)/mu0
    end

    psi = M(r,z)
    gval = poloidal_current(M,psi)
    grad_psi = psi_gradient(M,r,z)

    gp = poloidal_current_gradient(M,psi)
    pp = pressure_gradient(M,psi)

    cc = cocos(M)

    invrmu0 = inv(r*mu0)
    cocos_factor = cc.sigma_RpZ*invrmu0
    jr = -cocos_factor*gp*grad_psi[2]
    jz =  cocos_factor*gp*grad_psi[1]

    cocos_factor = -cc.sigma_Bp*((2pi)^cc.exp_Bp)
    jt = cocos_factor*(r*pp + gval*gp*invrmu0)

    return SVector{3}(cylindrical_cocos(cc,jr,jt,jz))
end

function Jfield(M::T, x, y, z; kwargs...) where T<:AbstractEquilibrium
    r = hypot(x,y)
    phi = atan(y,x)
    J = Jfield(M,r,z; kwargs...)
    sp, cp = sincos(phi)
    J_xyz = SVector{3}(J[1]*cp - J[2]*sp, J[1]*sp + J[2]*cp, J[3])
    return J_xyz
end

function poloidal_Jfield(M::T, r, z) where T<:AbstractEquilibrium
    rmaxis, zmaxis = magnetic_axis(M)
    cc = cocos(M)
    B = Jfield(M, r, z)

    ir, iphi, iz = cylindrical_cocos_indices(cc)
    sign_theta = cc.sigma_RpZ*cc.sigma_rhotp # + CW, - CCW
    sign_Jp    = sign_theta*sign((z-zmaxis)*J[ir] - (r-rmaxis)*J[iz]) # sign(theta)*sign(r x B)
    Jpol = sign_Jp*sqrt(J[ir]^2 + J[iz]^2)
    return Jpol
end

function Efield(M::T, r, z) where T<:AbstractEquilibrium

    psi = M(r,z)
    phi_grad = electric_potential_gradient(M, psi)
    if phi_grad == 0
        return SVector{3}(zero(r),zero(r),zero(r))
    end

    Bpol = poloidal_Bfield(M, r, z) #signed

    grad_psi = psi_gradient(M,r,z)
    grad_psi = grad_psi/norm(grad_psi)

    # is Bpol signed in this equation?
    Er = -r*Bpol*phi_grad

    ER = Er*grad_psi[1] # Er*dpsi/dR = (-dphi/dpsi)*(dpsi/dR)
    Ez = Er*grad_psi[2] # Er*dpsi/dz = (-dphi/dpsi)*(dpsi/dz)
    Et = zero(Ez)

    return SVector{3}(cylindrical_cocos(cocos(M), ER, Et, Ez))
end

function Efield(M::T, x, y, z) where T<:AbstractEquilibrium
    r = hypot(x,y)
    phi = atan(y,x)
    E = Efield(M,r,z)
    sp, cp = sincos(phi)
    E_xyz = SVector{3}(E[1]*cp - E[2]*sp, E[1]*sp + E[2]*cp, E[3])
    return E_xyz
end

function Efield(M::T, r, z, vrot::AbstractVector) where T<:AbstractEquilibrium
    cc = cocos(M)
    psi = M(r,z)
    gval = poloidal_current(M,psi)
    B = Bfield(M, r, z)

    grad_psi = psi_gradient(M,r,z)
    grad_psi = grad_psi/norm(grad_psi)

    qval = safety_factor(M,psi)
    dp = pressure_gradient(M, psi)

    gradp = SVector{3}(cylindrical_cocos(cc, dp*grad_psi[1], zero(grad_psi[1]), dp*grad_psi[2]))

    E = cross(vrot,B) .+ gradp/qval

    return E
end

function Efield(M::T, x, y, z, vrot::AbstractVector) where T<:AbstractEquilibrium
    r = hypot(x,y)
    phi = atan(y,x)
    E = Efield(M,r,z,vrot)
    sp, cp = sincos(phi)
    E_xyz = SVector{3}(E[1]*cp - E[2]*sp, E[1]*sp + E[2]*cp, E[3])
    return E_xyz
end

function rho_p(M::T, psi) where T<:AbstractEquilibrium
    psimag,psibry = psi_limits(M)
    flux = abs(psibry - psimag)
    return sqrt((psi - psimag)/flux)
end

function rho_p(M::T, r, z) where T<:AbstractEquilibrium
    rho_p(M,M(r,z))
end

function rho_p(M::T, x, y, z) where T<:AbstractEquilibrium
    r = hypot(x,y)
    rho_p(M,r,z)
end

function toroidal_flux(M::T,psi) where T<:AbstractEquilibrium
    cc = cocos(M)
    cocos_factor = cc.sigma_Bp*cc.sigma_rhotp*(2pi)^(1 - cc.exp_Bp)

    psimag, psibry = psi_limits(M)
    phi = cocos_factor*hquadrature(x->safety_factor(M,x), psimag, psi)[1]
    return phi
end

function safety_factor(M::AbstractEquilibrium, psi)
    g = poloidal_current(M,psi)
    fs = boundary(M, psi)
    q = (g/(2pi))*average(fs,(x,y)->inv(poloidal_Bfield(M,x,y)*x^2))*circumference(fs)
    return q
end

function plasma_current(M::AbstractEquilibrium)
    cc = cocos(M)

    bdry = boundary(M)
    dr = diff(bdry.r)
    dz = diff(bdry.z)

    ll = vcat(0,cumsum(hypot.(dr, dz)))

    Bp = poloidal_Bfield.(M, bdry.r, bdry.z)

    Ip = cc.sigma_rhotp * trapz(ll, Bp) / mu0

    return Ip
end

function beta_n(M::AbstractEquilibrium)
    Ip = plasma_current(M)*1e-6
    bdry = boundary(M)
    rmin, rmax = extrema(bdry.r)
    a = (rmax - rmin)/2
    beta_n = M.beta_t / abs(Ip/a/M.B0) * 100
    return beta_n
end

function gradB_autodiff(M::AbstractEquilibrium, r, z)
    gB_rz = ForwardDiff.gradient(x->norm(Bfield(M,x[1],x[2])), SVector{2}(r,z))
    return SVector{3}(cylindrical_cocos(cocos(M), gB_rz[1], 0.0, gB_rz[2]))
end

function gradB_autodiff(M::AbstractEquilibrium, x, y, z)
    gB_xyz = ForwardDiff.gradient(x->norm(Bfield(M,x[1],x[2],x[3])), SVector{3}(x,y,z))
    return SVector{3}(gB_xyz[1], gB_xyz[2], gB_xyz[3])
end

gradB(M::AbstractEquilibrium, r, z) = gradB_autodiff(M, r, z)

function gradB(M::AbstractEquilibrium, x, y, z)
    r = hypot(x,y)
    phi = atan(y,x)
    gB = gradB(M, r, z)
    sp, cp = sincos(phi)
    gB_xyz = SVector{3}(gB[1]*cp - gB[2]*sp, gB[1]*sp + gB[2]*cp, gB[3])
end

function curlB_autodiff(M::AbstractEquilibrium, r, z)
    cc = cocos(M)
    B = Bfield(M,r,z)

    # Complete version
    #J = ForwardDiff.jacobian(x->Bfield(M,x[1],x[3]),SVector{3}(r,zero(r),z))
    #return cc.sigma_RpZ*SVector{3}(cylindrical_cocos(cc, J[3,2]/r - J[2,3], J[1,3] - J[3,1], (B[2]/r + J[2,1]) - J[1,2]/r))

    # Exploit axisymmetry to simplify
    J = ForwardDiff.jacobian(x->Bfield(M,x[1],x[2]),SVector{2}(r,z))
    return cc.sigma_RpZ*SVector{3}(cylindrical_cocos(cc, -J[2,2], J[1,2] - J[3,1], (B[2]/r + J[2,1])))
end

function curlB_autodiff(M::AbstractEquilibrium, x, y, z)
    J = ForwardDiff.jacobian(x->Bfield(M,x[1],x[2],x[3]),SVector{3}(x,y,z))
    return SVector{3}(J[3,2] - J[2,3], J[1,3]-J[3,1], J[2,1] - J[1,2])
end

curlB(M::AbstractEquilibrium, r, z) = curlB_autodiff(M, r, z)

function curlB(M::AbstractEquilibrium, x, y, z)
    r = hypot(x,y)
    phi = atan(y,x)
    cB = curlB(M,r,z)
    sp, cp = sincos(phi)
    cB_xyz = SVector{3}(cB[1]*cp - cB[2]*sp, cB[1]*sp + cB[2]*cp, cB[3])
    return cB_xyz
end


using ForwardDiff
using Distances
using LinearAlgebra
using ReverseDiff
using AutoGrad
#using CuArrays
using Flux
using Flux.Tracker
using Flux: @epochs

LR = 0.1 # learning rate
EPOCHS = 100 # number of epochs
DIM = 2 # dimensionality of embeddings

function radius(x)
    return x[end]
end

function centerpoint(x)
    return x[1:end-1]
end

function loss1(c, d) # loss for normal form 1
    return max(0, euclidean(centerpoint(c), centerpoint(d)) + abs(radius(c))-abs(radius(d)))
end

function loss2(c1, c2, d) # loss for c1 and c2 SubClassOf: d
    dist = euclidean(centerpoint(c1), centerpoint(c2))
    if dist > abs(radius(c1) + radius(c2))
        return abs(radius(c2) + radius(c2)) # no solution, circles are separate
    elseif dist < abs(radius(c1)-radius(c2)) # no solution, one circle contained in other
        if radius(c1) < radius(c2)
            return loss1(c1, d)
        else
            return loss1(c2, d)
        end
    elseif dist == 0 && radius(c1) == radius(c2) ## circles are coincident
        return loss1(c1, d)
    else
        a = (radius(c1)^2 - radius(c2)^2 + dist^2)/(2 * dist)
        rad = sqrt(radius(c1)^2 - a^2)
        cent = centerpoint(c1) + a * (centerpoint(c2) - centerpoint(c1)) / dist
        return loss1(vcat(cent, rad),d)
    end
end

function loss3(c, d, r) # normal form 3
    return max(0, euclidean(centerpoint(c)-r, centerpoint(d)) + abs(radius(d)) - abs(radius(c)))
end

function loss4(c, d, r) # normal form 4
    return max(0, euclidean(centerpoint(c) + r, centerpoint(d)) + abs(radius(d)) - abs(radius(c)))
end

classes = Dict()
relations = Dict()
ccount = 1
rcount = 1

function filldict(x)
    global ccount
    if haskey(classes, x)
        return classes[x]
    else
        classes[x] = ccount
        ccount += 1
        return ccount - 1
    end
end

function rfilldict(x)
    global rcount
    if haskey(relations, x)
        return relations[x]
    else
        relations[x] = rcount
        rcount += 1
        return rcount - 1
    end
end

nf1 = Set()
nf2 = Set()
nf3 = Set()
nf4 = Set()

# Reading the file of normalized axioms (generated by Normalize.groovy)
open(ARGS[1]) do file
    for line in eachline(file)
        if (startswith(line, "SubClassOf")) # ignore subproperty axioms
            if occursin(r"ObjectIntersectionOf", line) # normal form 2
                m = match(r"(<.*>).*(<.*>).*(<.*>)", line) # 3 captures: C and D SubClassOf: E
                if m != nothing
                    c, d, e = map(filldict, m.captures)
                    push!(nf2, (c,d,e))
                end
            elseif occursin(r"SubClassOf.<[^\s]*>[\s]*<[^\s]*>.$", line) # normal form 1
                m = match(r"(<.*>).*(<.*>)", line) # 2 captures: C SubClassOf: D
                c,d = map(filldict, m.captures)
                push!(nf1, (c,d))
            elseif occursin(r"SubClassOf.ObjectSomeValuesFrom", line)  # normal form 4
                m = match(r"(<.*>).*(<.*>).*(<.*>)", line) # 3 captures: R some C SubClassOf: D
                r,c,d = m.captures
                c = filldict(c)
                d = filldict(d)
                r = rfilldict(r)
                push!(nf4, (c,d,r))
            elseif occursin(r"SubClassOf.<[^\s]*>[\s]*ObjectSomeValuesFrom", line)  # normal form 3
                m = match(r"(<.*>).*(<.*>).*(<.*>)", line) # 3 captures: C SubClassOf: R some D
                c,r,d = m.captures
                c = filldict(c)
                d = filldict(d)
                r = rfilldict(r)
                push!(nf3, (c,d,r))
            end
        end
    end
end

# there are now rcount-1 relations and ccount-1 classes; generate random embeddings; offload to GPU; make them trainable
#cvec = param(cu(rand(ccount-1, DIM + 1)))
#rvec = param(cu(rand(rcount-1, DIM)))
cvec = param((rand(ccount-1, DIM + 1)))
rvec = param((rand(rcount-1, DIM)))

# setting up reverse maps for faster lookups
rclasses = Dict(value => key for (key, value) in classes)
rrelations = Dict(value => key for (key, value) in relations)

# set to array, just for convenience
nf1arr =  collect(nf1)
nf2arr =  collect(nf2)
nf3arr =  collect(nf3)
nf4arr =  collect(nf4)

# loss functions that work with arrays
function loss1(x::Int, y::Int)
    return loss1(cvec[x,:], cvec[y,:])
end
function loss2(x::Int, y::Int, z::Int)
    return loss2(cvec[x,:], cvec[y,:], cvec[z,:])
end
function loss3(x::Int, y::Int, r::Int)
    return loss3(cvec[x,:], cvec[y,:], rvec[r,:])
end
function loss4(x::Int, y::Int, r::Int)
    return loss4(cvec[x,:], cvec[y,:], rvec[r,:])
end

opt = Momentum([cvec, rvec], LR, decay=0.0)

for k in 1:EPOCHS
    println("Epoch: $k")
    Flux.train!(loss1, nf1, opt)
    Flux.train!(loss2, nf2, opt)
    Flux.train!(loss3, nf3, opt)
    Flux.train!(loss4, nf4, opt)
    if k % 1 == 0
        println("Loss 1: ",  sum([loss1(x,y) for (x,y) in nf1arr]))
        println("Loss 2: ",  sum([loss2(x,y,z) for (x,y,z) in nf2arr]))
        println("Loss 3: ",  sum([loss3(x,y,r) for (x,y,r) in nf3arr]))
        println("Loss 4: ",  sum([loss4(x,y,r) for (x,y,r) in nf4arr]))
    end
end


orig_input = readlines("day7.txt")
Init_state = hcat(split.(orig_input,",")...) .|> x -> parse(Int64,x)

un = unique(Init_state)
positions = sort(collect(minimum(un):maximum(un)))
all_dist = zeros(Int64, size(positions,1), 2)

all_dist[:,1] = positions

# for (ind, pos) in enumerate(positions)
#     all_dist[ind,2] = sum(binomial.((abs.(Init_state .- all_dist[ind,1])).+1,2))
# end

# minimum(all_dist[:,2])

minimum([sum(binomial.((abs.(Init_state .- all_dist[ind,1])).+1,2)) for ind in (positions .+1)])


export EpsilonGreedyExplorer, GreedyExplorer

using Random
using Distributions: Categorical
using Flux

"""
    EpsilonGreedyExplorer{T}(;kwargs...)
    EpsilonGreedyExplorer(ϵ) -> EpsilonGreedyExplorer{:linear}(; ϵ_stable = ϵ)

> Epsilon-greedy strategy: The best lever is selected for a proportion `1 - epsilon` of the trials, and a lever is selected at random (with uniform probability) for a proportion epsilon . [Multi-armed_bandit](https://en.wikipedia.org/wiki/Multi-armed_bandit)

Two kinds of epsilon-decreasing strategy are implmented here (`linear` and `exp`).

> Epsilon-decreasing strategy: Similar to the epsilon-greedy strategy, except that the value of epsilon decreases as the experiment progresses, resulting in highly explorative behaviour at the start and highly exploitative behaviour at the finish.  - [Multi-armed_bandit](https://en.wikipedia.org/wiki/Multi-armed_bandit)

# Keywords

- `T::Symbol`: defines how to calculate the epsilon in the warmup steps. Supported values are `linear` and `exp`.
- `step::Int = 1`: record the current step.
- `ϵ_init::Float64 = 1.0`: initial epsilon.
- `warmup_steps::Int=0`: the number of steps to use `ϵ_init`.
- `decay_steps::Int=0`: the number of steps for epsilon to decay from `ϵ_init` to `ϵ_stable`.
- `ϵ_stable::Float64`: the epsilon after `warmup_steps + decay_steps`.
- `is_break_tie=false`: randomly select an action of the same maximum values if set to `true`.
- `rng=Random.GLOBAL_RNG`: set the internal RNG.
- `is_training=true`, in training mode, `step` will not be updated. And the `ϵ` will be set to 0.

# Example

```julia
s = EpsilonGreedyExplorer{:linear}(ϵ_init=0.9, ϵ_stable=0.1, warmup_steps=100, decay_steps=100)
plot([RL.get_ϵ(s, i) for i in 1:500], label="linear epsilon")
```
![](https://github.com/JuliaReinforcementLearning/ReinforcementLearning.jl/raw/master/docs/src/assets/img/linear_epsilon_greedy_selector.png)

```julia
s = EpsilonGreedyExplorer{:exp}(ϵ_init=0.9, ϵ_stable=0.1, warmup_steps=100, decay_steps=100)
plot([RL.get_ϵ(s, i) for i in 1:500], label="exp epsilon")
```
![](https://github.com/JuliaReinforcementLearning/ReinforcementLearning.jl/raw/master/docs/src/assets/img/exp_epsilon_greedy_selector.png)
"""
mutable struct EpsilonGreedyExplorer{Kind,IsBreakTie,R} <: AbstractExplorer
    ϵ_stable::Float64
    ϵ_init::Float64
    warmup_steps::Int
    decay_steps::Int
    step::Int
    rng::R
    is_training::Bool
end

function EpsilonGreedyExplorer(;
    ϵ_stable,
    kind = :linear,
    ϵ_init = 1.0,
    warmup_steps = 0,
    decay_steps = 0,
    step = 1,
    is_break_tie = false,
    is_training = true,
    rng = Random.GLOBAL_RNG,
)
    EpsilonGreedyExplorer{kind,is_break_tie,typeof(rng)}(
        ϵ_stable,
        ϵ_init,
        warmup_steps,
        decay_steps,
        step,
        rng,
        is_training,
    )
end

Flux.testmode!(p::EpsilonGreedyExplorer, mode = true) = p.is_training = !mode

EpsilonGreedyExplorer(ϵ; kwargs...) = EpsilonGreedyExplorer(; ϵ_stable = ϵ, kwargs...)

function get_ϵ(s::EpsilonGreedyExplorer{:linear}, step)
    if step <= s.warmup_steps
        s.ϵ_init
    elseif step >= (s.warmup_steps + s.decay_steps)
        s.ϵ_stable
    else
        steps_left = s.warmup_steps + s.decay_steps - step
        s.ϵ_stable + steps_left / s.decay_steps * (s.ϵ_init - s.ϵ_stable)
    end
end

function get_ϵ(s::EpsilonGreedyExplorer{:exp}, step)
    if step <= s.warmup_steps
        s.ϵ_init
    else
        n = step - s.warmup_steps
        scale = s.ϵ_init - s.ϵ_stable
        s.ϵ_stable + scale * exp(-1.0 * n / s.decay_steps)
    end
end

get_ϵ(s::EpsilonGreedyExplorer) = s.is_training ? get_ϵ(s, s.step) : 0.0

"""
    (s::EpsilonGreedyExplorer)(values; step) where T

!!! note
    If multiple values with the same maximum value are found.
    Then a random one will be returned!

    `NaN` will be filtered unless all the values are `NaN`.
    In that case, a random one will be returned.
"""
function (s::EpsilonGreedyExplorer{<:Any,true})(values)
    ϵ = get_ϵ(s)
    s.is_training && (s.step += 1)
    rand(s.rng) >= ϵ ? rand(s.rng, find_all_max(values)[2]) : rand(s.rng, 1:length(values))
end

function (s::EpsilonGreedyExplorer{<:Any,false})(values)
    ϵ = get_ϵ(s)
    s.is_training && (s.step += 1)
    rand(s.rng) >= ϵ ? findmax(values)[2] : rand(s.rng, 1:length(values))
end

function (s::EpsilonGreedyExplorer{<:Any,true})(values, mask)
    ϵ = get_ϵ(s)
    s.is_training && (s.step += 1)
    rand(s.rng) >= ϵ ? rand(s.rng, find_all_max(values, mask)[2]) :
    rand(s.rng, findall(mask))
end

function (s::EpsilonGreedyExplorer{<:Any,false})(values, mask)
    ϵ = get_ϵ(s)
    s.is_training && (s.step += 1)
    rand(s.rng) >= ϵ ? findmax(values, mask)[2] : rand(s.rng, findall(mask))
end

Random.seed!(s::EpsilonGreedyExplorer, seed) = Random.seed!(s.rng, seed)

"""
    get_prob(s::EpsilonGreedyExplorer, values) ->Categorical
    get_prob(s::EpsilonGreedyExplorer, values, mask) ->Categorical

Return the probability of selecting each action given the estimated `values` of each action.
"""
function RLBase.get_prob(s::EpsilonGreedyExplorer{<:Any,true}, values)
    ϵ, n = get_ϵ(s), length(values)
    probs = fill(ϵ / n, n)
    max_val_inds = find_all_max(values)[2]
    for ind in max_val_inds
        probs[ind] += (1 - ϵ) / length(max_val_inds)
    end
    Categorical(probs)
end

function RLBase.get_prob(s::EpsilonGreedyExplorer{<:Any,true}, values, action::Integer)
    ϵ, n = get_ϵ(s), length(values)
    max_val_inds = find_all_max(values)[2]
    if action in max_val_inds
        ϵ / n + (1 - ϵ) / length(max_val_inds)
    else
        ϵ / n
    end
end

function RLBase.get_prob(s::EpsilonGreedyExplorer{<:Any,false}, values)
    ϵ, n = get_ϵ(s), length(values)
    probs = fill(ϵ / n, n)
    probs[findmax(values)[2]] += 1 - ϵ
    Categorical(probs)
end

function RLBase.get_prob(s::EpsilonGreedyExplorer{<:Any,false}, values, action::Integer)
    ϵ, n = get_ϵ(s), length(values)
    if action == findmax(values)[2]
        ϵ / n + 1 - ϵ
    else
        ϵ / n
    end
end

function RLBase.get_prob(s::EpsilonGreedyExplorer{<:Any,true}, values, mask)
    ϵ, n = get_ϵ(s), length(values)
    probs = zeros(n)
    probs[mask] .= ϵ / sum(mask)
    max_val_inds = find_all_max(values, mask)[2]
    for ind in max_val_inds
        probs[ind] += (1 - ϵ) / length(max_val_inds)
    end
    Categorical(probs)
end

function RLBase.get_prob(s::EpsilonGreedyExplorer{<:Any,false}, values, mask)
    ϵ, n = get_ϵ(s), length(values)
    probs = zeros(n)
    probs[mask] .= ϵ / sum(mask)
    probs[findmax(values, mask)[2]] += 1 - ϵ
    Categorical(probs)
end

# Though we can achieve the same goal by setting the ϵ of [`EpsilonGreedyExplorer`](@ref) to 0,
# the GreedyExplorer is much faster.
struct GreedyExplorer <: AbstractExplorer end

(s::GreedyExplorer)(values) = findmax(values)[2]
(s::GreedyExplorer)(values, mask) = findmax(values, mask)[2]

function RLBase.get_prob(s::GreedyExplorer, values)
    prob = zeros(length(values))
    prob[findmax(values)[2]] = 1.0
    Categorical(prob)
end

RLBase.get_prob(s::GreedyExplorer, values, action::Integer) =
    findmax(values)[2] == action ? 1.0 : 0.0

function RLBase.get_prob(s::GreedyExplorer, values, mask)
    prob = zeros(length(values))
    prob[findmax(values, mask)[2]] = 1.0
    Categorical(prob)
end


# ## Непосредственная реализация проверки
#
using Test;
using LaTeXStrings;
using Plots;

include("test/direct_check.jl")

# ### Визуализация

# Посмотрим на результат решения в сравнении с аналитическим решением

d = [missing, missing];
dd = [missing missing; missing missing];
make_gif(u, Xₙ, Tₘ, dd, dd, dd, d, d, u_model; name = "dicrect_check.gif")

# Найдем абсолютную погрешность численного решения

err = u .- u_model
heatmap(Xₙ, Tₘ, err', xlabel=L"X_n", ylabel=L"T_m", title="Absolute Error", size=(1200, 800))


# # A model with homeowners and renters

using Test #jl
using Aiyagari
using QuantEcon, Parameters
using Plots #md

PKG_HOME = joinpath(dirname(pathof(Aiyagari)), "..")

# ## Setting up the state space

# Exogenous states (incomes)

z_grid = [0.5; 1.0; 1.5]
z_prob = [0.7 0.15 0.15;
          0.2 0.6 0.2;
          0.15 0.15 0.7]
z_MC = MarkovChain(z_prob, z_grid, :z)

exo = ExogenousStateSpace([z_MC])

# ## Define utility function for the case of two goods (consumption good and housing)

u(c; γ) = c^(1-γ) / (1-γ)

function u(c,h; ξ=0.8159, ρ=map(s -> (s-1)/s, 0.13), γ=2.0)
  C = (ξ * h^ρ + (1-ξ) * c^ρ)^(1/ρ)
  u(C, γ=γ)
end

include(joinpath(PKG_HOME, "test", "housing-simple-nlopt.jl"))
include(joinpath(PKG_HOME, "test", "renting-nlopt.jl"))

# ## Aggregate state

mutable struct HousingAS{T1,T2,T3,T4} <: AggregateState
  r::T1
  p::T2 # house price
  ρ::T3 # rent
  dist::T4 # the distribution over idiosynchratic states
end

function HousingAS(r, p, w_grid, exo, param; ρ=p * (param.δ + r))
  dist_proto = zeros((length(w_grid), length(exo)))
  HousingAS(r, p, ρ, dist_proto)
end

# # Case 1: Forever home owners

w_grid_own = LinRange(0.0, 3.0, 40)
endo_own = EndogenousStateSpace((w=w_grid_own,))

r_own = 0.15
p_own = 0.9
param_own  = (β = 0.9, θ = 0.9, δ = 0.1, h_thres = eps())
agg_state_own = HousingAS(r_own, p_own, w_grid_own, exo, param_own)

@unpack policies_full, val = solve_bellman(endo_own, exo, agg_state_own, param_own, Owner(), rtol=√eps())

dist = stationary_distribution(endo_own, exo, policies_full.w_next)

using DelimitedFiles #jl

#

plt_w = plot(w_grid_own, policies_full.w_next, title="wealth next period", xlab="wealth") #md
plt_m_zoomed = plot(w_grid_own[1:15], policies_full.m[1:15,:], title="mortgage (zoomed)", xlab="wealth") #md

plt_m = plot(w_grid_own, policies_full.m, title="mortgage", xlab="wealth") #md

plt_dist = plot(w_grid_own, dist, xlab="wealth", title="stationary distribution") #md

plot(plt_w, plt_dist, plt_m, plt_m_zoomed, legend=false) #md

# using DelimitedFiles #jl
#writedlm(joinpath(PKG_HOME, "test/matrices", "housing-simple-value.txt"), val) #src
#writedlm(joinpath(PKG_HOME, "test/matrices", "housing-simple-dist.txt"), dist) #src

@testset "regression test housing-simple" begin #jl
  value_test = readdlm(joinpath(PKG_HOME, "test/matrices", "housing-simple-value.txt")) #jl
  dist_test = readdlm(joinpath(PKG_HOME, "test/matrices", "housing-simple-dist.txt")) #jl
  @test all(val .== value_test) #jl
  @test maximum(abs, dist .- dist_test) < 1e-12 #jl
end #jl

# ## Equilibrium

# First, let's compute excess demand

function excess_demand(r, p)
  agg_state_own = HousingAS(r, p, w_grid_own, exo, param_own)

  @unpack policies_full, val = solve_bellman(endo_own, exo, agg_state_own, param_own, Owner(), rtol=√eps())

  dist = stationary_distribution(endo_own, exo, policies_full.w_next)

  (m = sum(dist .* policies_full.m), h=sum(dist .* policies_full.h))
end

excess_demand(r_own, p_own)

# TODO: Housing supply

# # Own big, rent small

# Now we want to give households to choice whether to buy or to rent

function combined_policies(policies_full, owner, own, rent=own; f_own =identity, f_rent=identity)
  f_own.(getproperty(policies_full[1], own)) .* owner .+ f_rent.(getproperty(policies_full[2], rent)) .* .!(owner)
end

function excess_demand(r, p, endo, exo, param, hh; maxiter_bellman=200)
  agg_state = HousingAS(r, p, endo.grids.w, exo, param, ρ= p * (param[1].δ + r))

  out = solve_bellman(endo, exo, agg_state, param, hh, maxiter=300, rtol=√eps())

  @unpack val, policy, policies_full, policy_hh = out
  owner = policy_hh[1] .== 1

  w_next_all = combined_policies(policies_full, owner, :w_next)

  dist = stationary_distribution(endo, exo, w_next_all)
  a_all = combined_policies(policies_full, owner, :m, :w_next, f_own = x -> -x)
  h_all = combined_policies(policies_full, owner, :h)

  (a = sum(dist .* a_all), h=sum(dist .* h_all), dist=dist, out...)
end

r = 0.01
p = 1.7
w_grid = LinRange(0.0, 4.0, 40)
endo = EndogenousStateSpace((w=w_grid,))

param_both = (β = 0.93, θ = 0.9, δ = 0.1, h_thres = 1.2)
param = [param_both, param_both]

@unpack val, policy, policies_full, policy_hh, a, h, dist = excess_demand(r, p, endo, exo, param, OwnOrRent())

@show r => a, p => h
owner = policy_hh[1] .== 1
value = val[1]

plt_own = plot(w_grid, owner, title="Who owns?") #md

w_next_all = combined_policies(policies_full, owner, :w_next)
h_all = combined_policies(policies_full, owner, :h)
c_all = combined_policies(policies_full, owner, :c)
a_all = combined_policies(policies_full, owner, :m, :w_next, f_own = x -> -x)

plt_h = plot(w_grid, h_all, legend=:false, title="House size", markerstrokewidth=0, xlab="wealth") #md

plt_a = plot(w_grid, a_all, title = "asset/mortgage") #md

plt_w = plot(w_grid, w_next_all, title = "wealth next period") #md

plt_dist = plot(w_grid, dist, title = "stationary distribution") #md

plot(plt_h, plt_a, plt_w, plt_dist, legend=false) #md
#- #md

#writedlm(joinpath(PKG_HOME, "test/matrices", "own-rent-value.txt"), value) #src
#writedlm(joinpath(PKG_HOME, "test/matrices", "own-rent-dist.txt"), dist) #src

@testset "own-rent" begin #jl
  value_test = readdlm(joinpath(PKG_HOME, "test/matrices", "own-rent-value.txt")) #jl
  dist_test = readdlm(joinpath(PKG_HOME, "test/matrices", "own-rent-dist.txt")) #jl
  @test all(value .≈ value_test) #jl
  @test maximum(abs, dist .- dist_test) < 1e-12 #jl
end #jl


using Base.Test

# choose the next block for redistribution
# @return the value and index of the block to distribute (1 based)
function choose_block(x)
  findmax(x)
end

# redistribute the max block
# @return the redistributed memory
function redistribute(x)
  val, index = choose_block(x)
  result = x
  result[index] = 0 # zero redistributed block
  len = size(x)[1]
  # we distribute one-by-one until nothing is left
  for i = 1:val
    index = (index == len) ? 1 : (index + 1) # wrap around at last index
    result[index] += 1
  end
  result
end

# count redistributions until cycle is detected
# @return the number of redistributions overall, and the cycle length
function count_until_cycle(x)
  seen = Dict()
  n = 0
  next = x
  # we use the string representation of the array as key to avoid aliasing problems
  while !haskey(seen, "$next")
    seen["$next"] = n
    next = redistribute(next)
    n += 1
  end
  (n, n - seen["$next"])
end

# tests for choose_block
@test choose_block([0, 2, 7, 0]) == (7,3)
# choose_block should choose the first one if multiple max-value blocks exist
@test choose_block([0, 7, 0, 7]) == (7,2)

# tests for redistribution
@test redistribute([0, 0, 0, 1]) == [1, 0, 0, 0]
@test redistribute([0, 2, 7, 0]) == [2, 4, 1, 2]

# tests for detecting cycles
@test count_until_cycle([0, 2, 7, 0]) == (5, 4)

# solve it
input = [11,11,13,7,0,15,5,5,4,4,1,1,7,1,15,11]
println(count_until_cycle(input))


using Test, Printf
using GenSPGL, JOLI

include("test_NormL1.jl")
include("test_spgl1.jl")


"""Bitstring function with a split-mode that splits the bitstring into
sign, exponent and significant bits."""
function Base.bitstring(x::Float16,mode::Symbol)
    if mode == :split
        bstr = bitstring(x)
        return "$(bstr[1]) $(bstr[2:6]) $(bstr[7:end])"
    else
        bitstring(x)
    end
end

"""Bitstring function with a split-mode that splits the bitstring into
sign, exponent and significant bits."""
function Base.bitstring(x::Float32,mode::Symbol)
    if mode == :split
        bstr = bitstring(x)
        return "$(bstr[1]) $(bstr[2:9]) $(bstr[10:end])"
    else
        bitstring(x)
    end
end

"""Bitstring function with a split-mode that splits the bitstring into
sign, exponent and significant bits."""
function Base.bitstring(x::Float64,mode::Symbol)
    if mode == :split
        bstr = bitstring(x)
        return "$(bstr[1]) $(bstr[2:12]) $(bstr[13:end])"
    else
        bitstring(x)
    end
end


@testset "Indexing" begin
    @testset "Int" begin
        leaf = Leaf(Tape(), 5 * [1, 1, 1, 1, 1])
        y = getindex(leaf, 1)
        @test unbox(y) == 5
        @test ∇(y, one(unbox(y)))[leaf] == [1, 0, 0, 0, 0]
    end

    @testset "Vector" begin
        x = Leaf(Tape(), 10 * [1, 1, 1])
        y = x[2:3]
        @test unbox(y) == [10, 10]
        @test ∇(y, oneslike(unbox(y)))[x] == [0, 1, 1]
    end

    @testset "Overlapping indices (#139)" begin
        x = Leaf(Tape(), 10 * [1, 1, 1])
        y = x[[2, 3, 3]]
        @test unbox(y) == [10, 10, 10]
        @test ∇(y, oneslike(unbox(y)))[x] == [0, 1, 2]
    end

    @testset "Ref" begin
        @test ref_equal(∇(getindex)(Ref(4))[1], Ref(1))
    end
end


#using Morbit
#using Test

# Two standard (easy) test problems with known Pareto Data
# Many combinations of constaints, heterogenity, descent method and sampling algorithm are used

#using Logging
#global_logger(SimpleLogger(stdout, Logging.Debug))
#%%
@testset "1In1OutProblem_SteepestDescent" begin

    x0 = [5.0];
    f1(x) = x[1]^2;

    # I) unconstrained
    opt_settings = AlgoConfig(
        max_iter = 30
    );
    ## 1) treat objective as cheap
    mop = MixedMOP();
    add_objective!( mop, f1, :cheap );
    x,fx = optimize(mop, x0; algo_config = opt_settings);

    @test x ≈ [ 0.0 ] atol=1e-2

    ## 2) treat objective as expensive
    opt_settings = AlgoConfig(
        max_iter = 30
    );
    mop = MixedMOP();
    add_objective!( mop, f1, :expensive )
    x,fx = optimize(mop, x0; algo_config = opt_settings)

    @test x ≈ [ 0.0 ] atol = 1e-2

    #II) constrained

    ## 1) treat objective as cheap
    opt_settings = AlgoConfig(
        max_iter = 10
    );
    mop = MixedMOP( [-8.0], [8.0] );
    add_objective!( mop, f1, :cheap)
    x,fx = optimize(mop, x0; algo_config = opt_settings)

    @test x ≈ [ 0.0 ] atol=0.1

    ## 2) treat objective as expensive
    opt_settings = AlgoConfig(
        max_iter = 30
    );
    mop = MixedMOP( [-8.0], [8.0] );
    add_objective!( mop, f1, :expensive);
    x,fx = optimize(mop, x0; algo_config = opt_settings);

    @test x ≈ [ 0.0 ] atol=.1
end

#%%
@testset "2In2OutProblem_Steepest_Descent" begin
    x0 = [ π , π^2 ]
    ub = 10 .* ones(2);
    lb = -ub;

    g1(x) = sum( (x .- 1.0).^2 );
    g2(x) = sum( (x .+ 1.0).^2 );

    # unconstrained, cheap
    @testset "unbound_cheap_cheap" begin
        opt_settings = AlgoConfig(
            Δ_0= 0.2,
            max_iter = 30,
        );
        mop = MixedMOP();
        add_objective!(mop, g1, :cheap);
        add_objective!(mop, g2, :cheap);
        x,fx = optimize(mop, x0; algo_config = opt_settings);

        @test x[1] ≈ x[2] atol = .1
    end

    # unconstrained, expensive
    @testset "unbound_exp_exp" begin
        opt_settings = AlgoConfig(
            max_iter = 30
        );
        mop = MixedMOP();
        add_objective!(mop, g1, :expensive);
        add_objective!(mop, g2, :expensive);
        x,fx = optimize(mop, x0; algo_config = opt_settings);

        @test x[1] ≈ x[2] atol = .1
    end

    # unconstrained, heterogenous
    @testset "unbound_exp_cheap" begin
        opt_settings = AlgoConfig(
            max_iter = 30
        );
        mop = MixedMOP();
        add_objective!(mop, g1, :expensive);
        add_objective!(mop, g2, :cheap);
        x,fx = optimize(mop, x0; algo_config = opt_settings);

        @test x[1] ≈ x[2] atol = .1
    end

    # constrained, cheap
    @testset "bound_cheap_cheap" begin
        opt_settings = AlgoConfig(
            max_iter = 30
        );
        mop = MixedMOP(lb, ub);
        add_objective!(mop, g1, :cheap);
        add_objective!(mop, g2, :cheap);
        x,fx = optimize(mop, x0; algo_config = opt_settings);

        @test x[1] ≈ x[2] atol = .1
    end

    # constrained, expensive
    @testset "bound_exp_exp" begin
        opt_settings = AlgoConfig(
            max_iter = 50,
        );
        mop = MixedMOP(lb,ub);
        add_objective!(mop, g1, :expensive);
        add_objective!(mop, g2, :expensive);
        x,fx = optimize(mop, x0; algo_config = opt_settings);

        @test x[1] ≈ x[2] atol = .1
    end

    # constrained, heterogenous
    @testset "bound_exp_cheap" begin
        opt_settings = AlgoConfig(
            max_iter = 50,
            Δ_max = .4,
            _eps_crit = 0.1,
            strict_acceptance_test = true,
        );
        mop = MixedMOP(lb,ub);
        add_objective!(mop, g1, RbfConfig(kernel=:multiquadric));
        add_objective!(mop, g2, :cheap);
        x,fx = optimize(mop, x0; algo_config = opt_settings);

        @test x[1] ≈ x[2] atol = .1
    end
end
#%%
#=
@testset "direct_search" begin
    # 1D1D
    x0 = [5.0];
    f1(x) = x[1]^2;

    for lb_ub ∈ [ (Float32[],Float32[]), ([-5.0], [5.0]), ([-10.0], [10.0])]
        lb, ub = lb_ub
        for type ∈ [:cheap, :expensive ]
            for ideal_point ∈ [Float32[], [0.0]]
                opt_settings = AlgoConfig(
                    max_iter = 20,
                    descent_method = :direct_search,
                    ideal_point = ideal_point,
                )
                mop = MixedMOP(lb,ub)
                Morbit.scale( mop, x0)
                add_objective!(mop, f1, type)
                x, fx = optimize(mop, x0; algo_config = opt_settings)
                @test x[end] ≈ 0.0 atol = .1
            end
        end
    end

    # 2D2D
    #=
    x0 = [ π , π^2 ]
    ub = 10 .* ones(2);
    lb = -ub;

    g1(x) = sum( (x .- 1.0).^2 );
    g2(x) = sum( (x .+ 1.0).^2 );
    for lb_ub ∈ [ ( [],[] ), ([-10.0,-100], [10.0, 10.0])]
        lb,ub = lb_ub
        for type_tuple ∈ [ (:cheap, :cheap), (:expensive, :expensive), (:expensive, :cheap) ]
            type1, type2 = type_tuple
            for ideal_point ∈ [ [], zeros(2) ]
                opt_settings = AlgoConfig(
                    max_iter = 50,
                    rbf_kernel = :multiquadric,
                    rbf_shape_parameter = cs -> let Δ = cs.iter_data.Δ; return 1/(10*Δ) end,
                    Δ_0= 0.5,
                    all_objectives_descent = true,
                    use_max_points = type_tuple ==  (:expensive, :expensive) ? true : false,
                    max_model_points = 12,
                )
                mop = MixedMOP(lb,ub)
                add_objective!(mop, g1, type1)
                add_objective!(mop, g2, type2)
                x, fx = optimize(mop, x0; algo_config = opt_settings)
                @test x[1] ≈ x[2] atol = .5
            end
        end
    end
    =#
end
=#


# example of a unit test in Julia
# runs a test for a certain module
# to run the tests, first open the package manager (`]` in REPL),
# activate the project if not done so and then enter `test`

# wrap all your tests and subgroups in a `@testset` block
@testset "BeesAlgorithm" begin
    using STMOZOO.BeesAlgorithm  # load YOUR module
    using STMOZOO.BeesAlgorithm: initialize_population, compute_objective, compute_fitness, employed_bee_phase, scouting_phase, onlooker_bee_phase, foodsource_info_prob, create_newsolution

    S = 24
    bounds_lower = [-100,-100,-100,-100]
    bounds_upper = [100,100,100,100]
    D = 4
    limit = D * (S/2)
    T = 500
    n = Int8(S/2)

    @testset "initialize_population" begin

        population = initialize_population(D, bounds_lower, bounds_upper, n)
        @test length(population) == n
        @test all([bounds_lower <= population[i] <= bounds_upper for i in 1:length(population)])

    end

    @testset "compute_objective" begin

        population = initialize_population(D, bounds_lower, bounds_upper, n)
        objectives = compute_objective(population,sphere)
        @test length(objectives) == length(population)

    end

    @testset "compute_fitness" begin

        population = initialize_population(D, bounds_lower, bounds_upper, n)
        objectives = compute_objective(population,sphere)
        fitness = compute_fitness(objectives)
        @test length(fitness) == length(objectives) == length(population)

    end

    @testset "foodsource_info_prob" begin

        population = initialize_population(D, bounds_lower, bounds_upper, n)
        objectives = compute_objective(population,sphere)
        fitness = compute_fitness(objectives)
        probab = foodsource_info_prob(fitness)
        @test length(probab) == length(fitness)
        @test all(0 .≤ probab .≤ 1)

    end

    @testset "create_newsolution" begin

        population = initialize_population(D, bounds_lower, bounds_upper, n)
        solution = [4, 0, 1, 4]
        new_solution = create_newsolution(solution, population, bounds_lower, bounds_upper)
        @test length(new_solution) == length(solution)
        @test sum(new_solution .==  solution) == D-1

    end

    @testset "employed_bee_phase" begin

        population = initialize_population(D, bounds_lower, bounds_upper, n)
        trial = zeros(size(population)[1])
        population_new_evolved, fitness_new_evolved, objective_new_evolved, newtrial = employed_bee_phase(population, bounds_lower, bounds_upper, trial,n ,sphere)

        @test length(population_new_evolved) == length(population)
        @test length(newtrial) == length(trial)
    end

    @testset "onlooker_bee_phase" begin

        population = initialize_population(D, bounds_lower, bounds_upper, n)
        trial = zeros(size(population)[1])
        population_new_evolved, fitness_new_evolved, objective_new_evolved, newtrial = onlooker_bee_phase(population, bounds_lower, bounds_upper, trial,n ,sphere)

        @test length(population_new_evolved) == length(population)
        @test length(newtrial) == length(trial)
    end

    @testset "scouting_phase" begin

        population = initialize_population(D, bounds_lower, bounds_upper, n)
        trial = ones(size(population)[1])
        objective= compute_objective(population,sphere)
        fitness = compute_fitness(objective)
        population_new_evolved, fitness_new_evolved, objective_new_evolved, newtrial  = scouting_phase(population, bounds_lower, bounds_upper,D, trial, fitness, objective, 0, sphere)

        @test length(objective) == length(fitness)
        @test length(population_new_evolved) == length(population)
        @test length(newtrial) == length(trial)
    end

    @testset "ArtificialBeeColonization" begin

        optimal_solution,populations,best_fitness_tracker = ArtificialBeeColonization(D, bounds_lower, bounds_upper, S, T, limit, sphere)

        @test length(optimal_solution) == length(bounds_lower) == length(bounds_upper) == D
        @test optimal_solution isa Vector
        @test optimal_solution >= bounds_lower
        @test optimal_solution <= bounds_upper

    end

end


# # [Overdetermined Beam](@ref overdetermined)
#
# This example shows how to predict the behavior of a beam which is clamped at one end and
# simply supported at the other end when subjected to a linear distributed load.
#
# ![](../assets/overdetermined-drawing.svg)
#
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`overdetermined.ipynb`](@__NBVIEWER_ROOT_URL__/examples/overdetermined.ipynb).
#-

using GXBeam, LinearAlgebra

nelem = 16

## create points
L = 1
x = range(0, L, length=nelem+1)
y = zero(x)
z = zero(x)
points = [[x[i],y[i],z[i]] for i = 1:length(x)]

## index of endpoints for each beam element
start = 1:nelem
stop = 2:nelem+1

## create compliance matrix for each beam element
EI = 1e7
compliance = fill(Diagonal([0, 0, 0, 0, 1/EI, 0]), nelem)

## create assembly
assembly = Assembly(points, start, stop, compliance=compliance)

## set prescribed conditions
prescribed_conditions = Dict(
    ## simply supported left endpoint
    1 => PrescribedConditions(uz=0),
    ## clamped right endpoint
    nelem+1 => PrescribedConditions(ux=0, uy=0, uz=0, theta_x=0, theta_y=0,
        theta_z=0)
)

## create distributed load
qmax = 1000
distributed_loads = Dict()
for i = 1:nelem
    distributed_loads[i] = DistributedLoads(assembly, i; s1=x[i],
        s2=x[i+1], fz = (s) -> qmax*s)
end

## perform a static analysis
system, converged = static_analysis(assembly;
    prescribed_conditions = prescribed_conditions,
    distributed_loads = distributed_loads,
    linear = true)

## post-process the results
state = AssemblyState(system, assembly;
    prescribed_conditions = prescribed_conditions)
#!jl nothing #hide

# An analytical solution to this over-determined problem can be found using the method of
# superposition using the analytical solutions to a cantilever with a linearly distributed
# load and a cantilever with an end load.

## construct analytical solution
x_a = range(0.0, L, length=1000)
w_a = @. qmax*(1-x_a)^2/(120*EI)*(4 - 8*(1-x_a) + 5*(1-x_a)^2 - (1-x_a)^3)
slope_a = @. -qmax*(1-x_a)/(120*EI)*(8 - 24*(1-x_a) + 20*(1-x_a)^2 - 5*(1-x_a)^3)
M_a = @. qmax/120*(8 - 48*(1-x_a) + 60*(1-x_a)^2 - 20*(1-x_a)^3)

theta_a = -atan.(slope_a)

## adjust coordinate system of analytical solution
M_a = -M_a
#!jl nothing #hide

# Plotting the results reveals that the analytical and computational solutions show
# excellent agreement.

using Plots
#md using Suppressor #hide
pyplot()
#!jl nothing #hide

#-

#md @suppress_err begin #hide
#md local x #hide

plot(
    xlim = (0.0, 1.0),
    xticks = 0.0:0.2:1.0,
    xlabel = "x (m)",
    ylabel = "Linear Deflection (m)",
    grid = false,
    overwrite_figure=false
    )

x = [assembly.points[ipoint][1] + state.points[ipoint].u[1] for ipoint =
    1:length(assembly.points)]
deflection = [state.points[ipoint].u[3] for ipoint = 1:length(assembly.points)]
plot!(x_a, w_a, label="Analytical")
scatter!(x, deflection, label="GXBeam")
#!nb plot!(show=true)
#md savefig("../assets/overdetermined-linear-deflection.svg") #hide
#md closeall() #hide
#md end #hide
#md nothing #hide

#md # ![](../assets/overdetermined-linear-deflection.svg)

#-

#md @suppress_err begin #hide
#md local x #hide

plot(
    xlim = (0.0, 1.0),
    xticks = 0.0:0.2:1.0,
    xlabel = "x (m)",
    ylabel = "Angular Deflection (rad)",
    grid = false,
    overwrite_figure=false
    )

x = [assembly.points[ipoint][1] + state.points[ipoint].u[1]
    for ipoint = 1:length(assembly.points)]
theta = [4*atan.(state.points[ipoint].theta[2]/4)
    for ipoint = 1:length(assembly.points)]
plot!(x_a, theta_a, label="Analytical")
scatter!(x, theta, label="GXBeam")
#!nb plot!(show=true)
#md savefig("../assets/overdetermined-angular-deflection.svg") #hide
#md closeall() #hide
#md end #hide
#md nothing #hide

#md # ![](../assets/overdetermined-angular-deflection.svg)

#-

#md @suppress_err begin #hide
#md local x #hide

plot(
    xlim = (0.0, 1.0),
    xticks = 0.0:0.2:1.0,
    xlabel = "x (m)",
    ylabel = "Bending Moment (\$Nm\$)",
    grid = false,
    overwrite_figure=false
    )

x = [assembly.elements[ielem].x[1] + state.elements[ielem].u[1] for ielem =
    1:length(assembly.elements)]
M = [state.elements[ielem].M[2] for ielem = 1:length(assembly.elements)]
plot!(x_a, M_a, label="Analytical")
scatter!(x, M, label="GXBeam")
#!nb plot!(show=true)
#md savefig("../assets/overdetermined-bending-moment.svg") #hide
#md closeall() #hide
#md end #hide
#md nothing #hide

#md # ![](../assets/overdetermined-bending-moment.svg)

# Note that we could have easily performed a nonlinear analysis for this problem by
# setting `linear=false`.


## weights
function getweight(P, i, j, P2, TEs, M, maxmag, flags) # Phase, index, neighbor, ...
    weight = 1.0
    if flags[1] weight *= (0.1 + 0.9phasecoherence(P, i, j)) end
    if flags[2] weight *= (0.1 + 0.9phasegradientcoherence(P, P2, TEs, i, j)) end
    if flags[3] weight *= (0.1 + 0.9phaselinearity(P, i, j)) end
    if !isnothing(M)
        small, big = minmax(M[i], M[j])
        if flags[4] weight *= (0.1 + 0.9magcoherence(small, big)) end
        if flags[5] weight *= (0.1 + 0.9magweight(small, maxmag)) end
        if flags[6] weight *= (0.1 + 0.9magweight2(big, maxmag)) end
    end
    return weight
end

phasecoherence(P, i, j) = 1 - abs(γ(P[i] - P[j]) / π)
phasegradientcoherence(P, P2, TEs, i, j) = max(0, 1 - abs(γ(P[i] - P[j]) - γ(P2[i] - P2[j]) * TEs[1] / TEs[2]))
magcoherence(small, big) = (small / big) ^ 2
magweight(small, maxmag) = 0.5 + 0.5min(1, small / (0.5 * maxmag))
magweight2(big, maxmag) = 0.5 + 0.5min(1, (0.5 * maxmag) / big) # too high magnitude is not good either (flow artifact)

phaselinearity(P, i, j, k) = max(0, 1 - abs(rem2pi(P[i] - 2P[j] + P[k], RoundNearest)/2))
function phaselinearity(P, i, j)
    neighbor = j - i
    h = i - neighbor
    k = j + neighbor
    if 0 < h && k <= length(P)
        return phaselinearity(P, h, i, j) * phaselinearity(P, i, j, k)
    else
        return 0.9
    end
end

"""
    calculateweights(wrapped; weights=:romeo, kwargs...)

Calculates weights for all edges.
size(weights) == [3, size(wrapped)...]

###  Optional keyword arguments:

- `weights`: Options are [`:romeo`] | `:romeo2` | `:romeo3` | `:bestpath`.
- `mag`: Additional mag weights are used.
- `mask`: Unwrapping is only performed inside the mask.
- `phase2`: A second reference phase image (possibly with different echo time).
   It is used for calculating the phasecoherence weight.
- `TEs`: The echo times of the phase and the phase2 images as a tuple (eg. (5, 10) or [5, 10]).

"""
function calculateweights(wrapped; weights=:romeo, kwargs...)
    weights = if weights isa AbstractArray{<:Real,4}
        weights
    elseif weights == :bestpath
        calculateweights_bestpath(wrapped; kwargs...)
    else
        calculateweights_romeo(wrapped, weights; kwargs...)
    end
    # these edges do not exist
    weights[1,end,:,:] .= 0
    weights[2,:,end,:] .= 0
    weights[3,:,:,end] .= 0
    return weights
end

calculateweights_romeo(wrapped, weights::AbstractArray{T,4}; kw...) where T = UInt8.(weights)
# defines default weights
function calculateweights_romeo(wrapped, weights::Symbol; kwargs...)
    flags = falses(6)
    if weights == :romeo
        flags[1:4] .= true
    elseif weights == :romeo2
        flags[[1,4]] .= true # phasecoherence, magcoherence
    elseif weights == :romeo3
        flags[[1,2,4]] .= true # phasecoherence, phasegradientcoherence, magcoherence
    elseif weights == :romeo4
        flags[1:4] .= true
    elseif weights == :romeo6
        flags[1:6] .= true # additional magweight, magweight2
    else
        throw(ArgumentError("Weight '$weights' not defined!"))
    end
    return calculateweights_romeo(wrapped, flags; kwargs...)
end

function calculateweights_romeo(wrapped, flags::AbstractArray{Bool,1}; type::Type{T}=UInt8, rescale=rescale, kwargs...) where T
    mask, P2, TEs, M, maxmag = parsekwargs(kwargs, wrapped)
    updateflags!(flags, wrapped, P2, TEs, M)
    stridelist = getdimoffsets(wrapped)
    weights = zeros(T, 3, size(wrapped)...)
    for dim in 1:3
        neighbor = stridelist[dim]
        for I in LinearIndices(wrapped)
            J = I + neighbor
            if mask[I] && checkbounds(Bool, wrapped, J)
                w = getweight(wrapped, I, J, P2, TEs, M, maxmag, flags)
                weights[dim + (I-1)*3] = rescale(w)
            end
        end
    end
    return weights
end

## utility functions
function parsekwargs(kwargs, wrapped)
    getval(key) = if haskey(kwargs, key) kwargs[key] else nothing end
    mag = getval(:mag)
    mask = getval(:mask)
    if !isnothing(mask)
        if !isnothing(mag)
            mag = mag .* mask
        end
    else
        mask = trues(size(wrapped))
    end
    maxmag = if !isnothing(mag)
        maximum(mag[isfinite.(mag)]) else nothing
    end
    return mask, getval(:phase2), getval(:TEs), mag, maxmag
end

function updateflags!(flags, P, P2, TEs, M)
    if isnothing(M)
        flags[4:6] .= false
    end
    if isnothing(P2) || isnothing(TEs)
        flags[2] = false
    end
end

# from: 1 is best and 0 worst
# to: 1 is best, NBINS is worst, 0 is not valid (not added to queue)
function rescale(w)
    if 0 ≤ w ≤ 1
        max(round(Int, (1 - w) * (NBINS - 1)), 1)
    else
        0
    end
end

## best path weights
# Abdul-Rahamn https://doi.org/10.1364/AO.46.006623

function calculateweights_bestpath(wrapped; kwargs...)
    scale(w) = UInt8.(min(max(round((1 - (w / 10)) * (NBINS - 1)), 1), 255))
    weights = scale.(getbestpathweight(wrapped))
    if haskey(kwargs, :mask) # apply mask to weights
        mask = kwargs[:mask]
        weights .*= reshape(mask, 1, size(mask)...)
    end
    weights
end

function getbestpathweight(φ)
    R = 1 ./ getD(φ)
    weight = zeros(3, size(R)...)
    for idim in 1:3
        n = getdimoffsets(R)[idim]
        @inbounds for i in 1:length(R)-n
            weight[idim + 3i] = R[i] + R[i+n]
        end
    end
    return weight
end

function getD(φ)
    directions = Iterators.product(-1:1,-1:1,-1:1)
    neighbors = unique(abs(sum(getdimoffsets(φ) .* d)) for d in directions if d != (0,0,0))
    D2 = zeros(size(φ))
    @inbounds for n in neighbors, i in 1+n:length(φ)-n
        D2[i] += (γ(φ[i-n] - φ[i]) - γ(φ[i] - φ[i+n])) ^ 2
    end
    return sqrt.(D2)
end


# We work with Rydberg units, length scale Bohr radius a_0, energy scale: Ry
using StaticArrays

###### constants ###########
const e0 = sqrt(2)  # electric charge
const me = 0.5  # electron mass
const dim = 3    # dimension (D=2 or 3, doesn't work for other D!!!)
const spin = 2  # number of spins

const rs = 1.0
const kF = (dim == 3) ? (9π / (2spin))^(1 / 3) / rs : sqrt(4 / spin) / rs
const EF = kF^2 / (2me)
const beta = 0.05
const β = beta/ EF
const mass2 = 0.01

const Weight = SVector{2,Float64}
const Base.abs(w::Weight) = abs(w[1]) + abs(w[2]) # define abs(Weight)
const INL, OUTL, INR, OUTR = 1, 2, 3, 4
# const Nf = (D==3) ?

println("rs=$rs, β=$β, kF=$kF, EF=$EF, mass2=$mass2")