# Utility #
# ========================== #
function evaluate_objective(objective::JuMP.GenericAffExpr, x::AbstractVector)
    val = objective.constant
    for (var, coeff) in objective.terms
        val += coeff*x[index(var).value]
    end
    return val
end

function invalidate_cache!(stochasticprogram::StochasticProgram)
    cache = problemcache(stochasticprogram)
    delete!(cache, :evp)
    delete!(cache, :dep)
    delete!(cache, :stage_1)
    cache = solutioncache(stochasticprogram)
    delete!(cache, :solution)
    return nothing
end

function remove_scenarios!(stochasticprogram::StochasticProgram, s::Integer = 2)
    remove_scenarios!(scenarioproblems(stochasticprogram, s))
    return nothing
end

function remove_decision_variables!(stochasticprogram::StochasticProgram, s::Integer)
    remove_decision_variables!(scenarioproblems(stochasticprogram, s))
    return nothing
end

function remove_subproblems!(stochasticprogram::StochasticProgram, s::Integer = 2)
    remove_subproblems!(scenarioproblems(stochasticprogram, s))
    return nothing
end

function transfer_model!(dest::StochasticProgram, src::StochasticProgram)
    empty!(dest.generator)
    merge!(dest.generator, src.generator)
    return dest
end

function copy_decision_objective!(src::JuMP.Model, dest::JuMP.Model, vars::Vector{DecisionRef})
    src_obj = objective_function(src)
    src_obj_sense = objective_sense(src)
    dest_obj_sense = objective_sense(dest)
    for var in vars
        src_var = decision_by_name(src, name(var))
        src_var === nothing && error("Cannot copy objective function. Variable $var not in src model.")
        coeff = JuMP._affine_coefficient(src_obj, src_var)
        if dest_obj_sense == src_obj_sense
            set_objective_coefficient(dest, var, coeff)
        else
            set_objective_coefficient(dest, var, -coeff)
        end
    end
    dest_obj = objective_function(dest)
    if dest_obj_sense == src_obj_sense
        set_objective_function(dest, dest_obj + constant(src_obj))
    else
        set_objective_function(dest, -dest_obj + constant(src_obj))
        set_objective_sense(dest, src_obj_sense)
    end
    return nothing
end

function supports_zero(types::Vector, provided_def::Bool)
    for vartype in types
        if !hasmethod(zero, (Type{vartype}, ))
            !provided_def && @warn "Zero not defined for $vartype. Cannot generate zero function."
            return false
        end
    end
    return true
end

function supports_expected(types::Vector, provided_def::Bool)
    for vartype in types
        if !hasmethod(+, (vartype, vartype))
            !provided_def && @warn "Addition not defined for $vartype. Cannot generate expectation function."
            return false
        end
        if !hasmethod(*, (Float64, vartype)) || Base.code_typed(*, (Float64, vartype))[1].second != vartype
            !provided_def && @warn "Scalar multiplication with Float64 not defined for $vartype. Cannot generate expectation function."
            return false
        end
    end
    return true
end

problemcache(stochasticprogram::StochasticProgram) = stochasticprogram.problemcache
solutioncache(stochasticprogram::StochasticProgram) = stochasticprogram.solutioncache
function get_problem(stochasticprogram::StochasticProgram, key::Symbol)
    haskey(stochasticprogram.problemcache, key)|| error("No $key in problem cache")
    return stochasticprogram.problemcache[key]
end

typename(dtype::UnionAll) = dtype.body.name.name
typename(dtype::DataType) = dtype.name.name

function add_subscript(src::AbstractString, subscript::Integer)
    return @sprintf("%s%s", src, unicode_subscript(subscript))
end
add_subscript(src::Symbol, subscript::Integer) = add_subscript(String(src), subscript)

function unicode_subscript(subscript::Integer)
    if subscript < 0
        error("$subscript is negative")
    end
    return join('₀'+d for d in reverse(digits(subscript)))
end

function extract_set(expr)
    set = NoSpecifiedConstraint()
    found = false
    quit = false
    new_expr = prewalk(expr) do x
        if @capture(x, var_Symbol in constrset_) && !found && !quit
            set = constrset
            found = true
            return :($var)
        elseif @capture(x, var_Symbol[ids__] in constrset_) && !found && !quit
            set = constrset
            found = true
            return :($(x.args[2]))
        elseif @capture(x, [ids__] in constrset_) && !found && !quit
            set = constrset
            found = true
            return :($(x.args[2]))
        elseif @capture(x, set = constrset_) && !found && !quit
            set = constrset
            found = true
            return :()
        elseif @capture(x, var_Symbol[ids__])
            # Break here to prevent indices from being filtered
            quit = true
            return x
        else
            return x
        end
    end
    return set, new_expr, found
end

function decision_variables_at_stage(stochasticprogram::StochasticProgram{N}, s::Integer) where N
    1 <= s <= N || error("Stage $s not in range 1 to $N.")
    return all_decisions(get_decisions(proxy(stochasticprogram, s), s))
end

function decision_constraints_at_stage(stochasticprogram::StochasticProgram{N}, s::Integer) where N
    1 <= s <= N || error("Stage $s not in range 1 to $N.")
    return mapreduce(vcat, MOI.get(proxy(stochasticprogram, s), MOI.ListOfConstraints())) do (F, S)
        if F <: SingleDecision || F <: AffineDecisionFunction || F <: QuadraticDecisionFunction
            return MOI.get(proxy(stochasticprogram, s), MOI.ListOfConstraintIndices{F,S}())
        end
        return CI[]
    end
end

function attach_mocks!(structure::DeterministicEquivalent)
    MOIU.attach_optimizer(structure.model)
    return nothing
end

function attach_mocks!(structure::VerticalStructure)
    MOIU.attach_optimizer(structure.first_stage)
    for s in subproblems(structure, 2)
        MOIU.attach_optimizer(s)
    end
    return nothing
end

function attach_mocks!(structure::HorizontalStructure)
    return nothing
end

function decision_index(model::MOI.ModelLike, index::MOI.VariableIndex)
    ci = CI{MOI.SingleVariable,SingleDecisionSet{Float64}}(index.value)
    if MOI.is_valid(model, ci)
        return MOI.get(model, DecisionIndex(), ci)
    end
    # Locate multiple decision set
    F = MOI.VectorOfVariables
    S = MultipleDecisionSet{Float64}
    for ci in MOI.get(model, MOI.ListOfConstraintIndices{F,S}())
        f = MOI.get(model, MOI.ConstraintFunction(), ci)
        i = something(findfirst(vi -> vi == index, f.variables), 0)
        if i != 0
            return MOI.get(model, DecisionIndex(), ci)
        end
    end
end

function decision_index(model::MOI.ModelLike, ci::MOI.ConstraintIndex)
    return MOI.get(model, DecisionIndex(), ci)
end

function MOI.get(src::MOIU.MockOptimizer, attr::ScenarioDependentModelAttribute)
    return MOI.get(src, attr.attr)
end

function MOI.get(src::MOIU.MockOptimizer, attr::ScenarioDependentVariableAttribute, index::MOI.VariableIndex)
    return MOI.get(src, attr.attr, index)
end

function MOI.get(src::MOIU.MockOptimizer, attr::ScenarioDependentConstraintAttribute, ci::MOI.ConstraintIndex)
    return MOI.get(src, attr.attr, ci)
end

function MOI.set(src::MOIU.MockOptimizer, attr::ScenarioDependentModelAttribute, value)
    return MOI.set(src, attr.attr, value)
end

function MOI.set(src::MOIU.MockOptimizer, attr::ScenarioDependentVariableAttribute, index::MOI.VariableIndex, value)
    return MOI.set(src, attr.attr, index, value)
end

function MOI.set(src::MOIU.MockOptimizer, attr::ScenarioDependentConstraintAttribute, ci::MOI.ConstraintIndex, value)
    return MOI.set(src, attr.attr, ci, value)
end
# ========================== #


INITIALIZE_DEFAULT(cb,t,u,integrator) = nothing

immutable ContinuousCallback{F1,F2,F3,F4,T,T2,I} <: DECallback
  condition::F1
  affect!::F2
  affect_neg!::F3
  initialize::F4
  idxs::I
  rootfind::Bool
  interp_points::Int
  save_positions::Tuple{Bool,Bool}
  abstol::T
  reltol::T2
end

ContinuousCallback(condition,affect!,affect_neg!;
                   initialize = INITIALIZE_DEFAULT,
                   idxs = nothing,
                   rootfind=true,
                   save_positions=(true,true),
                   interp_points=10,
                   abstol=1e-12,reltol=0) = ContinuousCallback(
                              condition,affect!,affect_neg!,initialize,
                              idxs,
                              rootfind,interp_points,
                              save_positions,abstol,reltol)

function ContinuousCallback(condition,affect!;
                   initialize = INITIALIZE_DEFAULT,
                   idxs = nothing,
                   rootfind=true,
                   save_positions=(true,true),
                   affect_neg! = affect!,
                   interp_points=10,
                   abstol=1e-12,reltol=0)

 ContinuousCallback(
            condition,affect!,affect_neg!,initialize,idxs,
            rootfind,interp_points,
            save_positions,abstol,reltol)

end

immutable DiscreteCallback{F1,F2,F3} <: DECallback
  condition::F1
  affect!::F2
  initialize::F3
  save_positions::Tuple{Bool,Bool}
end
DiscreteCallback(condition,affect!;
        initialize = INITIALIZE_DEFAULT,save_positions=(true,true)) = DiscreteCallback(condition,affect!,initialize,save_positions)

# DiscreteCallback(condition,affect!,save_positions) = DiscreteCallback(condition,affect!,save_positions)

immutable CallbackSet{T1<:Tuple,T2<:Tuple} <: DECallback
  continuous_callbacks::T1
  discrete_callbacks::T2
end

CallbackSet(callback::DiscreteCallback) = CallbackSet((),(callback,))
CallbackSet(callback::ContinuousCallback) = CallbackSet((callback,),())
CallbackSet() = CallbackSet((),())
CallbackSet(cb::Void) = CallbackSet()

# For Varargs, use recursion to make it type-stable
CallbackSet(callbacks::Union{DECallback,Void}...) = CallbackSet(split_callbacks((), (), callbacks...)...)

@inline split_callbacks(cs, ds) = cs, ds
@inline split_callbacks(cs, ds, c::Void, args...) = split_callbacks(cs, ds, args...)
@inline split_callbacks(cs, ds, c::ContinuousCallback, args...) = split_callbacks((cs..., c), ds, args...)
@inline split_callbacks(cs, ds, d::DiscreteCallback, args...) = split_callbacks(cs, (ds..., d), args...)
@inline function split_callbacks(cs, ds, d::CallbackSet, args...)
  split_callbacks((cs...,d.continuous_callbacks...), (ds..., d.discrete_callbacks...), args...)
end

initialize!(cb::CallbackSet,t,u,integrator::DEIntegrator) = initialize!(t,u,integrator,cb.continuous_callbacks...,cb.discrete_callbacks...)
initialize!(cb::CallbackSet{Tuple{},Tuple{}},t,u,integrator::DEIntegrator) = nothing
initialize!(t,u,integrator::DEIntegrator,c::DECallback,cs::DECallback...) = (c.initialize(c,t,u,integrator); initialize!(t,u,integrator,cs...))
initialize!(t,u,integrator::DEIntegrator,c::DECallback) = c.initialize(c,t,u,integrator)


# В данном файле собраны служебные функции, необходимые для:
# 1. обработки входных аргументов вызова скрипта;
# 2. чтения больших файлов в заданном формате с возможным пропуском закомментированных или ненужных строк  заголовка;
# 3. вывод числа типа Int с добавлением необходимого числа нулей в начале;
# 4. функция логирования событий.

# Функции обработки входящего аргумента скрипта, задающего набор степеней полиномов для вывода результатов аппроксимации.
# Данная функция обрабатывает параметр, заключенный в квадратные скобки
# и на основании указанных данных формирует массив элементов
# Пример:
# на вход подается строка "[1,2,3,4,5:5:20,30,40]"
# на выходе формируется массив Array{Int}[1,2,3,4,5,10,15,20,30,40]

function make_numberlist(arg_string)

    arg_split = split(strip(arg_string, ['[',']']), ",")

    arg_deg = Any[]
    for arg in arg_split
        s_arg = split(arg, ":")
        if length(s_arg) >1
            if length(s_arg) == 2
                st = parse(Int, s_arg[1])
                en = parse(Int, s_arg[2])
                push!(arg_deg, st:en)
            end
            if length(s_arg) == 3
                st = parse(Int, s_arg[1])
                sp = parse(Int, s_arg[2])
                en = parse(Int, s_arg[3])
                push!(arg_deg, st:sp:en)
            end
        else
            push!(arg_deg, parse(Int, arg))
        end
    end
    degs = Int[]
    for els in arg_deg
        if typeof(els) <: Integer
            push!(degs, els)
        else
            for el_in_collect in els
                push!(degs, el_in_collect)
            end
        end
    end
    return degs
end

# Данная функция предназначена для формирования набора окон для обработки
# На вход функции подается строка аргументов с параметрами окна и максимальный размер данных, на которые формируются окна
# Строка аргументов может быть задана тремя способами:
# M:N - одним окном;
# M:W:N - окнами с перекрытием в одну точку;
# M:W:S:N - окнами с заданным перекрытием.
# При этом всюду выше:
# M - начальная точка первого окна,
# N - последняя точка последнего окна (в качестве последней точки
# можно указать конец файла, для этого используется ключевое слово end),
# W - размер окна,
# S - перекрытие соседних окон.
# !!!ВАЖНОЕ ЗАМЕЧАНИЕ - 1!!!
# При формировании окон последнее окно формируется так, чтобы
# последняя точка окна совпадала с последней заданной точкой N,
# при этом перекрытие с предпоследним окном может быть больше, чем указано в исходных параметрах
# !!!ВАЖНОЕ ЗАМЕЧАНИЕ - 2!!!
# Если размер окна задан больше, чем количество доступных данных, то скрипт выдает ошибку в лог-файл и завершает работу
# **********************************************************# Примеры
# "1:10", 20 => [(1,10)]
# "1:10:19", 18 => [(1, 10), (9, 18)]
# "1:10:19", 20 => [(1, 10), (10, 19)]
# "1:10:20", 20 => [(1, 10), (10, 19), (11, 20)]
# "1:10:3:20", 20 => [(1, 10), (8, 17), (11, 20)]

function make_windows(arg_string, data_length)
    windows_array = Any[]
    arg = split(arg_string, ':')
    if length(arg) == 2 # Вариант, когда используется 1 окно
        start_index = parse(Int, arg[1])
        if arg[2] == "end"
            end_index = data_length
        else
            end_index =  parse(Int, arg[2])
            if end_index>data_length
                return windows_array
            end
        end
        windows_array = [(start_index, end_index)]
    elseif length(arg) == 3 # окна c перекрытием в 1 точку
        start_index = parse(Int, arg[1])
        step_size = parse(Int, arg[2])
        if arg[3] == "end"
            end_index = data_length
        else
            end_index =  parse(Int, arg[3])
            if end_index>data_length
                return windows_array
            end
        end
        for start_window = start_index:step_size-1:end_index - step_size
            push!(windows_array, (start_window, start_window + step_size-1))
        end
        if length(windows_array) ==0
            return windows_array
        end
        if windows_array[end][2] < end_index
            push!(windows_array, (end_index - step_size+1, end_index))
        end
    elseif length(arg) == 4 # окна с перекрытием
        start_index = parse(Int, arg[1])
        step_size = parse(Int, arg[2])
        overlap = parse(Int, arg[3])
        if arg[4] == "end"
            end_index = data_length
        else
            end_index =  parse(Int, arg[4])
            if end_index>data_length
                println("return")
                return windows_array
            end
        end
        println("read param: start$start_index - step$step_size - overlap$overlap - end$end_index")
        for start_window = start_index:(step_size - overlap):(end_index - step_size)
            push!(windows_array, (start_window, start_window + step_size-1))
        end
        if length(windows_array) ==0
            return windows_array
        end
        if windows_array[end][2] < end_index
            push!(windows_array, (end_index - step_size + 1, end_index))
        end
    else
        windows_array = [(1, data_length)]
    end
    return windows_array
end

# функция для чтения данных
function readbigfile(filename::AbstractString,           # имя входного файла
    Columns::Array{Int,1},     # номера стобцов, которые необходимо считать. Если указать 0 как один из номеров для считывания, то будет сформирован столбец с номерами строк (пропущенные при чтении строки не считаются)
    T::Type                    # тип данных, в который будут сконвертированы считанные столбцы
    ;
    skip_first_row = 0        # число строк в начале файла, которые необходимо пропустить
    )::Array{T,2}             # функция возвращает двумерный массив

    # Открываем файл для чтения
    in_file = readlines(filename)
    skip_row = 0
    for line in in_file
        if skip_first_row > skip_row
            skip_row=skip_row+1
            continue
        end
        # пропускаем строки, начинающиеся (В ПЕРВОЙ ПОЗИЦИИ!) c #
        if line[1] == '#'
            skip_row = skip_row+1
            continue
        end
        # if length(split(line)) < length(Columns)
        #     skip_row = skip_row+1
        #     continue
        # end
    end
    println("skipped rows: $skip_row")
    # Создаем нулевой двумерный массив типа T
    out_array = zeros(T, (length(in_file) - skip_row, length(Columns)))
    # Построчно обрабатываем файл
    num_out_array_row = 1
    for num_str = 1+skip_first_row:length(in_file)
        if in_file[num_str][1] == '#'
            continue
        end
#  пропускаем строки с недостаточным числом столбцов
        # if length(split(in_file[num_str])) < length(Columns)
        #     continue
        # end
        value_str = split(in_file[num_str])
        for num_col = 1:length(Columns)
            if Columns[num_col] == 0
                out_array[num_out_array_row,num_col] = num_out_array_row
            else
                out_array[num_out_array_row,num_col] = parse(T, value_str[Columns[num_col]])
            end
        end
        num_out_array_row = num_out_array_row + 1
    end
    # Возвращаем считанный массив
    return out_array
end

# Данная функция принимает на вход два числа:
# 1. целое положительное число, которое необходимо напечатать
# 2. минимальное число символов в выводимом числе.
# Если символов в числе меньше указанного, перед числом добавляются нули так, чтобы число с нулями стало необходимой минимальной длины.
# Если число длиннее, выводится просто указанное число
function print_with_zero(num::Integer, min_length_num::Integer)
    str_num = "$num"
    zero_in_begin = "0"^maximum([0,(min_length_num-length(str_num))])
    return "$(zero_in_begin)$(str_num)"
end

# Функция дописывает сообщение, указанное вторым аргументом функции в файл с именем указанным как первый аргумент функции,
# при этом указывается время записи
function add_log_messange(log_name, message)
    out_file = open(log_name, "a")
    println(out_file, "[$(now())] $(message)")
    flush(out_file)
    close(out_file)
end

# Функция для вывода двумерных и одномерных массивов с указанием разделителя между элементами
# На вход подается входной массив и могут быть указаны разделители между элементами строки (col_delims)
# и между строками (row_delims),
# если не указывать - по умолчанию между элементами табуляция и между строками перенос строки
# также может быть указан поток вывода, по умолчанию STDOUT
function println_array(InputArray::Array;col_delims = "\t", row_delims="\n", out_stream = STDOUT)
    dims = ndims(InputArray)
    if dims == 1
        println(out_stream,join(map(x->"$(x)",InputArray),col_delims))
    elseif dims == 2
        num_rows = size(InputArray)[1]
        for ind_row = 1:num_rows
            print(out_stream,join(map(x->"$(x)",InputArray[ind_row, :]),col_delims))
            print(out_stream,row_delims)
        end
        println(out_stream)
    else
        println("warning: Dims > 2")
    end
end


function OrientedPoint(arg0::Vector1D, arg1::jboolean, arg2::jdouble)
    return OrientedPoint((Vector1D, jboolean, jdouble), arg0, arg1, arg2)
end

function copy_self(obj::OrientedPoint)
    return jcall(obj, "copySelf", OrientedPoint, ())
end

function empty_hyperplane(obj::OrientedPoint)
    return jcall(obj, "emptyHyperplane", SubOrientedPoint, ())
end

function get_location(obj::OrientedPoint)
    return jcall(obj, "getLocation", Vector1D, ())
end

function get_offset(obj::OrientedPoint, arg0::JVector)
    return jcall(obj, "getOffset", jdouble, (JVector,), arg0)
end

function get_offset(obj::OrientedPoint, arg0::Point)
    return jcall(obj, "getOffset", jdouble, (Point,), arg0)
end

function get_tolerance(obj::OrientedPoint)
    return jcall(obj, "getTolerance", jdouble, ())
end

function is_direct(obj::OrientedPoint)
    return jcall(obj, "isDirect", jboolean, ())
end

function project(obj::OrientedPoint, arg0::Point)
    return jcall(obj, "project", Point, (Point,), arg0)
end

function revert_self(obj::OrientedPoint)
    return jcall(obj, "revertSelf", void, ())
end

function same_orientation_as(obj::OrientedPoint, arg0::Hyperplane)
    return jcall(obj, "sameOrientationAs", jboolean, (Hyperplane,), arg0)
end

function whole_hyperplane(obj::OrientedPoint)
    return jcall(obj, "wholeHyperplane", SubOrientedPoint, ())
end

function whole_space(obj::OrientedPoint)
    return jcall(obj, "wholeSpace", IntervalsSet, ())
end


using Coconet
using Flux
using BSON
using Plots
using MusicTransformer
using Dates

function train_model(;
    model = nothing,
    opt = nothing,
    losses = nothing,

    head_count = 8,
    latent_size = 32,
    word_size = 3,
    N = 6,

    sequence_length = 32,
    batch_size = 5,
    epochs = 1,
    set_length = nothing,
    max_regions = nothing,

    scramble_rate = 0.1,
    scramble_increase = 0.1,
    scramble_max = 0.9,
    graduation_rate = 0.5,

    dataset_name = "Tonnetz_Bach",
)

    difficulty_level = 1
    difficulties = [scramble_rate:scramble_increase:scramble_max...]

    pos_data = nothing

    ## Get model
    epoch_start = 1

    if model == nothing
        @info "Building model..."
        model = Transformer(N, head_count, word_size, latent_size)
        opt = Flux.ADAM()
    end

    model = model |> gpu
    ps = params(model)

    ## Data
    @info "Wrangling data"
    path = "datasets"
    dataset = BSON.load(joinpath(path, string(dataset_name, ".bson")))

    if set_length != nothing
        dataset[:train] = dataset[:train][1:set_length]
        dataset[:test] = dataset[:test][1:set_length]
    end

    ## Loss
    function loss(query_data, mask, y)
        l = y .- model(query_data, pos_data, mask=mask)
        l = l .* l
        sum(l)
    end

    function test_loss(set::Symbol)
        L_test = 0 # Total loss
        i = 0 # counter for the progress bar

        test = Coconet.get_batch(
            dataset,
            sequence_length,
            batch_size,
            set = set,
            max_regions = max_regions,
        )

        for batch in test
            L_batch = 0 # Loss per batch
            for song in batch
                L_song = 0 # loss per song
                song = gpu(song)
                for region in song
                    # data
                    mask = gpu(MusicTransformer.get_mask(region, scramble_rate))

                    query_data = region ./ (mask[1,1,:,:,:] .+ 1)
                    key_data = query_data
                    y = region
                    pos_data = gpu(get_positional_encoding(region))

                    # Loss per region
                    L_song += loss(query_data, mask, y)
                end

                # Average of song losses across batch
                L_batch += L_song

                # Progress bar
                if i % 4 == 0
                    if i % 16 == 0
                        print(' ')
                    end
                    print('/')
                end
                i += 1
            end

            # Losses across test set
            L_test += L_batch
        end

        print(string("> ", L_test, "\n"))
        return L_test
    end

    function ratio()
        @info string("Current Ratio")
        print(string("Train: ", losses[length(losses)][1]/losses[1][1], "\n"))
        print(string("Test: ", losses[length(losses)][2]/losses[1][2], "\n"))
    end

    ## Plot
    function save_to_plot(epoch)
        train = [loss[1] for loss in losses] ./ losses[1][1]
        test = [loss[2] for loss in losses] ./ losses[1][2]
        x = 0:1:(length(losses)-1)

        p = plot(x, train, label="Train")
        plot!(x, test, label="Test")

        ylims!(0, max(train..., test..., 1))
        xlims!(0, length(losses) - 1)

        xlabel!("Epochs")
        p = ylabel!("percent of loss_0")

        p = plot(
            p,
            size = (1280,720),
            xticks = x,
            yticks = 0:0.1:1,
        )

        savefig(p, string("checkpoint/loss_curve"))
        return p
    end

    # Get initial loss
    @info string("Initial Loss:\n")
    if losses == nothing
        print("Train: ")
        tr_l = test_loss(:train)

        print("Test: ")
        te_l = test_loss(:test)

        losses = [(tr_l, te_l)]
    else
        print(string("Train: ", losses[1][1],"\n"))
        print(string("Test: ", losses[1][2],"\n"))
    end
    ratio()
    best_loss = losses[length(losses)][2]
    last_improvement = best_loss / losses[1][2]

    # Train!
    for epoch = epoch_start:(epoch_start+epochs - 1)
        i = 0 # counter for the progress bar
        @info string("Epoch ", epoch)
        train = Coconet.get_batch(
            dataset,
            sequence_length,
            batch_size,
            set = :train,
            max_regions = max_regions,
        )

        L_train = 0
        batch_num = 1

        # Epoch
        for batch in train
            print(string("Batch ", batch_num," of ",length(train)))

            ∇_batch = [] # Gradient per batch
            L_batch = 0 # Loss per batch
            for song in batch
                ∇_song = [] # Gradient per song
                L_song = 0 # Loss per song
                song = gpu(song)
                for region in song
                    # data
                    mask = gpu(MusicTransformer.get_mask(region, scramble_rate))

                    query_data = region ./ (mask[1,1,:,:,:] .+ 1)
                    key_data = query_data
                    y = region
                    pos_data = gpu(get_positional_encoding(region))

                    L_region = 0

                    #time
                    t = now()

                    # Calculate gradient over region
                    ∇_region = gradient(ps) do
                        L_region = loss(query_data, mask, y)
                    end

                    #time
                    print(string("\n", Dates.canonicalize(now() - t), "\n"))

                    push!(∇_song, cpu(∇_region))
                    L_song += L_region
                end

                # Store song gradients
                push!(∇_batch, ∇_song)
                L_batch += L_song

                # Progress bar
                print('.')
            end
            print("\n")

            ∇_batch = gpu(∇_batch)
            # Combine gradients
            ∇_batch = sum([sum(∇_song) for ∇_song in ∇_batch])
            L_train += L_batch

            # Apply gradients (a.k.a. Train!)
            Flux.update!(opt, ps, ∇_batch)
            batch_num += 1
        end

        # Get loss
        L_test = test_loss(:test)
        push!(losses, (L_train, L_test))

        # Print current state
        ratio()

        #FUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU
        if isnan(L_test)
            print("FUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU")
            break
        end

        # Make it harder if it's doing too well
        train_ratio = losses[length(losses)][2] / losses[1][2]
        improvement = 1 - (train_ratio / last_improvement)

        if epoch > 0 && epoch % 2 == 0
            if difficulty_level < length(difficulties)
                difficulty_level += 1
                last_improvement = train_ratio
                @info string("Increasing scramble rate to ", difficulties[difficulty_level])
            end
        elseif improvement < 0
            last_improvment = train_ratio
            @warn string("Looks like they're struggling...")
        end

        # Save checkpoint
        if losses[length(losses)][2] < best_loss
            @info "Saving model"
            best_loss = losses[length(losses)][2]
            m = model |> cpu
            o = opt |> cpu
            BSON.@save string("checkpoint/my_model.bson") model=m opt=o epoch_start=epoch+1 losses=losses scramble_rate
        end

        # Plot losses
        save_to_plot(epoch)
    end
end


# Julia wrapper for header: /opt/intel/compilers_and_libraries_2017.1.132/linux/mkl/include/mkl_dnn.h
# Automatically generated using Clang.jl wrap_c, version 0.0.0

function dnnLayoutCreate_F32(pLayout,dimension,size,strides)
    check_dnnerror(ccall((:dnnLayoutCreate_F32,libmkl),dnnError_t,(Ptr{dnnLayout_t},Csize_t,Ptr{Csize_t},Ptr{Csize_t}),pLayout,dimension,size,strides))
end

function dnnLayoutCreateFromPrimitive_F32(pLayout,primitive,_type)
    check_dnnerror(ccall((:dnnLayoutCreateFromPrimitive_F32,libmkl),dnnError_t,(Ptr{dnnLayout_t},dnnPrimitive_t,dnnResourceType_t),pLayout,primitive,_type))
end

function dnnLayoutGetMemorySize_F32(layout)
    ccall((:dnnLayoutGetMemorySize_F32,libmkl),Csize_t,(dnnLayout_t,),layout)
end

function dnnLayoutCompare_F32(l1,l2)
    ccall((:dnnLayoutCompare_F32,libmkl),Cint,(dnnLayout_t,dnnLayout_t),l1,l2)
end

function dnnAllocateBuffer_F32(pPtr,layout)
    check_dnnerror(ccall((:dnnAllocateBuffer_F32,libmkl),dnnError_t,(Ptr{Ptr{Cvoid}},dnnLayout_t),pPtr,layout))
end

function dnnReleaseBuffer_F32(ptr)
    check_dnnerror(ccall((:dnnReleaseBuffer_F32,libmkl),dnnError_t,(Ptr{Cvoid},),ptr))
end

function dnnLayoutDelete_F32(layout)
    check_dnnerror(ccall((:dnnLayoutDelete_F32,libmkl),dnnError_t,(dnnLayout_t,),layout))
end

function dnnPrimitiveAttributesCreate_F32(attributes)
    check_dnnerror(ccall((:dnnPrimitiveAttributesCreate_F32,libmkl),dnnError_t,(Ptr{dnnPrimitiveAttributes_t},),attributes))
end

function dnnPrimitiveAttributesDestroy_F32(attributes)
    check_dnnerror(ccall((:dnnPrimitiveAttributesDestroy_F32,libmkl),dnnError_t,(dnnPrimitiveAttributes_t,),attributes))
end

function dnnPrimitiveGetAttributes_F32(primitive,attributes)
    check_dnnerror(ccall((:dnnPrimitiveGetAttributes_F32,libmkl),dnnError_t,(dnnPrimitive_t,Ptr{dnnPrimitiveAttributes_t}),primitive,attributes))
end

function dnnExecute_F32(primitive,resources)
    check_dnnerror(ccall((:dnnExecute_F32,libmkl),dnnError_t,(dnnPrimitive_t,Ptr{Ptr{Cvoid}}),primitive,resources))
end

function dnnExecuteAsync_F32(primitive,resources)
    check_dnnerror(ccall((:dnnExecuteAsync_F32,libmkl),dnnError_t,(dnnPrimitive_t,Ptr{Ptr{Cvoid}}),primitive,resources))
end

function dnnWaitFor_F32(primitive)
    check_dnnerror(ccall((:dnnWaitFor_F32,libmkl),dnnError_t,(dnnPrimitive_t,),primitive))
end

function dnnDelete_F32(primitive)
    check_dnnerror(ccall((:dnnDelete_F32,libmkl),dnnError_t,(dnnPrimitive_t,),primitive))
end

function dnnConversionCreate_F32(pConversion,from,to)
    check_dnnerror(ccall((:dnnConversionCreate_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnLayout_t,dnnLayout_t),pConversion,from,to))
end

function dnnConversionExecute_F32(conversion,from,to)
    check_dnnerror(ccall((:dnnConversionExecute_F32,libmkl),dnnError_t,(dnnPrimitive_t,Ptr{Cvoid},Ptr{Cvoid}),conversion,from,to))
end

function dnnSumCreate_F32(pSum,attributes,nSummands,layout,coefficients)
    check_dnnerror(ccall((:dnnSumCreate_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,dnnLayout_t,Ptr{Cfloat}),pSum,attributes,nSummands,layout,coefficients))
end

function dnnConcatCreate_F32(pConcat,attributes,nSrcTensors,src)
    check_dnnerror(ccall((:dnnConcatCreate_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{dnnLayout_t}),pConcat,attributes,nSrcTensors,src))
end

function dnnSplitCreate_F32(pSplit,attributes,nDstTensors,layout,dstChannelSize)
    check_dnnerror(ccall((:dnnSplitCreate_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,dnnLayout_t,Ptr{Csize_t}),pSplit,attributes,nDstTensors,layout,dstChannelSize))
end

function dnnScaleCreate_F32(pScale,attributes,dataLayout,alpha)
    check_dnnerror(ccall((:dnnScaleCreate_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Cfloat),pScale,attributes,dataLayout,alpha))
end

function dnnConvolutionCreateForward_F32(pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnConvolutionCreateForward_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnConvolutionCreateForwardBias_F32(pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnConvolutionCreateForwardBias_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnConvolutionCreateBackwardData_F32(pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnConvolutionCreateBackwardData_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnConvolutionCreateBackwardFilter_F32(pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnConvolutionCreateBackwardFilter_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnConvolutionCreateBackwardBias_F32(pConvolution,attributes,algorithm,dimension,dstSize)
    check_dnnerror(ccall((:dnnConvolutionCreateBackwardBias_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Ptr{Csize_t}),pConvolution,attributes,algorithm,dimension,dstSize))
end

function dnnGroupsConvolutionCreateForward_F32(pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnGroupsConvolutionCreateForward_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnGroupsConvolutionCreateForwardBias_F32(pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnGroupsConvolutionCreateForwardBias_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnGroupsConvolutionCreateBackwardData_F32(pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnGroupsConvolutionCreateBackwardData_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnGroupsConvolutionCreateBackwardFilter_F32(pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnGroupsConvolutionCreateBackwardFilter_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnGroupsConvolutionCreateBackwardBias_F32(pConvolution,attributes,algorithm,groups,dimension,dstSize)
    check_dnnerror(ccall((:dnnGroupsConvolutionCreateBackwardBias_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Csize_t,Ptr{Csize_t}),pConvolution,attributes,algorithm,groups,dimension,dstSize))
end

function dnnReLUCreateForward_F32(pRelu,attributes,dataLayout,negativeSlope)
    check_dnnerror(ccall((:dnnReLUCreateForward_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Cfloat),pRelu,attributes,dataLayout,negativeSlope))
end

function dnnReLUCreateBackward_F32(pRelu,attributes,diffLayout,dataLayout,negativeSlope)
    check_dnnerror(ccall((:dnnReLUCreateBackward_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,dnnLayout_t,Cfloat),pRelu,attributes,diffLayout,dataLayout,negativeSlope))
end

function dnnLRNCreateForward_F32(pLrn,attributes,dataLayout,kernel_size,alpha,beta,k)
    check_dnnerror(ccall((:dnnLRNCreateForward_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Csize_t,Cfloat,Cfloat,Cfloat),pLrn,attributes,dataLayout,kernel_size,alpha,beta,k))
end

function dnnLRNCreateBackward_F32(pLrn,attributes,diffLayout,dataLayout,kernel_size,alpha,beta,k)
    check_dnnerror(ccall((:dnnLRNCreateBackward_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,dnnLayout_t,Csize_t,Cfloat,Cfloat,Cfloat),pLrn,attributes,diffLayout,dataLayout,kernel_size,alpha,beta,k))
end

function dnnBatchNormalizationCreateForward_F32(pBatchNormalization,attributes,dataLayout,eps)
    check_dnnerror(ccall((:dnnBatchNormalizationCreateForward_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Cfloat),pBatchNormalization,attributes,dataLayout,eps))
end

function dnnBatchNormalizationCreateBackwardScaleShift_F32(pBatchNormalization,attributes,dataLayout,eps)
    check_dnnerror(ccall((:dnnBatchNormalizationCreateBackwardScaleShift_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Cfloat),pBatchNormalization,attributes,dataLayout,eps))
end

function dnnBatchNormalizationCreateBackwardData_F32(pBatchNormalization,attributes,dataLayout,eps)
    check_dnnerror(ccall((:dnnBatchNormalizationCreateBackwardData_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Cfloat),pBatchNormalization,attributes,dataLayout,eps))
end

function dnnPoolingCreateForward_F32(pPooling,attributes,op,srcLayout,kernelSize,kernelStride,inputOffset,borderType)
    check_dnnerror(ccall((:dnnPoolingCreateForward_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,dnnLayout_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pPooling,attributes,op,srcLayout,kernelSize,kernelStride,inputOffset,borderType))
end

function dnnPoolingCreateBackward_F32(pPooling,attributes,op,srcLayout,kernelSize,kernelStride,inputOffset,borderType)
    check_dnnerror(ccall((:dnnPoolingCreateBackward_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,dnnLayout_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pPooling,attributes,op,srcLayout,kernelSize,kernelStride,inputOffset,borderType))
end

function dnnInnerProductCreateForward_F32(pInnerProduct,attributes,dimensions,srcSize,outputChannels)
    check_dnnerror(ccall((:dnnInnerProductCreateForward_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{Csize_t},Csize_t),pInnerProduct,attributes,dimensions,srcSize,outputChannels))
end

function dnnInnerProductCreateForwardBias_F32(pInnerProduct,attributes,dimensions,srcSize,outputChannels)
    check_dnnerror(ccall((:dnnInnerProductCreateForwardBias_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{Csize_t},Csize_t),pInnerProduct,attributes,dimensions,srcSize,outputChannels))
end

function dnnInnerProductCreateBackwardData_F32(pInnerProduct,attributes,dimensions,srcSize,outputChannels)
    check_dnnerror(ccall((:dnnInnerProductCreateBackwardData_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{Csize_t},Csize_t),pInnerProduct,attributes,dimensions,srcSize,outputChannels))
end

function dnnInnerProductCreateBackwardFilter_F32(pInnerProduct,attributes,dimensions,srcSize,outputChannels)
    check_dnnerror(ccall((:dnnInnerProductCreateBackwardFilter_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{Csize_t},Csize_t),pInnerProduct,attributes,dimensions,srcSize,outputChannels))
end

function dnnInnerProductCreateBackwardBias_F32(pInnerProduct,attributes,dimensions,dstSize)
    check_dnnerror(ccall((:dnnInnerProductCreateBackwardBias_F32,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{Csize_t}),pInnerProduct,attributes,dimensions,dstSize))
end

function dnnLayoutCreate_F64(pLayout,dimension,size,strides)
    check_dnnerror(ccall((:dnnLayoutCreate_F64,libmkl),dnnError_t,(Ptr{dnnLayout_t},Csize_t,Ptr{Csize_t},Ptr{Csize_t}),pLayout,dimension,size,strides))
end

function dnnLayoutCreateFromPrimitive_F64(pLayout,primitive,_type)
    check_dnnerror(ccall((:dnnLayoutCreateFromPrimitive_F64,libmkl),dnnError_t,(Ptr{dnnLayout_t},dnnPrimitive_t,dnnResourceType_t),pLayout,primitive,_type))
end

function dnnLayoutGetMemorySize_F64(layout)
    ccall((:dnnLayoutGetMemorySize_F64,libmkl),Csize_t,(dnnLayout_t,),layout)
end

function dnnLayoutCompare_F64(l1,l2)
    ccall((:dnnLayoutCompare_F64,libmkl),Cint,(dnnLayout_t,dnnLayout_t),l1,l2)
end

function dnnAllocateBuffer_F64(pPtr,layout)
    check_dnnerror(ccall((:dnnAllocateBuffer_F64,libmkl),dnnError_t,(Ptr{Ptr{Cvoid}},dnnLayout_t),pPtr,layout))
end

function dnnReleaseBuffer_F64(ptr)
    check_dnnerror(ccall((:dnnReleaseBuffer_F64,libmkl),dnnError_t,(Ptr{Cvoid},),ptr))
end

function dnnLayoutDelete_F64(layout)
    check_dnnerror(ccall((:dnnLayoutDelete_F64,libmkl),dnnError_t,(dnnLayout_t,),layout))
end

function dnnPrimitiveAttributesCreate_F64(attributes)
    check_dnnerror(ccall((:dnnPrimitiveAttributesCreate_F64,libmkl),dnnError_t,(Ptr{dnnPrimitiveAttributes_t},),attributes))
end

function dnnPrimitiveAttributesDestroy_F64(attributes)
    check_dnnerror(ccall((:dnnPrimitiveAttributesDestroy_F64,libmkl),dnnError_t,(dnnPrimitiveAttributes_t,),attributes))
end

function dnnPrimitiveGetAttributes_F64(primitive,attributes)
    check_dnnerror(ccall((:dnnPrimitiveGetAttributes_F64,libmkl),dnnError_t,(dnnPrimitive_t,Ptr{dnnPrimitiveAttributes_t}),primitive,attributes))
end

function dnnExecute_F64(primitive,resources)
    check_dnnerror(ccall((:dnnExecute_F64,libmkl),dnnError_t,(dnnPrimitive_t,Ptr{Ptr{Cvoid}}),primitive,resources))
end

function dnnExecuteAsync_F64(primitive,resources)
    check_dnnerror(ccall((:dnnExecuteAsync_F64,libmkl),dnnError_t,(dnnPrimitive_t,Ptr{Ptr{Cvoid}}),primitive,resources))
end

function dnnWaitFor_F64(primitive)
    check_dnnerror(ccall((:dnnWaitFor_F64,libmkl),dnnError_t,(dnnPrimitive_t,),primitive))
end

function dnnDelete_F64(primitive)
    check_dnnerror(ccall((:dnnDelete_F64,libmkl),dnnError_t,(dnnPrimitive_t,),primitive))
end

function dnnConversionCreate_F64(pConversion,from,to)
    check_dnnerror(ccall((:dnnConversionCreate_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnLayout_t,dnnLayout_t),pConversion,from,to))
end

function dnnConversionExecute_F64(conversion,from,to)
    check_dnnerror(ccall((:dnnConversionExecute_F64,libmkl),dnnError_t,(dnnPrimitive_t,Ptr{Cvoid},Ptr{Cvoid}),conversion,from,to))
end

function dnnSumCreate_F64(pSum,attributes,nSummands,layout,coefficients)
    check_dnnerror(ccall((:dnnSumCreate_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,dnnLayout_t,Ptr{Cdouble}),pSum,attributes,nSummands,layout,coefficients))
end

function dnnConcatCreate_F64(pConcat,attributes,nSrcTensors,src)
    check_dnnerror(ccall((:dnnConcatCreate_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{dnnLayout_t}),pConcat,attributes,nSrcTensors,src))
end

function dnnSplitCreate_F64(pSplit,attributes,nDstTensors,layout,dstChannelSize)
    check_dnnerror(ccall((:dnnSplitCreate_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,dnnLayout_t,Ptr{Csize_t}),pSplit,attributes,nDstTensors,layout,dstChannelSize))
end

function dnnScaleCreate_F64(pScale,attributes,dataLayout,alpha)
    check_dnnerror(ccall((:dnnScaleCreate_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Cdouble),pScale,attributes,dataLayout,alpha))
end

function dnnConvolutionCreateForward_F64(pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnConvolutionCreateForward_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnConvolutionCreateForwardBias_F64(pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnConvolutionCreateForwardBias_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnConvolutionCreateBackwardData_F64(pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnConvolutionCreateBackwardData_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnConvolutionCreateBackwardFilter_F64(pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnConvolutionCreateBackwardFilter_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnConvolutionCreateBackwardBias_F64(pConvolution,attributes,algorithm,dimension,dstSize)
    check_dnnerror(ccall((:dnnConvolutionCreateBackwardBias_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Ptr{Csize_t}),pConvolution,attributes,algorithm,dimension,dstSize))
end

function dnnGroupsConvolutionCreateForward_F64(pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnGroupsConvolutionCreateForward_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnGroupsConvolutionCreateForwardBias_F64(pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnGroupsConvolutionCreateForwardBias_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnGroupsConvolutionCreateBackwardData_F64(pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnGroupsConvolutionCreateBackwardData_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnGroupsConvolutionCreateBackwardFilter_F64(pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType)
    check_dnnerror(ccall((:dnnGroupsConvolutionCreateBackwardFilter_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Csize_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pConvolution,attributes,algorithm,groups,dimension,srcSize,dstSize,filterSize,convolutionStrides,inputOffset,borderType))
end

function dnnGroupsConvolutionCreateBackwardBias_F64(pConvolution,attributes,algorithm,groups,dimension,dstSize)
    check_dnnerror(ccall((:dnnGroupsConvolutionCreateBackwardBias_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,Csize_t,Csize_t,Ptr{Csize_t}),pConvolution,attributes,algorithm,groups,dimension,dstSize))
end

function dnnReLUCreateForward_F64(pRelu,attributes,dataLayout,negativeSlope)
    check_dnnerror(ccall((:dnnReLUCreateForward_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Cdouble),pRelu,attributes,dataLayout,negativeSlope))
end

function dnnReLUCreateBackward_F64(pRelu,attributes,diffLayout,dataLayout,negativeSlope)
    check_dnnerror(ccall((:dnnReLUCreateBackward_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,dnnLayout_t,Cdouble),pRelu,attributes,diffLayout,dataLayout,negativeSlope))
end

function dnnLRNCreateForward_F64(pLrn,attributes,dataLayout,kernel_size,alpha,beta,k)
    check_dnnerror(ccall((:dnnLRNCreateForward_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Csize_t,Cdouble,Cdouble,Cdouble),pLrn,attributes,dataLayout,kernel_size,alpha,beta,k))
end

function dnnLRNCreateBackward_F64(pLrn,attributes,diffLayout,dataLayout,kernel_size,alpha,beta,k)
    check_dnnerror(ccall((:dnnLRNCreateBackward_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,dnnLayout_t,Csize_t,Cdouble,Cdouble,Cdouble),pLrn,attributes,diffLayout,dataLayout,kernel_size,alpha,beta,k))
end

function dnnBatchNormalizationCreateForward_F64(pBatchNormalization,attributes,dataLayout,eps)
    check_dnnerror(ccall((:dnnBatchNormalizationCreateForward_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Cdouble),pBatchNormalization,attributes,dataLayout,eps))
end

function dnnBatchNormalizationCreateBackwardScaleShift_F64(pBatchNormalization,attributes,dataLayout,eps)
    check_dnnerror(ccall((:dnnBatchNormalizationCreateBackwardScaleShift_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Cdouble),pBatchNormalization,attributes,dataLayout,eps))
end

function dnnBatchNormalizationCreateBackwardData_F64(pBatchNormalization,attributes,dataLayout,eps)
    check_dnnerror(ccall((:dnnBatchNormalizationCreateBackwardData_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnLayout_t,Cdouble),pBatchNormalization,attributes,dataLayout,eps))
end

function dnnPoolingCreateForward_F64(pPooling,attributes,op,srcLayout,kernelSize,kernelStride,inputOffset,borderType)
    check_dnnerror(ccall((:dnnPoolingCreateForward_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,dnnLayout_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pPooling,attributes,op,srcLayout,kernelSize,kernelStride,inputOffset,borderType))
end

function dnnPoolingCreateBackward_F64(pPooling,attributes,op,srcLayout,kernelSize,kernelStride,inputOffset,borderType)
    check_dnnerror(ccall((:dnnPoolingCreateBackward_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,dnnAlgorithm_t,dnnLayout_t,Ptr{Csize_t},Ptr{Csize_t},Ptr{Cint},dnnBorder_t),pPooling,attributes,op,srcLayout,kernelSize,kernelStride,inputOffset,borderType))
end

function dnnInnerProductCreateForward_F64(pInnerProduct,attributes,dimensions,srcSize,outputChannels)
    check_dnnerror(ccall((:dnnInnerProductCreateForward_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{Csize_t},Csize_t),pInnerProduct,attributes,dimensions,srcSize,outputChannels))
end

function dnnInnerProductCreateForwardBias_F64(pInnerProduct,attributes,dimensions,srcSize,outputChannels)
    check_dnnerror(ccall((:dnnInnerProductCreateForwardBias_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{Csize_t},Csize_t),pInnerProduct,attributes,dimensions,srcSize,outputChannels))
end

function dnnInnerProductCreateBackwardData_F64(pInnerProduct,attributes,dimensions,srcSize,outputChannels)
    check_dnnerror(ccall((:dnnInnerProductCreateBackwardData_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{Csize_t},Csize_t),pInnerProduct,attributes,dimensions,srcSize,outputChannels))
end

function dnnInnerProductCreateBackwardFilter_F64(pInnerProduct,attributes,dimensions,srcSize,outputChannels)
    check_dnnerror(ccall((:dnnInnerProductCreateBackwardFilter_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{Csize_t},Csize_t),pInnerProduct,attributes,dimensions,srcSize,outputChannels))
end

function dnnInnerProductCreateBackwardBias_F64(pInnerProduct,attributes,dimensions,dstSize)
    check_dnnerror(ccall((:dnnInnerProductCreateBackwardBias_F64,libmkl),dnnError_t,(Ptr{dnnPrimitive_t},dnnPrimitiveAttributes_t,Csize_t,Ptr{Csize_t}),pInnerProduct,attributes,dimensions,dstSize))
end


"""
This implements a cylindrical cutoff for the bond environments:
* A central bond (i,j) is within the cutoff if r_ij < rcutbond.
* A neighbour atom at position rkij (relative position to midpoint) is within
the environment if - after transformation to (r, θ, z) coordinates, it satisfies
`r <= rcutenv` and `abs(z) <= zcutenv`.

This struct implements the resulting filter under `env_filter`.
"""
struct CylindricalCutoff{T}
   rcutbond::T
   rcutenv::T
   zcutenv::T
end

env_filter(r, z, cutoff::CylindricalCutoff) =
      (r <= cutoff.rcutenv) && (abs(z) <= cutoff.zcutenv)

struct ACEBondPotential{TM} <: AbstractCalculator
   models::Dict{Tuple{AtomicNumber, AtomicNumber}, TM}
   cutoff::CylindricalCutoff{Float64}
end

struct ACEBondPotentialBasis{TM} <: JuLIP.MLIPs.IPBasis
   models::Dict{Tuple{AtomicNumber, AtomicNumber}, TM}  # model = basis
   inds::Dict{Tuple{AtomicNumber, AtomicNumber}, UnitRange{Int}}
   cutoff::CylindricalCutoff{Float64}
end

ACEBondCalc = Union{ACEBondPotential, ACEBondPotentialBasis}

function _get_basisinds(V::ACEBondPotential)
   inds = Dict{Tuple{AtomicNumber, AtomicNumber}, UnitRange{Int}}()
   zz = sort(collect(keys(V.models)))
   i0 = 0
   for z in zz
      mo = V.models[z]
      len = length(mo.basis)
      inds[z] = (i0+1):(i0+len)   # to generalize for general models
      i0 += len
   end
   return inds
end

_get_basisinds(V::ACEBondPotentialBasis) = V.inds

function basis(V::ACEBondPotential)
   models = Dict( [zz => model.basis for (zz, model) in V.models]... )
   inds = _get_basisinds(V)
   return ACEBondPotentialBasis(models, inds, V.cutoff)
end

# TODO:
#   - nparams
#   - get_params
#   - set_params!

Base.length(basis::ACEBondPotentialBasis) =
      sum(length(inds) for (_, inds) in basis.inds)

# --------------------------------------------------------

import JuLIP: energy, forces, virial
import ACE: evaluate, evaluate_d, grad_config

# overload the initiation of the bonds iterator to correctly extract the
# right cutoffs.
bonds(at::Atoms, calc::ACEBondCalc) =
         bonds( at, calc.cutoff.rcutbond,
                calc.cutoff.rcutbond/2 + calc.cutoff.zcutenv,
                (r, z) -> env_filter(r, z, calc.cutoff) )

_get_model(calc::ACEBondCalc, zi, zj) =
      calc.models[(min(zi, zj), max(zi,zj))]

function energy(calc::ACEBondPotential, at::Atoms)
   E = 0.0
   for (i, j, rrij, Js, Rs, Zs) in bonds(at, calc)
      # find the right ace model
      ace = _get_model(calc, at.Z[i], at.Z[j])
      # transform the euclidean to cylindrical coordinates
      env = eucl2cyl(rrij, at.Z[i], at.Z[j], Rs, Zs)
      # evaluate
      Eij = evaluate(ace, env)
      E += Eij.val
   end
   return E
end

function forces(calc::ACEBondPotential, at::Atoms)
   F = zeros(SVector{3, Float64}, length(at))
   for (i, j, rrij, Js, Rs, Zs) in bonds(at, calc)
      Zi, Zj = at.Z[i], at.Z[j]
      # find the right ace model
      ace = _get_model(calc, Zi, Zj)
      # transform the euclidean to cylindrical coordinates
      env = eucl2cyl(rrij, Zi, Zj, Rs, Zs)
      # evaluate
      dV_cyl = grad_config(ace, env)
      # transform back?
      dV_drrij, dV_dRs = rrule_eucl2cyl(rrij::SVector, Zi, Zj, Rs, Zs, dV_cyl)
      # assemble the forces
      F[i] += dV_drrij
      F[j] -= dV_drrij
      for (k, dv) in zip(Js, dV_dRs)
         F[k] -= dv
         F[i] += 0.5 * dv
         F[j] += 0.5 * dv
      end
   end
   return F
end

# site_virial(dV::AbstractVector{JVec{T1}}, R::AbstractVector{JVec{T2}}
#             ) where {T1, T2} =  (
#       length(R) > 0 ? (- sum( dVi * Ri' for (dVi, Ri) in zip(dV, R) ))
#                     : zero(JMat{fltype_intersect(T1, T2)})
#       )

function virial(calc::ACEBondPotential, at::Atoms{T}) where {T}
   V = zero(SMatrix{3, 3, T})
   for (i, j, rrij, Js, Rs, Zs) in bonds(at, calc)
      Zi, Zj = at.Z[i], at.Z[j]
      # find the right ace model for this bond
      ace = _get_model(calc, Zi, Zj)
      # transform the euclidean to cylindrical coordinates
      env = eucl2cyl(rrij, Zi, Zj, Rs, Zs)
      # evaluate
      dV_cyl = grad_config(ace, env)
      # transform back?
      dV_drrij, dV_dRs = rrule_eucl2cyl(rrij::SVector, Zi, Zj, Rs, Zs, dV_cyl)
      # assemble the virial
      #   dV_dRs contain derivative relative to midpoint
      #   dV_drrij contain derivative w.r.t. rrij
      V -= dV_drrij * rrij'
      for q = 1:length(dV_dRs)
         V -= dV_dRs[q] * Rs[q]'
      end
   end
   return V
end

function energy(basis::ACEBondPotentialBasis, at::Atoms)
   E = zeros(Float64, length(basis))
   Et = zeros(Float64, length(basis))
   for (i, j, rrij, Js, Rs, Zs) in bonds(at, basis)
      # find the right ace model
      ace = _get_model(basis, at.Z[i], at.Z[j])
      # transform the euclidean to cylindrical coordinates
      env = eucl2cyl(rrij, at.Z[i], at.Z[j], Rs, Zs)
      # evaluate
      ACE.evaluate!(Et, ace, ACE.ACEConfig(env))
      E += Et
   end
   return E
end

function forces(basis::ACEBondPotentialBasis, at::Atoms)
   F = zeros(SVector{3, Float64}, length(basis), length(at))
   for (i, j, rrij, Js, Rs, Zs) in bonds(at, basis)
      Zi, Zj = at.Z[i], at.Z[j]
      # find the right ace model
      ace = _get_model(basis, Zi, Zj)
      # transform the euclidean to cylindrical coordinates
      env = eucl2cyl(rrij, Zi, Zj, Rs, Zs)
      # evaluate
      dB_cyl = evaluate_d(ace, env)
      # transform back?
      dB_drrij, dB_dRs = rrule_eucl2cyl(rrij::SVector, Zi, Zj, Rs, Zs, dB_cyl)
      # assemble the forces
      F[:, i] += dB_drrij
      F[:, j] -= dB_drrij
      for n = 1:length(Js)
         k = Js[n]
         dv = dB_dRs[:, n]
         F[:, k] -= dv
         F[:, i] += 0.5 * dv
         F[:, j] += 0.5 * dv
      end
   end
   return F  # TODO: this is probably the wrong format
end

function virial(basis::ACEBondPotentialBasis, at::Atoms{T}) where {T}
   V = zeros(SMatrix{3, 3, T}, length(basis))
   for (i, j, rrij, Js, Rs, Zs) in bonds(at, basis)
      Zi, Zj = at.Z[i], at.Z[j]
      # find the right ace model
      ace = _get_model(basis, Zi, Zj)
      # transform the euclidean to cylindrical coordinates
      env = eucl2cyl(rrij, Zi, Zj, Rs, Zs)
      # evaluate
      dB_cyl = evaluate_d(ace, env)
      # transform back?
      dB_drrij, dB_dRs = rrule_eucl2cyl(rrij::SVector, Zi, Zj, Rs, Zs, dB_cyl)
      # assemble the virials
      for iB = 1:length(basis)
         V[iB] -= dB_drrij[iB] * rrij'
         for q = 1:length(Rs)
            V[iB] -= dB_dRs[iB, q] * Rs[q]'
         end
      end
   end
   return V  # TODO: double-check the format
end


import NMRDataSetup
import NMRSpectraSimulator
import NMRSpecifyRegions

include("../src/NMRCalibrate.jl")
import .NMRCalibrate

using LinearAlgebra
using FFTW
import PyPlot
import BSON
import JSON

import Statistics
import Random

# include("./helpers/final_helpers.jl")
# include("./helpers/solute_calibrate.jl")
# include("./helpers/loop_entries1.jl")

##### global constants.
#cs_config_path = "/home/roy/Documents/repo/NMRData/src/input/reduced_cs_config.txt"
SH_config_path = "/home/roy/Documents/repo/NMRData/input/SH_configs/select_compounds_SH_configs.json"
surrogate_config_path = "/home/roy/Documents/repo/NMRData/input/surrogate_configs/select_compounds_SH_configs.json"
fit_config_path = "/home/roy/Documents/repo/NMRData/input/fit_configs/calibrate_700MHz_type1_select_compounds.json"

# get mapping from molecule names to their spin system info json files.
H_params_path = "/home/roy/Documents/repo/NMRData/input/coupling_info"
dict_compound_to_filename = JSON.parsefile("/home/roy/Documents/repo/NMRData/input/compound_mapping/select_compounds.json")

# where the experiment bson file is located.
root_folder = "/home/roy/MEGAsync/outputs/NMR/experiments/NRC"
experiment_path = joinpath(root_folder, "NRC-glucose-2018") # second input is subfolder name.
load_path = joinpath(experiment_path, "experiment.bson")

# where the saved calibration project folder is for this experiment.
project_name = "NRC-glucose-2018"

save_base_folder = "/home/roy/MEGAsync/outputs/NMR/calibrate/NRC"
save_folder = joinpath(save_base_folder, project_name)
save_path = joinpath(save_folder, "proxy.bson")
isdir(save_path) || mkpath(save_folder)
#####

##### user input for loading data.

molecule_names = ["D-(+)-Glucose";]
w = [1.0;]

# molecule_names = ["D-(+)-Glucose"; "DSS"; "Ethanol"]
# w = [1.0; 1.0; 1.0]

# molecule_names = ["D-(+)-Glucose"; "DSS"; ]
# w = [1.0; 1.0;]

# note that 0.1% DSS is 0.0046 M = 4.6 mM.

##### end inputs for loading data.

##### user input for surrogate.
dummy_SSFID = NMRSpectraSimulator.SpinSysParamsType1(0.0) # level 2 model.

##### end inputs for surrogate.

println("Now on $(project_name)")

##### load data that was saved in store_exp.jl. See that script for details of the following commands.
load_path = joinpath(load_path)
dic = BSON.load(load_path)
s_t = dic[:s_t]
fs = dic[:fs]
SW = dic[:SW]
ν_0ppm = dic[:ν_0ppm]
α_0ppm = dic[:α_0ppm]
β_0ppm = dic[:β_0ppm]
λ_0ppm = dic[:λ_0ppm]

hz2ppmfunc = uu->(uu - ν_0ppm)*SW/fs
ppm2hzfunc = pp->(ν_0ppm + pp*fs/SW)

# start the first entry at the frequency corresponding to offset_Hz, which we set to 0.3 ppm.
offset_ppm = 0.3
offset_Hz = ν_0ppm - (ppm2hzfunc(offset_ppm)-ppm2hzfunc(0.0))
DFT_s = fft(s_t)
U_DFT, U_y, U_inds = NMRDataSetup.getwraparoundDFTfreqs(length(s_t), fs, offset_Hz)

S_U = DFT_s[U_inds]
P_y = hz2ppmfunc.(U_y)

val, ind = findmin( abs.(U_y .- ν_0ppm) )
Z = abs(S_U[ind])

y = S_U ./ Z

### end loading spectrum.
#@assert 1==2

####### compute surrogate.

# TODO document how to use the configuration files. then delete these default values from tutorial.
tol_coherence = 1e-2 # resonances are pairs of eigenvalues of the Hamiltonian that have quantum numbers that differ by -1. This is the tolerance away from -1 that is allowed.
α_relative_threshold = 0.05 # resonances with relative amplitude less than this factor compared to the maximum resonance in the spin group will be removed. Set to 0.0 to see every single resonance component.
Δc_partition_radius = 0.3 # determines how many resonances get grouped together. Larger number means less groups and thus more resonances per group.
λ0 = λ_0ppm

Δr_default = 1.0 # the samples used to build the surrogate is taken every `Δr` radian on the frequency axis. Decrease for improved accuracy at the expense of computation resources.
Δκ_λ_default = 0.05 # the samples used to build thes urrogate for κ_λ are taken at this sampling spacing. Decrease for improved accuracy at the expense of computation resources.
Δcs_max_scalar_default = 0.2 # In units of ppm. interpolation border that is added to the lowest and highest resonance frequency component of the mixture being simulated.
κ_λ_lb_default = 0.7 # interpolation lower limit for κ_λ.
κ_λ_ub_default = 1.5 # interpolation upper limit for κ_λ.

# get a surrogate where K_{n,i} is encouraged to be no larger than `early_exit_part_size`.
println("Timing: mag equivalence")
@time MEs = NMRSpectraSimulator.getmageqinfomixture(molecule_names,
    H_params_path,
    dict_compound_to_filename;
    unique_cs_atol = 1e-6)

println("Timing: setupmixtureproxies()")
@time mixture_params = NMRSpectraSimulator.setupmixtureSH(molecule_names,
    H_params_path, dict_compound_to_filename, fs, SW,
    ν_0ppm;
    config_path = SH_config_path,
    tol_coherence = tol_coherence,
    α_relative_threshold = α_relative_threshold,
    Δc_partition_radius = Δc_partition_radius)
As = mixture_params

ΩS_ppm = NMRSpectraSimulator.getPsnospininfo(As, hz2ppmfunc)
ΩS_ppm_sorted = sort(NMRSpectraSimulator.combinevectors(ΩS_ppm))

u_offset = 0.5 # in ppm.
u_min = ppm2hzfunc(ΩS_ppm_sorted[1] - u_offset)
u_max = ppm2hzfunc(ΩS_ppm_sorted[end] + u_offset)

println("Timing: fitproxies!()")
dummy_SSFID = NMRSpectraSimulator.SpinSysParamsType1(0.0)
Bs = NMRSpectraSimulator.fitproxies(As, dummy_SSFID, λ0;
    names = molecule_names,
    config_path = surrogate_config_path,
    Δcs_max_scalar_default = Δcs_max_scalar_default,
    κ_λ_lb_default = κ_λ_lb_default,
    κ_λ_ub_default = κ_λ_ub_default,
    u_min = u_min,
    u_max = u_max,
    Δr_default = Δr_default,
    Δκ_λ_default = Δκ_λ_default)

####### end mixture proxy.

### save.
BSON.bson(save_path,
As = As,
Bs = Bs,
surrogate_config_path = surrogate_config_path,
SH_config_path = SH_config_path,
u_min = u_min,
u_max = u_max,
y = y,
U_y = U_y,
SW = SW,
fs = fs,
ν_0ppm  = ν_0ppm,
λ0 = λ0,
molecule_names = molecule_names,
w = w)


#! format: off
########################### ElectricLoad ####################################

get_variable_multiplier(_, ::Type{<:PSY.ElectricLoad}, ::AbstractLoadFormulation) = -1.0

########################### ActivePowerVariable, ElectricLoad ####################################

get_variable_binary(::ActivePowerVariable, ::Type{<:PSY.ElectricLoad}, ::AbstractLoadFormulation) = false
get_variable_lower_bound(::ActivePowerVariable, d::PSY.ElectricLoad, ::AbstractLoadFormulation) = 0.0
get_variable_upper_bound(::ActivePowerVariable, d::PSY.ElectricLoad, ::AbstractLoadFormulation) = PSY.get_active_power(d)

########################### ReactivePowerVariable, ElectricLoad ####################################

get_variable_binary(::ReactivePowerVariable, ::Type{<:PSY.ElectricLoad}, ::AbstractLoadFormulation) = false

get_variable_lower_bound(::ReactivePowerVariable, d::PSY.ElectricLoad, ::AbstractLoadFormulation) = 0.0
get_variable_upper_bound(::ReactivePowerVariable, d::PSY.ElectricLoad, ::AbstractLoadFormulation) = PSY.get_reactive_power(d)

########################### ReactivePowerVariable, ElectricLoad ####################################

get_variable_binary(::OnVariable, ::Type{<:PSY.ElectricLoad}, ::AbstractLoadFormulation) = true

get_multiplier_value(::TimeSeriesParameter, d::PSY.ElectricLoad, ::StaticPowerLoad) = -1*PSY.get_max_active_power(d)
get_multiplier_value(::ReactivePowerTimeSeriesParameter, d::PSY.ElectricLoad, ::StaticPowerLoad) = -1*PSY.get_max_reactive_power(d)
get_multiplier_value(::TimeSeriesParameter, d::PSY.ElectricLoad, ::AbstractControllablePowerLoadFormulation) = PSY.get_max_active_power(d)

########################Objective Function##################################################
proportional_cost(cost::Nothing, ::OnVariable, ::PSY.ElectricLoad, ::AbstractControllablePowerLoadFormulation)=1.0
proportional_cost(cost::PSY.OperationalCost, ::OnVariable, ::PSY.ElectricLoad, ::AbstractControllablePowerLoadFormulation)=PSY.get_fixed(cost)

objective_function_multiplier(::VariableType, ::AbstractControllablePowerLoadFormulation)=OBJECTIVE_FUNCTION_NEGATIVE

variable_cost(::Nothing, ::PSY.ElectricLoad, ::ActivePowerVariable, ::AbstractControllablePowerLoadFormulation)=1.0
variable_cost(cost::PSY.OperationalCost, ::ActivePowerVariable, ::PSY.ElectricLoad, ::AbstractControllablePowerLoadFormulation)=PSY.get_variable(cost)

#! format: on

function get_default_time_series_names(
    ::Type{<:PSY.ElectricLoad},
    ::Type{<:Union{FixedOutput, AbstractLoadFormulation}},
)
    return Dict{Type{<:TimeSeriesParameter}, String}(
        ActivePowerTimeSeriesParameter => "max_active_power",
        ReactivePowerTimeSeriesParameter => "max_active_power",
    )
end

function get_default_attributes(
    ::Type{U},
    ::Type{V},
) where {U <: PSY.ElectricLoad, V <: Union{FixedOutput, AbstractLoadFormulation}}
    return Dict{String, Any}()
end

get_initial_conditions_device_model(
    ::OperationModel,
    ::DeviceModel{T, <:AbstractLoadFormulation},
) where {T <: PSY.ElectricLoad} = DeviceModel(T, StaticPowerLoad)

####################################### Reactive Power Constraints #########################
"""
Reactive Power Constraints on Controllable Loads Assume Constant power_factor
"""
function add_constraints!(
    container::OptimizationContainer,
    T::Type{<:ReactivePowerVariableLimitsConstraint},
    U::Type{<:ReactivePowerVariable},
    devices::IS.FlattenIteratorWrapper{V},
    ::DeviceModel{V, W},
    ::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.ElectricLoad, W <: AbstractControllablePowerLoadFormulation}
    time_steps = get_time_steps(container)
    constraint = add_constraints_container!(
        container,
        T(),
        V,
        [PSY.get_name(d) for d in devices],
        time_steps,
    )
    jump_model = get_jump_model(container)
    for t in time_steps, d in devices
        name = PSY.get_name(d)
        pf = sin(atan((PSY.get_max_reactive_power(d) / PSY.get_max_active_power(d))))
        reactive = get_variable(container, U(), V)[name, t]
        real = get_variable(container, ActivePowerVariable(), V)[name, t] * pf
        constraint[name, t] = JuMP.@constraint(jump_model, reactive == real)
    end
end

function add_constraints!(
    container::OptimizationContainer,
    T::Type{ActivePowerVariableLimitsConstraint},
    U::Type{<:VariableType},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    X::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.ControllableLoad, W <: DispatchablePowerLoad}
    add_parameterized_upper_bound_range_constraints(
        container,
        ActivePowerVariableTimeSeriesLimitsConstraint,
        U,
        ActivePowerTimeSeriesParameter,
        devices,
        model,
        X,
    )
    return
end

function add_constraints!(
    container::OptimizationContainer,
    T::Type{ActivePowerVariableLimitsConstraint},
    U::Type{<:VariableType},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    X::Type{<:PM.AbstractPowerModel},
) where {V <: PSY.ControllableLoad, W <: InterruptiblePowerLoad}
    add_parameterized_upper_bound_range_constraints(
        container,
        ActivePowerVariableTimeSeriesLimitsConstraint,
        U,
        ActivePowerTimeSeriesParameter,
        devices,
        model,
        X,
    )
    return
end

############################## FormulationControllable Load Cost ###########################
function objective_function!(
    container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    ::DeviceModel{T, U},
    ::Type{<:PM.AbstractPowerModel},
) where {T <: PSY.ControllableLoad, U <: DispatchablePowerLoad}
    add_variable_cost!(container, ActivePowerVariable(), devices, U())
    return
end

function objective_function!(
    container::OptimizationContainer,
    devices::IS.FlattenIteratorWrapper{T},
    ::DeviceModel{T, U},
    ::Type{<:PM.AbstractPowerModel},
) where {T <: PSY.ControllableLoad, U <: InterruptiblePowerLoad}
    add_proportional_cost!(container, OnVariable(), devices, U())
    return
end


module Simulations

Base.Experimental.@optlevel 3

export AbstractSimulation, CTSimulation
export simulate

using ..Monads
using ..CTImages
using ..Utils
using ProgressMeter, LinearAlgebra, IntervalSets
import Random, Distributions, StatsBase

abstract type AbstractSimulation end

function simulate end
function simulate_blocks end
function simulate_single end

include("simulations/sinogram_sampling.jl")

@inline function Random.rand(
    rng::Random.AbstractRNG,
    s::Random.SamplerTrivial{ORI{T}},
) where {T <: AbstractFloat}
    a, b = endpoints(s[])
    t = rand(rng, T)
    a * (one(T) - t) + b * t
end

end # module


module TrueRandom

using Random
import Random: rand, rand!, _rand52

import Base.show

using Base.Threads

using HTTP
using JSON3

include("UInt16RNG.jl")

export ANUQRNG

include("ANUQRNG.jl")

end


soma(x,y) = x+y

multiplicação(x,y) = x*y


__precompile__(true)

module SequenceAlignment
export levenshteinAlign
export levenshteinDistance
export ldn

include("levenshtein.jl")

end


module TestModelSearch

using Test
using MLJModels
import MLJBase.info

pca = info("PCA", pkg="MultivariateStats")
cnst = info("ConstantRegressor", pkg="MLJModels")

@test_throws ArgumentError info("Julia")

@info(ConstantRegressor) == cnst
@test info(Standardizer()) == info("Standardizer", pkg="MLJModels")

@testset "localmodels" begin
    tree = info("DecisionTreeRegressor")
    @test cnst in localmodels(modl=TestModelSearch)
    @test !(tree in localmodels(modl=TestModelSearch))
    import MLJDecisionTreeInterface.DecisionTreeRegressor
    @test tree in localmodels(modl=TestModelSearch)
end

@testset "models() and localmodels()" begin
    t(model) = model.is_pure_julia
    mods = models(t)
    @test pca in mods
    @test cnst in mods
    @test !(info("SVC") in mods)
    mods = localmodels(t, modl=TestModelSearch)
    @test cnst in mods
    @test !(pca in mods)
    u(model) = !(model.is_supervised)
    @test pca in models(u, t)
    @test !(cnst in models(u, t))
end

@testset "models(needle::Union{AbstractString,Regex}) and localmodels(needle::Union{AbstractString,Regex})" begin
    @test pca in models("PCA")
    @test pca ∉ models("PCA′")

    @test pca in models(r"PCA")
    @test pca in models(r"pca"i)
    @test pca ∉ models(r"PCA′")

    info("DecisionTreeRegressor") in localmodels("Decision"; modl = TestModelSearch)
    info("DecisionTreeRegressor") in localmodels(r"Decision"; modl = TestModelSearch)
end

end
true


const onevec = 1:1
const emptyvec = 1:0
@inline inbounds(itp::AbstractInterpolation, x...) = _inbounds.(bounds(itp), x)
_inbounds((l,u)::Tuple{Number,Number}, x::Number) = ifelse(l <= x <= u, onevec, emptyvec)
function _inbounds((l,u)::Tuple{Number,Number}, x::AbstractVector)
    ret = Int[]
    for i in eachindex(x)
        l <= x[i] <= u && push!(ret, i)
    end
    ret
end

function getindex!(dest, itp, xs...)
    for (i, x) in zip(eachindex(dest), Iterators.product(xs...))
        dest[i] = itp(x...)
    end
    return dest
end


module GEOTRACES

using OceanographyCruises
using NCDatasets
using Unitful
using Dates
using Match
using Measurements
using DataFrames

const metadatakeys=("lat", "lon", "depth", "cruise", "station", "date") # default

GEOTRACES_IDP21_DiscreteSamples_path() = get(ENV, "GEOTRACES_IDP21_PATH", joinpath(homedir(), "Data/GEOTRACES/GEOTRACES_IDP2021_v1/seawater/netcdf/GEOTRACES_IDP2021_Seawater_Discrete_Sample_Data_v1/GEOTRACES_IDP2021_Seawater_Discrete_Sample_Data_v1.nc"))

include("helper_functions.jl")

"""
    cruisetrack(cruise_name)

Construct GEOTRACES `CruiseTrack` from cruise name.
"""
function cruisetrack(ds::Dataset, cruise)
    ikeep = findall(list_of_cruises(ds) .== cruise)
    ikeep = CartesianIndices((1, length(ds[varname("lat")])))[ikeep]
    metadatakeys=("lat","lon","cruise","station","date")
    GEOTRACESmetadatakeys = varname.(metadatakeys)
    metadata = [metadatakeyvaluepair(ds[k], ikeep) for k in GEOTRACESmetadatakeys]
    m = (; metadata...)
    any(m.cruise .≠ cruise) && error("Not the right cruise!")
    stations = [Station(date=d,lat=y,lon=x,name=string(s)) for (d,y,x,s) in zip(m.date, m.lat, m.lon, m.station)]
    return CruiseTrack(name=cruise, stations=stations)
end

"""
    list_of_cruises()

List of GEOTRACES cruises.
"""
list_of_cruises(ds::Dataset) = reduce.(string, filter.(!=('\0'), eachcol(ds["metavar1"].var[:,:])))

"""
    list_of_stations()

List of GEOTRACES stations.
"""
list_of_stations(ds::Dataset) = string.(reduce.(string, filter.(!=('\0'), eachcol(ds["metavar2"].var[:,:]))))

"""
transect(tracer::String, cruise::String)

The `Transect` of observations of tracer `tracer` along cruise `cruise`.
"""
function transect(ds::Dataset, tracer::String; cruise::String, QCmax=1)
    ts = transects(ds, tracer, QCmax=QCmax)
    i = findall(ts.cruises .== cruise)
    length(i) == 1 ? i = i[1] : error("Multiple identical cruises is an issue!")
    return ts.transects[i]
end

"""
transects(tracer::String, cruise::String)

The `Transects` of observations of tracer `tracer`.
"""
function transects(obs::DataFrame, tracer::String)
    ts = Transect{eltype(getproperty(obs, tracer))}[]
    df_bycruise = groupby(obs, :cruise)
    for (cruise, cruise_df) in pairs(df_bycruise)
        df_byprofile = groupby(cruise_df, [:lat, :lon, :station, :date])
        profiles = [DepthProfile(station = Station(name=string(k.station), lat=k.lat, lon=k.lon, date=k.date),
                                 depths=ustrip.(v.depth),
                                 values=v[:,tracer]) for (k,v) in pairs(df_byprofile)]
        push!(ts, Transect(tracer=tracer, cruise=cruise.cruise, profiles=profiles))
    end
    return Transects(tracer=tracer, cruises=[k.cruise for k in keys(df_bycruise)], transects=ts)
end
function transects(ds::Dataset, tracer::String; QCmax=1)
    obs = observations(ds, tracer; metadatakeys, QCmax)
    return transects(obs, tracer)
end

#function cruisetrack(ds::Dataset, cruise::String)
#    Istations = findall(list_of_cruises(ds) .== cruise)
#    stations_names = list_of_stations(ds)
#    stations = Station[]
#
#    for istation in Istations
#        station = Station(name=stations_names[istation], lat=ds["latitude"][istation], lon=ds["longitude"][istation])
#        push!(stations, station)
#    end
#
#    return CruiseTrack(name=cruise, stations=stations)
#end

"""
    observations(tracer1, tracer2, tracer3, ...)

Returns the GEOTRACES observations of the given tracers.

### Example usage

```
x, y, z, ... = observations(tracer1, tracer2, tracer3, ...)
x = observations(tracer1)
```

### Quality control flags

By default, only observations with quality control flag of `1` are kept.
If you want more data, you must define keyword-argument `QCmax`.

```
x, y, z, ... = observations(tracer1, tracer2, tracer3, ...; QCmax = 2)
x = observations(tracer1; QCmax = 3)
```
"""
function observations(ds::Dataset, tracers::String...; metadatakeys=metadatakeys, QCmax=1)
    dfs = (observations(ds, tracer; metadatakeys, QCmax) for tracer in tracers)
    return innerjoin(dfs..., on=collect(metadatakeys))
end
function observations(ds::Dataset, tracer::String; metadatakeys=metadatakeys, QCmax=1)
    f = unitfunction(ds[varname(tracer)].attrib["units"])
    var = ds[varname(tracer)]
    v = var[:]
    qcv = ds[qcvarname(tracer)].var[:] .|> Char
    ikeep = findall(parse.(Int, qcv) .≤ QCmax)
    ikeep = CartesianIndices(size(qcv))[ikeep]
    GEOTRACESmetadatakeys = varname.(metadatakeys)
    metadata = (metadatakeyvaluepair(ds[k], ikeep) for k in GEOTRACESmetadatakeys)
    df = DataFrame(metadata..., Symbol(tracer)=>f.(float.(view(v, ikeep))))
    #mdf = attach_metadata(df, (GEOTRACESvarname=name(var),))
    return df
end

"""
    qualitycontrols(tracer1, tracer2, tracer3, ...)

Returns the GEOTRACES quality control flag of the given tracers.

GEOTRACES's `Char`s are converted to `Int` by this function.

From GEOTRACES:
```
1 = good quality
2 = not evaluated, not available or unknown quality
3 = questionable/suspect quality
4 = bad quality
9 = missing data
```

Note that there should be no `9` because `missing` data is skipped
by the `observations` function.
"""
function qualitycontrols(ds::Dataset, tracers::String...)
    vs = ((ds[varname(tracer)][:] for tracer in tracers)...,)
    ikeep = findall(i -> !any(ismissing.(getindex.(vs, i))), eachindex(vs[1]))
    qcvs = ((ds[qcvarname(tracer)].var[:] .|> Char for tracer in tracers)...,)
    return ((parse.(Int, view(qcv, ikeep)) for qcv in qcvs)...,)
end
function qualitycontrols(ds::Dataset, tracer::String)
    v = ds[varname(tracer)][:]
    ikeep = findall(!ismissing, v)
    qcv = ds[qcvarname(tracer)].var[:] .|> Char
    return parse.(Int, view(qcv, ikeep))
end

"""
    metadata(tracer1, tracer2, tracer3, ...)

Returns the GEOTRACES metadata for given tracers.
"""
function metadata(ds::Dataset, tracers::String...; metadatakeys=("lat", "lon", "depth"), QCmax=1)
    qcvs = ((ds[qcvarname(tracer)].var[:] .|> Char for tracer in tracers)...,)
    ikeep = findall(i -> all((parse.(Int, getindex.(qcvs, i))) .≤ QCmax), eachindex(qcvs[1]))
    ikeep = CartesianIndices(size(qcvs[1]))[ikeep]
    GEOTRACESmetadatakeys = varname.(metadatakeys)
    metadata = [metadatakeyvaluepair(ds[k], ikeep) for k in GEOTRACESmetadatakeys]
    namedmetadata = (; metadata...)
    return namedmetadata
end

# TODO treatment of uncertainty
# Not all variable have a "standard deviation" in GEOTRACES
# But I could create one from
# - the number of significant digits of each tracer per cruise?
# - the STD along each cruise?
# - the total STD?
function standarddeviations(ds::Dataset, tracer::String)
    _, s = observations(ds, varname(tracer), stdvarname(tracer))
    return s
end

function standarddeviations(ds::Dataset, tracers::String...)
    x_and_s = observations(ds, varname.(tracers)..., stdvarname.(tracers)...)
    return x_and_s[length(tracers)+1:end]
end

function observations_with_std(ds::Dataset, tracer::String)
    x, s = observations(ds, varname(tracer), stdvarname(tracer))
    u = unit(x[1])
    return @. (ustrip(x) ± ustrip(s)) * u
end
function observations_with_std(ds::Dataset, tracers::String...)
    x_and_s = observations(ds, varname.(tracers)..., stdvarname.(tracers)...)
    n = length(tracers)
    us = ((unit(x[1]) for x in x_and_s[1:n])...,)
    return (((ustrip.(x_and_s[i]) .± ustrip.(x_and_s[i+n])) * us[i] for i in 1:n)...,)
end

# open and close ds if not provided
for f in [:observations, :metadata, :transect, :transects, :list_of_cruises,
         :list_of_stations, :variable, :qualitycontrols, :cruisetrack,
         :standarddeviations, :matchingvariables, :observations_with_std]
    @eval begin
        $f(args...; kwargs...) =
        Dataset(GEOTRACES_IDP21_DiscreteSamples_path(), "r") do ds
             $f(ds, args...; kwargs...)
        end
    end
end

end # module


#################################
"""
    replaceunit!(fgmax::VisClaw.FGmax, unit::Symbol)
    replaceunit!(gauge::VisClaw.Gauge, unit::Symbol)
    replaceunit!(amrs::VisClaw.AMR, unit::Symbol)
    replaceunit!(track::VisClaw.Track, unit::Symbol)

Time unit converter
The second arg unit must be any one of :second, :minute, :hour, or :day.
"""
function replaceunit!(fgmax::VisClaw.FGmax, unit::Symbol)

    !haskey(timedict, unit) && error("Invalid specification of unit")
    !haskey(timedict, fgmax.unittime) && error("Invalid symbol in fgmax")

    ratio = timedict[fgmax.unittime]/timedict[unit]
    if abs(ratio - 1.0) < 1e-5; return fgmax; end

    # convert
    fgmax.unittime = unit

    fgmax.tD = ratio.*fgmax.tD
    fgmax.tarrival = ratio.*fgmax.tarrival
    !isempty(fgmax.tv) && (fgmax.tv = ratio.*fgmax.tv)
    if !isempty(fgmax.tM)
        fgmax.tM = ratio.*fgmax.tM
        fgmax.tMflux = ratio.*fgmax.tMflux
        fgmax.tDmin = ratio.*fgmax.tDmin
    end
    return fgmax
end
#################################

#################################
function replaceunit!(gauge::VisClaw.Gauge, unit::Symbol)
    !haskey(timedict, unit) && error("Invalid specification: unit")
    !haskey(timedict, gauge.unittime) && error("Invalid symbol in gauge")

    ratio = timedict[gauge.unittime]/timedict[unit]
    if abs(ratio - 1.0) < 1e-5; return gauge; end

    # convert
    gauge.unittime = unit
    gauge.time = ratio.*gauge.time
    # return value
    return gauge
end
#################################

#################################
function replaceunit!(amrs::VisClaw.AMR, unit::Symbol)
    !haskey(timedict, unit) && error("Invalid specification: unit")
    !haskey(timedict, amrs.unittime) && error("Invalid symbol in AMR")

    ratio = timedict[amrs.unittime]/timedict[unit]
    if abs(ratio - 1.0) < 1e-5; return amrs; end

    # convert
    amrs.unittime = unit
    amrs.timelap = ratio.*amrs.timelap
    # return value
    return amrs
end
#################################

#################################
function replaceunit!(track::VisClaw.Track, unit::Symbol)
    !haskey(timedict, unit) && error("Invalid specification: unit")
    !haskey(timedict, track.unittime) && error("Invalid symbol in AMR")

    ratio = timedict[track.unittime]/timedict[unit]
    if abs(ratio - 1.0) < 1e-5; return track; end

    # convert
    track.unittime = unit
    track.timelap = ratio.*track.timelap
    # return value
    return track
end
#################################


testdims(L::Integer) = ntuple(i -> 3 + i, Val(L))

randlike(x::Number) = rand(typeof(x))
randlike(A::AbstractArray{<:Number}) = rand(eltype(A), size(A)...)
randlike(A::AbstractArray{<:Number,0}) = rand!(zeros(eltype(A)))

function randlike(A::AbstractArray{<:AbstractArray{V,M},N}) where {V,M,N}
    B = Array{Array{V,M},N}(undef, size(A)...)
    for I in eachindex(B)
        B[I] = randlike(A[I])
    end
    return B
end

function randlike!(A::AbstractArray{<:AbstractArray})
    for I in eachindex(A)
        copyto!(A[I], randlike(A[I]))
    end
    return A
end

function randlike!(A::AbstractArray{<:Number})
    for I in eachindex(A)
        A[I] = randlike(A[I])
    end
    return A
end

macro qe(ex)
    Expr(:quote, MacroTools.striplines(esc(ex)))
end

macro test_inferred(ex)
    @qe begin
        $Test.@test (($Test.@inferred $ex); true)
    end
end

macro test_noalloc(ex)
    @qe begin
        local nbytes = $memory($BenchmarkTools.@benchmark $ex samples = 1 evals = 1)
        $iszero(nbytes) ? true : $error("Allocated $(BenchmarkTools.prettymemory(nbytes))")
    end
end

macro test_getindex(A, I)
    @qe begin
        if $I isa $Tuple
            @test_inferred $A[$I...]
            if $index_dimsum($I...) == ()
                $Test.@test $typeof($A[$I...]) == $eltype($A)
            else
                $Test.@test $size($A[$I...]) == $map($length, $index_shape($to_indices($A, $I)...))
                $Test.@test $eltype($A[$I...]) == $eltype($A)
            end
        else
            @test_inferred $A[$I]
            if $index_dimsum($I) == () && !($I isa AbstractArray{<:Any,0})
                $Test.@test $typeof($A[$I]) == $eltype($A)
            else
                $Test.@test $size($A[$I]) == $map($length, $index_shape($to_indices($A, ($I,))...))
                $Test.@test $eltype($A[$I]) == $eltype($A)
            end
        end
    end
end

macro test_setindex!(A, I)
    @qe begin
        if $I isa $Tuple
            local x = $randlike($A[$I...])
            @test_inferred $setindex!($A, x, $I...)
            $Test.@test $setindex!($A, x, $I...) === $A
            $Test.@test $A[$I...] == x
        else
            local x = $randlike($A[$I])
            @test_inferred $setindex!($A, x, $I)
            $Test.@test $setindex!($A, x, $I) === $A
            $Test.@test $A[$I] == x
        end
    end
end

macro test_getindices(A, Is)
    @qe begin
        $Test.@test $all($Is) do I
            if I isa $Tuple
                if $index_dimsum(I...) == ()
                    $typeof($A[I...]) == $eltype($A)
                else
                    $eltype($A[I...]) == $eltype($A) &&
                    $size($A[I...]) == $map($length, $index_shape($to_indices($A, I)...))
                end
            else
                #if $index_dimsum(I) == ()
                if $index_dimsum(I) == () && !(I isa AbstractArray{<:Any,0})
                    $typeof($A[I]) == $eltype($A)
                else
                    $eltype($A[I]) == $eltype($A) &&
                    $size($A[I]) == $map($length, $index_shape($to_indices($A, (I,))...))
                end
            end
        end
    end
end

macro test_setindices!(A, Is)
    @qe begin
        $Test.@test $all($Is) do I
            if I isa $Tuple
                x = $randlike($A[I...])
                $A[I...] != x && $setindex!($A, x, I...) === $A && $A[I...] == x
            else
                x = $randlike($A[I])
                $A[I] != x && $setindex!($A, x, I) === $A && $A[I] == x
            end
        end
    end
end

macro test_all_equal(A, B, Is)
    @qe begin
        $Test.@test all($Is) do I
            if I isa Tuple
                $A[I...] == $B[I...] #&& typeof($A[I...]) == typeof($B[I...])
            else
                $A[I] == $B[I] #&& typeof($A[I]) == typeof($B[I])
            end
        end
    end
end

macro test_copyto!(dest, src)
    @qe begin
        $randlike!($dest)
        @test_inferred $copyto!($dest, $src)
        $Test.@test $copyto!($dest, $src) === $dest
        $Test.@test $all(I -> $dest[I] == $src[I], $eachindex($dest, $src))
        $Test.@test $dest == $src

        if $length($dest) > 2
            $randlike!($dest)
            @test_inferred $copyto!($dest, 2, $src, 2, 1)
            $Test.@test $copyto!($dest, 2, $src, 2, 1) === $dest
            $Test.@test $dest[2] == $src[2] && $dest[1] != $src[1] && $dest[3:end] != $src[3:end]
        else
            $Test.@test_skip @test_inferred $copyto!($dest, 2, $src, 2, 1)
            $Test.@test_skip $copyto!($dest, 2, $src, 2, 1) === $dest
            $Test.@test_skip $dest[2] == $src[2] &&
                             $dest[1] != $src[1] &&
                             $dest[3:end] != $src[3:end]
        end
    end
end

macro test_array_attributes(A)
    @qe begin
        $Test.@test $eltype($A) === $typeof($first($A))
        $Test.@test $ndims($A) == $length($axes($A))
        $Test.@test $size($A) == $map($length, $axes($A))
        $Test.@test $length($A) == $prod($size($A))

        $Test.@test $all(1:$ndims($A)) do dim
            $size($A)[dim] == $size($A, dim)
        end
        $Test.@test $size($A, $ndims($A) + 1) == 1

        $Test.@test $all(1:$ndims($A)) do dim
            $axes($A)[dim] == $axes($A, dim)
        end
        $Test.@test $axes($A, $ndims($A) + 1) == $Base.OneTo(1)

        @test_inferred $eltype($A)
        @test_inferred $ndims($A)
        @test_inferred $axes($A)
        @test_inferred $size($A)
        @test_inferred $length($A)
    end
end

macro test_similar(A, T, dims)
    @qe begin
        @test_inferred $similar($A)
        $Test.@test $eltype($similar($A)) === $eltype($A)
        $Test.@test $size($similar($A)) == $size($A)

        @test_inferred $similar($A, $T)
        $Test.@test $eltype($similar($A, $T)) === $T
        $Test.@test $size($similar($A, $T)) == $size($A)

        @test_inferred $similar($A, $dims)
        $Test.@test eltype($similar($A, $dims)) === $eltype($A)
        $Test.@test size($similar($A, $dims)) == $dims

        @test_inferred $similar($A, $T, $dims)
        $Test.@test eltype($similar($A, $T, $dims)) === $T
        $Test.@test size($similar($A, $T, $dims)) == $dims
    end
end

function test_indexing_AB(AB::Function)
    @testset "LinearIndices" begin
        let (A, _) = AB()
            @test_getindex A first(LinearIndices(A))
            @test_setindex! A first(LinearIndices(A))
        end
        let (A, B) = AB()
            @test LinearIndices(A) === LinearIndices(B)
            @test_all_equal A B LinearIndices(A)
            @test_getindices A LinearIndices(A)
            @test_setindices! A LinearIndices(A)
        end
        let (A, B) = AB()
            @test A[LinearIndices(A)] == B[LinearIndices(A)]
            @test_getindex A LinearIndices(A)
            @test_setindex! A LinearIndices(A)
        end
        let (A, B) = AB()
            @test begin
                I = hcat(LinearIndices(A), LinearIndices(A))
                A[I] == B[I]
            end
            @test_getindex A hcat(LinearIndices(A), LinearIndices(A))
            #@test_setindex! A hcat(LinearIndices(A), LinearIndices(A))
        end
    end
    @testset "CartesianIndices" begin
        let (A, _) = AB()
            @test_getindex A first(CartesianIndices(A))
            @test_setindex! A first(CartesianIndices(A))
        end
        let (A, B) = AB()
            @test CartesianIndices(A) === CartesianIndices(B)
            @test_all_equal A B CartesianIndices(A)
            @test_getindices A CartesianIndices(A)
            @test_setindices! A CartesianIndices(A)
        end
        let (A, B) = AB()
            @test A[CartesianIndices(A)] == B[CartesianIndices(A)]
            @test_getindex A CartesianIndices(A)
            @test_setindex! A CartesianIndices(A)
        end
        let (A, B) = AB()
            @test begin
                I = hcat(CartesianIndices(A), CartesianIndices(A))
                A[I] == B[I]
            end
            @test_getindex A hcat(CartesianIndices(A), CartesianIndices(A))
            #@test_setindex! A CartesianIndices(A)
        end
    end
    @testset "trailing singleton" begin
        let (A, _) = AB()
            @test_getindex A (Tuple(first(CartesianIndices(A)))..., 1)
            @test_setindex! A (Tuple(first(CartesianIndices(A)))..., 1)
        end
        let (A, B) = AB()
            @test_all_equal A B [(Tuple(I)..., 1) for I in CartesianIndices(A)]
            @test_getindices A [(Tuple(I)..., 1) for I in CartesianIndices(A)]
            @test_setindices! A [(Tuple(I)..., 1) for I in CartesianIndices(A)]
        end
    end
    @testset "trailing colon" begin
        let (A, _) = AB()
            @test_getindex A (Tuple(first(CartesianIndices(A)))..., :)
            @test_setindex! A (Tuple(first(CartesianIndices(A)))..., :)
        end
        let (A, B) = AB()
            @test_all_equal A B [(Tuple(I)..., :) for I in CartesianIndices(A)]
            @test_getindices A [(Tuple(I)..., :) for I in CartesianIndices(A)]
            @test_setindices! A [(Tuple(I)..., :) for I in CartesianIndices(A)]
        end
    end
    @testset "dropped singleton" begin
        _, B = AB()
        if ndims(B) > 0 && size(B, ndims(B)) == 1
            let (A, _) = AB()
                @test_getindex A front(Tuple(first(CartesianIndices(A))))
                @test_setindex! A front(Tuple(first(CartesianIndices(A))))
            end
            let (A, B) = AB()
                @test_all_equal A B [front(Tuple(I)) for I in CartesianIndices(A)]
                @test_getindices A [front(Tuple(I)) for I in CartesianIndices(A)]
                @test_setindices! A [front(Tuple(I)) for I in CartesianIndices(A)]
            end
        end
    end
    @testset "logical" begin
        let (A, B) = AB()
            @test begin
                I = map(isodd, LinearIndices(A))
                A[I] == B[I]
            end
            @test_getindex A map(isodd, LinearIndices(A))
            @test_setindex! A map(isodd, LinearIndices(A))
        end
        let (A, B) = AB()
            @test begin
                I = Tuple(map(isodd, ax) for ax in axes(A))
                A[I...] == B[I...]
            end
            @test_getindex A Tuple(map(isodd, ax) for ax in axes(A))
            @test_setindex! A Tuple(map(isodd, ax) for ax in axes(A))
        end
    end
    @testset "single colon" begin
        A, B = AB()
        @test A[:] == B[:]
    end
end

function test_array_AB(AB::Function)
    @testset "array attributes" begin
        A, _ = AB()
        @test_array_attributes A

        @test_noalloc eltype($A)
        @test_noalloc ndims($A)
        @test_noalloc axes($A)
        @test_noalloc size($A)
        @test_noalloc length($A)
    end

    @testset "indexing" begin
        test_indexing_AB(AB)
    end

    @testset "copy/copyto!/equality" begin
        let (A1, _) = AB()
            A2 = copy(A1)
            @test A1 == A2
            randlike!(A1)
            @test A1 != A2
        end
        let (A, B) = AB()
            @test_copyto! A B
        end
        let (A, B) = AB()
            @test_copyto! B A
        end
        let A1 = first(AB()), A2 = first(AB())
            @test_copyto! A1 A2
        end
    end
end