text:
data, evts = loadtestdata("test_case_3a") #
f = @formula 0 ~ 1 + conditionA + continuousA # 1

# prepare data
data_r = reshape(data, (1, :))
data_r = vcat(data_r, data_r)#add second channel

#--------------------------#
## Mass Univariate Linear ##
#--------------------------#
data_e, times = Unfold.epoch(data = data_r, tbl = evts, τ = (-1.0, 1.9), sfreq = 20) # cut the data into epochs

# test manual pathway
uf = UnfoldLinearModel(Dict(Any => (f, times)))
designmatrix!(uf, evts; eventcolumn = "type")
fit!(uf, data_e)

@test typeof(uf.modelfit) == Unfold.LinearModelFit
@test !isempty(coef(uf.modelfit))

# test "automatic", non-dictionary call
m_mul = coeftable(fit(UnfoldLinearModel, f, evts, data_e, times))

@test m_mul[(m_mul.channel.==1).&(m_mul.time.==0.1), :estimate] ≈ [2, 3, 4]

# test Autodetection
@test Unfold.designToModeltype(Dict(Any => (@formula(0 ~ 1), 0:10))) == UnfoldLinearModel
@test Unfold.designToModeltype(Dict(Any => (@formula(0 ~ 1 + A), 0:10))) ==
      UnfoldLinearModel
@test Unfold.designToModeltype(
    Dict(Any => (@formula(0 ~ 1 + A), firbasis(τ = (-1, 1), sfreq = 20, name = "basisA"))),
) == UnfoldLinearModelContinuousTime
@test Unfold.designToModeltype(Dict(Any => (@formula(0 ~ 1 + (1 | test)), 0:10))) ==
      UnfoldLinearMixedModel
@test Unfold.designToModeltype(
    Dict(
        Any => (
            @formula(0 ~ 1 + (1 | test)),
            firbasis(τ = (-1, 1), sfreq = 20, name = "basisA"),
        ),
    ),
) == UnfoldLinearMixedModelContinuousTime

data_e_noreshape, times = Unfold.epoch(data = data, tbl = evts, τ = (-1.0, 1.9), sfreq = 20) # cut the data into epochs
m_mul_noreshape = coeftable(fit(UnfoldLinearModel, f, evts, data_e_noreshape, times))

@test m_mul_noreshape[
    (m_mul_noreshape.channel.==1).&(m_mul_noreshape.time.==0.1),
    :estimate,
] ≈ [2, 3, 4]
@test size(m_mul_noreshape)[1] == size(m_mul)[1] / 2

# Add Missing in Data
data_e_missing = data_e
data_e_missing[1, 25, end-5:end] .= missing
m_mul_missing = coeftable(Unfold.fit(UnfoldLinearModel, f, evts, data_e_missing, times))

@test m_mul_missing.estimate ≈ m_mul.estimate

# Special solver solver_lsmr_se with Standard Error
se_solver = solver = (x, y) -> Unfold.solver_default(x, y, stderror = true)
m_mul_se = coeftable(Unfold.fit(UnfoldModel, f, evts, data_e, times, solver = se_solver))
@test all(m_mul_se.estimate .≈ m_mul.estimate)
@test !all(isnothing.(m_mul_se.stderror))

# robust solver
rob_solver = (x, y) -> Unfold.solver_robust(x, y)#,rlmOptions=(initial_coef=zeros(3 *length(times)),))
data_outlier = copy(data_e)
data_outlier[:,31,1] .= 1000
m_mul_rob = coeftable(Unfold.fit(UnfoldModel, f, evts, data_outlier, times, solver = rob_solver))
ix = findall(m_mul_rob.time .≈ 0.5)
@test all(m_mul_rob.estimate .≈ m_mul.estimate)

m_mul_outlier = coeftable(Unfold.fit(UnfoldModel, f, evts, data_outlier, times))

#---------------------------------#
## Timexpanded Univariate Linear ##
#---------------------------------#
basisfunction = firbasis(τ = (-1, 1), sfreq = 20, name = "basisA")
m_tul = coeftable(fit(UnfoldModel, f, evts, data_r, basisfunction))

@test isapprox(
    m_tul[(m_tul.channel.==1).&(m_tul.time.==0.1), :estimate],
    [2, 3, 4],
    atol = 0.01,
)

# test without reshape, i.e. 1 channel vector e.g. size(data) = (1200,)
m_tul_noreshape = coeftable(fit(UnfoldModel, f, evts, data, basisfunction));
@test size(m_tul_noreshape)[1] == size(m_tul)[1] / 2

# Test under missing data
data_missing = Array{Union{Missing,Number}}(undef, size(data_r))
data_missing .= data_r
data_missing[4500:4600] .= missing

m_tul_missing = coeftable(fit(UnfoldModel, f, evts, data_missing, basisfunction))

@test isapprox(m_tul_missing.estimate, m_tul.estimate, atol = 1e-4)  # higher tol because we remove stuff

## Test multiple basisfunctions
b1 = firbasis(τ = (-1, 1), sfreq = 20, name = "basisA")
b2 = firbasis(τ = (-1, 1), sfreq = 20, name = "basisB")

f1 = @formula 0 ~ 1 + continuousA # 1
f2 = @formula 0 ~ 1 + continuousA # 1

# Fast-lane new implementation
res = coeftable(
    fit(
        UnfoldModel,
        Dict(0 => (f1, b1), 1 => (f2, b2)),
        evts,
        data_r,
        eventcolumn = "conditionA",
    ),
)

# slow manual
X1 = designmatrix(
    UnfoldLinearModelContinuousTime,
    f1,
    filter(x -> (x.conditionA == 0), evts),
    b1,
)
X2 = designmatrix(
    UnfoldLinearModelContinuousTime,
    f2,
    filter(x -> (x.conditionA == 1), evts),
    b2,
)
uf = UnfoldLinearModelContinuousTime(Dict(), X1 + X2, [])
@time fit!(uf, data_r)
tmp = coeftable(uf)

# test fast way & slow way to be identical
@test all(tmp.estimate .== res.estimate)

# runntime tests - does something explode?
for k = 1:4
    local f
    if k == 1
        f = @formula 0 ~ 1
    elseif k == 2
        f = @formula 0 ~ 1 + conditionA
    elseif k == 3
        f = @formula 0 ~ 0 + conditionA
    elseif k == 4
        f = @formula 0 ~ 1 + continuousA
    end

    fit(UnfoldModel, f, evts, data_e, times)
    fit(UnfoldModel, f, evts, data, basisfunction)
end

# Special solver solver_lsmr_se with Standard Error
se_solver = solver = (x, y) -> Unfold.solver_default(x, y, stderror = true)
m_tul_se = coeftable(fit(UnfoldModel, f, evts, data_r, basisfunction, solver = se_solver))
@test all(m_tul_se.estimate .== m_tul.estimate)
@test !all(isnothing.(m_tul_se.stderror))

#m_mul_se,m_mul_se = fit(UnfoldLinearModel,f,evts,data_e.+randn(size(data_e)).*5,times,solver=se_solver)
#plot(m_mul_se[m_mul_se.channel.==1,:],se=true)
#m_tul_se,m_tul_se = fit(UnfoldLinearModel,f,evts,data_r.+randn(size(data_r)).*5,basisfunction,solver=se_solver)
#plot(m_tul_se[m_tul_se.channel.==1,:],se=true)

##
data_long, evts_long = loadtestdata("test_case_1c") #
data_long = reshape(data_long, (1, :))

data_long = vcat(data_long, data_long)
f_long = @formula 0 ~ 1
basisfunction_long = firbasis(τ = (-1, 1), sfreq = 1000, name = "basisA")

@time designmatrix(UnfoldLinearModel, f_long, evts_long, basisfunction_long)
@time m_tul_long =
    coeftable(fit(UnfoldModel, f_long, evts_long, data_long, basisfunction_long))

@test isapprox(
    m_tul_long[(m_tul_long.channel.==1).&(m_tul_long.time.==0.1), :estimate],
    [2],
    atol = 0.0001,
)
# ~21s, test_case_1c, sfreq = 1000, 6000 events
#
## Older numbers for "designmatrix" only. But see "benchmark/benchmarkjl" for better benchmarks
# new version 7s-10s, dataset4, sfreq=1000, 1200stim,
# ~13s, test_case_1c, sfreq = 1000, 6000 events

###############################
##  Mixed Model tests
###############################
data, evts = loadtestdata("testCase3", dataPath = (@__DIR__) * "/data") #
append!(data, zeros(1000))
data = reshape(data, 1, :)
data = vcat(data, data)
data = data .+ 1 * randn(size(data)) # we have to add minimal noise, else mixed models crashes.
data_missing = Array{Union{Missing,Number}}(undef, size(data))
data_missing .= data

data_missing[4500:4600] .= missing

transform!(evts, :subject => categorical => :subject)

f = @formula 0 ~ 1 + condA + condB + (1 + condA + condB | subject)
#f  = @formula 0~1 + (1|subject)

# cut the data into epochs
# TODO This ignores subject bounds
data_e, times = Unfold.epoch(data = data, tbl = evts, τ = (-1.0, 1.9), sfreq = 10)
data_missing_e, times =
    Unfold.epoch(data = data_missing, tbl = evts, τ = (-1.0, 1.9), sfreq = 10)
evts_e, data_e = Unfold.dropMissingEpochs(copy(evts), data_e)
evts_missing_e, data_missing_e = Unfold.dropMissingEpochs(copy(evts), data_missing_e)

######################
##  Mass Univariate Mixed
@time m_mum = fit(
    UnfoldModel,
    f,
    evts_e,
    data_e,
    times,
    contrasts = Dict(:condA => EffectsCoding(), :condB => EffectsCoding()),
)
df = Unfold.coeftable(m_mum)
@test isapprox(
    df[(df.channel.==1).&(df.coefname.=="condA: 1").&(df.time.==0.0), :estimate],
    [5.618, 9.175],
    rtol = 0.1,
)

# with missing
@time m_mum = fit(
    UnfoldModel,
    f,
    evts_missing_e,
    data_missing_e,
    times,
    contrasts = Dict(:condA => EffectsCoding(), :condB => EffectsCoding()),
)
df = coeftable(m_mum)
@test isapprox(
    df[(df.channel.==1).&(df.coefname.=="condA: 1").&(df.time.==0.0), :estimate],
    [5.618, 9.175],
    rtol = 0.1,
)

# Timexpanded Univariate Mixed
f = @formula 0 ~ 1 + condA + condB + (1 + condA + condB | subject)
basisfunction = firbasis(τ = (-0.2, 0.3), sfreq = 10, name = "ABC")
@time m_tum = fit(
    UnfoldModel,
    f,
    evts,
    data,
    basisfunction,
    contrasts = Dict(:condA => EffectsCoding(), :condB => EffectsCoding()),
)
df = coeftable(m_tum)
@test isapprox(
    df[(df.channel.==1).&(df.coefname.=="condA: 1").&(df.time.==0.0), :estimate],
    [5.618, 9.175],
    rtol = 0.1,
)

# missing data in LMMs
# not yet implemented
Test.@test_broken m_tum = fit(
    UnfoldModel,
    f,
    evts,
    data_missing,
    basisfunction,
    contrasts = Dict(:condA => EffectsCoding(), :condB => EffectsCoding()),
)

evts.subjectB = evts.subject;
evts1 = evts[evts.condA.==0, :]
evts2 = evts[evts.condA.==1, :]

f0_lmm = @formula 0 ~ 1 + condB + (1 | subject) + (1 | subjectB)
@time m_tum = coeftable(fit(UnfoldModel, f0_lmm, evts, data, basisfunction))

f1_lmm = @formula 0 ~ 1 + condB + (1 | subject)
f2_lmm = @formula 0 ~ 1 + condB + (1 | subjectB)

b1 = firbasis(τ = (-0.2, 0.3), sfreq = 10, name = "A")
b2 = firbasis(τ = (-0.1, 0.3), sfreq = 10, name = "B")

X1_lmm = designmatrix(UnfoldLinearMixedModel, f1_lmm, evts1, b1)
X2_lmm = designmatrix(UnfoldLinearMixedModel, f2_lmm, evts2, b2)

r = fit(UnfoldLinearMixedModelContinuousTime, X1_lmm + X2_lmm, data);
df = coeftable(r)

@test isapprox(
    df[(df.channel.==1).&(df.coefname.=="condB").&(df.time.==0.0), :estimate],
    [18.21, 17.69],
    rtol = 0.1,
)

# Fast-lane new implementation
m = coeftable(
    fit(
        UnfoldModel,
        Dict(0 => (f1_lmm, b1), 1 => (f2_lmm, b2)),
        evts,
        data,
        eventcolumn = "condA",
    ),
)

if 1 == 0
    using WGLMakie, AlgebraOfGraphics
    m = mapping(
        :time,
        :estimate,
        color = :colname,
        layout_x = :colname,
        layout_y = :basisname,
    )

    AlgebraOfGraphics.data(df[df.channel.==2, :]) * visual(Lines) * m |> draw

end

## Condense check for multi channel, multi
@testset "LMM multi channel, multi basisfunction" begin
    data,evts = loadtestdata("testCase3", dataPath = (@__DIR__) * "/data")
    transform!(evts,:subject=>categorical=>:subject);
    data = vcat(data',data')

	bA0 = firbasis(τ=(-0.0,0.1),sfreq=10,name="bA0")
	bA1 = firbasis(τ=(0.1,0.2),sfreq=10,name="bA1")
	evts.subject2 = evts.subject
	fA0 = @formula 0~1+condB + zerocorr(1|subject)
	fA1  =@formula 0~1+condB + zerocorr(1|subject2)
	m = fit(UnfoldModel,
		Dict(0=>(fA0,bA0),
			 1=>(fA1,bA1)),
		evts,data,eventcolumn="condA")

	res = coeftable(m)

    @test all(last(.!isnothing.(res.group),8))
    @test all(last(res.coefname,8).=="(Intercept)")
end

##------

if 1 == 0
    # Fit mass-univariate 1st level for all subjects
    basisfunction = firbasis(τ = (-.1, 0.4), sfreq = 10, name = "A")
    resAll = DataFrame()
    f = @formula 0 ~ 1 + condA + condB
    for s in unique(evts.subject)
        ##
        from = minimum(evts[evts.subject.==s, :latency]) - 10
        to = maximum(evts[evts.subject.==s, :latency]) + 40
        to = min(to, size(data, 1))
        evts_s = evts[evts.subject.==s, :]
        evts_s.latency .-= from
        m = fit(
            UnfoldLinearModel,
            f,
            evts_s,
            data[from:to],
            basisfunction,
            contrasts = Dict(:condA => EffectsCoding(), :condB => EffectsCoding()),
        )
        m.results.subject = s
        append!(resAll, m.results)
    end

    results = resAll[resAll.coefname.=="(Intercept)", :]

    results[results.time.==0.1, :]
end

max_stars_repo_name: unfoldtoolbox/jlunfold

max_issues_repo_name: unfoldtoolbox/jlunfold

max_forks_repo_name: unfoldtoolbox/jlunfold


text:
import Base: union, union!, setdiff, setdiff!, intersect!, intersect

struct AVLSet{K} <: AbstractSet{K}
    tree::AVLTree{K,Nothing}
end

AVLSet() = AVLSet{Any}(AVLTree{Any,Nothing}())
AVLSet{K}() where {K} = AVLSet{K}(AVLTree{K,Nothing}())

function AVLSet(x::K) where {K <: AbstractVector}
    t = AVLTree{eltype(x),Nothing}()
    for i in x
        insert!(t, i, nothing)
    end
    return AVLSet{eltype(x)}(t)
end

Base.eltype(::Type{AVLSet{K}}) where {K} = K
Base.length(set::AVLSet) = length(set.tree)
Base.in(x::K, set::AVLSet{K}) where {K} = x in set.tree

function iterate(set::AVLSet{K}) where {K}
    ret = iterate(set.tree)
    if ret === nothing return nothing else return (ret[1][1], ret[2])  end
end

function iterate(set::AVLSet{K}, node::Node{K,Nothing}) where {K}
    ret = iterate(set.tree, node)
    if ret === nothing return nothing else return (ret[1][1], ret[2])  end
end

Base.push!(set::AVLSet{K}, item::K) where {K} = insert!(set.tree, item, nothing)

Base.delete!(set::AVLSet{K}, item) where {K} = delete!(set.tree, item)

Base.union(set::AVLSet{K}, sets...) where {K} = union!(deepcopy(set), sets...)

function Base.union!(set::AVLSet{K}, sets...) where {K}
    (key -> push!.(Ref(set), key)).(sets)
    return set
end

Base.setdiff(set::AVLSet{K}, sets...) where {K} = setdiff!(deepcopy(set), sets...)

function Base.setdiff!(set::AVLSet{K}, sets...) where {K}
    (key -> delete!.(Ref(set), key)).(sets)
    return set
end

Base.intersect(set::AVLSet{K}, s::AbstractSet) where {K} = intersect!(deepcopy(set), s)

function Base.intersect!(set::AVLSet{K}, s::AbstractSet) where {K}
    _set = collect(set)
    for key in _set
        if key ∉ s
            delete!(set, key)
        end
    end
    return set
end

ext: jl


text:
using GLMakie

const p1 = Node(Vector{Point2f0}(undef, 6000)) # 5 min
const y_label = Node("")

function plot_sin(x)
    y_label[]="sin"
    y = sin.(x)
    p1[] =  Point2f0.(x, y)
end

function plot_cos(x)
    y_label[]="cos"
    y = cos.(x)
    p1[] =  Point2f0.(x, y)
end

function main(gl_wait=true)
    fig=Figure()
    ax=Axis(fig[1, 1], xlabel = "time [s]", ylabel = y_label)
    x_1 = 0f0:0.1f0:10f0*pi
    x_2 = 0:0.1:5*pi
    plot_sin(x_1)
    po = lines!(p1)

    @async begin
        for i in 0:2
            plot_sin(x_1)
            autolimits!(ax)
            sleep(2)
            plot_cos(x_2)
            autolimits!(ax)
            sleep(2)
        end
    end

    gl_screen = display(fig)
    if gl_wait
        wait(gl_screen)
    end
    return nothing
end

main()


text:
# This file was generated, do not modify it. # hide
params(forest) |> pprint


text:
#=
em.jl

    Expectation-Maximization (EM) algorithm to estimate the hyper parameters of
    a linear Gaussian State Space model.

@author: Quint Wiersma <q.wiersma@vu.nl>

@date: 2022/02/09
=#

# Expectation-Maximization
struct EMState{Tv, Tm} <: EMOptimizerState
	ψ::Tv		    # current state
	ψ_prev::Tv	    # previous state
    ψ_prev_m::Tv    # M step update buffer
	Δ::Tv		    # change in state
    V_sum::Tm       # buffer for E step sum Vₜ
    V_0::Tm         # buffer for E step V₀
    V_1::Tm         # buffer for E step V₋₁
    V_01::Tm        # buffer for E step V₀₋₁
end

# Expectation-Conditional Maximization
struct ECMState{Tv, Tm} <: EMOptimizerState
	ψ::Tv		    # current state
	ψ_prev::Tv	    # previous state
	Δ::Tv		    # change in state
    V_sum::Tm       # buffer for E step sum Vₜ
    V_0::Tm         # buffer for E step V₀
    V_1::Tm         # buffer for E step V₋₁
    V_01::Tm        # buffer for E step V₀₋₁
end

"""
    em!(model, pen; init=NamedTuple(), ϵ=1e-4, max_iter=1000)

Expectation-Maximization (EM) algorithm to estimate the hyper parameters of a
linear Gaussian State Space model as defined by `model`, storing the results in
`model`.

#### Arguments
  - `model::StateSpaceModel`: state space model
  - `pen::NamedTuple`       : penalization parameters
  - `method::Symbol`        : filtering method
  - `init::NamedTuple`      : initial model parameters
  - `ϵ::Real`               : tolerance
  - `max_iter::Integer`     : max number of iterations

#### Returns
  - `ll::Real`  : log-likelihood value
"""
function em!(model::StateSpaceModel, pen::NamedTuple;
            method::Symbol=:univariate, init::NamedTuple=NamedTuple(),
            ϵ::Real=1e-4, max_iter::Integer=1000)
    # Initialize state space model and system
    sys= init!(model, init, method)

    # Get dims
    (n,T_len)= size(model.y)
    p= length(sys.a1)
    n_params= nparams(model)

    # Initialize parameters
    ψ= similar(model.y, n_params)
    get_params!(ψ, model)

    # Initialize state
    state= EMState(ψ, similar(ψ), similar(ψ), similar(ψ), similar(ψ, p, p),
                    similar(ψ, p, p), similar(ψ, p, p), similar(ψ, p, p))

    # Initialize filter and smoother
    T= eltype(model.y)  # type
    if method === :univariate
        filter= UnivariateFilter(n, p, T_len, T)
    elseif method === :collapsed
        filter= UnivariateFilter(p, p, T_len, T)
    elseif method === :multivariate
        filter= MultivariateFilter(n, p, T_len, T)
    elseif method === :woodbury
        filter= WoodburyFilter(n, p, T_len, T)
    else
        throw(ArgumentError("Invalid method name $(method)"))
    end
    smoother= Smoother(p, 1, T_len, T)

    # Initialize relative change
    rel_change= Inf
    # Initialize iteration counter
    iter= 1
    # EM algorithm
    while rel_change > ϵ && iter < max_iter
        # Store current parameters
        copyto!(state.ψ_prev, state.ψ)

        # E-step
        init_system!(sys, model)
        get_system!(sys, model, method)
        estep!(smoother, filter, sys)

        # M-step
        mstep!(state, model, smoother, pen)

        # Store change in state
        state.Δ.= state.ψ .- state.ψ_prev

        # Relative change
        rel_change= norm(state.Δ, Inf) * inv(1 + norm(state.ψ, Inf))

        # Update iteration counter
        iter+=1
    end

    return loglik(filter, model, method)
end

"""
    ecm!(model, pen; init=NamedTuple(), ϵ=1e-4, max_iter=1000)

Expectation-Conditional Maximization (ECM) algorithm to estimate the hyper
parameters of a linear Gaussian State Space model as defined by `model`, results
are stored in `model`.

#### Arguments
  - `model::StateSpaceModel`: state space model
  - `pen::NamedTuple`       : penalization parameters
  - `method::Symbol`        : filtering method
  - `init::NamedTuple`      : initial model parameters
  - `ϵ::Real`               : tolerance
  - `max_iter::Integer`     : max number of iterations

#### Returns
  - `ll::Real`  : log-likelihood value
"""
function ecm!(model::StateSpaceModel, pen::NamedTuple;
                method::Symbol=:univariate, init::NamedTuple=NamedTuple(),
                ϵ::Real=1e-4, max_iter::Integer=1000)
    # Initialize state space model and system
    sys= init!(model, init, method)

    # Get dims
    (n,T_len)= size(model.y)
    p= length(sys.a1)
    n_params= nparams(model)

    # Initialize parameters
    ψ= similar(model.y, n_params)
    get_params!(ψ, model)

    # Initialize state
    state= ECMState(ψ, similar(ψ), similar(ψ), similar(ψ, p, p), similar(ψ, p, p),
                    similar(ψ, p, p), similar(ψ, p, p))

    # Initialize filter and smoother
    T= eltype(model.y)  # type
    if method === :univariate
        filter= UnivariateFilter(n, p, T_len, T)
    elseif method === :collapsed
        filter= UnivariateFilter(p, p, T_len, T)
    elseif method === :multivariate
        filter= MultivariateFilter(n, p, T_len, T)
    elseif method === :woodbury
        filter= WoodburyFilter(n, p, T_len, T)
    else
        throw(ArgumentError("Invalid method name $(method)"))
    end
    smoother= Smoother(p, 1, T_len, T)

    # Initialize relative change
    rel_change= Inf
    # Initialize iteration counter
    iter= 1
    # EM algorithm
    while rel_change > ϵ && iter < max_iter
        # Store current parameters
        copyto!(state.ψ_prev, state.ψ)

        # E-step
        init_system!(sys, model)
        get_system!(sys, model, method)
        estep!(smoother, filter, sys)

        # Conditional M-step
        mstep!(state, model, smoother, pen)

        # Store change in state
        state.Δ.= state.ψ .- state.ψ_prev

        # Relative change
        rel_change= norm(state.Δ, Inf) * inv(1 + norm(state.ψ, Inf))

        # Update iteration counter
        iter+=1
    end

    return loglik(filter, model, method)
end


text:
using Flux
using Flux: @treelike
using Flux.Tracker: data
using LinearAlgebra: tril!

struct MultiheadAttention
    head::Int
    future::Bool
    iqproj::Dense
    ikproj::Dense
    ivproj::Dense
    oproj::Dense
    drop::Dropout
end

@treelike MultiheadAttention

"""
    MultiheadAttention(head::Int, is::Int, hs::Int, os::Int; future::Bool=true, pdrop = 0.1)

Multihead dot product Attention Layer, `head` is the number of head, `is` is the input size, `hs` is the hidden size of input projection layer of each head,
`os` is the output size. When `future` is `false`, the k-th token can't see tokens at > k. `pdrop` is the dropout rate.
"""
MultiheadAttention(head::Int,
                   is::Int,
                   hs::Int,
                   os::Int;
                   future::Bool=true, pdrop = 0.1) = MultiheadAttention(head,
                                                                        future,
                                                                        Dense(is, hs*head),
                                                                        Dense(is, hs*head),
                                                                        Dense(is, hs*head),
                                                                        Dense(hs*head, os),
                                                                        Dropout(pdrop),
                                                                        )

function Base.show(io::IO, mh::MultiheadAttention)
    hs = div(size(mh.iqproj.W)[1], mh.head)
    is = size(mh.iqproj.W)[end]
    os = size(mh.oproj.W)[1]

    print(io, "MultiheadAttention(")
    print(io, "head=$(mh.head), ")
    print(io, "head_size=$(hs), ")
    print(io, "$(is)=>$(os)")

    if mh.drop.active
        print(io, ", dropout=$(mh.drop.p))")
    else
        print(io, ")")
    end
end

function (mh::MultiheadAttention)(query::Abstract3DTensor{T},
                                  key::Abstract3DTensor{T},
                                  value::Abstract3DTensor{T};
                                  mask=nothing) where T
    qs = size(query)
    ks = size(key)
    vs = size(value)

    #size(ipq) == (h, q_seq_len, batch)
    ipq = @toNd mh.iqproj(query)
    ipk = @toNd mh.ikproj(key)
    ipv = @toNd mh.ivproj(value)

    h = size(ipq, 1)
    hs = div(h, mh.head)

    #size(ipq) == (hs, q_seq_len, head, batch)
    ipq = permutedims(reshape(ipq, hs, mh.head, qs[2], qs[3]), [1, 3, 2, 4])
    ipk = permutedims(reshape(ipk, hs, mh.head, ks[2], ks[3]), [1, 3, 2, 4])
    ipv = permutedims(reshape(ipv, hs, mh.head, vs[2], vs[3]), [1, 3, 2, 4])

    #size(ipq) == (hs, q_seq_len, head * batch)
    ipq = reshape(ipq, hs, qs[2], :)
    ipk = reshape(ipk, hs, ks[2], :)
    ipv = reshape(ipv, hs, vs[2], :)

    atten = attention(ipq,ipk,ipv;
                      mask=mask,
                      future=mh.future,
                      dropout=mh.drop)

    atten = permutedims(reshape(atten, hs, qs[2], mh.head, qs[3]), [1, 3, 2, 4]) #size(atten) == (hs, head, ql, b)
    atten = reshape(atten, h, qs[2], qs[3]) #size(atten) == (h, ql, b)

    out = @toNd mh.oproj(atten)
    out #size(out) == (h, q_seq_len, batch)
end

function (mh::MultiheadAttention)(query::AbstractMatrix{T},
                                  key::AbstractMatrix{T},
                                  value::AbstractMatrix{T};
                                  mask=nothing) where T
    # size(query) == (dims, seq_len)
    ipq = mh.iqproj(query)
    ipk = mh.ikproj(key)
    ipv = mh.ivproj(value)

    h = size(ipq)[1] #h == hs * head
    hs = div(h, mh.head)

    #size(hq) == (hs, seq_len, head)
    hq = permutedims(reshape(ipq, hs, mh.head, :), [1, 3, 2])
    hk = permutedims(reshape(ipk, hs, mh.head, :), [1, 3, 2])
    hv = permutedims(reshape(ipv, hs, mh.head, :), [1, 3, 2])

    atten = attention(hq, hk, hv;
                      mask=mask,
                      future=mh.future,
                      dropout=mh.drop)

    # size(atten) == (head*hs, seq_len)
    atten = reshape(permutedims(atten, [1, 3, 2]), h, :)

    mh.oproj(atten)
end

function attention(query::AbstractMatrix{T},
                   key::AbstractMatrix{T},
                   value::AbstractMatrix{T};
                   mask=nothing, future::Bool = false,
                   dropout=nothing) where T
    # size(query) == (dims, {q,k}_seq_len) == size(key) == size(value)
    # size(score) == (k_seq_len, q_seq_len)
    dk = size(key)[1]
    score = transpose(key) * query
    score = score ./ convert(T, sqrt(dk))

    if mask !== nothing
        @. mask = (1 - mask) * convert(T, -1e9)
        score = score .+ mask
    end

    if !future
        fmask = tril!(fill!(similar(score), convert(T, -1e9)), -1)
        score = score .+ fmask
    end

    score = softmax(score)
    dropout !== nothing && (score = dropout(score))
    value * score #size(return) == (dims, q_seq_len)
end

function attention(query::Abstract3DTensor{T},
                   key::Abstract3DTensor{T},
                   value::Abstract3DTensor{T};
                   mask=nothing, future::Bool = false,
                   dropout=nothing) where T
    #size(query) == (dims, {q,k}_seq_len, batch) == size(key) == size(value)
    #size(score) == (k_seq_len, q_seq_len, batch)
    dk = size(key, 1)
    score = batchedmul(key, query; transA = true)
    score = score ./ convert(T, sqrt(dk))

    s = size(score)

    if mask !== nothing
        #weird issue on @. mask = (1 - mask) * -1e9 which casue mask to be -Inf
        mask = (1 .- mask) .* convert(T, -1e9)
        ms = size(mask)
        #score = score .+ mask; use broadcast instead of repeat mask for head
        score = reshape(reshape(score, s[1:end-1]..., :, ms[end]) .+ reshape(mask, ms[1:end-1]..., 1, ms[end]), s)
    end

    if !future
        #without ... will cause data move back to cpu
        fmask = tril!(fill!(similar(score, s[1:end-1]...), convert(T, -1e9)), -1)
        score = score .+ fmask
    end

    score = @toNd softmax(score) #reshape(softmax(reshape(score, s[1], :)) , s)
    dropout !== nothing && (score = dropout(score))
    batchedmul(value, score) #size(return) == (dims, q_seq_len, batch)
end


text:
using DataFlow:Call, syntax, vertex, constant
include("ops.jl")
vcall(a...) = vertex(Call(), constant.(a)...)

#function creat_graph(a::Array{Any, 1})
struct graph
    input
    output
end

function get_basic_graph(a::Array{Any, 1})
    res = Array{graph, 1}()
    for x=1:length(a)
        if x==1
            temp = graph(new_type(:Input, Dict{Any, Any}()), a[x])
        else
            temp = graph(a[x-1], a[x])
        end
        push!(res, temp)
    end
    return res
end

function create_graph(a::Array{Any, 1})
    b = Array{Any, 1}()
    for x=1:length(a)
        push!(b, ops[a[x].layer_type](a[x].fields))
    end
    return b
end

vcall_chain(a...) = vertex(Call(), :Chain, constant.(a)...)

#Create an array of ops from Keras.load_layers
function get_ops(a::Array{Any, 1})
    res = Array{Any, 1}()
    for ele in a
        push!(res, ops[ele.layer_type](ele))
    end
    return res
end

#Chainify get_ops
function chainify(a::Array{Any, 1}, ip)
    res = ip
    for ele in a
        if typeof(ele) <: Tuple
            temp = vcall(:Chain, ele[1], ele[2])
        else
            temp = vcall(:Chain, ele)
        end
        res = vcall(temp, res)
    end
    return res
end

#function write_chainify(ll)
#    text="Chain("
#    for ele in ll
#        text= text*"$(ele.fields["name"]) ,"
#    end
#    text = text*")"
#    write!("(x,)-> begin \n $text \n end")
#end

function load(structure_file, weight_file, ip...)
    if check_modeltype(structure_file) == "Sequential"
        global weight = weights(weight_file)
        s = load_structure(structure_file)
        l = load_layers(s)
        go = get_ops(l)
        if isempty(ip)
            return go
        else
            return go(ip[1])
        end
    elseif check_modeltype(structure_file) == "Model"
        s = load_structure(structure_file)["layers"]
        filter!(x->x["class_name"]!="InputLayer", s)
        if isempty(ip)
            error("Cannot call non-Sequential models without input")
        else
            ll = load_layers(load_structure(structure_file)["layers"])
            return graphify(ll, structure_file, weight_file, ip[1])
        end
    end
end

function (m::Array{Any, 1})(x)
    return chainify(m, x) |> syntax |> eval
end

#"""
#Takes load_layer as input
#"""
#function graphify(m::Array{Any, 1})
#    global weight = weights("resnet.h5")
#    res = Dict{Any, Any}()
#    for ele in m
#        if ele.layer_type == :InputLayer
#            res[ele.fields["name"]] = :ip
#        else
#            res[ele.fields["name"]] = vcall(ops[ele.layer_type](ele), res[ele.input_nodes[1][1][1]])
#        end
#    end
#    return res
#end


text:
using Test
using FileIO

@testset "VulkanExamples.jl" begin
    examples_dir = joinpath(dirname(@__DIR__), "examples")
    @testset "Headless" begin
        include(joinpath(examples_dir, "headless", "headless.jl"))
        dst = joinpath(examples_dir, "headless", tempname() * ".png")
        try
            main(
                dst,
                Point2f[(-0.5, -0.5), (0.5, -0.5), (0.5, 0.5), (-0.5, 0.5)],
                [
                    RGBA{Float16}(0.0, 0.0, 1.0, 0.0),
                    RGBA{Float16}(0.0, 1.0, 0.0, 0.05),
                    RGBA{Float16}(1.0, 1.0, 1.0, 0.4),
                    RGBA{Float16}(1.0, 0.0, 0.0, 1.0),
                ],
            )
            @test stat(dst).size > 1
        finally
            ispath(dst) && rm(dst)
        end
    end
    @testset "Texture" begin
        include(joinpath(examples_dir, "texture", "texture_2d.jl"))
        dst = joinpath(examples_dir, "texture", tempname() * ".png")
        try
            main(
                dst,
                Point2f[(-1., -1.), (1., -1.), (-1., 1.), (1., 1.)],
                joinpath(examples_dir, "texture", "texture_2d.png"),
                width = 1024,
                height = 1024,
            )
            @test stat(dst).size > 1
        finally
            ispath(dst) && rm(dst)
        end
    end
end


text:
mutable struct DEOptions{SType,CType}
    saveat::SType
    save_on::Bool
    save_everystep::Bool
    callback::CType
end

mutable struct ODEInterfaceIntegrator{F,algType,uType,uPrevType,oType,SType,solType,P,CallbackCacheType} <: DiffEqBase.AbstractODEIntegrator{algType, true, uType, Float64}
    f::F
    u::uType
    uprev::uPrevType
    t::Float64
    tprev::Float64
    p::P
    opts::oType
    u_modified::Bool
    tdir::Float64
    sizeu::SType
    sol::solType
    eval_sol_fcn
    event_last_time::Int
    vector_event_last_time::Int
    callback_cache::CallbackCacheType
    alg::algType
    last_event_error::Float64
end

@inline function (integrator::ODEInterfaceIntegrator)(t,deriv::Type{Val{N}}=Val{0};idxs=nothing) where N
  @assert N==0 "ODEInterface does not support dense derivative"
  sol = integrator.eval_sol_fcn(t)
  return idxs == nothing ? sol : sol[idxs]
end


text:
function gen_cpp_wrapaccessors(modname::String, typename::String, fieldinfos; indent::Int = 4)
    indentstr = repeat(' ', indent)
    lines = [ "$modname.add_type<$typename>(\"$typename\")" ]
    for (T, V, n) in fieldinfos
        push!(lines, "$indentstr.method(\"$(typename)_get_$(n)\", [](const $(typename)& x) { return x.$(n); })")
        push!(lines, "$indentstr.method(\"$(typename)_set_$(n)\", []($(typename)& x, const $V v) { x.$(n) = v; })")
    end
    return join(lines, '\n') * ';'
end


text:
### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 7cfa427a-9251-11eb-3b66-d313a367496d
md"# Projects 2021"

# ╔═╡ 4987617e-9252-11eb-0351-e76aec0c1812
md"""## Groups

### French speaking

#### Group A

* Bryan Frédéric Fonteyn
* Mohamed Yacine Bouzaienne
* Lionel Eddy De Laet
* Louis-Henri Gustave Laquay
* Hugo Van Loo

#### Group B

* Victoria Nathalie Duhez
* Mazarine Catherine Remy-Paquay
* Alexandre Baudoin de Merode
* Florian Van Hoeck
* Noémie Nadia Van Roosbroeck

#### Group C

* Hugo Jules Carpiaux
* Adrien Dumont
* Paul Charles Lecocq
* Alina Marie Schütz

#### Combinations

* 5 groups of one person from group A and one person from group B
* 2 groups of 2 persons from group C

### Dutch speaking

#### Group A

* Cloë Jean Bouchez
* Mathias Van der Stricht
* Johannes Albrecht
* Joeri Stefan Winckelmans
* Toon Ivens
* Younes Sabori
* Cloë Theys
* Arthur Lots
* Lennart Ché Wanzeele
* Tim De Smet

#### Group B

* Wout Van den Bosch
* Pieterjan Frank Van Riet
* Emile Van Vaerenbergh
* Stan Vandekeere
* Robbe Stein	Berten
* Ben Dana De Groot
* Marion Monique Decaesteker-Baelen
* Bram Steven Janssens
* Tomas Ellen Oostvogels
* Pieter Van der Veken
* Lander Van Eeckhout

#### Combinations

* 9 groups of one person from group A and one person from group B
* 1 group of one person from group A and two persons from group B
"""

# ╔═╡ 6a896c5e-9253-11eb-2f93-a9faf3ba8385
md"""## Games

### Pen and paper games

* [Paper soccer](https://en.wikipedia.org/wiki/Paper_soccer)
* [Quantum tic-tac-toe](https://en.wikipedia.org/wiki/Quantum_tic-tac-toe)
* [Mastermind](https://en.wikipedia.org/wiki/Mastermind_(board_game))
* [Dots and boxes](https://en.wikipedia.org/wiki/Dots_and_Boxes)
* [Racetrack](https://en.wikipedia.org/wiki/Racetrack_(game))

### [Logic games](https://www.chiark.greenend.org.uk/~sgtatham/puzzles/)

* Palisade
* Singles
* Filling
* Unequal
* Tracks
* Range

### Smart games

* [Jungle hide and seek](https://www.smartgames.eu/nl/try-smartgames-online/1-speler-spellen/jungle-hide-seek)
* [Anti-virus](https://www.smartgames.eu/nl/try-smartgames-online/1-speler-spellen/anti-virus)
"""

# ╔═╡ 924df842-9259-11eb-380c-61f6d9928b7c
md"""## Daily Work

* Implement the game
* Include documentation
* Prepare a Pluto notebook with demo
* Explain the datastructures
* Detail the created functions
* Post your code and notebook on BelADL before 10th of May
"""

# ╔═╡ 0ed42ce4-925a-11eb-1555-ebecc574440a
md"""## Exam

* Implement a bot playing the game automatically
* Based on your own strategy
* Prepare a Pluto notebook with demo
* Explain the implemented strategy
* Detail the datastructures and the logic
* Post your code and notebook on BelADL before the day of the exam
"""

# ╔═╡ Cell order:
# ╟─7cfa427a-9251-11eb-3b66-d313a367496d
# ╟─4987617e-9252-11eb-0351-e76aec0c1812
# ╟─6a896c5e-9253-11eb-2f93-a9faf3ba8385
# ╟─924df842-9259-11eb-380c-61f6d9928b7c
# ╟─0ed42ce4-925a-11eb-1555-ebecc574440a


text:
# PiecewiseLinearContinuousFitting
# Piecewise Linear Continouous Interpolation Constraints
# Sparse L2 Optimal Fitting subject Continuity Constraints
# Integrated Square

using Test

# TODO what is this?
#@everywhere include(joinpath(dirname(@__FILE__),"..","src","jl"))

@everywhere using Polynomials, IterTools#, Plots

testN = 9
@everywhere t1 = time()
#maxlengths = Array{Int64,1}(undef, testN) ;        maxlengths[1] = 1
#times = Array{Float64,1}(undef, testN);            times[1] = 0.
#sizes = Array{Array{Int64,2},1}(undef, testN);     sizes[1] = reshape([1],(1,1))
#gr()
j = 9
#@everywhere function testi(j)
    N = j
    t = range(0, stop=2π, length=50)
    #g = Poly([0,0,1,-1])
    ran = 0.1*randn(N)
    z(t) = ran[floor(Int64, t/2pi*(N-1))+1]
    g = sin
    g2(t) = g(t) + z(t)
    ℓ = compute_transition_costs(g2, t);
    K = 10
    #@time I, Y, f = brute_force_search(ℓ, K-1);
    Λ_0 = [create_new_pwq(minimize_wrt_x2(ℓ[i, N])) for i in 1:N-1];
    Λ, t, _, _, _ = @timed pwq_dp_constrained(Λ_0, ℓ, N-1);
    println("$j, tot: $(time()-t1), $t")
    times = t
    sizes = reshape([isassigned(Λ, i) ? length(Λ[i]) : 0 for i in eachindex(Λ)], size(Λ))
    maxlengths = maximum(sizes)
    #return times, sizes, maxlengths
#end
@time result = pmap(testi, 300:testN)
times, sizes, maxlengths = collect(zip(result...))
times = [times...]; sizes=[sizes...]; maxlengths=[maxlengths...];

function plotall(Λ)
    M,N = size(Λ)
    p = plot(layout=(M-2,N-1), size=(1400,1000))
    k = 1
    for m = (M-2):-1:1
        for n = 1:N-1
            idx = (n-1)*M+m
            if isassigned(Λ,idx)
                x, y, x_all, y_all = get_vals(Λ[m,n])
                for i in eachindex(x_all)
                    plot!(p, y_all[i][[1,end]], x_all[i][[1,end]], l = 0, m=(1,:cross,:orange), subplot=k, lab="")
                    plot!(p, y_all[i], x_all[i], l=(1,:blue), lab="", subplot=k)
                end
                plot!(p, y, x, l=:red, subplot=k, lab="$m, $n, $k", ylims=(-4,4), xlims=(0,6))
            end
            k += 1
        end
    end
    return p
end

plot(times)
plot!(collect(500:512),times3*2)
plot!(((1:512).^3.3)*maximum(times)/300^3.3)
plot!(ylims=(0,1400))
plot!()
heatmap!(sizes[end])

plot!(maximum(times)*sum.(sizes)/maximum(sum.(sizes)))

using Plots
plot(t, g.(t), lab="g(t)")
for k = 7:7
    @time I2, y2, f2 = recover_solution(Λ[k, 1], 1, N)
    Y2, _ = find_optimal_y_values(ℓ, I2)

    println("Comparison: ", sqrt(f), " --- ", sqrt(f2))

    #find_optimal_y_values

    l = zeros(size(Λ))
    for i=1:size(l,1)
        for j=1:size(l,2)
            if isassigned(Λ,i,j)
                l[i,j] = length(Λ[i, j])
            end
        end
    end

    #x = range(0, stop=1, length=100)

    #closefig()
    plot!(t[I2],Y2, m=:circle, lab="k=$k")
end
plot!()
#using PyPlot
# #close()
# figure(1)
# plot(t, g.(t), "g", t[I2], Y2, "bo-")
#
# figure(2)
# plot(1:length(t)-1, l')


text:
#  Copyright 2016-2019 Eric S. Tellez <eric.tellez@infotec.mx>
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

import Base:
    push!, popfirst!, pop!, length, last, first, empty!

export Item, KnnResult, maxlength, covrad, NnResult

struct Item{T}
    objID::T
    dist::Float64
end

mutable struct KnnResult{T}
    k::Int
    pool::Vector{Item{T}}

    function KnnResult(T::Type, k)
        v = Vector{Item{T}}()
        sizehint!(v, k)
        new{T}(k, v)
    end

    KnnResult(k::Integer) = KnnResult(Int, k)
end

NnResult() = KnnResult(1)

"""
fix_order! fixes the sorted state of the array. It implements a kind of insertion sort
It is efficient due to the expected distribution of the items being inserted
(few smaller than the ones already inside)
"""
@inline function fix_order!(res::KnnResult{T}) where T
    arr::Vector{Item{T}} = res.pool

    item = arr[end]
    i = length(arr)
    @inbounds while i > 1
        if item.dist < arr[i-1].dist
            arr[i] = arr[i-1]
        else
            arr[i] = item
            return
        end
        i -= 1
    end
    arr[1] = item
end

"""
push! appends an item to the end of the result set
"""
push!(p::KnnResult{Int64}, objID::I, dist::F) where {I <: Union{Int32,Int16}, F <: Real} = push!(p, convert(Int64, objID), convert(Float64, dist))
push!(p::KnnResult{T}, objID::T, dist::F) where {T, F <: Union{Float16, Float32}} = push!(p, objID, convert(Float64, dist))

function push!(p::KnnResult{T}, objID::T, dist::Float64) where T
    if length(p.pool) < p.k
        # fewer items than the maximum capacity
        push!(p.pool, Item(objID, dist))
        fix_order!(p)
        return true
    end

    @inbounds last_item = p.pool[end]
    if dist >= last_item.dist
        # p.k == length(p.pool) but item.dist doesn't improve the pool's radius
        return false
    end

    # p.k == length(p.pool) but item.dist improves the result set
    @inbounds p.pool[end] = Item(objID, dist)
    fix_order!(p)
    return true
end

"""
return the first item of the result set, the closest item
"""
function first(p::KnnResult{T}) where T
    @inbounds return p.pool[1]
end

"""
returns the last item of the result set
"""
function last(p::KnnResult{T}) where T
    @inbounds return p.pool[end]
end

"""
apply popfirst!(p.pool), an O(length(p.pool)) operation
"""
function popfirst!(p::KnnResult{T}) where T
    popfirst!(p.pool)
end

"""
apply pop!(p), an O(1) operation
"""
function pop!(p::KnnResult{T}) where T
    return pop!(p.pool)
end

"""
length returns the number of items in the result set
"""
Base.length(p::KnnResult{T}) where T = length(p.pool)

"""
The maximum allowed cardinality (the k of knn)
"""
maxlength(p::KnnResult{T}) where T = p.k

"""
covrad returns the coverage radius of the result set; if length(p) < K then typemax(Float32) is returned
"""
function covrad(p::KnnResult{T})::Float64 where T
    return length(p.pool) < p.k ? typemax(Float64) : last(p).dist
end

function empty!(p::KnnResult{T}) where T
    resize!(p.pool, 0)
end

##### iterator interface
### KnnResult
function Base.iterate(p::KnnResult{T}) where T
    return length(p) == 0 ? nothing : (first(p), 2)
end

function Base.iterate(p::KnnResult{T}, state::Int) where T
    if state > length(p)
        return nothing
    end

    return p.pool[state], state + 1
end


text:
# Inside make.jl
push!(LOAD_PATH,"../src/")
using PotentialLearning
using Documenter

DocMeta.setdocmeta!(PotentialLearning, :DocTestSetup, :(using PotentialLearning); recursive=true)

makedocs(
        modules  = [PotentialLearning],
        authors = "CESMIX-MIT",
        repo="https://github.com/cesmix-mit/PotentialLearning.jl/blob/{commit}{path}#{line}",
        sitename="PotentialLearning.jl",
        format=Documenter.HTML(;
            prettyurls=get(ENV, "CI", "false") == "true",
            canonical="https://cesmix-mit.github.io/PotentialLearning.jl",
            assets=String[],
        ),
        pages = [
            "Home" => "index.md",
            "Functions" => "functions.md",
        ],
        doctest = true,
        linkcheck = true,
        strict = true,
)

deploydocs(;
    repo="github.com/cesmix-mit/PotentialLearning.jl",
    devbranch = "main",
)


text:
"""
Types used in the PMDP
"""
const Action = Float64
const Timestep = Int64

"""
State for the PMDP
"""
struct State{n_res}
    c::SVector{n_res,Int64}   # Capacity vector
    t::Timestep               # Timestep
    iₚ::Int64        # Requested product index
end

abstract type PMDP{State, Action} <: MDP{State, Action} end

const EMPTY_PRODUCT_USER_BUDGET = -1.
const REJECT_ACTION = floatmax()
const ϵ = 10^-12


text:
import Base.zeros
function zeros(path::ASCIIString, nz::Int64, nx::Int64, ot::Float64, dt::Float64, nt::Int64)
    fid = open(path, "w")
    write(fid, nz, nx, ot, dt, nt)
    write(fid, zeros(nz*nx*nt))
    close(fid)
    return nothing
end

import Base.ones
function ones(path::ASCIIString, nz::Int64, nx::Int64, ot::Float64, dt::Float64, nt::Int64)
    fid = open(path, "w")
    write(fid, nz, nx, ot, dt, nt)
    write(fid, ones(nz*nx*nt))
    close(fid)
    return nothing
end

function dotSrcSpt(path::ASCIIString)
    (nz, nx, ot, dt, nt) =  InfoSrcSpt(path)
    fid = open(path, "r")
    seek(fid, sizeof(Float64)*5)
    l = 0.0
    for it = 1 : nt
        p = read(fid, Float64, nz*nx)
        l = l + dot(p,p)
    end
    close(fid)
    return l
end

function dotSrcSpt(path::ASCIIString, W::Array{Float64,1})
    (nz, nx, ot, dt, nt) =  InfoSrcSpt(path)
    if nz*nx != length(W)
       error("size dismatch")
    end
    fid = open(path, "r")
    seek(fid, sizeof(Float64)*5)
    l = 0.0
    for it = 1 : nt
        p = read(fid, Float64, nz*nx)
        l = l + dot(p, W.*p)
    end
    close(fid)
    return l
end

function AddSrcSpt(path1::ASCIIString, path2::ASCIIString, alpha::Float64; item=2)
    (nz1, nx1, ot1, dt1, nt1) = InfoSrcSpt(path1)
    (nz , nx , ot , dt , nt ) = InfoSrcSpt(path2)
    if nz1 != nz || nx1 != nx || nt1 != nt
       error("size dismatch of those two srcspts")
    end
    fid1 = open(path1, "r+")
    fid2 = open(path2, "r" )
    head = sizeof(Float64) * 5
    lspt = sizeof(Float64) * nz * nx
    if item == 1
       for it = 1 : nt
           position = head + (it-1)*lspt
           seek(fid1, position)
           seek(fid2, position)
           p1 = read(fid1, Float64, nz*nx)
           p2 = read(fid2, Float64, nz*nx)
           p1 = alpha*p1 + p2
           seek(fid1, position)
           write(fid1, p1)
       end
    elseif item == 2
       for it = 1 : nt
           position = head + (it-1)*lspt
           seek(fid1, position)
           seek(fid2, position)
           p1 = read(fid1, Float64, nz*nx)
           p2 = read(fid2, Float64, nz*nx)
           p1 = p1 + alpha * p2
           seek(fid1, position)
           write(fid1, p1)
       end
    end
    close(fid1); close(fid2);
    return nothing
end

function CGjoint(ntw::Int64, path_x::ASCIIString, path_s::ASCIIString, path_p::ASCIIString, shot::Shot, fidMtx::FidMtx, pos_rec::Array{Int64,2}; niter=10, mu=0.1, tmax=0.5)
    cost = zeros(niter+1)
    nz = fidMtx.nz; nx = fidMtx.nx; dt = fidMtx.dt
    zeros(path_x, nz, nx, 0.0, dt, ntw)
    shot_r = copy(shot)
    MultiStepAdjoint(ntw, path_s, shot_r, fidMtx)
    cp(path_s, path_p, remove_destination=true)
    gamma = dotSrcSpt(path_s)
    cost0 = dot(vec(shot_r.d), vec(shot_r.d))
    cost[1] = 1.0
    println("iteration 0, objective 1")
    for k = 1 : niter
        shot_q = MultiStepForward(pos_rec, path_p, fidMtx, tmax=tmax)
        delta = dot(vec(shot_q.d), vec(shot_q.d)) + mu * dotSrcSpt(path_p)
        if delta <= 0
           println("waring: indefinite")
        end
        alpha = gamma / delta
        AddSrcSpt(path_x, path_p, alpha)
        shot_r.d = shot_r.d - alpha*(shot_q.d)
        rm(path_s)
        MultiStepAdjoint(ntw, path_s, shot_r, fidMtx)
        AddSrcSpt(path_s, path_x, -mu)
        gamma0 = copy(gamma)
        gamma  = dotSrcSpt(path_s)
        df = dot(vec(shot_r.d), vec(shot_r.d)); mnorm = mu*dotSrcSpt(path_x)
        ratio = mnorm / df
        cost1  = (df + mnorm) / cost0
        cost[k+1] = cost1
        println("iteration $k, objective $cost1, ratio $ratio")
        beta = gamma / gamma0
        AddSrcSpt(path_p, path_s, beta, item=1)
    end
    return cost
end

function CGjoint_GP(ntw::Int64, path_x::ASCIIString, path_s::ASCIIString, path_p::ASCIIString, shot::Shot, fidMtx::FidMtx, pos_rec::Array{Int64,2}; inner=10, outer=5, mu=0.1, tmax=0.5)
    obj = CGjoint(ntw, path_x, path_s, path_p, shot, fidMtx, pos_rec, tmax=tmax, mu=mu, niter=inner)
    rm(path_s); rm(path_p);
    nz = fidMtx.nz; nx = fidMtx.nx;
    W = zeros(nz, nx)
    for iter_o = 1 : outer
        W = GPW(path_x)
        CG_GP(W, ntw, path_x, path_s, path_p, shot, fidMtx, pos_rec, niter=inner, mu=mu, tmax=tmax)
        rm(path_s)
        rm(path_p)
    end
    return W
end

function GPW(path_x::ASCIIString; delta = 1e-10)
    (nz, nx, ot, dt, ntw) = InfoSrcSpt(path_x)
    fid = open(path_x, "r")
    seek(fid, sizeof(Float64)*5)
    W = zeros(nz*nx)
    for it = 1 : ntw
        p  = read(fid, Float64, nz*nx)
        W  = W + p.^2
    end
    W = sqrt(W)
    W = 1 ./ (W+delta)
    return W
end

function obtainS(path_s::ASCIIString, mu::Float64, W::Array{Float64,1}, path_x::ASCIIString)
    W = mu * W
    (nz1, nx1, ot, dt, ntw1) = InfoSrcSpt(path_x)
    (nz , nx , ot, dt, ntw ) = InfoSrcSpt(path_x)
    if nz1 != nz || nx1 != nx || ntw1 != ntw
       error("size dismatch")
    end
    fids = open(path_s, "r+");
    fidx = open(path_x, "r" );
    for it = 1 : ntw
        pos = sizeof(Float64)*(5 + nz*nx*(it-1))
        seek(fids, pos);
        seek(fidx, pos);
        ps = read(fids, Float64, nz*nx)
        px = read(fidx, Float64, nz*nx)
        ps = ps - (W.*px)
        seek(fids, pos)
        write(fids, ps)
    end
    return nothing
end

function CG_GP(W::Array{Float64,1}, ntw::Int64, path_x::ASCIIString, path_s::ASCIIString, path_p::ASCIIString, shot_b::Shot, fidMtx::FidMtx, pos_rec::Array{Int64,2}; niter=10, mu=0.1, tmax=0.5)
    obj = zeros(niter+1)
    nz = fidMtx.nz; nx = fidMtx.nx; dt = fidMtx.dt
    shot_r = shot_b - MultiStepForward(pos_rec, path_x, fidMtx, tmax=tmax)
    MultiStepAdjoint(ntw, path_s, shot_r, fidMtx)
    obtainS(path_s, mu, W, path_x)
    cp(path_s, path_p, remove_destination=true)
    gamma = dotSrcSpt(path_s)
    obj[1] = gamma
    println("iteration 0, objective $gamma")
    for k = 1 : niter
        shot_q = MultiStepForward(pos_rec, path_p, fidMtx, tmax=tmax)
        delta = dot(vec(shot_q.d), vec(shot_q.d)) + mu * dotSrcSpt(path_p, W)
        if delta <= 0
           println("waring: indefinite")
        end
        if delta == 0
           delta = eps()
        end
        alpha = gamma / delta
        AddSrcSpt(path_x, path_p, alpha)
        shot_r.d = shot_r.d - alpha*(shot_q.d)
        rm(path_s)
        MultiStepAdjoint(ntw, path_s, shot_r, fidMtx)
        obtainS(path_s, mu, W, path_x)
        gamma0 = gamma
        gamma  = dotSrcSpt(path_s)
        println("iteration $k, objective $gamma")
        beta = gamma / gamma0
        obj[k+1] = gamma
        AddSrcSpt(path_p, path_s, beta, item=1)
    end
    return nothing
end

max_stars_repo_name: WneleiGao/AcousticWaveLab

max_issues_repo_name: WneleiGao/AcousticWaveLab

max_forks_repo_name: WneleiGao/AcousticWaveLab


text:
using SmoothingSplines
using Test

using LinearAlgebra
using Random

#n=50
#X = rand(n) .* 3
#Y = 2 .* X.^2 - X .- randn(n)
#idx = sortperm(X)
#X = X[idx]
#Y = Y[idx]
#h = diff(X)

# first test some of the helper functions
n=5
Y = rand(n)
h = [1.0 for i=1:(n-1)]
X = [0.0; cumsum(h)]
λ = 1.0

ws1 = ones(Float64, length(Y))
ws2 = 1.0 .+ rand(n)
ws_dict = Dict("uniform weights"=> ws1, "random weights"=>ws2)
for (k,ws) in ws_dict
  println(k)
  Q = SmoothingSplines.ReinschQ(h)
  R = SmoothingSplines.ReinschR(h)
  Qfull = reshape(Float64[x for x in Q], size(Q))
  Rfull = reshape(Float64[x for x in R], size(R))
  QtYfull = Qfull'*Y
  QtY = mul!(fill(0.0, n-2), transpose(Q), Y)

  @test QtY ≈ QtYfull

  tmpQtQpRfull = Rfull + λ*Qfull'*Matrix(Diagonal(1.0 ./ ws))*Qfull
  tmpQtQpR =SmoothingSplines. QtQpR(h, λ, ws)
  @test vec(tmpQtQpR[3,:]) ≈ diag(tmpQtQpRfull)
  @test vec(tmpQtQpR[2,2:end]) ≈ diag(tmpQtQpRfull,1)
  @test vec(tmpQtQpR[1, 3:end]) ≈ diag(tmpQtQpRfull,2)

  γfull = tmpQtQpRfull\QtYfull
  γ = copy(QtY)
  SmoothingSplines.pbtrf!('U', 2, tmpQtQpR)
  SmoothingSplines.pbtrs!('U', 2, tmpQtQpR, γ)
  @test γ ≈ γfull

  gfull = Y - λ*Matrix(Diagonal(1.0./ws))*Qfull*γfull
  g = Y - λ*mul!(zero(Y), Q, γ)./ws
  @test g ≈ gfull

  fullfit = fit(SmoothingSpline, X, Y, λ, ws)
end
# now test the full algorithm, and predict functions
println("testing predict function")
Random.seed!(1)
n=50
X = rand(n) .* 3
Y = 2 .* X.^2 .- X .- randn(n)

spl = fit(SmoothingSpline, X, Y, 1.0)

@test predict(spl,X) ≈ predict(spl)

println("testing against smooth.spline on R cars dataset")
# real data examples
# cars dataset
# using RDatasets
# cars = dataset("datasets","cars")
# X = map(Float64,convert(Array,cars[:Speed]))
X = [4.0,4.0,7.0,7.0,8.0,9.0,10.0,10.0,10.0,11.0,11.0,12.0,12.0,12.0,
  12.0,13.0,13.0,13.0,13.0,14.0,14.0,14.0,14.0,15.0,15.0,15.0,16.0,
  16.0,17.0,17.0,17.0,18.0,18.0,18.0,18.0,19.0,19.0,19.0,20.0,20.0,
  20.0,20.0,20.0,22.0,23.0,24.0,24.0,24.0,24.0,25.0]
# Y = map(Float64,convert(Array,cars[:Dist]))
Y = [2.0,10.0,4.0,22.0,16.0,10.0,18.0,26.0,34.0,17.0,28.0,14.0,20.0,24.0,
  28.0,26.0,34.0,34.0,46.0,26.0,36.0,60.0,80.0,20.0,26.0,54.0,32.0,40.0,
  32.0,40.0,50.0,42.0,56.0,76.0,84.0,36.0,46.0,68.0,32.0,48.0,52.0,56.0,
  64.0,66.0,54.0,70.0,92.0,93.0,120.0,85.0]
# so that λ for this package is comparable to λ in smooth.spline
X = (X .- minimum(X))/(maximum(X)-minimum(X))

# compare to R.stats:
# attach(cars)
# cars.spl <- smooth.spline(speed, dist)
λ= 0.1112206 # cars.spl$lambda
# Rpred <- predict(cars.spl)$y
Rpred =  [ 1.657809, 11.682913, 15.064409, 18.482339, 21.947078, 25.465623, 29.050518,
    32.707552, 36.423468, 40.195168, 44.054188, 48.025219, 52.115627, 56.323874,
    60.673887, 69.808295, 74.533703, 79.313597, 84.103688]
cars_fit = fit(SmoothingSpline, X, Y, λ)
@test cars_fit.g ≈ Rpred rtol=1e-4


text:
#################### Random Walk Metropolis ####################

#################### Types and Constructors ####################

mutable struct RWMTune <: SamplerTune
  logf::Union{Function, Missing}
  scale::Union{Float64, Vector{Float64}}
  eligible::Vector{Symbol}
  proposal::SymDistributionType

  RWMTune() = new()

  function RWMTune(x::Vector, scale::Real, logf::Union{Function, Missing}, eligible::Vector{Symbol};
                   proposal::SymDistributionType=Normal)
    new(logf, Float64(scale), eligible, proposal)
  end

  function RWMTune(x::Vector, scale::Vector{T},
                  logf::Union{Function, Missing},
                  eligible::Vector{Symbol};
                  proposal::SymDistributionType=Normal) where {T<:Real}
    new(logf, convert(Vector{Float64}, scale), eligible, proposal)
  end
end

const RWMVariate = Sampler{RWMTune, T} where T

validate(v::RWMVariate{T}) where T<:AbstractArray{S} where S<:Union{Real, GeneralNode} = validate(v, v.tune.scale)

validate(v::RWMVariate{T}, scale::Float64) where T<:AbstractArray{S} where S<:Union{Real, GeneralNode} = v

function validate(v::RWMVariate{T}, scale::Vector) where T<:AbstractArray{S} where S<:Union{Real, GeneralNode}
  n = length(v)
  length(scale) == n ||
    throw(ArgumentError("length(scale) differs from variate length $n"))
  v
end

#################### Sampler Constructor ####################
"""
    RWM(params::ElementOrVector{Symbol},
                  scale::ElementOrVector{T}; args...) where {T<:Real})

Construct a `Sampler` object for RWM sampling. Parameters are assumed to be
continuous, but may be constrained or unconstrained.

Returns a `Sampler{RWMTune}` type object.

* `params`: stochastic node(s) to be updated with the sampler. Constrained parameters are mapped to unconstrained space according to transformations defined by the Stochastic `unlist()` function.

* `scale`: scaling value or vector of the same length as the combined elements of nodes `params` for the `proposal` distribution. Values are relative to the unconstrained parameter space, where candidate draws are generated.

* `args...`: additional keyword arguments to be passed to the `RWMVariate` constructor.
"""
function RWM(params::ElementOrVector{Symbol},
              scale::ElementOrVector{T}; args...) where {T<:Real}
  tune = RWMTune(Float64[], scale, logpdf!, Symbol[])
  Sampler(params, tune, Symbol[], false)
end

"""
    RWM(params::ElementOrVector{Symbol}, moves::Array{Symbol}; args...)

Construct the RWM sampler for Trees. If you set moves to :all it will use all
eligible moves to change the tree. These are currently:
NNI, SPR, Slide, Swing, :EdgeLength

Returns a `Sampler{RWMTune}` type object.
"""
function RWM(params::ElementOrVector{Symbol}, moves::ElementOrVector{Symbol}; kwargs...)
  eligible = [:NNI, :SPR, :Slide, :Swing, :EdgeLength]
  to_use = Symbol[]
  if moves == :all
    to_use = eligible
  else
    for i in moves
      !(i in eligible) &&
        throw("$i is not an eligible tree move. The list of eligible tree moves is $eligible")
      push!(to_use, i)
    end
  end

  tune = RWMTune(Float64[], logpdf!, 1.0, to_use)
  Sampler(params, tune, Symbol[], transform)
end

#################### Sampling Functions ####################

sample!(v::RWMVariate) = sample!(v, v.tune.logf)

"""
    sample!(v::RWMVariate, logf::Function, moves::Array{Symbol})

Propose a new tree by randomly performing a move from the ones specified in `moves`.

Returns `v` updated with simulated values and associated tuning parameters.
"""
function sample!(v::RWMVariate{T}, logf::Function; kwargs...) where T<:GeneralNode
  tree = v[1]
  tc = deepcopy(tree)

  move = rand(v.elegible)
  if move == :NNI
    NNI!(tree)
  elseif move == :SPR
    tree = SPR(tree)
  elseif move == :Slide
    slide!(tree)
  elseif move == :Swing
    swing!(tree)
  elseif move == :EdgeLength
    change_edge_length!(tree)
  else
    throw("Tree move not elegible ")
  end
  if rand() < exp(logf(tree) - logf(tc))
    v[1] = tree
  else
    v[1] = tc
  end
  v
end

"""
    sample!(v::RWMVariate, logf::Function)

Draw one sample from a target distribution using the RWM sampler. Parameters
are assumed to be continuous and unconstrained.

Returns `v` updated with simulated values and associated tuning parameters.
"""
function sample!(v::RWMVariate{T}, logf::Function; kwargs...) where T<:AbstractArray{<:Real}
  x = v + v.tune.scale .* rand(v.tune.proposal(0.0, 1.0), length(v))
  if rand() < exp(logf(x) - logf(v.value))
    v[:] = x
  end
  v
end

max_stars_repo_name: erathorn/MCPhylo

max_issues_repo_name: erathorn/MCPhylo

max_forks_repo_name: erathorn/MCPhylo


text:
using Test
using Expronicon

@testset "is_function" begin
    @test is_function(:(foo(x) = x))
    @test is_function(:(x -> 2x))
end

@testset "uninferrable_typevars" begin
    def = @expr JLKwStruct struct Inferable1{T}
        x::Constaint{T, <(2)}
    end

    @test uninferrable_typevars(def) == []

    def = @expr JLKwStruct struct Inferable2{T}
        x::Constaint{Float64, <(2)}
    end

    @test uninferrable_typevars(def) == [:T]

    def = @expr JLKwStruct struct Inferable3{T, N}
        x::Int
        y::N
    end
    @test uninferrable_typevars(def) == [:T]

    def = @expr JLKwStruct struct Inferable4{T, N}
        x::T
        y::N
    end
    @test uninferrable_typevars(def) == []

    def = @expr JLKwStruct struct Inferable5{T, N}
        x::T
        y::Float64
    end

    @test uninferrable_typevars(def) == [:T, :N]
    @test uninferrable_typevars(def; leading_inferable=false) == [:N]
end

@testset "has_plain_constructor" begin
    def = @expr JLKwStruct struct Foo1{T, N}
        x::Int
        y::N
        Foo1{T, N}(x, y) where {T, N} = new{T, N}(x, y)
    end

    @test has_plain_constructor(def) == true

    def = @expr JLKwStruct struct Foo2{T, N}
        x::T
        y::N
        Foo2(x, y) = new{typeof(x), typeof(y)}(x, y)
    end

    @test has_plain_constructor(def) == false

    def = @expr JLKwStruct struct Foo3{T, N}
        x::Int
        y::N
        Foo3{T}(x, y) where T = new{T, typeof(y)}(x, y)
    end

    @test has_plain_constructor(def) == false

    def = @expr JLKwStruct struct Foo4{T, N}
        x::T
        y::N
        Foo4{T, N}(x::T, y::N) where {T, N} = new{T, N}(x, y)
    end

    @test has_plain_constructor(def) == false
end

@testset "is_kw_function" begin
    @test is_kw_function(:(
        function foo(x::Int; kw=1)
        end
    ))

    ex = :(function (x::Int; kw=1) end)
    @test is_kw_function(ex)
    @test !is_kw_function(true)

    @test !is_kw_function(:(
        function foo(x::Int)
        end
    ))

    @test !is_kw_function(:(
        function (x::Int)
        end
    ))
end

@testset "JLFunction(ex)" begin
    jlfn = JLFunction()
    @test jlfn.name === nothing

    @test_expr JLFunction function foo(x::Int, y::Type{T}) where {T <: Real}
        return x
    end

    def = @test_expr JLFunction function (x, y)
        return 2
    end
    @test is_kw_function(def) == false

    def = @test_expr JLFunction function (x, y; kw=2)
        return "aaa"
    end
    @test is_kw_function(def) == true

    @test_expr JLFunction (x, y)->sin(x)

    # canonicalize head when it's a block
    @test_expr JLFunction function (x::Int; kw=1) end

    ex = :(struct Foo end)
    @test_throws AnalysisError JLFunction(ex)
    ex = :(@foo(2, 3))
    @test_throws AnalysisError split_function_head(ex)

    ex = :(Foo(; a = 1) = new(a))
    @test JLFunction(ex).kwargs[1] == Expr(:kw, :a, 1)

    @test_expr JLFunction function f(x::T; a=10)::Int where T
        return x
    end

    @test_expr JLFunction f(x::Int)::Int = x
end

@testset "JLStruct(ex)" begin
    @test JLField(;name=:x).name === :x
    @test JLField(;name=:x).type === Any
    @test JLStruct(;name=:Foo).name === :Foo

    ex = :(struct Foo
        x::Int
    end)

    jlstruct = JLStruct(ex)
    println(jlstruct)
    @test jlstruct.name === :Foo
    @test jlstruct.ismutable === false
    @test length(jlstruct.fields) == 1
    @test jlstruct.fields[1].name === :x
    @test jlstruct.fields[1].type === :Int
    @test jlstruct.fields[1].line isa LineNumberNode
    @test codegen_ast(jlstruct) == ex

    ex = :(mutable struct Foo{T, S <: Real} <: AbstractArray
        a::Float64

        function foo(x, y, z)
            new(1)
        end
    end)

    jlstruct = JLStruct(ex)
    println(jlstruct)
    @test jlstruct.ismutable == true
    @test jlstruct.name === :Foo
    @test jlstruct.typevars == Any[:T, :(S <: Real)]
    @test jlstruct.supertype == :AbstractArray
    @test jlstruct.misc[1] == ex.args[3].args[end]
    @test rm_lineinfo(codegen_ast(jlstruct)) == rm_lineinfo(ex)

    ex = quote
        """
        Foo
        """
        struct Foo
            "xyz"
            x::Int
            y

            Foo(x) = new(x)
            1 + 1
        end
    end
    ex = ex.args[2]
    jlstruct = JLStruct(ex)
    @test jlstruct.doc == "Foo\n"
    @test jlstruct.fields[1].doc == "xyz"
    @test jlstruct.fields[2].type === Any
    @test jlstruct.constructors[1].name === :Foo
    @test jlstruct.constructors[1].args[1] === :x
    @test jlstruct.misc[1] == :(1 + 1)
    ast = codegen_ast(jlstruct)
    @test ast.args[1] == GlobalRef(Core, Symbol("@doc"))
    @test ast.args[3] == "Foo\n"
    @test ast.args[4].head === :struct
    @test is_function(ast.args[4].args[end].args[end-1])
    println(jlstruct)

    @test_throws AnalysisError split_struct_name(:(function Foo end))
end

@testset "JLKwStruct" begin
    def = @expr JLKwStruct struct Trait end
    @test_expr prettify(codegen_ast_kwfn(def)) == quote
    end

    @test JLKwField(;name=:x).name === :x
    @test JLKwField(;name=:x).type === Any
    @test JLKwStruct(;name=:Foo).name === :Foo

    def = @expr JLKwStruct struct ConvertOption
        include_defaults::Bool=false
        exclude_nothing::Bool=false
    end

    @test_expr codegen_ast_kwfn(def, :create) == quote
        function create(::Type{S}; include_defaults = false, exclude_nothing = false) where S <: ConvertOption
            ConvertOption(include_defaults, exclude_nothing)
        end
    end

    def = @expr JLKwStruct struct Foo1{N, T}
        x::T = 1
    end
    println(def)

    @test_expr codegen_ast_kwfn(def, :create) == quote
        function create(::Type{S}; x = 1) where {N, T, S <: Foo{N, T}}
            Foo1{N, T}(x)
        end
        function create(::Type{S}; x = 1) where {N, S <: Foo{N}}
            Foo1{N}(x)
        end
    end

    @test_expr codegen_ast(def) == quote
        struct Foo1{N, T}
            x::T
        end
        function Foo1{N, T}(; x = 1) where {N, T}
            Foo1{N, T}(x)
        end
        function Foo1{N}(; x = 1) where N
            Foo1{N}(x)
        end
    end

    def = @expr JLKwStruct struct Foo2 <: AbstractFoo
        x = 1
        y::Int
    end

    @test_expr codegen_ast(def) == quote
        struct Foo2 <: AbstractFoo
            x
            y::Int
        end
        function Foo2(; x = 1, y)
            Foo2(x, y)
        end
    end

    ex = quote
        """
        Foo
        """
        mutable struct Foo
            "abc"
            a::Int = 1
            b

            Foo(x) = new(x)
            1 + 1
        end
    end
    ex = ex.args[2]
    jlstruct = JLKwStruct(ex)
    @test jlstruct.doc == "Foo\n"
    @test jlstruct.fields[1].doc == "abc"
    @test jlstruct.fields[2].name === :b
    @test jlstruct.constructors[1].name === :Foo
    @test jlstruct.misc[1] == :(1 + 1)
    println(jlstruct)

    def = @expr JLKwStruct struct Foo3
        a::Int = 1
        Foo3(;a = 1) = new(a)
    end

    @test_expr codegen_ast(def) == quote
        struct Foo3
            a::Int
            Foo3(; a = 1) = new(a)
        end
    end

    def = @expr JLKwStruct struct Potts{Q}
        L::Int
        beta::Float64=1.0
        neighbors::Neighbors = square_lattice_neighbors(L)
    end

    @test_expr codegen_ast_kwfn(def, :create) == quote
        function create(::Type{S}; L, beta = 1.0, neighbors = square_lattice_neighbors(L)) where {Q, S <: Potts{Q}}
            Potts{Q}(L, beta, neighbors)
        end
    end

    def = @expr JLKwStruct struct Flatten
        x = 1
        begin
            y = 1
        end
    end

    @test def.fields[1].name === :x
    @test def.fields[2].name === :y
end

@test sprint(print, AnalysisError("a", "b")) == "expect a expression, got b."

@testset "JLIfElse" begin
    jl = JLIfElse()
    jl[:(foo(x))] = :(x = 1 + 1)
    jl[:(goo(x))] = :(y = 1 + 2)
    jl.otherwise = :(error("abc"))
    println(jl)

    ex = codegen_ast(jl)
    dst = JLIfElse(ex)
    @test_expr dst[:(foo(x))] == :(x = 1 + 1)
    @test_expr dst[:(goo(x))] == :(y = 1 + 2)
    @test_expr dst.otherwise == :(error("abc"))
end

@testset "JLFor" begin
    ex = :(for i in 1:10, j in 1:20,
            k in 1:10
        1 + 1
    end)
    jl = JLFor(ex)
    println(jl)
    @test codegen_ast(jl) == ex

    jl = JLFor(;vars=[:x], iterators=[:itr], kernel=:(x + 1))
    ex = codegen_ast(jl)
    @test ex.head === :for
    @test ex.args[1].args[1] == :(x = itr)
    @test ex.args[2] == :(x+1)

    ex = :(for i in 1:10
        1 + 1
    end)
    jl = JLFor(ex)
    println(jl)
    @test jl.vars == [:i]
    @test jl.iterators == [:(1:10)]
end


text:
abstract type GraphicElement2D <: GraphicElement end

#--- 2D POINTS -----------------------------------------
#-------------------------------------------------------

struct Vec2{T<:Real}
    x::T
    y::T
end

Vec2(x::Real,y::Real) = Vec2(promote(x,y)...)
Vec2(t::Tuple{<:Real,<:Real}) = Vec2(t...)
Vec2(z::Complex) = Vec2(reim(z)...)
complex(V::Vec2) = V.x + im*V.y
abs(V::Vec2) = hypot(V.x,V.y)

string(v::Vec2) = "($(v.x),$(v.y))"

import Base.convert
convert(::Type{Vector},V::Vec2) = [V.x,V.y]

import LinearAlgebra.×
×(V::Vec2,W::Vec2) = V.x*W.y - W.x*V.y

struct PointShape
    name::Symbol
    size::Float64
end

const _DEFAULT_POINTSHAPE_KWARGS =
    OrderedDict(
        :name => :dot,
        :size => 0.025
    )

PointShape(name;kw...) = PointShape(name=name,kw...)
PointShape(;kw...) = PointShape(updatedvals(_DEFAULT_POINTSHAPE_KWARGS,kw)...)

convert(::Type{PointShape},s::Symbol) = PointShape(name=s)

"""
    Point2D(x::Real,y::Real; label="",pen=Pen())
    Point2D(P; label="",pen=Pen())

A graphics primitive representing a two-dimensional point.

`P` may be a 2-tuple of real numbers, a `Vec2`, or a `Complex`

# Examples
```julia-repl
julia> Point2D(3,-1;pen="DarkGreen")
Point2D(3,-1;pen=DarkGreen)
```
"""
struct Point2D <: GraphicElement2D
    P::Vec2
    label::AbstractString
    pen::Pen
    shape::PointShape
end

"""
    Point(x::Real,y::Real;kwargs...)
    Point(x::Real,y::Real,z::Real;kwargs...)
    Point(P;kwargs...)

Return a Point2D or Point3D object, as appropriate

# Examples
```julia-repl
julia> Point(1,2;pen=Pen(opacity=0.5),shape=:plus)
Point2D(1,2;pen=opacity(0.5))
```
"""
function Point end

const _DEFAULT_POINT2D_KWARGS =
    OrderedDict(
        :label => "",
        :pen => Pen(),
        :shape => PointShape()
    )

Point2D(v::Vec2;kwargs...) =
    Point2D(v,updatedvals(_DEFAULT_POINT2D_KWARGS,
                                process_pen_kwargs(kwargs))...)
Point(v::Vec2;kwargs...) = Point2D(v;kwargs...)
Point2D(P::Tuple{<:Real,<:Real};kwargs...) = Point2D(Vec2(P);kwargs...)
Point(P::Tuple{<:Real,<:Real};kwargs...) = Point2D(P;kwargs...)
Point2D(x::Real,y::Real;kwargs...) =
                Point2D(Vec2(x,y);kwargs...)
Point(x::Real,y::Real;kwargs...) = Point2D(x,y;kwargs...)
Point2D(z::Union{Real,Complex};kwargs...) = Point2D(reim(z)...;kwargs...)
Point(z::Union{Real,Complex};kwargs...) = Point2D(z;kwargs...)

function point_drawcommand(x::String,y::String,p::Pen,shape::PointShape,lbl::AbstractString)
    label = lbl == "" ? "" : "L=$(enclosequote(lbl)),"
    pen = isdefault(p) ? "" : ",p=$(string(p))"
    if shape.name == :dot
        "dot($label($x,$y)$pen);"
    elseif shape.name == :cross
        ϵ = shape.size / √(2)
        """
        draw($label($x,$y) + (-$ϵ,-$ϵ) -- ($x,$y) + ($ϵ,$ϵ)$pen);
        draw(($x,$y) + ($ϵ,-$ϵ) -- ($x,$y) + (-$ϵ,$ϵ)$pen);
        """
    elseif shape.name == :plus
        ϵ = shape.size
        """
        draw($label($x,$y) + (-$ϵ,0) -- ($x,$y) + ($ϵ,0)$pen);
        draw(($x,$y) + (0,-$ϵ) -- ($x,$y) + (0,$ϵ)$pen);
        """
    elseif shape.name == :circle
        ϵ = shape.size
        "draw($(label)circle(($x,$y), $ϵ$pen);"
    end
end
point_drawcommand(x::Real,y::Real,p::Pen,shape::PointShape,lbl::AbstractString) = point_drawcommand(string(x),string(y),p,shape,lbl)

function AsyString(P::Point2D)
    AsyString(point_drawcommand(P.P.x,P.P.y,P.pen,P.shape,P.label))
end

function Base.show(io::IO,P::Point2D)
    kwargs = kwargstring(P,_DEFAULT_POINT2D_KWARGS)
    print(io,"Point2D($(P.P.x),$(P.P.y)$kwargs)")
end

function PointCloud end

struct PointCloud2D <: GraphicElement2D
    points::Array{<:Vec2,1}
    pen::Pen
    shape::PointShape
end

const _DEFAULT_POINTCLOUD2D_KWARGS =
    OrderedDict(
        :pen => Pen(),
        :shape => PointShape()
    )

PointCloud2D(points::Array{<:Vec2,1};kwargs...) =
        PointCloud2D(points,updatedvals(_DEFAULT_POINTCLOUD2D_KWARGS,
                                process_pen_kwargs(kwargs))...)

PointCloud(points::Array{<:Vec2,1};kwargs...) = PointCloud2D(points;kwargs...)

PointCloud2D(points::Array{<:Tuple{<:Real,<:Real},1};kwargs...) =
                                PointCloud2D(map(Vec2,points);kwargs...)

PointCloud(points::Array{<:Tuple{<:Real,<:Real},1};kwargs...) =
                                    PointCloud2D(map(Vec2,points);kwargs...)

PointCloud2D(coords::Array{<:Real,2};kwargs...) =
            PointCloud2D([Vec2(coords[i,:]...) for i=1:size(coords,1)];kwargs...)

PointCloud2D(points::Array{<:RealOrComplex,1};kwargs...) =
                        PointCloud2D(map(Vec2,points);kwargs...)

PointCloud2D(x,y;kwargs...) = PointCloud2D(collect(zip(x,y));kwargs...)

PointCloud(points::Array{<:RealOrComplex,1};kwargs...) =
                                PointCloud2D(map(Vec2,points);kwargs...)

PointCloud(x,y;kwargs...) = PointCloud2D(x,y;kwargs...)

function PointCloud(coords::Array{<:Real,2};kwargs...)
    if size(coords,2) == 2
        PointCloud2D(coords;kwargs...)
    else
        PointCloud3D(coords;kwargs...)
    end
end

const _DEFAULT_POINTCLOUD_KWARGS =
    OrderedDict(
        :pen => Pen(),
        :shape => PointShape()
    )

function AsyString(P::PointCloud2D)
    coords = vcat([[p.x p.y] for p in P.points]...)
    AsyString("""
    file pathdata = input("pointcloud{IDENTIFIER}.csv");
    real[][] A = pathdata.csv().dimension(0,2);
    close(pathdata);
    for(int i=0; i<A.length; ++i){
        $(point_drawcommand("A[i][0]","A[i][1]",P.pen,P.shape,""))
    }
    """,
    Dict("pointcloud{IDENTIFIER}.csv"=>coords))
end

#--- 2D PATHS -----------------------------------------
#------------------------------------------------------

"""
    Path2D(points;label="",pen=Pen(),arrow=NoArrow(),spline=false)
    Path2D(x,y;label="",pen=Pen(),arrow=NoArrow(),spline=false)

A graphics primitive representing a two-dimensional path

`points` may be an Array of `Vec2`s, an Array of 2-tuples,
or an ``n × 2`` Array. Alternatively, iterables of coordinates
may be supplied separately as `x` and `y`

# Examples
```julia-repl
julia> Path2D([(0,0),(1,0),(1,1)];pen="MidnightBlue")
Path2D(<3 points>;pen=MidnightBlue)
```
"""
struct Path2D <: GraphicElement2D
    points::Array{<:Vec2,1}
    label::AbstractString
    pen::Pen
    arrow::Arrow
    spline::Bool
end

"""
    Path(A;kwargs)

Return a Path2D or Path3D object, as appropriate

# Examples
```julia-repl
julia> Path([1 2; 3 4])
Path2D(<2 points>)
```
"""
function Path end

const _DEFAULT_PATH2D_KWARGS =
    OrderedDict(
          :label=>"",
          :pen=>Pen(),
          :arrow=>NoArrow(),
          :spline=>false)

Path2D(points::Array{<:Vec2,1};kwargs...) =
        Path2D(points,updatedvals(_DEFAULT_PATH2D_KWARGS,
                                process_pen_kwargs(kwargs))...)

Path(points::Array{<:Vec2,1};kwargs...) = Path2D(points;kwargs...)

Path2D(points::Array{<:Tuple{<:Real,<:Real},1};kwargs...) =
                                Path2D(map(Vec2,points);kwargs...)

Path(points::Array{<:Tuple{<:Real,<:Real},1};kwargs...) =
                                    Path2D(map(Vec2,points);kwargs...)

Path2D(coords::Array{<:Real,2};kwargs...) =
            Path2D([Vec2(coords[i,:]...) for i=1:size(coords,1)];kwargs...)

Path2D(points::Array{<:RealOrComplex,1};kwargs...) =
                        Path2D(map(Vec2,points);kwargs...)

Path2D(x,y;kwargs...) = Path2D(collect(zip(x,y));kwargs...)

Path(points::Array{<:RealOrComplex,1};kwargs...) =
                                Path2D(map(Vec2,points);kwargs...)

Path(x,y;kwargs...) = Path2D(x,y;kwargs...)

function Path(coords::Array{<:Real,2};kwargs...)
    if size(coords,2) == 2
        Path2D(coords;kwargs...)
    else
        Path3D(coords;kwargs...)
    end
end

function AsyString(path::Path2D)
    label = path.label == "" ? "" : "L=$(enclosequote(path.label))"
    pathname = "p{IDENTIFIER}"
    coords = vcat([[p.x p.y] for p in path.points]...)
    spline = path.spline ? ".." : "--"
    arrow = path.arrow == NoArrow() ? "" : AsyString(path.arrow)
    if length(path.points) > 10
        AsyString("""
        file pathdata = input("path{IDENTIFIER}.csv");
        real[][] A = pathdata.csv().dimension(0,2);
        close(pathdata);
        guide $pathname;
        for(int i=0; i<A.length; ++i){
            $pathname = $pathname $spline (A[i][0],A[i][1]);
        }
        draw($(filterjoin(label,pathname,path.pen,arrow)));
        """,
        Dict("path{IDENTIFIER}.csv"=>coords))
    else
        AsyString("""
        path $pathname = $(join([string(p) for p in path.points],spline));
        draw($(filterjoin(label,pathname,path.pen,arrow)));
        """)
    end
end

function Base.show(io::IO,P::Path2D)
    l = length(P.points)
    s = l ≠ 1 ? "s" : ""
    kwargs = kwargstring(P,_DEFAULT_PATH2D_KWARGS)
    print(io,"""
    Path2D(<$l point$s>$kwargs)""")
end

import Base.map
function Base.map(f::Function,P::Path2D)
    Path(map(f,P.points);label=P.label,
                         pen=P.pen,
                         arrow=P.arrow,
                         spline=P.spline)
end

#--- 2D CIRCLES ---------------------------------------
#------------------------------------------------------

"""
    Circle2D(center,radius; pen::Pen,fillpen::Pen)

A graphics primitive representing a circle in the plane

`center` may be a `Vec2` or a 2-tuple of `Real`s or a
`Complex`

# Examples
```julia-repl
julia> Circle2D((0,0),1;pen="LightBlue",fillpen="red")
Circle2D((0,0),1;pen=LightBlue,fillpen=red)
```
"""
struct Circle2D <: GraphicElement2D
    center::Vec2
    radius::Real
    pen::Pen
    fillpen::Pen
    clip::Bool
end

const _DEFAULT_CIRCLE2D_KWARGS =
    OrderedDict(
        :pen=>Pen(),
        :fillpen=>NoPen(),
        :clip=>false
    )

function Circle2D(center::Vec2,
                  radius::Real;kwargs...)
    Circle2D(center,radius,updatedvals(_DEFAULT_CIRCLE2D_KWARGS,
                                    process_pen_kwargs(kwargs))...)
end

Circle(center::Vec2,radius::Real;kwargs...) =
                            Circle2D(center,radius;kwargs...)

function Circle2D(x::Real,y::Real,radius::Real;kwargs...)
    Circle2D(Vec2(x,y),radius;kwargs...)
end

Circle2D(z::RealOrComplex,radius::Real;kwargs...) =
    Circle2D(Vec2(z),radius;kwargs...)

Circle(z::RealOrComplex,radius::Real;kwargs...) =
            Circle2D(z,radius;kwargs...)

Circle2D(center::Tuple{<:Real,<:Real},radius::Real;kwargs...) =
    Circle2D(Vec2(center),radius;kwargs...)

Circle(center::Tuple{<:Real,<:Real},radius::Real;kwargs...) =
    Circle2D(Vec2(center),radius;kwargs...)

is_no_pen(P::Pen) = P.other == "NoPen"

function AsyString(C::Circle2D)
    if is_no_pen(C.pen) && is_no_pen(C.fillpen)
        return AsyString("")
    elseif is_no_pen(C.fillpen)
        drawcommand = "draw"
        penname = "p"
    elseif is_no_pen(C.pen)
        drawcommand = "fill"
        fillname = "p"
    else
        drawcommand = "filldraw"
        fillname = "fillpen"
        penname = "drawpen"
    end
    fillpen = is_no_pen(C.fillpen) ? "" : "$fillname=$(C.fillpen)"
    drawpen = (is_no_pen(C.pen) || isdefault(C.pen)) ? "" : "$penname=$(C.pen)"
    center = string(C.center)
    circle = "circle($center,$(C.radius))"
    if C.clip
        clipstatement = "clip($circle);\n"
    else
        clipstatement = ""
    end
    AsyString("""
    $clipstatement$drawcommand($(filterjoin(circle,fillpen,drawpen)));
    """)
end

function Base.show(io::IO,C::Circle2D)
    kwargs = kwargstring(C,_DEFAULT_CIRCLE2D_KWARGS)
    print(io,"""
    Circle2D(($(C.center.x),$(C.center.y)),$(C.radius)$kwargs)""")
end

#--- 2D POLYGONS -----------------------------------------
#------------------------------------------------------

"""
    Polygon2D(points;pen=Pen(),
                     fillpen=Pen(color="white"),
                     spline=false)

A graphics primitive representing a two-dimensional polygon

`points` may be an Array of `Vec2`s, an Array of 2-tuples,
or an ``n × 2`` Array. Alternatively, iterables of coordinates
may be supplied separately as `x` and `y`

# Examples
```julia-repl
julia> Polygon2D([(0,0),(1,0),(1,1)];pen="MidnightBlue")
Polygon2D(<3 points>;pen=MidnightBlue)
```
"""
struct Polygon2D <: GraphicElement2D
    points::Array{<:Vec2,1}
    pen::Pen
    fillpen::Pen
    spline::Bool
    clip::Bool
end

const _DEFAULT_POLYGON2D_KWARGS =
    OrderedDict(
        :pen=>Pen(),
        :fillpen=>NoPen(),
        :spline=>false,
        :clip=>false,
    )

Polygon2D(points::Array{<:Vec2,1};kwargs...) =
       Polygon2D(points,updatedvals(_DEFAULT_POLYGON2D_KWARGS,
                                            process_pen_kwargs(kwargs))...)

Polygon(points::Array{<:Vec2,1};kwargs...) = Polygon2D(points;kwargs...)

Polygon(points::Array{<:Tuple{<:Real,<:Real},1};kwargs...) = Polygon2D(points;kwargs...)

Polygon2D(points::Array{<:Tuple{<:Real,<:Real},1};kwargs...) =
                                Polygon2D(map(Vec2,points);kwargs...)

Polygon2D(coords::Array{<:Real,2};kwargs...) =
            Polygon2D([Vec2(coords[i,:]...) for i=1:size(coords,1)];kwargs...)

Polygon2D(points::Array{<:RealOrComplex,1};kwargs...) =
                        Polygon2D(map(Vec2,points);kwargs...)

Polygon(points::Array{<:RealOrComplex,1};kwargs...) =
                      Polygon2D(map(Vec2,points);kwargs...)

function counterclockwise(P::Polygon2D)
    v = P.points
    sum(p×q for (p,q) in zip(v,vcat(v[2:end],v[1:1]))) > 0
end

function Polygon(P::Path2D;kwargs...)
    lastindex = endof(P.points) - (P.points[1] == P.points[end] ? 1 : 0)
    Polygon(P.points[1:lastindex];kwargs...)
end

function box(a,b,c,d;kwargs...)
    Polygon([a b; c b; c d; a d];kwargs...)
end

"""
    Polygon(points;kwargs...)

Return a `Polygon2D` or a `Polygon3D` as appropriate

`points` may be an array of `Vec2`s or `Complex`es,
or an ``n × 2`` array of `Real`s

# Examples
```julia-repl
julia> Polygon([im,0,1])
Polygon2D(<3 points>)
```
"""
function Polygon(coords::Array{<:Real,2};kwargs...)
    if size(coords,2) == 2
        Polygon2D(coords;kwargs...)
    else
        Polygon3D(coords;kwargs...)
    end
end

function AsyString(P::Polygon2D)
    if is_no_pen(P.pen) && is_no_pen(P.fillpen)
        return AsyString("")
    elseif is_no_pen(P.fillpen)
        drawcommand = "draw"
        penname = "p"
    elseif is_no_pen(P.pen)
        drawcommand = "fill"
        fillname = "p"
    else
        drawcommand = "filldraw"
        fillname = "fillpen"
        penname = "drawpen"
    end
    fillpen = P.fillpen.other == "NoPen" ? "" : "$fillname=$(P.fillpen)"
    drawpen = (is_no_pen(P.pen) || isdefault(P.pen)) ? "" : "$penname=$(P.pen)"
    pathname = "p{IDENTIFIER}"
    coords = vcat([[p.x p.y] for p in P.points]...)
    spline = P.spline ? ".." : "--"
    if P.clip
        clipstatement = "clip($pathname);\n"
    else
        clipstatement = ""
    end
    if length(P.points) > 10
        AsyString("""
        file pathdata = input("path{IDENTIFIER}.csv");
        real[][] A = pathdata.csv().dimension(0,2);
        close(pathdata);
        guide $pathname;
        for(int i=0; i<A.length; ++i){
            $pathname = $pathname $spline (A[i][0],A[i][1]);
        }
        $pathname = $pathname $spline cycle;
        $clipstatement$drawcommand($(filterjoin(pathname,fillpen,drawpen)));
        """,
        Dict("path{IDENTIFIER}.csv"=>coords))
    else
        AsyString("""
        path $pathname = $(join([string(p) for p in P.points],spline))--cycle;
        $clipstatement$drawcommand($(filterjoin(pathname,fillpen,drawpen)));
        """)
    end
end

function Base.show(io::IO,P::Polygon2D)
    l = length(P.points)
    s = l ≠ 1 ? "s" : ""
    kwargs = kwargstring(P,_DEFAULT_POLYGON2D_KWARGS)
    print(io,"""
    Polygon2D(<$l point$s>$kwargs)""")
end

#--- GRAPHIC STRINGS ----------------------------------
#------------------------------------------------------

"""
    RawString(s::AbstractString)

Container for directly inserting Asymptote drawing
commands

# Examples
```julia-repl
julia> Plot([Circle((0,0),1),RawString2D("draw((0,0)--dir(20));")])
```
"""
struct RawString <: GraphicString
    s::AbstractString
end

function AsyString(S::RawString)
    AsyString(S.s)
end

function Base.show(io::IO,S::RawString)
    print(io,"RawString($(enclosequote(S.s)))")
end

struct Label2D <: GraphicElement2D
    s::AbstractString
    location::Vec2
    pen::Pen
    rotation::Float64
end

const _DEFAULT_LABEL2D_KWARGS =
    OrderedDict(
        :pen => Pen(),
        :rotation => 0.0
    )

Label2D(s::AbstractString,v::Vec2;kwargs...) =
    Label2D(s,v,updatedvals(_DEFAULT_LABEL2D_KWARGS,
                                process_pen_kwargs(kwargs))...)
Label(s::AbstractString,v::Vec2;kwargs...) = Label2D(s,v;kwargs...)
Label2D(s::AbstractString,
        P::Tuple{<:Real,<:Real};kwargs...) =
                                        Label2D(s,Vec2(P);kwargs...)
Label(s::AbstractString,P::Tuple{<:Real,<:Real};kwargs...) =
                                            Label2D(s,P;kwargs...)
Label2D(s,x::Real,y::Real;kwargs...) =
                Label2D(s,Vec2(x,y);kwargs...)
Label(s::AbstractString,x::Real,y::Real;kwargs...) = Label2D(s,x,y;kwargs...)
Label2D(s::AbstractString,z::Union{Real,Complex};kwargs...) =
                                            Label2D(s,reim(z)...;kwargs...)
Label(s::AbstractString,z::Union{Real,Complex};kwargs...) = Label2D(s,z;kwargs...)

function Base.show(io::IO,L::Label2D)
    kwargs = kwargstring(L,_DEFAULT_LABEL2D_KWARGS)
    print(io,"Label2D($(enclosequote(L.s)),($(L.location.x),$(L.location.y))$kwargs)")
end

#--- Pixel Maps ---------------------------------------

"""
    PixelMap(pixels, lowerleft, upperright; kwargs...)

A graphics primitive representing a two-dimensional array of
pixels situated in the box with given lower left and upper
right corners.

`pixels` is an array of `NamedColor`s, while `lowerleft`
and `upperright` are tuples.

# Examples
```julia-repl
julia> pixels = [x*NamedColor("blue") +
                    (1-x)*NamedColor("red")
                        for x in 0:0.1:1, y in 0:0.1:1]
julia> PixelMap(pixels, (0,0), (1,1))
PixelMap(<11×11>,[0,1]×[0,1])
```
"""
struct PixelMap <: GraphicElement2D
    pixels::Array{NamedColor,2}
    alpha::Array{<:Real,2}
    lowerleft::Tuple{Real,Real}
    upperright::Tuple{Real,Real}
    smooth::Bool
    antialias::Bool
end

const _DEFAULT_PIXELMAP_KWARGS =
    OrderedDict(
        :smooth => false,
        :antialias => false
    )

function PixelMap(pixels::Array{NamedColor,2},
                  lowerleft::Tuple{Real,Real},
                  upperright::Tuple{Real,Real};kwargs...)
    PixelMap(pixels,ones(size(pixels)...),lowerleft,upperright;kwargs...)
end

function PixelMap(pixels::Array{NamedColor,2},
                  alpha::Array{<:Real,2},
                  lowerleft::Tuple{Real,Real},
                  upperright::Tuple{Real,Real};
                  kwargs...)
    PixelMap(pixels,alpha,lowerleft,upperright,
                updatedvals(_DEFAULT_PIXELMAP_KWARGS,kwargs)...)
end

function Base.show(io::IO,P::PixelMap)
    m,n = size(P.pixels)
    a,b = P.lowerleft
    c,d = P.upperright
    print(io,"PixelMap(<$(m)×$(n)>,[$a,$c]×[$b,$d])")
end

function AsyString(P::PixelMap)
    r(N::NamedColor) = Float64(N.color.r)
    g(N::NamedColor) = Float64(N.color.g)
    b(N::NamedColor) = Float64(N.color.b)

    imagecommands = if P.smooth
        """
        pen[][] pixels = new pen[m][n];
        for(int i=0;i<m;++i){
          for(int j=0;j<n;++j){
              pixels[i][j] = (rgb(redvalues[i][j],
                                  greenvalues[i][j],
                                  bluevalues[i][j]) +
                                  opacity(alphavalues[i][j]));
          }
        }

        image(pixels,$(P.lowerleft),$(P.upperright),antialias=$(P.antialias));
        """
    else
        """
        for(int i=0;i<m;++i){
          for(int j=0;j<n;++j){
            if (alphavalues[i][j] > 0) {
              fill(box((a+i/m*(c-a),b+j/n*(d-b)),
                       (a+(i+1)/m*(c-a),b+(j+1)/n*(d-b))),
                       p=rgb(redvalues[i][j],
                               greenvalues[i][j],
                               bluevalues[i][j])+
                               opacity(alphavalues[i][j]));
            }
          }
        }
        """
    end
    AsyString("""
    int m = $(size(P.pixels,1));
    int n = $(size(P.pixels,2));
    file reds = input("red{IDENTIFIER}.csv");
    real[][] redvalues = reds.csv().dimension(m,n);
    close(reds);
    file blues = input("blue{IDENTIFIER}.csv");
    real[][] bluevalues = blues.csv().dimension(m,n);
    close(reds);
    file greens = input("green{IDENTIFIER}.csv");
    real[][] greenvalues = greens.csv().dimension(m,n);
    close(reds);
    file alphas = input("alpha{IDENTIFIER}.csv");
    real[][] alphavalues = alphas.csv().dimension(m,n);
    close(alphas);

    real a = $(Float64(P.lowerleft[1]));
    real b = $(Float64(P.lowerleft[2]));
    real c = $(Float64(P.upperright[1]));
    real d = $(Float64(P.upperright[2]));

    $imagecommands

    """,
    Dict("red{IDENTIFIER}.csv"=>r.(P.pixels),
         "green{IDENTIFIER}.csv"=>g.(P.pixels),
         "blue{IDENTIFIER}.csv"=>b.(P.pixels),
         "alpha{IDENTIFIER}.csv"=>P.alpha))
end

const _DEFAULT_HEATMAP_KWARGS =
    OrderedDict(
         :colors => ["MidnightBlue",
                     "LightSeaGreen",
                     "Yellow",
                     "Tomato"]
    )

"""
    heatmap(A; colors)
    heatmap(A, lowerleft, upperright; colors)
    heatmap(xs, ys, f::Function)

Plot a heatmap of the values stored in the matrix `A`, using the colormap represented by the vector `colors` of `NamedColors`. Place the resulting `PixelMap` according to the given `lowerleft` and `upperright` tuples (which default to (0,0) and size(A)).

# Examples
```julia-repl
julia> heatmap(0:10, 0:10, (x,y) -> x^2 + y^2)
```
"""
function heatmap(A::Array{<:Real,2};
                 colors=NamedColor.(_DEFAULT_HEATMAP_KWARGS[:colors]),
                 kwargs...)
    heatmap(A, (0,0), size(A); colors=colors, kwargs...)
end

function heatmap(A::Array{<:Real,2},
                 lowerleft::Tuple{Real,Real},
                 upperright::Tuple{Real,Real};
                 colors=NamedColor.(_DEFAULT_HEATMAP_KWARGS[:colors]),
                 kwargs...)
    m,M = extrema(A)
    C = [cmap(colors,(A[i,j]-m)/(M-m)) for i=1:size(A,1),j=1:size(A,2)]
    if :alpha in keys(kwargs)
        alpha = kwargs[:alpha]
        kwargs = [(a,b) for (a,b) in kwargs if a ≠ :alpha]
        Plot(PixelMap(C,alpha,lowerleft,upperright; kwargs...))
    else
        Plot(PixelMap(C,lowerleft,upperright; kwargs...))
    end
end

function heatmap(xs, ys, f::Function; kwargs...)
    heatmap([f(x,y) for x in xs, y in ys],
            (xs[1], ys[1]),
            (xs[end], ys[end]); kwargs...)
end

function cmap(colors::Array{NamedColor,1},r::Real)
    fpart, ipart = modf(1 + r*(length(colors)-1))
    if fpart == 0
        colors[Int(ipart)]
    else
        (1-fpart) * colors[Int(ipart)] + fpart * colors[Int(ipart)+1]
    end
end

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

#--- 2D PLOTS -----------------------------------------
#------------------------------------------------------

"""
    Plot2D(elements::Array{<:GraphicElement,1},
           options::Array{Any,1})

A container for a list of graphics primitives to draw,
together with drawing options

# Examples
```julia-repl
julia> Plot([Path([0 0; 1 1]),Polygon([exp(2*pi*im*k/5) for k=1:5])])
```
"""
struct Plot2D <: Plot
    elements::Array{<:GraphicElement,1}
    options::Array{Any,1}
end

Plot2D(;kwargs...) = Plot2D(GraphicElement[],kwargs...)

Plot2D(elements::Array{<:GraphicElement,1};kwargs...) =
        Plot2D(elements,collect(kwargs))

Plot2D(element::GraphicElement;kwargs...) = Plot2D([element];kwargs...)

function Plot2D(P::Plot2D;kwargs...)
    Plot2D(P.elements,updatedvals(
                    _DEFAULT_PLOT2D_KWARGS,kwargs))
end

Plot(P::Plot2D;kwargs...) = Plot2D(P;kwargs...)

const _DEFAULT_PLOT2D_KWARGS =
    OrderedDict(
         :axes => false,
         :axispen => Pen(),
         :xaxisarrow => NoArrow(),
         :yaxisarrow => NoArrow(),
         :axisarrow => NoArrow(),
         :xlabel => "",
         :ylabel => "",
         :packages => ["mathpazo"],
         :xmin => "-infinity",
         :xmax => "infinity",
         :ymin => "-infinity",
         :ymax => "infinity",
         :xticks => "NoTicks",
         :yticks => "NoTicks",
         :ticks => "NoTicks",
         :ignoreaspect => false,
         :width => _DEFAULT_WIDTH,
         :height => "",
         :bgcolor => NamedColor("white"),
         :bgfill => true,
         :border => 3,
         :pdf => true)

function AsyString(P::Plot2D)
    global _DEFAULT_PLOT2D_KWARGS
    D = deepcopy(_DEFAULT_PLOT2D_KWARGS)
    for k in keys(Dict(P.options))
        if !(k in keys(D))
            error("Unknown Plot2D option $k")
        end
    end
    merge!(D,Dict(P.options))
    for s in [:xmin,:xmax,:ymin,:ymax]
        if isa(D[s],Real)
            D[s] = string(D[s])
        end
    end
    if isa(D[:bgcolor],AbstractString)
        D[:bgcolor] = NamedColor(D[:bgcolor])
    end
    if isa(D[:axispen],AbstractString) || isa(D[:axispen],NamedColor)
        D[:axispen] = convert(Pen,D[:axispen])
    end
    if D[:axisarrow] ≠ _DEFAULT_PLOT2D_KWARGS[:axisarrow]
        D[:xaxisarrow] = D[:axisarrow]
        D[:yaxisarrow] = D[:axisarrow]
    end
    if D[:ticks] ≠ _DEFAULT_PLOT2D_KWARGS[:ticks]
        D[:xticks] = D[:ticks]
        D[:yticks] = D[:ticks]
    end
    pen = (is_no_pen(D[:axispen]) || isdefault(D[:axispen]) ?
                "" :
                """\n$(" "^9)p   = $(D[:axispen]),""")
    if D[:axes]
        axesstring = """
        xaxis(   L   = $(enclosequote(D[:xlabel])),
               arrow = $(AsyString(D[:xaxisarrow])),
               xmin  = $(D[:xmin]),
               xmax  = $(D[:xmax]),$pen
               ticks = $(D[:xticks]));
        yaxis(   L   = $(enclosequote(D[:ylabel])),
               arrow = $(AsyString(D[:yaxisarrow])),
               ymin  = $(D[:ymin]),
               ymax  = $(D[:ymax]),$pen
               ticks = $(D[:yticks]));
        """
    else
        axesstring = ""
    end

    shipout = D[:bgfill] ? "shipout(bbox(FillDraw($(D[:border]),fillpen=$(string(D[:bgcolor])),drawpen=invisible)));" :
        "shipout(bbox($(D[:border]),invisible));"

    pdf = D[:pdf] ? "settings.outformat=\"pdf\";" : ""
    ignoreaspect = D[:ignoreaspect] ? "IgnoreAspect" : ""
    packages = join("usepackage($(enclosequote(s)));"
                                    for s in D[:packages])

    asystrings = map(AsyString,P.elements)
    drawingcommands = join(replace(s.str,"{IDENTIFIER}"=>string(j))
                                for (j,s) in enumerate(asystrings))

    data = [A.data for A in asystrings]
    merged_data = merge([Dict(replace(k,"{IDENTIFIER}"=>j) => v
                for (k,v) in D) for (j,D) in enumerate(data)]...)

    AsyString("""
    import x11colors;
    import graph;
    import palette;

    $pdf

    $packages

    size($(filterjoin(D[:width],D[:height],ignoreaspect)));

    $drawingcommands

    $axesstring

    $shipout
    """,
    merged_data)
end

function +(P::Plot2D,Q::Plot2D)
    global _DEFAULT_PLOT2D_KWARGS
    elements = vcat(P.elements,Q.elements)
    options = copy(_DEFAULT_PLOT2D_KWARGS)
    merge!(options,Dict(P.options))
    merge!(options,Dict(Q.options))
    Plot2D(elements,collect(options))
end

# Convenience functions for plotting -------------------

const blue, red, green = [Colors.parse(Colors.Colorant,c) for
                                c in ("MidnightBlue","DarkRed","SeaGreen")]

"""
    plot(x,y;kwargs...)
    plot(y;kwargs...)

Return a graph of the path with ``x`` and ``y`` values
given by `x` and `y`

`x` defaults to `0:length(y)-1`. `kwargs` are applied to
the `Path2D` object representing the line or to the
containing `Plot2D`, as appropriate

    plot(xs::Vector{<:Vector{<:Real}},
         ys::Vector{<:Vector{<:Real}};
         kwargs...)

Multiple line graphs in the same figure

    plot(x,y,z;kwargs...)
    plot(z::Array{<:Real,2};kwargs...)

A graph of the surface with ``x``, ``y``, and ``z``
values `x`, `y`, and `z`

`x` defaults to `[i-1 for i=1:size(z,1),j=1:size(z,2)]` and
`y` defaults to `[j-1 for i=1:size(z,1),j=1:size(z,2)]`

# Examples
```julia-repl
plot(cumsum(randn(100)))
plot(rand(5,5))
```
"""
function plot(xs::Vector{<:Vector{<:Real}},
              ys::Vector{<:Vector{<:Real}};
              colors=Colors.distinguishable_colors(length(xs),
                        [blue,red,green],
                        lchoices=range(50,stop=100,length=15),
                        cchoices=range(50,stop=100,length=15),
                        hchoices=range(170,stop=340,length=15)),
              kwargs...)
    if length(xs) ≠ length(ys)
      error("List of x vectors and list of y vectors
             should have the same length")
    end
    global _DEFAULT_PATH2D_KWARGS
    pathkwargs, plotkwargs = splitkwargs(process_pen_kwargs(kwargs),
                                                _DEFAULT_PATH2D_KWARGS)
    penkwarg = [(a,b) for (a,b) in pathkwargs if a == :pen]
    if length(penkwarg) > 0
        ps = [[penkwarg[1]] for c in colors]
    else
        ps = [[(:pen,Pen(color=c,linewidth=1.5))] for c in colors]
    end
    if :ticks in keys(kwargs)
        xticks = kwargs[:ticks]
        yticks = kwargs[:ticks]
    else
        xticks = (:xticks in keys(kwargs)) ? kwargs[:xticks] : "Ticks(OmitTick(0))"
        yticks = (:yticks in keys(kwargs)) ? kwargs[:yticks] : "Ticks(OmitTick(0))"
    end
    Plot2D([Path2D(x,y;p...,pathkwargs...) for (x,y,p) in zip(xs,ys,ps)];
                   ignoreaspect=true,
                   axes=true,
                   xticks=xticks,
                   yticks=yticks,
                   plotkwargs...)
end

function plot(x,ys::Vector{<:Array{<:Real,1}};kwargs...)
    plot(fill(collect(x),length(ys)),ys;kwargs...)
end

function plot(ys::Vector{<:Array{<:Real,1}};kwargs...)
    @assert length(Set(map(length,ys))) == 1
    plot(collect(1:length(ys[1])),ys;kwargs...)
end

function plot(x::Union{Vector{<:Real},UnitRange},
              y::Union{Vector{<:Real},UnitRange};kwargs...)
    plot([collect(x)],[collect(y)];kwargs...)
end

function plot(y::Array{<:Real,1};kwargs...)
    plot(collect(0:length(y)-1),y;kwargs...)
end

function plot(fs,
              a::Real,
              b::Real;
              n::Integer=100,
              kwargs...)
    x = range(a,stop=b,length=n)
    plot(collect(x),[f.(x) for f in fs];kwargs...)
end

function plot(f::Function,a::Real,b::Real;kwargs...)
    plot([f],a,b;kwargs...)
end

function plot(f_or_fs,t::Tuple{<:Real,<:Real};kwargs...)
    plot(f_or_fs,t...;kwargs...)
end