text:
struct SineOperation{rank} <: Operation{rank}
    shape::Tuple{}
    children::Tuple{AbstractTensor}
    freeindices::Tuple{}
    function SineOperation(A::AbstractTensor{rank}) where rank
        new{rank}(A.shape, (A,), ())
    end
end

function Base.sin(s:: AbstractTensor{0})
    return SineOperation(Tensor(s))
end

struct CosineOperation{rank} <: Operation{rank}
    shape::Tuple{}
    children::Tuple{AbstractTensor}
    freeindices::Tuple{}
    function CosineOperation(A::AbstractTensor{rank}) where rank
        new{rank}(A.shape, (A,), ())
    end
end

function Base.cos(s::AbstractTensor{0})
    return CosineOperation(Tensor(s))
end

struct TangentOperation{rank} <: Operation{rank}
    shape::Tuple{}
    children::Tuple{AbstractTensor}
    freeindices::Tuple{}
    function TangentOperation(A::AbstractTensor{rank}) where rank
        new{rank}(A.shape, (A,),())
    end
end

function Base.tan(s::AbstractTensor{0})
    return TangentOperation(Tensor(s))
end

struct ArcsineOperation{rank} <: Operation{rank}
    shape::Tuple{}
    children::Tuple{AbstractTensor}
    freeindices::Tuple{}
    function ArcsineOperation(A::AbstractTensor{rank}) where rank
        if A isa ConstantTensor{T} where T<: Number && abs(A.value) > 1
            throw(DomainError(A, " asin is not defined for |x| > 1."))
        end
        new{rank}(A.shape, (A,),())
    end
end

function Base.asin(s::AbstractTensor{0})
    return ArcsineOperation(Tensor(s))
end

struct ArccosineOperation{rank} <: Operation{rank}
    shape::Tuple{}
    children::Tuple{AbstractTensor}
    freeindices::Tuple{}
    function ArccosineOperation(A::AbstractTensor{rank}) where rank
        if A isa ConstantTensor{T} where T<: Number && abs(A.value) > 1
            throw(DomainError(A, " acos is not defined for |x| > 1."))
        end
        new{rank}(A.shape, (A,),())
    end
end

function Base.acos(s::AbstractTensor{0})
    return ArccosineOperation(Tensor(s))
end

struct ArctangentOperation{rank} <: Operation{rank}
    shape::Tuple{}
    children::Tuple{AbstractTensor}
    freeindices::Tuple{}
    function ArctangentOperation(A::AbstractTensor{rank}) where rank
        new{rank}(A.shape, (A,),())
    end
end

function Base.atan(s::AbstractTensor{0})
    return ArctangentOperation(Tensor(s))
end

max_stars_repo_name: Sohva/TensorDSL

max_issues_repo_name: Sohva/TensorDSL

max_forks_repo_name: Sohva/TensorDSL


text:
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
using ForwardDiff
using Plots

f(x, y) = exp(-(x^2 + y^2))
u(x, y) = -ForwardDiff.derivative(Base.Fix2(f, y), x) # = f_x(x, y)
v(x, y) = -ForwardDiff.derivative(Base.Fix1(f, x), y) # = f_y(x, y)

xlim = (-1.7, 1.7)
ylim = (-1.7, 1.7)
xs = range(xlim...; length=200)
ys = range(ylim...; length=200)

c = 0.5
x = range(-1.5, 1.5; length=11)
y = range(-1.5, 1.5; length=11)
X, Y = reim(complex.(x', y)) # meshgrid
U, V = c*u.(x', y), c*v.(x', y)

heatmap(xs, ys, f)
quiver!(vec(X-U/2), vec(Y-V/2); quiver=(vec(U), vec(V)), color=:cyan)
plot!(; xlim, ylim, size=(450, 400))

# %%
using ForwardDiff
using Plots

f(x, y) = y * exp(-(x^2 + y^2))
u(x, y) = -ForwardDiff.derivative(Base.Fix2(f, y), x) # = f_x(x, y)
v(x, y) = -ForwardDiff.derivative(Base.Fix1(f, x), y) # = f_y(x, y)

xlim = (-2.2, 2.2)
ylim = (-2.2, 2.2)
xs = range(xlim...; length=200)
ys = range(ylim...; length=200)

c = 0.5
x = range(-2.0, 2.0; length=13)
y = range(-2.0, 2.0; length=13)
X, Y = reim(complex.(x', y)) # meshgrid
U, V = c*u.(x', y), c*v.(x', y)

heatmap(xs, ys, f)
quiver!(vec(X-U/2), vec(Y-V/2); quiver=(vec(U), vec(V)), color=:cyan)
plot!(; xlim, ylim, size=(450, 400))

# %%

ext: jl

max_stars_repo_path: 0015/Pkg/heatmap and quiver example.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_path: 0015/Pkg/heatmap and quiver example.jl

max_issues_repo_name: genkuroki/public

max_forks_repo_path: 0015/Pkg/heatmap and quiver example.jl

max_forks_repo_name: genkuroki/public


text:
module TestConverterLens

include("preamble.jl")
using StaticArrays

@testset begin
    obj = (x = ((0, 1, 2), "A"), y = "B");
    lens = (@lens _.x[1]) ∘ getting(SVector)
    @test get(obj, lens) === SVector(obj.x[1])
    @test set(obj, lens, SVector(3, 4, 5)) ===
        (x = ((3, 4, 5), "A"), y = "B")
end

@testset begin
    obj = (x = ((a = 0, b = 1, c = 2), "A"), y = "B");
    @testset for lens in [
        (@lens _.x[1]) ∘ getting(Base.splat(SVector))
        (@lens _.x[1]) ∘ getting(Tuple) ∘ getting(SVector)
    ]
        @test get(obj, lens) === SVector(obj.x[1]...)
        @test set(obj, lens, SVector(3, 4, 5)) ===
            (x = ((a = 3, b = 4, c = 5), "A"), y = "B")
    end
end

struct CustomType{T}
    a::T
    b::T
    c::T
end

CustomType(x::SVector) = CustomType(x...)
StaticArrays.SVector(obj::CustomType) = SVector(obj.a, obj.b, obj.c)

@testset begin
    obj = (x = (CustomType(0, 1, 2), "A"), y = "B");
    lens = (@lens _.x[1]) ∘ getting(SVector)
    @test get(obj, lens) === SVector(obj.x[1])
    @test set(obj, lens, SVector(3, 4, 5)) ===
        (x = (CustomType(3, 4, 5), "A"), y = "B")
end

end  # module


text:
function Grad(M, bound = "sym", order = 1)
    """
        grad - gradient, forward differences

          [gx,gy] = grad(M, options);
        or
          g = grad(M, options);

          options.bound = 'per' or 'sym'
          options.order = 1 (backward differences)
                        = 2 (centered differences)

          Works also for 3D array.
          Assme that the function is evenly sampled with sampling step 1.

          See also: div.

          Copyright (c) Gabriel Peyre
    """

    # retrieve number of dimensions
    nbdims = ndims(M)

    if bound == "sym"
        nx = size(M)[1]
        if order == 1
            fx = M[vcat(collect(2 : nx), nx), :] - M
        else
            fx = (M[vcat(collect(2 : nx), nx), :] - M[vcat(1, collect(1 : nx - 1)), :])./2.
            # boundary
            fx[1, :] = M[2, :] - M[1, :]
            fx[nx, :] = M[nx, :] - M[nx - 1, :]
        end

        if nbdims >= 2
            ny = size(M)[2]
            if order == 1
                fy = M[:, vcat(collect(2 : ny), ny)] - M
            else
                fy = (M[:, vcat(collect(2 : ny), ny)] - M[:, vcat(1, collect(1 : ny - 1))])./2.
                # boundary
                fy[:, 1] = M[:, 2] - M[:, 1]
                fy[:, ny] = M[:, ny] - M[:, ny - 1]
            end
        end

        if nbdims >= 3
            nz = size(M)[3]
            if order == 1
                fz = M[:, :, vcat(collect(2 : nz), nz)] - M
            else
                fz = (M[:, :, vcat(collect(2 : nz), nz)] - M[:, :, vcat(1, collect(1 : nz - 1))])./2.
                # boundary
                fz[:, :, 1] = M[:, :, 2] - M[:, :, 1]
                fz[:, :, ny] = M[:, :, nz] - M[:, :, nz - 1]
            end
        end
    else
        nx = size(M)[1]
        if order == 1
            fx = M[vcat(collect(2 : nx), 1), :] - M
        else
            fx = (M[vcat(collect(2 : nx), 1), :] - M[vcat(nx, collect(1 : nx - 1)), :])./2.
        end

        if nbdims >= 2
            ny = size(M)[2]
            if order == 1
                fy = M[:, vcat(collect(2 : ny), 1)] - M
            else
                fy = (M[:, vcat(collect(2 : ny), 1)] - M[:, vcat(ny, collect(1 : ny - 1))])./2.
            end
        end

        if nbdims >= 3
            nz = size(M)[3]
            if order == 1
                fz = M[:, :, vcat(collect(2 : nz), 1)] - M
            else
                fz = (M[:, :, vcat(collect(2 : nz), 1)] - M[:, :, vcat(nz, collect(1 : nz - 1))])./2.
            end
        end
    end

    if nbdims==2
        fx = cat(3, fx[:, :], fy[:, :])
    elseif nbdims==3
        fx = cat(4, (fx[:, :, :], fy[:, :, :], fz[:, :, :]))
    end

    fx[end, :, 1] = fx[end - 1, :, 1]
    fx[:, end, 2] = fx[:, end - 1, 2]

    return fx
end


text:
"""
# Arguments
- `S::Int` - Current Susceptible Population
- `I::Vector{Int}` - Current Infected Population
- `R::Int` - Current Recovered Population
- `Tests::Matrix{Int}` - Array for which people belonging to array element ``T_{i,j}`` are ``i-1`` days away
    from receiving positive test and have infection age ``j``
"""
struct DoubleCovidState <: CovidPOMDPs.CovidState
    S::Int # Current Susceptible Population
    I1::Vector{Int} # Current Infected Population
    I2::Vector{Int} # Current Infected Population
    R::Int # Current Recovered Population
    T1::Matrix{Int} # Rows: Days from receiving test result; Columns: Infection Age
    T2::Matrix{Int} # Rows: Days from receiving test result; Columns: Infection Age
    params::InfParams
    prev_action::CovidAction
end

CovidPOMDPs.CovidState(S,I1,I2,R,T1,T2,p,a) = DoubleCovidState(S,I1,I2,R,T1,T2,p,a)

CovidPOMDPs.statevars(s::DoubleCovidState) = (s.S, s.I1, s.I2, s.R, s.T1, s.T2)

CovidPOMDPs.susceptible(s::DoubleCovidState) = s.S
CovidPOMDPs.infected(s::DoubleCovidState) = sum(s.I1) + sum(s.I2)
CovidPOMDPs.recovered(s::DoubleCovidState) = s.R

Base.@kwdef struct DoubleCovidPOMDP{A} <: POMDP{DoubleCovidState, CovidPOMDPs.CovidAction, Int}
    "Delay (in days) between test being administered and result of test being received `(≥ 0)`"
    test_delay::Int = 1

    "Total population count `(> 0)`"
    N::Int = 10^6

    "POMDP discount factor `(γ ∈ [0,1])`"
    discount::Float64 = 0.95

    "Weight with which to penalize new infections `(≥ 0.0)`"
    inf_loss::Float64 = 1.0

    "Weight with which to penalize testing rate `(≥ 0.0)`"
    test_loss::Float64 = 1.0

    "Weight with which to penalize changes in testing rate `(≥ 0.0)`"
    testrate_loss::Float64 = 1.0

    "Number of days for which a testing policy must be held `(≥ 1)`"
    test_period::Int = 1

    "CovidPOMDP action space (default is continuous [0,1])"
    actions::A = CovidPOMDPs.CovidActionSpace()

    params::CovidPOMDPs.InfParams = DEFAULT_PARAMS
end

"""
Take given CovidPOMDP obj and return same CovidPOMDP obj only with test_period changed to 1
"""
function CovidPOMDPs.unity_test_period(pomdp::DoubleCovidPOMDP)
    return DoubleCovidPOMDP(
        pomdp.test_delay,
        pomdp.N,
        pomdp.discount,
        pomdp.inf_loss,
        pomdp.test_loss,
        pomdp.testrate_loss,
        1,
        pomdp.actions,
        pomdp.params
    )
end

function CovidPOMDPs.init_SIRT(pomdp::DoubleCovidPOMDP, rng::AbstractRNG=Random.GLOBAL_RNG)
    N = pomdp.N
    S, inf, R = floor.(Int, CovidPOMDPs.simplex_sample(3, Float64(N), rng))
    inf1, inf2 = floor.(Int, CovidPOMDPs.simplex_sample(2, Float64(inf), rng))

    horizon = INFECTION_HORIZON

    I1 = floor.(Int, CovidPOMDPs.simplex_sample(horizon, Float64(inf1), rng))
    I2 = floor.(Int, CovidPOMDPs.simplex_sample(horizon, Float64(inf2), rng))

    leftover = N - (S + sum(I1) + sum(I2) + R)
    R += leftover

    T1 = zeros(Int, pomdp.test_delay+1, horizon)
    T2 = zeros(Int, pomdp.test_delay+1, horizon)

    return S, I1, I2, R, T1, T2
end

function CovidPOMDPs.init_SIRT(pomdp::DoubleCovidPOMDP, inf::Int, rng::AbstractRNG=Random.GLOBAL_RNG)
    N = pomdp.N
    R = 0

    horizon = INFECTION_HORIZON
    inf1, inf2 = floor.(Int, simplex_sample(horizon, Float64(inf), rng))
    I1 = floor.(Int, CovidPOMDPs.simplex_sample(horizon, Float64(inf1), rng))
    I2 = floor.(Int, CovidPOMDPs.simplex_sample(horizon, Float64(inf2), rng))

    leftover = N - (sum(I1) + sum(I2))
    S = leftover

    T1 = zeros(Int, pomdp.test_delay+1, horizon)
    T2 = zeros(Int, pomdp.test_delay+1, horizon)

    return S, I1, I2, R, T1, T2
end


text:
"""
    int_F²(f)

Compute the time profile contribution to the ``b``-value given the time profile `f`. The
``b``-value is given by

```julia
b = γ^2 * g^2 * int_F²(f)
```

where `γ` is the gyromagnetic ratio of the water proton and `g` is the gradient amplitude.
"""
int_F²(f::TimeProfile{T}) where {T} = quadgk(s -> integral(f, s)^2, zero(T), echotime(f))
int_F²(f::PGSE) = f.δ^2 * (f.Δ - f.δ / 3)
int_F²(f::DoublePGSE) = 2f.δ^2 * (f.Δ - f.δ / 3)
int_F²(f::CosOGSE) = f.δ^3 / (2π * f.nperiod)^2
int_F²(f::SinOGSE) = 3f.δ^3 / (2π * f.nperiod)^2


text:
# AUTO GENERATED FILE - DO NOT EDIT

export vtk_slicerepresentation

"""
    vtk_slicerepresentation(;kwargs...)
    vtk_slicerepresentation(children::Any;kwargs...)
    vtk_slicerepresentation(children_maker::Function;kwargs...)

A SliceRepresentation component.
SliceRepresentation is responsible to convert a vtkPolyData into rendering
It takes the following set of properties:
  - colorBy: ['POINTS', ''],
  - pointSize: 1,
  - color: [1,1,1],
Keyword arguments:
- `children` (Array of a list of or a singular dash component, string or numbers | a list of or a singular dash component, string or number; optional)
- `id` (String; optional): The ID used to identify this component.
- `mapper` (Dict; optional): Properties to set to the mapper
- `actor` (Dict; optional): Properties to set to the slice/actor
- `property` (Dict; optional): Properties to set to the volume.property
- `colorMapPreset` (String; optional): Preset name for the lookup table color map
- `colorDataRange` (Array of Reals | String; optional): Data range use for the colorMap
- `iSlice` (Real; optional): index of the slice along i
- `jSlice` (Real; optional): index of the slice along j
- `kSlice` (Real; optional): index of the slice along k
- `xSlice` (Real; optional): index of the slice along x
- `ySlice` (Real; optional): index of the slice along y
- `zSlice` (Real; optional): index of the slice along z
"""
function vtk_slicerepresentation(; kwargs...)
        available_props = Symbol[:children, :id, :mapper, :actor, :property, :colorMapPreset, :colorDataRange, :iSlice, :jSlice, :kSlice, :xSlice, :ySlice, :zSlice]
        wild_props = Symbol[]
        return Component("vtk_slicerepresentation", "SliceRepresentation", "dash_vtk", available_props, wild_props; kwargs...)
end

vtk_slicerepresentation(children::Any; kwargs...) = vtk_slicerepresentation(;kwargs..., children = children)
vtk_slicerepresentation(children_maker::Function; kwargs...) = vtk_slicerepresentation(children_maker(); kwargs...)

max_stars_repo_name: ericbdaniels/dash-vtk

max_issues_repo_name: ericbdaniels/dash-vtk

max_forks_repo_name: ericbdaniels/dash-vtk


text:
abstract type Normalizer end

export Normalizer

export MeanNorm
include("./MeanNorm.jl")

export ZNorm
export BatchNorm0d, BatchNorm1d
include("./ZNorm.jl")

export wnorm
include("./weight-normer.jl")

max_stars_repo_name: sonosole/Delta

max_issues_repo_name: sonosole/Delta

max_forks_repo_name: sonosole/Delta


text:
function jt01b()
#
#  test polynomial suggested by Jenkins and Traub
#
   a = 10.0^(-10)
   p = [1;-1;-a^2;a^2]
   z = [a 1;-a 1;1 1]
   p, PolyZeros(z)
end


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

function validations(c::Collar, s::Survey, intervals::Intervals)
    # get initial info
    cfields = [c.holeid,c.x,c.y,c.z]
    sfields = [s.holeid,s.at,s.azm,s.dip]
    c.enddepth != nothing && push!(cfields,c.enddepth)
    it = intervals isa Interval ? [intervals] : intervals

    # read files
    dfc = c.file isa String ? CSV.read(c.file, DataFrame, select=cfields) : select(c.file,cfields)
    dfs = s.file isa String ? CSV.read(s.file, DataFrame, select=sfields) : select(s.file,sfields)
    dfi = []
    for f in it
        ffields = [f.holeid,f.from,f.to]
        d = f.file isa String ? CSV.read(f.file, DataFrame, select=ffields) : select(f.file,ffields)
        push!(dfi,d)
    end

    # output warning table
    out = DataFrame(TYPE = String[], FILE = String[], DESCRIPTION = String[])

    ## ERRORS
    # check for inappropriate nans
    nanc = [missing isa dtype for dtype in eltype.(eachcol(dfc))]
    nans = [missing isa dtype for dtype in eltype.(eachcol(dfs))]
    nani = [[missing isa dtype for dtype in eltype.(eachcol(df))] for df in dfi]
    # check for inappropriate duplicates
    dupc = findall(nonunique(dfc,[c.holeid]))
    dups = findall(nonunique(dfs,[s.holeid,s.at]))
    # check for typerrors; numeric: coords, enddepth, at, azm, dip from, to
    typc = [Number <: x for x in eltype.(eachcol(dfc))[2:end]]
    typs = [Number <: x for x in eltype.(eachcol(dfs))[2:end]]
    typi = [[Number <: x for x in eltype.(eachcol(df))[2:end]] for df in dfi]
    # check for overlaps
    ovlp = [overlaps(dfi[x],it[x]) for x in 1:length(it)]

    # error descriptions
    outnan(df,nan) = "Missing values in the column(s) $(names(df)[nan])"
    outdup(rows)   = replace("Duplicate values in the row(s) $rows","Any"=>"")
    outtyp(df,typ) = "Non-numeric values in the column(s) $(names(df)[2:end][typ])"
    outovl(rows)   = replace("Overlaps in the row(s) $rows","Any"=>"")
    fname(x)       = x isa String ? Base.Filesystem.basename(x) : @varname(x)

    # check for errors and add it to output table if it exists
    sum(nanc)>0    && push!(out,("Error",fname(c.file),outnan(dfc,nanc)))
    sum(nans)>0    && push!(out,("Error",fname(s.file),outnan(dfs,nans)))
    length(dupc)>0 && push!(out,("Error",fname(c.file),outdup(dupc)))
    length(dups)>0 && push!(out,("Error",fname(s.file),outdup(dups)))
    sum(typc)>0    && push!(out,("Error",fname(c.file),outtyp(dfc,typc)))
    sum(typs)>0    && push!(out,("Error",fname(s.file),outtyp(dfs,typs)))
    for x in 1:length(it)
      sum(nani[x])>0  && push!(out,("Error",fname(it[x].file),outnan(dfi[x],nani[x])))
      sum(typi[x])>0  && push!(out,("Error",fname(it[x].file),outtyp(dfi[x],typi[x])))
      sum(typi[x])==0 && length(ovlp[x])>0 && push!(out,("Error",fname(it[x].file),outovl(ovlp[x])))
    end

    # if some error exists, stop validations and return the errors
    size(out,1)>0 && (return sort!(out,[:TYPE,:DESCRIPTION,:FILE]))

    ## WARNINGS

    # check enddepth
    # check duplicate names and inform (maybe also if there is a preexisting x y z length)
    # bhid do not exist in collar table
    # survey do not exist
    # survey do not exist at zero
    # survey duplicated (same props)
    # survey single value
    # survey vector before and after too different
    # intervals list of absent collars
    # list of numeric interval values; if not, check for alpha values in table
    # warn if coordinate == -9999.9999 during fillxyz!
    # table columns: warn repeated names

    sort!(out,[:TYPE,:DESCRIPTION,:FILE])
end

# check for overlaps
function overlaps(df,pars)
    bh, f, t, ix = pars.holeid, pars.from, pars.to, :_IDX_
    n = size(df,1)
    df[!,ix] = collect(1:n)
    sort!(df,[bh,f,t])

    ovlps = []
    for i in 1:size(df,1)
        df[i,f] >= df[i,t] && push!(ovlps,df[i,ix])
        i==1 && continue
        df[i,bh] != df[i-1,bh] && continue
        df[i,f] < df[i-1,t] && append!(ovlps,[df[i-1,ix],df[i,ix]])
    end
    ovlps
end

"""
    exportwarns(dh::DrillHole, outname="errors")

Export errors and warnings identified during drill hole desurvey of `dh`
to the file `outname`.csv.
"""
exportwarns(dh::DrillHole, outname::String="errors") = CSV.write("$outname.csv", dh.warns)


text:
# ------------------------------------------
#    ACE Fitting

import IPFitting, ACE1pack, ACE1
using Dates, Base

export fit_params, fit_ace, make_ace_db, db_params, fit_ace_db, save_fit

"""
`fit_ace(params::Dict)` : function to set up and fit the least-squares
problem of "atoms' positions" -> "energy, forces, (virials)". Takes in a
dictionary with all the parameters. see `?fit_params` for details.

Returns `IP`, `lsqinfo`
"""
function fit_ace(params::Dict)

    db = _decide_how_to_get_db(params)
    IP, lsqinfo = fit_ace_db(db, params)

    return IP, lsqinfo
end

"""
`fit_params(; kwargs...)` : returns a dictionary containing all of the
parameters needed for making an ACE potential. All parameters are passed
as keyword argumts.

### Parameters
* `data` : data parameters, see `?data_params` for details (mandatory)
* `basis` : dictionary containing dictionaries that specify the basis used
in fitting. Usually just `ace_params` and `pair_params` (mandatory).
For example
```
basis = Dict(
    "ace" => ace_params(; kwargs...),
    "pair" => pair_params(; kwargs...))
````
keys of `basis` must correspond to one of the "type" of `basis_params`, see
`?basis_params` and values are corresponding parameters' dictionaries.
* `solver` : dictionary containing parameters that specify the solver for
least squares problem (mandatory). See `?solver_params`.
* `e0` : Dict{String, Float} containing reference values for isolated atoms'
energies (mandatory).
* `weights` : dictionary of `Dict("config_type" => Dict("E" => Float, "F => Float))``
entries specifying fitting weights. "default" is set to `1.0` for all of "E", "F",
and "V" weights.
* `P` : regularizer parameters (optional), see `?regularizer_params`.
* `ACE_fname = "ACE_fit.json"` : filename to save ACE to. Potential & info
do not get saved if `ACE_fname` isnothing() or is set to `""`. Files already _parse_entry
are renamed and not overwritten.
* `LSQ_DB_fname_stem = ""` : stem to save LsqDB to. Doesn't get saved if set to an empty
string (""). If the file is already present, but `fit_from_LSQ_DB` is set to false,
the old database is renamed, a new one constructed and saved under the given name.
* `fit_from_LSQ_DB = false`: whether to fit from a least squares database specified with
`LSQ_DB_fname_stem`. If `LSQ_DB_fname_stem * "_kron.h5"` file is not present, LsqDB is
constructed from scratch and saved.

"""
function fit_params(;
    data = nothing,
    basis = nothing,
    solver = nothing,
    e0 = nothing,
    weights = nothing,
    P = nothing,
    ACE_fname = "ACE_fit.json",
    LSQ_DB_fname_stem = "",
    fit_from_LSQ_DB = false)

    @assert !isnothing(data) "`data` is mandatory"
    @assert !isnothing(basis) "`basis` is mandatory"
    @assert !isnothing(solver) "`solver` is mandatory"
    @assert !isnothing(e0) "`e0` is mandatory"

    return Dict(
            "data" => data,
            "basis" => basis,
            "solver" => solver,
            "e0" => e0,
            "weights" => weights,
            "P" => P,
            "ACE_fname" => ACE_fname,
            "LSQ_DB_fname_stem" => LSQ_DB_fname_stem,
            "fit_from_LSQ_DB" => fit_from_LSQ_DB)
end

"""
`save_fit(fname, IP, lsqinfo)` : saves Dict("IP" => IP, "info" => lsqinfo) to fname.
If `fname` is already present, it is renamed and dictionary saved to `fname`.
"""
function save_fit(fname, IP, lsqinfo)
    # ENH: save to yace option
    if fname == "" || isnothing(fname)
        return
    end
    if isfile(fname)
        stem = replace(fname, ".json" => "")
        fnew =  stem * "." * String(rand('a':'z', 5)) * ".json"
        @warn("The file $fname already exists. It will be renamed to $fnew to avoid overwriting.")
        mv(fname, fnew)
    end
    @info("Saving ace fit to $(fname)")
    save_dict(fname, Dict("IP" => write_dict(IP), "info" => lsqinfo))
end

"""
`fit_ace_db(params::Dict)` : fits LsqDB with `params["LSQ_DB_fname_stem"], which must be already present.
`params["fit_from_LSQ_DB"]` must be set to true. See `?fit_params` for `params` specification,
of which `data` and `basis` aren't needed (are ignored).
"""
function fit_ace_db(params::Dict)
    @assert params["fit_from_LSQ_DB"]
    db = LsqDB(params["LSQ_DB_fname_stem"])
    IP, lsqinfo = fit_ace_db(db, params)
    return IP, lsqinfo
end

"""
`fit_ace_db(db::IPFitting.LsqDB, params::Dict)` : fits the given LsqDB. See `?fit_params` for
`params` specification, of which `data` and `basis` aren't needed (are ignored).
"""
function fit_ace_db(db::IPFitting.LsqDB, params::Dict)
    solver = ACE1pack.generate_solver(params["solver"])

    if !isnothing(params["P"])
        solver["P"] = ACE1pack.generate_regularizer(db.basis, params["P"])
    end

    if typeof(params["e0"]) == Dict{Any, Any}
        # sometimes gets read in (from yaml?) as Dict{Any, Any}
        # which gives StackOverflowError somewhere in OneBody
        params["e0"] = convert(Dict{String, Any}, params["e0"])
    end

    Vref = OneBody(params["e0"])
    weights = params["weights"]

    vals, solve_time, bytes, _, _ = @timed IPFitting.Lsq.lsqfit(
        db,
        Vref=Vref,
        solver=solver,
        weights=weights,
        error_table=true
    );

    IP = vals[1]
    lsqinfo = vals[2]
    solve_time = canonicalize(Dates.CompoundPeriod(Second(trunc(Int, solve_time))))
    bytes = Base.format_bytes(bytes)
    lsqinfo["solve_time"] = solve_time
    lsqinfo["solve_mem"] = bytes
    lsqinfo["fit_params"] = params

    save_fit(params["ACE_fname"], IP, lsqinfo)

    @info("Fitting errors")
    rmse_table(lsqinfo["errors"])
    @info("LsqDB solve time: $(solve_time)")
    @info("LsqDB memory: $(bytes)")

    return IP, lsqinfo
end

"""
`make_ace_db(params::Dict)` : makes a LsqDB from given parameters' dictionary.
For `params` see `?db_params`; parameters from `fit_params` also work, except
unnecessary entries will be ignored.

"""
function make_ace_db(params::Dict)
    data =  ACE1pack.read_data(params["data"])
    basis = [ACE1pack.generate_basis(basis_params) for (basis_name, basis_params) in params["basis"]]
    basis = JuLIP.MLIPs.IPSuperBasis(basis);
    db = LsqDB(params["LSQ_DB_fname_stem"], basis, data)
    return db
end

"""
db_params(; kwargs...)` : returns a dictionary containing all of the
parameters needed for making a LsqDB. All parameters are passed
as keyword argumts.

### Parameters
* `data` : data parameters, see `?data_params` for details (mandatory)
* `basis` : dictionary containing dictionaries that specify the basis used
in fitting. Usually just `ace_params` and `pair_params` (mandatory).
For example
```
basis = Dict{
    "ace" => ace_params(; kwargs...),
    "pair" => pair_params(; kwargs...)}
````
keys of `basis` must correspond to one of the "type" of `basis_params`, see
`?basis_params` and values are corresponding parameters' dictionaries.
* `LSQ_DB_fname_stem = ""` : stem to save LsqDB to. Doesn't get saved if set to an empty
string (""). If `LSQ_DB_fname_stem * "_kron.h5"` file is not present it gets renamed,
a new LsqDB is constructed and saved.
"""
function db_params(;
    data = nothing,
    basis = nothing,
    LSQ_DB_fname_stem = "")

    @assert !isnothing(data) "`data` is mandatory"
    @assert !isnothing(basis) "`basis` is mandatory"

    return Dict(
            "data" => data,
            "basis" => basis,
            "LSQ_DB_fname_stem" => LSQ_DB_fname_stem)
end

function _decide_how_to_get_db(params::Dict)
    if params["LSQ_DB_fname_stem"] == ""
        return make_ace_db(params)
    else
        db_fname = params["LSQ_DB_fname_stem"] * "_kron.h5"
        if ~isfile(db_fname)
            return make_ace_db(params)
        else
            if params["fit_from_LSQ_DB"]
                return LsqDB(params["LSQ_DB_fname_stem"])
            else
                @warn("Found $(db_fname), but wasn't asked to fit it. Making a new one and saving to $db_fname).")
                return make_ace_db(params)
            end
        end
    end
end


text:
#=
# This is currently slower than splitting the loop.
# See https://github.com/JuliaSIMD/LoopVectorization.jl/issues/284
function downconvert_and_correlate(
    signal::StructArray{Complex{T}},
    correlator::C,
    code,
    correlator_sample_shifts,
    carrier_frequency,
    sampling_frequency,
    start_phase,
    start_sample,
    num_samples
) where {T, C <: AbstractCorrelator}
    s_re = signal.re; s_im = signal.im
    accumulators = zero_accumulators(get_accumulators(correlator), signal)
    a_re = real.(accumulators)
    a_im = imag.(accumulators)
    @avx for i = start_sample:start_sample + num_samples - 1
        c_im, c_re = sincos(T(2π) * ((i - start_sample) * T(upreferred(carrier_frequency / Hz)) / T(upreferred(sampling_frequency / Hz)) + T(start_phase)))
        d_re = s_re[i] * c_re + s_im[i] * c_im
        d_im = s_im[i] * c_re - s_re[i] * c_im
        for j = 1:length(a_re)
            sample_shift = correlator_sample_shifts[j] - correlator_sample_shifts[1]
            a_re[j] += d_re * code[i + sample_shift]
            a_im[j] += d_im * code[i + sample_shift]
        end
    end
    accumulators_result = complex.(a_re, a_im)
    C(map(+, get_accumulators(correlator), accumulators_result))
end
=#

# CUDA Kernel
function downconvert_and_correlate_kernel(
    res_re,
    res_im,
    signal_re,
    signal_im,
    codes,
    code_frequency,
    correlator_sample_shifts,
    carrier_frequency,
    sampling_frequency,
    start_code_phase,
    carrier_phase,
    code_length,
    prn,
    num_samples,
    num_ants,
    num_corrs
)
    cache = @cuDynamicSharedMem(Float32, (2 * blockDim().x, num_ants, num_corrs))
    sample_idx   = 1 + ((blockIdx().x - 1) * blockDim().x + (threadIdx().x - 1))
    antenna_idx  = 1 + ((blockIdx().y - 1) * blockDim().y + (threadIdx().y - 1))
    corr_idx     = 1 + ((blockIdx().z - 1) * blockDim().z + (threadIdx().z - 1))
    iq_offset = blockDim().x
    cache_index = threadIdx().x - 1

    code_phase = accum_re = accum_im = dw_re = dw_im = carrier_re = carrier_im = 0.0f0
    mod_floor_code_phase = Int(0)

    if sample_idx <= num_samples && antenna_idx <= num_ants && corr_idx <= num_corrs
        # generate carrier
        carrier_im, carrier_re = CUDA.sincos(2π * ((sample_idx - 1) * carrier_frequency / sampling_frequency + carrier_phase))

        # downconvert with the conjugate of the carrier
        dw_re = signal_re[sample_idx, antenna_idx] * carrier_re + signal_im[sample_idx, antenna_idx] * carrier_im
        dw_im = signal_im[sample_idx, antenna_idx] * carrier_re - signal_re[sample_idx, antenna_idx] * carrier_im

        # calculate the code phase
        code_phase = code_frequency / sampling_frequency * ((sample_idx - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase

        # wrap the code phase around the code length e.g. phase = 1024 -> modfloorphase = 1
        mod_floor_code_phase = 1 + mod(floor(Int32, code_phase), code_length)

        # multiply elementwise with the code
        accum_re += codes[mod_floor_code_phase, prn] * dw_re
        accum_im += codes[mod_floor_code_phase, prn] * dw_im
    end

    cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] = accum_re
    cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] = accum_im

    ## Reduction
    # wait until all the accumulators have done writing the results to the cache
    sync_threads()

    i::Int = blockDim().x ÷ 2
    @inbounds while i != 0
        if cache_index < i
            cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 0 * iq_offset + i, antenna_idx, corr_idx]
            cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 1 * iq_offset + i, antenna_idx, corr_idx]
        end
        sync_threads()
        i ÷= 2
    end

    if (threadIdx().x - 1) == 0
        res_re[blockIdx().x, antenna_idx, corr_idx] += cache[1 + 0 * iq_offset, antenna_idx, corr_idx]
        res_im[blockIdx().x, antenna_idx, corr_idx] += cache[1 + 1 * iq_offset, antenna_idx, corr_idx]
    end
    return nothing
end

function downconvert_and_correlate_kernel_wrapper(
    system,
    signal,
    correlator,
    code_phase,
    carrier_phase,
    code_frequency,
    correlator_sample_shifts,
    carrier_frequency,
    sampling_frequency,
    signal_start_sample,
    num_samples_left,
    prn
)
    num_corrs = length(correlator_sample_shifts)
    num_ants = size(signal, 2)
    num_samples = size(signal, 1)
    block_dim_z = num_corrs
    block_dim_y = num_ants
    # keep num_corrs and num_ants in seperate dimensions, truncate num_samples accordingly to fit
    block_dim_x = prevpow(2, 1024 ÷ block_dim_y ÷ block_dim_z)
    threads = (block_dim_x, block_dim_y, block_dim_z)
    blocks = cld(size(signal, 1), block_dim_x)
    res_re = CUDA.zeros(Float32, blocks, block_dim_y, block_dim_z)
    res_im = CUDA.zeros(Float32, blocks, block_dim_y, block_dim_z)
    shmem_size = sizeof(ComplexF32)*block_dim_x*block_dim_y*block_dim_z
    @cuda threads=threads blocks=blocks shmem=shmem_size downconvert_and_correlate_kernel(
        res_re,
        res_im,
        signal.re,
        signal.im,
        system.codes,
        Float32(code_frequency),
        correlator_sample_shifts,
        Float32(carrier_frequency),
        Float32(sampling_frequency),
        Float32(code_phase),
        Float32(carrier_phase),
        size(system.codes, 1),
        prn,
        num_samples,
        num_ants,
        num_corrs
    )
    return sum(res_re .+ 1im*res_im, dims=1)
end


text:
# Hooks and event loop functionality

struct KeyEvent
    spec::AbstractString
end

import Base: ==
==(e1::KeyEvent, e2::KeyEvent) = e1.spec == e2.spec

struct CursorPosition
    pos::SVector{3,Float64}
end

_eventspec_str(s::AbstractString) = s
_eventspec_str(e::KeyEvent) = "key:$(e.spec)"

_argspec_str(s::AbstractString) = s
_argspec_str(::Type{Nothing}) = "null"
_argspec_str(::Type{CursorPosition}) = "cursor"

"""
    event_loop(callback::Function, [plotobj::DisplazWindow], event_list...)

Subscribe to a list of events, calling `callback` each time one is received.

Each event comes with some optional some state which is attached at the time
the event is triggered.  The events are specified as a list of `event=>state`
pairs.

Currently only `KeyEvent` is supported, with the possible arguments `Nothing` or
`CursorPosition`.  For example:

```
Displaz.event_loop(
        KeyEvent("c")=>Nothing,
        KeyEvent("p")=>CursorPosition
) do event, arg
    @show event, arg
    if event == KeyEvent("c")
        clearplot()
    end
end
```
"""
function event_loop(callback::Function, plotobj::DisplazWindow, event_list...)

    command = hook_command(plotobj.name, event_list...)

    open(command, "r") do event_stream
        handle_events(callback, event_stream)
    end

end

function hook_command(servername, event_list...)

    hookopts = [["-hook", _eventspec_str(e), _argspec_str(p)]
    for (e,p) in event_list]

    return `$_displaz_cmd -server $(servername) $(vcat(hookopts...))`
end

function handle_events(callback, event_stream)

    while !eof(event_stream)
        rawline = readline(event_stream)
        line = split(rawline)
        if length(line) < 2
            @warn("Unrecognized displaz hook string: \"$line\"")
            break
        end
        eventspec = line[1]
        if startswith(eventspec, "key:")
            event = KeyEvent(eventspec[5:end])
        else
            @warn("Unrecognized displaz event type: \"$eventspec\"")
            event = eventspec
        end
        argspec = line[2]
        if argspec == "null"
            arg = nothing
        elseif argspec == "cursor"
            arg = CursorPosition(map(s->parse(Float64,s), line[3:end]))
        else
            @warn("Unrecognized displaz hook payload: \"$argspec\"")
            arg = line[3:end]
        end
        callback(event, arg) != false || break
    end

end

event_loop(callback::Function, event_list...) = event_loop(callback, current(), event_list...)

# Wait until the given displaz event occurs
Base.wait(event::KeyEvent) = event_loop((e,a)->false, event=>Nothing)


text:
function helpers_test_unit()
    test_longest_consecutive()
    test_anneal_availability()
    test_get_day_bounds()
    test_get_longest_availability()
    test_sub_array()
    test_days_available_per_week()
    test_day_sum_coverage()
    test_week_sum_coverage()
    test_day_hours_scheduled()
    test_week_hours_scheduled()
end

function test_longest_consecutive()
    test_cases = [
        # (input, expected)
        ([1 1 1 1 1 0 0 1 1], 5),
        ([0 1 1 1 1 0 0 1 1], 4),
        (zeros(Int, 8), 0),
        (ones(Int, 10), 10),
        ([1 1 1 0 1 1 1 0 0 1 1], 3),
        ([1 1 0 1 1 1 0 0 1 1], 3),
        ([ones(Int, 12)], 12)
    ]

    for case in test_cases
        input, expected = case
        @test longest_consecutive(input) == expected
    end
end

function test_anneal_availability()
    test_cases = [
        # (avail, length, expected)
        (
            [1 1 1 1 1 0 0 1 1 0],
            5,
            [1 1 1 1 1 0 0 0 0 0]
        ),
        ([1 1 1 1 1 0 0 1 1], 5, [1 1 1 1 1 0 0 0 0]),
        ([1 1 1 1 1 0 1 1 1], 3, [1 1 1 1 1 0 1 1 1]),
        ([1 1 1 1], 4, [1 1 1 1]),
        ([1 1 1 1], 3, [1 1 1 1]),
        ([1 1 1 1 0 0 0 0 1 1 1 1 0 0 0], 4, [1 1 1 1 0 0 0 0 1 1 1 1 0 0 0]),
        ([1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0], 5, zeros(Int, 15)),
        ([0, 0, 1, 1, 1, 0, 1, 0], 3, [0, 0, 1, 1, 1, 0, 0, 0]),
        ([zeros(Int, 2), ones(Int, 10)], 10, [zeros(Int, 2), ones(Int, 10)]),
        (
            [[1, 1, 1, 1, 1, 0, 0, 1, 1, 0] [1, 1, 1, 1, 1, 0, 0, 1, 1, 0];],
            5,
            [[1, 1, 1, 1, 1, 0, 0, 0, 0, 0] [1, 1, 1, 1, 1, 0, 0, 0, 0, 0]]
        ),
    ]

    for case in test_cases
        avail, length, expected = case
        @test anneal_availability(avail, length) == expected
    end
end

function test_get_day_bounds()
    # Test a week where min requires you to assign full amount
    # to each day
    min, max = get_day_bounds(
        [8, 4, 3, 3, 3],
        [1, 1, 1, 1, 1],
        0,
        20,
        23,
        3
    )

    @test max == 8
    @test min == 7

    min, max = get_day_bounds(
        [4, 3, 3, 3],
        [1, 1, 1, 1],
        7,
        20,
        23,
        3
    )

    @test max == 4
    @test min == 4

    min, max = get_day_bounds(
        [3],
        [1],
        19,
        20,
        23,
        3
    )

    @test max == 3
    @test min == 3

    min, max = get_day_bounds(
        [5, 8, 8, 8],
        [1, 1, 0, 1],
        0,
        0,
        9,
        3,
    )

    @test max == 3
    @test min == 3

    min, max = get_day_bounds(
        [8, 8, 8],
        [1, 0, 1],
        3,
        0,
        9,
        3,
    )

    @test max == 3
    @test min == 3

    min, max = get_day_bounds(
        [8],
        [1],
        6,
        0,
        9,
        3,
    )

    @test max == 3
    @test min == 3

end

function test_get_longest_availability()
    day = 1
    employee = {
        "availability" => Array[
            [0, 1, 1, 1, 1, 0], # Don't trigger bounds
            [0, 1, 1, 0, 0, 0], # trigger min
            [1, 1, 1, 1, 1, 1], # trigger max
        ],
        "shift_time_min" => 3,
        "shift_time_max" => 5,
    }

    @test get_longest_availability(employee, 1) == 4
    @test get_longest_availability(employee, 2) == 0
    @test get_longest_availability(employee, 3) == employee["shift_time_max"]

end

function test_sub_array()
    array = [1, 2, 3, 4, 5]

    @test sub_array(array, 1, 5) == [1, 2, 3, 4, 5]
    @test sub_array(array, 2, 1) == [2, 3, 4, 5, 1]
    @test sub_array(array, 3, 2) == [3, 4, 5, 1, 2]
    @test sub_array(array, 1, 1) == [1]
end

function test_days_available_per_week()
    lenny = {
        "shift_time_min" => 4,
        "shift_time_max" => 6,
        "hours_min" => 16, # hours worked over simulation
        "hours_max" => 25, # hours worked over simulation
        "shift_count_min" => 4,
        "shift_count_max" => 5,
        "availability" => Array[
            [ones(Int, 7), zeros(Int, 5)],
            [ones(Int, 7), zeros(Int, 5)],
            [zeros(Int, 12)],
            [ones(Int, 7), zeros(Int, 5)],
            [zeros(Int, 5), ones(Int, 3), zeros(Int, 4)], # Annealed out
            [zeros(Int, 12)],
            [zeros(Int, 12)],
        ],
    }

    @test days_available_per_week(lenny) == 3
end

function test_day_sum_coverage()
    test_env = {
        "coverage"  => [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
        "cycle_length" => 12,
    }
    expected = 49
    @test day_sum_coverage(test_env) == expected
end

function test_week_sum_coverage()
    test_env = {
        "shift_time_min" => 3,
        "shift_time_max" => 8,
        "coverage"  => Array[
            # Monday
            [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
            # Tuesday
            [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
            # Wednesday
            [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
            # Thursday
            [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
            # Friday
            [3, 4, 5 * ones(Int, 6), 4 * ones(Int, 2), 2 * ones(Int, 2)],
            # Saturday
            [1, 2 * ones(Int, 11)],
            # Sunday
            [1, 2 * ones(Int, 11)],
        ],
        "time_between_coverage" => 12,
        "intershift" => 13, # Time periods between shifts   # only serve 1 shift/day
    }

    expected = 49*5 + 23*2
    @test week_sum_coverage(test_env) == expected
end
function test_day_hours_scheduled()
    schedule = {
        "bob" => {
            "start" => 8,
            "length" => 13,
        },
        "joe" => {
            "start" => 1,
            "length" => 30,
        }
    }

    expected = 43
    @test day_hours_scheduled(schedule) == expected
end
function test_week_hours_scheduled()
    schedule = {
        "bob" => [
            {
                "start" => 8,
                "length" => 13,
                "day" => 2,
            },
            {
                "start" => 0,
                "length" => 3,
                "day" => 2,
            },
        ],
        "billy" => [
            {
                "start" => 8,
                "length" => 19,
                "day" => 2,
            },
            {
                "start" => 0,
                "length" => 3,
                "day" => 2,
            },
        ],
    }

    expected = 38
    @test week_hours_scheduled(schedule) == expected
end

max_stars_repo_name: shdann/autoshift

max_issues_repo_name: shdann/autoshift

max_forks_repo_name: shdann/autoshift


text:
module BuildTree

    using PhylogeneticTrees, JuMP, Gurobi
    @static if VERSION < v"0.7.0-DEV.2005"
        using Base.Test
    else
        using Test
    end

    println("=============================")
    println("Testing node-based tree model")
    println("=============================")

    @testset "Solving a tree problem" begin

    	pd = PhylogeneticTrees.PopulationData(
    		"testdata/f3.Europe6.csv",
    		"testdata/f3.Europe6-covariance.csv"
    		)

        bt = PhylogeneticTrees.BinaryTree(3)
        tp = PhylogeneticTrees.NodeTreeProblem(pd, bt,
            binaryencoding = true, # this slows it down significantly but is just to test
            solver = GurobiSolver(OutputFlag = 0))
        PhylogeneticTrees.breaksymmetries(tp, rules = [:leftfirst, :alphabetize])
        @time solve(tp.model)
        # 12.707202 seconds (18.43 k allocations: 3.351 MiB, 0.11% gc time)

        leaves = PhylogeneticTrees.getnodes(bt, bt.depth)

        # test integrality
        for a in 1:pd.npop, u in leaves
            @test isapprox(getvalue(tp.assign[a,u,1]), 1) || isapprox(getvalue(tp.assign[a,u,1]), 0)
        end

        xval = [leaves[findfirst(round.(getvalue(tp.assign[a,:,1])) .> 0)] for a in 1:pd.npop]

        @test isapprox(round(getvalue(tp.f3formula[1,1,xval[1],xval[1],1,1])), 478)
        @test isapprox(round(getvalue(tp.f3formula[1,2,xval[1],xval[2],1,1])), 33)
        @test isapprox(round(getvalue(tp.f3formula[2,2,xval[2],xval[2],1,1])), 247)
        for a in 1:pd.npop, b in a:pd.npop, u in leaves, v in leaves
            u == tp.outgroupnode && continue
            v == tp.outgroupnode && continue
            xval[a] == u && continue
            xval[b] == v && continue
            @test isapprox(getvalue(tp.f3formula[a,b,u,v,1,1]), 0, atol=1e-4)
        end

        @test isapprox(getobjectivevalue(tp.model), 267.7057665969791)

        PhylogeneticTrees.removesolution(tp, getvalue(tp.assign))
        @time solve(tp.model)
        # 13.446876 seconds (35.10 k allocations: 2.273 MiB)
        @test isapprox(getobjectivevalue(tp.model), 267.70576659697684, atol=1e-4) # symmetric solution

    end

end


text:
default_combine(x, y) = y

merge!(out, a) = merge!(default_combine, out, a)
merge!(combine, out, a, b, others...) = merge!(merge!(combine, out, a), combine, b, others...)

function merge!(combine, out, a)
    for (k, v) ∈ pairs(a)
        if haskey(out, k)
            out[k] = combine(out[k], v)
        else
            out[k] = v
        end
    end
    return out
end

merge(a, b) = merge(default_combine, a, b)
merge(combine, a, b, c, others...) = merge(combine, merge(combine, a, b), c, others...)

function merge(combine, a, b)
    # TODO possibly make this work for different eltype, keytype of b, or immutable a, etc
    out = copy(a)
    merge!(combine, out, b)
    return out
end

merge(combine, a::Number, b::Number) = combine(a, b)

function merge(::typeof(default_combine), a::Base.OneTo, b::Base.OneTo)
    return Base.OneTo(max(last(a), last(b)))
end

function merge(::typeof(default_combine), a::AbstractVector, b::AbstractVector)
    if keys(a) ⊆ keys(b)
        return b
    else
        out = similar(a, promote_type(eltype(a), eltype(b)), merge(keys(a), keys(b)))
        out[keys(a)] = a
        out[keys(b)] = b
    end
    return out
end

function merge(::typeof(default_combine), a::Tuple, b::Tuple)
    # TODO this sucks...
    if length(a) <= length(b)
        return b
    else
        return ntuple(i -> i <= length(b) ? b[i] : a[i], a, b)
    end
end


text:
@testset "Constraint Sets" begin
model = Cartpole()
n,m = size(model)
N = 11
x,u = rand(model)
dt = 0.1
z = KnotPoint(x,u,dt)

#--- Generate some constraints
# Circle Constraint
xc = SA[1,1,1]
yc = SA[1,2,3]
r  = SA[1,1,1]
cir = CircleConstraint(n, xc, yc, r)

# Goal Constraint
xf = @SVector rand(n)
goal = GoalConstraint(xf)

# Linear Constraint
p = 5
A = @SMatrix rand(p,n+m)
b = @SVector rand(p)
lin = LinearConstraint(n,m,A,b, Inequality())

# Bound Constraint
xmin = -@SVector rand(n)
xmax = +@SVector rand(n)
umin = -@SVector rand(m)
umax = +@SVector rand(m)
bnd = BoundConstraint(n,m, x_min=xmin, x_max=xmax, u_min=umin, u_max=umax)

# Dynamics Constraint
dyn = TO.DynamicsConstraint(model, N)

#--- Create the Constraint List
cons = ConstraintList(n,m,N)
add_constraint!(cons, cir, 1:N)
add_constraint!(cons, goal, N)
add_constraint!(cons, lin, 2:N-1)
cons2 = copy(cons)
cons2_ = deepcopy(cons)
add_constraint!(cons, bnd, 1:N-1)
add_constraint!(cons, dyn, 1:N-1)
@test length(cons) == 5
@test length(cons2) == 3
@test cons[2] === cons2[2]
@test cons[2] !== cons2_[2]

end


text:
# a variety of rankings
#
# Please refer to http://en.wikipedia.org/wiki/Ranking#Strategies_for_assigning_rankings
# to see the definitions of a variety of ranking strategies
#
# The implementations here follow this wikipedia page.
#

function _check_randparams(rks, x, p)
    n = length(rks)
    length(x) == length(p) == n || raise_dimerror()
    return n
end

# ranking helper function: calls sortperm(x) and then ranking method f!
function _rank(f!, x::AbstractArray, R::Type=Int; sortkwargs...)
    rks = similar(x, R)
    ord = reshape(sortperm(vec(x); sortkwargs...), size(x))
    return f!(rks, x, ord)
end

# ranking helper function for arrays with missing values
function _rank(f!, x::AbstractArray{>: Missing}, R::Type=Int; sortkwargs...)
    inds = findall(!ismissing, vec(x))
    isempty(inds) && return missings(R, size(x))
    xv = disallowmissing(view(vec(x), inds))
    ordv = sortperm(xv; sortkwargs...)
    rks = missings(R, size(x))
    f!(view(rks, inds), xv, ordv)
    return rks
end

# Ordinal ranking ("1234 ranking") -- use the literal order resulted from sort
function _ordinalrank!(rks::AbstractArray, x::AbstractArray, p::IntegerArray)
    _check_randparams(rks, x, p)
    @inbounds for i in eachindex(p)
        rks[p[i]] = i
    end
    return rks
end

"""
    ordinalrank(x; lt=isless, by=identity, rev::Bool=false, ...)

Return the [ordinal ranking](https://en.wikipedia.org/wiki/Ranking#Ordinal_ranking_.28.221234.22_ranking.29)
("1234" ranking) of an array. Supports the same keyword arguments as the `sort` function.
All items in `x` are given distinct, successive ranks based on their position
in the sorted vector.
Missing values are assigned rank `missing`.
"""
ordinalrank(x::AbstractArray; sortkwargs...) =
    _rank(_ordinalrank!, x; sortkwargs...)

# Competition ranking ("1224" ranking) -- resolve tied ranks using min
function _competerank!(rks::AbstractArray, x::AbstractArray, p::IntegerArray)
    n = _check_randparams(rks, x, p)

    @inbounds if n > 0
        p1 = p[1]
        v = x[p1]
        rks[p1] = k = 1

        for i in 2:n
            pi = p[i]
            xi = x[pi]
            if xi != v
                v = xi
                k = i
            end
            rks[pi] = k
        end
    end

    return rks
end

"""
    competerank(x; lt=isless, by=identity, rev::Bool=false, ...)

Return the [standard competition ranking](http://en.wikipedia.org/wiki/Ranking#Standard_competition_ranking_.28.221224.22_ranking.29)
("1224" ranking) of an array. Supports the same keyword arguments as the `sort` function.
Equal (*"tied"*) items are given the same rank, and the next rank comes after a gap
that is equal to the number of tied items - 1.
Missing values are assigned rank `missing`.
"""
competerank(x::AbstractArray; sortkwargs...) =
    _rank(_competerank!, x; sortkwargs...)

# Dense ranking ("1223" ranking) -- resolve tied ranks using min
function _denserank!(rks::AbstractArray, x::AbstractArray, p::IntegerArray)
    n = _check_randparams(rks, x, p)

    @inbounds if n > 0
        p1 = p[1]
        v = x[p1]
        rks[p1] = k = 1

        for i in 2:n
            pi = p[i]
            xi = x[pi]
            if xi != v
                v = xi
                k += 1
            end
            rks[pi] = k
        end
    end

    return rks
end

"""
    denserank(x; lt=isless, by=identity, rev::Bool=false, ...)

Return the [dense ranking](http://en.wikipedia.org/wiki/Ranking#Dense_ranking_.28.221223.22_ranking.29)
("1223" ranking) of an array. Supports the same keyword arguments as the `sort` function.
Equal items receive the same rank, and the next subsequent rank is
assigned with no gap.
Missing values are assigned rank `missing`.
"""
denserank(x::AbstractArray; sortkwargs...) =
    _rank(_denserank!, x; sortkwargs...)

# Tied ranking ("1 2.5 2.5 4" ranking) -- resolve tied ranks using average
function _tiedrank!(rks::AbstractArray, x::AbstractArray, p::IntegerArray)
    n = _check_randparams(rks, x, p)

    @inbounds if n > 0
        v = x[p[1]]

        s = 1  # starting index of current range
        for e in 2:n # e is pass-by-end index of current range
            cx = x[p[e]]
            if cx != v
                # fill average rank to s : e-1
                ar = (s + e - 1) / 2
                for i = s : e-1
                    rks[p[i]] = ar
                end
                # switch to next range
                s = e
                v = cx
            end
        end

        # the last range
        ar = (s + n) / 2
        for i = s : n
            rks[p[i]] = ar
        end
    end

    return rks
end

"""
    tiedrank(x; lt=isless, by=identity, rev::Bool=false, ...)

Return the [tied ranking](http://en.wikipedia.org/wiki/Ranking#Fractional_ranking_.28.221_2.5_2.5_4.22_ranking.29),
also called fractional or "1 2.5 2.5 4" ranking,
of an array. Supports the same keyword arguments as the `sort` function.
Equal (*"tied"*) items receive the mean of the ranks they would
have been assigned under the ordinal ranking (see [`ordinalrank`](@ref)).
Missing values are assigned rank `missing`.
"""
tiedrank(x::AbstractArray; sortkwargs...) =
    _rank(_tiedrank!, x, Float64; sortkwargs...)


text:
@inline g(contact::Contact, x, q) = contact.Nx[SA[1; 2; 3]]' * (x + vrotate(contact.p,q) - contact.offset[SA[1; 2; 3]]) #c entfernung mit orientierung
@inline g(contact::Contact, state::State, Δt) = g(contact, getx3(state, Δt), getq3(state, Δt))

##Descrete Time Position Derivatives
@inline function ∂g∂pos(contact::Contact, x::AbstractVector, q::UnitQuaternion)
    X = contact.Nx                                                     #∂g∂x
    Q = contact.Nx * ((VLmat(q) * Lmat(UnitQuaternion(contact.p)) * Tmat()) + (VRᵀmat(q) * Rmat(UnitQuaternion(contact.p))))  #∂g∂q
    return X, Q                                                        #[∂g∂x ∂g∂q]
end

##Shurf und ShurD
#Shurf
@inline function schurf(contact::Contact, state::State, γ, s, μ, Δt)
    φ = g(contact, getx3(state, Δt), getq3(state, Δt))
    return ∂g∂ʳpos(contact, state)' * (γ / s * φ - μ / s)
end
#ShurD
@inline function schurD(contact::Contact,state::State, γ, s, Δt)
    return ∂g∂ʳpos(contact, state)' * γ/s * ∂g∂ʳvel(contact, state, Δt)
end

#Friction
@inline function calcFrictionForce!(mechanism::Mechanism{T}, friction::Contact, body::Body, γ) where T
    cf = friction.cf
    D = friction.D
    state = body.state

    d = state.d
    v = state.vsol[2]
    ω = state.ωsol[2]
    state.vsol[2] = szeros(T, 3)
    state.ωsol[2] = szeros(T, 3)
    dyn = dynamics(mechanism, body)
    state.vsol[2] = v
    ω = state.ωsol[2] = ω
    state.d = d

    b0 = D*dyn + friction.b # remove old friction force

    if norm(b0) > cf*γ
        friction.b = b0/norm(b0)*cf*γ
    else
        friction.b = b0
    end
    return
end


text:
#Copyright (c) 2015 Rex Computing and Isaac Yonemoto

#see LICENSE.txt

#this work was supported in part by DARPA Contract D15PC00135

#test-warlpiri-multiplication.jl

###############################################################################
##COMPLETE MULTIPLICATION TABLE

#column key
#none_ pfew_ pone_ psome ptwo_ pmany pall_ junk_ nnone n_few n_one nsome n_two nmany n_all
wmult = Utype[
#none_ * line
 none_ none_ none_ none_ none_ none_ junk_ junk_ none_ none_ none_ none_ none_ none_ junk_;
#pfew_ * line
 none_ pfew_ pfew_ pf_ps pf_ps pf_pm pall_ junk_ none_ n_few n_few ns_nf ns_nf nm_nf n_all;
#pone_ * line
 none_ pfew_ pone_ psome ptwo_ pmany pall_ junk_ none_ n_few n_one nsome n_two nmany n_all;
#psome * line
 none_ pf_ps psome ps_pm pmany pmany pall_ junk_ none_ ns_nf nsome nm_ns nmany nmany n_all;
#ptwo_ * line
 none_ pf_ps ptwo_ pmany pmany pmany pall_ junk_ none_ ns_nf n_two nmany nmany nmany n_all;
#pmany * line
 none_ pf_pm pmany pmany pmany pmany pall_ junk_ none_ nm_nf nmany nmany nmany nmany n_all;
#pall_ * line
 junk_ pall_ pall_ pall_ pall_ pall_ pall_ junk_ junk_ n_all n_all n_all n_all n_all n_all;
#junk_ * line
 junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_ junk_;
#nnone * line
 none_ none_ none_ none_ none_ none_ junk_ junk_ none_ none_ none_ none_ none_ none_ junk_;
#n_few * line
 none_ n_few n_few ns_nf ns_nf nm_nf n_all junk_ none_ pfew_ pfew_ pf_ps pf_ps pf_pm pall_;
#n_one * line
 none_ n_few n_one nsome n_two nmany n_all junk_ none_ pfew_ pone_ psome ptwo_ pmany pall_;
#nsome * line
 none_ ns_nf nsome nm_ns nmany nmany n_all junk_ none_ pf_ps psome ps_pm pmany pmany pall_;
#n_two * line
 none_ ns_nf n_two nmany nmany nmany n_all junk_ none_ pf_ps ptwo_ pmany pmany pmany pall_;
#nmany * line
 none_ nm_nf nmany nmany nmany nmany n_all junk_ none_ pf_pm pmany pmany pmany pmany pall_;
#n_all * line
 junk_ n_all n_all n_all n_all n_all n_all junk_ junk_ pall_ pall_ pall_ pall_ pall_ pall_
]

testop(*, wmult)

max_stars_repo_name: REX-Computing/unumjl

max_issues_repo_name: REX-Computing/unumjl

max_forks_repo_name: REX-Computing/unumjl


text:
using PyMNE
using PyCall
using Test

using PyCall: PyError

@testset "sklearn" begin
    PyMNE.install_sklearn()
    # requires sklearn
    @test PyMNE.preprocessing.ICA(method="fastica") isa PyObject
end

@testset "create_info and get_info" begin
    dat = zeros(1, 100)
    naive_info = PyMNE.mne.create_info([:a], 100)
    wrapped_info = PyMNE.create_info([:a], 100)
    raw = PyMNE.io.RawArray(dat, wrapped_info)
    @test raw.get_data() == dat
    @test get_info(raw) isa PyObject
    @test raw.info isa Dict
    @test_throws PyError PyMNE.io.RawArray(dat, naive_info)
end


text:
# This file is a part of SimilaritySearch.jl

using SimilaritySearch
export Kvp, k_near_and_far, fit, search, push!

struct Kvp{DataType<:AbstractVector,DistanceType<:PreMetric} <: AbstractSearchContext
    dist::DistanceType
    db::DataType
    refs::DataType
    sparsetable::Vector{Vector{Pair{Int32,Float32}}}
    ksparse::Int
end

Kvp(dist::PreMetric, db, refs, sparsetable, ksparse::Integer) =
    Kvp(dist, db, refs, sparsetable, ksparse)

Base.copy(kvp::Kvp;
        dist::PreMetric=kvp.dist,
        db::AbstractVector=kvp.db,
        refs::AbstractVector=kvp.refs,
        sparsetable::AbstractVector=kvp.sparsetable,
        ksparse::Integer=kvp.ksparse
    ) =
    Kvp(dist, db, refs, sparsetable, ksparse)

Base.string(p::Kvp) = "{Kvp: dist=$(p.dist), n=$(length(p.db)), refs=$(length(p.refs)), ksparse=$(p.ksparse)}"

"""
    k_near_and_far(dist::PreMetric, near::KnnResult, far::KnnResult, obj::T, refs::Vector{T}, k::Integer) where T

Searches for `k` near and far objects in the set of references
"""
function k_near_and_far(dist::PreMetric, near::KnnResult, far::KnnResult, obj::T, refs::Vector{T}, k::Integer) where T
    empty!(near, k)
    empty!(far, k)

    for refID in eachindex(refs)
        d = evaluate(dist, obj, refs[refID])
        push!(near, refID, d)
        push!(far, refID, -d)
    end

    row = [near[i] for i in eachindex(near)]

    for i in length(far):-1:1
        id_, dist_ = far[i]
        push!(row, id_ => -dist_)
    end

    row
end

"""
    Kvp(dist::PreMetric, db, refs, sparsetable, ksparse::Integer)
    Kvp(dist::PreMetric, db::AbstractVector, refs::AbstractVector, ksparse::Integer)
    Kvp(dist::PreMetric,
        db::AbstractVector;
        numpivots::Integer=ceil(Int, sqrt(length(db))),
        ksparse::Integer=ceil(Int, log2(length(db)))
        )
Creates a K vantage points index: a sparse pivot table storing only `ksparse` near and far pivots using references as pivots.

"""

function Kvp(dist::PreMetric, db::AbstractVector, refs::AbstractVector, ksparse::Integer)
    @info "Kvp, refs=$(typeof(db)), k=$(ksparse), numrefs=$(length(refs)), dist=$(dist)"
    near = KnnResult(ksparse)
    far = KnnResult(ksparse)

    sparsetable = [k_near_and_far(dist, near, far, db[i], refs, ksparse) for i in 1:length(db)]
    Kvp(dist, db, refs, sparsetable, ksparse)
end

function Kvp(dist::PreMetric,
        db::AbstractVector;
        numpivots::Integer=ceil(Int, sqrt(length(db))),
        ksparse::Integer=ceil(Int, log2(length(db)))
        )
    L = unique(rand(1:length(db), numpivots))
    Kvp(dist, db, db[L], ksparse)
end

"""
    search(kvp::Kvp, q::T, res::KnnResult) where T

Searches for `q` in the `kvp` index
"""
function search(kvp::Kvp, q::T, res::KnnResult) where T
    # d::Float64 = 0.0
    qI = [evaluate(kvp.dist, q, piv) for piv in kvp.refs]

    for i in eachindex(kvp.db)
        obj::T = kvp.db[i]
        objSparseRow = kvp.sparsetable[i]

        discarded = false
        @inbounds for (pivID, dop) in objSparseRow
            if abs(dop - qI[pivID]) > covrad(res)
                discarded = true
                break
            end
        end

        discarded && continue
        push!(res, i, evaluate(kvp.dist, q, obj))
    end

    res
end

##
## function push!(index::Kvp{T}, dist, obj::T) where T
##     push!(index.db, obj)
##     row = k_near_and_far(dist, obj, index.refs, index.k)
##     push!(index.sparsetable, row)
##     length(index.db)
## end


text:
# <: Real to make tracing easier. Maybe a bad idea?
struct Variable <: Expression
    name::Symbol
    value
    value_type::DataType
    subtype::Symbol
    diff::Union{AbstractOperator,Void}
    dependents::Vector{Variable}
    description::String
    flow::Bool
    domain::AbstractDomain
    size
    context
end

Variable(name,
         value = nothing,
         value_type = typeof(value);
         subtype::Symbol=:Variable,
         dependents::Vector{Variable} = Variable[],
         flow::Bool = false,
         description::String = "",
         domain = Reals(),
         size = nothing,
         context = nothing) =
         Variable(name,value,value_type,subtype,nothing,
                  dependents,description,flow,domain,size,context)
Variable(name,args...;kwargs...) = Variable(name,args...;subtype=:Variable,kwargs...)

Variable(name,x::Variable) = Variable(name,x.value,x.value_type,
                x.subtype,D,x.dependents,x.description,x.flow,x.domain,
                x.size,x.context)

Parameter(name,args...;kwargs...) = Variable(name,args...;subtype=:Parameter,kwargs...)
Constant(value::Number) = Variable(Symbol(value),value,typeof(value);subtype=:Constant)
Constant(name,args...;kwargs...) = Variable(name,args...;subtype=:Constant,kwargs...)
IndependentVariable(name,args...;kwargs...) = Variable(name,args...;subtype=:IndependentVariable,kwargs...)

function DependentVariable(name,args...;dependents = [],kwargs...)
    @assert !isempty(dependents)
    Variable(name,args...;subtype=:DependentVariable,dependents=dependents,kwargs...)
end

function StateVariable(name,args...;dependents = [],kwargs...)
    @assert !isempty(dependents)
    Variable(name,args...;subtype=:StateVariable,dependents=dependents,kwargs...)
end

function ControlVariable(name,args...;dependents = [],kwargs...)
    @assert !isempty(dependents)
    Variable(name,args...;subtype=:ControlVariable,dependents=dependents,kwargs...)
end

function JumpVariable(name,args...;dependents = [],kwargs...)
    @assert !isempty(dependents)
    Variable(name,args...;subtype=:JumpVariable,dependents=dependents,kwargs...)
end

function NoiseVariable(name,args...;dependents = [],kwargs...)
    @assert !isempty(dependents)
    Variable(name,args...;subtype=:NoiseVariable,dependents=dependents,kwargs...)
end

export Variable,Parameter,Constant,DependentVariable,IndependentVariable,JumpVariable,NoiseVariable,
       @Var, @DVar, @IVar, @Param, @Const

# Variables use isequal for equality since == is an Operation
function Base.:(==)(x::Variable,y::Variable)
    x.name == y.name && x.subtype == y.subtype && x.value == y.value &&
    x.value_type == y.value_type && x.diff == y.diff
end

function Base.:(==)(x::Variable,y::Number)
    x == Constant(y)
end

function Base.:(==)(x::Number,y::Variable)
    Constant(x) == y
end

function Base.Expr(x::Variable)
    if x.diff == nothing
        return :($(x.name))
    else
        return :($(Symbol("$(x.name)_$(x.diff.x.name)")))
    end
end

function Base.show(io::IO, A::Variable)
    if A.subtype == :Constant
        print(io,"Constant($(A.value))")
    else
        str = "$(A.subtype)($(A.name))"
        if A.value != nothing
            str *= ", value = " * string(A.value)
        end

        if A.diff != nothing
            str *= ", diff = " * string(A.diff)
        end

        print(io,str)
    end
end

extract_idv(eq) = eq.args[1].diff.x

function extract_elements(ops, targetmap, default = nothing)
    elems = Dict{Symbol, Vector{Variable}}()
    names = Dict{Symbol, Set{Symbol}}()
    if default == nothing
        targets =  unique(collect(values(targetmap)))
    else
        targets = [unique(collect(values(targetmap))), default]
    end
    for target in targets
        elems[target] = Vector{Variable}()
        names[target] = Set{Symbol}()
    end
    for op in ops
        extract_elements!(op, elems, names, targetmap, default)
    end
    Tuple(elems[target] for target in targets)
end
# Walk the tree recursively and push variables into the right set
function extract_elements!(op::AbstractOperation, elems, names, targetmap, default)
    for arg in op.args
        if arg isa Operation
                extract_elements!(arg, elems, names, targetmap, default)
        elseif arg isa Variable
            if default == nothing
                target = haskey(targetmap, arg.subtype) ? targetmap[arg.subtype] : continue
            else
                target = haskey(targetmap, arg.subtype) ? targetmap[arg.subtype] : default
            end
            if !in(arg.name, names[target])
                push!(names[target], arg.name)
                push!(elems[target], arg)
            end
        end
    end
end

# Build variables more easily
function _parse_vars(macroname, fun, x)
    ex = Expr(:block)
    lhss = Symbol[]
    # if parsing things in the form of
    # begin
    #     x
    #     y
    #     z = exp(2)
    # end
    x = flatten_expr!(x)
    for _var in x
        iscall = typeof(_var) <: Expr && _var.head == :call
        issym    = _var isa Symbol
        isassign = issym ? false : _var.head == :(=)
        @assert iscall || issym || isassign "@$macroname expects a tuple of expressions!\nE.g. `@$macroname x y z=1`"
        if iscall || issym
            if iscall
                dependents = :([$(_var.args[2:end]...)])
                var = _var.args[1]
            else
                dependents = Variable[]
                var = _var
            end
            lhs = var
            push!(lhss, lhs)
            expr = :( $lhs = $fun( Symbol($(String(lhs))) ,
                      dependents = $dependents))
        end
        if isassign
            iscall = typeof(_var.args[1]) <: Expr && _var.args[1].head == :call
            if iscall
                dependents = :([$(_var.args[1].args[2:end]...)])
                lhs = _var.args[1].args[1]
            else
                dependents = Variable[]
                lhs = _var.args[1]
            end
            rhs = _var.args[2]
            push!(lhss, lhs)
            expr = :( $lhs = $fun( Symbol($(String(lhs))) , $rhs,
                      dependents = $dependents))
        end
        push!(ex.args, expr)
    end
    push!(ex.args, Expr(:tuple, lhss...))
    ex
end

for funs in ((:DVar, :DependentVariable), (:IVar, :IndependentVariable),
             (:Var, :Variable),
             (:Param, :Parameter))
    @eval begin
        macro ($(funs[1]))(x...)
            esc(_parse_vars(String($funs[1]), $funs[2], x))
        end
    end
end

function _const_assign(x)
    ex = Expr(:block)
    lhss = Symbol[]
    x = flatten_expr!(x)
    for eq in x
        @assert eq isa Expr && eq.head == :(=) "@Const expects a tuple of assignments!\nE.g. `@Const D=t W=g`"
        lhs = eq.args[1]
        push!(lhss, lhs)
        rhs = eq.args[2]
        expr = :($lhs = Constant($rhs))
        push!(ex.args,  expr)
    end
    push!(ex.args, Expr(:tuple, lhss...))
    ex
end

macro Const(x...)
    esc(_const_assign(x))
end


text:
"
expands the bond dimension by adding random unitary vectors
"
@with_kw struct RandExpand<:Algorithm
    trscheme::TruncationScheme = truncdim(1)
end

function changebonds(state::InfiniteMPS,alg::RandExpand)
    #determine optimal expansion spaces around bond i
    pexp = PeriodicArray(map(1:length(state)) do i
        ACAR = _permute_front(state.AC[i])*_permute_tail(state.AR[i+1])
        AC2 = TensorMap(rand,eltype(ACAR),codomain(ACAR),domain(ACAR));

        #Calculate nullspaces for AL and AR
        NL = leftnull(state.AL[i])
        NR = rightnull(_permute_tail(state.AR[i+1]))

        #Use this nullspaces and SVD decomposition to determine the optimal expansion space
        intermediate = adjoint(NL)*AC2*adjoint(NR)
        (U,S,V) = tsvd(intermediate,trunc=alg.trscheme,alg=SVD())

        (NL*U,V*NR)
    end)

    newstate = copy(state);

    #do the actual expansion
    for i in 1:length(state)
        al = _permute_tail(catdomain(newstate.AL[i],pexp[i][1]))
        lz = TensorMap(zeros,_lastspace(pexp[i-1][1])',domain(al))
        newstate.AL[i] = _permute_front(catcodomain(al,lz))

        ar = _permute_front(catcodomain(_permute_tail(newstate.AR[i+1]),pexp[i][2]))
        rz = TensorMap(zeros,codomain(ar),space(pexp[i+1][2],1))
        newstate.AR[i+1] = catdomain(ar,rz)

        l = TensorMap(zeros,codomain(newstate.CR[i]),space(pexp[i][2],1))
        newstate.CR[i] = catdomain(newstate.CR[i],l)
        r = TensorMap(zeros,_lastspace(pexp[i][1])',domain(newstate.CR[i]))
        newstate.CR[i] = catcodomain(newstate.CR[i],r)

        newstate.AC[i] = newstate.AL[i]*newstate.CR[i]
    end

    return newstate
end

function changebonds(state::MPSMultiline,alg::RandExpand)
    #=
        todo : merge this with the MPSCentergauged implementation
    =#
    #determine optimal expansion spaces around bond i
    pexp = PeriodicArray(map(Iterators.product(1:size(state,1),1:size(state,2))) do (i,j)

        ACAR = _permute_front(state.AC[i-1,j])*_permute_tail(state.AR[i-1,j+1])
        AC2 = TensorMap(rand,eltype(ACAR),codomain(ACAR),domain(ACAR));

        #Calculate nullspaces for AL and AR
        NL = leftnull(state.AL[i,j])
        NR = rightnull(_permute_tail(state.AR[i,j+1]))

        #Use this nullspaces and SVD decomposition to determine the optimal expansion space
        intermediate = adjoint(NL)*AC2*adjoint(NR)
        (U,S,V) = tsvd(intermediate,trunc=alg.trscheme,alg=SVD())

        (NL*U,V*NR)
    end)

    newstate = copy(state);

    #do the actual expansion
    for (i,j) in Iterators.product(1:size(state,1),1:size(state,2))
        al = _permute_tail(catdomain(newstate.AL[i,j],pexp[i,j][1]))
        lz = TensorMap(zeros,_lastspace(pexp[i,j-1][1])',domain(al))

        newstate.AL[i,j] = _permute_front(catcodomain(al,lz))

        ar = _permute_front(catcodomain(_permute_tail(newstate.AR[i,j+1]),pexp[i,j][2]))
        rz = TensorMap(zeros,codomain(ar),space(pexp[i,j+1][2],1))
        newstate.AR[i,j+1] = catdomain(ar,rz)

        l = TensorMap(zeros,codomain(newstate.CR[i,j]),space(pexp[i,j][2],1))
        newstate.CR[i,j] = catdomain(newstate.CR[i,j],l)
        r = TensorMap(zeros,_lastspace(pexp[i,j][1])',domain(newstate.CR[i,j]))
        newstate.CR[i,j] = catcodomain(newstate.CR[i,j],r)

        newstate.AC[i,j] = newstate.AL[i,j]*newstate.CR[i,j]
    end

    return newstate
end

changebonds(state::Union{FiniteMPS,MPSComoving}, alg::RandExpand) = changebonds!(copy(state),alg)
function changebonds!(state::Union{FiniteMPS,MPSComoving},alg::RandExpand)
    for i in 1:(length(state)-1)
        ACAR = _permute_front(state.AC[i])*_permute_tail(state.AR[i+1])
        AC2 = TensorMap(rand,eltype(ACAR),codomain(ACAR),domain(ACAR));

        #Calculate nullspaces for left and right
        NL = leftnull(state.AC[i])
        NR = rightnull(_permute_tail(state.AR[i+1]))

        #Use this nullspaces and SVD decomposition to determine the optimal expansion space
        intermediate = adjoint(NL) * AC2 * adjoint(NR);
        (U,S,V) = tsvd(intermediate,trunc=alg.trscheme,alg=SVD())

        ar_re = V*NR;
        ar_le = TensorMap(zeros,codomain(state.AC[i]),space(V,1))

        (nal,nc) = leftorth(catdomain(state.AC[i],ar_le),alg=QRpos())
        nar = _permute_front(catcodomain(_permute_tail(state.AR[i+1]),ar_re));

        state.AC[i] = (nal,nc)
        state.AC[i+1] = (nc,nar)
    end

    state
end


text:
export Non_negative
export Map_0_1

 # All these functions return a mapping function and
 # the inverse of it
 # they are used to map the real numbers to non-negative real numbers

"""
    Non_negative()

Returns a function and a inverse function inverse function
to map numbers to non-negative numbers.
We use a parabola.

# Examples
```julia-repl
julia> p, p_inv = Non_negative()
(DeconvOptim.var"#5#7"(), DeconvOptim.var"#6#8"())

julia> x = [-1, 2, -3]
3-element Array{Int64,1}:
 -1
  2
 -3

julia> p(x)
3-element Array{Int64,1}:
 1
 4
 9

julia> p_inv(p(x))
3-element Array{Float64,1}:
 1.0
 2.0
 3.0
```
"""
function Non_negative()
    return x -> map(abs2, x) , (x -> sqrt.(x))
end

"""
    Map_0_1()

Returns a function and a inverse function inverse function
to map numbers to an interval between 0 and 1.

"""
function Map_0_1()
    f(x) = 1 .- exp.(.- x.^2)
    f_inv(y) = sqrt.(.- log.(1 .- y))
    return f, f_inv
end


text:
#=

Draw radar map in resulting world frame (after solve)

=#

##

using Distributed
# addprocs(10)

using Colors
using Caesar
@everywhere using Caesar

using DistributedFactorGraphs: listDataEntrySequence

import Rotations as _Rot
using CoordinateTransformations

# using DistributedFactorGraphs.DocStringExtensions
using Dates
using JSON2
# using BSON
using Serialization
using FixedPointNumbers
using Manifolds
using StaticArrays
using TensorCast
using DataStructures

using ImageMagick, FileIO
using Images
using ImageDraw
using ImageView
using ImageTransformations
using ImageFeatures

using VideoIO

using RoMEPlotting
Gadfly.set_default_plot_size(35cm,20cm)

using ProgressMeter

using Optim

## load the solved graph and reattach the datastores

dfg_datafolder = "/tmp/caesar/philos"

fg = loadDFG("$dfg_datafolder/results_5/x0_5_x280_stab")

##

ds = FolderStore{Vector{UInt8}}(:radar, "$dfg_datafolder/data/radar")
addBlobStore!(fg, ds)

ds = FolderStore{Vector{UInt8}}(:gps_fix, "$dfg_datafolder/data/gps")
addBlobStore!(fg, ds)

ds = FolderStore{Vector{UInt8}}(:lidar, "$dfg_datafolder/data/lidar")
addBlobStore!(fg, ds)

ds = FolderStore{Vector{UInt8}}(:camera, "$dfg_datafolder/data/camera")
addBlobStore!(fg, ds)

##

M = SpecialEuclidean(2)
e0 = identity_element(M)

##

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

include(joinpath(@__DIR__,"ImageOnlyStabilization.jl"))

##

lbls = sortDFG(ls(fg, tags=[:POSE;])) #[1:100]

_stabImageSeq(fg, lbls; filecount=20, threshold=0.2, kparg1=12, size=512, window=20, logentry=true, N=15, scale_r=0.8)
# _stabImageSeq(fg; filecount=11, threshold=0.2, kparg1=10, logentry=true)

##

saveDFG("$dfg_datafolder/results_5/x0_5_x280_stab", fg)

##

# count = 0
# let count = count
#   @showprogress 1 "sweep" for kp1=[12;], thr=[0.2;0.5;0.8;1.1;]
#     count += 1
#     _stabImageSeq(fg, lbls; filecount=count, threshold=thr, kparg1=kp1, logentry=true)
#   end
# end

##

# T = typeof(IMGS[1])
# IMGS_ = Vector{T}(undef, length(IMGS))
# IMGS_ .= IMGS
# writevideo("/tmp/caesar/philos/quickstab.mp4", IMGS_; fps=90)

##

# mosaicMatches(imgG[[1;3]]..., matches) |> imshow
# trans, matches = calcMatchMinimize(imgG[[1;2]]...); trans

##

# img1 = Gray.(imgs[1])
# img2 = Gray.(imgs[3])
# tr, matches = calcMatchMinimize(img1,img2; threshold=0.4, size=1024, window=50)
# mosaicMatches(img1, img2, matches) |> imshow

##

#


text:
module DynamicFactorModeling

######################
# Import packages
using LinearAlgebra
using Statistics
using Random
using Distributions
using PDMats
using PDMatsExtras
using ShiftedArrays
using Parameters
using Polynomials
using Kronecker

######################
# Include all package scripts
include("simulations/dgp.jl")
include("common_types/common_types.jl")
include("pca/pca_1level_estimator.jl")
include("pca/pca_2level_estimator.jl")
include("kim_nelson/kn_1level_estimator.jl")
include("kim_nelson/kn_2level_estimator.jl")
include("otrok_whiteman/ow_1level_estimator.jl")
include("otrok_whiteman/ow_2level_estimator.jl")
include("output_analysis/variance_decomposition.jl")
include("hyperparameter_est/hyperparam_est.jl")
include("test_scripts.jl")

######################
# Export package objects
export SSModel, HDFM, convertHDFMtoSS, simulateSSModel,
        DFMStruct, HDFMStruct,
        KN1LevelEstimator, KN2LevelEstimator,
        OW1LevelEstimator, OW2LevelEstimator,
        PCA1LevelEstimator, PCA2LevelEstimator,
        Hyperparam2LEstimator,
        vardecomp2level,
        sayhi
end


text:
using DataStructures
using NearestNeighbors

include("map.jl")

mutable struct Tree
    init_pose::Vector{Float64}
    params::parameters
    current_goal::Vector{Float64}
    all_nodes
    QrPrior::Stack{Int}
    Qr::Queue{Int}
    Qs::Queue{Int}
    # Passcode to prevent unneccessary looping
    passcode::Float64
    positions
    kdtree
    map::PlanningMap

    function Tree(init_pose::Vector{Float64}, params::parameters)
        println("Initializing RT-RRT* ...")
        root_node = TreeNode(1, init_pose, 0)
        QrPrior = Stack{Int}()
        Qr = Queue{Int}()
        Qs = Queue{Int}()
        passcode = rand()
        kdtree = KDTree(init_pose[1:2])
        positions = init_pose[1:2]
        map = PlanningMap(params)
        load_default_map(map, params)
        new(init_pose, params, [1, 1], [root_node], QrPrior, Qr, Qs, passcode, positions, kdtree, map)
    end
end

mutable struct TreeNode
    "ID in tree"
    id::Int64
    "Pose of the 2D robot"
    pose::Vector{Float64}
    "Parent Node"
    parent_id::Int64
    "Total cost = Parent Cost + edge_parent"
    totalcost::Float64
    "Edge cost"
    edge_cost::Float64
    "Child IDs"
    child_ids::Vector{Int64}
    "Passcode (has already been processed in rewire from root?)"
    passcode::Float64

    function TreeNode(id::Int64, pose::Vector{Float64}, parent_id::Int)
        new(id, pose, parent_id, 0, 0, [], 0)
    end
end

function steer(closest::TreeNode, pos::Vector{Float64}, params::parameters)
    delta = [pos[1] - closest.pose[1], pos[2] - closest.pose[2]]
    delta /= norm(delta)
    new_pose = [clamp(closest.pose[1] + delta[1] * params.nominal_step, params.x_range[1], params.x_range[2]),
                clamp(closest.pose[2] + delta[2] * params.nominal_step, params.y_range[1], params.y_range[2]),
                closest.pose[3] + atan(delta[2]/delta[1])]
    return new_pose
end

function near_nodes(this::Tree, node)
    "Naive Implementation, KD Trees will follow"
    r = 0.8
    nearnode_ids = inrange(this.kdtree, node.pose[1:2], r)
    return nearnode_ids
end

function expansion(this::Tree)
    start_time = time()
    while size(this.all_nodes, 1) <= this.params.max_tree_nodes
        sampled_pos = sample_point(this.params)
        # dist_, closest_ind = findmin(map(x->norm(this.all_nodes[x].pose[1:2] - sampled_pos), 1:size(this.all_nodes, 1)))
        # Nearest node obtained through KD Tree
        closest_ind, dist_ = knn(this.kdtree, sampled_pos, 1)
        closest_ind = closest_ind[1]
        new_pose = steer(this.all_nodes[closest_ind], sampled_pos, this.params)

        if !edge_collision(this.map, this.all_nodes[closest_ind].pose[1:2], new_pose[1:2])
            # Create new node
            new_node = TreeNode(size(this.all_nodes, 1) + 1, new_pose, closest_ind)
            new_node.edge_cost = calc_edge_cost(this.all_nodes[closest_ind].pose, new_pose)
            new_node.totalcost = this.all_nodes[closest_ind].totalcost + new_node.edge_cost
            push!(this.all_nodes[new_node.parent_id].child_ids, new_node.id)

            # nearnodes = near_nodes(this.all_nodes, new_node, this.params.gamma)
            nearnode_ids = near_nodes(this, new_node)
            c_min = new_node.totalcost
            new_parent_id = new_node.parent_id
            parent_changed = false
            for node_id in nearnode_ids
                c_new = this.all_nodes[node_id].totalcost + calc_edge_cost(this.all_nodes[node_id].pose, new_node.pose)
                if c_new < c_min
                    if !edge_collision(this.map, this.all_nodes[node_id].pose[1:2], new_node.pose[1:2])
                        new_parent_id = node_id
                        c_min = c_new
                        parent_changed = true
                    end
                end
            end

            push!(this.all_nodes, new_node)
            update_kdtree(this)

            if parent_changed
                change_parent(this, new_node.id, new_parent_id)
            else
                change_parent(this, new_node.id, new_node.parent_id)
            end

            push!(this.QrPrior, size(this.all_nodes, 1))
        end
    end
end

function rewire(this::Tree)
    start_time = time()
    while !(isempty(this.QrPrior) && isempty(this.Qr)) #&& (time() - start_time) <= this.params.RewireTime
        if !isempty(this.QrPrior)
            rewire_ind = pop!(this.QrPrior)
        else
            rewire_ind = dequeue!(this.Qr)
        end
        # nearnodes = near_nodes(this.all_nodes, this.all_nodes[rewire_ind], this.params.gamma)
        nearnode_ids = near_nodes(this, this.all_nodes[rewire_ind])
        # Boolean value for changing parent
        parent_changed = false

        for node_id in nearnode_ids
            # New costs when changing parent
            c_new = this.all_nodes[rewire_ind].totalcost + calc_edge_cost(this.all_nodes[rewire_ind].pose, this.all_nodes[node_id].pose)

            # update parent
            if c_new < this.all_nodes[node_id].totalcost
                if !edge_collision(this.map, this.all_nodes[node_id].pose[1:2], this.all_nodes[rewire_ind].pose[1:2])
                    change_parent(this, node_id, rewire_ind)
                    parent_changed = true
                    enqueue!(this.Qr, node_id)
                end
            end

            if parent_changed
                propagate_costs(this, node_id)
            end
        end
    end
end

function rewire_from_root(this::Tree)
    start_time = time()

    if isempty(this.Qs)
        enqueue!(this.Qs, 1)
    end
    while !isempty(this.Qs) && (time() - start_time) <= this.params.RewireTime
        rewire_ind = dequeue!(this.Qs)

        # nearnodes = near_nodes(this.all_nodes, this.all_nodes[rewire_ind], this.params.gamma)
        nearnode_ids = near_nodes(this, this.all_nodes[rewire_ind])

        for node_id in nearnode_ids
            # New costs when changing parent
            c_new = this.all_nodes[rewire_ind].totalcost + calc_edge_cost(this.all_nodes[rewire_ind].pose, this.all_nodes[node_id].pose)

            parent_changed = false
            # update parent
            if c_new < this.all_nodes[node_id].totalcost
                if !edge_collision(this.map, this.all_nodes[node_id].pose[1:2], this.all_nodes[rewire_ind].pose[1:2])
                    change_parent(this, node_id, rewire_ind)
                    parent_changed = true
                end
            end
            if this.all_nodes[node_id].passcode != this.passcode
                enqueue!(this.Qs, node_id)
                this.all_nodes[node_id].passcode = this.passcode
            end

            if parent_changed
                # Propagate costs
                propagate_costs(this, node_id)
            end
        end
    end
end

function propagate_costs(this::Tree, id::Int64)
    # Recursion
    if !isempty(this.all_nodes[id].child_ids)
        for child_id in this.all_nodes[id].child_ids
            # Update own costs
            this.all_nodes[child_id].totalcost = this.all_nodes[id].totalcost + this.all_nodes[child_id].edge_cost
            propagate_costs(this, child_id)
        end
    end
end

function change_parent(this::Tree, node_id, new_parent_id)
    # Update child id's of parent node
    setdiff!(this.all_nodes[this.all_nodes[node_id].parent_id].child_ids, node_id)
    this.all_nodes[node_id].parent_id = new_parent_id
    this.all_nodes[node_id].edge_cost = calc_edge_cost(this.all_nodes[new_parent_id].pose, this.all_nodes[node_id].pose)
    this.all_nodes[node_id].totalcost = this.all_nodes[new_parent_id].totalcost + this.all_nodes[node_id].edge_cost
    push!(this.all_nodes[new_parent_id].child_ids, node_id)
end

function update_costs(this::Tree, node_id::Int64)
    parent_id = this.all_nodes[node_id].parent_id
    this.all_nodes[node_id].edge_cost = calc_edge_cost(his.all_nodes[parent_id].pose, this.all_nodes[node_id].pose)
    this.all_nodes[nodes_id].totalcost = this.all_nodes[parent_id].totalcost + this.all_nodes[node_id].edge_cost
end

function update_kdtree(this::Tree)
    this.positions = cat(this.positions, this.all_nodes[end].pose[1:2], dims=2)
    this.kdtree = KDTree(this.positions)
end

function extract_goal_path(this::Tree)
    closest_ind, dist_ = knn(this.kdtree, this.current_goal, 1)
    ind = closest_ind[1]

    goal_path = []
    while true
        push!(goal_path, this.all_nodes[ind].pose)
        if this.all_nodes[ind].parent_id == 0
            break
        end
        ind = this.all_nodes[ind].parent_id
    end
    return goal_path
end

function update_root_node(this::Tree, dx::Float64, dy::Float64)
    new_pose = this.all_nodes[1].pose + [dx, dy, 0]
    new_pose[1] = clamp(new_pose[1], this.params.x_range[1], this.params.x_range[2])
    new_pose[2] = clamp(new_pose[2], this.params.y_range[1], this.params.y_range[2])
    if !edge_collision(this.map, this.all_nodes[1].pose[1:2], new_pose[1:2])
        this.all_nodes[1].pose = new_pose
        this.positions[:, 1] = new_pose[1:2]
        this.kdtree = KDTree(this.positions)

        # change costs of root node children and propagate
        for node_id in this.all_nodes[1].child_ids
            if !edge_collision(this.map, this.all_nodes[node_id].pose[1:2], this.all_nodes[1].pose[1:2])
                change_parent(this, node_id, 1)
                propagate_costs(this, node_id)
            else
                # Assign very high cost to colliding node
                this.all_nodes[node_id].totalcost = Inf

                # Propagate to avoid cycles in tree
                propagate_costs(this, node_id)
                nearnode_ids = near_nodes(this, this.all_nodes[node_id])

                for node_id_ in nearnode_ids
                    c_new = this.all_nodes[node_id_].totalcost + calc_edge_cost(this.all_nodes[node_id_].pose, this.all_nodes[node_id].pose)

                    if !edge_collision(this.map, this.all_nodes[node_id].pose[1:2], this.all_nodes[node_id_].pose[1:2]) && c_new < this.all_nodes[node_id].totalcost
                        change_parent(this, node_id, node_id_)
                        propagate_costs(this, node_id)
                        break
                    end
                end
            end

        end

        empty!(this.Qs)
        this.passcode = rand()
    end
end


text:
function hide(x) x end
function g() :: Int64 hide("foo") end
g() # <-- error: cannot convert

max_stars_repo_name: gabriel-doriath-dohler/pjulia

max_issues_repo_name: gabriel-doriath-dohler/pjulia

max_forks_repo_name: gabriel-doriath-dohler/pjulia


text:
using JLD
"""GraphSerializer is a type for custom serialization into JLD files.
It has no use except on disk. This type supports JLD.writeas(g::Graph)
and JLD.readas(gs::GraphSerializer). It is a form of Compressed Sparse Column format
of the adjacency matrix of a graph g.

Fields
======
vertices: the UnitRange of vertex numbers
ne: the number of edges in g
packed_adjlist: the concatenation of fadj into a 1D array
n_adjlist: the degree of each vertex. use cumsum(n_adjlist) to get the CSC offsets.

This type compacts the adjacency matrix into a 1D structure so that it can be represented on disk.
We do not use CSC/CSR for graphs because we need amortized O(1) insertion, and CSC provides O(n)
insertion because you have to shift entries in memory.

This format is fine for storage on disk because the data is not changing.
JLD.readas(gs::GraphSerializer) creates the adjacency list representation directly instead of calling add_edge!
repeatedly in an attempt to improve performance.

This type has not been tested with mmaped files or compression in JLD.
"""
type GraphSerializer
    vertices::UnitRange{Int}
    ne::Int
    packed_adjlist::Vector{Int}
    n_adjlist::Vector{Int} #stores the end offset into packed_adjlist
end

function JLD.writeas(g::Graph)
    n_adjlist = zeros(Int,nv(g))
    @assert sum(degree(g))/2 == ne(g)
    packed_adjlist = Array(Int, 2*ne(g))
    k = 0
    degree(g), sum(degree(g))
    for (i,lst) in enumerate(g.fadjlist)
        #create the offsets
        if i != 1
            n_adjlist[i] = n_adjlist[i-1] + length(lst)
        else
            n_adjlist[1] = length(lst)
        end

        # pack the neighbors
        for v in lst
            packed_adjlist[k+=1] = v
        end
    end
    GraphSerializer(g.vertices, g.ne, packed_adjlist, n_adjlist)
end

function JLD.readas(gs::GraphSerializer)
    n = length(gs.vertices)
    adj = Vector{Vector{Int}}(n)
    @assert length(adj) == n
    for i in gs.vertices
        if i == 1
            posbegin = 1
        else
            posbegin = gs.n_adjlist[i-1] +1
        end
        posend   = gs.n_adjlist[i]
        adj[i] = gs.packed_adjlist[posbegin:posend]
    end
    @assert sum(map(length, adj)) == 2gs.ne
    g = Graph(1:n, gs.ne, adj)
    return g
end

type Network{G,V,E}
    graph::G
    vprop::Vector{V}
    eprop::Dict{Edge,E}
end

import Base.==

==(n::Network, m::Network) = (n.graph == m.graph) && (n.vprop == m.vprop) && (n.eprop == m.eprop)


text:
using MolecularGraph
using ..ChemistryFeaturization.Codec

@testset "SpeciesFeatureDescriptor" begin
    # because if it doesn't work on caffeine then why bother...
    caffeine = smilestomol("CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
    ag = AtomGraph(caffeine)

    # because tea smells good...
    sfd = SpeciesFeatureDescriptor("isaromatic")
    @test all(get_value(sfd, ag) .== isaromatic(caffeine))
    @test length(encodable_elements(sfd)) == 118

    # test the other codec case...
    sfd = SpeciesFeatureDescriptor("lonepair")
    @test all(get_value(sfd, ag) .== lonepair(caffeine))
    encoded = encode(sfd, ag)
    @test all(encoded[1, :] .== encoded[5, :] .== 0)
    @test all(encoded[:, 8] .== encoded[:, 11] .== [0, 0, 0, 1, 0])

    # build one "from scratch"
    ag = AtomGraph(Float32.([0 1; 1 1]), ["H", "O"])
    num_nbs = g -> first.(length.(neighbors.(Ref(g), 1:nv(g))))
    codec = OneHotOneCold(true, [1, 2, 3, 4])
    categorical = true
    ee = ["O", "H"]
    sfd = SpeciesFeatureDescriptor{SimpleWeightedGraph,OneHotOneCold}(
        "num_neighbors",
        num_nbs,
        codec,
        categorical,
        ee,
    )
    @test get_value(sfd, ag) == [1, 2]
    @test encode(sfd, ag) == [1 0; 0 1; 0 0; 0 0]
    # add one to verify encodable_elements check happens
end


text:
mutable struct Out_Layer <: NN
    r1::Residual
    w::PType4 #weights
    b::PType4#biases
    ms::KnetMoment #batch norm moments
    bn_p::PType1 #batch norm parameters
    training::Bool
end

function (c::Out_Layer)(x::HGType)
    r1=c.r1(x)
    out1 = conv4(c.w,r1); myfree(r1);
    out2 = out1 .+ c.b; myfree(out1);
    bn=batchnorm(out2,c.ms,c.bn_p,training=c.training); myfree(out2);
    out3 = relu.(bn);

     myfree(bn)

    out3
end

function Out_Layer(in_dim::Int,atype=KnetArray)
    r1=Residual(in_dim,in_dim,atype)
    bn_p = bnparams(Float32,in_dim)
    bn_p = Param(convert(atype{Float32,1},bn_p))
    ms = bnmoments()

    w = xavier_normal(Float32,1,1,in_dim,in_dim)
    b = xavier_normal(Float32,1,1,in_dim,1)
    w = Param(convert(atype,w))
    b = Param(convert(atype,b))
    Out_Layer(r1,w,b,ms,bn_p,true)
end

function set_testing(f::Out_Layer,training=false)
    set_testing(f.r1,training)
    f.training=training
    nothing
end

struct FirstBlock <: NN
    c1::Conv1
    r1::Residual_skip
    p1::Pool
    r2::Residual
    r3::Residual_skip
end
FirstBlock(N::Int,atype=KnetArray)=FirstBlock(Conv1(3,64,7,2,3,atype),Residual_skip(64,128,atype),Pool(),Residual(128,128,atype),Residual_skip(128,N,atype))

function (f::FirstBlock)(x::HGType)
    c1=f.c1(x)
    r1=f.r1(c1); myfree(c1);
    p1=f.p1(r1); myfree(r1);
    r2=f.r2(p1); myfree(p1);
    r3=f.r3(r2); myfree(r2)

    r3
end

function set_testing(f::FirstBlock,training=false)
    set_testing(f.c1,training)
    set_testing(f.r1,training)
    set_testing(f.r2,training)
    set_testing(f.r3,training)
    nothing
end

struct Hourglass <: NN
    n::Int64
    up1::Residual
    pool1::Pool
    low1::Residual
    low2::Union{Hourglass,Residual} #This could be a residual or another Hourglass
    low3::Residual
    up2::Unpool
end

function Hourglass(n::Int,f_num::Int,atype=KnetArray)
    up1 = Residual(f_num,f_num,atype)
    pool1 = Pool()
    low1 = Residual(f_num,f_num,atype)
    if n > 1
        low2 = Hourglass(n-1,f_num,atype)
    else
        low2 = Residual(f_num,f_num,atype)
    end
    low3 = Residual(f_num,f_num,atype)
    up2 = Unpool()
    Hourglass(n,up1,pool1,low1,low2,low3,up2)
end

function (h::Hourglass)(x::HGType)

    up1 = h.up1(x)
    pool1 = h.pool1(x)
    low1 = h.low1(pool1); myfree(pool1);
    low2 = h.low2(low1); myfree(low1);
    low3 = h.low3(low2); myfree(low2);
    up2 = h.up2(low3); myfree(low3);
    out = up1 .+ up2

    myfree(up1); myfree(up2)

    out
end

function set_testing(f::Hourglass,training=false)
    set_testing(f.up1,training)
    set_testing(f.low1,training)
    set_testing(f.low2,training)
    set_testing(f.low3,training)
    nothing
end

struct HG2 <: NN
    nstack::Int64
    fb::FirstBlock
    hg::Array{Hourglass,1}
    o1::Array{Out_Layer,1}
    c1::Array{Conv0,1}
    merge_features::Array{Conv0,1}
    merge_preds::Array{Conv0,1}
end

#=
N = Number of Channels (64 default)
K = number of features for prediction
nstack = number of hourglasses
=#
function HG2(N::Int,K::Int,nstack::Int,atype=KnetArray)
    fb = FirstBlock(N,atype)

    hg=[Hourglass(4,N,atype) for i=1:nstack];
    o1=[Out_Layer(N,atype) for i=1:nstack];
    c1=[Conv0(N,K,1,1,0,atype) for i=1:nstack];
    merge_features=[Conv0(N,N,1,1,0,atype) for i=1:(nstack-1)]
    merge_preds=[Conv0(K,N,1,1,0,atype) for i=1:(nstack-1)]
    HG2(nstack,fb,hg,o1,c1,merge_features,merge_preds)
end

features(hg::HG2)=size(hg.merge_preds[1].w,3)

function (h::HG2)(x::HGType)
    temp=h.fb(x)

    preds=Array{typeof(temp),1}() #Can this be typed to be the same as input?
    temps=Array{typeof(temp),1}(undef,h.nstack) #Can this be typed to be the same as input?
    temps[1]=temp
    for i=1:h.nstack
        hg=h.hg[i](temps[i])
        features=h.o1[i](hg)
        pred=h.c1[i](features)
        push!(preds,pred)
        if i<h.nstack
            m_features = h.merge_features[i](features)
            m_preds = h.merge_preds[i](pred)
            temp1 = m_features + m_preds
            temps[i+1] = temp1 + temps[i]
            myfree(m_features); myfree(m_preds); myfree(temp1)
        end

        myfree(hg); myfree(features);
    end
    myfree(temp)
    for i=1:h.nstack
        myfree(temps[i])
    end

    preds
end

function (h::HG2)(x,y)
    preds=h(x)
    loss=0.0f0
    for i=1:h.nstack
        loss += pixel_mse(y,preds[i])
    end
    loss
end
(h::HG2)(d::Knet.Data) = mean(h(x,y) for (x,y) in d)

function set_testing(f::HG2,training=false)
    set_testing(f.fb,training)
    for i=1:f.nstack
        set_testing(f.hg[i],training)
        set_testing(f.o1[i],training)
    end
    nothing
end


text:
# vanderPol orbit diagram and winding number versus ω
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot

function stroboscopic_plot_winding(ds,ω,y_init,Ttr,nattr,nwind)
    period = 2π/ω
    ds.p[3] = period
    T = period*nattr   # no. of periods required for transient
    Ttr = period*ceil(Ttr/period)     # no. of periods used for averaging

    tr = trajectory(ds, T, y_init; Δt = period, Ttr)
    u1ret = tr[:, 1]

    # compute winding number
    Δt = period/20
    tr = trajectory(ds, nwind*period, y_init; Δt, Ttr)
    u1, u2 = columns(tr)
    θ = 0.0
    θ_old = atan(u2[1], u1[1])
    for (x, y) ∈ zip(u1, u2)
        θ_new = atan(y, x)
        θ += mod(θ_new - θ_old + π, 2π) - π
        θ_old = θ_new
    end
    W = abs(θ / nwind*2π)
    return tr[end], W, u1ret
end # of function

# -------- main -----

d = 5.0
a = 1.0
Ttr = 1000      # minimal transient time
nattr = 1000    # no. of points on the attractor
nwind = 100    # averaging of winding number
y_init = SVector(0.1, 0.1)
ds = Systems.vanderpol(y_init; μ = d, F = a)

# compute and plot orbit diagram
# ------------------------------
ω_min = 0.2
ω_max = 2.
nω =  100 # 800
ω_vec = LinRange(ω_min, ω_max, nω)
wind_no_vec = zeros(nω)  # winding numbers
u1s = []

pcol = COLORS[1]
for iω = 1:length(ω_vec)
    @show iω
    ω = ω_vec[iω]
    y_init, wind_no, u1 =  stroboscopic_plot_winding(ds,ω,y_init,Ttr,nattr,nwind)
    wind_no_vec[iω] = 1 / wind_no
    push!(u1s, u1)
end # of frequency loop

# zoom
ω_min = 0.8
ω_max = 1.0
ω_vec_zoom = LinRange(ω_min, ω_max, nω)
wind_no_vec_zoom = zeros(nω)  # winding numbers
u1s_zoom = []

for iω = 1:length(ω_vec_zoom)
    @show iω
    ω = ω_vec_zoom[iω]
    y_init, wind_no, u1 =  stroboscopic_plot_winding(ds,ω,y_init,Ttr,nattr,nwind)
    wind_no_vec_zoom[iω] = 1 / wind_no
    push!(u1s_zoom, u1)
end # of frequency loop

# %% plot
fig, axs = subplots(2,2; figsize = (figx, 1.5figy))

j = 0
for (uvec, ωvec) in zip((u1s, u1s_zoom), (ω_vec, ω_vec_zoom))
    j += 1
    for (i, ω) in enumerate(ωvec)
        axs[1,j].scatter(fill(ω, length(uvec[i])), uvec[i], s = 1, c = "C0", alpha = 0.2)
    end
    w = (wind_no_vec, wind_no_vec_zoom)[j]
    axs[2,j].plot(ω_vec, w; color = "C2")
    axs[1,j].set_xticklabels([])
    axs[1,j].set_xlim(ωvec[1], ωvec[end])
    axs[2,j].set_xlim(ωvec[1], ωvec[end])
end

axs[1,1].

subplot(2,2,1)
xlabel(L"\ω")
ylabel(L"u_1")

ax = gca()
ax.set_xlim([ω_min,ω_max])

# plot winding number
subplot(2,2,3)
plt.plot(ω_vec,wind_no_vec,c = "k")

xlabel(L"\ω")
ylabel(L"W")
ax = gca()
ax.set_xlim([ω_min,ω_max])

xlabel(L"\ω")
ylabel(L"u_1")

ax = gca()
ax.set_xlim([ω_min,ω_max])

# plot winding number
subplot(2,2,4)
plt.plot(ω_vec,wind_no_vec,c = "k")

xlabel(L"\ω")
ylabel(L"W")
ax = gca()
ax.set_xlim([ω_min,ω_max])

add_identifiers!(fig)
fig.tight_layout(pad=0.3)

# Output
# ------
# dname = "/Users/parlitz/ownCloud/NonlinearDynamicsTextbook/figure_generation/9/van_der_Pol/orbit_diagram/"
# fname = "van_der_Pol_orbit_diagram_vs_frequency_wind_no_v4_d=5_a=1_v5"
# ftype = ".png"

# savefig(dname*fname*ftype)

max_stars_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_issues_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_forks_repo_name: JuliaDynamics/NonlinearDynamicsTextbook


text:
# using Revise for auto-updating packages/functions
try
    using Revise
catch e
    @warn "Error initializing Revise" exception=(e, catch_backtrace())
end

max_stars_repo_name: nhejazi/mydots

max_issues_repo_name: nhejazi/mydotfiles

max_forks_repo_name: nhejazi/mydotfiles