mutable struct PlateUpdateState{U,V,W}
    score::Float64
    is_empty::Bool
    args::U
    nodes::V
    changes::W
    discard::DynamicAssignment
end

function update_existing_trace!(gen::Plate{T,U}, key::Int, state::PlateUpdateState, subtraces, retvals) where {T,U}
    node = haskey(state.nodes, key) ? state.nodes[key] : EmptyAssignment()
    if haskey(state.changes, key)
        kernel_change = state.changes[key]
    else
        kernel_change = nothing
    end
    subtrace::U = subtraces[key]
    prev_call = get_call_record(subtrace)
    kernel_args = get_args_for_key(state.args, key)
    (subtrace, _, kernel_discard) = update(
        gen.kernel, kernel_args, kernel_change, subtrace, node)
    set_internal_node!(state.discard, key, kernel_discard)
    call::CallRecord = get_call_record(subtrace)
    state.score += (call.score - prev_call.score)
    subtraces = assoc(subtraces, key, subtrace)
    retvals = assoc(retvals, key, call.retval::T)
    state.is_empty = state.is_empty && !has_choices(subtraces)
    (subtraces, retvals)
end

function assess_new_trace!(gen::Plate{T,U}, key::Int, state::PlateUpdateState, subtraces, retvals) where {T,U}
    node = haskey(state.nodes, key) ? state.nodes[key] : EmptyAssignment()
    kernel_args = get_args_for_key(state.args, key)
    subtrace::U = assess(gen.kernel, kernel_args, node)
    subtraces = push(subtraces, subtrace)
    call::CallRecord = get_call_record(subtrace)
    state.score += call.score
    retvals = push(retvals, call.retval::T)
    (subtraces, retvals)
end

function update_process_constraints!(nodes, to_visit, key::Int, node, new_length)
    if key <= new_length
        push!(to_visit, key)
        nodes[key] = node
    else
        error("Update did not consume constraints at key $key")
    end
end

function update_process_constraints!(nodes, to_visit, key, node, new_length)
    # key is something other than an Int
    error("Update did not consume constraints at key $key")
end

function _update(gen::Plate{T,U}, args, prev_trace::VectorTrace{T,U}, constraints, to_visit, changes) where {T,U}

    (new_length, prev_length) = get_prev_and_new_lengths(args, prev_trace)

    # collect constraints, indexed by key
    nodes = Dict{Int, Any}()
    for (key, node) in get_internal_nodes(constraints)
        update_process_constraints!(nodes, to_visit, key, node, new_length)
    end

    # discard subtraces of deleted applications
    discard = DynamicAssignment()
    for key=new_length+1:prev_length
        set_internal_node!(discard, key, get_assignment(prev_trace.subtraces[key]))
    end

    # collect initial state
    state = PlateUpdateState(prev_trace.call.score, prev_trace.is_empty, args, nodes, changes, discard)
    subtraces = prev_trace.subtraces
    retvals = prev_trace.call.retval

    # visit existing applications
    for key in to_visit
        (subtraces, retvals) = update_existing_trace!(gen, key, state, subtraces, retvals)
    end

    # visit each new application and assess the new trace
    for key=prev_length+1:new_length
        (subtraces, retvals) = assess_new_trace!(gen, key, state, subtraces, retvals)
    end

    call = CallRecord(state.score, retvals, args)
    trace = VectorTrace(subtraces, call, state.is_empty)

    # compute the weight
    prev_score = prev_trace.call.score
    weight = state.score - prev_score
    retchange = nothing
    (trace, weight, state.discard, retchange)
end

# No change to arguments
function update(gen::Plate{T,U}, args, change::NoChange, trace::VectorTrace{T,U}, constraints) where {T,U}
    to_visit = Set{Int}()
    changes = Dict{Int, T}()
    _update(gen, args, trace, constraints, to_visit, changes)
end

# Known change to arguments
function update(gen::Plate{T,U}, args, change::PlateChange, trace::VectorTrace{T,U}, constraints) where {T,U}

    (new_length, prev_length) = get_prev_and_new_lengths(args, trace)

    # calculate which existing applications to visit
    to_visit = Set{Int}()
    changes = Dict{Int, T}()
    for (i, kernel_change) in zip(change.changed_args, change.sub_changes)
        if i <= new_length
            push!(to_visit, i)
            changes[i] = kernel_change
        end
    end

    _update(gen, args, trace, constraints, to_visit, changes)
end

# Unknown change to arguments
function update(gen::Plate{T,U}, args, change::Nothing, trace::VectorTrace{T,U}, constraints) where {T,U}

    (new_length, prev_length) = get_prev_and_new_lengths(args, trace)

    # visit all existing applications that are preserved
    to_visit = Set(1:min(prev_length, new_length))
    changes = Dict{Int, Any}() # not used

    _update(gen, args, trace, constraints, to_visit, changes)
end


# Check the interactions of different Containers via `convert`

using TypeClasses
using Test
using Suppressor
using DataTypesBasic
splitln(str) = split(strip(str), "\n")

load_square(x, prefix="") = @ContextManager function (cont)
  println("$(prefix)before $x")
  result = cont(x*x)
  println("$(prefix)after $x")
  result
end

# Multiple Combinations
# =====================

# unfortunately this does not work, as the typeinference does not work
# and hence an Array of Any is constructed instead of an Array of ContextManager

# the version without the Array component works like a charm

h2() = @syntax_flatmap begin
  (a, b) = [(1, 2), (2, 3), (3, 4), (4, 5), (5, 6)]
  c = iftrue(a % 2 == 0) do
    a + b
  end
  d = load_square(c) # c*c
  @pure d
end

@test @suppress(h2()) == [25, 81]
@test splitln(@capture_out h2()) == ["before 5", "after 5", "before 9", "after 9"]

# Try
# ---

mytry(x) = x == 4 ? error("error") : x

h3() = @syntax_flatmap begin
  (a, b) = [(1, 2), (2, 3), (3, 4), (4, 5), (5, 6)]
  c = iftrue(a % 2 == 0) do
    a + b
  end
  d = load_square(c)  # c*c
  e = @Try mytry(d / a)
  @pure e + 1
end
@test @suppress(h3()) == [(2+3)^2/2 + 1, (4+5)^2/4 + 1]  # TODO not working because of bad type-inference
@test splitln(@capture_out h3()) == [
  "before 5",
  "after 5",
  "before 9",
  "after 9",
]

# Either
# ------

h4() = @syntax_flatmap begin
  (a, b) = [(1, 2), (2, 3), (3, 4), (4, 5), (5, 6)]
  c = iftrue(a % 2 == 0) do
    a + b
  end
  d = load_square(c) # c*c
  f = either("left", d > 4, :right)
  e = @Try "$a, $b, $c, $d, $f"
end
@test @suppress(h4()) == ["2, 3, 5, 25, right", "4, 5, 9, 81, right"]
@test splitln(@capture_out h4()) == [
  "before 5",
  "after 5",
  "before 9",
  "after 9",
]

# Writer + ContextManager
# ---------------------

writer_cm() = @syntax_flatmap begin
  i = Writer("hi", 4)
  l = load_square(i)
  Writer(" $l yes", l+i)
end

# importantly, writer does not work smoothly with other Contexts, but only by forgetting its accumulator
@test @suppress(writer_cm()) != Writer("hi 4 yes", 4*4+4)
@test @suppress(writer_cm()) == Writer("hi", 4*4+4)
@test splitln(@capture_out writer_cm()) == [
  "before 4",
  "after 4"
]

# combining ContextManager with other monads
# ------------------------------------------

# ContextManager as main Monad

cm2 = @syntax_flatmap begin
  i = load_square(4)
  v = [i, i+1, i+3]
  @pure v + 2
end
@test_throws MethodError @suppress(cm2(x -> x).value) == [6,7,9]
@test_throws MethodError splitln(@capture_out cm2(x -> x)) == ["before 4", "after 4"]

# ContextManager as sub monad

vector2() = @syntax_flatmap begin
  i = [1,2,3]
  c = load_square(i)
  @pure c + 2
end
@test @suppress(vector2()) == [1*1+2, 2*2+2, 3*3+2]
@test splitln(@capture_out vector2()) == [
  "before 1", "after 1", "before 2", "after 2", "before 3", "after 3"]

# alternating contextmanager and vector does not work, as there is no way to convert a vector to a contextmanager

multiplecm() = @syntax_flatmap begin
  i = [1,2,3]
  c = load_square(i, "i ")
  j = [c, c*c]
  c2 = load_square(j, "j ")
  @pure c + c2
end

@test_throws MethodError @suppress(multiplecm()) == [2, 2, 4, 6, 6, 12]
@test_throws MethodError splitln(@capture_out multiplecm()) == ["i before 1",
  "j before 1",
  "j after 1",
  "j before 1",
  "j after 1",
  "i after 1",
  "i before 2",
  "j before 2",
  "j after 2",
  "j before 4",
  "j after 4",
  "i after 2",
  "i before 3",
  "j before 3",
  "j after 3",
  "j before 9",
  "j after 9",
  "i after 3"]


export_plot(z, all_data, ϕ_all, filename, ylabel) = nothing
export_plot_snapshot(z, all_data, ϕ_all, filename, ylabel) = nothing

# using Requires
# @init @require Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" begin
#   using .Plots

"""
    plot_friendly_name(ϕ)

Get plot-friendly string, since many Unicode
characters do not render in plot labels.
"""
function plot_friendly_name(ϕ)
    s = ϕ
    s = replace(s, "ρ" => "rho")
    s = replace(s, "α" => "alpha")
    s = replace(s, "∂" => "partial")
    s = replace(s, "∇" => "nabla")
    return s
end

"""
    export_plot(z, all_data, ϕ_all, filename, ylabel)

Export plot of all variables, or all
available time-steps in `all_data`.
"""
function export_plot(z, all_data, ϕ_all, filename, ylabel)
    ϕ_all isa Tuple || (ϕ_all = (ϕ_all,))
    p = plot()
    for n in 0:(length(keys(all_data)) - 1)
        for ϕ in ϕ_all
            ϕ_string = String(ϕ)
            ϕ_name = plot_friendly_name(ϕ_string)
            ϕ_data = all_data[n][ϕ_string][:]
            plot!(ϕ_data, z, xlabel = ϕ_name, ylabel = ylabel)
        end
    end
    savefig(filename)
end

"""
    export_plot_snapshot(z, all_data, ϕ_all, filename, ylabel)

Export plot of all variables in `all_data`
"""
function export_plot_snapshot(z, all_data, ϕ_all, filename, ylabel)
    ϕ_all isa Tuple || (ϕ_all = (ϕ_all,))
    p = plot()
    for ϕ in ϕ_all
        ϕ_string = String(ϕ)
        ϕ_name = plot_friendly_name(ϕ_string)
        ϕ_data = all_data[ϕ_string][:]
        plot!(ϕ_data, z, xlabel = ϕ_name, ylabel = ylabel)
    end
    savefig(filename)
end

# end


"""
This file contains functions that extract information on the timestep of the Isca model output
"""

function supportedcalendars()
    return ["no_calendar","thirty_day"]
end

function thirty2threesixty(tvec::Vector{<:Real},units::AbstractString)

    @info "$(Dates.now()) - The Isca \"THIRTY_DAY\" calendar corresponds to the \"360_DAY\" calendar in CF convention.  Converting date vector into DateTime360Day type so that it is recognizable by CFTime.jl"

    return timeencode(tvec,units,calendar="360_day")

end

function retrievetime(fnc::AbstractString)

    @warn "$(Dates.now()) - Isca Calendars may not be recognized by CF conventions and therefore it is necessary to manually convert the time-information, where applicable, into CF conventions so that it can be parsed by CFTime.jl"

    ds = Dataset(fnc); t = ds["time"][:]; tattr = ds["time"].attrib;
    tdict = Dict{AbstractString,AbstractString}()
    tinfo = Dict{AbstractString,Any}()

    tdict["units"] = tattr["units"];
    tdict["long_name"] = tattr["long_name"];
    tdict["cartesian_axis"] = tattr["cartesian_axis"];
    tdict["calendar_type"] = tattr["calendar_type"]
    tdict["calendar"] = tattr["calendar"]
    tdict["bounds"] = tattr["bounds"]

    cal = tattr["calendar"]; supcals = supportedcalendars()
    if sum(lowercase(cal).==supcals) == 0
        @warn "$(Dates.now()) - $cal is currently not supported in IscaTools.jl, which means that some of the functionalities of IscaTools.jl cannot be used.  Please submit an issue or a PR."
    end

    if cal == "THIRTY_DAY"; t = thirty2threesixty(t,tdict["units"]) end

    tinfo["calendar"] = tattr["calendar"]
    tinfo["nhr"] = 0; tinfo["ndy"] = 0;
    tinfo["time"] = t; nt = length(t);
    tinfo["raw"] = ds["time"][:]

    if nt > 1
          tstep = tinfo["raw"][2] - tinfo["raw"][1]
    else; tstep = tinfo["raw"][1] * 2
    end

    tinfo["ncattribs"] = tdict;

    if cal == "NO_CALENDAR"

        unit = split(tattr["units"]," since ")[1];
        if occursin("day",unit)
            if tstep < 1;  tinfo["nhr"] = 1 / tstep end
            tinfo["ndy"] = length(t) * tstep;
        elseif occursin("hour",unit)
            if tstep < 24; tinfo["nhr"] = 24 / tstep end
            tinfo["ndy"] = length(t) * tstep / 24;
        elseif  occursin("month",unit)
            tinfo["ndy"] = length(t) * tstep * 30;
        else
            error("$(Dates.now()) - IscaTools.jl requires that if your calendar was set to \"NO_CALENDAR\", you need to express the frequency at which you save your data in units of hours or days.")
        end

    else

        ttot = t[end] - t[1] + tstep; nhr = 24 / Dates.value(Hour(tstep))
        tinfo["ndy"] = Day(ttot); if nhr > 1; tinfo["nhr"] = nhr; end

    end

    if round(Int,tinfo["nhr"]) != tinfo["nhr"]; tinfo["nhr"] = 0; end

    if cal == "NOLEAP" && tinfo["ndy"] !== 365
        @error("$(Dates.now()) - IscaTools.jl requires that if your calendar was set to \"NOLEAP\" or \"JULIAN\", your output data $(BOLD("MUST")) be saved every $(BOLD("year (360 days)")).")
    elseif cal == "JULIAN" && sum(tinfo["ndy"] .== [365,366]) == 0
        @error("$(Dates.now()) - IscaTools.jl requires that if your calendar was set to \"NOLEAP\" or \"JULIAN\", your output data $(BOLD("MUST")) be saved every $(BOLD("year (360 days)")).")
    elseif cal == "THIRTY_DAY" && sum(tinfo["ndy"] .== [30,360]) == 0
        @error("$(Dates.now()) - IscaTools.jl requires that if your calendar was set to \"THIRTY_DAY\", your output data $(BOLD("MUST")) be saved every $(BOLD("month (30 days)")) or $(BOLD("year (360 days)")).")
    elseif cal == "NO_CALENDAR" && sum(tinfo["ndy"] .== [30,360]) == 0
        @warn "$(Dates.now()) - IscaTools.jl has detected that your calendar was set to \"NO_CALENDAR\", and that you did not save data every $(BOLD("month (30 days)")), or $(BOLD("year (360 days)")).  Therefore, IscaTools.jl will perform time-averaging operations over the length of the $(BOLD("entire run ($(tinfo["ndy"]) days)"))."
    end

    if tinfo["ndy"] == 30; tinfo["isyear"] = false; else; tinfo["isyear"] = true; end

    close(ds); return tinfo

end

function retrieveruns!(itime::AbstractDict,iroot::AbstractDict)

    fol = glob("run*",iroot["raw"]); nfol = length(fol);
    itime["nruns"] = nfol;

    if isdir(iroot["spinup"])
          fol = glob("run*",iroot["spinup"]); nfol = length(fol);
          itime["nspin"] = nfol;
    else; itime["nspin"] = 0;
    end

    return

end


__precompile__(false)

module MagneticLocSuchowiak

using Plots
using DataFrames
using CSV
using MLJ
using Distances
using Statistics

@load XGBoostRegressor
@load NeuralNetworkRegressor
@load KNNRegressor
@load NuSVR
@load RandomForestRegressor pkg=DecisionTree
#@load LGBMRegressor

include("graphs.jl")
include("datareader.jl")
include("models.jl")

export createfullgraph, createcomparisongraphs, readdata, fitknn, testmachines, fitforest, fitnusvr, fitcustommodel

end


import DataFrames

"""
"""
function fix_column_types!(
        df::DataFrames.AbstractDataFrame;
        categorical_feature_names::AbstractVector{Symbol} = Symbol[],
        continuous_feature_names::AbstractVector{Symbol} = Symbol[],
        categorical_label_names::AbstractVector{Symbol} = Symbol[],
        continuous_label_names::AbstractVector{Symbol} = Symbol[],
        float_type::Type{<:AbstractFloat} = Float64,
        )::Nothing
    function make_categorical_column(
            old_vector::AbstractVector,
            )::Vector{Any}
        new_vector = Vector{Any}(undef, length(old_vector))
        for i = 1:length(old_vector)
            old_vector_ith_element = old_vector[i]
            if DataFrames.ismissing(old_vector_ith_element)
                new_vector[i] = DataFrames.missing
            else
                new_vector[i] = string(old_vector_ith_element)
            end
        end
        return new_vector
    end
    function make_continuous_column(
            old_vector::AbstractVector,
            )::Vector{Any}
        new_vector = Vector{Any}(undef, length(old_vector))
        for i = 1:length(old_vector)
            old_vector_ith_element = old_vector[i]
            if DataFrames.ismissing(old_vector_ith_element)
                new_vector[i] = DataFrames.missing
            else
                new_vector[i] = float_type(old_vector_ith_element)
            end
        end
        return new_vector
    end
    for x in categorical_feature_names
        transform_columns!(df, make_categorical_column, x)
    end
    for x in continuous_feature_names
        transform_columns!(df, make_continuous_column, x)
    end
    for x in categorical_label_names
        transform_columns!(df, make_categorical_column, x)
    end
    for x in continuous_label_names
        transform_columns!(df, make_continuous_column, x)
    end
    for x in DataFrames.names(df)
        transform_columns!(df, fix_type, x)
    end
    return nothing
end

"""
"""
function check_column_types(
        df::DataFrames.AbstractDataFrame;
        categorical_feature_names::AbstractVector{Symbol} = Symbol[],
        continuous_feature_names::AbstractVector{Symbol} = Symbol[],
        categorical_label_names::AbstractVector{Symbol} = Symbol[],
        continuous_label_names::AbstractVector{Symbol} = Symbol[],
        )::Nothing
    for column_name in DataFrames.names(df)
        column_eltype=eltype(
            collect(
                DataFrames.skipmissing(
                    df[column_name]
                    )
                )
            )
        if column_name in categorical_feature_names
            if column_eltype<:AbstractString
            else
                error(
                    string(
                        "Column \"",
                        column_name,
                        "\" has eltype() \"",
                        column_eltype,
                        "\". However, this column is categorical,",
                        "and therefore its eltype() must be <: ",
                        "AbstractString.",
                        )
                    )
            end
        elseif column_name in categorical_label_names
            if column_eltype<:AbstractString
            else
                error(
                    string(
                        "Column \"",
                        column_name,
                        "\" has eltype() \"",
                        column_eltype,
                        "\". However, this column is categorical,",
                        "and therefore its eltype() must be <: ",
                        "AbstractString.",
                        )
                    )
            end
        elseif column_name in continuous_feature_names
            if column_eltype<:AbstractFloat
            else
                error(
                    string(
                        "Column \"",
                        column_name,
                        "\" has eltype() \"",
                        column_eltype,
                        "\". However, this column is continuous,",
                        "and therefore its eltype() must be <:",
                        "AbstractFloat.",
                        )
                    )
            end
        elseif column_name in continuous_label_names
            if column_eltype<:AbstractFloat
            else
                error(
                    string(
                        "Column \"",
                        column_name,
                        "\" has eltype() \"",
                        column_eltype,
                        "\". However, this column is continuous,",
                        "and therefore its eltype() must be <:",
                        "AbstractFloat.",
                        )
                    )
            end
        else
            if column_eltype<:AbstractString
            elseif column_eltype<:AbstractFloat
                error(
                    string(
                        "Column \"",
                        column_name,
                        "\" has eltype() \"",
                        column_eltype,
                        "\". However, we only allow AbstractStrings and ",
                        "AbstractFloats. Use T <: AbstractString if",
                        "it is a categorical column. Use T <: ",
                        "AbstractFloat if it is a continuous column."
                        )
                    )
            else
            end
        end
    end
    return nothing
end


#
# Iterator to join left and right streams on position
#

#
# uses vector as heap for right iterator
#
# iterator state is
# (lelt_s, relt_s, rq, rqidx)
# where lelt_s is current elt from left and state
#       relt_s is element, state that just did not make it into rq
#       rq is heap of rows from right to be reported withh lelt_s[1], ordered by (Chrom, bpEnd)
#       rqidx is index into rq to be reported on next iteration
#

#
# How do we combine elts from left and right stream?
#
# If we use an inner join, can simply use mergex from utils.jl.
# If we use a left join, all columns from right stream can be missing,
# i.e. we need to declare their types as Union{Missing, coltype}
# implemented as mergeleftx(left, right) or mergeleftx(left, R) in utils.jl
#

export join

using Tables
using DataStructures: heappush!, heappop!

abstract type AbstractJoin{L,R} end

struct InnerJoin{L,R} <: AbstractJoin{L,R}
    left::L
    right::R
    window::Int64
    kind::Symbol
    lendcol::Int64
    rendcol::Int64
end

struct LeftJoin{L,R} <: AbstractJoin{L,R}
    left::L
    right::R
    window::Int64
    kind::Symbol
    lendcol::Int64
    rendcol::Int64
end

Base.IteratorEltype(::Type{<:AbstractJoin{L,R}}) where {L,R} = Base.HasEltype()

Base.eltype(::Type{InnerJoin{L,R}}) where {L,R} =
    Base.return_types(mergex, (eltype(L), eltype(R)))[1]

Base.eltype(::Type{LeftJoin{L,R}}) where {L,R} =
    Base.return_types(mergeleftx, (eltype(L), Type{eltype(R)}))[1]

Base.IteratorSize(::Type{<:AbstractJoin{L,R}}) where {L,R} = Base.SizeUnknown()

Tables.istable(j::AbstractJoin{L,R}) where {L,R} = true
Tables.rowaccess(j::AbstractJoin{L,R}) where {L,R} = true
Tables.schema(j::AbstractJoin{L,R}) where {L,R} = Tables.Schema(eltype(j))
Tables.rows(j::AbstractJoin{L,R}) where {L,R} = j

# depending on kind of join
# the interval to be used for overlap is different
# :snpsnp  (start, start)  (start, start)
# :snpseg  (start, start)  (start, end)
# :segsnp  (start, end)    (start, start)
# :segseg  (start, end)    (start, end)

lendcol = (snpsnp = 2, snpseg = 2, segsnp = 3, segseg = 3)
rendcol = (snpsnp = 2, snpseg = 3, segsnp = 2, segseg = 3)

"""
    join(left, right; kind = :snpsnp, leftjoin = false, window = 0)

    left |> join(right; kind = :snpsnp, leftjoin = false, window = 0)

Join genome ordered streams `left` and `right` on `(elt[1], elt[2])`.

# Arguments

- `leftjoin::Bool`: should left join be performed
- `kind::Symbol`: how should overlap be determined (:snpsnp, :snpseg, :segsnp, :segseg)
- `window::Int`: allow `window` base pairs difference in position
"""
function join(left, right; kind = :snpsnp, leftjoin = false, window = 0)
    @assert kind in (:snpsnp, :snpseg, :segsnp, :segseg) "Unknown join kind: $kind"

    leftjoin ?
        LeftJoin{typeof(left), typeof(right)}(left, right, window, kind, lendcol[kind], rendcol[kind]) :
        InnerJoin{typeof(left), typeof(right)}(left, right, window, kind, lendcol[kind], rendcol[kind])
end

join(right; kind = :snpsnp, leftjoin = false, window = 0) =
    left -> join(left, right; leftjoin = leftjoin, window = window, kind = kind)

# given the current lelt_s,
# update rq to contain all rows to be reported (maybe more)

function advance_right!(j::AbstractJoin, state)
    lelt_s, relt_s, rq, rqidx = state

    # compute new criteria based on position of left for rows to be included in rq
    lmin = (lelt_s[1][1], lelt_s[1][2] - j.window)
    lmax = (lelt_s[1][1], lelt_s[1][j.lendcol] + j.window)

    # remove elts from front of rq that end before lmin
    while length(rq) > 0 && (rq[1][1] < lmin)
	heappop!(rq)
    end

    # add elts from right to rq if they overlap (lmin, lmax)
    while relt_s ≠ nothing

        # relt after lelt
        relt_s[1][1] > lelt_s[1][1] && break
        relt_s[1][1] == lelt_s[1][1] && relt_s[1][2] > lelt_s[1][j.lendcol] + j.window && break
        # overlap
        if lelt_s[1][1] == relt_s[1][1] && lelt_s[1][2] - j.window <= relt_s[1][j.rendcol] && relt_s[1][2] <= lelt_s[1][j.lendcol] + j.window
	    heappush!(rq, ((relt_s[1][1], relt_s[1][j.rendcol]), relt_s[1]))
        end

	relt_s = iterate(j.right, relt_s[2])
    end

    # return new state
    return (lelt_s, relt_s, rq, 1)
end

function Base.iterate(j::AbstractJoin)

    # prime state for first iteration
    lelt_s = iterate(j.left)
    relt_s = iterate(j.right)

    lelt_s == nothing && return nothing

    # heap, ordered by (rChrom, rEnd)
    rq = Vector{Tuple{Tuple{String, Int64}, eltype(j.right)}}()

    state = advance_right!(j, (lelt_s, relt_s, rq, 1))

    # and iterate
    iterate(j, state)
end

function Base.iterate(j::InnerJoin{L,R}, state) where {L,R}

    lelt_s, relt_s, rq, rqidx = state

    while true

	if rqidx <= length(rq)
            # need to test for overlap again
            # it could be that left ends before right starts
	    relt = rq[rqidx][2]
            if lelt_s[1][1] == relt[1] && lelt_s[1][2] - j.window <= relt[j.rendcol] && relt[2] <= lelt_s[1][j.lendcol] + j.window
	        return mergex(lelt_s[1], relt), (lelt_s, relt_s, rq, rqidx + 1)
            else
                rqidx += 1
            end
	end

	if rqidx > length(rq)
	    lelt_s = iterate(j.left, lelt_s[2])
	    lelt_s == nothing && return nothing
	    lelt_s, relt_s, rq, rqidx = advance_right!(j, (lelt_s, relt_s, rq, rqidx))
	end

    end
end

function Base.iterate(j::LeftJoin{L,R}, state) where {L,R}

    lelt_s, relt_s, rq, rqidx = state

    while true

	if length(rq) == 0 && rqidx == 1
	    return mergeleftx(lelt_s[1], eltype(R)), (lelt_s, relt_s, rq, rqidx + 1)
	end

	if rqidx <= length(rq)
	    relt = rq[rqidx][2]
            if lelt_s[1][1] == relt[1] && lelt_s[1][2] - j.window <= relt[j.rendcol] && relt[2] <= lelt_s[1][j.lendcol] + j.window
	        return mergex(lelt_s[1], relt), (lelt_s, relt_s, rq, rqidx + 1)
            else
                rqidx += 1
            end
	end

	if rqidx > length(rq)
	    lelt_s = iterate(j.left, lelt_s[2])
	    lelt_s == nothing && return nothing
	    lelt_s, relt_s, rq, rqidx = advance_right!(j, (lelt_s, relt_s, rq, rqidx))
	end

    end
end


# Testing a normalizing flow for seismic denosing
# Authors: Ali Siahkoohi, alisk@gatech.edu
#          Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
#          Philipp Witte, pwitte3@gatech.edu
# Date: September 2020
"""
Example:
julia test_warm-start_unsupervised_hint.jl
"""

using DrWatson
@quickactivate :FastApproximateInference

using InvertibleNetworks
using JLD
using JLD2
using Random
using Statistics
using ArgParse
using BSON
using Printf
using PyPlot
using Seaborn
using Flux: gpu, cpu
using Flux.Data: DataLoader

set_style("whitegrid")
rc("font", family="serif", size=13)
font_prop = matplotlib.font_manager.FontProperties(
    family="serif",
    style="normal",
    size=13
)
sfmt=matplotlib.ticker.ScalarFormatter(useMathText=true)
sfmt.set_powerlimits((0, 0))

s = ArgParseSettings()
@add_arg_table! s begin
    "--max_epoch"
        help = "Maximum # of epochs"
        arg_type = Int
        default = 15
    "--epoch"
        help = "Epoch to load net params"
        arg_type = Int
        default = 10
    "--lr"
        help = "starting learning rate"
        arg_type = Float32
        default = 1f-4
    "--lr_step"
        help = "lr scheduler applied at evey lr_step epoch"
        arg_type = Int
        default = 5
    "--batchsize"
        help = "batch size"
        arg_type = Int
        default = 4
    "--n_hidden"
        help = "# hidden channels"
        arg_type = Int
        default = 64
    "--depth"
        help = "depth of the network"
        arg_type = Int
        default = 12
    "--sigma"
        help = "noise std"
        arg_type = Float32
        default = 2f-1
    "--sim_name"
        help = "simulation name"
        arg_type = String
        default = "warm-start-unsupervised-hint-seismic-denoising-256x256-0.3CS"
end
parsed_args = parse_args(s)

max_epoch = parsed_args["max_epoch"]
lr = parsed_args["lr"]
lr_step = parsed_args["lr_step"]
batchsize = parsed_args["batchsize"]
n_hidden = parsed_args["n_hidden"]
depth = parsed_args["depth"]
sigma = parsed_args["sigma"]
sim_name = parsed_args["sim_name"]

if parsed_args["epoch"] == -1
    parsed_args["epoch"] = parsed_args["max_epoch"]
end
epoch = parsed_args["epoch"]

# Loading the experiment—only network weights and training loss
Params, fval, exp_path = load_experiment(parsed_args; return_path=true)

Y_fixed = wload(exp_path)["Array(Y_obs)"]
X_fixed = wload(exp_path)["X_true"]
Zy_fixed = wload(exp_path)["Array(Zy_obs)"]
Zx = wload(exp_path)["Array(Z_train)"]
A = wload(exp_path)["A"]
pretrained_args = wload(exp_path)["pretrained_args"]

nx, ny, n_in = size(X_fixed)[1:3]

CH = NetworkConditionalHINT(n_in, n_hidden, depth)
CH_rev = NetworkConditionalHINT(n_in, n_hidden, depth)

put_params!(CH_rev, convert(Array{Any,1}, Params))
CH_rev = CH_rev |> gpu

Params0, fval0 = load_experiment(pretrained_args)
put_params!(CH, convert(Array{Any,1}, Params0))

# Testing data
test_size = size(Zx, 4)

# Now select single fixed sample from all Ys
Y_fixed = Y_fixed |> gpu
Zy_fixed = Zy_fixed |> gpu

CH_rev = reverse(CH_rev)

# Draw new Zx, while keeping Zy fixed
test_size = 1000
X_post0 = zeros(Float32, nx, ny, n_in, test_size)
X_post = zeros(Float32, nx, ny, n_in, test_size)

CH = CH |> gpu

test_loader = DataLoader(randn(Float32, nx, ny, n_in, test_size) |> gpu, batchsize=8, shuffle=false)

for (itr, zx) in enumerate(test_loader)
    counter = (itr - 1)*size(zx)[4] + 1

    zx = zx |> gpu

    X_post0[:, :, :, counter:(counter+size(zx)[4]-1)] = (CH.inverse(zx, repeat(Zy_fixed, 1, 1, 1, size(zx, 4)))[1] |> cpu)
    X_post[:, :, :, counter:(counter+size(zx)[4]-1)] = (CH_rev.forward(zx, repeat(Zy_fixed, 1, 1, 1, size(zx, 4)))[1] |> cpu)
    @printf("Sampling: [%d/%d]\r", counter, test_size)
end

X_post = wavelet_unsqueeze(X_post)
X_post0 = wavelet_unsqueeze(X_post0)

X_fixed = wavelet_unsqueeze(X_fixed)
Y_fixed = wavelet_unsqueeze(Y_fixed)

# Some stats
X_post_mean = mean(X_post; dims=4)
X_post_std = std(X_post; dims=4)
X_post_mean0 = mean(X_post0; dims=4)
X_post_std0 = std(X_post0; dims=4)

save_dict = @strdict max_epoch epoch lr lr_step batchsize n_hidden depth sigma
save_path = plotsdir(sim_name, savename(save_dict; digits=6))

# Plot the conditional mean estimate
fig = figure("x_cm", figsize=(5, 5))
imshow(
    X_fixed[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title("True model")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "true_model.png"), fig)
close(fig)

fig = figure("x_cm", figsize=(5, 5))
imshow(
    Y_fixed[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Observed data")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "observed_data.png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("x_cm", figsize=(5, 5))
imshow(
    X_post_mean[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Conditional mean — Equation (6)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "conditional_mean.png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("x_cm", figsize=(5, 5))
imshow(
    X_post_mean0[:, :, 1, 1], vmin=-4.0, vmax=4.0, aspect=1, cmap="Greys", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Conditional mean — Equation (3)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "pretrained-conditional_mean.png"), fig)
close(fig)

# Plot the pointwise standard deviation
fig = figure("x_std", figsize=(5, 5))
imshow(
    X_post_std0[:, :, 1, 1], vmin=.30,
    vmax=.42, aspect=1, cmap="OrRd", resample=true,
    interpolation="lanczos", filterrad=1
)
title("STD — Equation (3)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "pretrained-pointwise_std.png"), fig)
close(fig)

# Plot the pointwise standard deviation
fig = figure("x_std", figsize=(5, 5))
imshow(
    X_post_std[:, :, 1, 1], vmin=.13,
    vmax=.19, aspect=1, cmap="OrRd", resample=true,
    interpolation="lanczos", filterrad=1
)
title("STD — Equation (6)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "pointwise_std.png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("err", figsize=(5, 5))
imshow(
    2f0.*(X_fixed[:, :, 1, 1] - X_post_mean[:, :, 1, 1])./(abs.(X_fixed[:, :, 1, 1])  + abs.(X_post_mean[:, :, 1, 1])), vmin=-2.0, vmax=2.0, aspect=1, cmap="twilight_shifted", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Relative error — Equation (6)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "conditional_mean_err.png"), fig)
close(fig)

# Plot the conditional mean estimate
fig = figure("err", figsize=(5, 5))
imshow(
    2f0.*(X_fixed[:, :, 1, 1] - X_post_mean0[:, :, 1, 1])./(abs.(X_fixed[:, :, 1, 1])  + abs.(X_post_mean0[:, :, 1, 1])), vmin=-2.0, vmax=2.0, aspect=1, cmap="twilight_shifted", resample=true,
    interpolation="lanczos", filterrad=1
)
title("Relative error — Equation (3)")
colorbar(fraction=0.047, pad=0.01, format=sfmt)
grid(false)
ax = plt.gca()
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
safesave(joinpath(save_path, "pretrained-conditional_mean_err.png"), fig)
close(fig)


using Metal

@show devices()
dev = MtlDevice(1)

bufferSize = 128
bufferA = MtlArray{Float32,1}(undef, tuple(bufferSize), storage=Shared)
bufferB = MtlArray{Float32,1}(undef, tuple(bufferSize), storage=Shared)
bufferC = MtlArray{Float32,1}(undef, tuple(bufferSize), storage=Shared)

vecA = unsafe_wrap(Vector{Float32}, bufferA.buffer, tuple(bufferSize))
vecB = unsafe_wrap(Vector{Float32}, bufferB.buffer, tuple(bufferSize))
vecC = unsafe_wrap(Vector{Float32}, bufferC.buffer, tuple(bufferSize))

using Random
rand!.([vecA, vecB])
vecC .= 0.0

## Setup
src = read(dirname(pathof(Metal))*"/Metal/kernels/vadd.metal", String)
opts = MtlCompileOptions()
lib = MtlLibrary(dev, src, opts)

fun = MtlFunction(lib, "add_vectors")
pip_addfun = MtlComputePipelineState(dev, fun)
queue = global_queue(dev) #MtlCommandQueue(dev)

##
argbufenc = MtlArgumentEncoder(fun, 1)
argbuf1 = alloc(Cchar, dev, sizeof(argbufenc), storage=Shared)
MTL.assign_argument_buffer!(argbufenc, argbuf1, 1)
MTL.set_field!(argbufenc, 128, 1)
set_buffer!(argbufenc, bufferA.buffer, 0, 2)
content1 = (128, Base.unsafe_load(Base.convert(Ptr{Int}, MTL.content(argbuf1)),2))

argbufenc = MtlArgumentEncoder(fun, 2)
argbuf2 = alloc(Cchar, dev, sizeof(argbufenc), storage=Shared)
MTL.assign_argument_buffer!(argbufenc, argbuf2, 1)
MTL.set_field!(argbufenc, 128, 1)
set_buffer!(argbufenc, bufferB.buffer, 0, 2)
content2 = (128, Base.unsafe_load(Base.convert(Ptr{Int}, MTL.content(argbuf2)),2))

argbufenc = MtlArgumentEncoder(fun, 3)
argbuf3 = alloc(Cchar, dev, sizeof(argbufenc), storage=Shared)
MTL.assign_argument_buffer!(argbufenc, argbuf3, 1)
MTL.set_field!(argbufenc, 128, 1)
set_buffer!(argbufenc, bufferC.buffer, 0, 2)
content3 = (128, Base.unsafe_load(Base.convert(Ptr{Int}, MTL.content(argbuf3)),2))

rc1 = Base.RefValue(content1)
rc2 = Base.RefValue(content2)
rc3 = Base.RefValue(content3)
pt1 = Base.unsafe_convert(Ptr{typeof(content1)}, rc1)
pt2 = Base.unsafe_convert(Ptr{typeof(content1)}, rc2)
pt3 = Base.unsafe_convert(Ptr{typeof(content1)}, rc3)

cmd = MTL.commit!(queue) do cmdbuf
    MtlComputeCommandEncoder(cmdbuf) do enc
        MTL.use!(enc, bufferA.buffer, MTL.ReadWriteUsage)
        MTL.use!(enc, bufferB.buffer, MTL.ReadWriteUsage)
        MTL.use!(enc, bufferC.buffer, MTL.ReadWriteUsage)
        #MTL.set_buffers!(enc,
        #                        [argbuf1, argbuf2, argbuf3],
        #                        [0,0,0], 1:3)
        MTL.set_bytes!(enc, pt1, 16, 1)
        MTL.set_bytes!(enc, pt2, 16, 2)
        MTL.set_bytes!(enc, pt3, 16, 3)
        gridSize = MtSize(length(vecA), 1, 1)
        MTL.set_function!(enc, pip_addfun)
        threadGroupSize = min(length(vecA), pip_addfun.maxTotalThreadsPerThreadgroup)
        threadGroupSize = MTL.MtSize(threadGroupSize, 1, 1)
        MTL.append_current_function!(enc, gridSize, threadGroupSize)
    end
end

####

# Execute
wait(cmd)

@show vecC


# options checking specific to the Euler Equations

"""
  This function checks the options in the options dictionary after default
  values have been supplied and throws exceptions if unsupported/incompatable
  options are specified

  **Inputs**

   * opts: options dictionary
   * comm: MPI communicator

  **Outputs**

    none
"""
function checkOptions(opts, comm=MPI.COMM_WORLD)

  # this is a problem for nested physics
#  if opts["physics"] != PhysicsName
#    error("physics not specified as $PhysicsName, are you lost?")
#  end

  if opts["use_edge_res"]
    error("use_edge_res does not work correctly, do not use")
  end

  if opts["use_staggered_grid"]
    if opts["volume_integral_type"] != 2
      error("cannot use staggered grids with non entropy stable volume integrals")
    end
#=
    if opts["face_integral_type"] != 2
      error("cannot use staggered grids with non entropy stable face integrals")
    end
=#
    if opts["operator_type2"] == "SBPNone"
      error("must specify a real SBP operator for staggered grid")
    end
  end

  run_type = opts["run_type"]
  #TODO: what to do for jac_type 4?
  if opts["use_DG"] &&
     opts["Flux_name"] in keys(FluxDict_diff) &&
     (run_type == 5 || run_type == 40 || run_type == 20 || run_type == 50) &&
     opts["jac_type"] >= 2 #&&
#     !opts["use_src_term"]
    get!(opts, "calc_jac_explicit", true)
  else
    get!(opts, "calc_jac_explicit", false)
    if opts["calc_jac_explicit"]
      error("cannot calculate jacobian explicitly with given options")
    end
  end

  if opts["calc_jac_explicit"]
    get!(opts, "preallocate_jacobian_coloring", false)
  else
    get!(opts, "preallocate_jacobian_coloring", true)
  end

  get!(opts, "homotopy_function", "FirstOrderDissipation")
  get!(opts, "freeze_viscosity", false)
  get!(opts, "shock_capturing_variables", "IR")
  get!(opts, "sensor_pp_use_filtered", false)
  get!(opts, "sensor_pp_s0_offset", 0)

  return nothing
end

import PDESolver: _getSparsityPattern

function _getSparsityPattern(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts)

  if !opts["addShockCapturing"]
    return INVISCID
  else
    return getSparsityPattern(eqn.shock_capturing)
  end

end


using ACE1pack, Test, JuLIP, LazyArtifacts
using ACE1.Testing: println_slim
using JuLIP.Testing: print_tf

test_train_set = joinpath(artifact"TiAl_tiny_dataset", "TiAl_tiny.xyz")

data = Dict(
    "energy_key"   => "energy",
    "fname" => test_train_set,
    "virial_key"   => "virial")

@info("Quick test for filling in missing param entries with defaults")
data = ACE1pack.fill_defaults!(data, param_key = "data")
println_slim(@test "force_key" in collect(keys(data)))

@info("Test multitransform (it's a bit trickier)")
multitransform_params = Dict(
    "cutoffs" => Dict(
       ("C", "C") => (1, 2),
       ("C", "H") => (1, 3),
       ("H", "H") => (1, 3)),
    "transforms" => Dict(
        ("C", "C") => Dict("type"=> "polynomial"),
        ("C", "H") => Dict("type"=> "polynomial"),
        ("H", "H") => Dict("type" => "polynomial")),
    "type" => "multitransform")

out = fill_defaults!(multitransform_params, param_key="transform")
for (key, val) in out["transforms"]
    for extra_key in ["r0", "p"]
        print_tf(@test haskey(val, extra_key))
    end
end
println()

@info("Test filling in ace_basis without \"type\" among radial basis parameters")
ace_params = Dict(
    "species" => "something",
    "N" => 2,
    "maxdeg" => 2,
    "radial" => Dict(
        "rin" => 2.0,
        "rcut" => 7.0
    ))
out = fill_defaults!(ace_params, param_key="ace")

@info("Test parsing string of tuple into tuple of strings in basis")
ace_basis_params = Dict(
        "degree" => Dict(
            "type" => "sparseM",
            "Dn" => Dict("default"=> 1.0),
            "Dl" => Dict("default"=> 1.5),
            "Dd" => Dict(
                "default" => 10,
                "(4, C)" => 8)),
        "transform" => Dict(
            "cutoffs" => Dict("(C, C)" => "(1.0, 2.3)"),
            "transforms" => Dict(
                "(C, C)" => Dict(
                    "type" => "polynomial")),
            "type" => "multitransform"))

out = parse_ace_basis_keys(ace_basis_params)
print_tf(@test haskey(out["degree"]["Dd"], (4, "C")))
print_tf(@test haskey(out["transform"]["transforms"], ("C", "C")))
print_tf(@test haskey(out["transform"]["cutoffs"], ("C", "C")))
println()

@info("Test that all *params get filled in correctly on smallest allowed input.")

@info("fit_params")
minimal_params = Dict(
    "data" => Dict(
        "fname" => "something"),
    "basis" => Dict(
        "ace" => Dict(
            "species" => "something",
            "N" => 1,
            "maxdeg" => 1,
            "type" => "ace"
            ),
        "pair" => Dict(
            "species" => "sth",
            "maxdeg" => 1,
            "type" => "pair"
            ),),
    "solver" => Dict(
        "type" => "rrqr"),
    "e0" => "something")

filled_params = fill_defaults!(minimal_params)
for extra_key in ["ACE_fname", "LSQ_DB_fname_stem", "fit_from_LSQ_DB"]
    print_tf(@test haskey(filled_params, extra_key))
end
for extra_key in ["energy_key", "force_key", "virial_key"]
    print_tf(@test haskey(filled_params["data"], extra_key))
end
for extra_key in ["radial", "transform", "degree"]
    print_tf(@test haskey(filled_params["basis"]["ace"], extra_key))
end
for extra_key in ["rcut", "rin", "pcut", "pin", "transform"]
    print_tf(@test haskey(filled_params["basis"]["pair"], extra_key))
end
for extra_key in ["rrqr_tol"]
    print_tf(@test haskey(filled_params["solver"], extra_key))
end
println()

@info("The rest to be or not to be done...")


#### Structs for positions in the hamiltonian space

abstract type HMC_State end

mutable struct Tree_HMC_State{T<:GeneralNode} <: HMC_State
    x::T
    r::Vector{Float64}
    g::Vector{Float64}
    lf::Float64
    extended::Bool
    nni::Int64
end

function Tree_HMC_State(tree::T, r::Vector{Float64}, g::Vector{Float64}, lf::Float64)::Tree_HMC_State{T} where T<:GeneralNode
    Tree_HMC_State(tree, r, g, lf, false, 0)
end

mutable struct Array_HMC_State{T<:Array{<:Real}} <: HMC_State
    x::T
    r::Vector{Float64}
    g::Vector{Float64}
    lf::Float64
end

function transfer(s1::T)::T where T<:HMC_State
    T(deepcopy(s1.x), s1.r[:],s1.g[:],s1.lf, s1.extended, s1.nni)
end

function transfer(s1::T)::T where T<:Array_HMC_State
    T(s1.x[:], s1.r[:],s1.g[:],s1.lf)
end

hamiltonian(s::HMC_State) = s.lf - 0.5 * dot(s.r, s.r)

function transfer!(s1::HMC_State, s2::HMC_State)
    s1.x = deepcopy(s2.x)
    s1.r[:] .= s2.r[:]
    s1.g[:] .= s2.g[:]
    s1.lf = s2.lf
    s1.extended = s2.extended
    s1.nni = s2.nni
    nothing
end

function transfer!(s1::Array_HMC_State, s2::Array_HMC_State)
    s1.x[:] .= s2.x[:]
    s1.r[:] .= s2.r[:]
    s1.g[:] .= s2.g[:]
    s1.lf = s2.lf
    nothing
end

#### PNUTS Parameter Structs

struct NUTS_StepParams
    μ::Float64
    δ::Float64
    γ::Float64
    κ::Float64
    t0::Int
    τ::Float64
end

function update_step(p::NUTS_StepParams, μ::Float64)
    p1 = NUTS_StepParams(μ, p.δ, p.γ, p.κ, p.t0, p.τ)
    p1
end

mutable struct NUTSstepadapter
    m::Int
    s_bar::Float64
    x_bar::Float64
    params::NUTS_StepParams
    metro_acc_prob::Float64
    avg_nni::Float64

    function NUTSstepadapter(m, s_bar, x_bar, μ, δ, γ, κ, t0, τ)
        new(m, s_bar, x_bar, NUTS_StepParams(μ, δ, γ, κ, t0, τ), 0.0, 0.0)
    end

    function NUTSstepadapter(m, s_bar, x_bar, params::NUTS_StepParams)
        new(m, s_bar, x_bar, params, 0.0, 0.0)
    end
end

mutable struct NUTSMeta
    nni::Float64
    alpha::Float64
    nalpha::Float64
    accnni::Float64
end

NUTSMeta() = NUTSMeta(0.0,0.0,0.0,0.0)

function update!(x::NUTSMeta, y::NUTSMeta)
    x.nni += y.nni
    x.alpha += y.alpha
    x.nalpha += y.nalpha
    x.accnni += y.accnni
end


# TODO: Need to finish
#       Need to test
function fit_fs_tp!(init::StateFS,
                    cfs::ConstantsFS,
                    dfs::DataFS;
                    nmcmc::Int, nburn::Int,
                    # These args are for iMCMC:
                    batchprop::Float64=0.1,
                    batchsizes=nothing,
                    prior_thin::Int=2,
                    temper::Float64=1.0, mb_update_burn_prop=0.6, anneal=false,
                    min_temper::Float64=1.0,
                    # End of iMCMC args.
                    tfs::Union{Nothing, Vector{TunersFS}}=nothing,
                    monitors=[monitor1, monitor2],
                    fix::Vector{Symbol}=Symbol[],
                    thins::Vector{Int}=[2, nsamps_to_thin(10, nmcmc)],
                    thin_dden::Int=1,
                    printFreq::Int=0,
                    checkpoint=0,
                    computeDIC::Bool=false, computeLPML::Bool=false,
                    computedden::Bool=false,
                    sb_ibp::Bool=false,
                    use_repulsive::Bool=true,
                    Z_marg_lamgam::Float64=1.0,
                    Z_marg_lamgam_decay_rate::Float64=100.0,
                    Z_marg_lamgam_min::Float64=0.05,
                    verbose::Int=1,
                    time_updates=false,
                    seed::Int=-1)

  @assert 0 <= mb_update_burn_prop <= 1

  printMsg(iter::Int, msg::String) = if printFreq > 0 && iter % printFreq == 0
    print(msg)
  end

  # Set random seed if needed
  if seed >= 0
    Random.seed!(seed)
  end

  # Print temperature
  println("temperature: $(temper)")

  # Create Tuners if needed
  if tfs == nothing
    tfs = TunersFS(Tuners(dfs.data.y, cfs.constants.K),
                   init.theta,
                   dfs.X)
  end

  if verbose >= 1
    fixed_vars_str = join(fix, ", ")
    if fixed_vars_str == ""
      fixed_vars_str = "nothing"
    end
    println("fixing: $fixed_vars_str")
    println("Use stick-breaking IBP: $(sb_ibp)")
    println("Z_marg_lamgam: $(Z_marg_lamgam)")
    println("Z_marg_lamgam_decay_rate: $(Z_marg_lamgam_decay_rate)")
    println("Z_marg_lamgam_min: $(Z_marg_lamgam_min)")
    println("use_repulsive: $(use_repulsive)")
    flush(stdout)
  end

  @assert printFreq >= -1
  if printFreq == 0
    numPrints = 10
    printFreq = Int(ceil((nburn + nmcmc) / numPrints))
  end

  # Instantiate (but not initialize) CPO stream
  if computeLPML
    cpoStream = MCMC.CPOstream{Float64}()
  end

  # DIC
  if computeDIC
    dicStream = initDicStream(dfs.data)
    loglikeDIC(param::DICparam) = computeLoglikeDIC(dfs.data, param)

    convertStateToDicParam(s::State)::DICparam = let
      _convertStateToDicParam(s, cfs.constants, dfs.data)
    end
  end

  # Cache for data density
  dden = Matrix{Vector{Float64}}[]

  # Cache for loglikelihood
  ll = Float64[]

  # Update function
  function update!(state, iter::Int, out)
    # Whether or not to marginalize over lambda and gamma.
    # We want to do this more often at the beginning, and less at the end.
    zmarg = ((Z_marg_lamgam - Z_marg_lamgam_min) *
             exp(-iter/Z_marg_lamgam_decay_rate) +
             Z_marg_lamgam_min) > rand()

    # Iterations to do minibatch update
    iters_mb_update_all_params = round(Int, nburn * mb_update_burn_prop)

    # Calculate current temperature
    curr_temper = if anneal
      # Gradually cool the temperature so that it is 1 when
      # we reach iteration `iters_mb_update_all_params`, which
      # is less than `nburn`.
      itmax = iters_mb_update_all_params
      power = (itmax - iter + 1) / itmax
      clamp(temper^power, min_temper, Inf)
    else
      temper
    end
    if verbose > 0
      print(" -- Current temper: $(curr_temper)")
    end

    # Update state using trained prior
    minibatch_update_all_params = iter < iters_mb_update_all_params || !anneal
    update_via_trained_prior!(state, dfs, cfs, tfs, batchprop, prior_thin,
                              batchsizes=batchsizes,
                              fix=fix, use_repulsive=use_repulsive,
                              Z_marg_lamgam=zmarg, sb_ibp=sb_ibp,
                              time_updates=time_updates, temper=curr_temper,
                              minibatch_update_all_params=minibatch_update_all_params,
                              verbose=verbose-2)

    # Pull out inner componenets for convenience
    s = state.theta
    c = cfs.constants
    d = dfs.data

    # Append loglike
    append!(ll, compute_marg_loglike(s, c, d, 1.0) / curr_temper)

    if computedden && iter > nburn && (iter - nburn) % thin_dden == 0
      # NOTE: `datadensity(s, c, d)` returns an (I x J) matrix of vectors of
      # length g.
      append!(dden, [datadensity(s, c, d)])
    end

    if computeLPML && iter > nburn
      # Inverse likelihood for each data point
      like = [[compute_like(i, n, s, c, d)
               for n in 1:d.N[i]] for i in 1:d.I]

      # Update (or initialize) CPO
      MCMC.updateCPO(cpoStream, vcat(like...))

      # Add to printMsg
      printMsg(iter, " -- LPML: $(MCMC.computeLPML(cpoStream))")
    end

    if computeDIC && iter > nburn
      # Update DIC
      MCMC.updateDIC(dicStream, s, updateParams,
                     loglikeDIC, convertStateToDicParam)

      # Add to printMsg
      printMsg(iter, " -- DIC: $(MCMC.computeDIC(dicStream, loglikeDIC,
                                                 paramMeanCompute))")

      DICg = MCMC.DIC_gelman(ll[(nburn+1):end])
      printMsg(iter, " -- DIC_gelman: $(DICg)")
    end

    printMsg(iter, "\n")
    flush(stdout)
  end

  println("Running Gibbs sampler ...")
  samples, state = MCMC.gibbs(deepcopy(init), update!, monitors=monitors,
                              thins=thins, nmcmc=nmcmc, nburn=nburn,
                              printFreq=printFreq, loglike=ll,
                              printlnAfterMsg=false)

  out = Dict(:samples => samples,
             :lastState => state,
             :init => init,
             :ll => ll,
             :c => cfs,
             :d => dfs)

  if computeDIC || computeLPML
    LPML = computeLPML ? MCMC.computeLPML(cpoStream) : NaN
    Dmean, pD = computeDIC ? MCMC.computeDIC(dicStream, loglikeDIC,
                                             paramMeanCompute,
                                             return_Dmean_pD=true) : (NaN, NaN)

    DICg = MCMC.DIC_gelman(ll[(nburn+1):end])

    metrics = Dict(:LPML => LPML,
                   :DIC => Dmean + pD,
                   :Dmean => Dmean,
                   :pD => pD,
                   :DICg => DICg)
    println()
    println("metrics:")
    for (k, v) in metrics
      out[k] = v
      println("$k => $v")
    end
    flush(stdout)
  end

  if computedden
    out[:dden] = dden
  end

  out[:nburn] = nburn
  out[:nmcmc] = nmcmc
  out[:batchprop] = batchprop
  out[:batchsizes] = batchsizes
  out[:prior_thin] = prior_thin
  out[:temper] = temper
  out[:anneal] = anneal

  return out
end


@testset "simple workflow" begin
    env = StateTransformedEnv(CartPoleEnv{Float32}();state_mapping=deepcopy)
    policy = RandomPolicy(action_space(env))
    N_EPISODE = 10_000
    hook = TotalRewardPerEpisode()
    run(policy, env, StopAfterEpisode(N_EPISODE), hook)

    @test isapprox(sum(hook[]) / N_EPISODE, 21; atol=2)
end

@testset "multi agent" begin
    # https://github.com/JuliaReinforcementLearning/ReinforcementLearning.jl/issues/393
    rps = RockPaperScissorsEnv() |> SequentialEnv
    ma_policy = MultiAgentManager(
        (
            NamedPolicy(p => RandomPolicy())
            for p in players(rps)
        )...
    )

    run(ma_policy, rps, StopAfterEpisode(10))
end


#=
Create training file
=#

function set_up_training(han::Tracker_Handles,get_mean=true)
    woi=get_woi_array(han)
    set_up_training(han.nn,han.wt.vid_name,han.end_frame,woi,han.wt.pad_pos,han.frame_list,get_mean)
end

function set_up_training(nn,vid_name,max_frames,woi,pad_pos,frame_list,get_mean=true)

    (w,h,fps)=get_vid_dims(vid_name)

    if get_mean
        (mean_img,std_img)=mean_std_video_gpu(vid_name,max_frames)
        nn.norm.min_ref = 0
        nn.norm.max_ref = 255
        nn.norm.mean_img = mean_img
        nn.norm.std_img = std_img

        #Rotate and Reshape to 256 256
        nn.norm.mean_img = reshape(imresize(nn.norm.mean_img[:,:,1]',(256,256)),(256,256,1))
    end

    (new_woi, new_frame_list) = check_whiskers(woi,frame_list,max_frames)

    WT_reorder_whisker(new_woi,pad_pos)

    nn.labels = make_heatmap_labels(new_woi,pad_pos)
    nn.imgs = get_labeled_frames(vid_name,new_frame_list);

    #Normalize
    nn.imgs = StackedHourglass.normalize_new_images(nn.imgs,nn.norm.mean_img);

    (nn.imgs,nn.labels)=augment_images(nn.imgs,nn.labels);
end

function create_new_weights(nn)
    hg=HG2(size(nn.labels,1),size(nn.labels,3),4);
    load_hourglass(nn.weight_path,hg)
    StackedHourglass.change_hourglass(hg,size(nn.labels,1),1,size(nn.labels,3))
    nn.features=features(hg)
    nn.hg = hg
    nn.use_existing_weights=true
    nothing
end

#=
Run Training methods
=#

function make_training_batch(img,ll,batch_size=8)
    dtrn=minibatch(img,ll,batch_size,xtype=KnetArray,ytype=KnetArray)
end

function run_training(hg,trn::Knet.Data,this_opt,p,epochs=100,ls=Array{Float64,1}())

    total_length=length(trn) * epochs
    minimizer = Knet.minimize(hg,ncycle(trn,epochs),this_opt)
    last_update = 0.0
    count=0
    Gtk.set_gtk_property!(p, :fraction, 0)
    sleep(0.0001)

        for x in takenth(minimizer,1)
            push!(ls,x)
            count+=1
            complete=round(count/total_length,digits=2)
            if complete > last_update
                Gtk.set_gtk_property!(p, :fraction, complete)
                last_update = complete
                reveal(p,true)
            end
            sleep(0.0001)
        end

    ls
end

function run_training_no_gui(hg,trn::Knet.Data,this_opt,epochs=100,ls=Array{Float64,1}())

    total_length=length(trn) * epochs
    minimizer = Knet.minimize(hg,ncycle(trn,epochs),this_opt)

    for x in takenth(minimizer,1)
        push!(ls,x)
        sleep(0.0001)
    end

    ls
end

#=
Error Checking before training
=#

#Make sure that dimensions of model weights match label dimensions
function check_hg_features(nn)

    #Output features
    if size(nn.labels,3) != features(nn.hg)
        StackedHourglass.change_hourglass_output(nn.hg,size(nn.labels,1),size(nn.labels,3))
        nn.features = features(nn.hg)
    end
end

function check_whiskers(woi,frame_list,max_frames)

    #Make sure all of the frames actually have a tracked whisker
    empty=trues(length(woi))
    for i=1:length(woi)
        if empty_whiskers(woi[i])
            empty[i] = false
            println("You have an untraced whisker at index: ", i)
        end
    end

    #Last frames don't always work
    if (frame_list[end] == max_frames)
        empty[end] = false
    end

    new_woi = woi[empty]
    new_frame_list = frame_list[empty]

    (new_woi, new_frame_list)
end

function empty_whiskers(w)

    (length(w.x)==0)|(length(w.y)==0)

end


# All necessary variables (ie, setup vars) should be defined here, to be
# shared between any disparate processes
# push!(LOAD_PATH, "/home/ian/Grad/Research/Disks/DiskJockey/")
# push!(LOAD_PATH, "/n/home07/iczekala/DiskJockey/")
# push!(LOAD_PATH, "/pool/scout0/DiskJockey/")

module consonance

export pars, nframes, nframes_per_proc, nchan, vels, lams, vmax, moviedir
export vmax_d, vmax_g, vmax_v

using DiskJockey.constants
using DiskJockey.model

# For odyssey use
# homedir = "/n/home07/iczekala/DiskJockey/"
# scratchdir = "/scratch/"
# outdir = "/n/home07/iczekala/DiskJockey/output/movie/" #Where the images.out and images.png are stored

# For local use
moviedir = "/home/ian/Grad/Research/Disks/DiskJockey/movie/"

# For cfa CF use
# homedir = "/pool/scout0/DiskJockey/"
# scratchdir = "/pool/cf/iczekala/scratch/"
# outdir = "/pool/scout0/DiskJockey/output/movie/"

# How many frames per process?
nframes_per_proc = 50

global const nchan = 5
global const vels = linspace(-1.5, 1.5, nchan) # [km/s]
# CO 2-1 rest frame
lam0 = cc/230.538e9 * 1e4 # [microns]
# convert velocities to wavelengths
lams = lam0 * (vels/c_kms + 1)

M_star = 1.75 # [M_sun] stellar mass
r_c = 45. # [AU] characteristic radius
r_in = 1.0 # [AU]
r_cav = 2.0 # [AU]
delta = 1.0 # no change
T_10 = 115. # [K] temperature at 10 AU
q = 0.63 # temperature gradient exponent
gamma = 1.0 # surface density gradient
logM_CO = 0.2 # [M_earth] disk mass of CO
ksi = 0.14 # [km/s] microturbulence
dpc = 73. # [pc] distance
incl = 45. # [degrees] inclination
PA = 0. # [degrees] position angle
vel = 0.0 # [km/s]
mu_RA = 0.0 # [arcsec] centroid location
mu_DEC = 0.0 # [arcsec]

# Given a starting parameter, vary it by steps dp to the low bound, then high
# bound, then back to where we started

# In the new edition, keep it at the top and bottom values for some period
function smooth_vary(start, low, high, dp, npause)
    ndown1 = round(Integer, (start - low)/dp) + 1
    nup = round(Integer, (high - low)/dp) + 1
    ndown2 = round(Integer, (high - start)/dp) + 1

    return [linspace(start, low, ndown1)' low * ones(npause)' linspace(low, high, nup)' high * ones(npause)' linspace(high, start, ndown2)' start * ones(npause)']
end

# Create a master parameter list
# First, adjust in radius
# radiuses = smooth_vary(r_c, 25., 65., 0.5)
# nradiuses = length(radiuses)

# First, adjust in inclination
incls = smooth_vary(incl, 0., 90., 1., 15)
nincls = length(incls)

# then, adjust in mass
masses = smooth_vary(M_star, 1.0, 2.5, 0.02, 15)
nmasses = length(masses)

# nframes = nincls + nmasses
#
# # Now create a giant array of Parameters objects
# pars = Array(Parameters, nframes)
#
# for i=1:nframes
#     if i <= nincls
#         pars[i] = Parameters(M_star, r_c, r_in, r_cav, delta, T_10, q, gamma, 10^logM_CO, ksi, dpc, incls[i], PA, vel, mu_RA, mu_DEC)
#     else
#         pars[i] = Parameters(masses[i - nincls], r_c, r_in, r_cav, delta, T_10, q, gamma, 10^logM_CO, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC)
#     end
# end
#
# println("There are ", nframes, " frames to be generated.")
# println("Inclinations ", incls)
# println("Masses ", masses)
# println("Radiuses ", radiuses)

# Return a normalized instance that is symmetric about 0
# function scale(data)
#     s = maximum(abs(data))
#     return norm = plt.Normalize(vmin=-s, vmax=s, clip=false)
# end

# used as arg to imshow

# Poster plot material
nframes = 5
pars = [Parameters(1.0, r_c, r_in, r_cav, delta, T_10, q, gamma, 10^logM_CO, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC),
Parameters(1.75, r_c, r_in, r_cav, delta, T_10, q, gamma, 10^logM_CO, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC),
Parameters(2.50, r_c, r_in, r_cav, delta, T_10, q, gamma, 10^logM_CO, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC),
Parameters(1.75, r_c, r_in, r_cav, delta, T_10, q, gamma, 10^logM_CO, ksi, dpc, 0.0, PA, vel, mu_RA, mu_DEC),
Parameters(1.75, r_c, r_in, r_cav, delta, T_10, q, gamma, 10^logM_CO, ksi, dpc, 90.0, PA, vel, mu_RA, mu_DEC)
]

vmin_d = 0.0
vmax_d = 6.415994694015499e9

vmin_g = 0.0
vmax_g = 2.9207242096782996e11

vmin_v = 0.0
vmax_v = 45.0193854891226

end # module


# Test backward compatibility

### Model
nsrc = 2
model, model0, dm = setup_model(tti, viscoacoustic, nlayer)
q, srcGeometry, recGeometry = setup_geom(model; nsrc=nsrc)
dt = srcGeometry.dt[1]
nt = srcGeometry.nt[1]
nrec = length(recGeometry.xloc[1])

@testset "Backward compatibility" begin
    @timeit TIMEROUTPUT "Backward compatibility" begin
        info =  Info(prod(model.n), nt, nsrc)
        @test_logs (:warn, "Info is deprecated and will be removed in future versions") Info(prod(model.n), nt, nsrc)
        @test_logs (:warn, "judiModeling(info::Info, ar...; kw...) is deprecated, use judiModeling(ar...; kw...)") judiModeling(info, model)
        @test_logs (:warn, "judiModeling(info::Info, ar...; kw...) is deprecated, use judiModeling(ar...; kw...)") judiModeling(info, model; options=Options())
        @test_logs (:warn, "judiProjection(info::Info, ar...; kw...) is deprecated, use judiProjection(ar...; kw...)") judiProjection(info, recGeometry)
        @test_logs (:warn, "judiWavefield(info::Info, ar...; kw...) is deprecated, use judiWavefield(ar...; kw...)") judiWavefield(info, dt, nsrc, randn(nt, model.n...))

        @test_throws ArgumentError judiRHS(info, recGeometry, randn(Float32, nt, nrec))
        @test_throws ArgumentError judiLRWF(info, nsrc, randn(nt))

    end
end


module InfiniteOpt

# Import and export JuMP
import Reexport
Reexport.@reexport using JuMP

# Import the necessary packages.
import MathOptInterface
import Distributions
import DataStructures
import FastGaussQuadrature

# Make useful aliases (note we get MOI and MOIU from JuMP)
const JuMPC = JuMP.Containers
const MOIUC = MOIU.CleverDicts

# Import the Collections module
include("Collections/Collections.jl")

# Import core methods
include("datatypes.jl")
include("infinite_domains.jl")
include("scalar_parameters.jl")
include("array_parameters.jl")
include("variable_basics.jl")
include("infinite_variables.jl")
include("semi_infinite_variables.jl")
include("point_variables.jl")
include("finite_variables.jl")
include("expressions.jl")
include("measures.jl")

# Import and export MeasureToolbox
include("MeasureToolbox/MeasureToolbox.jl")
Reexport.@reexport using .MeasureToolbox

# import more core methods
include("derivatives.jl")
include("constraints.jl")
include("macros.jl")
include("objective.jl")
include("measure_expansions.jl")
include("derivative_evaluations.jl")
include("optimize.jl")
include("results.jl")
include("show.jl")
include("utilities.jl")
include("general_variables.jl")

# Import and export TranscriptionOpt
include("TranscriptionOpt/TranscriptionOpt.jl")
Reexport.@reexport using .TranscriptionOpt

# Deprecations introduced with v0.4.0
Base.@deprecate(ParameterBounds, DomainRestrictions)
Base.@deprecate(has_parameter_bounds, has_domain_restrictions)
Base.@deprecate(parameter_bounds, domain_restrictions)
Base.@deprecate(set_parameter_bounds, set_domain_restrictions)
Base.@deprecate(add_parameter_bounds, add_domain_restrictions)
Base.@deprecate(delete_parameter_bounds, delete_domain_restrictions)
Base.@deprecate(IntervalSet, IntervalDomain)
Base.@deprecate(UniDistributionSet, UniDistributionDomain)
Base.@deprecate(MultiDistributionSet, MultiDistributionDomain)
Base.@deprecate(CollectionSet, CollectionDomain)
Base.@deprecate(supports_in_set, supports_in_domain)
Base.@deprecate(collection_sets, collection_domains)
Base.@deprecate(infinite_set, infinite_domain)
Base.@deprecate(set_infinite_set, set_infinite_domain)
Base.@deprecate(InfiniteParameterFunction, ParameterFunction)

for op in (:has_parameter_bounds, :parameter_bounds, :set_parameter_bounds,
           :add_parameter_bounds, :delete_parameter_bounds)
    @eval begin
        function $op(::JuMP.AbstractVariableRef, args...; kwargs...)
            error("The use of parameter bounds (now called domain ",
                  "domain restrictions) has been discontinued for finite ",
                  "variables. The preferred syntax is to specify ",
                  "`DomainRestrictions` in connection with constraints.")
        end
    end
end

macro independent_parameter(model, args...)
    error("`@independent_parameter` has been deprecated in favor of ",
          "`@infinite_parameter`. ",
          "\n\nOld Syntax: `@independent_parameter(model, par[...] in set, ",
          "kwargs...)`\nNew Syntax: `@infinite_parameter(model, par[...] in ",
          "set, independent = true, kwargs...)`.")
end
macro dependent_parameters(model, args...)
    error("`@Dependent_parameters` has been deprecated in favor of ",
          "`@infinite_parameter`. ",
          "\n\nOld Syntax: `@dependent_parameters(model, par[...] in set, ",
          "kwargs...)`\nNew Syntax: `@infinite_parameter(model, par[...] in ",
          "set, kwargs...)`.")
end
macro infinite_variable(model, args...)
    error("`@infinite_variable` has been deprecated in favor of `@variable`. ",
          "\n\nOld Syntax: `@infinite_variable(model, var[idxs...](params...), ",
          "args..., kwargs...)`\nNew Syntax: `@variable(model, var[idxs...], ",
          "Infinite(params...), args..., kwargs...)`.")
end
macro point_variable(model, args...)
    error("`@point_variable` has been deprecated in favor of `@variable`. ",
          "\n\nOld Syntax: `@point_variable(model, ivref(param_values...), ",
          "var_expr, args..., kwargs...)`\nNew Syntax: `@variable(model, ",
          "var_expr, Point(ivref, param_values...), args..., kwargs...)`.")
end
macro hold_variable(model, args...)
    error("`@hold_variable` has been deprecated in favor of `@variable`.",
          "\n\nOld Syntax: `@hold_variable(model, var_expr, args..., ",
          "kwargs...)`\nNew Syntax: `@variable(model, var_expr, args..., ",
          "kwargs...)`.")
end
macro derivative_variable(model, args...)
    error("`@derivative_variable` has been deprecated in favor of `@variable`. ",
          "\n\nOld Syntax: `@derivative_variable(model, d(ivref)/d(pref), ",
          "var_expr, kwargs...)`\nNew Syntax: `@variable(model, var_expr, ",
          "Deriv(ivref, pref), kwargs...)`.")
end
macro BDconstraint(model, args...)
    error("`@BDconstraint` has been deprecated in favor of `@constraint`. ",
          "\n\nOld Syntax: `@BDconstraint(model, ref_expr(conditions...), ",
          "constr_expr, kwargs...)`\nNew Syntax: `@constraint(model, ref_expr, ",
          "constr_expr, DomainRestrictions(conditions...), kwargs...)`.")
end
macro set_parameter_bounds(args...)
    error("`@set_parameter_bounds` has been deprecated in favor of ",
          "`set_domain_restrictions`. Macro based modification is no longer ",
          "supported.")
end
macro add_parameter_bounds(args...)
    error("`@add_parameter_bounds` has been deprecated in favor of ",
          "`add_domain_restrictions`. Macro based modification is no longer ",
          "supported.")
end

# Define additional stuff that should not be exported
const _EXCLUDE_SYMBOLS = [Symbol(@__MODULE__), :eval, :include]

# Following JuMP, export everything that doesn't start with a _
for sym in names(@__MODULE__, all = true)
    sym_string = string(sym)
    if sym in _EXCLUDE_SYMBOLS || startswith(sym_string, "_") || startswith(sym_string, "@_")
        continue
    end
    if !(Base.isidentifier(sym) || (startswith(sym_string, "@") && Base.isidentifier(sym_string[2:end])))
        continue
    end
    @eval export $sym
end

end # module


module XORShift

import Base: rand, rand!

type XORShiftStar1024 <: AbstractRNG
    x00::UInt64
    x01::UInt64
    x02::UInt64
    x03::UInt64
    x04::UInt64
    x05::UInt64
    x06::UInt64
    x07::UInt64
    x08::UInt64
    x09::UInt64
    x10::UInt64
    x11::UInt64
    x12::UInt64
    x13::UInt64
    x14::UInt64
    x15::UInt64
    p::UInt64
end
XORShiftStar() = XORShiftStar1024(rand(UInt64, 16)..., 0)

const globalXORShiftStar1024 = XORShiftStar()

fromUInt64(::Type{UInt64}, u::UInt64) = u
fromUInt64(::Type{Float64}, u::UInt64) = 5.421010862427522e-20 * u

function rand(rng::XORShiftStar1024, ::Type{UInt64})
    @inbounds begin
        p  = rng.p
        s0 = getfield(rng, reinterpret(Int, p)+1)
        p = (p + 1) % 16
        s1 = getfield(rng, reinterpret(Int, p)+1)
        s1 $= s1 << 31
        s1 $= s1 >> 11
        s0 $= s0 >> 30
        s0 $= s1
        unsafe_store!(convert(Ptr{UInt64}, pointer_from_objref(rng))+8,
                      s0, reinterpret(Int, p)+1)
        rng.p = p
    end
    return s0 * 1181783497276652981
end

rand{T<:Number}(rng::XORShiftStar1024, ::Type{T}) = fromUInt64(T, rand(rng, UInt64))

@eval function rand!{T}(rng::XORShiftStar1024, a::AbstractArray{T})
    p = rng.p
    n = length(a)
    q = min((16 - p) % 16, n)
    r = (n - q) % 16
    @inbounds for i = 1:q
        a[i] = rand(rng, T)
    end
    $(Expr(:block, [
        let x = symbol(@sprintf("x%02d", k % 16))
            :($x = rng.$x)
        end
    for k = 0:15 ]...))
    r = q + (n-q) >> 4 << 4
    @inbounds for i = q+1:16:r
        $(Expr(:block, [
            let x0 = symbol(@sprintf("x%02d", k % 16)),
                x1 = symbol(@sprintf("x%02d", (k + 1) % 16))
                quote
                    s0 = $x0
                    s1 = $x1
                    s1 $= s1 << 31
                    s1 $= s1 >> 11
                    s0 $= s0 >> 30
                    s0 $= s1
                    $x1 = s0
                    a[i + $k] = fromUInt64(T, s0 * 1181783497276652981)
                end
            end
        for k = 0:15 ]...))
    end
    $(Expr(:block, [
        let x = symbol(@sprintf("x%02d", k % 16))
            :(rng.$x = $x)
        end
    for k = 0:15 ]...))
    rng.p = 0
    @inbounds for i = r+1:n
        a[i] = rand(rng, T)
    end
    return a
end

end # module


using BenchmarkTools

const BENCHMARKS_DIR = "../benchmark/"
cp(joinpath(BENCHMARKS_DIR, "benchmark_utils.jl"), joinpath(@__DIR__, "benchmark_utils.jl"), force=true)

benchmark_filepath(benchmark_name, benchmarks_dir=BENCHMARKS_DIR) =
    joinpath(benchmarks_dir, "benchmark_" * benchmark_name * ".jl")

@testset "Performance benchmarks" begin
    @info "Testing performance benchmarks..."

    @testset "Performance benchmark scripts" begin
        @info "  Testing performance benchmark scripts..."

        @testset "Static ocean benchmark" begin
            @info "    Running static ocean benchmarks..."

            replace_strings = [
                ("Ns = [(16, 16, 16), (32, 32, 32), (64, 64, 64), (128, 128, 128), (256, 256, 256)]",
                 "Ns = [(1, 1, 1)]")
            ]

            @test run_script(replace_strings, "static_ocean", benchmark_filepath("static_ocean"))
        end

        @testset "Channel benchmark" begin
            @info "    Running channel benchmarks..."

            replace_strings = [
                ("Ns = [(32, 32, 32), (64, 64, 64), (128, 128, 128), (256, 256, 256)]",
                 "Ns = [(1, 1, 1)]")
            ]

            @test run_script(replace_strings, "channel", benchmark_filepath("channel"))
        end

        @testset "Turbulence closures benchmark" begin
            @info "    Running turbulence closures benchmark..."

            replace_strings = [
                ("Ns = [(32, 32, 32), (256, 256, 128)]",
                 "Ns = [(1, 1, 1)]")
            ]

            @test run_script(replace_strings, "turbulence_closures", benchmark_filepath("turbulence_closures"))
        end

        @testset "Tracers benchmark" begin
            @info "    Running tracers benchmark..."

            replace_strings = [
                ("(32, 32, 32)", "(1, 1, 1)"),
                ("(256, 256, 128)", "(1, 1, 1)"),
                ("test_cases = [(0, 0), (0, 1), (0, 2), (1, 0), (2, 0), (2, 3), (2, 5), (2, 10)]",
                 "test_cases = [(0, 0), (2, 0), (2, 3)]")
            ]

            @test run_script(replace_strings, "tracers", benchmark_filepath("tracers"))
        end
    end

    @testset "Selected performance benchmarks" begin
        @info "  Running selected performance benchmarks..."

        for arch in archs
            sizes = [(16, 16, 16), (32, 32, 32)]
            for sz in sizes
                grid = RegularCartesianGrid(size=sz, extent=(1, 1, 1))
                model = IncompressibleModel(architecture=arch, grid=grid)

                @info "Benchmarking [$arch, Float64, $sz]..."
                b = @benchmark time_step!($model, 1)
                display(b)
                println()

                @test model isa IncompressibleModel
            end
        end
    end
end

rm(joinpath(@__DIR__, "benchmark_utils.jl"))


function Sim_Model(Cell,Input,Def,Tk,SList,SOC,A0,B0,C0,D0,t)
    """
    Function to simulate generated reduced-order models.

    Sim_Model(Cell,Input,Def,Tk,SOC,A,B,C,D)

    """
    #Determine time span and allocate arrays
    tlength = size(Input,1)

    A = Array{Float64}(undef,size(A0[1]))
    B = Array{Float64}(undef,size(B0[1]))
    C = Array{Float64}(undef,size(C0[1]))
    D = Array{Float64}(undef,size(D0[1]))

    #Selecting SS Models
    for i in 1:length(SList)
        if SOC == SList[i] # Fix this, round?
            A = A0[i]
            B = B0[i]
            C = C0[i]
            D = D0[i]
        end
    end

    #Capturing Indices
    tfstr = Array{String}(undef,0,1)
    for i in 1:length(Cell.Transfer.tfs)
        t1 = t2 = Array{String}(undef,0,1)
        for j in 1:length(Cell.Transfer.Locs[i])
            t1 ="$(Cell.Transfer.tfs[i])_$(Cell.Transfer.Elec[i])"
            t2 = [t2; t1]
        end
        tfstr = [tfstr; t2]
    end

    CeInd = findall(isequal("C_e_Na"), tfstr)
    ϕ_ẽInd = findall(isequal("Phi_e_Na"), tfstr)
    CsePosInd = findall(isequal("C_se_Pos"), tfstr)
    CseNegInd = findall(isequal("C_se_Neg"), tfstr)
    ϕ_sNegInd = findall(isequal("Phi_s_Neg"), tfstr)
    ϕ_sPosInd = findall(isequal("Phi_s_Pos"), tfstr)
    ϕ_sePosInd = findall(isequal("Phi_se_Pos"), tfstr)
    ϕ_seNegInd = findall(isequal("Phi_se_Neg"), tfstr)
    FluxNegInd = findall(isequal("Flux_Neg"), tfstr)
    FluxPosInd = findall(isequal("Flux_Pos"), tfstr)

    #@infiltrate cond = true

    CeNeg = Cell.Transfer.Locs[1].<=Cell.Neg.L
    CeNegOffset = Cell.Transfer.Locs[1].<Cell.Neg.L
    CeSep = Cell.Transfer.Locs[1].<=Cell.Neg.L+Cell.Sep.L
    CeSepOffset = Cell.Transfer.Locs[1].<Cell.Neg.L+Cell.Sep.L
    CePos = Cell.Transfer.Locs[1].<=Cell.Const.Ltot
    CeNegInd = findall(CeNeg .== 1)
    CeSepInd = findall(CeSep.-CeNegOffset .== 1)
    CePosInd = findall(CePos.-CeSepOffset .== 1)

    csegain_neg = C[CseNegInd[1][1],1] #First Column in C Array (zeros column)
    csegain_pos = C[CsePosInd[1][1],1] #First Column in C Array (zeros column)

    #Memory Allocation
    θ_neg = Array{Float64}(undef,tlength,1) .= 0.
    θ_pos = Array{Float64}(undef,tlength,1) .= 0.
    jeq_neg = Array{Float64}(undef,tlength,1) .= 0.
    jeq_pos = Array{Float64}(undef,tlength,1) .= 0.
    x = Array{Float64}(undef,tlength+1,size(A,1)) .= 0.
    y = Array{Float64}(undef,tlength,size(C,1)) .= 0.
    Cse_Neg = Array{Float64}(undef,tlength,size(CseNegInd,1)) .= 0.
    Cse_Pos = Array{Float64}(undef,tlength,size(CsePosInd,1)) .= 0.
    Ce = Array{Float64}(undef,tlength,size(CeInd,1)) .= Cell.Const.ce0
    η0 = Array{Float64}(undef,tlength,1) .= 0.
    η_neg = Array{Float64}(undef,tlength,size(FluxNegInd,1)) .= 0.
    ηL = Array{Float64}(undef,tlength,1) .= 0.
    η_pos = Array{Float64}(undef,tlength,size(FluxPosInd,1)) .= 0.
    ϕ_ẽ1 = Array{Float64}(undef,tlength,size(ϕ_ẽInd,1)) .= 0.
    ϕ_ẽ2 = Array{Float64}(undef,tlength,size(CeInd,1)) .= 0.
    ϕ_se_neg_0 = Array{Float64}(undef,tlength,1) .= 0. #Replace with length of ϕ_seNegInd @ zero
    jNeg = Array{Float64}(undef,tlength,size(FluxNegInd,1)) .= 0.
    j0 = Array{Float64}(undef,tlength,1) .= 0.
    jPos = Array{Float64}(undef,tlength,size(FluxPosInd,1)) .= 0.
    jL = Array{Float64}(undef,tlength,1) .= 0.
    j0_CC_neg = Array{Float64}(undef,tlength,1) .= 0.
    Rtot_neg = Array{Float64}(undef,tlength,1) .= 0.
    Rtot_pos = Array{Float64}(undef,tlength,1) .= 0.
    j0_CC_pos = Array{Float64}(undef,tlength,1) .= 0.
    Uocp_Neg = Array{Float64}(undef,tlength,1) .= 0.
    Uocp_Pos = Array{Float64}(undef,tlength,1) .= 0.
    Cell_V = Array{Float64}(undef,tlength,1) .= 0.
    ϕ_e = Array{Float64}(undef,tlength,size(CeInd,1)) .= 0.
    Cell_SOC = Array{Float64}(undef,tlength,1) .= 0
    Iapp = Array{Float64}(undef,tlength+1,1) .= 0

    #Defining SOC
    SOC_Neg = SOC * (Cell.Neg.θ_100-Cell.Neg.θ_0) + Cell.Neg.θ_0
    SOC_Pos = SOC * (Cell.Pos.θ_100-Cell.Pos.θ_0) + Cell.Pos.θ_0
    θ_neg[1] = SOC_Neg
    θ_pos[1] = SOC_Pos
    Cell_SOC[1] = (SOC_Neg-Cell.Neg.θ_0)/(Cell.Neg.θ_100-Cell.Neg.θ_0)

    #Loop through time
    #Compute dependent variables (voltage, flux, etc.)
    for i in 0:(tlength-1)
        cs_neg_avg = x[i+1,1] * csegain_neg + SOC_Neg * Cell.Neg.cs_max < 0. ? 0. : x[i+1,1] * csegain_neg + SOC_Neg * Cell.Neg.cs_max #Zero if < 0
        cs_pos_avg = x[i+1,1] * csegain_pos + SOC_Pos * Cell.Pos.cs_max < 0. ? 0. : x[i+1,1] * csegain_pos + SOC_Pos * Cell.Pos.cs_max #Zero if < 0

        if cs_neg_avg > Cell.Neg.cs_max
            cs_neg_avg = Cell.Neg.cs_max
        end

        if cs_pos_avg > Cell.Pos.cs_max
            cs_pos_avg = Cell.Pos.cs_max
        end

        #Reaction Rates
        if Cell.Const.CellTyp == "Doyle_94"
            k_neg = Cell.Neg.k_norm#/Cell.Neg.cs_max/Cell.Const.ce0^(1-Cell.Neg.α)
            k_pos = Cell.Pos.k_norm#/Cell.Pos.cs_max/Cell.Const.ce0^(1-Cell.Pos.α)
            jeq_neg[i+1] = k_neg*sqrt(cs_neg_avg*Cell.Const.ce0*(Cell.Neg.cs_max-cs_neg_avg))
            jeq_pos[i+1] = k_pos*sqrt(cs_pos_avg*Cell.Const.ce0*(Cell.Pos.cs_max-cs_pos_avg))
        else
            #jeq_neg = Cell.Neg.k_norm*Cell.Neg.cs_max*(Cell.Const.ce0*(cs_neg_avg/Cell.Neg.cs_max*(1-cs_neg_avg/Cell.Neg.cs_max)))^(1-Cell.Neg.α)
            #jeq_pos = Cell.Pos.k_norm*Cell.Pos.cs_max*(Cell.Const.ce0*(cs_pos_avg/Cell.Pos.cs_max*(1-cs_pos_avg/Cell.Pos.cs_max)))^(1-Cell.Pos.α)
            k_neg = Cell.Neg.k_norm
            k_pos = Cell.Pos.k_norm
            jeq_neg[i+1] = Cell.Neg.k_norm*(Cell.Const.ce0*cs_neg_avg*(Cell.Neg.cs_max-cs_neg_avg))^(1-Cell.Neg.α)
            jeq_pos[i+1] = Cell.Pos.k_norm*(Cell.Const.ce0*cs_pos_avg*(Cell.Pos.cs_max-cs_pos_avg))^(1-Cell.Pos.α)
        end

        θ_neg[i+1] = cs_neg_avg/Cell.Neg.cs_max
        θ_pos[i+1] = cs_pos_avg/Cell.Pos.cs_max
        Cell_SOC[i+1] = (θ_neg[i+1]-Cell.Neg.θ_0)/(Cell.Neg.θ_100-Cell.Neg.θ_0)

        javg_neg = Iapp[i+1]/(Cell.Neg.as*F*Cell.Neg.L*Cell.Const.CC_A)
        javg_pos = Iapp[i+1]/(Cell.Pos.as*F*Cell.Pos.L*Cell.Const.CC_A)

        Arr_Factor = ((1/Cell.Const.T_ref)-(1/Tk[i+1]))/R
        κneg = Cell.Const.κf(mean(Ce[i+1,CeNegInd]))*exp(Cell.Const.Ea_κ*Arr_Factor)
        κpos = Cell.Const.κf(mean(Ce[i+1,CePosInd]))*exp(Cell.Const.Ea_κ*Arr_Factor)
        κsep = Cell.Const.κf(mean(Ce[i+1,CeSepInd]))*exp(Cell.Const.Ea_κ*Arr_Factor)
        σ_neg = Cell.Neg.σ*exp(Cell.Const.Ea_κ*Arr_Factor)
        σ_pos = Cell.Pos.σ*exp(Cell.Const.Ea_κ*Arr_Factor)
        κ_eff_Neg = κneg*(Cell.Neg.ϵ_e^(Cell.Neg.κ_brug))
        κ_eff_Sep = κsep*(Cell.Sep.ϵ_e^(Cell.Sep.κ_brug))
        κ_eff_Pos = κpos*(Cell.Pos.ϵ_e^(Cell.Pos.κ_brug))
        σ_eff_Neg = σ_neg*Cell.Neg.ϵ_s^Cell.Neg.σ_brug #Effective Conductivity Neg
        σ_eff_Pos = σ_pos*Cell.Pos.ϵ_s^Cell.Pos.σ_brug #Effective Conductivity Pos

        #Resistances
        #Rtot_neg[i] = (Tk[i]*R)/(F^2*mean([abs(jeq_neg[i]) abs(javg_neg)]))+Cell.Neg.RFilm
        Rtot_neg[i+1] = (Tk[i+1]*R)/(F^2*sqrt(jeq_neg[i+1]^2+javg_neg^2/4))+Cell.Neg.RFilm
        #Rtot_pos[i] = (Tk[i]*R)/(F^2*mean([abs(jeq_pos[i]) abs(javg_pos)]))+Cell.Pos.RFilm
        Rtot_pos[i+1] = (Tk[i+1]*R)/(F^2*sqrt(jeq_pos[i+1]^2+javg_pos^2/4))+Cell.Pos.RFilm

        #Condensing Variable
        ν_neg = @. Cell.Neg.L*sqrt((Cell.Neg.as*(1/κ_eff_Neg+1/σ_eff_Neg))/Rtot_neg[i+1])
        ν_pos = @. Cell.Pos.L*sqrt((Cell.Pos.as*(1/κ_eff_Pos+1/σ_eff_Pos))/Rtot_pos[i+1])

        #Relinearise dependent on ν, σ, κ
        D = D_Linear(Cell, ν_neg, ν_pos, σ_eff_Neg, κ_eff_Neg, σ_eff_Pos, κ_eff_Pos, κ_eff_Sep)

        #Interpolate C & D Matrices
        C = interp(C0,SList,Cell_SOC[i+1])
        #D = interp(D0,SList,Cell_SOC[i+1])
        #SS Output
        y[i+1,:] = C*x[i+1,:] + D*Iapp[i+1]

        #Concentrations & Force electrode concentration maximum
        Cse_Neg[i+1,:] = (SOC_Neg.*Cell.Neg.cs_max .+ y[i+1,CseNegInd]) > ones(size(Cse_Neg,2))*Cell.Neg.cs_max ? ones(size(Cse_Neg,2))*Cell.Neg.cs_max : (SOC_Neg.*Cell.Neg.cs_max .+ y[i+1,CseNegInd])
        Cse_Pos[i+1,:] = (SOC_Pos.*Cell.Pos.cs_max .+ y[i+1,CsePosInd]) > ones(size(Cse_Pos,2))*Cell.Pos.cs_max ? ones(size(Cse_Pos,2))*Cell.Pos.cs_max : (SOC_Pos.*Cell.Pos.cs_max .+ y[i+1,CsePosInd])
        Ce[i+1,:] = @. Cell.Const.ce0 + y[i+1,CeInd]

        #Potentials
        Uocp_Neg[i+1] = Cell.Const.Uocp("Neg",Cse_Neg[i+1,1]/Cell.Neg.cs_max)
        Uocp_Pos[i+1] = Cell.Const.Uocp("Pos",Cse_Pos[i+1,1]/Cell.Pos.cs_max)
        ϕ_se_neg_0[i+1] = y[i+1, ϕ_seNegInd[1]] + Uocp_Neg[i+1] #Location 0
        ϕ_ẽ1[i+1,:] = y[i+1,ϕ_ẽInd]
        ϕ_ẽ2[i+1,:] = @. ((Tk[i+1]*2*R*(1-Cell.Const.tpf(Ce[i+1,:])))/F)*(log(Ce[i+1,:]/Ce[i+1,1]))
        ϕ_e[i+1,:] = @. [0; ϕ_ẽ1[i+1,:]]+ϕ_ẽ2[i+1,:]-ϕ_se_neg_0[i+1]

        #Flux
        jNeg[i+1,:] = y[i+1,FluxNegInd]
        j0[i+1] = y[i+1,FluxNegInd[1]]
        jPos[i+1,:] = y[i+1,FluxPosInd]
        jL[i+1] = y[i+1,FluxPosInd[1]]

        #j0_CC_neg[i+1] = findmax([eps(); (k_neg*(Cell.Neg.cs_max-Cse_Neg[i+1,1])^(1-Cell.Neg.α))*((Cse_Neg[i+1,1]^Cell.Neg.α)*(Ce[i+1,1]^(1-Cell.Neg.α)))])[1]
        j0_neg = findmax([ones(size(Cse_Neg,2))*eps() (k_neg.*(Cse_Neg[i+1,:].^Cell.Neg.α).*(Ce[i+1,1].^(1-Cell.Neg.α))).*(Cell.Neg.cs_max.-Cse_Neg[i+1,:]).^(1-Cell.Neg.α)], dims=2)[1]

        η0[i+1] = Tk[i+1]*2*R/F*asinh(j0[i+1]/(2*j0_neg[1]))
        η_neg[i+1,:] = @. (Tk[i+1]*2*R)/F*asinh((jNeg[i+1,:])/(2*j0_neg))

        #j0_CC_pos[i+1] = findmax([eps(); (k_pos*(Cell.Pos.cs_max-Cse_Pos[i+1,1])^(1-Cell.Pos.α))*((Cse_Pos[i+1,1]^Cell.Pos.α)*(Ce[i+1,end]^(1-Cell.Pos.α)))])[1]
        j0_pos = findmax([ones(size(Cse_Pos,2))*eps() (k_pos.*(Cse_Pos[i+1,:].^Cell.Pos.α).*(Ce[i+1,1].^(1-Cell.Pos.α))).*(Cell.Pos.cs_max.-Cse_Pos[i+1,:]).^(1-Cell.Pos.α)], dims=2)[1]
        ηL[i+1] = (Tk[i+1]*2*R)/F*asinh(jL[i+1]/(2*j0_pos[1]))
        η_pos[i+1,:] = @. (Tk[i+1]*2*R)/F*asinh(jPos[i+1,:]/(2*j0_pos))

        #Cell Voltage
        Cell_V[i+1] = @. (Uocp_Pos[i+1]-Uocp_Neg[i+1]) + (ηL[i+1]-η0[i+1]) + (ϕ_ẽ1[i+1,end]+ϕ_ẽ2[i+1,end]) + (Cell.Pos.RFilm*jL[i+1]-Cell.Neg.RFilm*j0[i+1])*F

        #ϕ_s
        ϕ_s_neg = y[i+1,ϕ_sNegInd]
        ϕ_s_pos = @. y[i+1,ϕ_sPosInd] + Cell_V[i+1]

        #Interpolate A Matrix
        A = interp(A0,SList,Cell_SOC[i+1])
        B = interp(B0,SList,Cell_SOC[i+1])

        #Update States
        x[i+2,:] = A*x[i+1,:] + B*Iapp[i+1]

        if Def =="Power"
            Iapp[i+2] = Input[i+1,1]/Cell_V[i+1]
        else
            Iapp[i+2] = Input[i+1,1]
        end

    end

return Cell_V, Ce, jNeg, jPos, Rtot_neg, Rtot_pos, η0, ηL, η_neg, η_pos, ϕ_ẽ1, ϕ_ẽ2, Uocp_Neg, Uocp_Pos, ϕ_e, Cse_Neg, Cse_Pos, Cell_SOC, jeq_neg, jeq_pos, j0, jL, t
end


function plotFunc_timecourse(Sim::Module)
    rc("figure",figsize = (12,8));
    rc("font",size = 18);
    rc("axes",linewidth = 1);
    rc("lines",linewidth = 2);

    subplots_adjust(wspace=0.5, hspace=0.5);

    subplot(2,2,1);
    plot(Sim.t,Sim.CBMa./maximum(Sim.CBMa).*100,"k");
    xlim(-0.5,9.5);
    xticks([0,3,6,9]);
    ylim(0,120);
    yticks([0,20,40,60,80,100]);
    xlabel("Time (min)");
    ylabel("CP668 activity \n (relative % max.)");

    subplot(2,2,2);
    plot(Sim.t,Sim.TAK1a./maximum(Sim.TAK1a).*100,"k");
    xlim(-0.5,9.5);
    xticks([0,3,6,9]);
    ylim(0,120);
    yticks([0,20,40,60,80,100]);
    xlabel("Time (min)");
    ylabel("TAK1 activity \n (relative % max.)");

    subplot(2,2,3);
    plot(Sim.t,Sim.IKKa./maximum(Sim.IKKa).*100,"k");
    xlim(-0.5,9.5);
    xticks([0,3,6,9]);
    ylim(0,120);
    yticks([0,20,40,60,80,100]);
    xlabel("Time (min)");
    ylabel("IKK activity \n (relative % max.)");

    subplot(2,2,4);
    plot(Sim.t,Sim.NFkBa./maximum(Sim.NFkBa).*100,"k");
    xlim(0,360);
    xticks([0,120,240,360]);
    ylim(0,120);
    yticks([0,20,40,60,80,100]);
    xlabel("Time (min)");
    ylabel("NFκB activity \n (relative % max.)");

    show();
end


function build_master(sub_problems::T where T <: Dict, tree::T where T <: AbstractTree, probabilities::Dict{AbstractTree,Float64}, solver, discount_factor::Float64,
  risk::Union{Risk,Array{Risk,1}}, sideconstraints)

	if typeof(solver) <: Tuple
		model = Model(solver[1])
	else
		model = Model(solver)
	end
	@objective(model,Min,0)

	model.ext[:columns] = Dict{AbstractTree,Array{Column,1}}()
	for node in collect(tree)
		model.ext[:columns][node]=Array{Column,1}()
	end
	leafs=Array{Leaf,1}()

	model.ext[:expansions] = Dict{AbstractTree,Dict{Symbol,Any}}()

	model.ext[:options]=Dict{Symbol,Tuple}()
	for (name,options) in sub_problems[tree].ext[:options]
		model.ext[:options][name]=options
	end

	form=nothing
	for (name,options) in model.ext[:options]
		if options[4]==:Bin
			if form==:continuous
				form=:mixed
				break
			else
				form=:binary
			end
		elseif options[4]==:Int
			form=:mixed
			break
		else
			if form==:binary
				form=:mixed
				break
			else
				form=:continuous
			end
		end
	end

	model.ext[:discrete_con]=Dict{AbstractTree,Dict{Int64,ConstraintRef}}()
	model.ext[:discrete_var]=Dict{AbstractTree,Dict{Int64,VariableRef}}()

	for (node,sp) in sub_problems
		if haskey(node.ext,:sum_max) || haskey(node.ext,:sum_min) || haskey(node.ext,:sum_set)
			sp.ext[:form]=:binarycolumns
		else
			sp.ext[:form]=form
		end
		if sp.ext[:form] ∉ [:binary,:binarycolumns]
			vars = all_variables(sp)
			if form==:continuous
				for var in vars
					if is_binary(var) || is_integer(var)
						sp.ext[:form]=:mixed
						break
					end
				end
			end

			if sp.ext[:form]==:mixed
				sp.ext[:discrete]=VariableRef[]
				sp.ext[:discrete_branch]=VariableRef[]
				ignore=VariableRef[]

				for (name,variable) in sp.ext[:expansions]
					if typeof(variable) <: AbstractArray
						for index in eachindex(variable)
							push!(ignore,variable[index])
						end
					else
						push!(ignore,variable)
					end
				end

				for var in vars
					if (is_binary(var) || is_integer(var)) && var ∉ ignore
						push!(sp.ext[:discrete],var)
						push!(sp.ext[:discrete_branch],var)
					end
				end
			end
		end
	end

	for (node,sp) in sub_problems
		if typeof(node)==Leaf
			push!(leafs,node)
		end
		model.ext[:expansions][node] = Dict{Symbol,Any}()

		for (name,variable) in sp.ext[:expansions]
			if sp.ext[:options][name][4]==:Bin
				model.ext[:expansions][node][name] = copy_variable!(model, variable, relaxbinary)
			elseif sp.ext[:options][name][4]==:Int
				model.ext[:expansions][node][name] = copy_variable!(model, variable, relaxinteger)
			else
				model.ext[:expansions][node][name] = copy_variable!(model, variable)
			end
			if model.ext[:options][name][5]!=nothing
				if typeof(variable) <: AbstractArray
					for i in eachindex(variable)
						set_lower_bound(model.ext[:expansions][node][name][i],model.ext[:options][name][5])
					end
				else
					set_lower_bound(model.ext[:expansions][node][name],model.ext[:options][name][5])
				end
			end
			if model.ext[:options][name][6]!=nothing
				if typeof(variable) <: AbstractArray
					for i in eachindex(variable)
						set_upper_bound(model.ext[:expansions][node][name][i],model.ext[:options][name][6])
					end
				else
					set_upper_bound(model.ext[:expansions][node][name],model.ext[:options][name][6])
				end
			end
		end
	end

	model.ext[:scenprofit_var] = Dict{Leaf,VariableRef}()
	model.ext[:scenprofit_con] = Dict{Leaf,ConstraintRef}()

	for leaf in leafs
		model.ext[:scenprofit_var][leaf] = @variable(model)
		model.ext[:scenprofit_con][leaf] = @constraint(model, 0 == model.ext[:scenprofit_var][leaf])

  		set_objective_coefficient(model, model.ext[:scenprofit_var][leaf], probabilities[leaf])
	end

	risk_objectives=AffExpr[]
	remain=1.0
	if typeof(risk)==Risk
		if (risk.α==1.0 || risk.λ==0.0) && risk.bound==nothing
			risk=[]
		else
			risk=[risk]
		end
	end
    for i in 1:length(risk)
		risk_objective=AffExpr(0.0)
		eta=@variable(model)
	    for leaf in leafs
			v = @variable(model)
			w = @variable(model)
			set_lower_bound(v,0.0)
			set_lower_bound(w,0.0)
			if risk[i].offset==nothing
				@constraint(model,v>=eta-model.ext[:scenprofit_var][leaf])
				@constraint(model,w>=model.ext[:scenprofit_var][leaf]-eta)
				add_to_expression!(risk_objective,(model.ext[:scenprofit_var][leaf])*probabilities[leaf])
			else
				@constraint(model,v>=eta-model.ext[:scenprofit_var][leaf]+risk[i].offset[leaf])
				@constraint(model,w>=model.ext[:scenprofit_var][leaf]-risk[i].offset[leaf]-eta)
				add_to_expression!(risk_objective,(model.ext[:scenprofit_var][leaf]-risk[i].offset[leaf])*probabilities[leaf])
			end
			add_to_expression!(risk_objective,probabilities[leaf]*(v+w/risk[i].α*(1-risk[i].α)))
		end
		remain-=risk[i].λ
		push!(risk_objectives,risk_objective)
	end

	for leaf in leafs
		nodes=history(leaf)
		for n in eachindex(nodes)
			node=nodes[n]
			sp=sub_problems[node]
			df=discount_factor^depth(node)
			for (name,variable) in sp.ext[:expansions]
				interval=max(1,n-sp.ext[:options][name][3]-sp.ext[:options][name][2]+1):n-sp.ext[:options][name][2]
				disc=Dict{Int64,Float64}()
				for i in interval
					disc[i]=df/discount_factor^(i-1)
				end
				if typeof(variable) <: AbstractArray
					for i in eachindex(variable)
						cost_coef=df*coef(sp.ext[:capitalcosts],variable[i])
						for j in interval
							cost_coef+=disc[j]*coef(sp.ext[:ongoingcosts],variable[i])
						end
						set_normalized_coefficient(model.ext[:scenprofit_con][leaf], model.ext[:expansions][node][name][i], cost_coef)
					end
				else
					cost_coef=df*coef(sp.ext[:capitalcosts],variable)
					for j in interval
						cost_coef+=disc[j]*coef(sp.ext[:ongoingcosts],variable)
					end
					set_normalized_coefficient(model.ext[:scenprofit_con][leaf], model.ext[:expansions][node][name], cost_coef)
				end
			end
		end
	end

	# create the cover constraints
	model.ext[:coverconstraint] = Dict{AbstractTree,Dict{Symbol,Any}}()
	for (node,sp) in sub_problems
		model.ext[:coverconstraint][node] = Dict{Symbol,Any}()
		past=history(node)
		for (name,variable) in sp.ext[:expansions]
			interval=sp.ext[:options][name][2]+1:min(sp.ext[:options][name][2]+sp.ext[:options][name][3],length(past))
			if typeof(variable) <: AbstractArray
				model.ext[:coverconstraint][node][name] = Dict()
				for i in eachindex(variable)
					if sp.ext[:options][name][1]==:shutdown
						model.ext[:coverconstraint][node][name][i] = @constraint(model, 0 >= sum(model.ext[:expansions][past[index]][name][i] for index in interval))
					elseif sp.ext[:options][name][1]==:expansion
						model.ext[:coverconstraint][node][name][i] = @constraint(model, 0 <= sum(model.ext[:expansions][past[index]][name][i] for index in interval))
					elseif sp.ext[:options][name][1]==:enforced
						model.ext[:coverconstraint][node][name][i] = @constraint(model, 0 == sum(model.ext[:expansions][past[index]][name][i] for index in interval))
					elseif sp.ext[:options][name][1]==:state
						if node.parent==nothing
							model.ext[:coverconstraint][node][name][i] = @constraint(model, 0 == -sp.ext[:options][name][7] + model.ext[:expansions][node][name][i])
						else
							model.ext[:coverconstraint][node][name][i] = @constraint(model, 0 == model.ext[:expansions][node][name][i] - model.ext[:expansions][node.parent][name][i])
						end
					end
				end
				# if typeof(node)==Leaf && !sp.ext[:options][name][4]
				# 	for i in eachindex(variable)
				# 		@constraint(model, sum(model.ext[:expansions][past][name][i] for past in history_function(node))<=1)
				# 	end
				# end
			else
				if sp.ext[:options][name][1]==:shutdown
					model.ext[:coverconstraint][node][name] = @constraint(model, 0 >= sum(model.ext[:expansions][past[index]][name] for index in interval))
				elseif sp.ext[:options][name][1]==:expansion
					model.ext[:coverconstraint][node][name] = @constraint(model, 0 <= sum(model.ext[:expansions][past[index]][name] for index in interval))
				elseif sp.ext[:options][name][1]==:enforced
					model.ext[:coverconstraint][node][name] = @constraint(model, 0 == sum(model.ext[:expansions][past[index]][name] for index in interval))
				elseif sp.ext[:options][name][1]==:state
					if node.parent==nothing
						model.ext[:coverconstraint][node][name] = @constraint(model, 0 == -sp.ext[:options][name][7] + model.ext[:expansions][node][name])
					else
						model.ext[:coverconstraint][node][name] = @constraint(model, 0 == model.ext[:expansions][node][name] - model.ext[:expansions][node.parent][name])
					end
				end
				# if typeof(node)==Leaf && !sp.ext[:options][name][4]
				# 	@constraint(model, sum(model.ext[:expansions][past][name] for past in history_function(node))<=1)
				# end
			end
		end
	end

	model.ext[:convexcombination] = Dict{AbstractTree,Array{ConstraintRef,1}}()
	model.ext[:mip]=false
	model.ext[:branch_cons]=ConstraintRef[]

	for node in keys(sub_problems)
		model.ext[:convexcombination][node]=ConstraintRef[]
		if haskey(node.ext,:sum_max)
			push!(model.ext[:convexcombination][node],@constraint(model, 0 <= node.ext[:sum_max]))
		end
		if haskey(node.ext,:sum_min)
			push!(model.ext[:convexcombination][node],@constraint(model, 0 >= node.ext[:sum_min]))
		end
		if haskey(node.ext,:sum_set)
			if length(model.ext[:convexcombination][node])>0
				error(":eq and :le/:ge specifed for node "*node.name)
			end
			push!(model.ext[:convexcombination][node],@constraint(model, 0 == node.ext[:sum_set]))
		end
		if length(model.ext[:convexcombination][node])==0
			push!(model.ext[:convexcombination][node],@constraint(model, 0 == 1))
		end
	end

	if typeof(sideconstraints) <: Function
		map(Main.eval,unpack_expansions(model.ext[:expansions])) #bring expansion variables into global scope
		sideconstraints(model,tree)
		map(Main.eval,clear_expansions(model.ext[:expansions]))
	end

	objective_fn=objective_function(model)*remain

	for i in 1:length(risk_objectives)
		objective_fn+=risk_objectives[i]*risk[i].λ
		if risk[i].bound!=nothing
			surplus=@variable(model)
			set_lower_bound(surplus,0)
			@constraint(model,risk_objectives[i]<=risk[i].bound+surplus)
			objective_fn+=risk[i].penalty*surplus
		end
	end

	set_objective_function(model,objective_fn)

	model
end

function solve_master_binary(judge::JuDGEModel,initial_time::Float64,termination::Termination,warm_starts::Bool,mp_callback,verbose)
	set=0
	current=nothing
	if (termination.allow_frac==:binary_solve || termination.allow_frac==:binary_solve_return_relaxation) #&& judge.log[end].int>0
		if verbose==2
			print("Solving master problem as MIP")
		end
		set=1
		if solve_binary(judge,termination.abstol,termination.reltol,warm_starts,mp_callback)
			current = ConvergenceState(judge.log[end].obj, judge.bounds.UB, judge.bounds.LB, time()-initial_time, judge.log[end].iter + 1, 0)
			if verbose>0
				if verbose==2
					overprint("")
				end
				display(current,relaxation=false)
			end
			push!(judge.log,current)
		else
			if verbose==2
				overprint("")
			end
		end
	end
	if termination.allow_frac==:binary_solve_return_relaxation && set==1
		remove_binary(judge)
		optimize!(judge.master_problem)
	end
	current
end

function solve_binary(judge::JuDGEModel,abstol::Float64,reltol::Float64,warm_starts::Bool,mp_callback)
	if typeof(judge.master_solver) <: Tuple
		JuMP.set_optimizer(judge.master_problem,judge.master_solver[2])
	end

	judge.master_problem.ext[:mip]=true

	for node in keys(judge.master_problem.ext[:expansions])
		if judge.sub_problems[node].ext[:form]==:binarycolumns
			for col in judge.master_problem.ext[:columns][node]
			 	if !haskey(col.node.ext,:col_max)
			 		set_binary(col.var)
				else
			 		set_integer(col.var)
			 	end
			 end
		 end

		for x in keys(judge.master_problem.ext[:expansions][node])
			var = judge.master_problem.ext[:expansions][node][x]
			if judge.sub_problems[judge.tree].ext[:options][x][4]==:Bin
				if typeof(var) <: JuMP.Containers.SparseAxisArray
					for key in keys(var.data)
						set_binary(var.data[key])
					end
				elseif typeof(var) <: AbstractArray
					for key in keys(var)
						set_binary(var[key])
					end
				else
					set_binary(var)
				end
			elseif judge.sub_problems[judge.tree].ext[:options][x][4]==:Int
				if typeof(var) <: JuMP.Containers.SparseAxisArray
					for key in keys(var.data)
						set_integer(var.data[key])
					end
				elseif typeof(var) <: AbstractArray
					for key in keys(var)
						set_integer(var[key])
					end
				else
					set_integer(var)
				end
			end
		end
		if judge.sub_problems[node].ext[:form]==:mixed
			for (i,var) in judge.master_problem.ext[:discrete_var][node]
				set_integer(var)
			end
		end
	end

	if typeof(mp_callback) <: Function
		mp_callback(judge,abstol,reltol)
		printright("Initialising")
	end
	optimize!(judge.master_problem)
	status=termination_status(judge.master_problem)
	if status!=MathOptInterface.INFEASIBLE_OR_UNBOUNDED && status!=MathOptInterface.INFEASIBLE && status!=MathOptInterface.DUAL_INFEASIBLE
		obj = objective_value(judge.master_problem)
		if obj<judge.bounds.UB
			judge.bounds.UB=obj
			if warm_starts
				vars=all_variables(judge.master_problem)
				set_start_value.(vars, JuMP.value.(vars))
			end
		end
		return true
	else
		return false
	end
end

function remove_binary(judge::JuDGEModel)
	# if judge.master_problem.ext[:binarycolumns]
	# 	for col in judge.master_problem.ext[:columns]
	# 		if !haskey(col.node.ext,:col_max)
	# 			unset_binary(col.var)
	# 		else
	# 			unset_integer(col.var)
	# 		end
	# 	end
	# end
	for node in keys(judge.master_problem.ext[:expansions])
		if judge.sub_problems[node].ext[:form]==:binarycolumns
			for col in judge.master_problem.ext[:columns][node]
			 	if !haskey(col.node.ext,:col_max)
			 		unset_binary(col.var)
				else
			 		unset_integer(col.var)
			 	end
			 end
		 end
		for x in keys(judge.master_problem.ext[:expansions][node])
			var = judge.master_problem.ext[:expansions][node][x]
			if judge.sub_problems[judge.tree].ext[:options][x][4]==:Bin
				if typeof(var) <: AbstractArray
					for key in eachindex(var)
						unset_binary(var[key])
					end
				else
					unset_binary(var)
				end
			elseif judge.sub_problems[judge.tree].ext[:options][x][4]==:Int
				if typeof(var) <: AbstractArray
					for key in eachindex(var)
						unset_integer(var[key])
					end
				else
					unset_integer(var)
				end
			end
		end
		if judge.sub_problems[node].ext[:form]==:mixed
			for (i,var) in judge.master_problem.ext[:discrete_var][node]
				unset_integer(var)
			end
		end
	end

	judge.master_problem.ext[:mip]=false

	if typeof(judge.master_solver) <: Tuple
		JuMP.set_optimizer(judge.master_problem,judge.master_solver[1])
	end
end


# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

module ParaxialAnalysis

using StaticArrays: promote_tuple_eltype
using ..OpticSim
using ..OpticSim.Geometry
using ..OpticSim.Repeat
using StaticArrays
using LinearAlgebra
import LazySets
using Unitful
using Unitful.DefaultSymbols

include("LensletAssembly.jl")
include("HeadEyeModel/HeadEyeModel.jl")

end #module
export ParaxialAnalysis