using GeneralizedGenerated
using Random: GLOBAL_RNG
using TupleVectors: chainvec

export rand
EmptyNTtype = NamedTuple{(),Tuple{}} where T<:Tuple

function Base.rand(rng::AbstractRNG, d::ConditionalModel, N::Int)
    r = chainvec(rand(rng, d), N)
    for j in 2:N
        @inbounds r[j] = rand(rng, d)
    end
    return r
end

Base.rand(d::ConditionalModel, N::Int) = rand(GLOBAL_RNG, d, N)

@inline function Base.rand(rng::AbstractRNG, c::ConditionalModel)
    m = Model(c)
    return _rand(getmoduletypencoding(m), m, argvals(c))(rng)
end

@inline function Base.rand(m::ConditionalModel)
    rand(GLOBAL_RNG, m)
end

@inline function Base.rand(rng::AbstractRNG, m::Model)
    return _rand(getmoduletypencoding(m), m, NamedTuple())(rng)
end

rand(m::Model) = rand(GLOBAL_RNG, m)

sourceRand(m::Model) = sourceRand()(m)
sourceRand(jd::ConditionalModel) = sourceRand(jd.model)

export sourceRand
function sourceRand()
    function(_m::Model)
        proc(_m, st::Assign)  = :($(st.x) = $(st.rhs))
        proc(_m, st::Sample)  = :($(st.x) = rand(_rng, $(st.rhs)))
        proc(_m, st::Return)  = :(return $(st.rhs))
        proc(_m, st::LineNumber) = nothing

        vals = map(x -> Expr(:(=), x,x),parameters(_m))

        wrap(kernel) = @q begin
            _rng -> begin
                $kernel
                $(Expr(:tuple, vals...))
            end
        end

        buildSource(_m, proc, wrap) |> MacroTools.flatten
    end
end

@gg function _rand(M::Type{<:TypeLevel}, _m::Model, _args)
    body = type2model(_m) |> sourceRand() |> loadvals(_args, NamedTuple())
    @under_global from_type(_unwrap_type(M)) @q let M
        $body
    end
end

@gg function _rand(M::Type{<:TypeLevel}, _m::Model, _args::NamedTuple{()})
    body = type2model(_m) |> sourceRand()
    @under_global from_type(_unwrap_type(M)) @q let M
        $body
    end
end


function read_atoms(fn::IO)
    sane_readline() = strip(strip_split(readline(fn), ":")[1], '\'')
    nspecies = parse(Int, sane_readline())
    species_counter = 0

    reading_atoms = false
    atname = :nothing
    natoms = 0
    atcounter = 0
    positions = Point3{Float64}[]
    bfcmts = Vec3{Float64}[]
    atoms = Tuple{Symbol,Point3{Float64},Vec3{Float64}}[]
    while species_counter < nspecies
        l = sane_readline()
        if !reading_atoms
            atname = Symbol(split(l, ".")[1])
            natoms = parse(Int, sane_readline())
            reading_atoms = true
        elseif atcounter < natoms
            sline = split(l)
            push!(positions, Point3{Float64}(parse.(Float64, sline[1:3])))
            push!(bfcmts, Vec3{Float64}(parse.(Float64, sline[4:end])))
            atcounter += 1
        elseif isempty(l)
            for (p, mag) in zip(positions, bfcmts)
                push!(atoms, (atname, p, mag))
            end
            atname          = :nothing
            reading_atoms   = false
            natoms          = 0
            atcounter       = 0
            positions       = Point3{Float64}[]
            bfcmts          = Vec3{Float64}[]
            species_counter += 1
        end
    end
    return atoms
end

function elk_parse_DFTU(atoms::Vector{Atom}, blocknames_flaglines::Dict{Symbol,Any})
    lines = pop!(blocknames_flaglines, Symbol("dft+u"))
    t_l = split(lines[1])
    species_start = 3
    if length(t_l) == 1
        dftu_type = parse(Int, t_l[1])
        inpdftu   = parse(Int, lines[2])
    else
        dftu_type, inpdftu = parse.(Int, t_l)
        species_start = 2
    end
    species_names = name.(unique(atoms))
    for line in lines[species_start:end]
        sline         = split(line)
        species_id, l = parse.(Int, sline[1:2])
        U, J0         = parse.(Float64, sline[3:4])
        for at in filter(x -> x.name == species_names[species_id], atoms)
            at.dftu = DFTU{Float64}(; U = U, J0 = J0, l = l)
        end
    end
    return dftu_type, inpdftu
end

"""
    elk_read_calculation(filename; execs=[Exec(exec="elk")], run=true, structure_name="noname")

Reads an Elk calculation file. The `ELK_EXEC` inside execs gets used to find which flags are allowed in this calculation file, and convert the read values to the correct Types.
Returns a `Calculation{Elk}` and the `Structure` that is found in the calculation.
"""
function elk_read_calculation(fn::String; exec = Exec(; exec = "elk"), run = true,
                              structure_name = "noname")
    blocknames_flaglines = Dict{Symbol,Any}()
    atoms = Atom[]
    calculations = Calculation[]
    dir, file = splitdir(fn)

    #TODO it's a bit messy that the atoms block is not handled like this, not sure if we want to clean this up
    #blocks that require specific parsing
    special_blocks = (:avec, :tasks, Symbol("dft+u"))

    open(fn, "r") do f
        sane_readline() = strip_split(readline(f), ':')[1]
        while !eof(f)
            line = sane_readline()
            isempty(line) || line[1] == '!' && continue
            if !isempty(line)
                blockname = Symbol(line)
                info = elk_block_info(blockname)
                if blockname == :atoms
                    atoms = read_atoms(f)
                elseif blockname == :wannierExtra
                    @warn "Please supply the .win file instead of wannierExtra block for extracting the Wannier Calculation."

                    wflags, wdata, ab, cb, proj_block = wan_read_calculation(f)
                    continue
                    # push!(calculations, Calculation{Wannier90}("wannier", dir, wflags, wdata, execs, run))
                elseif info !== nothing
                    blocklines = String[]
                    line = sane_readline()
                    while !isempty(line)
                        push!(blocklines,
                              strip(replace(strip(split(line, ':')[1]), "'" => ""), '.'))
                        line = sane_readline()
                    end
                    if !in(blockname, special_blocks)
                        blocknames_flaglines[blockname] = parse_block_flags(blockname,
                                                                            collect(Iterators.flatten(split.(filter(x -> !isempty(x),
                                                                                blocklines)))))
                    else
                        blocknames_flaglines[blockname] = filter(x -> !isempty(x),
                                                                 blocklines)
                    end
                end
            end
        end
    end
    # for i in
    #cell
    scale = haskey(blocknames_flaglines, :scale) ? last(pop!(blocknames_flaglines[:scale])) : 1.0
    cell  = uconvert.(Ang, Mat3(reshape(scale .* parse.(Float64, collect(Iterators.flatten(split.(pop!(blocknames_flaglines, :avec))))), 3, 3) .* 1a₀))'

    newats = [Atom(; name = x[1], element = x[1].element, position_cryst = x[2],
                   position_cart = cell * x[2], magnetization = x[3]) for x in atoms]
    #structure #TODO make positions be in lattices coordinates
    structure = Structure(cell, newats)

    #different tasks
    tasks_     = pop!(blocknames_flaglines, :tasks)
    should_run = map(x -> !occursin("!", x), tasks_)
    tasks      = map(x -> strip(strip(x, '!')), tasks_)

    #The wan tasks can prbably be skipped completely and added depending on the wannier calculation
    wan_tasks = String[]
    for (x, run) in zip(tasks, should_run)
        if x ∈ ["601", "602", "603", "604", "605"]
            push!(wan_tasks, x)
        else
            push!(calculations,
                  calculation_from_task(x, blocknames_flaglines, dir, execs, run))
        end
    end
    if haskey(blocknames_flaglines, :wannier) && !isempty(wan_tasks)
        flags = pop!(blocknames_flaglines, :wannier)
        # flags[:elk2wan_tasks] = wan_tasks
        push!(calculations,
              Calculation{Elk}(; name = "elk2wannier", flags = flags,
                               exec = exec, run = true))
    end
    for f in (:ngrid, :vkloff, :plot1d, :plot2d, :plot3d)
        haskey(blocknames_flaglines, f) && pop!(blocknames_flaglines, f)
    end

    #elk2wannier

    flags = Dict{Symbol,Any}()

    if haskey(blocknames_flaglines, Symbol("dft+u"))
        flags[:dftu], flags[:inpdftu] = elk_parse_DFTU(structure.atoms,
                                                       blocknames_flaglines)
    end

    for (k, v) in blocknames_flaglines
        merge!(flags, v)
    end

    for i in calculations
        if isempty(i.flags)
            i.flags = deepcopy(flags)
        end
    end
    return calculations, structure
end

function find_data(flags, blocknames_flaglines)
    data = InputData[]
    for f in flags
        if haskey(blocknames_flaglines, f)
            push!(data, InputData(f, :nothing, blocknames_flaglines[f]))
        end
    end
    return data
end

function calculation_from_task(task, blocknames_flaglines, dir, exec, run)
    if task ∈ ["0", "1", "2", "3", "5"]
        data = find_data((:ngridk, :vkloff), blocknames_flaglines)
    elseif task ∈ ["20", "21"]
        data = find_data((:plot1d, :plot2d, :plot3d), blocknames_flaglines)
    end
    return Calculation{Elk}(; name = task, data = data, exec = exec, run = run)
end

function parse(::Type{UnitRange{Int}}, l::AbstractString)
    for delim in [":", ",", "-"]
        if occursin(delim, l)
            return UnitRange(parse.(Int, split(l, delim))...)
        end
    end
end

function parse_block_flags(blockname::Symbol, lines::Vector{<:AbstractString})
    flags = Dict{Symbol,Any}()
    i = 1
    totlen = length(lines)
    for flaginfo in elk_block_info(blockname).flags
        flagname = flaginfo.name
        i > length(lines) && break
        typ = eltype(flaginfo)
        if typ == String || typ == Vector{String}
            flags[flagname] = lines[i]
            i += 1
        elseif typ == Symbol
            flags[flagname] = Symbol(lines[i])
            i += 1
        else
            eltyp = eltype(typ)
            # The assumption is that the infinite repeating values are at the end of a block
            if flaginfo == last(elk_block_info(blockname).flags)
                typelen = length(typ)
                parsed = parse.(eltyp, lines[i:end])
                flags[flagname] = typ[]
                @assert length(parsed) % typelen == 0
                for j in 1:typelen:length(parsed)
                    push!(flags[flagname], typ(parsed[j:j+typelen-1]...))
                end
            elseif typ <: UnitRange
                flags[flagname] = parse(typ, lines[i])
                i += 1
            else
                typelen = length(typ)
                parsed = parse.(eltyp, lines[i:i+typelen-1])
                flags[flagname] = length(parsed) == 1 ? parsed[1] : typ(parsed)
                i += typelen
            end
        end
    end
    return flags
end

function elk_write_structure(f, structure)
    write(f, "atoms\n\t")
    species = unique(map(x -> x.name, structure.atoms))
    nspecies = length(species)
    write(f, "$nspecies\n")
    for s in species
        write(f, "\t'$s.in'\n")
        ats = filter(x -> x.name == s, structure.atoms)
        write(f, "\t$(length(ats))\n")
        for a in ats
            write(f, "\t")
            for i in a.position_cryst
                write(f, "$(round(i, digits=8)) ")
            end
            for i in magnetization(a)
                write(f, "$i ")
            end
            write(f, "\n")
        end
        write(f, "\n")
    end
    write(f, "\n")
    write(f, "avec\n\t")

    for i in eachindex(structure.cell')
        c = ustrip(uconvert(a₀, structure.cell[i]))
        if i % 3 == 0
            write(f, "$c\n\t")
        else
            write(f, "$c ")
        end
    end
    return write(f, "\n")
end

function construct_flag_blocks(calculations::Vector{Calculation{Elk}})
    all_flags = Dict{Symbol,Any}()
    for i in calculations
        merge!(all_flags, i.flags)
        for d in data(i)
            merge!(all_flags, d.data)
        end
    end
    block_flags = Dict{Symbol,Any}()

    for b in ELK_CONTROLBLOCKS
        bvals = Any[]
        for f in b.flags
            if haskey(all_flags, f.name)
                push!(bvals, pop!(all_flags, f.name))
            end
        end
        if !isempty(bvals)
            block_flags[b.name] = bvals
        end
    end
    return block_flags
end

function elk_write_flag_val(f, val)
    write(f, "\t")
    if isa(val, UnitRange)
        write(f, "$(first(val))-$(last(val))")
    elseif isa(val, AbstractVector)
        for i in val
            write(f, "$i ")
        end
    else
        if isa(val, AbstractString)
            write(f, "'$val'")
        else
            write(f, "$val")
        end
    end
    return write(f, "\n")
end

function elk_write_DFTU(f::IO, structure::DFC.Structure, dftu_vals::Vector)
    write(f, "dft+u\n")
    for v in dftu_vals
        elk_write_flag_val(f, v)
    end
    species = unique(structure.atoms)
    for (i, s) in enumerate(species)
        write(f, "\t$i $(string(Elk, dftu(s)))\n")
    end
    return write(f, "\n")
end

function save(calculations::Vector{Calculation{Elk}}, structure::Structure, dir)
    tasks = map(x -> x.run ? "$(x.name)" : "!$(x.name)",
                filter(x -> x.name != "elk2wannier", calculations))
    elk2wan_calculation = getfirst(x -> x.name == "elk2wannier", calculations)
    if elk2wan_calculation != nothing && haskey(elk2wan_calculation.flags, :elk2wan_tasks)
        append!(tasks, pop!(elk2wan_calculation.flags, :elk2wan_tasks))
    end
    flags = construct_flag_blocks(calculations)
    open(joinpath(dir, "elk.in"), "w") do f
        write(f, "tasks\n")
        for t in tasks
            write(f, "\t$t\n")
        end
        write(f, "\n")

        elk_write_structure(f, structure)

        haskey(flags, Symbol("dft+u")) &&
            elk_write_DFTU(f, structure, pop!(flags, Symbol("dft+u")))

        #remaining generic blocks
        for (k, v) in flags
            write(f, "$k\n")
            for v_ in values(v)
                if isa(v_, Vector{<:AbstractVector})
                    for v__ in v_
                        elk_write_flag_val(f, v__)
                    end
                else
                    elk_write_flag_val(f, v_)
                end
            end
            write(f, "\n")
        end
    end
end


export main_reported, filter_to_relevant_cases, check_and_correct_missing_days, are_days_missing

function main_reported(fig_name, file_name, days_col, data_col, ylabel_R, ylabel_N, aggregation::String="country", pop_days::Int=0)
    aggregation = lowercase(aggregation)
    @assert aggregation ∈ ["county", "state", "country"] "$aggregation is not a valid aggregation level."
    # download full COVID-19 data set from the Robert Koch Institut (RKI)
    # this file is ≈50mb

    if !isfile(file_name)
        run(`curl https://www.arcgis.com/sharing/rest/content/items/f10774f1c63e40168479a1feb6c7ca74/data --output $file_name`)
    end

    # generation-time
    k_gen = 4

    # pre-processing
    df = CSV.File(file_name) |> DataFrame

    df_filter = filter_to_relevant_cases(df)

    # aggregate by country, state, or county
    if aggregation == "country"
        df_temp = aggregate_Germany(df_filter, days_col)
    elseif aggregation == "state"
        df_temp = aggregate_state(df_filter, days_col)
    elseif aggregation == "county"
        throw(error("Not yet supported."))
    end
    # convert date column
    df_temp = convert_date_column(df_temp, days_col)
    # standardize column
    df_cases = get_reference_data(df_temp, days_col = days_col, data_col = data_col, kind = "cases")
    # account for the fact that cases were not reported properly before March 2020
    df_cases = filter_to_relevant_days(df_cases, Date("2020-03-01"))
    # check for missing days
    df_cases = check_and_correct_missing_days(df_cases, kind="cases")
    # do the math
    N, R = compute_and_plot(df_cases, fig_name, k_gen, ylabel_R, ylabel_N, "reported number of cases", pop_days)
    # return values
    N, R, df_cases
end

function filter_to_relevant_days(df::DataFrame, day::Date)
    df[df.days .>= day, :]
end

function filter_to_relevant_days(d::Dict, day::Date)
    dd = deepcopy(d)
    for (name, df) in dd
        dd[name] = filter_to_relevant_days(df, day)
    end
    dd
end

function filter_to_relevant_cases(df::DataFrame)
    # we need to filter out all new cases that are *not* -1
    df_filtered = df[df.NeuerFall .!= -1, :]
end

function aggregate_Germany(df_filtered, days_col)
    df = combine(groupby(df_filtered, days_col), :AnzahlFall => sum) |> sort
    Dict("Germany"=>df)
end

function aggregate_state(df_filtered, days_col)
    gdf = groupby(df_filtered, :Bundesland)
    cases_dict = Dict{String, DataFrame}()

    umlaute = Dict("ä"=>"ae", "ö"=>"oe", "ü"=>"ue", "ß"=>"ss")
    for (i, df_state) in enumerate(gdf)
        # get name of current state
        state_name = unique(gdf[i].Bundesland) |> first
        for (key, val) in umlaute
            state_name = replace(state_name, key=>val)
        end
        # aggregate state data by days
        gdf_state = groupby(df_state, days_col)
        # sum over the number of cases & sort
        cases = combine(gdf_state, :AnzahlFall => sum) |> sort
        cases_dict[state_name] = cases
    end
    cases_dict
end

function convert_date_column(df::DataFrame, days_col)
    df[!, days_col] = Date.(df[!, days_col], Dates.DateFormat("yyyy/mm/dd H:M:S"))
    df
end

function convert_date_column(df::Dict, days_col)
    for (key, value) in df
        df[key] = convert_date_column(value, days_col)
    end
    df
end

function check_and_correct_missing_days(df::DataFrame; name::String="", kind::String)
    d_start, d_end = first(df.days), last(df.days)
    if are_days_missing(df.days, d_start, d_end)
        display("Filling non-existing values for $name.")
        d_ref = d_start:Day(1):d_end
        y_ref = fill_non_existing_entries(d_ref, df.days, df.cases)
        return DataFrame(Dict("days" => d_ref, "$kind" => y_ref))
    end
    return df
end

function check_and_correct_missing_days(d::Dict; kind::String)
    for (key, value) in d
        d[key] = check_and_correct_missing_days(value, name=key, kind=kind)
    end
    d
end

function are_days_missing(days::Vector{Date}, start_day::Date, end_day::Date)
    days_ref = start_day:Day(1):end_day
    days_ref != days
end


module EUregmort

using CSV, DataFrames, PyCall, PyPlot, Statistics

const shpreader = PyNULL()
const ccrs = PyNULL()

function __init__()
	copy!(shpreader, pyimport("cartopy.io.shapereader"))
	copy!(ccrs, pyimport("cartopy.crs"))
end

export nuts2ids, caprop_regcmp, caprop_regsexplot, fourp, fivep, caprop_mapplot, meanrate
mainpath = normpath(@__DIR__, "..")
datapath = normpath(mainpath, "data")
shpdatapath = normpath(datapath, "NUTS_2013_03M_SH", "data")
ycdr = CSV.read(normpath(datapath, "hlth_cd_ycdr2.csv"); missingstring = ":", rows_for_type_detect = 200)
nuts = CSV.read(normpath(datapath, "NUTS_AT_2013.csv"); rows_for_type_detect = 200)
nids = convert(Array, collect(skipmissing(nuts[:NUTS_ID])))
ageseq = ["Y_LT1"; "Y1-4"; "Y5-9"; "Y10-14"; "Y15-19";
	"Y20-24"; "Y25-29"; "Y30-34"; "Y35-39"; "Y40-44"; "Y45-49";
	"Y50-54"; "Y55-59"; "Y60-64"; "Y65-69"; "Y70-74"; "Y75-79";
	"Y80-84"; "Y85-89"; "Y90-94"; "Y_GE95"]
agesplitter(age) = split(age,  ['Y'; 'M'; '_'; '-'], keepempty = false)

function agealias(age)
	aspl = agesplitter(age)
	if aspl[1] == "TOTAL"
		alias = "alla åldrar"
	elseif aspl[1] == "LT1"
		alias = "0"
	elseif startswith(aspl[1], "LT")
		alias = "0\u2013$(parse(aspl[end][3:end])-1)"
	elseif startswith(aspl[1], "GE")
		alias = "$(aspl[end][3:end])\u2013"
	else
		alias = "$(aspl[1])\u2013$(aspl[end])"
	end

	if (startswith(age, "YM") || startswith(age, "M"))
		return "$alias genomsnitt över åldrar"
	else
		return alias
	end
end

sexaliases = Dict("F" => "kvinnor", "M" => "män", "T" => "alla")
causealiases = Dict("pop" => L"$10^5$ personår", "A-R_V-Y" => "alla orsaker")

function causealias(ca)
	if ca in keys(causealiases)
		return causealiases[ca]
	else
		return ca
	end
end

function nuts2ids(ctry)
	n2ids = map((x)->Symbol(x), filter((x)->(length(x)==4 && startswith(x, ctry)), nids))
	filter!((x) -> x in names(ycdr), n2ids)
end

dfarrmatch(col, arr) = map((x) -> in(x, arr), Vector(col))

function meanrate(sage, eage)
	sexframes = Dict()
	sind = indexin([sage], ageseq)[1]
	eind = indexin([eage], ageseq)[1]
	agelist = ageseq[sind:eind]
	if sage == "Y_LT1"
		if eage == "Y_GE95"
			agestr = "MTOTAL"
		else
			agestr = "YM_LT$(parse(agesplitter(eage)[end])+1)"
		end
	else
		if eage == "Y_GE95"
			agestr = "YM_GE$(agesplitter(sage)[1])"
		else
			agestr = "YM$(agesplitter(sage)[1])-$(agesplitter(eage)[end])"
		end
	end
	for sex in keys(sexaliases)
		ycdr_sub = ycdr[((ycdr[:sex].==sex) & (dfarrmatch(ycdr[:age], agelist))),
			 [4;6:size(ycdr)[2]]]
		ycdr_sub_agg = aggregate(ycdr_sub, :icd10, mean)
		aggregs = rename((x)->Symbol(replace(string(x), "_mean", "")),
			ycdr_sub_agg[:, 2:end])
		aggl = size(ycdr_sub_agg)[1]
		metaframe = DataFrame(unit = fill("RT", aggl), sex = fill(sex, aggl),
			age = fill(agestr, aggl),
			icd10 = ycdr_sub_agg[:icd10], time_geo = fill(2013, aggl))
		sexframes[sex] = hcat(metaframe, aggregs)
	end
	vcat(values(sexframes)...)
end

function ycdr_regcmp(nutsregs, sex, age, icd10, time_geo, inframe)
	inframe_sub = DataFrame(unit = inframe[:unit], sex = inframe[:sex],
		age = inframe[:age], icd10 = inframe[:icd10],
		time_geo = inframe[:time_geo])
	for reg in nutsregs
		inframe_sub[reg] = inframe[reg]
	end
	inframe_long = stack(inframe_sub, nutsregs)
	inframe_long[((inframe_long[:sex].==sex) .& (inframe_long[:age].==age)
		.& (inframe_long[:icd10].==icd10) .& (inframe_long[:time_geo].==time_geo)), :]
end

function caprop_regcmp(nutsregs, sex, age, ca1, ca2, time_geo, inframe)
	ca1frame = ycdr_regcmp(nutsregs, sex, age, ca1, time_geo, inframe)
	if ca2 == "pop"
		capropframe = ca1frame
	else
		ca2frame = ycdr_regcmp(nutsregs, sex, age, ca2, time_geo, inframe)
		capropframe = copy(ca1frame)
		capropframe[:value] = ca1frame[:value]./ca2frame[:value]
	end
	capropframe
end

function caprop_regsexplot(nutsregs, age, ca1; ca2 = "A-R_V-Y", time_geo = 2013,
	inframe = ycdr)
	fframe = caprop_regcmp(nutsregs, "F", age, ca1, ca2, time_geo, inframe)
	mframe = caprop_regcmp(nutsregs, "M", age, ca1, ca2, time_geo, inframe)
	scatter(fframe[:value], mframe[:value], alpha=0.3)
	for i in 1:size(nutsregs)[1]
		text(fframe[:value][i], mframe[:value][i], fframe[:variable][i],
			ha = "center", va = "center")
	end
	grid(1)
	xlabel(sexaliases["F"])
	ylabel(sexaliases["M"])
	title(*("Döda $(time_geo) $(causealias(ca1))/$(causealias(ca2)) ",
		"$(agealias(age))"))
end

perc_round(value) = replace("$(round(value; digits=4))", "." => ",")
fourp(prop) =
	[
	Dict("col" => "lightyellow", "value" => quantile(prop, 1/4));
	Dict("col" => "yellow", "value" => quantile(prop, 2/4));
	Dict("col" => "tomato", "value" => quantile(prop, 3/4));
	Dict("col" => "red", "value" => quantile(prop, 1))
	]
fivep(prop) =
	[
	Dict("col" => "lightyellow", "value" => quantile(prop, 1/5));
	Dict("col" => "yellow", "value" => quantile(prop, 2/5));
	Dict("col" => "orange", "value" => quantile(prop, 3/5));
	Dict("col" => "tomato", "value" => quantile(prop, 4/5));
	Dict("col" => "red", "value" => quantile(prop, 1))
	]

function caprop_mapplot(nutsregs, sex, age, ca1; ca2 = "A-R_V-Y", time_geo = 2013,
		percfunc = fivep,
		shapefname = normpath(shpdatapath, "NUTS_RG_03M_2013_3034.shp"),
		inframe = ycdr)
	pframe = caprop_regcmp(nutsregs, sex, age, ca1, ca2, time_geo, inframe)
	region_shp = shpreader[:Reader](shapefname)
	regstrings = map((x)->string(x), nutsregs)
	proj = ccrs[:LambertConformal](central_longitude = 10, central_latitude = 52,
		standard_parallels = (35,65), false_easting = 4000000,
		false_northing = 2800000, globe = ccrs[:Globe](ellipse = "GRS80"))
	ax = plt[:axes](projection = proj)
	prop = pframe[:value]
	propdict = Dict(zip(regstrings, prop))
	quantiles = percfunc(prop)
	boundlist = []
	fcolor = "blue"
	for region_rec in region_shp[:records]()
		nid = region_rec[:attributes]["NUTS_ID"]
		if nid in regstrings
			boundlist = vcat(boundlist, region_rec[:bounds])
			xmean = mean([boundlist[end][1];
				boundlist[end][3]])
			ymean = mean([boundlist[end][2];
				boundlist[end][4]])
			for quantile in quantiles
				if propdict[nid] <= quantile["value"]
					fcolor = quantile["col"]
					break
				end
			end
			ax[:add_geometries](region_rec[:geometry], proj,
				edgecolor = "black", facecolor = fcolor)
			ax[:annotate](nid, (xmean, ymean), ha = "center")
		end
	end
	xminimum = minimum([bound[1] for bound in boundlist])
	xmaximum = maximum([bound[3] for bound in boundlist])
	yminimum = minimum([bound[2] for bound in boundlist])
	ymaximum = maximum([bound[4] for bound in boundlist])
	ax[:set_xlim](xminimum, xmaximum)
	ax[:set_ylim](yminimum, ymaximum)
	percpatches = []
	perclabels = []
	for (i, quantile) in enumerate(quantiles)
		percpatch = matplotlib[:patches][:Rectangle]((0, 0), 1, 1,
			facecolor = quantile["col"])
		percpatches = vcat(percpatches, percpatch)
		if i == 1
			perclabel = *("\u2265", perc_round(minimum(prop)),
				"\n\u2264", perc_round(quantile["value"]))
		else
			perclabel = *("\u2264", perc_round(quantile["value"]))
		end
		perclabels = vcat(perclabels, perclabel)
	end
	legend(percpatches, perclabels, loc = "upper left",
		framealpha = 0.75, bbox_to_anchor=(1,1))
	title(*("Döda $(time_geo) $(causealias(ca1))/$(causealias(ca2)) ",
		"$(sexaliases[sex]) $(agealias(age))"))
	subplots_adjust(right = 0.8)
end

end # module


#!/usr/bin/env julia

using Weber

version = v"0.0.1"
sid,skip = @read_args("Duration Discrimination ($version).")

#======================================================================
Experiment Setup
======================================================================#
# 2AFC duration discrimination. Each stimulus consists of 2 short tones
# separated by given duration. The task is to indicate which of 2 intervals
# is longer.

const atten_dB = 30 # adjust to calibrate levels
const n_trials_per_block = 60
const n_blocks = 6

const tone_freq = 1kHz
const tone_length = 10ms
const standard_length = 100ms # from start of tone 1 to start of tone 2

const SOA = 900ms # how long to wait to starting play stimulus 2
const response_delay = 300ms # how long to wait before asking for a response
const trial_spacing = 500ms # how long to wait between trials

const adapter = levitt_adapter(first_delta=2,down=3,up=1,
                               big=0.5,little=0.05,min_delta=0,
                               max_delta=(SOA/2-tone_length)/tone_length)

#======================================================================#

experiment = Experiment(
  skip=skip,
  columns = [
    :sid => sid,
    :condition => "example",
    :version => version,
    :standard => standard_length,
    :block
  ]
)

isresponse(e) = iskeydown(e,key"p") || iskeydown(e,key"q")

beep = attenuate(ramp(tone(tone_freq,tone_length)),atten_dB)
beep_beep(len) = [beep;silence(len - tone_length);beep]

function trial_stimuli(first_longer,delta)
  standard = beep_beep(standard_length)
  signal = beep_beep(standard_length*(1+delta))
  if first_longer
    [signal;silence(SOA-duration(signal));standard]
  else
    [standard;silence(SOA-duration(standard));signal]
  end
end

function duration2AFC(adapter;keys...)
  first_longer = rand(Bool)
  resp = response(adapter,key"q" => "first_longer",key"p" => "second_longer",
                  correct=(first_longer? "first_longer" : "second_longer");
                  keys...)

  stim() = trial_stimuli(first_longer,delta(adapter))
  [moment(trial_spacing,play,stim),show_cross(),
   moment(SOA + standard_length + response_delay,display,
          "Was the first [Q] or second sound [P] longer?"),
   resp,await_response(isresponse)]
end

setup(experiment) do
  addbreak(moment(record,"start"),
           moment(play,@> tone(1kHz,1s) ramp attenuate(atten_dB)),
           moment(1s))

  addbreak(instruct("""

      On each trial, you will hear two sounds. Indicate which of the two sounds
  you heard as longer. Hit 'Q' if the first sound was longer, and 'P' if the
  second was longer.

  """))

  for block in 1:n_blocks
    @addtrials let adapter = adapter
      for trial in 1:n_trials_per_block
        addtrial(duration2AFC(adapter,block=block))
      end
      # define this string during experiment setup
      # when we know what block we're on...
      heading = "Block $block of $n_blocks"

      function report_threshold()
        mean,sd = estimate(adapter)
        thresh = round(mean,3)*standard_length
        thresh_sd = round(sd,3)*standard_length

        # define this string during run time when we know
        # what the threshold estimate is.
        "$heading: Threshold $thresh (SD: $thresh_sd)\n"*
        "Hit spacebar to continue..."
      end

      addbreak(moment(display,report_threshold,clean_whitespace=false),
               await_response(iskeydown(key":space:")))
    end
  end
end

# play a test tone, to verify levels
run(experiment)


############################################################
# joData # #################################################
############################################################

# Authors: Philipp Witte (pwitte@eos.ubc.ca), Henryk Modzelewski (hmodzelewski@eos.ubc.ca)
# Date: January 2017

export joData, joDataException, subsample, joData_to_SeisBlock

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

# structure for seismic data as an abstract vector
mutable struct joData{vDT<:Number} <: joAbstractLinearOperator{vDT,vDT}
	name::String
	m::Integer
	n::Integer
	nsrc::Integer
	geometry::Geometry
	data
end

mutable struct joDataException <: Exception
	msg :: String
end

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

## outer constructors

"""
    joData
	    name::String
        m::Integer
        n::Integer
        nsrc::Integer
        geometry::Geometry
        data

Abstract vector for seismic data. This vector-like structure contains the geometry and data for either\\
receiver data (shot records) or source data (wavelets).

Constructors
============

Construct vector from `Geometry` structure and cell array of shot records or wavelets. The `data` keyword\\
can also be a single (non-cell) array, in which case the data is the same for all source positions:

    joData(geometry, data)

Construct vector for observed data from `SeisIO.SeisBlock`. `segy_depth_key` is the `SeisIO` keyword \\
that contains the receiver depth coordinate:

    joData(SeisIO.SeisBlock; segy_depth_key="RecGroupElevation")

Construct vector for observed data from out-of-core data container of type `SeisIO.SeisCon`:

    joData(SeisIO.SeisCon; segy_depth_key="RecGroupElevation")

Examples
========

(1) Construct data vector from `Geometry` structure and a cell array of shot records:

    dobs = joData(rec_geometry, shot_records)

(2) Construct data vector for a seismic wavelet (can be either cell arrays of individual\\
wavelets or a single wavelet as an array):

    q = joData(src_geometry, wavelet)

(3) Construct data vector from `SeisIO.SeisBlock` object:

    using SeisIO
    seis_block = segy_read("test_file.segy")
	dobs = joData(seis_block; segy_depth_key="RecGroupElevation")

(4) Construct out-of-core data vector from `SeisIO.SeisCon` object (for large SEG-Y files):

    using SeisIO
    seis_container = segy_scan("/path/to/data/directory","filenames",["GroupX","GroupY","RecGroupElevation","SourceDepth","dt"])
    dobs = joData(seis_container; segy_depth_key="RecGroupElevation")

"""
function joData(geometry::Geometry,data::Array; vDT::DataType=Float32)
	vDT == Float32 || throw(joDataException("Domain type not supported"))

	# length of vector
	n = 1
	if typeof(geometry) == GeometryOOC
		nsrc = length(geometry.container)
		m = sum(geometry.nsamples)
	else
		nsrc = length(geometry.xloc)
		m = 0
		for j=1:nsrc
			m += length(geometry.xloc[j])*geometry.nt[j]
		end
	end
	dataCell = Array{Array}(nsrc)
	for j=1:nsrc
		dataCell[j] = data
	end
	return joData{Float32}("Seismic data vector",m,n,nsrc,geometry,dataCell)
end

# constructor if data is passed as a cell array
function joData(geometry::Geometry,data::Union{Array{Any,1},Array{Array,1}}; vDT::DataType=Float32)
	vDT == Float32 || throw(joDataException("Domain and range types not supported"))

	# length of vector
	if typeof(geometry) == GeometryOOC
		nsrc = length(geometry.container)
		m = sum(geometry.nsamples)
	else
		nsrc = length(geometry.xloc)
		m = 0
		for j=1:nsrc
			m += length(geometry.xloc[j])*geometry.nt[j]
		end
	end
	n = 1
	return joData{Float32}("Seismic data vector",m,n,nsrc,geometry,data)
end

############################################################
# constructors from SEGY files or out-of-core containers

# contructor for in-core data container
function joData(data::SeisIO.SeisBlock; segy_depth_key="RecGroupElevation", vDT::DataType=Float32)
	vDT == Float32 || throw(joDataException("Domain and range types not supported"))

	# length of data vector
	src = get_header(data,"FieldRecord")
 	nsrc = length(unique(src))

	numTraces = get_header(data,"TraceNumber")[end] - get_header(data,"TraceNumber")[1] + 1
	numSamples = get_header(data,"ns")[end]
	m = numTraces*numSamples
	n = 1

	# extract geometry from data container
	geometry = Geometry(data; key="receiver", segy_depth_key=segy_depth_key)

	# fill data vector with pointers to data location
	dataCell = Array{Array}(nsrc)
	full_data = convert(Array{Float32,2},data.data)
	for j=1:nsrc
		traces = find(src .== unique(src)[j])
		dataCell[j] = full_data[:,traces]
	end

	return joData{Float32}("Julia seismic data container",m,n,nsrc,geometry,dataCell)
end

# contructor for in-core data container and given geometry
function joData(geometry::Geometry, data::SeisIO.SeisBlock; vDT::DataType=Float32)
	vDT == Float32 || throw(joDataException("Domain and range types not supported"))

	# length of data vector
	src = get_header(data,"FieldRecord")
 	nsrc = length(unique(src))

	numTraces = get_header(data,"TraceNumber")[end] - get_header(data,"TraceNumber")[1] + 1
	numSamples = get_header(data,"ns")[end]
	m = numTraces*numSamples
	n = 1

	# fill data vector with pointers to data location
	dataCell = Array{Array}(nsrc)
	for j=1:nsrc
		traces = find(src .== unique(src)[j])
		dataCell[j] = convert(Array{Float32,2},data.data[:,traces])
	end

	return joData{Float32}("Julia seismic data container",m,n,nsrc,geometry,dataCell)
end

# contructor for out-of-core data container from single container
function joData(data::SeisIO.SeisCon; segy_depth_key="RecGroupElevation", vDT::DataType=Float32)
	vDT == Float32 || throw(joDataException("Domain and range types not supported"))

	# length of data vector
	nsrc = length(data)
	numTraces = 0
	for j=1:nsrc
		numTraces += Int((data.blocks[j].endbyte - data.blocks[j].startbyte)/(240 + data.ns*4))
	end
	m = numTraces*data.ns
	n = 1

	# extract geometry from data container
	geometry = Geometry(data; key="receiver", segy_depth_key=segy_depth_key)

	# fill data vector with pointers to data location
	dataCell = Array{SeisIO.SeisCon}(nsrc)
	for j=1:nsrc
		dataCell[j] = split(data,j)
	end

	return joData{Float32}("Julia seismic data container",m,n,nsrc,geometry,dataCell)
end

# contructor for single out-of-core data container and given geometry
function joData(geometry::Geometry, data::SeisIO.SeisCon; vDT::DataType=Float32)
	vDT == Float32 || throw(joDataException("Domain and range types not supported"))

	# length of data vector
	nsrc = length(data)
	numTraces = 0
	for j=1:nsrc
		numTraces += Int((data.blocks[j].endbyte - data.blocks[j].startbyte)/(240 + data.ns*4))
	end
	m = numTraces*data.ns
	n = 1

	# fill data vector with pointers to data location
	dataCell = Array{SeisIO.SeisCon}(nsrc)
	for j=1:nsrc
		dataCell[j] = split(data,j)
	end

	return joData{Float32}("Julia seismic data container",m,n,nsrc,geometry,dataCell)
end

# contructor for out-of-core data container from cell array of containers
function joData(data::Array{SeisIO.SeisCon,1}; segy_depth_key="RecGroupElevation", vDT::DataType=Float32)
	vDT == Float32 || throw(joDataException("Domain and range types not supported"))

	# length of data vector
	nsrc = length(data)
	numTraces = 0
	for j=1:nsrc
		numTraces += Int((data[j].blocks[1].endbyte - data[j].blocks[1].startbyte)/(240 + data.ns*4))
	end
	m = numTraces*data.ns
	n = 1

	# extract geometry from data container
	geometry = Geometry(data; key="receiver", segy_depth_key=segy_depth_key)

	# fill data vector with pointers to data location
	dataCell = Array{SeisIO.SeisCon}(nsrc)
	for j=1:nsrc
		dataCell[j] = data[j]
	end

	return joData{Float32}("Julia seismic data container",m,n,nsrc,geometry,dataCell)
end

# contructor for out-of-core data container from cell array of containers and given geometry
function joData(geometry::Geometry, data::Array{SeisIO.SeisCon}; vDT::DataType=Float32)
	vDT == Float32 || throw(joDataException("Domain and range types not supported"))

	# length of data vector
	nsrc = length(data)
	numTraces = 0
	for j=1:nsrc
		numTraces += Int((data[j].blocks[1].endbyte - data[j].blocks[1].startbyte)/(240 + data[j].ns*4))
	end
	m = numTraces*data[1].ns
	n = 1

	# fill data vector with pointers to data location
	dataCell = Array{SeisIO.SeisCon}(nsrc)
	for j=1:nsrc
		dataCell[j] = data[j]
	end

	return joData{Float32}("Julia seismic data container",m,n,nsrc,geometry,dataCell)
end

############################################################
## overloaded Base functions

# conj(jo)
conj{vDT}(a::joData{vDT}) =
	joData{vDT}("conj("*a.name*")",a.m,a.n,a.nsrc,a.geometry,a.data)

# transpose(jo)
transpose{vDT}(a::joData{vDT}) =
	joData{vDT}(""*a.name*".'",a.n,a.m,a.nsrc,a.geometry,a.data)

# ctranspose(jo)
ctranspose{vDT}(a::joData{vDT}) =
	joData{vDT}(""*a.name*"'",a.n,a.m,a.nsrc,a.geometry,a.data)

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

# +(joData, joData)
function +{avDT,bvDT}(a::joData{avDT}, b::joData{bvDT})
	size(a) == size(b) || throw(joDataException("dimension mismatch"))
	compareGeometry(a.geometry, b.geometry) == 1 || throw(joDataException("geometry mismatch"))
	c = deepcopy(a)
	c.data = a.data + b.data
	return c
end

# -(joData, joData)
function -{avDT,bvDT}(a::joData{avDT}, b::joData{bvDT})
	size(a) == size(b) || throw(joDataException("dimension mismatch"))
	compareGeometry(a.geometry, b.geometry) == 1 || throw(joDataException("geometry mismatch"))
	c = deepcopy(a)
	c.data = a.data - b.data
	return c
end

# *(joData, number)
function *{avDT}(a::joData{avDT},b::Number)
	c = deepcopy(a)
	c.data = c.data*b
	return c
end

# *(number, joData)
function *{bvDT}(a::Number,b::joData{bvDT})
	c = deepcopy(b)
	c.data = a*c.data
	return c
end

# *(joLinearFunction, joData)
function *{ADDT,ARDT,avDT}(A::joLinearFunction{ADDT,ARDT},v::joData{avDT})
	A.n == size(v,1) || throw(joDataException("shape mismatch"))
	jo_check_type_match(ADDT,avDT,join(["DDT for *(joLinearFunction,joData):",A.name,typeof(A),avDT]," / "))
	V = A.fop(v)
	jo_check_type_match(ARDT,eltype(V),join(["RDT from *(joLinearFunction,joData):",A.name,typeof(A),eltype(V)]," / "))
	return V
end

# *(joLinearOperator, joData)
function *{ADDT,ARDT,avDT}(A::joLinearOperator{ADDT,ARDT},v::joData{avDT})
	A.n == size(v,1) || throw(joDataException("shape mismatch"))
	jo_check_type_match(ADDT,avDT,join(["DDT for *(joLinearFunction,joData):",A.name,typeof(A),avDT]," / "))
	V = A.fop(v)
	jo_check_type_match(ARDT,eltype(V),join(["RDT from *(joLinearFunction,joData):",A.name,typeof(A),eltype(V)]," / "))
	return V
end

# /(joData, number)
function -{avDT}(a::joData{avDT},b::Number)
	c = deepcopy(a)
	c.data = c.data/b
	return c
end

# /(number, joData)
function -{bvDT}(a::Number,b::joData{bvDT})
	c = deepcopy(b)
	c.data = a./c.data
	return c
end

# vcat
function vcat{avDT,bvDT}(a::joData{avDT},b::joData{bvDT})
	typeof(a.geometry) == typeof(b.geometry) || throw(joDataException("Geometry type mismatch"))
	m = a.m + b.m
	n = 1
	nsrc = a.nsrc + b.nsrc

	if typeof(a.data) == Array{SeisIO.SeisCon,1} && typeof(b.data) == Array{SeisIO.SeisCon,1}
		data = Array{SeisIO.SeisCon}(nsrc)
	else
		data = Array{Array}(nsrc)
	end

	dt = Array{Any}(nsrc)
	nt = Array{Any}(nsrc)
	t = Array{Any}(nsrc)
	if typeof(data) == Array{SeisIO.SeisCon,1}
		nsamples = Array{Any}(nsrc)
	else
		xloc = Array{Any}(nsrc)
		yloc = Array{Any}(nsrc)
		zloc = Array{Any}(nsrc)
	end

	# Merge data sets and geometries
	for j=1:a.nsrc
		data[j] = a.data[j]
		if typeof(data) == Array{SeisIO.SeisCon,1}
			nsamples[j] = a.geometry.nsamples[j]
		else
			xloc[j] = a.geometry.xloc[j]
			yloc[j] = a.geometry.yloc[j]
			zloc[j] = a.geometry.zloc[j]
		end
		dt[j] = a.geometry.dt[j]
		nt[j] = a.geometry.nt[j]
		t[j] = a.geometry.t[j]
	end
	for j=a.nsrc+1:nsrc
		data[j] = b.data[j-a.nsrc]
		if typeof(data) == Array{SeisIO.SeisCon,1}
			nsamples[j] = b.geometry.nsamples[j-a.nsrc]
		else
			xloc[j] = b.geometry.xloc[j-a.nsrc]
			yloc[j] = b.geometry.yloc[j-a.nsrc]
			zloc[j] = b.geometry.zloc[j-a.nsrc]
		end
		dt[j] = b.geometry.dt[j-a.nsrc]
		nt[j] = b.geometry.nt[j-a.nsrc]
		t[j] = b.geometry.t[j-a.nsrc]
	end

	if typeof(data) == Array{SeisIO.SeisCon,1}
		geometry = GeometryOOC(data,dt,nt,t,nsamples,a.geometry.key,a.geometry.segy_depth_key)
	else
		geometry = Geometry(xloc,yloc,zloc,dt,nt,t)
	end
	nvDT = promote_type(avDT,bvDT)
	return joData{nvDT}(a.name,m,n,nsrc,geometry,data)
end

# dot product
function dot{avDT,bvDT}(a::joData{avDT}, b::joData{bvDT})
# Dot product for data containers
	size(a) == size(b) || throw(joDataException("dimension mismatch"))
	compareGeometry(a.geometry, b.geometry) == 1 || throw(joDataException("geometry mismatch"))
	dotprod = 0f0
	for j=1:a.nsrc
		dotprod += dot(vec(a.data[j]),vec(b.data[j]))
	end
	return dotprod
end

# norm
function norm{avDT}(a::joData{avDT}; p=2)
	x = 0.f0
	for j=1:a.nsrc
		x += norm(vec(a.data[j]),p)^p
	end
	return x^(1.f0/p)
end

# abs
function abs{avDT}(a::joData{avDT})
	b = deepcopy(a)
	for j=1:a.nsrc
		b.data[j] = abs(a.data[j])
	end
	return b
end

# Subsample data container
"""
    subsample(x,source_numbers)

Subsample seismic data vectors or matrix-free linear operators and extract the entries that correspond\\
to the shot positions defined by `source_numbers`. Works for inputs of type `joData`, `joModeling`, \\
`joProjection`, `joJacobian`, `Geometry`, `joRHS`, `joPDE`, `joPDEfull`.

Examples
========

(1) Extract 2 shots from `joData` vector:

    dsub = subsample(dobs,[1,2])

(2) Extract geometry for shot location 100:

    geometry_sub = subsample(dobs.geometry,100)

(3) Extract Jacobian for shots 10 and 20:

    Jsub = subsample(J,[10,20])

"""
function subsample{avDT}(a::joData{avDT},srcnum)
	geometry = subsample(a.geometry,srcnum)		# Geometry of subsampled data container
	return joData(geometry,a.data[srcnum];vDT=avDT)
end

getindex(x::joData, a) = subsample(x,a)

# Create SeisBlock from joData container to write to file
function joData_to_SeisBlock{avDT}(d::joData{avDT},q::joData{avDT}; source_depth_key="SourceSurfaceElevation", receiver_depth_key="RecGroupElevation")

	typeof(d.geometry) == GeometryOOC && (d.geometry = Geometry(d.geometry))
	typeof(q.geometry) == GeometryOOC && (q.geometry = Geometry(q.geometry))

	blocks = Array{Any}(d.nsrc)
	count = 0
	for j=1:d.nsrc

		# create SeisBlock
		blocks[j] = SeisBlock(d.data[j])
		numTraces = size(d.data[j],2)
		traceNumbers = convert(Array{Integer,1},count+1:count+numTraces)

		# set headers
		set_header!(blocks[j], "GroupX", convert(Array{Integer,1},round.(d.geometry.xloc[j]*1f3)))
		if length(d.geometry.yloc[j]) == 1
			set_header!(blocks[j], "GroupY", Int(round(d.geometry.yloc[j]*1f3)))
		else
			set_header!(blocks[j], "GroupY", convert(Array{Integer,1},round.(d.geometry.yloc[j]*1f3)))
		end
		set_header!(blocks[j], receiver_depth_key, convert(Array{Integer,1},round.(d.geometry.zloc[j]*1f3)))
		set_header!(blocks[j], "SourceX", Int(round.(q.geometry.xloc[j]*1f3)))
		set_header!(blocks[j], "SourceY", Int(round.(q.geometry.yloc[j]*1f3)))
		set_header!(blocks[j], source_depth_key, Int(round.(q.geometry.zloc[j]*1f3)))

		set_header!(blocks[j], "dt", Int(d.geometry.dt[j]*1f3))
		set_header!(blocks[j], "FieldRecord",j)
		set_header!(blocks[j], "TraceNumWithinLine", traceNumbers)
		set_header!(blocks[j], "TraceNumWithinFile", traceNumbers)
		set_header!(blocks[j], "TraceNumber", traceNumbers)
		set_header!(blocks[j], "ElevationScalar", -1000)
		set_header!(blocks[j], "RecSourceScalar", -1000)
		count += numTraces
	end

	# merge into single block
	fullblock = blocks[1]
	for j=2:d.nsrc
		fullblock = merge(fullblock,blocks[j])
		blocks[j] = []
	end
	return fullblock
end

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

# Overload base function for SeisIO objects

vec(x::SeisIO.SeisCon) = vec(x[1].data)
norm(x::SeisIO.IBMFloat32; kwargs...) = norm(convert(Float32,x); kwargs...)
dot(x::SeisIO.IBMFloat32, y::SeisIO.IBMFloat32) = dot(convert(Float32,x), convert(Float32,y))
dot(x::SeisIO.IBMFloat32, y::Float32) = dot(convert(Float32,x), y)
dot(x::Float32, y::SeisIO.IBMFloat32) = dot(x, convert(Float32,y))

# binary operations return dense arrays
+(x::SeisIO.SeisCon, y::SeisIO.SeisCon) = +(x[1].data,y[1].data)
+(x::SeisIO.IBMFloat32, y::SeisIO.IBMFloat32) = +(convert(Float32,x),convert(Float32,y))
+(x::SeisIO.IBMFloat32, y::Float32) = +(convert(Float32,x),y)
+(x::Float32, y::SeisIO.IBMFloat32) = +(x,convert(Float32,y))

-(x::SeisIO.SeisCon, y::SeisIO.SeisCon) = -(x[1].data,y[1].data)
-(x::SeisIO.IBMFloat32, y::SeisIO.IBMFloat32) = -(convert(Float32,x),convert(Float32,y))
-(x::SeisIO.IBMFloat32, y::Float32) = -(convert(Float32,x),y)
-(x::Float32, y::SeisIO.IBMFloat32) = -(x,convert(Float32,y))

+(a::Number, x::SeisIO.SeisCon) = +(a,x[1].data)
+(x::SeisIO.SeisCon, a::Number) = +(x[1].data,a)

-(a::Number, x::SeisIO.SeisCon) = -(a,x[1].data)
-(x::SeisIO.SeisCon, a::Number) = -(x[1].data,a)

*(a::Number, x::SeisIO.SeisCon) = *(a,x[1].data)
*(x::SeisIO.SeisCon, a::Number) = *(x[1].data,a)

/(a::Number, x::SeisIO.SeisCon) = /(a,x[1].data)
/(x::SeisIO.SeisCon, a::Number) = /(x[1].data,a)


# Mon 31 Dec 2012 01:39:13 PM EST
#
# NeHe Tut 12 - Rotate a textured cube
#
# Q - quit
# Up/Down - increase/decrease rotation of each cube about it's own x-axis
# Left/Right - increase/decrease rotation of each cube about it's own y-axis

# load necessary GLUT/OpenGL routines

global OpenGLver="1.0"
using OpenGL
using GLUT

# initialize variables

global window

global xrot = 0
global yrot = 0

global box  = 0
global top  = 0

boxcol      = [1.0 0.0 0.0;
               1.0 0.5 0.0;
               1.0 1.0 0.0;
               0.0 1.0 0.0;
               0.0 1.0 1.0]

topcol      = [0.5 0.0  0.0;
               0.5 0.25 0.0;
               0.5 0.5  0.0;
               0.0 0.5  0.0;
               0.0 0.5  0.5]

global tex  = Array(Uint32,1) # generating 1 texture

width       = 640
height      = 480

# load textures from images

function LoadGLTextures()
    global tex

    img, w, h = glimread(expanduser("~/.julia/GLUT/Examples/NeHe/tut12/cube.bmp"))

    glGenTextures(1,tex)
    glBindTexture(GL_TEXTURE_2D,tex[1])
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR_MIPMAP_NEAREST)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexImage2D(GL_TEXTURE_2D, 0, 3, w, h, 0, GL_RGB, GL_UNSIGNED_BYTE, img)
    gluBuild2DMipmaps(GL_TEXTURE_2D, 3, w, h, GL_RGB, GL_UNSIGNED_BYTE, img)
end

# function to init OpenGL context

function initGL(w::Integer,h::Integer)
    global box
    global top

    glViewport(0,0,w,h)
    LoadGLTextures()
    glClearColor(0.0, 0.0, 0.0, 0.5)
    glClearDepth(1.0)
    glDepthFunc(GL_LEQUAL)
    glEnable(GL_DEPTH_TEST)
    glShadeModel(GL_SMOOTH)
    glHint(GL_PERSPECTIVE_CORRECTION_HINT, GL_NICEST)

    # enable texture mapping
    glEnable(GL_TEXTURE_2D)

    #enable simple lighting
    glEnable(GL_LIGHT0)
    glEnable(GL_LIGHTING)
    glEnable(GL_COLOR_MATERIAL)

    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()

    gluPerspective(45.0,w/h,0.1,100.0)

    glMatrixMode(GL_MODELVIEW)

    # build the display lists

    box = glGenLists(2)

    glNewList(box, GL_COMPILE)
        glBegin(GL_QUADS)
            # Bottom Face
            glTexCoord(0.0, 1.0)
            glVertex(-1.0, -1.0, -1.0)
            glTexCoord(1.0, 1.0)
            glVertex(1.0, -1.0, -1.0)
            glTexCoord(1.0, 0.0)
            glVertex(1.0, -1.0,  1.0)
            glTexCoord(0.0, 0.0)
            glVertex(-1.0, -1.0,  1.0)

            # Front Face
            glTexCoord(1.0, 0.0)
            glVertex(-1.0, -1.0,  1.0)
            glTexCoord(0.0, 0.0)
            glVertex(1.0, -1.0,  1.0)
            glTexCoord(0.0, 1.0)
            glVertex(1.0,  1.0,  1.0)
            glTexCoord(1.0, 1.0)
            glVertex(-1.0,  1.0,  1.0)

            # Back Face
            glTexCoord(0.0, 0.0)
            glVertex(-1.0, -1.0, -1.0)
            glTexCoord(0.0, 1.0)
            glVertex(-1.0,  1.0, -1.0)
            glTexCoord(1.0, 1.0)
            glVertex(1.0,  1.0, -1.0)
            glTexCoord(1.0, 0.0)
            glVertex(1.0, -1.0, -1.0)

            # Right Face
            glTexCoord(0.0, 0.0)
            glVertex(1.0, -1.0, -1.0)
            glTexCoord(0.0, 1.0)
            glVertex(1.0,  1.0, -1.0)
            glTexCoord(1.0, 1.0)
            glVertex(1.0,  1.0,  1.0)
            glTexCoord(1.0, 0.0)
            glVertex(1.0, -1.0,  1.0)

            # Left Face
            glTexCoord(1.0, 0.0)
            glVertex(-1.0, -1.0, -1.0)
            glTexCoord(0.0, 0.0)
            glVertex(-1.0, -1.0,  1.0)
            glTexCoord(0.0, 1.0)
            glVertex(-1.0,  1.0,  1.0)
            glTexCoord(1.0, 1.0)
            glVertex(-1.0,  1.0, -1.0)
        glEnd()
    glEndList()

    top = uint32(box+1)

    glNewList(top, GL_COMPILE)
        glBegin(GL_QUADS)
            # Top Face
            glTexCoord(1.0, 1.0)
            glVertex(-1.0, 1.0, -1.0)
            glTexCoord(1.0, 0.0)
            glVertex(-1.0, 1.0,  1.0)
            glTexCoord(0.0, 0.0)
            glVertex(1.0, 1.0,  1.0)
            glTexCoord(0.0, 1.0)
            glVertex(1.0, 1.0, -1.0)
        glEnd()
    glEndList()
end

# prepare Julia equivalents of C callbacks that are typically used in GLUT code

function ReSizeGLScene(w::Int32,h::Int32)
    if h == 0
        h = 1
    end

    glViewport(0,0,w,h)

    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()

    gluPerspective(45.0,w/h,0.1,100.0)

    glMatrixMode(GL_MODELVIEW)

    return nothing
end

_ReSizeGLScene = cfunction(ReSizeGLScene, Void, (Int32, Int32))

function DrawGLScene()
    global tex
    global xrot
    global yrot
    global box
    global top

    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

    glBindTexture(GL_TEXTURE_2D, tex[1])

    for yloop = 1:5
	      for xloop = 1:yloop
            glLoadIdentity()

            glTranslate(1.4+2.8xloop-1.4yloop, ((6.0-yloop)*2.4)-7.0, -20.0)

            glRotate(45.0-(2.0yloop)+xrot, 1.0, 0.0, 0.0)
            glRotate(45.0+yrot, 0.0, 1.0, 0.0)

            glColor(boxcol[yloop,:])
            glCallList(box)

            glColor(topcol[yloop,:])
            glCallList(top)
        end
    end

    glutSwapBuffers()

    return nothing
end

_DrawGLScene = cfunction(DrawGLScene, Void, ())

function keyPressed(the_key::Char,x::Int32,y::Int32)
    if the_key == int('q')
        glutDestroyWindow(window)
    end

    return nothing # keyPressed returns "void" in C. this is a workaround for Julia's "automatically return the value of the last expression in a function" behavior.
end

_keyPressed = cfunction(keyPressed, Void, (Char, Int32, Int32))

function specialKeyPressed(the_key::Int32,x::Int32,y::Int32)
    global xrot
    global yrot

    if the_key == GLUT_KEY_UP
        xrot -=0.2
    elseif the_key == GLUT_KEY_DOWN
        xrot +=0.2
    elseif the_key == GLUT_KEY_LEFT
        yrot -=0.2
    elseif the_key == GLUT_KEY_RIGHT
        yrot +=0.2
    end

    return nothing # specialKeyPressed returns "void" in C-GLUT. this is a workaround for Julia's "automatically return the value of the last expression in a function" behavior.
end

_specialKeyPressed = cfunction(specialKeyPressed, Void, (Int32, Int32, Int32))

# run GLUT routines

glutInit()
glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_ALPHA | GLUT_DEPTH)
glutInitWindowSize(width, height)
glutInitWindowPosition(0, 0)

window = glutCreateWindow("NeHe Tut 12")

glutDisplayFunc(_DrawGLScene)
glutFullScreen()

glutIdleFunc(_DrawGLScene)
glutReshapeFunc(_ReSizeGLScene)
glutKeyboardFunc(_keyPressed)
glutSpecialFunc(_specialKeyPressed)

initGL(width, height)

glutMainLoop()


# Polar formulation
#
#
#
include("caches.jl")

"""
    PolarForm{T, IT, VT, MT}

Takes as input a [`PS.PowerNetwork`](@ref) network and
implement the polar formulation model associated to this network.
The structure `PolarForm` stores the topology of the network, as
well as the complete indexing used in the polar formulation.

A `PolarForm` structure can be instantiated both on the host `CPU()`
or directly on the device `CUDADevice()`.
"""
struct PolarForm{T, IT, VT, MT} <: AbstractFormulation where {T, IT, VT, MT}
    network::PS.PowerNetwork
    device::KA.Device
    # bounds
    x_min::VT
    x_max::VT
    u_min::VT
    u_max::VT
    # costs
    costs_coefficients::MT
    # Constant loads
    active_load::VT
    reactive_load::VT
    # Indexing of the PV, PQ and slack buses
    indexing::IndexingCache{IT}
    # struct
    topology::NetworkTopology{IT, VT}
    # Jacobian indexing
    mapx::IT
    mapu::IT
    # Hessian structures and indexing
    hessianstructure::HessianStructure
end

include("kernels.jl")
include("derivatives.jl")
include("Constraints/constraints.jl")
include("powerflow.jl")
include("objective.jl")

function PolarForm(pf::PS.PowerNetwork, device::KA.Device)
    if isa(device, KA.CPU)
        IT = Vector{Int}
        VT = Vector{Float64}
        M = SparseMatrixCSC
        AT = Array
    elseif isa(device, KA.CUDADevice)
        IT = CUDA.CuVector{Int64}
        VT = CUDA.CuVector{Float64}
        M = CUSPARSE.CuSparseMatrixCSR
        AT = CUDA.CuArray
    end

    nbus = PS.get(pf, PS.NumberOfBuses())
    npv = PS.get(pf, PS.NumberOfPVBuses())
    npq = PS.get(pf, PS.NumberOfPQBuses())
    nref = PS.get(pf, PS.NumberOfSlackBuses())
    ngens = PS.get(pf, PS.NumberOfGenerators())

    topology = NetworkTopology{IT, VT}(pf)
    # Get coefficients penalizing the generation of the generators
    coefs = convert(AT{Float64, 2}, PS.get_costs_coefficients(pf))
    # Move load to the target device
    pload = convert(VT, PS.get(pf, PS.ActiveLoad()))
    qload = convert(VT, PS.get(pf, PS.ReactiveLoad()))

    # Move the indexing to the target device
    idx_gen = PS.get(pf, PS.GeneratorIndexes())
    idx_ref = PS.get(pf, PS.SlackIndexes())
    idx_pv = PS.get(pf, PS.PVIndexes())
    idx_pq = PS.get(pf, PS.PQIndexes())
    # Build-up reverse index for performance
    pv_to_gen = similar(idx_pv)
    ref_to_gen = similar(idx_ref)
    ## We assume here that the indexing of generators is the same
    ## as in MATPOWER
    for i in 1:ngens
        bus = idx_gen[i]
        i_pv = findfirst(isequal(bus), idx_pv)
        if !isnothing(i_pv)
            pv_to_gen[i_pv] = i
        else
            i_ref = findfirst(isequal(bus), idx_ref)
            if !isnothing(i_ref)
                ref_to_gen[i_ref] = i
            end
        end
    end

    gidx_gen = convert(IT, idx_gen)
    gidx_ref = convert(IT, idx_ref)
    gidx_pv = convert(IT, idx_pv)
    gidx_pq = convert(IT, idx_pq)
    gref_to_gen = convert(IT, ref_to_gen)
    gpv_to_gen = convert(IT, pv_to_gen)

    # Bounds
    ## Get bounds on active power
    p_min, p_max = PS.bounds(pf, PS.Generators(), PS.ActivePower())
    p_min = convert(VT, p_min)
    p_max = convert(VT, p_max)
    ## Get bounds on voltage magnitude
    v_min, v_max = PS.bounds(pf, PS.Buses(), PS.VoltageMagnitude())
    v_min = convert(VT, v_min)
    v_max = convert(VT, v_max)
    ## Instantiate arrays
    nᵤ = nref + 2*npv
    nₓ = npv + 2*npq
    u_min = convert(VT, fill(-Inf, nᵤ))
    u_max = convert(VT, fill( Inf, nᵤ))
    x_min = convert(VT, fill(-Inf, nₓ))
    x_max = convert(VT, fill( Inf, nₓ))
    ## Bounds on v_pq
    x_min[npv+npq+1:end] .= v_min[gidx_pq]
    x_max[npv+npq+1:end] .= v_max[gidx_pq]
    ## Bounds on v_pv
    u_min[nref+1:nref+npv] .= v_min[gidx_pv]
    u_max[nref+1:nref+npv] .= v_max[gidx_pv]
    ## Bounds on v_ref
    u_min[1:nref] .= v_min[gidx_ref]
    u_max[1:nref] .= v_max[gidx_ref]
    ## Bounds on p_pv
    u_min[nref+npv+1:nref+2*npv] .= p_min[gpv_to_gen]
    u_max[nref+npv+1:nref+2*npv] .= p_max[gpv_to_gen]

    indexing = IndexingCache(gidx_pv, gidx_pq, gidx_ref, gidx_gen, gpv_to_gen, gref_to_gen)
    mappv = [i + nbus for i in idx_pv]
    mappq = [i + nbus for i in idx_pq]
    # Ordering for x is (θ_pv, θ_pq, v_pq)
    statemap = vcat(mappv, mappq, idx_pq)
    controlmap = vcat(idx_ref, idx_pv, idx_pv .+ nbus)
    hessianmap = vcat(statemap, idx_ref, idx_pv, idx_pv .+ 2*nbus)
    hessianstructure = HessianStructure(IT(hessianmap))
    return PolarForm{Float64, IT, VT, AT{Float64,  2}}(
        pf, device,
        x_min, x_max, u_min, u_max,
        coefs, pload, qload,
        indexing,
        topology,
        statemap, controlmap,
        hessianstructure
    )
end
# Convenient constructor
PolarForm(datafile::String, device) = PolarForm(PS.PowerNetwork(datafile), device)

array_type(polar::PolarForm) = array_type(polar.device)

# Getters
function get(polar::PolarForm, ::NumberOfState)
    npv = PS.get(polar.network, PS.NumberOfPVBuses())
    npq = PS.get(polar.network, PS.NumberOfPQBuses())
    return 2*npq + npv
end

function get(polar::PolarForm, ::NumberOfControl)
    nref = PS.get(polar.network, PS.NumberOfSlackBuses())
    npv = PS.get(polar.network, PS.NumberOfPVBuses())
    return nref + 2*npv
end

function get(polar::PolarForm, attr::PS.AbstractNetworkAttribute)
    return get(polar.network, attr)
end

# Setters
function setvalues!(polar::PolarForm, ::PS.ActiveLoad, values)
    @assert length(polar.active_load) == length(values)
    copyto!(polar.active_load, values)
end
function setvalues!(polar::PolarForm, ::PS.ReactiveLoad, values)
    @assert length(polar.reactive_load) == length(values)
    copyto!(polar.reactive_load, values)
end

## Bounds
function bounds(polar::PolarForm{T, IT, VT, MT}, ::State) where {T, IT, VT, MT}
    return polar.x_min, polar.x_max
end

function bounds(polar::PolarForm{T, IT, VT, MT}, ::Control) where {T, IT, VT, MT}
    return polar.u_min, polar.u_max
end

# Initial position
function initial(polar::PolarForm{T, IT, VT, MT}, X::Union{State,Control}) where {T, IT, VT, MT}
    # Load data from PowerNetwork
    pbus = real.(polar.network.sbus) |> VT
    qbus = imag.(polar.network.sbus) |> VT
    vmag = abs.(polar.network.vbus) |> VT
    vang = angle.(polar.network.vbus) |> VT

    npv = get(polar, PS.NumberOfPVBuses())
    npq = get(polar, PS.NumberOfPQBuses())
    nref = get(polar, PS.NumberOfSlackBuses())

    if isa(X, State)
        # build vector x
        dimension = get(polar, NumberOfState())
        x = xzeros(VT, dimension)
        x[1:npv] = vang[polar.network.pv]
        x[npv+1:npv+npq] = vang[polar.network.pq]
        x[npv+npq+1:end] = vmag[polar.network.pq]
        return x
    elseif isa(X, Control)
        dimension = get(polar, NumberOfControl())
        u = xzeros(VT, dimension)
        u[1:nref] = vmag[polar.network.ref]
        # u is equal to active power of generator (Pᵍ)
        # As P = Pᵍ - Pˡ , we get
        u[nref + 1:nref + npv] = vmag[polar.network.pv]
        u[nref + npv + 1:nref + 2*npv] = pbus[polar.network.pv] + polar.active_load[polar.network.pv]
        return u
    end
end

function get(form::PolarForm{T, IT, VT, MT}, ::PhysicalState) where {T, IT, VT, MT}
    nbus = PS.get(form.network, PS.NumberOfBuses())
    ngen = PS.get(form.network, PS.NumberOfGenerators())
    n_state = get(form, NumberOfState())
    gen2bus = form.indexing.index_generators
    return PolarNetworkState{VT}(nbus, ngen, n_state, gen2bus)
end

function get!(
    polar::PolarForm{T, IT, VT, MT},
    ::State,
    x::AbstractVector,
    buffer::PolarNetworkState
) where {T, IT, VT, MT}
    npv = get(polar, PS.NumberOfPVBuses())
    npq = get(polar, PS.NumberOfPQBuses())
    nref = get(polar, PS.NumberOfSlackBuses())
    # Copy values of vang and vmag into x
    # NB: this leads to 3 memory allocation on the GPU
    #     we use indexing on the CPU, as for some reason
    #     we get better performance than with the indexing on the GPU
    #     stored in the buffer polar.indexing.
    x[1:npv] .= @view buffer.vang[polar.network.pv]
    x[npv+1:npv+npq] .= @view buffer.vang[polar.network.pq]
    x[npv+npq+1:npv+2*npq] .= @view buffer.vmag[polar.network.pq]
end

function get!(
    polar::PolarForm{T, IT, VT, MT},
    ::Control,
    u::AbstractVector,
    buffer::PolarNetworkState,
) where {T, IT, VT, MT}
    npv = get(polar, PS.NumberOfPVBuses())
    npq = get(polar, PS.NumberOfPQBuses())
    nref = get(polar, PS.NumberOfSlackBuses())
    # build vector u
    nᵤ = get(polar, NumberOfControl())
    u[1:nref] .= @view buffer.vmag[polar.network.ref]
    u[nref + 1:nref + npv] .= @view buffer.vmag[polar.network.pv]
    u[nref + npv + 1:nref + 2*npv] .= @view buffer.pg[polar.indexing.index_pv_to_gen]
    return u
end

function init_buffer!(form::PolarForm{T, IT, VT, MT}, buffer::PolarNetworkState) where {T, IT, VT, MT}
    # FIXME: add proper getters in PowerSystem
    pbus = real.(form.network.sbus)
    qbus = imag.(form.network.sbus)
    vmag = abs.(form.network.vbus)
    vang = angle.(form.network.vbus)

    pg = get(form.network, PS.ActivePower())
    qg = get(form.network, PS.ReactivePower())

    copyto!(buffer.vmag, vmag)
    copyto!(buffer.vang, vang)
    copyto!(buffer.pg, pg)
    copyto!(buffer.qg, qg)
    copyto!(buffer.pinj, pbus)
    copyto!(buffer.qinj, qbus)
end

function direct_linear_solver(polar::PolarForm)
    is_cpu = isa(polar.device, KA.CPU)
    if is_cpu
        jac = jacobian_sparsity(polar, power_balance, State())
        return LinearSolvers.DirectSolver(jac)
    else
        # Factorization is not yet supported on the GPU
        return LinearSolvers.DirectSolver(nothing)
    end
end

function build_preconditioner(polar::PolarForm; nblocks=-1)
    jac = jacobian_sparsity(polar, power_balance, State())
    n = size(jac, 1)
    npartitions = if nblocks > 0
        nblocks
    else
        div(n, 32)
    end
    return LinearSolvers.BlockJacobiPreconditioner(jac, npartitions, polar.device)
end

function Base.show(io::IO, polar::PolarForm)
    # Network characteristics
    nbus = PS.get(polar.network, PS.NumberOfBuses())
    npv = PS.get(polar.network, PS.NumberOfPVBuses())
    npq = PS.get(polar.network, PS.NumberOfPQBuses())
    nref = PS.get(polar.network, PS.NumberOfSlackBuses())
    ngen = PS.get(polar.network, PS.NumberOfGenerators())
    nlines = PS.get(polar.network, PS.NumberOfLines())
    # Polar formulation characteristics
    n_states = get(polar, NumberOfState())
    n_controls = get(polar, NumberOfControl())
    print(io,   "Polar formulation model")
    println(io, " (instantiated on device $(polar.device))")
    println(io, "Network characteristics:")
    @printf(io, "    #buses:      %d  (#slack: %d  #PV: %d  #PQ: %d)\n", nbus, nref, npv, npq)
    println(io, "    #generators: ", ngen)
    println(io, "    #lines:      ", nlines)
    println(io, "giving a mathematical formulation with:")
    println(io, "    #controls:   ", n_controls)
    print(io,   "    #states  :   ", n_states)
end


export visualize

"""
visualize(states::AbstractArray; statenames = string.(1:length(states)), quantiles = (0.1, 0.99), aspect = (1,1,1), resolution = (1920, 1080), statistics = false, title = "Field = ")
# Description
Visualize 3D states
# Arguments
- `states`: Array{Array{Float64,3},1}. An array of arrays containing different fields
# Keyword Arguments
- `statenames`: Array{String,1}. An array of stringnames
- `aspect`: Tuple{Int64,Int64,Float64}. Determines aspect ratio of box for volumes
- `resolution`: Resolution of preliminary makie window
- `statistics`: boolean. toggle for displaying statistics
# Return
- `scene`: Scene. A preliminary scene object for manipulation
"""
function visualize(
    states::AbstractArray;
    statenames = string.(1:length(states)),
    units = ["" for i in eachindex(states)],
    aspect = (1, 1, 1),
    resolution = (1920, 1080),
    statistics = false,
    title = "Field = ",
    bins = 300,
)
    # Create scene
    scene, layout = layoutscene(resolution = resolution)
    lscene = layout[2:4, 2:4] = LScene(scene)
    width = round(Int, resolution[1] / 4) # make menu 1/4 of preliminary resolution

    # Create choices and nodes
    stateindex = collect(1:length(states))
    statenode = Node(stateindex[1])

    colorchoices = [:balance, :thermal, :dense, :deep, :curl, :thermometer]
    colornode = Node(colorchoices[1])

    if statistics
        llscene =
            layout[4, 1] = Axis(
                scene,
                xlabel = @lift(statenames[$statenode] * units[$statenode]),
                xlabelcolor = :black,
                ylabel = "pdf",
                ylabelcolor = :black,
                xlabelsize = 40,
                ylabelsize = 40,
                xticklabelsize = 25,
                yticklabelsize = 25,
                xtickcolor = :black,
                ytickcolor = :black,
                xticklabelcolor = :black,
                yticklabelcolor = :black,
            )
        layout[3, 1] = Label(scene, "Statistics", width = width, textsize = 50)
    end

    # x,y,z are for determining the aspect ratio of the box
    if (typeof(aspect) <: Tuple) & (length(aspect) == 3)
        x, y, z = aspect
    else
        x, y, z = size(states[1])
    end

    # Clim sliders
    upperclim_slider =
        Slider(scene, range = range(0, 1, length = 101), startvalue = 0.99)
    upperclim_node = upperclim_slider.value
    lowerclim_slider =
        Slider(scene, range = range(0, 1, length = 101), startvalue = 0.01)
    lowerclim_node = lowerclim_slider.value

    # Lift Nodes
    state = @lift(states[$statenode])
    statename = @lift(statenames[$statenode])
    clims = @lift((
        quantile($state[:], $lowerclim_node),
        quantile($state[:], $upperclim_node),
    ))
    cmap_rgb = @lift(to_colormap($colornode))
    titlename = @lift(title * $statename) # use padding and appropriate centering

    # Statistics
    if statistics
        histogram_node = @lift(histogram($state, bins = bins))
        xs = @lift($histogram_node[1])
        ys = @lift($histogram_node[2])
        pdf = GLMakie.AbstractPlotting.barplot!(
            llscene,
            xs,
            ys,
            color = :red,
            strokecolor = :red,
            strokewidth = 1,
        )
        @lift(GLMakie.AbstractPlotting.xlims!(llscene, extrema($state)))
        @lift(GLMakie.AbstractPlotting.ylims!(
            llscene,
            extrema($histogram_node[2]),
        ))
        vlines!(
            llscene,
            @lift($clims[1]),
            color = :black,
            linewidth = width / 100,
        )
        vlines!(
            llscene,
            @lift($clims[2]),
            color = :black,
            linewidth = width / 100,
        )
    end

    # Volume Plot
    volume!(
        lscene,
        0..x,
        0..y,
        0..z,
        state,
        camera = cam3d!,
        colormap = cmap_rgb,
        colorrange = clims,
    )
    # Camera
    cam = cameracontrols(scene.children[1])
    eyeposition = Float32[2, 2, 1.3]
    lookat = Float32[0.82, 0.82, 0.1]
    # Title
    supertitle =
        layout[1, 2:4] = Label(scene, titlename, textsize = 50, color = :black)

    # Menus
    statemenu = Menu(scene, options = zip(statenames, stateindex))
    on(statemenu.selection) do s
        statenode[] = s
    end

    colormenu = Menu(scene, options = zip(colorchoices, colorchoices))
    on(colormenu.selection) do s
        colornode[] = s
    end
    lowerclim_string = @lift(
        "lower clim quantile = " *
        @sprintf("%0.2f", $lowerclim_node) *
        ", value = " *
        @sprintf("%0.1e", $clims[1])
    )
    upperclim_string = @lift(
        "upper clim quantile = " *
        @sprintf("%0.2f", $upperclim_node) *
        ", value = " *
        @sprintf("%0.1e", $clims[2])
    )
    # depends on makie version, vbox for old, vgrid for new
    layout[2, 1] = vgrid!(
        Label(scene, "State", width = nothing),
        statemenu,
        Label(scene, "Color", width = nothing),
        colormenu,
        Label(scene, lowerclim_string, width = nothing),
        lowerclim_slider,
        Label(scene, upperclim_string, width = nothing),
        upperclim_slider,
    )
    layout[1, 1] = Label(scene, "Menu", width = width, textsize = 50)

    # Modify Axis
    axis = scene.children[1][OldAxis]
    # axis[:names][:axisnames] = ("↓ Zonal [m] ", "Meriodonal [m]↓ ", "Depth [m]↓ ")
    axis[:names][:axisnames] = ("↓", "↓ ", "↓ ")
    axis[:names][:align] =
        ((:left, :center), (:right, :center), (:right, :center))
    # need to adjust size of ticks first and then size of axis names
    axis[:names][:textsize] = (50.0, 50.0, 50.0)
    axis[:ticks][:textsize] = (00.0, 00.0, 00.0)
    # axis[:ticks][:ranges_labels].val # current axis labels
    xticks = collect(range(-0, aspect[1], length = 2))
    yticks = collect(range(-0, aspect[2], length = 6))
    zticks = collect(range(-0, aspect[3], length = 2))
    ticks = (xticks, yticks, zticks)
    axis[:ticks][:ranges] = ticks
    xtickslabels = [@sprintf("%0.1f", (xtick)) for xtick in xticks]
    xtickslabels[end] = "1e6"
    ytickslabels = ["", "south", "", "", "north", ""]
    ztickslabels = [@sprintf("%0.1f", (xtick)) for xtick in xticks]
    labels = (xtickslabels, ytickslabels, ztickslabels)
    axis[:ticks][:labels] = labels

    display(scene)
    # Change the default camera position after the fact
    # note that these change dynamically as the plot is manipulated
    return scene
end

# 2D visualization
function visualize(
    states::Array{Array{S, 2}, 1};
    statenames = string.(1:length(states)),
    units = ["" for i in eachindex(states)],
    aspect = (1, 1, 1),
    resolution = (2412, 1158),
    title = "Zonal and Temporal Average of ",
    xlims = (0, 1),
    ylims = (0, 1),
    bins = 300,
) where {S}
    # Create scene
    scene, layout = layoutscene(resolution = resolution)
    lscene =
        layout[2:4, 2:4] = Axis(
            scene,
            xlabel = "South to North [m]",
            xlabelcolor = :black,
            ylabel = "Depth [m]",
            ylabelcolor = :black,
            xlabelsize = 40,
            ylabelsize = 40,
            xticklabelsize = 25,
            yticklabelsize = 25,
            xtickcolor = :black,
            ytickcolor = :black,
            xticklabelcolor = :black,
            yticklabelcolor = :black,
            titlesize = 50,
        )
    width = round(Int, resolution[1] / 4) # make menu 1/4 of preliminary resolution

    # Create choices and nodes
    stateindex = collect(1:length(states))
    statenode = Node(stateindex[1])

    colorchoices = [:balance, :thermal, :dense, :deep, :curl, :thermometer]
    colornode = Node(colorchoices[1])

    interpolationlabels = ["contour", "heatmap"]
    interpolationchoices = [true, false]
    interpolationnode = Node(interpolationchoices[1])

    # Statistics
    llscene =
        layout[4, 1] = Axis(
            scene,
            xlabel = @lift(statenames[$statenode] * " " * units[$statenode]),
            xlabelcolor = :black,
            ylabel = "pdf",
            ylabelcolor = :black,
            xlabelsize = 40,
            ylabelsize = 40,
            xticklabelsize = 25,
            yticklabelsize = 25,
            xtickcolor = :black,
            ytickcolor = :black,
            xticklabelcolor = :black,
            yticklabelcolor = :black,
        )
    layout[3, 1] = Label(scene, "Statistics", width = width, textsize = 50)

    # Clim sliders
    upperclim_slider =
        Slider(scene, range = range(0, 1, length = 101), startvalue = 0.99)
    upperclim_node = upperclim_slider.value
    lowerclim_slider =
        Slider(scene, range = range(0, 1, length = 101), startvalue = 0.01)
    lowerclim_node = lowerclim_slider.value

    #ylims = @lift(range($lowerval, $upperval, length = $))
    # Lift Nodes
    state = @lift(states[$statenode])
    statename = @lift(statenames[$statenode])
    unit = @lift(units[$statenode])
    oclims = @lift((
        quantile($state[:], $lowerclim_node),
        quantile($state[:], $upperclim_node),
    ))
    cmap_rgb = colornode
    clims = @lift(
        $oclims[1] != $oclims[2] ? (minimum($oclims), maximum($oclims)) :
        (minimum($oclims) - 1, maximum($oclims) + 1)
    )
    xlims = Array(range(xlims[1], xlims[2], length = 4)) #collect(range(xlims[1], xlims[2], length = size(states[1])[1]))
    ylims = Array(range(ylims[1], ylims[2], length = 4)) #@lift(collect(range($lowerval], $upperval, length = size($state)[2])))
    # newrange = @lift(range($lowerval, $upperval, length = 4))
    # lscene.yticks = @lift(Array($newrange))
    titlename = @lift(title * $statename) # use padding and appropriate centering
    layout[1, 2:4] = Label(scene, titlename, textsize = 50)
    # heatmap
    heatmap1 = heatmap!(
        lscene,
        xlims,
        ylims,
        state,
        interpolate = interpolationnode,
        colormap = cmap_rgb,
        colorrange = clims,
    )

    # statistics
    histogram_node = @lift(histogram($state, bins = bins))
    xs = @lift($histogram_node[1])
    ys = @lift($histogram_node[2])
    pdf = GLMakie.AbstractPlotting.barplot!(
        llscene,
        xs,
        ys,
        color = :red,
        strokecolor = :red,
        strokewidth = 1,
    )
    @lift(GLMakie.AbstractPlotting.xlims!(llscene, extrema($state)))
    @lift(GLMakie.AbstractPlotting.ylims!(llscene, extrema($histogram_node[2])))
    vlines!(llscene, @lift($clims[1]), color = :black, linewidth = width / 100)
    vlines!(llscene, @lift($clims[2]), color = :black, linewidth = width / 100)

    # Menus
    statemenu = Menu(scene, options = zip(statenames, stateindex))
    on(statemenu.selection) do s
        statenode[] = s
    end

    colormenu = Menu(scene, options = zip(colorchoices, colorchoices))
    on(colormenu.selection) do s
        colornode[] = s
    end

    interpolationmenu =
        Menu(scene, options = zip(interpolationlabels, interpolationchoices))
    on(interpolationmenu.selection) do s
        interpolationnode[] = s
        heatmap1 = heatmap!(
            lscene,
            xlims,
            ylims,
            state,
            interpolate = s,
            colormap = cmap_rgb,
            colorrange = clims,
        )
    end

    newlabel = @lift($statename * " " * $unit)
    cbar = Colorbar(scene, heatmap1, label = newlabel)
    cbar.width = Relative(1 / 3)
    cbar.height = Relative(5 / 6)
    cbar.halign = :center
    # cbar.flipaxisposition = true
    # cbar.labelpadding = -350
    cbar.labelsize = 50

    lowerclim_string = @lift(
        "clim quantile = " *
        @sprintf("%0.2f", $lowerclim_node) *
        ", value = " *
        @sprintf("%0.1e", $clims[1])
    )
    upperclim_string = @lift(
        "clim quantile = " *
        @sprintf("%0.2f", $upperclim_node) *
        ", value = " *
        @sprintf("%0.1e", $clims[2])
    )

    # depends on makie version, vbox for old, vgrid for new
    layout[2, 1] = vgrid!(
        Label(scene, "State", width = nothing),
        statemenu,
        Label(
            scene,
            "plotting options",
            width = width,
            textsize = 30,
            padding = (0, 0, 10, 0),
        ),
        interpolationmenu,
        Label(scene, "Color", width = nothing),
        colormenu,
        Label(scene, lowerclim_string, width = nothing),
        lowerclim_slider,
        Label(scene, upperclim_string, width = nothing),
        upperclim_slider,
    )

    layout[2:4, 5] = vgrid!(
        Label(
            scene,
            "Color Bar",
            width = width / 2,
            textsize = 50,
            padding = (25, 0, 0, 00),
        ),
        cbar,
    )
    layout[1, 1] = Label(scene, "Menu", width = width, textsize = 50)
    display(scene)
    return scene
end


const DEFAULT_LIN_DISC_FORCE_OUTLIERS_TO_CLOSEST = true

"""
    LinearDiscretizer

A discretizer which encodes (typically) continuous values into discrete bins.
A univariate domain is divided into discrete by edges.
A value V will be encoded into bin B if V ∈ [Bₗ Bᵣ)
(or V ∈ [Bₗ Bᵣ] if B is the rightmost bin)

If force_outliers_to_closest is set to true, then values outside of binedges will be
shunted into the nearest bin. Otherwise an error is thrown.

TODO(tim):
    handle NA
    handle Nullable
    handle NaN
    ability to specify decoding behavior
"""
struct LinearDiscretizer{N<:Real, D<:Integer} <: AbstractDiscretizer{N,D}
    binedges :: Vector{N}   # list of bin edges, sorted smallest to largest
    nbins    :: Int
    i2d      :: Dict{Int,D} # maps bin index to discrete label
    d2i      :: Dict{D,Int} # maps discrete label to bin index

    force_outliers_to_closest :: Bool # if true, real values outside of the bin ranges will be forced to the nearest bin, otherwise will throw an error
end

function LinearDiscretizer( binedges::Vector{N}, i2d::Dict{Int,D};
    force_outliers_to_closest::Bool = DEFAULT_LIN_DISC_FORCE_OUTLIERS_TO_CLOSEST ) where {N<:Real, D<:Integer}

    length(binedges) > 1 || error("bin edges must contain at least 2 values")

    if N <: AbstractFloat # continuous values
        findfirst(i->binedges[i-1] ≥ binedges[i], 2:length(binedges)) == nothing ||
            error("Bin edges must be sorted in increasing order")
    else # for integers, bins of unit width require repeated values
        (findfirst(i->binedges[i-1] ≥ binedges[i], 2:length(binedges)-1) == nothing &&
         findfirst(i->binedges[i-1] > binedges[i], 2:length(binedges)) == nothing ) ||
            error("Bin edges must be sorted in increasing order")
    end

    all(haskey(i2d, i) for i in 1 : length(binedges)-1) || error("i2d must contain all necessary keys")

    d2i = Dict{D,Int}(v => k for (k,v) in i2d)

    return LinearDiscretizer{N,D}(binedges, length(binedges)-1, i2d, d2i, force_outliers_to_closest)
end
function LinearDiscretizer(binedges::Vector{N}, ::Type{D} = Int;
    force_outliers_to_closest::Bool = DEFAULT_LIN_DISC_FORCE_OUTLIERS_TO_CLOSEST ) where {N<:Real, D<:Integer}

    length(binedges) > 1 || error("bin edges must contain at least 2 values")

    if N <: AbstractFloat
        findfirst(i->binedges[i-1] ≥ binedges[i], 2:length(binedges)) == nothing ||
            error("Bin edges must be sorted in increasing order")
    else # for integers, bins of unit width require repeated values
        (findfirst(i->binedges[i-1] ≥ binedges[i], 2:length(binedges)-1) == nothing &&
         findfirst(i->binedges[i-1] > binedges[i], 2:length(binedges)) == nothing ) ||
            error("Bin edges must be sorted in increasing order")
    end

    i2d = Dict{Int,D}(i => convert(D, i) for i in 1 : length(binedges)-1)

    return LinearDiscretizer(binedges,i2d,force_outliers_to_closest=force_outliers_to_closest)
end
function LinearDiscretizer(arr::AbstractVector{N}, ::Type{D}=Int;
    force_outliers_to_closest::Bool = DEFAULT_LIN_DISC_FORCE_OUTLIERS_TO_CLOSEST) where {N<:Real, D<:Integer}
    LinearDiscretizer(convert(Vector{N}, arr), D, force_outliers_to_closest=force_outliers_to_closest)
end

function supports_encoding(ld::LinearDiscretizer{N,D}, x::N) where {N<:Real,D<:Integer}
    if ld.force_outliers_to_closest
        return true
    else
        return ld.binedges[1] ≤ x ≤ ld.binedges[end]
    end
end
supports_decoding(ld::LinearDiscretizer{N,D}, d::D) where {N<:Real,D<:Integer} = 1 ≤ d ≤ ld.nbins

function encode(ld::LinearDiscretizer{N,D}, x::N) where {N<:Real,D<:Integer}
    !isnan(x) || error("cannot encode NaN values")

    if x < ld.binedges[1]
        return ld.force_outliers_to_closest ? ld.i2d[1] : throw(BoundsError())
    elseif x > ld.binedges[end]
        return ld.force_outliers_to_closest ? ld.i2d[ld.nbins] : throw(BoundsError())
    else

        # run bisection search
        binedges = ld.binedges
        a, b = 1, length(binedges)
        va, vb = binedges[a], binedges[b]
        while b-a > 1
            c = div(a+b, 2)
            vc = binedges[c]
            if x < vc
                b, vb = c, vc
            else
                a, va = c, vc
            end
        end
        return ld.i2d[a]
    end
end
encode(ld::LinearDiscretizer{N,D}, x) where {N<:Real,D<:Integer} = encode(ld, convert(N, x))::D
function encode(ld::LinearDiscretizer{N,D}, data::AbstractArray) where {N<:Real,D<:Integer}
    arr = [encode(ld, x) for x in data]
    reshape(arr, size(data))
end

"""
There are several methods for decoding a LinearDiscretizer.
The default is SampleUniform, which uniformly samples from the bin.
SampleBinCenter returns the bin's center value.
"""
function decode(ld::LinearDiscretizer{N,D}, d::D, ::SampleUniform) where {N<:AbstractFloat,D<:Integer}
    ind = ld.d2i[d]
    lo  = ld.binedges[ind]
    hi  = ld.binedges[ind+1]
    convert(N, lo + rand()*(hi-lo))
end
function decode(ld::LinearDiscretizer{N,D}, d::D, ::SampleBinCenter) where {N<:AbstractFloat,D<:Integer}
    ind = ld.d2i[d]
    lo  = ld.binedges[ind]
    hi  = ld.binedges[ind+1]
    convert(N, (hi + lo)/2)
end
decode(ld::LinearDiscretizer{N,D}, d::D) where {N<:AbstractFloat,D<:Integer} = decode(ld, d, SAMPLE_UNIFORM)

function decode(ld::LinearDiscretizer{N,D}, d::D, ::SampleUniform) where {N<:Integer,D<:Integer}
    ind = ld.d2i[d]
    lo  = ld.binedges[ind]
    hi  = ld.binedges[ind+1]
    if hi != ld.binedges[end]
        retval = rand(lo:hi-1)::Int
    else
        retval = rand(lo:hi)::Int
    end
    convert(N, retval)
end
function decode(ld::LinearDiscretizer{N,D}, d::D, ::SampleBinCenter) where {N<:Integer,D<:Integer}
    ind = ld.d2i[d]
    lo = ld.binedges[ind]
    hi = ld.binedges[ind+1]
    convert(N, div(lo+hi,2))
end
decode(ld::LinearDiscretizer{N,D}, d::D) where {N<:Integer,D<:Integer} = decode(ld, d, SAMPLE_UNIFORM)

decode(ld::LinearDiscretizer{N,D}, d::I, method::AbstractSampleMethod=SAMPLE_UNIFORM) where {N<:Real,D<:Integer,I<:Integer} =
    decode(ld, convert(D,d), method)

function decode(ld::LinearDiscretizer{N,D}, data::AbstractArray{D}, ::AbstractSampleMethod=SAMPLE_UNIFORM) where {N<:Real,D<:Integer}
    arr = Vector{N}(undef, length(data))
    for (i,d) in enumerate(data)
        arr[i] = decode(ld, d)
    end
    reshape(arr, size(data))
end

Base.max(ld::LinearDiscretizer) = ld.binedges[end]
Base.min(ld::LinearDiscretizer) = ld.binedges[1]
function Base.extrema(ld::LinearDiscretizer{N,D}) where {N<:Real,D<:Integer}
    lo  = ld.binedges[1]
    hi  = ld.binedges[end]
    return (lo, hi)
end
function Base.extrema(ld::LinearDiscretizer{N,D}, d::D) where {N<:AbstractFloat,D<:Integer}
    ind = ld.d2i[d]
    lo  = ld.binedges[ind]
    hi  = ld.binedges[ind+1]
    return (lo, hi)
end
function Base.extrema(ld::LinearDiscretizer{N,D}, d::D) where {N<:Integer,D<:Integer}
    ind = ld.d2i[d]
    lo  = ld.binedges[ind]
    hi  = ld.binedges[ind+1]
    if hi == ld.binedges[end]
        (lo, hi)
    else
        (lo, hi-1)
    end
end
function totalwidth(ld::LinearDiscretizer)
    ex = extrema(ld)
    ex[2] - ex[1]
end

nlabels(ld::LinearDiscretizer) = ld.nbins
binedges(ld::LinearDiscretizer) = ld.binedges
bincenters(ld::LinearDiscretizer{N,D}) where {N<:AbstractFloat,D<:Integer} = (0.5*(ld.binedges[1:ld.nbins] + ld.binedges[2:end]))::Vector{Float64}
function bincenters(ld::LinearDiscretizer{N,D}) where {N<:Integer,D<:Integer}
    retval = Vector{Float64}(undef, ld.nbins)
    for i = 1 : length(retval)-1
        retval[i] = 0.5(ld.binedges[i+1]-1 + ld.binedges[i])
    end
    retval[end] = 0.5(ld.binedges[end] + ld.binedges[end-1])
    retval
end
binwidth(ld::LinearDiscretizer{N,D}, d::D) where {N<:Real,D<:Integer} = ld.binedges[d+1] - ld.binedges[d]
binwidths(ld::LinearDiscretizer{N,D}) where {N<:Real,D<:Integer} = ld.binedges[2:end] - ld.binedges[1:end-1]


primitives(vp::AbstractVolumePrimitive) = (vp,)
function primitives(csg::AbstractConstructiveGeometry)
    ps = []
    push!(ps, primitives(csg.a)...)
    push!(ps, primitives(csg.b)...)
    ps
end

@recipe function f(csg::AbstractConstructiveGeometry{T}; n_samples = 40, CSG_scale = missing) where {T}
    seriestype --> :csg
    ps = primitives(csg)
    if haskey(plotattributes, :seriestype) && plotattributes[:seriestype] == :samplesurface
        spacing::T = T((ismissing(CSG_scale) ? get_scale(csg) : CSG_scale)/n_samples)
    end
    @series begin
        label --> "CSG"
        if haskey(plotattributes, :seriestype)
            if plotattributes[:seriestype] == :samplesurface
                if occursin("GRBackend", string(typeof(plotattributes[:plot_object].backend)))
                    aspect_ratio --> 1.0
                end
                seriesalpha --> 0.2
                filter(p -> in(p,csg), sample(ps[1], spacing))
            else
                if !isClosedPrimitive(ps[1])
                    fillcolor := :white
                    fillalpha --> 0.2
                    if haskey(plotattributes, :seriestype)
                        if plotattributes[:seriestype] in [:csg, :wireframe]
                            linewidth := 1
                        end
                    end
                end
                ps[1]
            end
        end
    end
    for i in 2:length(ps)
        @series begin
            label := ""
            if haskey(plotattributes, :seriestype)
                if plotattributes[:seriestype] == :samplesurface
                    seriesalpha --> 0.2
                    filter(p -> in(p,csg), sample(ps[i], spacing))
                else
                    if !isClosedPrimitive(ps[i])
                        fillcolor := :white
                        fillalpha --> 0.2
                        if haskey(plotattributes, :seriestype)
                            if plotattributes[:seriestype] in [:csg, :wireframe]
                                linewidth := 1
                            end
                        end
                    end
                    ps[i]
                end
            end
        end
    end
end


# Routines for overlaying a hierarchical rectangular grid onto a gmsh CAD model

function gmsh_overlay_rectangular_grid(bb::NTuple{4, Float64},
                                       material::String,
                                       x::Vector{Vector{Float64}},
                                       y::Vector{Vector{Float64}})

    @info "Overlaying rectangular grid"
    # Ensure that the material is of the form "MATERIAL_xxxx"
    if (length(material) < 9)
        @error "material argument must of of the form 'MATERIAL_xxxxxx'"
    end
    if material[1:9] != "MATERIAL_"
        @error "material argument must of of the form 'MATERIAL_xxxxxx'"
    end

    model_dim_tags = gmsh.model.get_entities(2)
    grid_tags = gmsh_generate_rectangular_grid(bb, x, y; material = material)
    grid_dim_tags = [(2, tag) for tag in grid_tags]
    union_of_dim_tags = vcat(model_dim_tags, grid_dim_tags)
    groups = gmsh.model.get_physical_groups()
    names = [gmsh.model.get_physical_name(grp[1], grp[2]) for grp in groups]
    material_indices = findall(x->occursin("MATERIAL", x), names)
    # material hierarchy with the grid material at the bottom.
    material_hierarchy = names[material_indices]
    push!(material_hierarchy, material)
    out_dim_tags = gmsh_group_preserving_fragment(
                        union_of_dim_tags,
                        union_of_dim_tags;
                        material_hierarchy = material_hierarchy
                   )

    return out_dim_tags
end

function gmsh_overlay_rectangular_grid(bb::NTuple{4, Float64},
                                       material::String,
                                       nx::Vector{Int64},
                                       ny::Vector{Int64})

    @info "Overlaying rectangular grid"
    # Ensure that the material is of the form "MATERIAL_xxxx"
    if (length(material) < 9)
        @error "material argument must of of the form 'MATERIAL_xxxxxx'"
    end
    if material[1:9] != "MATERIAL_"
        @error "material argument must of of the form 'MATERIAL_xxxxxx'"
    end

    model_dim_tags = gmsh.model.get_entities(2)
    grid_tags = gmsh_generate_rectangular_grid(bb, nx, ny; material = material)
    grid_dim_tags = [(2, tag) for tag in grid_tags]
    union_of_dim_tags = vcat(model_dim_tags, grid_dim_tags)
    groups = gmsh.model.get_physical_groups()
    names = [gmsh.model.get_physical_name(grp[1], grp[2]) for grp in groups]
    material_indices = findall(x->occursin("MATERIAL", x), names)
    # material hierarchy with the grid material at the bottom.
    material_hierarchy = names[material_indices]
    push!(material_hierarchy, material)
    out_dim_tags = gmsh_group_preserving_fragment(
                        union_of_dim_tags,
                        union_of_dim_tags;
                        material_hierarchy = material_hierarchy
                   )

    return out_dim_tags
end


struct PFCombinatorTrace
    args::Tuple
    gen_fn::Any
    emission_choices::DynamicChoiceMap
    log_ml_est::Float64
end

Gen.get_args(tr::PFCombinatorTrace) = tr.args
Gen.get_retval(tr::PFCombinatorTrace) = nothing
Gen.get_gen_fn(tr::PFCombinatorTrace) = tr.gen_fn
Gen.get_score(tr::PFCombinatorTrace) = tr.log_ml_est
Gen.get_choices(tr::PFCombinatorTrace) = tr.emission_choices

struct ParticleFilterCombinator <: GenerativeFunction{Any,PFCombinatorTrace}
    combined::GenerativeFunction
    num_particles::Int
end

function ParticleFilterCombinator(init, dynamics, emission, num_particles::Int)

    @gen function kernel(t::Int, prev_state,
            dynamics_args::Tuple, emission_args::Tuple)
        state = @trace(dynamics(t+1, prev_state, dynamics_args...), :dynamics)
        @trace(emission(t+1, state, emission_args...), :emission)
        return state
    end

    chain = Unfold(kernel)

    @gen function combined(T::Int, init_args::Tuple, dynamics_args::Tuple, emission_args::Tuple)
        init_state = @trace(init(init_args...), :init)
        @trace(emission(1, init_state, emission_args...), :init_emission)
        @trace(chain(T, init_state, dynamics_args, emission_args...), :chain, @argdiff())
    end

    ParticleFilterCombinator(combined, num_particles)
end

function run_particle_filter(gen_fn::ParticleFilterCombinator, args::Tuple, choices::ChoiceMap, num_particles::Int)
    (T::Int, initial_args, dynamics_args, emission_args) = args
    init_obs = choicemap()
    set_submap!(init_obs, :init_emission, get_submap(choices, 1))
    init_args = (0, initial_args, dynamics_args, emission_args)
    state = initialize_particle_filter(gen_fn.combined,
        init_args, init_obs, num_particles)
    argdiff = UnfoldCustomArgDiff(false, false)
    for t=2:T
        new_args = (t-1, initial_args, dynamics_args, emission_args)
        maybe_resample!(state)
        obs = choicemap()
        set_submap!(obs, :chain => t-1 => :emission, get_submap(choices, t))
        particle_filter_step!(state, new_args, argdiff, obs)
    end
    log_ml_estimate(state)
end

function Gen.generate(gen_fn::ParticleFilterCombinator, args::Tuple, choices::ChoiceMap)
    log_ml_est = run_particle_filter(gen_fn, args, choices, gen_fn.num_particles)
    trace = PFCombinatorTrace(args, gen_fn, choices, log_ml_est)
    (trace, log_ml_est)
end

function Gen.update(trace::PFCombinatorTrace, args::Tuple, argdiff, choices::ChoiceMap)
    if !isempty(choices)
        error("Not implemented")
    end
    gen_fn = get_gen_fn(trace)
    new_log_ml_est = run_particle_filter(gen_fn, args, trace.emission_choices, gen_fn.num_particles)
    new_trace = PFCombinatorTrace(args, gen_fn, trace.emission_choices, new_log_ml_est)
    weight = new_log_ml_est - trace.log_ml_est
    (new_trace, weight, DefaultRetDiff(), EmptyChoiceMap())
end

function Gen.regenerate(trace::PFCombinatorTrace, args::Tuple, argdiff, selection::Selection)
    if !isempty(selection)
        error("Not implemented")
    end
    gen_fn = get_gen_fn(trace)
    new_log_ml_est = run_particle_filter(gen_fn, args, trace.emission_choices, gen_fn.num_particles)
    new_trace = PFCombinatorTrace(args, gen_fn, trace.emission_choices, new_log_ml_est)
    weight = new_log_ml_est - trace.log_ml_est
    (new_trace, weight, DefaultRetDiff(), EmptyChoiceMap())
end

function Gen.extend(::PFCombinatorTrace, ::Tuple, argdiff, ::ChoiceMap)
    error("Not implemented")
end

function Gen.choice_gradients(::PFCombinatorTrace, ::Selection)
    error("Not implemented")
end

function Gen.accumulate_param_gradients!(::PFCombinatorTrace,
    retgrad, scale_factor)
    error("Not implemented")
end

function Gen.propose(::ParticleFilterCombinator, args::Tuple)
    error("Not implemented")
end

function Gen.assess(::ParticleFilterCombinator, args::Tuple, choices::ChoiceMap)
    error("Not implemented")
end

function Gen.project(tr::PFCombinatorTrace, ::EmptySelection)
    # q(t; u=empty, x) = simulate forward from model
    # q(r; x, t) = run sMC
    # therefore, the result is SMC-ML-est / forward probability
    #tr.log_ml_est - tr.model_project
    0.
end


fname = :TRKsparse
c = Combi(hessian_sparse,PDataK,solve_modelKTR,preprocessKTR,decreaseKTR,TparamsKTR())
include("Template.jl")