using Oceananigans.Operators: index_and_interp_dependencies
using Oceananigans.Utils: tupleit, user_function_arguments

import Oceananigans: location

"""
    struct ContinuousBoundaryFunction{X, Y, Z, I, F, P, D, N, ℑ} <: Function

A wrapper for the user-defined boundary condition function `func` at location
`X, Y, Z`. `I` denotes the boundary-normal index (`I=1` at western boundaries,
`I=grid.Nx` at eastern boundaries, etc). `F, P, D, N, ℑ` are, respectively, the
user-defined function, parameters, field dependencies, indices of the field dependencies
in `model_fields`, and interpolation operators for interpolating `model_fields` to the
location at which the boundary condition is applied.
"""
struct ContinuousBoundaryFunction{X, Y, Z, I, F, P, D, N, ℑ} <: Function
                          func :: F
                    parameters :: P
            field_dependencies :: D
    field_dependencies_indices :: N
     field_dependencies_interp :: ℑ

    """ Returns a location-less wrapper for `func`, `parameters`, and `field_dependencies`."""
    function ContinuousBoundaryFunction(func::F, parameters::P, field_dependencies) where {F, P}
        field_dependencies = tupleit(field_dependencies)
        D = typeof(field_dependencies)
        return new{Nothing, Nothing, Nothing, Nothing, F, P, D, Nothing, Nothing}(func, parameters, field_dependencies, nothing, nothing)
    end

    function ContinuousBoundaryFunction{X, Y, Z, I}(func::F,
                                                    parameters::P,
                                                    field_dependencies::D,
                                                    field_dependencies_indices::N,
                                                    field_dependencies_interp::ℑ) where {X, Y, Z, I, F, P, D, ℑ, N}

        return new{X, Y, Z, I, F, P, D, N, ℑ}(func, parameters, field_dependencies, field_dependencies_indices, field_dependencies_interp)
    end
end

location(::ContinuousBoundaryFunction{X, Y, Z}) where {X, Y, Z} = X, Y, Z

#####
##### "Regularization" for IncompressibleModel setup
#####

"""
    regularize_boundary_condition(bc::BoundaryCondition{C, <:ContinuousBoundaryFunction},
                                  topo, loc, dim, I, prognostic_field_names) where C

Regularizes `bc.condition` for location `loc`, boundary index `I`, and `prognostic_field_names`,
returning `BoundaryCondition(C, regularized_condition)`.

The regularization of `bc.condition::ContinuousBoundaryFunction` requries

1. Setting the boundary location to `LX, LY, LZ`.
   The location in the boundary-normal direction is `Nothing`.

2. Setting the boundary-normal index `I` for indexing into `field_dependencies`.
   `I` is either `1` (for left boundaries) or
   `size(grid, n)` for a boundary in the `n`th direction where `n ∈ (1, 2, 3)` corresponds
   to `x, y, z`.

3. Determining the `indices` that map `model_fields` to `field_dependencies`.

4. Determining the `interps` functions that interpolate field_dependencies to the location
   of the boundary.
"""
function regularize_boundary_condition(bc::BoundaryCondition{C, <:ContinuousBoundaryFunction},
                                       topo, loc, dim, I, prognostic_field_names) where C

    boundary_func = bc.condition

    # Set boundary-normal location to Nothing:
    LX, LY, LZ = Tuple(i == dim ? Nothing : loc[i] for i = 1:3)

    indices, interps = index_and_interp_dependencies(LX, LY, LZ,
                                                     boundary_func.field_dependencies,
                                                     prognostic_field_names)

    regularized_boundary_func = ContinuousBoundaryFunction{LX, LY, LZ, I}(boundary_func.func,
                                                                          boundary_func.parameters,
                                                                          boundary_func.field_dependencies,
                                                                          indices, interps)

    return BoundaryCondition(C, regularized_boundary_func)
end

@inline boundary_index(i) = ifelse(i == 1, 1, i + 1) # convert near-boundary Center() index into boundary Face() index

#####
##### Kernel functions
#####

# Return ContinuousBoundaryFunction on east or west boundaries.
@inline function (bc::ContinuousBoundaryFunction{Nothing, LY, LZ, i})(j, k, grid, clock, model_fields) where {LY, LZ, i}
    args = user_function_arguments(i, j, k, grid, model_fields, bc.parameters, bc)

    i′ = boundary_index(i)

    return bc.func(ynode(Face(), LY(), LZ(), i′, j, k, grid),
                   znode(Face(), LY(), LZ(), i′, j, k, grid),
                   clock.time, args...)
end

# Return ContinuousBoundaryFunction on south or north boundaries.
@inline function (bc::ContinuousBoundaryFunction{LX, Nothing, LZ, j})(i, k, grid, clock, model_fields) where {LX, LZ, j}
    args = user_function_arguments(i, j, k, grid, model_fields, bc.parameters, bc)

    j′ = boundary_index(j)

    return bc.func(xnode(LX(), Face(), LZ(), i, j′, k, grid),
                   znode(LX(), Face(), LZ(), i, j′, k, grid),
                   clock.time, args...)
end

# Return ContinuousBoundaryFunction on bottom or top boundaries.
@inline function (bc::ContinuousBoundaryFunction{LX, LY, Nothing, k})(i, j, grid, clock, model_fields) where {LX, LY, k}
    args = user_function_arguments(i, j, k, grid, model_fields, bc.parameters, bc)

    k′ = boundary_index(k)

    return bc.func(xnode(LX(), LY(), Face(), i, j, k′, grid),
                   ynode(LY(), LY(), Face(), i, j, k′, grid),
                   clock.time, args...)
end

# Don't re-convert ContinuousBoundaryFunctions passed to BoundaryCondition constructor
BoundaryCondition(Classification::DataType, condition::ContinuousBoundaryFunction) = BoundaryCondition(Classification(), condition)

Adapt.adapt_structure(to, bf::ContinuousBoundaryFunction{LX, LY, LZ, I}) where {LX, LY, LZ, I} =
    ContinuousBoundaryFunction{LX, LY, LZ, I}(Adapt.adapt(to, bf.func),
                                              Adapt.adapt(to, bf.parameters),
                                              nothing,
                                              Adapt.adapt(to, bf.field_dependencies_indices),
                                              Adapt.adapt(to, bf.field_dependencies_interp))


function merge_groups(g::Symbol...)
    return unique(vcat(getindex.(Ref(parameter_groups), g)...))
end

# initialize parameter groups with all basic sets and then add combined ones
const parameter_groups = Dict(
    :standard => ["molec_id", "local_iso_id", "nu", "sw", "a", "elower", "gamma_air",
                "delta_air","gamma_self","n_air","n_self","gp","gpp"],

    :labels => ["statep", "statepp"],

    :voigt_air => ["gamma_air","delta_air","deltap_air","n_air"],
    :voigt_self => ["gamma_self","delta_self","deltap_self","n_self"],
    :voigt_h2 => ["gamma_H2","delta_H2","deltap_H2","n_H2"],
    :voigt_co2 => ["gamma_CO2","delta_CO2","n_CO2"],
    :voigt_he => ["gamma_He","delta_He","n_He"],
    :voigt_h2o => ["gamma_H2O","n_H2O"],
    :voigt_linemixing => ["y_air","y_self"],

    :sdvoigt_air => ["gamma_air","delta_air","deltap_air","n_air", "SD_air"],
    :sdvoigt_self => ["gamma_self","delta_self","deltap_self","n_self", "SD_self"],
    :sdvoigt_linemixing => ["Y_SDV_air_296","Y_SDV_self_296"],

    :ht_self => ["gamma_HT_0_self_50","n_HT_self_50","gamma_HT_2_self_50",
                "delta_HT_0_self_50","deltap_HT_self_50","delta_HT_2_self_50",
                "gamma_HT_0_self_150","n_HT_self_150","gamma_HT_2_self_150",
                "delta_HT_0_self_150","deltap_HT_self_150","delta_HT_2_self_150",
                "gamma_HT_0_self_296","n_HT_self_296","gamma_HT_2_self_296",
                "delta_HT_0_self_296","deltap_HT_self_296","delta_HT_2_self_296",
                "gamma_HT_0_self_700","n_HT_self_700","gamma_HT_2_self_700",
                "delta_HT_0_self_700","deltap_HT_self_700","delta_HT_2_self_700",
                "nu_HT_self","kappa_HT_self","eta_HT_self","Y_HT_self_296"],
    :ht_air => ["gamma_HT_0_air_296","n_HT_air_296","gamma_HT_2_air_296",
                "delta_HT_0_air_296","deltap_HT_air_296","delta_HT_2_air_296",
                "nu_HT_air","kappa_HT_air","eta_HT_air","Y_HT_air_296"]

)
parameter_groups[:voigt_all] = merge_groups(:voigt_air, :voigt_self, :voigt_h2, :voigt_co2, :voigt_he, :voigt_h2o, :voigt_linemixing)
parameter_groups[:sdvoigt_all] = merge_groups(:sdvoigt_air, :sdvoigt_self, :sdvoigt_linemixing)
parameter_groups[:ht_all] = merge_groups(:ht_self, :ht_air)

parameter_groups[:all] = merge_groups(:standard, :labels, :voigt_all, :sdvoigt_all, :ht_all)


module AIBECS

using Reexport
@reexport using SparseArrays
using SuiteSparse
using LinearAlgebra
using Dates
using Printf
using ForwardDiff
@reexport using UnPack
using FieldMetadata
using Flatten
using DataFrames
using Distributions
using JLD2
using CodecZlib
using NCDatasets
@reexport using OceanGrids
@reexport using Unitful
using Unitful: °
using RecipesBase
using UnitfulRecipes
using Interpolations
using Distances
using NearestNeighbors
@reexport using MetadataArrays
using Shapefile
using Bijectors
using ImageFiltering
using SciMLBase

include("CirculationGeneration.jl")
include("SinkingParticles.jl")
include("newTypes.jl")
include("CTKsolvers.jl")
include("Parameters.jl")
include("overload_solve.jl")
include("multiTracer.jl")

# Available circulations
include("OCIM0.jl")
include("OCIM1.jl")
include("OCIM2.jl")
include("OCIM2_48L.jl")
include("OCCA.jl")
include("Archer_etal_2000.jl")
include("TwoBoxModel.jl")
include("Primeau_2x2x2.jl")

# AWESOME OCIM data
include("AO.jl") # TODO talk about it to Seth

# Aeolian source data
include("aeolian_sources.jl")

# River source data
include("Rivers.jl")

# Groundwater source data
include("GroundWaters.jl")

# ETOPO data
include("ETOPO.jl")

# Recipes for plotting
include("plot_recipes.jl")

# Diagnostics
include("diagnostics.jl")

end # module


"""
    `poly(vertices, indices; kwargs...)`
    `poly(points; kwargs...)`
    `poly(shape; kwargs...)`

Plots a polygon based on the arguments given.
When vertices and indices are given, it functions similarly to `mesh`.
When points are given, it draws one polygon that connects all the points in order.
When a shape is given (essentially anything decomposable by `GeometryTypes`), it will plot `decompose(shape)`.

    poly(coordinates, connectivity; kwargs...)

Plots polygons, which are defined by
`coordinates` (the coordinates of the vertices) and
`connectivity` (the edges between the vertices).

## Theme
$(ATTRIBUTES)
"""
@recipe(Poly) do scene
    Theme(;
        color = theme(scene, :color),
        visible = theme(scene, :visible),
        strokecolor = RGBAf0(0,0,0,0),
        colormap = theme(scene, :colormap),
        colorrange = automatic,
        strokewidth = 0.0,
        shading = false,
        linestyle = nothing,
        overdraw = false,
        transparency = false,
    )
end
convert_arguments(::Type{<: Poly}, v::AbstractVector{<: AbstractVector{<: VecTypes}}) = (v,)
convert_arguments(::Type{<: Poly}, v::AbstractVector{<: VecTypes}) = ([convert_arguments(Scatter, v)[1]],)
convert_arguments(::Type{<: Poly}, v::AbstractVector{<: Union{Circle, Rectangle, HyperRectangle}}) = (v,)
convert_arguments(::Type{<: Poly}, args...) = ([convert_arguments(Scatter, args...)[1]],)
convert_arguments(::Type{<: Poly}, vertices::AbstractArray, indices::AbstractArray) = convert_arguments(Mesh, vertices, indices)

function plot!(plot::Poly{<: Tuple{Union{AbstractMesh, GeometryPrimitive}}})
    mesh!(
        plot, plot[1],
        color = plot[:color], colormap = plot[:colormap], colorrange = plot[:colorrange],
        shading = plot[:shading], visible = plot[:visible], overdraw = plot[:overdraw]
    )
    wireframe!(
        plot, plot[1],
        color = plot[:strokecolor], linestyle = plot[:linestyle],
        linewidth = plot[:strokewidth], visible = plot[:visible], overdraw = plot[:overdraw]
    )
end

function plot!(plot::Poly{<: Tuple{<: AbstractVector{P}}}) where P <: AbstractVector{<: VecTypes}
    polygons = plot[1]
    color_node = plot[:color]
    attributes = Attributes()
    meshes = lift(polygons) do polygons
        polys = Vector{Point2f0}[]
        for poly in polygons
            s = GeometryTypes.split_intersections(poly)
            append!(polys, s)
        end
        GLNormalMesh.(polys)
    end
    mesh!(plot, meshes,
        visible = plot.visible,
        shading = plot.shading,
        color = plot.color,
        colormap = plot.colormap,
        colorrange = plot.colorrange,
        overdraw = plot.overdraw,
        transparency = plot.transparency
    )
    outline = lift(polygons) do polygons
        line = Point2f0[]
        for poly in polygons
            append!(line, poly)
            push!(line, poly[1])
            push!(line, Point2f0(NaN))
        end
        line
    end
    lines!(
        plot, outline, visible = plot.visible,
        color = plot.strokecolor, linestyle = plot.linestyle,
        linewidth = plot.strokewidth,
        overdraw = plot.overdraw, transparency = plot.transparency
    )
end

function plot!(plot::Mesh{<: Tuple{<: AbstractVector{P}}}) where P <: AbstractMesh
    meshes = plot[1]
    color_node = plot[:color]
    attributes = Attributes(
        visible = plot[:visible], shading = plot[:shading]
    )
    if haskey(plot, :colormap)
        attributes[:colormap] = plot[:colormap]
    end
    if haskey(plot, :colorrange)
        attributes[:colorrange] = plot[:colorrange]
    end
    bigmesh = if color_node[] isa Vector && length(color_node[]) == length(meshes[])
        lift(meshes, color_node) do meshes, colors
            meshes = GeometryTypes.add_attribute.(GLNormalMesh.(meshes), to_color.(colors))
            merge(meshes)
        end
    else
        attributes[:color] = color_node
        lift(meshes) do meshes
            merge(GLPlainMesh.(meshes))
        end
    end
    mesh!(plot, attributes, bigmesh)
end

function plot!(plot::Poly{<: Tuple{<: AbstractVector{T}}}) where T <: Union{Circle, Rectangle, Rect}
    positions = plot[1]
    position = lift(positions) do rects
        Point.(minimum.(rects))
    end
    markersize = lift(positions, name = "markersize") do rects
        widths.(rects)
    end
    scatter!(
        plot, position,
        marker = T, markersize = markersize, transform_marker = true,
        marker_offset = Vec2f0(0),
        color = plot[:color],
        strokecolor = plot[:strokecolor],
        colormap = plot[:colormap],
        colorrange = plot[:colorrange],
        strokewidth = plot[:strokewidth], visible = plot[:visible]
    )
end
function data_limits(p::Poly{<: Tuple{<: AbstractVector{T}}}) where T <: Union{Circle, Rectangle, Rect}
    xyz = p.plots[1][1][]
    msize = p.plots[1][:markersize][]
    xybb = FRect3D(xyz)
    mwidth = FRect3D(xyz .+ msize)
    union(mwidth, xybb)
end

"""
    `arrows(points, directions; kwargs...)`
    `arrows(x, y, u, v)`
    `arrows(x::AbstractVector, y::AbstractVector, u::AbstractMatrix, v::AbstractMatrix)`
    `arrows(x, y, z, u, v, w)`

Plots arrows at the specified points with the specified components.
`u` and `v` are interpreted as vector components (`u` being the x
and `v` being the y), and the vectors are plotted with the tails at
`x`, `y`.

If `x, y, u, v` are `<: AbstractVector`, then each 'row' is plotted
as a single vector.

If `u, v` are `<: AbstractMatrix`, then `x` and `y` are interpreted as
specifications for a grid, and `u, v` are plotted as arrows along the
grid.

`arrows` can also work in three dimensions.

## Theme
$(ATTRIBUTES)
"""
@recipe(Arrows, points, directions) do scene
    theme = Theme(
        arrowhead = automatic,
        arrowtail = nothing,
        linecolor = :black,
        linewidth = 1,
        arrowsize = 0.3,
        linestyle = nothing,
        scale = Vec3f0(1),
        normalize = false,
        lengthscale = 1.0f0
    )
    # connect arrow + linecolor by default
    get!(theme, :arrowcolor, theme[:linecolor])
    theme
end

# For the matlab/matplotlib users
const quiver = arrows
const quiver! = arrows!
export quiver, quiver!

arrow_head(N, marker) = marker
arrow_head(N, marker::Automatic) = N == 2 ? '▲' : Pyramid(Point3f0(0, 0, -0.5), 1f0, 1f0)

scatterfun(N) = N == 2 ? scatter! : meshscatter!

convert_arguments(::Type{<: Arrows}, x, y, u, v) = (Point2f0.(x, y), Vec2f0.(u, v))
function convert_arguments(::Type{<: Arrows}, x::AbstractVector, y::AbstractVector, u::AbstractMatrix, v::AbstractMatrix)
    (vec(Point2f0.(x, y')), vec(Vec2f0.(u, v)))
end
convert_arguments(::Type{<: Arrows}, x, y, z, u, v, w) = (Point3f0.(x, y, z), Vec3f0.(u, v, w))

function plot!(arrowplot::Arrows{<: Tuple{AbstractVector{<: Point{N, T}}, V}}) where {N, T, V}
    @extract arrowplot (points, directions, lengthscale, arrowhead, arrowsize, arrowcolor)
    headstart = lift(points, directions, lengthscale) do points, directions, s
        map(points, directions) do p1, dir
            dir = arrowplot[:normalize][] ? StaticArrays.normalize(dir) : dir
            Point{N, Float32}(p1) => Point{N, Float32}(p1 .+ (dir .* Float32(s)))
        end
    end
    linesegments!(
        arrowplot, headstart,
        color = arrowplot[:linecolor], linewidth = arrowplot[:linewidth],
        linestyle = arrowplot[:linestyle],
    )
    scatterfun(N)(
        arrowplot,
        lift(x-> last.(x), headstart),
        marker = lift(x-> arrow_head(N, x), arrowhead), markersize = arrowsize,
        color = arrowcolor, rotations = directions
    )
end

"""
    `wireframe(x, y, z)`, `wireframe(positions)`, or `wireframe(mesh)`

Draws a wireframe, either interpreted as a surface or as a mesh.

## Theme
$(ATTRIBUTES)
"""
@recipe(Wireframe) do scene
    default_theme(scene, LineSegments)
end

function argument_conversion(::Type{Wireframe}, x::AbstractVector, y::AbstractVector, z::AbstractMatrix)
    (ngrid(x, y)..., z)
end

xvector(x::AbstractVector, len) = x
xvector(x::ClosedInterval, len) = range(minimum(x), stop=maximum(x), length=len)
xvector(x::AbstractMatrix, len) = x

yvector(x, len) = xvector(x, len)'
yvector(x::AbstractMatrix, len) = x

function plot!(plot::Wireframe{<: Tuple{<: Any, <: Any, <: AbstractMatrix}})
    points_faces = lift(plot[1:3]...) do x, y, z
        T = eltype(z); M, N = size(z)
        points = vec(Point3f0.(xvector(x, M), yvector(y, N), z))
        # Connect the vetices with faces, as one would use for a 2D Rectangle
        # grid with M,N grid points
        faces = decompose(Face{2, GLIndex}, SimpleRectangle(0, 0, 1, 1), (M, N))
        view(points, faces)
    end
    linesegments!(plot, Theme(plot), points_faces)
end

function plot!(plot::Wireframe{Tuple{T}}) where T
    points = lift(plot[1]) do g
        # get the point representation of the geometry
        indices = decompose(Face{2, GLIndex}, g)
        points = decompose(Point3f0, g)
        view(points, indices)
    end
    linesegments!(plot, Theme(plot), points)
end

function sphere_streamline(linebuffer, ∇ˢf, pt, h, n)
    push!(linebuffer, pt)
    df = normalize(∇ˢf(pt[1], pt[2], pt[3]))
    push!(linebuffer, normalize(pt .+ h*df))
    for k=2:n
        cur_pt = last(linebuffer)
        push!(linebuffer, cur_pt)
        df = normalize(∇ˢf(cur_pt...))
        push!(linebuffer, normalize(cur_pt .+ h*df))
    end
    return
end

"""
    StreamLines

TODO add function signatures
TODO add descripton

## Theme
$(ATTRIBUTES)
"""
@recipe(StreamLines, points, directions) do scene
    Theme(
        h = 0.01f0,
        n = 5,
        color = :black,
        linewidth = 1
    )
end

# typeof(([1, 2, 3] |> vec, [1, 2, 3] |> vec)) <: Tuple{<: AbstractVector{T}, <: AbstractVector{T}} where T
# true
# streamlines([1, 2, 3] |> vec, [1, 2, 3] |> vec)
# ERROR: Plotting for the arguments (::Array{Int64,1}, ::Array{Int64,1}) not defined for AbstractPlotting.streamlines. If you want to support those arguments, overload plot!(plot::AbstractPlotting.streamlines{ <: Tuple{Array{Int64,1},Array{Int64,1}}})

# function plot!(plot::StreamLines) where T
#     @extract plot (points, directions)
#     linebuffer = T[]
#     lines = lift(directions, points, plot[:h], plot[:n]) do ∇ˢf, origins, h, n
#         empty!(linebuffer)
#         for point in origins
#             sphere_streamline(linebuffer, ∇ˢf, point, h, n)
#         end
#         linebuffer
#     end
#     linesegments!(plot, Theme(plot), lines)
# end

"""
    Series - ?

TODO add function signatures
TODO add description

## Theme
$(ATTRIBUTES)
"""
@recipe(Series, series) do scene
    Theme(
        seriescolors = :Set1,
        seriestype = :lines
    )
end
convert_arguments(::Type{<: Series}, A::AbstractMatrix{<: Number}) = (A,)
function plot!(sub::Series)
    A = sub[1]
    colors = map_once(sub[:seriescolors], A) do colors, A
        cmap = to_colormap(colors)
        if size(A, 2) > length(cmap)
            @info("Colormap doesn't have enough distinctive values. Please consider using another value for seriescolors")
            cmap = interpolated_getindex.((cmap,), range(0, stop=1, length=M))
        end
        cmap
    end
    map_once(A, sub[:seriestype]) do A, stype
        empty!(sub.plots)
        N, M = size(A)
        for i = 1:M
            c = lift(getindex, colors, i)
            attributes = Theme(color = c)
            a_view = view(A, :, i)
            if stype in (:lines, :scatter_lines)
                lines!(sub, attributes, a_view)
            end
            # if stype in (:scatter, :scatter_lines)
            #     scatter!(subsub, attributes, a_view)
            # end
            # subsub
        end
    end
    labels = get(sub, :labels) do
        map(i-> "y $i", 1:size(A[], 2))
    end
    legend!(sub, copy(sub.plots), labels)
    sub
end

"""
    `annotations(strings::Vector{String}, positions::Vector{Point})`

Plots an array of texts at each position in `positions`.

## Theme
$(ATTRIBUTES)
"""
@recipe(Annotations, text, position) do scene
    default_theme(scene, Text)
end

"""
Pushes an updates to all listeners of `node`
"""
function notify!(node::Node)
    node[] = node[]
end

function plot!(plot::Annotations)
    position = plot[2]
    sargs = (
        plot[:model], plot[:font],
        plot[1], position,
        getindex.(plot, (:color, :textsize, :align, :rotation))...,
    )
    N = to_value(position) |> eltype |> length
    atlas = get_texture_atlas()
    combinedpos = [Point3f0(0)]; colors = RGBAf0[RGBAf0(0,0,0,0)]
    scales = Vec2f0[(0,0)]; fonts = NativeFont[to_font("Dejavu Sans")]
    rotations = Quaternionf0[Quaternionf0(0,0,0,0)]

    tplot = text!(plot, "",
        align = Vec2f0(0), model = Mat4f0(I),
        position = combinedpos, color = colors,
        textsize = scales, font = fonts, rotation = rotations
    ).plots[end]

    onany(sargs...) do model, font, args...
        if length(args[1]) != length(args[2])
            error("For each text annotation, there needs to be one position. Found: $(length(t)) strings and $(length(p)) positions")
        end
        io = IOBuffer();
        empty!(combinedpos); empty!(colors); empty!(scales); empty!(fonts); empty!(rotations)

        broadcast_foreach(1:length(args[1]), args...) do idx, text, startpos, color, tsize, alignment, rotation
            # the fact, that Font == Vector{FT_FreeType.Font} is pretty annoying for broadcasting.
            # TODO have a better Font type!
            f = to_font(font)
            f = isa(f, NativeFont) ? f : f[idx]
            c = to_color(color)
            rot = to_rotation(rotation)
            pos, s = layout_text(text, startpos, tsize, f, alignment, rot, model)
            print(io, text)
            n = length(pos)
            append!(combinedpos, pos)
            append!(scales, s)
            append!(colors, repeated(c, n))
            append!(fonts, repeated(f, n))
            append!(rotations, repeated(rot, n))
        end
        str = String(take!(io))
        # update string the signals
        tplot[1] = str
        tplot[:scales] = scales
        tplot[:color] = colors
        tplot[:rotation] = rotations
        # fonts shouldn't need an update, since it will get udpated when listening on string
        #
        return
    end
    # update one time in the beginning, since otherwise the above won't run
    notify!(sargs[1])
    plot
end

"""
    arc(origin, radius, start_angle, stop_angle; kwargs...)

This function plots a circular arc, centered at `origin` with radius `radius`,
from `start_angle` to `stop_angle`.
`origin` must be a coordinate in 2 dimensions (i.e., a `Point2`); the rest of the arguments must be
`<: Number`.

Examples:

`arc(Point2f0(0), 1, 0.0, π)`
`arc(Point2f0(1, 2), 0.3. π, -π)`

## Theme
$(ATTRIBUTES)
"""
@recipe(Arc, origin, radius, start_angle, stop_angle) do scene
    Theme(;
        default_theme(scene, Lines)...,
        resolution = 361,
    )
end

function plot!(p::Arc)
    args = getindex.(p, (:origin, :radius, :start_angle, :stop_angle, :resolution))
    positions = lift(args...) do origin, radius, start_angle, stop_angle, resolution
        map(range(start_angle, stop=stop_angle, length=resolution)) do angle
            origin .+ Point2f0((sin(angle), cos(angle)) .* radius)
        end
    end
    lines!(p, Theme(p), positions)
end

function AbstractPlotting.plot!(plot::Plot(AbstractVector{<: Complex}))
    plot[:axis, :labels] = ("Re(x)", "Im(x)")
    lines!(plot, lift(im-> Point2f0.(real.(im), imag.(im)), x[1]))
end

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

Plots a barplot; `y` defines the height.  `x` and `y` should be 1 dimensional.

## Theme
$(ATTRIBUTES)
"""
@recipe(BarPlot, x, y) do scene
    Theme(;
        fillto = 0.0,
        color = theme(scene, :color),
        colormap = theme(scene, :colormap),
        colorrange = automatic,
        marker = Rect,
        strokewidth = 0,
        strokecolor = :white,
        width = automatic
    )
end

function data_limits(p::BarPlot)
    xy = p.plots[1][1][]
    msize = p.plots[1][:markersize][]
    xstart, xend = first(msize)[1], last(msize)[1]
    xvals = first.(xy)
    # correct widths
    xvals[1] = xvals[1] - (xstart / 2)
    xvals[end] = xvals[end] + (xend / 2)
    xybb = FRect3D(xy)
    y = last.(msize) .+ last.(xy)
    bb = xyz_boundingbox(xvals, y)
    union(bb, xybb)
end

conversion_trait(::Type{<: BarPlot}) = PointBased()

function AbstractPlotting.plot!(p::BarPlot)
    pos_scale = lift(p[1], p[:fillto], p[:width]) do xy, fillto, hw
        # compute half-width of bars
        if hw === automatic
            hw = mean(diff(first.(xy))) # TODO ignore nan?
        end
        # make fillto a vector... default fills to 0
        positions = Point2f0.(first.(xy), Float32.(fillto))
        scales = Vec2f0.(abs.(hw), last.(xy))
        offset = Vec2f0.(hw ./ -2f0, 0)
        positions, scales, offset
    end

    scatter!(
        p, lift(first, pos_scale),
        marker = p[:marker], marker_offset = lift(last, pos_scale),
        markersize = lift(getindex, pos_scale, Node(2)),
        color = p[:color], colormap = p[:colormap], colorrange = p[:colorrange],
        transform_marker = true, strokewidth = p[:strokewidth],
        strokecolor = p[:strokecolor]
    )
end

function convert_arguments(P::PlotFunc, r::AbstractVector, f::Function)
    ptype = plottype(P, Lines)
    to_plotspec(ptype, convert_arguments(ptype, r, f.(r)))
end

function convert_arguments(P::PlotFunc, i::AbstractInterval, f::Function)
    convert_arguments(P, PlotUtils.adapted_grid(f, endpoints(i)), f)
end

to_tuple(t::Tuple) = t
to_tuple(t) = (t,)

function convert_arguments(P::PlotFunc, f::Function, args...; kwargs...)
    tmp = f(args...; kwargs...) |> to_tuple
    convert_arguments(P, tmp...)
end

"""
    scatterlines(xs, ys, [zs]; kwargs...)

Plots `lines` between sets of x and y coordinates provided,
as well as plotting those points using `scatter`.

## Theme
$(ATTRIBUTES)
"""
@recipe(ScatterLines) do scene
    merge(default_theme(scene, Scatter), default_theme(scene, Lines))
end

function plot!(p::Combined{scatterlines, <:NTuple{N, Any}}) where N
    plot!(p, Lines, Theme(p), p[1:N]...)
    plot!(p, Scatter, Theme(p), p[1:N]...)
end

"""
    band(x, ylower, yupper; kwargs...)

Plots a band from `ylower` to `yupper` along `x`.

## Theme
$(ATTRIBUTES)
"""
@recipe(Band, x, ylower, yupper) do scene
    Theme(;
        default_theme(scene, Mesh)...,
        color = RGBAf0(1.0,0,0,0.2)
    )
end

function band_connect(n)
    ns = 1:n-1
    ns2 = n+1:2n-1
    [GLTriangle.(ns, ns .+ 1, ns2); GLTriangle.(ns .+ 1, ns2 .+ 1, ns2)]
end

function plot!(plot::Band)
    coordinates = lift( (x, ylower, yupper) -> [Point2f0.(x, ylower); Point2f0.(x, yupper)], plot[1], plot[2], plot[3])
    connectivity = lift(x -> band_connect(length(x)), plot[1])
    mesh!(plot, coordinates, connectivity;
        color = plot[:color], colormap = plot[:colormap],
        colorrange = plot[:colorrange],
        shading = false, visible = plot[:visible]
    )
end

function fill_view(x, y1, y2, where::Nothing)
  x, y1, y2
end
function fill_view(x, y1, y2, where::Function)
  fill_view(x, y1, y2, where.(x, y1, y2))
end
function fill_view(x, y1, y2, bools::AbstractVector{<: Union{Integer, Bool}})
  view(x, bools), view(y1, bools), view(y2, bools)
end

"""
    fill_between!(x, y1, y2; where = nothing, scene = current_scene(), kw_args...)

fill the section between 2 lines with the condition `where`
"""
function fill_between!(x, y1, y2; where = nothing, scene = current_scene(), kw_args...)
  xv, ylow, yhigh = fill_view(x, y1, y2, where)
  band!(scene, xv, ylow, yhigh; kw_args...)
end

export fill_between!

"""
    contour(x, y, z)
Creates a contour plot of the plane spanning x::Vector, y::Vector, z::Matrix

## Theme
$(ATTRIBUTES)
"""
@recipe(Contour) do scene
    default = default_theme(scene)
    pop!(default, :color)
    Theme(;
        default...,
        color = nothing,
        colormap = theme(scene, :colormap),
        colorrange = AbstractPlotting.automatic,
        levels = 5,
        linewidth = 1.0,
        fillrange = false,
    )
end

"""
    contour3d(x, y, z)
Creates a 3D contour plot of the plane spanning x::Vector, y::Vector, z::Matrix,
with z-elevation for each level.

## Theme
$(ATTRIBUTES)
"""
@recipe(Contour3d) do scene
    default_theme(scene, Contour)
end

function contourlines(::Type{<: Contour}, contours, cols)
    result = Point2f0[]
    colors = RGBA{Float32}[]
    for (color, c) in zip(cols, Contours.levels(contours))
        for elem in Contours.lines(c)
            append!(result, elem.vertices)
            push!(result, Point2f0(NaN32))
            append!(colors, fill(color, length(elem.vertices) + 1))
        end
    end
    result, colors
end

function contourlines(::Type{<: Contour3d}, contours, cols)
    result = Point3f0[]
    colors = RGBA{Float32}[]
    for (color, c) in zip(cols, Contours.levels(contours))
        for elem in Contours.lines(c)
            for p in elem.vertices
                push!(result, Point3f0(p[1], p[2], c.level))
            end
            push!(result, Point3f0(NaN32))
            append!(colors, fill(color, length(elem.vertices) + 1))
        end
    end
    result, colors
end

to_levels(x::AbstractVector{<: Number}, cnorm) = x
function to_levels(n::Integer, cnorm)
    zmin, zmax = cnorm
    dz = (zmax - zmin) / (n + 1)
    range(zmin + dz; step = dz, length = n)
end
conversion_trait(::Type{<: Contour3d}) = SurfaceLike()
conversion_trait(::Type{<: Contour}) = SurfaceLike()
conversion_trait(::Type{<: Contour{<: Tuple{X, Y, Z, Vol}}}) where {X, Y, Z, Vol} = VolumeLike()
conversion_trait(::Type{<: Contour{<: Tuple{<: AbstractArray{T, 3}}}}) where T = VolumeLike()

function plot!(plot::Contour{<: Tuple{X, Y, Z, Vol}}) where {X, Y, Z, Vol}
    x, y, z, volume = plot[1:4]
    @extract plot (colormap, levels, linewidth, alpha)
    valuerange = lift(nan_extrema, volume)
    cliprange = replace_automatic!(plot, :colorrange) do
        valuerange
    end
    cmap = lift(colormap, levels, linewidth, alpha, cliprange, valuerange) do _cmap, l, lw, alpha, cliprange, vrange
        levels = to_levels(l, vrange)
        nlevels = length(levels)
        N = nlevels * 50
        iso_eps = nlevels * ((vrange[2] - vrange[1]) / N) # TODO calculate this
        cmap = to_colormap(_cmap)
        v_interval = cliprange[1] .. cliprange[2]
        # resample colormap and make the empty area between iso surfaces transparent
        map(1:N) do i
            i01 = (i-1) / (N - 1)
            c = AbstractPlotting.interpolated_getindex(cmap, i01)
            isoval = vrange[1] + (i01 * (vrange[2] - vrange[1]))
            line = reduce(levels, init = false) do v0, level
                (isoval in v_interval) || return false
                v0 || (abs(level - isoval) <= iso_eps)
            end
            RGBAf0(Colors.color(c), line ? alpha : 0.0)
        end
    end
    volume!(
        plot, x, y, z, volume, colormap = cmap, colorrange = cliprange, algorithm = 7,
        transparency = plot[:transparency],
        overdraw = plot[:overdraw]
    )
end

function color_per_level(color, colormap, colorrange, alpha, levels)
    color_per_level(to_color(color), colormap, colorrange, alpha, levels)
end

function color_per_level(color::Colorant, colormap, colorrange, alpha, levels)
    fill(color, length(levels))
end

function color_per_level(colors::AbstractVector, colormap, colorrange, alpha, levels)
    color_per_level(to_colormap(colors), colormap, colorrange, alpha, levels)
end

function color_per_level(colors::AbstractVector{<: Colorant}, colormap, colorrange, alpha, levels)
    if length(levels) == length(colors)
        return colors
    else
        # TODO resample?!
        error("For a contour plot, `color` with an array of colors needs to
        have the same length as `levels`.
        Found $(length(colors)) colors, but $(length(levels)) levels")
    end
end

function color_per_level(::Nothing, colormap, colorrange, a, levels)
    cmap = to_colormap(colormap)
    map(levels) do level
        c = interpolated_getindex(cmap, level, colorrange)
        RGBAf0(color(c), alpha(c) * a)
    end
end

function plot!(plot::T) where T <: Union{Contour, Contour3d}
    x, y, z = plot[1:3]
    if to_value(plot[:fillrange])
        plot[:interpolate] = true
        # TODO normalize linewidth for heatmap
        plot[:linewidth] = map(x-> x ./ 10f0, plot[:linewidth])
        heatmap!(plot, Theme(plot), x, y, z)
    else
        zrange = lift(nan_extrema, z)
        levels = lift(plot[:levels], zrange) do levels, zrange
            if levels isa AbstractVector{<: Number}
                return levels
            elseif levels isa Integer
                to_levels(levels, zrange)
            else
                error("Level needs to be Vector of iso values, or a single integer to for a number of automatic levels")
            end
        end
        replace_automatic!(plot, :colorrange) do
            lift(nan_extrema, levels)
        end
        args = @extract plot (color, colormap, colorrange, alpha)
        level_colors = lift(color_per_level, args..., levels)
        result = lift(x, y, z, levels, level_colors) do x, y, z, levels, level_colors
            t = eltype(z)
            # Compute contours
            xv, yv = to_vector(x, size(z,1), t), to_vector(y, size(z,2), t)
            contours = Contours.contours(xv, yv, z,  convert(Vector{eltype(z)}, levels))
            contourlines(T, contours, level_colors)
        end
        lines!(
            plot, lift(first, result);
            color = lift(last, result), linewidth = plot[:linewidth]
        )
    end
    plot
end

function AbstractPlotting.data_limits(x::Contour{<: Tuple{X, Y, Z}}) where {X, Y, Z}
    AbstractPlotting.xyz_boundingbox(to_value.((x[1], x[2]))...)
end

"""
    VolumeSlices

TODO add function signatures
TODO add descripton

## Theme
$(ATTRIBUTES)
"""
@recipe(VolumeSlices, x, y, z, volume) do scene
    Theme(
        colormap = theme(scene, :colormap),
        colorrange = nothing,
        alpha = 0.1,
        contour = Theme(),
        heatmap = Theme(),
    )
end

convert_arguments(::Type{<: VolumeSlices}, x, y, z, volume) = convert_arguments(Contour, x, y, z, volume)
function plot!(vs::VolumeSlices)
    @extract vs (x, y, z, volume)
    replace_automatic!(vs, :colorrange) do
        map(extrema, volume)
    end
    keys = (:colormap, :alpha, :colorrange)
    contour!(vs, vs[:contour], x, y, z, volume)
    planes = (:xy, :xz, :yz)
    hattributes = vs[:heatmap]
    sliders = map(zip(planes, (x, y, z))) do plane_r
        plane, r = plane_r
        idx = node(plane, Node(1))
        vs[plane] = idx
        hmap = heatmap!(vs, hattributes, x, y, zeros(length(x[]), length(y[]))).plots[end]
        on(idx) do i
            transform!(hmap, (plane, r[][i]))
            indices = ntuple(Val(3)) do j
                planes[j] == plane ? i : (:)
            end
            hmap[3][] = view(volume[], indices...)
        end
        idx
    end
    plot!(scene, vs, rest)
end

"""
    showlibrary(lib::Symbol)::Scene

Shows all colour gradients in the given library.
Returns a Scene with these colour gradients arranged
as horizontal colourbars.
"""
function showlibrary(lib::Symbol)::Scene

    cgrads = sort(PlotUtils.cgradients(lib))

    PlotUtils.clibrary(lib)

    showgradients(cgrads)

end

"""
    showgradients(
        cgrads::AbstractVector{Symbol};
        h = 0.0, offset = 0.2, textsize = 0.7,
        resolution = (800, length(cgrads) * 84)
    )::Scene

Plots the given colour gradients arranged as horizontal colourbars.
If you change the offsets or the font size, you may need to change the resolution.
"""
function showgradients(
        cgrads::AbstractVector{Symbol};
        h = 0.0,
        offset = 0.4,
        textsize = 0.7,
        resolution = (800, length(cgrads) * 84),
        monospace = true
    )::Scene

    scene = Scene(resolution = resolution)

    map(collect(cgrads)) do cmap

         c = to_colormap(cmap)

         cbar = image!(
             scene,
             range(0, stop = 10, length = length(c)),
             range(0, stop = 1, length = length(c)),
             reshape(c, (length(c),1)),
             show_axis = false
         )[end]

         cmapstr = monospace ? UnicodeFun.to_latex("\\mono{$cmap}") : string(cmap, ":")

         text!(
             scene,
             cmapstr,
             position = Point2f0(-0.1, 0.5 + h),
             align = (:right, :center),
             show_axis = false,
             textsize = textsize
         )

         translate!(cbar, 0, h, 0)

         h -= (1 + offset)

    end

    scene

end

"""
    timeseries(x::Node{{Union{Number, Point2}}})

Plots a sampled signal.
Usage:
```julia
signal = Node(1.0)
scene = timeseries(signal)
display(scene)
# @async is optional, but helps to continue evaluating more code
@async while isopen(scene)
    # aquire data from e.g. a sensor:
    data = rand()
    # update the signal
    signal[] = data
    # sleep/ wait for new data/ whatever...
    # It's important to yield here though, otherwise nothing will be rendered
    sleep(1/30)
end

```
"""
@recipe(TimeSeries, signal) do scene
    Theme(
        history = 100;
        default_theme(scene, Lines)...
    )
end

signal2point(signal::Number, start) = Point2f0(time() - start, signal)
signal2point(signal::Point2, start) = signal
signal2point(signal, start) = error(""" Signal needs to be of type Number or Point.
Found: $(typeof(signal))
""")

function AbstractPlotting.plot!(plot::TimeSeries)
    # normal plotting code, building on any previously defined recipes
    # or atomic plotting operations, and adding to the combined `plot`:
    points = Node(fill(Point2f0(NaN), plot.history[]))
    buffer = copy(points[])
    lines!(plot, points)
    start = time()
    on(plot.signal) do x
        points[][end] = signal2point(x, start)
        circshift!(buffer, points[], 1)
        buff_ref = buffer
        buffer = points[]
        points[] = buff_ref
        update!(parent(plot))
    end
    plot
end

"""
    streamplot(f::function, xinterval, yinterval;
        kwargs...)
f must either accept `f(::Point)` or `f(x::Number, y::Number)`.
f must return a Point2.
Example:
```julia
using MakieGallery, Makie
run_example("streamplot")
```
## Theme
$(ATTRIBUTES)
"""
@recipe(StreamPlot, f, limits) do scene
    Theme(
        stepsize = 0.01,
        gridsize = (32, 32, 32),
        colormap = theme(scene, :colormap),
        arrow_size = 0.03
    )
end

function convert_arguments(::Type{<: StreamPlot}, f::Function, xrange, yrange)
    xmin, xmax = extrema(xrange)
    ymin, ymax = extrema(yrange)
    return (f, Rect(xmin, ymin, xmax - xmin, ymax - ymin))
end
function convert_arguments(::Type{<: StreamPlot}, f::Function, xrange, yrange, zrange)
    xmin, xmax = extrema(xrange)
    ymin, ymax = extrema(yrange)
    zmin, zmax = extrema(zrange)
    mini = Vec3f0(xmin, ymin, zmin)
    maxi = Vec3f0(xmax, ymax, zmax)
    return (f, Rect(mini, maxi .- mini))
end

function convert_arguments(::Type{<: StreamPlot}, f::Function, limits::Rect)
    return (f, limits)
end

"""
Code adapted from an example implementation by Moritz Schauer (@mschauer)
from https://github.com/JuliaPlots/Makie.jl/issues/355#issuecomment-504449775
"""
function streamplot_impl(CallType, f, limits::Rect{N, T}, resolutionND, stepsize) where {N, T}
    resolution = to_ndim(Vec{N, Int}, resolutionND, last(resolutionND))
    mask = trues(resolution...)
    arrow_pos = Point{N, Float32}[]
    arrow_dir = Vec{N, Float32}[]
    line_points = Point{N, Float32}[]
    colors = Float64[]
    line_colors = Float64[]
    dt = Point{N, Float32}(stepsize)
    mini, maxi = minimum(limits), maximum(limits)
    r = ntuple(N) do i
        LinRange(mini[i], maxi[i], resolution[i] + 1)
    end
    apply_f(x0, P) = if P <: Point
        f(x0)
    else
        f(x0...)
    end
    for c in CartesianIndices(mask)
        x0 = Point(ntuple(N) do i
            first(r[i]) + (c[i] - 0.5) * step(r[i])
        end)
        if mask[c]
            point = apply_f(x0, CallType)
            if !(point isa Point2 || point isa Point3)
                error("Function passed to streamplot must return Point2 or Point3")
            end
            pnorm = norm(point)
            push!(arrow_pos, x0)
            push!(arrow_dir, point ./ pnorm)
            push!(colors, pnorm)
            mask[c] == false
            for d in (-1, 1)
                n_linepoints = 1
                x = x0
                push!(line_points, Point{N, Float32}(NaN), x)
                push!(line_colors, 0.0, pnorm)
                while x in limits && n_linepoints < 500
                    point = apply_f(x, CallType)
                    pnorm = norm(point)
                    x = x .+ d .* dt .* point ./ pnorm
                    if !(x in limits)
                        break
                    end
                    # WHAT? Why does point behave different from tuple in this
                    # broadcast
                    idx = CartesianIndex(searchsortedlast.(r, Tuple(x)))
                    if idx != c
                        if !mask[idx]
                            break
                        end
                        mask[idx] = false
                        c = idx
                    end
                    push!(line_points, x)
                    push!(line_colors, pnorm)
                    n_linepoints += 1
                end
            end
        end
    end
    return (
        arrow_pos,
        arrow_dir,
        line_points,
        colors,
        line_colors,
    )
end

function plot!(p::StreamPlot)
    data = lift(p.f, p.limits, p.gridsize, p.stepsize) do f, limits, resolution, stepsize
        P = if applicable(f, Point2f0(0)) || applicable(f, Point3f0(0))
            Point
        else
            Number
        end
        streamplot_impl(P, f, limits, resolution, stepsize)
    end
    lines!(
        p,
        lift(x->x[3], data), color = lift(last, data), colormap = p.colormap
    )
    N = ndims(p.limits[])
    scatterfun(N)(
        p,
        lift(first, data), markersize = p.arrow_size, marker = arrow_head(N, automatic),
        color = lift(x-> x[4], data), rotations = lift(x-> x[2], data),
        colormap = p.colormap,
    )
end

"""
    spy(x::Range, y::Range, z::AbstractSparseArray)
Visualizes big sparse matrices.
Usage:
```julia
N = 200_000
x = sprand(Float64, N, N, (3(10^6)) / (N*N));
spy(x)
# or if you want to specify the range of x and y:
spy(0..1, 0..1, x)
```
## Theme
$(ATTRIBUTES)
"""
@recipe(Spy, x, y, z) do scene
    Theme(
        marker = automatic,
        markersize = automatic,
        colormap = theme(scene, :colormap),
        colorrange = automatic,
        framecolor = :black,
        framesize = 1,
    )
end

function convert_arguments(::Type{<: Spy}, x::SparseArrays.AbstractSparseArray)
    (0..size(x, 1), 0..size(x, 2), x)
end
function convert_arguments(::Type{<: Spy}, x, y, z::SparseArrays.AbstractSparseArray)
    (x, y, z)
end
function calculated_attributes!(::Type{<: Spy}, plot)
end

function plot!(p::Spy)
    rect = lift(p.x, p.y) do x, y
        xe = extrema(x)
        ye = extrema(y)
        FRect2D((xe[1], ye[1]), (xe[2] - xe[1], ye[2] - ye[1]))
    end
    # TODO FastPixel isn't accepting marker size in data coordinates
    # but instead in pixel - so we need to fix that in GLMakie for consistency
    # and make this nicer when redoing unit support
    markersize = lift(p.markersize, rect, p.z) do msize, rect, z
        if msize === automatic
            widths(rect) ./ Vec2f0(size(z))
        else
            msize
        end
    end
    # TODO correctly align marker
    xycol = lift(rect, p.z, markersize) do rect, z, markersize
        x, y, color = SparseArrays.findnz(z)
        points = map(x, y) do x, y
            (((Point2f0(x, y) .- 1) ./ Point2f0(size(z) .- 1)) .*
            widths(rect) .+ minimum(rect))
        end
        points, color
    end
    replace_automatic!(p, :colorrange) do
        lift(extrema_nan ∘ SparseArrays.nonzeros, p.z)
    end
    marker = lift(p.marker) do x
        if x === automatic
            # If we currently use GLMakie, we can go super fast!
            BackendModule = parentmodule(typeof(AbstractPlotting.current_backend[]))
            if nameof(BackendModule) == :GLMakie
                BackendModule.FastPixel()
            else
                :rect
            end
        else
            x
        end
    end

    scatter!(
        p,
        lift(first, xycol), color = lift(last, xycol),
        marker = marker, markersize = markersize, colorrange = p.colorrange
    )

    lines!(p, rect, color = p.framecolor, linewidth = p.framesize)
end


synth_imgs = [
    "Float64_0" => rand(127, 257),
    "Float64_1" => rand(127, 257, 1),
    "Float64_2" => rand(127, 257, 2),
    "Float64_3" => rand(127, 257, 3),
    "Float64_4" => rand(127, 257, 4),
    "Bool_0" => rand(Bool, 127, 257),
    "Bool_1" => rand(Bool, 127, 257, 1),
    "Bool_2" => rand(Bool, 127, 257, 2),
    "Bool_3" => rand(Bool, 127, 257, 3),
    "Bool_4" => rand(Bool, 127, 257, 4),
    "UInt8_0" => rand(UInt8, 127, 257),
    "UInt8_1" => rand(UInt8, 127, 257, 1),
    "UInt8_2" => rand(UInt8, 127, 257, 2),
    "UInt8_3" => rand(UInt8, 127, 257, 3),
    "UInt8_4" => rand(UInt8, 127, 257, 4),
    "UInt16_0" => rand(UInt16, 127, 257),
    "UInt16_1" => rand(UInt16, 127, 257, 1),
    "UInt16_2" => rand(UInt16, 127, 257, 2),
    "UInt16_3" => rand(UInt16, 127, 257, 3),
    "UInt16_4" => rand(UInt16, 127, 257, 4),
    "N0f8_0" => rand(N0f8, 127, 257),
    "N0f8_1" => rand(N0f8, 127, 257, 1),
    "N0f8_2" => rand(N0f8, 127, 257, 2),
    "N0f8_3" => rand(N0f8, 127, 257, 3),
    "N0f8_4" => rand(N0f8, 127, 257, 4),
    "N0f16_0" => rand(N0f16, 127, 257),
    "N0f16_1" => rand(N0f16, 127, 257, 1),
    "N0f16_2" => rand(N0f16, 127, 257, 2),
    "N0f16_3" => rand(N0f16, 127, 257, 3),
    "N0f16_4" => rand(N0f16, 127, 257, 4),
    "Gray" => rand(Gray, 127, 257),
    "GrayA" => rand(GrayA, 127, 257),
    "RGB" => rand(RGB, 127, 257),
    "RGBA" => rand(RGBA, 127, 257),
    "Gray-N0f8" => rand(Gray{N0f8}, 127, 257),
    "GrayA-N0f8" => rand(GrayA{N0f8}, 127, 257),
    "RGB-N0f8" => rand(RGB{N0f8}, 127, 257),
    "RGBA-N0f8" => rand(RGBA{N0f8}, 127, 257),
    "Gray-N0f16" => rand(Gray{N0f16}, 127, 257),
    "GrayA-N0f16" => rand(GrayA{N0f16}, 127, 257),
    "RGB-N0f16" => rand(RGB{N0f16}, 127, 257),
    "RGBA-N0f16" => rand(RGBA{N0f16}, 127, 257),
    "Gray-N0f8" => rand(Gray{N0f8}, 127, 257),
    "GrayA-N0f8" => rand(GrayA{N0f8}, 127, 257),
    "Gray-Float64" => rand(Gray{Float64}, 127, 257),
    "GrayA-Float64" => rand(GrayA{Float64}, 127, 257),
    "RGB-Float64" => rand(RGB{Float64}, 127, 257),
    "RGBA-Float64" => rand(RGBA{Float64}, 127, 257),
]

edge_case_imgs = [
    ("BitArray_0", x -> _convert(Gray, N7f1, x), randn(127, 257) .> 0),
    ("BitArray_1", x -> _convert(Gray, N7f1, x), randn(127, 257, 1) .> 0),
    ("BitArray_2", x -> _convert(GrayA, N7f1, x), randn(127, 257, 2) .> 0),
    ("BitArray_3", x -> _convert(RGB, N7f1, x), randn(127, 257, 3) .> 0),
    ("BitArray_4", x -> _convert(RGBA, N7f1, x), randn(127, 257, 4) .> 0),
    ("N4f12_0", x -> _convert(Gray, N0f16, x), rand(N4f12, 127, 257)),
    ("N4f12_1", x -> _convert(Gray, N0f16, x), rand(N4f12, 127, 257, 1)),
    ("N4f12_2", x -> _convert(GrayA, N0f16, x), rand(N4f12, 127, 257, 2)),
    ("N4f12_3", x -> _convert(RGB, N0f16, x), rand(N4f12, 127, 257, 3)),
    ("N4f12_4", x -> _convert(RGBA, N0f16, x), rand(N4f12, 127, 257, 4)),
    ("BGR-N0f8", identity, rand(BGR{N0f8}, 127, 257)),
    ("BGR-N01f16", identity, rand(BGR{N0f16}, 127, 257)),
    ("BGR-Float64", identity, BGR{Float64}.(rand(BGR{N0f8}, 127, 257))),
    ("BGRA-N0f8", identity, rand(BGRA{N0f8}, 127, 257)),
    ("BGRA-N0f16", identity, rand(BGRA{N0f16}, 127, 257)),
    ("BGRA-Float64", identity, BGRA{Float64}.(rand(BGRA{N0f8}, 127, 257))),
    ("ABGR-N0f8", identity, rand(ABGR{N0f8}, 127, 257)),
    ("ABGR-N0f16", identity, rand(ABGR{N0f16}, 127, 257)),
    ("ABGR-Float64", identity, ABGR{Float64}.(rand(ABGR{N0f8}, 127, 257))),
    ("ARGB-N0f8", identity, rand(ARGB{N0f8}, 127, 257)),
    ("ARGB-N0f16", identity, rand(ARGB{N0f16}, 127, 257)),
    ("ARGB-Float64", identity, ARGB{Float64}.(rand(ARGB{N0f8}, 127, 257))),
]

@testset "synthetic images" begin
    for (case, image) in synth_imgs
        @testset "$(case)" begin
            expected = collect(PNGFiles._prepare_buffer(image))
            fpath = joinpath(PNG_TEST_PATH, "test_img_$(case).png")
            @testset "write" begin
                open(io->PNGFiles.save(io, image), fpath, "w") #test IO method
                @test PNGFiles.save(fpath, image) == 0
            end
            @testset "read" begin
                global read_in_pngf = PNGFiles.load(fpath)
                @test typeof(read_in_pngf) <: AbstractMatrix
            end
            @testset "compare" begin
                @test all(expected .≈ read_in_pngf)
            end
            global read_in_immag = _standardize_grayness(ImageMagick.load(fpath))
            @testset "$(case): ImageMagick read type equality" begin
                # The lena image is Grayscale saved as RGB...
                @test eltype(_standardize_grayness(read_in_pngf)) == eltype(read_in_immag)
            end
            @testset "$(case): ImageMagick read values equality" begin
                imdiff_val = imdiff(read_in_pngf, read_in_immag)
                onfail(@test imdiff_val < 0.01) do
                    PNGFiles._inspect_png_read(fpath)
                    _add_debugging_entry(fpath, case, imdiff_val)
                end
            end
            path, ext = splitext(fpath)
            newpath = path * "_new" * ext
            PNGFiles.save(newpath, read_in_pngf)
            @testset "$(case): Fpath save IO is idempotent" begin
                @test all(read_in_pngf .≈ PNGFiles.load(newpath))
                @test all(read_in_pngf .≈ open(io->PNGFiles.load(io), newpath))
            end
            open(io->PNGFiles.save(io, read_in_pngf), newpath, "w") #test IO method
            @testset "$(case): Stream save IO is idempotent" begin
                @test all(read_in_pngf .≈ PNGFiles.load(newpath))
                @test all(read_in_pngf .≈ open(io->PNGFiles.load(io), newpath))
            end
        end
    end

    for (case, func_in, image) in edge_case_imgs
        @testset "$(case)" begin
            fpath = joinpath(PNG_TEST_PATH, "test_img_$(case).png")
            @testset "write" begin
                @test PNGFiles.save(fpath, image) == 0
            end
            @testset "read" begin
                global read_in_pngf = PNGFiles.load(fpath)
                @test read_in_pngf isa Matrix
            end
            @testset "compare" begin
                @test all(read_in_pngf .== func_in(image))
            end
            global read_in_immag = _standardize_grayness(ImageMagick.load(fpath))
            @testset "$(case): ImageMagick read type equality" begin
                @test eltype(read_in_pngf) == eltype(read_in_immag)
            end
            @testset "$(case): ImageMagick read values equality" begin
                imdiff_val = imdiff(read_in_pngf, read_in_immag)
                onfail(@test imdiff_val < 0.01) do
                    PNGFiles._inspect_png_read(fpath)
                    _add_debugging_entry(fpath, case, imdiff_val)
                end
            end
            path, ext = splitext(fpath)
            newpath = path * "_new" * ext
            PNGFiles.save(newpath, read_in_pngf)
            @testset "$(case): IO is idempotent" begin
                @test imdiff(read_in_pngf, PNGFiles.load(newpath)) < 0.01
                @test imdiff(read_in_pngf, open(io->PNGFiles.load(io), newpath)) < 0.01
            end
        end
    end
end


# To-do: derive high-fidelity gradient
using DwsimOpt
using JuMP
using Ipopt
using FiniteDifferences
using LinearAlgebra
# using GAMS
# using MadNLP

## SIMULATION OPTIMIZATION PROBLEM DEFINITION
path2sim = "c:/Users/lfsfr/Desktop/DwsimOpt.jl/examples/PRICO/PRICO_composite.dwxmz"
path2dwsim = "C:/Users/lfsfr/AppData/Local/DWSIM7/"

include("c:/Users/lfsfr/Desktop/DwsimOpt.jl/src/SimOpt.jl")
sim_jl = py"sim"

# Declare the simulation optimization problem
py"""
# Import dwsim-python data exchange interface (has to be after Automation2)
from dwsimopt.py2dwsim import create_pddx, assign_pddx

# Assign DoF:
create_pddx( ["MR-1", "CompoundMassFlow", "Nitrogen", "kg/s"],    sim, element="dof" )
create_pddx( ["MR-1", "CompoundMassFlow", "Methane", "kg/s"],     sim, element="dof" )
create_pddx( ["MR-1", "CompoundMassFlow", "Ethane", "kg/s"],      sim, element="dof" )
create_pddx( ["MR-1", "CompoundMassFlow", "Propane", "kg/s"],     sim, element="dof" )
create_pddx( ["VALV-01", "OutletPressure", "Mixture", "Pa"],      sim, element="dof" )
create_pddx( ["COMP-4", "OutletPressure", "Mixture", "Pa"],       sim, element="dof" )

# Assign F
create_pddx( ["Sum_W", "EnergyFlow", "Mixture", "kW"], sim, element="fobj" )

# adding constraints (g_i <= 0):
g1 = create_pddx( ["MITA1-Calc", "OutputVariable", "mita", "°C"], sim, element="constraint", assign=False )
r = 1.0/3
assign_pddx( lambda: r*(3-g1[0]()) , ["MITA1-Calc", "OutputVariable", "mita", "°C"], sim, element="constraint" )
create_pddx( ["MSTR-27", "MassFraction", "Liquid", "x"], sim, element="constraint" )
create_pddx( ["MR-1", "MassFraction", "Liquid", "x"], sim, element="constraint" )
create_pddx( ["MSTR-03", "MassFraction", "Liquid", "x"], sim, element="constraint" )
create_pddx( ["MSTR-05", "MassFraction", "Liquid", "x"], sim, element="constraint" )

# decision variables bounds
# x0 = np.array( [1.07249112e-04, 1.93057244e-04, 2.16881626e-04, 7.44968925e-04, 3.43386624e+05, 6.39810411e+06] )
# x0 = np.array( [1.12083333e-04, 1.30420113e-04, 2.57333445e-04, 8.00266697e-04, 3.45000000e+05, 1.53000000e+02] )
x0 = np.array( [1.12083333e-04, 2.20416667e-04, 1.42216680e-04, 7.54165016e-04, 3.45000000e+05, 7.20000000e+06] )
x0 = np.array( [1.3*1.12083333e-04, 1.3*2.20416667e-04, 1.3*1.42216680e-04, 1.3*7.54165016e-04, 3.45000000e+05, 7.20000000e+06] )
bounds_raw = np.array( [0.5*np.asarray(x0), 1.5*np.asarray(x0)] )   # 50 % around base case

# regularizer calculation
regularizer = np.zeros(x0.size)
regularizer_b = np.zeros(x0.size)
for i in range(len(regularizer)):
    regularizer[i] = 10**(-1*math.floor(math.log(x0[i],10))) # regularizer for magnitude order of 1e0

# bounds regularized
bounds_reg = regularizer*bounds_raw + regularizer_b

# objective and constraints lambda definitions
f = lambda x: sim.calculate_optProblem(np.asarray(x)/regularizer)[0:sim.n_f]
g = lambda x: sim.calculate_optProblem(np.asarray(x)/regularizer)[sim.n_f:(sim.n_f+sim.n_g)]
"""

f = py"f"
g = py"g"
x0 = py"x0*regularizer"
searchSpace = py"bounds_reg"
dim = sim_jl.n_dof
op1 = optProblem(f, g, x0, searchSpace, dim, sim_jl, py"sim.n_dof", py"sim.n_f", py"sim.n_g")
op1.sim_jl.verbose = false;
save_sim() = py"""sim.interface.SaveFlowsheet(sim.flowsheet,$pwd()+"/examples/PRICO/PRICO_composite2.dwxmz",True)"""

# external functions and jacobian calculations
x0 = 1.0 * (op1.searchSpace[1, :] + op1.searchSpace[2, :]) ./ 2
# x0 = py"x0*regularizer"
iter = 0
function f_block(x)
    global iter
    iter += 1
    return [op1.f(x)[1]; op1.g(x)[:]]
end
first_time = true
h = 0.01
flag_diff_method = "2nd_central"
flag_FiniteDifference = false
function jac_block(x)
    # println("hi")
    global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp, first_time
    if norm(x0 - x) > 1e-6 || first_time == true
        f0 = f_block(x)
        println("Calculating jacobian of black-box functions at x=", x, " and f^{sim}=", f0)
        # flag_diff_method = "1st_forward"
        global h, flag_FiniteDifference, flag_diff_method
        tmp = zeros(6, length(x))
        if flag_FiniteDifference
            tmp = jacobian(central_fdm(2, 1), f_block, x)[1]
        else
            for i = 1:length(x)
                if flag_diff_method == "1st_forward"
                    tmp[:, i] = (f_block([x[1:i-1]; x[i] + h; x[i+1:end]]) - f0) / h
                elseif flag_diff_method == "2nd_central"
                    tmp[:, i] = (f_block([x[1:i-1]; x[i] + h; x[i+1:end]]) - f_block([x[1:i-1]; x[i] - h; x[i+1:end]])) / (2 * h)
                end
            end
        end
        println("∇f_block=", tmp)
        # tmp = jacobian(forward_fdm(2, 1), f_block, x)[1]
        ∇f_bkp = tmp[1, :]
        ∇g1_bkp = tmp[2, :]
        ∇g2_bkp = tmp[3, :]
        ∇g3_bkp = tmp[4, :]
        ∇g4_bkp = tmp[5, :]
        ∇g5_bkp = tmp[6, :]
        x0 = x
        first_time = false
        return tmp
    end
end

function setup_model()
    ## JuMP MODEL STUFF
    model = Model(optimizer_with_attributes(Ipopt.Optimizer))
    # model = Model(Ipopt.Optimizer)
    # model = Model(GAMS.Optimizer)
    # set_optimizer_attribute(model, GAMS.ModelType(), "NLP")
    # set_optimizer_attribute(model, "NLP", "CONOPT")
    # set_optimizer_attribute(model, "resLIM", 1200.0)
    # model = Model(() -> MadNLP.Optimizer(print_level = MadNLP.INFO, max_iter = 100))

    # VARIABLES
    i = 1
    @variable(model, op1.searchSpace[1, i]' <= x1 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 2
    @variable(model, op1.searchSpace[1, i]' <= x2 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 3
    @variable(model, op1.searchSpace[1, i]' <= x3 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 4
    @variable(model, op1.searchSpace[1, i]' <= x4 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 5
    @variable(model, op1.searchSpace[1, i]' <= x5 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)
    i = 6
    @variable(model, op1.searchSpace[1, i]' <= x6 <= op1.searchSpace[2, i]', start = (op1.searchSpace[1, i] + op1.searchSpace[2, i]) ./ 2)

    # OBJECTIVE FUNCTION
    fobj_bb(x1, x2, x3, x4, x5, x6) = op1.f([x1, x2, x3, x4, x5, x6])[1]
    function ∇fobj_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇fobj_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇f_tmp = ∇f_bkp
        # ∇f_tmp = grad(forward_fdm(2, 1), op1.f, [x1, x2, x3, x4, x5, x6, x7, x8])[1]
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇f_tmp[i]
        end
    end
    register(model, :fobj_bb, op1.dim, fobj_bb, ∇fobj_bb; autodiff = false)
    @NLobjective(model, Min, fobj_bb(x1, x2, x3, x4, x5, x6))

    # CONSTRAINTS
    g1(x) = op1.g(x)[1]
    g2(x) = op1.g(x)[2]
    g3(x) = op1.g(x)[3]
    g4(x) = op1.g(x)[4]
    g5(x) = op1.g(x)[5]
    g1_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[1]
    g2_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[2]
    g3_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[3]
    g4_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[4]
    g5_bb(x1, x2, x3, x4, x5, x6) = op1.g([x1, x2, x3, x4, x5, x6])[5]
    function ∇g1_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g1_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g1_tmp = ∇g1_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g1_tmp[i]
        end
    end
    function ∇g2_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g2_tmp = ∇g2_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g2_tmp[i]
        end
    end
    function ∇g3_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g3_tmp = ∇g3_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g3_tmp[i]
        end
    end
    function ∇g4_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g4_tmp = ∇g4_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g4_tmp[i]
        end
    end
    function ∇g5_bb(g::AbstractVector{T}, x1::T, x2::T, x3::T, x4::T, x5::T, x6::T) where {T}
        # println("I am in ∇g2_bb with x=", x1)
        global x0, ∇f_bkp, ∇g1_bkp, ∇g2_bkp, ∇g3_bkp, ∇g4_bkp, ∇g5_bkp
        jac_block([x1, x2, x3, x4, x5, x6])
        ∇g5_tmp = ∇g5_bkp
        for i = 1:length([x1, x2, x3, x4, x5, x6])
            g[i] = ∇g5_tmp[i]
        end
    end
    register(model, :g1_bb, op1.dim, g1_bb, ∇g1_bb; autodiff = false)
    @NLconstraint(model, g1_bb(x1, x2, x3, x4, x5, x6) <= 0)
    register(model, :g2_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g2_bb(x1, x2, x3, x4, x5, x6) == 0)
    register(model, :g3_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g3_bb(x1, x2, x3, x4, x5, x6) == 0)
    register(model, :g4_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g4_bb(x1, x2, x3, x4, x5, x6) == 0)
    register(model, :g5_bb, op1.dim, g2_bb, ∇g2_bb; autodiff = false)
    @NLconstraint(model, g5_bb(x1, x2, x3, x4, x5, x6) == 0)

    return model
end

model = setup_model()
function set_opt_params(model)
    set_optimizer_attribute(model, "tol", 1e-2)
    set_optimizer_attribute(model, "dual_inf_tol", 1e-2)
    set_optimizer_attribute(model, "constr_viol_tol", 1e-3)
    set_optimizer_attribute(model, "compl_inf_tol", 1e-2)
    set_optimizer_attribute(model, "acceptable_tol", 1e-2)
    set_optimizer_attribute(model, "acceptable_iter", 3)
    # set_optimizer_attribute(model, "mu_target", 1e-6)

    set_optimizer_attribute(model, "hessian_approximation", "limited-memory")

    set_optimizer_attribute(model, "max_iter", 50)
    set_optimizer_attribute(model, "max_cpu_time", 3600.0)
    set_optimizer_attribute(model, "print_level", 5)
    return model
end
model = set_opt_params(model)
JuMP.optimize!(model)

iter
println("Model used h=", h, " and tol=", get_optimizer_attribute(model, "tol"))
println("Initial value x0=", (op1.searchSpace[1, :] + op1.searchSpace[2, :]) ./ 2)
op1.searchSpace

@show JuMP.termination_status(model)
@show JuMP.primal_feasibility_report(model)
@show JuMP.dual_status(model)

# retrieve the objective value, corresponding x values and the status
println([value(model[:x1]), value(model[:x2]), value(model[:x3]), value(model[:x4]), value(model[:x5]), value(model[:x6])])
println(JuMP.objective_value(model))
println(JuMP.termination_status(model))

f_block(JuMP.value.([value(model[:x1]), value(model[:x2]), value(model[:x3]), value(model[:x4]), value(model[:x5]), value(model[:x6])]))


module PWAHelper

using amplitudes_compass

using DelimitedFiles
using Juno

export precalculate_compass_basis, read_precalc_basis
# export precalculate_compass_basis_txt, read_precalc_basis_txt
export get_npars
export get_parameter_map, make_pblock_inds
export extnd, extnd!, shrnk, cohsq, cohts, cohts!
export contract_to_intensity, get_intesity
export get_parameter_ranges, normalize_pars!
export calculate_integrals_from_precalc_basis

function precalculate_compass_basis(basis,fin,fout)
    io = open(fin,"r")
    Nd = read(io,Int32)
    variable_mat = Array{Float64}(undef,6,Nd)
    read!(io,variable_mat)
    #variable_mat = readdlm(fin); Nd = size(variable_mat,1)
    function_mat = fill(0.0im, Nd, length(basis))
    @progress for e in 1:Nd
        for (i,b) in enumerate(basis)
            function_mat[e,i] = b(@view(variable_mat[2:end,e])..., variable_mat[1,e]);
        end
    end
    io = open(fout,"w")
    write(io,trunc(Int64, Nd))
    write(io,trunc(Int64, length(basis)+length(basis)))
    write(io,[real(function_mat) imag(function_mat)])
    close(io)
end

function read_precalc_basis(fname)
    io = open(fname,"r")
    r=read(io,Int64)
    c=read(io,Int64)
    ld = Array{Float64}(undef,r,c)
    read!(io,ld)
    Nh = div(size(ld,2),2)
    return ld[:,1:Nh]+1im*ld[:,(Nh+1):end]
end

function get_parameter_map(block_inds, Nw)
    temp = []; numb = []
    for (i,bl) in enumerate(block_inds)
        # push false for the first parameter, true for others
        push!(temp,false,[true for i=1:(length(bl)-1)]...);
        push!(numb,bl...)
    end
    Tmap = fill(0,2,sum(temp .+ 1))
    count=1
    for (i,b) in enumerate(temp)
        Tmap[1,count] = numb[i]
        count+=1
        if b
            Tmap[2,count] = numb[i];
            count+=1
        end
    end
    return Tmap
end

function get_npars(block_inds)
    sum(2*length(bl)-1 for bl in block_inds)
end

function make_pblock_inds(block_inds)
    Nps = [2*length(bl)-1 for bl in block_inds]
    counter = 1
    pbls = [let v = collect(counter:(counter+Np-1))
                counter += Np
                v
            end for Np in Nps]
    return pbls
end

extnd(Ψ, tmap)  = [((tmap[2,i]==0) ? Ψ[tmap[1,i]] : 1im*Ψ[tmap[2,i]]) for i in 1:size(tmap,2)]
function extnd!(X, Ψ, tmap)
    for i in 1:size(tmap,2)
        @inbounds X[i] = (tmap[2,i]==0) ? Ψ[tmap[1,i]] : 1im*Ψ[tmap[2,i]]
    end
end
function shrnk(p, tmap)
    Nw = tmap[2,end] # work around
    outv = fill(0.0im,Nw)
    for i in 1:size(tmap,2)
        (tmap[1,i] != 0) && (outv[tmap[1,i]] += p[i]);
        (tmap[2,i] != 0) && (outv[tmap[2,i]] += 1.0im*p[i]);
    end
    outv
end

function cohsq(X, pblocks)
    sum(abs2, (sum(@view(X[bl]))) for bl in pblocks)
end
function cohts(X, pblocks)
    Np = length(X)
    vblocks = [let v = fill(0.0,Np)
        v[bl] .= 1.0
        v
    end for bl in pblocks]
    sum(sum(X[bl])*pl for (pl,bl) in zip(vblocks, pblocks))
end
function cohts!(X, pblocks)
    # Np = length(X)
    # vblocks = [
    #     let v = fill(0.0,Np)
    #         v[bl] .= 1.0
    #         v
    #     end for bl in pblocks]
    # sum(sum(X[bl])*pl for (pl,bl) in zip(vblocks, pblocks))
    for bl in pblocks
        @inbounds X[bl] .= sum(@view(X[bl]))
    end
end

function contract_to_intensity(ΨΨstar, block_inds)
    Tmap = get_parameter_map(block_inds, size(ΨΨstar,1))
    Np = size(Tmap,2)
    pblocks = make_pblock_inds(block_inds)
    COH = fill(0.0, Np, Np)
    for bl in pblocks
        COH[bl,bl] .= 1.0
    end
    get(v,w) = (v==0 || w==0) ? 0.0im : ΨΨstar[v,w]
    [COH[i,j]*( # meaning is ([1]-i[2]) * ([1]+i[2])
            get(Tmap[1,i],Tmap[1,j]) +
            get(Tmap[2,i],Tmap[2,j]) -
        1im*get(Tmap[2,i],Tmap[1,j]) +
        1im*get(Tmap[1,i],Tmap[2,j])
            ) for j in 1:Np, i in 1:Np]
end

function get_intesity(pars, form, block_inds)
    pblocks = make_pblock_inds(block_inds)
    BM = real.(contract_to_intensity(form,block_inds))
    Np = length(pars)
    sum(pars[i]*BM[i,j]*pars[j] for i=1:Np, j=1:Np)
end

function get_parameter_ranges(form, block_inds)
    Np = get_npars(block_inds)
    [let p = fill(0.0, Np)
        p[i] = 1.0
        1/sqrt(get_intesity(p, form, block_inds))
     end for i in 1:Np]
end

function normalize_pars!(pars, form, block_inds)
    Intens = get_intesity(pars, form, block_inds)
    pars ./=  sqrt(Intens)
    pars
end

end


#=
    many of the types are
    Used from https://github.com/lanl-ansi/Juniper.jl
=#

mutable struct RegisteredFunction
    s         :: Symbol
    dimension :: Integer
    f         :: Function
    gradf     :: Union{Nothing, Function}
    grad2f    :: Union{Nothing, Function}
    autodiff  :: Bool
end

# Options for the solver (more details like defaults in solver.jl)
mutable struct SolverOptions
    nl_solver                           :: Any# needs to be set
    log_levels                          :: Vector{Symbol}
    silent                              :: Bool
    atol                                :: Float64
    incumbent_constr                    :: Bool
    obj_epsilon                         :: Float64

    time_limit                          :: Float64
    mip_gap                             :: Float64

    best_obj_stop                       :: Float64
    solution_limit                      :: Int64
    all_solutions                       :: Bool

    list_of_solutions                   :: Bool
    mip_solver                          :: Any
    allow_almost_solved                 :: Bool
    allow_almost_solved_integral        :: Bool
    registered_functions                :: Union{Nothing, Vector{RegisteredFunction}}
    processors                          :: Int64
end

mutable struct SolutionObj
    solution    :: Vector{Float64}
    objval      :: Float64
end

#OriginalProblem
mutable struct OriginalProblem
    nl_solver           :: Any
    nl_solver_options   :: Vector{Pair}
    model               :: JuMP.Model
    relaxation_status   :: MOI.TerminationStatusCode
    relaxation_objval   :: Float64
    relaxation_solution :: Vector{Float64}
    status              :: MOI.TerminationStatusCode
    objval              :: Float64
    best_bound          :: Float64
    x                   :: Vector{JuMP.VariableRef}
    primal_start        :: Vector{Real}
    num_constr          :: Int64
    num_nl_constr       :: Int64
    num_q_constr        :: Int64
    num_l_constr        :: Int64
    num_var             :: Int64
    l_var               :: Vector{Float64}
    u_var               :: Vector{Float64}
    has_nl_objective    :: Bool
    nlp_evaluator       :: MOI.AbstractNLPEvaluator
    objective           :: Union{SVF, SAF, SQF, Nothing}
    disc2var_idx        :: Vector{Int64}
    var2disc_idx        :: Vector{Int64}
    var_type            :: Vector{Symbol}
    obj_sense           :: Symbol
    num_disc_var        :: Int64
    solution            :: Vector{Float64}
    soltime             :: Float64
    solutions           :: Vector{SolutionObj}
    nsolutions          :: Int64
    mip_solver          :: Any
    mip_solver_options  :: Vector{Pair}
    relaxation_time     :: Float64
    start_time          :: Float64
    # Info
    nintvars            :: Int64
    nbinvars            :: Int64
    options             :: SolverOptions
    nnodes              :: Int64

    nlp_constr_type     :: Vector{Symbol}

    OriginalProblem() = new()
end

#OriginalProblem
mutable struct OAdata
    mip_model                 :: JuMP.Model
    nlp_model                 :: JuMP.Model
    ref_nlp_model             :: JuMP.Model
    ref_mip_model             :: JuMP.Model
    ref_feasibility_model     :: JuMP.Model

    mip_x                     :: Vector{JuMP.VariableRef}
    nlp_x                     :: Vector{JuMP.VariableRef}
    ref_nlp_x                 :: Vector{JuMP.VariableRef}
    ref_mip_x                 :: Vector{JuMP.VariableRef}
    ref_feasibility_x         :: Vector{JuMP.VariableRef}

    num_l_constr              :: Int

    ref_l_var                 :: Vector{Float64}
    ref_u_var                 :: Vector{Float64}
    ref_var_type              :: Vector{Symbol}
    ref_num_var               :: Int
    ref_num_nl_constr         :: Int
    ref_nlp_constr_type       :: Vector{Symbol}

    obj_sense                 :: Symbol

    oa_status                 :: Symbol
    oa_started                :: Bool
    incumbent                 :: Vector{Float64}
    new_incumbent             :: Bool
    total_time                :: Float64
    obj_val                   :: Float64
    obj_bound                 :: Float64
    obj_gap                   :: Float64
    oa_iter                   :: Int64

    milp_sol_available        :: Bool
    ref_mip_solution          :: Vector{Float64}
    mip_infeasible            :: Bool
    ref_nlp_solution          :: Vector{Float64}
    nlp_infeasible            :: Bool
    ref_feasibility_solution  :: Vector{Float64}

    int_idx                   :: Vector{Int64}
    prev_ref_mip_solution     :: Vector{Float64}

    ref_linear_le_constraints ::Vector{Tuple{SAF, MOI.LessThan{Float64}}}
    ref_linear_ge_constraints ::Vector{Tuple{SAF, MOI.GreaterThan{Float64}}}
    ref_linear_eq_constraints ::Vector{Tuple{SAF, MOI.EqualTo{Float64}}}

    OAdata() = new()
end

"""
    is_type_correct(x, var_type, atol)
Check whether a variable x has the correct type
"""
function is_type_correct(x, var_type, atol)
    if var_type != :Cont
        if !isapprox(round(x)-x, 0; atol=atol)
            return false
        end
    end
    return true
end

"""
    are_type_correct(sol, types, disc2var_idx, atol)
Check whether all variables have the correct type
"""
function are_type_correct(sol, types, disc2var_idx, atol)
    for i in disc2var_idx
        if !isapprox(round(sol[i])-sol[i], 0; atol=atol)
            return false
        end
    end
    return true
end

"""
    all_reasonable_type_correct(sol, disc2var_idx, reasonable_idx, atol)
Check whether all reasonable variable have are discrete alreadyy
"""
function all_reasonable_type_correct(sol, disc2var_idx, reasonable_idx, atol)
    for i in reasonable_idx
        idx = disc2var_idx[i]
        if !isapprox(round(sol[idx])-sol[idx], 0; atol=atol)
            return false
        end
    end
    return true
end


module CommonFunctions

    export nbasis, nnodes, nodes, order, degree

    nbasis() = nothing
    nnodes() = nothing
    nodes()  = nothing
    order()  = nothing
    degree() = nothing

    export evaluate, evaluate!

    evaluate() = nothing
    evaluate!() = nothing

    export derivative, integral

    derivative() = nothing
    integral() = nothing

    export write_to_hdf5

    write_to_hdf5() = nothing

    export get_solution, get_solution!, set_solution!, reset!, cut_periodic_solution!

    get_solution() = nothing
    get_solution!() = nothing
    set_solution!() = nothing
    reset!() = nothing
    cut_periodic_solution!() = nothing

end


# This file is a part of Julia. License is MIT: https://julialang.org/license

module Random

using Base.dSFMT
using Base.GMP: Limb, MPZ
import Base: copymutable, copy, copy!, ==, hash

export srand,
       rand, rand!,
       randn, randn!,
       randexp, randexp!,
       bitrand,
       randstring,
       randsubseq, randsubseq!,
       shuffle, shuffle!,
       randperm, randperm!,
       randcycle, randcycle!,
       AbstractRNG, MersenneTwister, RandomDevice,
       GLOBAL_RNG, randjump

abstract type AbstractRNG end

abstract type FloatInterval{T<:AbstractFloat} end

struct CloseOpen{  T<:AbstractFloat} <: FloatInterval{T} end # interval [0,1)
struct Close1Open2{T<:AbstractFloat} <: FloatInterval{T} end # interval [1,2)

const FloatInterval_64 = FloatInterval{Float64}
const CloseOpen_64     = CloseOpen{Float64}
const Close1Open2_64   = Close1Open2{Float64}

CloseOpen(  ::Type{T}=Float64) where {T<:AbstractFloat} = CloseOpen{T}()
Close1Open2(::Type{T}=Float64) where {T<:AbstractFloat} = Close1Open2{T}()

const BitFloatType = Union{Type{Float16},Type{Float32},Type{Float64}}

function __init__()
    try
        srand()
    catch ex
        Base.showerror_nostdio(ex,
            "WARNING: Error during initialization of module Random")
    end
end

include("RNGs.jl")
include("generation.jl")
include("normal.jl")
include("misc.jl")

## rand & rand! & srand docstrings

"""
    rand([rng=GLOBAL_RNG], [S], [dims...])

Pick a random element or array of random elements from the set of values specified by `S`;
`S` can be

* an indexable collection (for example `1:n` or `['x','y','z']`),
* an `Associative` or `AbstractSet` object,
* a string (considered as a collection of characters), or
* a type: the set of values to pick from is then equivalent to `typemin(S):typemax(S)` for
  integers (this is not applicable to [`BigInt`](@ref)), and to ``[0, 1)`` for floating
  point numbers;

`S` defaults to [`Float64`](@ref).

# Examples
```julia-repl
julia> rand(Int, 2)
2-element Array{Int64,1}:
 1339893410598768192
 1575814717733606317

julia> rand(MersenneTwister(0), Dict(1=>2, 3=>4))
1=>2
```

!!! note
    The complexity of `rand(rng, s::Union{Associative,AbstractSet})`
    is linear in the length of `s`, unless an optimized method with
    constant complexity is available, which is the case for `Dict`,
    `Set` and `IntSet`. For more than a few calls, use `rand(rng,
    collect(s))` instead, or either `rand(rng, Dict(s))` or `rand(rng,
    Set(s))` as appropriate.
"""
rand

"""
    rand!([rng=GLOBAL_RNG], A, [S=eltype(A)])

Populate the array `A` with random values. If `S` is specified
(`S` can be a type or a collection, cf. [`rand`](@ref) for details),
the values are picked randomly from `S`.
This is equivalent to `copy!(A, rand(rng, S, size(A)))`
but without allocating a new array.

# Examples
```jldoctest
julia> rng = MersenneTwister(1234);

julia> rand!(rng, zeros(5))
5-element Array{Float64,1}:
 0.590845
 0.766797
 0.566237
 0.460085
 0.794026
```
"""
rand!

"""
    srand([rng=GLOBAL_RNG], seed) -> rng
    srand([rng=GLOBAL_RNG]) -> rng

Reseed the random number generator. If a `seed` is provided, the RNG will give a
reproducible sequence of numbers, otherwise Julia will get entropy from the system. For
`MersenneTwister`, the `seed` may be a non-negative integer or a vector of [`UInt32`](@ref)
integers. `RandomDevice` does not support seeding.

# Examples
```jldoctest
julia> srand(1234);

julia> x1 = rand(2)
2-element Array{Float64,1}:
 0.590845
 0.766797

julia> srand(1234);

julia> x2 = rand(2)
2-element Array{Float64,1}:
 0.590845
 0.766797

julia> x1 == x2
true
```
"""
srand

end # module


module FlyDi

export c, q, u, m_e, h, a0, Ry, kB,
    FastAdjust,
    h5read_pa, potential,
    grid_r, electrode_g, pa_g, potential_g, field_g, amp_field_g, grad_field_g,
    kinetic_energy, potential_energy,
    euler, leapfrog, rk4,
    traj, traj_last,
    flydi, flydi_last,
    mc_supersonic, mc_uniform

include("external.jl")
include("constants.jl")
include("fastadjust.jl")
include("traj.jl")
include("fly.jl")
include("montecarlo.jl")

end


#include("ConeMantle.jl")
#include("EllipsoidMantle.jl")
#include("EllipticalSurface.jl")
#include("Plane.jl")
#include("Polygon.jl")
#include("TorusMantle.jl")

@recipe function f(s::AbstractSurfacePrimitive{T}; n_arc = 40, n_vert_lines = 2, n_samples = 40) where {T}
    seriestype --> :csg
    l = get_label_name(s)
    if haskey(plotattributes, :seriestype)
        if plotattributes[:seriestype] == :csg
            @series begin
                label := ""
                linewidth := 0
                linecolor --> :white
                mesh(s, n_arc)
            end
            @series begin
                label --> l
                linecolor --> :black
                fillalpha := 1
                linewidth --> 2
                mesh(s, n_arc, n_vert_lines)
            end
        elseif plotattributes[:seriestype] == :mesh3d
            label --> l
            linecolor --> :white
            mesh(s, n_arc)
        elseif plotattributes[:seriestype] == :wireframe
            label --> l
            seriescolor --> :black
            fillalpha := 1
            linewidth --> 2
            mesh(s, n_arc, n_vert_lines)
        elseif plotattributes[:seriestype] == :samplesurface
            label --> l
            seriesalpha --> 0.4
            sample(s, extremum(s)/n_samples)
        else
            @warn "The only seriestypes which will return a plot are :csg, :wireframe, :mesh3d, and :samplesurface"
        end
    end
end

@recipe function f(::Type{Val{:samplesurface}}, x, y, z)
    seriescolor --> 1
    seriesalpha --> 0.2
    markerstrokewidth --> 0
    markersize --> 4
    seriestype := :scatter
    ()
end

@recipe function f(vp::AbstractVector{<:AbstractSurfacePrimitive})
    linecolor --> :black
    @series begin
        label --> "Faces"
        show_normal --> false
        vp[1]
    end
    if length(vp) > 1
        for p in vp[2:end]
            @series begin
                show_normal --> false
                label := nothing
                p
            end
        end
    end
end


module CLagrangianFESpaces

using Gridap
using Gridap.Helpers
using Gridap.CellValuesGallery
using Gridap.DOFBases: _length
using Gridap.ConformingFESpaces: _CellField
using Gridap.ConformingFESpaces: CellEqClass
using Gridap.BoundaryGrids: _setup_tags

import Gridap: num_free_dofs
import Gridap: num_diri_dofs
import Gridap: diri_tags
import Gridap: apply_constraints
import Gridap: apply_constraints_rows
import Gridap: apply_constraints_cols
import Gridap: celldofids
import Gridap: interpolate_values
import Gridap: interpolate_diri_values
import Gridap: CellField
import Gridap: CellBasis
import Gridap: Triangulation

export CLagrangianFESpace

struct CLagrangianFESpace{D,Z,T,S,B} <: FESpace{D,Z,T}
  grid::Grid{D,Z}
  node_to_label::Vector{Int}
  tag_to_labels::Vector{Vector{Int}}
  diritags::Vector{Int}
  dirimasks::Vector{B}
  nfreedofs::Int
  ndiridofs::Int
  node_and_comp_to_dof::Vector{S}
  cell_to_dofs::CellVector{Int}
  reffe::LagrangianRefFE{Z,T}
  cellbasis::CellBasis{Z,T}
  dirinodes::Vector{Int}
end

function CLagrangianFESpace(
  ::Type{T},
  grid::Grid{D,Z},
  node_to_label::AbstractVector{<:Integer},
  tag_to_labels::Vector{Vector{Int}},
  diritags::Vector{Int},
  dirimasks::Vector{B}) where {D,Z,T,B}

  node_and_comp_to_dof, nfreedofs, ndiridofs, dirinodes =
    _setup_node_and_comp_to_dof(
    T, node_to_label, tag_to_labels, diritags, dirimasks)

  cell_to_nodes = cells(grid)

  reffe = _setup_reffe(T,celltypes(grid),cellorders(grid))
  lnode_and_comp_to_ldof = reffe.dofbasis.node_and_comp_to_dof

  cell_to_dofs = _setup_cell_to_dofs(
    cell_to_nodes,
    node_and_comp_to_dof,
    lnode_and_comp_to_ldof)

  cellbasis = _setup_cellbasis(reffe,grid)

  S = _S(T)

  CLagrangianFESpace{D,Z,T,S,B}(
    grid,
    node_to_label,
    tag_to_labels,
    diritags,
    dirimasks,
    nfreedofs,
    ndiridofs,
    node_and_comp_to_dof,
    cell_to_dofs,
    reffe,
    cellbasis,
    dirinodes)

end

function CLagrangianFESpace(
  ::Type{T},
  model::DiscreteModel,
  order::Integer,
  diritags=Int[],
  dirimasks=nothing) where T
  CLagrangianFESpace(T,model,FaceLabels(model),order,diritags,dirimasks)
end

function CLagrangianFESpace(
  ::Type{T},
  model::DiscreteModel,
  facelabels::FaceLabels,
  order::Integer,
  diritags=Int[],
  dirimasks=nothing) where T

  _diri_tags = _setup_tags(facelabels,diritags)
  _diri_masks = _setup_masks(T,_diri_tags,dirimasks)

  grid, node_to_label, tag_to_labels = _setup_grid(model,facelabels,order)

  CLagrangianFESpace(
    T,grid,node_to_label,tag_to_labels,_diri_tags,_diri_masks)

end

num_free_dofs(fesp::CLagrangianFESpace) = fesp.nfreedofs

num_diri_dofs(fesp::CLagrangianFESpace) = fesp.ndiridofs

diri_tags(fesp::CLagrangianFESpace) = fesp.diritags

apply_constraints(
  ::CLagrangianFESpace, cellvec::CellVector, cellids::CellNumber) = cellvec

apply_constraints_rows(
  ::CLagrangianFESpace, cellmat::CellMatrix, cellids::CellNumber) = cellmat

apply_constraints_cols(
  ::CLagrangianFESpace, cellmat::CellMatrix, cellids::CellNumber) = cellmat

celldofids(fesp::CLagrangianFESpace) = fesp.cell_to_dofs

function interpolate_values(fesp::CLagrangianFESpace,fun::Function)
  _interpolate_values(fesp,fun)
end

function interpolate_values(fesp::CLagrangianFESpace{D,Z,T},val::T) where {D,Z,T}
  _interpolate_values(fesp,val)
end

function _interpolate_values(fesp,fun)
  zh = zero(fesp)
  fdof_to_val = free_dofs(zh)
  ddof_to_val = diri_dofs(zh)
  _fill_interpolated_vals!(fdof_to_val,ddof_to_val,fesp,fun)
  (fdof_to_val, ddof_to_val)
end

function interpolate_diri_values(fesp::CLagrangianFESpace, funs::Vector{<:Function})
  _interpolate_diri_values(fesp, funs)
end

function interpolate_diri_values(fesp::CLagrangianFESpace{D,Z,T},vals::Vector{T}) where {D,Z,T}
  _interpolate_diri_values(fesp, vals)
end

function _interpolate_diri_values(fesp, funs)

  T = value_type(fesp)
  E = eltype(T)
  ndiri = num_diri_dofs(fesp)
  ddof_to_val = zeros(E,ndiri)
  node_to_coords = points(fesp.grid)
  node_and_comp_to_dof = fesp.node_and_comp_to_dof
  node_to_label = fesp.node_to_label
  tag_to_labels = fesp.tag_to_labels
  dirinodes = fesp.dirinodes

  for (fun,diritag,dirimask) in zip(funs,fesp.diritags,fesp.dirimasks)
    _fill_diri_values!(
      ddof_to_val,
      fun,
      diritag,
      dirimask,
      dirinodes,
      node_to_label,
      tag_to_labels,
      node_to_coords,
      node_and_comp_to_dof)
  end

  ddof_to_val
end

function CellField(
  fesp::CLagrangianFESpace{D,Z,T},
  free_dofs::AbstractVector{E},
  diri_dofs::AbstractVector{E}) where {D,Z,T,E}

  _CellField(fesp, free_dofs, diri_dofs, T, E)
end

CellBasis(fesp::CLagrangianFESpace{D,Z,T}) where {D,Z,T} = fesp.cellbasis

Triangulation(fesp::CLagrangianFESpace) = Triangulation(fesp.grid)

# Helpers

function _setup_masks(T,diritags,dirimasks)
  if dirimasks === nothing
    ncomps = _length(T)
    ntags = length(diritags)
    t = tuple(fill(true,ncomps)...)
    return fill(t,ntags)
  else
    @assert length(diritags) == length(dirimasks)
    return dirimasks
  end
end

function _setup_node_and_comp_to_dof(
  ::Type{T}, node_to_label, tag_to_labels, diritags, dirimasks) where T

  S = _S(T)
  nnodes = length(node_to_label)
  node_and_comp_to_dof = zeros(S,nnodes)
  dirinodes = Int[]

  nf, nd = _fill_node_and_comp_to_dof!(
    node_and_comp_to_dof, dirinodes, node_to_label, tag_to_labels, diritags, dirimasks)

  (node_and_comp_to_dof, nf, nd, dirinodes)

end

function _fill_node_and_comp_to_dof!(
  node_and_comp_to_dof::Vector{S},
  dirinodes,
  node_to_label,
  tag_to_labels,
  diritags,
  dirimasks) where S

  nnodes = length(node_to_label)
  fdof = 1
  ddof = -1

  for node in 1:nnodes
    label = node_to_label[node]

    comp_to_dof, fdof, ddof, isdirinode =
      _compute_comp_to_dof(S,label,tag_to_labels,diritags,dirimasks,fdof,ddof)

    if isdirinode
      push!(dirinodes,node)
    end

    node_and_comp_to_dof[node] = comp_to_dof

  end

  (fdof-1,-(ddof+1))

end

function _compute_comp_to_dof(
  ::Type{S}, label, tag_to_labels, diritags, dirimasks, fdof, ddof) where S <: MultiValue

  comp_to_dof = zero(mutable(S))
  comps = eachindex(comp_to_dof)

  isdirinode = false
  for comp in comps
    dof, fdof, ddof = _compute_dof(label,tag_to_labels,diritags,dirimasks,comp,fdof,ddof)
    comp_to_dof[comp] = dof
    if dof < 0
      isdirinode = true
    end
  end

  (comp_to_dof, fdof, ddof, isdirinode)

end

function _compute_comp_to_dof(
  ::Type{S}, label, tag_to_labels, diritags, dirimasks, fdof, ddof) where S <: Real

  comp = 1
  dof, fdof, ddof =  _compute_dof(
    label,tag_to_labels,diritags,dirimasks,comp,fdof,ddof)

  isdirinode = dof < 0

  (dof, fdof, ddof, isdirinode)

end

function _compute_dof(label,tag_to_labels,diritags,dirimasks,comp,fdof,ddof)

  isfree = true
  for (i,diritag) in enumerate(diritags)

    if label in tag_to_labels[diritag]
      comp_to_mask = dirimasks[i]
      mask = comp_to_mask[comp]
      if mask
        dof = ddof
        ddof -= 1
        isfree = false
        break
      end
    end
  end

  if isfree
    dof = fdof
    fdof += 1
  end

  (dof,fdof,ddof)

end

_S(::Type{<:Real}) = Int

_S(::Type{MultiValue{S,T,N,L}}) where {S,T,N,L} = MultiValue{S,Int,N,L}

function _setup_reffe(ct,co)
  @notimplemented
end

function _setup_reffe(
  ::Type{T},
  ct::ConstantCellValue{<:NTuple{Z}},
  co::ConstantCellValue) where {T,Z}
  t = ct.value
  o = co.value
  p = Polytope(t)
  LagrangianRefFE(T,p,fill(o,Z))
end

function _setup_cellbasis(reffe,grid)
  lshapes = shfbasis(reffe)
  nc = ncells(grid)
  shb = ConstantCellValue(lshapes, nc)
  trian = Triangulation(grid)
  phi = CellGeomap(trian)
  basis = attachgeomap(shb,phi)
  basis
end

function _setup_cell_to_dofs(
  cell_to_nodes,
  node_and_comp_to_dof::Vector{S},
  lnode_and_comp_to_ldof) where S

  ncomp = _length(S)
  ncells = length(cell_to_nodes)
  nlnodes, ncomp = size(lnode_and_comp_to_ldof)
  nldofs = nlnodes * ncomp

  ndata = ncells*nldofs
  I = eltype(S)

  cell_to_dofs_data = zeros(I,ndata)

  _fill_cell_to_dofs!(
    cell_to_dofs_data,
    cell_to_nodes,
    node_and_comp_to_dof,
    lnode_and_comp_to_ldof)

  CellVectorFromDataAndStride(cell_to_dofs_data,nldofs)

end

function  _fill_cell_to_dofs!(
  cell_to_dofs_data,
  cell_to_nodes,
  node_and_comp_to_dof,
  lnode_and_comp_to_ldof)

  k = 0

  nlnodes, ncomps = size(lnode_and_comp_to_ldof)
  nldofs = nlnodes * ncomps

  for nodes in cell_to_nodes

    for (lnode,node) in enumerate(nodes)
      for comp in 1:ncomps
        ldof = lnode_and_comp_to_ldof[lnode,comp]
        dof = node_and_comp_to_dof[node][comp]
        cell_to_dofs_data[k+ldof] = dof
      end
    end

    k += nldofs

  end

end

function _fill_interpolated_vals!(fdof_to_val,ddof_to_val,fesp,fun)

  x = points(fesp.grid)
  fx = _eval_bc(fun,x)
  T = value_type(fesp)
  @assert eltype(fx) == T
  z = zero(T)
  comps = eachindex(z)

  _fill_interpolated_vals_kernel!(
    fdof_to_val,
    ddof_to_val,
    fesp.node_and_comp_to_dof,
    fx,
    comps)

end

_eval_bc(fun::Function,x) = fun.(x)

_eval_bc(val,x) = fill(val,length(x))

function  _fill_interpolated_vals_kernel!(
  fdof_to_val,
  ddof_to_val,
  node_and_comp_to_dof,
  node_and_comp_to_val,
  comps)

  nnodes = length(node_and_comp_to_dof)
  @assert nnodes == length(node_and_comp_to_val)

  for node in 1:nnodes
    comp_to_dof = node_and_comp_to_dof[node]
    comp_to_val = node_and_comp_to_val[node]
    for comp in comps
      dof = comp_to_dof[comp]
      val = comp_to_val[comp]
      if dof > 0
        fdof_to_val[dof] = val
      elseif dof < 0
        ddof_to_val[-dof] = val
      else
        @unreachable
      end
    end
  end

end

function _fill_diri_values!(
  ddof_to_val,
  fun,
  diritag,
  dirimask,
  dirinodes,
  node_to_label,
  tag_to_labels,
  node_to_coords,
  node_and_comp_to_dof)

  for node in dirinodes

    label = node_to_label[node]
    if label in tag_to_labels[diritag]
      x = node_to_coords[node]
      fx = _eval(fun,x)
      comp_to_dof = node_and_comp_to_dof[node]
      for comp in eachindex(fx)
        mask = dirimask[comp]
        if mask
          val = fx[comp]
          dof = comp_to_dof[comp]
          ddof = -dof
          ddof_to_val[ddof] = val
        end
      end
    end

  end

end

@inline _eval(fun::Function,x) = fun(x)

@inline _eval(val,x) = val

function _same_order(order,co)
  @notimplemented
end

function _same_order(order,co::ConstantCellValue)
  order == co.value
end

function _setup_grid(model,facelabels,order)

  _grid = Grid(model)

  if _same_order(order,cellorders(_grid))
    @notimplementedif order !=1
    grid = _grid
    node_dim = 0
    node_to_label = labels_on_dim(facelabels,node_dim)
  else
    grid, node_to_label = grid_from_model_and_order(model,order)
  end

  tag_to_labels = facelabels.tag_to_labels

  (grid, node_to_label, tag_to_labels)

end

function grid_from_model_and_order(
  model::DiscreteModel{D}, order::Integer) where D

  node_to_coords, dim_to_nface_to_nodes, node_to_label =
    _setup_nodes_to_coords(model,D,order)

  graph = GridGraph(model)

  dim_to_cell_to_nface = [connections(graph,D,i) for i in 0:D]

  offset = length.(dim_to_nface_to_nodes)
  for i in 2:length(offset)
    offset[i] += offset[i-1]
  end
  offset = tuple(offset[1:(end-1)]...)

  cell_to_faces = local_append(offset, dim_to_cell_to_nface...)

  face_to_nodes = append(dim_to_nface_to_nodes...)

  grid = Grid(model,D)
  reffe = _setup_reffe(order,celltypes(grid))

  cell_to_nodes = CellEqClass(cell_to_faces, face_to_nodes, reffe)

  cell_to_nodes_data, cell_to_nodes_ptrs = compress(cell_to_nodes)

  cell_to_extrusion = celltypes(grid)

  n = ncells(grid)
  cell_to_order = ConstantCellValue(order,n)

  grid = UnstructuredGrid(
    node_to_coords,
    cell_to_nodes_data,
    cell_to_nodes_ptrs,
    cell_to_extrusion,
    cell_to_order)

  (grid, node_to_label)

end

function _setup_nodes_to_coords(model,D,order)

  node_to_coords = Point{D,Float64}[]
  node_to_label = Int[]

  facelabels = FaceLabels(model)

  vertex_to_label = labels_on_dim(facelabels,0)

  grid = Grid(model,0)
  _push_nodes!(node_to_coords,points(grid))
  _push_nodes!(node_to_label,vertex_to_label)
  node_to_nodes = cells(grid)

  dim_to_nface_to_nodes = []
  push!(dim_to_nface_to_nodes,node_to_nodes)

  for d in 1:D

    grid = Grid(model,d)
    trian = Triangulation(grid)
    reffe = _setup_reffe(order,celltypes(grid))

    ilnodes = reffe.nfacenodes[end]
    linode_to_qcoords = reffe.dofbasis.nodes[ilnodes]
    n = ncells(grid)
    q = ConstantCellValue(linode_to_qcoords,n)
    phi = CellGeomap(trian)
    nface_ilnode_to_coords = evaluate(phi,q)

    nface_to_nodes_data = Int[]
    nface_to_nodes_ptrs = zeros(Int,n+1)

    nface_to_label = labels_on_dim(facelabels,d)

    _push_nface_nodes!(
      node_to_coords,
      node_to_label,
      nface_to_nodes_data,
      nface_to_nodes_ptrs,
      nface_ilnode_to_coords,
      nface_to_label)

    nface_to_nodes = CellVectorFromDataAndPtrs(
      nface_to_nodes_data, nface_to_nodes_ptrs)

    push!(dim_to_nface_to_nodes,nface_to_nodes)

  end

  (node_to_coords, dim_to_nface_to_nodes, node_to_label)

end

function _setup_reffe(order,ct::ConstantCellValue)
  e = ct.value
  p = Polytope(e)
  D = length(e)
  LagrangianRefFE(Float64,p,fill(order,D))
end

function _push_nodes!(node_to_coords,coords)
  for x in coords
    push!(node_to_coords,x)
  end
end

function _push_nface_nodes!(
  node_to_coords,
  node_to_label,
  nface_to_nodes_data,
  nface_to_nodes_ptrs,
  nface_ilnode_to_coords,
  nface_to_label)

  node = length(node_to_coords)+1
  for (nface,ilnode_to_coords) in enumerate(nface_ilnode_to_coords)

    label = nface_to_label[nface]
    nface_to_nodes_ptrs[nface+1] = length(ilnode_to_coords)

    for coords in ilnode_to_coords
      push!(node_to_coords,coords)
      push!(node_to_label,label)
      push!(nface_to_nodes_data,node)
      node += 1
    end
  end

  length_to_ptrs!(nface_to_nodes_ptrs)

end

end # module


__precompile__()
module Comp

import DataStructures: counter
using SeqHax.Utils
using ArgParse
using Bio.Seq

function add_args(argparser)
    add_common_args(argparser["comp"])
    @add_arg_table argparser["comp"] begin
        "input"
            help = "Input FASTQ read file(s) (may be gzipped)"
            required = true
            nargs = '+'
            arg_type = String
    end
end

function main(args)
    comp = []

    for readfile in args["input"]
        npair = foreach_read(readfile, quiet=args["quiet"]) do idx, read
            while length(comp) < length(read.seq)
                push!(comp, counter(DNANucleotide))
            end
            for (i, base) in enumerate(read.seq)
                if isambiguous(base)
                    push!(comp[i], DNA_N)
                else
                    push!(comp[i], base)
                end
            end
        end
        println(STDERR, "Finished '$readfile'")
    end
    println("cycle\tA\tC\tG\tT\tN")
    for (i, ctr) in enumerate(comp)
        A, C, G, T, N = [ctr[b] for b in [DNA_A, DNA_C, DNA_G, DNA_T, DNA_N]]
        println("$i\t$A\t$C\t$G\t$T\t$N")
    end
end

end # module Comp


using SRKGenerator
#Standard Driver Script

if isempty(ARGS)

  ## Called from test. Use a default setup
  imin = -12
  jmax = 1

  ## Parallelize correctly
  host = gethostname()
  if host == "ChrisRackBed"
    #numCores = 4
    #addprocs(numCores-length(workers()))
    ptx_str = joinpath(Pkg.dir("SRKGenerator"),"deps","integrationWin.ptx")
    cudaCores=[1664,640]
  end

  if host == "crackauc2"
    #numCores = 16
    #addprocs(numCores-length(workers()))
    ptx_str = joinpath(Pkg.dir("SRKGenerator"),"deps","integrationLinux.ptx")
    cudaCores = [2816]
  end

  if host == "ChrisRackTV"
    #numCores = 4
    #addprocs(numCores-length(workers()))
    ptx_str = joinpath(Pkg.dir("SRKGenerator"),"deps","integrationWin.ptx")
    cudaCores=[2560]
  end

else

  ## Called from commandline. Assume Linux.
  ptx_str = joinpath(Pkg.dir("SRKGenerator"),"deps","integrationLinux.ptx")
  ## Arg Form:
  cudaCores = [parse(Int,ARGS[1])]
  imin = parse(Int,ARGS[2])
  jmax = parse(Int,ARGS[3]); jmin = -jmax

end

resString = srk_optimize(:LN_AUGLAG_EQ,1/100,0,imin,1,-jmax,jmax,
                         rand_minmax = rand(1:3),
                         NLoptRandSeed=rand(1:Int(1e8)),gpuEnabled = true,
                         cudaCores=cudaCores,ptx_str=ptx_str)