##########
## ITMIX (of all versions loaders)
#########
#

#####
# ITMIX data
#####
"Translate part of the file-name to DataKind subtypes"
const dk_dict = Dict("RES"=>ThicknessData,
                     "margin"=>OutlineData,
                     "surface"=>DEMData,
                     "mb"=>BdotData,
                     "dhdt"=>DhdtData,
                     "speed"=>IVData,
                     "azimuth"=>nothing,
                     "velocity"=>IVData)

todo = 1
"Whether a glacier is calving, is an ice cap, and its RGI 4.0 number (if there is a single one)."
const iscalving_isicecap = Dict(:Academy=>(true,true,nothing),
                                :Austfonna=>(true,true,nothing),
                                :Devon=>(true,true,"RGI40-03.02443"), # one of the major outlet glaciers in the North
                                :Hellstugubreen=>(false,false,"RGI40-08.02182"),
                                :Mocho=>(false,false,nothing), # glacier complex
                                :SouthGlacier=>(false,false,"RGI40-01.17598"),
                                :Tasman=>(false,false,"RGI40-18.02342"),
                                :Washmawapta=>(false,false,"RGI40-02.04894"),
                                :Aqqutikitsoq=>(false,false,"RGI40-05.00524"), # not sure, size does not match well
                                :Brewster=>(false,false,"RGI40-18.01130"),
                                :Elbrus=>(false,false,nothing), # glacier complex
                                :NorthGlacier=>(false,false,"RGI40-01.18240"),
                                :Starbuck=>(true,false,nothing),
                                :Unteraar=>(false,false,"RGI40-11.01328"),
                                :Columbia=>(true,false,"RGI40-01.11506"),
                                :Freya=>(false,false,"RGI40-05.19983"),
                                :Kesselwandferner=>(false,false,"RGI40-11.00787"),
                                :Urumqi=>(false,false,"RGI40-13.27054"),
                                :Synthetic1=>(false,false,nothing),
                                :Synthetic2=>(false,false,nothing),
                                :Synthetic3=>(false,true,nothing),
                                :AustreGroenfjordbreen=>(false,false,nothing),
                                :ChhotaShigri=>(false,false,nothing),
                                )

# ITMIX glacier with IV data:
const itmix_glaciers_iv = [:Synthetic1,  :Unteraar,
                           :Synthetic2, :Synthetic3,
                           :Tasman, :Austfonna, :Devon,
                           :SouthGlacier, :NorthGlacier,
                           :Brewster, :Hellstugubreen]

# of which these only have point data (at GPS locations)
const itmix_glaciers_iv_point = [:NorthGlacier, :SouthGlacier, :Hellstugubreen, :Brewster]

function LoadPara(gid::ITMIXGlacier; use_glogem=false, pl_kwargs...)
    pl_kwargs = deepcopy(pl_kwargs)
    id = getid(gid)
    data_root = gid.version==1? joinpath(LoadPara(gid=gid).data_root,"ITMIX") : joinpath(LoadPara(gid=gid).data_root,"ITMIX2")
    pl = LoadPara(; gid=gid, data_root=data_root, pl_kwargs...)
    # set the sigmas of uncertainty for the data:
    for (k,v) in Dict(ThicknessData=>10.0, # maybe change to more
                      BdotData=>0.2,  # m/a, 1m/a is what Matthias says about GloGEM
                      DhdtData=>0.2,  # m/a, e.g. 1m/a Nuth et al. 2010, Table 1, for Svalbard
                      IVData=>10.0)
        pl.dataset_opts[k][:sigma]=v
    end
    updated_pm_defaults = Dict(
        :bandsize => 30,
        :window_width_smooth=>100, # can make a noisy tau
        :cap_at_floatation=>false,
        :dist_exp => 0.5,
        #:window => 0
    )
    if iscalving_isicecap[id][2]
        merge!(updated_pm_defaults, Dict(:h_at_top => -1, :shapeF => false))
    end
    if iscalving_isicecap[id][1]
        # relative flux lost through calving
        pl.dataset_opts[TerminusFluxData][:fluxkind] =
            FluxKind.rel_btilde_teminus # compare to prior expectation of btilde integral
        #pl.dataset_opts[TerminusFluxData][:terminus_flux] = SomeData(vals=1.0, sigma=0.5, range_abs=(0.1,3.0))
        pl.dataset_opts[TerminusFluxData][:terminus_flux] =
            SomeData(vals=NaN) # require this to be set individually
    else
        # approximately zero flux at terminus +/- 1% of max bdot flux accumulation
        pl.dataset_opts[TerminusFluxData][:fluxkind] = FluxKind.rel_bdot_max
        pl.dataset_opts[TerminusFluxData][:terminus_flux] = SomeData(vals=0.0, sigma=0.01)
    end

    if id==:Austfonna
        # This makes a coarser grid as the Austfonna data is on a
        # 50m grid!
        # with 50m DEM:
        #  - 400s without cache
        #  - 35s with cache
        # with 250m DEM:
        #  - 208s without cache
        #  - 35s  with cache
        if gid.version==1
            pl.dataset_opts[DataKind][:grid_downsample_step]=5
        end
        #pl.dataset_LOADERS[DhdtData] = SyntheticBenchLoader # something is off with Austfonna's dh/dt (or bdot)

        pl.dataset_opts[TerminusFluxData][:fluxkind] = FluxKind.rel_bdot_max # compare to prior expectation of bdot integral
        # Dowdeswell at al. 2008 says this:
        pl.dataset_opts[TerminusFluxData][:terminus_flux] = SomeData(vals=1/3,
                                                                     sigma=0.4,
                                                                     range_abs=(0.05, 3.0))
    elseif id==:AustreGroenfjordbreen
        # compare to prior expectation of btilde integral because rel_bdot_max is negative
        pl.dataset_opts[TerminusFluxData][:fluxkind] = FluxKind.rel_btilde_teminus
        pl.dataset_opts[TerminusFluxData][:terminus_flux] = SomeData(vals=0.0, sigma=0.01)
    elseif id==:Columbia
        # the size of (1040, 920) make it run slow, about 0.5s per forward run
        # reduced by step=2 (4x) it's down to 0.07, in-line with the others.
        pl.dataset_opts[DataKind][:grid_downsample_step] = 2

        pl.dataset_opts[BdotData][:smb_gradient] = 0.02 # maritime
        pl.dataset_opts[TerminusFluxData][:fluxkind] = FluxKind.rel_bdot_max # compare to prior expectation of max-bdot integral
                                                                             # Note, bdot intergral is negative at terminus
        pl.dataset_opts[TerminusFluxData][:terminus_flux] = SomeData(vals=0.5,
                                                                     sigma=0.3,
                                                                     range_abs=(0.1,3.0))
        merge!(updated_pm_defaults, Dict(:h_at_top => -1)) # more ice cap like at devide
    elseif id==:Aqqutikitsoq
        if gid.version==1
            pl.dataset_opts[DataKind][:grid_downsample_step]=2
        end
    elseif id==:Hellstugubreen
        # pl.dataset_LOADERS[DhdtData] = SyntheticBenchLoader
        # pl.dataset_opts[DhdtData][:ela_elevation] = 2000
    elseif id==:Elbrus
        if gid.version==1
            pl.dataset_opts[DataKind][:grid_downsample_step]=3
        else
            # the size of (601,581) make it run slow, about 0.12s per forward run
            # reduced by step=2 (4x) it's down to 0.025, in-line with the others.
            # (now on 60m resolution vs 30m resolution)
            pl.dataset_opts[DataKind][:grid_downsample_step] = 2
        end

        pl.dataset_LOADERS[BdotData] = SyntheticBenchLoader # original only covers a small fraction of the area
                                                            # TODO: this could be turned into elevation band data:
                                                             # pl.dataset_opts[BdotData][:vals_prior] = 0.0
        # pl.dataset_LOADERS[DhdtData] = SyntheticBenchLoader # doesn't make much sense to then not use this either.
        merge!(updated_pm_defaults, Dict(:h_at_top => -1, :shapeF => false))
    elseif id==:Unteraar
        # something is amiss with dhdt or bdot data as glacier stops at confluence:
        # pl.dataset_LOADERS[DhdtData] = SyntheticBenchLoader
        # pl.dataset_opts[DhdtData][:ela_elevation] = 2900
        pl.dataset_opts[IVData][:sigma] = 4.0 # Vogel et al 2012
        pl.dataset_opts[ThicknessData][:sigma] = 27.0 # Bauder etal 2003
    elseif id==:Freya
        # uses the SyntheticBenchLoader
        pl.dataset_opts[DhdtData][:ela_elevation] = 985 # gets negative fluxes in the middle otherwise
    elseif id==:Urumqi
        pl.dataset_opts[DhdtData][:ela_elevation] = 4200
    elseif id==:Tasman
        merge!(updated_pm_defaults, Dict(:bandsize => 35))
    elseif id==:Devon
        pl.dataset_opts[TerminusFluxData][:fluxkind] = FluxKind.rel_bdot_max # compare to prior expectation of bdot integral
        # https://www.igsoc.org/annals/46/a46a288.pdf says 1/4 calving and 3/4 melt
        pl.dataset_opts[TerminusFluxData][:terminus_flux] = SomeData(vals=1/4,
                                                                     sigma=0.1,
                                                                     range_abs=(0.1,3.0))
    elseif id==:Synthetic3
        merge!(updated_pm_defaults, Dict(:h_at_top => -1))
    elseif id==:Synthetic1 || id==:Synthetic2 || id==:Synthetic3
        # there are no "measurement" errors here as no noise was put on the data.
        # Note there is still the BITEModel-error in pm.sigma_*
        pl.dataset_opts[IVData][:sigma] = 0.0
        pl.dataset_opts[ThicknessData][:sigma] = 0.0
    elseif id==:ChhotaShigri
        # flux not zero at terminus
    elseif id==:Academy
        pl.dataset_opts[TerminusFluxData][:fluxkind] = FluxKind.rel_bdot_max
        # Dowdeswell at all 2008 says this:
        pl.dataset_opts[TerminusFluxData][:terminus_flux] = SomeData(vals=1/3,
                                                                     sigma=0.3,
                                                                     range_abs=(0.1,3.0))
    elseif id==:Mocho
        # has negative mass-balance at the top!
        merge!(updated_pm_defaults, Dict(:h_at_top => -1, :shapeF => false))
    elseif id==:Starbuck
        pl.dataset_opts[TerminusFluxData][:fluxkind] = FluxKind.rel_bdot_teminus # compare to prior expectation of bdot integral
        pl.dataset_opts[TerminusFluxData][:terminus_flux] = SomeData(vals=1.0,
                                                                     sigma=0.1,
                                                                     range_abs=(0.1,3.0))
        merge!(updated_pm_defaults, Dict(:h_at_top => -1))  # more ice cap like at devide
    end
    pl.dataset_opts[DEMData][:cropbox] = :from_DEM

    # update pm but keeping any explicitly set (via kwargs) options
    merge!((a,b)->a, pl.dataset_opts[ParaData][:pm], updated_pm_defaults)

    # GloGEM
    if use_glogem && iscalving_isicecap[gid.id][3]!=nothing
        error("glogem loader disabled for ITMIX2 runs")
        pl.dataset_LOADERS[BdotData] = GloGEMLoader
        pl.dataset_LOADERS[DhdtData] = SyntheticBenchLoader
    else
        if use_glogem
            error("Not using GloGEM for $gid")
        end
    end

    # Set loaders and files:
    @unpack dataset_LOADERS, dataset_files, dataset_opts = pl
    # check what files and thus data exist
    name = getname(gid)
    dir = joinpath(pl.data_root, name)
    zip = joinpath(pl.data_root, name*".zip")
    fls = Dict()
    if isfile(zip) # zip-archive
        zipfl = ZipFile.Reader(zip)
        toread = zipfl.files
    elseif isdir(dir)
        toread = readdir(dir)
        append!(toread, joinpath.("shapefiles", readdir(joinpath(dir, "shapefiles"))))
    else
        error("Cannot read $dir or $zip")
    end
    for fl in toread
        out = parse_ITMIX_filename(fl)
        if out!=nothing
            fls[out[1]] = out[2:end]
        end
    end
    isfile(zip) && close(zipfl)
    proj_ = fls[OutlineData][end]
    args = []

    # Set the LOADERS
    for DK in datakinds()
        if !haskey(dataset_LOADERS, DK)
            if haskey(fls,DK) # ITMIX file present
                fln, nr, nameinfile, date_range, proj = fls[DK]
                dataset_LOADERS[DK] = ITMIXLoader
                dataset_files[DK] = [joinpath(dir,fln)]
                dataset_opts[DK][:proj] = proj_
            else
                if DK==ParaData
                    dataset_LOADERS[DK] = ITMIXLoader
                    if iscalving_isicecap[Symbol(name)][2]
                        dataset_opts[DK][:iv_extrapolation_scaling] = :huss
                    end
                # elseif DK==BdotData # TODO: use GloGEM by default
                #     dataset_LOADERS[DK] = GloGEMLoader
                #     dataset_files[DK] = [joinpath(dir,fln)]
                elseif DK==IVData || DK==ThicknessData
                    dataset_LOADERS[DK] = NothingLoader
                else
                    # if there is no file use the synthetic loader
                    dataset_LOADERS[DK] = SyntheticBenchLoader
                end
            end
        else
            # custom set loader & other bits, leave as is
        end
    end

    # bug in shapefile
    if id==:Tasman
        dataset_files[OutlineData] = ["data/ITMIX/Tasman/01_margin_Tasman_1986_UTM-59.csv"]
    end
    return pl
end

# Outline
function load!(ds::DataSet{OutlineData, ITMIXLoader}, gid)
    opts = ds.opts
    proj = opts[:proj]
    # fl, fid = open_file(ds.files[1])
    # outline = readdlm(fl, ';', Float64)'
    # close(fid)
    @assert length(ds.files)==1
    if endswith(ds.files[1], ".shp")
        outline, splits = load_shapefile!(ds, :outline)
    else
        outline = @open_magic readdlm(fl, ';', F)'

        if size(outline,1)==3
            # TODO
            # # process boundary ID stuff:
            # # TODO: not sure this is right yet.
            # boundary_ind = outline[3,:]
            # xy = outline[1:2,:]
            # # all boundary_ind>1 treat as outcrops
            # outline = xy[:,boundary_ind.==boundary_ind[1]]
            # outcropsx = Float64[]
            # outcropsy = Float64[]
            # splits = Int[]
            # for i=2:1000
            #     inds = boundary_ind.==i
            #     length(inds)==0 && break
            #     append!(outcropsx, xy[1,inds])
            #     append!(outcropsy, xy[2,inds])
            # end
            # outline = xy
            outline = outline[1:2,:]
        end
        if outline[:,1] !=outline[:,end]
            # close polygon
            outline = hcat(outline, outline[:,1])
        end
        splits = VAWTools.find_poly_splits(outline)
    end
    # Fix Unteraar
    if getname(gid) == "Unteraar"
        outline = outline[:, splits[end-1]:end-1]
        # remove self-intersections
        outline = outline[:,[1:1483, 1487:1667, 1671:end;]]
        splits = [1,size(outline)[2]]
        @assert outline[:,1]==outline[:,end]
    end

    iscalving,isicecap,rgi = iscalving_isicecap[getid(gid)]
    @assert isa(iscalving, Bool) "Input whether calving or not in iscalving_dict"
    return Outline(gid, outline, splits),
           iscalving::Bool, isicecap::Bool, proj::String
end

"""
opts:
- grid : if specified use a different grid from the grid of the DEM as
  computational grid.
- cropbox -- used for cropping the DEM and other 2D data.  Needs to
  contain all of the outline. `F[xmin, xmax, ymin, ymax]`, defaults
  being calculated during DEM-load

Return:
- dem::Gridded
- glaciermask::Matrix{Bool} -- true on glacier
- landmask::Matrix{Bool} -- true on solid ground (i.e. false on glaciers, ice shelves, water)
- dist_land::Matrix{F} -- distance to closest land
- dist_land_maxdist -- maximum distance to which to calculate above
- cropbox::Vector{Float64} -- box around region of interest (contains glacier + some buffer)
                              `F[xmin, xmax, ymin, ymax]`
"""
function load!(ds::DataSet{DEMData, ITMIXLoader}, outline::Outline, pp::Phys
               )::Tuple{Gridded, Matrix{Bool}, Matrix{Bool}, Matrix{F}, F, Tuple, String}
    @assert ds.opts[:cropbox]==:from_DEM "ITMIX cropbox needs to be: `:from_DEM`"
    return load!(convert(DataSet{DEMData, :ASCIIGridLoader}, ds), outline, pp)
end
function load!(ds::DataSet{IVData, ITMIXLoader}, dem, glaciermask, target_proj,
               outline, cropbox, pp::Phys)::Tuple{SomeData, Matrix{Bool}}
    if splitext(ds.files[1])[2]==".txt"
        # IV at points
        tmp = @open_magic readdlm(fl, header=true)
        speed = F.(tmp[1][:,4])
        x = F.(tmp[1][:,end-2])
        y = F.(tmp[1][:,end-1])
        iv = Traj(x,y,speed)
        ivmask = Bool.(trues(0,0))
        return SomeData(vals=iv, sigma=ds.opts[:sigma]), ivmask
    elseif splitext(ds.files[1])[2]==".asc"
        # IV in 2D
        return load!(convert(DataSet{IVData, :ASCIIGridLoader}, ds), dem, glaciermask, target_proj,
                     outline, cropbox, pp)
    else
        error("Extension of file $file not supported.  Need .txt or .asc")
    end
end
function load!(ds::DataSet{ThicknessData, ITMIXLoader}, dem, glaciermask, target_proj, outline, cropbox,
               pp::Phys)::SomeData
    @unpack files, opts, daterange = ds
    @unpack gid = outline
    sigma = opts[:sigma]
    if gid.version==1
        zb = get_cache_ram!(ds, :zb, @open_magic_fn readdlm(fl; header=true)[1])
        x,y = zb[:,1], zb[:,2]
        zb = zb[:,3]
        izs  = Interp.interpolate((dem.x, dem.y), dem.v, Interp.Gridded(Interp.Linear()) )
        zs = [izs[xx,yy] for (xx,yy) in zip(x,y)]
        thick = zs-zb
        out = Traj(x, y, thick, zs)
        split_traj!(out, 1e3) # TODO adjust split-distance
    else
        zb = get_cache_ram!(ds, :zb, @open_magic_fn readdlm(fl; header=true, skipstart=16)[1])
        lines = unique(zb[:,1])
        x,y,thick,zs = zb[:,3], zb[:,4], zb[:,end], zb[:,5]
        splits = []
        for l in lines
            push!(splits, findfirst(zb[:,1], l):findlast(zb[:,1], l))
        end
        out = Traj(x, y, thick, zs, splits)
    end
    return SomeData(vals=out, sigma=sigma)
end
function load!(ds::DataSet{BdotData, ITMIXLoader}, dem, glaciermask, target_proj, outline, cropbox,
               pp::Phys)
    bdot = load_asciigrid!(ds, :bdot, cropbox)
    # scale to m ice / a
    scale_m_ice!(bdot, pp)
    if haskey(ds.opts, :vals_prior)
        return SomeData(vals=bdot, sigma=ds.opts[:sigma], vals_prior=ds.opts[:vals_prior])
    else
        return SomeData(vals=bdot, sigma=ds.opts[:sigma])
    end
end
function load!(ds::DataSet{DhdtData, ITMIXLoader},
               bdot, terminus_flux,
               dem, glaciermask, target_proj, outline, cropbox,
               iscalving, pp::Phys)
    return load!(convert(DataSet{DhdtData, :ASCIIGridLoader}, ds),
                 dem, glaciermask, target_proj, outline, cropbox, pp)
end

function load!(ds::DataSet{FslData, ITMIXLoader}, dem, glaciermask, target_proj, outline, cropbox,
               pp::Phys)
    error("FSL data not provided in ITMIX.  Use different loader, probably syntetic")
end
function load!(ds::DataSet{TempData, ITMIXLoader}, dem, glaciermask, target_proj, outline, cropbox,
               pp::Phys)
    error("Ice-temp data not provided in ITMIX.  Use different loader, probably syntetic")
end
function load!(ds::DataSet{TerminusFluxData, ITMIXLoader}, dem, glaciermask, iscalving,
               pp::Phys)
    error("Calving flux not provided in ITMIX.  Use different loader, probably synthetic")
end

#############################
# Helper

# Each of this folders contains a combination of the following files:
#
# - 01_margin_GGGGGG_YYYY_UTMZZZ.csv
# - 02_surface_GGGGGG_YYYY_UTMZZZ.asc  #
# - 03_RES_GGGGGG.txt                  # GPR file
# - 04_mb_GGGGGG_XXXX-YYYY_UTMZZZ.asc
# - 05_dhdt_GGGGGG_XXXX-YYYY_UTMZZZ.asc
# - 06_speed_GGGGGG_XXXX-YYYY_UTMZZZ.asc
# - 07_azimuth_GGGGGG_XXXX-YYYY_UTMZZZ.asc
# - 08_velocity_GGGGGG_XXXX-YYYY_UTMZZZ.txt
#
# where
#
# GGGGGG = name of the glacier (e.g. "Brewster")
# YYYY   = year the data are referring to (e.g. "2003")
# ZZZZ   = UTM zone of the coordinate system in which the data are projected (e.g. "32") - negative values indicate the southern hemisphere.
# Files 04-08 are mean values over the period "XXXX-YYYY" (e.g. "1997-2009")
#
# The file content is as follows:
# 01: Margin of the glacier (m).
# 02: DEM of the glacier (m a.s.l.) - Note that the DEM might be cropped with the glacier margin or not
# 04: Spatially distributed average (surface) mass balance of the glacier (m w.e./a)
# 05: Spatially distributed average ice thickness change rate (m/a)
# 06: Spatially distributed average surface flow speed of the glacier (m/a)
# 07: Spatially distributed average azimuth (="direction") of the surface flow (clockwise degrees from north)
# 08: Point observations of surface flow velocity (comprising speed and azimuth)
#
# .txt-files are plain text files.
# .csv-files are comma-separated-value files.
# .asc-files are ASCII-grid files (http://en.wikipedia.org/wiki/Esri_grid)
function parse_ITMIX_filename(fln::ZipFile.ReadableFile)
    error("zip files not supported currently")
    dir,fl = splitdir(fln.name)
    if contains(dir,"shapefiles")
        return nothing
    else
        return parse_ITMIX_filename(fl)
    end
end
function parse_ITMIX_filename(fln::String)
    fln=="" && return nothing
    fl, ext = splitext(splitdir(fln)[2])
    splits = find(c->c=='_', fl)
    length(splits)<2 && return nothing
    !(ext in [".txt",".asc",".shp"]) && return nothing # drop the ".csv" outlines in favor of shp
    nr = fl[1:splits[1]-1]
    !all(isnumber, nr) && return nothing
    nr = parse(Int,nr)
    dk = fl[splits[1]+1:splits[2]-1]
    # if dk in ["margin","surface"]
    #     name = fl[splits[2]+1:splits[3]-1]
    #     DK = dk=="margin" ? Outline : DEM
    #     y = parse(Int, fl[splits[3]+1:splits[4]-1])
    #     date_range = (Date(y),Date(y))
    # else
    DK = dk_dict[dk]
    if dk=="RES"
        name = fl[splits[2]+1:end]
        date_range = (Date(),Date())
    else
        name = fl[splits[2]+1:splits[3]-1]
        sp = split(fl[splits[3]+1:splits[4]-1],'-')
        if length(sp)==1
            y1 = parse(Int,sp[1])
            y2 = y1
        else
            y1,y2 = (parse(Int,s) for s in sp)
        end
        date_range = (Date(y1),Date(y2))
    end
    if dk!="RES"
        if name=="Unteraar"
            # Unteraar is not in proper UTM projection!
            zone = fl[splits[end]+1:end]
            proj = "close to UTM zone $zone, but no cigar"
        else
            zone = fl[splits[end]+1:end]
            proj = "+proj=utm +zone=$zone +ellps=WGS84 +datum=WGS84 +units=m +no_defs"
        end
    else
        proj = "unknown"
    end
    return DK, fln, nr, name, date_range, proj
end

"Read a agr file"
_read(name, dirs, num) = read_agr(_getfile(dirs, name, num), F, NA=FILL)

function generate_DEM_masks(gid::ITMIXGlacier, dem, outline, sealevel_cutoff=10.0)
    glaciermask, _ = maskit(dem, outline)
    # also DEM holes from glaciermask
    glaciermask[dem.v.==FILL] = false
    # remove points below sea level
    glaciermask[dem.v.<0] = false

    remove_slivers = true
    landmask = generate_landmask(glaciermask, dem, sealevel_cutoff, remove_slivers)
    return glaciermask, landmask
end

######################
# Special processing (aka hacks)
######################

function special_processing!(gid::ITMIXGlacier, outline, dem, glaciermask, landmask,
                             iv, ivmask, thickness, bdot, dhdt, fsl, temp,
                             pl::LoadPara, pp::Phys, pm::MPara, pn::Num)
    if getname(gid)=="Austfonna"
        # no easy way of figuring land vs sea out for Austfonna, set all to sea.
        landmask[:] = false
        ivmask[iv.vals.v.<=0] = false
    elseif getname(gid)=="Tasman" || getname(gid)=="Devon" || getname(gid)=="Unteraar"
        # not masked at all in the input data
        ivmask[iv.vals.v.<=0] = false
    end
    return nothing
end

###########################
## ITMIX2 Utilities
###########################
module ITMIX2
import ..BITEModel; const BM = BITEModel
import VAWTools
"Number of experiments per glacier"
const nexps = 16
using DataFrames, FileIO, CSVFiles
int2expnr(i) = i>9 ? Symbol("exp$i") : Symbol("exp0$i")

"""
    load_experiment_table(fln="data/ITMIX2/ITMIX2_experiments_key.txt")

Load ITMIX2_experiments_key.txt as a DataFrame
"""
function load_experiment_table(fln)
    tab = DataFrame(load(File(format"CSV", fln),
                         spacedelim=true,
                         header_exists=true,
                         skiplines_begin=11,
                         type_detect_rows=512,
                         #colparsers=[String,Int, fill(x->Bool(parse(Int,x)),16)..., Int]
                         ))
    tab[:glacier] = [BM.ITMIXGlacier(Symbol(gl), 2) for gl in tab[:glacier]]
    for i=1:nexps
        tab[int2expnr(i)] = Bool.(tab[int2expnr(i)])
    end
    # add exp00
    tab[:exp00] = tab[:exp01].*false
    return tab
end

"Get the priority of a glacier"
getpriority(gid) = getexps(gid)[:prio][1]
getpriority() = Dict(((gln, getexps(gln)[:prio][1]) for gln in BM.ITMIXnames))

"Get experiments for a glacier `gl`"
getexps(gid::BM.ITMIXGlacier) = DataFrames.filter(row->row[:glacier]==gid, tab)
getexps(name::Union{Symbol,String}) = getexps(BM.ITMIXGlacier(name,2))

"Get fitting radar-lines of experiment `i` for a glacier `gl`."
get_rlines(gid::BM.ITMIXGlacier, i) = i==0 ? Int[] : find(getexps(gid)[int2expnr(i)])
get_rlines(name::Union{Symbol,String}, i) =  i==0 ? Int[] : get_rlines(BM.ITMIXGlacier(name,2), i)

fln = joinpath(@__DIR__, "../../scripts/data/ITMIX2/ITMIX2_experiments_key.txt")
"The ITMIX2 experiment table"
const tab = isfile(fln) ? load_experiment_table(fln) : nothing

"The ITMIX2 glaciers sorted according to their priorities"
const priorities = tab==nothing ? nothing : [[BM.ITMIXGlacier(gln,2) for gln in keys(getpriority()) if getpriority()[gln]==i] for i=1:4]

# # iteration over the ITMIX 2 experiments
# """
#     experiments([gid])

# Return iterator for ITMIX2 experiments.

# ```
# for (gid, i, radarlines) in experiments([gid])
#    ...
# end
# ```
# """
# function experiments end

# experiments(gid) = ITMIX2Iterator(gid, getexps(gid))
# struct ITMIX2Iterator
#     gid::BM.ITMIXGlacier
# end
# Base.start(::ITMIX2Iterator) = 1
# Base.next(ii::ITMIX2Iterator, state) = ( (gid, int2expnr(state), getexp(ii.gid, state)), state+1)
# Base.done(ii::ITMIX2Iterator, state) = state==17

# experiments() = ITMIX2Iterator_All()
# struct ITMIX2Iterator_All end
# Base.start(::ITMIX2Iterator_All) = (1, BM.ITMIXv2_glaciers[1])
# function Base.next(ii::ITMIX2Iterator_All, state)
#     expnr, gid = state
#     item = (gid, int2expnr(expnr), getexp(gid, expnr))
#     state = if expnr==16 && gid!=BM.ITMIXv2_glaciers[end]
#         (1, BM.ITMIXv2_glaciers[findfirst(BM.ITMIXv2_glaciers, gid)+1])
#     else
#         (expnr+1, gid)
#     end

#     return item, state
# end
# Base.done(ii::ITMIX2Iterator_All, state) = state[1]==17

"""
    make_ITMIX2_glacier(gl, expnr::Number, rm_iv=false)
    make_ITMIX2_glacier(gl, tracks2keep, rm_iv=false)

Take the glacier with all radarlines loaded and return one which only has the
lines of that experiment loaded.  The second interface allows to specify which
tracks to keep and ditch the rest.

Also has the extra feature to make more experiments:
- exp0: no radar lines
- exp 17-32: no IV, otherwise as 1.16

Above two only make sense (in the sense that there is something to fit)
for glaciers with IV data.
"""
make_ITMIX2_glacier(gl, expnr::Number, rm_iv=false) =
    make_ITMIX2_glacier(gl, get_rlines(gl.gid, expnr), rm_iv, expnr)
function make_ITMIX2_glacier(gl, tracks, rm_iv=false, expnr=-1)
    new_h = if length(tracks)>0  # ITMIX2 experiments
        BM.SomeData(gl.h, vals=VAWTools.Traj(gl.h.vals, tracks))
    else # remove all radar-lines: name it exp00.  This is the ITMIX1 run.
        new_h = BM.SomeData(nothing)
    end
    misc = merge(gl.misc, Dict{Symbol,Any}(:expnr=>expnr, :rm_iv=>rm_iv))
    return if rm_iv
        # check that there is IV
        gl.iv isa BM.SomeData{Void} && error("Cannot remove IV from glacier $(BM.getname(gl)) as it has no IV.")
        # remove IV data
        BM.Glacier(gl, h=new_h, iv=BM.SomeData(nothing), misc=misc)
    else
        # normal run
        BM.Glacier(gl, h=new_h, misc=misc)
    end
end

##########################
## Saving
##########################
# needs to be saved on the same grid as the input data
#=
ITMIX2 - instructions for results submission

1) Results format and content
- Submitted results shall be distributed estimates of the ice thickness.
- The unit for the ice thickness is "meters".
- Results are to be submitted as Esri grids, i.e. in the same format of the input data.
- All results shall have same resolution, cellsize, extent, and coordinate system as the input grids.

2) Results nomenclature
- All results shall be contained in a single directory, called "results_[Yourname]" (e.g. "results_Farinotti")
- The directory shall be compressed (.zip) before submission. The name of the uploaded file will thus be: "results_[Yourname].zip"
- The individual files in the folder shall be named
   "thickness_[NameOfTheGlacier]_[Yourname]_exp[XX].asc"
  where
   [NameOfTheGlacier] = Name of the glacier, as defined by the input files (e.g. "NorthGlacier")
   [Yourname] = Your family name (e.g. "Farinotti")
   [XX] = Number of the experiment, as given in the "experiments_key"-file, the result is referring to (e.g. "02")
  A valid filename will thus look like, e.g., "thickness_NorthGlacier_Farinotti_exp02.asc".
- For a test case to be labelled as "considered", results for all 16 experiments will need to be submitted.
  Since there are three test cases defined as "compulsory", the submitted directory will thus contain 3*16=48 files at least.

3) Results upload
- To uplod your results, simply drag-and-drop your .zip-repository to the following address:
 https://drive.switch.ch/index.php/s/hovc5mJjSztoLVd
- When uploaded, please send a notification email to: daniel.farinotti@ethz.ch
=#

"""
    filepaths(gl::BM.Glacier, pl, expnr, rm_iv, dir)

Returns saving filenames including their full path.
The path is `joinpath(pl.data_root, dir).

Returns file names for h, h-sigma, iv, iv-sigma, and jld.
"""
function filepaths(gl::BM.Glacier, pl, expnr, rm_iv, dir)
    expnrs = expnr>9 ? Symbol("exp$expnr") : Symbol("exp0$expnr")
    glaciername = BM.getname(gl)
    fl = rm_iv ?
        joinpath(pl.data_root, dir, "thickness_$(glaciername)_Werder_$(expnrs)_rmIV") :
        joinpath(pl.data_root, dir, "thickness_$(glaciername)_Werder_$(expnrs)")
    # make the directory if it does not exist
    !isdir(pl.data_root) && error("Data-root dir does not exist: $(pl.data_root)")
    !isdir(joinpath(pl.data_root, dir)) && mkpath(joinpath(pl.data_root, dir))
    return fl*".asc", fl*"-sigma.asc", fl*"-iv.asc", fl*"-iv-sigma.asc", fl*".jld"
end

"""
    save_itmix2_run(hs2d, hs2d_sigma, gl, pl, iv2d=nothing, iv2d_sigma=nothing)

Saves the thickness map according to ITMIX2 specs.
"""
function save_itmix2_run(hs2d, hs2d_sigma, gl, pl, ivs2d=nothing, ivs2d_sigma=nothing;
                         dir="results_Werder")
    expnr, rm_iv = gl.misc[:expnr], gl.misc[:rm_iv]
    fln, flns, fln_iv, fln_ivs, fln_jld = filepaths(gl, pl, expnr, rm_iv, dir)

    downsample = get(pl.dataset_opts[BM.DataKind], :grid_downsample_step, 1)
    hs2d, hs2d_sigma = copy(hs2d), copy(hs2d_sigma)
    hs2d[.!gl.glaciermask] = BM.FILL
    hs2d_sigma[.!gl.glaciermask] = BM.FILL
    h = VAWTools.Gridded(gl.dem.x, gl.dem.y, hs2d)
    hs = VAWTools.Gridded(gl.dem.x, gl.dem.y, hs2d_sigma)
    if ivs2d!=nothing
        ivs2d, ivs2d_sigma = copy(ivs2d), copy(ivs2d_sigma)
        ivs2d[.!gl.glaciermask] = BM.FILL
        ivs2d_sigma[.!gl.glaciermask] = BM.FILL
        iv = VAWTools.Gridded(gl.dem.x, gl.dem.y, ivs2d)
        ivs = VAWTools.Gridded(gl.dem.x, gl.dem.y, ivs2d_sigma)
    end
    if downsample!=1
        # up-sample output
        dem_orig = VAWTools.read_agr(pl.dataset_files[BM.DEMData][1])
        h = VAWTools.upsample(h, dem_orig.x, dem_orig.y)
        hs = VAWTools.upsample(hs, dem_orig.x, dem_orig.y)
        if ivs2d!=nothing
            iv = VAWTools.upsample(iv, dem_orig.x, dem_orig.y)
            ivs = VAWTools.upsample(ivs, dem_orig.x, dem_orig.y)
        end
    end
    VAWTools.write_agr(h, fln, NA_g=BM.FILL, NA_agr=-9999.0)
    VAWTools.write_agr(hs, flns, NA_g=BM.FILL, NA_agr=-9999.0)
    if ivs2d!=nothing
        VAWTools.write_agr(iv, fln_iv, NA_g=BM.FILL, NA_agr=-9999.0)
        VAWTools.write_agr(ivs, fln_ivs, NA_g=BM.FILL, NA_agr=-9999.0)
    end
    # TODO make sub-folders with computer-name-date
    return fln, flns, fln_iv, fln_ivs, fln_jld
end
end #module ITMIX2


mutable struct RelaxedSequencingNode{T<:Real,U<:Int,V<:Int}<:Node
    uuid                    ::UUID
    length                  ::U
    state                   ::V
    allUp                   ::BitVector
    allDown                 ::BitVector
    someUp                  ::BitVector
    someDown                ::BitVector
    lengthToRoot            ::T
    lengthToTerminal        ::T
    bestSeq                 ::Vector{V}
    exact                   ::Bool
    arcsIn                  ::Vector{ExactArc}
    arcsOut                 ::Vector{ExactArc}
end

#change hashing of nodes to use uuids
Base.hash(node::RelaxedSequencingNode) = hash(node.uuid)
Base.:(==)(node1::RelaxedSequencingNode, node2::RelaxedSequencingNode) = node1.uuid == node2.uuid
Base.length(node::RelaxedSequencingNode) = node.length
#do UUIDs automatically
function RelaxedSequencingNode(pathLength::U,state::V,allUp::BitVector,allDown::BitVector,someUp::BitVector,someDown::BitVector,lengthToRoot::T,lengthToTerminal::T,bestSeq::Vector{V}, exact::Bool,arcsIn::Vector{ExactArc},arcsOut::Vector{ExactArc})where{T<:Real,U<:Int,V<:Int}
    return RelaxedSequencingNode(uuid4(), pathLength,state,allUp,allDown,someUp,someDown,lengthToRoot,lengthToTerminal,bestSeq, exact,arcsIn,arcsOut)
end

"""
    deepCopy(node::RelaxedSequencingNode)

Deep copy a node

# Arguments
- `node::RelaxedSequencingNode`: the node
"""
function deepCopy(node::RelaxedSequencingNode)
    return RelaxedSequencingNode(
        node.uuid, #uuid
        node.length, #path length
        node.state, #current state
        copy(node.allUp), #allUp
        copy(node.allDown), #allDown
        copy(node.someUp), #someUp
        copy(node.someDown), #someDown
        node.lengthToRoot, #to root
        node.lengthToTerminal, #to terminal
        copy(node.bestSeq), #best Seq
        node.exact, #exact
        deepCopy(node.arcsIn), #arcsin
        deepCopy(node.arcsOut) #arcsout
    )
end

Base.show(io::IO,node::RelaxedSequencingNode) = Base.print(
    io,
    "\n   pathLength: ", node.length, "\n",
    "   state: ", node.state, "\n",
    "   allUp: ", node.allUp, "\n",
    "   allDown: ", node.allDown, "\n",
    "   someUp: ", node.someUp, "\n",
    "   someDown: ", node.someDown, "\n",
    "   lengthToRoot: ", node.lengthToRoot, "\n",
    "   lengthToTerminal: ", node.lengthToTerminal, "\n",
    "   best Sequence: ", node.bestSeq, "\n",
    "   isExact: ", node.exact, "\n",
    "   arcsIn: ", showInArcOrigins(node.arcsIn), "\n",
    "   arcsOut: ", showOutArcAssignments(node.arcsOut), "\n",
)

function showOutArcAssignments(arcs::Vector{ExactArc})
    str = ""
    for arc in arcs
        str = string(str, arc.assignment, " ")
    end
    return str
end

function showInArcOrigins(arcs::Vector{ExactArc})
    str = ""
    for arc in arcs
        str = string(str, arc.origin.state, " ")
    end
    return str
end

"""
    createRootRelaxedSequencingNodeFromModel(model::SequencingModel)

Create a root node for a relaxed DD from the sequencing model

# Arguments
- `model::SequencingModel`: the sequencing problem to be evaluated
"""
function createRootRelaxedSequencingNodeFromModel(model::SequencingModel)
    restrictedNode = createRestrictedRootNode(model)
    relaxedNode = RelaxedSequencingNode(
        length(restrictedNode.path), #length
        last(restrictedNode.path), #current state
        falses(length(model)), #allUp
        copy(restrictedNode.visited), #allDown
        falses(length(model)), #someUp
        copy(restrictedNode.visited), #someDown
        0, #to root
        0, #to terminal
        copy(restrictedNode.path),
        true, #exact
        Vector{ExactArc}(), #arcsin
        Vector{ExactArc}() #arcsout
    )
    return relaxedNode, restrictedNode.domain
end

"""
    createPartialRelaxedNode(;length::T=0, state::U=-1, allUp::BitVector=BitVector(), allDown::BitVector=BitVector(), someUp::BitVector=BitVector(), someDown::BitVector=BitVector(), bestSeq::Vector{U}=Vector{Int}(), lengthToRoot::V=0, lengthToTerminal::V=0, exact::Bool=false, arcsIn::Vector{ExactArc}= Vector{ExactArc}(), arcsOut::Vector{ExactArc}= Vector{ExactArc}())where{T<:Int, U<:Int, V<:Real}

Create a node using default values for any combination of fields

# Arguments
- `length::T`: the length field
- `state::U`: the state field
- `allUp::BitVector`: the allUp field
- `allDown::BitVector`: the allDown field
- `someUp::BitVector`: the someUp field
- `someDown::BitVector`: the someDown field
- `bestSeq::Vector{Int}`: the bestSeq field
- `lengthToRoot::V`: the lengthToRoot field
- `lengthToTerminal::V`: the lengthToTerminal field
- `exact::Bool`: the exact field
- `arcsIn::Vector{ExactArc}`: the arcsIn field
- `arcsOut::Vector{ExactArc}`: the arcsOut field
"""
function createPartialRelaxedNode(;length::T=0, state::U=-1, allUp::BitVector=BitVector(), allDown::BitVector=BitVector(), someUp::BitVector=BitVector(), someDown::BitVector=BitVector(), bestSeq::Vector{U}=Vector{Int}(), lengthToRoot::V=0, lengthToTerminal::V=0, exact::Bool=false, arcsIn::Vector{ExactArc}= Vector{ExactArc}(), arcsOut::Vector{ExactArc}= Vector{ExactArc}())where{T<:Int, U<:Int, V<:Real}
    return RelaxedSequencingNode(
        length, #length
        state, #current state
        allUp, #allUp
        allDown, #allDown
        someUp, #someUp
        someDown, #someDown
        lengthToRoot, #to root
        lengthToTerminal, #to terminal
        bestSeq,
        exact, #exact
        arcsIn, #arcsin
        arcsOut #arcsout
    )
end

"""
    getParents(node::RelaxedSequencingNode)

Find and return the parents of a node

# Arguments
- `node::RelaxedSequencingNode`: the node whose parents to find
"""
function getParents(node::RelaxedSequencingNode)
    parents = Set{RelaxedSequencingNode}()
    for arc in node.arcsIn
        push!(parents, arc.origin)
    end
    return parents
end

"""
    getChildren(node::RelaxedSequencingNode)

Find and return the children
# Arguments
- `node::RelaxedSequencingNode`: the node whose children to find
"""
function getChildren(node::RelaxedSequencingNode)
    children = Set{RelaxedSequencingNode}()
    for arc in node.arcsOut
        push!(children, arc.destination)
    end
    return children
end

"""
    updateNodeDownVariables!(node::RelaxedSequencingNode,model::SequencingModel)

Update the state variable for allDown and someDown and lengthToRoot and exact

# Arguments
- `node::RelaxedSequencingNode`: the node to update
- `model::SequencingModel`: the sequencing problem being evaluated
"""
function updateNodeDownVariables!(node::RelaxedSequencingNode,model::SequencingModel)
    node.allDown = falses(length(model))
    node.someDown = falses(length(model))
    node.exact = true

    node.someDown[node.state] = true
    node.allDown[node.state] = true
    #update the best path variables
    bestArc, node.lengthToRoot = findBestArc(model, node.arcsIn)
    node.bestSeq = push!(copy(bestArc.origin.bestSeq), node.state)
    #union domains with parents and check for exactness
    for parent in getParents(node)
        node.someDown = node.someDown .| parent.someDown
        node.allDown = node.allDown .| parent.allDown
        if !parent.exact
            node.exact = false
        end
    end
    #check each all down to make sure its valid
    for i in 1:length(node.allDown)
        if node.allDown[i]
            for arc in node.arcsIn
                if arc.assignment != i
                    if !arc.origin.allDown[i]
                        node.allDown[i]=false
                        break
                    end
                end
            end
        end
    end
    #final check for exactness (first part is just a check to see if the second part should run)
    if node.exact && (node.someDown != node.allDown || node.length!=sum(node.allDown))
        node.exact = false
    end
end

"""
    updateNodeUpVariables!(node::RelaxedSequencingNode,model::SequencingModel)

Update the state variable for someUp, allUp, and lengthToTerminal

# Arguments
- `node::RelaxedSequencingNode`: the node to update
- `model::SequencingModel`: the sequencing problem being evaluated
"""
function updateNodeUpVariables!(node::RelaxedSequencingNode,model::SequencingModel)
    node.allUp = falses(length(model))
    node.someUp = falses(length(model))
    node.lengthToTerminal = node.arcsOut[1].value + node.arcsOut[1].destination.lengthToTerminal
    for arc in node.arcsOut
        node.someUp[arc.assignment] = true
        node.allUp[arc.assignment] = true
        #update the length to terminal
        if isBetterSolutionValue(model, arc.value + arc.destination.lengthToTerminal, node.lengthToTerminal)
            node.lengthToTerminal = arc.value + arc.destination.lengthToTerminal
        end
    end

    #union domains with children
    for child in getChildren(node)
        node.someUp = node.someUp .| child.someUp
        node.allUp = node.allUp .| child.allUp
    end
    #check each all up to make sure its valid
    for i in 1:length(node.allUp)
        if node.allUp[i]
            for arc in node.arcsOut
                if arc.assignment != i
                    if !arc.destination.allUp[i]
                        node.allUp[i]=false
                        break
                    end
                end
            end
        end
    end
end

"""
    createSquareStartDD(model::SequencingModel;rootNode::Union{Nothing, RelaxedSequencingNode}=nothing, masterDomain::BitVector=BitVector())

Create an initial dd from the model where all the arcs have exact values

# Arguments
- `model::SequencingModel`: the sequencing problem being evaluated
- `rootNode::Union{Nothing, RelaxedSequencingNode}`: a root node if not starting from the model
- `masterDomain::BitVector`: the domain of the root node if there is a root node, must be included if root node is included
"""
function createSquareStartDD(model::SequencingModel;rootNode::Union{Nothing, RelaxedSequencingNode}=nothing, masterDomain::BitVector=BitVector())
    if isnothing(rootNode)
        rootNode, masterDomain = createRootRelaxedSequencingNodeFromModel(model)
    end
    relaxedDD = Vector{Vector{RelaxedSequencingNode}}()
    push!(relaxedDD, [rootNode])

    #for each layer going down
    terminalNode = createPartialRelaxedNode(
        length = length(model),
        state=0,
        allUp=falses(length(masterDomain)),
        allDown=falses(length(masterDomain)),
        someUp=falses(length(masterDomain)),
        someDown=falses(length(masterDomain))
    )
    if hasEnd(model)
        lastLayerLoopIndex = length(model)-length(rootNode)
        terminalNode.state = model.seqEnd
    else
        lastLayerLoopIndex = length(model)-length(rootNode)+1
    end
    for layer in 2:lastLayerLoopIndex
        push!(relaxedDD,Vector{RelaxedSequencingNode}())
        #new node starts as a state and length
        for (index, inDomain) in enumerate(masterDomain)
            if inDomain
                newNode = createPartialRelaxedNode(length = layer, state=index)
                push!(relaxedDD[layer],newNode)
            end
        end
        #update domain based on precedence
        if model.precedenceConstraints
            for (i, inDomain) in enumerate(masterDomain)
                if !inDomain && (!hasStart(model) || i!=model.seqStart)
                    #add satisfied precedence constraints to the domain
                    if layer >= get(model.precedenceNumbers,i,0)
                        masterDomain[i] = true
                    end
                end
            end
        end
    end#layer iteration
    push!(relaxedDD, [terminalNode])

    #add all arcs except terminal arcs and update down values
    for layer in 1:length(relaxedDD)-2
        for parent in relaxedDD[layer]
            if layer != 1
                updateNodeDownVariables!(parent,model)
            end
            for child in relaxedDD[layer+1]
                if hasKey(model.objective, parent.state, child.state) && !parent.allDown[child.state]
                    newArc = ExactArc(parent, child, child.state, evaluateDecision(model.objective, parent.state, child.state))
                    push!(parent.arcsOut, newArc)
                    push!(child.arcsIn, newArc)
                end
            end
        end
        toRemove = Vector{RelaxedSequencingNode}()
        for child in relaxedDD[layer+1]
            if isempty(child.arcsIn)
                push!(toRemove, child)
            end
        end
        filter!(x->!(x in toRemove), relaxedDD[layer+1])
    end
    #add arcs to terminal
    if hasEnd(model)
        for parent in relaxedDD[length(relaxedDD)-1]
            if hasKey(model.objective, parent.state, model.seqEnd)
                newArc = ExactArc(parent, terminalNode, model.seqEnd, evaluateDecision(model.objective, parent.state, model.seqEnd))
                push!(parent.arcsOut, newArc)
                push!(terminalNode.arcsIn, newArc)
            end
            updateNodeDownVariables!(parent,model)
        end
    else
        for parent in relaxedDD[length(relaxedDD)-1]
            newArc = ExactArc(parent, terminalNode, -1, 0)
            push!(parent.arcsOut, newArc)
            push!(terminalNode.arcsIn, newArc)
            updateNodeDownVariables!(parent,model)
        end
    end
    updateNodeDownVariables!(terminalNode,model)

    #update all the up variables
    for layer in Iterators.reverse(1:length(relaxedDD)-1)
        for node in relaxedDD[layer]
            updateNodeUpVariables!(node,model)
        end
    end

    return relaxedDD
end


module LineNumbers

export SourceFile, compute_line, LineBreaking
import Base: getindex, setindex!, length

# Offsets are 0 based
struct SourceFile
    data::Vector{UInt8}
    offsets::Vector{UInt64}
end
length(file::SourceFile) = length(file.offsets)

function SourceFile(data)
    offsets = UInt64[0]
    buf = IOBuffer(data)
    local line
    while !eof(buf)
        line = readuntil(buf,'\n')
        !eof(buf) && push!(offsets, position(buf))
    end
    if !isempty(offsets) && line[end] == '\n'
        push!(offsets, position(buf))
    end
    SourceFile(Vector{UInt8}(data),offsets)
end

function compute_line(file::SourceFile, offset)
    ind = searchsortedfirst(file.offsets, offset)
    ind <= length(file.offsets) && file.offsets[ind] == offset ? ind : ind - 1
end

function getindex(file::SourceFile, line::Int)
    if line == length(file.offsets)
        return file.data[(file.offsets[end]+1):end]
    else
        # - 1 to skip the '\n'
        return file.data[(file.offsets[line]+1):(file.offsets[line+1]-1)]
    end
end
getindex(file::SourceFile, arr::AbstractArray) = [file[x] for x in arr]

# LineBreaking

"""
Indexing adaptor to map from a flat byte offset to a `[line][offset]` pair.
Optionally, off may specify a byte offset relative to which the line number and
offset should be computed
"""
struct LineBreaking{T}
    off::UInt64
    file::SourceFile
    obj::T
end

function indtransform(lb::LineBreaking, x::Int)
    offline = compute_line(lb.file, lb.off)
    line = compute_line(lb.file, x)
    lineoffset = lb.file.offsets[line]
    off = x - lineoffset + 1
    if lineoffset < lb.off
        off -= lb.off - lineoffset
    end
    (line - offline + 1), off
end

function getindex(lb::LineBreaking, x::Int)
    l, o = indtransform(lb, x)
    lb.obj[l][o]
end

function setindex!(lb::LineBreaking, y, x::Int)
    l, o = indtransform(lb, x)
    lb.obj[l][o] = y
end
function setindex!(lb::LineBreaking, y, x::AbstractArray)
    for i in x
        lb[i] = y
    end
end

end # module


using DataDeps

register(DataDep(
  "hill-valley",
  "http://archive.ics.uci.edu/ml/datasets/Hill-Valley",
  "http://archive.ics.uci.edu/ml/machine-learning-databases/hill-valley/Hill_Valley_without_noise_Training.data",
  "af232ff5258abeb5b1c61974648fd00bc0d37a545de33a2c3493692fc20330d5",
  post_fetch_method=(path -> begin
    UCIData.process_dataset(path,
        target_index=101,
        feature_indices=1:100,
        skipto=2,
    )
  end),
))


"""
`DataValueArray{T, N}` is an efficient alternative to `Array{DataValue{T}, N}`.
"""
struct DataValueArray{T,N} <: AbstractArray{DataValue{T},N}
    values::Array{T,N}
    isna::Array{Bool,N}

    function DataValueArray{T,N}(d::AbstractArray{T, N}, m::AbstractArray{Bool, N}) where {T,N}
        if size(d) != size(m)
            msg = "values and missingness arrays must be the same size"
            throw(ArgumentError(msg))
        end
        new(d, m)
    end
end

const DataValueVector{T} = DataValueArray{T, 1}
const DataValueMatrix{T} = DataValueArray{T, 2}


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Tests related to ECEF to ECEF transformations using satellite state vector.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# Get the current EOP Data.
#
# TODO: The EOP data obtained from IERS website does not match the values in the
# examples in [1]. However, this should be enough, because 1) the individual
# functions at the low level are tested using the same values of [1], and 2) the
# difference is smaller than 30 cm.

eop_iau1980  = read_iers_eop("./eop_IAU1980.txt",  Val(:IAU1980))
eop_iau2000a = read_iers_eop("./eop_IAU2000A.txt", Val(:IAU2000A))

# File: ./src/transformations/sv_ecef_to_ecef.jl
# ==============================================

# Functions: sv_ecef_to_ecef
# --------------------------

# ITRF <=> PEF
# ============

################################################################################
#                                  IAU-76/FK5
################################################################################

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-15: Performing IAU-76/FK5 reduction.
#
# According to this example and Table 3-6, using:
#
#   UTC    = April 6, 2004, 07:51:28.386009
#   r_itrf = -1033.4793830    i + 7901.2952754    j + 6380.3565958    k [km]
#   v_itrf =    -3.225636520  i -    2.872451450  j +    5.531924446  k [km/s]
#
# one gets:
#
#   r_pef  = -1033.47503130   i + 7901.30558560   j + 6380.34453270   k [km]
#   v_pef  =    -3.2256327470 i -    2.8724425110 j +    5.5319312880 k [km/s]
#
################################################################################

@testset "Function sv_ecef_to_ecef ITRF <=> PEF" begin
    JD_UTC = date_to_jd(2004, 4, 6, 7, 51, 28.386009)

    ## ITRF => PEF
    ## ===========

    r_itrf  = [-1033.4793830; 7901.2952754; 6380.3565958]
    v_itrf  = [-3.225636520; -2.872451450; +5.531924446]
    sv_itrf = orbsv(JD_UTC, r_itrf, v_itrf)
    sv_pef  = sv_ecef_to_ecef(sv_itrf, ITRF(), PEF(), JD_UTC, eop_iau1980)

    @test sv_pef.t === JD_UTC

    @test sv_pef.r[1] ≈ -1033.47503130 atol=1e-4
    @test sv_pef.r[2] ≈ +7901.30558560 atol=1e-4
    @test sv_pef.r[3] ≈ +6380.34453270 atol=1e-4

    @test sv_pef.v[1] ≈ -3.2256327470  atol=1e-7
    @test sv_pef.v[2] ≈ -2.8724425110  atol=1e-7
    @test sv_pef.v[3] ≈ +5.5319312880  atol=1e-7

    ## PEF => ITRF
    ## ===========

    r_pef   = [-1033.47503130; 7901.30558560; 6380.34453270]
    v_pef   = [-3.2256327470; -2.8724425110; +5.5319312880]
    sv_pef  = orbsv(JD_UTC, r_pef, v_pef)
    sv_itrf = sv_ecef_to_ecef(sv_pef, PEF(), ITRF(), JD_UTC, eop_iau1980)

    @test sv_itrf.t === JD_UTC

    @test sv_itrf.r[1] ≈ -1033.4793830 atol=1e-4
    @test sv_itrf.r[2] ≈ +7901.2952754 atol=1e-4
    @test sv_itrf.r[3] ≈ +6380.3565958 atol=1e-4

    @test sv_itrf.v[1] ≈ -3.225636520  atol=1e-7
    @test sv_itrf.v[2] ≈ -2.872451450  atol=1e-7
    @test sv_itrf.v[3] ≈ +5.531924446  atol=1e-7

end

################################################################################
#                                IAU-2006/2010
################################################################################

## ITRF <=> TIRS
## =============

################################################################################
#                                 Test Results
################################################################################
#
# Scenario 01
# ===========
#
# Example 3-14: Performing an IAU-2000 reduction [1, p. 220]
#
# According to this example and Table 3-6, using:
#
#   UTC    = April 6, 2004, 07:51:28.386009
#   r_itrf = -1033.4793830    i + 7901.2952754    j + 6380.3565958    k [km]
#   v_itrf =    -3.225636520  i -    2.872451450  j +    5.531924446  k [km/s]
#
# one gets the following:
#
#   r_tirs  = -1033.47503120   i + 7901.30558560   j + 6380.34453270   k [km]
#   v_tirs  =    -3.2256327470 i -    2.8724425110 j +    5.5319312880 k [km/s]
#
################################################################################

@testset "Function sv_ecef_to_ecef ITRF <=> TIRS" begin
    JD_UTC = date_to_jd(2004, 4, 6, 7, 51, 28.386009)

    ## ITRF => TIRS
    ## ===========

    r_itrf = [-1033.4793830; 7901.2952754; 6380.3565958]
    v_itrf = [-3.225636520; -2.872451450; +5.531924446]
    sv_itrf = orbsv(JD_UTC, r_itrf, v_itrf)
    sv_tirs = sv_ecef_to_ecef(sv_itrf, ITRF(), TIRS(), JD_UTC, eop_iau2000a)

    @test sv_tirs.t === JD_UTC

    @test sv_tirs.r[1] ≈ -1033.47503120 atol=1e-4
    @test sv_tirs.r[2] ≈ +7901.30558560 atol=1e-4
    @test sv_tirs.r[3] ≈ +6380.34453270 atol=1e-4

    @test sv_tirs.v[1] ≈ -3.2256327470  atol=1e-7
    @test sv_tirs.v[2] ≈ -2.8724425110  atol=1e-7
    @test sv_tirs.v[3] ≈ +5.5319312880  atol=1e-7

    ## TIRS => ITRF
    ## ===========

    r_tirs  = [-1033.47503120; 7901.30558560; 6380.34453270]
    v_tirs  = [-3.2256327470; -2.8724425110; +5.5319312880]
    sv_tirs = orbsv(JD_UTC, r_tirs, v_tirs)
    sv_itrf = sv_ecef_to_ecef(sv_tirs, TIRS(), ITRF(), JD_UTC, eop_iau2000a)

    @test sv_itrf.t === JD_UTC

    @test sv_itrf.r[1] ≈ -1033.4793830 atol=1e-4
    @test sv_itrf.r[2] ≈ +7901.2952754 atol=1e-4
    @test sv_itrf.r[3] ≈ +6380.3565958 atol=1e-4

    @test sv_itrf.v[1] ≈ -3.225636520  atol=1e-7
    @test sv_itrf.v[2] ≈ -2.872451450  atol=1e-7
    @test sv_itrf.v[3] ≈ +5.531924446  atol=1e-7

end


# load_constraints.jl
# Gets city information from manhattan.jld object and travel time information from Travel time CSV
# Authored by Arthur J Delarue on 6/8/15

function saveManhattan(pb::TaxiProblem, name::AbstractString; compress=false)
  save("Cities/Saved/$name.jld", "pb", pb, compress=compress)
end

function loadManhattan(name::AbstractString)
  pb = load("Cities/Saved/$name.jld","pb")
  return pb
end

function saveRoadTimes(times::SparseMatrixCSC{Float64,Int}, name::AbstractString; compress=false)
  save("Outputs/$name.jld", "times", times, compress=compress)
end

function loadRoadTimes(name::AbstractString)
  times = load("Outputs/$name.jld","times")
  return times
end

MANHATTAN_NODES = 6134
# Speed limits in km/h from OpenStreetMap/src/speeds.jl
URBAN_SPEEDS = [95,72,48,32,22,12,8,5]

function loadCityGraph(;fromScratch::Bool=false, useShortestPaths::Bool=true,fixSpeeds=false)
	"""
	Load Manhattan object and return it
	"""
	println("**** Loading city graph ****")
	if fromScratch
		manhattan = Manhattan(sp=useShortestPaths)
		# if useShortestPaths
		# 	sp = manhattan.sp
		# else
		# 	sp = parallelShortestPathsAuto(manhattan.network, manhattan.roadTime, manhattan.roadCost)
		# 	manhattan.sp = sp
		# end
		if fixSpeeds
			# Make sure the speeds are in URBAN_SPEEDS
			# If not match speed to closest actual speed
			for i=1:MANHATTAN_NODES,j=1:MANHATTAN_NODES
				if manhattan.distances[i,j] != 0
					# Don't forget to convert from m/s to km/h
					if manhattan.roadTime[i,j] != 0
						speed = 3.6 * manhattan.distances[i,j]/manhattan.roadTime[i,j]
					else
						speed = 3.6 * manhattan.distances[i,j]/0.1
					end
					road_type = indmin(abs(URBAN_SPEEDS - speed))
					if manhattan.roadTime[i,j] == 0
						println(manhattan.roadTime[i,j], "\t", manhattan.distances[i,j], "\t", speed, "\t", URBAN_SPEEDS[road_type])
					end
					manhattan.roadTime[i,j] = manhattan.distances[i,j]/URBAN_SPEEDS[road_type]
				end
			end
			shortestPaths!(manhattan)
		end
		saveManhattan(manhattan, "full_manhattan")
	else
		manhattan = loadManhattan("full_manhattan")
	end
	return manhattan
end

function loadTravelTimeData(;radius::Int=40, times::AbstractString= "1214", min_rides::Int=4, num_nodes::Int=MANHATTAN_NODES, preprocess::Bool=false, num_clusters::Int64 = 50, sampleSize::Int64=10000)
	# Hardcoded for simplicity. Map of Manhattan not likely to change anytime soon
	"""
	Return array of travel times between pairs of nodes A,B obtained from NYC dataset
	"""
	println("**** Loading travel times ****")
	# Load Manhattan highway nodes
	highwayNodes = load("Cities/Saved/highwayNodes.jld", "highwayNodes")
	if preprocess
		dataFile = "../Travel_times/travel_times_r$(radius)_wd_$(times)_clust$(num_clusters)_rides$(sampleSize)_minr$(min_rides).csv"
	else
		dataFile = "../Travel_times/travel_times_r$(radius)_wd_$(times).csv"
	end
	println("-- Loading from $dataFile...")
	data = readcsv(dataFile)
	nodePairs = data[:,1]
	average = data[:,2]
	minRides = data[:,5]

	# Put travel time data in an array
	travelTimes = zeros(num_nodes,num_nodes)
	for (i, element) in enumerate(nodePairs)
		node = split(element, ";")
		# Remember +1 because Python indexes from 0 but Julia from 1
		src = int(node[1]) + 1
		dest = int(node[2]) + 1
		if minRides[i] >= min_rides && average[i] >= 60 && src != dest && !(src in highwayNodes) && !(dest in highwayNodes)
			travelTimes[src, dest] = average[i]
		end
	end
	return travelTimes
end

function loadNewTravelTimeData(;trainOrTest::AbstractString="training", radius::Int=40, times::AbstractString="1214", num_nodes::Int=MANHATTAN_NODES, preprocess::Bool=false, num_clusters::Int64=50, sampleSize::Int64=10000, min_rides::Int=4, saveTestingMatrix::Bool=false, loadTestingMatrixDirectly=false)
	"""
	Return array of travel times between pairs of nodes A,B obtained from NYC dataset, as well as number of rides for each node pair.
	"""
	println("**** Loading $(trainOrTest) data ****")
	# Load Manhattan highway nodes
	highwayNodes = load("Cities/Saved/highwayNodes.jld", "highwayNodes")
	if preprocess
		dataFile = "../Travel_times/$(trainOrTest)_r$(radius)_wd_$(times)_clust$(num_clusters)_rides$(sampleSize)_minr$(min_rides).csv"
	else
		dataFile = "../Travel_times/$(trainOrTest)_r$(radius)_wd_$(times).csv"
	end
	# Special case: load testing matrix directly
	if preprocess == false && loadTestingMatrixDirectly == true
		println("-- Loading directly from JLD archive")
		travelTimes = load("../Travel_times/$(trainOrTest)_r$(radius)_wd_$(times).jld", "travelTimes")
		numRides = load("../Travel_times/$(trainOrTest)_r$(radius)_wd_$(times).jld", "numRides")
		return travelTimes, numRides
	end
	println("-- Loading from $dataFile...")
	data = readcsv(dataFile)
	nodePairs = data[:,1]
	averageTime = data[:,2]
	minRides = data[:,5]
	travelTimes = zeros(num_nodes,num_nodes)
	numRides = zeros(Int, (num_nodes, num_nodes))
	for (i, element) in enumerate(nodePairs)
		node = split(element, ";")
		src = int(node[1]) + 1
		dest = int(node[2]) + 1
		if src != dest && !(src in highwayNodes) && !(dest in highwayNodes) && averageTime[i] >= 60
			travelTimes[src, dest] = averageTime[i]
			numRides[src, dest] = minRides[i]
		end
	end
	if preprocess == false && saveTestingMatrix == true
		println("-- Saving to JLD file for speed")
		save("../Travel_times/$(trainOrTest)_r$(radius)_wd_$(times).jld", "travelTimes", travelTimes, "numRides", numRides)
	end
	return travelTimes, numRides
end

function chooseConstraints(travelTimes::Array{Float64, 2}, numRides::Array{Int, 2}; num_nodes::Int = MANHATTAN_NODES, sample_size::Int = 5000)
	"""
	Given a matrix of travel times and number of rides, randomly selects sample_size nonzero entries
	"""
	newTravelTimes = zeros(num_nodes, num_nodes)
	newNumRides = zeros(Int, (num_nodes, num_nodes))
	for i = 1:sample_size
		startNode = rand(1:num_nodes)
		endNode = rand(1:num_nodes)
		while true
			# Keep them if not added yet
			if numRides[startNode, endNode] > 0 && newNumRides[startNode, endNode] == 0
				break
			end
			startNode = rand(1:num_nodes)
			endNode = rand(1:num_nodes)
		end
		newNumRides[startNode, endNode] = numRides[startNode, endNode]
		newTravelTimes[startNode, endNode] = travelTimes[startNode, endNode]
	end
	return newTravelTimes, newNumRides
end


const turing_key_map = Dict(
    "hamiltonian_energy" => "energy",
    "hamiltonian_energy_error" => "energy_error",
    "is_adapt" => "tune",
    "max_hamiltonian_energy_error" => "max_energy_error",
    "nom_step_size" => "step_size_nom",
    "numerical_error" => "diverging",
)
const stan_key_map = Dict(
    "accept_stat__" => "acceptance_rate",
    "divergent__" => "diverging",
    "energy__" => "energy",
    "lp__" => "lp",
    "n_leapfrog__" => "n_steps",
    "stepsize__" => "step_size",
    "treedepth__" => "tree_depth",
)
const stats_key_map = merge(turing_key_map, stan_key_map)

"""
    reshape_values(x::AbstractArray) -> AbstractArray

Convert from `MCMCChains` variable values with dimensions `(ndraw, nchain, size...)` to
ArviZ's expected `(nchain, ndraw, size...)`.
"""
reshape_values(x::AbstractArray{T,N}) where {T,N} = permutedims(x, (2, 1, (3:N)...))

headtail(x) = x[1], x[2:end]

function split_locname(name)
    name = replace(name, r"[\[,]" => '.')
    name = replace(name, ']' => "")
    name, loc = headtail(split(name, '.'))
    isempty(loc) && return name, ()
    loc = tryparse.(Int, loc)
    Nothing <: eltype(loc) && return name, ()
    return name, tuple(loc...)
end

function varnames_locs_dict(loc_names, loc_str_to_old)
    vars_to_locs = Dict()
    for loc_name in loc_names
        var_name, loc = split_locname(loc_name)
        if var_name ∉ keys(vars_to_locs)
            vars_to_locs[var_name] = ([loc_str_to_old[loc_name]], [loc])
        else
            push!(vars_to_locs[var_name][1], loc_str_to_old[loc_name])
            push!(vars_to_locs[var_name][2], loc)
        end
    end
    # ensure that elements are ordered in the same order as they would be iterated
    for loc_name_locs in values(vars_to_locs)
        perm = sortperm(loc_name_locs[2]; by=CartesianIndex)
        permute!(loc_name_locs[1], perm)
        permute!(loc_name_locs[2], perm)
    end
    return vars_to_locs
end

function attributes_dict(chns::Chains)
    info = delete(chns.info, :hashedsummary)
    return Dict{String,Any}((string(k), v) for (k, v) in pairs(info))
end

function section_dict(chns::Chains, section)
    ndraws, _, nchains = size(chns)
    loc_names_old = getfield(chns.name_map, section) # old may be Symbol or String
    loc_names = string.(loc_names_old)
    loc_str_to_old = Dict(
        name_str => name_old for (name_str, name_old) in zip(loc_names, loc_names_old)
    )
    vars_to_locs = varnames_locs_dict(loc_names, loc_str_to_old)
    vars_to_arrays = Dict{String,Array}()
    for (var_name, names_locs) in vars_to_locs
        loc_names, locs = names_locs
        sizes = reduce((a, b) -> max.(a, b), locs)
        ndim = length(sizes)
        var_views = (@view(chns.value[:, n, :]) for n in loc_names)
        oldarr = reshape_values(replacemissing(convert(Array, cat(var_views...; dims=3))))
        if iszero(ndim)
            arr = dropdims(oldarr; dims=3)
        else
            arr = Array{Union{typeof(NaN),eltype(oldarr)}}(undef, nchains, ndraws, sizes...)
            fill!(arr, NaN)
            for i in eachindex(locs)
                arr[:, :, locs[i]...] = oldarr[:, :, i]
            end
        end
        vars_to_arrays[var_name] = arr
    end
    return vars_to_arrays
end

function chains_to_dict(
    chns::Chains; ignore=String[], section=:parameters, rekey_fun=identity
)
    section in sections(chns) || return Dict()
    chns_dict = section_dict(chns, section)
    removekeys!(chns_dict, ignore)
    return rekey_fun(chns_dict)
end

"""
    convert_to_inference_data(obj::Chains; group = :posterior, kwargs...) -> InferenceData

Convert the chains `obj` to an [`InferenceData`](@ref) with the specified `group`.

Remaining `kwargs` are forwarded to [`from_mcmcchains`](@ref).
"""
function convert_to_inference_data(chns::Chains; group=:posterior, kwargs...)
    group = Symbol(group)
    group === :posterior && return from_mcmcchains(chns; kwargs...)
    return from_mcmcchains(; group => chns)
end

@doc doc"""
    from_mcmcchains(posterior::MCMCChains.Chains; kwargs...) -> InferenceData
    from_mcmcchains(; kwargs...) -> InferenceData
    from_mcmcchains(
        posterior::MCMCChains.Chains,
        posterior_predictive::Any,
        predictions::Any,
        log_likelihood::Any;
        kwargs...
    ) -> InferenceData

Convert data in an `MCMCChains.Chains` format into an [`InferenceData`](@ref).

Any keyword argument below without an an explicitly annotated type above is allowed, so long
as it can be passed to [`convert_to_inference_data`](@ref).

# Arguments

- `posterior::MCMCChains.Chains`: Draws from the posterior

# Keywords

- `posterior_predictive::Any=nothing`: Draws from the posterior predictive distribution or
    name(s) of predictive variables in `posterior`
- `predictions::Any=nothing`: Out-of-sample predictions for the posterior.
- `prior::Any=nothing`: Draws from the prior
- `prior_predictive::Any=nothing`: Draws from the prior predictive distribution or name(s)
    of predictive variables in `prior`
- `observed_data::Dict{String,Array}=nothing`: Observed data on which the `posterior` is
    conditional. It should only contain data which is modeled as a random variable. Keys are
    parameter names and values.
- `constant_data::Dict{String,Array}=nothing`: Model constants, data included in the model
    which is not modeled as a random variable. Keys are parameter names and values.
- `predictions_constant_data::Dict{String,Array}=nothing`: Constants relevant to the model
     predictions (i.e. new `x` values in a linear regression).
- `log_likelihood::Any=nothing`: Pointwise log-likelihood for the data. It is recommended
     to use this argument as a dictionary whose keys are observed variable names and whose
     values are log likelihood arrays.
- `log_likelihood::String=nothing`: Name of variable in `posterior` with log likelihoods
- `library=MCMCChains`: Name of library that generated the chains
- `coords::Dict{String,Vector}=Dict()`: Map from named dimension to named indices
- `dims::Dict{String,Vector{String}}=Dict()`: Map from variable name to names of its
    dimensions
- `eltypes::Dict{String,DataType}=Dict()`: Apply eltypes to specific variables. This is used
    to assign discrete eltypes to discrete variables.

# Returns

- `InferenceData`: The data with groups corresponding to the provided data
"""
from_mcmcchains

function from_mcmcchains(
    posterior,
    posterior_predictive,
    predictions,
    log_likelihood;
    library=MCMCChains,
    eltypes=Dict(),
    kwargs...,
)
    kwargs = convert(Dict, merge((; dims=Dict()), kwargs))
    library = string(library)
    rekey_fun = d -> rekey(d, stats_key_map)

    # Convert chains to dicts
    if posterior === nothing
        post_dict = nothing
        stats_dict = nothing
    else
        post_dict = convert_to_eltypes(chains_to_dict(posterior), eltypes)
        stats_dict = chains_to_dict(posterior; section=:internals, rekey_fun=rekey_fun)
        stats_dict = enforce_stat_eltypes(stats_dict)
        stats_dict = convert_to_eltypes(stats_dict, Dict("is_accept" => Bool))
    end

    all_idata = InferenceData()
    for (group, group_data) in [
        :posterior_predictive => posterior_predictive,
        :predictions => predictions,
        :log_likelihood => log_likelihood,
    ]
        group_data === nothing && continue
        if group_data isa Union{Symbol,String}
            group_data = [string(group_data)]
        end
        if group_data isa Union{AbstractVector{Symbol},NTuple{N,Symbol} where {N}}
            group_data = map(string, group_data)
        end
        if group_data isa Union{AbstractVector{String},NTuple{N,String} where {N}}
            group_data = popsubdict!(post_dict, group_data)
        end
        group_dataset = if group_data isa Chains
            convert_to_dataset(group_data; library=library, eltypes=eltypes, kwargs...)
        else
            convert_to_dataset(group_data; library=library, kwargs...)
        end
        setattribute!(group_dataset, "inference_library", library)
        concat!(all_idata, InferenceData(; group => group_dataset))
    end

    attrs_library = Dict("inference_library" => library)
    if posterior === nothing
        attrs = attrs_library
    else
        attrs = merge(attributes_dict(posterior), attrs_library)
    end
    kwargs = convert(Dict, merge((; attrs=attrs, dims=Dict()), kwargs))
    post_idata = _from_dict(post_dict; sample_stats=stats_dict, kwargs...)
    concat!(all_idata, post_idata)
    return all_idata
end
function from_mcmcchains(
    posterior=nothing;
    posterior_predictive=nothing,
    predictions=nothing,
    prior=nothing,
    prior_predictive=nothing,
    observed_data=nothing,
    constant_data=nothing,
    predictions_constant_data=nothing,
    log_likelihood=nothing,
    library=MCMCChains,
    eltypes=Dict(),
    kwargs...,
)
    kwargs = convert(Dict, merge((; dims=Dict(), coords=Dict()), kwargs))

    all_idata = from_mcmcchains(
        posterior,
        posterior_predictive,
        predictions,
        log_likelihood;
        library=library,
        eltypes=eltypes,
        kwargs...,
    )

    if prior !== nothing
        pre_prior_idata = convert_to_inference_data(
            prior;
            posterior_predictive=prior_predictive,
            library=library,
            eltypes=eltypes,
            kwargs...,
        )
        prior_idata = rekey(
            pre_prior_idata,
            Dict(
                :posterior => :prior,
                :posterior_predictive => :prior_predictive,
                :sample_stats => :sample_stats_prior,
            ),
        )
        concat!(all_idata, prior_idata)
    end

    for (group, group_data) in [
        :observed_data => observed_data,
        :constant_data => constant_data,
        :predictions_constant_data => predictions_constant_data,
    ]
        group_data === nothing && continue
        group_data = convert_to_eltypes(group_data, eltypes)
        group_dataset = convert_to_constant_dataset(group_data; library=library, kwargs...)
        concat!(all_idata, InferenceData(; group => group_dataset))
    end

    return all_idata
end

"""
    from_cmdstan(posterior::Chains; kwargs...) -> InferenceData

Call [`from_mcmcchains`](@ref) on output of `CmdStan`.
"""
function from_cmdstan(posterior::Chains; kwargs...)
    return from_mcmcchains(posterior; library="CmdStan", kwargs...)
end


immutable RectangularBinGeometry <: Gadfly.GeometryElement
    default_statistic::Gadfly.StatisticElement

    function RectangularBinGeometry(
            default_statistic::Gadfly.StatisticElement=Gadfly.Stat.identity())
        new(default_statistic)
    end
end

const rectbin = RectangularBinGeometry

function histogram2d(; xbincount=nothing, xminbincount=3, xmaxbincount=150,
                       ybincount=nothing, yminbincount=3, ymaxbincount=150)
    RectangularBinGeometry(
        Gadfly.Stat.histogram2d(xbincount=xbincount,
                                xminbincount=xminbincount,
                                xmaxbincount=xmaxbincount,
                                ybincount=ybincount,
                                yminbincount=yminbincount,
                                ymaxbincount=ymaxbincount))
end

function default_statistic(geom::RectangularBinGeometry)
    geom.default_statistic
end

function element_aesthetics(::RectangularBinGeometry)
    [:x, :y, :xmin, :xmax, :ymin, :ymax, :color]
end

# Render rectangular bin (e.g., heatmap) geometry.
#
# Args:
#   geom: rectbin geometry
#   theme: the plot's theme
#   aes: some aesthetics
#
# Returns
#   A compose form.
#
function render(geom::RectangularBinGeometry, theme::Gadfly.Theme, aes::Gadfly.Aesthetics)

    default_aes = Gadfly.Aesthetics()
    default_aes.color = PooledDataArray(RGB{Float32}[theme.default_color])
    aes = inherit(aes, default_aes)

    if aes.x === nothing && (aes.xmin === nothing || aes.xmax === nothing)
        error("Geom.rectbin requires either x or both xmin and xmax be defined.")
    end

    if aes.y === nothing && (aes.ymin === nothing || aes.ymax === nothing)
        error("Geom.rectbin requires either y or both ymin and ymax be defined.")
    end

    if aes.xmin != nothing && length(aes.xmin) != length(aes.xmax)
        error("Geom.rectbin requires that xmin and xmax be of equal length.")
    end

    if aes.ymin != nothing && length(aes.ymin) != length(aes.ymax)
        error("Geom.rectbin requires that ymin and ymax be of equal length.")
    end

    nx = aes.xmin === nothing ? length(aes.x) : length(aes.xmin)
    ny = aes.ymin === nothing ? length(aes.y) : length(aes.ymin)

    if nx != ny
        error("""Geom.rectbin requires an equal number of x (or xmin/xmax) and
                 y (or ymin/ymax) values.""")
    end
    n = nx

    if aes.xmin === nothing
        xmin = [x - 0.5 for x in aes.x]
        xwidths = [1.0*cx - theme.bar_spacing]
    else
        xmin = aes.xmin
        xwidths = [(x1 - x0)*cx - theme.bar_spacing
                   for (x0, x1) in zip(aes.xmin, aes.xmax)]
    end

    if aes.ymin === nothing
        ymin = [y - 0.5 for y in aes.y]
        ywidths = [1.0*cy - theme.bar_spacing]
    else
        ymin = aes.ymin
        ywidths = [(y1 - y0)*cy - theme.bar_spacing
                   for (y0, y1) in zip(aes.ymin, aes.ymax)]
    end

    if length(aes.color) == n
        cs = aes.color
    else
        cs = Array(RGB{Float32}, n)
        for i in 1:n
            cs[i] = aes.color[((i - 1) % length(aes.color)) + 1]
        end
    end

    allvisible = true
    for c in cs
        if c == nothing
            allvisible = false
            break
        end
    end

    if !allvisible
        visibility = cs .!= nothing
        cs = cs[visibility]
        xmin = xmin[visibility]
        xmax = xmax[visibility]
        ymin = ymin[visibility]
        ymax = ymax[visibility]
    end

    return compose!(
        context(),
        rectangle(xmin, ymin, xwidths, ywidths),
        fill(cs),
        stroke(nothing),
        svgclass("geometry"),
        svgattribute("shape-rendering", "crispEdges"))
end


using Test
using Boltzmann
using BenchmarkTools

"""
Runs a series of basic smoke and benchmark tests on the interface of any RBM.

Default methods tested:

* `fit(model, X)`
* `transform(model, X)`
* `generate(model, X)`

Args:
* model::Union{AbstractRBM, Net} - the model to test

Optional:
* opts::Dict - the context passed when running fit. (default=DEFAULT_CONTEXT)
* n_obs::Int - total number of observations in generated dataset (default=1000)

Returns: BenchmarkGroup
"""
function benchmark!(rbm::AbstractRBM{T,V,H}, suite::BenchmarkGroup; opts::Dict=DefaultContext(), input_size=-1, debug=false, n_obs=(1000, 5000, 10000)) where {T,V,H}
    tname = String(typeof(rbm).name)
    suite[tname, String(T.name)] = BenchmarkGroup()

    if debug && isa(ctx[:reporter], TestReporter)
        ctx[:reporter].log = true
    end

    for m in n_obs
        ctx = deepcopy(opts)
        n_hid, n_vis = size(rbm.W)

        if input_size > 0
            n_vis = input_size
        end

        DX = generate_dataset(T, n_vis; n_obs=n_obs)
        SX = generate_dataset(T, n_vis; n_obs=n_obs, sparsity=0.3)

        for X in (DX, SX)
            suite[tname]["fit", size(X), String(typeof(X).name)] = @benchmarkable fit($rbm, $X, $ctx)
            suite[tname]["transform", size(X), String(typeof(X).name)] = @benchmarkable transform($rbm, $X)
            suite[tname]["generate", size(X), String(typeof(X).name)] = @benchmarkable generate($rbm, $X)
        end
    end

    return suite
end

#=
Available test groups are:
* :unit,
* :integration,
* :acceptance,
* :benchmark,
=#
const DEFAULT_TEST_GROUPS = Set([:unit, :integration])
TEST_GROUPS = DEFAULT_TEST_GROUPS

# If the JULIA_TEST_GROUPS ENV variable is set
# then use the inte
if haskey(ENV, "JULIA_TEST_GROUPS")
    TEST_GROUPS = Set(
        Symbol.(split(ENV["JULIA_TEST_GROUPS"]))
    )
end

tests = [
    "utils",
    "rbm",
    "nets",
    "conditional",
]

for t in tests
    println("\n\nRunning $t tests...\n")
    include(string("test", t, ".jl"))
end


module SVG

export Inline, Class, Svg, Polyline, Circle, Line, write, bounds, bounds_viewbox, scaled, points2circles

import Base.==

dp(n; digits=2) = round(n; digits)

abstract type SvgObject end

abstract type Style end

"""
	Inline
Inline style object
# Properties
- `attribs` Dict of CSS attributes
# Constructor
	Inline(;fill="none", stroke="black", stroke_width=1)
All keyword arguments
"""
struct Inline <: Style
	attribs
	function Inline(pairs=[];fill="none", stroke="black", stroke_width=1)
		attribs = Dict()
		attribs["fill"] = fill
		attribs["stroke"] = stroke
		attribs["stroke-width"] = stroke_width
		for (k,v) in pairs
			attribs[k] = v
		end
		new(attribs)
	end
end

"""
	Class
CSS class name Object
# Constructor
	Class(class::String)
"""
struct Class <: Style
	class::String
end

"""
	Polyline
Polyline object
# Properties
- `xs` Vector of X coordinates, one per vertex
- `ys` Vector of Y coordinates, one per vertex
# Constructors
Quite the choice.
- Construct with separate vectors or vectors of 2d points.
- Supply an optional transformation function `fx`, `fy`, `fxy`
- Or empty or with pre-allocated vectors of `xs`/`ys`
- Supply `Style` using kwarg
```
Polyline(; style::Style=Inline())
Polyline(npoints::Int; style::Style=Inline())
Polyline(xs::Vector, ys::Vector; style::Style=Inline())
Polyline(xs::Vector, ys::Vector, fx::Function, fy::Function; style::Style=Inline())
Polyline(xys::Vector; style::Style=Inline())
Polyline(xys::Vector, fxy::Function; style::Style=Inline())
Polyline(xys::Vector, fx::Function, fy::Function; style::Style=Inline())
```
"""
struct Polyline <: SvgObject
	xs
	ys
	style::Style
	Polyline(; style::Style=Inline()) = new(Vector{Float64}(), Vector{Float64}(), style)
	Polyline(npoints::Int; style::Style=Inline()) = new(Vector{Float64}(undef, npoints), Vector{Float64}(undef, npoints), style)
	Polyline(xs::Vector, ys::Vector; style::Style=Inline()) = new(xs, ys, style)
	Polyline(xs::Vector, ys::Vector, fx::Function, fy::Function; style::Style=Inline()) = Polyline(map(fx, xs), map(fy, ys); style)
	Polyline(xys::Vector; style::Style=Inline()) = Polyline(map(xy->xy[1], xys), map(xy->xy[2], xys); style)
	Polyline(xys::Vector, fxy::Function; style::Style=Inline()) = Polyline(map(xy->fxy(xy), xys); style)
	Polyline(xys::Vector, fx::Function, fy::Function; style::Style=Inline()) = Polyline(map(xy->fx(xy[1]), xys), map(xy->fy(xy[2]), xys); style)
end

==(p1::Polyline, p2::Polyline) = p1.xs == p2.xs && p1.ys == p2.ys && p1.style.attribs == p2.style.attribs

scaled(p::Polyline, fx, fy) = Polyline(p.xs, p.ys, fx, fy; style=p.style)

"""
	Line
# Properties
- `x1`, `y1` `x2`, `y2` - Endpoints of the line
- `style` - Style object for this line
# Constructor
	Line(x1, y1, x2, y2; style::Style=Inline())
"""
struct Line <: SvgObject
	x1
	y1
	x2
	y2
	style::Style
	Line(x1, y1, x2, y2; style::Style=Inline()) = new(x1, y1, x2, y2, style)
end

scaled(l::Line, fx, fy) = Line(fx(l.x1), fy(l.y1), fx(l.x2), fy(l.y2); style=l.style)

"""
	Circle
Circle object
# Properties
- `x`, `y`, `r` - position and radius
- `style` = Style object
Radius is invariant under the module supplied scaling operations
"""
struct Circle <: SvgObject
	x
	y
	r
	style::Style
	Circle(x, y, r; style::Style=Inline()) = new(x, y, r, style)
end

scaled(c::Circle, fx, fy) = Circle(fx(c.x), fy(c.y), c.r; style=c.style)

"""
	Svg
Svg object
# Properties
- `objects` - Vector of SVG objects
# Constructors
Call either empty or with Vector of existing objects
- `Svg()`
- `Svg(objs::SvgObject)`
"""
struct Svg
	objects::Vector{SvgObject}
	Svg() = new(Vector{SvgObject}())
	Svg(objs) = new(objs)
end
"""
	scaled(s::Svg, w, h; digits, flip_y)
	scaled(s::Svg, fx, fy)
	scaled(s::SvgObject, fx, fy)
Scale either the whole Svg or individual Objects
- `s` either the whole Svg or individual objects
- `fx`, `fy` - functions to scale the positional co-ordinates
- `w`, `h` - scale the entire Svg up to the smallest of `w` / `h` and maintain aspect ratio
- `digits::Int` - round to this many digits
- `flip_y::Bool` - flip vertically
"""
scaled(s::Svg, fx::Function, fy::Function) = Svg(map(o->scaled(o, fx, fy), s.objects))
function scaled(svg, width::Real, height::Real; flip_y=false, digits=2)
    xmin, ymin, xmax, ymax = bounds(svg)
    xmx = xmax - xmin
    ymx = ymax - ymin
    scale = min(width, height) / min(xmx, ymx)
    fx = x -> round(scale * (x - xmin); digits)
    fy = y -> round(flip_y ? height - scale * (y - ymin) : scale * (y - ymin); digits)
	scaled(svg, fx, fy)
end

"""
	bounds(svg::Svg)
	bounds(svg::SvgObject)
Return the bounding box of the entire `Svg` or an individual `SvgObject`
# Returns
NamedTuple(xmin, ymin, xmax, ymax)
"""
function bounds(svg::Svg)
	xmin, ymin, xmax, ymax = Inf, Inf, -Inf, -Inf
	for b in map(bounds, svg.objects)
		if b.xmin < xmin
			xmin = b.xmin
		end
		if b.ymin < ymin
			ymin = b.ymin
		end
		if b.xmax > xmax
			xmax = b.xmax
		end
		if b.ymax > ymax
			ymax = b.ymax
		end
	end
	(;xmin, ymin, xmax, ymax)
end

"""
	bounds_viewbox(svg::Svg)
	bounds_viewbox(bounds::NamedTuple)
Create a viewbox string using either the Svg or a set of bounds
"""
bounds_viewbox(svg::Svg) = bounds_viewbox(bounds(svg))
bounds_viewbox(bounds::NamedTuple) = "$(bounds.xmin) $(bounds.ymin) $(bounds.xmax) $(bounds.ymax)"

"""
	write(fn::String, svg::Svg, width, height; viewbox="", inhtml=false, digits=2, objwrite_fn=write_objs)
	write(io::IO, svg::Svg, width, height; viewbox="", inhtml=false, digits=2, objwrite_fn=write_objs)
write the Svg to a file or IO object.
# Arguments
- `fn` - write to filename
- `io` - write to IO Object
- `svg` - the Svg to write
- `width`, `height` - the width / height of the SVG object to supply
- `viewbox` - specify a particular viewbox: "xmin ymin xmax ymax"
- `inhtml` - Boolean flag whether to wrap the SVG in an HTML document for viewing in a browser.
- `digits` - round all SVG Object co-ordinates to this many digits
- `objwrite_fn` - a callback so the caller can modify the objects as they are written
"""
Base.write(fn::String, svg::Svg, width, height; viewbox="", inhtml=false, stylesheet="", digits=2, objwrite_fn=write_objs) = open(fn, "w+") do io write(io, svg, width, height; viewbox, inhtml, stylesheet, digits, objwrite_fn) end

write_objs(io::IO, svg::Svg) = foreach(o->write(io, o), svg.objects)

function Base.write(io::IO, svg::Svg, width, height; viewbox="", inhtml=false, stylesheet="", digits=2, objwrite_fn=write_objs)
	open_html(io; inhtml)
	open_svg(io, width, height; viewbox, stylesheet)
	objwrite_fn(io, svg)
	close_svg(io)
	close_html(io; inhtml)
end

function open_svg(io, width, height; viewbox="", stylesheet="")
	if viewbox != ""
		viewbox = " viewBox=\"$viewbox\""
	end
	println(io, """<svg width="$width" height="$height"$viewbox xmlns="http://www.w3.org/2000/svg">""")
	if stylesheet != ""
		println(io, "")
	end
end
close_svg(io) = println(io, "</svg>")
open_html(io; inhtml) = if inhtml println(io, "
") end
close_html(io; inhtml) = if inhtml println(io, "
") end

function Base.write(io::IO, s::Inline)
	print(io, "style=\"")
	for (k,v) in s.attribs
		print(io, "$k:$v; ")
	end
	print(io, "\"")
end
Base.write(io::IO, s::Class) = print(io, "class=\"$(s.class)\"")

function Base.write(io::IO, p::Polyline)
	print(io, "<polyline points=\"")
	foreach(i->print(io, p.xs[i], ",", p.ys[i], " "), 1:length(p.xs))
	print(io, "\" ")
	write(io, p.style)
	println(io, " />")
end

bounds(p::Polyline) = (xmin=minimum(p.xs), ymin=minimum(p.ys), xmax=maximum(p.xs), ymax=maximum(p.ys))

function Base.write(io::IO, L::Line)
	print(io, "<line x1=\"", L.x1, "\" y1=\"", L.y1, "\" x2=\"", L.x2, "\" y2=\"", L.y2, "\" ")
	write(io, L.style)
	println(io, " />")
end

bounds(l::Line) = (xmin=min(l.x1,l.x2), ymin=min(l.y1,l.y2), xmax=max(l.x1,l.x2), ymax=max(l.y1,l.y2))

function Base.write(io::IO, c::Circle)
	print(io, "<circle cx=\"", c.x, "\" cy=\"", c.y, "\" r=\"", c.r, "\" ")
	write(io, c.style)
	println(io, " />")
end

"""
	points2circles(xs, ys, r::Real)
	points2circles(xys, r::Real)
	points2circles(xs, ys, rs::Vector{Real})
	points2circles(xys, rs::Vector{Real})
Helper function to create a collection of points of a given radius to a collection of Circle objects
# Arguments
- `xs`, `ys` Vectors or x/y corrdinates
- `xys` Vector of 2D points
- `r` Radius
- `rs` Vector of Radii
"""
points2circles(xs, ys, r::Real=1) = map(i->Circle(xs[i], ys[i], r), 1:length(xs))
points2circles(xs, ys, rs::Vector{Real}) = map(i->Circle(xs[i], ys[i], r[i]), 1:length(xs))
points2circles(xys, r::Real=1) = map(i->Circle(xys[i][1], xys[i][2], r), 1:length(xys))
points2circles(xys, rs::Vector{Real}) = map(i->Circle(xys[i][1], xys[i][2], r[i]), 1:length(xys))

bounds(c::Circle) = (xmin=c.x-c.r, ymin=c.y-c.r, xmax=c.x+c.r, ymax=c.y+c.r)

###
end


"""

    calib_wrapper(param, st_hydro, prec, epot, q_obs, q_sim)

Wrapper function required for calibrating hydrological routing model.

"""

function calib_wrapper(param, st_hydro, prec, epot, q_obs, q_sim, states_sim, warmup, force_states)

    # Assign parameter values

    assign_param(st_hydro, param)

    # Initilize model states

    init_states(st_hydro)

    # Run model

    ntimes = size(prec, 2)

    for itime in 1:ntimes

        get_input(st_hydro, prec, epot, itime)

        run_timestep(st_hydro)

        q_sim[itime] = st_hydro.q_sim

        if ~isempty(states_sim)

            states_hydro = get_states(st_hydro)

            states_sim[:, itime] = vcat(states_hydro)

        end

    end

    # Compute penatly

    if force_states == true && ~isempty(states_sim)

        penalty = zeros(size(states_sim, 1))

        for istate = 1:size(states_sim, 1)

            slope = cov(states_sim[istate, warmup:end], warmup:ntimes) / var(warmup:ntimes)

            penalty[istate] = 2 ./ (1  + exp(-(10 * slope).^2)) - 1

        end

        penalty = maximum(penalty)

        # penalty = mean(abs(slope))

    else

        penalty = 0.0

    end

    # Compute performance measure

    perf_measure = 1.0 - nse(q_sim[warmup:end], q_obs[warmup:end]) + penalty

    return perf_measure

end

"""

    calib_wrapper(param, st_snow, st_hydro, tair, prec, epot, q_obs)

Wrapper function required for calibrating snow and hydrological routing model.

"""

function calib_wrapper(param, st_snow, st_hydro, tair, prec, epot,
                       q_obs, q_sim, states_sim, warmup, force_states)

    # Assign parameter values

    assign_param(st_snow, param[1:length(st_snow.param)])
    assign_param(st_hydro, param[length(st_snow.param)+1:end])

    # Initilize model states

    init_states(st_snow)
    init_states(st_hydro)

    # Run model

    ntimes = size(prec, 2)

    for itime in 1:ntimes

        get_input(st_snow, prec, tair, itime)

        run_timestep(st_snow)

        get_input(st_snow, st_hydro, epot, itime)

        run_timestep(st_hydro)

        q_sim[itime] = st_hydro.q_sim

        states_snow = get_states(st_snow)
        states_hydro = get_states(st_hydro)

        states_sim[:, itime] = vcat(states_snow, states_hydro)

    end

    # Compute penatly

    if force_states == true

        penalty = zeros(size(states_sim, 1))

        for istate = 1:size(states_sim, 1)

            slope = cov(states_sim[istate, warmup:end], warmup:ntimes) / var(warmup:ntimes)

            penalty[istate] = 2 ./ (1  + exp(-(10 * slope).^2)) - 1

        end

        penalty = maximum(penalty)

        # penalty = mean(abs(slope))

    else

        penalty = 0.0

    end

    # Compute performance measure

    perf_measure = 1.0 - nse(q_sim[warmup:end], q_obs[warmup:end]) + penalty

    return perf_measure

end

"""

    run_model_calib(st_hydro::Hydro, prec, epot, q_obs)

Run calibration of hydrological routing model, for example HBV.

"""

function run_model_calib(st_hydro::Hydro, prec, epot, q_obs;
                         verbose = :silent, warmup = 3*365, force_states = false)

    # Get parameter range

    param_range = get_param_range(st_hydro)

    # Allocate output array

    states_hydro = Vann.get_states(st_hydro)

    ntimes = size(prec, 2)
    nstates = length(states_hydro)

    q_sim = zeros(ntimes)

    states_sim = zeros(nstates, ntimes)

    # Run calibration

    calib_wrapper_tmp(param) = calib_wrapper(param, st_hydro, prec, epot, q_obs,
                                             q_sim, states_sim, warmup, force_states)

    res = bboptimize(calib_wrapper_tmp; SearchRange = param_range, TraceMode = verbose)

    param_hydro = best_candidate(res)

    return param_hydro

end

"""

    run_model_calib(st_snow::Snow, st_hydro::Hydro, tair, prec, epot, q_obs)

Run calibration of snow and hydrological routing model.

"""

function run_model_calib(st_snow::Snow, st_hydro::Hydro, tair, prec, epot, q_obs;
                         verbose = :silent, warmup = 3*365, force_states = false)

    # Get parameter range

    param_range_snow  = get_param_range(st_snow)
    param_range_hydro = get_param_range(st_hydro)

    param_range = vcat(param_range_snow, param_range_hydro)

    # Allocate output array

    states_snow = Vann.get_states(st_snow)
    states_hydro = Vann.get_states(st_hydro)

    ntimes = size(prec, 2)
    nstates = length(vcat(states_snow, states_hydro))

    q_sim = zeros(ntimes)
    states_sim = zeros(nstates, ntimes)

    # Run calibration

    calib_wrapper_tmp(param) = calib_wrapper(param, st_snow, st_hydro, tair, prec, epot,
                                             q_obs, q_sim, states_sim, warmup, force_states)

    res = bboptimize(calib_wrapper_tmp; SearchRange = param_range, TraceMode = verbose)

    # Extract parameters for snow and hydrological routing model

    param = best_candidate(res)

    param_snow = param[1:length(st_snow.param)]
    param_hydro = param[length(st_snow.param)+1:end]

    return param_snow, param_hydro

end

"""

    kge(q_sim, q_obs)

Compute modified Kling-Gupta efficiency

"""

function kge(q_sim, q_obs)

    # ikeep = q_obs .!= -999.

    if all(isnan, q_sim) || all(isnan, q_obs)

        kge = NaN

    else

        ikeep = !isnan(q_obs)

        q_sim = q_sim[ikeep]
        q_obs = q_obs[ikeep]

        r = cor(q_sim, q_obs)

        beta = mean(q_sim) / mean(q_obs)

        gamma = (std(q_sim) / mean(q_sim)) / (std(q_obs) / mean(q_obs))

        kge = 1 - sqrt( (r-1)^2 + (beta-1)^2 + (gamma-1)^2 )

    end

    return kge

end

"""

    nse(q_sim, q_obs)

Compute Nash-Sutcliffe efficiency

"""

function nse(q_sim, q_obs)

    # ikeep = q_obs .!= -999.

    ikeep = !isnan(q_obs)

    q_sim = q_sim[ikeep]
    q_obs = q_obs[ikeep]

    ns = 1 - sum((q_sim - q_obs).^2) / sum((q_obs - mean(q_obs)).^2)

end

#=
### Snow and hydrological model

# Run model and compute performance measure

function calib_wrapper(param::Vector, grad::Vector, st_snow, st_hydro, date, tair, prec, epot, q_obs)

  # Assign parameter values

  assign_param(st_snow, param[1:length(st_snow.param)]);
  assign_param(st_hydro, param[length(st_snow.param)+1:end]);

  # Initilize model states

  init_states(st_snow);
  init_states(st_hydro);

  # Run model

  ntimes = size(prec, 2);

  q_sim = zeros(Float64, ntimes);

  for itime in 1:ntimes

    get_input(st_snow, prec, tair, date, itime);

    run_timestep(st_snow);

    get_input(st_snow, st_hydro, epot, itime);

    run_timestep(st_hydro);

    q_sim[itime] = st_hydro.q_sim

  end

  return(kge(q_sim, q_obs));

end

# Run calibration

function run_model_calib(st_snow::Snow, st_hydro::Hydro, date, tair, prec, epot, q_obs)

  # Get parameter range

  param_range_snow  = get_param_range(st_snow);
  param_range_hydro = get_param_range(st_hydro);

  param_range = vcat(param_range_snow, param_range_hydro);

  # Lower and upper bounds for parameters

  param_lower = [param_range[i][1] for i in eachindex(param_range)];
  param_upper = [param_range[i][2] for i in eachindex(param_range)];

  # Starting point

  param_start = (param_lower + param_upper) / 2;

  # Wrapper function

  calib_wrapper_tmp(param::Vector, grad::Vector) = calib_wrapper(param::Vector, grad::Vector, st_snow, st_hydro, date, tair, prec, epot, q_obs);

  # Perform global optimization

  opt_global = Opt(:GN_ESCH, length(param_start));

  max_objective!(opt_global, calib_wrapper_tmp);

  maxeval!(opt_global, 5000);

  lower_bounds!(opt_global, param_lower);
  upper_bounds!(opt_global, param_upper);

  (min_func, best_global, ret_nlopt) = optimize(opt_global, param_start);

  # Check that optimal value is inside bounds

  for iparam = 1:length(best_global)

    if best_global[iparam] < param_lower[iparam] + 0.2*(param_upper[iparam] - param_lower[iparam])
      best_global[iparam] = param_lower[iparam] + 0.2*(param_upper[iparam] - param_lower[iparam]);
    end

    if best_global[iparam] > param_upper[iparam] - 0.2*(param_upper[iparam] - param_lower[iparam]);
      best_global[iparam] = param_upper[iparam] - 0.2*(param_upper[iparam] - param_lower[iparam]);
    end

  end

  # Perform local optimization

  opt_local = Opt(:LN_NELDERMEAD, length(param_start));

  max_objective!(opt_local, calib_wrapper_tmp);

  lower_bounds!(opt_local, param_lower);
  upper_bounds!(opt_local, param_upper);

  (min_func, best_local, ret_nlopt) = optimize(opt_local, best_global);

  # Return best parameter values

  return(best_local)

end

### Hydrological model

# Run model and compute performance measure

function calib_wrapper(param::Vector, grad::Vector, st_hydro, prec, epot, q_obs)

  # Assign parameter values

  assign_param(st_hydro, param);

  # Initilize model states

  init_states(st_hydro);

  # Run model

  ntimes = size(prec, 2);

  q_sim = zeros(Float64, ntimes);

  for itime in 1:ntimes

    get_input(st_hydro, prec, epot, itime);

    q_sim[itime] = run_timestep(st_hydro);

  end

  return(kge(q_sim, q_obs));

end

# Run calibration

function run_model_calib(st_hydro::Hydro, prec, epot, q_obs)

  # Get parameter range

  param_range = get_param_range(st_hydro);

  # Lower and upper bounds for parameters

  param_lower = [param_range[i][1] for i in eachindex(param_range)];
  param_upper = [param_range[i][2] for i in eachindex(param_range)];

  # Starting point

  param_start = (param_lower + param_upper) / 2;

  # Wrapper function

  calib_wrapper_tmp(param::Vector, grad::Vector) = calib_wrapper(param::Vector, grad::Vector, st_hydro, prec, epot, q_obs);

  # Perform global optimization

  opt_global = Opt(:GN_ESCH, length(param_start));

  max_objective!(opt_global, calib_wrapper_tmp);

  maxeval!(opt_global, 5000);

  lower_bounds!(opt_global, param_lower);
  upper_bounds!(opt_global, param_upper);

  (min_func, best_global, ret_nlopt) = optimize(opt_global, param_start);

  # Check that optimal value is inside bounds

  for iparam = 1:length(best_global)

    if best_global[iparam] < param_lower[iparam] + 0.2*(param_upper[iparam] - param_lower[iparam])
      best_global[iparam] = param_lower[iparam] + 0.2*(param_upper[iparam] - param_lower[iparam]);
    end

    if best_global[iparam] > param_upper[iparam] - 0.2*(param_upper[iparam] - param_lower[iparam]);
      best_global[iparam] = param_upper[iparam] - 0.2*(param_upper[iparam] - param_lower[iparam]);
    end

  end

  # Perform local optimization

  opt_local = Opt(:LN_NELDERMEAD, length(param_start));

  max_objective!(opt_local, calib_wrapper_tmp);

  lower_bounds!(opt_local, param_lower);
  upper_bounds!(opt_local, param_upper);

  (min_func, best_local, ret_nlopt) = optimize(opt_local, best_global);

  # Return best parameter values

  return(best_local)

end

=#


"""
A simple common interface to diverse graphics packages backends useful for
drawings in SDD algorithms.
"""
module SDDGraphics

using Reexport

include("canvas.jl")

export
  canvassize,
  width,
  height,
  rectregion,
  xlims,
  ylims,
  xmin,
  ymin,
  xmax,
  ymax,
  square,
  insiderectregion,
  pointtopixel,
  pixeltopoint,
  pointtocanvas,
  canvastopoint

include("colors.jl")
include("coloringfunctions.jl")

export
  colortype,
  color,
  bgcolor,
  fgcolor,
  colorscheme,
  colormap,
  colorarray,
  updatecolorarray,
  colorinterpolation,
  colorinterpolationbg,
  colorinterpolationfg,
  coloringfunction,
  RadialColoringFunction, RadialCF,
  AngleColoringFunction, AngleCF

include("drawing.jl")

#ToDo!: Documentation.
"""
    configure(; kwargs...)

Shortcut function to configure graphics.

#### Arguments
- `canvassize=(w,h)` or `canvas=(w,h)`: Set the canvas size.
"""
  function configure(; kwargs...)
    for (k,v) in kwargs
      if k == :canvassize || k == :canvas
        canvassize(v...)
      elseif k == :width || k == :w
        width(v)
      elseif k == :height || k == :h
        height(v)
      elseif k == :xmin
        xmin(v)
      elseif k == :xmax
        xmax(v)
      elseif k == :ymin
        ymin(v)
      elseif k == :ymax
        ymax(v)
      elseif k == :xlims
        xlims(v...)
      elseif k == :ylims
        width(v...)
      elseif k == :rectregion || k == :rect
        rectregion(v...)
      elseif k == :color
        color(v)
      elseif k == :bgcolor || k == :bgc
        bgcolor(v)
      elseif k == :fgcolor || k == :fgc
        fgcolor(v)
      elseif k == :colormap || k == :cm || k == :colorscheme || k == :cs
        colorscheme(v)
      elseif k == :colormapinv || k == :cminv || k == :colorschemeinv || k == :csinv
        colorscheme(v,true)
      elseif k == :coloringfunction || k == :cf
        coloringfunction(v)
      elseif k == :axes
        axes(v)
      elseif k == :pointsize || k == :ps
        pointsize(v)
      elseif k == :linewidth || k == :lw
        linewidth(v)
      elseif k == :style
        style(v)
      #elseif k == :dash
      #  dash(v)
      end
    end
  end

export
  pointsize,
  linewidth,
  style, # :stroke, :fill, :fillstroke, :fillgrad
  #dash
  configure,
  drawingkind,
  newdrawing,
  drawpixel,
  drawpoint,
  drawline,
  #drawray,
  drawlinesegment,
  drawcircle,
  drawarc,
  #drawpath, #?
  #drawbox, #?
  #drawrect, #?
  axes,
  drawing

include("backends/images.jl") # Default backend
include("backends.jl")

export
  backend,
  backends,
  supported

end