text:
# AVIcodec_type
import Base: read, read!, show, close, eof, isopen, seek, seekstart

export read, read!, read_raw, read_raw!, pump, openvideo, opencamera, playvideo,
    viewcam, play, gettime, skipframe, skipframes, counttotalframes,
    raw_frame_size, out_frame_size, raw_pixel_format, out_bytes_size,
    out_frame_eltype

const ReaderBitTypes = Union{UInt8, UInt16}
const ReaderNormedTypes = Normed{T} where T<: ReaderBitTypes
const ReaderGrayTypes = Gray{T} where T<: ReaderNormedTypes
const ReaderRgbTypes = RGB{T} where T<:ReaderNormedTypes
const ReaderElTypes = Union{ReaderGrayTypes, ReaderRgbTypes, ReaderNormedTypes,
                            ReaderBitTypes}
const PermutedArray{T,N,perm,iperm,AA <: StridedArray} =
    Base.PermutedDimsArrays.PermutedDimsArray{T,N,perm,iperm,AA}
const VidArray{T,N} = Union{StridedArray{T,N},PermutedArray{T,N}}

const VidRawBuff = StridedArray{UInt8}

const VIO_ALIGN = 32

abstract type StreamContext end

# An audio-visual input stream/file
mutable struct AVInput{I}
    io::I
    format_context::AVFormatContextPtr
    avio_context::AVIOContextPtr
    avio_ctx_buffer_size::UInt
    packet::AVPacketPtr

    unknown_indices::Vector{Int}
    video_indices::Vector{Int}
    audio_indices::Vector{Int}
    data_indices::Vector{Int}
    subtitle_indices::Vector{Int}
    attachment_indices::Vector{Int}

    listening::Set{Int}
    stream_contexts::Dict{Int, StreamContext}

    isopen::Bool
    finished::Bool
end

function show(io::IO, avin::AVInput)
    println(io, "AVInput(", avin.io, ", ...), with")
    (len = length(avin.video_indices))      > 0 && println(io, "  $len video stream(s)")
    (len = length(avin.audio_indices))      > 0 && println(io, "  $len audio stream(s)")
    (len = length(avin.data_indices))       > 0 && println(io, "  $len data stream(s)")
    (len = length(avin.subtitle_indices))   > 0 && println(io, "  $len subtitle stream(s)")
    (len = length(avin.attachment_indices)) > 0 && println(io, "  $len attachment stream(s)")
    (len = length(avin.unknown_indices))    > 0 && println(io, "  $len unknown stream(s)")
end

const TRANSCODE = true
const NO_TRANSCODE = false

mutable struct VideoReader{transcode, T<:GraphType, I} <: StreamContext
    avin::AVInput{I}
    stream_index0::Int

    codec_context::AVCodecContextPtr

    frame_queue::Vector{Vector{UInt8}}
    frame_graph::T
    graph_input_occupied::Bool
    bits_per_result_pixel::Int

    flush::Bool
    finished::Bool
end

"""
    load(filename::String, args...; kwargs...)

Load video file `filename` into memory as vector of image arrays, setting `args` and `kwargs` on the `openvideo` process.
"""
function load(filename::String, args...; kwargs...)
    openvideo(filename, args...; kwargs...) do io
        collect(io)
    end
end

show(io::IO, vr::VideoReader) = print(io, "VideoReader(...)")

function iterate(r::VideoReader, state = 0)
    eof(r) && return
    return read(r), state + 1
end

is_finished(r::VideoReader) = r.finished

IteratorSize(::Type{<:VideoReader}) = Base.SizeUnknown()

IteratorEltype(::Type{<:VideoReader}) = Base.EltypeUnknown()

convert(::Type{Rational{T}}, r::AVRational) where T = Rational{T}(r.num, r.den)
convert(::Type{Rational}, r::AVRational) = Rational(r.num, r.den)
convert(::Type{AVRational}, r::Rational) = AVRational(numerator(r), denominator(r))

# Pump input for data
function pump(avin::AVInput)
    while avin.isopen && !avin.finished
        sigatomic_begin()
        ret = av_read_frame(avin.format_context, avin.packet)
        sigatomic_end()
        if ret < 0
            avin.finished = true
            break
        end
        stream_index = avin.packet.stream_index
        if stream_index in avin.listening
            # Decode the packet, and check if the frame is complete
            r = avin.stream_contexts[stream_index]
            decode(r, avin.packet)
            av_packet_unref(avin.packet)

            # If the frame is complete, we're done
            frame_is_queued(r) && return Int(stream_index)
        else
            # If we're not listening to this stream, skip it
            av_packet_unref(avin.packet)
        end
    end
    if avin.finished
        for (stream_index, r) in pairs(avin.stream_contexts)
            if !r.finished
                decode(r, C_NULL) # Flush packet
                r.flush = true
                frame_is_queued(r) && return stream_index
            end
        end
    end
    return -1
end

pump(r::StreamContext) = pump(r.avin)

function pump_until_frame(r, err = false)
    while !frame_is_queued(r)
        idx = pump(r.avin)
        idx == r.stream_index0 && break
        if idx == -1
            err ? throw(EOFError()) : return false
        end
    end
    return true
end

# This will point to _read_packet, but is set in __init__()
const read_packet = Ref{Ptr{Cvoid}}(C_NULL)

function _read_packet(avin_ptr::Ptr, pbuf::Ptr{UInt8}, buf_size::Cint)
    avin = unsafe_pointer_to_objref(avin_ptr)
    out = unsafe_wrap(Array, pbuf, (buf_size,))
    convert(Cint, readbytes!(avin.io, out))
end

function AVIOContextPtr(avin::AVInput{<:IO})
    avin_ptr = pointer_from_objref(avin)
    AVIOContextPtr(avin.avio_ctx_buffer_size, avin_ptr, read_packet[])
end

function open_avinput(avin::AVInput, io::IO, input_format=C_NULL, options=C_NULL)

    !isreadable(io) && error("IO not readable")

    # These allow control over how much of the stream to consume when
    # determining the stream type
    # TODO: Change these defaults if necessary, or allow user to set
    #av_opt_set(avin.format_context "probesize", "100000000", 0)
    #av_opt_set(avin.format_context, "analyzeduration", "1000000", 0)

    # Allocate the io buffer used by AVIOContext
    # Must be done with av_malloc, because it could be reallocated
    avio_context = AVIOContextPtr(avin)
    avin.avio_context = avio_context
    avin.format_context.pb = @preserve avio_context unsafe_convert(Ptr{AVIOContext}, avio_context)

    # "Open" the input
    ret = avformat_open_input(avin.format_context, C_NULL, input_format, options)
    if ret != 0
        error("Unable to open input")
    end

    nothing
end

function open_avinput(avin::AVInput, source::AbstractString,
                      input_format=C_NULL, options=C_NULL)
    ret = avformat_open_input(avin.format_context, source, input_format, options)
    ret != 0 && error("Could not open $source")
    nothing
end

function AVInput(
        source::T, input_format=C_NULL, options=C_NULL;
        avio_ctx_buffer_size=4096 # recommended value by FFMPEG documentation
    ) where T <: Union{IO,AbstractString}

    # Register all codecs and formats

    packet = AVPacketPtr()
    format_context = AVFormatContextPtr()
    avio_context = AVIOContextPtr(C_NULL)

    # Allocate this object (needed to pass into AVIOContext in open_avinput)
    avin = AVInput{T}(source, format_context, avio_context,
                      avio_ctx_buffer_size, packet,
                      [Int[] for i = 1:6]..., Set(Int[]),
                      Dict{Int, StreamContext}(), false, false)

    # Set up the format context and open the input, based on the type of source
    open_avinput(avin, source, input_format, options)
    avin.isopen = true

    # Get the stream information
    ret = avformat_find_stream_info(avin.format_context, C_NULL)
    ret < 0 && error("Unable to find stream information")

    for i = 1:format_context.nb_streams
        codec_type = format_context.streams[i].codecpar.codec_type
        target_array = codec_type == AVMEDIA_TYPE_VIDEO      ? avin.video_indices      :
                       codec_type == AVMEDIA_TYPE_AUDIO      ? avin.audio_indices      :
                       codec_type == AVMEDIA_TYPE_DATA       ? avin.data_indices       :
                       codec_type == AVMEDIA_TYPE_SUBTITLE   ? avin.data_indices       :
                       codec_type == AVMEDIA_TYPE_ATTACHMENT ? avin.attachment_indices :
                                                               avin.unknown_indices
        push!(target_array, i - 1)
        # Set the stream to discard all packets. Individual StreamContexts can
        # re-enable streams.
        format_context.streams[i].discard = AVDISCARD_ALL
    end
    avin
end

get_stream(avin::AVInput, stream_index0) = avin.format_context.streams[stream_index0 + 1]
get_stream(r::VideoReader) = get_stream(r.avin, r.stream_index0)

function VideoReader(avin::AVInput{I}, video_stream = 1;
                     transcode::Bool = true,
                     target_format::Union{Nothing, Cint} = nothing,
                     pix_fmt_loss_flags = 0,
                     target_colorspace_details = nothing,
                     allow_vio_gray_transform = true,
                     swscale_options::OptionsT = (;),
                     sws_color_options::OptionsT = (;),
                     thread_count::Union{Nothing, Int} = Sys.CPU_THREADS) where I
    bad_px_type = transcode && target_format !== nothing &&
        !is_pixel_type_supported(target_format)
    bad_px_type && error("Unsupported pixel format $target_format")

    1 <= video_stream <= length(avin.video_indices) ||
        error("video stream $video_stream not found")

    stream_index0 = avin.video_indices[video_stream]

    # Find decoder
    stream = get_stream(avin, stream_index0)
    codec = avcodec_find_decoder(stream.codecpar.codec_id)
    check_ptr_valid(codec, false) || error("Failed to find decoder")

    # Tell the demuxer to retain packets from this stream
    stream.discard = AVDISCARD_DEFAULT

    # Create decoder context
    codec_context = AVCodecContextPtr(codec) # Allocates
    codec_context.thread_count = thread_count === nothing ? codec_context.thread_count : thread_count
    # Transfer parameters to decoder context
    ret = avcodec_parameters_to_context(codec_context, stream.codecpar)
    ret < 0 && error("Could not copy the codec parameters to the decoder")
    codec_context.pix_fmt < 0 && error("Unknown pixel format")

    # Open the decoder
    sigatomic_begin()
    lock(VIO_LOCK)
    ret = avcodec_open2(codec_context, codec, C_NULL)
    unlock(VIO_LOCK)
    sigatomic_end()
    ret < 0 && error("Could not open codec")

    if target_format === nothing # automatically determine format
        dst_pix_fmt, pix_fmt_loss = _vio_determine_best_pix_fmt(
            codec_context.pix_fmt; loss_flags = pix_fmt_loss_flags
        )
    else
        dst_pix_fmt = target_format
    end

    if target_colorspace_details === nothing
        colorspace_details = get(VIO_DEFAULT_TRANSFER_COLORSPACE_DETAILS,
                                 dst_pix_fmt,
                                 VIO_DEFAULT_COLORSPACE_DETAILS)
    else
        colorspace_details = target_colorspace_details
    end
    bit_depth, padded_bits_per_pixel = pix_fmt_to_bits_per_pixel(dst_pix_fmt)
    width = codec_context.width
    height = codec_context.height

    use_vio_gray_transform = transcode &&
        codec_context.pix_fmt == dst_pix_fmt && allow_vio_gray_transform &&
        dst_pix_fmt in VIO_GRAY_SCALE_TYPES &&
        colorspace_details.color_range != codec_context.color_range
    if use_vio_gray_transform
        frame_graph = GrayTransform()
        set_basic_frame_properties!(frame_graph.dstframe, width, height,
                                    dst_pix_fmt)
        frame_graph.src_depth = frame_graph.dst_depth = bits_per_pixel
        frame_graph.dstframe.color_range = colorspace_details.color_range
    elseif ! transcode || codec_context.pix_fmt == dst_pix_fmt
        frame_graph = AVFramePtr()
    else
        frame_graph = SwsTransform(width, height, codec_context.pix_fmt,
                                   codec_context.color_primaries,
                                   codec_context.color_range, dst_pix_fmt,
                                   colorspace_details.color_primaries,
                                   colorspace_details.color_range,
                                   sws_color_options, swscale_options)
        set_basic_frame_properties!(frame_graph.dstframe, width, height,
                                    dst_pix_fmt)
    end

    vr = VideoReader{transcode, typeof(frame_graph), I}(avin, stream_index0,
                                                        codec_context,
                                                        Vector{UInt8}[],
                                                        frame_graph, false,
                                                        padded_bits_per_pixel,
                                                        false, false)

    push!(avin.listening, stream_index0)
    avin.stream_contexts[stream_index0] = vr

    vr
end

VideoReader(s::Union{IO, AbstractString}, args...; kwargs...) =
    VideoReader(AVInput(s), args...; kwargs... )

# a convenience function for getting the aspect ratio
function aspect_ratio(f::VideoReader)
    aspect_num = f.codec_context.sample_aspect_ratio.num
    aspect_den = f.codec_context.sample_aspect_ratio.den
    aspect_ratio = (aspect_den == 0) ? Cint(0) // Cint(1) :
        aspect_num // aspect_den
    # if the stored aspect ratio is nonsense then we default to one. OBS, this
    # might still be wrong for some videos and an unnecessary test for most
    if iszero(aspect_ratio) || isnan(aspect_ratio) || isinf(aspect_ratio)
        fixed_aspect = 1//1
    else
        fixed_aspect = aspect_ratio
    end
    fixed_aspect
end

framerate(f::VideoReader) = f.codec_context.time_base.den // f.codec_context.time_base.num
height(f::VideoReader) = f.codec_context.height
width(f::VideoReader) = f.codec_context.width

# Does not check input size, meant for internal use only
function stash_graph_input!(imgbuf, r::VideoReader, align = VIO_ALIGN)
    frame = graph_input_frame(r)
    av_image_copy_to_buffer(imgbuf, length(imgbuf), Ref(frame.data),
                            Ref(frame.linesize), r.codec_context.pix_fmt,
                            r.codec_context.width, r.codec_context.height,
                            align)
    return imgbuf
end

stash_graph_input(r, align = VIO_ALIGN) =
    stash_graph_input!(Vector{UInt8}(undef, out_bytes_size(r, align)), r, align)

function unpack_stashed_planes!(r::VideoReader, imgbuf)
    frame = graph_input_frame(r)
    av_make_writable(frame)
    av_image_fill_arrays(frame, frame.linesize, imgbuf, r.codec_context.pix_fmt,
                         r.codec_context.width, r.codec_context.height,
                         VIO_ALIGN)
    return r
end

check_send_packet_return(r) = (r < 0 && r != -Libc.EAGAIN) && error("Could not send packet")

function decode(r::VideoReader, packet)
    # Do we already have a complete frame that hasn't been consumed?
    r.finished && return
    if graph_blocked(r)
        push!(r.frame_queue, stash_graph_input(r))
        remove_graph_input!(r)
    end
    if !r.flush
        pret = avcodec_send_packet(r.codec_context, packet)
        check_send_packet_return(pret)
    end
    fret = avcodec_receive_frame(r.codec_context, graph_input_frame(r))
    if fret == 0
        r.graph_input_occupied = true
    elseif fret == VIO_AVERROR_EOF
        r.finished = true
    elseif fret != -Libc.EAGAIN
        error("Decoding error $fret")
    end
    if !r.finished && !r.flush && pret == -Libc.EAGAIN
        pret2 = avcodec_send_packet(r.codec_context, packet)
        check_send_packet_return(pret2)
    end
    return
end

graph_input_frame(r::VideoReader) = graph_input_frame(r.frame_graph)
graph_output_frame(r::VideoReader) = graph_output_frame(r.frame_graph)

unsupported_retrieval_format(fmt) = error("Unsupported format $(fmt)")

function drop_frame!(r::VideoReader)
    if isempty(r.frame_queue)
        graph_blocked(r) && remove_graph_input!(r)
    else
        if r.graph_input_occupied
            push!(r.frame_queue, stash_graph_input(r))
            remove_graph_input!(r)
        end
        popfirst!(r.frame_queue)
    end
    r
end

function drop_frames!(r::VideoReader)
    empty!(r.frame_queue)
    remove_graph_input!(r)
end

n_queued_frames(r::VideoReader) = length(r.frame_queue) + r.graph_input_occupied

function fill_graph_input!(r::VideoReader)
    pump_until_frame(r)

    # Get the oldest frame stored by VideoReader

    # If we have an unprocessed frames in the queue
    if !isempty(r.frame_queue)
        # But also a decoded frame
        if graph_blocked(r)
            # Then move the decoded frame to the back of the queue
            push!(r.frame_queue, stash_graph_input(r))
        end
        unpack_stashed_planes!(r, popfirst!(r.frame_queue))
        r.graph_input_occupied = true
    end
end

execute_graph!(r::VideoReader) = exec!(r.frame_graph)

function transfer_graph_output!(buf, r::VideoReader)
    bytes_per_pixel, nrem = divrem(r.bits_per_result_pixel, 8)
    nrem > 0 && error("Unsupported pixel size")
    transfer_frame_to_img_buf!(buf, graph_output_frame(r), bytes_per_pixel)
end

# Converts a grabbed frame to the correct format (RGB by default)
function _retrieve!(r::VideoReader, buf)
    fill_graph_input!(r)
    execute_graph!(r)
    transfer_graph_output!(buf, r)
    remove_graph_input!(r)
    return buf
end

function retrieve!(r::VideoReader{TRANSCODE}, buf::VidArray{T}
                   ) where T <: ReaderElTypes
    if !out_img_check(r, buf)
        throw(ArgumentError("Buffer is the wrong size, element type, or stride"))
    end
    _retrieve!(r, buf)
end

function _retrieve_raw!(r, buf::VidRawBuff, align = VIO_ALIGN)
    fill_graph_input!(r)
    stash_graph_input!(buf, r, align)
    remove_graph_input!(r)
    return buf
end

function retrieve_raw!(r, buf::VidRawBuff, align = VIO_ALIGN)
    if !raw_buff_check(buf, r, align)
        throw(ArgumentError("Buffer is the wrong size or stride"))
    end
    _retrieve_raw!(r, buf, align)
end

# Converts a grabbed frame to the correct format (RGB by default)
# TODO: type stability?
function retrieve(r::VideoReader{TRANSCODE}) # true=transcode
    pix_fmt = out_frame_format(r)
    if !is_pixel_type_supported(pix_fmt)
        unsupported_retrieveal_format(pix_fmt)
    end
    elt = VIO_PIX_FMT_DEF_ELTYPE_LU[pix_fmt]
    width, height = out_frame_size(r)
    buf = PermutedDimsArray(Matrix{elt}(undef, width, height), (2, 1))

    _retrieve!(r, buf)
end

retrieve!(r::VideoReader{NO_TRANSCODE}, buf::VidRawBuff, args...) =
    retrieve_raw!(r, buf, args...)

function retrieve_raw(r::VideoReader, align = VIO_ALIGN)
    imgbuf = Vector{UInt8}(undef, out_bytes_size(r, align))
    retrieve_raw!(r, imgbuf, align)
    imgbuf, align
end

retrieve(r::VideoReader{NO_TRANSCODE}, align::Integer = 1) =
    retrieve_raw(r, align)[1]

# Utility functions

open(filename::AbstractString) = AVInput(filename) # not exported

"""
    openvideo(file[, video_stream = 1]; <keyword arguments>) -> reader
    openvideo(f, ...)

Open `file` and create an object to read and decode video stream number
`video_stream`. `file` can either be a `AVInput` created by `VideoIO.open`,
the name of a file as an `AbstractString`, or instead an `IO` object. However,
support for `IO` objects is incomplete, and does not currently work with common
video containers such as `*.mp4` files.

Frames can be read from the `reader` with `read` or `read!`, or alternatively by
using the iterator interface provided for `reader`. To close the `reader`,
simply use `close`. Seeking within the reader can be accomplished using `seek`,
`seekstart`. Frames can be skipped with `skipframe`, or `skipframes`. The
current time in the video stream can be accessed with `gettime`. Details about
the frame dimension can be found with `out_frame_size`. The total number of
frames can be found with `counttotalframes`.

If called with a single argument function as the first argument, the `reader`
will be passed to the function, and will be closed once the call returns whether
or not an error occurred.

The decoder options and conversion to Julia arrays is controlled by the keyword
arguments listed below.

# Keyword arguments
- `transcode::Bool = true`: Determines whether decoded frames are transferred
    into a Julia matrix with easily interpretable element type, or instead
    returned as raw byte buffers.
- `target_format::Union{Nothing, Cint} = nothing`: Determines the target pixel
    format that decoded frames will be transformed into before being transferred
    to an output array. This can either by a `VideoIO.AV_PIX_FMT_*` value
    corresponding to a FFmpeg
    [`AVPixelFormat`]
    (https://ffmpeg.org/doxygen/4.1/pixfmt_8h.html#a9a8e335cf3be472042bc9f0cf80cd4c5),
    and must then also be a format supported by the VideoIO, or instead
    `nothing`, in which case the format will be automatically chosen by FFmpeg.
    This list of currently supported pixel formats, and the matrix element type
    that each pixel format corresponds with, are elements of
    `VideoIO.VIO_PIX_FMT_DEF_ELTYPE_LU`.
- `pix_fmt_loss_flags = 0`: Loss flags to control how transfer pixel format is
    chosen. Only valid if `target_format = nothing`. Flags must correspond to
    FFmpeg
    [loss flags]
    (http://ffmpeg.org/doxygen/2.5/pixdesc_8h.html#a445e6541dde2408332c216b8d0accb2d).
- `target_colorspace_details = nothing`: Information about the color space
    of output Julia arrays. If `nothing`, then this will correspond to a
    best-effort interpretation of `Colors.jl` for the corresponding element
    type. To override these defaults, create a `VideoIO.VioColorspaceDetails`
    object using the appropriate `AVCOL_` definitions from FFmpeg, or use
    `VideoIO.VioColorspaceDetails()` to use the FFmpeg defaults. To avoid
    rescaling limited color range data (mpeg) to full color range output (jpeg),
    then set this to `VideoIO.VioColorspaceDetails()` to avoid additional
    scaling by `sws_scale`.
- `allow_vio_gray_transform = true`: Instead of using `sws_scale` for gray data,
    use a more accurate color space transformation implemented in `VideoIO` if
    `allow_vio_gray_gransform = true`. Otherwise, use `sws_scale`.
- `swscale_options::OptionsT = (;)`: A `Namedtuple`, or `Dict{Symbol, Any}` of
    options for the swscale object used to perform color space scaling. Options
    must correspond with options for FFmpeg's
    [scaler](https://ffmpeg.org/ffmpeg-all.html#Scaler-Options) filter.
- `sws_color_options::OptionsT = (;)`: Additional keyword arguments passed to
    [sws_setColorspaceDetails]
    (http://ffmpeg.org/doxygen/2.5/group__libsws.html#ga541bdffa8149f5f9203664f955faa040).
- `thread_count::Union{Nothing, Int} = Sys.CPU_THREADS`: The number of threads the codec is
    allowed to use or `nothing` for default codec behavior. Defaults to `Sys.CPU_THREADS`.
"""
openvideo(s::Union{IO, AbstractString, AVInput}, args...; kwargs...) =
    VideoReader(s, args...; kwargs...)

function openvideo(f, args...; kwargs...)
    r = openvideo(args...; kwargs...)
    try
        f(r)
    finally
        close(r)
    end
end

"""
    read(reader::VideoReader) -> frame::PermutedDimsArray{Matrix{T}}

Allocate a matrix `frame` and fill it with the next video frame from `reader`,
advancing the video stream. For details about the frame matrix, see `read!`.
"""
read(r::VideoReader, args...) = retrieve(r, args...)

"""
    read_raw(reader::VideoReader[, align = VideoIO.VIO_ALIGN]) ->
        (buff::Vector{UInt8}, align::Int)

Allocate a `Vector{UInt8}` and return the bytes and scanline memory alignment of
the next untransformed frame of `reader`, advancing the video stream.
"""
read_raw(r::VideoReader, args...) = retrieve_raw(r, args...)

"""
    read!(reader::VideoReader, buf::Union{PermutedArray{T,N}, Array{T,N}}) where T<:ReaderElTypes

Read a frame from `reader` into array `buf`, which can either be a regular array
or a permuted view of a regular array, and advance the video stream. `buf` must
be scanline-major, meaning that adjacent pixels on the same horizontal line of
the video frame are also adjacent in memory. The element type of `buf` must be
one of the supported element types, any key of
`VideoIO.VIO_DEF_ELTYPE_PIX_FMT_LU`, and must match the pixel format chosen when
`reader` was created. For any supported `Gray` element type, VideoIO will also
support the corresponding `Normed` or `Unsigned` element types as well. The size
of `buf` must correspond to the size of `out_frame_size`. If `buf` is a regular
array, then frame width must correspond to the number of rows of `buf`, and
conversely if `buf` is a `PermutedDimsArray` view, frame width must correspond
to the number of columns in the view.
"""
read!(r::VideoReader, buf::AbstractArray{T}) where {T <: ReaderElTypes} =
    retrieve!(r, buf)

"""
    read_raw!(reader, buf[, align = VideoIO.VIO_ALIGN])

Read the next frame in the video stream described by `reader`, and transfer the
untransformed bytes into `buf`, which has scanline memory alignment `align`.
`buf` must be `out_bytes_size` long.
"""
read_raw!(r::VideoReader, buf, args...) = retrieve_raw!(r, buf, args...)

isopen(avin::AVInput{I}) where {I <: IO} = isopen(avin.io)
isopen(avin::AVInput) = avin.isopen
isopen(r::VideoReader) = isopen(r.avin)

"""
    raw_frame_size(reader) -> (width, height)

Return the dimensions of the raw frames returned by `reader`.
"""
raw_frame_size(r::VideoReader) = (r.codec_context.width, r.codec_context.height)

"""
    out_frame_size(reader) -> (width, height)

Return frame dimensions, `(width, height)` of frames returned by `read(reader)`.
`read` will return frames as `PermutedDimsArray` views of Julia matrices with
the frame height spanning the rows of the matrix, and frame width spanning the
columns. The underlying array has the opposite orientation.
"""
function out_frame_size end

out_frame_size(r::VideoReader{<:Any, AVFramePtr}) = raw_frame_size(r)
out_frame_size(t::SwsTransform) = (t.dstframe.width, t.dstframe.height)
out_frame_size(r::VideoReader{TRANSCODE, SwsTransform}) =
    out_frame_size(r.frame_graph)

"""
    raw_pixel_format(reader) -> Cint

Return the pixel format for frames of `reader`. The return value will be a
`Cint` corresponding to `VideoIO.AV_PIX_FMT_*`, which corresponds to a
[`AVPixelFormat`]
(https://ffmpeg.org/doxygen/4.1/pixfmt_8h.html#a9a8e335cf3be472042bc9f0cf80cd4c5).
This will determine the format of byte arrays returned by `read_raw`, and
`read_raw!`.
"""
raw_pixel_format(r::VideoReader) = r.codec_context.pix_fmt

out_frame_format(r::VideoReader{<:Any, AVFramePtr}) = raw_pixel_format(r)
out_frame_format(t::SwsTransform) = t.dstframe.format
out_frame_format(t::GrayTransform) = t.dstframe.format
out_frame_format(r::VideoReader{TRANSCODE, SwsTransform}) = out_frame_format(r.frame_graph)
out_frame_format(r::VideoReader{TRANSCODE, GrayTransform}) = out_frame_format(r.frame_graph)

"""
    out_frame_eltype(reader) -> T

Return the Julia eltype `T` of matrices returned by `read(reader)`.
"""
out_frame_eltype(r::VideoReader) = VIO_PIX_FMT_DEF_ELTYPE_LU[out_frame_format(r)]

out_img_size_check(r, buf) = all(size(buf) .>= out_frame_size(r))
out_img_size_check(r, buf::PermutedDimsArray) = out_img_size_check(r, parent(buf))

out_img_eltype_check(::Type{T}, p) where T <: UInt8 = p == AV_PIX_FMT_GRAY8
out_img_eltype_check(::Type{T}, p) where T <: UInt16 = p == AV_PIX_FMT_GRAY10LE
out_img_eltype_check(::Type{X}, p) where {T, X <: Normed{T}} =
    out_img_eltype_check(T, p)
out_img_eltype_check(::Type{X}, p) where {T, X <: Gray{T}} =
    out_img_eltype_check(T, p)
out_img_eltype_check(::Type{T}, p) where T<:RGB{N0f8} = p == AV_PIX_FMT_RGB24
out_img_eltype_check(::Type{T}, p) where T<:RGB{N6f10} = p == AV_PIX_FMT_GBRP10LE

out_img_eltype_check(fmt::Integer, ::AbstractArray{T}) where T =
    out_img_eltype_check(T, fmt)
out_img_eltype_check(r, buf) = out_img_eltype_check(out_frame_format(r), buf)

stride_check(buf) = stride(buf, 1) == 1
stride_check(buf::PermutedDimsArray) = stride_check(parent(buf))

out_img_check(r, buf) = out_img_size_check(r, buf) &
    out_img_eltype_check(r, buf) & stride_check(buf)

"""
    out_bytes_size(reader[, align = VideoIO.VIO_ALIGN])

Return the number required to store the raw frame data for frames of `reader`
with memory alignment for each scanline `align`. Input bye buffers to
`read_raw!` must be this length.
"""
function out_bytes_size end

function out_bytes_size(fmt, width, height, align = VIO_ALIGN)
    align > 0 || throw(ArgumentError("align must be greater than zero"))
    sz = av_image_get_buffer_size(fmt, width, height, align)
    sz < 0 && error("Could not determine the buffer size")
    sz
end
out_bytes_size(r, args...) =
    out_bytes_size(raw_pixel_format(r), out_frame_size(r)..., args...)

out_bytes_size_check(buf, r, args...) = sizeof(buf) == out_bytes_size(r, args...)

raw_buff_check(buf, r, args...) = out_bytes_size_check(buf, r, args...) &
    stride_check(buf)

graph_blocked(r) = r.graph_input_occupied # TODO: make sure the last frame was made available
frame_is_queued(r::StreamContext) = !isempty(r.frame_queue) || graph_blocked(r)
frame_is_queued(avin::AVInput) = mapreduce(frame_is_queued, |, values(avin.stream_contexts), init = false)

function remove_graph_input!(r::VideoReader)
    av_frame_unref(graph_input_frame(r))
    r.graph_input_occupied = false
end

seconds_to_timestamp(s, time_base::Rational) = round(Int, s / time_base)
seconds_to_timestamp(s, time_base::AVRational) = seconds_to_timestamp(s, convert(Rational, time_base))

get_video_stream_index(s::VideoReader) =
    findfirst(x -> x == s.stream_index0, s.avin.video_indices)

"""
    gettime(s::VideoReader)

Return timestamp of the last returned frame, if available. Else, return `0.0`.
"""
function gettime(s::VideoReader)
    # No frame has been read
    t = position(s)
    t == nothing && return 0.0
    return Float64(t)
end

# To be called for all stream contexts following a seek of AVInput
function reset_file_position_information!(r::VideoReader)
    avcodec_flush_buffers(r.codec_context)
    drop_frames!(r)
    r.flush = false
    r.finished = false
end

get_frame_presentation_time(stream, frame) =
    convert(Rational, stream.time_base) * convert(Rational, frame.pts)

function get_frame_period_timebase(r::VideoReader)
    # This is probably not valid for many codecs, frame period in timebase
    # units
    stream = get_stream(s)
    time_base = convert(Rational, stream.time_base)
    time_base == 0 && return nothing
    frame_rate = convert(Rational, av_stream_get_r_frame_rate(stream))
    frame_period_timebase = round(Int64,  1 / (frame_rate * time_base))
    frame_period_timebase
end

"""
    seek(reader::VideoReader, seconds)

Seeks into the parent `AVInput` using this video stream's index. See [`seek`]
for `AVInput`.
"""
function seek(r::VideoReader, seconds::Number)
    !isopen(r) && throw(ErrorException("Video input stream is not open!"))
    video_stream_idx = get_video_stream_index(r)
    seek(r.avin, seconds, video_stream_idx)
    return r
end

function seek_trim(r::VideoReader, seconds::Number)
    stream = get_stream(r)
    time_base = convert(Rational, stream.time_base)
    time_base == 0 && error("No time base")
    target_pts = seconds_to_timestamp(seconds, time_base)
    frame_rate = convert(Rational, av_stream_get_r_frame_rate(stream))
    frame_period_timebase = round(Int64,  1 / (frame_rate * time_base))
    gotframe = pump_until_frame(r, false)
    # If advancing another frame would still leave us before the target

    frame = graph_input_frame(r)
    while gotframe && frame.pts != AV_NOPTS_VALUE &&
          frame.pts + frame_period_timebase <= target_pts
        drop_frame!(r)
        gotframe = pump_until_frame(r, false)
    end
end

"""
    seek(avin::AVInput, seconds::AbstractFloat, video_stream::Integer=1)

Seek through the container format `avin` so that the next frame returned by
the stream indicated by `video_stream` will have a timestamp greater than or
equal to `seconds`.
"""
function seek(avin::AVInput{T}, seconds::Number, video_stream::Integer=1) where T <: AbstractString
    stream_index0 = avin.video_indices[video_stream]
    stream = get_stream(avin, stream_index0)
    time_base = convert(Rational, stream.time_base)
    time_base == 0 && error("No time base for stream")
    pts = seconds_to_timestamp(seconds, time_base)
    ret = avformat_seek_file(avin.format_context, stream_index0, typemin(Int),
                             pts, typemax(Int), 0)
    ret < 0 && throw(ErrorException("Could not seek in stream"))
    avin.finished = false
    for r in values(avin.stream_contexts)
        reset_file_position_information!(r)
        seek_trim(r, seconds)
    end
    return avin
end
"""
    seekstart(reader::VideoReader)

Seek to time zero of the parent `AVInput` using `reader`'s stream index. See
`seekstart` for `AVInput` objects.
"""
seekstart(s::VideoReader, args...) = seek(s, 0, args...)

"""
    seekstart(avin::AVInput{T}, video_stream_index=1) where T <: AbstractString

Seek to time zero of AVInput object.
"""
seekstart(avin::AVInput{<:AbstractString}, args...) = seek(avin, 0, args...)

seekstart(avin::AVInput{<:IO}, args...) =
    throw(ErrorException("Sorry, Seeking is not supported from IO streams"))

"""
    skipframe(s::VideoReader; throwEOF=true)

Skip the next frame. If End of File is reached, EOFError thrown if throwEOF=true.
Otherwise returns true if EOF reached, false otherwise.
"""
function skipframe(s::VideoReader; throwEOF=true)
    gotframe = pump_until_frame(s, throwEOF)
    gotframe && drop_frame!(s)
    return !gotframe
end

"""
    skipframes(s::VideoReader, n::Int; throwEOF=true) -> n

Skip the next `n` frames. If End of File is reached and `throwEOF=true`,
a `EOFError` will be thrown. Returns the number of frames that were skipped.
"""
function skipframes(s::VideoReader, n::Int; throwEOF=true)
    for i in 1:n
        skipframe(s, throwEOF=throwEOF) && return i - 1
    end
    return n
end

"""
    counttotalframes(reader) -> n::Int

Count the total number of frames in the video by seeking to start, skipping
through each frame, and seeking back to the start.

For a faster alternative that relies on video container metadata, try `get_number_frames`.
"""
function counttotalframes(s::VideoReader)
    seekstart(s)
    n = 0
    while true
        skipframe(s, throwEOF = false) && break
        n += 1
    end
    seekstart(s)
    return n
end

function eof(avin::AVInput)
    !isopen(avin) && return true
    frame_is_queued(avin) && return false
    allfinished = mapreduce(is_finished, &, values(avin.stream_contexts),
                            init = true)
    allfinished && return true
    got_frame = pump(avin) != -1
    return !got_frame
end

eof(r::VideoReader) = eof(r.avin)

close(r::VideoReader) = close(r.avin)

# Free AVIOContext object when done
function close(avin::AVInput)
    avin.isopen || return
    avin.isopen = false

    empty!(avin.stream_contexts)
    empty!(avin.listening)
    if check_ptr_valid(avin.format_context, false)
        # Replace the existing object in avin with a null pointer. The finalizer
        # for AVFormatContextPtr will close it and clean up its memory
        sigatomic_begin()
        avformat_close_input(avin.format_context) # This is also done in the finalizer,
        # but closing the input here means users won't have to wait for the GC to run
        # before trying to remove the file
        avin.format_context = AVFormatContextPtr(Ptr{AVFormatContext}(C_NULL))
        sigatomic_end()
    end

    if check_ptr_valid(avin.avio_context, false)
        # Replace the existing object in avin with a null pointer. The finalizer
        # will close it and clean up its memory
        sigatomic_begin()
        avin.avio_context = AVIOContextPtr(Ptr{AVIOContext}(C_NULL))
        sigatomic_end()
    end
    nothing
end

function position(r::VideoReader)
    # this field is not indended for public use, should we be using it?
    last_pts = r.codec_context.pts_correction_last_pts
    # At the beginning, or just seeked
    last_pts == AV_NOPTS_VALUE && return nothing
        # Try to figure out the time of the oldest frame in the queue
    stream = get_stream(r)
    time_base = convert(Rational, stream.time_base)
    time_base == 0 && return nothing
    nqueued = n_queued_frames(r)
    nqueued == 0 && return last_pts * time_base
    frame_rate = convert(Rational, av_stream_get_r_frame_rate(stream))
    last_returned_pts = last_pts - round(Int, nqueued / (frame_rate * time_base))
    return last_returned_pts * time_base
end

### Camera Functions

# These are set in __init__()
const DEFAULT_CAMERA_FORMAT = Ref{Any}()
const CAMERA_DEVICES = String[]
const DEFAULT_CAMERA_DEVICE = Ref{String}()
const DEFAULT_CAMERA_OPTIONS = AVDict()

function get_camera_devices(ffmpeg, idev, idev_name)
    camera_devices = String[]

    read_vid_devs = false
    lines = FFMPEG.exe(`-list_devices true -f $idev -i $idev_name`,collect=true)

    for line in lines
        if occursin("video devices",line)
            read_vid_devs = true
            continue
        elseif occursin("audio devices",line) || occursin("exit", line) || occursin("error", line)
            read_vid_devs = false
            continue
        end

        if read_vid_devs
            m = match(r"""\[.*"(.*)".?""", line)
            if m != nothing
                push!(camera_devices, m.captures[1])
            end

            # Alternative format (TODO: could be combined with the regex above)
            m = match(r"""\[.*\] \[[0-9]\] (.*)""", line)
            if m != nothing
                push!(camera_devices, m.captures[1])
            end
        end
    end

    return camera_devices
end

function opencamera(
        device=DEFAULT_CAMERA_DEVICE[],
        format=DEFAULT_CAMERA_FORMAT[],
        options = DEFAULT_CAMERA_OPTIONS,
        args...;
        kwargs...
    )
    camera = AVInput(device, format, options)
    VideoReader(camera, args...; kwargs...)
end


text:
"""
```
    MidwayEqualization <: AbstractHistogramAdjustmentAlgorithm
    MidwayEqualization(; nbins = 256, minval = 0, maxval = 1)

    adjust_histogram([T,] img_sequence, f::MidwayEqualization(nbins = 256, edges = nothing))
    adjust_histogram!([out_sequence,] img_sequence, f::MidwayEqualization(nbins = 256, edges = nothing))
```

Gives a pair of images the same histogram whilst maintaining as
much as possible their previous grey level dynamics.

# Details

The purpose of midway histogram equalization is to transform the intensities in
a pair of images so that the intensities distribute according to a common
"midway" distribution. The histogram representing the common distribution is
chosen so that the original  gray level dynamics of the images are preserved as
much as possible. If one interprets histograms as piecewise-constant models of
probability density functions (see [`build_histogram`](@ref
build_histogram(::AbstractArray, ::Integer, ::Union{Real,AbstractGray},
::Union{Real,AbstractGray}))), then the midway histogram equalization task can
be modeled as the problem of transforming one probability distribution into
another (see [`adjust_histogram`](@ref adjust_histogram(::Matching,::AbstractArray, ::AbstractArray, ::Integer))).
It turns out that the solution to this transformation problem involves the
cumulative and inverse cumulative distribution functions of the source and
"midway" probability density functions. In particular, let the random variables ``X_i \\thicksim p_{x_i} \\; (i = 1,2)``,
and ``Z \\thicksim p_{z}``  represent an intensity in the first, second and
"midway" image respectively, and let

```math
 S_{X_i}(x) = \\int_0^{x}p_{x_i}(w)\\mathrm{d} w \\; \\quad \\text{and} \\quad
 T_{Z}(x) = \\frac{2}{\\frac{1}{S_{X_1}(x)} + \\frac{1}{S_{X_2}(x)}}
```
represent the cumulative distribution functions of the two input images, and
their *harmonic mean*, respectively. Then the sought-after mapping
``Q_{X_i}(\\cdot)`` ``(i = 1,2)`` such that ``Q_{X_i}(x) \\thicksim p_{z} `` is
given by

```math
Q_{X_i}(x) =  T_{Z}^{-1}\\left( S_{X_i}(x) \\right),
```

where ``T_{Z}^{-1}(y) = \\operatorname{min} \\{ x \\in \\mathbb{R} : y \\leq
T_{Z}(x) \\}`` is the inverse cumulative distribution function of ``T_{Z}(x)``.

# Options

Various options for the parameters of the `adjust_histogram` function and
`MidwayEqualization` types are described in more detail below.

## Choices for `img_sequence`

The function `adjust_histogram` expects a length-2 `Vector` of images (the pair
of images) and returns a length-2 `Vector` of modified images.  The  function
can handle a variety of input types. The type of the returned image matches the
input type.

For colored images, the inputs are converted to
[YIQ](https://en.wikipedia.org/wiki/YIQ)  type and the distributions of the Y
channels are transformed according to a "midway" distribution. The modified Y
channel is then combined with the I and Q channels and the resulting image
converted to the same type as the input.

## Choices for `nbins`

You can specify the total number of bins in the histogram. If you do not
specify the number of bins then a default value of 256 bins is utilized.

## Choices for `edges`

If you do not designate the number of bins, then you have the option to directly
stipulate how the intervals will be divided by specifying a `AbstractRange`
type.

# Example

```julia
using Images, TestImages, ImageView, ImageContrastAdjustment

img = testimage("mandril_gray")

# The same image but with different intensitiy distributions
img1 = adjust_histogram(img, GammaCorrection(gamma = 2))
img2 = adjust_histogram(img, GammaCorrection(gamma = 1.2))

# Midway histogram equalization will transform these two images so that their
# intensity distributions are almost identical.
img_sequence = adjust_histogram([img1, img2], MidwayEqualization(nbins = 256))
img1o = first(img_sequence)
img2o = last(img_sequence)
```

# References
1. T. Guillemot and J. Delon, “*Implementation of the Midway Image Equalization*,” Image Processing On Line, vol. 5, pp. 114–129, Jun. 2016. [doi:10.5201/ipol.2016.140](https://doi.org/10.5201/ipol.2016.140)
"""
Base.@kwdef struct MidwayEqualization{T₁ <: Union{Integer, Nothing},
                					  T₂ <: Union{AbstractRange, Nothing}} <: AbstractHistogramAdjustmentAlgorithm
    nbins::T₁ = 256
    edges::T₂ = nothing
end

function (f::MidwayEqualization)(out_sequence::Vector{<:GenericGrayImage}, in_sequence::Vector{<:GenericGrayImage})
    length(out_sequence) == 2 || error("Please supply a length-2 output vector to store the pair of images.")
    length(in_sequence) == 2 || error("Please supply a length-2 input vector storing the image pair.")
    out1 = first(out_sequence)
    out2 = last(out_sequence)
    in1 = first(in_sequence)
    in2 = last(in_sequence)
    out1 .= in1
    out2 .= in2
    edges, pdf1, pdf2 = isnothing(f.edges) ? construct_pdfs(out1, out2, f.nbins) : construct_pdfs(out1, out2, f.edges)
    midway_pdf = zero(pdf1)
    cdf1 = cumsum(pdf1)
    cdf2 = cumsum(pdf2)
    # midway_cdf is the harmonic mean between cdf1 and cdf2.
    midway_cdf =  similar(cdf1)
    for i in eachindex(cdf1)
        if cdf1[i] == 0 || cdf2[i] == 0
            midway_cdf[i] = 0
        else
            midway_cdf[i] = 2 / (1/cdf1[i] + 1/cdf2[i])
        end
    end
    cdf2pdf!(midway_pdf, midway_cdf)
    index₁ = firstindex(out_sequence)
    index₂ = lastindex(out_sequence)
    out_sequence[index₁]  = match_pdf!(out1, edges, pdf1, midway_pdf)
    out_sequence[index₂]  = match_pdf!(out2, edges, pdf2, midway_pdf)
end

function (f::MidwayEqualization)(out_sequence::Vector{<:AbstractArray{<:Color3}}, in_sequence::Vector{<:AbstractArray{<:Color3}})
    length(out_sequence) == 2 || error("Please supply a length-2 output vector to store the pair of images.")
    length(in_sequence) == 2 || error("Please supply a length-2 input vector storing the image pair.")
    out1 = first(out_sequence)
    out2 = last(out_sequence)
    in1 = first(in_sequence)
    in2 = last(in_sequence)

    T₁ = eltype(out1)
    T₂ = eltype(out2)

    out1 .= in1
    out2 .= in2

    in_yiq1 = convert.(YIQ, in1)
    in_yiq1_view = channelview(in_yiq1)
    in_yiq2 = convert.(YIQ, in2)
    in_yiq2_view = channelview(in_yiq2)
    in_yiq_view_sequence = [view(in_yiq1_view,1,:,:), view(in_yiq2_view,1,:,:)]

    out_yiq1 = convert.(YIQ, out1)
    out_yiq1_view = channelview(out_yiq1)
    out_yiq2 = convert.(YIQ, out2)
    out_yiq2_view = channelview(out_yiq2)
    out_yiq_view_sequence = [view(out_yiq1_view,1,:,:), view(out_yiq2_view,1,:,:)]

    adjust_histogram!(out_yiq_view_sequence, in_yiq_view_sequence, f)

    index₁ = firstindex(out_sequence)
    index₂ = lastindex(out_sequence)
    out_sequence[index₁] .= convert.(T₁, out_yiq1)
    out_sequence[index₂] .= convert.(T₂, out_yiq2)
end

(f::MidwayEqualization)(out_sequence::Vector{<:GenericGrayImage}, img_sequence::Vector{<:AbstractArray{<:Color3}}) =
    f(out_sequence, map(img -> of_eltype(Gray, img), img_sequence))

ext: jl


text:
if ! isdefined(:LazyMaps)
    include(joinpath("..", "src", "LazyMaps.jl"))
end

module LazyMapsTests

using LazyMaps, Base.Test

const DEBUG = false

function identicalvalues(A, B, ::Type{Val{:array}})
    for i in eachindex(A, B)
        A[i] !== B[i] && return false
    end
    return true
end

function identicalvalues(A::AbstractArray, B::LazyMap,
                         ::Type{Val{:collection}})
    i = 0
    for b in B
        i += 1
        A[i] !== b && return false
    end
    return true
end

a = randn(11)
b = randn(5,6,7)
c = view(b, :, 2:3, 3:6)
d = Dict("A" => 33, "B" => 100, "C" => pi)

@testset "LazyMaps" begin

    @testset "arrays" begin
        @test identicalvalues(cos.(a), lazymap(cos, a), Val{:array})
        @test identicalvalues(cos.(b), lazymap(cos, b), Val{:array})
        @test identicalvalues(cos.(c), lazymap(cos, c), Val{:array})
    end

    @testset "collections" begin
        @test identicalvalues(sin.(a), lazymap(sin, a), Val{:collection})
        @test identicalvalues(sin.(b), lazymap(sin, b), Val{:collection})
        @test identicalvalues(sin.(c), lazymap(sin, c), Val{:collection})
    end

    #@testset "dictionaries" begin
    #    @test identicalvalues([repr(pair) for pair in d],
    #                          lazymap(repr, d), Val{:collection})
    #end
end

end # module


text:
function stretch_linearly(x, A, B, a, b)
    (x-A) * ((b-a) / (B-A)) + a
end

function extract_string(buffer)
    first_nul = findfirst(isequal('\0'), buffer) - 1
    buffer[1:first_nul]
end


text:
using JLD2, NetCDF, Statistics, LinearAlgebra, Plots
include("kernels.jl")
include("gaussian_process.jl")
include("../les/get_les_data.jl")
include("../les/custom_avg.jl")

using Plots; pyplot()

filename = "general_strat_16_profiles.jld2"
data = get_les_data(filename);

V = data.T
t = data.t
Nz, Nt = size(V)

D = 16 # gridpoints
zavg = custom_avg(data.z, D) # compress z vector to D values
vavg = [custom_avg(V[:,j], D) for j in 1:Nt] # compress variable array to D values per time
x = vavg[1:(Nt-1)] # (v₀, v₁, ... ,v_(Nt-1)) (Nt-1)-length array of D-length inputs
y = vavg[2:Nt]     # (v₁, v₂, ... ,v_Nt    ) (Nt-1)-length array of D-length targets

# reserve 25% of data for training, but across the entire time interval
total_set = 1:(Nt-1)
training_set = 1:4:(Nt-1)
verification_set = setdiff(total_set, training_set)

n = length(training_set) # num. training points
x_train = x[training_set]
y_train = y[training_set]

# function get_mean_loss(kernel_obj)
#     𝒢 = construct_gpr(x_train, y_train, kernel_obj)
#     return mean_log_marginal_loss(y_train,𝒢)
# end

# get_mean_loss(SquaredExponentialKernelI(1000.0,1.0))
# get_mean_loss(Matern12I(1000.0,1.0))
# get_mean_loss(Matern32I(1000.0,1.0))
# get_mean_loss(Matern52I(1000.0,1.0))

# σs = (1.0:0.1:2)*1.0
# γs = (500:20:1000)*1.0
# Z = [get_mean_loss(SquaredExponentialKernelI(γ,σ)) for γ in γs, σ in σs]
# hplot_exponential = surface(σs,γs,Z,
#                     xlabel="signal variance",
#                     ylabel="squared length scale",
#                     title="log marginal likelihood")

"""
surface plot- hyperparameter landscape
"""
function get_mean_loss(γ,σ)
    kernel = SquaredExponentialKernelI(γ,σ)
    𝒢 = construct_gpr(x_train, y_train, kernel)
    return mean_log_marginal_loss(y_train,𝒢)
end

σs = (0.1:0.1:2)*1.0
γs = (1:5:1000)*1.0
plot(γs,σs,get_mean_loss,st=:surface,c=cgrad([:red,:blue]),camera=(-30,30))

"""
surface plot- hyperparameter landscape
"""
Nt = length(data.t)
set = 1:(Nt-2)
gpr_prediction = similar(y[total_set])
gpr_prediction[1] = x[1] # starting
Nt = length(y[total_set])

# function get_mean_max_error(γ,σ)
#     kernel = SquaredExponentialKernelI(γ,σ)
#     𝒢 = construct_gpr(x_train, y_train, kernel)
#
#
#     for i in 1:(Nt-2)
#         gpr_prediction[i+1] = prediction([gpr_prediction[i]], 𝒢)
#     end
#
#     maxes = zeros(Nt-1)
#     for i in 1:500
#         exact    = y[i+1]
#         predi    = gpr_prediction[i+1]
#         # println(exact)
#         # println(predi)
#         maxes[i] = maximum((exact - predi).^2)
#     end
#
#     return maximum(maxes)
# end
#
# σs = (1.0:0.1:1.5)*1.0
# γs = (500:50:1000)*1.0
# plot(γs,σs,get_mean_max_error,st=:surface,c=cgrad([:red,:blue]),camera=(-30,30))

max_stars_repo_name: adelinehillier/LearnConvection

max_issues_repo_name: adelinehillier/LearnConvection

max_forks_repo_name: adelinehillier/LearnConvection


text:
export frac_iso, is_frac_iso

"""
    frac_iso(G,H)

Find a doubly stochastic matrix `S` so that `AS=SB` where `A` and `B`
are the adjacency matrices of `G` and `H` respectively.
"""
function frac_iso(G::SimpleGraph, H::SimpleGraph)
    if deg(G) != deg(H) || degdeg(G) != degdeg(H) # quick basic check before LP
        error("The graphs are not fractionally isomorphic")
    end
    A = adjacency(G)
    B = adjacency(H)
    return frac_iso(A, B)
end

function frac_iso(A::Matrix, B::Matrix)
    n, c = size(A)
    nn, cc = size(B)
    @assert n == c && nn == cc "Matrices must be square"
    @assert n == nn "Matrices must have the same size"

    m = Model(get_solver())

    @variable(m, S[1:n, 1:n])

    for i = 1:n
        for k = 1:n
            @constraint(
                m,
                sum(A[i, j] * S[j, k] for j = 1:n) == sum(S[i, j] * B[j, k] for j = 1:n)
            )
        end
        @constraint(m, sum(S[i, j] for j = 1:n) == 1)
        @constraint(m, sum(S[j, i] for j = 1:n) == 1)
    end

    for i = 1:n
        for j = 1:n
            @constraint(m, S[i, j] >= 0)
        end
    end

    optimize!(m)
    status = Int(termination_status(m))

    if status != 1
        error("The graphs are not fractionally isomorphic")
    end

    return value.(S)

end

"""
    is_frac_iso(G,H)
Test if two graphs are fractionally isomorphic.
"""
function is_frac_iso(G::SimpleGraph, H::SimpleGraph)::Bool
    try
        f = frac_iso(G, H)
        return true
    catch
        return false
    end
    false
end

max_stars_repo_name: scheinerman/SimpleGraphAlgorithms

max_issues_repo_name: scheinerman/SimpleGraphAlgorithms

max_forks_repo_name: scheinerman/SimpleGraphAlgorithms


text:
export default_if_empty

import Base: show

"""
    default_if_empty(value::T)

Creates a `default_if_empty` operator, which emits a given value if the source Observable completes
without emitting any next value, otherwise mirrors the source Observable.

```jldoctest
using Rocket

source = completed(Int) |> default_if_empty(0)

subscribe!(source, logger())
;

# output
[LogActor] Data: 0
[LogActor] Completed
```
See also: [`AbstractOperator`](@ref), [`InferableOperator`](@ref), [`logger`](@ref), [`map`](@ref)
"""
default_if_empty(value::T) where T = DefaultIfEmptyOperator{T}(value)

struct DefaultIfEmptyOperator{T} <: InferableOperator
    value :: T
end

operator_right(operator::DefaultIfEmptyOperator{T}, ::Type{L}) where { L, T } = Union{L, T}

function on_call!(::Type{L}, ::Type{Union{L, T}}, operator::DefaultIfEmptyOperator{T}, source) where { L, T }
    return proxy(Union{L, T}, source, DefaultIfEmptyProxy{Union{L, T}}(convert(Union{L, T}, operator.value)))
end

struct DefaultIfEmptyProxy{L} <: ActorProxy
    default :: L
end

actor_proxy!(::Type, proxy::DefaultIfEmptyProxy{L}, actor::A) where { L, A } = DefaultIfEmptyActor{L, A}(actor, false, proxy.default)

mutable struct DefaultIfEmptyActor{L, A} <: Actor{L}
    actor      :: A
    is_emitted :: Bool
    default    :: L
end

function on_next!(actor::DefaultIfEmptyActor{L}, data::L) where L
    actor.is_emitted = true
    next!(actor.actor, data)
end

function on_error!(actor::DefaultIfEmptyActor, err)
    actor.is_emitted = true
    error!(actor.actor, err)
end

function on_complete!(actor::DefaultIfEmptyActor)
    if !actor.is_emitted
        next!(actor.actor, actor.default)
    end
    complete!(actor.actor)
end

Base.show(io::IO, ::DefaultIfEmptyOperator{T}) where T = print(io, "DefaultIfEmptyOperator($T)")
Base.show(io::IO, ::DefaultIfEmptyProxy{L})    where L = print(io, "DefaultIfEmptyProxy($L)")
Base.show(io::IO, ::DefaultIfEmptyActor{L})    where L = print(io, "DefaultIfEmptyActor($L)")


text:
using Spec
using Base.Threads: @spawn
export re!, re, re_all!, Replica, ReplicaAlg

using InferenceBase

"Replica Exchange (Parallel Tempering)"
struct ReplicaAlg end
const Replica = ReplicaAlg()

"swap `v[i]` and `v[j]`"
function swap!(v, i, j)
  temp = v[i]
  v[i] = v[j]
  v[j] = temp
end

struct PreSwap end
struct PostSwap end

"Swap adjacent chains"
function swap_contexts!(rng, logenergys, states)
  for i in length(states):-1:2
    j = i - 1

    # Evaluate energy of state i at temperature i
    @show E_i_i = logenergys[i](states[i])
    @show E_i_j = logenergys[i](states[j])
    @show E_j_i = logenergys[j](states[i])
    @show E_j_j = logenergys[j](states[j])

    probaccept = (E_i_j + E_j_i) - (E_i_i + E_j_j)

    doswap = log(rand(rng)) < probaccept
    if doswap
      println("did swap", i, j)
      swap!(states, i, j)
    else
      println("did not swap", i, j)
    end
  end
end

"`every(m)` -- `f` where `f(i) is true every i steps"
every(m) = i -> i % m == 0

"""
  `re!(rng,
        logenergys,
        samples_per_swap,
        num_swaps,
        samples,
        states,
        sim_chain_keep_n,
        sim_chain_keep_last = last ∘ sim_chain_keep_n,
        swap_contexts! = swap_contexts!)`

Replica Exchange Markov Chain Monte Carlo

Replica exhange:
- There are `n` different contexts, where a different context means a different target density
- There arae `A` different algorithms.  These may be actually different algorithms
  such as HMC vs SSMH, or different parameterizations of the same algorithm
- The user provides some subset of the cross product, i.e. `(ctx1, alg1), (ctx2, alg2)`
- runs `nreplicas = length(algs)` MCMC chains in parallel
- Each alg is run in a different context.
  - The most common form of a context is a temperature
- We assume there is a ground context which is the true model we wish to sample from
-

# Arguments
- `rng`: AbstractRng used to sample proposals in MH loop
- `logenergys`: collection of `n` logenergys to sample from
- `samples_per_swap` : number of samples drawn between each exchange
- `num_swaps`: number of swaps (num_samples = `samples_per_swap` * `num_swaps`)
- `states`: initial states
- `sim_chain_keep_n` : algorithm to take `samples_per_swap` mcmc steps and return all n
  - should support `sim_chain_keep_n(logenergy, init_state, n_samples)``
- `sim_chain_keep_1` : algorithm to take `samples_per_swap` mcmc steps and return only last
- `samples`: Vector of samples to mutate
- `swap_contexts!` function that swaps contexts every `num_swaps`, i.e. does exchange

# Returns
- `n` samples drawn from ground state
"""
function re!(rng,
             logenergys,
             samples_per_swap,
             num_swaps,
             states,
             samples,
             sim_chain_keep_n,
             sim_chain_keep_last = last ∘ sim_chain_keep_n,
             swap_contexts! = swap_contexts!,
             cb = tonothing)
  nreplicas = length(states)

  GROUNDID = 1
  for num_swap = 1:num_swaps
    for i = 1:nreplicas
      @show i, :, logenergys[i](states[GROUNDID])
      if i == GROUNDID
        # @show typeof(logenergys)
        # @assert false
        # If we're ground state, return swap_every samples
        lb = (num_swap - 1) * samples_per_swap + 1
        ub = lb + samples_per_swap - 1
        @inbounds samples[lb:ub] = sim_chain_keep_n(rng,
                                                    logenergys[GROUNDID],
                                                    states[GROUNDID],
                                                    samples_per_swap,
                                                    i)
        @inbounds states[GROUNDID] = samples[ub]
      else
        # If not ground state just return last sample
        @inbounds states[i] = sim_chain_keep_last(rng, logenergys[i], states[i], samples_per_swap, i)
      end
    end
    swap_contexts!(rng, logenergys, states)
  end
  samples
end
# @pre length(logenergys) == length(states) "Must have one density per initial state"

"Non-mutating re!"
re(rng,
   logenergys,
   samples_per_swap,
   num_swaps,
   states,
   sim_chain_keep_n,
   sim_chain_keep_last = last ∘ sim_chain_keep_n,
  swap_contexts! = swap_contexts!,
  cb = tonothing) =
  re!(rng,
      logenergys,
      samples_per_swap,
      num_swaps,
      states,
      Vector{eltype(states)}(undef, num_swaps * samples_per_swap),
      sim_chain_keep_n,
      sim_chain_keep_last,
      swap_contexts!,
      cb)

function re_all!(rng,
                 logenergys,
                 samples_per_swap,
                 num_swaps,
                 states,
                 simulate_n,
                 sim_chain_keep_last,
                 evaluate,
                 samples = [Vector{typeof(states[i])}(undef, num_swaps*samples_per_swap) for i = 1:length(states)])
            #  swap = every(div(n, 10)))
  # @pre length(ctxs) == length(algs)
  nreplicas = length(states)
  if length(logenergys) != length(states)
    error("length(logenergys) != length(states)")
  end

  GROUNDID = 1
  for num_swap = 1:num_swaps
    for i = 1:nreplicas
      # If we're ground state, return swap_every samples
      lb = (num_swap - 1) * samples_per_swap + 1
      ub = lb + samples_per_swap - 1
      samples[i][lb:ub] = simulate_n(logenergys[i],
                                               states[i],
                                               samples_per_swap)
      states[i] = samples[i][ub]
    end
    swap_contexts!(rng, logenergys, states, evaluate)
  end
  samples
end

# Questions:
# Can there be more than one ground context?
# Is it true that from the ground context we need many samples and from
# everythinge else we need 1 sample?
# -- No, we only need samples from the ground context, we need scores
# From the others
# We need to be able to compute:
  # In some context i, density of some other point

# How should burn in / thinning work?
# Choice of representation:
## - alg(ctx, init_state)
## - apply(alg, ctx, init_state)
## - userprovidedfunc(alg, ctx, init_state)
## Do we need both Algs and Ctxs??

max_stars_repo_name: zenna/expect

max_issues_repo_name: zenna/expect

max_forks_repo_name: zenna/expect


text:
export
    LogisticCircuit,
    LogisticLeafNode, LogisticInnerNode,
    LogisticLiteralNode, Logistic⋀Node, Logistic⋁Node,
    num_classes, num_parameters_per_class

#####################
# Infrastructure for logistic circuit nodes
#####################

"Root of the logistic circuit node hierarchy"
abstract type LogisticCircuit <: LogicCircuit end

"""
A logistic leaf node
"""
abstract type LogisticLeafNode <: LogisticCircuit end

"""
A logistic inner node
"""
abstract type LogisticInnerNode <: LogisticCircuit end

"""
A logistic literal node
"""
mutable struct LogisticLiteralNode <: LogisticLeafNode
    literal::Lit
    data
    counter::UInt32
    LogisticLiteralNode(l) = begin
        new(l, nothing, 0)
    end
end

"""
A logistic conjunction node (And node)
"""
mutable struct Logistic⋀Node <: LogisticInnerNode
    children::Vector{<:LogisticCircuit}
    data
    counter::UInt32
    Logistic⋀Node(children) = begin
        new(convert(Vector{LogisticCircuit}, children), nothing, 0)
    end
end

"""
A logistic disjunction node (Or node)
"""
mutable struct Logistic⋁Node <: LogisticInnerNode
    children::Vector{<:LogisticCircuit}
    thetas::Matrix{Float32}
    data
    counter::UInt32
    Logistic⋁Node(children, class::Int) = begin
        new(convert(Vector{LogisticCircuit}, children), init_array(Float32, length(children), class), nothing, 0)
    end
end

#####################
# traits
#####################

import LogicCircuits.GateType # make available for extension
@inline GateType(::Type{<:LogisticLiteralNode}) = LiteralGate()
@inline GateType(::Type{<:Logistic⋀Node}) = ⋀Gate()
@inline GateType(::Type{<:Logistic⋁Node}) = ⋁Gate()

#####################
# methods
#####################

import LogicCircuits: children # make available for extension
@inline children(n::LogisticInnerNode) = n.children
@inline num_classes(n::Logistic⋁Node) = size(n.thetas)[2]

@inline num_parameters(c::LogisticCircuit) = sum(n -> num_children(n) * num_classes(n), ⋁_nodes(c))
@inline num_parameters_per_class(c::LogisticCircuit) = sum(n -> num_children(n), ⋁_nodes(c))

"Get the parameters associated with a or node"
params(n::Logistic⋁Node) = n.thetas

#####################
# constructors and conversions
#####################

function LogisticCircuit(circuit::LogicCircuit, classes::Int)
    f_con(n) = error("Cannot construct a logistic circuit from constant leafs: first smooth and remove unsatisfiable branches.")
    f_lit(n) = compile(LogisticCircuit, literal(n))
    f_a(n, cn) = Logistic⋀Node(cn)
    f_o(n, cn) = Logistic⋁Node(cn, classes)
    foldup_aggregate(circuit, f_con, f_lit, f_a, f_o, LogisticCircuit)
end

compile(::Type{<:LogisticCircuit}, l::Lit) =
    LogisticLiteralNode(l)

function compile(::Type{<:LogisticCircuit}, classes, circuit::LogicCircuit)
    LogisticCircuit(circuit, classes)
end

import LogicCircuits: fully_factorized_circuit #extend

function fully_factorized_circuit(::Type{<:LogisticCircuit}, n::Int; classes::Int)
    ff_logic_circuit = fully_factorized_circuit(PlainLogicCircuit, n)
    compile(LogisticCircuit, classes, ff_logic_circuit)
end

function check_parameter_integrity(circuit::LogisticCircuit)
    for node in or_nodes(circuit)
        @assert all(θ -> !isnan(θ), node.thetas) "There is a NaN in one of the log_probs"
    end
    true
end


text:
@testset "stream" begin

let
    s = CuStream()
    synchronize(s)
    let s2 = CuStream()
        @test s != s2
    end

    let s3 = CuStream(CUDAdrv.STREAM_NON_BLOCKING)
        @test s != s3
    end

    let s4 = CuStream(CUDAdrv.STREAM_DEFAULT | CUDAdrv.STREAM_NON_BLOCKING)
        @test s != s4
    end

    synchronize(CuDefaultStream())
end

end


text:
module InvolutiveMCMC

# using Distributions
# using Random
using InferenceBase: keepall

export imcmc!, imcmc

"`imcmc!(rng, p, Vₓ, p_vₓ, f, x, xs; n)` like `imcmc`, except mutates `xs`"
function imcmc!(rng, p, Vₓ, p_vₓ, f, x, xs; n, keep = keepall)
  logjoint(x, v) = p(x) + p_vₓ(v, x)
  i = 1
  s = 1
  while s <= n
    v = Vₓ(rng, x)
    x_n, v_n = f(x, v)
    ratio = logjoint(x_n, v_n) - logjoint(x, v)
    if log(rand(rng)) < ratio
      x = x_n
    end
    if keep(i)
      @inbounds xs[s] = x
      s += 1
    end
    i += 1
  end
  xs
end

"""
`imcmc(rng, p, Vₓ, f, x; n)`

Involutive MCMC

# Input:
- `rng`: RNG object
- `p(x)`: target logdensity `p(x) -> Real`
- `Vₓ(rng, x)`:  conditional random variable - samples from `v` given `x` where
    `v` is auxiliary variable
- `f(x, v) -> (x', v')`: an involution, i.e., `f(f((x, v))) == (x, v)`
- `x`: initial value of `x`
- `n`: number of samples

# Returns
- Set of samples drawn approimately according to `p`
"""
imcmc(rng, p, Vₓ, p_vₓ, f, x::X; n, keep = keepall) where X =
  imcmc!(rng, p, Vₓ, p_vₓ, f, x, Vector{X}(undef, n); n = n, keep = keep)

end # module


text:
# Inside make.jl
push!(LOAD_PATH,"../src/")
using NotebookToLaTeX
using Documenter
makedocs(
         sitename="NotebookToLaTeX.jl",
         modules=[NotebookToLaTeX],
         pages=[
                "Home" => "index.md"
               ])
deploydocs(;
    repo="github.com/davibarreira/NotebookToLaTeX.jl",
)


text:
# -*- coding: utf-8 -*-
#
# Calculate the derivatives of an image using the Scharr operator
# of kernal size 3x3.
#
# References:
# [1] https://en.wikipedia.org/wiki/Sobel_operator
# [2] http://docs.opencv.org/doc/tutorials/imgproc/imgtrans/sobel_derivatives/sobel_derivatives.html
#
# Aaron LI
# 2015/07/20
#

# Calculate the derivatives of the image using the Scharr operator
# img - input image
# dx  - order of the derivative x
# dy  - order of the derivative y
function scharr(img, dx, dy)
    rows, cols = size(img)
    img_d = float(img)
    (isa(dx, Int) && dx >= 0) || error("dx should be an integer >= 0")
    (isa(dy, Int) && dy >= 0) || error("dy should be an integer >= 0")
    # Scharr operator
    Gy = [-3.0 -10.0 -3.0; 0.0 0.0 0.0; 3.0 10.0 3.0];
    Gx = Gy'
    # calculate the derivatives using convolution
    for i = 1:dx
        img_d = conv2(img_d, Gx)
    end
    for i = 1:dy
        img_d = conv2(img_d, Gy)
    end
    # FIXME: 'conv2' will increase the image size
    rows_d, cols_d = size(img_d)
    return img_d[(div(rows_d-rows, 2)+1):(div(rows_d-rows, 2)+rows), (div(cols_d-cols, 2)+1):(div(cols_d-cols, 2)+cols)]
end

max_stars_repo_name: liweitianux/atoolbox

max_issues_repo_name: liweitianux/atoolbox

max_forks_repo_name: liweitianux/atoolbox


text:
#MDDatasets: Register dataset operations
#-------------------------------------------------------------------------------

#==NOTE
DataF1 cannot represent matrices.  Element-by-element operations will therefore
be the default.  There is not need to use the "." operator versions.
==#

#==NOTE:
	:cummax, :cummin, :cumprod, :cumsum, maxabs, are high-level functions... not sure
	if these should be ported, or new names/uses should be found
==#

#==Support basic math operations
===============================================================================#

#Unary operators
const _operators1 = Symbol[:-, :+, :!]

for op in _operators1; @eval begin #CODEGEN--------------------------------------

#op Index:
Base.$op(i::Index) = Index($op(i.v))

#op DataF1:
Base.$op(d::DataF1) = DataF1(d.x, ($op).(d.y))

#Everything else:
Base.$op(d::DataMD) = broadcastMD(CAST_BASEOP1, Base.$op, d)

end; end #CODEGEN---------------------------------------------------------------

const _operators2 = Symbol[:-, :+, :/, :*, :^, :>, :<, :>=, :<=, :!=, :(==)]

for op in _operators2; @eval begin #CODEGEN--------------------------------------

#Index op Index:
Base.$op(i1::Index, i2::Index) = Index(($op).(i1.v, i2.v))

#DataF1 op DataF1:
Base.$op(d1::DataF1, d2::DataF1) = apply(Base.$op, d1, d2)

#DataF1 op Number:
Base.$op(d::DataF1, n::Number) = DataF1(d.x, ($op).(d.y, n))

#Number op DataF1:
Base.$op(n::Number, d::DataF1) = DataF1(d.x, ($op).(n, d.y))

#Everything else:
Base.$op(d1::DataMD, d2::DataMD) = broadcastMD(CAST_BASEOP2, Base.$op, d1, d2)
Base.$op(d1::Number, d2::DataMD) = broadcastMD(CAST_BASEOP2, Base.$op, d1, d2)
Base.$op(d1::DataMD, d2::Number) = broadcastMD(CAST_BASEOP2, Base.$op, d1, d2)

end; end #CODEGEN---------------------------------------------------------------

#==Support for 1-argument Base.functions
===============================================================================#

#==TODO:
Check out statistical stuff

isempty, isfinite, isinf, isinteger, isnan, isposdef, isreal

fft, etc

Mapping ??
map, mapreduce, mapreducedim, mapslices

#Number converters
#Bool(): map(Bool, x)
round
...

#other
#eachindex
rand
==#

#1-argument Base.functions:
const _basefn1 = [:(Base.$fn) for fn in [
	:abs, :abs2, :angle,
	:imag, :real, :exponent,
	:exp, :exp2, :exp10, :expm1,
	:log, :log10, :log1p, :log2,
	:ceil, :floor,
	:asin, :asind, :asinh, :acos, :acosd, :acosh,
	:atan, :atand, :atanh, :acot, :acotd, :acoth,
	:asec, :asecd, :asech, :acsc, :acscd, :acsch,
	:sin, :sind, :sinh, :cos, :cosd, :cosh,
	:tan, :tand, :tanh, :cot, :cotd, :coth,
	:sec, :secd, :sech, :csc, :cscd, :csch,
	:sinpi, :cospi,
	:sinc, :cosc, #cosc: d(sinc)/dx
	:deg2rad, :rad2deg,

	#I think broadcast operations will allow user to specify low & high as arguments.
	#so: really a function of 1 DataMD argument - and 2 "parameters":
	:clamp, #Needs low and high.
]]

#Custom functions of 1 argument that operate on base types:
const _custbasefn1 = [
	#Also function of 1 DataMD argument, but 2 other "parameters":
	:sanitize,
]

for fn in vcat(_basefn1,_custbasefn1); @eval begin #CODEGEN---------------------

#fn(DataF1)
$fn(d::DataF1, args...; kwargs...) = DataF1(d.x, ($fn).(d.y, args...; kwargs...))

#Everything else:
$fn(d::DataMD, args...; kwargs...) = broadcastMD(CAST_BASEOP1, $fn, d, args...; kwargs...)

end; end #CODEGEN---------------------------------------------------------------

#zeros & one function... must use "fill" on functions:
const _onezerofn1 = [:(Base.$fn) for fn in [
	:zeros, :ones
]]

#Define old-ish "ones"/"zeros" syntax to simplify definition of ones/zeros:
Base.ones(::MDCUST, a) = fill(Base.one(eltype(a)), size(a))
Base.zeros(::MDCUST, a) = fill(Base.zero(eltype(a)), size(a))

for fn  in vcat(_onezerofn1); @eval begin #CODEGEN---------------------

#fn(DataF1)
$fn(d::DataF1) = DataF1(d.x, $fn(MDCUST(), d.y))

#Everything else:
$fn(d::DataMD) = broadcastMD(CAST_BASEOP2, $fn, MDCUST(), d)

end; end #CODEGEN---------------------------------------------------------------

#==Support for 2-argument Base.functions
===============================================================================#

const _basefn2 = [:(Base.$fn) for fn in [
	:max, :min,
	:atan, :hypot,
]]

for fn in _basefn2; @eval begin #CODEGEN----------------------------------------

#fn(DataF1, DataF1):
$fn(d1::DataF1, d2::DataF1) = apply($fn, d1, d2)

#Everything else:
$fn(d1::DataMD, d2::DataMD) = broadcastMD(CAST_BASEOP2, $fn, d1, d2)
$fn(d1::Number, d2::DataMD) = broadcastMD(CAST_BASEOP2, $fn, d1, d2)
$fn(d1::DataMD, d2::Number) = broadcastMD(CAST_BASEOP2, $fn, d1, d2)

end; end #CODEGEN---------------------------------------------------------------

#==Support reducing/collpasing Base.functions:
===============================================================================#
const _baseredfn1 = [:(Base.$fn) for fn in [
	:maximum, :minimum,
	:prod, :sum,
]]

for fn in _baseredfn1; @eval begin #CODEGEN-------------------------------------

#fn(DataF1):
$fn(d::DataF1) = ($fn)(d.y)

#Everything else:
$fn(d::DataMD) = broadcastMD(CAST_BASEOPRED1, $fn, d)

end; end #CODEGEN---------------------------------------------------------------

#==Support reducing/collpasing Statistics.functions:
===============================================================================#
const _statredfn1 = [:(Statistics.$fn) for fn in [
	:mean, :median, :middle,
	:std, :var,
]]

for fn in _statredfn1; @eval begin #CODEGEN-------------------------------------

#fn(DataF1):
$fn(d::DataF1) = ($fn)(d.y)

#Everything else:
$fn(d::DataMD) = broadcastMD(CAST_BASEOPRED1, ($fn), d)

end; end #CODEGEN---------------------------------------------------------------

#==Support 1-argument functions of DataF1
===============================================================================#
const _custfn1 = [
	:clip, :xval, :sample,
	:delta, :xshift, :xscale,
	:deriv, :iinteg,
	:xcross, :measperiod, :measfreq, :measduty,
	:measck2q, :measrise, :measfall,
]

for fn in _custfn1; @eval begin #CODEGEN---------------------------------------

$fn(d::DataMD, args...; kwargs...) = broadcastMD(CAST_MD1, $fn, d, args...; kwargs...)

end; end #CODEGEN---------------------------------------------------------------

#Support reducing/collpasing funcitons of DataF1
#-------------------------------------------------------------------------------
const _custredfn1 = [
	:xcross1, :integ, :xmin, :xmax, :value,
]

for fn in _custredfn1; @eval begin #CODEGEN-------------------------------------

$fn(d::DataMD, args...; kwargs...) = broadcastMD(CAST_MDRED1, $fn, d, args...; kwargs...)

end; end #CODEGEN---------------------------------------------------------------

#==Support 2-argument functions of DataF1
===============================================================================#
const _custfn2 = [
	:yvsx,
	:measdelay,
	:measck2q,
]

for fn in _custfn2; @eval begin #CODEGEN----------------------------------------

$fn(d1::DataMD, d2::DataMD, args...; kwargs...) =
	broadcastMD(CAST_MD2, $fn, d1, d2, args...; kwargs...)

end; end #CODEGEN---------------------------------------------------------------

#==Special cases
===============================================================================#
#2nd argument does not have to be DataMD:
#-------------------------------------------------------------------------------
ycross(d1::DataMD, d2, args...; kwargs...) =
	broadcastMD(CastType(DataF1, 1, Number, 2), ycross, d1, d2, args...; kwargs...)
ycross1(d1::DataMD, d2, args...; kwargs...) =
	broadcastMD(CastTypeRed(DataF1, 1, Number, 2), ycross1, d1, d2, args...; kwargs...)

#First argument is ::DS{}
#-------------------------------------------------------------------------------
measdelay(ds::DS, d1::DataMD, d2::DataMD, args...; kwargs...) =
	broadcastMD(CastType(DataF1, 2, DataF1, 3), measdelay, ds, d1, d2, args...; kwargs...)
xcross(ds::DS, d::DataMD, args...; kwargs...) =
	broadcastMD(CastType(DataF1, 2), xcross, ds, d, args...; kwargs...)
xcross1(ds::DS, d::DataMD, args...; kwargs...) =
	broadcastMD(CastTypeRed(DataF1, 2), xcross1, ds, d, args...; kwargs...)

#Last line


text:
using RadeonProRender, GeometryBasics, Colors
const RPR = RadeonProRender
using Colors: N0f8

# Create a scene
context = RPR.Context()
scene = RPR.Scene(context)
matsys = RPR.MaterialSystem(context, 0)
camera = RPR.Camera(context)
lookat!(camera, Vec3f(3), Vec3f(0), Vec3f(0, 0, 1))
set!(scene, camera)

# create some x,y,z data
function create_mesh(DN = 512, dphi=pi/200.0f0, dtheta=dphi)
    phi = 0.0f0:dphi:(pi + dphi * 1.5f0)
    theta = (0.0f0:dtheta:(2.0f0 * pi + dtheta * 1.5f0))'

    m0 = 4.0f0;
    m1 = 3.0f0;
    m2 = 2.0f0;
    m3 = 3.0f0;
    m4 = 6.0f0;
    m5 = 2.0f0;
    m6 = 6.0f0;
    m7 = 4.0f0;
    a = @. sin(m0 * phi) ^ m1
    b = @. cos(m2 * phi) ^ m3
    c = @. sin(m4 * theta) ^ m5
    d = @. cos(m6 * theta) ^ m7
    r = @. a + b + c + d;

    x = @. r * sin(phi) * cos(theta)
    y = @. r * cos(phi)
    z = @. r * sin(phi) * sin(theta)

    xyz = Point3f[Point3f(x[i, j], y[i, j], z[i, j]) for i in 1:size(x, 1), j in 1:size(x, 2)]
    r = Tesselation(Rect2f((0, 0), (1, 1)), size(x))
    # decomposing a rectangle into uv and triangles is what we need to map the z coordinates on
    # since the xyz data assumes the coordinates to have the same neighouring relations
    # like a grid
    faces = decompose(GLTriangleFace, r)
    uv = decompose_uv(r)
    # with this we can beuild a mesh
    return GeometryBasics.Mesh(meta(vec(xyz), uv=uv), faces), xyz
end

raw_mesh, xyz = create_mesh()
surfacemesh = RPR.Shape(context, raw_mesh)
push!(scene, surfacemesh)

# create layered material
tex = RPR.MaterialNode(matsys, RPR.RPR_MATERIAL_NODE_IMAGE_TEXTURE)
base = RPR.MaterialNode(matsys, RPR.RPR_MATERIAL_NODE_DIFFUSE)
top = RPR.MaterialNode(matsys, RPR.RPR_MATERIAL_NODE_REFLECTION)

# Diffuse color
set!(base, RPR.RPR_MATERIAL_INPUT_COLOR, tex);
set!(base, RPR.RPR_MATERIAL_INPUT_ROUGHNESS, 0.12, 0.0, 0.0, 1.0)
set!(top, RPR.RPR_MATERIAL_INPUT_COLOR, 0.1, 0.3, 0.9, 1.0)

# Create layered shader
layered = RPR.layeredshader(matsys, base, top)

colorim = map(xyz) do xyz
    return RGBA{N0f8}(clamp(abs(xyz[1]), 0, 1), clamp(abs(xyz[2]), 0, 1), clamp(abs(xyz[3]), 0, 1), 1.0)
end

color = RPR.Image(context, colorim)
set!(tex, RPR.RPR_MATERIAL_INPUT_DATA, color)

set!(surfacemesh, layered)
set!(context, scene)

ibl = RPR.EnvironmentLight(context)
image_path = joinpath(@__DIR__, "studio026.exr")
set!(ibl, context, image_path)
set!(scene, ibl)

set_standard_tonemapping!(context)

fb_size = (800, 600)
frame_buffer = RPR.FrameBuffer(context, RGBA, fb_size)
frame_bufferSolved = RPR.FrameBuffer(context, RGBA, fb_size)
frame_buffer_post = RPR.FrameBuffer(context, RGBA, fb_size)
set!(context, RPR.RPR_AOV_COLOR, frame_buffer)

begin
    clear!(frame_buffer)
    clear!(frame_buffer_post)
    RPR.rprContextSetParameterByKey1u(context, RPR.RPR_CONTEXT_ITERATIONS, 100)
    RPR.render(context)
    RPR.rprContextResolveFrameBuffer(context, frame_buffer, frame_bufferSolved, true)
    RPR.save(frame_bufferSolved, "surface_mesh.png")
    RPR.rprContextResolveFrameBuffer(context, frame_buffer, frame_buffer_post, false)
    RPR.save(frame_buffer_post, "surface_mesh_postprocess.png")
end


text:
# ENV["SHAPE_ENCODING"] = "UTF-8"
function writeORG(table, fn::AbstractString; layer_name::AbstractString="data",
    geom_column::Symbol=:geom, crs::GFT.GeoFormat=GFT.EPSG(4326), driver::Union{Nothing,AbstractString}=nothing)

    rows = Tables.rows(table)
    sch = Tables.schema(rows)

    # Geom type can't be inferred from here
    geom_type = AG.getgeomtype(rows[1][geom_column])

    # Find driver
    _, extension = splitext(fn)
    if extension in keys(drivermapping)
        driver = AG.getdriver(drivermapping[extension])
    elseif driver !== nothing
        driver = AG.getdriver(driver)
    else
        error("Couldn't determine driver for $extension. Please provide one of $(keys(drivermapping))")
    end

    # Figure out attributes
    fields = Vector{Tuple{Symbol,DataType}}()
    for (name, type) in zip(sch.names, sch.types)
        if type != AG.IGeometry && name != geom_column
            push!(fields, (Symbol(name), type))
        end
    end

    options = ["ENCODING=UTF-8", "SHAPE_ENCODING=UTF-8"];
    AG.create(
        fn,
        driver=driver,
        options=options
    ) do ds
        AG.createlayer(
            name=layer_name,
            geom=geom_type,
            spatialref=AG.importCRS(crs),
            options=options
        ) do layer
            for (name, type) in fields
                AG.addfielddefn!(layer, String(name), fieldmapping[type])
            end
            for row in rows
                AG.createfeature(layer) do feature
                    AG.setgeom!(feature, getproperty(row, geom_column))
                    for (name, type) in fields
                        AG.setfield!(feature, AG.findfieldindex(feature, name), getproperty(row, name))
                    end
                end
            end
            AG.copy(layer, dataset=ds, name=layer_name, options=options)
        end
    end
    fn
end

export writeORG


text:
export make_engine_editor, show_engine_editor

using SQLite
using DataFrames
using Gtk
using Rockets

"""
    show_engine_editor

Show GTK window with engine data by loading database in `data/rockets.db`.
"""
function show_engine_editor()
    enginesdata = load_engines()
    win = make_engine_editor(enginesdata)
    showall(win)
    return win
end

"""
    make_engine_editor(enginesdata::DataFrame)

Show GTK window with engine data stored in `enginesdata`.
"""
function make_engine_editor(enginesdata::DataFrame)
    path =     path = joinpath(datadir, "engineui.glade")
    builder = GtkBuilder(filename=path);
    w = builder["enginewin"];

    enginelistview = builder["enginelist"];

    # Will get some dummy data, we will immediately replace
    enginestore = builder["enginestore"];
    updatestore!(enginestore, enginesdata)

    enginename_entry = builder["engine_name"];
    mass_entry = builder["mass"];
    thrust_sl_entry = builder["thrust_sl"];
    thrust_vac_entry = builder["thrust_vac"];
    isp_sl_entry = builder["isp_sl"];
    isp_vac_entry = builder["isp_vac"];

    addbtn = builder["add"];
    removebtn = builder["remove"];
    reloadbtn = builder["reload"];
    savebtn = builder["save"];

    # Deal with selection and model update
    selection = GAccessor.selection(enginelistview)
    chid = signal_connect(selection, "changed") do widget
        # selection = GAccessor.selection(enginelistview)
        if hasselection(selection)
            current = selected(selection)
            enginename = enginestore[current, 1]

            i = findfirst(==(enginename), enginesdata.Engine)

            if !isnothing(i)
                set_gtk_property!(enginename_entry, :text, enginesdata.Engine[i])
                set_gtk_property!(mass_entry, :text, enginesdata.Mass[i])
                set_gtk_property!(thrust_sl_entry, :text, enginesdata.ThrustSL[i])
                set_gtk_property!(thrust_vac_entry, :text, enginesdata.ThrustVac[i])
                set_gtk_property!(isp_sl_entry, :text, enginesdata.IspSL[i])
                set_gtk_property!(isp_vac_entry, :text, enginesdata.IspVac[i])
            end
        end
    end

    # Adding row. Note we keep ID (addid) in case we copy paste this into REPL
    # Allows you to disconnect: signal_handler_disconnect(addbtn, addid)
    addid = signal_connect(addbtn, "clicked") do widget
        name = get_gtk_property(enginename_entry, :text, String)

        # Cannot get text as numbers. Need to make that conversion ourselves
        enginename = get_gtk_property(enginename_entry, :text, String)
        masstxt = get_gtk_property(mass_entry, :text, String)
        thrust_sltxt = get_gtk_property(thrust_sl_entry, :text, String)
        thrust_vactxt = get_gtk_property(thrust_vac_entry, :text, String)
        isp_sltxt = get_gtk_property(isp_sl_entry, :text, String)
        isp_vactxt = get_gtk_property(isp_vac_entry, :text, String)

        mass = parse(Float64, masstxt)
        thrust_sl = parse(Float64, thrust_sltxt)
        thrust_vac = parse(Float64, thrust_vactxt)
        isp_sl = parse(Float64, isp_sltxt)
        isp_vac = parse(Float64, isp_vactxt)

        push!(enginesdata, (
            Engine = enginename,
            Mass = mass,
            ThrustSL = thrust_sl,
            ThrustVac = thrust_vac,
            IspSL = isp_sl,
            IspVac = isp_vac
        ))

        push!(enginestore, (name,))
    end

    rmid = signal_connect(removebtn, "clicked") do widget
        if hasselection(selection)
            current = selected(selection)
            enginename = enginestore[current, 1]

            i = findfirst(==(enginename), enginesdata.Engine)

            if !isnothing(i)
                # Remove from DataFrame
                deleteat!(enginesdata, i)

                # Put on this on main GUI task. I think.
                # GUI stuff cannot be run on other tasks
                Gtk.@sigatom begin
                    deleteat!(enginestore, current)

                    set_gtk_property!(enginename_entry, :text, "")
                    set_gtk_property!(mass_entry, :text, "")
                    set_gtk_property!(thrust_sl_entry, :text, "")
                    set_gtk_property!(thrust_vac_entry, :text, "")
                    set_gtk_property!(isp_sl_entry, :text, "")
                    set_gtk_property!(isp_vac_entry, :text, "")

                    display(enginesdata)
                end
            end
        end
    end

    # Reloading
    reloadid = signal_connect(reloadbtn, "clicked") do widget
        Gtk.@sigatom begin
            updatestore!(enginestore, enginesdata)
        end
    end

    saveid = signal_connect(savebtn, "clicked") do widget
        # This is just to demonstrate different ways of working
        # It is not ideal to  delete a whole table and then add all the data
        path = joinpath(datadir, "rockets.db")
        db = SQLite.DB(path)

        # Remove table Engines. Don't throw exception if the table
        # is not there (ifexists=true).
        stmt = SQLite.drop!(db, "Engines", ifexists = true)

        # Create Engines table from scratch by loading data
        # from DataFrame
        SQLite.load!(enginesdata, db, "Engines")

        DBInterface.close!(db)
    end

    showall(w)
    return w
end

"Put engine names in `enginesdata` into  `enginestore` used by tree view"
function updatestore!(enginestore::GtkListStore, enginesdata::DataFrame)
    empty!(enginestore)
    for name in enginesdata.Engine
       push!(enginestore, (name,))
    end
end


text:
"""
Data structure describing the FEM degrees of freedom.

isfree[nd]     is true if nd is free, false is nd is fixed.
freenode[k]    is the gmsh node number of the kth free node.
fixednode[k]   is the gmsh node number of the kth fixed node.
node2free[nd]  = k if nd = freenode[k], -1 otherwise.
node2fixed[nd] = k if nd = fixednode[k], -1 otherwise.
"""
struct DoF
    isfree     :: Vector{Bool}
    freenode   :: Vector{Integer}
    fixednode  :: Vector{Integer}
    node2free  :: Vector{Integer}
    node2fixed :: Vector{Integer}
end

struct VariationalProblem
    mesh         :: Mesh
    dof          :: DoF
    essential_bc :: Vector{String}
    bilin_form   :: Vector{Any}
    lin_functnl  :: Vector{Any}
    ufixed       :: Vector{Float64}
end

struct EigenProblem
    mesh           :: Mesh
    dof            :: DoF
    essential_bc   :: Vector{String}
    LHS_bilin_form :: Vector{Any}
    RHS_bilin_form :: Vector{Any}
end

function VariationalProblem(mesh::Mesh,
                            essential_bc::Array{String})
    for name in essential_bc
        if !(name in keys(mesh.nodes_of))
            error("$name: unknown physical name")
        end
    end
    dof = degrees_of_freedom(mesh, essential_bc)
    bilin_form   = Any[]
    lin_functnl  = Any[]
    nofixed = length(dof.fixednode)
    # Zero boundary condition by default.
    ufixed = zeros(nofixed)
    return VariationalProblem(mesh, dof, essential_bc,
                              bilin_form, lin_functnl, ufixed)
end

# Use in case of pure Neumann bc.
VariationalProblem(mesh::Mesh) = VariationalProblem(mesh, String[])

function EigenProblem(mesh::Mesh, essential_bc::Array{String})
    for name in essential_bc
        if !(name in keys(mesh.nodes_of))
            error("$name: unknown physical name")
        end
    end
    dof = degrees_of_freedom(mesh, essential_bc)
    LHS_bilin_form   = Any[]
    RHS_bilin_form   = Any[]
    nofixed = length(dof.fixednode)
    return EigenProblem(mesh, dof, essential_bc,
                        LHS_bilin_form, RHS_bilin_form)
end

function assign_bdry_vals!(vp::VariationalProblem,
                           name::String, g::Float64)
    if !(name in vp.essential_bc)
        error("$name: not listed in essential_bc")
    end
    for nd in vp.mesh.nodes_of[name]
        i = vp.dof.node2fixed[nd]
        vp.ufixed[i] = g
    end
end

function assign_bdry_vals!(vp::VariationalProblem,
                           name::String, g::Function)
    if !(name in vp.essential_bc)
        error("$name: not listed in essential_bc")
    end
    for nd in vp.mesh.nodes_of[name]
        i = vp.dof.node2fixed[nd]
        x = vp.mesh.coord[:,nd]
        vp.ufixed[i] = g(x)
    end
end

"""
Returns the DoF object for the given mesh and essential boundary
conditions.
"""
function degrees_of_freedom(mesh::Mesh, essential_bc::Array{String})
    nonodes = size(mesh.coord, 2)
    isfree = Array{Bool}(undef, nonodes)
    fill!(isfree, true)
    for name in essential_bc
        if !(name in values(mesh.physname))
            error("Unknown physical region $name")
        end
        elm = mesh.elms_of[name]
        for k = 1:size(elm, 2)
            for p = 1:size(elm, 1)
                isfree[elm[p,k]] = false
            end
        end
    end
    nofree = 0
    nofixed = 0
    for nd = 1:nonodes
        if isfree[nd]
            nofree += 1
        else
            nofixed += 1
        end
    end
    kfree = 0
    kfixed = 0
    freenode  = zeros(Integer, nofree)
    fixednode = zeros(Integer, nofixed)
    for nd = 1:nonodes
        if isfree[nd]
            kfree += 1
            freenode[kfree] = nd
        else
            kfixed += 1
            fixednode[kfixed] = nd
        end
    end
    node2free  = zeros(Integer, nonodes)
    fill!(node2free, -1)
    for k = 1:nofree
        nd = freenode[k]
        node2free[nd] = k
    end
    node2fixed = zeros(Integer, nonodes)
    fill!(node2fixed, -1)
    for k = 1:nofixed
        nd = fixednode[k]
        node2fixed[nd] = k
    end
    return DoF(isfree, freenode, fixednode, node2free, node2fixed)
end

function assembled_linear_system(vp::VariationalProblem)
    mesh = vp.mesh
    dof = vp.dof
    nofree = length(dof.freenode)
    nofixed = length(dof.fixednode)
    A = assembled_matrix(vp)
    b = zeros(nofree)
    assembled_vector!(b, vp)
    if nofixed > 0
        b .= b - A[:,nofree+1:end] * vp.ufixed
    end
    return A[:,1:nofree], b
end

function assembled_eigenproblem_matrices(ep::EigenProblem)
    mesh = ep.mesh
    dof = ep.dof
    nofree = length(dof.freenode)
    nofixed = length(dof.fixednode)
    A = sparse(Int64[], Int64[], Float64[], nofree, nofree)
    B = sparse(Int64[], Int64[], Float64[], nofree, nofree)
    for (name, elm_mat!, coef) in ep.LHS_bilin_form
        next = assembled_matrix(name, elm_mat!, coef, mesh, dof)
        A += next[:,1:nofree]
    end
    for (name, elm_mat!, coef) in ep.RHS_bilin_form
        next = assembled_matrix(name, elm_mat!, coef, mesh, dof)
        B += next[:,1:nofree]
    end
    return A, B
end

function element_matrix(elm_A::Matrix{Float64}, z::Matrix{Float64},
                        coef::Union{Float64,Function}, k::Int64,
                        elm_mat!::Function)
    elm_mat!(elm_A, z, coef)
end

function element_matrix(elm_A::Matrix{Float64}, z::Matrix{Float64},
                        coef::Vector{Float64}, k::Int64, elm_mat!::Function)
    elm_mat!(elm_A, z, coef[k])
end

function assembled_matrix(vp::VariationalProblem)
    mesh = vp.mesh
    dof = vp.dof
    nofree = length(dof.freenode)
    nofixed = length(dof.fixednode)
    A = sparse(Int64[], Int64[], Float64[], nofree, nofree+nofixed)
    for (name, elm_mat!, coef) in vp.bilin_form
        next = assembled_matrix(name, elm_mat!, coef, mesh, dof)
        A += next
    end
    return A
end

function assembled_matrix(name::String, elm_mat!::Function,
                          coef::Union{Float64, Function, Vector{Float64}},
                          mesh::Mesh, dof::DoF)
    # name = physical name of sub-domain
    # coef    = constant coefficient
    # coef(x) = value of coef at x
    # coef[j] = value of coef at centroid of jth element
    nofree = length(dof.freenode)
    nofixed = length(dof.fixednode)
    isfree = dof.isfree
    I = Int64[]
    J = Int64[]
    V = Float64[]
    elm = mesh.elms_of[name]
    nonodes_per_elm = size(elm, 1)
    elm_A = zeros(nonodes_per_elm, nonodes_per_elm)
    z = zeros(size(mesh.coord,1), nonodes_per_elm)
    for k = 1:size(elm, 2)
        for i = 1:size(z,1), p = 1:nonodes_per_elm
            z[i,p] = mesh.coord[i,elm[p,k]]
        end
        element_matrix(elm_A, z, coef, k, elm_mat!)
        for p = 1:nonodes_per_elm
            if !isfree[elm[p,k]]
                continue
            end
            for q = 1:nonodes_per_elm
                push!(I, dof.node2free[elm[p,k]])
                if isfree[elm[q,k]]
                    push!(J, dof.node2free[elm[q,k]])
                else
                    push!(J, nofree + dof.node2fixed[elm[q,k]])
                end
                push!(V, elm_A[p,q])
            end
        end
    end
    A = sparse(I, J, V, nofree, nofree+nofixed)
    return A
end

function assembled_vector!(loadvec::Vector{Float64}, vp::VariationalProblem)
    mesh = vp.mesh
    dof = vp.dof
    nofree = length(dof.freenode)
    nofixed = length(dof.fixednode)
    @assert length(loadvec) == nofree
    fill!(loadvec, 0.0)
    for (name, elm_vec!, f) in vp.lin_functnl
        assembled_vector!(loadvec, name, elm_vec!, f, mesh, dof)
    end
end

function assembled_vector(name::String, elm_vec!::Function,
                          f::Function, mesh::Mesh, dof::DoF)
    nofree = length(dof.freenode)
    loadvec = zeros(nofree)
    assembled_vector!(loadvec, name, elm_vec!, f, mesh, dof)
    return loadvec
end

function assembled_vector!(loadvec::Vector{Float64},
                           name::String, elm_vec!::Function,
                           f::Function, mesh::Mesh, dof::DoF)
    nofree = length(dof.freenode)
    @assert length(loadvec) == nofree
    nofixed = length(dof.fixednode)
    isfree = dof.isfree
    elm = mesh.elms_of[name]
    nonodes_per_elm = size(elm, 1)
    elm_v = zeros(nonodes_per_elm)
    z = zeros(size(mesh.coord,1), nonodes_per_elm)
    for k = 1:size(elm, 2)
        for p = 1:nonodes_per_elm
            for i = 1:size(z,1)
                z[i,p] = mesh.coord[i,elm[p,k]]
            end
        end
        elm_vec!(elm_v, z, f)
        for p = 1:nonodes_per_elm
            if isfree[elm[p,k]]
                i = dof.node2free[elm[p,k]]
                loadvec[i] += elm_v[p]
            end
        end
    end
end

"""
Returns the vector of nodal values, including both the free
and fixed nodes.
"""
function complete_soln(ufree::Vector{Float64}, vp::VariationalProblem)
    dof = vp.dof
    nofree = length(dof.freenode)
    nofixed = length(dof.fixednode)
    u = zeros(nofree+nofixed)
    for k = 1:nofree
        nd = dof.freenode[k]
        u[nd] = ufree[k]
    end
    for k = 1:nofixed
        nd = dof.fixednode[k]
        u[nd] = vp.ufixed[k]
    end
    return u
end

function complete_soln!(u::Vector{Float64}, ufree::Vector{Float64},
                        ufixed::Vector{Float64}, dof::DoF)
    nofree = length(dof.freenode)
    nofixed = length(dof.fixednode)
    @assert length(ufree) == nofree
    @assert length(ufixed) == nofixed
    for k = 1:nofree
        nd = dof.freenode[k]
        u[nd] = ufree[k]
    end
    for k = 1:nofixed
        nd = dof.fixednode[k]
        u[nd] = ufixed[k]
    end
end

function complete_soln(vfree::Matrix{Float64}, ep::EigenProblem)
    dof = ep.dof
    nofree = length(dof.freenode)
    nofixed = length(dof.fixednode)
    nev = size(vfree, 2)
    v = zeros(nofree+nofixed, nev)
    for j = 1:nev, k = 1:nofree
        nd = dof.freenode[k]
        v[nd,j] = vfree[k,j]
    end
    return v
end


text:
module CommunalHelperCustomSummitGemManager

using ..Ahorn, Maple

@mapdef Entity "CommunalHelper/CustomSummitGemManager" SummitGemManager(
    x::Integer,
    y::Integer,
    gemIds::String="",
    melody::String="",
    heartOffset::String="0.0,0.0",
    nodes::Array{Tuple{Integer,Integer},1}=Tuple{Integer,Integer}[],
)

const placements = Ahorn.PlacementDict(
    "Summit Gem Manager (Communal Helper)" => Ahorn.EntityPlacement(
        SummitGemManager,
    ),
)

Ahorn.nodeLimits(entity::SummitGemManager) = 0, -1

const gemSprite = "collectables/summitgems/0/gem00"

const smallGemSprites = String["collectables/summitgems/$i/small00" for i in 0:7]

function Ahorn.selection(entity::SummitGemManager)
    x, y = Ahorn.position(entity)

    nodes = get(entity.data, "nodes", ())

    res = Ahorn.Rectangle[Ahorn.Rectangle(x - 15, y - 15, 30, 30)]

    for node in nodes
        nx, ny = node
        push!(res, Ahorn.getSpriteRectangle(gemSprite, nx, ny))
    end

    hx, hy = getHeartOffset(entity)
    if hx + hy != 0
        push!(res, Ahorn.getSpriteRectangle("collectables/heartGem/ghost00", x + hx, y + hy))
    end

    return res
end

# Cruor please don't yell at me
function Ahorn.Selection.applyMovement!(target::SummitGemManager, ox, oy, node=0)
    if node == 0
        target.data["x"] += ox
        target.data["y"] += oy
    else
        nodes = get(target.data, "nodes", ())

        if length(nodes) >= node
            nodes[node] = nodes[node] .+ (ox, oy)
        elseif node == length(nodes) + 1
            heartOffset = getHeartOffset(target)
            (hx, hy) = heartOffset .+ (ox, oy)
            target.data["heartOffset"] = string(hx, ",", hy)
        end
    end

    return true
end

# not currently working
function Ahorn.deleted(target::SummitGemManager, node::Int64)
    nodes = get(target.data, "nodes", ())

    if node == length(nodes) + 1
        target.data["heartOffset"] = "0.0,0.0"
    end
end

function getHeartOffset(entity::SummitGemManager)
    heartOffset = get(entity.data, "heartOffset", "")
    if length(heartOffset) > 3
        heartOffset = (split(heartOffset, ","))
        heartOffset = tryparse.(Float64, heartOffset)
        if all(f -> isnothing(f), heartOffset)
            return (0, 0)
        end
        return (heartOffset[1], heartOffset[2])
    end
    return (0, 0)
end

function Ahorn.renderSelectedAbs(ctx::Ahorn.Cairo.CairoContext, entity::SummitGemManager)
    x, y = Ahorn.position(entity)

    nodes = get(entity.data, "nodes", ())

    for node in nodes
        nx, ny = node

        Ahorn.drawLines(ctx, Tuple{Number, Number}[(x, y), (nx, ny)], Ahorn.colors.selection_selected_fc)
    end

    hx, hy = getHeartOffset(entity)
    if hx + hy != 0
        Ahorn.drawArrow(ctx, x, y, x + hx, y + hy, Ahorn.colors.selection_selected_fc; headLength=6, headTheta=pi/8)
    end
end

function Ahorn.renderAbs(ctx::Ahorn.Cairo.CairoContext, entity::SummitGemManager)
    x, y = Ahorn.position(entity)

    hx, hy = getHeartOffset(entity)
    if hx + hy != 0
        Ahorn.drawSprite(ctx, "collectables/heartGem/ghost00", x + hx, y + hy)
    end

    Ahorn.Cairo.save(ctx)
    Ahorn.set_antialias(ctx, 1)
    Ahorn.set_line_width(ctx, 1)
    Ahorn.Cairo.arc(ctx, x, y, 20, 0, pi * 2)
    Ahorn.Cairo.set_source_rgba(ctx, 0.0, 0.0, 0.0, 0.4)
    Ahorn.Cairo.fill_preserve(ctx)
    Ahorn.Cairo.set_source_rgba(ctx, 1.0, 1.0, 1.0, 1.0)
    Ahorn.Cairo.stroke(ctx)
    Ahorn.restore(ctx)

    for i in 1:length(smallGemSprites)
        gx, gy = Base.Math.cospi((i / length(smallGemSprites)) * 2) * 14,
        Base.Math.sinpi((i / length(smallGemSprites)) * 2) * 14
        Ahorn.drawSprite(ctx, smallGemSprites[i], x + gx, y + gy)
    end

    nodes = get(entity.data, "nodes", ())

    for node in nodes
        nx, ny = node
        Ahorn.drawSprite(ctx, gemSprite, nx, ny)
    end
end

end