using NNlib: DenseConvDims

# descriptor

mutable struct ConvDesc
    ptr::cudnnConvolutionDescriptor_t
end

unsafe_free!(cd::ConvDesc) = cudnnDestroyConvolutionDescriptor(cd.ptr)

Base.unsafe_convert(::Type{cudnnConvolutionDescriptor_t}, cd::ConvDesc)=cd.ptr

function cdsize(w, nd)
    isa(w, Integer) && return Cint[fill(w,nd)...]
    length(w) == nd && return Cint[reverse(w)...]
    length(w) == 2*nd && return Cint[reverse(w[nd+1:end])...]
    throw(DimensionMismatch())
end

pdsize(w, nd)=Cint[reverse(psize(w,nd))...]
function psize(w, nd)
    isa(w, Integer) && return Cint[fill(w,nd)...]
    length(w) == nd && return w
    length(w) == 2*nd && return w[1:nd]
    throw(DimensionMismatch())
end

Base.cconvert(::Type{cudnnConvolutionMode_t}, x::Bool) = x ? CUDNN_CROSS_CORRELATION : CUDNN_CONVOLUTION

function ConvDesc(T, N, padding, stride, dilation, mode)
    cd = Ref{cudnnConvolutionDescriptor_t}()
    cudnnCreateConvolutionDescriptor(cd)
    if version() >= v"4"
        cudnnSetConvolutionNdDescriptor(cd[],N,cdsize(padding,N),cdsize(stride,N),cdsize(dilation,N),mode,cudnnDataType(T))
    elseif version() >= v"3"
        cudnnSetConvolutionNdDescriptor_v3(cd[],N,cdsize(padding,N),cdsize(stride,N),cdsize(dilation,N),mode,cudnnDataType(T))
    else
        cudnnSetConvolutionNdDescriptor(cd[],N,cdsize(padding,N),cdsize(stride,N),cdsize(dilation,N),mode)
    end
    this = ConvDesc(cd[])
    finalizer(unsafe_free!, this)
    return this
end

function ConvDesc(T, cdims::DenseConvDims)
    pd = NNlib.padding(cdims)
    if !all(pd[1:2:end] .== pd[2:2:end])
        @warn("CuDNN does not support asymmetric padding; defaulting to symmetric choice")
    end
    return ConvDesc(T, NNlib.spatial_dims(cdims), pd[1:2:end], NNlib.stride(cdims),
                       NNlib.dilation(cdims), NNlib.flipkernel(cdims))
end

# wrappers

function cudnnConvolutionBiasActivationForward(y::CuArray{T,N}, x::CuArray{T,N}, w::CuArray{T,N}, bias::CuArray{T,N};
                                               alpha1=1, workspace=CU_NULL, workspace_size=0,
                                               algo=0, alpha2=0, padding=0, stride=1, dilation=1, mode=0,
                                               activationMode=CUDNN_ACTIVATION_IDENTITY, activationCoeff=0.0,
                                               activationReluNanOpt=CUDNN_NOT_PROPAGATE_NAN) where {T,N}
    cd = ConvDesc(T, N-2, padding, stride, dilation, mode)
    ad = ActivationDesc(activationMode, T(activationCoeff), activationReluNanOpt)
    cudnnConvolutionBiasActivationForward(handle(), Ref(T(alpha1)),TensorDesc(x),x,FilterDesc(w),w,cd,cudnnConvolutionFwdAlgo_t(algo),workspace,
        workspace_size,Ref(T(alpha2)),TensorDesc(bias),bias,ad,TensorDesc(y),y)
    return y
end

function cudnnConvolutionForward(y::CuArray{T,N}, x::CuArray{T,N}, w::CuArray{T,N},
                                 cdims::DenseConvDims; algo=0, alpha=1, beta=0) where {T,N}
    @workspace size=@argout(
            cudnnGetConvolutionForwardWorkspaceSize(
                handle(), TensorDesc(x),
                FilterDesc(w), ConvDesc(T, cdims),
                TensorDesc(y),
                cudnnConvolutionFwdAlgo_t(algo),
                out(Ref{Csize_t}()))
        )[] workspace->begin
            cudnnConvolutionForward(
                handle(), Ref(T(alpha)), TensorDesc(x), x, FilterDesc(w), w,
                ConvDesc(T,cdims), cudnnConvolutionFwdAlgo_t(algo), workspace,
                sizeof(workspace), Ref(T(beta)), TensorDesc(y), y)
        end
    return y
end

function cudnnConvolutionBackwardData(dx::CuArray{T,N}, w::CuArray{T,N}, dy::CuArray{T,N},
                                      cdims::DenseConvDims; algo=0, alpha=1, beta=0) where {T,N}
    @workspace size=@argout(
            cudnnGetConvolutionBackwardDataWorkspaceSize(
                handle(), FilterDesc(w),
                TensorDesc(dy), ConvDesc(T, cdims), TensorDesc(dx),
                cudnnConvolutionBwdDataAlgo_t(algo),
                out(Ref{Csize_t}()))
        )[] workspace->begin
            cudnnConvolutionBackwardData(
                handle(), Ref(T(alpha)), FilterDesc(w), w,
                TensorDesc(dy), dy, ConvDesc(T, cdims),
                cudnnConvolutionBwdDataAlgo_t(algo),
                workspace, sizeof(workspace),
                Ref(T(beta)), TensorDesc(dx), dx)
        end
    return dx
end

function cudnnConvolutionBackwardFilter(dw::CuArray{T,N}, x::CuArray{T,N}, dy::CuArray{T,N},
                                        cdims::DenseConvDims; algo=0, alpha=1, beta=0) where {T,N}
    @workspace size=@argout(
            cudnnGetConvolutionBackwardFilterWorkspaceSize(
                handle(), TensorDesc(x),
                TensorDesc(dy),
                ConvDesc(T, cdims),
                FilterDesc(dw),
                cudnnConvolutionBwdFilterAlgo_t(algo),
                out(Ref{Csize_t}()))
        )[] workspace->begin
            cudnnConvolutionBackwardFilter(
                handle(), Ref(T(alpha)), TensorDesc(x), x,
                TensorDesc(dy), dy, ConvDesc(T, cdims),
                cudnnConvolutionBwdFilterAlgo_t(algo), workspace,
                sizeof(workspace), Ref(T(beta)), FilterDesc(dw), dw)
        end
    return dw
end

function cudnnConvolutionBackwardBias(db::CuArray{T,N}, dy::CuArray{T,N}; alpha=1, beta=0) where {T,N}
    cudnnConvolutionBackwardBias(handle(), Ref(T(alpha)), TensorDesc(dy), dy, Ref(T(beta)), TensorDesc(db), db)
    return db
end


# This is too much work to manually implement.
# Let's resort to using the existing logic for plus.

Add = SU.Add
add_t = SU.add_t
makeadd = SU.makeadd
promote_symtype = SU.promote_symtype

add_t(f, a::SSet, b::SSet) = SSet

function union(a::SSet, b::SSet)
    if a isa Add
        coeff, dict = makeadd(1, 0, b)
        @show symtype(a)
        T = promote_symtype(union, symtype(a), symtype(b))
        @show T
        return Add(add_t(a, b), union(a.coeff, coeff), _merge(union, a.dict, dict, filter=_iszero))
    elseif b isa Add
        return union(b, a)
    end
    Add(SSet, makeadd(1, 0, a, b)...)
end

union(a::SSet) = a

union(a::Add, b::Add) = Add(add_t(a,b),
                        union(a.coeff, b.coeff),
                        _merge(union, a.dict, b.dict, filter=_iszero))


import Mosek

function run_rmp_mosek(frmp)
    task = Mosek.maketask()
    Mosek.putstreamfunc(task, Mosek.MSK_STREAM_LOG, print)

    Mosek.readdata(task, frmp)

    # Set parameters
    Mosek.putintparam(task, Mosek.MSK_IPAR_NUM_THREADS, 1)
    Mosek.putintparam(task, Mosek.MSK_IPAR_PRESOLVE_USE, 0)
    Mosek.putintparam(task, Mosek.MSK_IPAR_INTPNT_BASIS, 0)

    # Solve
    t = @elapsed Mosek.optimize(task)

    println("Solve time: $t")
    println("Barrier iterations: ", Mosek.getintinf(task, Mosek.MSK_IINF_INTPNT_ITER))

    return nothing
end

if abspath(PROGRAM_FILE) == @__FILE__
    run_rmp_mosek(ARGS[1])
end


using OrthographicProjection
using FileManager
using Common
# using Visualization

txtpotreedirs = raw"C:\Users\marte\Downloads\point-cloud-private" #"C:/Users/marte/Documents/GEOWEB/wrapper_file/directory.txt"
potreedirs = OrthographicProjection.get_potree_dirs(txtpotreedirs)
metadata = CloudMetadata(potreedirs[1])
output = raw"C:\Users\marte\Downloads\point-cloud-private.las"
bbin = metadata.tightBoundingBox
#bbin = "C:/Users/marte/Documents/GEOWEB/wrapper_file/JSON/volume_COLOMBELLA.json"
# bbin = Volume([200., 200., 200.], [458309.223, 4493974.624, 199.400], [0.000, 0.000, 0.000])
bbin = Volume([37.332, 31.454, 30.835],[295485.136, 4781263.156, 285.114],[0.000, 0.000, -0.685])
model = getmodel(bbin)
epsg = nothing #2049
@time OrthographicProjection.segment(txtpotreedirs, output, model; epsg = epsg)


"""
This function creates a PowerSystems ThermalStandard unit.
"""
function create_PSY_generator(gen::ThermalGenEMIS{<: BuildPhase}, sys::PSY.System)

    tech = get_tech(gen)
    base_power = get_maxcap(gen)

    buses = PSY.get_components(PSY.Bus, sys)
    project_bus = PSY.Bus[]

    for b in buses
        if PSY.get_number(b) == get_bus(tech)
            push!(project_bus, b)
        end
    end

    PSY_gen =  PSY.ThermalStandard(
        get_name(gen), # name
        true,   # available
        true,   # status
        project_bus[1], # bus
        get_maxcap(gen) / base_power,    # active power
        1.0,                # reactive power
        get_maxcap(gen) / base_power,    # rating
        (min = get_mincap(gen) / base_power, max = get_maxcap(gen) / base_power), # active power limits
        nothing,      # reactive power limits
        (up = get_ramp_limits(tech)[:up] / (base_power * 60), down = get_ramp_limits(tech)[:down] / (base_power * 60)), # ramp limits
        get_operation_cost(tech), # operation cost
        base_power, # base power
        get_time_limits(tech), # up and down time limits
        PSY.PrimeMovers.Symbol(get_type(tech)), # primemover
        PSY.ThermalFuels.Symbol(get_fuel(tech)), # fuel type
    )
    return PSY_gen
end

"""
This function creates a PowerSystems RenewableDispatch unit.
"""
function create_PSY_generator(gen::RenewableGenEMIS{<: BuildPhase}, sys::PSY.System)
    tech = get_tech(gen)
    base_power = get_maxcap(gen)

    if get_type(tech) == "WT"
        primemover = PSY.PrimeMovers.WT
    elseif get_type(tech) == "PVe"
        primemover = PSY.PrimeMovers.PVe
    end

    buses = PSY.get_components(PSY.Bus, sys)
    project_bus = PSY.Bus[]

    for b in buses
        if PSY.get_number(b) == get_bus(tech)
            push!(project_bus, b)
        end
    end

    PSY_gen =  PSY.RenewableDispatch(
        get_name(gen),  # name
        true,           # available
        project_bus[1], # bus
        get_maxcap(gen) / base_power, # active power
        1.0,             # reactive power
        get_maxcap(gen) / base_power,    # rating
        primemover,     # primemover
        nothing,        # reactivepower limits
        1.0,            # power factor
        get_operation_cost(tech),
        base_power, # base power
    )
    return PSY_gen
end

"""
This function creates a PowerSystems GenericBattery unit.
"""
function create_PSY_generator(gen::BatteryEMIS{<: BuildPhase}, sys::PSY.System)
    tech = get_tech(gen)
    base_power = get_maxcap(gen)

    buses = PSY.get_components(PSY.Bus, sys)
    project_bus = PSY.Bus[]

    for b in buses
        if PSY.get_number(b) == get_bus(tech)
            push!(project_bus, b)
        end
    end

    PSY_gen =  PSY.GenericBattery(
        get_name(gen),  # name
        true,           # available
        project_bus[1], # bus
        PSY.PrimeMovers.Symbol(get_type(tech)), # primemover
        get_soc(tech) / (get_maxcap(gen) * base_power), #initial state of charge
        (min = get_storage_capacity(tech)[:min] / (get_maxcap(gen) * base_power), max = get_storage_capacity(tech)[:max] / (get_maxcap(gen) * base_power)), # state of charge limits
        get_maxcap(gen) / base_power, # rating
        get_maxcap(gen) / base_power, # active power
        (min = get_input_active_power_limits(tech)[:min] / base_power, max = get_input_active_power_limits(tech)[:max] / base_power), # input active power limits
        (min = get_output_active_power_limits(tech)[:min] / base_power, max = get_output_active_power_limits(tech)[:max] / base_power), # output active power limits
        get_efficiency(tech), # efficiency
        1.0,             # reactive power
        nothing,      # reactive power limits
        base_power, # base power
    )
    return PSY_gen
end


#///////////////////////////////////////
#// File Name: rs_sac_controller_test.jl
#// Author: Haruki Nishimura (hnishimura@stanford.edu)
#// Date Created: 2021/02/26
#// Description: Test code for src/rs_sac_controller.jl
#///////////////////////////////////////

using DataStructures
using LinearAlgebra
using Random
using RobotOS

@testset "RSSAC Controller Test" begin
# Cost parameters
Cep = Matrix(1.0I, 2, 2);
Cu = Matrix(1.0I, 2, 2);
β_pos = 0.6;
β_col = 0.4;
α_col = 100.0;
λ_col = 1.0;
σ_risk = 1.0;
cost_param = CostParameter(Cep, Cu, β_pos, α_col, β_col, λ_col, σ_risk);
# Simulation parameters
conf_file_name = "config.json";
test_data_name = "eth_test.pkl";
test_scene_id = 0;
start_time_idx = 50;
device = py"torch".device("cpu");
incl_robot_name = false;
scene_param = TrajectronSceneParameter(conf_file_name, test_data_name,
                                       test_scene_id, start_time_idx,
                                       incl_robot_name);
scene_loader = TrajectronSceneLoader(scene_param, verbose=false);
num_samples = 50;
prediction_steps = 12;
use_robot_future = false;
deterministic = false;
rng_seed_py = 1;
predictor_param = TrajectronPredictorParameter(prediction_steps, num_samples,
                                               use_robot_future, deterministic,
                                               rng_seed_py);
traj_predictor = TrajectronPredictor(predictor_param,
                                     scene_loader.model_dir,
                                     scene_loader.param.conf_file_name,
                                     device, verbose=false);
initialize_scene_graph!(traj_predictor, scene_loader);
ado_states = fetch_ado_positions!(scene_loader, return_full_state=true);
ado_positions = reduce_to_positions(ado_states);

dtc = 0.01;
sim_param = SimulationParameter(scene_loader, traj_predictor, dtc, cost_param);
# Initial conditions
ep = [-1., 3.5];
ev = [1., 1.];
t_init = Time(1.0);
e_init = RobotState([ep; ev], t_init);
ap_dict_init = convert_nodes_to_str(ado_positions);
t_last_m = Time(0.8);
w_init = WorldState(e_init, ap_dict_init, t_last_m);
# Target Trajectory
u_nominal_base = [0.0, 0.0];
u_array_base = [u_nominal_base for ii = 1:480];
wp1 = WayPoint2D([0.0, 0.0], Time(1.0));
wp2 = WayPoint2D([0.0, 0.0], Time(5, 8e8));
target_trajectory = Trajectory2D([wp1, wp2]);
# ControlParameter
eamax = 5.0;
tcalc = 0.1;
dtexec = [0.05, 0.01, 0.02];
dtr = 0.1;
nominal_search_depth = 2;
constraint_time = nothing;
u_nominal_cand = append!([[0.0, 0.0]],
                         [round.([a*cos(deg2rad(θ)), a*sin(deg2rad(θ))], digits=5)
                          for a = [2., 4.] for θ = 0.:45.:(360. - 45.)])             # nominal control candidate value [ax, ay] [m/s^2]
cnt_param = ControlParameter(eamax, tcalc, dtexec, dtr, u_nominal_base,
                             u_nominal_cand, nominal_search_depth,
                             constraint_time=constraint_time);
# Prediction Dict
prediction_dict = sample_future_ado_positions!(traj_predictor,
                                               ado_states);
num_controls = length(cnt_param.u_nominal_cand)^nominal_search_depth;
for key in keys(prediction_dict)
    prediction_dict[key] = repeat(prediction_dict[key], outer=(num_controls, 1, 1));
end
@test prediction_dict["PEDESTRIAN/25"][1:50, :, :] == prediction_dict["PEDESTRIAN/25"][51:100, :, :]

# # Helper functions test
# convert_to_schedule test
u_schedule = convert_to_schedule(w_init.t, u_array_base, sim_param);

@test length(u_schedule) == 480
@test all(collect(values(u_schedule)) .== [[0.0, 0.0]])
@test haskey(u_schedule, Time(1.0))
@test haskey(u_schedule, Time(1.01))
@test haskey(u_schedule, Time(5.79))

# get_nominal_u_arrays test: get nominal u_arrays from u_schedule and nominal control candidates
#u_nominal_mod_init_time = w_init.t + Duration(cnt_param.tcalc);
#u_nominal_mod_final_time = u_nominal_mod_init_time +
#                           Duration(sim_param.dto) -
#                           Duration(sim_param.dtc);
schedule_before = u_schedule;
u_arrays = get_nominal_u_arrays(u_schedule, sim_param, cnt_param);
@test all(u_arrays[1][1:end] .== [[0.0, 0.0]])
@test all([all(u_arrays[ii][1:10] .== [[0.0, 0.0]]) &&
           all(u_arrays[ii][11:50] .== [u_nominal_cand[div(ii - 1, length(cnt_param.u_nominal_cand)) + 1]]) &&
           all(u_arrays[ii][51:90] .== [u_nominal_cand[mod(ii - 1, length(cnt_param.u_nominal_cand)) + 1]]) &&
           all(u_arrays[ii][91:end] .== [[0.0, 0.0]])
           for ii = 2:length(cnt_param.u_nominal_cand)])
@test u_schedule == schedule_before

# get_robot_present_and_future_test
begin
    e_state = RobotState([-5.0, 0.0, 0.0, 1.0]);
    u_nominal = [[0.0, 0.0] for ii = 1:Int64(round(prediction_steps*sim_param.dto/dtc))];
    u_s = convert_to_schedule(e_state.t, u_nominal, sim_param);
    rpf = get_robot_present_and_future(e_state, u_s, sim_param, cnt_param);
    @test size(rpf) == (17^nominal_search_depth, 13, 6)
    @test all(isapprox.(cumsum(rpf[:, 1:end-1, 5:6], dims=2).*sim_param.dto .+ rpf[:, 1:1, 3:4], rpf[:, 2:end, 3:4], atol=1e-6))

end

# Get simulation result and best nominal u_array
sim_result, u_nominal_array = simulate(w_init, u_arrays, target_trajectory,
                                       prediction_dict, sim_param, cnt_param.constraint_time);

# get_control_coeffs_test
coeff_matrix, coeff_matrix_constraint = get_control_coeffs(sim_result, sim_param, cnt_param);
#@test size(coeff_matrix) == (2, 20);
@test size(coeff_matrix) == (2, 480);
@test isnothing(coeff_matrix_constraint)

test_time_id = 11;
den_test = sum(exp.(σ_risk.*sim_result.sampled_total_costs));
num_test = sum(exp.(σ_risk.*sim_result.sampled_total_costs).*
               [sim_result.e_costate_array[:, test_time_id, jj] for jj = 1:sim_param.num_samples]);

@test transition_control_coeff(sim_result.e_state_array[test_time_id])'*
    (num_test./den_test) ≈ coeff_matrix[:, test_time_id];

# get_control_schedule test
control_schedule_array = get_control_schedule(sim_result, u_nominal_array, coeff_matrix,
                                              sim_param, cnt_param);
#@test length(control_schedule_array) == Int64((tcalc+dtr)/dtc);
@test length(control_schedule_array) == Int64(round(sim_param.dto*sim_param.prediction_steps/dtc, digits=5));
@test maximum(map(s -> norm(vec(s.u)), control_schedule_array)) < eamax ||
      maximum(map(s -> norm(vec(s.u)), control_schedule_array)) ≈ eamax;
@test !any(isnan.(map(s -> s.cost, control_schedule_array)));

# determine_control_time test
control_chosen = determine_control_time(sim_result,
                                        control_schedule_array, sim_param,
                                        cnt_param);
@test t_init + Duration(cnt_param.tcalc) + Duration(maximum(cnt_param.dtexec)) <= control_chosen.t;
#@test control_chosen.t <= t_init + Duration(cnt_param.tcalc) + Duration(cnt_param.dtr)
@test control_chosen.t <= t_init + Duration(sim_param.dto*sim_param.prediction_steps);
@test norm(control_chosen.u) < eamax || norm(control_chosen.u) ≈ eamax;

# sac_control_update_test
tcalc_actual, best_schedule, sim_result =
    sac_control_update(w_init, u_schedule, target_trajectory, prediction_dict,
                       sim_param, cnt_param);
perturbation_times = [k for (k,v) in best_schedule if v==control_chosen.u];
if length(perturbation_times) > 1
    @test best_schedule[control_chosen.t - Duration(dtc)] == control_chosen.u
    @test best_schedule[control_chosen.t] != control_chosen.u
    @test in(to_sec(perturbation_times[end] - perturbation_times[1]) + sim_param.dtc, cnt_param.dtexec)
end

# # Main functions test
u_schedule = convert_to_schedule(w_init.t, u_array_base, sim_param);

controller = RSSACController(traj_predictor, u_schedule, sim_param, cnt_param);
scene_loader = TrajectronSceneLoader(scene_param, verbose=false);
traj_predictor = TrajectronPredictor(predictor_param,
                                     scene_loader.model_dir,
                                     scene_loader.param.conf_file_name,
                                     device, verbose=false);
initialize_scene_graph!(traj_predictor, scene_loader);
ado_states = fetch_ado_positions!(scene_loader, return_full_state=true);
ado_positions = reduce_to_positions(ado_states);

schedule_prediction!(controller, ado_states);
wait(controller.prediction_task);
@test istaskdone(controller.prediction_task);
@test controller.prediction_dict_tmp["PEDESTRIAN/25"][1:50, :, :] ==
      controller.prediction_dict_tmp["PEDESTRIAN/25"][51:100, :, :]

schedule_control_update!(controller, w_init, target_trajectory);
@test !isnothing(controller.prediction_dict);
@test isnothing(controller.prediction_task);
wait(controller.control_update_task);
@test istaskdone(controller.control_update_task);

ado_positions_str = convert_nodes_to_str(ado_positions);
latest_ado_pos_dict = deepcopy(ado_positions_str);
ado_1 = collect(keys(latest_ado_pos_dict))[1];
ado_2 = collect(keys(latest_ado_pos_dict))[2];
prediction_array_ado_2 = copy(controller.prediction_dict[ado_2]);
latest_ado_pos_dict[ado_2] += [0.5, 0.5];
pop!(latest_ado_pos_dict, ado_1);
latest_ado_pos_dict["PEDESTRIAN/999"] = [1.0, 1.0];
adjust_old_prediction!(controller, ado_positions_str, latest_ado_pos_dict);
@test all(controller.prediction_dict[ado_2][:, :, 1] .- prediction_array_ado_2[:, :, 1] .≈ 0.5);
@test all(controller.prediction_dict[ado_2][:, :, 2] .- prediction_array_ado_2[:, :, 2] .≈ 0.5);
@test !in(controller.prediction_dict, ado_1);
@test size(controller.prediction_dict["PEDESTRIAN/999"]) == (num_samples*num_controls, sim_param.prediction_steps, 2);
@test all(controller.prediction_dict["PEDESTRIAN/999"][:, :, 1] .== 1.0)
@test all(controller.prediction_dict["PEDESTRIAN/999"][:, :, 2] .== 1.0)

u_1 = control!(controller, Time(1.0));
@test u_1 == [0.0, 0.0];
@test !isnothing(controller.sim_result);
@test !isnothing(controller.tcalc_actual);
#@test !isnothing(controller.u_init_time);
#@test !isnothing(controller.u_last_time);
#@test !isnothing(controller.u_value);
@test isnothing(controller.control_update_task);

# # constraint_time tests
constraint_time = 0.1;
cnt_param = ControlParameter(eamax, tcalc, dtexec, dtr, u_nominal_base,
                             u_nominal_cand, nominal_search_depth,
                             constraint_time=constraint_time);
# Get simulation result and best nominal u_array
sim_result, u_nominal_array = simulate(w_init, u_arrays, target_trajectory,
                                       prediction_dict, sim_param, cnt_param.constraint_time);
# get_control_coeffs_test
coeff_matrix, coeff_matrix_constraint = get_control_coeffs(sim_result, sim_param, cnt_param);
#@test size(coeff_matrix) == (2, 20);
@test size(coeff_matrix) == (2, 480);
@test !isnothing(coeff_matrix_constraint)
@test size(coeff_matrix_constraint) == (2, 10);

u_array_constraint, cost_array_constraint = solve_multi_qcqp(u_nominal_array,
                                                             coeff_matrix,
                                                             coeff_matrix_constraint,
                                                             sim_param, cnt_param);
@test all([norm(u) <= cnt_param.eamax for u in u_array_constraint])
control_schedule_array_constraint =
    get_control_schedule(sim_result, u_nominal_array, coeff_matrix,
                         sim_param, cnt_param, coeff_matrix_constraint);
@test all([s.u == u for (s, u) in zip(control_schedule_array_constraint[1:10], u_array_constraint)])
@test all([s.cost == cost for (s, cost) in zip(control_schedule_array_constraint[1:10], cost_array_constraint)])
@test all([s.t == s_constraint.t for (s, s_constraint) in zip(control_schedule_array[1:10], control_schedule_array_constraint[1:10])])
@test all([s.u == s_constraint.u for (s, s_constraint) in zip(control_schedule_array[11:20], control_schedule_array_constraint[11:20])])
@test all([s.cost == s_constraint.cost for (s, s_constraint) in zip(control_schedule_array[11:20], control_schedule_array_constraint[11:20])])
@test all([s.t == s_constraint.t for (s, s_constraint) in zip(control_schedule_array[11:20], control_schedule_array_constraint[11:20])])
# sac_control_update_test
tcalc_actual, best_schedule, sim_result =
    sac_control_update(w_init, u_schedule, target_trajectory, prediction_dict,
                       sim_param, cnt_param);
end


const nbstandata = [
Dict(
  "K" => 4,
  "V" => 10,
  "M" => 200,
  "N" => 1955,
  "z" => [
    1, 1, 2, 1, 1, 3, 2, 3, 1, 1, 1, 1, 2, 1, 1,
    2, 4, 4, 1, 1, 1, 1, 4, 1, 2, 1, 1, 1, 3, 1, 3,
    1, 3, 3, 3, 1, 2, 1, 1, 1, 2, 1, 1, 1, 4, 3, 2,
    1, 2, 3, 1, 1, 1, 1, 2, 2, 1, 2, 1, 1, 1, 1, 1,
    1, 3, 1, 3, 1, 1, 1, 2, 1, 1, 1, 2, 1, 2, 1, 2,
    1, 3, 3, 2, 1, 4, 1, 1, 3, 4, 1, 2, 2, 1, 2, 3,
    1, 2, 3, 4, 1, 2, 3, 3, 3, 1, 2, 4, 3, 3, 4, 1,
    1, 2, 1, 1, 2, 1, 1, 3, 2, 4, 1, 1, 2, 1, 3, 1,
    1, 2, 2, 4, 2, 2, 1, 1, 2, 4, 1, 2, 2, 2, 1, 2,
    2, 1, 3, 1, 3, 2, 2, 2, 1, 1, 3, 1, 1, 1, 1, 1,
    1, 4, 1, 1, 1, 3, 2, 4, 1, 1, 1, 3, 1, 2, 1, 1,
    1, 1, 2, 1, 1, 1, 3, 4, 2, 3, 3, 2, 1, 4, 1, 3,
    1, 2, 1, 1, 1, 1, 3, 4, 1
  ],
  "w" => [7, 9, 5, 2, 9, 1, 1, 9, 6, 9, 1, 10, 6, 7, 2,
    3, 3, 2, 9, 3, 8, 10, 4, 3, 6, 2, 7, 1, 7, 7,
    1, 6, 7, 7, 1, 7, 7, 4, 7, 7, 4, 1, 7, 6, 6, 3,
    10, 5, 5, 4, 5, 5, 3, 5, 6, 3, 6, 5, 5, 3, 3,
    3, 7, 9, 2, 2, 9, 8, 5, 8, 9, 1, 7, 1, 7, 7, 1,
    1, 7, 7, 7, 4, 7, 4, 6, 4, 3, 9, 6, 7, 9, 9, 9,
    9, 6, 3, 7, 1, 1, 7, 1, 4, 7, 8, 1, 9, 9, 6, 8,
    2, 7, 7, 6, 9, 1, 7, 3, 7, 3, 3, 3, 4, 3, 9, 3,
    3, 3, 1, 4, 9, 10, 1, 6, 7, 6, 1, 7, 7, 1, 3,
    1, 7, 7, 6, 9, 1, 7, 3, 7, 7, 7, 6, 3, 3, 2, 8,
    3, 2, 3, 3, 1, 2, 3, 2, 2, 1, 3, 2, 9, 9, 2, 10,
    2, 1, 2, 2, 6, 2, 1, 1, 2, 2, 7, 1, 5, 1, 7, 7,
    7, 8, 3, 5, 7, 9, 7, 9, 7, 7, 1, 6, 1, 10, 7,
    7, 7, 9, 7, 6, 7, 9, 7, 9, 7, 7, 6, 7, 1, 1, 7,
    7, 1, 7, 2, 2, 1, 4, 2, 5, 7, 2, 1, 9, 4, 1, 7,
    7, 9, 1, 10, 7, 9, 2, 2, 2, 3, 3, 2, 4, 3, 4,
    4, 4, 4, 7, 1, 6, 7, 1, 7, 7, 7, 1, 7, 9, 4, 3,
    1, 9, 1, 7, 7, 9, 7, 1, 7, 9, 5, 5, 9, 5, 5, 4,
    8, 1, 9, 7, 6, 7, 7, 7, 9, 6, 8, 5, 5, 2, 9, 5,
    5, 1, 5, 1, 4, 9, 9, 7, 2, 7, 7, 7, 9, 5, 3, 5,
    5, 5, 6, 5, 6, 1, 5, 6, 3, 6, 5, 4, 6, 4, 3, 5,
    5, 5, 7, 7, 7, 7, 7, 7, 5, 1, 1, 7, 1, 7, 9, 9,
    2, 4, 2, 3, 3, 4, 3, 1, 1, 7, 1, 4, 7, 8, 7, 6,
    1, 1, 9, 1, 10, 3, 7, 7, 1, 1, 7, 2, 1, 7, 7,
    9, 4, 3, 3, 2, 3, 3, 2, 9, 3, 9, 7, 9, 7, 4, 1,
    1, 1, 1, 1, 7, 7, 3, 7, 1, 2, 7, 7, 7, 9, 1, 1,
    7, 1, 9, 7, 1, 7, 9, 7, 7, 6, 7, 1, 2, 2, 2, 1,
    9, 2, 2, 1, 1, 2, 6, 5, 5, 5, 6, 5, 6, 3, 2, 5,
    5, 3, 5, 5, 3, 3, 8, 9, 3, 3, 3, 9, 4, 3, 10,
    3, 2, 4, 6, 7, 10, 1, 10, 1, 9, 7, 7, 7, 3, 7,
    4, 9, 4, 4, 4, 3, 3, 3, 6, 3, 2, 2, 6, 7, 6, 5,
    5, 10, 3, 6, 5, 5, 5, 1, 9, 8, 8, 7, 7, 1, 7,
    7, 9, 1, 3, 1, 5, 10, 7, 7, 6, 7, 7, 7, 7, 7,
    7, 1, 7, 3, 1, 9, 7, 7, 7, 7, 8, 10, 7, 2, 7,
    5, 7, 4, 7, 7, 9, 5, 7, 6, 9, 1, 6, 4, 3, 2, 3,
    2, 4, 4, 3, 10, 2, 3, 3, 3, 9, 3, 3, 9, 3, 2,
    4, 3, 3, 6, 9, 3, 3, 7, 7, 1, 6, 9, 3, 7, 9, 3,
    7, 9, 9, 3, 4, 3, 4, 3, 9, 4, 6, 5, 1, 7, 9, 1,
    7, 4, 1, 7, 7, 7, 6, 7, 7, 7, 9, 7, 1, 6, 9, 7,
    9, 1, 9, 3, 7, 9, 7, 3, 7, 7, 7, 9, 3, 6, 7, 8,
    7, 7, 7, 7, 4, 3, 7, 7, 9, 7, 1, 1, 1, 7, 1, 7,
    9, 1, 7, 6, 5, 5, 10, 3, 3, 6, 1, 7, 7, 1, 9,
    7, 1, 3, 9, 7, 7, 5, 5, 6, 3, 5, 5, 1, 5, 5, 5,
    6, 3, 7, 8, 8, 8, 7, 4, 9, 9, 6, 7, 7, 7, 7, 1,
    2, 7, 1, 4, 9, 7, 7, 6, 1, 7, 7, 7, 7, 7, 7, 9,
    7, 9, 6, 7, 9, 3, 8, 3, 9, 9, 2, 3, 3, 3, 2, 9,
    3, 7, 10, 4, 7, 1, 7, 1, 9, 7, 7, 7, 4, 7, 1,
    4, 7, 7, 9, 1, 7, 7, 1, 7, 5, 2, 4, 4, 7, 7, 7,
    7, 7, 9, 3, 4, 4, 1, 3, 3, 4, 3, 3, 9, 9, 3, 2,
    7, 7, 7, 7, 7, 7, 4, 7, 7, 8, 3, 6, 2, 7, 10,
    3, 2, 7, 8, 3, 7, 1, 7, 1, 7, 8, 9, 1, 7, 7, 7,
    7, 7, 7, 3, 2, 9, 3, 2, 9, 7, 5, 5, 5, 6, 10,
    8, 5, 9, 6, 1, 5, 1, 5, 7, 7, 3, 2, 4, 3, 7, 4,
    9, 2, 2, 7, 1, 2, 5, 3, 2, 3, 4, 7, 7, 9, 7, 9,
    1, 1, 6, 9, 5, 9, 9, 2, 6, 2, 10, 2, 9, 1, 1,
    3, 8, 7, 7, 7, 1, 7, 7, 7, 9, 9, 7, 7, 4, 9, 1,
    4, 8, 5, 5, 5, 5, 6, 2, 1, 3, 5, 1, 2, 4, 7, 2,
    2, 2, 2, 2, 1, 9, 1, 10, 1, 6, 4, 3, 7, 7, 1,
    7, 3, 2, 6, 4, 3, 2, 8, 3, 3, 2, 8, 2, 7, 6, 3,
    8, 1, 10, 2, 10, 4, 3, 3, 3, 7, 9, 7, 7, 2, 7,
    4, 3, 2, 3, 5, 3, 3, 9, 9, 8, 5, 6, 8, 3, 5, 5,
    3, 1, 1, 7, 1, 7, 7, 7, 7, 7, 1, 7, 7, 3, 2, 3,
    3, 2, 3, 2, 2, 6, 5, 5, 5, 1, 2, 2, 5, 5, 10,
    7, 8, 7, 9, 7, 7, 1, 1, 8, 7, 7, 2, 3, 3, 7, 8,
    9, 5, 6, 3, 3, 5, 3, 5, 5, 5, 5, 5, 3, 1, 6, 3,
    5, 3, 5, 5, 5, 3, 5, 5, 5, 3, 5, 5, 7, 9, 4, 1,
    7, 1, 7, 3, 3, 5, 3, 6, 3, 3, 3, 2, 3, 3, 3, 2,
    1, 2, 2, 2, 9, 9, 3, 2, 3, 5, 6, 9, 3, 10, 5,
    5, 5, 5, 6, 6, 5, 2, 2, 2, 3, 2, 1, 2, 2, 1, 5,
    9, 1, 1, 7, 7, 8, 9, 7, 4, 1, 6, 7, 1, 4, 1, 3,
    3, 3, 2, 2, 7, 3, 2, 3, 2, 2, 4, 10, 2, 2, 3,
    3, 8, 1, 10, 8, 7, 7, 6, 7, 1, 7, 8, 9, 7, 7,
    7, 7, 7, 1, 1, 1, 1, 3, 7, 7, 1, 7, 7, 4, 3, 3,
    9, 3, 7, 4, 3, 7, 7, 7, 7, 9, 7, 9, 7, 8, 4, 1,
    1, 3, 10, 3, 5, 3, 5, 6, 6, 1, 3, 3, 2, 3, 2,
    7, 5, 7, 4, 2, 2, 1, 2, 9, 1, 1, 6, 2, 1, 7, 7,
    3, 10, 1, 7, 7, 3, 1, 1, 1, 2, 7, 6, 1, 9, 7,
    2, 2, 10, 4, 10, 2, 6, 7, 7, 7, 4, 5, 3, 5, 5,
    5, 6, 8, 6, 7, 7, 3, 4, 7, 7, 9, 4, 7, 10, 6,
    1, 7, 7, 7, 7, 1, 7, 1, 2, 1, 2, 9, 2, 3, 3, 3,
    10, 3, 3, 3, 9, 3, 3, 3, 3, 3, 3, 3, 3, 4, 2,
    7, 1, 7, 4, 10, 2, 2, 1, 8, 4, 8, 3, 3, 8, 9,
    3, 3, 4, 3, 3, 3, 2, 6, 2, 3, 4, 3, 3, 2, 7, 7,
    7, 6, 7, 7, 10, 4, 7, 3, 7, 1, 7, 7, 9, 9, 7,
    1, 7, 4, 9, 1, 9, 7, 7, 3, 2, 3, 9, 3, 9, 3, 2,
    2, 7, 8, 6, 2, 2, 1, 2, 2, 4, 2, 6, 7, 3, 1, 9,
    4, 7, 3, 3, 1, 3, 7, 3, 3, 9, 2, 3, 10, 3, 10,
    3, 3, 3, 3, 3, 4, 2, 2, 3, 3, 7, 6, 7, 7, 7, 7,
    7, 3, 3, 5, 6, 5, 9, 4, 7, 9, 3, 3, 9, 3, 2, 9,
    7, 6, 1, 7, 7, 3, 7, 1, 8, 7, 7, 3, 5, 3, 5, 5,
    6, 3, 2, 7, 7, 6, 7, 7, 1, 7, 6, 7, 1, 5, 5, 5,
    7, 5, 3, 3, 5, 2, 3, 4, 9, 3, 3, 3, 2, 2, 1, 3,
    2, 4, 2, 3, 2, 9, 3, 4, 3, 3, 7, 4, 3, 7, 7, 2,
    1, 3, 1, 7, 7, 7, 1, 7, 4, 7, 1, 7, 7, 9, 7, 6,
    7, 5, 7, 1, 7, 8, 6, 6, 3, 5, 6, 3, 5, 5, 5, 5,
    5, 6, 6, 5, 3, 5, 5, 7, 4, 7, 1, 1, 4, 9, 7, 7,
    7, 7, 7, 2, 7, 7, 1, 1, 7, 1, 3, 1, 6, 3, 2, 9,
    1, 9, 9, 7, 8, 7, 6, 1, 1, 7, 7, 7, 7, 9, 4, 9,
    7, 1, 6, 8, 1, 2, 2, 1, 9, 7, 2, 1, 2, 9, 9, 6,
    7, 1, 7, 7, 8, 7, 10, 7, 7, 7, 7, 7, 1, 1, 4,
    1, 7, 7, 1, 1, 1, 1, 7, 7, 4, 7, 5, 5, 6, 5, 10,
    5, 5, 5, 5, 5, 5, 6, 6, 5, 2, 3, 10, 3, 3, 4,
    3, 4, 2, 3, 1, 1, 2, 10, 1, 2, 7, 7, 9, 4, 6,
    7, 7, 5, 1, 8, 7, 6, 1, 9, 1, 3, 2, 3, 8, 1, 7,
    7, 8, 5, 7, 3, 8, 8, 9, 1, 9, 9, 1, 9, 2, 9, 6,
    5, 3, 6, 6, 9, 5, 4, 9, 3, 5, 5, 1, 5, 5, 6, 7,
    7, 8, 7, 1, 7, 7, 7, 3, 7, 6, 4, 2, 2, 3, 3, 2,
    3, 3, 4, 3, 3, 3, 3, 3, 4, 7, 6, 7, 9, 7, 1, 7,
    4, 7, 7, 7, 9, 1, 4, 9, 7, 7, 7, 1, 7, 7, 7, 1,
    6, 7, 7, 1, 1, 3, 7, 7, 7, 6, 7, 3, 5, 9, 7, 7,
    7, 1, 9, 7, 7, 1, 3, 3, 1, 3, 3, 3, 4, 2, 2, 3,
    3, 3, 6, 1, 3, 7, 3, 4, 7, 7, 9, 4, 6, 10, 9,
    7, 6, 7, 1, 2, 1, 1, 7, 1, 1, 7, 7, 7, 7, 7, 9,
    1, 7, 7, 7, 9, 9, 8, 10, 5, 4, 5, 5, 3, 5, 3,
    1, 6, 8, 5, 1, 4, 2, 1, 2, 2, 2, 1, 2, 6, 2, 4,
    9, 3, 3, 9, 10, 3, 3, 9, 6, 6, 3, 8, 5, 3, 5,
    3, 1, 7, 4, 2, 6, 3, 10, 7, 8, 9, 7, 7, 7, 7,
    7, 2, 1, 2, 2, 2, 1, 5, 7, 2, 2, 1, 2, 9, 4, 1,
    7, 7, 2, 4, 7, 6, 6, 10, 5, 6, 5, 5, 8, 9, 5,
    5, 7, 1, 6, 9, 7, 7, 7, 3, 3, 7, 3, 2, 1, 3, 4,
    7, 1, 7, 6, 7, 7, 7, 7, 4, 7, 5, 9, 6, 9, 4, 7,
    1, 7, 7, 1, 1, 7, 7, 1, 7, 7, 6, 7, 7, 1, 7, 8,
    7, 7, 9, 7, 9, 7, 2, 4, 9, 9, 5, 5, 4, 5, 1, 3,
    3, 7, 5, 2, 1, 7, 5, 5, 9, 9, 2, 2, 2, 5, 5, 1,
    7, 1, 7, 2, 4, 7, 7, 9, 1, 7, 9],
  "doc" => [1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 3,
    3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4,
    4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6,
    6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7,
    7, 7, 7, 7, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9,
    9, 9, 9, 9, 9, 9, 9, 10, 10, 10, 10, 10, 10, 10,
    10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
    11, 11, 11, 11, 11, 12, 12, 12, 12, 12, 12, 12, 13,
    13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 14, 14, 14,
    14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15,
    15, 15, 15, 15, 15, 15, 15, 16, 16, 16, 16, 16, 16,
    16, 16, 16, 17, 17, 17, 17, 17, 17, 17, 17, 17, 17,
    17, 17, 17, 18, 18, 18, 18, 18, 18, 18, 19, 19, 19,
    19, 19, 19, 19, 19, 19, 19, 19, 20, 20, 20, 20, 20,
    20, 20, 20, 21, 21, 21, 21, 21, 21, 21, 21, 21, 21,
    21, 21, 21, 21, 21, 21, 21, 22, 22, 22, 23, 23, 23,
    23, 23, 23, 23, 23, 23, 23, 23, 23, 23, 24, 24, 24,
    24, 24, 24, 24, 25, 25, 25, 25, 25, 25, 25, 25, 25,
    25, 25, 25, 25, 26, 26, 26, 26, 26, 27, 27, 27, 27,
    27, 27, 27, 27, 27, 27, 27, 27, 28, 28, 28, 28, 28,
    28, 29, 29, 29, 29, 29, 29, 29, 30, 30, 30, 30, 30,
    30, 30, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31,
    32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 33, 33, 33,
    33, 33, 34, 34, 34, 34, 34, 34, 34, 34, 34, 35, 35,
    35, 35, 35, 35, 35, 36, 36, 36, 36, 36, 36, 36, 36,
    36, 36, 36, 36, 36, 36, 37, 37, 37, 37, 37, 37, 37,
    38, 38, 38, 38, 38, 38, 38, 38, 38, 38, 38, 39, 39,
    39, 39, 39, 39, 39, 39, 40, 40, 40, 40, 40, 40, 40,
    41, 41, 41, 41, 41, 41, 41, 41, 42, 42, 42, 42, 42,
    42, 42, 42, 42, 43, 43, 43, 43, 43, 43, 43, 43, 43,
    43, 43, 43, 43, 43, 43, 43, 44, 44, 44, 44, 44, 44,
    44, 44, 44, 45, 45, 45, 45, 45, 45, 45, 45, 45, 45,
    46, 46, 46, 46, 46, 46, 46, 46, 46, 46, 46, 46, 46,
    46, 46, 47, 47, 47, 47, 47, 47, 47, 47, 47, 47, 47,
    47, 47, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48, 48,
    48, 48, 48, 49, 49, 49, 49, 49, 49, 49, 49, 49, 49,
    50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 50, 51, 51,
    51, 51, 51, 51, 51, 51, 51, 51, 51, 51, 52, 52, 52,
    52, 52, 52, 52, 52, 52, 52, 52, 52, 52, 52, 52, 53,
    53, 53, 53, 53, 53, 53, 53, 53, 53, 53, 53, 53, 53,
    54, 54, 54, 54, 54, 54, 54, 54, 54, 55, 55, 55, 55,
    55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 56, 56, 56,
    56, 56, 56, 56, 56, 56, 56, 56, 57, 57, 57, 57, 57,
    57, 57, 57, 57, 58, 58, 58, 58, 58, 58, 58, 58, 58,
    58, 58, 59, 59, 59, 59, 59, 59, 59, 59, 59, 59, 59,
    60, 60, 60, 60, 60, 60, 60, 60, 60, 60, 60, 60, 61,
    61, 61, 61, 61, 61, 61, 61, 61, 62, 62, 62, 62, 62,
    62, 62, 62, 62, 62, 63, 63, 63, 63, 63, 63, 63, 64,
    64, 64, 64, 64, 64, 64, 64, 65, 65, 65, 65, 65, 65,
    65, 65, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 67,
    67, 67, 67, 67, 67, 67, 67, 67, 67, 67, 67, 67, 67,
    67, 67, 67, 68, 68, 68, 68, 68, 68, 68, 68, 69, 69,
    69, 69, 69, 69, 69, 69, 69, 69, 69, 70, 70, 70, 70,
    70, 70, 70, 70, 70, 70, 71, 71, 71, 71, 71, 71, 71,
    71, 71, 71, 71, 71, 71, 72, 72, 72, 72, 72, 72, 72,
    72, 72, 72, 73, 73, 73, 73, 73, 73, 73, 73, 73, 73,
    73, 73, 73, 74, 74, 74, 74, 74, 74, 74, 74, 74, 75,
    75, 75, 75, 75, 75, 75, 75, 75, 75, 76, 76, 76, 76,
    76, 76, 76, 76, 76, 76, 76, 76, 76, 76, 77, 77, 77,
    77, 77, 77, 77, 77, 77, 77, 77, 78, 78, 78, 78, 78,
    78, 78, 78, 78, 78, 78, 78, 79, 79, 79, 79, 79, 79,
    80, 80, 81, 81, 81, 82, 82, 82, 82, 82, 82, 82, 82,
    82, 82, 82, 83, 83, 83, 83, 83, 83, 83, 83, 83, 83,
    83, 83, 83, 83, 83, 83, 83, 83, 84, 84, 84, 84, 84,
    84, 84, 84, 84, 85, 85, 85, 85, 85, 85, 85, 85, 85,
    85, 86, 86, 86, 86, 86, 86, 86, 87, 87, 87, 87, 87,
    87, 87, 87, 87, 87, 87, 88, 88, 88, 88, 88, 88, 88,
    88, 88, 88, 89, 89, 89, 89, 89, 89, 89, 89, 89, 89,
    90, 90, 90, 90, 90, 90, 90, 90, 90, 90, 90, 91, 91,
    91, 91, 91, 91, 91, 91, 92, 92, 92, 92, 92, 92, 92,
    92, 92, 92, 92, 92, 92, 92, 92, 92, 93, 93, 93, 93,
    93, 93, 94, 94, 94, 94, 94, 94, 94, 94, 94, 94, 95,
    95, 95, 95, 95, 95, 96, 96, 96, 96, 96, 96, 96, 96,
    96, 96, 96, 96, 96, 97, 97, 97, 97, 97, 97, 97, 97,
    98, 98, 98, 98, 98, 99, 99, 99, 99, 100, 100, 100,
    100, 100, 100, 100, 100, 100, 100, 100, 100, 101, 101,
    101, 101, 102, 102, 102, 102, 102, 102, 102, 102, 102,
    102, 103, 103, 103, 103, 103, 103, 103, 103, 103, 103,
    104, 104, 104, 104, 104, 104, 104, 104, 104, 105, 105,
    105, 105, 105, 105, 105, 106, 106, 106, 106, 106, 106,
    106, 106, 106, 106, 106, 106, 107, 107, 107, 107, 107,
    107, 107, 107, 107, 108, 108, 108, 108, 108, 109, 109,
    109, 109, 109, 109, 109, 109, 110, 110, 110, 110, 110,
    110, 110, 110, 110, 110, 110, 111, 111, 111, 111, 111,
    112, 112, 112, 112, 112, 112, 112, 112, 112, 113, 113,
    113, 113, 113, 113, 113, 113, 113, 113, 113, 113, 113,
    113, 113, 113, 113, 113, 114, 114, 114, 114, 114, 114,
    114, 114, 114, 114, 114, 114, 115, 115, 115, 115, 115,
    115, 115, 115, 115, 115, 115, 115, 115, 115, 116, 116,
    116, 116, 116, 117, 117, 117, 117, 117, 117, 118, 118,
    118, 118, 118, 118, 118, 118, 118, 118, 119, 119, 119,
    119, 119, 119, 119, 119, 120, 120, 120, 120, 120, 120,
    120, 120, 121, 121, 121, 121, 121, 121, 121, 121, 121,
    121, 122, 122, 122, 122, 122, 122, 122, 122, 123, 123,
    123, 123, 123, 123, 123, 123, 123, 123, 124, 124, 124,
    124, 124, 124, 125, 125, 125, 125, 126, 126, 126, 126,
    126, 126, 126, 126, 127, 127, 127, 127, 127, 127, 127,
    127, 127, 127, 127, 127, 128, 128, 128, 128, 128, 128,
    128, 128, 128, 128, 129, 129, 129, 129, 129, 129, 129,
    129, 129, 129, 129, 130, 130, 130, 130, 130, 130, 130,
    130, 130, 130, 131, 131, 131, 131, 131, 131, 131, 131,
    132, 132, 132, 132, 132, 132, 132, 133, 133, 133, 133,
    133, 133, 133, 133, 133, 133, 133, 133, 133, 133, 134,
    134, 134, 134, 134, 134, 134, 134, 134, 134, 134, 134,
    134, 134, 134, 135, 135, 135, 135, 135, 135, 135, 135,
    135, 135, 136, 136, 136, 136, 136, 136, 136, 137, 137,
    137, 137, 137, 137, 137, 137, 137, 137, 137, 137, 138,
    138, 138, 138, 138, 138, 138, 139, 139, 139, 139, 139,
    139, 140, 140, 140, 140, 140, 140, 140, 140, 141, 141,
    141, 141, 141, 141, 141, 141, 141, 141, 142, 142, 142,
    142, 142, 142, 143, 143, 143, 143, 144, 144, 144, 144,
    144, 144, 144, 144, 144, 144, 145, 145, 145, 145, 145,
    145, 145, 145, 145, 145, 145, 145, 146, 146, 146, 146,
    146, 146, 146, 146, 147, 147, 147, 147, 147, 147, 147,
    147, 147, 147, 148, 148, 148, 148, 148, 148, 148, 148,
    148, 149, 149, 149, 149, 149, 149, 150, 150, 150, 150,
    150, 150, 150, 150, 150, 150, 150, 150, 150, 150, 151,
    151, 151, 151, 151, 151, 152, 152, 152, 152, 152, 152,
    152, 152, 152, 152, 153, 153, 153, 153, 153, 153, 153,
    153, 153, 153, 153, 154, 154, 154, 154, 154, 154, 154,
    154, 154, 154, 154, 154, 154, 154, 154, 154, 154, 155,
    155, 155, 155, 155, 155, 155, 155, 156, 156, 157, 157,
    157, 157, 157, 158, 158, 158, 158, 158, 158, 158, 159,
    159, 159, 159, 159, 159, 159, 159, 159, 159, 159, 159,
    159, 160, 160, 160, 160, 160, 160, 160, 160, 160, 160,
    160, 160, 161, 161, 161, 161, 161, 161, 161, 161, 161,
    162, 162, 162, 162, 163, 163, 163, 163, 163, 163, 163,
    163, 163, 164, 164, 164, 164, 164, 164, 164, 164, 164,
    164, 164, 164, 164, 164, 165, 165, 165, 165, 165, 165,
    165, 165, 165, 165, 165, 165, 165, 165, 165, 165, 166,
    166, 166, 166, 166, 166, 166, 166, 166, 167, 167, 167,
    167, 167, 167, 168, 168, 168, 168, 168, 168, 168, 168,
    168, 168, 168, 168, 168, 169, 169, 169, 169, 169, 169,
    169, 169, 169, 170, 170, 170, 170, 170, 170, 170, 170,
    170, 170, 170, 170, 170, 171, 171, 171, 171, 171, 171,
    171, 171, 171, 171, 171, 171, 171, 171, 171, 171, 171,
    172, 172, 172, 172, 172, 172, 172, 172, 172, 172, 172,
    173, 173, 173, 173, 173, 173, 173, 173, 173, 173, 173,
    173, 173, 173, 174, 174, 174, 174, 174, 174, 174, 174,
    174, 174, 174, 174, 174, 174, 174, 174, 175, 175, 175,
    175, 175, 175, 175, 175, 175, 175, 175, 175, 176, 176,
    176, 176, 176, 176, 176, 176, 176, 176, 177, 177, 177,
    177, 177, 177, 177, 177, 177, 178, 178, 178, 178, 178,
    178, 178, 178, 178, 178, 178, 178, 179, 179, 179, 179,
    179, 179, 179, 179, 179, 179, 180, 180, 180, 180, 180,
    180, 180, 180, 181, 181, 181, 181, 181, 181, 181, 181,
    181, 181, 181, 181, 181, 181, 181, 181, 182, 182, 182,
    182, 182, 182, 182, 182, 182, 182, 182, 182, 182, 183,
    183, 183, 183, 183, 183, 183, 183, 183, 183, 183, 184,
    184, 184, 184, 184, 184, 184, 184, 184, 185, 185, 186,
    186, 186, 186, 186, 187, 187, 187, 187, 187, 187, 187,
    188, 188, 188, 188, 188, 188, 188, 188, 188, 189, 189,
    189, 189, 189, 189, 189, 189, 189, 189, 189, 189, 190,
    190, 190, 190, 190, 190, 190, 190, 190, 190, 191, 191,
    191, 191, 191, 191, 191, 191, 191, 192, 192, 192, 192,
    192, 192, 192, 193, 193, 193, 193, 193, 193, 193, 193,
    194, 194, 194, 194, 194, 194, 194, 195, 195, 195, 195,
    195, 195, 195, 195, 195, 195, 195, 195, 196, 196, 196,
    196, 196, 196, 196, 196, 196, 196, 196, 196, 197, 197,
    197, 197, 197, 197, 197, 197, 197, 197, 197, 198, 198,
    198, 198, 198, 198, 198, 198, 198, 198, 198, 198, 198,
    198, 199, 199, 199, 199, 199, 199, 199, 200, 200, 200,
    200, 200, 200, 200, 200, 200, 200, 200, 200],
  "alpha" => [1, 1, 1, 1],
  "beta" => [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1]
)
]


"""
    FullFeaturedCPNN(;
        graphChain::Flux.Chain
        nodeChain::Flux.Chain
        globalChain::Flux.Chain
        outputChain::Flux.Dense
    )

This structure is here to provide a flexible way to create a nn model which respect this approach:
Making modification on the graph, then extract one node feature and modify it.
"""
Base.@kwdef struct FullFeaturedCPNN <: NNStructure
    graphChain::Flux.Chain = Flux.Chain()
    nodeChain::Flux.Chain = Flux.Chain()
    globalChain::Flux.Chain = Flux.Chain()
    outputChain::Union{Flux.Dense, Flux.Chain} = Flux.Chain()
end

# Enable the `|> gpu` syntax from Flux
Flux.@functor FullFeaturedCPNN

# TODO make it possible to use global features
function (nn::FullFeaturedCPNN)(states::BatchedDefaultTrajectoryState)
    @assert !isnothing(states.allValuesIdx)

    variableIdx = states.variableIdx
    batchSize = length(variableIdx)
    allValuesIdx = states.allValuesIdx
    actionSpaceSize = size(allValuesIdx, 1)
    mask = device(states) == Val(:gpu) ? CUDA.zeros(Float32, 1, actionSpaceSize, batchSize) : zeros(Float32, 1, actionSpaceSize, batchSize) # this mask will replace `reapeat` using broadcasted `+`

    # chain working on the graph(s) with the GNNs
    featuredGraph = nn.graphChain(states.fg)
    nodeFeatures = featuredGraph.nf # FxNxB
    globalFeatures = featuredGraph.gf # GxB

    # Extract the features corresponding to the varibales
    variableIndices = nothing
    Zygote.ignore() do
        variableIndices = Flux.unsqueeze(CartesianIndex.(variableIdx, 1:batchSize), 1)
    end
    variableFeatures = nodeFeatures[:, variableIndices] # Fx1xB
    variableFeatures = reshape(nn.nodeChain(RL.flatten_batch(variableFeatures)), :, 1, batchSize) # F'x1xB

    # Extract the features corresponding to the values
    valueIndices = nothing
    Zygote.ignore() do
        valueIndices = CartesianIndex.(allValuesIdx, repeat(transpose(1:batchSize); outer=(actionSpaceSize, 1)))
    end
    valueFeatures = nodeFeatures[:, valueIndices] # FxAxB
    valueFeatures = reshape(nn.nodeChain(RL.flatten_batch(valueFeatures)), :, actionSpaceSize, batchSize) # F'xAxB

    finalFeatures = nothing
    if sizeof(globalFeatures) != 0

        # Extract the global features
        globalFeatures = reshape(nn.globalChain(globalFeatures), :, 1, batchSize) # G'x1xB

        # Prepare the input of the outputChain
        finalFeatures = vcat(
            variableFeatures .+ mask, # F'xAxB
            globalFeatures .+ mask, # G'xAxB
            valueFeatures,
        ) # (F'+G'+F')xAxB
        finalFeatures = RL.flatten_batch(finalFeatures) # (F'+G'+F')x(A+B)
    else
        # Prepare the input of the outputChain
        finalFeatures = vcat(
            variableFeatures .+ mask, # F'xAxB
            valueFeatures,
        ) # (F'+F')xAxB
        finalFeatures = RL.flatten_batch(finalFeatures) # (F'+F')x(A+B)
    end

    # output layer
    predictions = nn.outputChain(finalFeatures) # Ox(A+B)
    output = reshape(predictions, actionSpaceSize, batchSize) # OxAxB

    return output
end


# requires:

primitive type EdgeType 8 end

const EDGETYPE_TO_UINT8 = fill( 0x07, (4,4) )
      EDGETYPE_TO_UINT8[4,2] = 0x00 # 'SL' = 0x00; Tx Start
      EDGETYPE_TO_UINT8[4,3] = 0x01 # 'SR' = 0x01; Tandem Last Exon
      EDGETYPE_TO_UINT8[2,4] = 0x02 # 'LS' = 0x02; Tandem First Exon
      EDGETYPE_TO_UINT8[3,4] = 0x03 # 'RS' = 0x03; Tx End
      EDGETYPE_TO_UINT8[2,3] = 0x04 # 'LR' = 0x04; Intron 5'->3' SS
      EDGETYPE_TO_UINT8[2,2] = 0x05 # 'LL' = 0x05; Alt- 5'SS
      EDGETYPE_TO_UINT8[3,3] = 0x06 # 'RR' = 0x06; Alt- 3'SS

const INDEX_TO_EDGETYPE = transpose(reshape([[0x03,0x01,0x03,0x03];
                                             [0x06  for _ in 1:4 ];
                                             [0x05,0x04,0x05,0x02];
                                             [0x00  for _ in 1:4 ] ], (4,4)))

const INDEX_TO_EDGETYPE_NODE = transpose(reshape([[0x03  for _ in 1:4 ];
                                                  [0x06  for _ in 1:4 ];
                                                  [0x05  for _ in 1:4 ];
                                                  [0x00  for _ in 1:4 ] ], (4,4)))

function Base.convert( ::Type{EdgeType}, one::UInt8, two::UInt8 )
   @assert( 5 <= one <= 7 && 5 <= one <= 7 )
   EDGETYPE_TO_UINT8[one-3,two-3]
end

Base.convert( ::Type{EdgeType}, one::DNA, two::DNA ) =
              Base.convert( EdgeType, convert(UInt8, trailing_zeros(one)), convert(UInt8, trailing_zeros(two)) )
Base.convert( ::Type{EdgeType}, edge::UInt8 ) = reinterpret(EdgeType, edge)
Base.convert( ::Type{UInt8}, edge::EdgeType ) = reinterpret(UInt8, edge)
Base.convert( ::Type{I}, edge::EdgeType) where {I <: Integer} = Base.convert(I, Base.convert(UInt8, edge))

const EDGETYPE_SL = convert(EdgeType, 0b000)
const EDGETYPE_SR = convert(EdgeType, 0b001)
const EDGETYPE_LS = convert(EdgeType, 0b010)
const EDGETYPE_RS = convert(EdgeType, 0b011)

const EDGETYPE_LR = convert(EdgeType, 0b100)
const EDGETYPE_LL = convert(EdgeType, 0b101)
const EDGETYPE_RR = convert(EdgeType, 0b110)

istxstart( edge::EdgeType ) = edge == EDGETYPE_SL || edge == EDGETYPE_LS ? true : false
istxstop(  edge::EdgeType ) = edge == EDGETYPE_SR || edge == EDGETYPE_RS ? true : false

isexonstart( edge::EdgeType ) = edge in (EDGETYPE_LR, EDGETYPE_RR, EDGETYPE_SR, EDGETYPE_LS, EDGETYPE_SL) ? true : false
isexonstop(  edge::EdgeType ) = edge in (EDGETYPE_LR, EDGETYPE_LL, EDGETYPE_SR, EDGETYPE_LS, EDGETYPE_RS) ? true : false
ishardedge(  edge::EdgeType ) = edge in (EDGETYPE_LR, EDGETYPE_LS, EDGETYPE_SR) ? true : false

# Function takes coordinate types for node boundaries
# and returns an EdgeType
function get_edgetype( minidx::Int, secidx::Int, isnode::Bool, strand::Bool )
   if isnode
      ret = INDEX_TO_EDGETYPE_NODE[minidx,secidx]
   else
      ret = INDEX_TO_EDGETYPE[minidx,secidx]
   end
   ret = EdgeType(ret)
   if !strand
      ret = invert_edgetype( ret )
   end
   ret
end

# Invert EdgeTypes using xor
# 'LR' does not need to be inverted
function invert_edgetype( edge::EdgeType )
   if edge == EdgeType(0x04) # 'LR'
      return edge
   else
      return convert(EdgeType, convert(UInt8, edge) ⊻ 0b011)
   end
end

struct GCBiasParams
   length::Int64           # length offset
   width::Float64          # gc bin size
end

const ExpectedGC = Vector{Float16}

function ExpectedGC( seq::BioSequence{A}; gc_param = GCBiasParams(50, 0.05) ) where A <: BioSequences.Alphabet
   bins = zeros(Float16, Int(div(1.0, gc_param.width)+1))
   for i in 1:length(seq)-gc_param.length+1
      gc = Int(div(gc_content(seq[i:i+gc_param.length-1]), gc_param.width)+1)
      @inbounds bins[gc] += 1.0
   end
   # normalize columns
   binsum = sum(bins)
   binsum == 0.0 && (return bins)
   for i in 1:length(bins)
      @fastmath bins[i] /= binsum
   end
   return bins
end

# This holds a representation of the splice graph
# which is a directed multigraph
struct SpliceGraph{K}
   nodeoffset::Vector{CoordInt} # SG offset
   nodecoord::Vector{CoordInt}  # Genome offset
   nodelen::Vector{CoordInt}
   edgetype::Vector{EdgeType}
   edgeleft::Vector{SGKmer{K}}
   edgeright::Vector{SGKmer{K}}
   annopath::Vector{IntSet}
   annoname::Vector{String}
   annobias::Vector{ExpectedGC}
   seq::SGSequence
end
# All positive strand oriented sequences--->
# Node array: txStart| 1 |   2   | 3 |    4    |5| 6 |txEnd
# Edge array:        1   2       3   4         5 6   7
# Node coord:  chr   100 200     300 400      500 600 700

#= empty constructor
SpliceGraph(k::Int) = SpliceGraph( Vector{CoordInt}(), Vector{CoordInt}(),
                                   Vector{CoordInt}(), Vector{EdgeType}(),
                                   Vector{SGKmer{k}}(),Vector{SGKmer{k}}(), dna"" )
=#

# Main constructor
# Build splice graph here.
function SpliceGraph( gene::RefGene, genome::SGSequence, k::Int )
   # splice graph variables
   nodeoffset = Vector{CoordInt}()
   nodecoord  = Vector{CoordInt}()
   nodelen    = Vector{CoordInt}()
   edgetype   = Vector{EdgeType}()
   seq        = DNASequence()

   strand = gene.info.strand  # Bool

   # initialize iterators
   a,alen = 1,length(gene.acc)
   d,dlen = 1,length(gene.don)
   s,slen = 1,length(gene.txst)
   p,plen = 1,length(gene.txen)

   idx = [a, d, s, p]

   # return a tuple containing the min coordinate's idx,val
   # each of 4 categories, txstart, donor, acceptor, txend
   function getmin_ind_val( gene::RefGene, idx; def=typemax(CoordInt) )
      retarr = CoordInt[ get(gene.txen, idx[1], def),
                         get(gene.acc,  idx[2], def),
                         get(gene.don,  idx[3], def),
                         get(gene.txst, idx[4], def) ]
      indmin( retarr ), min( retarr... )
   end

   while( idx[1] <= alen || idx[2] <= dlen || idx[3] <= slen || idx[4] <= plen )
      # iterate through donors, and acceptors
      # left to right. '-' strand gets revcomp and unshift
      minidx,minval = getmin_ind_val( gene, idx )
      idx[minidx] += 1
      secidx,secval = getmin_ind_val( gene, idx )

      # last coordinate in the gene:
      # if secidx == 1 && get(gene.txst, idx[secidx], Inf) == Inf
      if secval == typemax(CoordInt)
         termedge = strand ? EdgeType(0x03) : EdgeType(0x00)
         stranded_push!(edgetype, termedge, strand)
         break
      end

      # now should we make a node?
      if issubinterval( gene.exons, Interval{CoordInt}(minval,secval) ) ||
         (minidx == 2 && minval in gene.novelacc) ||
         (secidx == 3 && secval in gene.noveldon)
         #(usebam && isretainedintron( )) # TODO (not yet implemented)
         leftadj  = (minidx == 1 || minidx == 3) && minval != secval ? 1 : 0
         rightadj = (secidx == 2 || secidx == 4) && minval != secval ? 1 : 0
         nodesize = Int(secval - minval) - leftadj - rightadj + 1
         nodeseq  = copy(genome[(Int(minval)+leftadj):(Int(secval)-rightadj)]) # collect slice
         edge     = get_edgetype( minidx, secidx, true, strand ) # determine EdgeType
         pushval  = minval + leftadj
         thridx = 0
      else # don't make a node, this is a sequence gap, make edge and inc+=2
         idx[secidx] += 1 #skip ahead again
         thridx,thrval = getmin_ind_val( gene, idx )
         rightadj = (thridx == 2 || thridx == 4) && secval != thrval ? 1 : 0
         nodesize = Int(thrval - secval) - rightadj + 1
         nodeseq  = copy(genome[Int(secval):(Int(thrval)-rightadj)])
         edge     = get_edgetype( minidx, secidx, false, strand )
         pushval  = secval
      end

      if strand
         seq *= nodeseq
      else # '-' strand
         seq = reverse_complement(nodeseq) * seq
      end
      stranded_push!(nodecoord, pushval,  strand)
      stranded_push!(nodelen,   nodesize, strand)
      stranded_push!(edgetype,  edge,     strand)

   end # end while

   # calculate node offsets-->
   curoffset = 1
   for n in nodelen
      push!(nodeoffset, curoffset)
      curoffset += n
   end

   eleft  = Vector{SGKmer{k}}(length(edgetype))
   eright = Vector{SGKmer{k}}(length(edgetype))

   paths = build_paths_edges( nodecoord, nodelen, gene )
   names = map( x->x.info.name, gene.reftx )
   expgc = map( x->ExpectedGC(path_to_seq(x, nodeoffset, nodelen, seq)), paths )

   return SpliceGraph( nodeoffset, nodecoord, nodelen, edgetype, eleft, eright, paths, names, expgc, seq )
end

# re-orient - strand by using unshift! instead of push!
function stranded_push!( collection, value, strand::Bool )
   if strand
      push!( collection, value )
   else # '-' strand
      unshift!( collection, value )
   end
end

# This function looks specifically for an intersection
# between one interval, and the interval tree, such that
# some interval in the tree fully contains the sub interval
#  itree contains: (low,             high)
#  subint must be:     (low+x, high-y)
#  where x and y are >= 0
# Returns: bool
function issubinterval( itree::T, subint::I ) where {T <: IntervalTree,
                                         I <: AbstractInterval}
   for i in intersect(itree, subint)
      if i.first <= subint.first && i.last >= subint.last
         return true
      end
   end
   false
end

# hacky extension of Base.get for tuples which isn't supported natively
# because tuples aren't really a type of collection, fair enough
function Base.get( collection::T, key::I, def ) where {T <: Tuple, I <: Integer}
   if 1 <= key <= length(collection)
      return collection[key]
   else
      return def
   end
end

# Functions for iso.jl annotated edges feature
# Take a RefTx and produce an IntSet through the nodes of a SpliceGraph
# returns: IntSet
function build_annotated_path( nodecoord::Vector{CoordInt},
                               nodelen::Vector{CoordInt},
                               tx::RefTx, strand::Bool )
   path = IntSet()
   # this may be `poor form`, but 256 is too big for default!
   resize!(path.bits, 64) # Deprecated:  = zeros(UInt32,64>>>5)
   for i in 1:length(tx.acc)
      ind = collect(searchsorted( nodecoord, tx.acc[i], rev=!strand ))[end]
      push!( path, ind )
      cur = ind + (strand ? 1 : -1)
      while 1 <= cur <= length(nodecoord) &&
            nodecoord[cur] <= tx.don[i]
         push!( path, cur )
         cur = cur + (strand ? 1 : -1 )
      end
   end
   path
end

function build_paths_edges( nodecoord::Vector{CoordInt},
                            nodelen::Vector{CoordInt},
                            gene::RefGene )
   paths = Vector{IntSet}()
   for tx in gene.reftx
      curpath = build_annotated_path( nodecoord, nodelen, tx, gene.info.strand )
      push!( paths, curpath )
   end
   paths
end

function path_to_seq( path::IntSet, nodeoffset::Vector{CoordInt},
                      nodelen::Vector{CoordInt}, seq::SGSequence )
   SGSequence(join( map( x->seq[nodeoffset[x]:(nodeoffset[x]+nodelen[x]-1)], collect(path) ) ))
end


#=
The main module for Finch.
=#
module Finch

export init_finch, set_language, set_custom_gen_target, dendro, set_solver, set_stepper, set_specified_steps, set_matrix_free,
        reformat_for_stepper, reformat_for_stepper_fv,
        add_mesh, output_mesh, add_boundary_ID, add_test_function,
        add_initial_condition, add_boundary_condition, add_reference_point, set_rhs, set_lhs, set_lhs_surface, set_rhs_surface,
        set_symexpressions, set_assembly_loops
export build_cache_level, build_cache, build_cache_auto
export sp_parse
export generate_code_layer, generate_code_layer_surface, generate_code_layer_fv
export Variable, add_variable, VariableTransform
export Coefficient, add_coefficient
export Parameter, add_parameter
export Indexer
export FVInfo

### Module's global variables ###
# configuration, output files, etc.
config = nothing;
prob = nothing;
project_name = "unnamedProject";
output_dir = pwd();
language = 0;
gen_framework = 0;
generate_external = false;
solver = nothing;
codegen_params = nothing;
#log
use_log = false;
log_file = "";
log_line_index = 1;
#mesh
mesh_data = nothing;    # The basic element information as read from a MSH file or generated here.
elemental_order = [];   # Determines the order of the elemental loops
needed_grid_types = [false, false, false]; # [CG, DG, FV] true if that type is needed
# FEM specific
grid_data = nothing;    # The full collection of nodes(including internal nodes) and other mesh info in the actual DOF ordering.
dg_grid_data = nothing; # This DG version is only made if using mixed CG/DG. Otherwise it is in grid_data.
geo_factors = nothing;  # Geometric factors
refel = nothing;        # Reference element(s?)
# FVM specific
fv_info = nothing;      # Finite volume info
fv_refel = nothing;     # Reference element for FV
parent_maps = nothing;  # For high order FV
fv_grid = nothing;      # Similar to grid_data, but only first order CG elements, and includes ghost info when partitioned.
fv_geo_factors = nothing;# Geometric factors for FV

#problem variables
var_count = 0;
variables = [];
coefficients = [];
parameters = [];
test_functions = [];
indexers = [];
variable_transforms = [];
#generated functions
genfunc_count = 0;
genfunctions = [];
callback_functions = [];
#rhs
linears = [];
face_linears = [];
#lhs
bilinears = [];
face_bilinears = [];
#assembly loop functions
assembly_loops = [];
#symbolic layer
symexpressions = [[],[],[],[]];
# string versions of generated code for printing
code_strings = [[],[],[],[],[]];
#time stepper
time_stepper = nothing;
specified_dt = 0;
specified_Nsteps = 0;
use_specified_steps = false;

use_cachesim = false;

#handles for custom code gen functions
custom_gen_funcs = [];

###############################################################
include("finch_includes.jl");
include("macros.jl");
include("finch_interface.jl"); # included here after globals are defined
###############################################################

config = Finch_config(); # These need to be initialized here
prob = Finch_prob();

# This has already been initialized on loading, so this is more of a reset function
function init_finch(name="unnamedProject")
    global config = Finch_config();
    global prob = Finch_prob();
    global project_name = name;
    global language = JULIA;
    global framework = 0;
    global generate_external = false;
    global codegen_params = nothing;
    global log_file = "";
    global use_log = false;
    global log_line_index = 1;
    global mesh_data = nothing;
    global needed_grid_types = [false, false, false];
    global grid_data = nothing;
    global dg_grid_data = nothing;
    global geo_factors = nothing;
    global fv_info = nothing;
    global refel = nothing;
    global fv_refel = nothing;
    global elemental_order = [];
    global parent_maps = nothing;
    global fv_grid = nothing;
    global fv_geo_factors = nothing;
    global var_count = 0;
    global variables = [];
    global coefficients = [];
    global parameters = [];
    global test_functions = [];
    global indexers = [];
    global variable_transforms = [];
    global genfunc_count = 0;
    global genfunctions = [];
    global callback_functions = [];
    global linears = [];
    global bilinears = [];
    global face_linears = [];
    global face_bilinears = [];
    global assembly_loops = [];
    global symexpressions = [[],[],[],[]];
    global code_strings = [[],[],[],[],[]];
    global time_stepper = nothing;
    global specified_dt = 0;
    global specified_Nsteps = 0;
    global use_specified_steps = false;
    global use_cachesim = false;

    # check for MPI and initialize
    if @isdefined(MPI)
        MPI.Init();
        config.num_procs = MPI.Comm_size(MPI.COMM_WORLD);
        if config.num_procs < 2
            # If only one process, ignore MPI.
            # MPI.Finalize();
        else
            config.use_mpi = true;
            config.proc_rank = MPI.Comm_rank(MPI.COMM_WORLD);
        end
    end

    # check for thread availability
    config.num_threads = Threads.nthreads();
    if (config.num_threads > 1 || config.num_procs > 1) && config.proc_rank == 0
        println("Initialized with "*string(config.num_procs)*" processes and "*string(config.num_threads)*" threads per proc.");
    end
end

# Sets the code generation target
function set_included_gen_target(lang, framework, dirpath, name; head="")
    #TODO
    println("Premade generation targets currently under construction. Sorry. Reverting to Julia target.");
    return;

    global language = lang;
    global gen_framework = framework;
    global output_dir = dirpath;
    global project_name = name;
    global generate_external = true;
    init_codegenerator(dirpath, name, head);

    # Need to set these three functions

    set_generation_target(get_external_language_elements, generate_external_code_layer, generate_external_files);
end

# Setting a custom target requires three functions
# 1. get_external_language_elements() - file extensions, comment chars etc.
# 2. generate_external_code_layer(var, entities, terms, lorr, vors) - Turns symbolic expressions into code
# 3. generate_external_files(var, lhs_vol, lhs_surf, rhs_vol, rhs_surf) - Writes all files based on generated code
function set_custom_gen_target(lang_elements, code_layer, file_maker, dirpath, name; head="", params=nothing)
    global language = CUSTOM_GEN_TARGET;
    global gen_framework = CUSTOM_GEN_TARGET;
    global output_dir = dirpath;
    global project_name = name;
    global generate_external = true;
    init_code_generator(dirpath, name, head);
    set_generation_target(lang_elements, code_layer, file_maker);
end

# Set parameters used by code generation. This is specific to the target.
function set_codegen_parameters(params)
    global codegen_params = params;
end

# Set the needed discretizations and backend
function set_solver(stype, backend)
    config.linalg_backend = backend;
    if typeof(stype) <: Array
        config.solver_type = MIXED;
        global solver = MixedSolver;
        for s in stype
            if s == CG
                global needed_grid_types[1] = true;
            elseif s == DG
                global needed_grid_types[2] = true;
            elseif s == FV
                global needed_grid_types[3] = true;
            end
        end
    elseif stype == DG
        config.solver_type = stype;
        global solver = DGSolver;
        global needed_grid_types[2] = true;
    elseif stype == CG
        config.solver_type = stype;
        global solver = CGSolver;
        global needed_grid_types[1] = true;
    elseif stype == FV
        config.solver_type = stype;
        global solver = FVSolver;
        global needed_grid_types[3] = true;
    end
    solver.init_solver();

    if backend == PETSC_SOLVER
        # Initialize using the first available library.
        # This should have been set up beforehand.
        if length(PETSc.petsclibs) == 0
            printerr(
"No PETSc library found. Please build PETSc first.
If you have a preferred PETSc library or the one supplied with PETSc.jl
is causing trouble, do this:
\$ export JULIA_PETSC_LIBRARY=/path/to/your/petsc_lib.so
julia> ]build PETSc", fatal=true);
        end
        petsclib = PETSc.petsclibs[1];
        PETSc.initialize(petsclib)
    end
end

# Sets the time stepper
function set_stepper(type, cfl)
    if typeof(type) <: Array
        first_stepper = Stepper(type[1], cfl);
        second_stepper = Stepper(type[2], cfl);

        global time_stepper = [first_stepper, second_stepper];
        global config.stepper = MIXED_STEPPER;
        log_entry("Set mixed time stepper to ["*type[1]*", "*type[2]*"]", 1);
    else
        global time_stepper = Stepper(type, cfl);
        global config.stepper = type;
        log_entry("Set time stepper to "*type, 1);
    end
end

# Time steps can be chosen automatically or specified. This is for manual specifying.
function set_specified_steps(dt, steps)
    global specified_dt = dt;
    global specified_Nsteps = steps;
    global use_specified_steps = true;
    prob.time_dependent = true;
    prob.end_time = dt*steps;
    log_entry("Set time stepper values to dt="*string(dt)*", Nsteps="*string(steps), 1);
end

# Uses matrix free iterative solver with the given max iterations and error tolerance.
function set_matrix_free(max, tol)
    config.linalg_matrixfree = true;
    config.linalg_matfree_max = max;
    config.linalg_matfree_tol = tol;
end

# Adds a mesh and builds the full grid and reference elements.
function add_mesh(mesh; partitions=0)
    global mesh_data = mesh;
    global refel;
    global fv_refel;
    global grid_data;
    global dg_grid_data;
    global fv_grid;
    global geo_factors;
    global fv_geo_factors;
    global fv_info;
    global elemental_order;

    # If no method has been specified, assume CG
    if !(needed_grid_types[1] || needed_grid_types[2] || needed_grid_types[3])
        set_solver(CG, DEFAULT_SOLVER);
    end

    if partitions==0
        np = config.num_procs; # By default each process gets a partition
        config.num_partitions = np;
        config.partition_index = config.proc_rank;
    else
        # If other partitioning strategies are desired.
        # More than one proc may be assigned to each mesh partition.
        # They are grouped by partition number: 0,0,0,1,1,1,...,p,p,p
        np = partitions;
        config.num_partitions = np;
        config.partition_index = Int(floor(((config.proc_rank+0.5) * np) / config.num_procs));
    end
    if config.use_mpi && np > 1
        if config.proc_rank == 0
            # Partition the mesh and send the partition info to each proc to build their subgrids.
            # For now, partition into the number of procs. This will be modified in the future.
            epart = get_element_partitions(mesh, np);

        else # other ranks recieve the partition info
            epart = fill(Cint(-1), mesh.nel);

        end
        # Broadcast partition info to all
        MPI.Bcast!(epart, 0, MPI.COMM_WORLD);

        # Each proc will build a subgrid containing only their elements and ghost neighbors.
        # Make a Grid struct for each needed type
        if needed_grid_types[1] # CG
            (refel, grid_data) = partitioned_grid_from_mesh(mesh_data, epart, grid_type=CG, order=config.basis_order_min);
        end
        if needed_grid_types[2] # DG
            if needed_grid_types[1] # CG also needed
                (refel, dg_grid_data) = partitioned_grid_from_mesh(mesh_data, epart, grid_type=DG, order=config.basis_order_min);
            else
                (refel, grid_data) = partitioned_grid_from_mesh(mesh_data, epart, grid_type=DG, order=config.basis_order_min);
            end
        end
        if needed_grid_types[3] # FV
            (fv_refel, fv_grid) = partitioned_grid_from_mesh(mesh_data, epart, grid_type=FV, order=1);
            # If only FV is used, also set grid_data and refel to this. Just in case.
            if !(needed_grid_types[1] || needed_grid_types[2])
                grid_data = fv_grid;
                refel = fv_refel;
            end
        end

    else
        # Make a Grid struct for each needed type
        if needed_grid_types[1] # CG
            (refel, grid_data) = grid_from_mesh(mesh_data, grid_type=CG, order=config.basis_order_min);
        end
        if needed_grid_types[2] # DG
            if needed_grid_types[1] # CG also needed
                (refel, dg_grid_data) = grid_from_mesh(mesh_data, grid_type=DG, order=config.basis_order_min);
            else
                (refel, grid_data) = grid_from_mesh(mesh_data, grid_type=DG, order=config.basis_order_min);
            end
        end
        if needed_grid_types[3] # FV
            (fv_refel, fv_grid) = grid_from_mesh(mesh_data, grid_type=FV, order=1);
            # If only FV is used, also set grid_data and refel to this. Just in case.
            if !(needed_grid_types[1] || needed_grid_types[2])
                grid_data = fv_grid;
                refel = fv_refel;
            end
        end
    end

    # regular parallel sided elements or simplexes have constant Jacobians, so only store one value per element.
    if ((config.geometry == SQUARE && config.mesh_type == UNIFORM_GRID) ||
        config.dimension == 1 ||
        (config.dimension == 2 && refel.Nfaces == 3) ||
        (config.dimension == 3 && refel.Nfaces == 4) )
        constantJ = true;
    else
        constantJ = false;
    end
    do_faces = false;
    do_vol = false;
    if config.solver_type == DG || config.solver_type == FV || config.solver_type == MIXED
        do_faces = true;
    end
    if config.solver_type == FV || config.solver_type == MIXED
        do_vol = true;
    end

    # FE version is always made?
    geo_factors = build_geometric_factors(refel, grid_data, do_face_detj=do_faces, do_vol_area=do_vol, constant_jacobian=constantJ);
    if needed_grid_types[3] # FV
        fv_geo_factors = build_geometric_factors(fv_refel, fv_grid, do_face_detj=do_faces, do_vol_area=do_vol, constant_jacobian=constantJ);
        fv_info = build_FV_info(fv_grid);
    end

    log_entry("Added mesh with "*string(mesh_data.nx)*" vertices and "*string(mesh_data.nel)*" elements.", 1);
    if np > 1
        e_count = zeros(Int, np);
        for ei=1:length(epart)
            e_count[epart[ei]+1] += 1;
        end
        log_entry("Number of elements in each partition: "*string(e_count), 2);
        log_entry("Full grid has "*string(length(grid_data.global_bdry_index))*" nodes.", 2);

    else
        log_entry("Full grid has "*string(size(grid_data.allnodes,2))*" nodes.", 2);
    end

    # Set elemental loop ordering to match the order from the grid for now.
    elemental_order = 1:grid_data.nel_owned;
end

# Write the mesh to a MSH file
function output_mesh(file, format)
    if config.proc_rank == 0
        write_mesh(file, format, mesh_data);
        log_entry("Wrote mesh data to file.", 1);
    end
end

function set_parent_and_child(p_maps, c_grid, order)
    global parent_maps = p_maps;
    global fv_grid = c_grid;
    dim = config.dimension;
    nfaces = size(fv_grid.element2face,1);
    global fv_refel = build_refel(dim, 1, nfaces, config.elemental_nodes);
    global fv_geo_factors = build_geometric_factors(fv_refel, fv_grid, do_face_detj=true, do_vol_area=true, constant_jacobian=true);
    global fv_info = build_FV_info(fv_grid, order);
    log_entry("Created parent-child grid with "*string(size(c_grid.allnodes,2))*" nodes and "*string(size(c_grid.loc2glb,2))*" elements.", 2);

    # If CELL variables exist, resize their values
    N = size(fv_grid.loc2glb, 2);
    for i=1:length(variables)
        if variables[i].location == CELL
            if variables[i].type == SCALAR
                val_size = (1, N);
            elseif variables[i].type == VECTOR
                val_size = (config.dimension, N);
            elseif variables[i].type == TENSOR
                val_size = (config.dimension*config.dimension, N);
            elseif variables[i].type == SYM_TENSOR
                val_size = (Int((config.dimension*(config.dimension+1))/2), N);
            elseif variables[i].type == VAR_ARRAY
                if typeof(variables[i].indexer) <: Array
                    comps = 1;
                    for j=1:length(variables[i].indexer)
                        comps *= length(variables[i].indexer[j].range);
                    end
                elseif typeof(variables[i].indexer) == Indexer
                    comps = length(variables[i].indexer.range);
                else
                    comps = 1;
                end
                val_size = (comps, N);
            end

            variables[i].values = zeros(val_size);
        end
    end
end

# Maybe remove. This is in grid.jl
function add_boundary_ID(bid, on_bdry)
    add_boundary_ID_to_grid(bid, on_bdry, grid_data);
end

# Defines a test function symbol by creating a special coefficient object.
function add_test_function(v, type)
    varind = length(test_functions) + 1;
    # make SymType
    dim = config.dimension;
    components = 1;
    if type == SCALAR
        components = 1;
    elseif type == VECTOR
        components = dim;
    elseif type == TENSOR
        components = dim*dim;
    elseif type == SYM_TENSOR
        if dim == 1
            components = 1;
        elseif dim == 2
            components = 3;
        elseif dim == 3
            components = 6;;
        elseif dim == 4
            components = 10;
        end
    elseif type == VAR_ARRAY
        components = 1;
    end
    symvar = sym_var(string(v), type, components);

    push!(test_functions, Finch.Coefficient(v, symvar, varind, type, NODAL, [], false););
    log_entry("Set test function symbol: "*string(v)*" of type: "*type, 2);
end

# Adds a variable and allocates everything associated with it.
function add_variable(var)
    global var_count += 1;

    # adjust values arrays
    if grid_data === nothing
        N = 1;
    else
        if var.location == CELL
            if fv_grid === nothing
                N = size(grid_data.loc2glb, 2);
            else
                N = size(fv_grid.loc2glb, 2);
            end
        else
            N = size(grid_data.allnodes,2);
        end
    end

    if var.type == SCALAR
        val_size = (1, N);
        comps = 1;
    elseif var.type == VECTOR
        val_size = (config.dimension, N);
        comps = config.dimension;
    elseif var.type == TENSOR
        val_size = (config.dimension*config.dimension, N);
        comps = config.dimension*config.dimension;
    elseif var.type == SYM_TENSOR
        val_size = (Int((config.dimension*(config.dimension+1))/2), N);
        comps = val_size[1];
    elseif var.type == VAR_ARRAY
        if typeof(var.indexer) <: Array
            comps = 1;
            for i=1:length(var.indexer)
                comps *= length(var.indexer[i].range);
            end
        elseif typeof(var.indexer) == Indexer
            comps = length(var.indexer.range);
        else
            comps = 1;
        end
        val_size = (comps, N);
    end
    var.total_components = comps;

    if language == JULIA || language == 0
        var.values = zeros(val_size);
    end

    # make symbolic layer variable symbols
    var.symvar = sym_var(string(var.symbol), var.type, var.total_components);

    global variables = [variables; var];

    global linears = [linears; nothing];
    global bilinears = [bilinears; nothing];
    global face_linears = [face_linears; nothing];
    global face_bilinears = [face_bilinears; nothing];
    global assembly_loops = [assembly_loops; nothing];
    global symexpressions[1] = [symexpressions[1]; nothing];
    global symexpressions[2] = [symexpressions[2]; nothing];
    global symexpressions[3] = [symexpressions[3]; nothing];
    global symexpressions[4] = [symexpressions[4]; nothing];
    global code_strings[1] = [code_strings[1]; ""];
    global code_strings[2] = [code_strings[2]; ""];
    global code_strings[3] = [code_strings[3]; ""];
    global code_strings[4] = [code_strings[4]; ""];
    global code_strings[5] = [code_strings[5]; ""];

    log_entry("Added variable: "*string(var.symbol)*" of type: "*var.type*", location: "*var.location, 2);
end

# Adds a coefficient with either constant value or some generated function of (x,y,z,t)
function add_coefficient(c, type, location, val, nfuns, element_array=false)
    global coefficients;
    # The values of c will have the same array structure as val
    if typeof(val) <: Array
        vals = Array{Any}(undef,size(val));
    else
        vals = [];
    end
    if nfuns == 0 # constant values
        vals = val;
        if length(vals) == 1 && !(typeof(vals) <: Array)
            vals = [val];
        end

    else # genfunction values
        if typeof(val) <: Array
            ind = length(genfunctions) - nfuns + 1;
            for i=1:length(val)
                if typeof(val[i]) == String
                    vals[i] = genfunctions[ind];
                    ind += 1;
                else
                    vals[i] = val[i];
                end
            end
        else
            push!(vals, genfunctions[end]);
        end
    end
    # If element_array, the last index in size(vals) should be per element.
    # c[n1, nel] for n1 components by nel elements.
    if element_array
        sz = size(vals);
        components = 0;
        for ind=1:length(sz)-1
            components += sz[ind];
        end
        if components == 0
            components = 1;
        end
    else
        components = length(vals);
    end

    symvar = sym_var(string(c), type, components);

    index = length(coefficients) + 1;
    push!(coefficients, Coefficient(c, symvar, index, type, location, vals, element_array));

    if element_array
        log_entry("Added coefficient "*string(c)*" : (array of elemental values)", 2);
    else
        log_entry("Added coefficient "*string(c)*" : "*string(val), 2);
    end

    return coefficients[end];
end

# Adds a parameter entity.
function add_parameter(p, type, val)
    index = length(parameters);
    push!(parameters, Parameter(p, index, type, val));

    log_entry("Added parameter "*string(p)*" : "*string(val), 2);

    return parameters[end];
end

# Adds an Indexer entity
function add_indexer(indexer)
    push!(indexers, indexer);
    if length(indexer.range) < 3
        log_entry("Added indexer "*string(indexer.symbol)*" : "*string(indexer.range));
    else
        log_entry("Added indexer "*string(indexer.symbol)*" : ["*string(indexer.range[1])*", ... "*string(indexer.range[end])*"]");
    end

end

# Defines a variable transform
function add_variable_transform(var1, var2, func)
    push!(variable_transforms, VariableTransform(var1, var2, func));
    log_entry("Added transform "*string(func)*": "*string(var1)*" -> "*string(var2));
    return variable_transforms[end];
end

# Needed to insert parameter expressions into the weak form expressions.
function swap_parameter_xyzt(ex)
    if typeof(ex) == Symbol
        if ex === :x
            return :parameterCoefficientForx ;
        elseif ex === :y
            return :parameterCoefficientFory ;
        elseif ex === :z
            return :parameterCoefficientForz ;
        elseif ex === :t
            return :parameterCoefficientFort ;
        end
    elseif typeof(ex) == Expr
        for i=1:length(ex.args)
            ex.args[i] = swap_parameter_xyzt(ex.args[i]); # Recursively swap
        end
    end
    return ex;
end

# Sets initial condition for the given variable, but does not initialize.
function add_initial_condition(varindex, ex, nfuns)
    global prob;
    while length(prob.initial) < varindex
        prob.initial = [prob.initial; nothing];
    end
    if typeof(ex) <: Array
        vals = [];
        ind = length(genfunctions) - nfuns + 1;
        for i=1:length(ex)
            if typeof(ex[i]) == String
                push!(vals, genfunctions[ind]);
                ind += 1;
            else
                push!(vals, ex[i]);
            end
        end
        prob.initial[varindex] = vals;
    else
        var_components = size(variables[varindex].values,1);
        if typeof(ex) == String
            prob.initial[varindex] = fill(genfunctions[end], var_components);
        else
            prob.initial[varindex] = fill(ex, var_components);
        end
    end

    if length(prob.initial[varindex]) < 10
        log_entry("Initial condition for "*string(variables[varindex].symbol)*" : "*string(prob.initial[varindex]), 2);
    else
        log_entry("Initial condition for "*string(variables[varindex].symbol)*" : "*string(prob.initial[varindex][1:6])*" (truncated for printing)", 2);
    end
    # hold off on initializing till solve or generate is determined.
end

# Evaluates all available intitial conditions setting the values for variables.
function eval_initial_conditions()
    dim = config.dimension;

    # build initial conditions
    for vind=1:length(variables)
        if !(variables[vind].ready) && vind <= length(prob.initial)
            if !(prob.initial[vind] === nothing)
                if variables[vind].location == CELL && !(fv_grid === nothing)
                    # Need to use the fv_grid instead of grid_data
                    this_grid_data = fv_grid;
                    this_geo_factors = fv_geo_factors;
                    this_refel = fv_refel;
                else
                    this_grid_data = grid_data;
                    this_geo_factors = geo_factors;
                    this_refel = refel;
                end

                # Evaluate at nodes
                nodal_values = zeros(length(prob.initial[vind]), size(this_grid_data.allnodes,2));
                for ci=1:length(prob.initial[vind])
                    for ni=1:size(this_grid_data.allnodes,2)
                        if typeof(prob.initial[vind][ci]) <: Number
                            nodal_values[ci,ni] = prob.initial[vind][ci];
                        elseif dim == 1
                            nodal_values[ci,ni] = prob.initial[vind][ci].func(this_grid_data.allnodes[ni],0,0,0);
                        elseif dim == 2
                            nodal_values[ci,ni] = prob.initial[vind][ci].func(this_grid_data.allnodes[1,ni],this_grid_data.allnodes[2,ni],0,0);
                        elseif dim == 3
                            nodal_values[ci,ni] = prob.initial[vind][ci].func(this_grid_data.allnodes[1,ni],this_grid_data.allnodes[2,ni],this_grid_data.allnodes[3,ni],0);
                        end
                    end
                end

                # compute cell averages using nodal values if needed
                if variables[vind].location == CELL
                    nel = size(this_grid_data.loc2glb, 2);
                    for ei=1:nel
                        # e = elemental_order[ei];
                        e = ei;
                        glb = this_grid_data.loc2glb[:,e];
                        vol = this_geo_factors.volume[e];
                        detj = this_geo_factors.detJ[e];

                        for ci=1:length(prob.initial[vind])
                            variables[vind].values[ci,e] = detj / vol * (this_refel.wg' * this_refel.Q * (nodal_values[ci,glb][:]))[1];
                        end
                    end
                else
                    variables[vind].values = nodal_values;
                end

                variables[vind].ready = true;
                log_entry("Built initial conditions for: "*string(variables[vind].symbol));
            end
        end
    end
end

# Sets boundary condition for a variable and BID.
function add_boundary_condition(var, bid, type, ex, nfuns)
    global prob;
    # make sure the arrays are big enough
    if size(prob.bc_func, 1) < var_count || size(prob.bc_func, 2) < bid
        if !(grid_data===nothing)
            nbid = length(grid_data.bids);
        else
            # For some targets no grid is created
            nbid = 1;
        end

        tmp1 = fill("", var_count, nbid);
        tmp2 = Array{Any,2}(undef, (var_count, nbid)); # need to keep this array open to any type
        fill!(tmp2, 0);
        tmp3 = zeros(Int, (var_count, nbid));

        for i=1:size(prob.bc_func,1)
            for j=1:size(prob.bc_func,2)
                tmp1[i,j] = prob.bc_type[i,j];
                tmp2[i,j] = prob.bc_func[i,j];
                tmp3[i,j] = prob.bid[i,j];
            end
        end
        prob.bc_type = tmp1;
        prob.bc_func = tmp2;
        prob.bid = tmp3;
    end

    # Add this boundary condition to the struct
    valstr = "";
    if typeof(ex) <: Array
        vals = [];
        valstr = "[";
        ind = length(genfunctions) - nfuns + 1;
        for i=1:length(ex)
            if typeof(ex[i]) <: Number
                push!(vals, ex[i]);
                valstr *= string(ex[i]);
            elseif typeof(ex[i]) == String
                push!(vals, genfunctions[ind]);
                valstr *= genfunctions[ind].name;
                ind += 1;
            else # callback
                push!(vals, ex[i]);
                valstr *= ex[i].name;
            end
            if i < length(ex)
                valstr *= ", ";
            end
        end
        valstr *= "]";
        prob.bc_func[var.index, bid] = vals;
    else
        var_components = size(var.values,1);
        if typeof(ex) <: Number
            prob.bc_func[var.index, bid] = fill(ex, var_components);
            valstr = string(ex);
        elseif typeof(ex) == String
            prob.bc_func[var.index, bid] = fill(genfunctions[end], var_components);
            valstr = genfunctions[end].name;
        else # callback
            prob.bc_func[var.index, bid] = fill(ex, var_components);
            valstr = ex.name;
        end
    end
    prob.bc_type[var.index, bid] = type;
    prob.bid[var.index, bid] = bid;

    log_entry("Boundary condition: var="*string(var.symbol)*" bid="*string(bid)*" type="*type*" val="*valstr, 2);
end

# A reference point is a single point with a defined value.
# It's treated like a Dirichlet boundary with one boundary node.
function add_reference_point(var, pos, val)
    global prob;
    # make sure the array is big enough
    if size(prob.ref_point,1) < var_count
        tmp = Array{Any,2}(undef, (var_count, 3));
        for i=1:size(prob.ref_point,1)
            tmp[i,:] = prob.ref_point[i,:];
        end
        for i=(size(prob.ref_point,1)+1):var_count
            tmp[i,1] = false;
            tmp[i,2] = [0,0];
            tmp[i,3] = [0];
        end
        prob.ref_point = tmp;
    end
    if typeof(pos) <: Number
        pos = pos*ones(config.dimension);
    end
    if typeof(val) <: Number
        val = val*ones(length(var.symvar));
    end

    # Find the closest vertex to pos
    # The stored pos is actually the index into glbvertex pointing to the closest vertex
    ind = [1,1];
    mindist = 12345;
    nel = size(grid_data.loc2glb,2);
    for ei=1:nel
        for i=1:size(grid_data.glbvertex,1)
            d = 0;
            for comp=1:length(pos)
                d = d + abs(grid_data.allnodes[comp, grid_data.glbvertex[i, ei]] - pos[comp]);
            end
            if d<mindist
                ind = [i,ei];
                mindist = d;
            end
        end
    end

    prob.ref_point[var.index, 1] = true;
    prob.ref_point[var.index, 2] = ind;
    prob.ref_point[var.index, 3] = val;

    log_entry("Reference point: var="*string(var.symbol)*" position="*string(pos)*" value="*string(val), 2);
end

# adds a callback function to the
function add_callback_function(f)
    push!(callback_functions, f);
end

# Sets the RHS(linear) function for the given variable(s).
function set_rhs(var, code="")
    global linears;
    if language == 0 || language == JULIA
        if typeof(var) <:Array
            for i=1:length(var)
                linears[var[i].index] = genfunctions[end];
                code_strings[2][var[i].index] = code;
            end
        else
            linears[var.index] = genfunctions[end];
            code_strings[2][var.index] = code;
        end

    else # external generation
        if typeof(var) <:Array
            for i=1:length(var)
                linears[var[i].index] = code;
            end
        else
            linears[var.index] = code;
        end
    end
end

# Sets the LHS(bilinear) function for the given variable(s).
function set_lhs(var, code="")
    global bilinears;
    if language == 0 || language == JULIA
        if typeof(var) <:Array
            for i=1:length(var)
                bilinears[var[i].index] = genfunctions[end];
                code_strings[1][var[i].index] = code;
            end
        else
            bilinears[var.index] = genfunctions[end];
            code_strings[1][var.index] = code;
        end

    else # external generation
        if typeof(var) <:Array
            for i=1:length(var)
                bilinears[var[i].index] = code;
            end
        else
            bilinears[var.index] = code;
        end
    end
end

# Sets the surface RHS(linear) function for the given variable(s).
function set_rhs_surface(var, code="")
    global face_linears;
    if language == 0 || language == JULIA
        if typeof(var) <:Array
            for i=1:length(var)
                face_linears[var[i].index] = genfunctions[end];
                code_strings[4][var[i].index] = code;
            end
        else
            face_linears[var.index] = genfunctions[end];
            code_strings[4][var.index] = code;
        end

    else # external generation
        if typeof(var) <:Array
            for i=1:length(var)
                face_linears[var[i].index] = code;
            end
        else
            face_linears[var.index] = code;
        end
    end
end

# Sets the surface LHS(bilinear) function for the given variable(s).
function set_lhs_surface(var, code="")
    global face_bilinears;
    if language == 0 || language == JULIA
        if typeof(var) <:Array
            for i=1:length(var)
                face_bilinears[var[i].index] = genfunctions[end];
                code_strings[3][var[i].index] = code;
            end
        else
            face_bilinears[var.index] = genfunctions[end];
            code_strings[3][var.index] = code;
        end

    else # external generation
        if typeof(var) <:Array
            for i=1:length(var)
                face_bilinears[var[i].index] = code;
            end
        else
            face_bilinears[var.index] = code;
        end
    end
end

function set_symexpressions(var, ex, lorr, vors)
    # If the ex is empty, don't set anything
    if ex == [] || ex == [[]]
        return;
    end

    if typeof(var) <:Array
        for i=1:length(var)
            set_symexpressions(var[i], ex, lorr, vors);
        end

    else
        global symexpressions;
        if lorr == LHS
            if vors == "volume"
                ind = 1;
            else
                ind = 2;
            end
        else # rhs
            if vors == "volume"
                ind = 3;
            else
                ind = 4;
            end
        end
        symexpressions[ind][var.index] = ex;
    end
end

function set_assembly_loops(var, code="")
    global assembly_loops;
    if language == 0 || language == JULIA
        if typeof(var) <:Array
            for i=1:length(var)
                assembly_loops[var[i].index] = genfunctions[end];
                code_strings[5][var[i].index] = code;
            end
        else
            assembly_loops[var.index] = genfunctions[end];
            code_strings[5][var.index] = code;
        end

    else # external generation
        if typeof(var) <:Array
            for i=1:length(var)
                assembly_loops[var[i].index] = code;
            end
        else
            assembly_loops[var.index] = code;
        end
    end
end

function finalize()
    finalize_finch();
end

end # module


# Wrapper to CARET library.
module CaretWrapper

importall Orchestra.Types
importall Orchestra.Util

using PyCall
@pyimport rpy2.robjects as RO
@pyimport rpy2.robjects.packages as RP
@pyimport rpy2.robjects.numpy2ri as N2R
N2R.activate()
RP.importr("caret")

export CRTLearner,
       fit!,
       transform!

# Convert vector to R equivalent.
vector_to_r{T<:Int}(vector::Vector{T}) = RO.IntVector(vector)
vector_to_r{T<:Real}(vector::Vector{T}) = RO.FloatVector(vector)
vector_to_r{T<:Bool}(vector::Vector{T}) = RO.BoolVector(vector)
vector_to_r{T<:String}(vector::Vector{T}) = RO.StrVector(vector)
function vector_to_r(vector::Vector{Any})
  vec_eltype = infer_eltype(vector)
  if vec_eltype == Any
    error("Cannot handle R conversion for vector with differing element types.")
  end

  return vector_to_r(convert(Vector{vec_eltype}, vector))
end
vector_to_r(vector::Vector) = error(
  "Cannot handle R conversion for $(typeof(vector))."
)

# Builds R dataframe out of dataset.
# Returns (dataframe, label_factor_levels).
function dataset_to_r_dataframe(
  instances::Matrix, labels=nothing)

  # Build dataframe
  df_dict = Dict()
  for col in 1:size(instances, 2)
    df_dict["X$col"] = vector_to_r(instances[:, col])
  end

  if labels != nothing
    r_labels = RO.FactorVector(labels)
    df_dict["Y"] = r_labels
    return (RO.DataFrame(df_dict), r_labels[:levels])
  else
    return (RO.DataFrame(df_dict), nothing)
  end
end

# CARET wrapper that provides access to all learners.
#
# Options for the specific CARET learner is to be passed
# in `options[:impl_options]` dictionary.
type CRTLearner <: Learner
  model
  options

  function CRTLearner(options=Dict())
    default_options = {
      # Output to train against
      # (:class).
      :output => :class,
      :learner => "svmLinear",
      :impl_options => Dict()
    }
    new(nothing, nested_dict_merge(default_options, options))
  end
end

function fit!(crtw::CRTLearner, instances::Matrix, labels::Vector)
  impl_options = crtw.options[:impl_options]
  crtw.model = Dict()
  crtw.model[:learner] = crtw.options[:learner]

  # Build R dataframe out of dataset
  (r_dataset_df, label_factors) = dataset_to_r_dataframe(instances, labels)

  # Assign label factors
  crtw.model[:label_factors] = collect(label_factors)

  # Train model
  caret_formula = RO.Formula("Y ~ .")
  r_fit_control = pycall(RO.r[:trainControl], PyObject,
    method = "none"
  )
  if isempty(impl_options)
    r_model = pycall(RO.r[:train], PyObject,
      caret_formula,
      method = crtw.model[:learner],
      data = r_dataset_df,
      trControl = r_fit_control,
      tuneLength = 1
    )
  else
    r_model = pycall(RO.r[:train], PyObject,
      caret_formula,
      method = crtw.model[:learner],
      data = r_dataset_df,
      trControl = r_fit_control,
      tuneGrid = RO.DataFrame(impl_options)
    )
  end
  crtw.model[:r_model] = r_model
end

function transform!(crtw::CRTLearner, instances::Matrix)
  (r_instance_df, _) = dataset_to_r_dataframe(instances)
  predictions = collect(pycall(RO.r[:predict], PyObject,
    crtw.model[:r_model],
    newdata = r_instance_df
  ))
  label_factors = crtw.model[:label_factors]
  predictions = map(x -> label_factors[x], predictions)
  return predictions
end

end # module


using JuMP

using SCIP
function scip_optimizer_factory()
    with_optimizer(SCIP.Optimizer)#, limits_time=2 * 60, heuristics_alns_priority=1000000, heuristics_subnlp_priority=1000000)
    # with_optimizer(SCIP.Optimizer,
    #     limits_gap=0.05, limits_time=10 * 60 * 60, display_verblevel=5,
    #     display_width=120, history_valuebased=true, lp_threads=10, branching_preferbinary=true, lp_scaling=false,
    #     branching_allfullstrong_priority=536870911, heuristics_multistart_freq=20, heuristics_multistart_onlynlps=false,
    #     heuristics_mpec_priority=536870911)#heuristics_subnlp_priority=536870911)#, nlp_solver="ipopt", heuristics_nlpdiving_priority=536870911)#,;
end

function scip_optimize_hook(model::JuMP.Model)
    caching_optimizer = backend(model)
    if caching_optimizer.mode == MOIU.AUTOMATIC && caching_optimizer.state == MOIU.EMPTY_OPTIMIZER
        MOIU.attach_optimizer(caching_optimizer)
    end
    mscip = caching_optimizer.optimizer.model.mscip
    SCIP.SCIPsetEmphasis(mscip, SCIP.SCIP_PARAMEMPHASIS_FEASIBILITY, true)#false); println()
    SCIP.SCIPsetPresolving(mscip, SCIP.SCIP_PARAMSETTING_AGGRESSIVE, true)#false); println()
    SCIP.SCIPsetHeuristics(mscip, SCIP.SCIP_PARAMSETTING_AGGRESSIVE, true)#false)
    # SCIP.set_parameter(mscip, "heuristics/alns/freq", 5)
    # SCIP.set_parameter(mscip, "heuristics/subnlp/freq", 5)
    # SCIP.set_parameter(mscip, "heuristics/mpec/freq", 5)
    # SCIP.set_parameter(mscip, "heuristics/feaspump/freq", 5)
    SCIP.set_parameter(mscip, "heuristics/rens/freq", -1)
    SCIP.set_parameter(mscip, "reoptimization/enable", true)
    SCIP.set_parameter(mscip, "limits/time", 2 * 60)
    # SCIP.set_parameter(mscip, "limits/time", 2 * 60 * 60)
    ret = MOI.optimize!(caching_optimizer)
    @assert caching_optimizer.optimizer.model.mscip === mscip
    return ret
end

using AmplNLWriter
AmplNLWriter.setdebug(true)

function bonmin_optimizer_factory()
    with_optimizer(AmplNLWriter.Optimizer, "/home/twan/code/bonmin/Bonmin-1.8.7/build/bin/bonmin")
end

function couenne_optimizer_factory()
    with_optimizer(AmplNLWriter.Optimizer, "/home/twan/code/couenne/couenne")
end

using Ipopt
# ENV["JULIA_IPOPT_LIBRARY_PATH"] = "/home/twan/code/ipopt/CoinIpopt/build/lib"
# ENV["JULIA_IPOPT_EXECUTABLE_PATH"] = "/home/twan/code/ipopt/CoinIpopt/build/bin"
# import Pkg; Pkg.build("Ipopt")

function ipopt_optimizer_factory()
    with_optimizer(Ipopt.Optimizer)
end

function complementarity_reformulate!(model::JuMP.Model)
    for var in JuMP.all_variables(problem.model)
        if JuMP.is_binary(var)
            @constraint model var * (1 - var) == 0
            JuMP.unset_binary(var)
        end
    end
end

# using Alpine
# function alpine_optimizer_factory()
#     with_optimizer(Alpine.Optimizer, nlp_optimizer=Ipopt.Optimizer(print_level=0), mip_optimizer=Gurobi.Optimizer(OutputFlag=0))
# end

# using Juniper
# function juniper_optimizer_factory()
#     params = Dict{Symbol,Any}()
#     params[:nl_solver] = with_optimizer(Ipopt.Optimizer, print_level=0)
#     # params[:mip_solver] = with_optimizer(Gurobi.Optimizer, OutputFlag=0)
#     params[:mip_solver] = with_optimizer(CPLEX.Optimizer, CPX_PARAM_SCRIND=0)
#     # Note to self: if you get "Cannot set bounds because variable is of type: BINARY", use LinQuadOptInterface master (needs #91)
#     params[:feasibility_pump] = false
#     with_optimizer(Juniper.Optimizer, params)
# end

using BARON
function baron_optimizer_factory()
    with_optimizer(BARON.Optimizer;
         CplexLibName="/opt/ibm/ILOG/CPLEX_Studio128/cplex/bin/x86-64_linux/libcplex1280.so",
         threads=Sys.CPU_THREADS ÷ 2, MaxTime=5 * 60 * 60, PrTimeFreq=5.,
         AllowFilterSD=1, AllowFilterSQP=1, AllowIpopt=1#=, NumLoc=20, LocRes=1=#)
end

using CPLEX
function cplex_optimizer_factory_continuous()
    # https://www.ibm.com/support/knowledgecenter/en/SSSA5P_12.9.0/ilog.odms.cplex.help/CPLEX/Parameters/topics/OptimalityTarget.html

    # Parameters tried:
    # CPX_PARAM_OPTIMALITYTARGET = 2: CPX_OPTIMALITYTARGET_FIRSTORDER - not available for MIP
    # CPX_PARAM_FPHEUR = 1 - not available for MIQCP
    # CPX_PARAM_DEPIND = 3 # dependency checking: turn on at beginning and at end of preprocessing
    # CPX_PARAM_MIQCPSTRAT = 1 - solve QCPs as subproblems
    # CPXPARAM_MIP_Cuts_LiftProj = 3 - Generate lift-and-project cuts very aggressively
    # CPX_PARAM_MIPEMPHASIS = CPLEX.CPX_MIPEMPHASIS_HIDDENFEAS - applies considerable additional effort toward finding high quality feasible solutions that are difficult to locate

    # Parameters to try:
    # CPX_PARAM_MIQCPSTRAT
    # CPX_MIPEMPHASIS_FEASIBILITY
    # barrier-related options
    with_optimizer(CPLEX.Optimizer,
        CPX_PARAM_OPTIMALITYTARGET = CPLEX.CPX_OPTIMALITYTARGET_AUTO,
        CPX_PARAM_DEPIND = 3,
        CPXPARAM_MIP_Cuts_LiftProj = 3)
end

function cplex_optimizer_factory()
    with_optimizer(CPLEX.Optimizer)
end

using Gurobi
function gurobi_optimizer_factory()
    with_optimizer(Gurobi.Optimizer)
end