text:
using Base.Test
import ArchGDAL; const AG = ArchGDAL

AG.registerdrivers() do
    AG.read("ospy/data4/aster.img") do ds
        @testset "version 1" begin
            band = AG.getband(ds, 1)
            count = 0
            total = 0
            buffer = Array{AG.getdatatype(band)}(AG.getblocksize(band)..., 1)
            for (cols,rows) in AG.windows(band)
                AG.rasterio!(ds, buffer, Cint[1], rows-1, cols-1)
                data = buffer[1:length(cols),1:length(rows)]
                count += sum(data .> 0)
                total += sum(data)
            end
            @test total / count ≈ 76.33891347095299
            @test total / (AG.height(ds) * AG.width(ds)) ≈ 47.55674749653172
        end

        @testset "version 2" begin
            band = AG.getband(ds, 1)
            count = 0
            total = 0
            buffer = Array{AG.getdatatype(band)}(AG.getblocksize(band)..., 1)
            for (cols,rows) in AG.windows(band)
                AG.read!(ds, buffer, Cint[1], rows-1, cols-1)
                data = buffer[1:length(cols),1:length(rows)]
                count += sum(data .> 0)
                total += sum(data)
            end
            @test total / count ≈ 76.33891347095299
            @test total / (AG.height(ds) * AG.width(ds)) ≈ 47.55674749653172
        end

        @testset "version 3" begin
            band = AG.getband(ds, 1)
            count = 0
            total = 0
            buffer = Array{AG.getdatatype(band)}(AG.getblocksize(band)...)
            for (cols,rows) in AG.windows(band)
                AG.read!(ds, buffer, 1, rows, cols)
                data = buffer[1:length(cols),1:length(rows)]
                count += sum(data .> 0)
                total += sum(data)
            end
            @test total / count ≈ 76.33891347095299
            @test total / (AG.height(ds) * AG.width(ds)) ≈ 47.55674749653172
        end

        @testset "version 4" begin
            band = AG.getband(ds, 1)
            count = 0
            total = 0
            buffer = Array{AG.getdatatype(band)}(AG.getblocksize(band)...)
            for (cols,rows) in AG.windows(band)
                AG.read!(band, buffer, rows, cols)
                data = buffer[1:length(cols),1:length(rows)]
                count += sum(data .> 0)
                total += sum(data)
            end
            @test total / count ≈ 76.33891347095299
            @test total / (AG.height(ds) * AG.width(ds)) ≈ 47.55674749653172
        end

        @testset "version 5" begin
            band = AG.getband(ds, 1)
            count = 0
            total = 0
            xbsize, ybsize = AG.getblocksize(band)
            buffer = Array{AG.getdatatype(band)}(ybsize, xbsize)
            for ((i,j),(nrows,ncols)) in AG.blocks(band)
                # AG.rasterio!(ds,buffer,Cint[1],i,j,nrows,ncols)
                # AG.read!(band, buffer, j, i, ncols, nrows)
                AG.readblock!(band, j, i, buffer)
                data = buffer[1:nrows, 1:ncols]
                count += sum(data .> 0)
                total += sum(data)
            end
            @test total / count ≈ 76.33891347095299
            @test total / (AG.height(ds) * AG.width(ds)) ≈ 47.55674749653172
        end

        @testset "version 6" begin
            band = AG.getband(ds, 1)
            buffer = Array{AG.getdatatype(band)}(AG.width(ds), AG.height(ds), 1)
            AG.rasterio!(ds, buffer, Cint[1])
            count = sum(buffer .> 0)
            total = sum(buffer)
            @test total / count ≈ 76.33891347095299
            @test total / (AG.height(ds) * AG.width(ds)) ≈ 47.55674749653172
        end

        @testset "version 7" begin
            band = AG.getband(ds, 1)
            buffer = Array{AG.getdatatype(band)}(AG.width(ds), AG.height(ds))
            AG.read!(band, buffer)
            count = sum(buffer .> 0)
            total = sum(buffer)
            @test total / count ≈ 76.33891347095299
            @test total / (AG.height(ds) * AG.width(ds)) ≈ 47.55674749653172
        end

        @testset "version 8" begin
            band = AG.getband(ds, 1)
            buffer = Array{AG.getdatatype(band)}(AG.width(ds), AG.height(ds))
            AG.read!(ds, buffer, 1)
            count = sum(buffer .> 0)
            total = sum(buffer)
            @test total / count ≈ 76.33891347095299
            @test total / (AG.height(ds) * AG.width(ds)) ≈ 47.55674749653172
        end

        @testset "version 9" begin
            band = AG.getband(ds, 1)
            buffer = Array{AG.getdatatype(band)}(AG.width(ds), AG.height(ds), 1)
            AG.read!(ds, buffer, Cint[1])
            count = sum(buffer .> 0)
            total = sum(buffer)
            @test total / count ≈ 76.33891347095299
            @test total / (AG.height(ds) * AG.width(ds)) ≈ 47.55674749653172
        end

        @testset "version 10" begin
            band = AG.getband(ds, 1)
            buffer = Array{AG.getdatatype(band)}(AG.width(ds), AG.height(ds), 3)
            AG.read!(ds, buffer)
            count = sum(buffer[:,:,1] .> 0)
            total = sum(buffer[:,:,1])
            @test total / count ≈ 76.33891347095299
            @test total / (AG.height(ds) * AG.width(ds)) ≈ 47.55674749653172
        end
    end
end

# Untested
# writeblock!(rb::RasterBand, xoffset::Integer, yoffset::Integer, buffer)
# read!(rb::RasterBand, buffer::Array{Real,2}, xoffset::Integer, yoffset::Integer, xsize::Integer, ysize::Integer)
# read!(dataset::Dataset, buffer::Array{T,2}, i::Integer, xoffset::Integer, yoffset::Integer, xsize::Integer, ysize::Integer)
# read!(dataset::Dataset, buffer::Array{T,3}, indices::Vector{Cint}, xoffset::Integer, yoffset::Integer, xsize::Integer, ysize::Integer)

# read{U <: Integer}(rb::RasterBand, rows::UnitRange{U}, cols::UnitRange{U})
# read(dataset::Dataset, indices::Vector{Cint})
# read(dataset::Dataset)
# read{T <: Integer}(dataset::Dataset, indices::Vector{T}, xoffset::Integer, yoffset::Integer, xsize::Integer, ysize::Integer)
# read{U <: Integer}(dataset::Dataset, i::Integer, rows::UnitRange{U}, cols::UnitRange{U})
# read{U <: Integer}(dataset::Dataset, indices::Vector{Cint}, rows::UnitRange{U}, cols::UnitRange{U})update!{T <: Real}(rb::RasterBand, buffer::Array{T,2}) =

# write!(rb::RasterBand, buffer::Array{T,2}, rows::UnitRange{U}, cols::UnitRange{U})
# write!(dataset::Dataset, buffer::Array{T,2}, i::Integer)
# write!(dataset::Dataset, buffer::Array{T,3}, indices::Vector{Cint})
# write!(dataset::Dataset, buffer::Array{T,3}, indices::Vector{Cint}, xoffset::Integer, yoffset::Integer, xsize::Integer, ysize::Integer)
# write!(dataset::Dataset, buffer::Array{T,2}, i::Integer, rows::UnitRange{U}, cols::UnitRange{U})
# write!(dataset::Dataset, buffer::Array{T,3}, indices::Vector{Cint}, rows::UnitRange{U}, cols::UnitRange{U})

# function rasterio!(dataset::Dataset,
#                            buffer::Array{$T, 3},
#                            bands::Vector{Cint},
#                            xoffset::Integer,
#                            yoffset::Integer,
#                            xsize::Integer,
#                            ysize::Integer,
#                            access::GDALRWFlag=GF_Read,
#                            pxspace::Integer=0,
#                            linespace::Integer=0,
#                            bandspace::Integer=0,
#                            extraargs=Ptr{GDAL.GDALRasterIOExtraArg}(C_NULL))
#             (dataset == C_NULL) && error("Can't read invalid rasterband")
#             xbsize, ybsize, zbsize = size(buffer)
#             nband = length(bands); @assert nband == zbsize
#             result = ccall((:GDALDatasetRasterIOEx,GDAL.libgdal),GDAL.CPLErr,
#                            (Dataset,GDAL.GDALRWFlag,Cint,Cint,Cint,Cint,
#                             Ptr{Void},Cint,Cint,GDAL.GDALDataType,Cint,
#                             Ptr{Cint},GDAL.GSpacing,GDAL.GSpacing,GDAL.GSpacing,
#                             Ptr{GDAL.GDALRasterIOExtraArg}),dataset,access,
#                             xoffset,yoffset,xsize,ysize,pointer(buffer),xbsize,
#                             ybsize,$GT,nband,pointer(bands),pxspace,linespace,
#                             bandspace,extraargs)
#             @cplerr result "Access in DatasetRasterIO failed."
#             buffer
#         end


text:
function preset_a_best(codecs)
    list = (251, 250, 249, 140, 139)
    for c ∈ list
        c ∈ codecs && return c
    end
end

function preset_limited_best(codecs; res_limit=4320)
    resolutions = (4320,2160,1440,1080,720,480,360,240,144)
    for r ∈ resolutions
        r > res_limit && continue
        cs_r = CODECS_EACH_RES[r]
        i_end = lastindex(cs_r)
        for i ∈ 1:i_end
            c = cs_r[i]
            if c ∉ codecs
                continue
            elseif (r ≤ 480) || (name(c) != "av1")
                return c
            else
                for j ∈ i+1:i_end
                    c_next = cs_r[j]
                    if c_next ∉ codecs
                        continue
                    elseif fps(c_next) != "30"
                        return c
                    else
                        return c_next
                    end
                end
                return c
            end
        end
    end
end

ext: jl


text:
D = 5

open("81.jl.out") do f
  for line in readlines(f)
    words = filter(x-> !isempty(x), strip.(split(line)))
    nw = length(words)
    for idx in 1:length(words)
      randd = rand(1:D)
      d = min(D, max(idx-1,1), nw-idx)
      print(words[idx], "\t")
      for i in -d:d
        i == 0 && continue
        if idx+i < 1 || idx+i > nw
          print("N/A\t")
        else
          print(words[idx+i], "\t")
        end
      end
      println()
    end
  end
end


text:
####
#### Manipulating keys (tuples of Symbol)
####
#### These used as of type parameters, so care is taken to allow @pure functions.
####

"""
Type for keys, used internally.
"""
const Key = Symbol

const Keys = Tuple{Vararg{Symbol}}

"""
$(SIGNATURES)

Check if a `key` is in `itr`.
"""
Base.@pure function key_in(key::Key, keys::Keys)
    for k in keys
        key ≡ k && return true
    end
    false
end

Base.@pure function key_issubset(a::Keys, b::Keys)
    for k in a
        key_in(k, b) || return false
    end
    true
end

Base.@pure function key_setdiff(a::Keys, b::Keys)
    kept = Symbol[]
    for k in a
        key_in(k, b) || push!(kept, k)
    end
    (kept..., )
end


text:
module DataWriters
using FileIO
using CSVFiles
using FeatherFiles
using DataFrames
using HDF5
using JLD
using Feather
using Parquet
using Dates

export DataWriter, fit!, transform!
export Transformer
export datawriterrun

using TSML
using TSML.TSMLTypes
import TSML.TSMLTypes.fit! # to overload
import TSML.TSMLTypes.transform! # to overload
using TSML.Utils

using TSML.DataReaders
using TSML.DataReaders: FILEFMT, DATEFMT

mutable struct DataWriter <: Transformer
    model
    args

    function DataWriter(args=Dict())
        default_args=Dict(
            :filename => "",
            :dateformat => DATEFMT,
            :impl_args => Dict()
        )
        new(nothing,mergedict(default_args,args))
    end
end

function fit!(dtr::DataWriter,x::T=[],y::Vector=[]) where {T<:Union{DataFrame,Vector,Matrix}}
    fname = dtr.args[:filename]
    fmt = dtr.args[:dateformat]
    (fname != "" && fmt != "") || error("missing filename or date format")
    dtr.model = dtr.args
end

function transform!(dtr::DataWriter,x::T=[]) where {T<:Union{DataFrame,Vector,Matrix}}
    fullname = dtr.args[:filename]
    fmt = dtr.args[:dateformat]
    (fullname != "" && fmt != "") || error("missing filename or date format")
    fname = basename(fullname)
    fname != "" || error("filename is empty")
    fn,ext=split(fname,".")
    ext != "" || error("no filename extension format")
    ext in keys(FILEFMT)  || error("extension not recognized "*ext)
    # dispatch based on extension
    writefmt(FILEFMT[ext],fullname,x,fmt)
end

function writefmt(::T,fname::String,data::S,datefmt::String) where {T<:Union{Val{:csv},Val{:feather},Val{:parquet}},S<:Union{Matrix,DataFrame}}
    df = data |> DataFrame
    ncol(df) == 2 || error("dataframe should have only two columns: Date,Value")
    rename!(df,names(df)[1]=>:Date,names(df)[2]=>:Value)
    eltype(df[:Date]) <: DateTime || error("Date format error")
    df[:Date] = Dates.format.(df[:Date],datefmt)
    df |> save(fname)
end

function writefmt(atype::Union{Val{:hdf5},Val{:jld}},fname::String, data::T,datefmt::String) where {T<:Union{Matrix,DataFrame}}
    df = data |> DataFrame
    ncol(df) == 2 || error("dataframe should have only two columns: Date,Value")
    rename!(df,names(df)[1]=>:Date,names(df)[2]=>:Value)
    eltype(df[:Date]) <: DateTime || error("Date format error")
    df[:Date] = Dates.format.(df[:Date],datefmt)
    if atype == Val{:hdf5}
        fileopen = h5open
    else
        fileopen = jldopen
    end
    ldate = df[:Date]
    lvalue = df[:Value]
    fileopen(fname,"w") do file
        write(file,"dateval/date",ldate)
        write(file,"dateval/value",lvalue)
    end
end

function datawriterrun()
    fname = joinpath(dirname(pathof(TSML)),"../data/testdateval.csv")
    lcsv=DataReader(Dict(:filename=>fname))
    fit!(lcsv)
    dateval=transform!(lcsv)
    @show "csv"
    @show first(dateval,2)
    csvname = replace(fname,"test"=>"out")
    wcsv = DataWriter(Dict(:filename=>csvname))
    fit!(wcsv)
    transform!(wcsv,dateval)
    # check hdf5
    hdf5name = replace(fname,"csv"=>"h5")
    lhdf5 = DataWriter(Dict(:filename=>hdf5name))
    fit!(lhdf5)
    transform!(lhdf5,dateval)
    whdf5 = DataReader(Dict(:filename=>hdf5name))
    fit!(whdf5)
    res = transform!(whdf5)
    @show "hdf5"
    @show first(res,2)
    # check feather
    feathername = replace(fname,"csv"=>"feather")
    lfeather = DataWriter(Dict(:filename=>feathername))
    fit!(lfeather)
    transform!(lfeather,dateval)
    wfeather = DataReader(Dict(:filename=>feathername))
    fit!(wfeather)
    res = transform!(wfeather)
    @show "feather"
    @show first(res,2)
    # check jld
    jldname = replace(fname,"csv"=>"jld")
    ljld = DataWriter(Dict(:filename=>jldname))
    fit!(ljld)
    transform!(ljld,dateval)
    wjld = DataReader(Dict(:filename=>jldname))
    fit!(wjld)
    res = transform!(wjld)
    @show "jld"
    @show first(res,2)
    return nothing
end

end


text:
using Test, Mill, Flux
using Mill: p_map, inv_p_map, r_map, inv_r_map
import Mill: bagnorm

@testset "basic aggregation functionality" begin
    W = [1, 1/2, 1/2, 1/8, 1/3, 13/24]
    X = Matrix{Float64}(reshape(1:12, 2, 6))
    bags = BAGS[1]
    @test SegmentedMean(2)(X, bags) ≈ [1.0 4.0 9.0; 2.0 5.0 10.0]
    @test SegmentedSum(2)(X, bags) ≈ length.(bags)' .* SegmentedMean(2)(X, bags)
    @test SegmentedMax(2)(X, bags) ≈ [1.0 5.0 11.0; 2.0 6.0 12.0]
    @test SegmentedMeanMax(2)(X, bags) ≈ cat(SegmentedMean(2)(X, bags), SegmentedMax(2)(X, bags), dims=1)
    @test SegmentedMean(2)(X, bags, W) ≈ [1.0 4.0 236/24; 2.0 5.0 260/24]
    @test SegmentedSum(2)(X, bags, W) ≈ [bagnorm(W, b) for b in bags]' .* SegmentedMean(2)(X, bags, W)
    @test SegmentedMax(2)(X, bags, W) ≈ SegmentedMax(2)(X, bags)
end

@testset "matrix weights" begin
    W = abs.(randn(6))
    W_mat = vcat(W, 2*W)
    X = Matrix{Float64}(reshape(1:12, 2, 6))
    for bags in BAGS
        @test SegmentedSum(2)(X, bags, W_mat) ≈ SegmentedSum(2)(X, bags, W)
        @test SegmentedMean(2)(X, bags, W_mat) ≈ SegmentedMean(2)(X, bags, W)
        @test SegmentedMax(2)(X, bags, W_mat) ≈ SegmentedMax(2)(X, bags, W)
    end
end

@testset "pnorm functionality" begin
    dummy = randn(2)
    for t = 1:10
        a, b, c, d, ρ1, ρ2, c1, c2, w1, w2 = randn(10)
        p1, p2 = p_map(ρ1), p_map(ρ2)
        w1 = abs(w1)
        w2 = abs(w2)
        bags = ScatteredBags([[1,2]])
        agg = SegmentedPNorm([ρ1, ρ2], [c1, c2], dummy)
        @test agg([a b; c d], bags) ≈ [
                                     (1/2*(abs(a-c1)^p1 + abs(b-c1)^p1))^(1/p1);
                                     (1/2*(abs(c-c2)^p2 + abs(d-c2)^p2))^(1/p2)
                                    ]
        @test agg([a b; c d], bags) ≈ [
                                     (1/2*(abs(a-c1)^p1 + abs(b-c1)^p1))^(1/p1);
                                     (1/2*(abs(c-c2)^p2 + abs(d-c2)^p2))^(1/p2)
                                    ]
        @test agg([a b; c d], bags, [w1, w2]) ≈ [
                                               (1/(w1+w2)*(w1*abs(a-c1)^p1 + w2*abs(b-c1)^p1))^(1/p1);
                                               (1/(w1+w2)*(w1*abs(c-c2)^p2 + w2*abs(d-c2)^p2))^(1/p2)
                                              ]
        for bags in BAGS
            X = randn(2, 6)
            agg = SegmentedPNorm(inv_p_map.([1+1e-16, 1+1e-16]), [0, 0], dummy)
            @test agg(X, bags) ≈ SegmentedMean(2)(abs.(X), bags)
            agg = SegmentedPNorm(inv_p_map.([2, 2]), [0, 0], dummy)
            @test agg(X, bags) ≈ hcat([sqrt.(sum(X[:, b] .^ 2, dims=2) ./ length(b)) for b in bags]...)
        end
    end
end

@testset "lse functionality" begin
    dummy = randn(2)
    for t = 1:10
        a, b, c, d, ρ1, ρ2 = randn(6)
        r1, r2 = r_map(ρ1), r_map(ρ2)
        @test SegmentedLSE([ρ1, ρ2], dummy)([a b; c d], ScatteredBags([[1,2]])) ≈ [
                                                                               1/r1*log(1/2*(exp(a*r1)+exp(b*r1)));
                                                                               1/r2*log(1/2*(exp(c*r2)+exp(d*r2)))
                                                                              ]
        for bags in BAGS
            X = randn(2, 6)
            W = abs.(randn(6))
            r1, r2 = randn(2)
            # doesn't use weights
            @test all(SegmentedLSE([ρ1, ρ2], dummy)(X, bags) .== SegmentedLSE([ρ1, ρ2], dummy)(X, bags, W))
            # the bigger value of r, the closer we are to the real maximum
            @test isapprox(SegmentedLSE([100, 100], dummy)(X, bags), SegmentedMax(2)(X, bags), atol=0.1)
        end
    end
end

@testset "pnorm numerical stability" begin
    k, d = 10, 5
    dummy, c = randn(d), zeros(d)
    b = AlignedBags([1:k])
    ρ1 = inv_p_map.(ones(d))
    ρ2 = randn(d)
    p2 = p_map(ρ2)
    Z = 1e5 .+ 1e3 .* randn(d, k)
    W = abs.(randn(k)) .+ 1e-2
    for X in [Z, -Z, randn(d, k)]
        @test SegmentedPNorm(ρ1, c, dummy)(X, b) ≈ sum(abs.(X); dims=2) ./ k
        @test SegmentedPNorm(ρ1, c, dummy)(X, b, W) ≈ sum(W' .* abs.(X); dims=2) / sum(W)
        @test SegmentedPNorm(ρ1, c, dummy)(X, b, W) ≈ sum(W' .* abs.(X); dims=2) / sum(W)
        for i in 1:k
            @test SegmentedPNorm(ρ1, c, dummy)(repeat(X[:, i], 1, k), b) ≈ abs.(X[:, i])
            @test SegmentedPNorm(ρ2, c, dummy)(repeat(X[:, i], 1, k), b) ≈ abs.(X[:, i])
            @test SegmentedPNorm(ρ1, c, dummy)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
            @test SegmentedPNorm(ρ1, c, dummy)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
            @test SegmentedPNorm(ρ2, c, dummy)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
            @test SegmentedPNorm(ρ2, c, dummy)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
        end
    end
end

@testset "lse numerical stability" begin
    k, d = 10, 5
    dummy = randn(d)
    b = AlignedBags([1:k])
    ρ1 = inv_r_map.(1e15 .+ 1e5 .* randn(d))
    ρ2 = inv_r_map.(abs.(1e5 .* randn(d)))
    Z = 1e5 .+ 1e3 .* randn(d, k)
    W = abs.(randn(k)) .+ 1e-2
    for X in [Z, -Z, randn(d, k)]
        @test SegmentedLSE(ρ1, dummy)(X, b) ≈ maximum(X; dims=2)
        # doesn't use weights
        @test SegmentedLSE(ρ1, dummy)(X, b, W) ≈ maximum(X; dims=2)

        # implementation immune to underflow not available yet
        @test_skip @test SegmentedLSE(ρ2, dummy)(X, b) ≈ sum(X; dims=2) ./ k
        # doesn't use weights
        @test_skip @test SegmentedLSE(ρ2, dummy)(X, b, W) ≈ sum(X; dims=2) ./ k

        for i in 1:k
            @test SegmentedLSE(randn(d), dummy)(repeat(X[:, i], 1, k), b) ≈ X[:, i]
            # doesn't use weights
            @test SegmentedLSE(randn(d), dummy)(repeat(X[:, i], 1, k), b, W) ≈ X[:, i]
        end
    end
end

@testset "missing values" begin
    dummy = randn(2)
    C = randn(2)
    for bags in [AlignedBags([0:-1]), AlignedBags([0:-1, 0:-1, 0:-1])]
        @test SegmentedMean(C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedSum(C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedMax(C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedLSE(dummy, C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedPNorm(dummy, dummy, C)(missing, bags) == repeat(C, 1, length(bags))
    end

    # default values C are indeed filled in
    for bags in vcat(BAGS2)
        idcs = isempty.(bags.bags)
        l = maximum(maximum.(bags.bags[.!idcs]))
        X = randn(2, l)
        @test SegmentedMean(C)(X, bags)[:, idcs] == repeat(C, 1, sum(idcs))
        @test SegmentedSum(C)(X, bags)[:, idcs] == repeat(C, 1, sum(idcs))
        @test SegmentedMax(C)(X, bags)[:, idcs] == repeat(C, 1, sum(idcs))
        @test SegmentedLSE(dummy, C)(X, bags)[:, idcs] == repeat(C, 1, sum(idcs))
        @test SegmentedPNorm(dummy, dummy, C)(X, bags)[:, idcs] == repeat(C, 1, sum(idcs))
    end
end


text:
# runtests.jl

include("ndgrid-oneto.jl")
include("ndgrid-unitr.jl")
include("ndgrid-range.jl")
include("ndgrid-avect.jl")
include("timer.jl")


text:
# AEdgeMap and AVertexMap defined her for include precedence

"""
    abstract type AEdgeMap{T} end

Type representing an abstract edge map with value type `T`.
"""
abstract type AEdgeMap{T} end

"""
    abstract type AVertexMap{T} end

Type representing an abstract vertex map with value type `T`.
"""
abstract type AVertexMap{T} end

"""
    abstract type AIndexedEdge <: AEdge end

Edge types with unique indexes, accessed by [`idx`](@ref)
"""
abstract type AIndexedEdge <: AEdge end

"""
    idx(e::AIndexedEdge)

Returns the index of edge `e`.
"""
idx(e::AIndexedEdge) = error("Method not defined")

"""
    abstract type ANetwork <: AGraph end

An abstract graph with the additional possibility to attach properties to vertices and edges.
"""
abstract type ANetwork <: AGraph end

"""
    abstract type ADiNetwork <: ADiGraph end

An abstract directed graph with the additional possibility to attach properties to vertices and edges.
"""
abstract type ADiNetwork <: ADiGraph end

const ANetOrDiNet = Union{ANetwork, ADiNetwork}

### GRAPH
"""
    set_graph_property!(g, name, x)

Set the property `name` to value `x` to `g`. Creates the property if it doesn't exist.
[`gprop!`](@ref) can be conveniently used as a short form of this function.

**Example**
```julia
g = Network(10, 20)
set_graph_property!(g, "label", "My Network")
# or equivalently
gprop!(g, "label", "My Network")
```
"""
set_graph_property!(g::ANetOrDiNet, name::AbstractString, x) = set_graph_property!(g.props, name, x)

"""See [`set_graph_property!`](@ref)"""
gprop! = set_graph_property!

"""
    rem_graph_property!(g, name)

Remove the property `name` from `g`.

[`rem_gprop!`](@ref) is the short form of this function.
"""
rem_graph_property!(g::ANetOrDiNet, name::AbstractString) = rem_graph_property!(g.props, name)

"""See [`rem_graph_property!`](@ref)"""
rem_gprop! = rem_graph_property!

"""
    graph_property(g, name)

Return the property `name` of `g`.

    graph_property(g)

Returns a dictionary with elements `property_name => property_value`

[`gprop`](@ref) is the short form of this function.
"""
graph_property(g::ANetOrDiNet, name::AbstractString) = graph_property(g.props, name)
graph_property(g::ANetOrDiNet) = graph_property(g.props)

"""See [`graph_property`](@ref)"""
gprop = graph_property

"""
    has_graph_property(g, name)

Check if network  `g` has a graph property named `name`.

[`has_gprop`](@ref) is the short form of this function.
"""
has_graph_property(g::ANetOrDiNet, name::AbstractString) = has_graph_property(g.props, name)

"""See [`has_graph_property`](@ref)"""
has_gprop = has_graph_property

### EDGE
"""
    add_edge_property!(g, name, T)
    add_edge_property!(g, name, emap)

Add the edge property  `name` to `g`.

If a type `T` is given as an input, an edge map with valtype `T` is created and
stored into `g`.

As an alternative, an existing edge map `emap` can be stored into `g`.

[`eprop!`](@ref) is the short form of this function.

**Example**
```julia
g = random_regular_graph(10, 3, Network)

add_edge_property!(g, "weight", Float64)
# or equivalently
eprop!(g, "weight", Float64)
```
"""
add_edge_property!(g::ANetOrDiNet, name::AbstractString, ::Type{T}) where {T} = add_edge_property!(g.props, name, EdgeMap(g, T))
add_edge_property!(g::ANetOrDiNet, name::AbstractString, emap::AEdgeMap) = add_edge_property!(g.props, name, emap)
add_edge_property!(g::ANetOrDiNet, name::AbstractString, data) = add_edge_property!(g, name, EdgeMap(g, data))

"""See [`add_edge_property!`](@ref)"""
eprop! = add_edge_property!

"""
    rem_edge_property!(g, name)

Remove the edge property  `name` from `g`.

[`rem_eprop!`](@ref) is the short form of this function.
"""
rem_edge_property!(g::ANetOrDiNet, name::AbstractString) = rem_edge_property!(g.props, name)

"""See [`rem_edge_property!`](@ref)"""
rem_eprop! = rem_edge_property!

"""
    edge_property(g, name)

Return an edge map corresponding to property `name` of edges in `g`.

    edge_property(g)

Returns a dictionary with elements `property_name => edge_map`.

    edge_property(g, e)

Returns a dictionary of the form `name => val` containing all the properties
associated to edge `e`.

[`eprop`](@ref) is the short form of this function.
"""
edge_property(g::ANetOrDiNet, name::AbstractString) = edge_property(g.props, name)
edge_property(g::ANetOrDiNet) = edge_property(g.props)
edge_property(g::ANetOrDiNet, e::AEdge) = edge_property(g.props, e)

"""See [`edge_property`](@ref)"""
eprop = edge_property

"""
    has_edge_property(g, name)
    has_edge_property(g, name, e)

Check if network  `g` has an edge property named `name`.
The second method checks also if edge `e` has an assigned value for
that property.

[`has_eprop`](@ref) is the short form of this function.
"""
has_edge_property(g::ANetOrDiNet, name::AbstractString) = has_edge_property(g.props, name)
has_edge_property(g::ANetOrDiNet, name::AbstractString, e::AEdge) = has_edge_property(g.props, name, e)

"""See [`has_edge_property`](@ref)"""
has_eprop = has_edge_property

####################

"""
    add_vertex_property!(g, name, T)
    add_vertex_property!(g, name, vmap)

Add the vertex property  `name` to `g`.

If a type `T` is given as an input, a vertex map with valtype `T` is created and
stored into `g`.

As an alternative, an existing vertex map `vmap` can be stored into `g`.

[`vprop!`](@ref) is the short form of this function.
"""
add_vertex_property!(g::ANetOrDiNet, name::AbstractString, ::Type{T}) where {T} = add_vertex_property!(g.props, name, VertexMap(g, T))
add_vertex_property!(g::ANetOrDiNet, name::AbstractString, vmap::AVertexMap) = add_vertex_property!(g.props, name, vmap)
add_vertex_property!(g::ANetOrDiNet, name::AbstractString, data) = add_vertex_property!(g, name, VertexMap(g, data))

"""See [`add_vertex_property!`](@ref)"""
vprop! = add_vertex_property!

"""
    rem_vertex_property!(g, name)

Remove the vertex property  `name` from `g`.

[`rem_vprop!`](@ref) is the short form of this function.
"""
rem_vertex_property!(g::ANetOrDiNet, name::AbstractString) = rem_vertex_property!(g.props, name)

"""See [`rem_vertex_property!`](@ref)"""
rem_vprop! = rem_vertex_property!

"""
    vertex_property(g, name)

Return an vertex map corresponding to property `name` of vertices in `g`.

    vertex_property(g)

Returns a dictionary with elements `property_name => vertex_map`.

    vertex_property(g, v)

Returns a dictionary of the form `name => val` containing all the properties
associated to vertex `v`.
"""
vertex_property(g::ANetOrDiNet, name::AbstractString) = vertex_property(g.props, name)
vertex_property(g::ANetOrDiNet) = vertex_property(g.props)
vertex_property(g::ANetOrDiNet, v::Integer) = vertex_property(g.props, v)

"""See [`vertex_property`](@ref)"""
vprop = vertex_property

"""
    has_vertex_property(g, name, v)

Check if network  `g` has a vertex property named `name`.
The second method checks also if vertex `v` has an assigned value for
that property.

[`has_vprop`](@ref) is the short form of this function.
"""
has_vertex_property(g::ANetOrDiNet, name::AbstractString) = has_vertex_property(g.props, name)
has_vertex_property(g::ANetOrDiNet, name::AbstractString, v::Integer) =
                                    has_vertex_property(g.props, name, v)

"""See [`has_vertex_property`](@ref)"""
has_vprop = has_vertex_property

# TODO export
gprop_names(g::ANetOrDiNet) = collect(keys(gprop(g)))
vprop_names(g::ANetOrDiNet) = collect(keys(vprop(g)))
eprop_names(g::ANetOrDiNet) = collect(keys(eprop(g)))


text:
# FIXME see whether we can remove the broadcast calls required for 0.7

println("\n\n\nStarting runtests.jl $(join(ARGS, " ")) ...")

using FunctionalData, FunctionalDataUtils, Test #, Colors
using Statistics

macro shouldtestset(a,b) length(ARGS) < 1 || ARGS[1] == a ?  :(@testset $a $b) : nothing end
macro shouldtestset2(a,b) length(ARGS) < 2 || ARGS[2] == a ?  :(@testset $a $b) : nothing end

@shouldtestset "machinelearning" begin
    @shouldtestset2 "randbase" begin
        # @test_equal size(randbase(2,10)) (2,10)
        # @test_almostequal sqrt(sum(randbase(2,10).^2, 2)) [1. 1.]'
        # @test_almostequal (@p randbase 2 10  | (.^) _ 2 |  sum _ 2) [1. 1.]'
    end
    @shouldtestset2 "randproj" begin
    end
    @shouldtestset2 "loocv" begin
    end
end

@shouldtestset "computing" begin
    @shouldtestset2 "fasthash" begin
        @test fasthash(1) == "20eae029a26a15420f9ed6d7a9999e823f251b9d44adea5a504e1d8a4f7c5512"
        @test fasthash(1.) == "cb5415fc791f65b70b303678e3601d60ac527cc4cdd3a10d1f85b7423b270112"
        @test fasthash('a') == "94923c7a4899a5baf82f29551f5e4652e3b7ee50699bbd0b4ddba1538e5963b2"
        @test fasthash("a") == "ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb"
        @test fasthash(["a"]) == "da3811154d59c4267077ddd8bb768fa9b06399c486e1fc00485116b57c9872f5"
        @test fasthash([1]) == "7c9fa136d4413fa6173637e883b6998d32e1d675f88cddff9dcbcf331820f4b8"
        @test fasthash(UInt8[1]) == "4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459a"
        @test fasthash(Int8[1]) == "4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459a"
        @test fasthash(UInt16[1]) == "47dc540c94ceb704a23875c11273e16bb0b8a87aed84de911f2133568115f254"
        @test fasthash(Int16[1]) == "47dc540c94ceb704a23875c11273e16bb0b8a87aed84de911f2133568115f254"
        @test fasthash(UInt32[1]) == "67abdd721024f0ff4e0b3f4c2fc13bc5bad42d0b7851d456d88d203d15aaa450"
        @test fasthash(Int32[1]) == "67abdd721024f0ff4e0b3f4c2fc13bc5bad42d0b7851d456d88d203d15aaa450"
        @test fasthash(UInt64[1]) == "7c9fa136d4413fa6173637e883b6998d32e1d675f88cddff9dcbcf331820f4b8"
        @test fasthash(Int64[1]) == "7c9fa136d4413fa6173637e883b6998d32e1d675f88cddff9dcbcf331820f4b8"
        @test fasthash(Any[1,"test"]) == "b59f20a7e3513ef702971ec5cbd288f19d760df46d49d5d27edb82fe73fc257b"
        data = Any[1,1., 'a', "a", rand(10,200), zeros(1000,10000), Dict("a" => 1, 2 => "b")]
        r = [fasthash(x) for x in data]
        @test fasthash(1:2) == "48adce367d9daa9d76033a4693b20d7fab23d357bf97e49a036f2dd32b1afff4"
        @test fasthash(1:3:10) == "dc29cadf1701b9023b1170a8101b40928770dbbd967e2740afa4986499522ca3"
    end
end

@shouldtestset "computervision" begin
    @shouldtestset2 "warp" begin
        img = rand(10,20,3)
        a = rand(3:9,2,4)
        b = copy(a)
        b[:,end] = 10
        warp(img,a,b)
        @test size(warp(img,a,b))  ==  (10,20,3)
        @test size(warp(img,a,b,(8,9)))  ==  (8,9,3)

    end
    @shouldtestset2 "iimg" begin
        a_ = [1 2 3]
        @test size(iimg(a_)) == (1,3)
        @test typeof(iimg(a_)) == Array{Float64,2}
        @test iimg(a_) == [1 3 6]
        @test iimg([1 2 3; 4 5 6]) == [1.0 3.0 6.0; 5.0 12.0 21.0]
        @test iimg(ones(1,2,3)) == float(cat([1 2],[2 4],[3 6],dims=3))
    end

    @shouldtestset2 "interp3" begin
        a = [o for m in 1:2, n in 1:2, o in 1:2]
        @test interp3(a,1.,1.,1.5) == 1.5
        @test interp3(a,2.,2.,1.5) == 1.5
        a = [n for m in 1:2, n in 1:2, o in 1:2]
        @test interp3(a,1.,1.5,1.) == 1.5
        @test interp3(a,2.,1.5,2.) == 1.5
        a = [m for m in 1:2, n in 1:2, o in 1:2]
        @test interp3(a,1.5,1.,1.) == 1.5
        @test interp3(a,1.5,2.,2.) == 1.5
    end

    @shouldtestset2 "interp3with01coords" begin
        a = [o for m in 1:2, n in 1:2, o in 1:2]
        @test interp3with01coords(a,0.,0.,0.5) == 1.5
        @test interp3with01coords(a,1.,1.,0.5) == 1.5
        a = [n for m in 1:2, n in 1:2, o in 1:2]
        @test interp3with01coords(a,0.,0.5,0.) == 1.5
        @test interp3with01coords(a,1.,0.5,1.) == 1.5
        a = [m for m in 1:2, n in 1:2, o in 1:2]
        @test interp3with01coords(a,0.5,0.,0.) == 1.5
        @test interp3with01coords(a,0.5,1.,1.) == 1.5
    end

    @shouldtestset2 "resize" begin
        @test resize(ones(Int,1,1), [1, 1]) == ones(Int, 1, 1)
        @test resize(ones(Float64,1,1), 1) == ones(Float64, 1, 1)
        @test resize(ones(Float64,1,1,1), [1, 1, 1]) == ones(Float64, 1, 1, 1)
        @test resize(ones(Int,1,1), [1, 10]) == ones(Int, 1, 10)
        @test maximum(abs.(resize([1. 2.], [1, 4])-[1. 1.33333 1.66666 2.])) < 0.01
        @test resize([1. 3.], [1, 3]) == [1. 2. 3.]
    end

    @shouldtestset2 "bwlabel" begin
        img = [1 2 3]
        @test bwlabel(img) == img

        img = [1 2 3 0 5 6]
        @test bwlabel(img) == [1 2 3 1 5 6]
        @test bwlabel(img,4) == [1 2 3 4 5 6]

        img = [1 1 0 5 2 2; 3 3 0 0 4 4]
        bwlabel(img)
    end

    @shouldtestset2 "border" begin
        img = zeros(5,7)
        img[3,4] = 1
        @test border(img)  ==  img
        img[2:4,3:6] = 1
        b = copy(img)
        b[3,4:5] = 0
        @test border(img)  ==  b
    end

    @shouldtestset2 "bwdist" begin
        a = [0 0 0 1 0 0]
        @test bwdist(a)  ==  [3 2 1 0 1 2]
        @test bwdist([1 10;2 12], [1 10; 1 10])  ==  [1,2]
    end

    @shouldtestset2 "rle" begin
        data = [1 1 2 2 2 3 3 3 3]
        @test unrle(rle(data)) == data
        data = rand(1:5, 1, 1000)
        @test unrle(rle(data)) == data
        for t = [Int16, Int32, Int64, Float32, Float64]
            for s1 = 1:10, s2 = 1:10
                data = rand(t, s1, s2)
                @test unrle(rle(data)) == data
            end
        end
    end

    @shouldtestset2 "resizeminmax" begin
        f(a) = size(resizeminmax(a, [30,40], [130, 140]))
        @test f(ones(1,1)) == (30,40)
        @test f(ones(100,200)) == (85,140)
        @test f(ones(10,200)) == (30,140)
        @test f(ones(50,10)) == (125,40)
        @test f(ones(50,60)) == (50,60)
        @test f(ones(1,600)) == (30,140)
        @test f(ones(600,1)) == (130,40)
    end

    @shouldtestset2 "blocks" begin
        a = ones(16,32)
        ind = stridedblocksub(a, 8; keepshape = true)
        @test size(ind) == (2,4)
        ind = stridedblocksub(a, 8)
        @test len(ind) == 8
        ind =  stridedblocksub(a, 8, 1)
        @test len(ind) == 9*25
        a = flatten(Any[ones(2,2) 2*ones(2,2); 3*ones(2,2) 4*ones(2,2)])
        subs = stridedblocksub(a, 2)
    end

    @shouldtestset2 "meshgrid" begin
        @test meshgrid(2,3)  ==  [1 2 1 2 1 2; 1 1 2 2 3 3]
        @test meshgrid(1:2,1:3)  ==  [1 2 1 2 1 2; 1 1 2 2 3 3]
        @test meshgrid(2:3,2:4)  ==  [2 3 2 3 2 3; 2 2 3 3 4 4]
        @test meshgrid(2:3,[2, 0, 4])  ==  [2 3 2 3 2 3; 2 2 0 0 4 4]
        @test centeredmeshgrid(2:3,2:4)  ==  [0 1 0 1 0 1; -1 -1 0 0 1 1]
    end

    # FIXME reenable once MultivariateStats works on 0.7
    # @shouldtestset2 "inpointcloud" begin
    #     l = row(linspace(0,1,100))
    #     o = row(ones(100))
    #     shape2d = [l o flipdim(l,2) 0*o; 0*o l o flipdim(l,2)]*10 .+ [2.5;5.5]
    #     shouldbe2d = zeros(30,30)
    #     shouldbe2d[3:12,6:15] = 1
    #     coords = meshgrid(30,30)
    #     result2d = reshape(map(coords, x->inpointcloud(x, shape2d)), 30,30);
    #     @test result2d  ==  shouldbe2d

    #     plate1 = [meshgrid(1:10,1:10); zeros(1,100)]
    #     plate2 = [meshgrid(1:10,1:10); 10*ones(1,100)]
    #     plate3 = plate1[[1, 3, 2],:]
    #     plate4 = plate2[[1, 3, 2],:]
    #     plate5 = plate1[[3, 1, 2],:]
    #     plate6 = plate2[[3, 1, 2],:]
    #     shape3d = [plate1 plate2 plate3 plate4 plate5 plate6] .+ [2.5;5.5;5.5]

    #     shouldbe3d = zeros(30,30,30)
    #     shouldbe3d[3:12,6:15,6:15] = 1
    #     coords = meshgrid(30,30,30)
    #     result3d = reshape(map(coords, x->inpointcloud(x, shape3d)), 30,30,30);
    #     result3d == shouldbe3d
    #     @test sum(result3d-shouldbe3d)  ==  38
    # end

    @shouldtestset2 "sampler" begin
        image = rand(200,100)
        sampler = makeSampler(image, 32)
        coords = asint(broadcast(+, 50, 50*randn(2,100)))
        @test size(@p map coords sampler) == (32,32,100)
        out = zeros(1024,1)
        pos = col([50,50])
        @test size(sample!(out, pos, sampler)) == (1024,1)

        sampler = makeSampler(image, 32, col = true)
        @test size(sample(pos, sampler)) == (1024,1)

        sampler = makeSampler(image, 32, col = true, normmeanstd = true)
        r = sample(pos, sampler)
        @test abs(mean(r))  < 1e-5
        @test abs(std(r)-1) < 1e-5

        image = [1 2 3 4 5; 10 20 30 40 50; 11 22 33 44 55]
        sampler = makeSampler(image, 3, centered = false)
        @test sample(col([1,2]), sampler) == [2 3 4; 20 30 40; 22 33 44]
        sampler = makeSampler(image, 3, centered = true)
        @test sample(col([1,2]), sampler) == [1 2 3; 10 20 30; 11 22 33]
        @test sample(col([1,3]), sampler) == [2 3 4; 20 30 40; 22 33 44]

        image = ones(20,10)
        sampler = makeSampler(image, 3, centered = false)
        sample(sampler, col([1 1]))
        sample(sampler,siz(image))

        sampler = makeSampler(image, 3, centered = true)
        sample(sampler, col([1 1]))
        sample(sampler,siz(image))
    end
end

# @shouldtestset "graphics" begin
#     @shouldtestset2 "jetcolormap" begin
#         r = jetcolormap(10)
#         @test all(broadcast(>=, r, 0)) == true
#         @test all(broadcast(<=, r, 1)) == true
#     end

#     @shouldtestset2 "jetcolors" begin
#         @test jetcolors(1:3) ≈ [0.0 0.5198019801980198 0.5396039603960396; 0.0 1.0 0.0; 0.5396039603960396 0.4801980198019802 0.0]
#         @test jetcolors(1:3,2,4) ≈ [0.0 0.0 0.5198019801980198; 0.0 0.0 1.0; 0.5396039603960396 0.5396039603960396 0.4801980198019802]
#         # @test jetcolorants(1:3) == Any[RGB{Float64}(0.0,0.0,0.5396039603960396),RGB{Float64}(0.5198019801980198,1.0,0.4801980198019802),RGB{Float64}(0.5396039603960396,0.0,0.0)] # disabled until Colors works on 0.7
#     end

#     @shouldtestset2 "asimagesc" begin
#         img = [1 2 3 4 5 6 7 8 9]
#         r = asimagescrgb(img)
#         @test size(r,3) == 3
#     end

#     @shouldtestset2 "blocksvisu" begin
#         @test blocksvisu([1 2 3]) == [1 0 3; 0 0 0; 2 0 0]
#         @test blocksvisu([1 2 3; 1 2 3; 1 2 3; 1 2 3]) == [1 1 0 3 3; 1 1 0 3 3; 0 0 0 0 0; 2 2 0 0 0; 2 2 0 0 0]
#     end

#     @shouldtestset2 "overlaygradient" begin
#         img = rand(100,200)
#         sp = zeros(size(img))
#         sp[25:75, 50:150] = 1
#         @test size(overlaygradient(img, sp)) == (100,200,3)
#     end
# end

@shouldtestset "numerical" begin
    @shouldtestset2 "nanfunctions" begin
        @test nanmedian([1 2 3], 1)  ==  [1.0 2.0 3.0]
        @test isequal(nanmedian([1 2 NaN], 1),  [1.0 2.0 NaN]) == true
        @test nanmedian([1 2 3], 2)  ==  row([2.0])
        @test nanmedian([1 2 NaN], 2)  ==  row([1.5])
        @test nanmedian([1, 2, NaN])  ==  1.5
    end
    @shouldtestset2 "distance" begin
        @test distance(1,1)  ==  0
        @test distance(1.,1.)  ==  0.
        @test distance([1,2,3],[2,2,3])  ==  1
        @test distance([0 1; 0 1],[0 1; 0 1])  ==  [0. sqrt(2); sqrt(2) 0.]
    end
    @shouldtestset2 "norms" begin
        @test normsum([1 2 1])  ==  [1/4 1/2 1/4]
        @test norm01([1 2 3])  ==  [0 1/2 1]
        @test normeuclid([1 1])  ==  broadcast(/, 1, [sqrt(2) sqrt(2)])
        @test normmeanstd([1 2 3])  ==  [-1 0 1]
        @test normquantile(collect(0:10))  ≈  asfloat64(collect(0:10))/10
        @test normquantile(collect(0:2:20))  ≈  asfloat64(collect(0:10))/10
    end
    @shouldtestset2 "extrema" begin
        @test (@p maximum [1 2 3; 4 5 0] (x->x[2]))  ==  [2 5]'
        @test (@p minimum [1 2 3; 4 5 0] (x->x[2]))  ==  [3 0]'
    end
    @shouldtestset2 "valuemap" begin
        @test valuemap([0 1 2 3], [10 20 30])  ==  [0. 10. 20. 30.]
    end
    @shouldtestset2 "clamp" begin
        @test clamp(1,1,1)  ==  1
        @test clamp(0,1,2)  ==  1
        @test clamp(2,1,2)  ==  2
        @test clamp([-1 2 3 10; -5 10 11  100],[1 10],[3 15])  ==  [1 2 3 3; 10 10 11 15]
        @test clamp(Float32[-1 2 3 10; -5 10 11  100],[1 10],[3 15])  ==  [1 2 3 3; 10 10 11 15]
        @test clamp([1,2],[3,4],[5,6])   ==  [3,4]
        @test clamp([0 1 11; -5 3 20], rand(10,15))  ==  [1 1 10; 1 3 15]
        @test clamp([0 1 11; -5 3 20; -4 3 5], rand(10,15,4))  ==  [1 1 10; 1 3 15; 1 3 4]
    end
    @shouldtestset2 "rand" begin
        @test size(randsiz(siz(rand(2,3)))) == (2,3)
        @test eltype(randsiz([2 3]', Float32)) == Float32
    end
    @shouldtestset2 "plus" begin
        @test (@p plus zeros(2,3) ones(2,1))  ==  ones(2,3)
        @test (@p minus zeros(2,3) ones(2,1))  ==  -ones(2,3)
        @test (@p times ones(2,3) zeros(2,1))  ==  zeros(2,3)
        @test (@p divby ones(2,3) 2*ones(2,1))  ==  0.5*ones(2,3)
    end
end

@shouldtestset "utils" begin
    @shouldtestset2 "histdict" begin
        d = [Dict(:a => 1), Dict(:a => 1), Dict(:a => 10)]
        r = Dict(1 => 2, 10 => 1)
        @test histdict(d, :a)  ==  r
        @test histdict([1,1,10])  ==  r
    end
end

println("done!\n")


text:
# MHD equations
# ∇⋅u = 0 in fluid
# u⋅∇(u) -ν*Δ(u) + (1/ρ)*∇(p) - (1/ρ)*(j×B) = (1/ρ)*f in fluid
# ∇⋅j = 0 in fluid and solid
# j + σ*∇(φ) - σ*(u×B) = 0 in fluid
# j + σ*∇(φ) = 0 in solid
# solving for u,p,j,φ for a given B,ν,ρ,σ
#
# One can provide characteristic quantities
# u0,B0,L
#
# To introduce these change of variables
# u = u0*ū, B = B0*B̄, j = σ*u0*B0*j̄, φ = u0*B0*L*φ̄ (I am not sure about L in φ)
# p = ρ*u0^2*p̄, f = (ρ*u0^2/L)*f̄ (option 1)
# p = σ*u0*B0*L*p̄, f = (σ*u0*B0^2)*f̄ (option 2)
#
# To solve the following equations in a scaled domain according to L
# ∇⋅ū = 0 in fluid
# ∇⋅j̄ = 0 in fluid and solid
# j̄ + ∇(φ̄) - ū×B̄ = 0 in fluid
# j̄ + ∇(φ̄) = 0 in solid
# ū⋅∇(ū) -(1/Re)*Δ(ū) + ∇(p̄) - N*(j̄×B̄) = f̄ in fluid (Option 1, CFD)
# (1/N)*ū⋅∇(ū) - (1/Ha^2)*Δ(ū) + ∇(p̄) - (j̄×B̄) = f̄ in fluid (Option 2,MHD)
#
# with
#  Re = u0*L/ν
#  Ha = B0*L*sqrt(σ/(ρ*ν))
#  N = Ha^2/Re

# In order to account for both options, the code solves these equations
#
# ∇⋅ū = 0 in fluid
# ∇⋅j̄ = 0 in fluid and solid
# j̄ + σ̄*∇(φ̄) - σ̄*(ū×B̄) = 0 in fluid
# j̄ + σ̄*∇(φ̄) = 0 in solid
# α*ū⋅∇(ū) - β*Δ(ū) + ∇(p̄) - γ*(j̄×B̄) = f̄ in fluid
#
# α = 1, β = (1/Re), γ = N (option 1,CFD)
# α = (1/N), β = (1/Ha^2), γ = 1 (option 2,MHD)
# σ̄ is 1, but is included just in case you want to solve
# the un-scaled equations directly
#
# Some boundary conditions
# (for simplicity we drop the bars, i.e. ū is simply u)
#
# Velocity bc
# u = u (imposed strongly)
#
# Charge bc
# j = j (imposed strongly in the normal direction to the boundary)
#
# Traction bC
# # n⋅∇(u) - p*n = t
#
# Potential bc
# φ = φ (Imposed weakly)
#
# Thin wall bc
# j⋅n + cw*n⋅∇(j)⋅n = jw (imposed weakly via a penalty of value τ)

"""
    add_default_params(params::Dict{Symbol}) -> full_params

Return a new `Dict` with the contents of `params` plus
default values for the optional keys not provided in `params`.
It also checks the validity of the main parameter dictionary `params`.
"""
function add_default_params(_params)
  mandatory = Dict(
    :ptimer=>false,
    :debug=>false,
    :solver=>false,
    :matrix_type=>false,
    :vector_type=>false,
    :model=>true,
    :k=>false,
    :fluid=>true,
    :solid=>false,
    :bcs=>true,
    :check_valid=>false,
    :solver_postpro=>false,
  )
  _check_mandatory(_params,mandatory,"")
  optional = Dict(
    :solid=>nothing,
    :ptimer=>default_ptimer(_params[:model]),
    :debug=>false,
    :solver=>NLSolver(show_trace=true,method=:newton),
    :solver_postpro => (x->nothing),
    :matrix_type=>SparseMatrixCSC{Float64,Int},
    :vector_type=>Vector{Float64},
    :k=>2,
    :check_valid=>true,
  )
  params = _add_optional(_params,mandatory,optional,_params,"")
  _check_unused(params,mandatory,params,"")
  # Process sub-params
  params[:fluid] = params_fluid(params)
  if params[:solid] !== nothing
    params[:solid] = params_solid(params)
  end
  params[:bcs] = params_bcs(params)
  params
end

default_ptimer(model) = PTimer(get_part_ids(sequential,1))
default_ptimer(model::GridapDistributed.DistributedDiscreteModel) = PTimer(get_part_ids(model.models))

"""
Valid keys for `params[:fluid]` are the following.

# Mandatory keys
- `domain`: Domain where to solve the MHD problem.
  The domain is represented with a `Triangulation` or with a `Integer`/`String`
  tag in the underlying discrete model.
  It can also be a `Gridap.DiscreteModel` or `GridapDistributed.DistributedDiscreteModel`
  if `params[:solid]` is `nothing`. In this case,
  one needs to guarantee `params[:model]===params[:fluid][:domain]`.
-  `:α`: Value of the parameter `α`.
-  `:β`: Value of the parameter `β`.
-  `:γ`: Value of the parameter `γ`.
-  `:B`: Value of the parameter `B`.

# Optional keys
-  `:f=>VectorValue(0,0,0)`: Value of the parameter `f`.
-  `:σ=>1`: Value of the parameter `σ`.
"""
function params_fluid(params::Dict{Symbol,Any})
  mandatory = Dict(
   :domain=>true,
   :α=>true,
   :β=>true,
   :γ=>true,
   :B=>true,
   :f=>false,
   :σ=>false,
  )
  optional = Dict(:σ=>1,:f=>VectorValue(0,0,0))
  fluid = _check_mandatory_and_add_optional(params[:fluid],mandatory,optional,params,"[:fluid]")
  fluid
end

"""
Valid keys for `params[:solid]` are the following

# Mandatory keys
- `domain`: Domain occupied by the solid.
  The domain is represented either with a `Triangulation` or with a `Integer`/`String`
  tag in the underlying discrete model.
# Optional keys
-  `:σ=>1`: Value of the parameter `σ`.
"""
function params_solid(params::Dict{Symbol,Any})
  mandatory = Dict(
   :domain=>true,
   :σ=>false,
  )
  optional = Dict(:σ=>1)
  solid = _check_mandatory_and_add_optional(params[:solid],mandatory,optional,params,"[:solid]")
  solid
end

"""
Valid keys for `params[:bcs]` are the following

# Mandatory keys
- `:u`: A `Dict` defining strong Dirichlet conditions on the fluid velocity.
   See  [`params_bcs_u`](@ref) for further details.
- `:j`: A `Dict` defining strong Dirichlet conditions on the charge current.
   See  [`params_bcs_j`](@ref) for further details.

# Optional keys
- `:φ=>[]`: A `Dict` or a vector of `Dict`s.
   Each `Dict` defines a weak boundary condition for the potential.
   See  [`params_bcs_φ`](@ref) for further details.
- `:t=>[]`: A `Dict` or a vector of `Dict`s.
  Each `Dict` defines a boundary traction on the fluid.
   See  [`params_bcs_t`](@ref) for further details.
- `:thin_wall=>[]`: A `Dict` or a vector of `Dict`s.
  Each Dict defines a thin wall law.
   See  [`params_bcs_thin_wall`](@ref) for further details.
"""
function params_bcs(params)
  mandatory = Dict(
   :u=>true,
   :j=>true,
   :φ=>false,
   :t=>false,
   :thin_wall=>false,
   :f => false,
   :B => false,
  )
  optional = Dict(
   :φ=>[],
   :t=>[],
   :thin_wall=>[],
   :f =>[],
   :B =>[],
  )
  bcs = _check_mandatory_and_add_optional(params[:bcs],mandatory,optional,params,"[:bcs]")
  # Sub params
  bcs[:u] = params_bcs_u(params)
  bcs[:j] = params_bcs_j(params)
  if bcs[:φ] !== optional[:φ]
    bcs[:φ] = params_bcs_φ(params)
  end
  if bcs[:t] !== optional[:t]
    bcs[:t] = params_bcs_t(params)
  end
  if bcs[:thin_wall] !== optional[:thin_wall]
    bcs[:thin_wall] = params_bcs_thin_wall(params)
  end
  bcs
end

"""
Valid keys for `params[:bcs][:u]` are the following

# Mandatory keys
- `:tags`: Dirichlet tags where to impose strong boundary conditions for the fluid velocity.

# Optional keys
- `:values => zero_values(prams[:bcs][:u][:tags])`: The fluid velocity value or function
   to be imposed at each of the given tags.
"""
function params_bcs_u(params::Dict{Symbol,Any})
  mandatory = Dict(
   :tags=>true,
   :values=>false,
  )
  _check_mandatory(params[:bcs][:u],mandatory,"[:bcs][:u]")
  optional = Dict(
    :values=>zero_values(params[:bcs][:u][:tags]),
  )
  u = _add_optional(params[:bcs][:u],mandatory,optional,params,"[:bcs][:u]")
  _check_unused(u,mandatory,params,"[:bcs][:u]")
  u
end

zero_values(tags) = VectorValue(0,0,0)
zero_values(tags::Vector) = map(zero_values,tags)

"""
Valid keys for `params[:bcs][:j]` are the following

# Mandatory keys
- `:tags`: Dirichlet tags where to impose strong boundary conditions for the charge current in normal direction.

# Optional keys
- `:values => zero_values(params[:bcs][:j][:tags])`: The charge current value or function
   to be imposed at each of the given tags.
"""
function params_bcs_j(params::Dict{Symbol,Any})
  mandatory = Dict(
   :tags=>true,
   :values=>false,
  )
  _check_mandatory(params[:bcs][:j],mandatory,"[:bcs][:j]")
  optional = Dict(
    :values=>zero_values(params[:bcs][:j][:tags]),
  )
  j = _add_optional(params[:bcs][:j],mandatory,optional,params,"[:bcs][:j]")
  _check_unused(j,mandatory,params,"[:bcs][:j]")
  j
end

"""
Valid keys for the dictionaries in `params[:bcs][:φ]` are the following.

# Mandatory keys
- `:domain`: Domain where to impose the potential weakly.
  The domain is represented either with a `Triangulation` or with a `Integer`/`String`
  tag in the underlying discrete model.
- `:value`: Value of the electric potential to be imposed weakly.
"""
function params_bcs_φ(params::Dict{Symbol,Any})
  mandatory = Dict(
   :domain=>true,
   :value=>true,
  )
  optional = Dict()
  _check_mandatory_and_add_optional_weak(params[:bcs][:φ],mandatory,optional,params,"[:bcs][:φ]")
end

"""
Valid keys for the dictionaries in `params[:bcs][:t]` are the following.

# Mandatory keys
- `:domain`: Domain where to impose the fluid boundary traction weakly.
  The domain is represented either with a `Triangulation` or with a `Integer`/`String`
  tag in the underlying discrete model.
- `:value`: Value of the fluid traction to be imposed weakly.
"""
function params_bcs_t(params::Dict{Symbol,Any})
  mandatory = Dict(
   :domain=>true,
   :value=>true,
  )
  optional = Dict()
  _check_mandatory_and_add_optional_weak(params[:bcs][:t],mandatory,optional,params,"[:bcs][:t]")
end

"""
Valid keys for the dictionaries in `params[:bcs][:thin_wall]` are the following.

The thin wall law is

    j⋅n + cw*n⋅∇(j)⋅n = jw

where `j` is unknown, `n`  is the boundary outward normal, and `cw,jw` are parameters.
The thin wall law is imposed weakly via a penalty parameter `τ`.

# Mandatory keys
- `:domain`: Domain where to impose the thin wall law.
  The domain is represented either with a `Triangulation` or with a `Integer`/`String`
  tag in the underlying discrete model.
- `:cw`: Value of the parameter `cw`.
- `:τ`: Value of the parameter `τ`.

# Optional keys
- `:jw=>0`: Value of the parameter `jw`.
"""
function params_bcs_thin_wall(params::Dict{Symbol,Any})
  mandatory = Dict(
   :domain=>true,
   :cw=>true,
   :τ=>true,
   :jw=>false,
  )
  optional = Dict(:jw=>0)
  _check_mandatory_and_add_optional_weak([:bcs][:thin_wall],mandatory,optional,params,"[:bcs][:thin_wall]")
end

function _check_mandatory(params,mandatory,p)
  if !isa(params,Dict{Symbol})
    error("The params$p has to be a Dict{Symbol}")
  end
  for key in keys(mandatory)
    if mandatory[key] && !haskey(params,key)
      error("Key :$key is a mandatory key in params$p, but it is not provided.")
    end
  end
end

function _add_optional(_subparams,mandatory,optional,params,p)
  # New dict
  subparams = Dict{Symbol,Any}()
  merge!(subparams,_subparams)
  # Check that we have computed defaults for all optionals
  for key in keys(mandatory)
    if !mandatory[key] && !haskey(optional,key)
      error("Internal error")
    end
  end
  # Set default args
  for key in keys(optional)
    if !haskey(subparams,key)
      subparams[key] = optional[key]
    end
  end
  subparams
end

function _check_unused(subparams,mandatory,params,p)
  if params[:check_valid]
    for key in keys(subparams)
      if !haskey(mandatory,key)
        error("Key :$key is not a valid key in params$p. Set params[:check_valid] = false to ignore invalid keys.")
      end
    end
  end
end

function _check_mandatory_and_add_optional(_subparams,mandatory,optional,params,p)
  _check_mandatory(_subparams,mandatory,p)
  subparams = _add_optional(_subparams,mandatory,optional,params,p)
  _check_unused(subparams,mandatory,params,p)
  subparams
end

function _check_mandatory_and_add_optional_weak(t,mandatory,optional,params,p)
  function _params_bcs(_t)
    error("The value params$p has to be a Dict{Symbol} or a vector of Dict{Symbol}.")
  end
  function _params_bcs(_t::AbstractVector)
    map(i->_params_bcs(_t[i],"[$i]"),1:length(_t))
  end
  function _params_bcs(_t::Dict{Symbol},i="")
    _check_mandatory_and_add_optional(_t,mandatory,optional,params,"$p$i")
  end
  r = _params_bcs(t)
  isa(r,Dict) ? [r] : r
end

"""
    main(params::Dict{Symbol}) -> xh, full_params

Solve the MHD equations. Function `main` takes a dictionary `params`
that contains several parameters defining a MHD problem (see below)
and returns `xh` and `full_params`.
Object `xh` represents the solution of the MHD problem and it is of type `Gridap.MultiField.MultiFieldFEFunction`,
which can be unpacked to get access to the different fields of the MHD solution
(fluid velocity, fluid pressure, charge current, and electric potential respectively).
One can further post process these quantities using the tools provided by Gridap.
On the other hand `full_params` is a copy of `params` augmented with the default values
used in the computation.

In the fluid domain, the equations being solved are

     ∇⋅j = 0
     j + σ*∇(φ) - σ*(u×B) = 0
     ∇⋅u = 0
     α*u⋅∇(u) - β*Δ(u) + ∇(p) - γ*(j×B) = f

where `u,p,j,φ` are the unknowns and `α,β,γ,σ,B,f` are (possibly spatially-dependent) parameters.  In the solid domain, we solve

     ∇⋅j = 0
     j + σ*∇(φ) = 0

where `j,φ` are the unknowns and `σ` is a (possibly spatially-dependent) parameter.
These equations are augmented with suitable boundary conditions (see below).

The MHD problem is customized by setting the following keys in `params`.

# Mandatory keys
- `:model`: The finite element mesh representing the union of the fluid and solid domains. It should be either a `Gridap.DiscreteModel`
or a `GridapDistributed.DistributedDiscreteModel`
- `:fluid`: A `Dict` defining the fluid domain and fluid parameters.
   See [`params_fluid`](@ref) for further details.
- `:bcs`: A `Dict` defining the boundary conditions and other external loads.
  See [`params_bcs`](@ref) for further details.

# Optional keys
- `:solid => nothing`:
  A `Dict` defining the solid domain and solid parameters.
  If not provided or set to `nothing` the solid domain is not taken into account.
  See [`params_solid`](@ref) for further details.
- `:k => 2`:
  Maximum interpolation order (i.e., the order used for the fluid velocity).
- `:solver => default_solver()`:
  Nonlinear solver to compute the solution.
    It should be an instance of some type implementing the `NonlinearSolver` interface of Gridap.
- `:debug => false`:
  If true, setup the problem, but do not solve it. Otherwise, solve it.
- `:check_valid => true`: If `true`, check that all given keys are valid. Otherwise, silently ignore invalid keys.
- `:matrix_type => SparseMatrixCSC{Float64,Int}`:
   Matrix type to assemble the problem.
- `:vector_type => Vector{Float64}`:
  Vector type to assemble the problem.
- `:ptimer => default_ptimer(params[:model])`:
  Instance of `PTimer` used to monitor times. New time measurements are added to the given timer.
"""
function main(_params::Dict)

  params = add_default_params(_params)

  t = params[:ptimer]
  tic!(t;barrier=true)

  # ReferenceFEs
  k::Int = params[:k]
  T = Float64
  model = params[:model]
  D = num_cell_dims(model)
  reffe_u = ReferenceFE(lagrangian,VectorValue{D,T},k)
  reffe_p = ReferenceFE(lagrangian,T,k-1;space=:P)
  reffe_j = ReferenceFE(raviart_thomas,T,k-1)
  reffe_φ = ReferenceFE(lagrangian,T,k-1)

  # Test spaces
  Ωf = _fluid_mesh(model,params[:fluid][:domain])
  V_u = TestFESpace(Ωf,reffe_u;dirichlet_tags=params[:bcs][:u][:tags])
  V_p = TestFESpace(Ωf,reffe_p;conformity=p_conformity(Ωf))
  V_j = TestFESpace(model,reffe_j;dirichlet_tags=params[:bcs][:j][:tags])
  V_φ = TestFESpace(model,reffe_φ;conformity=:L2)
  V = MultiFieldFESpace([V_u,V_p,V_j,V_φ])

  # Trial spaces
  # TODO improve for parallel computations
  tic!(t;barrier=true)
  z = zero(VectorValue{D,Float64})
  u_bc = params[:bcs][:u][:values]
  j_bc = params[:bcs][:j][:values]
  U_u = u_bc == z ? V_u : TrialFESpace(V_u,u_bc)
  U_j = j_bc == z ? V_j : TrialFESpace(V_j,j_bc)
  U_p = TrialFESpace(V_p)
  U_φ = TrialFESpace(V_φ)
  U = MultiFieldFESpace([U_u,U_p,U_j,U_φ])
  toc!(t,"fe_spaces")

  tic!(t;barrier=true)
  if params[:debug]
    Random.seed!(1234)
    vt = get_vector_type(U)
    free_ids = get_free_dof_ids(U)
    free_vals = _rand(vt,free_ids)
    xh = FEFunction(U,free_vals)
  else
    res, jac = weak_form(params,k)
    Tm = params[:matrix_type]
    Tv = params[:vector_type]
    assem = SparseMatrixAssembler(Tm,Tv,U,V)
    op = FEOperator(res,jac,U,V,assem)
    solver = params[:solver]
    xh = zero(U)
    xh,cache = solve!(xh,solver,op)
    solver_postpro = params[:solver_postpro]
    solver_postpro(cache)
  end
  toc!(t,"solve")

  xh, params
end

function _fluid_mesh(
  model,domain::Union{Gridap.DiscreteModel,GridapDistributed.DistributedDiscreteModel})
  msg = "params[:fluid][:domain] is a discrete model, but params[:fluid][:domain]===params[:model] is not true."
  @assert model === domain msg
  domain
end

function _fluid_mesh(
  model,
  domain::Union{Gridap.Triangulation,GridapDistributed.DistributedTriangulation})
  domain
end

function _fluid_mesh( model, domain)
  Interior(model,tags=domain)
end

function _interior(
  model,domain::Union{Gridap.DiscreteModel,GridapDistributed.DistributedDiscreteModel})
  Interior(domain)
end

function _interior(
  model,
  domain::Union{Gridap.Triangulation,GridapDistributed.DistributedTriangulation})
  domain
end

function _interior( model, domain)
  Interior(model,tags=domain)
end

function _boundary(
  model,
  domain::Union{Gridap.Triangulation,GridapDistributed.DistributedTriangulation})
  domain
end

function _boundary( model, domain)
  Boundary(model,tags=domain)
end

function _rand(vt::Type{<:Vector{T}},r::AbstractUnitRange) where T
  rand(T,length(r))
end

function _rand(vt::Type{<:PVector{T,A}},ids::PRange) where {T,A}
  values = map_parts(ids.partition) do partition
    Tv = eltype(A)
    _rand(Tv,1:num_lids(partition))
  end
  PVector(values,ids)
end

function p_conformity(poly::Polytope)
  if is_simplex(poly)
    conf = :H1
  elseif is_n_cube(poly)
    conf = :L2
  else
    @unreachable "unsupported cell topology"
  end
  conf
end

function p_conformity(Ω::Triangulation)
  reffes = get_reffes(Ω)
  @assert length(reffes) == 1
  reffe = first(reffes)
  poly = get_polytope(reffe)
  p_conformity(poly)
end

function p_conformity(Ω::GridapDistributed.DistributedTriangulation)
  p = map_parts(Ω.trians) do Ω
    reffes = get_reffes(Ω)
    @assert length(reffes) == 1
    reffe = first(reffes)
    poly = get_polytope(reffe)
    poly
  end
  poly = get_part(p) # We assume same polytope in all parts
  p_conformity(poly)
end

function p_conformity(model::DiscreteModel)
  Ω = Interior(model)
  p_conformity(Ω)
end

function p_conformity(model::GridapDistributed.DistributedDiscreteModel)
  Ω = Interior(model)
  p_conformity(Ω)
end

function add_defaults!(params,defaults)
  for (key,val) in defaults
    if !haskey(params,key)
      params[key] = val
    elseif isa(val,AbstractDict)
      @assert isa(params[key],AbstractDict)
      add_defaults!(params[key],val)
    end
  end
end

function weak_form(params,k)

  fluid = params[:fluid]

  Ωf = _interior(params[:model],fluid[:domain])
  dΩf = Measure(Ωf,2*k)

  solid = params[:solid]
  if solid !== nothing
    Ωs = _interior(params[:model],solid[:domain])
    dΩs = Measure(Ωs,2*k)
    σs = solid[:σ]
  end

  α = fluid[:α]
  β = fluid[:β]
  γ = fluid[:γ]
  f = fluid[:f]
  B = fluid[:B]
  σf = fluid[:σ]

  bcs = params[:bcs]

  params_φ = []
  for i in 1:length(bcs[:φ])
    φ_i = bcs[:φ][i][:value]
    Γ = _boundary(params[:model],bcs[:φ][i][:domain])
    dΓ = Measure(Γ,2*k)
    n_Γ = get_normal_vector(Γ)
    push!(params_φ,(φ_i,n_Γ,dΓ))
  end

  params_thin_wall = []
  for i in 1:length(bcs[:thin_wall])
    τ_i = bcs[:thin_wall][i][:τ]
    cw_i = bcs[:thin_wall][i][:cw]
    jw_i = bcs[:thin_wall][i][:jw]
    Γ = _boundary(params[:model],bcs[:thin_wall][i][:domain])
    dΓ = Measure(Γ,2*k)
    n_Γ = get_normal_vector(Γ)
    push!(params_thin_wall,(τ_i,cw_i,jw_i,n_Γ,dΓ))
  end

  if length(bcs[:t]) != 0
    error("Boundary tranction not yet implemented")
  end

  params_f = []
  for i in 1:length(bcs[:f])
    f_i = bcs[:f][i][:value]
    Ω_i = _interior(params[:model],bcs[:f][i][:domain])
    dΩ_i = Measure(Ω_i,2*k)
    push!(params_f,(f_i,dΩ_i))
  end

  params_B = []
  for i in 1:length(bcs[:B])
    B_i = bcs[:B][i][:value]
    Ω_i = _interior(params[:model],bcs[:B][i][:domain])
    dΩ_i = Measure(Ω_i,2*k)
    push!(params_f,(γ,B_i,dΩ_i))
  end

  function a(x,dy)
    r = a_mhd(x,dy,β,γ,B,σf,dΩf)
    for p in params_thin_wall
      r = r + a_thin_wall(x,dy,p...)
    end
    for p in params_B
      r = r + a_B(x,dy,p...)
    end
    if solid !== nothing
      r = r + a_solid(x,dy,σs,dΩs)
    end
    r
  end

  function ℓ(dy)
    r = ℓ_mhd(dy,f,dΩf)
    for p in params_φ
      r = r + ℓ_φ(dy,p...)
    end
    for p in params_thin_wall
      r = r + ℓ_thin_wall(dy,p...)
    end
    for p in params_f
      r = r + ℓ_f(dy,p...)
    end
    r
  end

  function c(x,dy)
    r = c_mhd(x,dy,α,dΩf)
    r
  end

  function dc(x,dx,dy)
    r = dc_mhd(x,dx,dy,α,dΩf)
    r
  end

  res(x,dy) = c(x,dy) + a(x,dy) - ℓ(dy)
  jac(x,dx,dy) = dc(x,dx,dy) + a(dx,dy)

  res, jac
end

conv(u,∇u) = (∇u')⋅u

function a_solid(x,dy,σ,dΩ)
  u, p, j, φ = x
  v_u, v_p, v_j, v_φ = dy
  ∫( j⋅v_j - σ*φ*(∇⋅v_j) + (∇⋅j)*v_φ)dΩ
end

function a_mhd(x,dy,β,γ,B,σ,dΩ)
  u, p, j, φ = x
  v_u, v_p, v_j, v_φ = dy
  ∫(
    β*(∇(u)⊙∇(v_u)) - p*(∇⋅v_u) -(γ*(j×B)⋅v_u) +
    (∇⋅u)*v_p +
    j⋅v_j - σ*φ*(∇⋅v_j) - σ*(u×B)⋅v_j +
    (∇⋅j)*v_φ ) * dΩ
end

function ℓ_mhd(dy,f,dΩ)
  v_u, v_p, v_j, v_φ = dy
  ∫( v_u⋅f )*dΩ
end

function c_mhd(x,dy,α,dΩ)
  u, p, j, φ = x
  v_u, v_p, v_j, v_φ = dy
  ∫( α*v_u⋅(conv∘(u,∇(u))) ) * dΩ
end

function dc_mhd(x,dx,dy,α,dΩ)
  u, p, j, φ = x
  du , dp , dj , dφ  = dx
  v_u, v_p, v_j, v_φ = dy
  ∫( α*v_u⋅( (conv∘(u,∇(du))) + (conv∘(du,∇(u))) ) ) * dΩ
end

function ℓ_φ(dy,φ,n_Γ,dΓ)
  v_u, v_p, v_j, v_φ = dy
  ∫( -(v_j⋅n_Γ)*φ )*dΓ
end

function ℓ_f(dy,f,dΩ)
  v_u, v_p, v_j, v_φ = dy
  ∫( v_u⋅f )*dΩ
end

function ℓ_thin_wall(dy,τ,cw,jw,n_Γ,dΓ)
  v_u, v_p, v_j, v_φ = dy
  ∫( τ*(v_j⋅n_Γ)*jw ) * dΓ
end

function a_thin_wall(x,dy,τ,cw,jw,n_Γ,dΓ)
  u, p, j, φ = x
  v_u, v_p, v_j, v_φ = dy
  ∫( τ*((v_j⋅n_Γ)*(j⋅n_Γ) + cw*(v_j⋅n_Γ)*(n_Γ⋅(∇(j)⋅n_Γ))) )*dΓ
end

function a_B(x,dy,γ,B,dΩ)
  u, p, j, φ = x
  v_u, v_p, v_j, v_φ = dy
  ∫( -(γ*(j×B)⋅v_u) - (u×B)⋅v_j )*dΩ
end


text:
"""
```
write_standard_packet(m, input_type, cond_type,
    output_vars = [:forecastobs, :forecastpseudo, :shockdecobs, :shockdecpseudo];
    sections = [:estimation, :forecast]
    forecast_string = "", outdir = joinpath(saveroot(m), \"Packets\", spec(m), subspec(m)))
```

Write standard estimation and forecast result packet to `outdir`.
"""
function write_standard_packet(m::AbstractModel, input_type::Symbol, cond_type::Symbol,
                               output_vars::Vector{Symbol} = [:forecastobs, :forecastpseudo,
                                                              :shockdecobs, :shockdecpseudo];
                               sections::Vector{Symbol} = [:estimation, :forecast],
                               forecast_string::String = "",
                               outdir::String = joinpath(saveroot(m), "Packets", spec(m), subspec(m)),
                               purpose::String = "")
    @assert issubset(sections, [:estimation, :forecast, :irf]) "Section specified in `section` kwarg is not supported. Must be a subset of [:estimation, :forecast, :irf]."

    # Title and authors
    title = "Estimation and Forecasting Results. \\\\ " * description(m)
    authors = "User"

    # Compute file name
    cdvt_str = cond_type == :none ? "" : "_cdvt=" * cond_vintage(m)
    fn = joinpath(outdir, "results_cond=" * string(cond_type) * cdvt_str * "_vint=" * data_vintage(m) * ".tex")
    isdir(dirname(fn)) || mkpath(dirname(fn))

    # Write packet
    open(fn, "w") do fid
        write_preamble(fid, title, authors)
        write_spec_section(fid, m, purpose = purpose)
        if :estimation in sections
            write_estimation_section(fid, m)
        end
        if :forecast in sections
            write_forecast_section(fid, m, input_type, cond_type, setdiff(output_vars, [:irfstates, :irfobs, :irfpseudo]),
                                   forecast_string = forecast_string)
        end
        if :irf in sections
            write_irf_section(fid, m, input_type, cond_type,
                              output_vars, forecast_string = forecast_string)
        end
        write_postamble(fid)
    end
    println("Wrote " * fn)
end

"""
```
plot_standard_packet(m, input_type, cond_type,
    output_vars = [:forecastobs, :forecastpseudo, :shockdecobs, :shockdecpseudo];
    sections = [:estimation, :forecast]
    forecast_string = "", hist_start-date = 0001-01-01)
```

Plot parameter prior/posterior histograms to `figurespath(m, \"estimate\")` and
forecasts/shock decompositions to `figurespath(m, \"forecast\')`.
"""
function plot_standard_packet(m::AbstractModel, input_type::Symbol, cond_type::Symbol,
                              output_vars::Vector{Symbol} = [:forecastobs, :forecastpseudo,
                                                              :shockdecobs, :shockdecpseudo];
                              sections::Vector{Symbol} = [:estimation, :forecast],
                              forecast_string::String = "",
                              hist_start_date::Date = Date("0001-01-01", "yyyy-mm-dd"))
    @assert issubset(sections, [:estimation, :forecast, :irf]) "Section specified in `section` kwarg is not supported. Must be a subset of [:estimation, :forecast, :irf]."
    if :estimation in sections
        plot_prior_posterior(m)
    end
    if :forecast in sections
        for output_var in setdiff(output_vars, [:irfstates, :irfobs, :irfpseudo])
            make_forecast_plots(m, input_type, cond_type, output_var,
                                forecast_string = forecast_string,
                                hist_start_date = hist_start_date)
        end
    end
    if :irf in sections
        plot_irf_section(m, input_type, cond_type, output_vars;
                         forecast_string = forecast_string)
    end
end

"""
```
write_spec_section(fid, m; purpose = "")

write_spec_section(fid, m1, m2; purpose = "")
```

Write user, model, and data specification section.
"""
function write_spec_section(fid::IOStream, m::AbstractModel; purpose::String = "")
    @printf fid "\n\n"
    @printf fid "\\section{Specification}\n"
    @printf fid "\n"

    @printf fid "\\subsection{User}\n"
    @printf fid "\n"
    @printf fid "\\begin{itemize}\n"
    @printf fid "  \\item Packet creator: %s\n" string(ENV["USER"])
    @printf fid "  \\item Creation time: %s\n" Dates.format(now(), "U d, YYYY at H:MM")
    if !isempty(purpose)
        @printf fid "\\item Purpose: %s\n" purpose
    end
    @printf fid "\\end{itemize}\n"
    @printf fid "\n"

    @printf fid "\\subsection{Model}\n"
    @printf fid "\n"
    @printf fid "\\begin{itemize}\n"
    @printf fid "  \\item Model: %s\n" m.spec
    @printf fid "  \\item Subspec: %s\n" m.subspec
    @printf fid "  \\item Number of states: %d\n" n_states_augmented(m)
    @printf fid "  \\item Number of shocks: %d\n" n_shocks_exogenous(m)
    @printf fid "  \\item Number of observables: %d\n" n_observables(m)
    @printf fid "  \\item Number of parameters: %d\n" n_parameters(m)
    @printf fid "  \\item Number of anticipated shocks: %d\n" n_anticipated_shocks(m)
    @printf fid "\\end{itemize}\n"
    @printf fid "\n"

    @printf fid "\\subsection{Data}\n"
    @printf fid "\n"
    @printf fid "\\begin{itemize}\n"
    @printf fid "  \\item Input data directory: \\path{%s}\n" get_setting(m, :dataroot)
    @printf fid "  \\item Output data directory: \\path{%s}\n" get_setting(m, :saveroot)
    @printf fid "  \\item Data vintage: %s\n" get_setting(m, :data_vintage)
    @printf fid "  \\item Dataset ID: %d\n" get_setting(m, :data_id)
    @printf fid "  \\item Conditional data vintage: %s\n" get_setting(m, :cond_vintage)
    @printf fid "  \\item Conditional dataset ID: %d\n" get_setting(m, :cond_id)
    @printf fid "  \\item Use population forecast: %s\n" get_setting(m, :use_population_forecast)
    @printf fid "  \\item Last quarter of data: %d-Q%d\n" Dates.year(date_mainsample_end(m)) Dates.quarterofyear(date_mainsample_end(m))
    @printf fid "  \\item Last quarter of conditional data: %d-Q%d\n" Dates.year(date_conditional_end(m)) Dates.quarterofyear(date_conditional_end(m))
    @printf fid "\\end{itemize}\n"
end

"""
```
write_estimation_section(fid, m; plotroot = "")
```

Write parameter moment tables and prior/posterior plots.
"""
function write_estimation_section(fid::IOStream, m::AbstractModel;
                                  plotroot::String = "")
    @printf fid "\n\n"
    @printf fid "\\clearpage\n"
    @printf fid "\\section{Estimation}\n"
    @printf fid "\n"

    @printf fid "\\subsection{Histograms}\n"
    write_estimation_plots(fid, m, plotroot = plotroot)
    @printf fid "\n"

    @printf fid "\\clearpage\n"
    @printf fid "\\subsection{Moments}\n"
    @printf fid "\n"

    @printf fid "\\input{%s}\n" tablespath(m, "estimate", "moments.tex")
end

"""
```
write_estimation_plots(fid, m; plotroot = "")
```

Write LaTeX code displaying plots of `m` `product`s to the `IOStream` `fid`. If
`plotroot` is not specified, plots from `figurespath(m, \"estimate\")` are used.
"""
function write_estimation_plots(fid::IOStream, m::AbstractModel;
                                plotroot::String = "")
    if isempty(plotroot)
        plotroot = figurespath(m, "estimate")
    end
    base = DSGE.filestring_base(m)
    filestr = DSGE.filestring(base, String[])

    @printf fid "\n"
    @printf fid "\\begin{longtable}{cc}\n"
    left = true
    for param in m.parameters
        if param.fixed
            continue
        else
            @printf fid "\\includegraphics[width=0.45\\textwidth]{%s/prior_posterior_%s%s.pdf} %s\n" plotroot DSGE.detexify(param.key) filestr (left ? "&" : "\\\\")
            left = !left
        end
    end

    @printf fid "\\end{longtable}\n"
end

"""
```
write_forecast_section(fid, m, input_type, cond_type,
    output_vars = [:forecastobs, :forecastpseudo, :shockdecobs, :shockdecpseudo];
    forecast_string = "". plotroot = "")
```

Write forecast specification and plots for the given `output_vars`. If `plotroot`
is not specified, plots from `figurespath(m, \"forecast\")` are used.
"""
function write_forecast_section(fid::IOStream, m::AbstractModel, input_type::Symbol, cond_type::Symbol,
                                output_vars::Vector{Symbol} = [:forecastobs, :forecastpseudo,
                                                               :shockdecobs, :shockdecpseudo];
                                forecast_string::String = "",
                                plotroot = "")
    @printf fid "\n\n"
    @printf fid "\\clearpage\n"
    @printf fid "\\section{Forecast}\n"
    @printf fid "\n"

    @printf fid "\\subsection{Specification}\n"
    @printf fid "\n"
    @printf fid "\\begin{itemize}\n"
    @printf fid "  \\item Forecast type: %s\n" input_type
    @printf fid "  \\item Forecast conditional type: %s\n" cond_type
    if !isempty(forecast_string)
        fc_string = if occursin("_", forecast_string)
            replace(forecast_string, "_" => "\\_")
        else
            forecast_string
        end
        @printf fid "  \\item Forecast identifying string: %s\n" fc_string
    end
    @printf fid "  \\item Estimation used: \\path{%s}\n" DSGE.get_forecast_input_file(m, input_type)
    @printf fid "\\end{itemize}\n"

    products = map(get_product, output_vars)

    if :forecast in products
        @printf fid "\\clearpage\n"
        @printf fid "\\subsection{Forecasts}\n"

        for output_var in Base.filter(x -> get_product(x) == :forecast, output_vars)
            write_forecast_plots(fid, m, input_type, cond_type, output_var,
                                 forecast_string = forecast_string,
                                 plotroot = plotroot)
            @printf fid "\n"
        end
    end

    n_output_vars = 2 * (:forecastobs in output_vars) + 4 * (:forecastpseudo in output_vars)
    if n_output_vars > 0
        @printf fid "\\clearpage\n"
        @printf fid "\\subsection{Mean Forecast Estimates}\n"
        write_forecast_table(fid, m, cond_type, output_vars; forecast_string = forecast_string)
    end

    if :shockdec in products
        @printf fid "\\clearpage\n"
        @printf fid "\\subsection{Shock Decompositions}\n"

        for output_var in Base.filter(x -> get_product(x) == :shockdec, output_vars)
            write_forecast_plots(fid, m, input_type, cond_type, output_var,
                                 forecast_string = forecast_string,
                                 plotroot = plotroot)
            @printf fid "\n"
        end
    end
end

"""
```
make_forecast_plots(m, input_type, cond_type, output_var;
    forecast_string = "", plotroot = "", hist_start_date = 0001-01-01)
```

Generate all `output_var` plots for the forecast of `m` specified by the
`input_type`, `cond_type`, and optional `forecast_string`. If `plotroot` is not
specified, plots are saved to `figurespath(m, \"forecast\")`.
"""
function make_forecast_plots(m::AbstractModel, input_type::Symbol, cond_type::Symbol, output_var::Symbol;
                             forecast_string::String = "",
                             plotroot::String = "",
                             hist_start_date = Date("0001-01-01", "yyyy-mm-dd"))

    # Output directory
    if isempty(plotroot)
        plotroot = figurespath(m, "forecast")
    end

    # Set class-specific variables
    class = get_class(output_var)
    product = get_product(output_var)

    if class == :obs
        varnames = collect(keys(m.observables))
        hist_label = "Data"
    elseif class == :pseudo
        varnames = collect(keys(m.pseudo_observables))
        hist_label = "Mean Estimate"
    elseif class == :states
        varnames = collect(keys(m.endogenous_states))
        hist_label = "Mean Estimate"
    else
        error("Unsupported class: " * string(class))
    end

    # Added min() so forecasts that are starting from a year before 2007 will plot properly
    if hist_start_date != Date("0001-01-01", "yyyy-mm-dd")
        start_date = hist_start_date
    else
        start_date = min(DSGE.quartertodate("2007-Q1"), date_mainsample_end(m) - Dates.Year(5))
    end
    if haskey(m.settings, :date_forecast_end)
        end_date = get_setting(m, :date_forecast_end)

        # check forecast horizon was long enough
        end_date -= Dates.Month(3 * get_setting(m, :forecast_horizons))
        if end_date > get_setting(m, :date_forecast_start)
            error("forecast horizons is not long enough for the requested forecast end date")
        else
            end_date = get_setting(m, :date_forecast_end)
        end
    elseif haskey(m.settings, :forecast_horizons)
        end_date = get_setting(m, :date_forecast_start) # this will become the end date
        end_date += Dates.Month((get_setting(m, :forecast_horizons)-1) * 3)
    else
        end_date = DSGE.quartertodate("2020-Q1")
    end
    forecast_label = "Mean Forecast"
    tick_size = 2

    # Make product-specific plot
    product = get_product(output_var)

    if product == :forecast
        plot_history_and_forecast(m, varnames, class, input_type, cond_type,
                                  forecast_string = forecast_string,
                                  bdd_and_unbdd = true,
                                  plotroot = plotroot,
                                  start_date = start_date,
                                  names = Dict(:hist => hist_label, :forecast => forecast_label),
                                  end_date = end_date,
                                  tick_size = tick_size,
                                  legend = :bottomleft)
    elseif product == :shockdec
        groups = DSGE.shock_groupings(m)
        plot_shock_decomposition(m, varnames, class, input_type, cond_type,
                                 forecast_string = forecast_string,
                                 plotroot = plotroot,
                                 start_date = start_date,
                                 end_date = end_date,
                                 hist_label = "",
                                 forecast_label = "",
                                 tick_size = tick_size,
                                 legend = :bottomleft)
    else
        error("Unsupported product: " * string(product))
    end
end

"""
```
write_forecast_plots(fid, m, input_type, cond_type, output_var;
    forecast_string = "", plotroot = "")
```

Write LaTeX code displaying plots of `m` `product`s to the `IOStream` `fid`. If
`plotroot` is not specified, plots from `figurespath(m, \"forecast\")` are used.
"""
function write_forecast_plots(fid::IOStream, m::AbstractModel,
                              input_type::Symbol, cond_type::Symbol, output_var::Symbol;
                              forecast_string::String = "",
                              plotroot::String = "")
    if isempty(plotroot)
        plotroot = figurespath(m, "forecast")
    end
    base = DSGE.filestring_base(m)
    addl = DSGE.get_forecast_filestring_addl(input_type, cond_type, forecast_string = forecast_string)
    filestr = DSGE.filestring(base, addl)

    class = get_class(output_var)
    if class == :obs
        section_title = "Observables"
        dict = m.observables
    elseif class == :pseudo
        section_title = "Pseudo-Observables"
        dict = m.pseudo_observables
    else
        error("Unsupported class: " * string(class))
    end

    product = get_product(output_var)

    @printf fid "\n"
    @printf fid "\\subsubsection{%s}\n" section_title
    @printf fid "\\begin{longtable}{c}\n"
    for var in keys(dict)
        @printf fid "\\includegraphics[height=0.4\\textheight]{%s/%s_%s%s.pdf} \\\\\n" plotroot product DSGE.detexify(var) filestr
    end
    @printf fid "\\end{longtable}\n"
end

"""
```
write_forecast_table(fid, m, cond_type, output_vars;
    forecast_string = "")
```

Write the standard packet's forecast table listing forecasts 3 years out.

### Inputs

- `fid::IOStream`: file to write to
- `m::AbstractModel`: model object for current Snapshot forecast
- `cond_type::Symbol`: type of conditional forecast
- `output_vars`: which type of forecast variables to print out

### Keyword Arguments

- `forecast_str::String`: string to use
"""
function write_forecast_table(fid::IOStream, m::AbstractModel, cond_type::Symbol,
                              output_vars::Vector{Symbol};
                              forecast_string::String = "")

    # Infer dates
    first_fcast_year = Dates.year(date_forecast_start(m))
    fcast_years = [first_fcast_year, first_fcast_year + 1, first_fcast_year + 2, first_fcast_year + 3]
    fcast_dates = map(year -> Date(year, 12, 31), fcast_years)

    curr_month = month_label(m)

    if curr_month != "Jan" && curr_month != "Feb"
        curr_year = "20"*data_vintage(m)[1:2]
    else
        # Because the first forecast of the year still forecasts the previous year.
        curr_year = "20"*string(parse(data_vintage(m)[1:2])-1)
    end

    # Open output file, write header
    year2 = string(Meta.parse(curr_year)+1)
    year3 = string(Meta.parse(curr_year)+2)
    year4 = string(Meta.parse(curr_year)+3)

    write(fid, "\\begin{table}[h!]\n")
    write(fid, "\\centering\n")
    write(fid, "\\begin{tabular}{l | c | c | c | c }\n")
    write(fid, "& " * "$(curr_year)" * " & " * "$(year2)" * " & " * "$(year3)" * " & " * "$(year4)" * " \\\\\n")
    write(fid, "\\hline\n\\hline\n")

    n_output_vars = 2 * (:forecastobs in output_vars) + 4 * (:forecastpseudo in output_vars)
    if :forecastobs in output_vars
        write(fid, print_variable_means(m, cond_type,
                                        :forecast4qobs, :obs_gdp, ["GDP growth ", "(Q4/Q4)"],
                                        fcast_dates, false, forecast_string = forecast_string))
        last_row = n_output_vars == 2 ? true : false
        write(fid, print_variable_means(m, cond_type,
                                        :forecast4qobs, :obs_corepce,
                                        ["Core PCE inflation ", "(Q4/Q4)"],
                                     fcast_dates, last_row, forecast_string = forecast_string))
    end
    if :forecastpseudo in output_vars
        write(fid, print_variable_means(m, cond_type,
                                        :forecast4qpseudo, :π_t, ["Inflation ",
                                                                     "(Q4)"],
                                        fcast_dates, false, forecast_string = forecast_string))
        write(fid, print_variable_means(m, cond_type,
                                        :forecast4qpseudo, :LongRunInflation,
                                        ["Long Run Inflation ", "(Q4)"],
                                        fcast_dates, false, forecast_string = forecast_string))
        write(fid, print_variable_means(m, cond_type,
                                        :forecast4qpseudo, :NaturalRate,
                                        ["Real natural rate", "of interest (Q4)"],
                                        fcast_dates, false, forecast_string = forecast_string))
        write(fid, print_variable_means(m, cond_type,
                                        :forecast4qpseudo, :OutputGap,
                                        ["Output gap", "(Q4)"],
                                        fcast_dates, true, forecast_string = forecast_string))
    end

    write(fid, "\\caption{Mean Forecast Estimates}\n")
    write(fid, "\\end{table}")
end

function print_variable_means(m::AbstractModel, cond_type::Symbol, output_var::Symbol,
                              varname::Symbol, vardesc::Vector{String}, fcast_dates::Vector{Date},
                              last_row::Bool; forecast_string::String = "")
    n_fcast_dates = length(fcast_dates)

    mb_curr = read_mb(m, :full, cond_type, output_var; forecast_string = forecast_string)

    # Index out Q4 values for current forecast
    all_dates   = mb_curr.metadata[:date_inds]

    # The following works because all_dates is an ordered dictionary so can just treat keys as a array and take their indices.
    # The commented out part stopped working because Julia7 doesn't like indexing with a Date key
    table_dates = findall(x -> x in fcast_dates, collect(keys(all_dates))) #map(date -> all_dates[date], fcast_dates)
    cur_val     = round.(mb_curr.means[:, [:date, varname]][table_dates, 2], digits = 1)

    row1_string = vardesc[1]*repeat(" ", 34-length(vardesc[1])) * "& "
    for i in 1:n_fcast_dates
        if i != n_fcast_dates
            row1_string = string(row1_string, cur_val[i]," & ")
        else
            row1_string = string(row1_string, cur_val[i]," \\\\\n")
        end
    end
    if last_row
        row2_string = vardesc[2]*repeat(" ", 34 - length(vardesc[2]))*repeat("&     ", n_fcast_dates)*"\\\\\n\\end{tabular}"
    else
        row2_string = vardesc[2]*repeat(" ", 34 - length(vardesc[2]))*repeat("&     ", n_fcast_dates)*"\\\\\n\\hline\n"
    end

    return row1_string*row2_string
end

"""
```
write_irf_section(fid, m, input_type, cond_type, forecast_string = "", plotroot = "")
```

Write impulse responses. If `plotroot`
is not specified, plots from `figurespath(m, \"forecast\")` are used.
"""
function write_irf_section(fid::IOStream, m::AbstractModel, input_type::Symbol, cond_type::Symbol,
                           output_vars::Vector{Symbol}; forecast_string::String = "",
                           plotroot = "")
    @printf fid "\n\n"
    @printf fid "\\clearpage\n"
    @printf fid "\\section{Impulse Responses}\n"
    @printf fid "\n"

    @printf fid "\\subsection{Plots}\n"
    if :irfstates in output_vars
        write_irf_plots(fid, m, input_type, cond_type, :irfstates, plotroot = plotroot,
                        forecast_string = forecast_string)
    end
    if :irfobs in output_vars
        write_irf_plots(fid, m, input_type, cond_type, :irfobs, plotroot = plotroot,
                        forecast_string = forecast_string)
    end
    if :irfpseudo in output_vars
        write_irf_plots(fid, m, input_type, cond_type, :irfpseudo, plotroot = plotroot,
                        forecast_string = forecast_string)
    end
    @printf fid "\n"
end

"""
```
write_irf_plots(fid, m, input_type, cond_type, output_var; forecast_string = "", plotroot = "")
```

Plot impulse responses. If `plotroot`
is not specified, plots from `figurespath(m, \"forecast\")` are used.
"""
function write_irf_plots(fid::IOStream, m::AbstractModel, input_type::Symbol, cond_type::Symbol,
                         output_var::Symbol; forecast_string::String = "", plotroot::String = "")
    if isempty(plotroot)
        plotroot = figurespath(m, "forecast")
    end
    base = DSGE.filestring_base(m)
    addl = DSGE.get_forecast_filestring_addl(input_type, cond_type,
                                             forecast_string = forecast_string)
    filestr = DSGE.filestring(base, addl)

    # Find shocks to remove from list of shocks to compute IRFs
    ant_sh  = [Symbol("rm_shl" * string(i)) for i = 1:n_anticipated_shocks(m)]
    zero_shocks = Vector{Symbol}(undef,0)
    for k in setdiff(keys(m.exogenous_shocks),
                     [Symbol("rm_shl" * string(i)) for i = 1:n_anticipated_shocks(m)])
        if k == :rm_sh && m[:σ_r_m].value == 0. && m[:σ_r_m].fixed
            zero_shocks = vcat(zero_shocks, :rm_sh)
        elseif k == :zp_sh && m[:σ_z_p].value == 0. && m[:σ_z_p].fixed
            zero_shocks = vcat(zero_shocks, :zp_sh)
        elseif k != :rm_sh && k != :zp_sh
            sig_name = Symbol("σ_" * replace(string(k), "_sh" => ""))
            if m[sig_name].value == 0. && m[sig_name].fixed
                zero_shocks = vcat(zero_shocks, k)
            end
        end
    end
    shocks = setdiff(setdiff(collect(keys(m.exogenous_shocks)), ant_sh), zero_shocks)

    class   = get_class(output_var)
    product = get_product(output_var)
    if class == :obs
        section_title = "Observables"
        ant_obs = [Symbol("obs_nominalrate" * string(i)) for i = 1:n_anticipated_shocks(m)]
        rowvars = setdiff(keys(m.observables), ant_obs)
    elseif class == :pseudo
        section_title = "Pseudo-Observables"
        rowvars = keys(m.pseudo_observables)
    elseif class == :states
        section_title = "Endogenous States"
        rowvars = keys(m.endogenous_states)
    else
        error("Unsupported class: " * string(class))
    end

    @printf fid "\n"
    @printf fid "\\subsubsection{%s}\n" section_title
    for var in shocks
        @printf fid "\\begin{figure}[h!]\n"
        @printf fid "\\centering\n"
        @printf fid "\\begin{longtable}{cccc}\n"
        loc = 1 # location within row
        for k in rowvars
            if loc < 4
                endstr = "&"
                loc += 1
            else
                endstr = "\\\\"
                loc = 1
            end
            @printf fid "\\includegraphics[width=0.24\\textwidth]{%s/%s_%s_%s%s.pdf} %s\n" plotroot product DSGE.detexify(var) DSGE.detexify(k) filestr endstr
        end
        @printf fid "\\end{longtable}\n"
        caption_string = if occursin("_", string(var))
            replace(string(DSGE.detexify(var)), "_" => "\\_")
        else
            string(DSGE.detexify(var))
        end
        @printf fid "\\caption{%s}\n" caption_string
        @printf fid "\\end{figure}\n"
    end
    @printf fid "\\clearpage\n"
end
"""
```
plot_irf_section(fid, m, input_type, cond_type; forecast_string = "", plotroot = "", ncols = 4)
```

Plot impulse responses. If `plotroot`
is not specified, plots from `figurespath(m, \"forecast\")` are used.
We assume the default number of columns for the joint plot is 3.
"""
function plot_irf_section(m::AbstractModel, input_type::Symbol, cond_type::Symbol,
                          output_vars::Vector{Symbol};
                          forecast_string::String = "", plotroot::String = "", ncols::Int64 = 4)
    # Create file name related strings
    if isempty(plotroot)
        plotroot = figurespath(m, "forecast")
    end
    base = DSGE.filestring_base(m)
    addl = DSGE.get_forecast_filestring_addl(input_type, cond_type,
                                             forecast_string = forecast_string)
    filestr = DSGE.filestring(base, addl)

    # Find shocks to remove from list of shocks to compute IRFs
    ant_obs = [Symbol("obs_nominalrate" * string(i)) for i = 1:n_anticipated_shocks(m)]
    ant_sh  = [Symbol("rm_shl" * string(i)) for i = 1:n_anticipated_shocks(m)]
    zero_shocks = Vector{Symbol}(undef,0)
    for k in setdiff(keys(m.exogenous_shocks),
                     [Symbol("rm_shl" * string(i)) for i = 1:n_anticipated_shocks(m)])
        if k == :rm_sh && m[:σ_r_m].value == 0. && m[:σ_r_m].fixed
            zero_shocks = vcat(zero_shocks, :rm_sh)
        elseif k == :zp_sh && m[:σ_z_p].value == 0. && m[:σ_z_p].fixed
            zero_shocks = vcat(zero_shocks, :zp_sh)
        elseif k != :rm_sh && k != :zp_sh
            sig_name = Symbol("σ_" * replace(string(k), "_sh" => ""))
            if m[sig_name].value == 0. && m[sig_name].fixed
                zero_shocks = vcat(zero_shocks, k)
            end
        end
    end

    shocks = setdiff(setdiff(collect(keys(m.exogenous_shocks)), ant_sh), zero_shocks)

    for output_var in intersect(output_vars, [:irfstates, :irfobs, :irfpseudo])
        class = get_class(output_var)
        if class == :obs
            vars = setdiff(collect(keys(m.observables)), ant_obs)
        elseif class == :pseudo
            vars = collect(keys(m.pseudo_observables))
        elseif class == :states
            vars = collect(keys(m.endogenous_states))
        end
        nvars = length(vars)
        nrows = convert(Int, ceil(nvars/ncols))

        for shock in shocks
            # Create individual plots
            plots = plot_impulse_response(m, shock, vars, class, input_type, cond_type,
                                          forecast_string = forecast_string)

            # Plot all variables together
            varplots = if mod(nvars, ncols) > 0
                empty = plot(grid = false, foreground_color_axis = :white, ticks = [])
                n_empty = ncols - mod(nvars, ncols)
                vcat(collect(values(plots)), repeat([empty], n_empty))
            else
                collect(values(plots))
            end
            p = plot(varplots..., layout = grid(nrows, ncols), size = (650, 775))

            # Save
            fn = joinpath(plotroot, "irf_" * string(DSGE.detexify(shock)) * "_" * string(class))
            fn *= filestr * ".pdf"
            DSGE.save_plot(p, fn)
        end
    end
end


text:
using Test
using WavePropBase
using WavePropBase.Geometry
using WavePropBase.Interpolation
using StaticArrays

@testset "HyperRectangle tests" begin
    low_corner  = SVector(0.0,0.0)
    high_corner = SVector(1.0,2.0)
    mid         = (low_corner + high_corner)/2
    rec = HyperRectangle(low_corner,high_corner)
    @test mid == center(rec)
    @test high_corner ∈ rec
    @test low_corner ∈ rec
    @test mid ∈ rec
    @test !in(high_corner + SVector(1,1),rec)
    rec1, rec2 = split(rec)
    @test low_corner ∈ rec1
    @test high_corner ∈ rec2
    @test !(low_corner ∈ rec2)
    @test !(high_corner ∈ rec1)
    @test diameter(rec) == sqrt(1^2 + 2^2)
    @test radius(rec) == sqrt(1^2 + 2^2)/2
    # bbox
    pts = SVector{2,Float64}[]
    for x=-1:0.1:1
        for y=-1:0.1:1
            push!(pts,SVector(x,y))
        end
    end
    @test HyperRectangle(pts)      == HyperRectangle(SVector(-1.,-1),SVector(1,1.))
    @test HyperRectangle(pts,true) == HyperRectangle(SVector(-1.,-1),SVector(1,1.))
    pts = SVector{2,Float64}[]
    for x=-1:0.1:1
        for y=-1:0.1:2
            push!(pts,SVector(x,y))
        end
    end
    @test HyperRectangle(pts)      == HyperRectangle(SVector(-1.,-1),SVector(1,2.))
    @test HyperRectangle(pts,true) == HyperRectangle(SVector(-1.5,-1),SVector(3/2,2.))
    rec1 = HyperRectangle(SVector(0,0),SVector(1,1))
    rec2 = HyperRectangle(SVector(2,0),SVector(3,1))
    @test distance(rec1,rec2) ≈ 1
    x  = SVector(0.5,0.5)
    @test distance(x,rec1) ≈ distance(rec1,x) ≈ 0
    x  = SVector(0.0,2.0)
    @test distance(x,rec1) ≈ distance(rec1,x) ≈ 1
    x  = SVector(2.0,2.0)
    @test distance(x,rec1) ≈ distance(rec1,x) ≈ √2
    rec2 = HyperRectangle(SVector(2,2),SVector(3,3))
    distance(rec1,rec2) ≈ sqrt(2)
end

@testset "HyperCube tests" begin
    low_corner  = SVector(1.5,-2.3)
    side = 2.2
    high_corner = low_corner .+ side
    mid         = (low_corner + high_corner)/2
    rec = HyperCube(low_corner,side)
    @test mid == center(rec)
    @test high_corner ∈ rec
    @test low_corner ∈ rec
    @test mid ∈ rec
    @test !in(high_corner + SVector(1,1),rec)
    rec1, rec2 = split(rec)
    @test low_corner ∈ rec1
    @test high_corner ∈ rec2
    @test !(low_corner ∈ rec2)
    @test !(high_corner ∈ rec1)
    @test diameter(rec) == side*sqrt(length(low_corner))
    @test radius(rec) == side*sqrt(length(low_corner))/2
    #=
    # bbox
    pts = SVector{2,Float64}[]
    for x=-1:0.1:1
        for y=-1:0.1:1
            push!(pts,SVector(x,y))
        end
    end
    @test HyperRectangle(pts)      == HyperRectangle(SVector(-1.,-1),SVector(1,1.))
    @test HyperRectangle(pts,true) == HyperRectangle(SVector(-1.,-1),SVector(1,1.))
    pts = SVector{2,Float64}[]
    for x=-1:0.1:1
        for y=-1:0.1:2
            push!(pts,SVector(x,y))
        end
    end
    @test HyperRectangle(pts)      == HyperRectangle(SVector(-1.,-1),SVector(1,2.))
    @test HyperRectangle(pts,true) == HyperRectangle(SVector(-1.5,-1),SVector(3/2,2.))
    =#
    rec1 = HyperCube(SVector(0,0),2)
    rec2 = HyperCube(SVector(3,0),1)
    @test distance(rec1,rec2) ≈ 1
    x  = SVector(0.5,0.5)
    @test distance(x,rec1) ≈ distance(rec1,x) ≈ 0
    x  = SVector(0.0,3.0)
    @test distance(x,rec1) ≈ distance(rec1,x) ≈ 1
    x  = SVector(3.0,3.0)
    @test distance(x,rec1) ≈ distance(rec1,x) ≈ √2
    rec2 = HyperCube(SVector(3,3),1)
    distance(rec1,rec2) ≈ sqrt(2)
end

max_stars_repo_name: WaveProp/WavePropBase

max_issues_repo_name: WaveProp/WavePropBase

max_forks_repo_name: WaveProp/WavePropBase


text:
# Pulse model for Bandit Arm

struct Pulse <: BanditArmBase
    step::Int64             # For tracking internally
    period::Int64           # Period of the pulse
    changePoint::Int64      # The point at which the signal goes high
    highDuration::Int64     # TIme to stay HIGH
    isRestless::Bool        # Is this a restless arm

    function Pulse( _period, _cp, _hd, _ir = true )
        new(
            0,
            _period,
            _cp,
            _hd,
            _ir
        )
    end
end

function pull!( arm::Pulse )
    arm.step    = arm.step + 1
    if arm.step > arm.period
        arm.step = 1
    end

    return ((arm.step>=arm.changePoint)&&(arm.step<=arm.changePoint+arm.highDuration)) ? 1 : 0;
end

function tick!( arm::Pulse )
    if arm.isRestless
        arm.step = arm.step + 1
        if arm.step > arm.period
            arm.step = 1
        end
    end
    nothing
end

function reset!( arm::Pulse )
    arm.step    = 0
    nothing
end