text:
# sample script to reproduce
# ERROR: fatal error in type inference (type bound)

# it seems related to Union{mseSplit,sseSplit}
# (search the code for that expression)
# I am quite sure that this did not happen on 0.7beta2

using Distributed
using Pkg
pkg"develop https://github.com/kafisatz/DecisionTrees.jl#master"
pkg"activate DecisionTrees.jl"
Distributed.@everywhere using DecisionTrees

# if the above does not work, maybe you need to clone the git
# git clone --depth=50 --branch=master https://github.com/kafisatz/DecisionTrees.jl.git
# cd kafisatz/DecisionTrees.jl

@assert VERSION >= v"0.7.0-rc2.0"

Distributed.@everywhere import CSV
Distributed.@everywhere import DataFrames
Distributed.@everywhere import DataFrames: DataFrame

##############################
# Read the data
##############################

dr = "C:\\Users\\bernhard.konig\\Documents\\ASync\\home\\Code\\Julia\\DecisionTrees.jl"
isdir(dr) && cd(dr)
pwd()

@warn("""This assumes pwd() == Pkg.dir("DecisionTrees.jl")""")
datafile = joinpath("data", "freMTPL2", "freMTPL2.csv")
if Sys.islinux()
    # trying to unzip the file...
    pathToZip = string(splitext(datafile)[1], ".zip")
    unzipLoc = splitdir(pathToZip)[1]
    cmd = `unzip $(pathToZip) -d $(unzipLoc)`
    run(cmd)
end

@assert isfile(datafile) "You need to unzip the file freMTPL2.zip in the folder data/freMTPL2"

@time fullData = CSV.read(datafile, rows_for_type_detect=100000, allowmissing=:none, categorical=false);

# some data manipulation
areasSorted = sort(unique(fullData[:Area]))
AreaInteger = map(x->findall((in)([x]), areasSorted)[1], fullData[:Area])
fullData[:AreaInteger] = AreaInteger
    # correct for unreasonable observations
for i = 1:size(fullData, 1)
    fullData[:ClaimNb][i] = min(4, fullData[:ClaimNb][i])
    fullData[:Exposure][i] = min(1.0, fullData[:Exposure][i])
end

# set independent variables
selected_explanatory_vars = ["Area","AreaInteger","VehPower","VehAge","DrivAge","BonusMalus","VehBrand","VehGas","Density","Region"]

##############################
# Prepare the data
##############################
dtmtable, sett, dfprepped = prepare_dataframe_for_dtm!(fullData, keycol="IDpol", trnvalcol="trnTest", numcol="ClaimNb", denomcol="Exposure", weightcol="Exposure", independent_vars=selected_explanatory_vars);

# run a few models

# this runs as expected
updateSettings!(sett, minWeight=-0.03, model_type="build_tree", calculatePoissonError=true)
resultingFiles, resM = dtm(dtmtable, sett)

# this runs as expected
updateSettings!(sett, crit="sse", model_type="build_tree")
resultingFiles, resM = dtm(dtmtable, sett)

# this one fails!
updateSettings!(sett, crit="sse", model_type="boosted_tree")
resultingFiles, resM = dtm(dtmtable, sett)

#=
# this is the error I get

"""
ERROR: fatal error in type inference (type bound)
Stacktrace:
 [1] _typed_hcat(::Type{Float64}, ::Tuple{Array{Int64,1},Array{Int64,1},Array{Float64,1},Array{Float64,1},Array{Float64,1},Array{Float64,1},Array{Float64,1},Array{Float64,1}}) at .\abstractarray.jl:1255
 [2] typed_hcat at .\abstractarray.jl:1243 [inlined]
 [3] hcat at C:\cygwin\home\Administrator\buildbot\worker\package_win64\build\usr\share\julia\stdlib\v0.7\SparseArrays\src\sparsevector.jl:1063 [inlined]
 [4] macro expansion at .\logging.jl:307 [inlined]
 [5] macro expansion at .\util.jl:156 [inlined]
 [6] macro expansion at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\boosting.jl:89 [inlined]
 [7] macro expansion at .\util.jl:156 [inlined]
 [8] macro expansion at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\helper_functions.jl:1469 [inlined]
 [9] boosted_tree(::DecisionTrees.DTMTable, ::ModelSettings) at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\boosting.jl:82
 [10] run_model_actual(::DecisionTrees.DTMTable, ::ModelSettings, ::String) at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\runmodel.jl:467
 [11] #dtm#110 at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\runmodel.jl:798 [inlined]
 [12] dtm(::DecisionTrees.DTMTable, ::ModelSettings) at C:\Users\bernhard.konig\Documents\ASync\home\Code\Julia\DecisionTrees.jl\src\runmodel.jl:798
 [13] top-level scope at none:0
 """ =#


text:
module Day10
export day10

function readInput(inputfile="inputs/input10.txt")
    vals = sort!([parse(Int, line) for line in eachline(inputfile)])
    vcat([0], vals, maximum(vals)+3)
end

function day10(inputfile="inputs/input10.txt")
    values = readInput(inputfile)
    (day10Part1(values), day10Part2(values))
end

function day10Part1(values)
    diff1 = 0
    diff3 = 0
    for (a, b) in zip(values, values[2:end])
        if b-a == 1
            diff1 += 1
        elseif b-a == 3
            diff3 += 1
        end
        if b-a > 3
            print(a," - ", b)
        end
    end
    diff1 * diff3
end

function day10Part2(values)
    combinations = zeros(Int, size(values))
    combinations[end-2:end] .= 1
    # the last adpater (n-1) can only connect directly to our device (n)
    # that means the second to last one (n-2) can only connect to the last (n-1)

    for index in reverse(1:length(combinations)-3)
        val = values[index]
        for δ in 1:3
            if values[index+δ] - val <= 3
                combinations[index] += combinations[index+δ]
            else
                break
            end
        end
    end
    combinations[1]
end

end #module

ext: jl


text:
# # Worst case risk analysis
# Generate data for worst-case risk analysis.
using Random

Random.seed!(2);
n = 5;
r = abs.(randn(n, 1)) / 15;
Sigma = 0.9 * rand(n, n) .- 0.15;
Sigma_nom = Sigma' * Sigma;
Sigma_nom .-= (maximum(Sigma_nom) - 0.9)

#-

# Form and solve portfolio optimization problem.
# Here we minimize risk while requiring a 0.1 return.
using Convex, SCS
w = Variable(n);
ret = dot(r, w);
risk = sum(quadform(w, Sigma_nom));
problem = minimize(risk, [sum(w) == 1, ret >= 0.1, norm(w, 1) <= 2])
solve!(problem, SCS.Optimizer; silent_solver = true)
wval = vec(evaluate(w))

#-

# Form and solve worst-case risk analysis problem.
Sigma = Semidefinite(n);
Delta = Variable(n, n);
risk = sum(quadform(wval, Sigma));
problem = maximize(
    risk,
    [
        Sigma == Sigma_nom + Delta,
        diag(Delta) == 0,
        abs(Delta) <= 0.2,
        Delta == Delta',
    ],
);
solve!(problem, SCS.Optimizer; silent_solver = true)
println(
    "standard deviation = ",
    round(sqrt(wval' * Sigma_nom * wval), sigdigits = 2),
);
println(
    "worst-case standard deviation = ",
    round(sqrt(evaluate(risk)), sigdigits = 2),
);
println("worst-case Delta = ");
println(round.(evaluate(Delta), sigdigits = 2));


text:
function second2hms(x::Second)::Dict{DataType, Int}
    ps = Dates.canonicalize(Dates.CompoundPeriod(x))
    a = Dict{DataType, Int}(k => 0 for k in [Hour, Minute, Second])
    ts = [Day, Week, Month, Year]
    for p in ps.periods
        if typeof(p) in ts
            a[Hour] += Hour(p).value
        else
            a[typeof(p)] += p.value
        end
    end
    return a
end
shorten(s::String, k::Int) = length(s) > 2k + 1 ? s[1:k]*"…"*s[(end-k + 1):end] : s
function shorten(vfs::OrderedSet{String}, m)
    nmax = maximum(map(length, vfs))
    n = min(m, nmax) - 1
    while n < nmax
        n += 1
        if allunique(shorten(vf, n) for vf in vfs)
            break
        end
    end
    return OrderedDict(shorten(vf, n) => vf for vf in vfs)
end

function openit(f::String)
    isfile(f) || systemerror("$f not found", true)
    cmd = if is_windows()
        `explorer $f`
    elseif is_linux()
        `xdg-open $f`
    elseif is_apple()
        try
            `open -a "quicktime player" $f`
        catch
            `open $f`
        end
    else
        error("Unknown OS")
    end
    return run(ignorestatus(cmd))
        #stream, proc = open(cmd)
        # return proc
        # try to see if you can kill the spawned process (closing the movie player). this will be useful for testing this, and for managing shit once the user is done (not sure if all the players automatically close when the user quits julia)
end
function findshortfile(v::String, d::Dict{String, String})::String
    for k in keys(d)
        d[k] == v && return k
    end
    error("Couldn't find $v in $d")
end
function validargs(_, y, m, d, rest...)
    0 < m < 13 || return Nullable{ArgumentError}(ArgumentError("Month: $m out of range (1:12)"))
    0 < d < daysinmonth(y,m)+1 || return Nullable{ArgumentError}(ArgumentError("Day: $d out of range (1:$(daysinmonth(y,m)))"))
    return Nullable{ArgumentError}()
end


text:
module IntegralEquations2D

using QuadGK, HCubature, FastGaussQuadrature,
    SpecialFunctions, StaticArrays,
    IntervalSets, DomainSets,
    DomainIntegrals,
    BasisFunctions, FrameFun, GridArrays,
    CardinalBSplines, CompactTranslatesDict,
    LinearAlgebra,
    RecipesBase

import Base:
    Threads

import DomainSets:
    Domain,
    EuclideanDomain,
    applymap,
    jacobian,
    parameterization

import BasisFunctions:
    AbstractOperator,
    GridSampling,
    ProjectionSampling,
    SamplingOperator,
    measure,
    integral,
    period,
    matrix,
    domain

import DomainIntegrals:
    Singularity,
    NoSingularity,
    CurveSingularity,
    SingularDiagonal,
    LogSingPoint

export measure

const Vec2{T} = SVector{2,T}
const Vec3{T} = SVector{3,T}

include("util/common.jl")

include("domains/param.jl")
include("domains/polygons.jl")
include("domains/smooth.jl")
include("domains/recipes.jl")

include("assembly/refinable.jl")
include("assembly/quadrature.jl")
include("assembly/sampling.jl")

include("basis/splines.jl")

include("operator/kernel.jl")
include("operator/operator.jl")
include("operator/bie.jl")
include("operator/bem.jl")

include("assembly/assembly.jl")
include("assembly/field.jl")

include("solve/solve.jl")

include("applications/helmholtz.jl")
include("applications/laplace.jl")

end # module


text:
using SeisPGVPGA
using Test

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


text:
struct Segmenter{S} end
function segment(::Segmenter{2}, Q̄::AbstractFloat, H̄::AbstractFloat)
    Q̄s = (0.75*Q̄, 0.25*Q̄)
    μ = H̄/(Q̄s[1] + 0.95*Q̄s[2])
    μs = (μ, 0.95*μ)
    return Q̄s, μs
end
function segment_percentage(::Segmenter{2}, segment::Integer)
    if segment == 1
        return 0.75
    elseif segment == 2
        return 0.25
    else
        error("Segment out of bounds.")
    end
end


text:
module TestProblems

import ..generateChild, ..gen_ref_dirs, ..norm, ..get_non_dominated_solutions

include("box-constrained.jl")
include("constrained.jl")
include("multiobjective.jl")

"""
    get_problem(problem)

Returns a 3-tuple with the objective function, the bounds and 100 Pareto solutions for
multi-objective optimization problems or the optimal solutions for (box)constrained optimization
problems.

Here, `problem` can be one of the following symbols:

Single-objective:
- `:sphere`
- `:discus`
- `:rastrigin`

Multi-objective:
- `:ZDT1`
- `:ZDT2`
- `:ZDT3`
- `:ZDT4`
- `:ZDT6`

Many-objective:
- `:DTLZ1`
- `:DTLZ2`
- `:DTLZ3`
- `:DTLZ4`
- `:DTLZ5`
- `:DTLZ6`

## Example

```jldoctest
julia> import Metaheuristics: TestProblems, optimize

julia> f, bounds, pareto_solutions = TestProblems.get_problem(:ZDT3);

julia> bounds
2×30 Array{Float64,2}:
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0     1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0

julia> pareto_solutions
                           F space
          ┌────────────────────────────────────────┐
        1 │⢅⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
          │⠈⢢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
          │⠀⠀⠙⠒⠀⠀⠀⠀⡆⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
          │⠀⠀⠀⠀⠀⠀⠀⠀⠘⡄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
          │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠸⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
          │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠱⠄⠀⠀⠀⠀⠀⠀⡄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
          │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢱⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
   f_2    │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠬⡦⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
          │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠘⠂⠀⠀⠀⠀⠀⠀⢢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
          │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠸⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
          │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢧⡀⠀⠀⠀⠀⠀⠀⢀⠀⠀⠀│
          │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⠀⠀│
          │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⡆⠀⠀│
          │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⠀⠀│
       -1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
          └────────────────────────────────────────┘
          0                                      0.9
                             f_1

```
"""
function get_problem(problem)
    expr = Meta.parse( string(problem) * "()")
    return eval(expr)
end

end


text:
#
#
#       None of this probably works!!!
#
#
"""
0 = <x|H - E0|x'>v(x') + <x|H - E0|p>v(p)
0 = <x|H - E0|x'>v(x') + <x|H|p>v(p)
A(x,x')v(x') = -H(x,p)v(p)

here, x is outside the reference space, and p is inside

Ax=b

works for one root at a time
"""
function tucker_cepa_solve!(ref_vector::TuckerState, ci_vector::TuckerState, cluster_ops, clustered_ham; tol=1e-5)
#={{{=#
    fold!(ref_vector)
    fold!(ci_vector)
    sig = deepcopy(ref_vector)
    zero!(sig)
    build_sigma!(sig, ref_vector, cluster_ops, clustered_ham)
    e0 = dot(ref_vector, sig)
    size(e0) == (1,1) || error("Only one state at a time please")
    e0 = e0[1,1]
    @printf(" Reference Energy: %12.8f\n",e0)

    x_vector = deepcopy(ci_vector)
    #
    # now remove reference space from ci_vector
    for (fock,configs) in ref_vector
        if haskey(x_vector, fock)
            for (config,coeffs) in configs
                if haskey(x_vector[fock], config)
                    delete!(x_vector[fock], config)
                end
            end
        end
    end

    b = deepcopy(x_vector)
    zero!(b)
    build_sigma!(b, ref_vector, cluster_ops, clustered_ham)
    bv = -get_vector(b)

    function mymatvec(v)
        unfold!(x_vector)
        set_vector!(x_vector, v)
        fold!(x_vector)
        sig = deepcopy(x_vector)
        zero!(sig)
        build_sigma!(sig, x_vector, cluster_ops, clustered_ham)
        unfold!(x_vector)
        unfold!(sig)

        sig_out = get_vector(sig)
        sig_out .-= e0*get_vector(x_vector)
        return sig_out
    end
    dim = length(x_vector)
    Axx = LinearMap(mymatvec, dim, dim)
    #Axx = LinearMap(mymatvec, dim, dim; issymmetric=true, ismutating=false, ishermitian=true)

    x, solver = cg!(get_vector(x_vector), Axx,bv,log=true)

    set_vector!(x_vector, x)

    sig = deepcopy(ref_vector)
    zero!(sig)
    build_sigma!(sig,x_vector, cluster_ops, clustered_ham)
    ecorr = dot(sig,ref_vector)
    size(ecorr) == (1,1) || error(" Dimension Error")
    ecorr = ecorr[1]

    zero!(ci_vector)
    add!(ci_vector, ref_vector)
    add!(ci_vector, x_vector)

    #x, info = linsolve(Hmap,zeros(size(v0)))
    return ecorr+e0, x
end#=}}}=#

"""
    get_foi(v::TuckerState, clustered_ham, q_spaces; nroots=1, nbody=2)
Compute the first-order interacting space as defined by clustered_ham

e.g.,
 1(p') 3(r) 4(q's)  *  v[(1,1),(1,1),(1,1),(1,1)][1:1,1:1,1:1,1:1]  =>  v[(2,1), (1,1), (0,1), (1,1)][1:N, 1:1, 1:N, 1:N]
"""
function get_foi(v::TuckerState, clustered_ham, q_spaces; nroots=1, nbody=2)
    println(" Prepare empty TuckerState spanning the FOI of input")#={{{=#
    foi = deepcopy(v)
    na = 0
    nb = 0

    for (fock, tconfigs) in v

        na = sum([f[1] for f in fock])
        nb = sum([f[2] for f in fock])
        for (fock_trans, terms) in clustered_ham

            #
            # new fock sector configuration
            new_fock = fock + fock_trans

            #
            # check that each cluster doesn't have too many/few electrons
            ok = true
            for ci in v.clusters
                if new_fock[ci.idx][1] > length(ci) || new_fock[ci.idx][2] > length(ci)
                    ok = false
                end
                if new_fock[ci.idx][1] < 0 || new_fock[ci.idx][2] < 0
                    ok = false
                end
            end
            ok == true || continue

            if haskey(v.data, new_fock) == false
                add_fockconfig!(foi, new_fock)
            end

            #
            # find the cluster state index ranges (TuckerConfig) reached by Hamiltonian
            for (tconfig, coeffs) in tconfigs
                for term in terms
                    new_tconfig = deepcopy(tconfig)

                    length(term.clusters) <= nbody || continue

                    #
                    # for current term, expand index ranges for active clusters
                    for cidx in 1:length(term.clusters)
                        ci = term.clusters[cidx]
                        #tmp_spaces[ci.idx] = q_spaces[ci.idx][new_fock[ci.idx]]
                        #start = 1
                        #stop  = size(bases[ci.idx][new_fock[ci.idx]])[2]
                        #new_tconfig[ci.idx] = start:stop
                        new_tconfig[ci.idx] = q_spaces[ci.idx][new_fock[ci.idx]]
                    end

                    #
                    # remove any previously defined TuckerConfigs which are subspaces of the new TuckerConfig
                    for key in keys(foi[new_fock])
                        if is_subset(key, new_tconfig)
                            delete!(foi[new_fock], key)
                        end
                    end

                    #
                    # determine if new TuckerConfig is a subspace of any previously defined TuckerConfigs
                    is_new = true
                    for key in keys(foi[new_fock])
                        if is_subset(new_tconfig, key)
                            is_new = false
                        end
                    end
                    if is_new == true
                        foi[new_fock][new_tconfig] = zeros(length.(new_tconfig)..., nroots)
                    end
                end
            end
        end
    end
    prune_empty_fock_spaces!(foi)
    return foi
#=}}}=#
end

"""
    get_nbody_tucker_space(v::TuckerState, p_spaces, q_spaces; nroots=1, nbody=2)
Get a vector dimensioned according to the n-body Tucker scheme
- `v::TuckerState` = reference P-space vector
- `p_spaces` = `Vector{ClusterSubspace}` denoting all the cluster P-spaces
- `q_spaces` = `Vector{ClusterSubspace}` denoting all the cluster Q-spaces
- `nbody`    = n-body order
"""
function get_nbody_tucker_space(v::TuckerState, p_spaces, q_spaces, na, nb; nroots=1, nbody=2)
    clusters = v.clusters
    println(" Prepare empty TuckerState spanning the n-body Tucker space with nbody = ", nbody)#={{{=#
    ci_vector = deepcopy(v)
    if nbody >= 1
        for ci in clusters
            tmp_spaces = copy(p_spaces)
            tmp_spaces[ci.idx] = q_spaces[ci.idx]
            FermiCG.add!(ci_vector, FermiCG.TuckerState(clusters, tmp_spaces, na, nb))
        end
    end
    if nbody >= 2
        for ci in clusters
            for cj in clusters
                ci.idx < cj.idx || continue
                tmp_spaces = copy(p_spaces)
                tmp_spaces[ci.idx] = q_spaces[ci.idx]
                tmp_spaces[cj.idx] = q_spaces[cj.idx]
                FermiCG.add!(ci_vector, FermiCG.TuckerState(clusters, tmp_spaces, na, na))
            end
        end
    end
    if nbody >= 3
        for ci in clusters
            for cj in clusters
                for ck in clusters
                    ci.idx < cj.idx || continue
                    cj.idx < ck.idx || continue
                    tmp_spaces = copy(p_spaces)
                    tmp_spaces[ci.idx] = q_spaces[ci.idx]
                    tmp_spaces[cj.idx] = q_spaces[cj.idx]
                    tmp_spaces[ck.idx] = q_spaces[ck.idx]
                    FermiCG.add!(ci_vector, FermiCG.TuckerState(clusters, tmp_spaces, na, na))
                end
            end
        end
    end
    if nbody >= 4
        for ci in clusters
            for cj in clusters
                for ck in clusters
                    for cl in clusters
                        ci.idx < cj.idx || continue
                        cj.idx < ck.idx || continue
                        ck.idx < cl.idx || continue
                        tmp_spaces = copy(p_spaces)
                        tmp_spaces[ci.idx] = q_spaces[ci.idx]
                        tmp_spaces[cj.idx] = q_spaces[cj.idx]
                        tmp_spaces[ck.idx] = q_spaces[ck.idx]
                        tmp_spaces[cl.idx] = q_spaces[cl.idx]
                        FermiCG.add!(ci_vector, FermiCG.TuckerState(clusters, tmp_spaces, na, nb))
                    end
                end
            end
        end
    end
    return ci_vector
#=}}}=#
end

max_stars_repo_name: nmayhall-vt/FermiCG

max_issues_repo_name: nmayhall-vt/FermiCG

max_forks_repo_name: nmayhall-vt/FermiCG


text:
function print_julia_version()
    jit = get_jit(BOOT_COMPILER_REF[])
    addr = lookup(jit, "print_julia_version")
    return ccall(pointer(addr), Cvoid, ())
end


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

using Base: @propagate_inbounds

# SIMD-compatible KBN-summation

@inline function kbn_add(a::NTuple{2, Real}, b::Real)
    s = a[1] + b
    c = a[2] + ifelse(
        abs(a[1]) >= abs(b),
        (a[1] - s) + b,
        (b - s) + a[1]
    )
    (s, c)
end

@inline function kbn_add(a::NTuple{2, Real}, b::NTuple{2, Real})
    s = a[1] + b[1]
    c = ifelse(
        abs(a[1]) > abs(b[1]),
        (((a[1] - s) + b[1]) + b[2]) + a[2],
        (((b[1] - s) + a[1]) + a[2]) + b[2]
    )
    (s, c)
end

doc"""
    OnlineMvMean{T<:AbstractFloat} <: AbstractVector{T}

Multi-variate mean implemented via Kahan-Babuška-Neumaier summation.
"""
mutable struct OnlineMvMean{T<:AbstractFloat} <: AbstractVector{T}
    m::Int
    sum_w::Double{T}
    S::Vector{T}
    C::Vector{T}

    OnlineMvMean{T}(m::Integer) where {T<:AbstractFloat} =
        new{T}(m, zero(Double{T}), zeros(T, m), zeros(T, m))
end

export OnlineMvMean

OnlineMvMean(m::Integer) = OnlineMvMean{Float64}(m::Integer)

@inline Base.size(omn::OnlineMvMean) = size(omn.S)

@propagate_inbounds function Base.getindex(omn::OnlineMvMean{T}, idxs::Integer...) where {T}
    T((Single(omn.S[idxs...]) + Single(omn.C[idxs...])) / omn.sum_w)
end

function Base.merge!(target::OnlineMvMean, others::OnlineMvMean...)
    for x in others
        target.m != x.m && throw(ArgumentError("can't merge OnlineMvMean instances with different size"))
    end
    for x in others
        target.sum_w += x.sum_w
        target_S = target.S; target_C = target.C
        x_S = x.S; x_C = x.C
        @assert eachindex(target_S) == eachindex(x.S) == eachindex(target_C) == eachindex(x.C)  # TODO: Use exception instead of assert
        @inbounds @simd for i in eachindex(target_S)
            target_S[i], target_C[i] = kbn_add(
                (target_S[i], target_C[i]),
                (x.S[i], x.C[i])
            )
        end
    end
    target
end

@inline function push_contiguous!(
    omn::OnlineMvMean{T}, data::AbstractArray,
    start::Integer, weight::Real = one(T)
) where {T}
    @assert _iscontiguous(data)  # TODO: Use exception instead of assert

    # Workaround for lack of promotion between, e.g., Float32 and Double{Float64}
    weight_conv = T(weight)

    m = omn.m
    S = omn.S
    C = omn.C

    idxs = Base.OneTo(m)

    dshft = Int(start) - 1

    @assert idxs == indices(S, 1) == indices(C, 1)  # TODO: Use exception instead of assert
    checkbounds(data, idxs + dshft)

    omn.sum_w += weight_conv

    @inbounds @simd for i in idxs
        x = weight_conv * data[i + dshft]
        S[i], C[i] = kbn_add((S[i], C[i]), x)
    end

    omn
end

doc"""
    OnlineMvCov{T<:AbstractFloat,W} <: AbstractMatrix{T}

Implementation based on variance calculation Algorithms of Welford and West.

`W` must either be `Weights` (no bias correction) or one of `AnalyticWeights`,
`FrequencyWeights` or `ProbabilityWeights` to specify the desired bias
correction method.
"""

mutable struct OnlineMvCov{T<:AbstractFloat,W} <: AbstractMatrix{T}
    m::Int
    n::Int64
    sum_w::Double{T}
    sum_w2::Double{T}
    Mean_X::Vector{T}
    New_Mean_X::Vector{T}
    S::Matrix{T}

    OnlineMvCov{T,W}(m::Integer) where {T<:AbstractFloat,W} =
        new{T,W}(
            m, zero(Int64), zero(Double{T}), zero(Double{T}),
            zeros(T, m), zeros(T, m), zeros(T, m, m)
        )
end

export OnlineMvCov

OnlineMvCov(m::Integer) = OnlineMvCov{Float64, ProbabilityWeights}(m::Integer)

@inline Base.size(ocv::OnlineMvCov) = size(ocv.S)

@propagate_inbounds function Base.getindex(ocv::OnlineMvCov{T, Weights}, idxs::Integer...) where {T}
    sum_w = ocv.sum_w
    ifelse(
        sum_w > zero(T),
        T(ocv.S[idxs...] / sum_w),
        T(NaN)
    )
end

@propagate_inbounds function Base.getindex(ocv::OnlineMvCov{T, AnalyticWeights}, idxs::Integer...) where {T}
    sum_w = ocv.sum_w
    sum_w2 = ocv.sum_w2
    d = sum_w - sum_w2 / sum_w
    ifelse(
        sum_w > zero(T) && d > zero(T),
        T(ocv.S[idxs...] / d),
        T(NaN)
    )
end

@propagate_inbounds function Base.getindex(ocv::OnlineMvCov{T, FrequencyWeights}, idxs::Integer...) where {T}
    sum_w = ocv.sum_w
    ifelse(
        sum_w > one(T),
        T(ocv.S[idxs...] / (sum_w - one(T))),
        T(NaN)
    )
end

@propagate_inbounds function Base.getindex(ocv::OnlineMvCov{T, ProbabilityWeights}, idxs::Integer...) where {T}
    n = ocv.n
    sum_w = ocv.sum_w
    ifelse(
        n > one(T) && sum_w > zero(T),
        T(ocv.S[idxs...] * n / ((n - one(T)) * sum_w)),
        T(NaN)
    )
end

function Base.merge!(target::OnlineMvCov{T,W}, others::OnlineMvCov{T,W}...) where {T,W}
    # TODO: Improve implementation

    for x in others
        target.m != x.m && throw(ArgumentError("can't merge OnlineMvCov instances with different size"))
    end
    m = target.m
    for x in others
        idxs = Base.OneTo(m)
        scl = (target.sum_w * x.sum_w) / (target.sum_w + x.sum_w)
        for j in idxs
            for i in idxs
                target.S[i, j] = scl *
                    (target.Mean_X[i] - x.Mean_X[i]) *
                    (target.Mean_X[j] - x.Mean_X[j]) +
                    target.S[i, j] + x.S[i, j]
            end
        end

        d = inv(T(target.sum_w + x.sum_w))
        for i in idxs
            target.Mean_X[i] = (target.Mean_X[i] * target.sum_w + x.Mean_X[i] * x.sum_w) * d
            target.New_Mean_X[i] = target.Mean_X[i]
        end

        target.n += x.n
        target.sum_w += x.sum_w
        target.sum_w2 += x.sum_w2
    end
    target
end

@inline function push_contiguous!(
    ocv::OnlineMvCov{T,W}, data::AbstractArray,
    start::Integer,
    weight::Real = one(T)
) where {T,W}
    @assert _iscontiguous(data)  # TODO: Use exception instead of assert

    # Ignore zero weights (can't be handled)
    if weight ≈ 0
        return ocv
    end

    # Workaround for lack of promotion between, e.g., Float32 and Double{Float64}
    weight_conv = T(weight)

    m = ocv.m
    n = ocv.n
    sum_w = ocv.sum_w
    sum_w2 = ocv.sum_w2
    Mean_X = ocv.Mean_X
    New_Mean_X = ocv.New_Mean_X
    S = ocv.S

    idxs = Base.OneTo(m)

    dshft = Int(start) - 1

    @assert idxs == indices(Mean_X, 1) == indices(New_Mean_X, 1) == indices(S, 1) == indices(S, 2)  # TODO: Use exception instead of assert
    checkbounds(data, idxs + dshft)

    n += one(n)
    sum_w += weight_conv
    sum_w2 += weight_conv^2

    weight_over_sum_w = T(weight_conv / sum_w)

    @inbounds @simd for i in idxs
        x = data[i + dshft]
        mean_X = Mean_X[i]
        New_Mean_X[i] = muladd(x - mean_X, weight_over_sum_w, mean_X)
    end

    @inbounds for j in idxs
        new_dx_j = data[j + dshft] - New_Mean_X[j]
        j_offs = sub2ind(S, 0, j)
        @assert sub2ind(S, last(idxs), j) == last(idxs) + j_offs  # TODO: Use exception instead of assert
        @simd for i in idxs
            dx_i = data[i + dshft] - Mean_X[i]
            S[i + j_offs] = muladd(dx_i, weight_conv * new_dx_j, S[i + j_offs])
        end
    end

    @inbounds @simd for i in idxs
        Mean_X[i] = New_Mean_X[i]
    end

    ocv.n = n
    ocv.sum_w = sum_w
    ocv.sum_w2 = sum_w2

    ocv
end

doc"""
    BasicMvStatistics{T<:Real,W}

`W` must either be `Weights` (no bias correction) or one of `AnalyticWeights`,
`FrequencyWeights` or `ProbabilityWeights` to specify the desired bias
correction method.
"""

mutable struct BasicMvStatistics{T<:Real,W}
    m::Int
    mean::OnlineMvMean{T}
    cov::OnlineMvCov{T,W}
    maximum::Vector{T}
    minimum::Vector{T}

    BasicMvStatistics{T,W}(
        mean::OnlineMvMean{T},
        cov::OnlineMvCov{T,W},
        maximum::Vector{T},
        minimum::Vector{T}
    ) where {T<:Real,W} = new(mean, cov, maximum, minimum)

    BasicMvStatistics{T,W}(m::Integer) where {T<:Real,W} =
        new(m, OnlineMvMean{T}(m), OnlineMvCov{T,W}(m), fill(typemin(T), m), fill(typemax(T), m))
end

export BasicMvStatistics

function Base.merge!(target::BasicMvStatistics, others::BasicMvStatistics...)
    for x in others
        merge!(target.mean, x.mean)
        merge!(target.cov, x.cov)
        target.maximum = max.(target.maximum, x.maximum)
        target.minimum = min.(target.minimum, x.minimum)
    end
    target
end

function push_contiguous!(
    stats::BasicMvStatistics{T,W}, data::AbstractArray,
    start::Integer,
    weight::Real = one(T)
) where {T,W}
    @assert _iscontiguous(data)  # TODO: Use exception instead of assert

    push_contiguous!(stats.mean, data, start, weight)
    push_contiguous!(stats.cov, data, start, weight)

    max_v = stats.maximum
    min_v = stats.minimum

    m = stats.m
    idxs = Base.OneTo(m)
    dshft = Int(start) - 1

    @assert idxs == indices(max_v, 1) == indices(min_v, 1)  # TODO: Use exception instead of assert
    checkbounds(data, idxs + dshft)

    @inbounds @simd for i in idxs
        x = data[i + dshft]
        max_v[i] = max(max_v[i], x)
        min_v[i] = min(min_v[i], x)
    end

    stats
end

const OnlineMvStatistic = Union{OnlineMvMean, OnlineMvCov, BasicMvStatistics}

@inline Base.push!(target::OnlineMvStatistic, data::AbstractVector, weight::Real = 1) =
    push_contiguous!(target, data, first(linearindices(data)), weight)

function Base.append!(target::OnlineMvStatistic, data::AbstractMatrix, vardim::Integer = 1)
    if (vardim == 1)
        throw(ArgumentError("vardim == $vardim not supported (yet)"))
    elseif (vardim == 2)
        @assert target.m == size(data, 1)  # TODO: Use exception instead of assert
        @assert _iscontiguous(data)  # TODO: Use exception instead of assert
        @inbounds for i in indices(data, 2)
            push_contiguous!(target, data, sub2ind(data, 1, i))
        end
    else
        throw(ArgumentError("Value of vardim must be 2, not $vardim"))
    end

    target
end

function Base.append!(target::OnlineMvStatistic, data::AbstractMatrix, weights::AbstractVector, vardim::Integer = 1)
    if (vardim == 1)
        throw(ArgumentError("vardim == $vardim not supported (yet)"))
    elseif (vardim == 2)
        @assert _iscontiguous(data)  # TODO: Use exception instead of assert
        @assert target.m == size(data, 1)  # TODO: Use exception instead of assert
        @assert indices(data, 2) == indices(weights, 1)  # TODO: Use exception instead of assert
        @inbounds for i in indices(data, 2)
            push_contiguous!(target, data, sub2ind(data, 1, i), weights[i])
        end
    else
        throw(ArgumentError("Value of vardim must be 2, not $vardim"))
    end

    target
end

Base.merge(x::S, others::S...) where {S <: OnlineMvStatistic} = merge!(deepcopy(x), others...)


text:
using Queueing
using Base.Test
using Distributions
using DataStructures

λ = 3
μ = 4
ρ = λ / μ
expected_total_time = ρ / (λ*(1 - ρ))
expected_wait_time = expected_total_time - 1 /μ
results = MM1(λ, μ, 100000)
@test_approx_eq_eps expected_wait_time results.average_wait_time 0.05
@test_approx_eq_eps expected_total_time results.average_system_time 0.05

K = 10
props = QueueProperties(Exponential(1/λ), Exponential(1/μ), 1, K)
a = SimulationArgs(100000, 0, [QueueNode(1, props, true, [])])
λ_1 = 1 - ((1-ρ)*ρ^K)/(1-ρ^(K+1))
N = (ρ*(1-(K+1)*ρ^K + K*ρ^(K+1))) / ((1-ρ)*(1-ρ^(K+1)))
expected_total_time = N / (λ*λ_1)
results = simulate(a)
@test_approx_eq_eps expected_total_time results.average_system_time 0.05
agg_results = aggregate_simulate(a, 5)
@test_approx_eq_eps expected_total_time agg_results.average_system_time 0.05

props = QueueProperties(Exponential(1/λ), Exponential(1/μ), 1, -1)
prob = 0.5
a = QueueNode(1, props, false, [])
b = QueueNode(2, props, true, [QueueEdge(a.queue_id, prob)])
jackson_network_args = SimulationArgs(100000, 0, [a, b])
λ_1 = λ * prob
expected_total_time = (1/λ) * ((λ / (μ - λ)) + (λ_1 / (μ - λ_1)))
results = simulate(jackson_network_args)
@test_approx_eq_eps expected_total_time results.average_system_time 0.05


text:
mutable struct PeriodicValues{T}
    values          ::Vector{Tuple{DateTime, T}}
    period          ::Period
    last_dt         ::DateTime
end

function periodic_collector(
    type        ::Type{T},
    period      ::Period
)::Tuple{Function, PeriodicValues{T}} where T
    values = Vector{Tuple{DateTime, T}}()
    pv = PeriodicValues{T}(values, period, DateTime(0))

    @inline function collector(dt::DateTime, value::T)::Nothing
        if (dt - pv.last_dt) >= period
            push!(values, (dt, value))
            pv.last_dt = dt
        end
        nothing
    end

    return collector, pv
end

# ----------------------------------------------------------

mutable struct PredicateValues{T}
    values          ::Vector{Tuple{DateTime, T}}
    last_dt         ::DateTime
    last_value      ::T
    data            ::Any               # user-defined data for use in predicate
end

function predicate_collector(
    type            ::Type{T},
    predicate       ::Function,
    init_value      ::T
)::Tuple{Function, PredicateValues{T}} where T
    values = Vector{Tuple{DateTime, T}}()
    pv = PredicateValues{T}(values, DateTime(0), init_value, nothing)

    @inline function collector(dt::DateTime, value::T)::Nothing
        if predicate(pv, dt, value)
            push!(values, (dt, value))
            pv.last_dt = dt
            pv.last_value = value
        end
        nothing
    end

    return collector, pv
end

# ----------------------------------------------------------

mutable struct MinValue{T}
    dt         ::DateTime
    min_value  ::T
end

function min_value_collector(type::Type{T})::Tuple{Function, MinValue{T}} where T
    mv = MinValue{T}(DateTime(0), 1e50)

    @inline function collector(dt::DateTime, value::T)::Nothing
        if value < mv.min_value
            mv.min_value = value
            mv.dt = dt
        end
        nothing
    end

    return collector, mv
end

# ----------------------------------------------------------

mutable struct MaxValue{T}
    dt         ::DateTime
    max_value  ::T
end

function max_value_collector(type::Type{T})::Tuple{Function, MaxValue{T}} where T
    mv = MaxValue{T}(DateTime(0), typemin(T))

    @inline function collector(dt::DateTime, value::T)::Nothing
        if value > mv.max_value
            mv.max_value = value
            mv.dt = dt
        end
        nothing
    end

    return collector, mv
end

# ----------------------------------------------------------

@enum DrawdownMode Percentage PnL

mutable struct DrawdownValues
    values          ::Vector{Tuple{DateTime, Price}}
    mode            ::DrawdownMode
    max_equity      ::Price
    last_dt         ::DateTime
    data            ::Any               # user-defined data for use in predicate
end

function drawdown_collector(mode::DrawdownMode, predicate::Function)::Tuple{Function, DrawdownValues}
    dv = DrawdownValues(
        Vector{Tuple{DateTime, Price}}(),
        mode,
        -1e50,
        DateTime(0),
        nothing)

    @inline function collector(dt::DateTime, equity::Price)::Nothing
        # keep track of max equity value
        dv.max_equity = max(dv.max_equity, equity)

        # should collect new drawdown value?
        if predicate(dv, dt, equity)
            drawdown = min(0.0, equity - dv.max_equity)
            if mode == Percentage::DrawdownMode
                drawdown /= dv.max_equity
            end
            push!(dv.values, (dt, drawdown))
            dv.last_dt = dt
        end

        nothing
    end

    return collector, dv
end

# ----------------------------------------------------------


text:
using BenchmarkTools
using Unitful, UnitfulAstro
using Transits.Orbits: KeplerianOrbit, flip,
                       relative_position, compute_aor, compute_rho, compute_b,
                       _star_position, _planet_position,
                       stringify_units

# Setup python env
import Pkg
ENV["PYTHON"] = ""
Pkg.build("PyCall")
using PyCall, Conda
Conda.add(["batman-package"]; channel="conda-forge")

# Constants
const G_nom = 2942.2062175044193 # Rsun^3/Msun/d^2
const MsunRsun_to_gcc = ustrip(u"g/cm^3", 1.0u"Msun/Rsun^3")

function compute_r(orbit, t)
    pos = relative_position.(orbit, t)
    r = map(pos) do arr
        hypot(arr[1], arr[2])
    end
    return r
end

# Convert vector of vectors -> matrix
as_matrix(pos) = reinterpret(reshape, Float64, pos) |> permutedims

# Tests from:
# https://github.com/exoplanet-dev/exoplanet/blob/main/tests/orbits/keplerian_test.py

@testset "KeplerianOrbit: sky coords" begin
    # Comparison coords from `batman`
    py"""
    import numpy as np
    from batman import _rsky

    def sky_coords():
        t = np.linspace(-100, 100, 1_000)

        t0, period, a, e, omega, incl = (
            x.flatten()
            for x in np.meshgrid(
                np.linspace(-5.0, 5.0, 2),
                np.exp(np.linspace(np.log(5.0), np.log(50.0), 3)),
                np.linspace(50.0, 100.0, 2),
                np.linspace(0.0, 0.9, 5),
                np.linspace(-np.pi, np.pi, 3),
                np.arccos(np.linspace(0, 1, 5)[:-1]),
            )
        )

        r_batman = np.empty((len(t), len(t0)))

        for i in range(len(t0)):
            r_batman[:, i] = _rsky._rsky(
                t, t0[i], period[i], a[i], incl[i], e[i], omega[i], 1, 1
            )

        m = r_batman < 100.0

        return {
            "m_sum" : m.sum().item(), # Save native Int format
            "r_batman" : r_batman,
            "m" : m,
            "t" : t,
            "t0" : t0,
            "period" : period,
            "a" : a,
            "e" : e,
            "omega" : omega,
            "incl" : incl,
        }

    """
    sky_coords = py"sky_coords"()

    # Create comparison orbits from Transits.jl
    orbits = [
        KeplerianOrbit(
            aR_star = sky_coords["a"][i],
            P = sky_coords["period"][i],
            incl = sky_coords["incl"][i],
            t0 = sky_coords["t0"][i],
            ecc = sky_coords["e"][i],
            Omega = 0.0,
            omega = sky_coords["omega"][i],
        )
        for i in 1:length(sky_coords["t0"])
    ]

    # Compute coords
    t = sky_coords["t"]
    x = Matrix{Float64}(undef, length(sky_coords["t"]), length(sky_coords["t0"]))
    y = similar(x)
    z = similar(x)
    for (col, orbit) in enumerate(orbits)
        for (row, _t) in enumerate(t)
            pos = relative_position(orbit, _t)
            x[row, col] = pos[1]
            y[row, col] = pos[2]
            z[row, col] = pos[3]
        end
    end

    # Compare
    m = sky_coords["m"]
    r = hypot.(x, y)
    r_Transits = r[m]
    r_batman = sky_coords["r_batman"][m]

    @test sum(m) > 0
    @test allclose(r_Transits, r_batman, atol=2e-5)
    @test all(>(0), z[m])
    no_transit = @. (z[!(m)] < 0) | (r[!(m)] > 2)
    @test all(no_transit)
end

@testset "KeplerianOrbit: construction performance" begin
    b_rho_star = @benchmark KeplerianOrbit(
        rho_star = 2.0,
        R_star = 0.5,
        period = 2.0,
        ecc = 0.0,
        t0 = 0.0,
        incl = π / 2.0,
        Omega = 0.0,
        omega = 0.0,
    )

    b_rho_star_units = @benchmark KeplerianOrbit(
        rho_star = 2.0u"g/cm^3",
        R_star = 0.5u"Rsun",
        period = 2.0u"d",
        ecc = 0.0,
        t0 = 0.0u"d",
        incl = 90.0u"°",
        Omega = 0.0u"°",
        omega = 0.0u"°",
    )

    b_aR_star = @benchmark KeplerianOrbit(
        aR_star = 7.5,
        P = 2.0,
        incl = π / 2.0,
        t0 = 0.0,
        ecc = 0.0,
        Omega = 0.0,
        omega = 0.0,
    )

    b_aR_star_units = @benchmark KeplerianOrbit(
        aR_star = 7.5,
        P = 2.0u"d",
        incl = 90.0u"°",
        t0 = 0.0u"d",
        ecc = 0.0,
        Omega = 0.0u"°",
        omega = 0.0u"°",
    )

    # Units
    @test median(b_rho_star_units.times) ≤ 100_000 # ns
    @test b_rho_star_units.allocs ≤ 500
    @test median(b_aR_star_units.times) ≤ 100_000   # ns
    @test b_aR_star_units.allocs ≤ 500

    # TODO: Look into this

    if Base.VERSION < v"1.7"
        @test b_rho_star.allocs == b_rho_star.memory == 0
        @test b_aR_star.allocs == b_aR_star.memory == 0
        @test median(b_rho_star.times) ≤ 500 # ns
        @test median(b_aR_star.times) ≤ 500 # ns
    else
        # TODO: investigate performance regression
        @test median(b_rho_star.times) ≤ 20_000 # ns
        @test median(b_aR_star.times) ≤ 20_000 # ns
    end
end

@testset "KeplerianOrbit: orbital elements" begin
    orbit = KeplerianOrbit(
        cos_omega=√(2)/2, sin_omega=√(2)/2,
        period=2.0, t0=0.0, b=0.01, M_star=1.0, R_star=1.0, ecc=0.0,
    )
    @test orbit.omega == atan(orbit.sin_omega, orbit.cos_omega)

    orbit = KeplerianOrbit(
        aR_star = 1.0,
        period = 2.0,
        t0 = 0.0,
    )
    @test iszero(orbit.b)

    orbit = KeplerianOrbit(
        period=2.0, t0=0.0, M_star=1.0, R_star=1.0,
        ecc=0.01, omega=0.1, r=0.01,
    )
    @test isnothing(orbit.duration)
end

@testset "KeplerianOrbit: valid inputs" begin
    @test_throws ArgumentError("`b` must also be provided for a circular orbit if `duration given`") KeplerianOrbit(
        duration=0.01,
        period=2.0, t0=0.0, R_star=1.0
    )
    @test_throws ArgumentError("`r` must also be provided if `duration` given") KeplerianOrbit(
        duration=0.01, b=0.0,
        period=2.0, t0=0.0, R_star=1.0
    )
    #@test_throws ArgumentError("Only `ω`, or `cos_ω` and `sin_ω` can be provided") KeplerianOrbit(
    #    omega=0.0, cos_omega=1.0, sin_omega=0.0,
    #    period=2.0, t0=0.0, b=0.0, M_star=1.0, R_star=1.0, ecc=0.0,
    #)
    #@test_throws ArgumentError("`ω` must also be provided if `ecc` specified") KeplerianOrbit(
    #    rho_star=2.0, R_star=0.5, period=2.0, t0=0.0, incl=π/2.0, Omega=0.0, ecc=0.0,
    #)
    @test_throws ArgumentError("Only `incl`, `b`, or `duration` can be given") KeplerianOrbit(
        incl=π/2.0, b=0.0, duration=1.0,
        period=2.0, t0=0.0, R_star=1.0, r=0.01,
    )
    @test_throws ArgumentError("Please specify either `t0` or `tp`") KeplerianOrbit(
        b=0.0, period=2.0, R_star=1.0, M_star=1.0,
    )
    @test_throws ArgumentError("Please only specify one of `t0` or `tp`") KeplerianOrbit(
        b=0.0, period=2.0, R_star=1.0, M_star=1.0, t0=0.0, tp=1.0,
    )
    @test_throws ArgumentError("At least `a` or `P` must be specified") KeplerianOrbit(
        b=0.0, R_star=1.0, M_star=1.0,
    )
    @test_throws ArgumentError("If both `a` and `P` are given, `rho_star` or `M_star` cannot be defined") KeplerianOrbit(
        rho_star=2.0,
        R_star=0.5, a=7.5, period=2.0, t0=0.0, incl=π/2.0, Omega=0.0, omega=0.0, ecc=0.0,
    )
    @test_throws ArgumentError("If both `a` and `P` are given, `rho_star` or `M_star` cannot be defined") KeplerianOrbit(
        M_star=1.0,
        R_star=0.5, a=7.5, period=2.0, t0=0.0, incl=π/2.0, Omega=0.0, omega=0.0, ecc=0.0,
    )
    @test_throws ArgumentError("Must provide exactly two of: `rho_star`, `R_star`, or `M_star` if rho_star not implied") KeplerianOrbit(
        R_star=0.5,
        period=2.0, t0=0.0, incl=π/2.0, Omega=0.0, omega=0.0, ecc=0.0,
    )
    @test_throws ArgumentError("Must provide exactly two of: `rho_star`, `R_star`, or `M_star` if rho_star not implied") KeplerianOrbit(
        M_star=0.5,
        period=2.0, t0=0.0, incl=π/2.0, Omega=0.0, omega=0.0, ecc=0.0,
    )
    @test_throws ArgumentError("Must provide exactly two of: `rho_star`, `R_star`, or `M_star` if rho_star not implied") KeplerianOrbit(
        M_star=0.5, R_star=0.5, rho_star=0.5,
        period=2.0, t0=0.0, incl=π/2.0, Omega=0.0, omega=0.0, ecc=0.0,
    )
end

@testset "KeplerianOrbit: implied inputs" begin
    # R_star ≡ 1.0 R⊙ if not specified
    orbit_no_R_star = KeplerianOrbit(
        rho_star=2.0, period=2.0, t0=0.0,
        incl=π/2.0, Omega=0.0, omega=0.0, ecc=0.0
    )
    @test orbit_no_R_star.R_star == one(orbit_no_R_star.a)

    # Compute M_tot if `a` and `period` given
    orbit_a_period = KeplerianOrbit(
        a=1.0, period=1.0, t0=0.0,
        incl=π/2.0, Omega=0.0, omega=0.0, ecc=0.0
    )
    @test orbit_a_period.M_planet + orbit_a_period.M_star == 4.0 * π^2 / G_nom

    # Compute `R_star` from `M_star`
    orbit_M_star = KeplerianOrbit(
        M_star=4.0*π, rho_star = 1.0, period = 2.0, t0 = 0.0,
        incl = π / 2.0, Omega = 0.0, omega = 0.0, ecc = 0.0
    )
    @test orbit_M_star.R_star == 3.0^(1/3)
end

@testset "KeplerianOrbit: small star" begin
    # Sample model from `Transits.jl`
    orbit = KeplerianOrbit(
        R_star = 0.189,
        M_star = 0.151,
        period =  0.4626413,
        t0 = 0.2,
        b = 0.5,
        ecc = 0.1,
        omega = 0.1,
    )

    # Comparison coords from `batman`
    py"""
    import numpy as np
    from batman import _rsky

    def small_star(period, t0, aR_star, incl, ecc, omega):
        t = np.linspace(0, period, 500)
        r_batman = _rsky._rsky(
            t,
            t0,
            period,
            aR_star,
            incl,
            ecc,
            omega,
            1,
            1
        )

        m = r_batman < 100.0

        return {
            "t": t,
            "r_batman": r_batman,
            "m": m,
        }
    """
    small_star = py"small_star"(orbit.period, orbit.t0, orbit.aR_star, orbit.incl, orbit.ecc, orbit.omega)

    # Compare
    t = small_star["t"]
    r_batman = small_star["r_batman"]
    m = small_star["m"]
    r = compute_r(orbit, t)
    @test sum(m) > 0
    @test allclose(r_batman[m], r[m], atol=2e-5)
end

@testset "KeplerianOrbit: impact" begin
    # Sample model from `Transits.jl`
    orbit = KeplerianOrbit(
        R_star = 0.189,
        M_star = 0.151,
        P = 0.4626413,
        t0 = 0.2,
        b = 0.5,
        ecc = 0.8,
        omega = 0.1,
    )

    pos = relative_position.(orbit, orbit.t0)
    @test hypot(pos[1], pos[2]) ≈ orbit.b
end

@testset "KeplerianOrbit: flip" begin
    orbit = KeplerianOrbit(
        M_star = 1.3,
        R_star = 1.1,
        t0 = 0.5,
        period = 100.0,
        ecc = 0.3,
        incl = 0.25 * π,
        omega = 0.5,
        Omega = 1.0,
        M_planet = 0.1,
    )

    orbit_flipped = flip(orbit, 0.7)

    t = range(0, 100; length=1_000)

    u_star = as_matrix(_star_position.(orbit, orbit.R_star, t))
    u_planet_flipped = as_matrix(_planet_position.(orbit_flipped, orbit.R_star, t))
    for i in 1:3
        @test allclose(u_star[:, i], u_planet_flipped[:, i], atol=1e-5)
    end

    u_planet = as_matrix(_planet_position.(orbit, orbit.R_star, t))
    u_star_flipped = as_matrix(_star_position.(orbit_flipped, orbit.R_star, t))
    for i in 1:3
        @test allclose(u_planet[:, i], u_star_flipped[:, i], atol=1e-5)
    end
end

@testset "KeplerianOrbit: flip circular" begin
    t = range(0, 100; length=1_000)

    orbit = KeplerianOrbit(
        M_star = 1.3,
        M_planet = 0.1,
        R_star = 1.0,
        P = 100.0,
        t0 = 0.5,
        incl = 45.0,
        ecc = 0.0,
        omega = 0.5,
        Omega = 1.0
    )
    orbit_flipped = flip(orbit, 0.7)

    u_star = as_matrix(_star_position.(orbit, orbit.R_star, t))
    u_planet_flipped = as_matrix(_planet_position.(orbit_flipped, orbit.R_star, t))
    for i in 1:3
        @test allclose(u_star[:, i], u_planet_flipped[:, i], atol=1e-5)
    end

    u_planet = as_matrix(_planet_position.(orbit, orbit.R_star, t))
    u_star_flipped = as_matrix(_star_position.(orbit_flipped, orbit.R_star, t))
    for i in 1:3
        @test allclose(u_planet[:, i], u_star_flipped[:, i], atol=1e-5)
    end
end

@testset "KeplerianOrbit: compute_aor" begin
    duration = 0.12
    period = 10.1235
    b = 0.34
    r = 0.06
    R_star = 0.7
    aor = compute_aor(duration, period, b, r=r)

    for orbit in [
            KeplerianOrbit(
                period=period, t0=0.0, b=b, a=R_star * aor, R_star=R_star
            ),
            KeplerianOrbit(
                period = period,
                t0 = 0.0,
                b = b,
                duration = duration,
                R_star = R_star,
                r = r,
            ),
        ]

        x, y, z = _planet_position(orbit, R_star, 0.5*duration)
        @test allclose(hypot(x, y), 1.0 + r)
        x, y, z = _planet_position(orbit, R_star, -0.5*duration)
        @test allclose(hypot(x, y), 1.0 + r)
        x, y, z = _planet_position(orbit, R_star, period + 0.5*duration)
        @test allclose(hypot(x, y), 1.0 + r)
    end
end

@testset "KeplerianOrbit: compute_rho" begin
    @test isnothing(compute_rho(0.0, nothing))
end

@testset "KeplerianOrbit: compute_b" begin
    @test iszero(compute_b(-1.0, 1.0, 10.0, 2.0, 10.0, 1.0, 0.0))
end

@testset "KeplerianOrbit: stringify_units" begin
    #@test stringify_units(1u"Rsun", "Rsun") == "1 R⊙"
    @test stringify_units(1, "R⊙") == "1.0000 R⊙"
end

@testset "KeplerianOrbit: unit conversions" begin
    orbit = KeplerianOrbit(a=12.0, t0=0.0, b=0.0, R_star=1.0, M_star=1.0, M_planet=0.01, r=0.01)
    rho_planet_1 = orbit.rho_planet*u"Msun/Rsun^3" |> u"g/cm^3"
    rho_planet_2 = orbit.rho_planet * MsunRsun_to_gcc
    @test rho_planet_1.val == rho_planet_2

    orbit = KeplerianOrbit(a=12.0u"Rsun", t0=0.0u"d", b=0.0, R_star=1.0u"Rsun", M_star=1.0u"Msun")
    @test isnan(orbit.rho_planet)

    orbit = KeplerianOrbit(a=12.0u"Rsun", t0=0.0u"d", b=0.0, R_star=1.0u"Rsun", M_planet=0.01u"Msun", M_star=1.0u"Msun", r=0.01)
    rho_planet = orbit.rho_planet |> u"g/cm^3"
    rho_star = orbit.rho_star |> u"g/cm^3"
    @test rho_planet.val == orbit.rho_planet.val*MsunRsun_to_gcc
    @test rho_star.val == orbit.rho_star.val*MsunRsun_to_gcc
end

@testset "KeplerianOrbit: aliased kwargs" begin
    orbit_standard_1 = KeplerianOrbit(a=12.0, t0=0.0, b=0.0, R_star=1.0, M_star=1.0, M_planet=0.01, r=0.01)
    orbit_kwarg_alias_1 = KeplerianOrbit(a=12.0, t0=0.0, b=0.0, Rs=1.0, Ms=1.0, Mp=0.01, RpRs=0.01)
    @test orbit_standard_1 === orbit_kwarg_alias_1

    orbit_standard_2 = KeplerianOrbit(
        rho_star = 2.0, R_star = 0.5, period = 2.0, ecc = 0.0, t0 = 0.0,
        incl = π / 2.0, Omega = 0.0, omega = 0.0,
    )
    orbit_kwarg_alias_2 = KeplerianOrbit(
        ρ_star = 2.0, Rs = 0.5, P = 2.0, e = 0.0, t0 = 0.0,
        incl = π / 2.0, Ω = 0.0, ω = 0.0,
    )
    @test orbit_standard_2 === orbit_kwarg_alias_2
end


text:
struct SupportSample
    s::Any
    x::Int
end

struct CategoricalWithSupport <: Distribution{SupportSample} end

"""
    categorical(probs::AbstractArray{U, 1}) where {U <: Real}

Given a vector of probabilities `probs` where `sum(probs) = 1`, and a vector of support, sample an element  `support[i]` from the support with probability `probs[i]`.
"""
const categorical_with_support = CategoricalWithSupport()

function logpdf(::CategoricalWithSupport, ss::SupportSample, probs::AbstractArray{U,1},support) where {U <: Real}
    log(probs[ss.x])
end

function logpdf_grad(::CategoricalWithSupport, ss::SupportSample, probs::AbstractArray{U,1},support)  where {U <: Real}
    grad = zeros(length(probs))
    grad[ss.x] = 1.0 / probs[ss.x]
    (nothing, grad)
end

function random(::CategoricalWithSupport, probs::AbstractArray{U,1},support) where {U <: Real}
    idx = rand(Distributions.Categorical(probs))

    SupportSample(support[idx],idx)
end

(::CategoricalWithSupport)(probs,support) = random(CategoricalWithSupport(), probs,support)

has_output_grad(::CategoricalWithSupport) = false
has_argument_grads(::CategoricalWithSupport) = (true,)

export categorical_with_support

max_stars_repo_name: rameshputalapattu/Gen

max_issues_repo_name: rameshputalapattu/Gen

max_forks_repo_name: rameshputalapattu/Gen


text:
@testset "Test resampling  methods" begin
    @testset "Test that upsample and downsample is reversible" begin
        for dim = 1:3
            for _ in 1:5
                s_small = ntuple(_ -> rand(1:13), dim)
                s_large = ntuple(i -> max.(s_small[i], rand(10:16)), dim)

                x = randn(Float32, (s_small))
                @test x == resample(x, s_small)
                @test Float32.(x) ≈ Float32.(resample(resample(x, s_large), s_small))
                @test x ≈ resample_by_FFT(resample_by_FFT(x, s_large), s_small)
                @test Float32.(x) ≈ Float32.(resample_by_RFFT(resample_by_RFFT(x, s_large), s_small))
                @test x ≈ FourierTools.resample_by_1D(FourierTools.resample_by_1D(x, s_large), s_small)
                x = randn(ComplexF32, (s_small))
                @test x ≈ resample(resample(x, s_large), s_small)
                @test x ≈ resample_by_FFT(resample_by_FFT(x, s_large), s_small)
                @test x ≈ resample_by_FFT(resample_by_FFT(real(x), s_large), s_small) + 1im .* resample_by_FFT(resample_by_FFT(imag(x), s_large), s_small)
                @test x ≈ FourierTools.resample_by_1D(FourierTools.resample_by_1D(x, s_large), s_small)
            end
        end

    end

    @testset "Test that different resample methods are consistent" begin
        for dim = 1:3
            for _ in 1:5
                s_small = ntuple(_ -> rand(1:13), dim)
                s_large = ntuple(i -> max.(s_small[i], rand(10:16)), dim)

                x = randn(Float32, (s_small))
                @test ≈(FourierTools.resample(x, s_large), FourierTools.resample_by_1D(x, s_large))
            end
        end
    end

    @testset "Test that complex and real routine produce same result for real array" begin
        for dim = 1:3
            for _ in 1:5
                s_small = ntuple(_ -> rand(1:13), dim)
                s_large = ntuple(i -> max.(s_small[i], rand(10:16)), dim)

                x = randn(Float32, (s_small))
                @test Float32.(resample(x, s_large)) ≈ Float32.(real(resample(ComplexF32.(x), s_large)))
                @test FourierTools.resample_by_1D(x, s_large) ≈ real(FourierTools.resample_by_1D(ComplexF32.(x), s_large))
            end
        end
    end

    @testset "Tests that resample_by_FFT is purely real" begin
        function test_real(s_1, s_2)
            x = randn(Float32, (s_1))
            y = resample_by_FFT(x, s_2)
            @test all(( imag.(y) .+ 1 .≈ 1))
            y = FourierTools.resample_by_1D(x, s_2)
            @test all(( imag.(y) .+ 1 .≈ 1))
        end

        for dim = 1:3
            for _ in 1:5
                s_1 = ntuple(_ -> rand(1:13), dim)
                s_2 = ntuple(i -> rand(1:13), dim)
                test_real(s_1, s_2)
            end
        end

        test_real((4, 4),(6, 6))
        test_real((4, 4),(6, 7))
        test_real((4, 4),(9, 9))
        test_real((4, 5),(9, 9))
        test_real((4, 5),(9, 8))
        test_real((8, 8),(6, 7))
        test_real((8, 8),(6, 5))
        test_real((8, 8),(4, 5))
        test_real((9, 9),(4, 5))
        test_real((9, 9),(4, 5))
        test_real((9, 9),(7, 8))
        test_real((9, 9),(6, 5))

    end

    @testset "Sinc interpolation based on FFT" begin

    function test_interpolation_sum_fft(N_low, N)
	    x_min = 0.0
	    x_max = 16π

	    xs_low = range(x_min, x_max, length=N_low+1)[1:N_low]
	    xs_high = range(x_min, x_max, length=N)[1:end-1]
	    f(x) = sin(0.5*x) + cos(x) + cos(2 * x) + sin(0.25*x)
	    arr_low = f.(xs_low)
	    arr_high = f.(xs_high)

	    xs_interp = range(x_min, x_max, length=N+1)[1:N]
	    arr_interp = resample(arr_low, N)
	    arr_interp2 = FourierTools.resample_by_1D(arr_low, N)

        @test ≈(arr_interp[2*N ÷10: N*8÷10], arr_high[2* N ÷10: N*8÷10], rtol=0.05)
        @test ≈(arr_interp2[2*N ÷10: N*8÷10], arr_high[2* N ÷10: N*8÷10], rtol=0.05)
    end

    test_interpolation_sum_fft(128, 1000)
    test_interpolation_sum_fft(129, 1000)
    test_interpolation_sum_fft(120, 1531)
    test_interpolation_sum_fft(121, 1211)
    end

    @testset "Upsample2 compared to resample" begin
    for sz in ((10,10),(5,8,9),(20,5,4))
        a = rand(sz...)
        @test ≈(upsample2(a),resample(a,sz.*2))
        @test ≈(upsample2_abs2(a),abs2.(resample(a,sz.*2)))
        a = rand(ComplexF32, sz...)
        @test ≈(upsample2(a),resample(a,sz.*2))
        @test ≈(upsample2_abs2(a),abs2.(resample(a,sz.*2)))
        s2 = (d == 2 ? sz[d]*2 : sz[d] for d in 1:length(sz))
        @test ≈(upsample2(a, dims=(2,)),resample(a,s2))
        @test ≈(upsample2_abs2(a, dims=(2,)),abs2.(resample(a,s2)))
    end
    end

    @testset "Downsampling based on frequency cutting" begin
    function test_resample(N_low, N)
	    x_min = 0.0
	    x_max = 16π

	    xs_low = range(x_min, x_max, length=N_low+1)[1:N_low]
	    f(x) = sin(0.5*x) + cos(x) + cos(2 * x) + sin(0.25*x)
	    arr_low = f.(xs_low)

	    xs_interp = range(x_min, x_max, length=N+1)[1:N]
	    arr_interp = resample(arr_low, N)

	    xs_interp_s = range(x_min, x_max, length=N+1)[1:N]

        arr_ds = resample(arr_interp, (N_low,) )
        @test ≈(arr_ds, arr_low)
        @test eltype(arr_low) === eltype(arr_ds)
        @test eltype(arr_interp) === eltype(arr_ds)
    end

    test_resample(128, 1000)
    test_resample(128, 1232)
    test_resample(128, 255)
    test_resample(253, 254)
    test_resample(253, 1001)
    test_resample(99, 100101)

    end

    @testset "FFT resample in 2D" begin

        function test_2D(in_s, out_s)
            x = range(-10.0, 10.0, length=in_s[1] + 1)[1:end-1]
            y = range(-10.0, 10.0, length=in_s[2] + 1)[1:end-1]'
    	    arr = abs.(x) .+ abs.(y) .+ sinc.(sqrt.(x .^2 .+ y .^2))
    	    arr_interp = resample(arr[1:end, 1:end], out_s);
    	    arr_ds = resample(arr_interp, in_s)
            @test arr_ds ≈ arr
        end

        test_2D((128, 128), (150, 150))
        test_2D((128, 128), (151, 151))
        test_2D((129, 129), (150, 150))
        test_2D((129, 129), (151, 151))

        test_2D((150, 128), (151, 150))
        test_2D((128, 128), (151, 153))
        test_2D((129, 128), (150, 153))
        test_2D((129, 128), (129, 153))

        x = range(-10.0, 10.0, length=129)[1:end-1]
        x2 = range(-10.0, 10.0, length=130)[1:end-1]
        x_exact = range(-10.0, 10.0, length=2049)[1:end-1]
        y = x'
        y2 = x2'
        y_exact = x_exact'
        arr = abs.(x) .+ abs.(y) .+sinc.(sqrt.(x .^2 .+ y .^2))
        arr2 = abs.(x) .+ abs.(y) .+sinc.(sqrt.(x .^2 .+ y .^2))
        arr_exact = abs.(x_exact) .+ abs.(y_exact) .+ sinc.(sqrt.(x_exact .^2 .+ y_exact .^2))
        arr_interp = resample(arr[1:end, 1:end], (131, 131));
        arr_interp2 = resample(arr[1:end, 1:end], (512, 512));
        arr_interp3 = resample(arr[1:end, 1:end], (1024, 1024));
        arr_ds = resample(arr_interp, (128, 128))
        arr_ds2 = resample(arr_interp, (128, 128))
        arr_ds23 = resample(arr_interp2, (512, 512))
        arr_ds3 = resample(arr_interp, (128, 128))

        @test ≈(arr_ds3, arr)
        @test ≈(arr_ds2, arr)
        @test ≈(arr_ds, arr)
        @test ≈(arr_ds23, arr_interp2)

    end

    @testset "FFT resample 2D for a complex signal" begin

        function test_2D(in_s, out_s)
        	x = range(-10.0, 10.0, length=in_s[1] + 1)[1:end-1]
        	y = range(-10.0, 10.0, length=in_s[2] + 1)[1:end-1]'
        	f(x, y) = 1im * (abs(x) + abs(y) + sinc(sqrt(x ^2 + y ^2)))
        	f2(x, y) =  abs(x) + abs(y) + sinc(sqrt((x - 5) ^2 + (y - 5)^2))

        	arr = f.(x, y) .+ f2.(x, y)
        	arr_interp = resample(arr[1:end, 1:end], out_s);
        	arr_ds = resample(arr_interp, in_s)

            @test eltype(arr) === eltype(arr_ds)
            @test eltype(arr_interp) === eltype(arr_ds)
            @test imag(arr) ≈ imag(arr_ds)
            @test real(arr) ≈ real(arr_ds)
        end

        test_2D((128, 128), (150, 150))
        test_2D((128, 128), (151, 151))
        test_2D((129, 129), (150, 150))
        test_2D((129, 129), (151, 151))

        test_2D((150, 128), (151, 150))
        test_2D((128, 128), (151, 153))
        test_2D((129, 128), (150, 153))
        test_2D((129, 128), (129, 153))
    end

    @testset "FFT resample in 2D for a purely imaginary signal" begin
        function test_2D(in_s, out_s)
        	x = range(-10.0, 10.0, length=in_s[1] + 1)[1:end-1]
        	y = range(-10.0, 10.0, length=in_s[2] + 1)[1:end-1]'
        	f(x, y) = 1im * (abs(x) + abs(y) + sinc(sqrt(x ^2 + y ^2)))

        	arr = f.(x, y)
        	arr_interp = resample(arr[1:end, 1:end], out_s);
        	arr_ds = resample(arr_interp, in_s)

            @test imag(arr) ≈ imag(arr_ds)
            @test all(real(arr_ds) .< 1e-13)
            @test all(real(arr_interp) .< 1e-13)
        end

        test_2D((128, 128), (150, 150))
        test_2D((128, 128), (151, 151))
        test_2D((129, 129), (150, 150))
        test_2D((129, 129), (151, 151))

        test_2D((150, 128), (151, 150))
        test_2D((128, 128), (151, 153))
        test_2D((129, 128), (150, 153))
        test_2D((129, 128), (129, 153))

    end

    @testset "test select_region_ft" begin
        x = [1,2,3,4]
        @test select_region_ft(ffts(x), (5,)) == ComplexF64[-1.0 + 0.0im, -2.0 - 2.0im, 10.0 + 0.0im, -2.0 + 2.0im, -1.0 + 0.0im]
        x = [3.1495759241275225 0.24720770605505335 -1.311507800204285 -0.3387627167144301; -0.7214121984874265 -0.02566249380406308 0.687066447881175 -0.09536748694092163; -0.577092696986848 -0.6320809680268722 -0.09460071173365793 0.7689715736798227; 0.4593837753047561 -1.0204193548690512 -0.28474772376166907 1.442443602597533]
        @test select_region_ft(ffts(x), (7, 7)) == ComplexF64[0.0 + 0.0im 0.0 + 0.0im 0.0 + 0.0im 0.0 + 0.0im 0.0 + 0.0im 0.0 + 0.0im 0.0 + 0.0im; 0.0 + 0.0im 0.32043577156395486 + 0.0im 2.321469443190397 + 0.7890379226962572im 0.38521287113798636 + 0.0im 2.321469443190397 - 0.7890379226962572im 0.32043577156395486 + 0.0im 0.0 + 0.0im; 0.0 + 0.0im 1.3691035744780353 + 0.16703621316206385im 2.4110077589815555 - 0.16558718095884828im 2.2813159163314163 - 0.7520360306228049im 7.47614366018844 - 4.139633109911205im 1.3691035744780353 + 0.16703621316206385im 0.0 + 0.0im; 0.0 + 0.0im 0.4801675770812479 + 0.0im 3.3142445917764407 - 3.2082400832669373im 1.6529948781166373 + 0.0im 3.3142445917764407 + 3.2082400832669373im 0.4801675770812479 + 0.0im 0.0 + 0.0im; 0.0 + 0.0im 1.3691035744780353 - 0.16703621316206385im 7.47614366018844 + 4.139633109911205im 2.2813159163314163 + 0.7520360306228049im 2.4110077589815555 + 0.16558718095884828im 1.3691035744780353 - 0.16703621316206385im 0.0 + 0.0im; 0.0 + 0.0im 0.32043577156395486 + 0.0im 2.321469443190397 + 0.7890379226962572im 0.38521287113798636 + 0.0im 2.321469443190397 - 0.7890379226962572im 0.32043577156395486 + 0.0im 0.0 + 0.0im; 0.0 + 0.0im 0.0 + 0.0im 0.0 + 0.0im 0.0 + 0.0im 0.0 + 0.0im 0.0 + 0.0im 0.0 + 0.0im]
    end

end


text:
"

 author: Atsushi Sakai
"

using Test
using JuMP
const __MAIN__ = length(PROGRAM_FILE)!=0 && occursin(PROGRAM_FILE, @__FILE__)

using Ipopt
m = Model(solver=IpoptSolver())

function main()
    println(PROGRAM_FILE," start!!")

    A = [1.0 0.0 -2.0;
         1.0 2.0 0.0;
         0.0 0.0 3.0]
    A = A*A'

    display(A)
    q = [3.0,4.0,-5.0]
    println("")
    display(q)
    println("")

    xmin = [-3.0, -5.0, -2.0]
    xmax = [1.0, 1.0, 1.0]

    @variable(m, x[1:3])

    @objective(m, Min, 0.5*x'*A*x+q'*x)

    @constraint(m, xmin .<= x .<= xmax)

    print(m)
    status = solve(m)

    println("Optimal Solutions:")
    println("x = ", getvalue(x))

    println(PROGRAM_FILE," Done!!")
end

if __MAIN__
    @time main()
end


text:
# # RecipesPipeline
module RecipesPipeline

import RecipesBase
import RecipesBase: @recipe, @series, RecipeData, is_explicit
import PlotUtils # tryrange and adapted_grid
using Dates
using NaNMath

export recipe_pipeline!
# Plots relies on these:
export SliceIt,
    DefaultsDict,
    Formatted,
    AbstractSurface,
    Surface,
    Volume,
    is3d,
    is_surface,
    needs_3d_axes,
    group_as_matrix,
    reset_kw!,
    pop_kw!,
    scale_func,
    inverse_scale_func,
    unzip,
    dateformatter,
    datetimeformatter,
    timeformatter
# API
export warn_on_recipe_aliases,
    splittable_attribute,
    split_attribute,
    process_userrecipe!,
    get_axis_limits,
    is_axis_attribute,
    type_alias,
    plot_setup!,
    slice_series_attributes!,
    process_sliced_series_attributes!

include("utils.jl")
include("api.jl")
include("series.jl")
include("group.jl")
include("user_recipe.jl")
include("type_recipe.jl")
include("plot_recipe.jl")
include("series_recipe.jl")
include("recipes.jl")

"""
    recipe_pipeline!(plt, plotattributes, args)

Recursively apply user recipes, type recipes, plot recipes and series recipes to build a
list of `Dict`s, each corresponding to a series. At the beginning, `plotattributes`
contains only the keyword arguments passed in by the user. Then, add all series to the plot
object `plt` and return it.
"""
function recipe_pipeline!(plt, plotattributes, args)
    @nospecialize
    plotattributes[:plot_object] = plt

    # --------------------------------
    # "USER RECIPES"
    # --------------------------------

    # process user and type recipes
    kw_list = _process_userrecipes!(plt, plotattributes, args)

    # --------------------------------
    # "PLOT RECIPES"
    # --------------------------------

    # The "Plot recipe" acts like a series type, and is processed before the plot layout
    # is created, which allows for setting layouts and other plot-wide attributes.
    # We get inputs which have been fully processed by "user recipes" and "type recipes",
    # so we can expect standard vectors, surfaces, etc.  No defaults have been set yet.

    kw_list = _process_plotrecipes!(plt, kw_list)

    # --------------------------------
    # Plot/Subplot/Layout setup
    # --------------------------------

    plot_setup!(plt, plotattributes, kw_list)

    # At this point, `kw_list` is fully decomposed into individual series... one KW per
    # series. The next step is to recursively apply series recipes until the backend
    # supports that series type.

    # --------------------------------
    # "SERIES RECIPES"
    # --------------------------------

    _process_seriesrecipes!(plt, kw_list)

    # --------------------------------
    # Return processed plot object
    # --------------------------------

    return plt
end

include("precompile_includer.jl")

end


text:
## Variation on Lazy macros to use _ syntax

_pipe(f, d::AbstractDataset, by; kwargs...) =  IndexedTables.groupby(f, d, by; kwargs...)
_pipe(f, d::AbstractDataset; kwargs...) = f(d)
_pipe(f, d::Columns, args...; kwargs...) = _pipe(f, _table(d), args...; kwargs...)
_pipe(f, args...; kwargs...) = d::Union{AbstractDataset, Columns} -> _pipe(f, d, args...; kwargs...)

_pipe_chunks(f, d::Dataset) = f(d)
_pipe_chunks(f, d::DDataset) = fromchunks(delayedmap(f, d.chunks))
_pipe_chunks(f, d::Columns) = _pipe_chunks(f, _table(d))
_pipe_chunks(f) = d::Union{AbstractDataset, Columns}  -> _pipe_chunks(f, d)

"""
`@apply(args...)`

Concatenate a series of operations. Non-macro operations from JuliaDB, are supported via
 the `_` curryfication syntax. A second optional argument is used for grouping:

```jldoctest apply
julia> t = table([1,2,1,2], [4,5,6,7], [0.1, 0.2, 0.3,0.4], names = [:x, :y, :z]);

julia> @apply t begin
          @where :x >= 2
          @transform {:x+:y}
          sort(_, :z)
       end
Table with 2 rows, 4 columns:
x  y  z    x + y
────────────────
2  5  0.2  7
2  7  0.4  9

julia> @apply t :x flatten=true begin
          @transform {w = :y + 1}
          sort(_, :w)
       end
Table with 4 rows, 4 columns:
x  y  z    w
────────────
1  4  0.1  5
1  6  0.3  7
2  5  0.2  6
2  7  0.4  8
```
"""
macro apply(args...)
    esc(Expr(:call, :(JuliaDBMeta._pipe), thread(args[end]), replace_keywords(args[1:end-1])...))
end

"""
`@applychunked(args...)`

Split the table into chunks, apply the processing pipeline separately to each chunk and return the result
as a distributed table.

```jldoctest applychunked
julia> t = table([1,2,1,2], [4,5,6,7], [0.1, 0.2, 0.3,0.4], names = [:x, :y, :z], chunks = 2);

julia> @applychunked t begin
          @where :x >= 2
          @transform {:x+:y}
          sort(_, :z)
       end
Distributed Table with 2 rows in 2 chunks:
x  y  z    x + y
────────────────
2  5  0.2  7
2  7  0.4  9
```
"""
macro applychunked(args...)
    esc(Expr(:call, :(JuliaDBMeta._pipe_chunks), thread(args[end]), args[1:end-1]...))
end

function thread(ex)
     if isexpr(ex, :block)
          thread(rmlines(ex).args...)
     elseif isa(ex, Expr)
          us = findall(ex.args .== :(_))
          i = gensym()
          ex.args[us] .= i
          isempty(us) ? ex : Expr(:(->), i, ex)
     else
          ex
     end
end

thread(exs...) = mapreduce(thread, (ex1, ex2) -> Expr(:call, Symbol(∘), ex2, ex1), exs)


text:
# this script does plotting for our paper
# ------------------------
# Including (in order):
# 1. MODEL EVALUATION: fertility experiments
#    1. UE-BMI under baseline scenario
#    2. Baseline compared with accounting results
#    3. Fertility experiments and UE-BMI
#    4. Summary table for benchmark
#    5. correlations (during plotting years) between econ vars & aging population ratio
# 2. MODEL EVALUATION: contribution rate increases
#    1. Rising of firm’s contribution rate to UE-BMI
#    2. ζ elasticity of four economic variables
# 3. MODEL EVALUATION: contract reforms on individual accounts
#    1. Reforms of individual accounts of UE-BMI
#  -----------------------
# NOTE: all figures are output to "pwd()/output/"

include("$(pwd())/main.jl") # if not run main.jl yet, run it first

# -----------------------------------
# SECTION introduction
# -----------------------------------
# NOTE: set pwd() as the root directory of the project to correctly cite paths!
   # --------- for plotting
   import EasyPlot
   using PyPlot
   # --------- for analysis & I/O
   import Statistics
   import DataFrames, CSV

# -----------------------------------
# SECTION: baseline & fertility experiments plotting
# -----------------------------------
# 0. define expr, collection; then run Fertiltiy Experiments
include("$(pwd())/scripts/subsections_PlotForPaper/sub_01_Def2RunFertilityExperiments.jl")

# 1. baseline plotting (no channel analysis)
   # Figure: Baseline compared with accounting results (BenchmarkCpAccount.pdf)
   # Figure: UE-BMI under baseline scenario (BenchProfile.pdf)
include("$(pwd())/scripts/subsections_PlotForPaper/sub_02_Plot4Baseline.jl")

# 2. fertility experiments plotting (no channel analysis)
   # NOTE: including def, expr, eval
   # Figure: Fertility shocks and UE-BMI (FertilityShockProfile.pdf)
   # Table: Summary of baseline simulation (SummaryBench.csv)
include("$(pwd())/scripts/subsections_PlotForPaper/sub_03_Plot4Baseline&FertilityExperiments.jl")

# 3. compute cor( econ var, aging population ratio )
   # No output file
include("$(pwd())/scripts/subsections_PlotForPaper/sub_04_Plot4BaselineCor.jl")

# 4. channel analysis in steady states
   # NOTE: based on the decomposition of gap
   # NOTE: including def, expr, eval
   # Figure: Effect decomposition in steady state 2010 parameters (Channel_DecomposeInSState.pdf)
include("$(pwd())/scripts/subsections_PlotForPaper/sub_05_Plot4ChannelAnalysis.jl")

# -----------------------------------
# SECTION: contribution rates increase
# -----------------------------------
# 0. define expr, collection; then get the econ vars under lattice of zeta values
include("$(pwd())/scripts/subsections_PlotForPaper/sub_06_Def&RunContributionRateRise.jl")

# 1. the four economic vars under different zeta values, plotting
   # Figure: Rising of firm’s contribution rate to UE-BMI (ContributionRateRise.pdf)
include("$(pwd())/scripts/subsections_PlotForPaper/sub_07_Plot4ContributionRateRise.jl")

# 2. define expr, collection; then compute the zeta elasticities of diff econ vars
   # NOTE: requires "sub_06_Def&RunContributionRateRise.jl" run
include("$(pwd())/scripts/subsections_PlotForPaper/sub_08_Def&RunElasticityOfZeta.jl")

# 3. elasticity analysis on zeta, plotting
   # Figure: ζ elasticity of four economic variables (ContributionRateRise_Elasticity.pdf)
include("$(pwd())/scripts/subsections_PlotForPaper/sub_09_Plot4ZetaElasticity.jl")

# -----------------------------------
# SECTION: UE-BMI contract reforms on individual accounts
# -----------------------------------
# 0. define expr, collection; then get the econ vars under diff scenarios
include("$(pwd())/scripts/subsections_PlotForPaper/sub_10_Def&RunIndiAcctReform.jl")

# 1. the four economic vars under different scenarios, plotting
   # Figure: Reforms of individual accounts of UE-BMI (IndiAcctReform.pdf)
include("$(pwd())/scripts/subsections_PlotForPaper/sub_11_Plot4IndiAcctReform.jl")

# 2. crowding out caused by cancelling individual accounts, plotting
   # Figure: Crowding out effects (CrowdingOutEffect.pdf)
include("$(pwd())/scripts/subsections_PlotForPaper/sub_12_Plot4CrowdingOutEffect.jl")

#


text:
using LinearAlgebra

"""
    Plane(c::Vector{Float64}, a::Vector{Float64})

Construct a `Plane` from a point `c` which lies in the plane and a unit normal `a`.
"""
struct Plane
    c::Vector{Float64}
    a::Vector{Float64}
    Plane(c,a) = begin
        if length(c) != 3 || length(a) != 3
            error("3-vectors please")
        else
            new(c,a/norm(a))
        end
    end
end

"""
    Cylinder(R::Float64, c::Vector{Float64}, a::Vector{Float64})

Construct a `Cylinder` with radius `R` and unit axis `a` which passes through point `c`.
"""
struct Cylinder
    R::Real          # radius
    c::Vector{Float64}  # center point
    a::Vector{Float64}  # axis (unit) (NOTE: c + h*a yields a point on the opposite cap => this is an inward normal axis)
    Cylinder(R,c,a) = begin
        if length(c) != 3 || length(a) != 3
            error("3-vectors please")
        else
            new(R,c,a/norm(a))
        end
    end
end

"""
    Paraboloid(b::Float64, c::Vector{Float64}, a::Vector{Float64})

Construct a `Paraboloid` at vertex `c` with growth rate `b` and unit axis `a`.
"""
struct Paraboloid
    b::Real
    c::Vector{Float64}
    a::Vector{Float64}
    Paraboloid(b,c,a) = begin
        if length(c) != 3 || length(a) != 3
            error("3-vectors please")
        else
            new(b,c,a/norm(a))
        end
    end
end

"""
    Hyperboloid(R::Float64, b::Float64, c::Vector{Float64}, a::Vector{Float64})

Construct a `Hyperboloid` of one sheet centered at `c` with minimum radius `R`, radial growth rate `b`, and unit axis `a`.
"""
struct Hyperboloid
    R::Real
    b::Real
    c::Vector{Float64}
    a::Vector{Float64}
    Hyperboloid(R,b,c,a) = begin
        if length(c) != 3 || length(a) != 3
            error("3-vectors please")
        else
            new(R,b,c,a/norm(a))
        end
    end
end

"""
    Quadric(Q::Matrix{Float64})

Construct a `Quadric` surface from a 4x4 symmetric quadric matrix.
"""
struct Quadric
    Q::Matrix{Float64}
    Quadric(Q) = begin
        if issymmetric(Q)
            new(Q)
        else
            error("quadric matrix must be symmetric")
        end
    end
end

"""
    TruncatedQuadric(q::Quadric, p::Vector{Plane}, caps)

Construct a `TruncatedQuadric` surface by slicing a non-reducible `Quadric` q with two planes `p`. Planes may not contain quadric axis. Argument `caps=true` closes the volume between the planes, `caps=false` defines an open-ended tube bounded by the `Quadric` and the planes.
"""
struct TruncatedQuadric
    q::Quadric
    p::Vector{Plane}
    caps::Bool
    TruncatedQuadric(q,p,caps) = begin
        Qr = q.Q[1:3,1:3]
        if all(eigvals(Qr) .!= 0)
            for plane in p
                if all(q.Q[1:3,1:3]*plane.a .== 0)
                    error("planes must be oblique to quadric axis")
                end
            end
        end
        new(q,p,caps)
    end
end

@doc raw"""
    Quadric(s::Plane)

Construct a `Quadric` from a `Plane` (with attributes axis = `a` = ``\boldsymbol{a}``; and point in plane = `c` = ``\boldsymbol{c}``) using the definition:

``\boldsymbol{Q}_s = \begin{bmatrix} \boldsymbol{0} & \frac{1}{2}\boldsymbol{a} \\ \frac{1}{2}\boldsymbol{a}^\mathsf{T} & -\boldsymbol{a}^\mathsf{T}\boldsymbol{c} \end{bmatrix}``
"""
function Quadric(s::Plane)
    Q = [zeros(3,3)      0.5*s.a;
         0.5*s.a'        -s.a'*s.c]
    return Quadric(Q)
end

@doc raw"""
    Quadric(s::Cylinder)

Construct a `Quadric` from a `Cylinder` (which has attributes axis = `a` = ``\boldsymbol{a}``; point on axis = `c` = ``\boldsymbol{c}``; and radius = `R` = ``R``) using the definitions:

``\boldsymbol{Q}_r = \boldsymbol{a}\boldsymbol{a}^\mathsf{T} - \boldsymbol{1}``

``\boldsymbol{Q}_c = \begin{bmatrix} \boldsymbol{Q}_r & -\boldsymbol{Q}_r\boldsymbol{c} \\ -\boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r & R^2 + \boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r\boldsymbol{c} \end{bmatrix}``
"""
function Quadric(s::Cylinder)
    Q = [s.a*s.a' - I            (I - s.a*s.a')*s.c;
         ((I - s.a*s.a')*s.c)'   s.R^2 + s.c'*(s.a*s.a' - I)*s.c]
    return Quadric(Q)
end

@doc raw"""
    Quadric(s::Paraboloid)

Construct a `Quadric` from a `Paraboloid` (which has attributes axis = `a` = ``\boldsymbol{a}``; point on axis = `c` = ``\boldsymbol{c}``; and quadratic parameter = `b` = ``b``) using the definitions:

``\boldsymbol{Q}_r = \boldsymbol{a}\boldsymbol{a}^\mathsf{T} - \boldsymbol{1}``

``\boldsymbol{Q}_p = \begin{bmatrix} \boldsymbol{Q}_r & \frac{b^2}{2}\boldsymbol{a} - \boldsymbol{Q}_r\boldsymbol{c} \\ \frac{b^2}{2}\boldsymbol{a}^\mathsf{T} - \boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r & -b^2\boldsymbol{a}^\mathsf{T}\boldsymbol{c} + \boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r\boldsymbol{c} \end{bmatrix}``
"""
function Quadric(s::Paraboloid)
    Q = [s.a*s.a' - I   s.b^2/2*s.a + (I - s.a*s.a')*s.c;
         (s.b^2/2*s.a + (I - s.a*s.a')*s.c)'    -s.b^2*s.a'*s.c + s.c'*(s.a*s.a' - I)*s.c]
    return Quadric(Q)
end

@doc raw"""
    Quadric(s::Hyperboloid)

Construct a `Quadric` from a `Hyperboloid` (which has attributes axis = `a` = ``\boldsymbol{a}``; point on axis = `c` = ``\boldsymbol{c}``; neck radius = `R` = ``R``; and slope parameter = `b` = ``b``) using the definitions:

``\gamma = 1 + \frac{R^2}{b^2}``

``\boldsymbol{Q}_r = \gamma\boldsymbol{a}\boldsymbol{a}^\mathsf{T} - \boldsymbol{1}``

``\boldsymbol{Q}_h = \begin{bmatrix} \boldsymbol{Q}_r & -\boldsymbol{Q}_r\boldsymbol{c} \\ -\boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r & R^2 + \boldsymbol{c}^\mathsf{T}\boldsymbol{Q}_r\boldsymbol{c} \end{bmatrix}``
"""
function Quadric(s::Hyperboloid)
    γ = (1 + s.R^2/s.b^2)
    Q = [γ*s.a*s.a' - I     (I - γ*s.a*s.a')*s.c;
         ((I - γ*s.a*s.a')*s.c)'  s.R^2 + s.c'*(γ*s.a*s.a' - I)*s.c]
    return Quadric(Q)
end

"""
    TruncatedQuadric(s::Cylinder, p::Vector{Plane}, caps=true)

Construct a `TruncatedQuadric` surface from a `Cylinder` and set of `Plane`s by converting `Cylinder` directly to a `Quadric`.
"""
TruncatedQuadric(s::Cylinder, p::Vector{Plane}, caps=true) = TruncatedQuadric(Quadric(s), p, caps)

"""
    TruncatedQuadric(s::Paraboloid, p::Vector{Plane}, caps=true)

Construct a `TruncatedQuadric` surface from a `Paraboloid` and set of `Plane`s by converting `Paraboloid` directly to a `Quadric`.
"""
TruncatedQuadric(s::Paraboloid, p::Vector{Plane}, caps=true) = TruncatedQuadric(Quadric(s), p, caps)

"""
    TruncatedQuadric(s::Hyperboloid, p::Vector{Plane}, caps=true)

Construct a `TruncatedQuadric` surface from a `Hyperboloid` and set of `Plane`s by converting `Hyperboloid` directly to a `Quadric`.
"""
TruncatedQuadric(s::Hyperboloid, p::Vector{Plane}, caps=true) = TruncatedQuadric(Quadric(s), p, caps)

"""
    TruncatedQuadric(q::Quadric, h1::Vector{Float64}, h2::Vector{Float64}, caps=true)

A convenience constructor for a `TruncatedQuadric`. Specify `h1` and `h2` as points along the axis of `q`, and the constructor will return `q` truncated by planes passing through `h1` and `h2` sharing a normal vector equal to the axis of `q`.
"""
function TruncatedQuadric(q::Quadric, h1::Vector{Float64}, h2::Vector{Float64}, caps=true)
    s = changerepresentation(q)
    if typeof(s) == Plane
        display("TruncatedQuadric just a bunch of planes")
    end

    if abs(((h1 - h2)/norm(h1 - h2))'*s.a/norm(s.a)) == 1
        p1 = Plane(h1, s.a)
        p2 = Plane(h2, s.a)
    else
        error("h1, h2 must lie on Quadric axis")
    end

    return TruncatedQuadric(q, [p1, p2], caps)
end

@doc raw"""
    changerepresentation(q::Quadric)

Converts a `Quadric` to a `Plane`, `Cylinder`, `Paraboloid`, or `Hyperboloid` based on its matrix structure.

!!! note
    This function is not fast. Minimize conversion between explicit surface and quadric matrix in heavy-lifting code.
"""
function changerepresentation(q::Quadric)
    ε = 1e-15

    Qr = q.Q[1:3, 1:3]
    qd = q.Q[1:3, end]
    q0 = q.Q[end, end]
    E = eigvals(Qr)
    if all(E .== 0)
        # plane
        a = 2*qd                    # Qr is zero, qd is parallel to plane normal
        c = -q0*a                   # q0 = -a'*c: choose c to lie on a (free choice)
        return Plane(c, a)

    elseif all(abs.(E) .>= ε)
        # hyperboloid
        v = eigvecs(Qr)
        a = v[:,end]                # axis is eigenvector for largest eigenvalue
        c = (-Qr)\qd                # center of hyperboloid
        γ = tr(Qr + I)              # Qr = γ*a*a' - I, and trace(a*a') == 1
        R = sqrt(q0 - c'*Qr*c)      # q0 = sum(c*c' .* Qr) + R^2
        b = 1/sqrt(γ - 1)*R         # γ = 1 + (R/b)^2
        return Hyperboloid(R, b, c, a)

    else
        v = eigvecs(Qr)
        a = v[:,end]        # axis is eigenvector for largest eigenvalue

        if rank(q.Q) == 4
            # paraboloid
            a = sign(qd'*a)*a       # correct for antiparallel axis
            b = sqrt(2*qd'*a)

            cb1 = v[:,1]'*qd
            cb2 = v[:,2]'*qd
            cb3 = -(q0 + cb1^2 + cb2^2)/b^2

            c = v*[cb1; cb2; cb3]

            return Paraboloid(b, c, a)

        elseif rank(q.Q) == 3
            # cylinder
            c = (-Qr)\qd
            badI = findfirst(isnan.(c) .| isinf.(c))
            if !isnothing(badI)
                c[badI] = 1.0
            end
            R = sqrt(q0 - c'*Qr*c)
            return Cylinder(R, c, a)

        else
            error("unclassified shape")

        end
    end
end

"""
    changerepresentation(s::Plane)

Convert a `Plane` to `Quadric` representation via the `Quadric(s::Plane)` constructor.
"""
changerepresentation(s::Plane) = Quadric(s)

"""
    changerepresentation(s::Cylinder)

Convert a `Cylinder` to `Quadric` representation via the `Quadric(s::Cylinder)` constructor.
"""
changerepresentation(s::Cylinder) = Quadric(s)

"""
    changerepresentation(s::Paraboloid)

Convert a `Paraboloid` to `Quadric` representation via the `Quadric(s::Paraboloid)` constructor.
"""
changerepresentation(s::Paraboloid) = Quadric(s)

"""
    changerepresentation(s::Hyperboloid)

Convert a `Hyperboloid` to `Quadric` representation via the `Quadric(s::Hyperboloid)` constructor.
"""
changerepresentation(s::Hyperboloid) = Quadric(s)

"""
    normal(q::Quadric, r::Vector{Float64})

Returns unit normal vector to a `Quadric` at position `r` on the surface.

!!! warning
    This will return a normal vector even if `r` does not fall on the surface. It is up to *you* to input an `r` you know to lie on the surface of `q`.
"""
function normal(q::Quadric, r::Vector{Float64})
    n = 2*q.Q*[r; 1]
    n = n[1:3]
    n = n/norm(n)
    return n
end

"""
    normal(s::Plane, r::Vector{Float64})

Returns unit normal vector to a `Plane`. This will return the `Plane`'s normal vector regardless of the `r` supplied; it takes the `r` argument nonetheless to maintain the same pattern as `normal()` functions for nonreducible `Quadric`s.
"""
function normal(s::Plane, r::Vector{Float64})
    return s.a
end

"""
    normal(s::Cylinder, r::Vector{Float64})

Returns unit normal vector to a `Cylinder` at position `r` on the surface.

!!! warning
    This will return a normal vector even if `r` does not fall on the surface. It is up to *you* to input an `r` you know to lie on the surface of `s`.
"""
function normal(s::Cylinder, r::Vector{Float64})
    n = 2*(s.a - (r - s.c)./norm(r - s.c))
    return n/norm(n)
end

"""
    normal(s::Paraboloid, r::Vector{Float64})

Returns unit normal vector to a `Paraboloid` at position `r` on the surface.

!!! warning
    This will return a normal vector even if `r` does not fall on the surface. It is up to *you* to input an `r` you know to lie on the surface of `s`.
"""
function normal(s::Paraboloid, r::Vector{Float64})
    n = s.b^2*s.a + 2*(s.a - (r - s.c)./norm(r - s.c))
    return n/norm(n)
end

"""
    normal(s::Hyperboloid, r::Vector{Float64})

Returns normal vector to a `Hyperboloid` at position `r` on the surface.

!!! warning
    This will return a normal vector even if `r` does not fall on the surface. It is up to *you* to input an `r` you know to lie on the surface of `q`.
"""
function normal(s::Hyperboloid, r::Vector{Float64})
    n = 2*((1 + s.R^2/s.b^2)*s.a - (r - s.c)./norm(r - s.c))
    return n/norm(n)
end

function classify(q::Quadric)

    Qr = q.Q[1:3,1:3]

    E = eigvals(Qr)

    p = sum(E .> 0)
    n = sum(E .< 0)
    z = sum(E .== 0)

    # options:
    #   [p, n, z]
    #   [0, 0, 3]:  plane
    #   [0, 1, 2]:  parabolic cylinder, parallel planes
    #   [0, 2, 1]:  imaginary cylinder
    #   [0, 3, 0]:  imag. ellipsoid
    #   [1, 0, 2]:  parabolic cylinder, parallel planes
    #   [1, 1, 1]:  hyperbolic paraboloid, hyperbolic cylinder
    #   [1, 2, 0]:  hyperboloid 2 sheet
    #   [2, 0, 1]:  elliptic paraboloid, elliptic cylinder
    #   [2, 1, 0]:  hyperboloid 1 sheet, cone @ origin
    #   [3, 0, 0]:  ellipsoid, imag. context
    shapes0 = [
        "none"  "none"  "none"  "ellipsoid";
        "none"  "none"  "hyperboloid"   "none";
        "none"  "hyperboloid"  "none"   "none";
        "ellipsoid" "none"  "none"  "none"
    ]
    shapes1 = [
        "none"  "none"  "elliptic cylinder/paraboloid"    "none";
        "none"  "hyperbolic paraboloid/cylinder"  "none"  "none";
        "elliptic cylinder/paraboloid" "none"   "none"  "none";
        "none"  "none"  "none"  "none"
    ]
    shapes2 = [
        "none"  "parabolic cylinder/planes" "none"  "none";
        "parabolic cylinder/planes" "none"  "none"  "none";
        "none"  "none"  "none"  "none";
        "none"  "none"  "none"  "none"
    ]
    shapes3 = [
        "plane"  "none"  "none"  "none";
        "none"   "none"  "none"  "none";
        "none"   "none"  "none"  "none";
        "none"   "none"  "none"  "none";
    ]

    shapes = cat(shapes0, shapes1, shapes2, shapes3, dims=3)

    shapename = shapes[p+1, n+1, z+1]
    return shapename
end


text:
using Plots
using LinearAlgebra

function angular_position(points)
    cpoints = points[1] + im * points[2]
    return rad2deg(angle(cpoints))
end

function perturbed_circle(point, meanradius, amplitude, frequency)
    angle = angular_position(point)
    interfaceposition = meanradius  + amplitude * sind(angle * frequency)
    distance = interfaceposition - norm(point)
    return distance
end

xrange = -1:1e-2:1
meanradius = 0.5
amplitude = 0.1
frequency = 5
contour(
    xrange,
    xrange,
    (x, y) ->
        perturbed_circle([x, y], meanradius, amplitude, frequency),
        levels=[0.0],
        aspect_ratio=:equal
)


text:
using LoopVectorization, SnoopCompile
tinf = @snoopi_deep include(joinpath(pkgdir(LoopVectorization), "test", "runtests.jl"))

ttot, pcs = SnoopCompile.parcel(tinf);

lv, pcslv = last(pcs);
@assert lv === LoopVectorization

blacklist = (
  :_turbo_!,
  :vmaterialize!,
  :vmaterialize,
  Symbol("vreduce##kw"),
  :_vreduce_dims!,
  :vreduce,
  :vmapreduce,
  :SIMDMapBack,
  :launch_thread_vmap!,
  :_turbo_loopset_debug,
  :all_dense,
  :sigmoid_fast,
  :rrule,
  :add_broadcast!,
  :create_mrefs!,
  :avx_config_val,
  :subsetview,
  :ifelsepartial,
  :tanh_fast,
  :check_args,
  :relu,
  :init_dual,
)
filteredmethods = filter(m -> !Base.sym_in(m[2].def.name, blacklist), last(pcslv));
length(filteredmethods);

SnoopCompile.write(
  "/tmp/precompile_loopvec",
  [LoopVectorization => (sum(first, filteredmethods), filteredmethods)],
)

# pc = SnoopCompile.parcel(tinf; blacklist=["vmaterialize", "vmaterialize!", "vreduce", "Base.Broadcast.materialize", "_vreduce_dims!", "vmapreduce"])
# pcs = pc[:LoopVectorization]
# open(joinpath(pkgdir, "src", "precompile.jl"), "w") do io
#     println(io, """
#     function _precompile_()
#         ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
#     """)
#     for stmt in sort(pcs)
#         println(io, "    ", stmt)
#     end
#     println(io, "end")
# end


text:
#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

################################## INSTALLATION ####################################################
# (1) Enter Package Manager via     ]
# (2) Install FMI via               add FMI       or   add "https://github.com/ThummeTo/FMI.jl"
# (3) Install FMIFlux via           add FMIFlux   or   add "https://github.com/ThummeTo/FMIFlux.jl"
################################ END INSTALLATION ##################################################

# this example covers creation and training am ME-NeuralFMUs

using FMI
using FMIFlux
using Flux
using DifferentialEquations: Tsit5
import Plots

modelFMUPath = joinpath(dirname(@__FILE__), "../model/SpringPendulum1D.fmu")
realFMUPath = joinpath(dirname(@__FILE__), "../model/SpringFrictionPendulum1D.fmu")

t_start = 0.0
t_step = 0.01
t_stop = 5.0
tData = collect(t_start:t_step:t_stop)

myFMU = fmiLoad(realFMUPath)
fmiInstantiate!(myFMU; loggingOn=false)
fmiSetupExperiment(myFMU, t_start, t_stop)

fmiEnterInitializationMode(myFMU)
fmiExitInitializationMode(myFMU)

x0 = fmiGetContinuousStates(myFMU)

realSimData = fmiSimulate(myFMU, t_start, t_stop; recordValues=["mass.s", "mass.v", "mass.f", "mass.a"], saveat=tData, setup=false)
fmiUnload(myFMU)

fmiPlot(realSimData)

myFMU = fmiLoad(modelFMUPath)

fmiInstantiate!(myFMU; loggingOn=false)
fmuSimData = fmiSimulate(myFMU, t_start, t_stop; recordValues=["mass.s", "mass.v", "mass.a"], saveat=tData)

posData = fmi2SimulationResultGetValues(realSimData, "mass.s")
velData = fmi2SimulationResultGetValues(realSimData, "mass.v")

# loss function for training
function losssum()
    solution = problem(x0, t_start)

    tNet = collect(data[1] for data in solution.u)
    posNet = collect(data[2] for data in solution.u)
    #velNet = collect(data[3] for data in solution.u)

    Flux.Losses.mse(posData, posNet) #+ Flux.Losses.mse(velData, velNet)
end

# callback function for training
global iterCB = 0
function callb()
    global iterCB += 1

    if iterCB % 10 == 1
        avg_ls = losssum()
        @info "Loss: $(round(avg_ls, digits=5))   Avg displacement in data: $(round(sqrt(avg_ls), digits=5))"
    end
end

# NeuralFMU setup
numStates = fmiGetNumberOfStates(myFMU)

net = Chain(inputs -> fmiDoStepME(myFMU, inputs),
            Dense(numStates, 16, tanh),
            Dense(16, 16, tanh),
            Dense(16, numStates))

problem = ME_NeuralFMU(myFMU, net, (t_start, t_stop), Tsit5(); saveat=tData)
solutionBefore = problem(x0, t_start)
fmiPlot(problem)

# train it ...
p_net = Flux.params(problem)

optim = ADAM()
Flux.train!(losssum, p_net, Iterators.repeated((), 1000), optim; cb=callb) # Feel free to increase training steps or epochs for better results

###### plot results mass.s
solutionAfter = problem(x0, t_start)
fig = Plots.plot(xlabel="t [s]", ylabel="mass position [m]", linewidth=2,
    xtickfontsize=12, ytickfontsize=12,
    xguidefontsize=12, yguidefontsize=12,
    legendfontsize=12, legend=:bottomright)
Plots.plot!(fig, tData, fmi2SimulationResultGetValues(fmuSimData, "mass.s"), label="FMU", linewidth=2)
Plots.plot!(fig, tData, posData, label="reference", linewidth=2)
Plots.plot!(fig, tData, collect(data[2] for data in solutionAfter.u), label="NeuralFMU", linewidth=2)
Plots.savefig(fig, "exampleResult_s.pdf")

fmiUnload(myFMU)


text:
## This file has methods to stably divide the jacobi singularity

function dirichlet_divide_singularity{T<:Number}(b::Bool,v::Vector{T})
    n=length(v)
    w=zeros(T,n-1)
    s=b?1:-1

    if n==1
        return w
    elseif n==2
        w[1]=-s*v[2]
        return w
    end

    w[n-1]=-2s*v[n]

    for k=n-2:-1:2
        @inbounds w[k]=-2s*(v[k+1] - .5w[k+1])
    end

    @inbounds w[1]=-s*(v[2] - .5w[2])

    w
end

# dirichletrange_divide_singularity divides by (1-x^2)
# where the coefficients are ChebyshevDirichlet{1,1}
# assume that v[1] == v[2] ==0.
# the output coefficients are Chebyshev

function dirichlet_divide_singularity{T<:Number}(v::Vector{T})
    n=length(v)
    if n ≤ 2    # assumes v[1]==v[2]==0 which is deleted
        return zeros(T,1)
    end

    w=zeros(T,n-2)

    w[n-2]=-4v[n]
    if n>3
        w[n-3]=-4v[n-1]
    end

    for k=n-4:-1:2
        @inbounds w[k]=-4*(v[k+2] - .25w[k+2])
    end

    w[1]=-2v[3]
    if n≥5
        w[1]+=0.5w[3]
    end

    w
end

divide_singularity(s::Bool,f::Fun)=Fun(dirichlet_divide_singularity(s,Fun(f,s?ChebyshevDirichlet{0,1}:ChebyshevDirichlet{1,0}).coefficients),
                                                                    Chebyshev(domain(f)))

divide_singularity(f::Fun)=Fun(dirichlet_divide_singularity(Fun(f,ChebyshevDirichlet{1,1}).coefficients),Chebyshev(domain(f)))

function divide_singularity(s::Tuple{Int,Int},f::Fun)
    if s[1]>0 && s[2]>0
        divide_singularity((s[1]-1,s[2]-1),divide_singularity(f))
    elseif s[1]>0
        divide_singularity((s[1]-1,s[2]),divide_singularity(false,f))
    elseif s[2]>0
        divide_singularity((s[1],s[2]-1),divide_singularity(true,f))
    else
        @assert s[1]==s[2]==0
        f
    end
end


text:
# imports

using JLD

# %% Create the network

include("../../Modules/Networks/net_res.jl")

n = Network([1200,64,64,2])

# %% load display_spikes

display_spikes = load("Julia/CartPole/Heatmaps/disp_spikesDVSv2.jld", "display_spikes")

# %% main sim

spike_hist = Array{Int64}[]

@inbounds for sec in 0:9
    @time @inbounds for msec in 1:1000
        time = 1000*sec+msec
        train = false
        if msec%10 == 0
            train = true
        end
        if time < 4000 && msec%20 == 0
            step_ = div(time,20) + 1
            input_spikes = display_spikes[step_]
            spiked = step!(n, input_spikes, train)
        else
            spiked = step!(n,train)
        end
        push!(spike_hist,spiked)
    end
end

# %% create a heatmap

using Plots
gr()
function heatmap_spikes(spikes::Array{Array{Int64}}, n_interval::Array{Int64},layer_bound::Array{Int64};
                        window::Int64 = 50, save::Bool = false)
    mat_spikes = zeros(n_interval[2],length(spikes))
    for time in 1:length(spikes)
        col = zeros(n_interval[2])
        if time > window
            for i in time-window:time
                indices = filter(x->n_interval[1]<= x <= n_interval[2] ,spikes[i])
                col[indices] .+= 1
            end
            col ./= window
            col ./= maximum(col)
        else
            for i in 1:time
                indices = filter(x->n_interval[1]<= x <= n_interval[2] ,spikes[i])
                col[indices] .+= 1
            end
            col ./= time
            col ./= maximum(col)
        end
        mat_spikes[:,time] .= col
    end
    display_mat = mat_spikes[n_interval[1]:end,:]
    heatmap(display_mat)
    plot!(title = "Network activity neuron $(n_interval[1]) to $(n_interval[2])",
          xlabel = "Time (msec)",
          ylabel = "Neuron")
    for i in layer_bound
        display(plot!([i], seriestype = "hline", color = :white, legend = false))
    end
    if save == true
        savefig("Julia/CartPole/Heatmaps/heatmapDVSv2.png")
    end
end

# %%

heatmap_spikes(spike_hist, [1200,1330], [64,128], save = true)

max_stars_repo_name: SuReLI/RL-STDP

max_issues_repo_name: SuReLI/RL-STDP

max_forks_repo_name: SuReLI/RL-STDP