text:
#### Elementwise manipulations (scaling/clamping/type conversion) ####

# This file exists primarily to handle conversions for display and
# saving to disk. Both of these operations require Ufixed-valued
# elements, but with display we always want to convert to 8-bit
# whereas saving can handle 16-bit.
# We also can't trust that user images are clamped properly.
# Finally, this supports adjustable contrast limits.

# Structure of MapInfo subtype definitions:
#   - type definition
#   - constructors for scalars
#   - constructors for AbstractArrays
#   - similar  (syntax: similar(mapi, ToType, FromType))
#   - implementation of map() for scalars
#   - implementation of map() for AbstractArrays
# map(mapi::MapInfo{T}, x) should return an object of type T (for x not an array)
# map1(mapi::MapInfo{T}, x) is designed to allow T<:ColorValue to work on
#    scalars x::Fractional

# Dispatch-based elementwise manipulations
abstract MapInfo{T}
eltype{T}(mapi::MapInfo{T}) = T

## MapNone
# At most, this does conversions

immutable MapNone{T} <: MapInfo{T}; end

# Constructors
MapNone{T}(::Type{T}) = MapNone{T}()
MapNone{T}(val::T) = MapNone{T}()
MapNone{T}(A::AbstractArray{T}) = MapNone{T}()

similar{T}(mapi::MapNone, ::Type{T}, ::Type) = MapNone{T}()

# Implementation
map{T<:Real}(mapi::MapNone{T}, val::Real) = convert(T, val)
map{C<:ColorValue}(mapi::MapNone{C}, c::ColorValue) = convert(C, c)
map{C<:ColorValue,T}(mapi::MapNone{AlphaColorValue{C,T}}, c::AlphaColorValue) = convert(AlphaColorValue{C,T}, c)
map1(mapi::Union(MapNone{RGB24}, MapNone{ARGB32}), b::Bool) = ifelse(b, 0xffuf8, 0x00uf8)
map1(mapi::Union(MapNone{RGB24},MapNone{ARGB32}), val::Fractional) = convert(Ufixed8, val)
map1{CT<:ColorType}(mapi::MapNone{CT}, val::Fractional) = convert(eltype(CT), val)

map{T<:ColorType}(mapi::MapNone{T}, img::AbstractImageIndexed{T}) = convert(Image{T}, img)
map{C<:ColorType}(mapi::MapNone{C}, img::AbstractImageDirect{C}) = img  # ambiguity resolution
map{T}(mapi::MapNone{T}, img::AbstractArray{T}) = img

## BitShift
# This is really a "saturating bitshift", for example
#    map(BitShift{Uint8,7}(), 0xf0ff) == 0xff rather than 0xe1 even though 0xf0ff>>>7 == 0x01e1

immutable BitShift{T,N} <: MapInfo{T} end
BitShift{T}(::Type{T}, n::Int) = BitShift{T,n}()  # note that this is not type-stable

similar{S,T,N}(mapi::BitShift{S,N}, ::Type{T}, ::Type) = BitShift{T,N}()

# Implementation
immutable BS{N} end
_map{T<:Unsigned,N}(::Type{T}, ::Type{BS{N}}, val::Unsigned) = (v = val>>>N; tm = oftype(val, typemax(T)); convert(T, ifelse(v > tm, tm, v)))
_map{T<:Ufixed,N}(::Type{T}, ::Type{BS{N}}, val::Ufixed) = reinterpret(T, _map(FixedPointNumbers.rawtype(T), BS{N}, reinterpret(val)))
map{T<:Real,N}(mapi::BitShift{T,N}, val::Real) = _map(T, BS{N}, val)
map{T<:Real,N}(mapi::BitShift{Gray{T},N}, val::Gray) = Gray(_map(T, BS{N}, val.val))
map1{N}(mapi::Union(BitShift{RGB24,N},BitShift{ARGB32,N}), val::Unsigned) = _map(Uint8, BS{N}, val)
map1{N}(mapi::Union(BitShift{RGB24,N},BitShift{ARGB32,N}), val::Ufixed) = _map(Ufixed8, BS{N}, val)
map1{CT<:ColorType,N}(mapi::BitShift{CT,N}, val::Ufixed) = _map(eltype(CT), BS{N}, val)

## Clamp types
# The Clamp types just enforce bounds, but do not scale or offset

# Types and constructors
abstract AbstractClamp{T} <: MapInfo{T}
immutable ClampMin{T,From} <: AbstractClamp{T}
    min::From
end
ClampMin{T,From}(::Type{T}, min::From) = ClampMin{T,From}(min)
ClampMin{T}(min::T) = ClampMin{T,T}(min)
immutable ClampMax{T,From} <: AbstractClamp{T}
    max::From
end
ClampMax{T,From}(::Type{T}, max::From) = ClampMax{T,From}(max)
ClampMax{T}(max::T) = ClampMax{T,T}(max)
immutable ClampMinMax{T,From} <: AbstractClamp{T}
    min::From
    max::From
end
ClampMinMax{T,From}(::Type{T}, min::From, max::From) = ClampMinMax{T,From}(min,max)
ClampMinMax{T}(min::T, max::T) = ClampMinMax{T,T}(min,max)
immutable Clamp{T} <: AbstractClamp{T} end  # specialized for clamping colorvalues (e.g., 0 to 1 for RGB, also fractional)
Clamp{T}(::Type{T}) = Clamp{T}()

similar{T,F}(mapi::ClampMin, ::Type{T}, ::Type{F}) = ClampMin{T,F}(convert(F, mapi.min))
similar{T,F}(mapi::ClampMax, ::Type{T}, ::Type{F}) = ClampMax{T,F}(convert(F, mapi.max))
similar{T,F}(mapi::ClampMinMax, ::Type{T}, ::Type{F}) = ClampMin{T,F}(convert(F, mapi.min), convert(F, mapi.max))
similar{T,F}(mapi::Clamp, ::Type{T}, ::Type{F}) = Clamp{T}()

# Implementation
map{T<:Real,F<:Real}(mapi::ClampMin{T,F}, val::F) = convert(T, max(val, mapi.min))
map{T<:Real,F<:Real}(mapi::ClampMax{T,F}, val::F) = convert(T, min(val, mapi.max))
map{T<:Real,F<:Real}(mapi::ClampMinMax{T,F}, val::F) = convert(T,min(max(val, mapi.min), mapi.max))
map{T<:Fractional,F<:Real}(mapi::ClampMin{Gray{T},F}, val::F) = convert(Gray{T}, max(val, mapi.min))
map{T<:Fractional,F<:Real}(mapi::ClampMax{Gray{T},F}, val::F) = convert(Gray{T}, min(val, mapi.max))
map{T<:Fractional,F<:Real}(mapi::ClampMinMax{Gray{T},F}, val::F) = convert(Gray{T},min(max(val, mapi.min), mapi.max))
map{T<:Fractional,F<:Fractional}(mapi::ClampMin{Gray{T},F}, val::Gray{F}) = convert(Gray{T}, max(val, mapi.min))
map{T<:Fractional,F<:Fractional}(mapi::ClampMax{Gray{T},F}, val::Gray{F}) = convert(Gray{T}, min(val, mapi.max))
map{T<:Fractional,F<:Fractional}(mapi::ClampMinMax{Gray{T},F}, val::Gray{F}) = convert(Gray{T},min(max(val, mapi.min), mapi.max))
map{T<:Fractional,F<:Fractional}(mapi::ClampMin{Gray{T},Gray{F}}, val::Gray{F}) = convert(Gray{T}, max(val, mapi.min))
map{T<:Fractional,F<:Fractional}(mapi::ClampMax{Gray{T},Gray{F}}, val::Gray{F}) = convert(Gray{T}, min(val, mapi.max))
map{T<:Fractional,F<:Fractional}(mapi::ClampMinMax{Gray{T},Gray{F}}, val::Gray{F}) = convert(Gray{T},min(max(val, mapi.min), mapi.max))
map1{T<:Union(RGB24,ARGB32),F<:Fractional}(mapi::ClampMin{T,F}, val::F) = convert(Ufixed8, max(val, mapi.min))
map1{T<:Union(RGB24,ARGB32),F<:Fractional}(mapi::ClampMax{T,F}, val::F) = convert(Ufixed8, min(val, mapi.max))
map1{T<:Union(RGB24,ARGB32),F<:Fractional}(mapi::ClampMinMax{T,F}, val::F) = convert(Ufixed8,min(max(val, mapi.min), mapi.max))
map1{CT<:ColorType,F<:Fractional}(mapi::ClampMin{CT,F}, val::F) = convert(eltype(CT), max(val, mapi.min))
map1{CT<:ColorType,F<:Fractional}(mapi::ClampMax{CT,F}, val::F) = convert(eltype(CT), min(val, mapi.max))
map1{CT<:ColorType,F<:Fractional}(mapi::ClampMinMax{CT,F}, val::F) = convert(eltype(CT), min(max(val, mapi.min), mapi.max))

map{To<:Real}(::Clamp{To}, val::Real) = clamp01(To, val)
map{To<:Real}(::Clamp{Gray{To}}, val::AbstractGray) = Gray(clamp01(To, val.val))
map{To<:Real}(::Clamp{Gray{To}}, val::Real) = Gray(clamp01(To, val))
map1{CT<:AbstractRGB}(::Clamp{CT}, val::Real) = clamp01(eltype(CT), val)
map1{CT<:AbstractRGB,T}(::Clamp{AlphaColor{CT,T}}, val::Real) = clamp01(T, val)
map1{CT<:AbstractRGB,T}(::Clamp{AlphaColorValue{CT,T}}, val::Real) = clamp01(T, val)

# Also available as a stand-alone function
clamp01{T}(::Type{T}, x::Real) = convert(T, min(max(x, zero(x)), one(x)))
clamp01(x::Real) = clamp01(typeof(x), x)
for CV in subtypes(AbstractRGB)
    @eval begin
        clamp01{T}(x::$CV{T}) = clamp01($CV{T}, x)
        clamp01{To}(::Type{$CV{To}}, x::AbstractRGB) = $CV{To}(clamp01(To, x.r), clamp01(To, x.g), clamp01(To, x.b))
        clamp01{T}(x::AlphaColor{$CV{T},T}) = clamp01(AlphaColor{$CV{T},T}, x)
        clamp01{T}(x::AlphaColorValue{$CV{T},T}) = clamp01(AlphaColorValue{$CV{T},T}, x)
        clamp01{To,CV2<:AbstractRGB}(::Type{AlphaColor{$CV{To},To}}, x::AlphaColor{CV2}) =
            AlphaColor{$CV{To},To}(clamp01(To, x.c.r), clamp01(To, x.c.g), clamp01(To, x.c.b), clamp01(To, x.alpha))
        clamp01{To,CV2<:AbstractRGB}(::Type{AlphaColorValue{$CV{To},To}}, x::AlphaColorValue{CV2}) =
            AlphaColorValue{$CV{To},To}(clamp01(To, x.c.r), clamp01(To, x.c.g), clamp01(To, x.c.b), clamp01(To, x.alpha))
        clamp01{To,CV2<:AbstractRGB}(::Type{AlphaColor{$CV{To},To}}, x::AlphaColorValue{CV2}) =
            AlphaColor{$CV{To},To}(clamp01(To, x.c.r), clamp01(To, x.c.g), clamp01(To, x.c.b), clamp01(To, x.alpha))
        clamp01{To,CV2<:AbstractRGB}(::Type{AlphaColorValue{$CV{To},To}}, x::AlphaColor{CV2}) =
            AlphaColorValue{$CV{To},To}(clamp01(To, x.c.r), clamp01(To, x.c.g), clamp01(To, x.c.b), clamp01(To, x.alpha))
    end
end

# clamp is generic for any colorspace; this version does the right thing for any RGB type
clamp(x::AbstractRGB) = clamp01(x)
for CV in subtypes(AbstractRGB)
    @eval begin
        clamp{T}(x::AlphaColorValue{$CV{T},T}) = clamp01(x)
        clamp{T}(x::AlphaColor{$CV{T},T}) = clamp01(x)
    end
end

## ScaleMinMax
# This clamps, subtracts the min value, then scales

immutable ScaleMinMax{To,From,S<:FloatingPoint} <: MapInfo{To}
    min::From
    max::From
    s::S

    function ScaleMinMax(min, max, s)
        min >= max && error("min must be smaller than max")
        new(min, max, s)
    end
end

ScaleMinMax{To,From}(::Type{To}, min::From, max::From, s::FloatingPoint) = ScaleMinMax{To,From,typeof(s)}(min, max, s)
ScaleMinMax{To<:Union(Fractional,ColorType),From}(::Type{To}, mn::From, mx::From) = ScaleMinMax(To, mn, mx, 1.0f0/(convert(Float32, mx)-convert(Float32, mn)))

# ScaleMinMax constructors that take AbstractArray input
ScaleMinMax{To,From<:Real}(::Type{To}, img::AbstractArray{From}, mn::Real, mx::Real) = ScaleMinMax(To, convert(From,mn), convert(From,mx), 1.0f0/(convert(Float32, convert(From, mx))-convert(Float32,convert(From, mn))))
ScaleMinMax{To,From<:Real}(::Type{To}, img::AbstractArray{Gray{From}}, mn::Real, mx::Real) = ScaleMinMax(To, convert(From,mn), convert(From,mx), 1.0f0/(convert(Float32, convert(From,mx))-convert(Float32, convert(From,mn))))
ScaleMinMax{To,From<:Real,R<:Real}(::Type{To}, img::AbstractArray{From}, mn::Gray{R}, mx::Gray{R}) = ScaleMinMax(To, convert(From,mn.val), convert(From,mx.val), 1.0f0/(convert(Float32, convert(From,mx.val))-convert(Float32, convert(From,mn.val))))
ScaleMinMax{To,From<:Real,R<:Real}(::Type{To}, img::AbstractArray{Gray{From}}, mn::Gray{R}, mx::Gray{R}) = ScaleMinMax(To, convert(From,mn.val), convert(From,mx.val), 1.0f0/(convert(Float32, convert(From,mx.val))-convert(Float32, convert(From,mn.val))))
ScaleMinMax{To}(::Type{To}, img::AbstractArray) = ScaleMinMax(To, img, minfinite(img), maxfinite(img))
ScaleMinMax{To,CV<:AbstractRGB}(::Type{To}, img::AbstractArray{CV}) = (imgr = reinterpret(eltype(CV), img); ScaleMinMax(To, minfinite(imgr), maxfinite(imgr)))

similar{T,F,To,From,S}(mapi::ScaleMinMax{To,From,S}, ::Type{T}, ::Type{F}) = ScaleMinMax{T,F,S}(convert(F,mapi.min), convert(F.mapi.max), mapi.s)

# Implementation
function map{To<:Real,From<:Union(Real,Gray)}(mapi::ScaleMinMax{To,From}, val::From)
    t = ifelse(val  < mapi.min, zero(From), ifelse(val  > mapi.max, mapi.max-mapi.min, val -mapi.min))
    convert(To, mapi.s*t)
end
function map{To<:Real,From<:Union(Real,Gray)}(mapi::ScaleMinMax{To,From}, val::Union(Real,ColorType))
    map(mapi, convert(From, val))
end
function map1{To<:Union(RGB24,ARGB32),From<:Real}(mapi::ScaleMinMax{To,From}, val::From)
    t = ifelse(val  < mapi.min, zero(From), ifelse(val  > mapi.max, mapi.max-mapi.min, val -mapi.min))
    convert(Ufixed8, mapi.s*t)
end
function map1{To<:ColorType,From<:Real}(mapi::ScaleMinMax{To,From}, val::From)
    t = ifelse(val  < mapi.min, zero(From), ifelse(val  > mapi.max, mapi.max-mapi.min, val -mapi.min))
    convert(eltype(To), mapi.s*t)
end
function map1{To<:Union(RGB24,ARGB32),From<:Real}(mapi::ScaleMinMax{To,From}, val::Union(Real,ColorType))
    map1(mapi, convert(From, val))
end
function map1{To<:ColorType,From<:Real}(mapi::ScaleMinMax{To,From}, val::Union(Real,ColorType))
    map1(mapi, convert(From, val))
end

## ScaleSigned
# Multiplies by a scaling factor and then clamps to the range [-1,1].
# Intended for positive/negative coloring

immutable ScaleSigned{T, S<:FloatingPoint} <: MapInfo{T}
    s::S
end
ScaleSigned{T}(::Type{T}, s::FloatingPoint) = ScaleSigned{T, typeof(s)}(s)

ScaleSigned{T}(::Type{T}, img::AbstractArray) = ScaleSigned(T, 1.0f0/maxabsfinite(img))
ScaleSigned(img::AbstractArray) = ScaleSigned(Float32, img)

similar{T,To,S}(mapi::ScaleSigned{To,S}, ::Type{T}, ::Type) = ScaleSigned{T,S}(mapi.s)

map{T}(mapi::ScaleSigned{T}, val::Real) = convert(T, clamppm(mapi.s*val))
function map{C<:Union(RGB24, RGB{Ufixed8})}(mapi::ScaleSigned{C}, val::Real)
    x = clamppm(mapi.s*val)
    g = Ufixed8(abs(x))
    ifelse(x >= 0, C(g, zero(Ufixed8), g), C(zero(Ufixed8), g, zero(Ufixed8)))
end

clamppm(x::Real) = ifelse(x >= 0, min(x, one(x)), max(x, -one(x)))

## ScaleAutoMinMax
# Works only on whole arrays, not values

immutable ScaleAutoMinMax{T} <: MapInfo{T} end
ScaleAutoMinMax{T}(::Type{T}) = ScaleAutoMinMax{T}()
ScaleAutoMinMax() = ScaleAutoMinMax{Ufixed8}()

similar{T}(mapi::ScaleAutoMinMax, ::Type{T}, ::Type) = ScaleAutoMinMax{T}()

# Conversions to RGB{T}, RGBA{T}, RGB24, ARGB32,
# for grayscale, AbstractRGB, and abstract ARGB inputs.
# This essentially "vectorizes" map over a single pixel's color channels using map1
for SI in (MapInfo, AbstractClamp)
    for ST in subtypes(SI)
        ST.abstract && continue
        ST == ScaleSigned && continue  # ScaleSigned gives an RGB from a scalar, so don't "vectorize" it
        @eval begin
            # Grayscale and GrayAlpha inputs
            map(mapi::$ST{RGB24}, g::Gray) = map(mapi, g.val)
            map(mapi::$ST{RGB24}, g::Real) = (x = map1(mapi, g); convert(RGB24, RGB{Ufixed8}(x,x,x)))
            function map(mapi::$ST{RGB24}, g::FloatingPoint)
                if isfinite(g)
                    x = map1(mapi, g)
                    convert(RGB24, RGB{Ufixed8}(x,x,x))
                else
                    RGB24(0)
                end
            end
            map{T}(mapi::$ST{RGB24}, g::GrayAlpha{T}) = map(mapi, g.c.val)
            map(mapi::$ST{ARGB32}, g::Gray) = map(mapi, g.val)
            function map(mapi::$ST{ARGB32}, g::Real)
                x = map1(mapi, g)
                convert(ARGB32, ARGB{Ufixed8}(x,x,x,0xffuf8))
            end
            function map{T}(mapi::$ST{ARGB32}, g::GrayAlpha{T})
                x = map1(mapi, g.c.val)
                convert(ARGB32, ARGB{Ufixed8}(x,x,x,map1(mapi, g.alpha)))
            end
        end
        for O in (:RGB, :BGR)
            @eval begin
                map{T}(mapi::$ST{$O{T}}, g::Gray) = map(mapi, g.val)
                function map{T}(mapi::$ST{$O{T}}, g::Real)
                    x = map1(mapi, g)
                    $O{T}(x,x,x)
                end
            end
        end
        for OA in (:RGBA, :ARGB, :BGRA)
            exAlphaGray = ST == MapNone ? :nothing : quote
                function map{T,S}(mapi::$ST{$OA{T}}, g::GrayAlpha{S})
                    x = map1(mapi, g.c.val)
                    $OA{T}(x,x,x,map1(mapi, g.alpha))
                end  # avoids an ambiguity warning with MapNone definitions
            end
            @eval begin
                map{T}(mapi::$ST{$OA{T}}, g::Gray) = map(mapi, g.val)
                function map{T}(mapi::$ST{$OA{T}}, g::Real)
                    x = map1(mapi, g)
                    $OA{T}(x,x,x)
                end
                $exAlphaGray
            end
        end
        @eval begin
            # AbstractRGB and abstract ARGB inputs
            map(mapi::$ST{RGB24}, rgb::AbstractRGB) =
                convert(RGB24, RGB{Ufixed8}(map1(mapi, rgb.r), map1(mapi, rgb.g), map1(mapi, rgb.b)))
            map{C<:AbstractRGB, TC}(mapi::$ST{RGB24}, argb::AbstractAlphaColorValue{C,TC}) =
                convert(RGB24, RGB{Ufixed8}(map1(mapi, argb.c.r), map1(mapi, argb.c.g),
                                            map1(mapi, argb.c.b)))
            map{C<:AbstractRGB, TC}(mapi::$ST{ARGB32}, argb::AbstractAlphaColorValue{C,TC}) =
                convert(ARGB32, ARGB{Ufixed8}(map1(mapi, argb.c.r), map1(mapi, argb.c.g),
                                              map1(mapi, argb.c.b), map1(mapi, argb.alpha)))
            map(mapi::$ST{ARGB32}, rgb::AbstractRGB) =
                convert(ARGB32, ARGB{Ufixed8}(map1(mapi, rgb.r), map1(mapi, rgb.g), map1(mapi, rgb.b)))
        end
        for O in (:RGB, :BGR)
            @eval begin
                map{T}(mapi::$ST{$O{T}}, rgb::AbstractRGB) =
                    $O{T}(map1(mapi, rgb.r), map1(mapi, rgb.g), map1(mapi, rgb.b))
                map{T,C<:AbstractRGB, TC}(mapi::$ST{$O{T}}, argb::AbstractAlphaColorValue{C,TC}) =
                    $O{T}(map1(mapi, argb.c.r), map1(mapi, argb.c.g), map1(mapi, argb.c.b))
            end
        end
        for OA in (:RGBA, :ARGB, :BGRA)
            @eval begin
                map{T, C<:AbstractRGB, TC}(mapi::$ST{$OA{T}}, argb::AbstractAlphaColorValue{C,TC}) =
                    $OA{T}(map1(mapi, argb.c.r), map1(mapi, argb.c.g),
                            map1(mapi, argb.c.b), map1(mapi, argb.alpha))
                map{T}(mapi::$ST{$OA{T}}, argb::ARGB32) = map(mapi, convert(RGBA{Ufixed8}, argb))
                map{T}(mapi::$ST{$OA{T}}, rgb::AbstractRGB) =
                    $OA{T}(map1(mapi, rgb.r), map1(mapi, rgb.g), map1(mapi, rgb.b))
                map{T}(mapi::$ST{$OA{T}}, rgb::RGB24) = map(mapi, convert(RGB{Ufixed8}, argb))
            end
        end
    end
end

# # Apply to any ColorType
# map(f::Callable, x::ColorValue) = f(x)
# map(mapi, x::ColorValue) = map(mapi, convert(RGB, x))
# map{C<:ColorValue, TC}(f::Callable, x::AbstractAlphaColorValue{C, TC}) = f(convert(ARGB, x))
# map{C<:ColorValue, TC}(mapi, x::AbstractAlphaColorValue{C, TC}) = map(mapi, convert(ARGB, x))

## Fallback definitions of map() for array types

function map{T}(mapi::MapInfo{T}, img::AbstractArray)
    out = similar(img, T)
    map!(mapi, out, img)
end

map{C<:ColorType,R<:Real}(mapi::MapNone{C}, img::AbstractImageDirect{R}) = mapcd(mapi, img)  # ambiguity resolution
map{C<:ColorType,R<:Real}(mapi::MapInfo{C}, img::AbstractImageDirect{R}) = mapcd(mapi, img)
function mapcd{C<:ColorType,R<:Real}(mapi::MapInfo{C}, img::AbstractImageDirect{R})
    # For this case we have to check whether color is defined along an array axis
    cd = colordim(img)
    if cd > 0
        dims = setdiff(1:ndims(img), cd)
        out = similar(img, C, size(img)[dims])
        map!(mapi, out, img, TypeConst{cd})
    else
        out = similar(img, C)
        map!(mapi, out, img)
    end
    out   # note this isn't type-stable
end

function map{T<:ColorType}(mapi::MapInfo{T}, img::AbstractImageIndexed)
    out = Image(Array(T, size(img)), properties(img))
    map!(mapi, out, img)
end

map!{T,T1,T2,N}(mapi::MapInfo{T1}, out::AbstractArray{T,N}, img::AbstractArray{T2,N}) =
    _map_a!(mapi, out, img)
@ngenerate N typeof(out) function _map_a!{T,T1,T2,N}(mapi::MapInfo{T1}, out::AbstractArray{T,N}, img::AbstractArray{T2,N})
    mi = take(mapi, img)
    dimg = data(img)
    dout = data(out)
    size(dout) == size(dimg) || throw(DimensionMismatch())
    @nloops N i dout begin
        @inbounds @nref(N, dout, i) = map(mi, @nref(N, dimg, i))
    end
    out
end

take(mapi::MapInfo, img::AbstractArray) = mapi
take{T}(mapi::ScaleAutoMinMax{T}, img::AbstractArray) = ScaleMinMax(T, img)

# Indexed images (colormaps)
map!{T,T1,N}(mapi::MapInfo{T}, out::AbstractArray{T,N}, img::AbstractImageIndexed{T1,N}) =
    _mapindx!(mapi, out, img)
@ngenerate N typeof(out) function _mapindx!{T,T1,N}(mapi::MapInfo{T}, out::AbstractArray{T,N}, img::AbstractImageIndexed{T1,N})
    dimg = data(img)
    dout = data(out)
    cmap = map(mapi, img.cmap)
    @nloops N i dout begin
        @inbounds @nref(N, dout, i) = cmap[@nref(N, dimg, i)]
    end
    out
end

# For when color is encoded along dimension CD
# NC is the number of color channels
# This is a very flexible implementation: color can be stored along any dimension, and it handles conversions to
# many different colorspace representations.
for (CT, NC) in ((Union(AbstractRGB,RGB24), 3), (Union(RGBA,ARGB,ARGB32), 4), (Union(GrayAlpha,AGray32), 2))
    for N = 1:4
        N1 = N+1
        @eval begin
function map!{T<:$CT,T1,T2,CD}(mapi::MapInfo{T}, out::AbstractArray{T1,$N}, img::AbstractArray{T2,$N1}, ::Type{TypeConst{CD}})
    mi = take(mapi, img)
    dimg = data(img)
    dout = data(out)
    # Set up the index along the color axis
    # We really only need dimension CD, but this will suffice
    @nexprs $NC k->(@nexprs $N1 d->(j_k_d = k))
    # Loop over all the elements in the output, performing the conversion on each color component
    @nloops $N i dout d->(d<CD ? (@nexprs $NC k->(j_k_d = i_d)) : (@nexprs $NC k->(j_k_{d+1} = i_d))) begin
        @inbounds @nref($N, dout, i) = @ncall $NC T k->(map1(mi, @nref($N1, dimg, j_k)))
    end
    out
end
        end
    end
end

#### MapInfo defaults
# Each "client" can define its own methods. "clients" include Ufixed, RGB24/ARGB32, and ImageMagick

const bitshiftto8 = ((Ufixed10, 2), (Ufixed12, 4), (Ufixed14, 6), (Ufixed16, 8))

# typealias GrayType{T<:Fractional} Union(T, Gray{T})
typealias GrayArray{T<:Fractional} Union(AbstractArray{T}, AbstractArray{Gray{T}})
# note, though, that we need to override for AbstractImage in case the "colorspace" property is defined differently

# mapinfo{T<:Union(Real,ColorType)}(::Type{T}, img::AbstractArray{T}) = MapNone(img)
mapinfo{T<:Ufixed}(::Type{T}, img::AbstractArray{T}) = MapNone(img)
mapinfo{T<:FloatingPoint}(::Type{T}, img::AbstractArray{T}) = MapNone(img)

# Grayscale methods
mapinfo(::Type{Ufixed8}, img::GrayArray{Ufixed8}) = MapNone{Ufixed8}()
mapinfo(::Type{Gray{Ufixed8}}, img::GrayArray{Ufixed8}) = MapNone{Gray{Ufixed8}}()
mapinfo(::Type{GrayAlpha{Ufixed8}}, img::AbstractArray{GrayAlpha{Ufixed8}}) = MapNone{GrayAlpha{Ufixed8}}()
for (T,n) in bitshiftto8
    @eval mapinfo(::Type{Ufixed8}, img::GrayArray{$T}) = BitShift{Ufixed8,$n}()
    @eval mapinfo(::Type{Gray{Ufixed8}}, img::GrayArray{$T}) = BitShift{Gray{Ufixed8},$n}()
    @eval mapinfo(::Type{GrayAlpha{Ufixed8}}, img::AbstractArray{GrayAlpha{$T}}) = BitShift{GrayAlpha{Ufixed8},$n}()
end
mapinfo{T<:Ufixed,F<:FloatingPoint}(::Type{T}, img::GrayArray{F}) = ClampMinMax(T, zero(F), one(F))
mapinfo{T<:Ufixed,F<:FloatingPoint}(::Type{Gray{T}}, img::GrayArray{F}) = ClampMinMax(Gray{T}, zero(F), one(F))
mapinfo{T<:FloatingPoint, R<:Real}(::Type{T}, img::AbstractArray{R}) = MapNone(T)

mapinfo(::Type{RGB24}, img::Union(AbstractArray{Bool}, BitArray)) = MapNone{RGB24}()
mapinfo(::Type{ARGB32}, img::Union(AbstractArray{Bool}, BitArray)) = MapNone{ARGB32}()
mapinfo{F<:Fractional}(::Type{RGB24}, img::GrayArray{F}) = ClampMinMax(RGB24, zero(F), one(F))
mapinfo{F<:Fractional}(::Type{ARGB32}, img::AbstractArray{F}) = ClampMinMax(ARGB32, zero(F), one(F))

# Color->Color methods
mapinfo(::Type{RGB{Ufixed8}}, img) = MapNone{RGB{Ufixed8}}()
mapinfo(::Type{RGBA{Ufixed8}}, img) = MapNone{RGBA{Ufixed8}}()
for (T,n) in bitshiftto8
    @eval mapinfo(::Type{RGB{Ufixed8}}, img::AbstractArray{RGB{$T}}) = BitShift{RGB{Ufixed8},$n}()
    @eval mapinfo(::Type{RGBA{Ufixed8}}, img::AbstractArray{RGBA{$T}}) = BitShift{RGBA{Ufixed8},$n}()
end
mapinfo{F<:Fractional}(::Type{RGB{Ufixed8}}, img::AbstractArray{RGB{F}}) = Clamp(RGB{Ufixed8})
mapinfo{F<:Fractional}(::Type{RGBA{Ufixed8}}, img::AbstractArray{RGBA{F}}) = Clamp(RGBA{Ufixed8})

# Color->RGB24/ARGB32
mapinfo(::Type{RGB24}, img::AbstractArray{RGB24}) = MapNone{RGB24}()
mapinfo(::Type{ARGB32}, img::AbstractArray{ARGB32}) = MapNone{ARGB32}()
for C in tuple(subtypes(AbstractRGB)..., Gray)
    @eval mapinfo(::Type{RGB24}, img::AbstractArray{$C{Ufixed8}}) = MapNone{RGB24}()
    @eval mapinfo(::Type{ARGB32}, img::AbstractArray{$C{Ufixed8}}) = MapNone{ARGB32}()
    for (T, n) in bitshiftto8
        @eval mapinfo(::Type{RGB24}, img::AbstractArray{$C{$T}}) = BitShift{RGB24, $n}()
        @eval mapinfo(::Type{ARGB32}, img::AbstractArray{$C{$T}}) = BitShift{ARGB32, $n}()
    end
    @eval mapinfo{F<:FloatingPoint}(::Type{RGB24}, img::AbstractArray{$C{F}}) = ClampMinMax(RGB24, zero(F), one(F))
    @eval mapinfo{F<:FloatingPoint}(::Type{ARGB32}, img::AbstractArray{$C{F}}) = ClampMinMax(ARGB32, zero(F), one(F))
    for AC in subtypes(AbstractAlphaColorValue)
        length(AC.parameters) == 2 || continue
        @eval mapinfo(::Type{ARGB32}, img::AbstractArray{$AC{$C{Ufixed8},Ufixed8}}) = MapNone{ARGB32}()
        @eval mapinfo(::Type{RGB24}, img::AbstractArray{$AC{$C{Ufixed8},Ufixed8}}) = MapNone{RGB24}()
        for (T, n) in bitshiftto8
            @eval mapinfo(::Type{ARGB32}, img::AbstractArray{$AC{$C{$T},$T}}) = BitShift{ARGB32, $n}()
            @eval mapinfo(::Type{RGB24}, img::AbstractArray{$AC{$C{$T},$T}}) = BitShift{RGB24, $n}()
        end
        @eval mapinfo{F<:FloatingPoint}(::Type{ARGB32}, img::AbstractArray{$AC{$C{F},F}}) = ClampMinMax(ARGB32, zero(F), one(F))
        @eval mapinfo{F<:FloatingPoint}(::Type{RGB24}, img::AbstractArray{$AC{$C{F},F}}) = ClampMinMax(RGB24, zero(F), one(F))
    end
end

mapinfo{CT<:ColorType}(::Type{RGB24},  img::AbstractArray{CT}) = MapNone{RGB24}()
mapinfo{CT<:ColorType}(::Type{ARGB32}, img::AbstractArray{CT}) = MapNone{ARGB32}()

# Uint32 conversions will use ARGB32 for images that have an alpha channel,
# and RGB24 when not
mapinfo{CV<:Union(Fractional,ColorValue)}(::Type{Uint32}, img::AbstractArray{CV}) = mapinfo(RGB24, img)
mapinfo{CV<:AbstractAlphaColorValue}(::Type{Uint32}, img::AbstractArray{CV}) = mapinfo(ARGB32, img)
mapinfo(::Type{Uint32}, img::Union(AbstractArray{Bool},BitArray)) = mapinfo(RGB24, img)
mapinfo(::Type{Uint32}, img::AbstractArray{Uint32}) = MapNone{Uint32}()

# ImageMagick client is defined in io.jl

# Backwards-compatibility
uint32color(img) = map(mapinfo(Uint32, img), img)
uint32color!(buf, img::AbstractArray) = map!(mapinfo(Uint32, img), buf, img)
uint32color!(buf, img::AbstractArray, mi::MapInfo) = map!(mi, buf, img)
uint32color!{T,N}(buf::Array{Uint32,N}, img::AbstractImageDirect{T,N}) =
    map!(mapinfo(Uint32, img), buf, img)
uint32color!{T,N,N1}(buf::Array{Uint32,N}, img::AbstractImageDirect{T,N1}) =
    map!(mapinfo(Uint32, img), buf, img, TypeConst{colordim(img)})
uint32color!{T,N}(buf::Array{Uint32,N}, img::AbstractImageDirect{T,N}, mi::MapInfo) =
    map!(mi, buf, img)
uint32color!{T,N,N1}(buf::Array{Uint32,N}, img::AbstractImageDirect{T,N1}, mi::MapInfo) =
    map!(mi, buf, img, TypeConst{colordim(img)})

sc(img::AbstractArray) = map(ScaleMinMax(Ufixed8, img), img)
sc(img::AbstractArray, mn::Real, mx::Real) = map(ScaleMinMax(Ufixed8, img, mn, mx), img)

for (fn,T) in ((:float32, Float32), (:float64, Float64), (:ufixed8, Ufixed8),
               (:ufixed10, Ufixed10), (:ufixed12, Ufixed12), (:ufixed14, Ufixed14),
               (:ufixed16, Ufixed16))
    @eval begin
        function $fn{C<:ColorType}(A::AbstractArray{C})
            newC = eval(C.name.name){$T}
            convert(Array{newC}, A)
        end
        $fn{C<:ColorType}(img::AbstractImage{C}) = shareproperties(img, $fn(data(img)))
    end
end

ufixedsc{T<:Ufixed}(::Type{T}, img::AbstractImageDirect) = map(mapinfo(T, img), img)
ufixed8sc(img::AbstractImageDirect) = ufixedsc(Ufixed8, img)


text:
# export random_behaviour_policy

# include("../../GVF.jl")
using StatsBase

function create_persistant_policy(action)
    next_action_pi(state_t; rng=Random.GLOBAL_RNG) = 1
    prob_pi(state_t, action_t) = if (action_t == action) 1 else 0 end
    pi = GeneralValueFunction.Policy(prob_pi, next_action_pi)
    return pi
end

random_behaviour_next_action(state_t; rng=Random.GLOBAL_RNG) = rand(rng, 1:4)
random_behaviour_probability(state_t, action_t) = 0.25
random_behaviour_policy = GeneralValueFunction.Policy(
    random_behaviour_probability, random_behaviour_next_action)

var_random_behaviour_next_action(state_t; rng=Random.GLOBAL_RNG) = StatsBase.sample(rng, 1:4, Weights([0.9/3, 0.9/3, 0.1, 0.9/3]))
var_random_behaviour_probability(state_t, action_t) = if (action_t == 3) 0.1 else 0.9/3 end
var_random_behaviour_policy = GeneralValueFunction.Policy(
    var_random_behaviour_probability, var_random_behaviour_next_action)

function next_action_variant(state_t; rng=Random.GLOBAL_RNG)
    high_prob = 0.45
    low_prob = 0.05
    # UP, RIGHT, DOWN, LEFT
    if state_t == [3,2]
        # Left Top
        act = StatsBase.sample(rng, 1:4, Weights([high_prob, low_prob, low_prob, high_prob]))
    elseif state_t == [3,9]
        # Right Top
        act = StatsBase.sample(rng, 1:4, Weights([high_prob, high_prob, low_prob, low_prob]))
    elseif state_t == [10,2]
        # Left Down
        act = StatsBase.sample(rng, 1:4, Weights([low_prob, low_prob, high_prob, high_prob]))
    elseif state_t == [10,9]
        # Right Down
        act = StatsBase.sample(rng, 1:4, Weights([low_prob, high_prob, high_prob, low_prob]))
    else
        act = rand(rng, 1:4)
    end
    return act
end

function probability_variant(state_t, action_t)
    high_prob = 0.45
    low_prob = 0.05
    # UP, RIGHT, DOWN, LEFT
    if state_t == [3,2]
        # Left Top
        if action_t == 1 || action_t == 4
            return high_prob
        else
            return low_prob
        end
    elseif state_t == [3,9]
        # Right Top
        if action_t == 1 || action_t == 2
            return high_prob
        else
            return low_prob
        end
    elseif state_t == [10,2]
        # Left Down
        if action_t == 3 || action_t == 4
            return high_prob
        else
            return low_prob
        end
    elseif state_t == [10,9]
        # Right Down
        if action_t == 2 || action_t == 3
            return high_prob
        else
            return low_prob
        end
    else
        return 0.25
    end
    return act
end

variant_behaviour_policy = GeneralValueFunction.Policy(
    probability_variant, next_action_variant)

function next_action_more_variant(state_t; rng=Random.GLOBAL_RNG)
    high_prob = 0.95/3
    low_prob = 0.05
    # UP, RIGHT, DOWN, LEFT
    # Assume target policy is going down.
    # Generated with gen_states(1143139448, 25)
    states = [[6, 7], [4, 8], [6, 9], [3, 11], [2, 10],
              [6, 2], [8, 1], [5, 2], [11, 8], [1, 2],
              [10, 1], [1, 5], [4, 7], [10, 11], [2, 9],
              [6, 5], [9, 4], [3, 6], [4, 1], [3, 2],
              [7, 5], [3, 9], [5, 6], [10, 9], [2, 1]]

    if state_t in states
        act = StatsBase.sample(rng, 1:4, Weights([high_prob, high_prob, low_prob, high_prob]))
    else
        act = rand(rng, 1:4)
    end
    return act
end

function gen_states(seed, num)
    mt = Random.MersenneTwister(seed)
    states = []
    while length(states) < num
        new_state = rand(mt, 1:11, 2)
        if !(new_state in states)
            append!(states, [new_state])
        end
    end
    return states
end

function find_same(states)
    s=0
    for p in collect(Iterators.product(states,states))
        s += if (p[1] == p[2]) 1 else 0 end
    end
    return (s - length(states))/2
end

function probability_more_variant(state_t, action_t)
    high_prob = 0.95/3
    low_prob = 0.05
    # UP, RIGHT, DOWN, LEFT
    # Assume target policy is going Down!
    # Generated with gen_states(1143139448, 25)
    states = [[6, 7], [4, 8], [6, 9], [3, 11], [2, 10],
              [6, 2], [8, 1], [5, 2], [11, 8], [1, 2],
              [10, 1], [1, 5], [4, 7], [10, 11], [2, 9],
              [6, 5], [9, 4], [3, 6], [4, 1], [3, 2],
              [7, 5], [3, 9], [5, 6], [10, 9], [2, 1]]

    if state_t in states
        # Assume
        if action_t != 3
            return high_prob
        else
            return low_prob
        end
    else
        return 0.25
    end
    return act
end

more_variant_behaviour_policy = GeneralValueFunction.Policy(
    probability_more_variant, next_action_more_variant)

function next_action_super_variant(state_t; rng=Random.GLOBAL_RNG)
    high_prob = 0.95/3
    low_prob = 0.05
    # UP, RIGHT, DOWN, LEFT
    # Assume target policy is going down.
    # Generated with gen_states(1143139448, 25)
    # states = [[6, 7], [4, 8], [6, 9], [3, 11], [2, 10],
    #           [6, 2], [8, 1], [5, 2], [11, 8], [1, 2],
    #           [10, 1], [1, 5], [4, 7], [10, 11], [2, 9],
    #           [6, 5], [9, 4], [3, 6], [4, 1], [3, 2],
    #           [7, 5], [3, 9], [5, 6], [10, 9], [2, 1]]
    states = [[6, 7], [4, 8], [6, 9], [3, 11], [2, 10],
              [6, 2], [8, 1], [5, 2], [11, 8], [1, 2],
              [10, 1], [1, 5], [4, 7], [10, 11], [2, 9],
              [6, 5], [9, 4], [3, 6], [4, 1], [3, 2],
              [7, 5], [3, 9], [5, 6], [10, 9], [2, 1],
              [6, 3], [8, 11], [9, 5], [11, 3], [8, 2],
              [5, 9], [7, 11], [10, 10], [9, 8], [1, 10],
              [7, 9], [1, 7], [3, 1], [7, 10], [7, 7],
              [10, 2], [10, 8], [1, 9], [11, 9], [1, 6],
              [11, 10], [8, 8], [10, 5], [5, 11], [2, 2]]
    if state_t in states
        act = StatsBase.sample(rng, 1:4, Weights([high_prob, high_prob, low_prob, high_prob]))
    else
        act = rand(rng, 1:4)
    end
    return act
end

function probability_super_variant(state_t, action_t)
    high_prob = 0.95/3
    low_prob = 0.05
    # UP, RIGHT, DOWN, LEFT
    # Assume target policy is going Down!
    # Generated with gen_states(1143139448, 25)
    # states = [[6, 7], [4, 8], [6, 9], [3, 11], [2, 10],
    #           [6, 2], [8, 1], [5, 2], [11, 8], [1, 2],
    #           [10, 1], [1, 5], [4, 7], [10, 11], [2, 9],
    #           [6, 5], [9, 4], [3, 6], [4, 1], [3, 2],
    #           [7, 5], [3, 9], [5, 6], [10, 9], [2, 1]]
    states = [[6, 7], [4, 8], [6, 9], [3, 11], [2, 10],
              [6, 2], [8, 1], [5, 2], [11, 8], [1, 2],
              [10, 1], [1, 5], [4, 7], [10, 11], [2, 9],
              [6, 5], [9, 4], [3, 6], [4, 1], [3, 2],
              [7, 5], [3, 9], [5, 6], [10, 9], [2, 1],
              [6, 3], [8, 11], [9, 5], [11, 3], [8, 2],
              [5, 9], [7, 11], [10, 10], [9, 8], [1, 10],
              [7, 9], [1, 7], [3, 1], [7, 10], [7, 7],
              [10, 2], [10, 8], [1, 9], [11, 9], [1, 6],
              [11, 10], [8, 8], [10, 5], [5, 11], [2, 2]]
    if state_t in states
        # Assume
        if action_t != 3
            return high_prob
        else
            return low_prob
        end
    else
        return 0.25
    end
    return act
end

super_variant_behaviour_policy = GeneralValueFunction.Policy(
    probability_super_variant, next_action_super_variant)

function next_action_rand_more_variant(states, state_t; rng=Random.GLOBAL_RNG)
    high_prob = 0.95/3
    low_prob = 0.05

    if state_t in states
        act = StatsBase.sample(rng, 1:4, Weights([high_prob, high_prob, low_prob, high_prob]))
    else
        act = rand(rng, 1:4)
    end
    return act
end

function probability_rand_more_variant(states, state_t, action_t)
    high_prob = 0.95/3
    low_prob = 0.05

    if state_t in states
        # Assume
        if action_t != 3
            return high_prob
        else
            return low_prob
        end
    else
        return 0.25
    end
    return act
end

mutable struct super_random
    dist
    function super_random(env_shape, num_actions; rng=Random.GLOBAL_RNG, ϵ=0.05)
        d = rand(rng, num_actions, env_shape[1], env_shape[2])
        d = clamp.(d, ϵ, 1)
        d ./= sum(d;dims=1)
        new(d)
    end
end

function next_action_random(strg::super_random, state_t; rng=Random.GLOBAL_RNG)
    high_prob = 0.95/3
    low_prob = 0.05
    # UP, RIGHT, DOWN, LEFT
    state_prob = strg.dist[:, state_t[1], state_t[2]]
    a = rand(rng)
    act = -1
    act_sum = 0
    for i in 1:length(state_prob)
        act_sum += state_prob[i]
        if act_sum < a
            return i
        end
    end
end

function probability_random(strg::super_random, state_t, action_t)
    state_prob = strg.dist[:, state_t[1], state_t[2]]
    return state_prob[action_t]
end

function create_random(env_size, num_actions; rng=Random.GLOBAL_RNG, strg=nothing)
    if strg == nothing
        strg = super_random(env_size, num_actions; rng=rng)
    end
    p(state_t, action_t) = probability_random(strg, state_t, action_t)
    act(state_t; rnng=rng) = next_action_random(strg, state_t; rng=rnng)
    return GeneralValueFunction(p, act)
end

# super_variant_behaviour_policy = GeneralValueFunction.Policy(
#     probability_super_variant, next_action_super_variant)

# Changing behaviour policy

mutable struct changing_behavior
    freq
    step
end

function next_action_cp(CB::changing_behavior, state_t; rng=Random.GLOBAL_RNG)

end

function probability_cp(CB::changing_behavior, state_t, action_t; rng=Random.GLOBAL_RNG)

end

function create_changing_behavior_policy(freq)
    changing_behavior(freq, 0)

end

ext: jl


text:
# LinearAlgebra.cross couldn't be applied to 2d vector
function twod_cross(x::AbstractVector, y::AbstractVector)
    @argcheck length(x) == length(y) == 2
    x[1] * y[2] - x[2] * y[1]
end

function moving_angle(x::AbstractVector, y::AbstractVector)
    @argcheck length(x) == length(y) == 2
    cos_theta = dot(x, y) / (norm(x) * norm(y))
    cos_theta = round(cos_theta, digits=6)
    # if cos_theta > 1.0
    #     println(x, y)
    # end
    theta = rad2deg(acos(cos_theta))
    theta = ifelse(twod_cross(x, y) < 0, -theta, theta)
    theta
end

# Naive implemation is 2times faster than using ifelse function
function spiff(x::AbstractArray, y::AbstractArray)
    @argcheck length(x) == length(y)
    total_len = length(x)

    frac_x = [modf(i)[1] for i in x]
    int_x = [modf(i)[2] for i in x]
    frac_y = [modf(i)[1] for i in y]
    int_y = [modf(i)[2] for i in y]

    x_plus = sort(frac_x[frac_x.>=0.5])
    x_minus = sort(frac_x[frac_x.<0.5])
    y_plus = sort(frac_y[frac_y.>=0.5])
    y_minus = sort(frac_y[frac_y.<0.5])

    x_plus_len, x_minus_len = length(x_plus), length(x_minus)
    y_plus_len, y_minus_len = length(y_plus), length(y_minus)

    corrected_x = zero(frac_x)
    corrected_y = zero(frac_y)

    @inbounds for i = 1:total_len
        # if frac_x[i] == 0
        #     continue
        # end
        # if frac_y[i] == 0
        #     continue
        # end
        #         corrected_x[i] = ifelse(
        #             frac_x[i] >= 0.5,
        #             int_x[i] + (sum(x_plus .<= frac_x[i])/x_plus_len)/2,
        #             int_x[i] - (sum(x_minus .> frac_x[i])/x_minus_len)/2
        #         )
        #          corrected_y[i] = ifelse(
        #             frac_x[i] >= 0.5,
        #             int_y[i] + (sum(y_plus .<= frac_y[i])/y_plus_len)/2,
        #             int_y[i] - (sum(y_minus .> frac_y[i])/y_minus_len)/2
        #         )
        if frac_x[i] >= 0.5
            tmp = sum(x_plus .<= frac_x[i]) / x_plus_len
            corrected_x[i] = int_x[i] + tmp / 2
        else
            tmp = sum((x_minus) .> (frac_x[i])) / x_minus_len
            corrected_x[i] = int_x[i] - tmp / 2
        end
        if frac_y[i] >= 0.5
            tmp = sum(y_plus .<= frac_y[i]) / y_plus_len
            corrected_y[i] = int_y[i] + tmp / 2
        else
            tmp = sum((y_minus) .> (frac_y[i])) / y_minus_len
            corrected_y[i] = int_y[i] - tmp / 2
        end
    end
    return corrected_x, corrected_y
end

function anisotropy_delta_t(
        df::DataFrames.DataFrame;
        maxt::Integer = 45,
        localization_error::Float64 = 0.03,
        num_of_resample = 50,
    )
    result = DataFrame(
        delta_t = Real[],
        cell_num = String[],
        fw = Real[],
        bw = Real[],
        n = Integer[],
        anisotropy = Float64[],
        std = Float64[],
        corrected_fw = Real[],
        corrected_bw = Real[],
        corrected_n = Integer[],
        corrected_anisotropy = Float64[],
        corrected_std = Float64[]
    )
    for t = 1:maxt
        tmp_data = df[
            (df.delta_t .== t) .& (df.dis_aft .>= localization_error) .&
            (df.dis_bfr .>= localization_error), :relative_angle]
        tmp_data = convert(Array, tmp_data)
        corrected_tmp_data = df[
            (df.delta_t .== t) .& (df.corrected_dis_aft .>= localization_error) .&
            (df.corrected_dis_bfr .>= localization_error), :corrected_angle]
        corrected_tmp_data = convert(Array, corrected_tmp_data)
        cell_num = df[1, :cell_num]

        fw, bw, n, anisotropy = calculate_fw_bw(tmp_data)
        corrected_fw, corrected_bw, corrected_n, corrected_anisotropy =
            calculate_fw_bw(corrected_tmp_data)

        boot_std = []
        corrected_boot_std = []
        for _ in 1:num_of_resample
            boot_df = tmp_data[
                rand(1:length(tmp_data), div(length(tmp_data), 2))]
            corrected_boot_df = corrected_tmp_data[
                rand(1:length(corrected_tmp_data), div(length(corrected_tmp_data), 2))]
            _, _, _, tmp_anisotropy = calculate_fw_bw(boot_df)
            _, _, _, corrected_tmp_anisotropy = calculate_fw_bw(corrected_boot_df)
            append!(boot_std, tmp_anisotropy)
            append!(corrected_boot_std, corrected_tmp_anisotropy)
        end

        push!(result, [
                t, cell_num, fw, bw, n, anisotropy, std(boot_std),
                corrected_fw, corrected_bw, corrected_n, corrected_anisotropy, std(corrected_boot_std)
                ])
    end
    return result
end

function calculate_fw_bw(data::AbstractArray)
    fw = size(data[-30 .<= data .<= 30, :], 1)
    bw = size(data[(150 .<= data) .|  (data .<= -150), :], 1)
    n = fw + bw
    anisotropy = bw / fw
    return fw, bw, n, anisotropy
end

function plot_anisotropy(grouped_df; maxt = 45, save_fig = false)
    for i = 1:length(grouped_df)
        result = []
        tmp_df = grouped_df[i]
        for t = 1:maxt
            a = tmp_df[tmp_df.delta_t.==t, :]
            push!(result, [t, mean(a.anisotropy), mean(a.std)])
        end

        result = reduce(vcat, result')
        genotype = tmp_df[1, :Genotype]
        plot(result[:, 1] * (1 / 45), result[:, 2], label = "$genotype")
        fill_between(
            result[:, 1] * (1 / 45),
            result[:, 2] .- result[:, 3],
            result[:, 2] .+ result[:, 3],
            alpha = 0.2,
        )
    end
    hlines([1], 0, 1, linestyle="--", color="k")

    xlim(0, maxt*(1/45))
    xlabel("Δt (sec)", fontsize = 14, family = "Arial")
    xticks(fontsize = 12)

    ylim(-1, )
    ylabel("Anisotropy", fontsize = 14, family = "Arial")
    yticks(fontsize = 12)

    legend(bbox_to_anchor = (1.02, 1.0), loc = 2, borderaxespad = 0.0, fontsize = 14)
    if save_fig
        savefig("anisotropy_dia0.25.png", bbox_inches = "tight", dpi = 800)
    end
end

# function anisotropy_mean_displacement(
#         df::AbstractDataFrame,
#         bin_array::Array{Float64,1}
#     )
#     result = DataFrame(
#         left = Float64[],
#         right = Float64[],
#         anisotropy = Float64[],
#         fw = Float64[],
#         bw = Float64[],
#         n = Int[]
#     )
#     for i in 1:length(bin_array)-1
#         tmp_array = df[df.bin .== i, :relative_angle]
#         fw, bw, n, anisotropy = SPT.calculate_fw_bw(tmp_array)
#         push!(result, [
#                 bin_array[i], bin_array[i+1], anisotropy, fw, bw, n
#             ])
#     end
#     return result
# end

function anisotropy_mean_displacement(
        df::AbstractDataFrame,
        bin_array::Array{Float64,1};
        num_resample::Integer = 50,
    )
    result = DataFrame(
        left = Float64[],
        right = Float64[],
        fw = Float64[],
        bw = Float64[],
        n = Int[],
        anisotropy = Float64[],
        std = Float64[],
        corrected_fw = Real[],
        corrected_bw = Real[],
        corrected_n = Integer[],
        corrected_anisotropy = Float64[],
        corrected_std = Float64[]
    )
    for i in 1:length(bin_array)-1
        tmp_data = df[df.bin .== i, :relative_angle]
        fw, bw, n, anisotropy = SPT.calculate_fw_bw(tmp_data)

        corrected_tmp_data = df[df.bin .== i, :corrected_angle]
        corrected_fw, corrected_bw, corrected_n, corrected_anisotropy =
            calculate_fw_bw(corrected_tmp_data)

        boot_std = []
        corrected_boot_std = []
        for _ in 1:num_resample
            boot_df = tmp_data[
                rand(1:length(tmp_data), div(length(tmp_data), 2))]
            corrected_boot_df = corrected_tmp_data[
                rand(1:length(corrected_tmp_data), div(length(corrected_tmp_data), 2))]
            _, _, _, tmp_anisotropy = calculate_fw_bw(boot_df)
            _, _, _, corrected_tmp_anisotropy = calculate_fw_bw(corrected_boot_df)
            append!(boot_std, tmp_anisotropy)
            append!(corrected_boot_std, corrected_tmp_anisotropy)
        end
        push!(result, [
                bin_array[i], bin_array[i+1],
                fw, bw, n, anisotropy, std(boot_std),
                corrected_fw, corrected_bw, corrected_n, corrected_anisotropy, std(corrected_boot_std)
        ])
    end
    return result
end

max_stars_repo_name: SosUts/spt

max_issues_repo_name: SosUts/spt

max_forks_repo_name: SosUts/spt


text:
islinked(program::GLuint) = glGetProgramiv(program, GL_LINK_STATUS) == GL_TRUE

const AttributeTuple = NamedTuple{(:name, :location, :T, :size), Tuple{Symbol, GLint, GLenum,  GLint}}
const UniformTuple = AttributeTuple
const INVALID_ATTRIBUTE = GLint(-1)
const INVALID_UNIFORM   = GLint(-1)

function setup_uniforms(program::GLuint)
    info = UniformTuple[]
    nuniforms = glGetProgramiv(program, GL_ACTIVE_UNIFORMS)
    for i=1:nuniforms
        name, typ, size = glGetActiveUniform(program, i-1)
        loc = glGetUniformLocation(program, name)
        push!(info, (name = name, location = loc, T = typ, size = size))
    end
    return info
end

function setup_attributes(program::GLuint)
    info = AttributeTuple[]
    nattribs = glGetProgramiv(program, GL_ACTIVE_ATTRIBUTES)
    for i=1:nattribs
        name, typ, size = glGetActiveAttrib(program, i-1)
        loc = glGetAttribLocation(program, name)
        push!(info, (name = name, location = loc, T = typ, size = size))
    end
    return info
end

abstract type AbstractProgram end
mutable struct Program <: AbstractProgram
    id        ::GLuint
    shaders   ::Vector{Shader}
    uniforms  ::Vector{UniformTuple}
    attributes::Vector{AttributeTuple}
    context   ::AbstractContext
    function Program(shaders::Vector{Shader}, fragdatalocation::Vector{Tuple{Int, String}})
        # Remove old shaders
        exists_context()
        program = glCreateProgram()::GLuint
        glUseProgram(program)
        #attach new ones
        foreach(shaders) do shader
            glAttachShader(program, shader.id)
        end

        #Bind frag data
        for (location, name) in fragdatalocation
            glBindFragDataLocation(program, location, ascii(name))
        end

        #link program
        glLinkProgram(program)
        if !islinked(program)
            for shader in shaders
                write(stdout, shader.source)
                println("---------------------------")
            end
            @error "program $program not linked. Error in: \n $(join(map(x-> string(x.name), shaders))), or, \n $(getinfolog(program))"
        end

        # generate the link locations
        uniforms = setup_uniforms(program)
        attribs  = setup_attributes(program)
        prog = new(program, shaders, uniforms, attribs, current_context())
        finalizer(free!, prog)
        prog
    end
end

Program(shaders::Vector{Shader}) = Program(shaders, Tuple{Int, String}[])

#REVIEW: This is a bit redundant seen as there is `load(source)` from FilIO for shaders but ok
function Program(sh_string_typ...)
    shaders = Shader[]
    for (source, typ) in sh_string_typ
        push!(shaders, Shader(gensym(), typ, Vector{UInt8}(source)))
    end
    Program(shaders)
end

attributes(program::Program) = program.attributes
uniforms(program::Program)   = program.uniforms
uniform_names(program::Program) = [x.name for x in program.uniforms]

attribute(program::Program, name::Symbol) =
    getfirst(x -> x.name == name, program.attributes)
uniform(program::Program, name::Symbol) =
    getfirst(x -> x.name == name, program.uniforms)

function attribute_location(program::Program, name::Symbol)
    att = attribute(program, name)
    return att != nothing ? att.location : INVALID_ATTRIBUTE
end

function uniform_location(program::Program, name::Symbol)
    u = uniform(program, name)
    return u != nothing ? u.location : INVALID_UNIFORM
end

function attribute_type(program::Program, name::Symbol)
    att = attribute(program, name)
    return att != nothing ? att.T : Nothing
end

function uniform_type(program::Program, name::Symbol)
    u = uniform(program, name)
    return u != nothing ? u.T : Nothing
end

function attribute_size(program::Program, name::Symbol)
    att = attribute(program, name)
    return att != nothing ? att.size : INVALID_ATTRIBUTE
end

function uniform_size(program::Program, name::Symbol)
    u = uniform(program, name)
    return u != nothing ? u.size : INVALID_UNIFORM
end

bind(program::Program) = glUseProgram(program.id)
unbind(program::AbstractProgram) = glUseProgram(0)

#REVIEW: Not sure if this is the best design decision
#REVIEW: Naming?
function set_uniform(program::Program, name::Symbol, vals::Tuple)
    loc = uniform_location(program, name)
    if loc != INVALID_UNIFORM
        gluniform(loc, vals...)
    end
end
function set_uniform(program::Program, name::Symbol, val)
    loc = uniform_location(program, name)
    if loc != INVALID_UNIFORM
        gluniform(loc, val)
    end
end

function Base.show(io::IO, p::Program)
    println(io, "Program: $(p.id)")
    println(io, "Shaders:")
    for shader in p.shaders
        println(io, shader)
    end
    println(io, "attributes:")
    for a in p.attributes
        println(io, "   ", a.name, "::", GLENUM(a.T).name)
    end
    println(io, "uniforms:")
    for u in p.uniforms
        println(io, "   ", u.name, "::", GLENUM(u.T).name)
    end
end

function infolog(program::Program)
    # Get the maximum possible length for the descriptive error message
    maxlength = GLint[0]
    glGetProgramiv(program.id, GL_INFO_LOG_LENGTH, maxlength)
    maxlength = first(maxlength)
    # Return the text of the message if there is any
    if maxlength > 0
        buffer = zeros(GLchar, maxlength)
        sizei = GLsizei[0]
        glGetProgramInfoLog(program.id, maxlength, sizei, buffer)
        length = first(sizei)
        return unsafe_string(pointer(buffer), length)
    else
        return "success"
    end
end

# display program's information
function program_info(p::Program)
    result = Dict{Symbol, Any}()
    program = p.id
    # check if name is really a program
    result[:program] = program
    # Get the shader's name
    result[:shaders] = glGetAttachedShaders(program)
    for shader in result[:shaders]
        result[Symbol("shader_type_$shader")] = GLENUM(convert(GLenum, glGetShaderiv(shader, GL_SHADER_TYPE))).name
    end
    # Get program info
    result[:program_seperable] = glGetProgramiv(program, GL_PROGRAM_SEPARABLE)
    result[:binary_retrievable_hint] = glGetProgramiv(program, GL_PROGRAM_BINARY_RETRIEVABLE_HINT)
    result[:link_status] = glGetProgramiv(program, GL_LINK_STATUS)
    result[:validate_status] = glGetProgramiv(program, GL_VALIDATE_STATUS)
    result[:delete_status] = glGetProgramiv(program, GL_DELETE_STATUS)
    result[:active_attributes] = glGetProgramiv(program, GL_ACTIVE_ATTRIBUTES)
    result[:active_uniforms] = glGetProgramiv(program, GL_ACTIVE_UNIFORMS)
    result[:active_uniform_blocks] = glGetProgramiv(program, GL_ACTIVE_UNIFORM_BLOCKS)
    result[:active_atomic_counter_buffers] = glGetProgramiv(program, GL_ACTIVE_ATOMIC_COUNTER_BUFFERS)
    result[:transform_feedback_buffer_mode] = glGetProgramiv(program, GL_TRANSFORM_FEEDBACK_BUFFER_MODE)
    result[:transform_feedback_varyings] = glGetProgramiv(program, GL_TRANSFORM_FEEDBACK_VARYINGS)
    result
end

function free!(x::Program)
    if !is_current_context(x.context)
        return x
    end
    try
        glDeleteProgram(x.id)
    catch e
        free_handle_error(e)
    end
    return
end

###############################################################################333

mutable struct LazyProgram <: AbstractProgram
    sources::Vector
    data::Dict
    compiled_program::Union{Program, Nothing}
end
LazyProgram(sources...; data...) = LazyProgram(Vector(sources), Dict(data), nothing)

function Program(lazy_program::LazyProgram)
    fragdatalocation = get(lazy_program.data, :fragdatalocation, Tuple{Int, String}[])
    shaders = haskey(lazy_program.data, :arguments) ? Shader.(lazy_program.sources, Ref(lazy_program.data[:arguments])) : Shader.()
    return Program([shaders...], fragdatalocation)
end
function bind(program::LazyProgram)
    iscompiled_orcompile!(program)
    bind(program.compiled_program)
end

function iscompiled_orcompile!(program::LazyProgram)
    if program.compiled_program == nothing
        program.compiled_program = Program(program)
    end
end

####################################################################################
# freeing

# OpenGL has the annoying habit of reusing id's when creating a new context
# We need to make sure to only free the current one

# display the values for a uniform in a named block
function uniform_in_block_info(p::Program, blockName, uniName)
    result = Dict{Symbol, Any}()
    program = p.id

    result[:index] = glGetUniformBlockIndex(program, blockName)
    if (index == GL_INVALID_INDEX)
        println("$uniName is not a valid uniform name in block $blockName")
    end
    result[:bindIndex] = glGetActiveUniformBlockiv(program, index, GL_UNIFORM_BLOCK_BINDING)
    result[:bufferIndex] = glGetIntegeri_v(GL_UNIFORM_BUFFER_BINDING, result[:bindIndex])
    result[:uniIndex] = glGetUniformIndices(program, uniName)
    uniIndex = result[:uniIndex]
    result[:uniType] = glGetActiveUniformsiv(program, uniIndex, GL_UNIFORM_TYPE)
    result[:uniOffset] = glGetActiveUniformsiv(program, uniIndex, GL_UNIFORM_OFFSET)
    result[:uniSize] = glGetActiveUniformsiv(program, uniIndex, GL_UNIFORM_SIZE)
    result[:uniArrayStride] = glGetActiveUniformsiv(program, uniIndex, GL_UNIFORM_ARRAY_STRIDE)
    result[:uniMatStride] = glGetActiveUniformsiv(program, uniIndex, GL_UNIFORM_MATRIX_STRIDE)
    result
end


text:
using DataStructures
import Base: isless

abstract type Event end
abstract type State end

# Captures an event and the time it takes place
struct TimedEvent
    event::Event
    time::Float64
end

# Comparison of two timed events - this will allow us to use them in a heap/priority-queue
isless(te1::TimedEvent, te2::TimedEvent) = te1.time < te2.time

"""
    new_timed_events = process_event(time, state, event)

Generate an array of 0 or more new `TimedEvent`s based on the current `event` and `state`.
"""
function process_event end # This defines a function with zero methods (to be added later)

# Generic events that we can always use

"""
    EndSimEvent()

Return an event that ends the simulation.
"""
struct EndSimEvent <: Event end

function process_event(time::Float64, state::State, es_event::EndSimEvent)
    println("Ending simulation at time $time.")
    return []
end

"""
    LogStateEvent()

Return an event that prints a log of the current simulation state.
"""
struct LogStateEvent <: Event end

function process_event(time::Float64, state::State, ls_event::LogStateEvent)
    println("Logging state at time $time.")
    println(state)
    return []
end

"""
The main simulation function gets an initial state and an initial event
that gets things going. Optional arguments are the maximal time for the
simulation, times for logging events, and a call-back function.
"""
function simulate(init_state::State, init_timed_event::TimedEvent
                    ;
                    max_time::Float64 = 10.0,
                    log_times::Vector{Float64} = Float64[],
                    callback = (time, state) -> nothing)

    # The event queue
    priority_queue = BinaryMinHeap{TimedEvent}()

    # Put the standard events in the queue
    push!(priority_queue, init_timed_event)
    push!(priority_queue, TimedEvent(EndSimEvent(), max_time))
    for log_time in log_times
        push!(priority_queue, TimedEvent(LogStateEvent(), log_time))
    end

    # initilize the state
    state = deepcopy(init_state)
    time = 0.0

    # Callback at simulation start
    callback(time, state)

    # The main discrete event simulation loop - SIMPLE!
    while true
        # Get the next event
        timed_event = pop!(priority_queue)

        # Advance the time
        time = timed_event.time

        # Act on the event
        new_timed_events = process_event(time, state, timed_event.event)

        # If the event was an end of simulation then stop
        if timed_event.event isa EndSimEvent
            break
        end

        # The event may spawn 0 or more events which we put in the priority queue
        for nte in new_timed_events
            push!(priority_queue,nte)
        end

        # Callback for each simulation event
        callback(time, state)
    end
end
"""
This function is designed to stich_steps of a discrete event curve.
"""
function stich_steps(epochs, values)
    n = length(epochs)
    new_epochs  = [epochs[1]]
    new_values = [values[1]]
    for i in 2:n
        push!(new_epochs, epochs[i])
        push!(new_values, values[i-1])
        push!(new_epochs, epochs[i])
        push!(new_values, values[i])
    end
    return (new_epochs, new_values)
end;


text:
using Statistics
using StatsBase: sample
using XGBoost
using EvoTrees
using BenchmarkTools

# prepare a dataset
features = rand(Int(1.25e6), 100)
# features = rand(100, 10)
X = features
Y = rand(size(X, 1))
𝑖 = collect(1:size(X,1))

# train-eval split
𝑖_sample = sample(𝑖, size(𝑖, 1), replace = false)
train_size = 0.8
𝑖_train = 𝑖_sample[1:floor(Int, train_size * size(𝑖, 1))]
𝑖_eval = 𝑖_sample[floor(Int, train_size * size(𝑖, 1))+1:end]

X_train, X_eval = X[𝑖_train, :], X[𝑖_eval, :]
Y_train, Y_eval = Y[𝑖_train], Y[𝑖_eval]

#######################
# xgboost
#######################
num_round = 100
param = ["max_depth" => 5,
         "eta" => 0.05,
         "objective" => "reg:linear",
         "print_every_n" => 5,
         "subsample" => 0.5,
         "colsample_bytree" => 0.5,
         "tree_method" => "hist",
         "max_bin" => 32]
metrics = ["rmse"]
@time xgboost(X_train, num_round, label = Y_train, param = param, metrics=metrics, silent=1);
@time dtrain = DMatrix(X_train, label = Y_train)
@time model_xgb = xgboost(dtrain, num_round, param = param, metrics=metrics, silent=1);
@time pred_train = XGBoost.predict(model_xgb, X_train)

# train model
params1 = EvoTreeRegressor(
    loss=:linear, metric=:none,
    nrounds=100,
    λ = 0.0, γ=0.0, η=0.05,
    max_depth = 6, min_weight = 1.0,
    rowsample=0.5, colsample=0.5, nbins=32)

# for 100k: 410.477 ms (44032 allocations: 182.68 MiB)
# for 1.25e6: 6.964114 seconds (6.05 M allocations: 2.350 GiB, 2.82% gc time)
# for 1.25e6 no eval: 6.200 s (44330 allocations: 2.19 GiB)
# for 1.25e6 mse with eval data: 6.321 s (45077 allocations: 2.19 GiB)
@time model, cache = init_evotree(params1, X_train, Y_train);
@time grow_evotree!(model, cache);
@time model = fit_evotree(params1, X_train, Y_train);
@btime model = fit_evotree(params1, X_train, Y_train);
@time pred_train = EvoTrees.predict(model, X_train)

@time model = fit_evotree(params1, X_train, Y_train, X_eval=X_eval, Y_eval=Y_eval, print_every_n=9999, early_stopping_rounds=9999);
@btime model = fit_evotree(params1, X_train, Y_train, X_eval=X_eval, Y_eval=Y_eval, print_every_n=9999, early_stopping_rounds=9999);

@time model = fit_evotree(params1, X_train, Y_train, early_stopping_rounds=10);
@time model = fit_evotree(params1, X_train, Y_train, print_every_n=2);

# @time model = grow_gbtree(X_train, Y_train, params1, X_eval = X_eval, Y_eval = Y_eval, print_every_n = 5);
# @btime model = grow_gbtree($X_train, $Y_train, $params1, X_eval = $X_eval, Y_eval = $Y_eval);
@time pred_train = predict(model, X_train)

#############################
# agaricus
#############################
function readlibsvm(fname::String, shape)
    dmx = zeros(Float32, shape)
    label = Float32[]
    fi = open(fname, "r")
    cnt = 1
    for line in eachline(fi)
        line = split(line, " ")
        push!(label, parse(Float64, line[1]))
        line = line[2:end]
        for itm in line
            itm = split(itm, ":")
            dmx[cnt, parse(Int, itm[1]) + 1] = parse(Int, itm[2])
        end
        cnt += 1
    end
    close(fi)
    return (dmx, label)
end

# we use auxiliary function to read LIBSVM format into julia Matrix
train_X, train_Y = readlibsvm("data/agaricus.txt.train", (6513, 126))
test_X, test_Y   = readlibsvm("data/agaricus.txt.test", (1611, 126))

#-------------Basic Training using XGBoost-----------------
# note: xgboost naturally handles sparse input
# use sparse matrix when your feature is sparse(e.g. when you using one-hot encoding vector)
# model parameters can be set as parameters for ```xgboost``` function, or use a Vector{String} / Dict()
num_round = 100
# you can directly pass Julia's matrix or sparse matrix as data,
# by calling xgboost(data, num_round, label=label, training-parameters)
metrics = ["logloss"]
@time bst = xgboost(train_X, num_round, label = train_Y, eta = 0.1, max_depth = 3, metrics = metrics, silent=0, objective = "binary:logistic")
features_xgb = XGBoost.importance(bst)

X_train = Float64.(train_X)
Y_train = Float64.(train_Y)
params1 = EvoTreeRegressor(
    loss=:logistic, metric=:logloss,
    nrounds=100,
    λ = 0.0, γ=0.0, η=0.1,
    max_depth = 4, min_weight = 1.0,
    rowsample=1.0, colsample=1.0, nbins=250)

@time model = fit_evotree(params1, X_train, Y_train, print_every_n=50);
@time pred_train = EvoTrees.predict(model, X_train)
features_evo = importance(model, 1:size(X_train,2))
sort(collect(values(features_evo)))


text:
# Extracting these closures into types improves type stability a lot.
# Essentially, I have reinveted [zippers](https://wiki.haskell.org/Zipper) ;)

abstract type Context{N, C} end

struct LeftContext{N, C} <: Context{N, C}
    conditions::Dict{Variable{N}, Float64}
    threshold::Float64
    right::DecisionTree{N, C}
    parent::Context{N, C}
end

function (context::LeftContext{N, C})(t::DecisionTree{N, C}) where {N, C}
    context.parent(Decision{N, C}(context.conditions, context.threshold, t, context.right))
end

struct RightContext{N, C} <: Context{N, C}
    conditions::Dict{Variable{N}, Float64}
    threshold::Float64
    left::DecisionTree{N, C}
    parent::Context{N, C}
end

function (context::RightContext{N, C})(t::DecisionTree{N, C}) where {N, C}
    context.parent(Decision{N, C}(context.conditions, context.threshold, context.left, t))
end

struct NoContext{N, C} <: Context{N, C} end

function (context::NoContext{N, C})(t::DecisionTree{N, C}) where {N, C}
    t
end

struct Cont{N, C}
    context::Context{N, C}
    chunk::DecisionTree{N, C}
end

function (continuation::Cont{N, C})(k::Function) where {N, C}
    newchunk = k(continuation.chunk)::DecisionTree{N, C}
    tree = continuation.context(newchunk)::DecisionTree{N, C}
    return tree, continuation.chunk
end

function randsplit_impl{N, C}(rng::AbstractRNG, t::Classification{N, C}, n::Int,
                                parent_context::Context{N, C}, continuation::Cont{N, C})
    n += 1
    if rand() ≤ 1/n
        return Cont{N, C}(parent_context, t), n
    else
        return continuation, n
    end
end

function randsplit_impl{N, C}(rng::AbstractRNG, t::Decision{N, C}, n,
                                parent_context::Context{N, C}, continuation::Cont{N, C})
    c1, n = randsplit_impl(rng, t.iftrue, n,
                             LeftContext{N, C}(t.conditions, t.threshold, t.iffalse, parent_context),
                             continuation)
    c2, n = randsplit_impl(rng, t.iffalse, n,
                             RightContext{N, C}(t.conditions, t.threshold, t.iftrue, parent_context),
                             c1)

    n += 1
    if rand() ≤ 1/n
        return Cont{N, C}(parent_context, t), n
    else
        return c2, n
    end
end

"""
    randsplit(action,[ rng,] t::DecisionTree)

Select uniformly at random a point to split `t`, then replace the subtree `s` by `action(s)`.
Returns both the new tree and the replaced subtree.
"""
function randsplit{N, C}(action, rng::AbstractRNG, t::DecisionTree{N, C})
    # Uses reservoir samling of continuations, see: https://stackoverflow.com/a/3272490/1346276.
    default_context = NoContext{N, C}()
    default_continuation = Cont{N, C}(default_context, t)
    cont, _ = randsplit_impl(rng, t, 0, default_context, default_continuation)
    cont(action)
end

randsplit{N, C}(action, t::DecisionTree{N, C}) = randsplit(action, Base.GLOBAL_RNG, t)

"""
    randchild([rng,] t::DecisionTree)

Select uniformly at random a subtree of `t`.
"""
randchild{N, C}(rng::AbstractRNG, t::DecisionTree{N, C}) = randsplit(identity, rng, t)[2]
randchild{N, C}(t::DecisionTree{N, C}) = randchild(Base.GLOBAL_RNG, t)

max_stars_repo_name: phipsgabler/soft-computing-final

max_issues_repo_name: phipsgabler/soft-computing-final

max_forks_repo_name: phipsgabler/soft-computing-final


text:
module TutorialsBuilder

# NOTE: while the clever_cd from PkgPage can be reused,
# you certainly don't want a "clear=true" always serve.
#

end # module


text:
module CMDPs

export generate_sim_q, run_study
export RewardvsN, RewardvsNResult
export TuneDPW, TuneDPWResult, report
export TuneCBTS, TuneCBTSResult

using POMDPs, POMDPToolbox, MCTS
using Parameters, DataFrames
using Walk1DMDP, VDPTag2
using POMDPModels #MountainCar, InvertedPendulum
using LightDarkPOMDPs
using RecipesBase

include("walk1d.jl")
include("vdptag.jl")
include("inverted.jl")
include("mountaincar.jl")
include("lightdark2d.jl")
include("sim.jl")
include("studies.jl")
include("tuning.jl")
include("observers.jl")

end # module


text:
using Distributions
using Flux
using Plots

mutable struct HMCState
    position::Union{Number, Array{Float64,1}}
    momentum::Union{Number, Array{Float64,1}}
    step_size::Float64
    num_steps::Int64
end

function HMCState(q, p, ssize, nsteps)
    position = q
    momentum = p
    step_size = ssize
    num_steps = nsteps
    HMCState(position, momentum, step_size, num_steps)
end

function leapfrog!(z::HMCState, gradU::Function, gradK::Function)
    step_size = z.step_size
    num_steps = z.num_steps
    q = z.position
    p = z.momentum

    # new state according to leapfrom for hamilton dynamics
    p -= step_size .* gradU.(q) / 2
    for jump in 1:1:num_steps-1
        q += step_size .* gradK.(p)
        p -= step_size .* gradU.(q)
    end
    q += step_size .* gradK.(p)
    p -= step_size .* gradU.(q) / 2 # the second half for the kinetic variables

    z.position = q
    z.momentum = -p
    return nothing
end

function sg_hmc!(state::HMCState, U::Function, gradU::Function,
                K::Function, gradK::Function; masses=Float64[])

    current_q = state.position
    d = length(current_q)
    masses = isempty(masses) ? 1.0 : masses
    state.momentum =  d > 1 ? rand(MvNormal(zeros(d), Diagonal(masses))) : rand(Normal(0, sqrt(masses)))
    current_p = state.momentum

    # new state according to leapfrom for hamilton dynamics
    leapfrog!(state, gradU, gradK)
    p = state.momentum
    q = state.position
    # Evaluate potential and kinetic energies at start and end of trajectory
    current_U = U(current_q)
    current_K = sum(current_p.^2 ./ (2*masses))
    proposed_U = U(q)
    proposed_K = sum(p.^2 / (2*masses))

    if  rand(Float64) .< exp(current_U - proposed_U + current_K - proposed_K)
        state.position = q
    else
        state.position = current_q
    end
    return nothing
end

function test_hmc(n=100000)
    samples = zeros(n)
    q, p, step_size, num_steps = 0.0, 1.0, 0.1, 10
    x = HMCState(q, p, step_size, num_steps)
    metric = 1.0
    U(x) = -logpdf(Normal(12.5,2.2), x)
    K(x) = x' * inv(metric) * x
    ∇U(x) = gradient(U, x)[1]     # numerical gradient
    ∇K(x) = gradient(K, x)[1]     # numerical gradient
    for i in 1:n
        sg_hmc!(x, U, ∇U, K, ∇K, masses=metric)
        samples[i] = x.position
    end
    return samples
end

@time samples = test_hmc()

histogram(samples)
mean(samples)
std(samples)

max_stars_repo_name: cmerkatas/HMC

max_issues_repo_name: cmerkatas/HMC

max_forks_repo_name: cmerkatas/HMC


text:
### CONCRETE TYPE: ITERATIVE PROX EVALUATION

struct IndAffineIterative{R <: Real, T <: RealOrComplex{R}, M <: AbstractMatrix{T}, V <: AbstractVector{T}} <: IndAffine
    A::M
    b::V
    res::V
    maxit::Integer
    tol::R
    function IndAffineIterative{R, T, M, V}(A::M, b::V) where {R <: Real, T <: RealOrComplex{R}, M <: AbstractMatrix{T}, V <: AbstractVector{T}}
        if size(A,1) > size(A,2)
            error("A must be full row rank")
        end
        normrowsinv = 1 ./ vec(sqrt.(sum(abs2.(A); dims=2)))
        A = normrowsinv.*A # normalize rows of A
        b = normrowsinv.*b # and b accordingly
        new(A, b, similar(b), 1000, 1e-8)
    end
end

is_prox_accurate(f::IndAffineIterative) = false
is_cone(f::IndAffineIterative) = norm(f.b) == 0.0

IndAffineIterative(A::M, b::V) where {R <: Real, T <: RealOrComplex{R}, M <: AbstractMatrix{T}, V <: AbstractVector{T}} = IndAffineIterative{R, T, M, V}(A, b)

function (f::IndAffineIterative{R, T, M, V})(x::V) where {R, T, M, V}
    mul!(f.res, f.A, x)
    f.res .= f.b .- f.res
    # the tolerance in the following line should be customizable
    if norm(f.res, Inf) <= f.tol
        return R(0)
    end
    return typemax(R)
end

function prox!(y::V, f::IndAffineIterative{R, T, M, V}, x::V, gamma::R=R(1)) where {R, T, M, V}
    # Von Neumann's alternating projections
    m = size(f.A, 1)
    y .= x
    for k = 1:f.maxit
        maxres = R(0)
        for i = 1:m
            resi = (f.b[i] - dot(f.A[i,:], y))
            y .= y + resi*f.A[i,:] # no need to divide: rows of A are normalized
            absresi = resi > 0 ? resi : -resi
            maxres = absresi > maxres ? absresi : maxres
        end
        if maxres < f.tol
            break
        end
    end
    return R(0)
end

function prox_naive(f::IndAffineIterative, x::AbstractArray{T,1}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    y = x + f.A'*((f.A*f.A')\(f.b - f.A*x))
    return y, R(0)
end


text:
#########################
# choice map interface #
#########################

"""
    schema = get_address_schema(::Type{T}) where {T <: ChoiceMap}

Return the (top-level) address schema for the given choice map.
"""
function get_address_schema end

"""
    submap = get_submap(choices::ChoiceMap, addr)

Return the sub-assignment containing all choices whose address is prefixed by addr.

It is an error if the assignment contains a value at the given address. If
there are no choices whose address is prefixed by addr then return an
`EmptyChoiceMap`.
"""
function get_submap end

"""
    value = get_value(choices::ChoiceMap, addr)

Return the value at the given address in the assignment, or throw a KeyError if
no value exists. A syntactic sugar is `Base.getindex`:

    value = choices[addr]
"""
function get_value end

"""
    key_submap_iterable = get_submaps_shallow(choices::ChoiceMap)

Return an iterable collection of tuples `(key, submap::ChoiceMap)` for each top-level key
that has a non-empty sub-assignment.
"""
function get_submaps_shallow end

"""
    has_value(choices::ChoiceMap, addr)

Return true if there is a value at the given address.
"""
function has_value end

"""
    key_submap_iterable = get_values_shallow(choices::ChoiceMap)

Return an iterable collection of tuples `(key, value)` for each
top-level key associated with a value.
"""
function get_values_shallow end

"""
    abstract type ChoiceMap end

Abstract type for maps from hierarchical addresses to values.
"""
abstract type ChoiceMap end

"""
    Base.isempty(choices::ChoiceMap)

Return true if there are no values in the assignment.
"""
function Base.isempty(::ChoiceMap)
    true
end

get_submap(choices::ChoiceMap, addr) = EmptyChoiceMap()
has_value(choices::ChoiceMap, addr) = false
get_value(choices::ChoiceMap, addr) = throw(KeyError(addr))
Base.getindex(choices::ChoiceMap, addr) = get_value(choices, addr)

function _has_value(choices::T, addr::Pair) where {T <: ChoiceMap}
    (first, rest) = addr
    submap = get_submap(choices, first)
    has_value(submap, rest)
end

function _get_value(choices::T, addr::Pair) where {T <: ChoiceMap}
    (first, rest) = addr
    submap = get_submap(choices, first)
    get_value(submap, rest)
end

function _get_submap(choices::T, addr::Pair) where {T <: ChoiceMap}
    (first, rest) = addr
    submap = get_submap(choices, first)
    get_submap(submap, rest)
end

function _print(io::IO, choices::ChoiceMap, pre, vert_bars::Tuple)
    VERT = '\u2502'
    PLUS = '\u251C'
    HORZ = '\u2500'
    LAST = '\u2514'
    indent_vert = vcat(Char[' ' for _ in 1:pre], Char[VERT, '\n'])
    indent_vert_last = vcat(Char[' ' for _ in 1:pre], Char[VERT, '\n'])
    indent = vcat(Char[' ' for _ in 1:pre], Char[PLUS, HORZ, HORZ, ' '])
    indent_last = vcat(Char[' ' for _ in 1:pre], Char[LAST, HORZ, HORZ, ' '])
    for i in vert_bars
        indent_vert[i] = VERT
        indent[i] = VERT
        indent_last[i] = VERT
    end
    indent_vert_str = join(indent_vert)
    indent_vert_last_str = join(indent_vert_last)
    indent_str = join(indent)
    indent_last_str = join(indent_last)
    key_and_values = collect(get_values_shallow(choices))
    key_and_submaps = collect(get_submaps_shallow(choices))
    n = length(key_and_values) + length(key_and_submaps)
    cur = 1
    for (key, value) in key_and_values
        print(io, indent_vert_str)
        print(io, (cur == n ? indent_last_str : indent_str) * "$(repr(key)) : $value\n")
        cur += 1
    end
    for (key, submap) in key_and_submaps
        print(io, indent_vert_str)
        print(io, (cur == n ? indent_last_str : indent_str) * "$(repr(key))\n")
        _print(io, submap, pre + 4, cur == n ? (vert_bars...,) : (vert_bars..., pre+1))
        cur += 1
    end
end

function Base.print(io::IO, choices::ChoiceMap)
    _print(io, choices, 0, ())
end

# assignments that have static address schemas should also support faster
# accessors, which make the address explicit in the type (Val(:foo) instaed of
# :foo)
function static_get_value end
function static_get_submap end

function _fill_array! end
function _from_array end

"""
    arr::Vector{T} = to_array(choices::ChoiceMap, ::Type{T}) where {T}

Populate an array with values of choices in the given assignment.

It is an error if each of the values cannot be coerced into a value of the
given type.

# Implementation

To support `to_array`, a concrete subtype `T <: ChoiceMap` should implement
the following method:

    n::Int = _fill_array!(choices::T, arr::Vector{V}, start_idx::Int) where {V}

Populate `arr` with values from the given assignment, starting at `start_idx`,
and return the number of elements in `arr` that were populated.
"""
function to_array(choices::ChoiceMap, ::Type{T}) where {T}
    arr = Vector{T}(undef, 32)
    n = _fill_array!(choices, arr, 0)
    @assert n <= length(arr)
    resize!(arr, n)
    arr
end

"""
    choices::ChoiceMap = from_array(proto_choices::ChoiceMap, arr::Vector)

Return an assignment with the same address structure as a prototype
assignment, but with values read off from the given array.

The order in which addresses are populated is determined by the prototype
assignment. It is an error if the number of choices in the prototype assignment
is not equal to the length the array.

# Implementation

To support `from_array`, a concrete subtype `T <: ChoiceMap` should implement
the following method:

    (n::Int, choices::T) = _from_array(proto_choices::T, arr::Vector{V}, start_idx::Int) where {V}

Return an assignment with the same address structure as a prototype assignment,
but with values read off from `arr`, starting at position `start_idx`, and the
number of elements read from `arr`.
"""
function from_array(proto_choices::ChoiceMap, arr::Vector)
    (n, choices) = _from_array(proto_choices, arr, 0)
    if n != length(arr)
        error("Dimension mismatch: $n, $(length(arr))")
    end
    choices
end

"""
    choices = Base.merge(choices1::ChoiceMap, choices2::ChoiceMap)

Merge two assignments.

It is an error if the assignments both have values at the same address, or if
one assignment has a value at an address that is the prefix of the address of a
value in the other assignment.
"""
function Base.merge(choices1::ChoiceMap, choices2::ChoiceMap)
    choices = DynamicChoiceMap()
    for (key, value) in get_values_shallow(choices1)
        choices.leaf_nodes[key] = value
    end
    for (key, node1) in get_submaps_shallow(choices1)
        node2 = get_submap(choices2, key)
        node = merge(node1, node2)
        choices.internal_nodes[key] = node
    end
    for (key, value) in get_values_shallow(choices2)
        if haskey(choices.leaf_nodes, key)
            error("choices1 has leaf node at $key and choices2 has leaf node at $key")
        end
        if haskey(choices.internal_nodes, key)
            error("choices1 has internal node at $key and choices2 has leaf node at $key")
        end
        choices.leaf_nodes[key] = value
    end
    for (key, node) in get_submaps_shallow(choices2)
        if haskey(choices.leaf_nodes, key)
            error("choices1 has leaf node at $key and choices2 has internal node at $key")
        end
        if !haskey(choices.internal_nodes, key)
            # otherwise it should already be included
            choices.internal_nodes[key] = node
        end
    end
    return choices
end

function Base.:(==)(a::ChoiceMap, b::ChoiceMap)
    for (addr, value) in get_values_shallow(a)
        if !has_value(b, addr) || (get_value(b, addr) != value)
            return false
        end
    end
    for (addr, value) in get_values_shallow(b)
        if !has_value(a, addr) || (get_value(a, addr) != value)
            return false
        end
    end
    for (addr, submap) in get_submaps_shallow(a)
        if submap != get_submap(b, addr)
            return false
        end
    end
    for (addr, submap) in get_submaps_shallow(b)
        if submap != get_submap(a, addr)
            return false
        end
    end
    return true
end

function Base.isapprox(a::ChoiceMap, b::ChoiceMap)
    for (addr, value) in get_values_shallow(a)
        if !has_value(b, addr) || !isapprox(get_value(b, addr), value)
            return false
        end
    end
    for (addr, value) in get_values_shallow(b)
        if !has_value(a, addr) || !isapprox(get_value(a, addr), value)
            return false
        end
    end
    for (addr, submap) in get_submaps_shallow(a)
        if !isapprox(submap, get_submap(b, addr))
            return false
        end
    end
    for (addr, submap) in get_submaps_shallow(b)
        if !isapprox(submap, get_submap(a, addr))
            return false
        end
    end
    return true
end

export ChoiceMap
export get_address_schema
export get_submap
export get_value
export has_value
export get_submaps_shallow
export get_values_shallow
export static_get_value
export static_get_submap
export to_array, from_array

######################
# static assignment #
######################

struct StaticChoiceMap{R,S,T,U} <: ChoiceMap
    leaf_nodes::NamedTuple{R,S}
    internal_nodes::NamedTuple{T,U}
    isempty::Bool
end

function StaticChoiceMap{R,S,T,U}(leaf_nodes::NamedTuple{R,S}, internal_nodes::NamedTuple{T,U}) where {R,S,T,U}
    is_empty = length(leaf_nodes) == 0 && all(isempty(n) for n in internal_nodes)
    StaticChoiceMap(leaf_nodes, internal_nodes, is_empty)
end

function StaticChoiceMap(leaf_nodes::NamedTuple{R,S}, internal_nodes::NamedTuple{T,U}) where {R,S,T,U}
    is_empty = length(leaf_nodes) == 0 && all(isempty(n) for n in internal_nodes)
    StaticChoiceMap(leaf_nodes, internal_nodes, is_empty)
end

# invariant: all internal_nodes are nonempty

function get_address_schema(::Type{StaticChoiceMap{R,S,T,U}}) where {R,S,T,U}
    keys = Set{Symbol}()
    for (key, _) in zip(R, S.parameters)
        push!(keys, key)
    end
    for (key, _) in zip(T, U.parameters)
        push!(keys, key)
    end
    StaticAddressSchema(keys)
end

function Base.isempty(choices::StaticChoiceMap)
    choices.isempty
end

get_values_shallow(choices::StaticChoiceMap) = pairs(choices.leaf_nodes)
get_submaps_shallow(choices::StaticChoiceMap) = pairs(choices.internal_nodes)
has_value(choices::StaticChoiceMap, addr::Pair) = _has_value(choices, addr)
get_value(choices::StaticChoiceMap, addr::Pair) = _get_value(choices, addr)
get_submap(choices::StaticChoiceMap, addr::Pair) = _get_submap(choices, addr)

# NOTE: there is no static_has_value because this is known from the static
# address schema

## has_value ##

function has_value(choices::StaticChoiceMap, key::Symbol)
    haskey(choices.leaf_nodes, key)
end

## get_submap ##

function get_submap(choices::StaticChoiceMap, key::Symbol)
    if haskey(choices.internal_nodes, key)
        choices.internal_nodes[key]
    elseif haskey(choices.leaf_nodes, key)
        throw(KeyError(key))
    else
        EmptyChoiceMap()
    end
end

function static_get_submap(choices::StaticChoiceMap, ::Val{A}) where {A}
    choices.internal_nodes[A]
end

## get_value ##

function get_value(choices::StaticChoiceMap, key::Symbol)
    choices.leaf_nodes[key]
end

function static_get_value(choices::StaticChoiceMap, ::Val{A}) where {A}
    choices.leaf_nodes[A]
end

# convert from any other schema that has only Val{:foo} addresses
function StaticChoiceMap(other::ChoiceMap)
    leaf_keys_and_nodes = collect(get_values_shallow(other))
    internal_keys_and_nodes = collect(get_submaps_shallow(other))
    if length(leaf_keys_and_nodes) > 0
        (leaf_keys, leaf_nodes) = collect(zip(leaf_keys_and_nodes...))
    else
        (leaf_keys, leaf_nodes) = ((), ())
    end
    if length(internal_keys_and_nodes) > 0
        (internal_keys, internal_nodes) = collect(zip(internal_keys_and_nodes...))
    else
        (internal_keys, internal_nodes) = ((), ())
    end
    StaticChoiceMap(
        NamedTuple{leaf_keys}(leaf_nodes),
        NamedTuple{internal_keys}(internal_nodes),
        isempty(other))
end

"""
    choices = pair(choices1::ChoiceMap, choices2::ChoiceMap, key1::Symbol, key2::Symbol)

Return an assignment that contains `choices1` as a sub-assignment under `key1`
and `choices2` as a sub-assignment under `key2`.
"""
function pair(choices1::ChoiceMap, choices2::ChoiceMap, key1::Symbol, key2::Symbol)
    StaticChoiceMap(NamedTuple(), NamedTuple{(key1,key2)}((choices1, choices2)),
        isempty(choices1) && isempty(choices2))
end

"""
    (choices1, choices2) = unpair(choices::ChoiceMap, key1::Symbol, key2::Symbol)

Return the two sub-assignments at `key1` and `key2`, one or both of which may be empty.

It is an error if there are any top-level values, or any non-empty top-level
sub-assignments at keys other than `key1` and `key2`.
"""
function unpair(choices::ChoiceMap, key1::Symbol, key2::Symbol)
    if !isempty(get_values_shallow(choices)) || length(collect(get_submaps_shallow(choices))) > 2
        error("Not a pair")
    end
    a = get_submap(choices, key1)
    b = get_submap(choices, key2)
    (a, b)
end

# TODO make it a generated function?
function _fill_array!(choices::StaticChoiceMap, arr::Vector{T}, start_idx::Int) where {T}
    if length(arr) < start_idx + length(choices.leaf_nodes)
        resize!(arr, 2 * length(arr))
    end
    for (i, value) in enumerate(choices.leaf_nodes)
        arr[start_idx + i] = value
    end
    idx = start_idx + length(choices.leaf_nodes)
    for (i, node) in enumerate(choices.internal_nodes)
        n_written = _fill_array!(node, arr, idx)
        idx += n_written
    end
    idx - start_idx
end

@generated function _from_array(proto_choices::StaticChoiceMap{R,S,T,U}, arr::Vector{V}, start_idx::Int) where {R,S,T,U,V}
    leaf_node_keys = proto_choices.parameters[1]
    leaf_node_types = proto_choices.parameters[2].parameters
    internal_node_keys = proto_choices.parameters[3]
    internal_node_types = proto_choices.parameters[4].parameters

    # leaf nodes
    leaf_node_refs = Expr[]
    for (i, key) in enumerate(leaf_node_keys)
        expr = Expr(:ref, :arr, Expr(:call, :(+), :start_idx, QuoteNode(i)))
        push!(leaf_node_refs, expr)
    end

    # internal nodes
    internal_node_blocks = Expr[]
    internal_node_names = Symbol[]
    for (key, typ) in zip(internal_node_keys, internal_node_types)
        node = gensym()
        push!(internal_node_names, node)
        push!(internal_node_blocks, quote
            (n_read, $node::$typ) = _from_array(proto_choices.internal_nodes.$key, arr, idx)
            idx += n_read
        end)
    end

    quote
        n_read = $(QuoteNode(length(leaf_node_keys)))
        leaf_nodes_field = NamedTuple{R,S}(($(leaf_node_refs...),))
        idx::Int = start_idx + n_read
        $(internal_node_blocks...)
        internal_nodes_field = NamedTuple{T,U}(($(internal_node_names...),))
        choices = StaticChoiceMap{R,S,T,U}(leaf_nodes_field, internal_nodes_field)
        (idx - start_idx, choices)
    end
end

@generated function Base.merge(choices1::StaticChoiceMap{R,S,T,U},
                               choices2::StaticChoiceMap{W,X,Y,Z}) where {R,S,T,U,W,X,Y,Z}

    # unpack first assignment type parameters
    leaf_node_keys1 = choices1.parameters[1]
    leaf_node_types1 = choices1.parameters[2].parameters
    internal_node_keys1 = choices1.parameters[3]
    internal_node_types1 = choices1.parameters[4].parameters
    keys1 = (leaf_node_keys1..., internal_node_keys1...,)

    # unpack second assignment type parameters
    leaf_node_keys2 = choices2.parameters[1]
    leaf_node_types2 = choices2.parameters[2].parameters
    internal_node_keys2 = choices2.parameters[3]
    internal_node_types2 = choices2.parameters[4].parameters
    keys2 = (leaf_node_keys2..., internal_node_keys2...,)

    # leaf vs leaf collision is an error
    colliding_leaf_leaf_keys = intersect(leaf_node_keys1, leaf_node_keys2)
    if !isempty(colliding_leaf_leaf_keys)
        error("choices1 and choices2 both have leaf nodes at key(s): $colliding_leaf_leaf_keys")
    end

    # leaf vs internal collision is an error
    colliding_leaf_internal_keys = intersect(leaf_node_keys1, internal_node_keys2)
    if !isempty(colliding_leaf_internal_keys)
        error("choices1 has leaf node and choices2 has internal node at key(s): $colliding_leaf_internal_keys")
    end

    # internal vs leaf collision is an error
    colliding_internal_leaf_keys = intersect(internal_node_keys1, leaf_node_keys2)
    if !isempty(colliding_internal_leaf_keys)
        error("choices1 has internal node and choices2 has leaf node at key(s): $colliding_internal_leaf_keys")
    end

    # internal vs internal collision is not an error, recursively call merge
    colliding_internal_internal_keys = (intersect(internal_node_keys1, internal_node_keys2)...,)
    internal_node_keys1_exclusive = (setdiff(internal_node_keys1, internal_node_keys2)...,)
    internal_node_keys2_exclusive = (setdiff(internal_node_keys2, internal_node_keys1)...,)

    # leaf nodes named tuple
    leaf_node_keys = (leaf_node_keys1..., leaf_node_keys2...,)
    leaf_node_types = map(QuoteNode, (leaf_node_types1..., leaf_node_types2...,))
    leaf_node_values = Expr(:tuple,
        [Expr(:(.), :(choices1.leaf_nodes), QuoteNode(key))
            for key in leaf_node_keys1]...,
        [Expr(:(.), :(choices2.leaf_nodes), QuoteNode(key))
            for key in leaf_node_keys2]...)
    leaf_nodes = Expr(:call,
        Expr(:curly, :NamedTuple,
            QuoteNode(leaf_node_keys),
            Expr(:curly, :Tuple, leaf_node_types...)),
        leaf_node_values)

    # internal nodes named tuple
    internal_node_keys = (internal_node_keys1_exclusive...,
                          internal_node_keys2_exclusive...,
                          colliding_internal_internal_keys...)
    internal_node_values = Expr(:tuple,
        [Expr(:(.), :(choices1.internal_nodes), QuoteNode(key))
            for key in internal_node_keys1_exclusive]...,
        [Expr(:(.), :(choices2.internal_nodes), QuoteNode(key))
            for key in internal_node_keys2_exclusive]...,
        [Expr(:call, :merge,
                Expr(:(.), :(choices1.internal_nodes), QuoteNode(key)),
                Expr(:(.), :(choices2.internal_nodes), QuoteNode(key)))
            for key in colliding_internal_internal_keys]...)
    internal_nodes = Expr(:call,
        Expr(:curly, :NamedTuple, QuoteNode(internal_node_keys)),
        internal_node_values)

    # construct assignment from named tuples
    Expr(:call, :StaticChoiceMap, leaf_nodes, internal_nodes)
end

export StaticChoiceMap
export pair, unpair

#######################
# dynamic assignment #
#######################

struct DynamicChoiceMap <: ChoiceMap
    leaf_nodes::Dict{Any,Any}
    internal_nodes::Dict{Any,Any}
    function DynamicChoiceMap(leaf_nodes::Dict{Any,Any}, internal_nodes::Dict{Any,Any})
        new(leaf_nodes, internal_nodes)
    end
end

# invariant: all internal nodes are nonempty

"""
    struct DynamicChoiceMap <: ChoiceMap .. end

A mutable map from arbitrary hierarchical addresses to values.

    choices = DynamicChoiceMap()

Construct an empty map.

    choices = DynamicChoiceMap(tuples...)

Construct a map containing each of the given (addr, value) tuples.
"""
function DynamicChoiceMap()
    DynamicChoiceMap(Dict(), Dict())
end

function DynamicChoiceMap(tuples...)
    choices = DynamicChoiceMap()
    for (addr, value) in tuples
        choices[addr] = value
    end
    choices
end

"""
    choices = DynamicChoiceMap(other::ChoiceMap)

Copy a choice map, returning a mutable choice map.
"""
function DynamicChoiceMap(other::ChoiceMap)
    choices = DynamicChoiceMap()
    for (addr, val) in get_values_shallow(other)
        choices[addr] = val
    end
    for (addr, submap) in get_submaps_shallow(other)
        set_submap!(choices, addr, DynamicChoiceMap(submap))
    end
    choices
end

"""
    choices = choicemap()

Construct an empty mutable choice map.
"""
function choicemap()
    DynamicChoiceMap()
end

"""
    choices = choicemap(tuples...)

Construct a mutable choice map initialized with given address, value tuples.
"""
function choicemap(tuples...)
    DynamicChoiceMap(tuples...)
end

get_address_schema(::Type{DynamicChoiceMap}) = DynamicAddressSchema()

get_values_shallow(choices::DynamicChoiceMap) = choices.leaf_nodes

get_submaps_shallow(choices::DynamicChoiceMap) = choices.internal_nodes

has_value(choices::DynamicChoiceMap, addr::Pair) = _has_value(choices, addr)

get_value(choices::DynamicChoiceMap, addr::Pair) = _get_value(choices, addr)

get_submap(choices::DynamicChoiceMap, addr::Pair) = _get_submap(choices, addr)

function get_submap(choices::DynamicChoiceMap, addr)
    if haskey(choices.internal_nodes, addr)
        choices.internal_nodes[addr]
    elseif haskey(choices.leaf_nodes, addr)
        throw(KeyError(addr))
    else
        EmptyChoiceMap()
    end
end

has_value(choices::DynamicChoiceMap, addr) = haskey(choices.leaf_nodes, addr)

get_value(choices::DynamicChoiceMap, addr) = choices.leaf_nodes[addr]

function Base.isempty(choices::DynamicChoiceMap)
    isempty(choices.leaf_nodes) && isempty(choices.internal_nodes)
end

# mutation (not part of the assignment interface)

"""
    set_value!(choices::DynamicChoiceMap, addr, value)

Set the given value for the given address.

Will cause any previous value or sub-assignment at this address to be deleted.
It is an error if there is already a value present at some prefix of the given address.

The following syntactic sugar is provided:

    choices[addr] = value
"""
function set_value!(choices::DynamicChoiceMap, addr, value)
    delete!(choices.internal_nodes, addr)
    choices.leaf_nodes[addr] = value
end

function set_value!(choices::DynamicChoiceMap, addr::Pair, value)
    (first, rest) = addr
    if haskey(choices.leaf_nodes, first)
        # we are not writing to the address directly, so we error instead of
        # delete the existing node.
        error("Tried to create assignment at $first but there was already a value there.")
    end
    if haskey(choices.internal_nodes, first)
        node = choices.internal_nodes[first]
    else
        node = DynamicChoiceMap()
        choices.internal_nodes[first] = node
    end
    node = choices.internal_nodes[first]
    set_value!(node, rest, value)
end

"""
    set_submap!(choices::DynamicChoiceMap, addr, submap::ChoiceMap)

Replace the sub-assignment rooted at the given address with the given sub-assignment.
Set the given value for the given address.

Will cause any previous value or sub-assignment at the given address to be deleted.
It is an error if there is already a value present at some prefix of address.
"""
function set_submap!(choices::DynamicChoiceMap, addr, new_node)
    delete!(choices.leaf_nodes, addr)
    delete!(choices.internal_nodes, addr)
    if !isempty(new_node)
        choices.internal_nodes[addr] = new_node
    end
end

function set_submap!(choices::DynamicChoiceMap, addr::Pair, new_node)
    (first, rest) = addr
    if haskey(choices.leaf_nodes, first)
        # we are not writing to the address directly, so we error instead of
        # delete the existing node.
        error("Tried to create assignment at $first but there was already a value there.")
    end
    if haskey(choices.internal_nodes, first)
        node = choices.internal_nodes[first]
    else
        node = DynamicChoiceMap()
        choices.internal_nodes[first] = node
    end
    set_submap!(node, rest, new_node)
end

Base.setindex!(choices::DynamicChoiceMap, value, addr) = set_value!(choices, addr, value)

function _fill_array!(choices::DynamicChoiceMap, arr::Vector{T}, start_idx::Int) where {T}
    if length(arr) < start_idx + length(choices.leaf_nodes)
        resize!(arr, 2 * length(arr))
    end
    leaf_keys_sorted = sort(collect(keys(choices.leaf_nodes)))
    internal_node_keys_sorted = sort(collect(keys(choices.internal_nodes)))
    for (i, key) in enumerate(leaf_keys_sorted)
        arr[start_idx + i] = choices.leaf_nodes[key]
    end
    idx = start_idx + length(choices.leaf_nodes)
    for key in internal_node_keys_sorted
        n_written = _fill_array!(get_submap(choices, key), arr, idx)
        idx += n_written
    end
    idx - start_idx
end

function _from_array(proto_choices::DynamicChoiceMap, arr::Vector{T}, start_idx::Int) where {T}
    @assert length(arr) >= start_idx
    choices = DynamicChoiceMap()
    leaf_keys_sorted = sort(collect(keys(proto_choices.leaf_nodes)))
    internal_node_keys_sorted = sort(collect(keys(proto_choices.internal_nodes)))
    for (i, key) in enumerate(leaf_keys_sorted)
        choices.leaf_nodes[key] = arr[start_idx + i]
    end
    idx = start_idx + length(choices.leaf_nodes)
    for key in internal_node_keys_sorted
        (n_read, node) = _from_array(get_submap(proto_choices, key), arr, idx)
        idx += n_read
        choices.internal_nodes[key] = node
    end
    (idx - start_idx, choices)
end

export DynamicChoiceMap
export choicemap
export set_value!
export set_submap!

#######################################
## vector combinator for assignments #
#######################################

# TODO implement LeafVectorChoiceMap, which stores a vector of leaf nodes

struct InternalVectorChoiceMap{T} <: ChoiceMap
    internal_nodes::Vector{T}
    is_empty::Bool
end

function vectorize_internal(nodes::Vector{T}) where {T}
    is_empty = all(map(isempty, nodes))
    InternalVectorChoiceMap(nodes, is_empty)
end

# note some internal nodes may be empty

get_address_schema(::Type{InternalVectorChoiceMap}) = VectorAddressSchema()

Base.isempty(choices::InternalVectorChoiceMap) = choices.is_empty
has_value(choices::InternalVectorChoiceMap, addr::Pair) = _has_value(choices, addr)
get_value(choices::InternalVectorChoiceMap, addr::Pair) = _get_value(choices, addr)
get_submap(choices::InternalVectorChoiceMap, addr::Pair) = _get_submap(choices, addr)

function get_submap(choices::InternalVectorChoiceMap, addr::Int)
    if addr > 0 && addr <= length(choices.internal_nodes)
        choices.internal_nodes[addr]
    else
        EmptyChoiceMap()
    end
end

function get_submaps_shallow(choices::InternalVectorChoiceMap)
    ((i, choices.internal_nodes[i])
     for i=1:length(choices.internal_nodes)
     if !isempty(choices.internal_nodes[i]))
end

get_values_shallow(::InternalVectorChoiceMap) = ()

function _fill_array!(choices::InternalVectorChoiceMap, arr::Vector{T}, start_idx::Int) where {T}
    idx = start_idx
    for key=1:length(choices.internal_nodes)
        n = _fill_array!(choices.internal_nodes[key], arr, idx)
        idx += n
    end
    idx - start_idx
end

function _from_array(proto_choices::InternalVectorChoiceMap{U}, arr::Vector{T}, start_idx::Int) where {T,U}
    @assert length(arr) >= start_idx
    nodes = Vector{U}(undef, length(proto_choices.internal_nodes))
    idx = start_idx
    for key=1:length(proto_choices.internal_nodes)
        (n_read, nodes[key]) = _from_array(proto_choices.internal_nodes[key], arr, idx)
        idx += n_read
    end
    choices = InternalVectorChoiceMap(nodes, proto_choices.is_empty)
    (idx - start_idx, choices)
end

export InternalVectorChoiceMap
export vectorize_internal

####################
# empty assignment #
####################

struct EmptyChoiceMap <: ChoiceMap end

Base.isempty(::EmptyChoiceMap) = true
get_address_schema(::Type{EmptyChoiceMap}) = EmptyAddressSchema()
get_submaps_shallow(::EmptyChoiceMap) = ()
get_values_shallow(::EmptyChoiceMap) = ()

_fill_array!(::EmptyChoiceMap, arr::Vector, start_idx::Int) = 0
_from_array(::EmptyChoiceMap, arr::Vector, start_idx::Int) = (0, EmptyChoiceMap())

export EmptyChoiceMap

max_stars_repo_name: dendisuhubdy/Gen

max_issues_repo_name: kumsh/Gen

max_forks_repo_name: kumsh/Gen


text:
import Base.Docs: Binding, MultiDoc, DocStr

get_multidoc(__module__::Module, b::Binding) = getindex(Docs.meta(__module__), b)

function get_sigs(__module__::Module, b::Binding)
    m = get_multidoc(__module__, b)
    m.order
end

function submodules(__module__::Module)
    itr = Iterators.filter(x->(Base.moduleroot(x) == __module__ && x != __module__), Docs.modules)
    collect(itr)
end

import Base: replace!

function gotkey(d::IdDict, k)
    for each in keys(d)
        k == each && return each
    end
    missing
end

"""
    replace!(__module__, binding, str, sig)
This is a muted version of [`Docs.doc!`](@ref), it replace the original docstring with
a new docstring `str` to the docsystem of `__module__` for `binding` and signature `sig`.
"""
function replace!(__module__::Module, b::Binding, str::DocStr, @nospecialize sig = Union{})
    Docs.initmeta(__module__)
    m = get!(Docs.meta(__module__), b, MultiDoc())
    sigm=gotkey(m.docs, sig)
    if ismissing(sigm)
        @error "Cannot find original docs for `$b :: $sig` in module `$(__module__)`"
    end

    old = m.docs[sigm]
    m.docs[sigm] = str

    # inherit the rest of data
    for (k, v) in old.data
        str.data[k] = v
    end

    str.data[:binding] = b
    str.data[:typesig] = sigm
    return b
end