# Plotting Options
using Printf
using Plots
using LaTeXStrings
using BSON
pgfplotsx()

includet("./models/dense_rren.jl")

theme(:default)

ms = 10
lw = 2.0
line_alpha = 0.8
c = palette(:default);

plot_args = (xaxis = :log10, seriestype = :scatter)
xaxis = (L"\gamma", (0, 1000), :log, font(20, "Courier"))

function plot_gamma_vs_perf!(data; kwargs...)
    plot!([data.gamma_lower], [mean.([data.val[1]["nrmse"]])]; kwargs...)
end

function plot_gamma_vs_perf!(data, gamma_upper; kwargs...)
    plot!([data.gamma_lower], [mean.([data.val[1]["nrmse"]])]; kwargs..., plot_args...)
    plot!([gamma_upper, gamma_upper], [0, 1.0];  color=kwargs[:color], label=nothing, ls=:dashdot)
end

function format!()
    tfont = font(22, "Courier")
    lfont = font(26, "Courier")

    plot!(;xlabel=L"\underline{\gamma}", xtickfont=tfont, labelfontsize=24, legend=(0.61, 0.99))
    plot!(;ylabel="NRMSE", ytickfont=tfont, legendfontsize=14, framestyle=:box)
end

# ## Plot data F16 Vibration Test
# mf16_ren_4p0 = load_sys_id_test("./results/sys_id/f16/bounded_ren_75_150_4.0")[2]
# mf16_ren_6p0 = load_sys_id_test("./results/sys_id/f16/bounded_ren_75_150_6.0")[2]
# mf16_ren_10p0 = load_sys_id_test("./results/sys_id/f16/bounded_ren_75_150_10.0")[2]
# mf16_ren_20p0 = load_sys_id_test("./results/sys_id/f16/bounded_ren_75_150_20.0")[2]
# mf16_ren_40p0 = load_sys_id_test("./results/sys_id/f16/bounded_ren_75_150_40.0")[2]
# mf16_ren_60p0 = load_sys_id_test("./results/sys_id/f16/bounded_ren_75_150_60.0")[2]
# mf16_ren_100p0 = load_sys_id_test("./results/sys_id/f16/bounded_ren_75_150_100.0")[2]
# mf16_ren_stable = load_sys_id_test("./results/sys_id/f16/stable_ren_75_150")[2]

# # ## Plot data F16 Vibration Test
# mf16_bounded_4p0 = load_sys_id_test("./results/sys_id_cdc2021/f16/bounded_rren_75_150_4.0")[2]
# mf16_bounded_6p0 = load_sys_id_test("./results/sys_id_cdc2021/f16/bounded_rren_75_150_6.0")[2]
# mf16_bounded_10p0 = load_sys_id_test("./results/sys_id_cdc2021/f16/bounded_rren_75_150_10.0")[2]
# mf16_bounded_20p0 = load_sys_id_test("./results/sys_id_cdc2021/f16/bounded_rren_75_150_20.0")[2]
# mf16_bounded_40p0 = load_sys_id_test("./results/sys_id_cdc2021/f16/bounded_rren_75_150_40.0")[2]
# mf16_bounded_60p0 = load_sys_id_test("./results/sys_id_cdc2021/f16/bounded_rren_75_150_60.0")[2]
# mf16_bounded_100p0 = load_sys_id_test("./results/sys_id_cdc2021/f16/bounded_rren_75_150_100.0")[2]
# mf16_stable = load_sys_id_test("./results/sys_id_cdc2021/f16/stable_rren_75_150")[2]

# mf16_ff_4p0 = load_sys_id_test("./results/sys_id/f16/bounded_ffren_75_150_4.0")[2]
# mf16_ff_6p0 = load_sys_id_test("./results/sys_id/f16/bounded_ffren_75_150_6.0")[2]
mf16_ff_10p0 = load_sys_id_test("./results/sys_id/f16/bounded_ffren_75_150_10.0")[2]
mf16_ff_20p0 = load_sys_id_test("./results/sys_id/f16/bounded_ffren_75_150_20.0")[2]
mf16_ff_40p0 = load_sys_id_test("./results/sys_id/f16/bounded_ffren_75_150_40.0")[2]
# mf16_ff_60p0 = load_sys_id_test("./results/sys_id/f16/bounded_ffren_75_150_60.0")[2]
# mf16_ff_100p0 = load_sys_id_test("./results/sys_id/f16/bounded_ffren_75_150_100.0")[2]
mf16_ff_stable = load_sys_id_test("./results/sys_id/f16/stable_ffren_75_150")[2]

# mf16_rrnn_4p0 = load_sys_id_test("./results/sys_id/f16/robust_rnn_lipschitz75_150_4.0")[2]
# mf16_rrnn_6p0 = load_sys_id_test("./results/sys_id/f16/robust_rnn_lipschitz75_150_6.0")[2]
# mf16_rrnn_10p0 = load_sys_id_test("./results/sys_id/f16/robust_rnn_lipschitz75_150_10.0")[2]
# mf16_rrnn_20p0 = load_sys_id_test("./results/sys_id/f16/robust_rnn_lipschitz75_150_20.0")[2]
# mf16_rrnn_40p0 = load_sys_id_test("./results/sys_id/f16/robust_rnn_lipschitz75_150_40.0")[2]
# mf16_rrnn_60p0 = load_sys_id_test("./results/sys_id/f16/robust_rnn_lipschitz75_150_60.0")[2]
# mf16_rrnn_100p0 = load_sys_id_test("./results/sys_id/f16/robust_rnn_lipschitz75_150_100.0")[2]
# mf16_rrnn_stable = load_sys_id_test("./results/sys_id/f16/robust_rnn_stable75_150")[2]

mf16_lstm = load_sys_id_test("./results/sys_id/f16/lstm_170")[2]
mf16_rnn = load_sys_id_test("./results/sys_id/f16/rnn_340")[2]

# gr()
pgfplotsx()
c = palette(:default);
lw = 1.5
plot_font = "Computer Modern"

default(fontfamily=plot_font, labelfontsize=22, legendfontsize=14, tickfontsize=16,
        linewidth=lw, framestyle=:box, label=nothing, grid=true)

p = plot()

plot_gamma_vs_perf!(mf16_ff_10p0, 10.0; label=L"\textrm{R-aREN } \gamma < 10", color=c[2], m=:diamond, msize=ms)
plot_gamma_vs_perf!(mf16_ff_20p0, 20.0; label=L"\textrm{R-aREN }{\gamma <20}", color=c[3], m=:diamond, msize=ms)
plot_gamma_vs_perf!(mf16_ff_40p0, 40.0; label=L"\textrm{R-aREN }{\gamma <40}", color=c[4], m=:diamond, msize=ms)
plot_gamma_vs_perf!(mf16_ff_stable; label=L"\textrm{C-aREN }", color=c[7], m=:diamond, msize=ms)

# Data comes from python impletmentation of Robust RNN
# Data can be found at https://github.com/imanchester/RobustRNN
# branch f16_exp - results/f16
plot!([9.045], [0.5610889]; label=L"\textrm{Robust RNN }{\gamma<10}", color=c[2], m=:dtriangle, msize=ms)
plot!([18.546], [0.475547]; label=L"\textrm{Robust RNN }{\gamma<20}", color=c[3], m=:dtriangle, msize=ms)
plot!([38.846], [0.325848]; label=L"\textrm{Robust RNN }{\gamma<40}", color=c[4], m=:dtriangle, msize=ms)
plot!([148.340], [0.243539]; label=L"\textrm{Robust RNN }{\gamma<\infty}", color=c[7], m=:dtriangle, msize=ms)
# plot!([49.896], [0.335553]; label=L"\textrm{Robust RNN }{\gamma<\infty}", color=c[7], m=:dtriangle, msize=ms)

plot_gamma_vs_perf!(mf16_lstm; label="LSTM", color=c[1], m=:rect, msize=ms)
plot_gamma_vs_perf!(mf16_rnn; label="RNN", color=c[1], m=:circle, msize=ms)

plot!(;xlabel=L"\underline{\gamma}", ylabel="NRMSE", legend=(0.59, 0.98))
ylims!(0.1, 0.60)
minor = [j*10.0^i for i=(0-1):(3+1) for j=2:9]
plot!(; xticks=((10, 100, 1000), ("10", "100", "1000")), minorticks=minor)

savefig(p, "./results/sys_id/f16.pdf")

# plot_gamma_vs_perf!(mf16_lstm; label="LSTM", color=c[1], m=:rect, msize=ms)
# plot_gamma_vs_perf!(mf16_rnn; label="RNN", color=c[1], m=:circle, msize=ms)

# ylims!(0.18, 0.75)
# savefig(p, "./results/sys_id/f16.png")

mean(mf16_ff_stable.val[1]["nrmse"][:])
mean(mf16_bounded_100p0.val[1]["nrmse"][:])

# mf16_bounded_10p0.gamma_lower
# mf16_bounded_20p0.gamma_lower
# mf16_bounded_40p0.gamma_lower
# mf16_bounded_60p0.gamma_lower
# mf16_bounded_100p0.gamma_lower

# mf16_ff_10p0.gamma_lower
# mf16_ff_20p0.gamma_lower
# mf16_ff_40p0.gamma_lower
# mf16_ff_60p0.gamma_lower
# mf16_ff_100p0.gamma_lower
# mf16_ff_stable.gamma_lower

# # Example outputs for f16 dataset & 91.0
# function make_plots()
#     plot_array = Any[]
#     range = 3500:3750
#     for output = 1:3
#         for batch = 1:3
#             y = extract(mf16_bounded_20p0.val[1]["y"], output, batch)[range]
#             yest1 = extract(mf16_bounded_4p0.val[1]["yest"], output, batch)[range]
#             # yest2 = extract(mf16_bounded_10p0.val[1]["yest"], output, batch)[range]
#             yest3 = extract(mf16_bounded_20p0.val[1]["yest"], output, batch)[range]
#             yest6 = extract(mf16_bounded_100p0.val[1]["yest"], output, batch)[range]
#             yest7 = extract(mf16_stable.val[1]["yest"], output, batch)[range]
#             yest8 = extract(mf16_lstm.val[1]["yest"], output, batch)[range]

#             push!(plot_array, plot())
#             plot!(y; c=:black, lw=3)
#             plot!(yest1; c=c[1],lw=1.5)
#             # plot!(yest2; c=c[2],lw=1.5)
#             plot!(yest3; c=c[2], lw=1.5)
#             plot!(yest7; c=c[3], lw=1.5)
#             plot!(yest8; c=c[4], lw=1.5,  legend=nothing)
#         end
#     end
#     return plot(plot_array..., layouts=(3, 3))
# end
# p = make_plots()
# ylims!(-4, 4)
# savefig(p, "./results/sys_id/f16_outputs.png")

# Plot training loss
# c1 = cgrad([:orange, :blue], LinRange(0, 1, 6))
c1 = palette(:thermal, 5)

p = plot(mf16_ff_4p0.tloss, label=L"\textrm{R-aREN }{\gamma < 4}", c=c1[4])
plot!(mf16_ff_10p0.tloss, label=L"\textrm{R-aREN }{\gamma < 10}", c=c1[3])
plot!(mf16_ff_40p0.tloss, label=L"\textrm{R-aREN }{\gamma < 40}", c=c1[2])
plot!(mf16_ff_stable.tloss, label=L"\textrm{C-aREN }", c=c1[1])

plot!(mf16_lstm.tloss, label="LSTM", lw=1.5, c=c[1])
plot!(mf16_rnn.tloss, label="RNN", lw=1.5, c=c[3])
plot!(;xlabel="Epochs", ylabel="Loss" ,yaxis=:log10, legend = (0.66, 0.98))
ylims!(0.5, 11.1)

plot!(;yticks=((1, 10), ("1", "10")))
savefig(p, "./results/sys_id/f16_training.pdf")

# Plot Wiener Hammerstein Results
wh_ff_1p0 = load_sys_id_test("./results/sys_id/wiener_hammerstein/bounded_ffren_40_100_1.0")[2]
wh_ff_1p5 = load_sys_id_test("./results/sys_id/wiener_hammerstein/bounded_ffren_40_100_1.5")[2]
wh_ff_2p0 = load_sys_id_test("./results/sys_id/wiener_hammerstein/bounded_ffren_40_100_2.0")[2]
wh_ff_2p5 = load_sys_id_test("./results/sys_id/wiener_hammerstein/bounded_ffren_40_100_2.5")[2]
wh_ff_3p0 = load_sys_id_test("./results/sys_id/wiener_hammerstein/bounded_ffren_40_100_3.0")[2]
wh_ff_3p5 = load_sys_id_test("./results/sys_id/wiener_hammerstein/bounded_ffren_40_100_3.5")[2]
wh_ff_5p0 = load_sys_id_test("./results/sys_id/wiener_hammerstein/bounded_ffren_40_100_5.0")[2]
wh_ff_8p0 = load_sys_id_test("./results/sys_id/wiener_hammerstein/bounded_ffren_40_100_8.0")[2]
wh_ff_15p0 = load_sys_id_test("./results/sys_id/wiener_hammerstein/bounded_ffren_40_100_15.0")[2]
wh_ff_stable = load_sys_id_test("./results/sys_id/wiener_hammerstein/stable_ffren_40_100")[2]
wh_lstm = load_sys_id_test("./results/sys_id/wiener_hammerstein/lstm_100")[2]
wh_rnn = load_sys_id_test("./results/sys_id/wiener_hammerstein/rnn_200")[2]

# gr()
pgfplotsx()

p = plot()
plot_gamma_vs_perf!(wh_ff_1p5, 1.5; label=L"\textrm{R-aREN } {\gamma < 1.5}", color=c[2], m=:diamond, msize=ms)
plot_gamma_vs_perf!(wh_ff_2p5, 2.5; label=L"\textrm{R-aREN } {\gamma < 2.5}", color=c[3], m=:diamond, msize=ms)
plot_gamma_vs_perf!(wh_ff_3p5, 3.5; label=L"\textrm{R-aREN } {\gamma < 3.5}", color=c[4], m=:diamond, msize=ms)
plot_gamma_vs_perf!(wh_ff_stable; label=L"\textrm{C-aREN }", color=c[7], m=:diamond, msize=ms)
plot_gamma_vs_perf!(wh_lstm; label="LSTM", color=c[1], m=:utriangle, msize=ms)
plot_gamma_vs_perf!(wh_rnn; label="RNN", color=c[1], m=:circle, msize=ms)

ylims!(0.25, 0.5)
xlims!(1.3, 1000)
plot!(;xlabel=L"\underline{\gamma}", ylabel="NRMSE", legend=(0.64, 0.43))
plot!(; xticks=((10, 100, 1000), ("10", "100", "1000")))
savefig(p, "./results/sys_id/wiener_hammerstein.pdf")

# Plot adversarial perterbations for f16 example
function calc_pert(data; pert_size=0.05)
    model = data.model
    u1 = data.u1
    u2 = data.u2

    du = u2 - u1
    Δu = pert_size * du / norm(du)

    x0 = init_state(model, 1)
    y1 = model(x0, u1)[1]
    y2 = model(x0, u1 + Δu)[1]

    Δy = y2 - y1
    println("Sensitivity: ", norm(Δy) / norm(Δu))

    return Δy, Δu
end

pgfplotsx()
c = palette(:default);
lw = 1.5
plot_font = "Computer Modern"

default(fontfamily=plot_font, labelfontsize=22, legendfontsize=14, tickfontsize=16,
        linewidth=lw, framestyle=:box, label=nothing, grid=true)

ffren_10_dat = BSON.load("./results/sys_id/f16_adversarial_perturbations/ffren_10.bson")[:data]
ffren_40_dat = BSON.load("./results/sys_id/f16_adversarial_perturbations/ffren_40.bson")[:data]
rnn_dat = BSON.load("./results/sys_id/f16_adversarial_perturbations/rnn.bson")[:data]
lstm_dat = BSON.load("./results/sys_id/f16_adversarial_perturbations/lstm.bson")[:data]

rnn_dy, rnn_du = calc_pert(rnn_dat)
lstm_dy, lstm_du = calc_pert(lstm_dat)
ffren_40_dy, ffren_40_du = calc_pert(ffren_40_dat)
ffren_10_dy, ffren_10_du = calc_pert(ffren_10_dat)

k = 1
p=plot(;ylabel=L"\Delta y_1")
plot!(extract(rnn_dy, k, 1); label=L"\textrm{RNN}", c=c[4])
plot!(extract(lstm_dy, k, 1); label=L"\textrm{LSTM}", c=c[3])
plot!(extract(ffren_40_dy, k, 1); label=L"\textrm{R-aREN} \gamma<40", c=c[2])
plot!(extract(ffren_10_dy, k, 1); label=L"\textrm{R-aREN} \gamma<10", c=c[1])
plot!(;xlabel="Time Steps", legend=(0.02, 0.98))
savefig(p, "./results/sys_id/adv_pert.pdf")

# Zoomed in version of the above plot
plot_range = 1000:1500
p=plot(;ylabel=L"\Delta y_1")
plot!(plot_range, extract(rnn_dy[plot_range], k, 1); label=L"\textrm{RNN}", c=c[4])
plot!(plot_range, extract(lstm_dy[plot_range], k, 1); label=L"\textrm{LSTM}", c=c[3])
plot!(plot_range, extract(ffren_40_dy[plot_range], k, 1); label=L"\textrm{R-aREN} \gamma<40", c=c[2])
plot!(plot_range, extract(ffren_10_dy[plot_range], k, 1); label=L"\textrm{R-aREN} \gamma<10", c=c[1])
plot!(;xlabel="Time Steps", legend=(0.01, 0.98))
savefig(p, "./results/sys_id/adv_pert_zoomed.pdf")


function header(f::IOStream)
    sdf_magic = simple_str(read(f, 4))
    endianness = read(f, Int32)
    file_version = read(f, Int32)
    file_revision = read(f, Int32)
    code_name = simple_str(read(f, ID_LENGTH))
    first_block_location = read(f, Int64)
    summary_location = read(f, Int64)
    summary_size = read(f, Int32)
    nblocks = read(f, Int32)
    block_header_length = read(f, Int32)
    step = read(f, Int32)
    time = read(f, Float64)
    jobid1 = read(f, Int32)
    jobid2 = read(f, Int32)
    string_length = read(f, Int32)
    code_io_version = read(f, Int32)

    Header(
        sdf_magic,
        endianness,
        file_version,
        file_revision,
        code_name,
        first_block_location,
        summary_location,
        summary_size,
        nblocks,
        block_header_length,
        step,
        time,
        jobid1,
        jobid2,
        string_length,
        code_io_version,
    )
end

header(filename::AbstractString) = open(header, filename)

function Base.read(f, block::BlockHeader{T,D,BLOCKTYPE_CONSTANT}) where {T,D}
    val = read(f, T)
    ConstantBlockHeader(block, val)
end

function Base.read(f, block::BlockHeader{T,N,BLOCKTYPE_PLAIN_MESH}) where {T,N}
    dims = Array{Int32, 1}(undef, N)

    mults, labels, units, geometry, minval, maxval = read_mesh_common!(f,N)
    read!(f, dims)

    PlainMeshBlockHeader(
        block,
        mults,
        labels,
        units,
        geometry,
        minval,
        maxval,
        dims,
    )
end

function Base.read(f, block::BlockHeader{T,N,BLOCKTYPE_POINT_MESH}) where {T,N}
    mults, labels, units, geometry, minval, maxval = read_mesh_common!(f,N)
    np = read(f, Int64)

    PointMeshBlockHeader(
        block,
        mults,
        labels,
        units,
        geometry,
        minval,
        maxval,
        np,
    )
end

function Base.read(f, block::BlockHeader{T,D,BLOCKTYPE_PLAIN_VARIABLE}) where {T,D}
    dims = Array{Int32, 1}(undef, D)

    mult, units, mesh_id = read_variable_common!(f)
    read!(f, dims)
    stagger = Stagger(read(f, Int32))

    PlainVariableBlockHeader(block, mult, units, mesh_id, Tuple(dims), stagger)
end

function Base.read(f, block::BlockHeader{T,D,BLOCKTYPE_POINT_VARIABLE}) where {T,D}
    mult, units, mesh_id = read_variable_common!(f)
    np = read(f, Int64)

    PointVariableBlockHeader(block, mult, units, mesh_id, np)
end

function Base.read(f, block::BlockHeader{T,D,BLOCKTYPE_RUN_INFO}) where {T,D}
    code_version = read(f, Int32)
    code_revision = read(f, Int32)
    commit_id = simple_str(read(f, ID_LENGTH))
    sha1sum = simple_str(read(f, ID_LENGTH))
    compile_machine = simple_str(read(f, ID_LENGTH))
    compile_flags = simple_str(read(f, ID_LENGTH))
    defines = read(f, Int64)
    compile_date = read(f, Int32)
    run_date = read(f, Int32)
    io_date = read(f, Int32)
    # TODO save these fields

    RunInfoBlockHeader(block)
end

function Base.read(f, block::BlockHeader{T,D,BLOCKTYPE_CPU_SPLIT}) where {T,D}
    CPUSplitBlockHeader(block)
end

function read_variable_common!(f)
    mult = read(f, Float64)
    units = simple_str(read(f, ID_LENGTH))
    mesh_id = simple_str(read(f, ID_LENGTH))

    mult, units, mesh_id
end

function read_mesh_common!(f, n)
    minval = Array{Float64, 1}(undef, n)
    maxval = Array{Float64, 1}(undef, n)

    mults = ntuple(n) do i
        read(f, Float64)
    end

    labels = ntuple(n) do i
        simple_str(read(f, ID_LENGTH))
    end
    units = ntuple(n) do i
        simple_str(read(f, ID_LENGTH))
    end
    geometry = Geometry(read(f, Int32))
    read!(f, minval)
    read!(f, maxval)

    mults, labels, units, geometry, minval, maxval
end

simple_str(s) = replace(String(rstrip(String(s))), "\0" => "")


import Base.show
using CUDA: CuArray

# There is a data loader PR: https://github.com/FluxML/Flux.jl/pull/450
mutable struct DataSetIterator
    raw_x::AbstractArray
    raw_y::AbstractArray
    index::Int
    batch_size::Int
    nbatch::Int
end
function DataSetIterator(x,y,batch_size)
    DataSetIterator(x, y, 1, batch_size,
                    convert(Int, floor(size(x)[end]/batch_size)))
end

mutable struct CuDataSetIterator
    gpu_x::Array{CuArray}
    gpu_y::Array{CuArray}
    index::Int
    batch_size::Int
    nbatch::Int
end

function merge_dses(dses)
    # merge datasets
    # 1. merge raw_x and raw_y
    xdim = length(size(dses[1].raw_x))
    ydim = length(size(dses[1].raw_y))
    @show xdim, ydim
    # 4
    raw_x = cat([ds.raw_x for ds in dses]..., dims=xdim)
    # 3
    raw_y = cat([ds.raw_y for ds in dses]..., dims=ydim)
    # 2. assert batch_size
    batch_size = dses[1].batch_size
    # 4. construct new ds
    DataSetIterator(raw_x, raw_y, batch_size)
end

function Base.convert(::Type{CuDataSetIterator}, ds::DataSetIterator)
    CuDataSetIterator(ds.raw_x, ds.raw_y, ds.batch_size)
end

function CuDataSetIterator(ds::DataSetIterator)
    CuDataSetIterator(ds.raw_x, ds.raw_y, ds.batch_size)
end

function CuDataSetIterator(x,y,batch_size)
    xdims = collect(repeated(:, length(size(x)) - 1))
    ydims = collect(repeated(:, length(size(y)) - 1))

    N = size(x)[end]
    nbatch = Int(floor(N/batch_size))

    # reshuffle
    indices = shuffle(1:size(x)[end])
    x = x[xdims...,indices]
    y = y[ydims...,indices]

    gpu_x = map(1:nbatch) do i
        start = (i-1) * batch_size + 1
        stop = i * batch_size
        x[xdims..., start:stop]
    end
    gpu_y = map(1:nbatch) do i
        start = (i-1) * batch_size + 1
        stop = i * batch_size
        y[ydims..., start:stop]
    end
    # FIXME .|>
    gpu_x = gpu_x |> gpu
    gpu_y = gpu_y |> gpu
    CuDataSetIterator(gpu_x, gpu_y, 1, batch_size, nbatch)
end

Base.show(io::IO, x::DataSetIterator) = begin
    println(io, "DataSetIterator:")
    println(io, "  batch size: ", x.batch_size)
    println(io, "  number of batches: ", x.nbatch)
    println(io, "  x data shape: ", size(x.raw_x))
    println(io, "  y data shape: ", size(x.raw_y))
    print(io, "  current index: ", x.index)
end

Base.show(io::IO, x::CuDataSetIterator) = begin
    println(io, "CuDataSetIterator:")
    println(io, "  batch size: ", x.batch_size)
    println(io, "  number of batches: ", x.nbatch)
    println(io, "  x data shape: ", size(x.gpu_x[1]))
    println(io, "  y data shape: ", size(x.gpu_y[1]))
    print(io, "  current index: ", x.index)
end

function next_batch!(ds::DataSetIterator)
    if ds.index > ds.nbatch
        ds.index = 1
    end

    xdims = collect(repeated(:, length(size(ds.raw_x)) - 1))
    ydims = collect(repeated(:, length(size(ds.raw_y)) - 1))

    if ds.index == 1
        indices = shuffle(1:size(ds.raw_x)[end])
        ds.raw_x = ds.raw_x[xdims...,indices]
        ds.raw_y = ds.raw_y[ydims...,indices]
    end

    a = (ds.index - 1) * ds.batch_size + 1
    b = ds.index * ds.batch_size

    ds.index += 1
    # FIXME use view instead of copy
    # FIXME properly slice with channel-last format
    return ds.raw_x[xdims...,a:b], ds.raw_y[ydims...,a:b]
end

function next_batch!(ds::CuDataSetIterator)
    if ds.index > ds.nbatch
        ds.index = 1
    end
    ds.index += 1
    return ds.gpu_x[ds.index-1], ds.gpu_y[ds.index-1]
end

_next_batch_step = 0
function next_batch!(dses::Array{T, N}
                     where T<:Union{DataSetIterator, CuDataSetIterator}
                     where N)
    # keep an internal step?
    global _next_batch_step
    _next_batch_step += 1
    next_batch!(dses[_next_batch_step % length(dses) + 1])
end

"""
1. shuffle and divide *all* data into batches
2. move *all* data to GPU
3. each time next_batch! is called, a pointer to the GPU array is returned
"""
function move_to_gpu!(ds::DataSetIterator)
    xdims = collect(repeated(:, length(size(ds.raw_x)) - 1))
    ydims = collect(repeated(:, length(size(ds.raw_y)) - 1))

    N = size(ds.raw_x)[end]

    # reshuffle
    indices = shuffle(1:size(ds.raw_x)[end])
    ds.raw_x = ds.raw_x[xdims...,indices]
    ds.raw_y = ds.raw_y[ydims...,indices]

    gpu_x = map(1:Int(floor(N/ds.batch_size))) do i
        start = (i-1) * ds.batch_size + 1
        stop = i * ds.batch_size
        ds.raw_x[xdims..., start:stop]
    end
    gpu_y = map(1:Int(floor(N/ds.batch_size))) do i
        start = (i-1) * ds.batch_size + 1
        stop = i * ds.batch_size
        ds.raw_y[ydims..., start:stop]
    end
    ds.gpu_x = gpu_x |> gpu
    ds.gpu_y = gpu_y |> gpu
    ds.gpu_index = 1
    nothing
end


module NDPrototype
using LightGraphs
using OrderedCollections
using Unrolled

include("utils.jl")
include("edge_coloring.jl")
include("component_functions.jl")
include("network_structure.jl")
include("construction.jl")
include("coreloop.jl")

end


function Base.map!(
    op::UnaryUnion, C::GBArray, A::GBArray;
    mask = nothing, accum = nothing, desc = nothing
)
    mask, accum = _handlenothings(mask, accum)
    desc = _handledescriptor(desc; in1=A)
    op = getoperator(op, eltype(A))
    accum = getaccum(accum, eltype(C))
    libgb.GrB_Matrix_apply(C, mask, accum, op, parent(A), desc)
    return C
end

Base.map!(op, C::GBArray, A::GBArray; mask = nothing, accum = nothing, desc = nothing) =
    map!(UnaryOp(op), C, A; mask, accum, desc)

function Base.map!(
    op::UnaryUnion, A::GBArray;
    mask = nothing, accum = nothing, desc = nothing
)
    return map!(op, A, A; mask, accum, desc)
end

Base.map!(op, A::GBArray; mask = nothing, accum = nothing, desc = nothing) =
    map!(UnaryOp(op), A; mask, accum, desc)

function Base.map(
    op::UnaryUnion, A::GBArray;
    mask = nothing, accum = nothing, desc = nothing
)
    t = inferoutputtype(A, op)
    return map!(op, similar(A, t), A; mask, accum, desc)
end

Base.map(op, A::GBArray; mask = nothing, accum = nothing, desc = nothing) =
    map(UnaryOp(op), A; mask, accum, desc)

function Base.map!(
    op::BinaryUnion, C::GBArray, x, A::GBArray;
    mask = nothing, accum = nothing, desc = nothing
)
    mask, accum = _handlenothings(mask, accum)
    desc = _handledescriptor(desc; in2=A)
    op = getoperator(op, optype(eltype(A), typeof(x)))
    accum = getaccum(accum, eltype(C))
    libgb.scalarmatapply1st[optype(typeof(x), eltype(A))](C, mask, accum, op, x, parent(A), desc)
    return C
end

function Base.map!(
    op, C::GBArray, x, A::GBArray;
    mask = nothing, accum = nothing, desc = nothing
)
    map!(BinaryOps.BinaryOp(op), C, x, A; mask, accum, desc)
end

function Base.map!(
    op::BinaryUnion, x, A::GBArray;
    mask = nothing, accum = nothing, desc = nothing
)
    return map!(op, A, x, A; mask, accum, desc)
end

function Base.map!(
    op, x, A::GBArray;
    mask = nothing, accum = nothing, desc = nothing
)
    map!(BinaryOps.BinaryOp(op), x, A; mask, accum, desc)
end

function Base.map(
    op::BinaryUnion, x, A::GBArray;
    mask = nothing, accum = nothing, desc = nothing
)
    t = inferoutputtype(A, op)
    return map!(op, similar(A, t), x, A; mask, accum, desc)
end

"""
    map(op::Union{Function, AbstractUnaryOp}, A::GBArray; kwargs...)::GBArray
    map(op::Union{Function, AbstractBinaryOp}, A::GBArray, x; kwargs...)::GBArray
    map(op::Union{Function, AbstractBinaryOp}, x, A::GBArray, kwargs...)::GBArray

Transform a GBArray by applying `op` to each element.

UnaryOps and single argument functions apply elementwise in the usual fashion.
BinaryOps and two argument functions require the additional argument `x` which is
    substituted as the first or second operand of `op` depending on its position.

# Arguments
- `op::Union{Function, AbstractUnaryOp, AbstractBinaryOp}`
- `A::GBArray`
- `x`: Position dependent argument to binary operators.

# Keywords
- `mask::Union{Nothing, GBVecOrMat} = nothing`: optional mask.
- `accum::Union{Nothing, AbstractBinaryOp} = nothing`: binary accumulator operation
    where `C[i,j] = accum(C[i,j], T[i,j])` where T is the result of this function before accum is applied.
- `desc::Union{Nothing, Descriptor} = nothing`
"""
function Base.map(
    op, x, A::GBArray;
    mask = nothing, accum = nothing, desc = nothing
)
    map(BinaryOps.BinaryOp(op), x, A; mask, accum, desc)
end

function Base.map!(
    op::BinaryUnion, C::GBArray, A::GBArray, x;
    mask = nothing, accum = nothing, desc = nothing
)
    mask, accum = _handlenothings(mask, accum)
    desc = _handledescriptor(desc; in1=A)
    op = getoperator(op, optype(eltype(A), typeof(x)))
    accum = getaccum(accum, eltype(C))
    libgb.scalarmatapply2nd[optype(typeof(x), eltype(A))](C, mask, accum, op, parent(A), x, desc)
    return C
end

function Base.map!(
    op, C::GBArray, A::GBArray, x;
    mask = nothing, accum = nothing, desc = nothing
)
    map!(BinaryOps.BinaryOp(op), C, A, x; mask, accum, desc)
end

function Base.map!(
    op::BinaryUnion, A::GBArray, x;
    mask = nothing, accum = nothing, desc = nothing
)
    return map!(op, A, A, x; mask, accum, desc)
end

function Base.map!(
    op, A::GBArray, x;
    mask = nothing, accum = nothing, desc = nothing
)
    map!(BinaryOps.BinaryOp(op), A, A, x; mask, accum, desc)
end

function Base.map(
    op::BinaryUnion, A::GBArray, x;
    mask = nothing, accum = nothing, desc = nothing
)
    t = inferoutputtype(A, op)
    return map!(op, similar(A, t), A, x; mask, accum, desc)
end

function Base.map(
    op, A::GBArray, x;
    mask = nothing, accum = nothing, desc = nothing
)
    map(BinaryOps.BinaryOp(op), A, x; mask, accum, desc)
end

Base.:+(x::valid_union, u::GBArray; mask = nothing, accum = nothing, desc = nothing) =
    map(BinaryOps.PLUS, x, u; mask, accum, desc)
Base.:+(u::GBArray, x::valid_union; mask = nothing, accum = nothing, desc = nothing) =
    map(BinaryOps.PLUS, u, x; mask, accum, desc)

Base.:*(x::valid_union, u::GBArray; mask = nothing, accum = nothing, desc = nothing) =
    map(BinaryOps.TIMES, x, u; mask, accum, desc)
Base.:*(u::GBArray, x::valid_union; mask = nothing, accum = nothing, desc = nothing) =
    map(BinaryOps.TIMES, u, x; mask, accum, desc)

Base.:-(x::valid_union, u::GBArray; mask = nothing, accum = nothing, desc = nothing) =
    map(BinaryOps.MINUS, x, u; mask, accum, desc)
Base.:-(u::GBArray, x::valid_union; mask = nothing, accum = nothing, desc = nothing) =
    map(BinaryOps.MINUS, u, x; mask, accum, desc)

Base.:-(u::GBArray) = map(UnaryOps.AINV, u)

"""
    mask!(C::GBArray, A::GBArray, mask::GBArray)

Apply a mask to matrix `A`, storing the results in C.

"""
function mask!(C::GBArray, A::GBArray, mask::GBArray; structural = false, complement = false)
    desc = Descriptor()
    structural && (desc.structural_mask=true)
    complement && (desc.complement_mask=true)
    map!(identity, C, A; mask, desc)
    return C
end

"""
    mask(A::GBArray, mask::GBArray)

Apply a mask to matrix `A`.
"""
function mask(A::GBArray, mask::GBArray; structural = false, complement = false)
    return mask!(similar(A), A, mask; structural, complement)
end


using .Glimpse
const Gl = Glimpse

using ..Structures
using ..DFControl

function Gl.Diorama(str::Structure)
    dio = Diorama(; name = Symbol("DFControl Structure"))
    add_structure!(dio, str)
    dio[Gl.PointLight].data[1] = Gl.PointLight(; diffuse = 1.5f0, ambient = 1.5f0)
    Gl.center_cameras(dio)
    return dio
end
function add_structure!(dio::Diorama, str::Structure; polyhedra::Vector{Symbol} = Symbol[])
    Entity(dio,
           Gl.assemble_wire_axis_box(; position = zero(Point3f0), x = Structures.ustrip(Structures.a(str)),
                                     y = Structures.ustrip(Structures.b(str)), z = Structures.ustrip(Structures.c(str)),
                                     color = Gl.BLACK)...)
    for at in str.atoms
        origin = Point3f0(Structures.ustrip(at.position_cart)...)
        color = RGB{Float32}(at.element.color...)
        Entity(dio,
               Gl.assemble_sphere(origin; color = color,
                                  radius = Float32(0.2 + at.element.atomic_weight / 400.0f0))...,
               Gl.Text(; str = String(at.name), font_size = 2))
        if sum(at.magnetization) != 0
            Entity(dio,
                   Gl.assemble_arrow(origin, origin + Point3f0(at.magnetization...);
                                     color = color, radius_ratio = 1.5f0,
                                     thickness = 0.05f0)...)
        end
    end
    for el in map(x->x.element, polyhedra)
        for at in str[el]
            distances = sort(distance.((at,), str.atoms))
            count = length(filter(i -> distances[i+1] - distances[i] < 1e-1angstrom),
                           1:length(distances)-1)
            poly = polyhedron(str, count)
        end
    end
    return dio
end


const IMAGENET_MEANS = SVector{3,Float32}(.485, 0.456, 0.406)
const IMAGENET_STDS = SVector{3,Float32}(0.229, 0.224, 0.225)

"""
    ImageTensor{N} <: Block

Block for N+1-dimensional arrays representing an N-dimensional
image with the color channels expanded.

"""
struct ImageTensor{N} <: Block
    nchannels::Int
end

function checkblock(block::ImageTensor{N}, a::AbstractArray{T,M}) where {M,N,T}
    # Tensor has dimensionality one higher and color channels need to be the same
    return (N + 1 == M) && (size(a, M) == block.nchannels)
end

"""
    ImagePreprocessing([; kwargs...]) <: Encoding

Encodes `Image`s by converting them to a common color type `C`,
expanding the color channels and normalizing the channel values.
Additionally, apply pixel-level augmentations passed in as `augmentations`
during `Training`.

Currently works with 2D images only, but this constraint will be
removed in the future.

Encodes
- `Image{2}` -> `ImageTensor{3}`

## Keyword arguments

- `augmentations::`[`DataAugmentation.Transform`](#): Augmentation to apply to every image
    before preprocessing. See [`augs_lighting`](#)
- `buffered = true`: Whether to use inplace transformations. Reduces memory usage.
- `means::SVector = IMAGENET_MEANS`: mean value of each color channel.
- `stds::SVector = IMAGENET_STDS`: standard deviation of each color channel.
- `C::Type{<:Colorant} = RGB{N0f8}`: color type to convert images to.
- `T::Type{<:Real} = Float32`: element type of output

"""
struct ImagePreprocessing{P,N,C <: Color{P,N},T <: Number} <: Encoding
    buffered::Bool
    augmentations::DataAugmentation.Transform
    stats::Tuple{SVector{N},SVector{N}}
    tfms::Dict{Context,DataAugmentation.Transform}
end

function ImagePreprocessing(;
        means::SVector{N}=IMAGENET_MEANS,
        stds::SVector{N}=IMAGENET_STDS,
        augmentations=Identity(),
        C::Type{<:Color{U,N}}=RGB{N0f8},
        T=Float32,
        buffered=true) where {N,U}
    # TODO: tensor of type T
    stats = means, stds
    basetfm = ToEltype(C) |> ImageToTensor{T}() |> Normalize(means, stds)
    if buffered
        tfms = Dict(
            Training() => BufferedThreadsafe(augmentations |> basetfm),
            Validation() => BufferedThreadsafe(basetfm),
            # Inference transform is not buffered since it can have
            # varying sizes
            Inference() => basetfm,
        )
    else
        tfms = Dict(
            Training() => augmentations |> basetfm,
            Validation() => basetfm,
            Inference() => basetfm,
        )
    end

    return ImagePreprocessing{U,N,C,T}(buffered, augmentations, stats, tfms)
end

colorchannels(C::Type{<:Color{T,N}}) where {T,N} = N

function encodedblock(ip::ImagePreprocessing{P,M,C}, ::Image{N}) where {P,M,C,N}
    return ImageTensor{N}(colorchannels(C))
end

function encode(ip::ImagePreprocessing, context, block::Image, data)
    return copy(apply(ip.tfms[context], DataAugmentation.Image(data)) |> itemdata)
end

decodedblock(::ImagePreprocessing, ::ImageTensor{N}) where N = Image{N}()

function decode(ip::ImagePreprocessing, context, block::ImageTensor, data)
    means, stds = ip.stats
    return copy(DataAugmentation.tensortoimage(DataAugmentation.denormalize(data, means, stds)))
end

# Augmentation helper

"""
    augs_lighting([; intensity = 0.2, p = 0.75])

Helper to create a set of lighting transformations for image data. With
probability `p`, applies [`AdjustBrightness`](#)`(intensity)` and
[`AdjustContrast`](#)`(intensity)`.
"""
function augs_lighting(;intensity=0.2, p=0.75)
    return Maybe(AdjustBrightness(intensity), p) |> Maybe(AdjustContrast(intensity), p)
end

# Pretty-printing

function Base.show(io::IO, ip::ImagePreprocessing{P,N,C,T}) where {P,N,C,T}
    show(io, ShowTypeOf(ip))
    fields = (
        buffered = ShowLimit(ip.buffered, limit=80),
        augmentations = ShowLimit(ip.augmentations, limit=80),
    )
    show(io, ShowProps(fields, new_lines=true))
end


struct NumberOfObjectives <: MOI.AbstractModelAttribute end

function MOI.set(model::Optimizer, ::NumberOfObjectives, n::Integer)
    set_intattr!(model.inner, "NumObj", n)
    return
end

function MOI.get(model::Optimizer, ::NumberOfObjectives)
    return get_intattr(model.inner, "NumObj")
end

struct MultiObjectiveFunction <: MOI.AbstractModelAttribute
    index::Int
end

function c_api_setobjectiven(
    model,
    index::Cint,
    priority::Cint,
    weight::Cdouble,
    abstol::Cdouble,
    reltol::Cdouble,
    name::String,
    constant::Cdouble,
    lnz::Cint,
    lind::Vector{Cint},
    lval::Vector{Cdouble}
)
    ret = @grb_ccall(
        setobjectiven,
        Cint,
        (
            Ptr{Cvoid}, Cint, Cint, Cdouble, Cdouble, Cdouble, Ptr{Cchar},
            Cdouble, Cint, Ptr{Cint}, Ptr{Cdouble}
        ),
        model, index, priority, weight, abstol, reltol, name, constant, lnz, lind, lval
    )
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
end

function MOI.set(
    model::Optimizer,
    attr::MultiObjectiveFunction,
    f::MOI.ScalarAffineFunction
)
    num_vars = length(model.variable_info)
    obj = zeros(Float64, num_vars)
    for term in f.terms
        column = _info(model, term.variable_index).column
        obj[column] += term.coefficient
    end
    indices, coefficients = _indices_and_coefficients(model, f)
    _update_if_necessary(model)
    c_api_setobjectiven(
        model.inner,
        Cint(attr.index - 1),
        Cint(0),
        1.0,
        0.0,
        0.0,
        "",
        f.constant,
        Cint(length(indices)),
        Cint.(indices) .- Cint(1),
        coefficients
    )
    _require_update(model)
    return
end

struct MultiObjectivePriority <: MOI.AbstractModelAttribute
    index::Int
end

function MOI.set(
    model::Gurobi.Optimizer, attr::MultiObjectivePriority, priority::Int
)
    set_int_param!(model.inner, "ObjNumber", attr.index - 1)
    set_intattr!(model.inner, "ObjNPriority", priority)
    _require_update(model)
    return
end

function MOI.get(
    model::Gurobi.Optimizer, attr::MultiObjectivePriority
)
    _update_if_necessary(model)
    set_int_param!(model.inner, "ObjNumber", attr.index - 1)
    return get_intattr(model.inner, "ObjNPriority")
end

struct MultiObjectiveWeight <: MOI.AbstractModelAttribute
    index::Int
end

function MOI.set(
    model::Gurobi.Optimizer, attr::MultiObjectiveWeight, weight::Float64
)
    set_int_param!(model.inner, "ObjNumber", attr.index - 1)
    set_dblattr!(model.inner, "ObjNWeight", weight)
    _require_update(model)
    return
end

function MOI.get(
    model::Gurobi.Optimizer, attr::MultiObjectiveWeight
)
    _update_if_necessary(model)
    set_int_param!(model.inner, "ObjNumber", attr.index - 1)
    return get_dblattr(model.inner, "ObjNWeight")
end

struct MultiObjectiveValue <: MOI.AbstractModelAttribute
    index::Int
end

function MOI.get(
    model::Gurobi.Optimizer, attr::MultiObjectiveValue
)
    set_int_param!(model.inner, "ObjNumber", attr.index - 1)
    return get_dblattr(model.inner, "ObjNVal")
end


"""

The exported names are
$(EXPORTS)
"""
module Visualization

using DocStringExtensions
using WriteVTK
using WriteVTK.VTKCellTypes: VTK_VERTEX
using WriteVTK.VTKCellTypes: VTK_LINE
using WriteVTK.VTKCellTypes: VTK_TRIANGLE
using WriteVTK.VTKCellTypes: VTK_QUAD
using WriteVTK.VTKCellTypes: VTK_TETRA
using WriteVTK.VTKCellTypes: VTK_HEXAHEDRON
using WriteVTK.VTKCellTypes: VTK_QUADRATIC_TRIANGLE
using WriteVTK.VTKCellTypes: VTK_QUADRATIC_TETRA
using WriteVTK.VTKCellTypes: VTK_QUADRATIC_EDGE
using WriteVTK.VTKCellTypes: VTK_QUADRATIC_HEXAHEDRON
#using WriteVTK.VTKCellTypes: VTK_BIQUADRATIC_HEXAHEDRON
using WriteVTK.VTKCellTypes: VTK_QUADRATIC_QUAD
using WriteVTK.VTKCellTypes: VTK_BIQUADRATIC_QUAD
using WriteVTK.VTKCellTypes: VTK_WEDGE
using WriteVTK.VTKCellTypes: VTK_PYRAMID
using WriteVTK.VTKCellTypes: VTK_LAGRANGE_QUADRILATERAL
using WriteVTK.VTKCellTypes: VTK_LAGRANGE_TRIANGLE
using WriteVTK.VTKCellTypes: VTK_LAGRANGE_HEXAHEDRON
using WriteVTK.VTKCellTypes: VTK_LAGRANGE_TETRAHEDRON

using Gridap.Helpers
using Gridap.TensorValues
using Gridap.Fields
using Gridap.Arrays
using Gridap.Polynomials
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.ReferenceFEs: _find_unique_with_indices
using Gridap.ReferenceFEs: SerendipityPolytope
using Gridap.FESpaces
using FillArrays
using Gridap.CellData

import Gridap.Geometry: get_reffes
import Gridap.Geometry: get_cell_type
import Gridap.Geometry: get_node_coordinates
import Gridap.Geometry: get_cell_nodes

import AbstractTrees

export writevtk
export createvtk
export write_vtk_file
export print_op_tree
export visualization_data

include("VisualizationData.jl")
include("Vtk.jl")

include("PrintOpTrees.jl")

end # module


using LinearAlgebra: norm
import Base: ==
const BLI = BLInvariants
using JuLIP: decode_dict

BondLengthDesc(transform::String, cutoff) =
         BondLengthDesc(AnalyticTransform(transform), fcut_analyse(cutoff))

# -------------- IO -------------------
Dict(D::BondLengthDesc) = Dict( "__id__"    =>  "BondLengthDesc",
                                "transform" =>  Dict(D.transform),
                                "cutoff"    =>  Dict(D.cutoff) )

BondLengthDesc(D::Dict) = BondLengthDesc( decode_dict(D["transform"]),
                                          decode_dict(D["cutoff"]) )

==(D1::BondLengthDesc, D2::BondLengthDesc) =
      ( (D1.transform == D2.transform) && (D1.cutoff == D2.cutoff) )

Base.convert(::Val{:BondLengthDesc}, D::Dict) = BondLengthDesc(D)

# ------------- Interface Code ---------------

tdegrees(::BondLengthDesc, vN::Val{N}) where {N} = BLI.tdegrees(vN)

@inline ricoords(D::BondLengthDesc, Rs, J) = edge_lengths(Rs, J)

@inline gradri2gradR!(desc::BondLengthDesc, dVsite, dV_dr, Rs, J, r) =
   _grad_len2pos!(dVsite, dV_dr, Rs, J, r)

@inline fcut(D::BondLengthDesc, r) = fcut(D.cutoff, r)
@inline fcut_d(D::BondLengthDesc, r) = fcut_d(D.cutoff, r)

@inline skip_simplex(D::BondLengthDesc, r) = (maximum(r) > cutoff(D.cutoff))

@inline invariants(D::BondLengthDesc, r) = BLI.invariants(transform.(Ref(D), r))

@inline function invariants_ed(D::BondLengthDesc, r)
   x = transform.(Ref(D), r)
   I1, I2, DI1, DI2 = BLI.invariants_ed(x)
   x_d = transform_d.(Ref(D), r)
   return I1, I2, _sdot(x_d, DI1), _sdot(x_d, DI2)
end

# ---------------- inner kernels for the edge length descriptor --------------

"""
`edge_lengths(Rs::AbstractVector, J::SVector{K, Int})`

Take collection `Rs::AbstractVector{JVec}` representing a simplex
and convert into an ordered vector of bondlengths, e.g.,
```
[ r01, r02, r03, r12, r13, r23 ]
[ r01, r02, r03, r04, r12, r13, r14, r23, r24, r34 ]
```

* `Rs` : collection of all R-vectors within a site energy computation
* `J` : indices of the vectors involved in the current nbody computation
"""
@generated function edge_lengths(Rs::AbstractVector, J::SVector{K, Int}) where K
   # note K = N-1 ]
   code = Expr[]
   idx = 0
   for n = 1:K
      idx += 1
      push_str!(code, "s_$idx = norm(Rs[J[$n]])")  # r_0n
   end
   for n = 1:K-1, m = (n+1):K
      idx += 1
      push_str!(code, "s_$idx = norm(Rs[J[$n]] - Rs[J[$m]])")   # r_nm
   end
   s_str = "s = @SVector [s_1"
   if idx > 1; s_str *= prod(", s_$i" for i = 2:idx); end
   push_str!(code, s_str * "]")
   quote
      $(Expr(:meta, :inline))
      @inbounds $(Expr(:block, code...))
      return s
   end
end

"""
convert ∇V (where ∇ is the gradient w.r.t. bond-lengths) into forces, i.e., into
∇Vsite (where ∇ is the gradient w.r.t. positions)

J : neighbour (sub-) indices
"""
@generated function _grad_len2pos!(dVsite, dV, Rs, J::SVector{K, Int}, r) where {K}
   # K is the number of neighbours, i.e. N = K+1 counting also the center atom
   # length(dV) == length(s) == length(S) == K * (K+1)/2
   # ------
   code = Expr[]
   idx = 0
   # the first K entries of dV, s, S are the |Ri - 0|
   for k = 1:K
      idx += 1   # idx == k of course
      push!(code, :( dVsite[J[$k]] += dV[$idx] * (Rs[J[$k]]/r[$idx]) ))
   end
   # the remaining ones are |R_i - R_j|
   for n = 1:K-1, m = (n+1):K
      idx += 1
      push!(code, :( S = (Rs[J[$n]]-Rs[J[$m]]) / r[$idx] ))
      push!(code, :( dVsite[J[$n]] += dV[$idx] * S ))
      push!(code, :( dVsite[J[$m]] -= dV[$idx] * S ))
   end
   quote
      $(Expr(:meta, :inline))
      @inbounds $(Expr(:block, code...))
      return dVsite
   end
end

# =============== Experimental Code for Cluster-based summation!!!!!!

ClusterBLDesc(transform::String, cutoff) =
         ClusterBLDesc(AnalyticTransform(transform), fcut_analyse(cutoff))

# -------------- IO -------------------
Dict(D::ClusterBLDesc) = Dict( "__id__"    =>  "ClusterBLDesc",
                                "transform" =>  Dict(D.transform),
                                "cutoff"    =>  Dict(D.cutoff) )

ClusterBLDesc(D::Dict) = ClusterBLDesc( decode_dict(D["transform"]),
                                        decode_dict(D["cutoff"]) )

==(D1::ClusterBLDesc, D2::ClusterBLDesc) =
      ( (D1.transform == D2.transform) && (D1.cutoff == D2.cutoff) )

Base.convert(::Val{:ClusterBLDesc}, D::Dict) = ClusterBLDesc(D)

ClusterBLDesc(D::BondLengthDesc) = ClusterBLDesc(D.transform, D.cutoff)

# ------------- Interface Code ---------------

tdegrees(::ClusterBLDesc, vN::Val{N}) where {N} = BLI.tdegrees(vN)

@inline ricoords(D::ClusterBLDesc, Rs, J) = edge_lengths(Rs, J)

@inline gradri2gradR!(desc::ClusterBLDesc, dVsite, dV_dr, Rs, J, r) =
   _grad_len2pos!(dVsite, dV_dr, Rs, J, r)

@inline fcut(D::ClusterBLDesc, r) = fcut(D.cutoff, r)
@inline fcut_d(D::ClusterBLDesc, r) = fcut_d(D.cutoff, r)

@inline skip_simplex(D::ClusterBLDesc, r) = (maximum(r) > cutoff(D.cutoff))

@inline invariants(D::ClusterBLDesc, r) = BLI.invariants(transform.(Ref(D), r))

@inline function invariants_ed(D::ClusterBLDesc, r)
   x = transform.(Ref(D), r)
   I1, I2, DI1, DI2 = BLI.invariants_ed(x)
   x_d = transform_d.(Ref(D), r)
   return I1, I2, _sdot(x_d, DI1), _sdot(x_d, DI2)
end


# TODO: parse expressions with MacroTools

"""
    @cfunc ex

Define a "calculation" function with an "adapter" methods used to parse args from model.
"calculate" functions take arguments from system state
and calculate "rate"s determining the probability weight of the reaction be selected.

# Example

For function definition:
```julia
@cfunc @inline growth_c(r, x::Vector) = r * x # function to calculate "growth rate"
```
This macro creates two methods, an "adapter" method
```julia
growth_c(args::NamedTuple) = growth_c(args.r, args.x)
```
and the origin method
```julia
@inline growth_c(r, x::Vector) = r * x
```

!!! info
    For function definition with other macros, put those macros after this macro.

!!! warning
    The argument name `t` is reserved for time.
    Don't use those variable names for other usage
"""
macro cfunc(ex)
    fname, fargs = parse_func(ex, (:t,))
    adapter = :(Base.@propagate_inbounds $fname(t, args::NamedTuple) = $fname($(fargs...)))
    return esc(Expr(:block, adapter, ex))
end

"""
    @ufunc ex

Define a "update" function with an "adapter" methods used to parse args from model.
"update" functions take arguments from system state and update system state.

# Example

For function definition:
```julia
@ufunc Base.@propagate_inbounds growth_u!(ind::CartesianIndex{1}, x) = x[ind] += 1
```
This macro creates two methods, an "adapter" method
```julia
growth_u!(ind, args::NamedTuple) = growth_u(ind, args.x)
```
and the origin method
```julia
Base.@propagate_inbounds growth_u!(ind::CartesianIndex{1}, x) = x[ind] += 1
```

!!! info
    For function definition with other macros, put those macros after this macro.

!!! warning
    The argument name `t` is reserved for time,
    arguments name `rng` is reserved for random number generator,
    and the argument name `ind` is preserved for index of "reaction".
    Don't use those variable names for other usage
"""
macro ufunc(ex)
    fname, fargs = parse_func(ex, (:t, :ind, :rng))
    adapter = :(Base.@propagate_inbounds $fname(rng, t, ind, args::NamedTuple) =
        $fname($(fargs...)))
    return esc(Expr(:block, adapter, ex))
end

"""
    @reaction name ex

Define a [`Reaction`](@ref) with the given `name` and `ex`.

```julia
@reaction growth begin
    r * x # "calculate" expression
    begin
        i = ind[1]
        x[i] += 1 # "update" expression
    end
end
```
will create a `Reaction` with a "calculation" function:
```julia
@cfunc Base.@propagate_inbounds growth_c(r, x) = r * x
```
and an "update" function:
```julia
@ufunc Base.@propagate_inbounds growth_u!(ind, x) = begin
    i = ind[1]
    x[i] += 1
end
```
where arguments of functions were collected from given expression automatically.

!!! note
    If there are global variables, this macro may can not collect arguments correctly.
    Especially, for functions which accept function and types as arguments,
    those functions and types may also be collect as an arguments.
    Thus, these variables must be marked as global variables by `global` before use them,
    even functions and types.
    Besides, type annotation expressions like `x::Vector{Int}`,
    types `Vector` and `Int` will not be collected.
    Avoid to defined your reaction with a type arguments for type annotation.

!!! warning
    The expression follow the same name preserve rule as `@cfunc` and `@ufunc`,
    don't use those variable names for other usage.

This macro cannot parse macro expression.
Thus in some cases, arguments of expression with macro may not be collected correctly.
To avoid this, define reaction with anonymous functions may helpfully:
```julia
@reaction reaction begin
    (A, B) -> @einsum C[i, k] := A[i, j] * B[j, k] # where C, i, j, k should not be collected
    begin
        # do something
    end
end
```
in this case, only arguments of anonymous function (`A` and `B`) will be collected.

Besides, there are another macro [`@quickloop`](@ref), which is similar to `@einsum`,
but it don't follow the einstein convenience which just generate nest loops out of the expression.
Because the `@quickloop` generate a function and arguments are provided,
indexing variables and return variable will not be collected.
"""
macro reaction(name::Symbol, block::Expr)
    block = rmlines(block)
    if block.head == :block && length(block.args) == 2
        cargs, cbody = parse_args_body(block.args[1])
        uargs, ubody = parse_args_body(block.args[2])
        return esc(gen_reaction(name, cargs, cbody, uargs, ubody))
    else
        throw(ArgumentError("can't parse expression"))
    end
end

function gen_reaction(name, cargs, cbody, uargs, ubody)
    cname = Symbol("$(name)_c")
    uname = Symbol("$(name)_u!")
    return Expr(
        :block,
        :(Base.@propagate_inbounds $cname(t, args::NamedTuple) =
            $cname($((_warparg(arg, (:t,)) for arg in cargs)...))),
        :($cname($(cargs...)) = $cbody),
        :(Base.@propagate_inbounds $uname(rng, t, ind, args::NamedTuple) =
            $uname($((_warparg(arg, (:t, :ind, :rng)) for arg in uargs)...))),
        :($uname($(uargs...)) = $ubody),
        Expr(:(=), name, :(Reaction($cname, $uname))),
    )
end

# extract name and arguments of given function expression
function parse_func(ex::Expr, preserve)
    head = ex.head
    if head in (:function, :(=), :where) # unpack function expression
        return parse_func(ex.args[1], preserve)
    elseif head == :macrocall # unpack macro expression
        return parse_func(ex.args[end], preserve)
    elseif head == :call # the args of :call expression is the function name and the arguments
        fname = ex.args[1]
        fargs = map(arg -> parse_arg(arg, preserve), ex.args[2:end])
        return fname, fargs
    else
        error("Unknown expr")
    end
end

# unpack argument names
parse_arg(arg::Symbol, preserve) = _warparg(arg, preserve)
parse_arg(arg::Expr, preserve) = _warparg(arg.args[1]::Symbol, preserve)
function _warparg(sym::Symbol, preserve)
    if sym in preserve
        return sym
    else
        return Expr(:., :args, QuoteNode(sym))
    end
end

# collect arguments from given expression
function parse_args_body(ex)
    if ex.head == :(->)
        return unbox_args!(Set{Symbol}(), ex.args[1]), unblock(ex.args[2])
    elseif ex.head == :macrocall && ex.args[1] == Symbol("@quickloop")
        args = filter(!isline, ex.args)
        return _quickloop(args[2:end]...)
    else
        return collectargs!(Set{Symbol}(), Set{Symbol}(), ex), ex
    end
end

function unbox_args!(args, ex::Expr)
    if ex.head == :(::)
        push!(args, ex.args[1])
    elseif ex.head == :tuple
        for arg in ex.args
            unbox_args!(args, arg)
        end
    else
        error("can't parse expression")
    end
    return args
end
unbox_args!(args, s::Symbol) = (push!(args, s), args)

# `args` is a list of symbols will be treat as arguments name
# `vars` is a list of symbols which are names of internal variable or global variables
collectargs!(args, _, _) = args # ignore when third argument is neither a symbol nor a expr
function collectargs!(args, vars, sym::Symbol) # symbols neither in syms or iv will be collected
    if !in(sym, args) && !in(sym, vars) && sym != :(:)
        push!(args, sym)
    end
    return args
end
function collectargs!(args, vars, _ex::Expr)
    ex = shortdef(_ex) # force convert function definition to short form
    head = ex.head
    if head == :(=) # assignment or function definition
        arg1 = ex.args[1]
        if (arg1 isa Symbol || (arg1 isa Expr && arg1.head == :tuple)) # assignment to one or multi new symbol
            collectargs!(args, vars, ex.args[2]) # must collect exprs right of equal firstly #2
            collectargs!(vars, args, ex.args[1]) # collect internal variables
            return args
        elseif arg1.head == :call # local function definition as short form
            push!(vars, arg1.args[1]) # local function name will be internal variable
            func_args = Set{Symbol}() # local function arguments
            collectargs!(func_args, Set{Symbol}(), arg1) # collect local function arguments
            collectargs!(args, union(vars, func_args), ex.args[2]) # collect variables used in local function
            return args
        else
            collectargs!(args, vars, ex.args[1]) # collect left side of equal
            collectargs!(args, vars, ex.args[2]) # collect right side of equal
            return args
        end
    elseif head == :-> # anonymous function
        func_args = Set{Symbol}() # anonymous function arguments
        collectargs!(func_args, Set{Symbol}(), ex.args[1]) # collect function arguments
        collectargs!(args, union(vars, func_args), ex.args[2]) # collect variables used in anonymous function
        return args
    elseif head == :. # for x.y only collect x, y will be ignore
        return collectargs!(args, vars, ex.args[1])
    elseif head == :global # global variables
        return union!(vars, ex.args)
    elseif head == :(::) # type annotation
        return collectargs!(args, vars, ex.args[1]) # collect left side of :: but ignore right side
    elseif head == :ref # ignore begin and end in reference
        collectargs!(args, vars, ex.args[1])
        foreach(
            arg -> collectargs!(args, TempUnion(vars, (:begin, :end)), arg),
            ex.args[2:end],
        )
    elseif head in (:call, :macrocall) # function and macro will not be treat as a args
        foreach(arg -> collectargs!(args, vars, arg), ex.args[2:end])
    else # in other cases, collect all arguments
        foreach(arg -> collectargs!(args, vars, arg), ex.args)
    end
    return args
end

# union two collections
# but push! to the first collection
struct TempUnion{S,T}
    s::S
    t::T
end
Base.in(item, s::TempUnion) = (item in s.s) || (item in s.t)
Base.push!(s::TempUnion, item) = push!(s.s, item)

"""
@quickloop ex opts...

Generate a function to calculate given expression with nested loops.

# Example
```julia
@quickloop x[i] * m[i, j] * y[j]
```
will generate a function
```julia
(m, y, x)->begin
    ret = Array{...}(undef, size(m, 1), size(m, 2)) # expression for type inference is omitted
    @fastmath @inbounds for j = axes(m, 2), i = axes(m, 1)
        ret[i, j] = x[i] * m[i, j] * y[j]
    end
    return ret
end
```
where the size of return array is determined automatically.
The size can be specified by
```
@quickloop z[j, i] := x[i] * m[i, j] * y[j]
```
which will generate a function
```julia
(m, y, x)->begin
    z = Array{...}(undef, size(m, 2), size(m, 1)) # expression for type inference is omitted
    @fastmath @inbounds for i = axes(m, 1), j = axes(m, 2)
        z[j, i] = x[i] * m[i, j] * y[j]
    end
    return z
end
```
where the return value is named by `z` and the size of return array is
`(size(m, 2), size(m, 1))` instead of `(size(m, 1), size(m, 2))`.

# Options
Options should be specified like `@quickloop ex opt=true/false` or `@quickloop ex opt`,
if the value of option is omitted, it will be treated as `true`.

- `avx` or `turbo`: use `LoopVectorization` to accelerate the loop, default is `false`,
    unless `LoopVectorization` have be loaded.
- `inbounds`: use `@inbounds` to accelerate the loop, default is `true`, this option is ignored
    if `avx` or `turbo` is enabled.
- `fastmath`: use `@fastmath` to accelerate the loop, default is `true`, this option is ignored
    if `avx` or `turbo` is enabled.
- `offset`: if `true`, input arrays will be treat as offset arrays, default is `false`.

!!! note
    The order of function arguments generated by this macro is undetermined.
    Thus, it's recommender use this macro inside of `@reaction`,
    which detect the order of function arguments automatically.
"""
macro quickloop(args...)
    args, fn_body = _quickloop(args...)
    fn_args = Expr(:tuple, args...)
    return esc(Expr(:(->), fn_args, fn_body))
end

function _quickloop(ex, opts...)
    lhs, rhs = _split_eq(ex)
    inds, inds_ax = collectaxes(lhs, rhs)
    fn_args = collectargs!(Set{Symbol}(), Set{Symbol}(inds), rhs)
    fn_body = gen_loop(inds, inds_ax, fn_args, lhs, rhs, opts...)
    return fn_args, fn_body
end

function _split_eq(ex)
    if ex.head == :(=) || ex.head == :(:=)
        return ex.args[1], ex.args[2]
    else
        return Expr(:ref, gensym(:ret)), ex
    end
end

function collectaxes(lhs, rhs)
    # Dict of index symbol and dict of array and dim of the index
    inds_ax = Dict{Symbol,Dict{Symbol,Vector{Int}}}()
    collectaxes!(rhs, inds_ax)
    if length(lhs.args) == 1
        inds = collect(keys(inds_ax))
        max_pos(ind) = maximum(Iterators.flatten(values(inds_ax[ind])))
        sort!(inds; by=max_pos)
        append!(lhs.args, inds)
    else
        inds = lhs.args[2:end]
        for ind in keys(inds_ax)
            if !in(ind, inds)
                throw(ArgumentError(
                    "index symbol `$ind` don't occur in left hand side `$lhs`"
                ))
            end
        end
    end
    return inds, inds_ax
end

function collectaxes!(ex::Expr, inds_ax)
    if ex.head == :ref
        array_name = ex.args[1]
        for i in 2:length(ex.args)
            ind = ex.args[i]
            if !haskey(inds_ax, ind)
                inds_ax[ind] = Dict{Symbol,Vector{Int}}()
            end
            if !haskey(inds_ax[ind], array_name)
                inds_ax[ind][array_name] = Vector{Int}()
            end
            push!(inds_ax[ex.args[i]][array_name], i - 1)
        end
    else
        for arg in ex.args
            collectaxes!(arg, inds_ax)
        end
    end
    return inds_ax
end
collectaxes!(::Any, axs) = axs

function gen_loop(inds, inds_ax, fn_args, lhs, rhs, opts...)
    if isempty(inds) # scalar expression
        return Expr(:return, rhs)
    end
    mode = copy(DEFAULT_LOOP_MODE)
    for ex in opts
        opt, flag = _parse_opt(ex)
        if opt == :turbo
            mode[1] = flag
        elseif opt == :avx
            mode[1] = flag
        elseif opt == :inbounds
            mode[2] = flag
        elseif opt == :fastmath
            mode[3] = flag
        # mode[4] is used to determine of check axes of given array is matched,
        # and it's removed because of the introduction of `ArrayInterface.indices`
        # this mode[4] maybe reused in the future
        elseif opt == :offset
            mode[5] = flag
        else
            error("unknown loop mode: $opt")
        end
    end
    ret_name = lhs.args[1]
    ret_size = Expr[]
    ret_type = _infer_ret_type(inds, inds_ax, fn_args, rhs)
    loop_head = Expr(:block)
    for (i, ind) in enumerate(inds)
        array_name, dims = first(inds_ax[ind])
        dim = first(dims)
        push!(ret_size, :(size($array_name, $dim)))
        array_tuple = Expr(:tuple, ret_name)
        dim_tuple = Expr(:tuple, i)
        for (a, ds) in inds_ax[ind]
            for d in ds
                push!(array_tuple.args, a)
                push!(dim_tuple.args, d)
            end
        end
        ax = Expr(:call, :indices, array_tuple, dim_tuple)
        pushfirst!(loop_head.args, Expr(:(=), ind, ax))
    end
    ret_init = :($ret_name = Array{$ret_type}(undef, $(ret_size...)))
    rhs_offset = if mode[5]
        MacroTools.postwalk(rhs) do ex
            if ex isa Expr && ex.head == :ref
                ex′ = deepcopy(ex)
                for i in 2:length(ex.args)
                    ex′.args[i] = :(axes($(ex.args[1]), $(i-1))[$(ex.args[i])])
                end
                return ex′
            else
                return ex
            end
        end
    else
        rhs
    end
    loop_body = Expr(:block, Expr(:(=), lhs, rhs_offset))
    loop = Expr(:for, loop_head, loop_body)
    if mode[1]
        loop = :(@turbo $loop)
    else
        if mode[2]
            loop = :(@inbounds $loop)
        end
        if mode[3]
            loop = :(@fastmath $loop)
        end
    end
    body = Expr(:block, ret_init, loop, Expr(:return, ret_name))
    return MacroTools.flatten(body)
end

# turbo, inbounds, fastmath, [removed], offset, ...
const DEFAULT_LOOP_MODE = Bool[0, 1, 1, 0, 0, 0, 0, 0]

_parse_opt(opt::Symbol) = opt, true
function _parse_opt(opt::Expr)
    if opt.head == :(=)
        return opt.args[1]::Symbol, opt.args[2]::Bool
    else
        throw(ArgumentError("options should be `opt` or `opt = value::Bool`"))
    end
end

function _infer_ret_type(inds, inds_ax, fn_args, ex)
    ex′ = Expr(:block, ex)
    for ind in inds
        array_name, dims = first(inds_ax[ind])
        dim = first(dims)
        assignment = :($ind = firstindex($array_name, $dim))
        pushfirst!(ex′.args, assignment)
    end
    args = collect(fn_args)
    promote_op_fn = Expr(:(->), Expr(:tuple, args...), ex′)
    args_type = map(arg -> :(typeof($arg)), args)
    ret_type = :(Base.promote_op($promote_op_fn, $(args_type...)))
    return ret_type # expression for inferring the return type
end

"""
    @reaction_eq name param eq [opts...]

Generate a reaction with the given name, parameters and equation.

# Example
```julia
@reaction growth r X[i] --> 2X[i]
```
is equivalent to
```julia
@reaction growth begin
    r * X
    begin
        i = ind[1]
        X[i] += 1
    end
end
```
"""
macro reaction_eq(name::Symbol, param, eq::Expr, opts...)
    lhs_counts, dif_count = parse_eq(eq)
    inds, cargs, cbody = gen_rate(name, param, lhs_counts, opts...)
    update_ex = gen_update(inds, dif_count)
    uargs, ubody = parse_args_body(update_ex)
    return esc(gen_reaction(name, cargs, cbody, uargs, ubody))
end

function parse_eq(eq::Expr)
    if Meta.isexpr(eq, :(-->), 2)
        return _parse_eq(eq.args[1], eq.args[2])
    elseif Meta.isexpr(eq, :call, 3) && eq.args[1] === :(→)
        return _parse_eq(eq.args[2], eq.args[3])
    else
        throw(ArgumentError("unknown expression: $eq, expected `-->` or `→`"))
    end
end
function _parse_eq(lhs, rhs)
    lhs_counts = countsyms(lhs)
    rhs_counts = countsyms(rhs)
    dif_counts = Dict{Any,Int}()
    syms = union(keys(lhs_counts), keys(rhs_counts))
    for sym in syms
        if !haskey(dif_counts, sym)
            dif_counts[sym] = 0
        end
        rhs_count = get(rhs_counts, sym, 0)
        lhs_count = get(lhs_counts, sym, 0)
        dif_counts[sym] = rhs_count - lhs_count
    end
    return lhs_counts, dif_counts
end

function gen_rate(name::Symbol, param, counts, opts...)
    ex = Expr(:call, :*, param)
    for (sym, n) in counts
        if n != 1
            push!(ex.args, :($sym^$n))
        else
            push!(ex.args, sym)
        end
    end
    lhs = Expr(:ref, Symbol(name, :_rate))
    inds, inds_ax = collectaxes(lhs, ex)
    fn_args = collectargs!(Set{Symbol}(), Set{Symbol}(inds), ex)
    fn_body = gen_loop(inds, inds_ax, fn_args, lhs, ex, opts...)
    return inds, fn_args, fn_body
end

function gen_update(inds, counts)
    update = quote
        $(Expr(:tuple, inds...)) = ind.I
    end
    for (sym, n) in counts
        if n != 0
            push!(update.args, :($(_maybe_ref(sym)) += $n))
        end
    end
    return update
end

_maybe_ref(sym::Symbol) = Expr(:ref, sym)
_maybe_ref(sym::Expr) = sym

function countsyms(ex)
    count_dict = Dict{Any, Int}()
    return countsyms!(count_dict, ex, 1)
end

function countsyms!(count_dict, ex::Expr, n::Int)
    if ex.head === :ref
        if haskey(count_dict, ex)
            count_dict[ex] += n
        else
            count_dict[ex] = n
        end
    elseif ex.head === :call && ex.args[1] == :+
        for arg in ex.args[2:end]
            countsyms!(count_dict, arg, n)
        end
    elseif ex.head === :call && ex.args[1] == :* && length(ex.args) == 3 && ex.args[2] isa Int
        countsyms!(count_dict, ex.args[3], ex.args[2] * n)
    else
        throw(ArgumentError("unexpected expression: $ex"))
    end
    return count_dict
end
function countsyms!(count_dict, sym::Symbol, n::Int)
    if haskey(count_dict, sym)
        count_dict[sym] += n
    elseif sym != :∅
        count_dict[sym] = n
    end
    return count_dict
end
countsyms!(count_dict, ::Any, ::Int) = count_dict


using Images, TestImages
using Base.Test

img = testimage("autumn_leaves")
@assert colorspace(img) == "RGBA"
@assert ndims(img) == 3
@assert colordim(img) == 1
@assert eltype(img) == Uint16

img = testimage("cameraman")
@assert colorspace(img) == "Gray"
@assert ndims(img) == 2
@assert colordim(img) == 0
@assert eltype(img) == Uint8

img = testimage("earth_apollo17")
@assert colorspace(img) == "RGB"
@assert ndims(img) == 3
@assert colordim(img) == 1
@assert eltype(img) == Uint8

img = testimage("fabio")
@assert colorspace(img) == "Gray"
@assert ndims(img) == 2
@assert colordim(img) == 0
@assert eltype(img) == Uint8

img = testimage("house")
@assert colorspace(img) == "GrayAlpha"
@assert ndims(img) == 3
@assert colordim(img) == 1
@assert eltype(img) == Uint8

img = testimage("jetplane")
@assert colorspace(img) == "GrayAlpha"
@assert ndims(img) == 3
@assert colordim(img) == 1
@assert eltype(img) == Uint8

img = testimage("lighthouse")
@assert colorspace(img) == "RGB"
@assert ndims(img) == 3
@assert colordim(img) == 1
@assert eltype(img) == Uint8

img = testimage("mandrill")
@assert colorspace(img) == "RGB"
@assert ndims(img) == 3
@assert colordim(img) == 1
@assert eltype(img) == Uint8

img = testimage("moonsurface")
@assert colorspace(img) == "Gray"
@assert ndims(img) == 2
@assert colordim(img) == 0
@assert eltype(img) == Uint8

img = testimage("mountainstream")
@assert colorspace(img) == "RGB"
@assert ndims(img) == 3
@assert colordim(img) == 1
@assert eltype(img) == Uint8


"""
Thermdynamic data collection

Function names list (order by appearance in code):
    metacyc
    metacycIllParsed
    collect_formation_dg (internal eQuilibrator setup)
    collect_formation_dg (overload without internal eQuilibrator setup)
    collect_reaction_dg
    reaction_dg_bounds
    collect_dGr_bounds
    reaction_bounds
"""

"""
    $(SIGNATURES)
Returns reaction string with a prefixed "metacyc.compound:".
"""
function metacyc(str)
    eQuilibrator._parse_reaction_string(str, "metacyc.compound")
end

"""
    $(SIGNATURES)
Returns modified string with prefixed "metacyc.compound:". If the MetaCyc-based model was created via the Python-based "Moped" module, metabolite and reaction names may have artifacts ("__45__" and "__45__" substrings) that should be replaced by a hyphen or periodt, respectively, before ΔG estimations.
"""
function metacycIllParsed(str)
    nstr = replace(str, "__45__46__45__" => "-46-")
    nstr = replace(nstr, "__46__" => ".")
    nstr = replace(nstr, "__45__" => "-")
    nstr = replace(nstr, "__43__" => "+")
    eQuilibrator._parse_reaction_string(nstr, "metacyc.compound")
end

"""
    collect_formation_dg(
        database,
        db_id;
        dGf_dict = OrderedDict(),
        ph::Float64 = 7.0,
        pmg::Float64 = 2.0,
        i_strengh::Quantity = 100.0u"mM",
        temp::Quantity = 25u"°C",
    )

Collects ΔG of formation for all metabolites in `database` via Component Contribution. The data bank source must be passed by `db_id` (Options: `bigg`, `kegg`, `metanex`, `chebi`). Alternatively a dictionary of ΔGs of formation `dGf_dict` can be passed to extend it by new entries for the database metabolites.
Keyword arguments setup eQuilibrator conditions.

# Example
```
collect_formation_dg(
    ecoli_core,
    bigg
)
```
"""
function collect_formation_dg(
    database,
    db_id;
    dGf_dict = OrderedDict(),
    met_prefix = "",
    ph::Float64 = 7.0,
    pmg::Float64 = 2.0,
    i_strengh::Quantity = 100.0u"mM",
    temp::Quantity = 25u"°C",
)
    println("eQuilibrator is initialized. Please wait...")

    equil = eQuilibrator.Equilibrator(
        pH = ph,
        pMg = pmg,
        ionic_strength = i_strengh,
        temperature = temp,
    )

    n = length(metabolites(database))
    p = Progress(n, 1)
    for met in metabolites(database)
        if met ∉ keys(dGf_dict)
            rxn_str = string(" = ", met[length(met_prefix)+1:end-2])
            try
                @suppress begin
                    dGf_dict[met] =
                        standard_dg_prime(equil, db_id(rxn_str); balance_warn = false)
                end
            catch
                dGf_dict[met] = "NA"
            end
        end
        next!(p)
    end
    return dGf_dict
end

"""
    collect_formation_dg(
        database,
        equil,
        db_id;
        met_prefix = "",
        dGf_dict = OrderedDict(),
    )

Collects ΔG of formation for all metabolites in `database` via Component Contribution. The data bank source must be passed by `db_id` (Options: `bigg`, `kegg`, `metanex`, `chebi`). An eQuilibrator setup has to be passed by `equil` (see also: https://github.com/stelmo/Equilibrator.jl). Alternatively a dictionary of ΔGs of formation `dGf_dict` can be passed to extend it by new entries for the database metabolites.

# Example
```
collect_formation_dg(
    ecoli_core,
    eQuilibrator.Equilibrator(ionic_strength=150.0u"mM"),
    bigg
)
```
"""
function collect_formation_dg(
    database,
    equil,
    db_id;
    met_prefix = "",
    dGf_dict = OrderedDict(),
)
    println("ΔG of formation are collected.")
    n = length(metabolites(database))
    p = Progress(n, 1)
    for met in metabolites(database)
        if met ∉ keys(dGf_dict)
            rxn_str = string(" = ", met[length(met_prefix)+1:end-2])
            try
                @suppress begin
                    dGf_dict[met] =
                        standard_dg_prime(equil, db_id(rxn_str); balance_warn = false)
                end
            catch
                dGf_dict[met] = measurement(1000000u"kJ/mol", 10000000u"kJ/mol") #setting infeasible value
            end
        end
        next!(p)
    end
    return dGf_dict
end

"""
function collect_reaction_dg(
    database,
    equil,
    db_id;
    met_prefix = "",
    dgr_dict = OrderedDict(),
)
Collects standard transformed ΔG of reactions for every reaction in an SBML model. The eQuilibrator setup (see `eQuilibrator.Equilibrator()`) is passed as the `equil` argument. The source database (e.g. BiGG) is passed by `db_id`. If a dictionary with ΔGs of reaction already exist and should be expand, pass the dictionary to the `dgr_dict` keyword argument.
"""
function collect_reaction_dg(
    database,
    equil,
    db_id;
    met_prefix = "",
    dgr_dict = OrderedDict(),
)
    println("Standard ΔG of reactions are collected.")
    n = length(metabolites(database))
    p = Progress(n, 1)
    for rxn in reactions(database)
        if rxn ∉ keys(dgr_dict)
            rxn_string = stoichiometry_string(
                reaction_stoichiometry(database, rxn),
                format_id = x -> db_id(x[length(met_prefix)+1:end-2]),
            )
            try
                @suppress begin
                    dgr_dict[rxn] = standard_dg_prime(equil, rxn_string, balance_warn = false)
                end
            catch
                dgr_dict[rxn] = measurement(1000000u"kJ/mol", 10000000u"kJ/mol") #setting infeasible value
            end
        end
        next!(p)
    end
    return dgr_dict
end

"""
    collect_dGr_bounds(
        database;
        return_opts::Int = 3,
        ph::Float64 = 7.0,
        pmg::Float64 = 2.0,
        i_strength::Quantity = 100.0u"mM",
        temp::Quantity = 25u"°C",
        db_id = bigg,
        met_suffix = "",
    )

DEV NOTE: Necessary for MinFlux/MinRxn procedure, currently work in progress. Function might change severely.
Collects ΔG of reaction for every reaction in `database` via Component Contribution.
Returns dictionary of ΔG of reactions (`return_opts = 1`), the calculated bounds for MinFlux/MinRxn (`return_opts = 2`) or both (`return_opts = 3`).
Kwargs setup eQuilibrator conditions and return.
"""
function collect_dGr_bounds(
    database;
    return_opts::Int = 3,
    ph::Float64 = 7.0,
    pmg::Float64 = 2.0,
    i_strength::Quantity = 100.0u"mM",
    temp::Quantity = 25u"°C",
    db_id = bigg,
    met_prefix = "",
)
    println("eQuilibrator is initialized. Please wait...")
    equil = eQuilibrator.Equilibrator(
        pH = ph,
        pMg = pmg,
        ionic_strength = i_strength,
        temperature = temp,
    )

    R = 8.31446261815324u"J/(K*mol)"
    R = uconvert(u"kJ/(K*mol)", R)

    println("Calculating standard ΔG of reaction:")
    dGr_standard_dict = OrderedDict()
    n = length(reactions(database))
    p = Progress(n, 1)
    for rxn in reactions(database)
        rxn_string = stoichiometry_string(
            reaction_stoichiometry(database, rxn),
            format_id = x -> x[length(met_prefix)+1:end-2],
        )
        try
            @suppress begin
                dGr_standard_dict[rxn] =
                    standard_dg_prime(equil, db_id(rxn_string), balance_warn = false)
            end
        catch
            dGr_standard_dict[rxn] = "NA"
        end
        next!(p)
    end

    println("Calculating flux bounds:")
    dGr_bounds = OrderedDict()
    n = length(reactions(database))
    p2 = Progress(n, 1)

    for (rxn, dG) in dGr_standard_dict
        try
            @suppress begin
                reactants =
                    [abs(v) for (k, v) in reaction_stoichiometry(database, rxn) if v < 0]
                products =
                    [abs(v) for (k, v) in reaction_stoichiometry(database, rxn) if v > 0]
                qmax =
                    prod(fill(0.1, length(products)) .^ products) /
                    prod(fill(1 * 10^-6, length(reactants)) .^ reactants)
                qmin =
                    prod(fill(1 * 10^-6, length(products)) .^ products) /
                    prod(fill(0.1, length(reactants)) .^ reactants)
                u_calc = R * uconvert(u"K", temp) * log(qmax)
                l_calc = R * uconvert(u"K", temp) * log(qmin)
                println(typeof(dG))
                dGu = dG.val + u_calc.val
                dGl = dG.val + l_calc.val
                dGr_bounds[rxn] = (dGl, dGu)
            end
        catch
            dGr_bounds[rxn] = "NA"
        end
        next!(p2)
    end
    if return_opts == 3 # Returns both ΔG of reaction dictionary and the calculated bounds
        return dGr_standard_dict, dGr_bounds
    elseif return_opts == 2
        return dGr_bounds
    elseif return_opts == 1
        return dGr_standard_dict
    else
        println("Return options not specified. Switched to complete return.")
        return dGr_standard_dict, dGr_bounds
    end
end

function reduce_model(database, dg_dictionary)
    for rxn in reactions(database)
        if rxn ∉ keys(dg_dictionary)
            database = remove_reaction(database, rxn)
        end
    end
    return database
end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ a1e77490-ba51-11eb-3cd8-1d82e8b5f7e4
begin
	using Pkg
	Pkg.activate(joinpath(@__DIR__, ".."))
end

# ╔═╡ 9bb18432-873d-4979-9e2c-a71e1007b2d7
using Agents, InteractiveDynamics, CairoMakie, PlutoUI

# ╔═╡ f3d42d7d-3e92-4f78-9284-689e2ac50017
using Setfield

# ╔═╡ 6ca2d3d4-d9ac-4e9d-9b1b-e8a66f00c7e6
using Random

# ╔═╡ d4ed1cd4-89a1-49f3-8c18-19215b4c2e8e
using DataFrames

# ╔═╡ 89d0deb2-13c1-45ca-b2ec-ee9704ae6e15
using Base.Threads

# ╔═╡ 3097bf32-e7ba-4b36-8674-3ddd37e1f395
mutable struct BPH <: AbstractAgent
	id::Int
	pos::Dims{2}
	energy
	age
	nb_reproduce
end

# ╔═╡ d073c77d-390c-437b-8ede-630f07e5d9d9
# module MyModel
# using Ag
# @agent BPH GridAgent{2} begin
# 	energy
# 	age
# 	nb_reproduce
# end
# end

# ╔═╡ 4abaea9e-d7ff-4d00-9433-0dd6e3d6da8d
# BPH = MyModel.BPH

# ╔═╡ d948d629-73bc-48b2-9275-51c0bd82a29a
function calculate_positions(pos, nb_flower, field_size)
	positions = if pos === :corner
		Iterators.product(1:5, 1:5)
	elseif pos === :border
		_, field_height = field_size
		Iterators.product(1:2, 1:field_height)
	elseif pos === :random
		field_width, field_height = field_size
		flower_start = (field_width ÷ 2 - nb_flower ÷ 2)
		Iterators.product(1:flower_start, 1:field_height)
	end
	collect(positions)
end

# ╔═╡ 60c02b02-92d3-4c58-b9f1-6e71dfaa28fc
function init_model(nb_bph_init, nb_flower, init_position, field_size;
		energy_miss = 0.025,
		age_init = 168,
		age_reproduce = 504,
		age_old = 600,
		age_die = 720,
		pr_move = 0.99,
		pr_reproduce = 0.07,
		pr_egg_death = 0.0025,
		pr_old_death = 0.04,
		offspring_max = 12,
		offspring_min = 5,
		energy_max = 1.0,
		energy_transfer = 0.1,
		energy_consume = 0.025,
		energy_move = 0.2,
		energy_reproduce = 0.8,
		seed = 1)
	field_width, field_height = field_size
	properties = (
		age_init = 168,
		age_reproduce = 504,
		age_old = 600,
		age_die = 720,
		pr_move = 0.99,
		pr_reproduce = 0.7,
		pr_egg_death = 0.0025,
		pr_old_death = 0.04,
		offspring_max = 12,
		offspring_min = 5,
		energy_miss = 0.025,
		energy_max = 1.0,
		energy_transfer = 0.1,
		energy_consume = 0.025,
		energy_move = 0.2,
		energy_reproduce = 0.8,
		flower_start = field_width ÷ 2 - nb_flower ÷ 2,
		flower_end = field_width ÷ 2 + nb_flower ÷ 2,
		food = fill(0.4, field_size)
	)
	field = GridSpace(field_size; periodic=false)
	rng = MersenneTwister(seed)
	scheduler = Schedulers.randomly
	model = ABM(BPH, field; properties, rng, scheduler)

	init_positions = calculate_positions(init_position, nb_flower, field_size)
	for _ in 1:nb_bph_init
		id = nextid(model)
		pos = rand(model.rng, init_positions)
		energy = rand(model.rng, 0.4:0.01:0.6)
		bph = BPH(id, pos, energy, rand(model.rng, 0:299), rand(model.rng, 1:21))
		add_agent_pos!(bph, model)
	end

	return model
end

# ╔═╡ cca251da-24fe-4b97-87db-337cc2d9490d
function bph_reflex!(bph, model)
	# older
	bph.age = bph.age + 1

	# move around
	directions = Iterators.product([-1, -2, 1, 2], [-1, -2, 1, 2])
	if bph.age > model.age_init && rand(model.rng) < model.pr_move
		walk!(bph, rand(model.rng, directions |> collect), model)
	end

	# loss energy normally
	if bph.age ≥ model.age_init
		bph.energy = bph.energy - model.energy_consume
	end

	# loss energy in flower
	x, y = bph.pos
	if model.flower_start ≤ x < model.flower_end
		bph.energy -= model.energy_miss
	elseif model.food[x, y] > 0 && bph.age ≥ model.age_init
		transfer = min(model.energy_transfer, model.food[x, y])
		model.food[x, y] -= transfer
		bph.energy = min(bph.energy + transfer, model.energy_max)
	end

	# reproduce
	if (bph.age ≥ model.age_reproduce &&
		bph.energy ≥ model.energy_reproduce &&
		bph.nb_reproduce < 21 &&
		rand(model.rng) ≤ model.pr_reproduce)
		#
		bph.nb_reproduce = bph.nb_reproduce + 1
		nb_offspring = rand(model.rng, model.offspring_min:model.offspring_max)
		for _ in 1:nb_offspring
			id = nextid(model)
			bph = BPH(id, bph.pos, 0.4, 0, 0)
			add_agent_pos!(bph, model)
		end
	end

	# die
	if bph.energy ≤ 0 ||
			(bph.age ≥ model.age_die) ||
			(model.age_die ≥ bph.age ≥ model.age_old && rand(model.rng) ≤ model.pr_old_death)
			(bph.age < model.age_init && rand(model.rng) ≤ model.pr_egg_death)
		kill_agent!(bph, model)
	end
end

# ╔═╡ af4042f1-7679-4cc5-ad20-b23937ffa40b
function rice_grow!(model)
	alive = @. model.food > 0
	@. model.food = min(model.food + model.food * 0.008 * alive , 1.0)
end

# ╔═╡ e8c37ce6-9376-4215-ae47-a3e3df26276a
bph_color(bph) = if bph.age < 168
	RGBf0(0, 0, 255)
else
	RGBf0(255, 0, 0)
end

# ╔═╡ 7e042ab7-6d6e-4df0-a9ef-b521fdf9f1de
rice_color(model) = map(positions(model)) do pos
	x, y = pos
	if model.flower_start ≤ x ≤ model.flower_end
		# RGBf0(255, 255, 0)
		-1
	else
		# RGBAf0(0, 255, 0, model.food[x, y])
		max(model.food[x, y], 0)
	end
end

# ╔═╡ 9ff14f09-e64b-4898-935a-c425ef1d3eaf
plotkwargs = (ac = bph_color, heatarray = rice_color)

# ╔═╡ 456e997d-c424-4647-9ba9-cd98a4c43ed2
heatkwargs = (nan_color = RGBf0(1.0, 1.0, 0),
	colormap = [RGBAf0(0, 1.0, 0, i) for i in 0:0.01:1]
)

# ╔═╡ 6f4dc104-c0f1-4c9f-af54-5de98a6e62e9
md"""
Generate Video $(@bind should_video PlutoUI.CheckBox())

Collect Data $(@bind should_data PlutoUI.CheckBox())

Flower $(@bind flower PlutoUI.Slider(12:20; show_value=true))
"""

# ╔═╡ 2da972ea-1e4b-4096-81b2-1e86b4eb4a10
@bind seed PlutoUI.NumberField(1:2000)

# ╔═╡ 48fb77fb-ab35-4a4b-948b-ce3599f7ec60
if should_video
	vid = "bph.mp4"
	let model = init_model(200, flower, :corner, (100, 100); seed = seed)
		abm_video(vid, model, bph_reflex!, rice_grow!; frames=1200 - 1, framerate=10, heatkwargs, plotkwargs...)
		LocalResource(vid)
	end
end

# ╔═╡ 149636d6-95b3-4ead-a607-f37c4de5ed52
1234

# ╔═╡ f0ddf1cc-12e5-4444-a5f4-d4085b8056b3
# @time max_times = map(1:100) do seed
# 	bph(x) = true
# 	model = init_model(200, 16, :corner, (100, 100); seed = seed)
# 	adf, _ = run!(model, bph_reflex!, rice_grow!, 2880; adata = [(bph, count)])
# 	seed => maximum(adf.step)
# end

# ╔═╡ 9e235da3-3f0f-42de-af15-10d940b3304c
@bind rerun_btn PlutoUI.Button("Re-run")

# ╔═╡ 3492f6a0-a867-4403-b32f-4dff1f706142
begin
	rerun_btn
	model = init_model(200, flower, :corner, (100, 100); seed)
end

# ╔═╡ 063e3c58-fb8f-46d7-8f37-677af66598bf
md"## Run & Collect data"

# ╔═╡ 09cf52d0-d6f5-4dad-ad43-68aea0f44cdd
# adata = [x -> x.pos, x -> x.age]

# ╔═╡ 736ae5e6-3c3f-4e24-9b93-6fc578d378db
adata = let bph_A1(bph) = bph.age ≥ model.age_init && bph.pos[1] < model.flower_start
	bph_E1(bph) = bph.age < model.age_init && bph.pos[1] < model.flower_start
	bph_A2(bph) = bph.age ≥ model.age_init && bph.pos[1] ≥ model.flower_end
	bph_E2(bph) = bph.age < model.age_init && bph.pos[1] ≥ model.flower_end
	[(bph_E1, count), (bph_E2, count), (bph_A1, count), (bph_A2, count)]
end

# ╔═╡ ff33a882-4092-42e5-a272-043cc9c2dbf0
# max_step = maximum(adfa.step)

# ╔═╡ bc177477-8bb5-4821-9076-517b07c0c0f7
# GC.gc()

# ╔═╡ 6835bd93-4824-4768-a575-912fdc4a94ee
let model = init_model(200, flower, :random, (100, 100))
	GC.gc()
	PlutoUI.with_terminal() do
		@time run!(model, bph_reflex!, rice_grow!, 2880; adata=adata);
	end
end

# ╔═╡ 8a7027e6-7959-43e9-b0b6-245ffbeeec74
adf, mdf = run!(model, bph_reflex!, rice_grow!, 2880; adata=adata);

# ╔═╡ 85e8ab56-3345-4e31-a6f0-cff58c42ac6c
# adf2 = rename(adf, [:step, :id, :pos, :age]);

# ╔═╡ 1471bab5-baab-48f3-8f13-975b2aabcd7b
adf

# ╔═╡ 4a9553e5-9229-47ec-8248-e72be094fa34

# ╔═╡ 088a6bab-0be7-4f0e-96bf-7497b6dd7209
let f = Figure()
	Axis(f[1, 1])
	step = adf.step[adf.count_bph_A1 .> 0]
	bph1 = adf.count_bph_A1[adf.count_bph_A1 .> 0]
	lines!(step, bph1)
	lines!(adf.step ./ 24, adf.count_bph_A1)
	lines!(adf.step ./ 24, adf.count_bph_A2)
	f
end

# ╔═╡ 96336b96-44d5-4356-ab9d-ab8cbd9ea1bb
let f = Figure()
	Axis(f[1, 1])
	@threads for seed = 1:5
		model = init_model(20, 18, :corner, (100, 100); seed = seed);
		adf, mdf = run!(model, bph_reflex!, rice_grow!, 2880; adata=adata);
		bph1 = adf.count_bph_A1[1:1700]
		steps = adf.step[1:1700]
		lines!(steps, bph1)
		# lines!(adf.step ./ 24, adf.count_bph_A1)
		# lines!(adf.step ./ 24, adf.count_bph_A2)
	end
	f
end

# ╔═╡ 67366c81-42fd-407c-bf21-1fe04cc4dda0
# adfa = let adult(age) = age .≥ model.age_init
# 	egg(age) = age .< model.age_init
# 	field1(pos) = getindex.(pos, 1) .< model.flower_start
# 	field2(pos) = getindex.(pos, 1) .≥ model.flower_end
# 	combine(groupby(adf2, :step),
# 		[:pos, :age] => ((p, a) -> count(egg(a) .& field1(p))) => :e1,
# 		[:pos, :age] => ((p, a) -> count(adult(a) .& field1(p))) => :a1,
# 		[:pos, :age] => ((p, a) -> count(egg(a) .& field2(p))) => :e2,
# 		[:pos, :age] => ((p, a) -> count(adult(a) .& field2(p))) => :a2,
# 	)
# end;

# ╔═╡ 5a8a235a-64ff-41a3-981b-684a1424fa07
# let f = Figure()
# 	Axis(f[1, 1])
# 	# lines!(adfa.step / 24, adfa.e1)
# 	# lines!(adfa.step / 24, adfa.e2)
# 	lines!(adfa.step, adfa.a1)
# 	lines!(adfa.step, adfa.a2)
# 	f
# end

# ╔═╡ d6502588-0196-4953-86f0-5244b519f9ee

# ╔═╡ Cell order:
# ╠═a1e77490-ba51-11eb-3cd8-1d82e8b5f7e4
# ╠═9bb18432-873d-4979-9e2c-a71e1007b2d7
# ╠═f3d42d7d-3e92-4f78-9284-689e2ac50017
# ╠═3097bf32-e7ba-4b36-8674-3ddd37e1f395
# ╠═d073c77d-390c-437b-8ede-630f07e5d9d9
# ╠═4abaea9e-d7ff-4d00-9433-0dd6e3d6da8d
# ╠═d948d629-73bc-48b2-9275-51c0bd82a29a
# ╠═60c02b02-92d3-4c58-b9f1-6e71dfaa28fc
# ╠═6ca2d3d4-d9ac-4e9d-9b1b-e8a66f00c7e6
# ╠═3492f6a0-a867-4403-b32f-4dff1f706142
# ╠═cca251da-24fe-4b97-87db-337cc2d9490d
# ╠═af4042f1-7679-4cc5-ad20-b23937ffa40b
# ╠═e8c37ce6-9376-4215-ae47-a3e3df26276a
# ╠═7e042ab7-6d6e-4df0-a9ef-b521fdf9f1de
# ╠═9ff14f09-e64b-4898-935a-c425ef1d3eaf
# ╠═456e997d-c424-4647-9ba9-cd98a4c43ed2
# ╠═48fb77fb-ab35-4a4b-948b-ce3599f7ec60
# ╠═6f4dc104-c0f1-4c9f-af54-5de98a6e62e9
# ╠═2da972ea-1e4b-4096-81b2-1e86b4eb4a10
# ╠═149636d6-95b3-4ead-a607-f37c4de5ed52
# ╠═f0ddf1cc-12e5-4444-a5f4-d4085b8056b3
# ╟─9e235da3-3f0f-42de-af15-10d940b3304c
# ╟─063e3c58-fb8f-46d7-8f37-677af66598bf
# ╠═09cf52d0-d6f5-4dad-ad43-68aea0f44cdd
# ╠═736ae5e6-3c3f-4e24-9b93-6fc578d378db
# ╠═ff33a882-4092-42e5-a272-043cc9c2dbf0
# ╠═bc177477-8bb5-4821-9076-517b07c0c0f7
# ╠═6835bd93-4824-4768-a575-912fdc4a94ee
# ╠═8a7027e6-7959-43e9-b0b6-245ffbeeec74
# ╠═d4ed1cd4-89a1-49f3-8c18-19215b4c2e8e
# ╠═85e8ab56-3345-4e31-a6f0-cff58c42ac6c
# ╠═1471bab5-baab-48f3-8f13-975b2aabcd7b
# ╠═4a9553e5-9229-47ec-8248-e72be094fa34
# ╠═088a6bab-0be7-4f0e-96bf-7497b6dd7209
# ╠═89d0deb2-13c1-45ca-b2ec-ee9704ae6e15
# ╠═96336b96-44d5-4356-ab9d-ab8cbd9ea1bb
# ╠═67366c81-42fd-407c-bf21-1fe04cc4dda0
# ╠═5a8a235a-64ff-41a3-981b-684a1424fa07
# ╠═d6502588-0196-4953-86f0-5244b519f9ee


module Molecules

using ForwardDiff
using LinearAlgebra

export @Orbital_str, @orbital_str,
    chemicalpotential, hamiltonian, makeaxes, molecularorbitals, overlap,
    selfenergy, slaterkoster, total_energy, vibrationalcoupling

const ⊗ = kron
const σ0 = I(2)
const σ = (
    x = [0.0 1.0; 1.0 0.0],
    y = [0.0 -1.0im; 1.0im 0.0],
    z = [1.0 0.0; 0.0 -1.0]
)

struct Orbital{N,L,M} end
macro Orbital_str(str::String)
    reg = r"^([0-9\*]+)([a-zA-Z]+)(?:_(?:\{(.+)\}|([a-zA-Z\*]+)))?$"
    m = match(reg, str)
    m === nothing && error("Invalid orbital")
    N = m[1] === nothing ? error("Invalid orbital") : m[1] == "*" ? nothing : parse(Int, m[1])
    L = m[2] === nothing ? error("Invalid orbital") : Symbol(m[2])
    M = m[3] === nothing ? m[4] === nothing ? Nothing : Symbol(m[4]) : Symbol(m[3])
    M === Symbol("*") && (M = nothing)

    N === nothing && M === nothing && return Orbital{<:Any,L,<:Any}
    N === nothing && return Orbital{<:Any,L,M}
    M === nothing && return Orbital{N,L,<:Any}
    return Orbital{N,L,M}
end

macro orbital_str(str::String)
    reg = r"^([0-9]+)([a-zA-Z]+)(?:_(?:\{(.+)\}|([a-zA-Z]+)))?$"
    m = match(reg, str)
    m === nothing && error("Invalid orbital")
    N = m[1] === nothing ? error("Invalid orbital") : parse(Int, m[1])
    L = m[2] === nothing ? error("Invalid orbital") : Symbol(m[2])
    M = m[3] === nothing ? m[4] === nothing ? Nothing : Symbol(m[4]) : Symbol(m[3])
    return Orbital{N,L,M}()
end

function Base.show(io::IO, ::Orbital{N,L,M}) where {N,L,M}
    if M === Nothing
        print(io, "orbital\"$N$L\"")
    elseif match(r"^[a-zA-Z]$", string(M)) !== nothing
        print(io, "orbital\"$N$(L)_$M\"")
    else
        print(io, "orbital\"$N$(L)_{$M}\"")
    end
end

function shortform(::Orbital{N,L,M}) where {N,L,M}
    l = L
    m = M == Nothing ? "" : replace(string(M), r"[_^+\-{}]" => s"")
    return Symbol("$l$m")
end

mutable struct Atom{SYMBOL,Z,N}
    position
    axes
    valence::Union{Int,Missing}
    orbitals

    function Atom{SYMBOL,Z,N}(position, axes, valence, orbitals::Orbital...) where {SYMBOL,Z,N}
        return new{SYMBOL,Z,N}(
            position,
            axes,
            valence,
            orbitals
        )
    end
end

function Atom{SYMBOL,Z,N}(position, orbitals::Orbital...; valence=missing, axes=makeaxes(0., 0.)) where {SYMBOL,Z,N}
    return Atom{SYMBOL,Z,N}(position, axes, valence, orbitals...)
end

Base.:(==)(atom1::Atom, atom2::Atom) = false
function Base.:(==)(atom1::T, atom2::T) where {T <: Atom}
    return (
        atom1.position == atom2.position &&
        atom1.orbitals == atom2.orbitals &&
        isequal(atom1.valence, atom2.valence)
    )
end

Base.getindex(atom::Atom, i::Integer) = collect(atom.orbitals)[i]

position(atom::Atom) = atom.position
atomic_positions(atoms::AbstractVector{<:Atom}) = position.(atoms)

struct Bond{A1<:Atom,A2<:Atom}
    atom1::A1
    atom2::A2

    function Bond{A1,A2}(atom1::A1, atom2::A2) where {A1<:Atom,A2<:Atom}
        if hash(atom1) > hash(atom2)
            return new{A1,A2}(atom1, atom2)
        else
            return new{A2,A1}(atom2, atom1)
        end
    end
end
Bond(atom1::A1, atom2::A2) where {A1, A2} = Bond{A1,A2}(atom1, atom2)

mutable struct Molecule
    position
    axes
    atoms::Vector{Atom}
    bonds::Set{Bond}
end

Molecule(atoms=Vector{Atom}(), bonds=Set{Bond}()) = Molecule(
    [0., 0., 0.],
    makeaxes(0., 0.),
    collect(Atom, atoms),
    Set{Bond}(bonds))

function Base.:(==)(mol1::Molecule, mol2::Molecule)
    return (
        mol1.position == mol2.position &&
        length(mol1.atoms) == length(mol2.atoms) &&
        all(atom1 -> any(atom2 -> atom1 == atom2, mol2.atoms), mol1.atoms) &&
        all(atom2 -> any(atom1 -> atom1 == atom2, mol1.atoms), mol2.atoms)
    )
end

Base.push!(mol::Molecule, atoms::Atom...) = push!(mol.atoms, atoms...)
Base.push!(mol::Molecule, bonds::Bond...) = push!(mol.bonds, bonds...)

Base.eachindex(mol::Molecule) = eachindex(collect(mol.atoms))
Base.iterate(mol::Molecule) = iterate(collect(mol.atoms))
Base.iterate(mol::Molecule, i) = iterate(collect(mol.atoms), i)
Base.length(mol::Molecule) = length(mol.atoms)
Base.getindex(mol::Molecule, i::Integer) = collect(mol.atoms)[i]
Base.lastindex(mol::Molecule) = lastindex(collect(mol.atoms))

function indices(mol::Molecule, atom::Atom)
    N0 = 1
    for atom′ in mol.atoms
        N1 = N0 + countorbitals(atom′)
        if atom === atom′
            return N0:(N1-1)
        end
        N0 = N1
    end
    return nothing
end

function indices(mol::Molecule, atom::Atom, orbital::Orbital)
    N0 = 1
    for atom′ in mol.atoms
        N1 = N0 + countorbitals(atom′)
        if atom === atom′
            for (i, orbital′) in enumerate(atom.orbitals)
                if orbital == orbital′
                    return (N0+i-1):(N0+i-1)
                end
            end
            return nothing
        end
        N0 = N1
    end
    return nothing
end

Base.Vector{T}(mol::Molecule) where {T} = normalize!(fill(oneunit(T), countorbitals(mol)))

function Base.Vector{T}(mol::Molecule, atom::Atom) where {T}
    u = fill(zero(T), countorbitals(mol))
    N = countorbitals(atom)
    u[indices(mol, atom)] .= 1 / sqrt(N)
    return u
end

function Base.Vector{T}(mol::Molecule, n::Integer) where {T}
    atom = mol[n]
    return Vector{T}(mol, atom)
end

function Base.Vector{T}(mol::Molecule, atom::Atom, orbital::Orbital) where {T}
    u = fill(zero(T), countorbitals(mol))
    u[indices(mol, atom, orbital)] .= 1
    return u
end

function Base.Vector{T}(mol::Molecule, atom::Atom, s::Integer) where {T}
    orbital = atom[s]
    return Vector{T}(mol, atom, orbital)
end

function Base.Vector{T}(mol::Molecule, n::Integer, s::Integer) where {T}
    atom = mol[n]
    orbital = atom[s]
    return Vector{T}(mol, atom, orbital)
end

function centre_of_mass(f, mol::Molecule)
    return sum(mass(f, atom)*position(atom) for atom in mol) / mass(f, mol)
end

function balance_position!(f, mol::Molecule)
    p0 = centre_of_mass(f, mol)
    for atom in mol.atoms
        atom.position -= p0
    end
    return mol
end

##################################################
# Chemical properties
##################################################

protonnumber(::Atom{SYMBOL,Z,N}) where {SYMBOL,Z,N} = Z
neutronnumber(::Atom{SYMBOL,Z,N}) where {SYMBOL,Z,N} = N
chemical_symbol(::Atom{SYMBOL,Z,N}) where {SYMBOL,Z,N} = SYMBOL
chemical_symbols(atoms::AbstractVector{<:Atom}) = chemical_symbol.(atoms)

countorbitals(atom::Atom)::Int = length(atom.orbitals)
countorbitals(mol::Molecule)::Int = sum(map(countorbitals, mol.atoms))

mass(f, atom::Atom) = f(chemical_symbol(atom))
mass(f, mol::Molecule) = sum(map(atom -> mass(f, atom), mol))

function electronvalence(atom::Atom)
    !ismissing(atom.valence) && return atom.valence
    n = protonnumber(atom)
    n <= 2 && return n
    n <= 10 && return n - 2
    n <= 18 && return n - 10 # Not very generic ...
    n > 18 && error("atom numbers greater than 18 are currently not supported")
end
"""
    valencecount(mol)

Return the sum of electron valences of each atom in `mol`.

"""
valencecount(mol::Molecule, oxidation=0) = sum(electronvalence(atom) for atom in mol.atoms) - oxidation

function molecularorbitals(H)
    vals, vecs = eigen(H)
    s = sortperm(vals)
    return vals[s], vecs[:, s]
end

function chemicalpotential(mol::Molecule, levels, oxidation=0)
    n = valencecount(mol, oxidation)
    μ = real(levels[ceil(Int, n/2)] + levels[ceil(Int, (n + 1)/2)])/2
    Δ = real(levels[ceil(Int, (n + 1)/2)] - levels[ceil(Int, n/2)])/2
    return μ, Δ
end

bond_length(bond::Bond) = norm(position(bond.atom2) - position(bond.atom1))

function total_energy(mol::Molecule, levels; oxidation = 0)
    n = valencecount(mol, oxidation)
    return sum(i -> real(levels[i]), 1:ceil(Int, n/2))
end

##################################################
# Hamiltonians
##################################################

include("hamiltonian.jl")
include("overlap.jl")

include("unitcell.jl")

##################################################
# Hamiltonian/overlap schemes
##################################################

include("slaterkoster.jl")

##################################################
# Elements
##################################################

chemical_symbol_to_atom_type(x::Symbol) = chemical_symbol_to_atom_type(Val(x))

const Hydrogen = Atom{:H, 1, 0}
Hydrogen(position; valence=1, kwargs...) = Hydrogen(position, orbital"1s"; valence=valence, kwargs...)
chemical_symbol_to_atom_type(::Val{:H}) = Hydrogen

const Deuterium = Atom{:D, 1, 1}
const Tritium = Atom{:T, 1, 2}

const Carbon = Atom{:C, 6, 6}
Carbon(position; valence=4, kwargs...) = Carbon(position, orbital"2s", orbital"2p_x", orbital"2p_y", orbital"2p_z"; valence=valence, kwargs...)
chemical_symbol_to_atom_type(::Val{:C}) = Carbon

const Carbon13 = Atom{:C, 6, 7}
const Nitrogen = Atom{:N, 7, 7}
const Oxygen = Atom{:O, 8, 8}

##################################################
# Coordinates
##################################################

function makeaxes(θ = 0.0, ϕ = 0.0, ρ = 0.0)
    R1 = LinearAlgebra.Givens(2, 1, cos(ρ), sin(ρ))
    R2 = LinearAlgebra.Givens(3, 1, cos(ϕ), sin(ϕ))
    R3 = LinearAlgebra.Givens(2, 1, cos(θ), sin(θ))
    return R3*(R2*(R1*Matrix(I, 3, 3)))
end

##################################################
# Basis change
##################################################

localbasis(atom::Atom) = localbasis(atom.axes, atom.orbitals...)

function localbasis(axes::AbstractMatrix{T}, ::Orbital"*s") where {T}
    return fill(one(T), 1, 1)
end

function localbasis(axes::AbstractMatrix{T}, ::Orbital"*s", ::Orbital"*p_x", ::Orbital"*p_y", ::Orbital"*p_z") where {T}
    Z = zero(T) / oneunit(T)
    return [one(T) fill(Z, 1, 3)
            fill(Z, 3, 1) axes]
end

function localbasis(axes::AbstractMatrix{T}, ::Orbital"*p_x", ::Orbital"*p_y") where {T}
    Z = zero(T) / oneunit(T)
    return axes[1:2,1:2]
end

##################################################
# Position operator
##################################################

_position(::Val{:x}, atom::Atom) = atom.position[1]
_position(::Val{:y}, atom::Atom) = atom.position[2]
_position(::Val{:z}, atom::Atom) = atom.position[3]

position(x::Symbol, args...) = position(Val(x), args...)

position(::Val, s1::Orbital, s2::Orbital) = s1 == s2 ? 1.0 : 0.0

function position(val::Val, atom::Atom)
    N = countorbitals(atom)
    U = localbasis(atom)
    X = fill(0.0im, N, N)
    for (j, s2) in enumerate(atom.orbitals), (i, s1) in enumerate(atom.orbitals)
        X[i, j] = _position(val, atom) * position(val, s1, s2)
    end
    return U'X*U
end

function position(val::Val, mol::Molecule)
    N = countorbitals(mol)
    X = fill(0.0im, N, N)

    i = 1
    for atom in mol.atoms
        Ni = countorbitals(atom)
        X[i:i+Ni-1, i:i+Ni-1] = position(val, atom)
        i += Ni
    end

    return X
end

# function position(val::Val, mol::Molecule, n)
#     N = countorbitals(mol)
#     X = fill(0.0im, N, N)
#     inds = indices(mol, mol[n])
#     X[inds, inds] = position(val, mol[n])
#     return X
# end

##################################################
# Angular momentum
##################################################

angularmomentum(vec::AbstractVector, args...) =
    vec[1]*angularmomentum(:x, args...) +
    vec[2]*angularmomentum(:y, args...) +
    vec[3]*angularmomentum(:z, args...)

angularmomentum(x::Symbol, args...) = angularmomentum(Val(x), args...)

angularmomentum(::Val, ::Orbital"*s", ::Orbital"*s") = 0.0
angularmomentum(::Val, ::Orbital"*p_{*}", ::Orbital"*s") = 0.0
angularmomentum(::Val, ::Orbital"*s", ::Orbital"*p_{*}") = 0.0
angularmomentum(::Val, ::Orbital"*p_{*}", ::Orbital"*p_{*}") = 0.0
angularmomentum(::Val{:x}, ::Orbital"*p_y", ::Orbital"*p_z") = -im
angularmomentum(::Val{:x}, ::Orbital"*p_z", ::Orbital"*p_y") = im
angularmomentum(::Val{:y}, ::Orbital"*p_z", ::Orbital"*p_x") = -im
angularmomentum(::Val{:y}, ::Orbital"*p_x", ::Orbital"*p_z") = im
angularmomentum(::Val{:z}, ::Orbital"*p_x", ::Orbital"*p_y") = -im
angularmomentum(::Val{:z}, ::Orbital"*p_y", ::Orbital"*p_x") = im

function angularmomentum(val::Val, atom::Atom)
    N = countorbitals(atom)
    U = localbasis(atom)
    L = fill(0.0im, N, N)
    for (j, s2) in enumerate(atom.orbitals), (i, s1) in enumerate(atom.orbitals)
        L[i, j] = angularmomentum(val, s1, s2)
    end
    return U'L*U
end

function angularmomentum(val::Val, mol::Molecule)
    N = countorbitals(mol)
    L = fill(0.0im, N, N)

    i = 1
    for atom in mol.atoms
        Ni = countorbitals(atom)
        L[i:i+Ni-1, i:i+Ni-1] = angularmomentum(val, atom)
        i += Ni
    end

    return L
end

function angularmomentum(val::Val, mol::Molecule, atom::Atom)
    N = countorbitals(mol)
    L = fill(0.0im, N, N)
    inds = indices(mol, atom)
    L[inds, inds] = angularmomentum(val, atom)
    return L
end

function angularmomentum(val::Val, mol::Molecule, n::Int)
    N = countorbitals(mol)
    L = fill(0.0im, N, N)
    inds = indices(mol, mol[n])
    L[inds, inds] = angularmomentum(val, mol[n])
    return L
end

##################################################
# Symmetries
##################################################

abstract type AbstractSymmetry end
struct Mirror{S} <: AbstractSymmetry end

parity(::Mirror{:X}, ::Orbital"*s") = 1.0
parity(::Mirror{:X}, ::Orbital"*p_*") = 1.0
parity(::Mirror{:X}, ::Orbital"*p_x") = -1.0

parity(::Mirror{:Y}, ::Orbital"*s") = 1.0
parity(::Mirror{:Y}, ::Orbital"*p_*") = 1.0
parity(::Mirror{:Y}, ::Orbital"*p_y") = -1.0

parity(::Mirror{:Z}, ::Orbital"*s") = 1.0
parity(::Mirror{:Z}, ::Orbital"*p_*") = 1.0
parity(::Mirror{:Z}, ::Orbital"*p_z") = -1.0

(s::Mirror)(a::T) where {T <: Atom} = T(s(a.position), a.orbitals...; valence=a.valence, axes=a.axes)

(s::Mirror{:X})(v::T) where {T <: AbstractVector} = convert(T, [-v[1], v[2], v[3]])

##################################################
# Distribution functions
##################################################

# In atomic units, room temperature is 300K / 315774K = 9.5e-4
nB(E, T = 9.5e-4) = 1/(exp(E/T) - 1)
nF(E, T = 9.5e-4) = 1/(exp(E/T) + 1)

##################################################
# Vibrations
##################################################

"""
    vibrationalcoupling(mol, atom, ω0, dr)

Calculate the numerical derivative of the `hamiltonian(mol)` w.r.t. a
displacement along `dr` of the position of `atom`.  `ω0` is the
characteristic frequency of the atomic vibration.

"""
function vibrationalcoupling(mol::Molecule, atom::Atom, ω0, dr)
    m = Molecules.mass(atom)
    H0 = hamiltonian(mol)
    atom.position += dr
    H′ = hamiltonian(mol)
    atom.position -= dr
    return (H′ - H0) / norm(dr) * √(1 / (2 * m * ω0))
end

##################################################
# Plots recipes
##################################################

include("recipe.jl")
include("blender.jl")
include("xyz.jl")

end #module