text:
#########################################################################
#
#   rdiff differentiation function
#
#########################################################################

######### expression version   ################
# TODO : break this huge function in smaller blocks

"""
Generates an expression calculating the derivative(s) for a given expression

    rdiff( ex::Expr; kwargs... )

Arguments:

- ex: is a Julia Expression containing the code to derive
- outsym: (default = nothing) is the symbol of the variable within ``ex`` containing the expression output (the result whose derivatives are needed). This variable must evaluate to a ``Real``. If not specified, ``outsym`` defaults to ``nothing`` which signals to ``rdiff`` that the last statement is the result of interest for derivation.
- order: (default = 1) is an integer indicating the derivation order (1 for 1st order, etc.). Order 0 is allowed and will produce an expression that is a processed version of ``ex`` with some variables names rewritten and possibly some optimizations.
- init: (multiple keyword arguments) is one or several symbol / DataType pairs used to indicate for which variable a derivative is needed and how they should be interpreted. By default the generated expression will yield the derivative for each variable given unless the variable is listed in the ``ignore`` argument.
- evalmod: (default=Main) module where the expression is meant to be evaluated. External variables and functions should be evaluable in this module.
- debug: (default=false) indicates if ``rdiff`` should dump the graph of the generating expression, instead of returning the expression itself.
- allorders: (default=true) indicates whether to generate the code for all orders up to ``order`` (true) or only the last order.
- ignore: (default=[]) do not differentiate against the listed variables, useful if you are not interested in having the derivative of one of several variables in ``init``.

Usage:
```julia
julia> rdiff( :(x^3) , x=Float64)  # first order
:(begin
    (x^3,3 * x^2.0)
    end)

julia> rdiff( :(x^3) , order = 3, x=Float64)  # orders up to 3
:(begin
        (x^3,3 * x^2.0,2.0 * (x * 3),6.0)
    end)

julia> rdiff( :(sin(x)) , order=10, x=Float64)  # derivatives up to order 10
:(begin
        _tmp1 = sin(x)
        _tmp2 = cos(x)
        _tmp3 = -_tmp1
        _tmp4 = -_tmp2
        _tmp5 = -_tmp3
        (_tmp1,_tmp2,_tmp3,_tmp4,_tmp5,_tmp2,_tmp3,_tmp4,_tmp5,_tmp2,_tmp3)
    end)

julia> ex = :(p^3+y)
julia> rdiff( ex , p=Float64, y=Float64, ignore=:y) # derive for p, not for y
```
"""
function rdiff(ex;
               outsym    = nothing,
               order::Int= 1,
               evalmod   = Main,
               debug     = false,
               allorders = true,
               ignore    = Symbol[],
               params...)

    # format parameters in `params`
    paramsym  = Symbol[ e[1] for e in params]
    paramval  = [ e[2] for e in params]
    pardict   = Dict(zip(paramsym, paramval))

    # list of variables to differentiate against
    paramdiff = if isa(ignore, Symbol)
                    setdiff(paramsym, [ignore;])
                elseif isa(ignore, Union{Vector,Tuple}) && eltype(ignore)==Symbol
                    setdiff(paramsym, ignore)
                else
                    error("'ignore' arg should be a Symbol or Vector/Tuple of Symbol")
                end

    # controls
    length(paramsym) >= 1 ||
        error("There should be at least one variable specified, none found")

    length(paramdiff) >= 1 ||
        error("There should be at least one variable for differentiation, none found")

    order <= 1 ||
        length(paramdiff) == 1 ||
        error("Only one differentiation variable allowed for order >= 2")

    order <= 1 ||
        (paramval[1] <: Vector && (eltype(paramval[1]) <: Real)) ||
        (paramval[1] <: Real) ||
        error("Param should be a real or vector for order >= 2")

    g = tograph(ex, evalmod)

    hassym(g.seti, outsym) ||
        error("can't find output var $( outsym==nothing ? "" : outsym)")

    # reduce to variable of interest
    g.seti = NSMap([getnode(g.seti, outsym)], [ outsym ])

    g |> splitnary! |> prune! |> simplify!
    calc!(g, params=pardict, emod=evalmod)

    ov = getnode(g.seti, outsym).val
    ov <: Real || error("output var should be a Real, $ov found")

    voi = Any[ outsym ]

    if order == 0
        # do nothing

    elseif order == 1
        dg = reversegraph(g, getnode(g.seti, outsym), paramdiff)
        append!(g.nodes, dg.nodes)

        for p in paramdiff
            nn = getnode(dg.seti, dprefix(p))  # find the exit node of deriv of p
            ns = newvar("_dv")
            g.seti[nn] = ns
            push!(voi, ns)
        end

        g |> splitnary! |> prune! |> simplify!

    elseif order > 1 && pardict[paramdiff[1]] <: Real
        for i in 1:order
            dg = reversegraph(g, getnode(g.seti, voi[i]), paramdiff)
            append!(g.nodes, dg.nodes)
            nn = collect(nodes(dg.seti))[1]  # only a single node produced
            ns = newvar("_dv")
            g.seti[nn] = ns
            push!(voi, ns)

            g |> splitnary! |> prune! |> simplify!

            calc!(g, params=pardict, emod=evalmod)
        end

    elseif order > 1 && pardict[paramdiff[1]] <: Vector &&
                        eltype(pardict[paramdiff[1]]) <: Real
        # do first order as usual
        dg = reversegraph(g, getnode(g.seti, outsym), paramdiff)
        append!(g.nodes, dg.nodes)
        ns = newvar(:_dv)
        g.seti[ collect(nodes(dg.seti))[1] ] = ns
        push!(voi, ns)

        g |> splitnary! |> prune! |> simplify!

        # now order 2 to n
        for i in 2:order
            # launch derivation on a single value of the preceding
            #   derivation vector
            no = getnode(g.seti, voi[i])
            si = newvar(:_idx)
            ni = addnode!(g, NExt(si))
            ns = addnode!(g, NRef(:getidx, [ no, ni ]))

            calc!(g, params=Dict(zip([paramsym; si], [paramval; Int64])), emod=evalmod)
            dg = reversegraph(g, ns, paramdiff)

            #### We will now wrap dg in a loop scanning all the elements of 'no'
            # first create ext nodes to make dg a complete subgraph
            dg2 = ExNode[]
            nmap = Dict()
            for n in dg.nodes  # n = dg.nodes[2]
                for (j, np) in enumerate(n.parents)  # j,np = 1, n.parents[1]
                    if haskey(nmap, np) # already remapped
                        n.parents[j] = nmap[np]

                    elseif np == ni # it's the loop index
                        nn = NExt(si)
                        push!(dg2, nn)
                        dg.exti[nn] = si
                        n.parents[j] = nn
                        nmap[np] = nn

                    elseif np == ns # it's the selected element of the deriv vector
                        # create 'no' ref if needed
                        if !haskey(nmap, no)
                            sn = newvar()
                            nn = NExt(sn)
                            push!(dg2, nn)
                            dg.exti[nn] = sn
                            dg.exto[no] = sn
                            nmap[no] = nn
                        end

                        nn = NRef(:getidx, [ nmap[no], nmap[ni] ])
                        push!(dg2, nn)
                        nmap[ns] = nn

                    elseif !(np in dg.nodes) # it's not in dg (but in g)
                        sn = newvar()
                        nn = NExt(sn)
                        push!(dg2, nn)
                        dg.exti[nn] = sn
                        dg.exto[np] = sn
                        n.parents[j] = nn
                        nmap[np] = nn

                    end
                end

                # update onodes in for loops
                if isa(n, NFor)
                    g2 = n.main[2]
                    for (o,s) in g2.exto
                        if haskey(nmap, o)
                            g2.exto[ nmap[o] ] = s  # replace
                        end
                    end
                end
            end
            append!(dg.nodes, dg2)
            # dg |> prune! |> simplify!

            # create for loop node
            nf = addnode!(g, NFor(Any[ si, dg ] ) )

            # create param size node
            nsz = addgraph!( :( length( x ) ), g, Dict( :x => getnode(g.exti, paramdiff[1]) ) )

            # create (n-1)th derivative size node
            ndsz = addgraph!( :( sz ^ $(i-1) ), g, Dict( :sz => nsz ) )

            # create index range node
            nid = addgraph!( :( 1:dsz ),  g, Dict( :dsz => ndsz ) )
            push!(nf.parents, nid)

            # pass size node inside subgraph
            sst = newvar()
            inst = addnode!(dg, NExt(sst))
            dg.exti[inst] = sst
            dg.exto[nsz]  = sst
            push!(nf.parents, nsz)

            # create result node (alloc in parent graph)
            nsa = addgraph!( :( zeros( $( Expr(:tuple, [:sz for j in 1:i]...) ) ) ),
                            g, Dict( :sz => nsz ) )
            ssa = newvar()
            insa = addnode!(dg, NExt(ssa))
            dg.exti[insa] = ssa
            dg.exto[nsa]  = ssa
            push!(nf.parents, nsa)

            # create result node update (in subgraph)
            nres = addgraph!( :( res[ ((sidx-1)*st+1):(sidx*st) ] = dx ; res ), dg,
                              Dict(:res  => insa,
                                   :sidx => nmap[ni],
                                   :st   => inst,
                                   :dx   => collect(dg.seti)[1][1] ) )
            dg.seti = NSMap([nres], [ssa])

            # create exit node for result
            nex = addnode!(g, NIn(ssa, [nf]))
            dg.seto = NSMap([nex], [ssa])

            # update parents of for loop
            append!( nf.parents, setdiff(collect( nodes(dg.exto)), nf.parents[2:end]) )

            ns = newvar(:_dv)
            g.seti[nex] = ns
            push!(voi, ns)

            g |> splitnary! |> prune! |> simplify!
        end

    else
        error("[rdiff] order / variable types inconsistent, check documentation")
    end

    if !allorders  # only keep the last derivative
        if order == 1 # potentially multiple diffs, issue #32
            voi = voi[2:end]
        else
            voi = [voi[end]]
        end
    end

    if length(voi) > 1  # create tuple if multiple variables
        voin = map( s -> getnode(g.seti, s), voi )
        nf = addnode!(g, NConst(tuple))
        exitnode = addnode!(g, NCall(:call, [nf, voin...]))
    else
        exitnode = getnode(g.seti, voi[1])
    end
    g.seti = NSMap( [exitnode], [nothing])  # make this the only exitnode of interest

    g |> splitnary! |> prune! |> simplify!

    resetvar()
    debug ? g : tocode(g, evalmod)
end


text:
using Distributed
using Test
addprocs(2)
@everywhere using Random
@everywhere using NeuralQuantumState

const energy_AFH_6_spins = -11.211102550927972

@testset "NeuralQuantumState.jl" begin
    @everywhere Random.seed!(1234321)
    NetSettings = NETSETTINGS(modelname = "U_afh",
        repetitions = 600,
        n = 6,
        α = 3,
        mag0 = true,
        γ_decay = 0.997,
        mc_trials = 200,
        writetofile = true,
        init_therm_steps = 100,
        therm_steps = 50,
        stat_samples = 2000,
        use_meter=false)

    energy = run_NQS(NetSettings)
    @test isapprox(energy[end], energy_AFH_6_spins, rtol=1e-3)

    NetSettings = NETSETTINGS(NetSettings, mag0=false) # Now without mag0
    energy = run_NQS(NetSettings)
    @test isapprox(energy[end], energy_AFH_6_spins, rtol=1e-2)
end

ext: jl


text:
using Documenter, MASW

makedocs(;
    modules=[MASW],
    format=Documenter.HTML(),
    pages=[
        "Home" => "index.md",
    ],
    repo="https://github.com/kura-okubo/MASW.jl/blob/{commit}{path}#L{line}",
    sitename="MASW.jl",
    authors="kurama",
    assets=String[],
)

deploydocs(;
    repo="github.com/kura-okubo/MASW.jl",
)


text:
using Random
Random.seed!(101)

@testset "Values tests" begin
  @testset "String values" begin
    @test rand(String) == "Tw2GszEb"
    @test rand(String) == "3U3vvin8"
    @test rand(String, 20) == "fhJmWB8TAixQ057e4eSF"
  end

  @testset "Pair values" begin
    @test rand(Pair{String,String}) == ("Tw2GszEb" => "3U3vvin8")
  end
end


text:
"""
Provides functions to interact with Cartesian ASCII ephemeris data.

This is BYOE (bring your own ephemeris) module! Check out
JPL Horizons and
[wrapper scripts](https://github.com/cadojo/Halo-Explorations/blob/main/scripts/ephemeris/fetch_ephemeris.sh)
to learn how to fetch ephemeris files.
"""
module Ephemeris

# Module Exports
export loadascii, Interpolator, interpolator

# Module Dependencies
using StaticArrays
using Interpolations
using CSV, DataFrames
using Unitful, UnitfulAstro, UnitfulAngles

using ..AstrodynamicsCore

# Module source code
include("LoadASCII.jl")

end


text:
function near02()
#
#  test polynomial suggested by Z. Zeng
#
    e = 0.01;
    p = poly([(1-e)*ones(20);ones(20);-0.5*[1;1;1;1;1]]);
    z = [1-e 20; 1 20; -0.5 5];
    p, PolyZeros(z)
end


text:
## Run test with the Julia module gaiaClustering
## @02.09.2018

## ... directories
rootdir = ENV["GAIA_ROOT"]
wdir    = "$rootdir/products"

push!(LOAD_PATH,"$rootdir/run/src")
using gaiaClustering

cd(wdir)

voname = "NGC 2516-1.0deg.vot"

data       = read_votable(voname)
df         = filter_data(data)
dfcart     = add_cartesian(df)

blck       = [[1,2,3],[4,5], [6,7,8]]
wghtblck   = [3.,2.,1.5]
norm = "normal"
dfcartnorm = normalization_PerBlock(dfcart, blck, wghtblck , norm)

println("## Test done...")

max_stars_repo_name: bosscha/gaia-shock

max_issues_repo_name: bosscha/gaia-shock

max_forks_repo_name: bosscha/gaia-shock


text:
"""
@addfunction f

Using `@addfunction` will write in a new function into the
current namespace at read-time.

# Examples
```julia-repl
addfunction lognormalconstraint(x, μ, σ) = (x-μ)^2/2/σ
```
"""
macro addfunction(f)
  :($f)
end

## Is this obsolete?
macro dataset( name, df )
  :( $(esc(name)) = $df )
end

"""
@add_dataset(name, df::DataFrame)

Introduce a new dataset into the namespace at read-time.
"""
function add_dataset(name, df::DataFrame)
  eval(:( $name = $df ) )
end

function add_dataset(df::DataFrame)
  eval(:( $(Symbol(df))=$df ))
end

function add_dataset(name::String, df::DataFrame)
  eval(:( $(Symbol(name)) = $df ))
end

function rename_dataset(name::Symbol, other::Symbol)
  eval(:( $(name) = $(other) ))
end

function add_array(name::Symbol, arr::Array{Float64})
  eval(:( $(name) = $arr ))
end

function add_array(name::String, arr::Array{Float64})
  eval(:( $(Symbol(name)) = $arr ))
end

"""
@add_function(name, function)

Take a user function and pass ownership to a symbol internal
to the Batman module.
"""
function add_function(name, func)
  eval(:( $name = $func ))
end

function add_function(fun)
  eval(:( $(Symbol(fun))=$fun ))
end

function add_function(name::String, func)
  eval(:( $(Symbol(name))=$func ))
end


text:
# Victoria Park

using Caesar, RoME
using RoMEPlotting
using Distributions

# load all the model data
# d = odometry information
# f = laser scanner detections
# MM = multi-modal individual id references
# MMr = reworked to map to only one previous feature
# examplefolder, datafolder
include(joinpath(dirname(@__FILE__),"loadVicPrkData.jl"))

# p = drawFeatTrackers( d[1], f[1] );

T=30 # 1400
fg = RoME.emptyFactorGraph();
idx = appendFactorGraph!(fg, d, f, toT=T, lcmode=:mmodal, MM=MMr);

drawPosesLandms(fg)

Graphs.plot(fg.g)

vc = startdefaultvisualization()

visualizeallposes!(vc, fg, drawlandms=false)

# Evaluate the likelihood of a point on the marginal belief of some variable
# note the dimensions must match
function evalLikelihood(fg::G, sym::Symbol, point::Vector{Float64}) where G <: AbstractDFG
  p = getVertKDE(fg, sym)
  Ndim(p) == length(point) ? nothing : error("point (dim=$(length(point))) must have same dimension as belief (dim=$(Ndim(p)))")
  evaluateDualTree(p, (point')')[1]
end

# Evaluate the likelihood of an Array{2} of points on the marginal belief of some variable
# note the dimensions must match
function evalLikelihood(fg::G, sym::Symbol, points::Array{Float64,2}) where G <: AbstractDFG
  p = getVertKDE(fg, sym)
  Ndim(p) == size(points,1) ? nothing : error("points (dim=$(size(points,1))) must have same dimension as belief (dim=$(Ndim(p)))")
  evaluateDualTree(p, (points))
end

evalLikelihood(fg, :l1, [1.0;0.2])

evalLikelihood(fg, :l1, zeros(2,3))

tree = prepBatchTree!(fg, drawpdf=true);
@time inferOverTree!(fg,tree, N=100);
# @save "fgT1400v09err.jld" fg

# include("examples/dev/ISAMRemoteSolve.jl")
# T=1400
# gtvals = doISAMSolve(d,f,toT=T, savejld=true, MM=MMr)

# @load "test/data/isamdict.jld"
# @load "test/data/fgu.jld"

# include("examples/dev/ISAMRemoteSolve.jl")

T=1400
fg = emptyFactorGraph();
idx = appendFactorGraph!(fg, d, f, toT=T, lcmode=:mmodal, MM=MMr);
tree = prepBatchTree!(fg, ordering=:qr,drawpdf=true);
@time inferOverTree!(fg,tree, N=100);
@save "fgT1400v06.jld" fg
FG = FactorGraph[]
TREE= BayesTree[]
push!(FG,deepcopy(fg));
push!(TREE,deepcopy(tree));
for i in 1:2
   @time inferOverTreeR!(fg,tree, N=100);
   push!(FG,deepcopy(fg));
   push!(TREE,deepcopy(tree));
end
for i in 1:2
   @time inferOverTree!(fg,tree, N=100);
   push!(FG,deepcopy(fg));
   push!(TREE,deepcopy(tree));
end

  # drawCompPosesLandm(fg,isamdict, fgu, lbls=false,drawunilm=false)

@load "results/fgT1400_Nightly7_28_perErr0.0.jld"
MMMAP0 = deepcopy(fg2)
UMMLE0d = deepcopy(isamdict)
UMMLE0fg = deepcopy(fgu)
rangesMM0_0 = rangeCompAllPoses(UMMLE0d, UMMLE0fg, MMMAP0)

@load "results/fgT1400_Friday7_29_perErr0.01.jld"
MMMAP1 = deepcopy(fg2)
UMMLE1d = deepcopy(isamdict)
UMMLE1fg = deepcopy(fgu)
rangesMM0_1 = rangeCompAllPoses(UMMLE1d, UMMLE1fg, MMMAP0)
rangesMM0_MM1 = rangeCompAllPoses(MMMAP0, MMMAP1)

@load "results/fgT1400_Friday7_29_perErr0.05.jld"
MMMAP5 = deepcopy(fg2)
UMMLE5d = deepcopy(isamdict)
UMMLE5fg = deepcopy(fgu)
rangesMM0_5 = rangeCompAllPoses(UMMLE5d, UMMLE5fg, MMMAP0)
rangesMM0_MM5 = rangeCompAllPoses(MMMAP0, MMMAP5)

@load "results/fgT1400_Friday7_29_perErr0.1.jld"
MMMAP10 = deepcopy(fg2)
UMMLE10d = deepcopy(isamdict)
UMMLE10fg = deepcopy(fgu)
rangesMM0_10 = rangeCompAllPoses(UMMLE10d, UMMLE10fg, MMMAP0)
rangesMM0_MM10 = rangeCompAllPoses(MMMAP0, MMMAP10)

@load "results/fgT1400_Friday7_29_perErr0.15.jld"
MMMAP15 = deepcopy(fg2)
UMMLE15d = deepcopy(isamdict)
UMMLE15fg = deepcopy(fgu)
rangesMM0_15 = rangeCompAllPoses(UMMLE15d, UMMLE15fg, MMMAP0)
rangesMM0_MM15 = rangeCompAllPoses(MMMAP0, MMMAP15)

@load "results/fgT1400_Friday7_29_perErr0.2.jld"
MMMAP20 = deepcopy(fg2)
UMMLE20d = deepcopy(isamdict)
UMMLE20fg = deepcopy(fgu)
rangesMM0_20 = rangeCompAllPoses(UMMLE20d, UMMLE20fg, MMMAP0)
rangesMM0_MM20 = rangeCompAllPoses(MMMAP0, MMMAP20)

ptsize = 1.8pt
# draw all
lyrs = []
push!(lyrs, Gadfly.layer(y=rangesMM0_0, Geom.line))

push!(lyrs, Gadfly.layer(y=rangesMM0_MM1, Geom.line, Theme(default_color=colorant"black")))
push!(lyrs, Gadfly.layer(y=rangesMM0_1, Geom.line, Theme(default_color=colorant"green")))

# push!(lyrs, Gadfly.layer(y=rangesMM0_5, Geom.line, Theme(default_color=colorant"yellow")))
# push!(lyrs, Gadfly.layer(y=rangesMM0_MM5, Geom.line, Theme(default_color=colorant"red")))

push!(lyrs, Gadfly.layer(y=rangesMM0_10, Geom.point, Theme(default_color=colorant"magenta", default_point_size=ptsize)))
push!(lyrs, Gadfly.layer(y=rangesMM0_MM10, Geom.point, Theme(default_color=colorant"black", default_point_size=ptsize)))

# push!(lyrs, Gadfly.layer(y=rangesMM0_15, Geom.line, Theme(default_color=colorant"cyan")))
# push!(lyrs, Gadfly.layer(y=rangesMM0_MM15, Geom.line, Theme(default_color=colorant"red")))

push!(lyrs, Gadfly.layer(y=rangesMM0_20, Geom.point, Theme(default_color=colorant"blue", default_point_size=ptsize)))
push!(lyrs, Gadfly.layer(y=rangesMM0_MM20, Geom.point, Theme(default_color=colorant"black", default_point_size=ptsize)))

push!(lyrs, Coord.Cartesian(xmin=0,xmax=length(rangesMM0_0),
                            ymin=0,ymax=350))

push!(lyrs,
  Guide.xlabel("Individual poses")
)
push!(lyrs,
  Guide.ylabel("| X - X' |,  [meters]")
)
push!(lyrs,
Guide.manual_color_key("", ["MM-MAP", "MLE ~0%", "MLE 1%", "MLE 10%", "MLE 20%"], ["black", "deepskyblue", "green", "magenta", "blue"])
)
push!(lyrs,
Theme(key_position=:top)
)

lpl = Gadfly.plot(lyrs...)
draw(PDF("poseErrs.pdf",15cm,8cm),lpl)
lpl

pl = drawCompPosesLandm(MMMAP10,UMMLE10d, UMMLE10fg, lbls=false,drawunilm=false)
pl.coord = Coord.Cartesian(xmin=-120,xmax=250,ymin=-100,ymax=250)
draw(PNG("VicPrk10.png",15cm,12cm),pl)

vpl = vstack(pl,lpl);
draw(PDF("test.pdf",15cm,24cm),vpl)


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<:Color to work on
#    scalars x::Fractional

# Dispatch-based elementwise manipulations
"""
`MapInfo{T}` is an abstract type that encompasses objects designed to
perform intensity or color transformations on pixels.  For example,
before displaying an image in a window, you might need to adjust the
contrast settings; `MapInfo` objects provide a means to describe these
transformations without calculating them immediately.  This delayed
execution can be useful in many contexts.  For example, if you want to
display a movie, it would be quite wasteful to have to first transform
the entire movie; instead, `MapInfo` objects allow one to specify a
transformation to be performed on-the-fly as particular frames are
displayed.

You can create your own custom `MapInfo` objects. For example, given a
grayscale image, you could color "saturated" pixels red using

```jl
immutable ColorSaturated{C<:AbstractRGB} <: MapInfo{C}
end

Base.map{C}(::ColorSaturated{C}, val::Union{Number,Gray}) = ifelse(val == 1, C(1,0,0), C(val,val,val))

imgc = map(ColorSaturated{RGB{U8}}(), img)
```

For pre-defined types see `MapNone`, `BitShift`, `ClampMinMax`, `ScaleMinMax`,
`ScaleAutoMinMax`, and `ScaleSigned`.
"""
abstract MapInfo{T}
eltype{T}(mapi::MapInfo{T}) = T

## MapNone
"`MapNone(T)` is a `MapInfo` object that converts `x` to have type `T`."
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}(mapi::MapNone{T}, val::Union{Number,Colorant}) = convert(T, val)
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<:Colorant}(mapi::MapNone{CT}, val::Fractional) = convert(eltype(CT), val)

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

## BitShift
"""
`BitShift{T,N}` performs a "saturating rightward bit-shift" operation.
It is particularly useful in converting high bit-depth images to 8-bit
images for the purpose of display.  For example,

```
map(BitShift(UFixed8, 8), 0xa2d5uf16) === 0xa2uf8
```

converts a `UFixed16` to the corresponding `UFixed8` by discarding the
least significant byte.  However,

```
map(BitShift(UFixed8, 7), 0xa2d5uf16) == 0xffuf8
```

because `0xa2d5>>7 == 0x0145 > typemax(UInt8)`.

When applicable, the main advantage of using `BitShift` rather than
`MapNone` or `ScaleMinMax` is speed.
"""
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<:Colorant,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}
"""
`ClampMin(T, minvalue)` is a `MapInfo` object that clamps pixel values
to be greater than or equal to `minvalue` before converting to type `T`.

See also: `ClampMax`, `ClampMinMax`.
"""
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)
"""
`ClampMax(T, maxvalue)` is a `MapInfo` object that clamps pixel values
to be less than or equal to `maxvalue` before converting to type `T`.

See also: `ClampMin`, `ClampMinMax`.
"""
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, minvalue, maxvalue)` is a `MapInfo` object that clamps
pixel values to be between `minvalue` and `maxvalue` before converting
to type `T`.

See also: `ClampMin`, `ClampMax`, and `Clamp`.
"""
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)
"""
`Clamp(C)` is a `MapInfo` object that clamps color values to be within
gamut.  For example,

```
map(Clamp(RGB{U8}), RGB(1.2, -0.4, 0.6)) === RGB{U8}(1, 0, 0.6)
```
"""
immutable Clamp{T} <: AbstractClamp{T} end
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<:Colorant,F<:Fractional}(mapi::ClampMin{CT,F}, val::F) = convert(eltype(CT), max(val, mapi.min))
map1{CT<:Colorant,F<:Fractional}(mapi::ClampMax{CT,F}, val::F) = convert(eltype(CT), min(val, mapi.max))
map1{CT<:Colorant,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{P<:TransparentRGB}(::Clamp{P}, val::Real) = clamp01(eltype(P), 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)
clamp01(x::Colorant) = clamp01(typeof(x), x)
clamp01{Cdest<:AbstractRGB   }(::Type{Cdest}, x::AbstractRGB)    = (To = eltype(Cdest);
    Cdest(clamp01(To, red(x)), clamp01(To, green(x)), clamp01(To, blue(x))))
clamp01{Pdest<:TransparentRGB}(::Type{Pdest}, x::TransparentRGB) = (To = eltype(Pdest);
    Pdest(clamp01(To, red(x)), clamp01(To, green(x)), clamp01(To, blue(x)), clamp01(To, alpha(x))))

# clamp is generic for any colorspace; this version does the right thing for any RGB type
clamp(x::Union{AbstractRGB, TransparentRGB}) = clamp01(x)

## ScaleMinMax
"""
`ScaleMinMax(T, min, max, [scalefactor])` is a `MapInfo` object that
clamps the image at the specified `min`/`max` values, subtracts the
`min` value, scales the result by multiplying by `scalefactor`, and
finally converts to type `T`.  If `scalefactor` is not specified, it
defaults to scaling the interval `[min,max]` to `[0,1]`.

Alternative constructors include `ScaleMinMax(T, img)` for which
`min`, `max`, and `scalefactor` are computed from the minimum and
maximum values found in `img`.

See also: `ScaleMinMaxNaN`, `ScaleAutoMinMax`, `MapNone`, `BitShift`.
"""
immutable ScaleMinMax{To,From,S<:AbstractFloat} <: 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::AbstractFloat) = ScaleMinMax{To,From,typeof(s)}(min, max, s)
ScaleMinMax{To<:Union{Fractional,Colorant},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<:Union{Real,AbstractGray},From<:Union{Real,AbstractGray}}(mapi::ScaleMinMax{To,From}, val::From)
    g = gray(val)
    t = ifelse(g < mapi.min, zero(From), ifelse(g > mapi.max, mapi.max-mapi.min, g-mapi.min))
    convert(To, mapi.s*t)
end
function map{To<:Union{Real,AbstractGray},From<:Union{Real,AbstractGray}}(mapi::ScaleMinMax{To,From}, val::Union{Real,Colorant})
    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<:Colorant,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,Colorant})
    map1(mapi, convert(From, val))
end
function map1{To<:Colorant,From<:Real}(mapi::ScaleMinMax{To,From}, val::Union{Real,Colorant})
    map1(mapi, convert(From, val))
end

## ScaleSigned
"""
`ScaleSigned(T, scalefactor)` is a `MapInfo` object designed for
visualization of images where the pixel's sign has special meaning.
It multiplies the pixel value by `scalefactor`, then clamps to the
interval `[-1,1]`. If `T` is a floating-point type, it stays in this
representation.  If `T` is an `AbstractRGB`, then it is encoded as a
magenta (positive)/green (negative) image, with the intensity of the
color proportional to the clamped absolute value.
"""
immutable ScaleSigned{T, S<:AbstractFloat} <: MapInfo{T}
    s::S
end
ScaleSigned{T}(::Type{T}, s::AbstractFloat) = 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<:AbstractRGB}(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
"""
`ScaleAutoMinMax(T)` constructs a `MapInfo` object that causes images
to be dynamically scaled to their specific min/max values, using the
same algorithm for `ScaleMinMax`. When displaying a movie, the min/max
will be recalculated for each frame, so this can result in
inconsistent contrast scaling.
"""
immutable ScaleAutoMinMax{T} <: MapInfo{T} end
ScaleAutoMinMax{T}(::Type{T}) = ScaleAutoMinMax{T}()
ScaleAutoMinMax() = ScaleAutoMinMax{UFixed8}()

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

## NaN-nulling mapping
"""
`ScaleMinMaxNaN(smm)` constructs a `MapInfo` object from a
`ScaleMinMax` object `smm`, with the additional property that `NaN`
values map to zero.

See also: `ScaleMinMax`.
"""
immutable ScaleMinMaxNaN{To,From,S} <: MapInfo{To}
    smm::ScaleMinMax{To,From,S}
end

"""
`Clamp01NaN(T)` or `Clamp01NaN(img)` constructs a `MapInfo` object
that clamps grayscale or color pixels to the interval `[0,1]`, sending
`NaN` pixels to zero.
"""
immutable Clamp01NaN{T} <: MapInfo{T} end

Clamp01NaN{T}(A::AbstractArray{T}) = Clamp01NaN{T}()

# Implementation
similar{T,F,To,From,S}(mapi::ScaleMinMaxNaN{To,From,S}, ::Type{T}, ::Type{F}) = ScaleMinMaxNaN{T,F,S}(similar(mapi.smm, T, F))
similar{T}(mapi::Clamp01NaN, ::Type{T}, ::Type) = Clamp01NaN{T}()

Base.map{To}(smmn::ScaleMinMaxNaN{To}, g::Number) = isnan(g) ? zero(To) : map(smmn.smm, g)
Base.map{To}(smmn::ScaleMinMaxNaN{To}, g::Gray) = isnan(g) ? zero(To) : map(smmn.smm, g)

function Base.map{T<:RGB}(::Clamp01NaN{T}, c::AbstractRGB)
    r, g, b = red(c), green(c), blue(c)
    if isnan(r) || isnan(g) || isnan(b)
        return T(0,0,0)
    end
    T(clamp(r, 0, 1), clamp(g, 0, 1), clamp(b, 0, 1))
end
function Base.map{T<:Union{Fractional,Gray}}(::Clamp01NaN{T}, c::Union{Fractional,AbstractGray})
    g = gray(c)
    if isnan(g)
        return T(0)
    end
    T(clamp(g, 0, 1))
end

# 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::AbstractFloat)
                if isfinite(g)
                    x = map1(mapi, g)
                    convert(RGB24, RGB{UFixed8}(x,x,x))
                else
                    RGB24(0)
                end
            end
            map{G<:Gray}(mapi::$ST{RGB24}, g::TransparentColor{G}) = map(mapi, gray(g))
            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{G<:Gray}(mapi::$ST{ARGB32}, g::TransparentColor{G})
                x = map1(mapi, gray(g))
                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,G<:Gray}(mapi::$ST{$OA{T}}, g::TransparentColor{G})
                    x = map1(mapi, gray(g))
                    $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, red(rgb)), map1(mapi, green(rgb)), map1(mapi, blue(rgb))))
            map{C<:AbstractRGB, TC}(mapi::$ST{RGB24}, argb::TransparentColor{C,TC}) =
                convert(RGB24, RGB{UFixed8}(map1(mapi, red(argb)), map1(mapi, green(argb)),
                                            map1(mapi, blue(argb))))
            map{C<:AbstractRGB, TC}(mapi::$ST{ARGB32}, argb::TransparentColor{C,TC}) =
                convert(ARGB32, ARGB{UFixed8}(map1(mapi, red(argb)), map1(mapi, green(argb)),
                                              map1(mapi, blue(argb)), map1(mapi, alpha(argb))))
            map(mapi::$ST{ARGB32}, rgb::AbstractRGB) =
                convert(ARGB32, ARGB{UFixed8}(map1(mapi, red(rgb)), map1(mapi, green(rgb)), map1(mapi, blue(rgb))))
        end
        for O in (:RGB, :BGR)
            @eval begin
                map{T}(mapi::$ST{$O{T}}, rgb::AbstractRGB) =
                    $O{T}(map1(mapi, red(rgb)), map1(mapi, green(rgb)), map1(mapi, blue(rgb)))
                map{T,C<:AbstractRGB, TC}(mapi::$ST{$O{T}}, argb::TransparentColor{C,TC}) =
                    $O{T}(map1(mapi, red(argb)), map1(mapi, green(argb)), map1(mapi, blue(argb)))
            end
        end
        for OA in (:RGBA, :ARGB, :BGRA)
            @eval begin
                map{T, C<:AbstractRGB, TC}(mapi::$ST{$OA{T}}, argb::TransparentColor{C,TC}) =
                    $OA{T}(map1(mapi, red(argb)), map1(mapi, green(argb)),
                            map1(mapi, blue(argb)), map1(mapi, alpha(argb)))
                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, red(rgb)), map1(mapi, green(rgb)), map1(mapi, blue(rgb)))
                map{T}(mapi::$ST{$OA{T}}, rgb::RGB24) = map(mapi, convert(RGB{UFixed8}, argb))
            end
        end
    end
end

# # Apply to any Colorant
# map(f::Callable, x::Color) = f(x)
# map(mapi, x::Color) = map(mapi, convert(RGB, x))
# map{C<:Color, TC}(f::Callable, x::TransparentColor{C, TC}) = f(convert(ARGB, x))
# map{C<:Color, TC}(mapi, x::TransparentColor{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<:Colorant,R<:Real}(mapi::MapNone{C}, img::AbstractImageDirect{R}) = mapcd(mapi, img)  # ambiguity resolution
map{C<:Colorant,R<:Real}(mapi::MapInfo{C}, img::AbstractImageDirect{R}) = mapcd(mapi, img)
function mapcd{C<:Colorant,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<:Colorant}(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)
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())
    for I in eachindex(dout, dimg)
        @inbounds dout[I] = map(mi, 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)
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)
    for I in eachindex(dout, dimg)
        @inbounds dout[I] = cmap[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{AGray,GrayA,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,Colorant}}(::Type{T}, img::AbstractArray{T}) = MapNone(img)
"""
`mapi = mapinf(T, img)` returns a `MapInfo` object that is deemed
appropriate for converting pixels of `img` to be of type `T`. `T` can
either be a specific type (e.g., `RGB24`), or you can specify an
abstract type like `Clamp` and it will return one of the `Clamp`
family of `MapInfo` objects.

You can define your own rules for `mapinfo`.  For example, the
`ImageMagick` package defines methods for how pixels values should be
converted before saving images to disk.
"""
mapinfo{T<:UFixed}(::Type{T}, img::AbstractArray{T}) = MapNone(img)
mapinfo{T<:AbstractFloat}(::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{GrayA{UFixed8}}, img::AbstractArray{GrayA{UFixed8}}) = MapNone{GrayA{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{GrayA{UFixed8}}, img::AbstractArray{GrayA{$T}}) = BitShift{GrayA{UFixed8},$n}()
end
mapinfo{T<:UFixed,F<:AbstractFloat}(::Type{T}, img::GrayArray{F}) = ClampMinMax(T, zero(F), one(F))
mapinfo{T<:UFixed,F<:AbstractFloat}(::Type{Gray{T}}, img::GrayArray{F}) = ClampMinMax(Gray{T}, zero(F), one(F))
mapinfo{T<:AbstractFloat, 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)
    C == RGB24 && continue
    @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<:AbstractFloat}(::Type{RGB24}, img::AbstractArray{$C{F}}) = ClampMinMax(RGB24, zero(F), one(F))
    @eval mapinfo{F<:AbstractFloat}(::Type{ARGB32}, img::AbstractArray{$C{F}}) = ClampMinMax(ARGB32, zero(F), one(F))
    for AC in subtypes(TransparentColor)
        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<:AbstractFloat}(::Type{ARGB32}, img::AbstractArray{$AC{$C{F},F}}) = ClampMinMax(ARGB32, zero(F), one(F))
        @eval mapinfo{F<:AbstractFloat}(::Type{RGB24}, img::AbstractArray{$AC{$C{F},F}}) = ClampMinMax(RGB24, zero(F), one(F))
    end
end

mapinfo{CT<:Colorant}(::Type{RGB24},  img::AbstractArray{CT}) = MapNone{RGB24}()
mapinfo{CT<:Colorant}(::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,Color,AbstractGray}}(::Type{UInt32}, img::AbstractArray{CV}) = mapinfo(RGB24, img)
mapinfo{CV<:TransparentColor}(::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}()

# Clamping mapinfo client. Converts to RGB and uses UFixed, clamping
# floating-point values to [0,1].
mapinfo{T<:UFixed}(::Type{Clamp}, img::AbstractArray{T}) = MapNone{T}()
mapinfo{T<:AbstractFloat}(::Type{Clamp}, img::AbstractArray{T}) = ClampMinMax(UFixed8, zero(T), one(T))
let handled = Set()
for ACV in (Color, AbstractRGB)
    for CV in subtypes(ACV)
        (length(CV.parameters) == 1 && !(CV.abstract)) || continue
        CVnew = CV<:AbstractGray ? Gray : RGB
        @eval mapinfo{T<:UFixed}(::Type{Clamp}, img::AbstractArray{$CV{T}}) = MapNone{$CVnew{T}}()
        @eval mapinfo{CV<:$CV}(::Type{Clamp}, img::AbstractArray{CV}) = Clamp{$CVnew{UFixed8}}()
        CVnew = CV<:AbstractGray ? Gray : BGR
        AC, CA       = alphacolor(CV), coloralpha(CV)
        if AC in handled
            continue
        end
        push!(handled, AC)
        ACnew, CAnew = alphacolor(CVnew), coloralpha(CVnew)
        @eval begin
            mapinfo{T<:UFixed}(::Type{Clamp}, img::AbstractArray{$AC{T}}) = MapNone{$ACnew{T}}()
            mapinfo{P<:$AC}(::Type{Clamp}, img::AbstractArray{P}) = Clamp{$ACnew{UFixed8}}()
            mapinfo{T<:UFixed}(::Type{Clamp}, img::AbstractArray{$CA{T}}) = MapNone{$CAnew{T}}()
            mapinfo{P<:$CA}(::Type{Clamp}, img::AbstractArray{P}) = Clamp{$CAnew{UFixed8}}()
        end
    end
end
end
mapinfo(::Type{Clamp}, img::AbstractArray{RGB24}) = MapNone{RGB{UFixed8}}()
mapinfo(::Type{Clamp}, img::AbstractArray{ARGB32}) = MapNone{BGRA{UFixed8}}()

# 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)})

"""
```
imgsc = sc(img)
imgsc = sc(img, min, max)
```

Applies default or specified `ScaleMinMax` mapping to the image.
"""
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<:Colorant}(A::AbstractArray{C})
            newC = eval(C.name.name){$T}
            convert(Array{newC}, A)
        end
        $fn{C<:Colorant}(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:
using BasisMatrices,LinearAlgebra,Parameters,Optim,QuantEcon,DataFrames,Gadfly,SparseArrays,Arpack
using Roots

@with_kw mutable struct HHModel
    #Preference Parameters
    γ::Float64 = 1. #Risk aversion
    β::Float64 = 0.985 #Discount Rate

    #Prices
    r̄::Float64 = .01 #quarterly
    w̄::Float64 = 1.

    #Asset Grid Parameters
    a̲::Float64 = 0. #Borrowing Constraint
    a̅::Float64 = 600. #Upper Bound on assets
    Na::Int64 = 100 #Number of grid points for splines

    #Income Process
    ρ_ϵ::Float64 = 0.9923 #calibrated to quarterly wage regressions
    σ_ϵ::Float64 = 0.0983
    Nϵ::Int64 = 7
    ϵ::Vector{Float64} = zeros(0)
    Π::Matrix{Float64} = zeros(0,0)

    #Solution
    k::Int = 2 #type of interpolation
    Vf::Vector{Interpoland} = Interpoland[]
    cf::Vector{Interpoland} = Interpoland[]

    #Extra
    EΦ′::SparseMatrixCSC{Float64,Int64} = spzeros(0,0)
    Φ::SparseMatrixCSC{Float64,Int64} = spzeros(0,0)
end;

"""
    U(HH::HHModel,c)
"""
function U(HH::HHModel,c)
    γ = HH.γ
    if γ == 1
        return log.(c)
    else
        return (c.^(1-γ))./(1-γ)
    end
end

"""
    setupgrids_shocks!(HH::HHModel, curv=1.7)

Set up non-linear grids for interpolation
"""
function setupgrids_shocks!(HH::HHModel, curv=1.7)
    @unpack a̲,a̅,Na,ρ_ϵ,σ_ϵ,Nϵ,k,r̄,w̄,β = HH
    #Compute grid on A
    agrid = (a̅-a̲).*LinRange(0,1,Na).^curv .+ a̲

    #Store markov chain
    mc = rouwenhorst(Nϵ,ρ_ϵ,σ_ϵ)
    HH.Π = Π = mc.p
    HH.ϵ = exp.(mc.state_values)

    #First guess of interpolation functions
    abasis = Basis(SplineParams(agrid,0,k))
    a = nodes(abasis)[1]

    Vf = HH.Vf = Vector{Interpoland}(undef,Nϵ)
    cf = HH.cf = Vector{Interpoland}(undef,Nϵ)
    for s in 1:Nϵ
        c = @. r̄*a + w̄*HH.ϵ[s]
        V = U(HH,c)./(1-β)

        Vf[s]= Interpoland(abasis,V)
        cf[s]= Interpoland(abasis,c)
    end

    #Expectations of 1st derivative of Basis functions
    HH.EΦ′ = kron(Π,BasisMatrix(abasis,Direct(),nodes(abasis)[1],[1]).vals[1])
    HH.Φ = kron(Matrix{Float64}(I,Nϵ,Nϵ),BasisMatrix(abasis,Direct()).vals[1])
end;

"""
computeoptimalconsumption(HH::HHModel,V)

Computes optimal savings using endogenous grid method.
"""
function computeoptimalconsumption(TM::TaxModel,Vcoefs::Array{Float64},ty,j)::Vector{Interpoland}
    @unpack a̲,EΦ′,η,ϵ,r,w,γ,β,ns,σ,jr,lgrid,nl,k,κ0,κ1,κ2,τp,τc,ȳ,ψ,Π = TM

    a′grid = nodes(TM.Vf[1,1,1].basis)[1]
    na = length(a′grid)
    EV_a = reshape(EΦ′*Vcoefs,:,ns) #Compute expectations of V'(a',s') using matrix multiplication
    v_max = zeros(na,ns)
    c = zeros(na,ns)
    a = zeros(na,ns)
    l = zeros(na,ns)

    if j < jr
        cl = zeros(nl,na,ns)
        al = zeros(nl,na,ns)
        for l in 1:nl # Find optimal consumption for each l in lgrid
            earn = w*ϵ[ty,j]*η*lgrid[l]
            lab_tax = tax_gs(κ0,κ1,κ2,earn)
            ss_tax =  τp*min.(earn,ȳ)
            # Consumption from Euler Equation, Implied assets today from Budget Constraint
            cl[l,:,:] = (β*ψ[j]/(γ*(1-lgrid[l]))^((1-σ)*(1-γ)).*EV_a).^(1/(γ*(1-σ)-1))
            al[l,:,:] = ((a′grid .+ (1+τc)*c[l] .+ ss_tax' .+ lab_tax' .- earn')/(1+r*(1-τk))) - Tr
            for apr in 1:na, s in 1:ns
                val[l,apr,s] = U(TM,c[l,apr,s],lgrid[l]) + β*ψ[j]*sum(map(eta -> TM.Vf[eta,ty,j+1](apr) * Π[s,eta],1:ns))
            end
        end
        for apr in 1:na, s in 1:ns
           v_max[apr,s],ind_max[apr,s] = findmax(val[:,apr,s])
           c[apr,s] = cl[ind_max[apr,s],apr,s]
           a[apr,s] = al[ind_max[apr,s],apr,s]
           l[apr,s] = lgrid[ind_max[apr,s]]
        end
    end

    cf = Vector{Interpoland}(undef,ns)#implied policy rules for each productivity
    for s in 1:ns
        #if a[1,s]> a̲
        #    c̲ = (earn - ss_tax + (1+r*(1-τk))*(a + Tr) - lab_tax - a̲)/(1+τc)
        #    cf[s]= Interpoland(Basis(SplineParams([a̲; a[:,s]],0,k)),[c̲;c[:,s]])
        #else
            cf[s]= Interpoland(Basis(SplineParams(a[:,s],0,k)),c[:,s])
        #end
    end
    return cf
end;

"""
iteratebellman_newton!(AM::AiyagriModel,Vcoefs)

Updates the coefficients of the value function using newton's method
"""
function iteratebellman_newton!(TM::TaxMod)
    @unpack β,ϵ,Π,r̄,w̄,Nϵ,Vf,Φ  ,ns,ty,j = TM

    agrid = nodes(Vf[1,1,1].basis)[1]
    Na = length(agrid)
    Vcoefs = zeros(ns)
    cf = Array{Interpoland}(undef,ns,nty,J);

    for ty in 1:nty, j in 1:J
        Vcoefs = vcat([Vf[s,ty,j].coefs for s in 1:ns]...)::Vector{Float64}
        cf[:,ty,j] = computeoptimalconsumption(TM,Vcoefs,ty,j) #Compute optimal consumption function

        c = zeros(Na*Nϵ)
        EΦ = spzeros(Na*Nϵ,Na*Nϵ)
        for s in 1:Nϵ
            for s′ in 1:Nϵ
                c[(s-1)*Na+1:s*Na] = cf[s](agrid) #compute consumption at gridpoints
                a′ = (1+r̄)*agrid .+ ϵ[s]*w̄ .- c[(s-1)*Na+1:s*Na] #asset choice
                #Compute expectation of basis functions at a′
                EΦ[(s-1)*Na+1:s*Na,(s′-1)*Na+1:s′*Na] = Π[s,s′]*BasisMatrix(Vf[s].basis,Direct(),a′).vals[1][:]
            end
        end

        Jac = β.*EΦ .- Φ
        res = U(HH,c) .+ Jac*Vcoefs
        Vcoefs′ = Vcoefs - Jac\res #newtons method
        for s in 1:Nϵ
            Vf[s].coefs .= Vcoefs′[1+(s-1)*Na:s*Na]
        end
    end
    return norm(res,Inf)
end;

"""
    iteratebellman_time!(AM::AiyagriModel,Vcoefs)

Updates the coefficients of the value function using time iteration of the bellman equation
"""
function iteratebellman_time!(HH::HHModel)
    @unpack β,ϵ,Π,r̄,w̄,Nϵ,Vf,Φ = HH
    Vcoefs = vcat([Vf[s].coefs for s in 1:Nϵ]...)::Vector{Float64}
    agrid = nodes(Vf[1].basis)[1]
    Na = length(agrid)

    cf = computeoptimalconsumption(HH,Vcoefs) #Compute optimal consumption function
    c = zeros(Na*Nϵ)
    EΦ = spzeros(Na*Nϵ,Na*Nϵ)
    for s in 1:Nϵ
        for s′ in 1:Nϵ
            c[(s-1)*Na+1:s*Na] = cf[s](agrid) #compute consumption at gridpoints
            a′ = (1+r̄)*agrid .+ ϵ[s]*w̄ .- c[(s-1)*Na+1:s*Na] #asset choice
            #Compute expectation of basis functions at a′
            EΦ[(s-1)*Na+1:s*Na,(s′-1)*Na+1:s′*Na] = Π[s,s′]*BasisMatrix(Vf[s].basis,Direct(),a′).vals[1][:]
        end
    end

    res = U(HH,c) .+ β.*EΦ*Vcoefs - Φ*Vcoefs
    Vcoefs′ = Φ\(U(HH,c) .+ β.*EΦ*Vcoefs)
    for s in 1:Nϵ
        Vf[s].coefs .= Vcoefs′[1+(s-1)*Na:s*Na]
    end
    return norm(res,Inf)
end;

"""
    solvebellman!(HH::HHModel)

Solves the bellman equation for given some initial
value function V.
"""
function solvebellman!(HH::HHModel,tol=1e-8)
    #increases stability to iterate on the time dimension a few times
    diff = 1.
    for _ in 1:5
        iteratebellman_time!(HH)
    end
    while diff > tol
        #then use newtons method
        diff = iteratebellman_newton!(HH)
    end
    Vcoefs = vcat([HH.Vf[s].coefs for s in 1:HH.Nϵ]...)::Vector{Float64}
    HH.cf = computeoptimalconsumption(HH,Vcoefs)
end;

"""
    setupgrids_shocks!(AM::AiyagariModel)

Setup the grids and shocks for the aiyagari model
"""
function setupgrids_shocks!(AM::AiyagariModel,curv=2.)
    @unpack HH,Ia,N̄= AM
    @unpack a̲,a̅,Nϵ = HH
    setupgrids_shocks!(HH)
    #Normalize so that average labor supply is 1
    πstat = real(eigs(HH.Π',nev=1)[2])
    πstat ./= sum(πstat)
    HH.ϵ = HH.ϵ./dot(πstat,HH.ϵ)*N̄
    #Grid for distribution
    agrid = (a̅-a̲).*LinRange(0,1,Ia).^curv .+ a̲
    AM.z̄ = hcat(kron(ones(Nϵ),agrid),kron(1:Nϵ,ones(Ia)))
    AM.ω̄ = ones(Ia*Nϵ)/(Ia*Nϵ)
end;

@with_kw mutable struct AiyagariModel
    HH::HHModel = HHModel()

    #Production Parameters
    α::Float64 = 0.3
    δ::Float64 = 0.025
    Θ̄::Float64 = 1.

    #Moments to match/prices
    W̄::Float64 = 1.
    R̄::Float64 = 1.01
    K2Y::Float64 = 10.2 #capital to output ratio
    N̄::Float64 = 1.

    #Distribution Parameters
    Ia::Int = 1000 #Number of gridpoints for distribution
    z̄::Matrix{Float64} = zeros(0,0) #Gridpoints for the state variables
    ω̄::Vector{Float64} = zeros(0) #Fraction of agents at each grid level
    H::SparseMatrixCSC{Float64,Int64} = spzeros(Ia,Ia) #Transition matrix
end;

"""
    find_stationarydistribution!(AM::AiyagariModel,V)

Computes the stationary distribution
"""
function find_stationarydistribution!(AM::AiyagariModel)
    @unpack Ia,z̄,HH,W̄,R̄ = AM
    @unpack ϵ,Π,Nϵ,cf,a̲,a̅ = HH

    a = z̄[1:Ia,1] #grids are all the same for all shocks
    c = hcat([cf[s](a) for s in 1:Nϵ]...) #consumption policy IaxNϵ
    a′ = R̄.*a .+ W̄.*ϵ' .- c #create a IaxNϵ grid for the policy rules

    #make sure we don't go beyond bounds.  Shouldn't bind if bmax is correct
    a′ = max.(min.(a′,a̅),a̲)

    Qa = BasisMatrix(Basis(SplineParams(a,0,1)),Direct(),reshape(a′,Ia*Nϵ)).vals[1]
    Q = spzeros(Ia*Nϵ,Ia*Nϵ)
    for s in 1:Nϵ
        Q[1+(s-1)*Ia:s*Ia,:] = kron(reshape(Π[s,:],1,:),Qa[1+(s-1)*Ia:s*Ia,:])
    end

    AM.H = Q'
    AM.ω̄ .= real(eigs(AM.H;nev=1)[2])[:]
    AM.ω̄ ./= sum(AM.ω̄) #normalize eigenvector
end;

"""
    capital_supply_demand(AM,R)

Compute the supply and demand for capital which prevails in
the stationary distribution at a given interest rate.
"""
function capital_supply_demand(AM,R)
    @unpack Θ̄,α,N̄,δ = AM
    AM.R̄ = R
    AM.HH.r̄ = R-1

    Y2K = (R-1+δ)/α
    K2N = (Y2K/Θ̄)^(1/(α-1))
    AM.W̄ = AM.HH.w̄ = (1-α)*Θ̄*K2N^α
    KD = K2N * N̄

    solvebellman!(AM.HH)
    find_stationarydistribution!(AM)
    KS = dot(AM.ω̄,AM.z̄[:,1])

    return [KS,KD]
end;

AM = AiyagariModel()
AM.HH.β = 0.99
setupgrids_shocks!(AM)

setupgrids_shocks!(HH,3.)
solvebellman!(HH)
setupgrids_shocks!(HH,3.)
@time solvebellman!(HH);

Rgrid = LinRange(1.,1.007,10)
KSKD = hcat([capital_supply_demand(AM,R) for R in Rgrid]...)
plot(layer(y=Rgrid,x=KSKD[1,:],color=["Capital Supplied"],Geom.line),
     layer(y=Rgrid,x=KSKD[2,:],color=["Capital Demanded"],Geom.line),
     Guide.ylabel("Gross Interest Rate"), Guide.xlabel("Capital"))

AM.R̄ = 1.01 #target a quarterly interest rate of 1%

function calibratesteadystate!(AM)
    @unpack Θ̄,α,N̄,K2Y,R̄ = AM
    AM.HH.r̄ = R̄ - 1
    Y2K = 1/K2Y
    AM.δ = α*Y2K + 1 - R̄ #matching capital to output ratio and interest rate gives depreciation rate
    K2N = (Y2K/Θ̄)^(1/(α-1)) #relationship between capital to output and capital to labor
    K̄ = K2N*N̄
    AM.W̄ = AM.HH.w̄ = (1-α)*Θ̄*K2N^α

    setupgrids_shocks!(AM)
    function βres(β)
        AM.HH.β=β
        solvebellman!(AM.HH)
        find_stationarydistribution!(AM)
        return dot(AM.ω̄,AM.z̄[:,1]) -K̄
    end

    Q̄ = 1/R̄
    fzero(βres,Q̄^2,Q̄^1.2)
end

AM = AiyagariModel()
calibratesteadystate!(AM)
plot(x=AM.z̄[AM.z̄[:,1].>1.,1],y=AM.ω̄[AM.z̄[:,1].>1.],Geom.bar,
    Guide.xlabel("Capital"), Guide.ylabel("Density"))


text:
"""
    explicit_euler_riemann(mu, beta, p0, h, t0, t1, a0, a1)

Integrates a population without migration using Explicit Euler-Riemann.
mu is a mortality function of age and time. Beta is a birth rate of
age and time. p0 is an initial population, as a function of age.
h is the step size in both age and time. (t0, t1) are time limits.
(a0, a1) are age limits. This returns a single array in age and time
with the value of the population.

You must have h*mu < 1.
"""
function explicit_euler_riemann(mu, beta, p0, h, t0, t1, a0, a1)
    a = a0:h:a1
    t = t0:h:t1
    m = length(a)
    n = length(t)

    u = zeros(h, m, n)
    u[:, 1] = p0(a, t0)

    for i in 1:(n - 1)
        u[:, i + 1] = (1 .- h .* mu(a, t[i])) .* u[:, i]
        u[1, i + 1] = h * sum(beta(a, t[i]) .* u[1:m, i])
    end
    u
end

"""
    implicit_euler_riemann(mu, beta, p0, h, t0, t1, a0, a1)

Integrates a population without migration using Implicit Euler-Riemann.
mu is a mortality function of age and time. Beta is a birth rate of
age and time. p0 is an initial population, as a function of age.
h is the step size in both age and time. (t0, t1) are time limits.
(a0, a1) are age limits. This returns a single array in age and time
with the value of the population.

You must have h*mu < 1.
"""
function implicit_euler_riemann(mu, beta, p0, h, t0, t1, a0, a1)
    a = a0:h:a1
    t = t0:h:t1
    m = length(a)
    n = length(t)

    u = zeros(h, m, n)
    u[:, 1] = p0(a, t0)

    for i in 1:(n - 1)
        u[:, i + 1] = u[:, i] ./ (1 .+ h .* mu(a, t[i]))
        u[1, i + 1] = h * sum(beta(a, t[i]) .* u[1:m, i])
    end
    u
end

"""
Solve the problem for q(a, t) = p(a, t) / survival.
Then multiply the survival back in order to get p(a, t).
"""
function euler_riemann_survival(s, beta, p0, h, t0, t1, a0, a1)
    a = a0:h:a1
    t = t0:h:t1
    m = length(a)
    n = length(t)

    u = zeros(h, m, n)
    u[:, 1] = p0(a, t0)

    for i in 1:(n - 1)
        birth = sum(beta(a, t[i]) .* s[a, t[i]] .* u[1:m, i])
        for j in 1:min(m, n - i)
            u[j, i + j] = birth
        end
        u[:, i] .*= s[a, t[i]]
    end
    u[:, n] .*= s[a, t[n]]
    u
end


text:
import Flux.testmode!
import Base: show, deepcopy
using Flux: @treelike
include("resnet.jl")

struct NeuralNet
  base_net::Chain
  value::Chain
  policy::Chain
end

function NeuralNet(env::T; tower_height::Int = 19) where T <: GameEnv
  N, num_planes = env.N, env.planes

  res_block() = ResidualBlock([256,256,256], [3,3], [1,1], [1,1])
  # 19 residual blocks
  tower = [res_block() for i = 1:tower_height]
  base_net = Chain(Conv((3,3), 2num_planes+1=>256, pad=(1,1)),
                   BatchNorm(256, relu),
                   tower...) |> gpu

  value = Chain(Conv((1,1), 256=>1),
	        BatchNorm(1, relu), x->reshape(x, N^2, :),
                Dense(N^2, 256, relu),
	        Dense(256, 1, tanh)) |> gpu

  policy = Chain(Conv((1,1), 256=>2),
                 BatchNorm(2, relu), x->reshape(x, 2N^2, size(x, 4)),
                 Dense(2N^2, env.action_space), softmax) |> gpu

  NeuralNet(base_net, value, policy)
end

function Base.show(io::IO, nn::NeuralNet)
  println("Base: ", nn.base_net)
  println("Policy: ", nn.policy)
  println("Value: ", nn.value)
end

#=
function deepcopy(nn::NeuralNet)
  base_net = deepcopy(nn.base_net)
  value = deepcopy(nn.value)
  policy = deepcopy(nn.policy)

  return NeuralNet(;base_net = base_net, value = value, policy = policy)
end
=#

function testmode!(nn::NeuralNet, val::Bool=true)
  testmode!(nn.base_net, val)
  testmode!(nn.policy, val)
  testmode!(nn.value, val)
end

function (nn::NeuralNet)(input::Vector{T}, train::Bool = false) where T <: Position
  nn_in = cat(dims=4, get_feats.(input)...) |> gpu

  !train && testmode!(nn)

  common_out = nn.base_net(nn_in)
  π, val = nn.policy(common_out), nn.value(common_out)

  !train && testmode!(nn, false)

  return π, val
end

function (nn::NeuralNet)(input::Position)
  p, v = nn([input])
  return p[:, 1], v[1]
end

loss_π(π, p) = crossentropy(p, π; weight = 0.01f0)

loss_value(z, v) = 0.01f0 * mse(z, v)

function loss_reg(nn::NeuralNet)
  sum_sqr(x) = sum([sum(i.^2) for i in x])
  0.0001f0 * sum_sqr(params(nn))
end

function _train(nn::NeuralNet, input_data::T, opt; epochs = 1) where T<:NamedTuple{(:positions, :policy, :values), Tuple{Vector{Position}, Matrix{Float32}, Vector{Int}}}
  positions = input_data.positions
  π = input_data.policy |> gpu
  z = input_data.values |> gpu

  data_size = length(positions)

  for i = 1:epochs, j = 1:32:data_size
    p, v = nn(positions[j:j+31], true)
    loss = loss_π(π[:, j:j+31], p) + loss_value(z[j:j+31], v) + loss_reg(nn)
    back!(loss)
    loss_avg += loss.data
    update!(opt, params(nn))
  end

  return loss_avg / epochs
end

function evaluate(env::GameEnv, black_net::NeuralNet, white_net::NeuralNet;
		  num_games = 400, ro = 800, verbose::Bool=false)
  games_won = 0

  testmode!(black_net)
  testmode!(white_net)

  black = MCTSPlayer(env, black_net, num_readouts = ro, two_player_mode = true)
  white = MCTSPlayer(env, white_net, num_readouts = ro, two_player_mode = true)

  for i = 1:num_games
    num_move = 0  # The move number of the current game

    initialize_game!(black)
    initialize_game!(white)

    while true
      active = num_move % 2 == true ? white : black
      inactive = num_move % 2 == true ? black : white

      current_readouts = N(active.root)
      readouts = active.num_readouts

      while N(active.root) < current_readouts + readouts
        tree_search!(active)
      end

      # First, check the roots for hopeless games.
      if should_resign(active)  # Force resign
        set_result!(active, -active.root.position.to_play, true)
        set_result!(inactive, -active.root.position.to_play, true)
	break
      end

      move = pick_move(active)
      play_move!(active, move)
      play_move!(inactive, move)
      num_move += 1

      if is_done(active)
        winner = result(active.root.position)
        set_result!(active,winner,false)
	set_result!(inactive, winner, false)
	break
      end
    end
    games_won += result(black.root.position) == BLACK
  end

  testmode!(black_net, false)
  testmode!(white_net, false)

  verbose && print("Won $games_won / $num_games. Win rate: $(games_won/num_games). ")

  return games_won / num_games ≥ 0.55
end


text:
# very simple implementation of iterating through all the reads of a BAM file

struct BamState
	io::IO
	zio::IO
	read::BamRead
end

BamState(f::BamFile, io::IO, zio::IO) = BamState(io,zio,BamRead(f))
function BamState(f::BamFile, io::IO)
	zio = ZlibInflateInputStream(io,reset_on_end=true)
	skip_bam_header(zio) # just skip header for now
	BamState(f, io, zio)
end
BamState(f::BamFile) = BamState(f,open(f.filename))

# Base.start(f::BamFile) = BamState(f)
# Base.next(::BamFile, s::BamState) = (nextread!(s.zio, s.read), s)
# Base.done(::BamFile, s::BamState) = eof(s.zio)

function iterate(f::BamFile, s::BamState=BamState(f))
	eof(s.zio) && return nothing
	(nextread!(s.zio, s.read), s)
end
Base.eltype(::Type{BamFile}) = BamRead


text:
export MvNormalMeanCovariance

import Distributions: logdetcov, distrname, sqmahal, sqmahal!, AbstractMvNormal
import LinearAlgebra: diag, Diagonal, dot
import Base: ndims, precision, length, size, prod

struct MvNormalMeanCovariance{ T <: Real, M <: AbstractVector{T}, P <: AbstractMatrix{T} } <: AbstractMvNormal
    μ :: M
    Σ :: P
end

function MvNormalMeanCovariance(μ::AbstractVector{ <: Real }, Σ::AbstractMatrix{ <: Real })
    T = promote_type(eltype(μ), eltype(Σ))
    return MvNormalMeanCovariance(convert(AbstractArray{T}, μ), convert(AbstractArray{T}, Σ))
end

function MvNormalMeanCovariance(μ::AbstractVector{ <: Integer}, Σ::AbstractMatrix{ <: Integer })
    return MvNormalMeanCovariance(float.(μ), float.(Σ))
end

function MvNormalMeanCovariance(μ::AbstractVector, σ::AbstractVector)
    return MvNormalMeanCovariance(μ, matrix_from_diagonal(promote_type(eltype(μ), eltype(σ)), σ))
end

function MvNormalMeanCovariance(μ::AbstractVector{T}) where T
    return MvNormalMeanCovariance(μ, convert(AbstractArray{T}, ones(length(μ))))
end

Distributions.distrname(::MvNormalMeanCovariance) = "MvNormalMeanCovariance"

function weightedmean(dist::MvNormalMeanCovariance)
    z = fastcholesky(cov(dist))
    return z \ mean(dist)
end

function weightedmean_invcov(dist::MvNormalMeanCovariance)
    W = precision(dist)
    xi = W * mean(dist)
    return (xi, W)
end

weightedmean_precision(dist::MvNormalMeanCovariance) = weightedmean_invcov(dist)

Distributions.mean(dist::MvNormalMeanCovariance)      = dist.μ
Distributions.var(dist::MvNormalMeanCovariance)       = diag(cov(dist))
Distributions.cov(dist::MvNormalMeanCovariance)       = dist.Σ
Distributions.invcov(dist::MvNormalMeanCovariance)    = cholinv(dist.Σ)
Distributions.std(dist::MvNormalMeanCovariance)       = cholsqrt(cov(dist))
Distributions.logdetcov(dist::MvNormalMeanCovariance) = chollogdet(cov(dist))

Distributions.sqmahal(dist::MvNormalMeanCovariance, x::AbstractVector) = sqmahal!(similar(x), dist, x)

function Distributions.sqmahal!(r, dist::MvNormalMeanCovariance, x::AbstractVector)
    μ = mean(dist)
    for i in 1:length(r)
        @inbounds r[i] = μ[i] - x[i]
    end
    return dot(r, invcov(dist), r) # x' * A * x
end

Base.eltype(::MvNormalMeanCovariance{T})     where T = T
Base.precision(dist::MvNormalMeanCovariance)         = invcov(dist)
Base.length(dist::MvNormalMeanCovariance)            = length(mean(dist))
Base.ndims(dist::MvNormalMeanCovariance)             = length(dist)
Base.size(dist::MvNormalMeanCovariance)              = (length(dist), )

Base.convert(::Type{ <: MvNormalMeanCovariance }, μ::AbstractVector, Σ::AbstractMatrix) = MvNormalMeanCovariance(μ, Σ)

function Base.convert(::Type{ <: MvNormalMeanCovariance{T} }, μ::AbstractVector, Σ::AbstractMatrix) where { T <: Real }
    return MvNormalMeanCovariance(convert(AbstractArray{T}, μ), convert(AbstractArray{T}, Σ))
end

vague(::Type{ <: MvNormalMeanCovariance }, dims::Int) = MvNormalMeanCovariance(zeros(dims), fill(huge, dims))

prod_analytical_rule(::Type{ <: MvNormalMeanCovariance }, ::Type{ <: MvNormalMeanCovariance }) = ProdAnalyticalRuleAvailable()

function Base.prod(::ProdPreserveType, left::MvNormalMeanCovariance, right::MvNormalMeanCovariance)
    invcovleft  = invcov(left)
    invcovright = invcov(right)
    Σ = cholinv(invcovleft + invcovright)
    μ = Σ * (invcovleft * mean(left) + invcovright * mean(right))
    return MvNormalMeanCovariance(μ, Σ)
end

function Base.prod(::ProdAnalytical, left::MvNormalMeanCovariance, right::MvNormalMeanCovariance)
    xi_left, W_left = weightedmean_precision(left)
    xi_right, W_right = weightedmean_precision(right)
    return MvNormalWeightedMeanPrecision(xi_left + xi_right, W_left + W_right)
end

function Base.prod(::ProdAnalytical, left::MvNormalMeanCovariance{T1}, right::MvNormalMeanCovariance{T2}) where { T1 <: LinearAlgebra.BlasFloat, T2 <: LinearAlgebra.BlasFloat }

    # start with parameters of left
    xi, W = weightedmean_precision(left)

    # update W
    W_right = precision(right)
    W .+= W_right

    # update xi without allocating another vector
    T = promote_type(T1, T2)
    xi = convert(AbstractVector{T}, xi)
    W  = convert(AbstractMatrix{T}, W)
    xi = BLAS.gemv!('N', one(T), convert(AbstractMatrix{T}, W_right), convert(AbstractVector{T}, mean(right)), one(T), xi)

    return MvNormalWeightedMeanPrecision(xi, W)
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

module StorageManagementClient

using Random
using Dates
using Swagger
import Swagger: field_name, property_type, hasproperty, validate_property, SwaggerApi, SwaggerModel
import Base: convert, propertynames

include("modelincludes.jl")

include("api_EncryptionScopesApi.jl")
include("api_LocationUsageApi.jl")
include("api_ManagementPoliciesApi.jl")
include("api_ObjectReplicationPoliciesApi.jl")
include("api_OperationsApi.jl")
include("api_PrivateEndpointConnectionsApi.jl")
include("api_PrivateLinkResourcesApi.jl")
include("api_SkusApi.jl")
include("api_StorageAccountsApi.jl")

# export models
export convert, AccountSasParameters
export convert, ActiveDirectoryProperties
export convert, AzureFilesIdentityBasedAuthentication
export convert, BlobRestoreParameters
export convert, BlobRestoreRange
export convert, BlobRestoreStatus
export convert, CheckNameAvailabilityResult
export convert, CustomDomain
export convert, DateAfterCreation
export convert, DateAfterModification
export convert, Dimension
export convert, Encryption
export convert, EncryptionScope
export convert, EncryptionScopeKeyVaultProperties
export convert, EncryptionScopeListResult
export convert, EncryptionScopeProperties
export convert, EncryptionService
export convert, EncryptionServices
export convert, Endpoints
export convert, ErrorResponse
export convert, GeoReplicationStats
export convert, IPRule
export convert, Identity
export convert, KeyVaultProperties
export convert, ListAccountSasResponse
export convert, ListServiceSasResponse
export convert, ManagementPolicy
export convert, ManagementPolicyAction
export convert, ManagementPolicyBaseBlob
export convert, ManagementPolicyDefinition
export convert, ManagementPolicyFilter
export convert, ManagementPolicyProperties
export convert, ManagementPolicyRule
export convert, ManagementPolicySchema
export convert, ManagementPolicySnapShot
export convert, MetricSpecification
export convert, NetworkRuleSet
export convert, ObjectReplicationPolicies
export convert, ObjectReplicationPolicy
export convert, ObjectReplicationPolicyFilter
export convert, ObjectReplicationPolicyProperties
export convert, ObjectReplicationPolicyRule
export convert, Operation
export convert, OperationDisplay
export convert, OperationListResult
export convert, OperationProperties
export convert, PrivateEndpoint
export convert, PrivateEndpointConnection
export convert, PrivateEndpointConnectionListResult
export convert, PrivateEndpointConnectionProperties
export convert, PrivateEndpointConnectionProvisioningState
export convert, PrivateEndpointServiceConnectionStatus
export convert, PrivateLinkResource
export convert, PrivateLinkResourceListResult
export convert, PrivateLinkResourceProperties
export convert, PrivateLinkServiceConnectionState
export convert, Resource
export convert, Restriction
export convert, RoutingPreference
export convert, SKUCapability
export convert, ServiceSasParameters
export convert, ServiceSpecification
export convert, Sku
export convert, SkuInformation
export convert, SkuName
export convert, StorageAccount
export convert, StorageAccountCheckNameAvailabilityParameters
export convert, StorageAccountCreateParameters
export convert, StorageAccountInternetEndpoints
export convert, StorageAccountKey
export convert, StorageAccountListKeysResult
export convert, StorageAccountListResult
export convert, StorageAccountMicrosoftEndpoints
export convert, StorageAccountProperties
export convert, StorageAccountPropertiesCreateParameters
export convert, StorageAccountPropertiesUpdateParameters
export convert, StorageAccountRegenerateKeyParameters
export convert, StorageAccountUpdateParameters
export convert, StorageSkuListResult
export convert, TagFilter
export convert, Tier
export convert, TrackedResource
export convert, Usage
export convert, UsageListResult
export convert, UsageName
export convert, VirtualNetworkRule

# export operations
export convert, EncryptionScopesApi, LocationUsageApi, ManagementPoliciesApi, ObjectReplicationPoliciesApi, OperationsApi, PrivateEndpointConnectionsApi, PrivateLinkResourcesApi, SkusApi, StorageAccountsApi

export check_required, field_name, property_type, hasproperty, propertynames, validate_property, convert

end


text:
using Test
using GrammianAngularFields
using StatsBase

@testset "trigonometric identities" begin
    θ = 2/3*π
    @test sin(-θ) ≈ -sin(θ)
    @test cos(-θ) ≈ cos(θ)
    @test sin(2θ) ≈ 2*sin(θ)*cos(θ)
    @test cos(2θ) ≈ cos(θ)^2 - sin(θ)^2
end;

@testset "transformation checks" begin
    x = GrammianAngularFields.normalize(collect(1:10))
    @test StatsBase.mean(x) ≈ 0
    @test StatsBase.stdm(x, 0) ≈ 1
    x = GrammianAngularFields.standardize(collect(1:10))
    @test maximum(x) ≈ 1
    @test minimum(x) ≈ 0
    x = GrammianAngularFields.standardize(collect(1:10), 3, 5)
    @test maximum(x) ≈ 5
    @test minimum(x) ≈ 3
    x = GrammianAngularFields.grammify(collect(1:3))
    @test x ≈ [1.0 -1.83697e-16 -1.0; -1.83697e-16 -1.0 6.12323e-17; -1.0 6.12323e-17 1.0]
end;


text:
export build_sparsity, get_variables, get_locctrcliques, get_maxcliques, collect_cliquesvars

"""
    momentmat_param, localizingmat_param = build_sparsity(relax_ctx, problem)

    Build the sparsitty pattern and variables decomposition for laying out the moment or SOS hierarchy
"""
function build_sparsity(relax_ctx::RelaxationContext, problem::Problem, max_cliques::Dict{String, Set{Variable}})

    ((relax_ctx.issparse == false) && (length(max_cliques) > 1)) && error(LOGGER, "build_sparsity(): Relaxation is not sparse, one clique is expected (not $(length(max_cliques)))")

    # Build localizing constraints order and variable set.
    localizingmat_param = Dict{String, Tuple{Set{String}, Int}}()
    for (ctrname, ctr) in problem.constraints
        ctrtype = get_cstrtype(ctr)
        ctrcliques = get_locctrcliques(ctr.p, max_cliques)

        if ctrtype == :ineqdouble
            ctrname_lo, ctrname_up = get_cstrname(ctrname, ctrtype)
            di_lo, ki_lo = relax_ctx.di[ctrname_lo], relax_ctx.ki[ctrname_lo]
            di_up, ki_up = relax_ctx.di[ctrname_up], relax_ctx.ki[ctrname_up]

            if relax_ctx.hierarchykind==:Complex
                localizingmat_param[ctrname_lo] = (ctrcliques, di_lo-ki_lo)
                localizingmat_param[ctrname_up] = (ctrcliques, di_up-ki_up)
            elseif relax_ctx.hierarchykind==:Real
                localizingmat_param[ctrname_lo] = (ctrcliques, di_lo-ceil(ki_lo/2))
                localizingmat_param[ctrname_up] = (ctrcliques, di_up-ceil(ki_up/2))
            else
                error(LOGGER, "build_sparsity(): Unknown relaxation kind $(relax_ctx.hierarchykind). Should be `:Real` or `:Complex`")
            end
        else # :ineqlo, :ineqhi, :eq
            di, ki = relax_ctx.di[get_cstrname(ctrname, ctrtype)], relax_ctx.ki[get_cstrname(ctrname, ctrtype)]
            if relax_ctx.hierarchykind == :Complex
                localizingmat_param[get_cstrname(ctrname, ctrtype)] = (ctrcliques, di-ki)
            elseif relax_ctx.hierarchykind == :Real
                localizingmat_param[get_cstrname(ctrname, ctrtype)] = (ctrcliques, di-ceil(ki/2))
            else
                error(LOGGER, "build_sparsity(): Unknown relaxation kind $(relax_ctx.hierarchykind). Should be `:Real` or `:Complex`")
            end
        end
    end

    # Build moment constraints order and variable set.
    momentmat_param = Dict{String, Int}()
    for (cliquename, cliquevars) in max_cliques
        cur_d::Int = -1
        for (ctrname, (ctrcliques, _)) in localizingmat_param
            if length(ctrcliques) == 1 && cliquename == first(ctrcliques)
                cur_d = max(cur_d, relax_ctx.di[ctrname])
            end
        end
        if cur_d == -1 # Clique does not match a full constraint
            cur_d = minimum(values(relax_ctx.di))
        end
        momentmat_param[cliquename] = cur_d
    end

    return momentmat_param, localizingmat_param
end

"""
pvars = get_variables(p)

Collect all variables appearing in `p`.
"""
function get_variables(p::Polynomial)
    pvars = Set{Variable}()
    for (expo, coeff) in p
        for (var, deg) in expo
            push!(pvars, var)
        end
    end
    return pvars
end

"""
    locctrcliques = get_locctrcliques(p, max_cliques)

    Find a minimal set of cliques gathering all variables from polynomial `p`.
"""
function get_locctrcliques(p::Polynomial, max_cliques::Dict{String, Set{Variable}})
    ctrvars = get_variables(p)

    # Build constraint variables to cliques dict
    var_to_cliques = Dict{Variable, Set{String}}()
    for (clique, cliquevars) in max_cliques
        for var in intersect(cliquevars, ctrvars)
            if !haskey(var_to_cliques, var)
                var_to_cliques[var] = Set{String}()
            end
            push!(var_to_cliques[var], clique)
        end
    end

    def_cliques = Set{String}()
    unaffected_vars = Set{Variable}(keys(var_to_cliques))

    keepon = true
    i = 0
    while keepon
        i += 1
        # Find which variables appear in one clique only, remove them from unaffected_vars
        for var in unaffected_vars
            cliques = var_to_cliques[var]
            if length(cliques) == 1
                push!(def_cliques, first(cliques))
                delete!(unaffected_vars, var)
            end
        end

        # Remove variables involved in at least one clique previously selected
        for var in unaffected_vars
            inter = intersect(var_to_cliques[var], def_cliques)
            if !isempty(inter)
                var_to_cliques[var] = Set{String}([first(inter)]) # NOTE: proper way to choose in inter here ?
                delete!(unaffected_vars, var)
            end
        end

        # Hopefully all variables are treated that way. Else repeat this process by choosing a clique. Again, which one ?
        if length(unaffected_vars) != 0
            # warn(LOGGER, "get_locctrcliques(): length(unaffected_vars) = $(length(unaffected_vars))") # TODO: better logging system...
            cliques_from_unaffvar = Dict{String, Int}()
            for var in unaffected_vars
                for clique in var_to_cliques[var]
                    # haskey(cliques_from_unaffvar, clique) || (cliques_from_unaffvar[clique] = 0)
                    # cliques_from_unaffvar[clique] += 1
                    addindex!(cliques_from_unaffvar, 1, clique)
                end
            end
            cur_clique, cur_pop = first(cliques_from_unaffvar)[1], first(cliques_from_unaffvar)[2]
            for (clique, pop) in cliques_from_unaffvar
                if pop > cur_pop
                    cur_clique = clique
                    cur_pop = pop
                end
            end
            push!(def_cliques, cur_clique)
        else
            keepon = false
        end
    end

    return def_cliques
end

function get_maxcliques(relax_ctx, problem)
    vars = Set{Variable}([Variable(name, kind) for (name, kind) in problem.variables])
    return Dict{String, Set{Variable}}("clique1"=>vars)
end

"""
    vars, blocname = collect_cliquesvars(clique_keys, max_cliques)

    Collect variables of `cliques_keys` cliques, described in `max_cliques`
"""
function collect_cliquesvars(clique_keys::Set{String}, max_cliques::Dict{String, Set{Variable}})
    # Collect variables involved in constraint
    vars = Set{Variable}()
    blocname = ""
    for clique_key in clique_keys
        union!(vars, max_cliques[clique_key])
        blocname = blocname*clique_key*"_"
    end
    return vars, blocname[1:end-1]
end

function print(io::IO, max_cliques::Dict{String, Set{Variable}})
    for cliquename in sort(collect(keys(max_cliques)))
        vars = max_cliques[cliquename]

        print(io, "$cliquename = ")
        for var in vars print(io, "$var, ") end
        @printf(io, "\b\b \n")
    end
end

#################################################################################
## Old stuff
"""
    SparsityPattern

    Type for storing and working on sparsitty patterns.
"""
abstract type SparsityPattern end

"""
    sparsity_pattern = compute_sparsitypattern(problem, di, ki)
    Compute the sparsity_pattern corresponding to the given partial orders.
"""
function compute_sparsitypattern(problem::Problem, relax_ctx)
    println("\n=== compute_sparsitypattern(problem::Problem, relax_ctx)")
    println("Compute the sparsity_pattern corresponding to the given partial orders.")
    println("-> Nb vertex / edges:      xx / xx")
    println("-> Order by vertex:        xx / xx (mean/std)")
    sparsity_pattern = SparsityPattern()
    return sparsity_pattern::SparsityPattern
end

"""
    compute_chordalextension!(sparsity_pattern)
    Compute the chordal extension on the provided sparsity_pattern.
"""
function compute_chordalextension!(sparsity_pattern::SparsityPattern)
    println("\n=== compute_chordalextension!(sparsity_pattern::SparsityPattern)")
    println("Compute the chordal extension on the provided sparsity_pattern.")
    println("-> Nb added edges:                 xx")
    println("-> Nb vertex / edges:              xx / xx")
    println("-> Order by vertex:                xx / xx (mean/std)")
end

"""
    maxcliques = compute_maxcliques(sparsity_pattern)
    Compute a `Array{Set{Variable}}` describing the maximum cliques on the provided sparsity_pattern.
"""
function compute_maxcliques(sparsity_pattern::SparsityPattern)
    println("\n=== compute_maxcliques(sparsity_pattern::SparsityPattern)")
    println("Compute a `Array{Set{Variable}}` describing the maximum cliques on the provided sparsity_pattern.")
    println("-> Nb cliques / nb vertex:         xx / xx")
    println("-> Nb vertex by clique:            xx / xx (mean/std)")
    maxcliques = Array{Set{Variable}}()
    return maxcliques::Array{Set{Variable}}
end

"""
    varsbycstr = compute_varsbycstr(problem)
    Compute a `Dict{String, Set{Variable}}` providing the set of variables involved in each constraint.
"""
function compute_varsbycstr(problem::Problem)
    println("\n=== compute_varsbycstr(problem::Problem)")
    println("Compute a `Dict{String, Set{Variable}}` providing the set of variables involved in each constraint.")
    println("-> Nb poly constraints:                xx")
    println("-> Nb variables by poly constraint:    xx / xx (mean/std)")
    varsbycstr = Dict{String, Set{Variable}}()
    return varsbycstr::Dict{String, Set{Variable}}
end

"""
    cliquevarsbycstr = compute_varsbycstr(sparsity_pattern, max_cliques, varsbycstr)
    Compute a `Dict{String, Set{Variable}}` providing the set of variables involved in the SDP localizing matrix corresponding to each constraint.
"""
function compute_varsbycstr(sparsity_pattern, max_cliques, varsbycstr)
    println("\n=== compute_varsbycstr(sparsity_pattern, max_cliques, varsbycstr)")
    println("Compute a `Dict{String, Set{Variable}}` providing the set of variables involved in the SDP localizing matrix corresponding to each constraint.")
    println("-> Nb SDP constraints:                 xx")
    println("-> Nb variables by SDP constraint:     xx / xx (mean/std)")
    cliquevarsbycstr = Dict{String, Set{Variable}}()
    return cliquevarsbycstr::Dict{String, Set{Variable}}
end

"""
    orderbyclique = compute_cliqueorders(sparsity_pattern, di, varsbycstr, max_cliques)
    Compute a `Array{Int}` providing the relaxation order corresponding to each clique.
"""
function compute_cliqueorders(sparsity_pattern, varsbycstr, max_cliques, relax_ctx)
    println("\n=== compute_cliqueorders(sparsity_pattern, varsbycstr, max_cliques, relax_ctx)")
    println("Compute a `Array{Int}` providing the relaxation order corresponding to each clique.")
    println("-> Nb cliques:                         xx")
    println("-> Relaxation order by clique:         xx / xx (mean/std)")
    orderbyclique = Dict{Int, Int}()
    return orderbyclique::Dict{Int, Int}
end


text:
module FEM_P2

using SparseArrays

@enum BdryCond Dirichlet Neumann

export Dirichlet, Neumann, Mesh
export assemble_matrix, assemble_vector, Neumann_bc_vector
export deriv_times_deriv!, func_times_func!
export evaluate_from_ndvals

struct Mesh
    node    :: Vector{Float64}
    elm     :: Matrix{Int64}
    left    :: BdryCond
    right   :: BdryCond
    ndof    :: Int64
    leftnd  :: Int64
    rightnd :: Int64
end

# Element stiffness matrix = MAT1 / (3hm)
const MAT1 = [ 7.0  -8.0   1.0
              -8.0  16.0  -8.0
               1.0  -8.0   7.0 ]

# Element mass matrix = (hm/30) * MAT2
const MAT2 = [  4.0   2.0  -1.0
                2.0  16.0   2.0
               -1.0   2.0   4.0 ]

Psi1(ξ) = 2 * ( ξ - 1/2 ) * ( ξ - 1 )
Psi2(ξ) = 4 * ξ * ( 1 - ξ )
Psi3(ξ) = 2 * ξ * ( ξ - 1/2 )
const PSI = ( Psi1, Psi2, Psi3 )

dPsi1(ξ) = 4ξ - 3
dPsi2(ξ) = 4 * ( 1 - 2ξ )
dPsi3(ξ) = 4ξ - 1
const DPSI = [ dPsi1(0.0) dPsi2(0.0) dPsi3(0.0)
               dPsi1(0.5) dPsi2(0.5) dPsi3(0.5)
               dPsi1(1.0) dPsi2(1.0) dPsi3(1.0) ]

const SIMPSON_WT = [ 1/6, 4/6, 1/6 ]

function Mesh(pts::Vector{Float64}, bc::Tuple{BdryCond,BdryCond})
    M = length(pts) - 1
    h = diff(pts)
    node = Vector{Float64}(undef, 2M+1)
    elm = Matrix{Int64}(undef, 3, M)
    if bc[1] == Dirichlet
        for m = 1:M
            node[2m-1] = pts[m] + h[m]/2
            node[2m]   = pts[m+1]
        end
        node[2M+1] = pts[1]
        leftnd = 2M+1
        rightnd = 2M
        elm[:,1] .= [ 2M+1, 1, 2]
        for m = 2:M
            elm[:,m] .= [ 2m-2, 2m-1, 2m ]
        end
        if bc[2] == Dirichlet
            ndof = 2M-1
        else
            ndof = 2M
        end
    else
        node[1] = pts[1]
        leftnd = 1
        rightnd = 2M+1
        for m = 1:M
            node[2m]   = pts[m] + h[m]/2
            node[2m+1] = pts[m+1]
            elm[:,m] .= [ 2m-1, 2m, 2m+1 ]
        end
        if bc[2] == Dirichlet
            ndof = 2M
        else
            ndof = 2M+1
        end
    end
    return Mesh(node, elm, bc[1], bc[2], ndof, leftnd, rightnd)
end

function assemble_matrix(mesh::Mesh, elm_mat!::Function,
                         coef::Union{Float64,Function})
    node, elm, ndof = mesh.node, mesh.elm, mesh.ndof
    M = size(elm, 2)
    Am = zeros(3, 3)
    nodem = zeros(3)
    I = Int64[]
    J = Int64[]
    V = Float64[]
    for m = 1:M
        for p = 1:3
            nodem[p] = node[elm[p,m]]
        end
        elm_mat!(Am, nodem, coef)
        for p = 1:3
            if elm[p,m] > ndof
                continue
            end
            for q = 1:3
                push!(I, elm[p,m])
                push!(J, elm[q,m])
                push!(V, Am[p,q])
            end
        end
    end
    return sparse(I, J, V, ndof, 2M+1)
end

function assemble_vector(mesh::Mesh, f::Function)
    node, elm, ndof = mesh.node, mesh.elm, mesh.ndof
    M = size(mesh.elm, 2)
    F = zeros(ndof)
    Fm = zeros(3)
    nodem = zeros(3)
    I = Int64[]
    V = Float64[]
    for m = 1:M
        for p = 1:3
            nodem[p] = node[elm[p,m]]
        end
        # Use Simpson's rule
        hm = nodem[3] - nodem[1]
        for p = 1:3
            Fm[p] = hm * SIMPSON_WT[p] * f(nodem[p])
        end
        for p = 1:3
            if elm[p,m] > ndof
                continue
            end
            F[elm[p,m]] += Fm[p]
        end
    end
    return F
end

function Neumann_bc_vector(mesh::Mesh, γ::Vector{Float64})
    left, right, ndof = mesh.left, mesh.right, mesh.ndof
    leftnd, rightnd = mesh.leftnd, mesh.rightnd
    count = 0
    if left == Neumann
        count += 1
    end
    if right == Neumann
        count += 1
    end
    if length(γ) != count
        ArgumentError("Wrong number of Neumann boundary values")
    end
    I = zeros(Int64, count)
    V = zeros(Float64, count)
    k = 0
    if left == Neumann
        k += 1
        I[k] = leftnd
        V[k] = γ[k]
    end
    if right == Neumann
        k += 1
        I[k] = rightnd
        V[k] = γ[k]
    end
    return sparsevec(I, V, ndof)
end

function deriv_times_deriv!(A::Matrix{Float64}, node::Vector{Float64},
                            coef::Float64)
    hm = node[3] - node[1]
    A .= MAT1
    A .*= coef/(3hm)
end

function deriv_times_deriv!(A::Matrix{Float64}, node::Vector{Float64},
                            coef::Function)
    hm = node[3] - node[1]
    coefs = coef.(node)
    for q = 1:3
        for p = 1:q-1
            A[p,q] = A[q,p]
        end
        for p = q:3
            s = 0.0
            for j = 1:3
                s += SIMPSON_WT[j] * coefs[j] * DPSI[j,q] * DPSI[j,p]
            end
            A[p,q] = s / hm
        end
    end
end

function func_times_func!(A::Matrix{Float64}, node::Vector{Float64},
                          coef::Float64)
    hm = node[3] - node[1]
    A .= MAT2
    A .*= hm / 30
end

function func_times_func!(A::Matrix{Float64}, node::Vector{Float64},
                          coef::Function)
    hm = node[3] - node[1]
    for p = 1:3
        A[p,p] = hm * SIMPSON_WT[p] * coef(node[p])
    end
end

function evaluate_from_ndvals(mesh::Mesh, Und::Vector{Float64},
                              r::Int64=10)
    node, elm, ndof = mesh.node, mesh.elm, mesh.ndof
    left, right = mesh.left, mesh.right
    leftnd, rightnd = mesh.leftnd, mesh.rightnd
    M = size(elm, 2)
    x = zeros(r*M+1)
    U = zeros(r*M+1)
    x[1] = node[leftnd]
    U[1] = Und[elm[1,1]]
    for m = 1:M
        lo = node[elm[1,m]]
        hi = node[elm[3,m]]
        hm = hi - lo
        step = hm / r
        for k = 1:r
            j = (m-1) * r + 1 + k
            x[j] = lo + k * step
            s = 0.0
            for p = 1:3
                ξ = k / r
                s += Und[elm[p,m]] * PSI[p](ξ)
            end
            U[j] = s
        end
    end
    return x, U
end

end # module

max_stars_repo_name: billmclean/ComputationalMathsNotes

max_issues_repo_name: billmclean/ComputationalMathsNotes

max_forks_repo_name: billmclean/ComputationalMathsNotes