text:
#///////////////////////////////////////
#// File Name: rs_sac_controller.jl
#// Author: Haruki Nishimura (hnishimura@stanford.edu)
#// Date Created: 2021/01/29
#// Description: Risk Sensitive SAC Controller
#///////////////////////////////////////

using DataStructures
using LinearAlgebra
using Printf
using Random
using RobotOS
import Convex: Variable, norm, quadform, minimize, dot, solve!
#using SCS
using ECOS

struct ControlParameter <: Parameter
    eamax::Float64  # Maximum abolute value of acceleration
    tcalc::Float64  # Allocated control computation time
    dtexec::Vector{Float64} # Control insertion duration candidates
    dtr::Float64    # Replanning time interval
    u_nominal_base::Vector{Float64} # baseline nominal control
    u_nominal_cand::Vector{Vector{Float64}} # Array of nominal control candidates (including u_nominal_base)
    nominal_search_depth::Int64
    improvement_threshold::Float64
    constraint_time::Union{Nothing, Float64}

    function ControlParameter(eamax::Float64, tcalc::Float64, dtexec::Vector{Float64},
                              dtr::Float64, u_nominal_base::Vector{Float64},
                              u_nominal_cand::Vector{Vector{Float64}},
                              nominal_search_depth::Int64,
                              improvement_threshold::Float64=0.0;
                              constraint_time::Union{Nothing, Float64}=nothing)
        @assert all(norm.(u_nominal_cand) .<= eamax)
        @assert in(u_nominal_base, u_nominal_cand)
        @assert tcalc <= dtr
        @assert nominal_search_depth >= 1
        @assert improvement_threshold >= 0.0
        @assert isnothing(constraint_time) || constraint_time > 0.0
        return new(eamax, tcalc, dtexec, dtr, u_nominal_base, u_nominal_cand,
                   nominal_search_depth, improvement_threshold,
                   constraint_time)
    end
end

mutable struct RSSACController
    sim_param::SimulationParameter
    cnt_param::ControlParameter
    predictor::Predictor

    u_schedule::OrderedDict{Time, Vector{Float64}}
    prediction_dict::Union{Nothing, Dict{String, Array{Float64, 3}}}
    sim_result::Union{Nothing, SimulationResult}
    tcalc_actual::Union{Nothing, Float64}
    #u_init_time::Union{Nothing, Time}
    #u_last_time::Union{Nothing, Time}
    #u_value::Union{Nothing, Vector{Float64}}

    prediction_task::Union{Nothing, Task}
    control_update_task::Union{Nothing, Task}

    u_schedule_tmp::Union{Nothing, OrderedDict{Time, Vector{Float64}}}
    prediction_dict_tmp::Union{Nothing, Dict{String, Array{Float64, 3}}}
    sim_result_tmp::Union{Nothing, SimulationResult}
    tcalc_actual_tmp::Union{Nothing, Float64}
    #u_init_time_tmp::Union{Nothing, Time}
    #u_last_time_tmp::Union{Nothing, Time}
    #u_value_tmp::Union{Nothing, Vector{Float64}}
end

function RSSACController(predictor::Predictor,
                         u_schedule::OrderedDict{Time, Vector{Float64}},
                         sim_param::SimulationParameter,
                         cnt_param::ControlParameter)
    return RSSACController(sim_param, cnt_param, predictor, u_schedule,
                           nothing, nothing, nothing, #nothing, nothing, nothing,
                           nothing, nothing,
                           nothing, nothing, nothing, nothing)#, nothing, nothing,
                           #nothing)
end

# # main control functions below
function control!(controller::RSSACController,
                  current_time::Time;
                  nominal_control::Bool=false,
                  log::Union{Nothing, Vector{Tuple{Time, String}}}=nothing)
    if !isnothing(controller.prediction_task) &&
        istaskdone(controller.prediction_task)
        if !isnothing(log)
            msg = "New prediction is available to the controller."
            push!(log, (current_time, msg))
        end
        controller.prediction_dict = copy(controller.prediction_dict_tmp);
        controller.prediction_task = nothing;
    end

    if !isnothing(controller.control_update_task) &&
        !isnothing(controller.u_schedule_tmp) &&
        istaskdone(controller.control_update_task)
        if !isnothing(log)
            msg = "New SAC control is available to the controller."
            push!(log, (current_time, msg))
        end
        controller.sim_result = deepcopy(controller.sim_result_tmp);
        controller.tcalc_actual = controller.tcalc_actual_tmp;
        #controller.u_init_time = controller.u_init_time_tmp;
        #controller.u_last_time = controller.u_last_time_tmp;
        #controller.u_value = copy(controller.u_value_tmp);
        controller.control_update_task = nothing;
        # update u_schedule as well.
        #=
        if controller.u_init_time < current_time
            if controller.u_last_time < current_time
                # This means that the sac control update was not performed in tcalc time.
                # In this case, just skip this part so the controller uses the nominal control.
                t = @sprintf "Time %.2f" round(to_sec(current_time), digits=5)
                @warn "$(t) [sec]: SAC control command is too old. Ignoring this SAC control."
            else
                # Rewrite u_init_time to make sure it is not in the past.
                controller.u_init_time = current_time;
            end
        end
        =#
        for time in keys(controller.u_schedule_tmp)
            if  !haskey(controller.u_schedule, time);
                continue;
            else
                controller.u_schedule[time] =
                    controller.u_schedule_tmp[time]
            end
        end
        controller.u_schedule_tmp = nothing;
    end

    if nominal_control
        pop!(controller.u_schedule, current_time);
        u = controller.cnt_param.u_nominal_base;
    else
        u = pop!(controller.u_schedule, current_time);
    end
    if !isnothing(log)
        msg = "Control: $(u) is applied to the system."
        push!(log, (current_time, msg))
    end

    # Insert new schedule
    new_control_time = collect(keys(controller.u_schedule))[end] +
                       Duration(controller.sim_param.dtc);
    @assert !haskey(controller.u_schedule, new_control_time)
    push!(controller.u_schedule, new_control_time => controller.cnt_param.u_nominal_base)

    return u
end

function adjust_old_prediction!(controller::RSSACController,
                                previous_ado_pos_dict::Dict{String, Vector{Float64}},
                                latest_ado_pos_dict::Dict{String, Vector{Float64}})
    @assert keys(controller.prediction_dict) == keys(previous_ado_pos_dict)

    ado_agents_removed = setdiff(keys(previous_ado_pos_dict),
                                 keys(latest_ado_pos_dict));
    ado_agents_added = setdiff(keys(latest_ado_pos_dict),
                               keys(previous_ado_pos_dict));
    keys_list = collect(keys(controller.prediction_dict))
    for key in keys_list
        if in(key, ado_agents_removed)
            # If this ado_agent is removed in latest_ado_pos_dict, then remove
            # it from prediction_dict.
            pop!(controller.prediction_dict, key)
        else
            if typeof(controller.predictor) != OraclePredictor # skip this step for OraclePredictor as predictions are perfect.
                # If this ado_agent is existent in both previous and latest
                # ado_pos_dict, then reuse and adjust the previous prediction.
                (diff_x, diff_y) = latest_ado_pos_dict[key] -
                                   previous_ado_pos_dict[key];
                controller.prediction_dict[key][:, :, 1] .+= diff_x;
                controller.prediction_dict[key][:, :, 2] .+= diff_y;
            end
        end
    end
    for key in ado_agents_added
        # Treat this newly added ado_agent as a static obstacle
        # (until the new prediction becomes available to the controller)
        num_controls = length(controller.cnt_param.u_nominal_cand)^controller.cnt_param.nominal_search_depth
        controller.prediction_dict[key] =
            zeros(controller.sim_param.num_samples*num_controls,
                  controller.sim_param.prediction_steps, 2);
        controller.prediction_dict[key][:, :, 1] .+= latest_ado_pos_dict[key][1];
        controller.prediction_dict[key][:, :, 2] .+= latest_ado_pos_dict[key][2];
    end
    @assert keys(controller.prediction_dict) == keys(latest_ado_pos_dict)
end

function schedule_prediction!(controller::RSSACController,
                              ado_pos_dict::Dict,
                              previous_ado_pos_dict::Union{Nothing, Dict{String, Vector{Float64}}}=nothing,
                              e_init::Union{Nothing, RobotState}=nothing);
    if !isnothing(previous_ado_pos_dict)
        adjust_old_prediction!(controller, previous_ado_pos_dict,
                               convert_nodes_to_str(reduce_to_positions(ado_pos_dict)));
    end
    if typeof(controller.predictor) == TrajectronPredictor
        controller.prediction_task = @task begin
            if controller.predictor.param.use_robot_future
                @assert !isnothing(e_init)
                robot_present_and_future=
                    get_robot_present_and_future(e_init,
                                                 controller.u_schedule,
                                                 controller.sim_param,
                                                 controller.cnt_param);

            else
                robot_present_and_future = nothing;
            end
            controller.prediction_dict_tmp =
                sample_future_ado_positions!(controller.predictor, ado_pos_dict,
                                             robot_present_and_future);
            if !controller.predictor.param.use_robot_future
                num_controls = length(controller.cnt_param.u_nominal_cand)^controller.cnt_param.nominal_search_depth;
                # each value has to be (num_samples*num_controls, prediction_steps, 2) array
                for key in keys(controller.prediction_dict_tmp)
                    controller.prediction_dict_tmp[key] =
                        repeat(controller.prediction_dict_tmp[key],
                               outer=(num_controls, 1, 1));
                end
            end
            # GC.gc();
        end
    elseif typeof(controller.predictor) == OraclePredictor
        controller.prediction_task = @task begin
            controller.prediction_dict_tmp =
                sample_future_ado_positions!(controller.predictor, ado_pos_dict);
            num_controls = length(controller.cnt_param.u_nominal_cand)^controller.cnt_param.nominal_search_depth;
            # each value has to be (num_samples*num_controls, prediction_steps, 2) array
            for key in keys(controller.prediction_dict_tmp)
                controller.prediction_dict_tmp[key] =
                    repeat(controller.prediction_dict_tmp[key],
                           outer=(num_controls, 1, 1));
            end
            # GC.gc();
        end
    elseif typeof(controller.predictor) == GaussianPredictor
        controller.prediction_task = @task begin
            controller.prediction_dict_tmp =
                sample_future_ado_positions!(controller.predictor,
                                             convert_nodes_to_str(ado_pos_dict));
            num_controls = length(controller.cnt_param.u_nominal_cand)^controller.cnt_param.nominal_search_depth;
            # each value has to be (num_samples*num_controls, prediction_steps, 2) array
            for key in keys(controller.prediction_dict_tmp)
                controller.prediction_dict_tmp[key] =
                    repeat(controller.prediction_dict_tmp[key],
                           outer=(num_controls, 1, 1));
            end
        end
    else
        @error "Type of controller.predictor: $(typeof(controller.predictor)) not supported."
    end
    schedule(controller.prediction_task);
end

function schedule_control_update!(controller::RSSACController,
                                  w_init::WorldState,
                                  target_trajectory::Trajectory2D;
                                  log::Union{Nothing, Vector{Tuple{Time, String}}}=nothing)
    if !isnothing(controller.prediction_task) &&
        istaskdone(controller.prediction_task)
        if !isnothing(log)
            msg = "New prediction is available to the controller."
            push!(log, (w_init.t, msg))
        end
        controller.prediction_dict = copy(controller.prediction_dict_tmp);
        controller.prediction_task = nothing;
    end
    if !isnothing(log)
        msg = "New SAC control computation is scheduled."
        push!(log, (w_init.t, msg))
    end
    u_schedule_copied = copy(controller.u_schedule);
    controller.control_update_task = @task begin
        controller.tcalc_actual_tmp,
        controller.u_schedule_tmp, controller.sim_result_tmp =
            sac_control_update(w_init, u_schedule_copied,
                               target_trajectory,
                               controller.prediction_dict,
                               controller.sim_param, controller.cnt_param);
    end
    schedule(controller.control_update_task);
end

# # Helper functions below
# convert u_array (array) to u_schedule (ordereddict)
@inline function convert_to_schedule(t_init::Time,
                                     u_array::Vector{Vector{Float64}},
                                     sim_param::SimulationParameter)
    u_schedule = OrderedDict{Time, Vector{Float64}}();
    plan_horizon = length(u_array)*sim_param.dtc;
    t = t_init;
    for ii = 1:length(u_array)
        u_schedule[t] = u_array[ii];
        t += Duration(sim_param.dtc);
    end
    return u_schedule
end

@inline function get_nominal_u_arrays(u_schedule::OrderedDict{Time, Vector{Float64}},
                                      modification_init_time::Time,
                                      modification_final_time::Time,
                                      cnt_param::ControlParameter,
                                      dto::Float64)
    @assert issorted(u_schedule);
    ordered_times = collect(keys(u_schedule))

    num_controls = length(cnt_param.u_nominal_cand);
    u_arrays = Vector{Vector{Vector{Float64}}}(undef, num_controls^cnt_param.nominal_search_depth);

    @assert cnt_param.u_nominal_cand[1] == [0.0, 0.0];
    u_arrays[1] = copy(collect(values(u_schedule)));
    if cnt_param.dtexec == [0.0] # SAC is not used, so set this nominal to zero
        fill!(u_arrays[1], [0.0, 0.0])
    end

    for ii = 2:length(u_arrays)
        @inbounds u_arrays[ii] = copy(collect(values(u_schedule)));
        cnt_idx_array = reverse(digits(ii - 1, base=num_controls, pad=cnt_param.nominal_search_depth) .+ 1);
        mod_init_time_tmp = modification_init_time;
        mod_final_time_tmp = modification_final_time;
        for idx in cnt_idx_array
            init_idx = findfirst(x -> x == mod_init_time_tmp, ordered_times);
            final_idx = findfirst(x -> x == mod_final_time_tmp, ordered_times);
            if !(isnothing(init_idx) || isnothing(final_idx))
                for jj = init_idx:final_idx
                    @inbounds u_arrays[ii][jj] = copy(cnt_param.u_nominal_cand[idx]);
                end
                mod_init_time_tmp += Duration(dto);
                mod_final_time_tmp += Duration(dto);
            end
        end
    end
    return u_arrays
end

@inline function get_nominal_u_arrays(u_schedule::OrderedDict{Time, Vector{Float64}},
                                      sim_param::SimulationParameter,
                                      cnt_param::ControlParameter)
    # make sure nominal control modification only happens after tcalc [s]
    # to not override the sac control from previous iteration.
    curr_time = minimum(keys(u_schedule));
    u_nominal_mod_init_time = curr_time + Duration(cnt_param.tcalc);
    u_nominal_mod_final_time = u_nominal_mod_init_time +
                               Duration(sim_param.dto) -
                               Duration(sim_param.dtc);
    return get_nominal_u_arrays(u_schedule, u_nominal_mod_init_time,
                                u_nominal_mod_final_time, cnt_param,
                                sim_param.dto)
end

# for Trajectron robot-future-conditoinal models
@inline function get_robot_present_and_future(e_init::RobotState,
                                      u_schedule::OrderedDict{Time, Vector{Float64}},
                                      sim_param::SimulationParameter,
                                      cnt_param::ControlParameter)
    # compute u_arrays from u_schedule and nominal control candidates
    u_arrays = get_nominal_u_arrays(u_schedule, sim_param, cnt_param);
    # process u_arrays
    u_array_gpu = cu(process_u_arrays(u_arrays));
    # ego state simulation
    ex_array_gpu = simulate_forward(e_init, u_array_gpu, sim_param)
    slice_factor = Int64(round(sim_param.dto/sim_param.dtc, digits=5));
    ex_array_cpu = collect(ex_array_gpu[:, 1:slice_factor:end, :])
    @assert size(ex_array_cpu, 2) == sim_param.prediction_steps + 1

    # (num_controls, 1 + prediction_steps, 6) array where the last dimension
    # is [pos_x, pos_y, vel_x, vel_y, acc_x, acc_y]. accelerations are computed
    # by finite-differencing velocity snapshots.
    robot_present_and_future = Array{Float64, 3}(undef, size(ex_array_cpu, 1),
                                                        size(ex_array_cpu, 2),
                                                        6);
    robot_present_and_future[:, :, 1:4] = ex_array_cpu;
    acc = diff(ex_array_cpu[:, :, 3:4], dims=2)./sim_param.dto
    robot_present_and_future[:, 1:end-1, 5:6] = acc;
    # pad last acceleration
    robot_present_and_future[:, end, 5:6] = robot_present_and_future[:, end-1, 5:6]

    return robot_present_and_future
end

@inline function get_control_coeffs(sampled_total_costs::Vector{Float64},
                                    e_state_array::Vector{RobotState},
                                    e_costate_array::Array{Float64, 3},
                                    valid_timesteps::Int64,
                                    sim_param::SimulationParameter,
                                    cnt_param::ControlParameter)
    # a trick to prevent underflow or overflow in computing sum of exponentials
    total_cost_extreme = nothing;
    if sim_param.cost_param.σ_risk >= 0.0
        total_cost_extreme = maximum(sampled_total_costs);
    else
        total_cost_extreme = minimum(sampled_total_costs);
    end
    coeffs = sim_param.cost_param.σ_risk.*(sampled_total_costs .- total_cost_extreme);

    # denominator
    den = sum(exp.(coeffs), dims=1);
    num = sum(exp.(coeffs).*
              [e_costate_array[:, ii, jj] for jj = 1:sim_param.num_samples, ii = 1:valid_timesteps],
              dims=1);

    coeff_array = transpose.(num./den);
    coeff_array .*= reshape(transition_control_coeff.(e_state_array[1:valid_timesteps]),
                           1, valid_timesteps);
    coeff_matrix = hcat(transpose.(coeff_array)...)

    return coeff_matrix
end

@inline function get_control_coeffs(sim_result::SimulationResult,
                            sim_param::SimulationParameter,
                            cnt_param::ControlParameter)
    # timesteps between current and (first measurement time + tcalc)
    #valid_timesteps = Int64((cnt_param.dtr + cnt_param.tcalc)/sim_param.dtc);
    sim_horizon = sim_param.prediction_steps*sim_param.dto;
    valid_timesteps = Int64(round(sim_horizon/sim_param.dtc, digits=5));
    coeff_matrix = get_control_coeffs(sim_result.sampled_total_costs,
                                      sim_result.e_state_array,
                                      sim_result.e_costate_array,
                                      valid_timesteps, sim_param, cnt_param);

    if !isnothing(cnt_param.constraint_time)
        valid_timesteps_constraint = size(sim_result.e_costate_array_constraint, 2) - 1;
        @assert valid_timesteps_constraint <= valid_timesteps
        coeff_matrix_constraint =
            get_control_coeffs(sim_result.sampled_total_costs_constraint,
                               sim_result.e_state_array,
                               sim_result.e_costate_array_constraint,
                               valid_timesteps_constraint, sim_param, cnt_param)
    else
        coeff_matrix_constraint = nothing
    end

    return coeff_matrix, coeff_matrix_constraint
end

struct ControlSchedule
    u::Vector{Float64}
    cost::Float64
    t::Time
end

@inline function solve_multi_qcqp(u_nominal_array::Vector{Vector{Float64}},
                                  coeff_matrix::Matrix{Float64},
                                  coeff_matrix_constraint::Matrix{Float64},
                                  sim_param::SimulationParameter,
                                  cnt_param::ControlParameter)
    u = Variable(2);
    cost_array = Vector{Float64}(undef, size(coeff_matrix_constraint, 2));
    u_array = Vector{Vector{Float64}}(undef, size(coeff_matrix_constraint, 2));
    Cu = sim_param.cost_param.Cu;
    for ii = 1:size(coeff_matrix_constraint, 2)
        coeff_vec = coeff_matrix[:, ii];
        u_nominal_vec = u_nominal_array[ii];
        coeff_vec_constraint = coeff_matrix_constraint[:, ii];
        if all(isapprox.(coeff_vec_constraint, 0.0, atol=1e-4))
            u_array[ii] = u_nominal_vec;
            cost_array[ii] = 0.0;
        else
            objective = 0.5*quadform(u, Cu) - dot(coeff_vec, (u - u_nominal_vec)) - 0.5*u_nominal_vec'*Cu*u_nominal_vec;
            problem = minimize(objective);
            problem.constraints += 0.5*quadform(u, Cu) - dot(coeff_vec_constraint, (u - u_nominal_vec)) -
                                   0.5*u_nominal_vec'*Cu*u_nominal_vec <= 0.0
            problem.constraints += norm(u, 2) <= cnt_param.eamax
            #solve!(problem, SCSSolver(verbose=false));
            solve!(problem, () -> ECOS.Optimizer(verbose=false));
            u_array[ii] = Float64[u.value[1], u.value[2]];
            cost_array[ii] = problem.optval;
        end
    end
    return u_array, cost_array
end

@inline function get_control_schedule(sim_result::SimulationResult,
                              u_nominal_array::Vector{Vector{Float64}},
                              coeff_matrix::Matrix{Float64},
                              sim_param::SimulationParameter,
                              cnt_param::ControlParameter,
                              coeff_matrix_constraint::Union{Nothing, Matrix{Float64}}=nothing)
    Cu = sim_param.cost_param.Cu;
    Cu_inv = inv(sim_param.cost_param.Cu);
    control_schedule_array = Vector{ControlSchedule}(undef,
                                                     size(coeff_matrix, 2));
    t_current = sim_result.e_state_array[1].t;

    if isnothing(coeff_matrix_constraint)
        time = t_current;
        for t = 1:size(coeff_matrix, 2)
            coeff_vec = coeff_matrix[:, t];
            u = -Cu_inv*coeff_vec;
            eamax = cnt_param.eamax;
            if norm(u) > eamax
                u = u./norm(u).*eamax; # This is valid as long as Cu is diagonal.
            end
            #cost = 1/2*u'*Cu*u + dot(u, coeff_vec);
            cost = 1/2*u'*Cu*u + dot((u - u_nominal_array[t]), coeff_vec) -
                   1/2*u_nominal_array[t]'*Cu*u_nominal_array[t];
            control_schedule_array[t] = ControlSchedule(u, cost, time);
            time += Duration(sim_param.dtc);
        end
    else
        u_array_constraint, cost_constraint =
            solve_multi_qcqp(u_nominal_array,
                             coeff_matrix,
                             coeff_matrix_constraint,
                             sim_param, cnt_param);
        time = t_current;
        for t = 1:size(coeff_matrix_constraint, 2)
            control_schedule_array[t] =
                ControlSchedule(u_array_constraint[t], cost_constraint[t], time);
                time += Duration(sim_param.dtc);
        end
        for t = size(coeff_matrix_constraint, 2)+1:size(coeff_matrix, 2)
            coeff_vec = coeff_matrix[:, t];
            u = -Cu_inv*coeff_vec;
            eamax = cnt_param.eamax;
            if norm(u) > eamax
                u = u./norm(u).*eamax; # This is valid as long as Cu is diagonal.
            end
            #cost = 1/2*u'*Cu*u + dot(u, coeff_vec);
            cost = 1/2*u'*Cu*u + dot((u - u_nominal_array[t]), coeff_vec) -
                   1/2*u_nominal_array[t]'*Cu*u_nominal_array[t];
            control_schedule_array[t] = ControlSchedule(u, cost, time);
                time += Duration(sim_param.dtc);
        end
    end
    return control_schedule_array
end

@inline function determine_control_time(
                                sim_result::SimulationResult,
                                control_schedule_array::Vector{ControlSchedule},
                                sim_param::SimulationParameter,
                                cnt_param::ControlParameter)
    # if dtexec is smaller than dtc, then use dtexec == dtc.
    if maximum(cnt_param.dtexec) < sim_param.dtc
        dtexec = sim_param.dtc;
    else
        dtexec = maximum(cnt_param.dtexec);
    end
    t_current = sim_result.e_state_array[1].t;
    t_allowed_min = t_current +
                    Duration(cnt_param.tcalc) +
                    Duration(dtexec);
    t_allowed_max = t_current +
                    Duration(sim_param.prediction_steps*sim_param.dto);
    if to_nsec(t_allowed_max - t_allowed_min) < 0.
        throw(ErrorException("t_min: $(to_sec(t_allowed_min)) > t_max:
                             $(to_sec(t_allowed_max)). There is no feasible
                             control insertion time."))
    end
    control_schedule_filtered =
        control_schedule_array[[t_allowed_min <= s.t <= t_allowed_max for s in
                                control_schedule_array]];
    control_chosen =
        control_schedule_filtered[findmin([s.cost for s in
                                  control_schedule_filtered])[2]];
    return control_chosen
end

@inline function perturbation_candidates(current_time::Time,
                                 control_chosen::ControlSchedule, u_nominal_array::Vector{Vector{Float64}},
                                 dtexec::Vector{Float64}, sim_param::SimulationParameter)
    u_arrays = Vector{Vector{Vector{Float64}}}();
    for dte in dtexec
        u_array = copy(u_nominal_array);
        final_idx = Int64(round(to_nsec(control_chosen.t - current_time)/
                                to_nsec(Duration(sim_param.dtc)),
                                digits=5));
        if dte <= sim_param.dtc
            if dte == 0.0
                push!(u_arrays, u_array)
            else
                u_array[final_idx] = control_chosen.u*(dte / sim_param.dtc);
                push!(u_arrays, u_array);
            end
        else
            start_time = control_chosen.t - Duration(dte);
            start_idx = Int64(round(to_nsec(start_time - current_time)/
                                    to_nsec(Duration(sim_param.dtc)),
                                    digits=5)) + 1
            for ii = start_idx:final_idx
                u_array[ii] = control_chosen.u
            end
            push!(u_arrays, u_array);
        end
    end
    return u_arrays
end

function get_perturbed_schedule(w_init::WorldState,
                                sim_result::SimulationResult,
                                u_perturb_arrays::Vector{Vector{Vector{Float64}}},
                                target_trajectory::Trajectory2D,
                                prediction_dict::Dict{String, Array{Float64, 3}},
                                sim_param::SimulationParameter,
                                cnt_param::ControlParameter)
    # process u_arrays
    u_array = process_u_arrays(u_perturb_arrays);
    u_array_gpu = cu(u_array);

    # ego state simulation
    ex_array_gpu = simulate_forward(w_init.e_state, u_array_gpu, sim_param);

    # process prediction_dict for further computation
    # note that each value in prediction_dict has to be (num_samples*num_controls, prediction_steps, 2) array
    num_controls = length(u_perturb_arrays)
    ap_array_gpu = repeat(sim_result.ap_array_gpu, outer=(num_controls, 1, 1, 1));
    measurement_schedule = sim_result.measurement_schedule;

    # ratio between dto and dtc
    expansion_factor = Int64(sim_param.dto/sim_param.dtc);
    # number of nominal control candidates
    num_controls = size(ex_array_gpu, 1);

    # total simulation timesteps
    total_timesteps = size(ex_array_gpu, 2);
    # timesteps between current and first measurement time
    init_timesteps =
        Int64(to_nsec(measurement_schedule[1] - w_init.t)/
              (sim_param.dtc*1e9));
    # timesteps between final measurement time and final e_state time
    remaining_timesteps = total_timesteps -
                          (init_timesteps + expansion_factor*(sim_param.prediction_steps - 1));
    # time_idx_ap_array
    time_idx_ap_array = Vector{Int64}(undef, total_timesteps);
    time_idx_ap_array[1:init_timesteps] .= 1;
    time_idx_ap_array[init_timesteps+1:end-remaining_timesteps] = repeat(2:1:sim_param.prediction_steps,
                                                                         inner=expansion_factor);
    time_idx_ap_array[end-remaining_timesteps+1:end] .= sim_param.prediction_steps + 1;

    # control_idx_ex_array
    control_idx_ex_array = repeat(1:1:num_controls, inner=sim_param.num_samples)

    time_idx_ap_array_gpu = CuArray{Int32}(time_idx_ap_array);
    control_idx_ex_array_gpu = CuArray{Int32}(control_idx_ex_array);

    # get target_pos array
    target_pos_array = get_target_pos_array(ex_array_gpu, w_init, target_trajectory, sim_param)
    target_pos_array_gpu = cu(target_pos_array);

    # compute costs
    cost_result = compute_costs(ex_array_gpu, u_array_gpu, ap_array_gpu,
                                time_idx_ap_array_gpu, control_idx_ex_array_gpu,
                                target_pos_array_gpu, sim_param.cost_param);

    # integrate costs
    total_cost_per_control_sample, risk_per_control = integrate_costs(cost_result, sim_param);
    if !isnothing(cnt_param.constraint_time)
        ~, risk_per_control_constraint = integrate_costs(cost_result, sim_param,
                                                         cnt_param.constraint_time);
        valid_indices = findall(risk_per_control_constraint .<= sim_result.risk_constraint);
        if isempty(valid_indices)
            time = @sprintf "Time %.2f" round(to_sec(w_init.t), digits=5)
            @info "$(time) [sec]: SAC did not find feasible improvement"
            nominal_schedule = convert_to_schedule(w_init.t, u_perturb_arrays[1], sim_param) # revert back to nominal control (assuming that 0.0 is the first element in dtexec)
            return sim_result.risk_val, nominal_schedule, 1
        end
    else
        valid_indices = 1:length(risk_per_control);
    end

    min_risk, control_idx = findmin(risk_per_control[valid_indices])
    if min_risk - sim_result.risk_val >= -cnt_param.improvement_threshold || length(valid_indices) == 1 # revert back to nominal control (assuming that 0.0 is the first element in dtexec)
        # because it's a stochastic algorithm you want to have some margin here.
        nominal_schedule = convert_to_schedule(w_init.t, u_perturb_arrays[1], sim_param)
        return sim_result.risk_val, nominal_schedule, 1
        #time = @sprintf "Time %.2f" round(to_sec(w_init.t), digits=5)
        #@info "$(time) [sec]: SAC did not improve risk"
    else
        best_schedule = convert_to_schedule(w_init.t, u_perturb_arrays[control_idx], sim_param)
        #time = @sprintf "Time %.2f" round(to_sec(w_init.t), digits=5)
        #@info "$(time) [sec]: SAC improved risk by $(min_risk - sim_result.risk_val)"
        return min_risk, best_schedule, control_idx
    end
end

@inline function sac_control_update(w_init::WorldState,
                            u_schedule::OrderedDict{Time, Vector{Float64}},
                            target_trajectory::Trajectory2D,
                            prediction_dict::Dict{String, Array{Float64, 3}},
                            sim_param::SimulationParameter,
                            cnt_param::ControlParameter)
    # dtexec has to be divisible by dtc, or smaller than dtc.
    @assert all([(to_nsec(Duration(dte)) %
            to_nsec(Duration(sim_param.dtc)) == 0) ||
            dte < sim_param.dtc for dte in cnt_param.dtexec])
    tcalc_actual =
        @elapsed u_arrays = get_nominal_u_arrays(u_schedule, sim_param,
                                                 cnt_param);
    tcalc_actual +=
        @elapsed sim_result, u_nominal_array =
            simulate(w_init, u_arrays, target_trajectory,
                     prediction_dict, sim_param, cnt_param.constraint_time);
    tcalc_actual +=
        @elapsed coeff_matrix, coeff_matrix_constraint =
            get_control_coeffs(sim_result, sim_param, cnt_param);
    tcalc_actual +=
        @elapsed control_schedule_array = get_control_schedule(sim_result,
                                                               u_nominal_array,
                                                               coeff_matrix,
                                                               sim_param,
                                                               cnt_param,
                                                               coeff_matrix_constraint);
    tcalc_actual +=
        @elapsed control_chosen =  determine_control_time(sim_result,
                                                          control_schedule_array,
                                                          sim_param, cnt_param);
    tcalc_actual +=
        @elapsed u_perturb_arrays = perturbation_candidates(w_init.t,
                                               control_chosen, u_nominal_array,
                                               cnt_param.dtexec, sim_param)
    tcalc_actual +=
        @elapsed min_risk, best_schedule, control_idx =
            get_perturbed_schedule(w_init, sim_result,
                                   u_perturb_arrays, target_trajectory,
                                   prediction_dict, sim_param, cnt_param);
    if tcalc_actual >= cnt_param.tcalc
        # In this case there is not much time left for the controller to insert control for dtexec seconds.
        # In simulation, throw the following warning message and override.
        time = @sprintf "Time %.2f" round(to_sec(w_init.t), digits=5)
        @warn "$(time) [sec]: SAC computation took $(round(tcalc_actual, digits=3)) [sec], which exceeds the maximum computation time allowed."
    end
    if min_risk >= sim_result.risk_val
    #    time = @sprintf "Time %.2f" round(to_sec(w_init.t), digits=5)
    #    @info "$(time) [sec]: SAC computation did not improve the best nominal risk value.
    else
    #    time = @sprintf "Time %.2f" round(to_sec(w_init.t), digits=5)
    #    @info "$(time) [sec]: SAC improved risk by $(min_risk - sim_result.risk_val)"
    end

    return tcalc_actual, best_schedule, sim_result
end


text:
# import TriangleMesh
import Triangulate

Base.@kwdef mutable struct FCFV_Mesh
    type   ::Union{String, Missing}          = missing
    nel    ::Union{Int64,  Missing}          = missing
    nf     ::Union{Int64,  Missing}          = missing
    nv     ::Union{Int64,  Missing}          = missing
    nn_el  ::Union{Int64,  Missing}          = missing
    nf_el  ::Union{Int64,  Missing}          = missing
    xv     ::Union{Vector{Float64}, Missing} = missing # node x coordinate
    yv     ::Union{Vector{Float64}, Missing} = missing # node y coordinate
    xf     ::Union{Vector{Float64}, Missing} = missing # face x coordinate
    yf     ::Union{Vector{Float64}, Missing} = missing # face y coordinate
    bc     ::Union{Vector{Int64},   Missing} = missing # face y coordinate
    xc     ::Union{Vector{Float64}, Missing} = missing # cent x coordinate
    yc     ::Union{Vector{Float64}, Missing} = missing # cent y coordinate
    e2v    ::Union{Matrix{Int64},   Missing} = missing # cell 2 node numbering
    e2f    ::Union{Matrix{Int64},   Missing} = missing # cell 2 face numbering
    vole   ::Union{Vector{Float64}, Missing} = missing # volume of element
    n_x    ::Union{Matrix{Float64}, Missing} = missing # normal 2 face x
    n_y    ::Union{Matrix{Float64}, Missing} = missing # normal 2 face y
    dA     ::Union{Matrix{Float64}, Missing} = missing # face length
    e2e    ::Union{Matrix{Int64},   Missing} = missing # element 2 element numbering
    f2e    ::Union{Matrix{Int64},   Missing} = missing # face 2 element numbering
    dA_f   ::Union{Matrix{Float64}, Missing} = missing # face 2 element numbering
    vole_f ::Union{Matrix{Float64}, Missing} = missing # volume of element
    n_x_f  ::Union{Matrix{Float64}, Missing} = missing # normal 2 face x
    n_y_f  ::Union{Matrix{Float64}, Missing} = missing # normal 2 face y
    # ---- mat props ---- #
    ke     ::Union{Vector{Float64}, Missing} = missing # diffusion coefficient
    phase  ::Union{Vector{Float64}, Missing} = missing # phase
    # phase_f::Union{Matrix{Float64}, Missing} = missing # normal 2 face x
end

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

function MakeTriangleMesh( nx, ny, xmin, xmax, ymin, ymax, inclusion, R, BC=[2; 1; 1; 1;], area = ((xmax-xmin)/nx)*((ymax-ymin)/ny), no_pts_incl = Int64(floor(1.0*pi*R/sqrt(((xmax-xmin)/nx)^2+((ymax-ymin)/ny)^2)))  )

    regions  = Array{Float64}(undef,4,0)
    holes    = Array{Float64}(undef,2,0)
    dx       = (xmax-xmin)/nx
    dy       = (ymax-ymin)/ny
    pts_l    = 1;
    pts_u    = 0;

    # 1. Four corners of the domain
    px     = [xmin; xmax; xmax; xmin]
    py     = [ymin; ymin; ymax; ymax]
    sx     = [ 1; 2; 3; 4; ]
    sy     = [ 2; 3; 4; 1; ]
    st     = BC  # segment markers == boundary flags

    no_pts = size(px,1);
    pts_l  = pts_l+no_pts;
    pts_u  = pts_u+no_pts;
    # Region 1
    h1     = [xmin+1e-13; ymin+1e-13; 1.0; 0.0]

    if inclusion==1
        # 2. Inclusion
        theta0       = collect(LinRange(0,2*pi,no_pts_incl+1));
        theta        = theta0[1:end-1] # do not include last point (periodic 0 == 2pi)
        xx           = cos.(theta);
        yy           = sin.(theta);
        center_x     = (xmax + xmin)/2.0
        center_y     = (ymax + ymin)/2.0
        X            = center_x .+ R*xx;
        Y            = center_y .+ R*yy;
        no_pts       = length(X);
        st1          = 3*ones(1,no_pts);
        pts_u        = pts_u + no_pts;
        sx1          = collect(pts_l:pts_u)
        sy1          = collect(pts_l+1:pts_u+1)
        sy1[end]     = pts_l # Periodic
        h2           = [0.0; 0.0; 2.0; 0.0] # Region 2
        for i=1:no_pts_incl
            px   = push!(px, X[i])
            py   = push!(py, Y[i])
            sx   = push!(sx, sx1[i])
            sy   = push!(sy, sy1[i])
            st   = push!(st, st1[i])
        end
        regions = hcat(h1,h2)
    end

    p       = hcat(px, py)         # points
    s       = hcat(sx, sy)         # segments
    p       = p'
    s       = s'
    st      = st[:]

    # Triangulation
    mesh        = FCFV_Mesh()
    mesh.type   = "UnstructTriangles"
    nvert_el = 3

    triin=Triangulate.TriangulateIO()
    triin.pointlist=Matrix{Cdouble}(vcat(px',py'))
    triin.segmentlist=Matrix{Cint}(vcat(sx',sy'))
    triin.segmentmarkerlist=Vector{Int32}(st)
    triin.regionlist=Matrix{Cdouble}(regions)
    astring     = @sprintf("%0.10lf", area)
    (trimesh, vorout)=Triangulate.triangulate("vQDpenq33o2IAa$(astring)", triin)  #
    mesh.nel    = size(trimesh.trianglelist,2)
    e2v         = trimesh.trianglelist[1:3,:]
    mesh.nv     = maximum(e2v)
    e2f         = trimesh.trianglelist[4:6,:] .- mesh.nv
    mesh.nf     = maximum(e2f)
    mesh.xv     = trimesh.pointlist[1,1:mesh.nv]
    mesh.yv     = trimesh.pointlist[2,1:mesh.nv]
    mesh.xf     = trimesh.pointlist[1,mesh.nv+1:end]
    mesh.yf     = trimesh.pointlist[2,mesh.nv+1:end]
    mesh.bc     = trimesh.pointmarkerlist[mesh.nv+1:end]
    mesh.phase  = trimesh.triangleattributelist[:]

    # domain   = TriangleMesh.Polygon_pslg(size(p,2), p, 0, Array{Int64}(undef,2,0), 0, Array{Float64}(undef,2,0),  size(s,2), s, st, 0, holes, size(regions,2), regions, 1.0)
    # astring  = @sprintf("%0.10lf", area)
    # switches = "vQDpenq33o2IAa$(astring)"  #QDpeq33o2Aa0.01

    # println("Arguments to Triangle: ", switches)
    # trimesh_0  = TriangleMesh.create_mesh(domain, switches)
    # # vertices per element
    # mesh.nel    = trimesh.n_cell
    # e2v         = trimesh.cell[1:3,:]
    # mesh.nv     = maximum(e2v)
    # e2f         = trimesh.cell[4:6,:] .- mesh.nv
    # mesh.nf     = maximum(e2f)
    # mesh.xv     = trimesh.point[1,1:mesh.nv]
    # mesh.yv     = trimesh.point[2,1:mesh.nv]
    # mesh.xf     = trimesh.point[1,mesh.nv+1:end]
    # mesh.yf     = trimesh.point[2,mesh.nv+1:end]
    # mesh.bc     = trimesh.point_marker[mesh.nv+1:end]
    # if length(trimesh.triangle_attribute)>0
    #     println(minimum(trimesh.triangle_attribute))
    #     println(maximum(trimesh.triangle_attribute))
    #     mesh.phase  = trimesh.triangle_attribute
    # else
    #     mesh.phase  = ones(trimesh.n_cell)
    # end

    mesh.ke     = ones(Float64,mesh.nel)
    nel  = mesh.nel
    vole = zeros(nel)
    xc   = zeros(nel)
    yc   = zeros(nel)

    @tturbo for iel=1:nel
        # Compute volumes of triangles - use vertices coordinates
        x1 = mesh.xv[e2v[1,iel]]
        y1 = mesh.yv[e2v[1,iel]]
        x2 = mesh.xv[e2v[2,iel]]
        y2 = mesh.yv[e2v[2,iel]]
        x3 = mesh.xv[e2v[3,iel]]
        y3 = mesh.yv[e2v[3,iel]]
        a         = sqrt((x1-x2)^2 + (y1-y2)^2)
        b         = sqrt((x2-x3)^2 + (y2-y3)^2)
        c         = sqrt((x1-x3)^2 + (y1-y3)^2)
        s         = 1/2*(a+b+c)
        vole[iel] = sqrt(s*(s-a)*(s-b)*(s-c))
        xc[iel]   = 1.0/3.0*(x1+x2+x3)
        yc[iel]   = 1.0/3.0*(y1+y2+y3)
    end

    mesh.e2v    = e2v'
    mesh.e2f    = e2f'
    mesh.nn_el  = 3
    mesh.nf_el  = 3
    mesh.xc     = xc
    mesh.yc     = yc
    mesh.vole   = vole

    nodeA = [2 3 1]
    nodeB = [3 1 2]
    nodeC = [1 2 3]

    # Compute normal to faces
    mesh.n_x = zeros(Float64,mesh.nel,mesh.nf_el)
    mesh.n_y = zeros(Float64,mesh.nel,mesh.nf_el)
    mesh.dA  = zeros(Float64,mesh.nel,mesh.nf_el)
    mesh.e2e = zeros(  Int64,mesh.nel,mesh.nf_el)
    mesh.ke  =  ones(Float64,mesh.nel)

    # Receive arrays from Triangulate
    vorodeges = zeros(Int64,size(vorout.edgelist))
    vorodeges[1,:] .= vorout.edgelist[1,:]
    vorodeges[2,:] .= vorout.edgelist[2,:]
    seglist   = zeros(Int64,size(trimesh.edgelist))
    seglist[1,:]   .= trimesh.edgelist[1,:]
    seglist[2,:]   .= trimesh.edgelist[2,:]

     # Assemble FCFV elements
     @tturbo for iel=1:mesh.nel

        for ifac=1:mesh.nf_el

            nodei  = mesh.e2f[iel,ifac]

            # Neighbouring element for this face
            ele1 = vorodeges[1,nodei]
            ele2 = vorodeges[2,nodei]
            iel2 = (iel==ele1) * ele2 + (iel==ele2) * ele1;
            mesh.e2e[iel,ifac] = iel2;

            # println(" element: ", iel, " connects: ", iel2)

            # Vertices
            vert1  = mesh.e2v[iel,nodeA[ifac]]
            vert2  = mesh.e2v[iel,nodeB[ifac]]
            vert3  = mesh.e2v[iel,nodeC[ifac]]
            bc     = mesh.bc[nodei]
            dx     = (mesh.xv[vert1] - mesh.xv[vert2] );
            dy     = (mesh.yv[vert1] - mesh.yv[vert2] );
            dAi    = sqrt(dx^2 + dy^2);

            # Face normal
            n_x  = -dy/dAi
            n_y  =  dx/dAi

            # Third vector
            v_x  = mesh.xf[nodei] - mesh.xc[iel]
            v_y  = mesh.yf[nodei] - mesh.yc[iel]

            # Check wether the normal points outwards
            dot                 = n_x*v_x + n_y*v_y
            mesh.n_x[iel,ifac]  = ((dot>=0.0)*n_x - (dot<0.0)*n_x)
            mesh.n_y[iel,ifac]  = ((dot>=0.0)*n_y - (dot<0.0)*n_y)
            mesh.dA[iel,ifac]   = dAi
        end
    end

    ############# FOR THE PSEUDO-TRANSIENT PURPOSES ONLY #############

    # Create face to element numbering
    mesh.f2e    = zeros(Int,mesh.nf,2)
    mesh.vole_f = zeros(Float64,mesh.nf,2)
    mesh.n_x_f  = zeros(Float64,mesh.nf,2)
    mesh.n_y_f  = zeros(Float64,mesh.nf,2)
    mesh.dA_f   = zeros(Float64,mesh.nf,2)

    # # Loop through field names and fields: standard
    # for fname in fieldnames(typeof(trimesh))
    #     println("Field name: ", fname)
    #     println("Content   : ", getfield(trimesh, fname))
    # end

    # Loop over edges and uses Voronoï diagram to get adjacent cells
    @tturbo for ifac=1:mesh.nf
        mesh.f2e[ifac,1] = vorodeges[1,ifac]
        mesh.f2e[ifac,2] = vorodeges[2,ifac]
        act1 = vorodeges[1,ifac] > 0
        act2 = vorodeges[2,ifac] > 0
        iel1 = (act1==1) * vorodeges[1,ifac] + (act1==0) * 1
        iel2 = (act2==1) * vorodeges[2,ifac] + (act2==0) * 1
        # Compute face length
        vert1  = seglist[1,ifac]
        vert2  = seglist[2,ifac]
        xf     = 0.5*(mesh.xv[vert1] + mesh.xv[vert2])
        yf     = 0.5*(mesh.yv[vert1] + mesh.yv[vert2])
        dx     = (mesh.xv[vert1] - mesh.xv[vert2] );
        dy     = (mesh.yv[vert1] - mesh.yv[vert2] );
        dAi    = sqrt(dx^2 + dy^2);
        mesh.dA_f[ifac,1] =  dAi
        mesh.dA_f[ifac,2] =  dAi
        # Volumes
        mesh.vole_f[ifac,1] = (act1==1) * mesh.vole[iel1]
        mesh.vole_f[ifac,2] = (act2==2) * mesh.vole[iel2]
        # Compute face normal
        n_x  = -dy/dAi
        n_y  =  dx/dAi
        # Third vector  (w.r.t. element 1)
        v_x  = xf - mesh.xc[iel1]
        v_y  = yf - mesh.yc[iel1]
        # Check wether the normal points outwards
        dot                 = n_x*v_x + n_y*v_y
        mesh.n_x_f[ifac,1]  = (act1==1) * ((dot>=0.0)*n_x - (dot<0.0)*n_x)
        mesh.n_y_f[ifac,1]  = (act1==1) * ((dot>=0.0)*n_y - (dot<0.0)*n_y)
        # Take the negative for the second element
        mesh.n_x_f[ifac,2]  = -mesh.n_x_f[ifac,1]
        mesh.n_y_f[ifac,2]  = -mesh.n_y_f[ifac,1]
    end
    ############# FOR THE PSEUDO-TRANSIENT PURPOSES ONLY #############

    return mesh

    end

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

function MakeQuadMesh( nx, ny, xmin, xmax, ymin, ymax, inclusion, R, BC=[2; 1; 1; 1;] )
    # Structure
    mesh        = FCFV_Mesh()
    # Generates a 2D rectangular mesh of nx*ny cells
    ncell      = nx*ny
    # Twice resolution mesh
    nx2 = 2*nx + 1
    ny2 = 2*ny + 1
    # 1D axis
    dx,dy = (xmax-xmin)/nx, (xmax-xmin)/ny
    xc  = collect(LinRange(xmin+dx/2.0,xmax-dx/2.0,nx))
    yc  = collect(LinRange(ymin+dy/2.0,ymax-dy/2.0,ny))
    x2  = LinRange(xmin,xmax,nx2)
    y2  = LinRange(ymin,ymax,ny2)
    # 2D mesh: fake ngrid for P-T space ;)
    x2d   = repeat(x2, 1, length(y2))'
    y2d   = repeat(y2, 1, length(x2))
    nodes = zeros(Int64, size(x2d))
    # Number active dofs (midfaces) - Face node numbering
    inum = 0
    for i=1:nx2
        for j=1:ny2
            if (mod(j,2)==1 && mod(i,2)==0) || (mod(j,2)==0 && mod(i,2)==1)
                inum += 1
                nodes[i,j] = inum
            end
        end
    end
    nface = inum;
    # Number vertices
    inum = 0
    for i=1:nx2
        for j=1:ny2
            if (mod(j,2)==1 && mod(i,2)==1) || (mod(j,2)==1 && mod(i,2)==1)
                inum += 1
                nodes[i,j] = inum
            end
        end
    end
    # Face nodes
    xf = zeros(nface)
    yf = zeros(nface)
    tf = zeros(Int64,nface) # node type - 1 boundary
    for i=1:nx2
        for j=1:ny2
            if (mod(j,2)==1 && mod(i,2)==0) || (mod(j,2)==0 && mod(i,2)==1)
                xf[nodes[i,j]] = x2d[i,j]
                yf[nodes[i,j]] = y2d[i,j]
                if ( j==1 ) # WEST
                    tf[nodes[i,j]] = BC[4]
                end
                if ( j==ny2 ) # EAST
                    tf[nodes[i,j]] = BC[2]
                end
                if ( i==nx2 ) # NORTH
                    tf[nodes[i,j]] = BC[3]
                end
                if ( i==1 ) # SOUTH
                    tf[nodes[i,j]] = BC[1] # set Neumann at the South
                end
            end
        end
    end
    # Vertices nodes
    xn = zeros( (nx+1)*(ny+1) )
    yn = zeros( (nx+1)*(ny+1) )
    tn = zeros(Int64, (nx+1)*(ny+1) ) # node type - 1 boundary
    for i=1:nx+1
        for j=1:ny+1
            ii = i + (j-1)*(ny+1)
                xn[ii] = xmin + (i-1)*dx
                yn[ii] = ymin + (j-1)*dy
                if (i==1 || i==nx+1 || j==1 || j==ny+1)
                    tn[ii] = 1
                end
        end
    end
    # Cell 2 face numbering - for matrix connectivity
    face = zeros(Int64, 4, ncell)
    e2e  = zeros(Int64, 4, ncell)
    @tturbo for i=1:nx
        for j=1:ny
            k  = j + (i-1)*ny
            jc = i+1 + (i-1)*1
            ic = j+1 + (j-1)*1
            face[1,k] =  nodes[ic,jc-1] # South
            e2e[1,k]  =  ((i-1) > 0) * (j + (i-1-1)*ny)  + ((i-1) <= 0) * -1

            face[2,k] =  nodes[ic-1,jc] # West
            e2e[2,k]  =  (j-1 > 0) * (j-1 + (i-1)*ny)  + ((j-1) <= 0) * -1

            face[3,k] =  nodes[ic+1,jc] # East
            e2e[3,k]  =  (j+1 <= ny ) * (j+1 + (i-1)*ny)  + (j+1 > ny) * -1

            face[4,k] =  nodes[ic,jc+1] # North
            e2e[4,k]  =  (i+1 <=nx) * (j + (i-1+1)*ny) + (i+1 >nx) * -1
        end
    end

    # Cell 2 vertices - used for visualisation of quads
    vert = zeros(Int64, 4, ncell)
    @tturbo for i=1:nx
        for j=1:ny
            k  = j + (i-1)*ny
            jc = i+1 + (i-1)*1
            ic = j+1 + (j-1)*1
            vert[1,k] =  nodes[ic-1,jc+1] # South
            vert[2,k] =  nodes[ic-1,jc-1] # East
            vert[3,k] =  nodes[ic+1,jc-1] # North
            vert[4,k] =  nodes[ic+1,jc+1] # West
        end
    end
    # Centroids
    xc = zeros(ncell)
    yc = zeros(ncell)
    w  = 0.25
    @tturbo for iel=1:ncell
        tempx = 0
        tempy = 0
        for j=1:4
            tempx += w * xf[face[j,iel]]
            tempy += w * yf[face[j,iel]]
        end
        xc[iel] = tempx
        yc[iel] = tempy
    end

    # Fill structure
    mesh.type   = "Quadrangles"
    mesh.nel    = ncell
    mesh.nf     = nface
    mesh.nf_el  = 4
    mesh.nf_el  = 4
    mesh.nv     = (nx+1)*(ny+1)
    mesh.xv     = xn
    mesh.yv     = yn
    mesh.xf     = xf
    mesh.yf     = yf
    mesh.e2v    = vert' # cell nodes - not dofs! Dofs are in 'cell'
    mesh.e2f    = face'
    mesh.xc     = xc
    mesh.yc     = yc
    mesh.vole   = dx*dy*ones(ncell)
    mesh.bc     = tf
    mesh.ke     =  ones(Float64,mesh.nel)
    mesh.e2e    = e2e'

    nodeA = [2 1 3 4]
    nodeB = [3 2 4 1]

    # phase
    mesh.phase = ones(mesh.nel)
    if inclusion==1
        @tturbo for iel=1:mesh.nel
            x               = mesh.xc[iel]
            y               = mesh.yc[iel]
            out             = (x^2 + y^2)>R^2
            mesh.phase[iel] = (out==1)*1.0 + (out!=1)*2.0
        end
    end

    # Compute normal to faces
    mesh.n_x = zeros(Float64,mesh.nel,mesh.nf_el)
    mesh.n_y = zeros(Float64,mesh.nel,mesh.nf_el)
    mesh.dA  = zeros(Float64,mesh.nel,mesh.nf_el)

    # Get the number of element that are on each side of a face
    face2element   = zeros(Int64,2,mesh.nf)
    face2vertices  = zeros(Int64,2,mesh.nf)
    elem           = collect(1:mesh.nel)
    @tturbo for iel=1:mesh.nel
        for ifac=1:mesh.nf_el
            nodei  = mesh.e2f[iel,ifac]
            vert1  = mesh.e2v[iel,nodeA[ifac]]
            vert2  = mesh.e2v[iel,nodeB[ifac]]
            # println("element: ",  iel, " neighbour: ", mesh.e2e[iel,ifac])
            # if mesh.e2e[iel,ifac] == -1 n0+=1 end
            face2element[1,nodei]  = elem[iel]
            face2element[2,nodei]  = mesh.e2e[iel,ifac]
            face2vertices[1,nodei] = vert1
            face2vertices[2,nodei] = vert2
        end
    end

    # Flag interface faces
    for iel=1:mesh.nel
        type1 = Int64(mesh.phase[iel])
        for ifac=1:mesh.nf_el
            nodei = mesh.e2f[iel,ifac]
            iel2  = mesh.e2e[iel,ifac]
            ind   = (iel2>0) * iel2 + (iel2<0) * iel
            type2 = Int64(mesh.phase[ind])
            flag  = mesh.bc[nodei]
            yes   = (type1 != type2)
            mesh.bc[nodei] = (yes==0) * flag + (yes==1) * 3
        end
    end

     # Assemble FCFV elements
    @tturbo for iel=1:mesh.nel

        # println("element: ",  iel)

        for ifac=1:mesh.nf_el

            nodei  = mesh.e2f[iel,ifac]

            # Vertices
            vert1  = mesh.e2v[iel,nodeA[ifac]]
            vert2  = mesh.e2v[iel,nodeB[ifac]]
            bc     = mesh.bc[nodei]
            dx     = abs(mesh.xv[vert1] - mesh.xv[vert2] );
            dy     = abs(mesh.yv[vert1] - mesh.yv[vert2] );
            dAi    = sqrt(dx^2 + dy^2);

            # if iel==1
            #     println(bc)
            #     @printf("face node, x = %2.2e y = %2.2e\n", mesh.xf[nodei], mesh.yf[nodei])
            #     @printf("vert1    , x = %2.2e y = %2.2e\n", mesh.xv[vert1], mesh.yv[vert1])
            #     @printf("vert2    , x = %2.2e y = %2.2e\n", mesh.xv[vert2], mesh.yv[vert2])
            # end

            # Face normal
            n_x  =  dy/dAi
            n_y  = -dx/dAi

            # Third vector
            v_x  = mesh.xf[nodei] - mesh.xc[iel]
            v_y  = mesh.yf[nodei] - mesh.yc[iel]

            # Check wether the normal points outwards
            dot                 = n_x*v_x + n_y*v_y
            mesh.n_x[iel,ifac]  = (dot>=0.0)*n_x - (dot<0.0)*n_x
            mesh.n_y[iel,ifac]  = (dot>=0.0)*n_y - (dot<0.0)*n_y
            mesh.dA[iel,ifac]   = dAi
        end
    end

    ############# FOR THE PSEUDO-TRANSIENT PURPOSES ONLY #############

    # Create face to element numbering
    mesh.f2e    = zeros(Int,mesh.nf,2)
    mesh.vole_f = zeros(Float64,mesh.nf,2)
    mesh.n_x_f  = zeros(Float64,mesh.nf,2)
    mesh.n_y_f  = zeros(Float64,mesh.nf,2)
    mesh.dA_f   = zeros(Float64,mesh.nf,2)

    # # Loop through field names and fields: standard
    # for fname in fieldnames(typeof(trimesh))
    #     println("Field name: ", fname)
    #     println("Content   : ", getfield(trimesh, fname))
    # end

    # Loop over edges and uses Voronoï diagram to get adjacent cells
    @tturbo for ifac=1:mesh.nf
        mesh.f2e[ifac,1] = face2element[1,ifac]
        mesh.f2e[ifac,2] = face2element[2,ifac]
        act1 = face2element[1,ifac] > 0
        act2 = face2element[2,ifac] > 0
        iel1 = (act1==1) * face2element[1,ifac] + (act1==0) * 1
        iel2 = (act2==1) * face2element[2,ifac] + (act2==0) * 1
        # Compute face length
        vert1  = face2vertices[1,ifac]
        vert2  = face2vertices[2,ifac]
        xf     = 0.5*(mesh.xv[vert1] + mesh.xv[vert2])
        yf     = 0.5*(mesh.yv[vert1] + mesh.yv[vert2])
        dx     = (mesh.xv[vert1] - mesh.xv[vert2] );
        dy     = (mesh.yv[vert1] - mesh.yv[vert2] );
        dAi    = sqrt(dx^2 + dy^2);
        mesh.dA_f[ifac,1] =  dAi
        mesh.dA_f[ifac,2] =  dAi
        # Volumes
        mesh.vole_f[ifac,1] = (act1==1) * mesh.vole[iel1]
        mesh.vole_f[ifac,2] = (act2==2) * mesh.vole[iel2]
        # Compute face normal
        n_x  = -dy/dAi
        n_y  =  dx/dAi
        # Third vector  (w.r.t. element 1)
        v_x  = xf - mesh.xc[iel1]
        v_y  = yf - mesh.yc[iel1]
        # Check wether the normal points outwards
        dot                 = n_x*v_x + n_y*v_y
        mesh.n_x_f[ifac,1]  = (act1==1) * ((dot>=0.0)*n_x - (dot<0.0)*n_x)
        mesh.n_y_f[ifac,1]  = (act1==1) * ((dot>=0.0)*n_y - (dot<0.0)*n_y)
        # Take the negative for the second element
        mesh.n_x_f[ifac,2]  = -mesh.n_x_f[ifac,1]
        mesh.n_y_f[ifac,2]  = -mesh.n_y_f[ifac,1]
    end
    ############# FOR THE PSEUDO-TRANSIENT PURPOSES ONLY #############

    return mesh
end

ext: jl


text:
# generic functions for distribution fitting

function suffstats{D<:Distribution}(dt::Type{D}, xs...)
    argtypes = tuple(D, map(typeof, xs)...)
    error("suffstats is not implemented for $argtypes.")
end

fit_mle{D<:UnivariateDistribution}(dt::Type{D}, x::AbstractArray) = fit_mle(D, suffstats(D, x))
fit_mle{D<:UnivariateDistribution}(dt::Type{D}, x::AbstractArray, w::AbstractArray) = fit_mle(D, suffstats(D, x, w))

fit_mle{D<:MultivariateDistribution}(dt::Type{D}, x::AbstractMatrix) = fit_mle(D, suffstats(D, x))
fit_mle{D<:MultivariateDistribution}(dt::Type{D}, x::AbstractMatrix, w::AbstractArray) = fit_mle(D, suffstats(D, x, w))

fit{D<:Distribution}(dt::Type{D}, x) = fit_mle(D, x)
fit{D<:Distribution}(dt::Type{D}, args...) = fit_mle(D, args...)


text:
using Test
using TightBindingLattice

using TightBindingLattice: parse_expr

@testset "parse_expr" begin
    @test 1 == parse_expr(1)
    @test 1 == parse_expr("1")
    @test 1.5 == parse_expr("1.5")
    @test 1.5im == parse_expr("1.5im")
    @test 1.5im == parse_expr("1.5i")
    @test [1,2,3] == parse_expr("[1,2,3]")
    @test [1,2,3im] == parse_expr([1, "2", "3i"])

    @test parse_expr("2") == 2
    @test parse_expr("i") == im
    @test parse_expr("cis(2)") == cis(2)
    @test parse_expr("exp(2i)") == exp(2im)
    @test parse_expr("[10, π, exp(1)]") == [10, π, exp(1)]
    @test parse_expr("[10 π]") == [10 π]
    @test parse_expr("[1 2; 3 4]") == [1 2; 3 4]
    @test parse_expr("(1, 2.0)") == (1, 2.0)

    @test parse_expr("2+3") == 2+3
    @test parse_expr("2-3") == 2-3
    @test parse_expr("2*3") == 2*3
    @test parse_expr("2/3") == 2/3
    @test parse_expr("2//3") == 2//3
    @test parse_expr("2\\3") == 2\3
    @test parse_expr("2^3") == 2^3

    for unary_ftn in [cos, sin, tan, exp, cis, cospi, sinpi, sqrt, log, abs, abs2, sign, conj, conj, real, imag, angle]
        @test parse_expr("$unary_ftn(3.0 + 5.0i)") == unary_ftn(3.0 + 5.0im)
    end
    @test_throws ErrorException parse_expr("j")
    @test_throws ErrorException parse_expr("true ? π : -1")
end

@testset "cleanup" begin
    using TightBindingLattice: cleanup_number
    @test cleanup_number(42, 1E-8) == 42
    @test cleanup_number(1.5 + 1E-12, 1E-8) == 1.5
    @test cleanup_number([1.0 + 1E-12, 2.0 - 1E-12, 3.0], 1E-8) == [1.0, 2.0, 3.0]
    @test cleanup_number(0.1234567, 1E-8) == 0.1234567
    @test cleanup_number(-1//3, 1E-8) == -1//3
end


text:
module ModelStorage

module Sessions

using Stipple
using Genie.Sessions

export init_from_storage

function init()
  Genie.Sessions.init()
end

function init_from_storage(m::Type{T})::T where T
  model_id = Symbol(m)
  instance = Genie.Sessions.get(model_id, nothing)

  model = if instance !== nothing
    Stipple.init(m; channel = getchannel(instance))
  else
    m |> Stipple.init
  end

  session = Genie.Sessions.set!(model_id, model)

  on(model.isready) do isready
    isready || return

    if instance !== nothing
      for f in fieldnames(typeof(model))
        setfield!(model, f, getfield(instance, f))
      end

      sleep(0.1)

      push!(model)
    end
  end

  for f in fieldnames(typeof(model))
    if isa(getfield(model, f), Reactive)
      on(getfield(model, f)) do _
        Genie.Sessions.set!(session, model_id, model)
      end
    end
  end

  model
end

end # module Sessions

end # module ModelStorage


text:
function cutborders!(array::Array{T,4},smooth_time,position) where T <: Real

    if position == 1
        array = array[:,:,:,1+div((smooth_time-1),2):end]
    elseif position == 2
        array = array[:,:,:,1+div((smooth_time-1),2):end-div((smooth_time-1),2)]
    elseif position == 3
        array = array[:,:,:,1:end - div((smooth_time-1),2)]
    end
    return array
end

function cutborders!(array::Array{T,3},smooth_time,position) where T <: Real
 if position == 1
        array = array[:,:,1+div((smooth_time-1),2):end]
    elseif position == 2
        array = array[:,:,1+div((smooth_time-1),2):end-div((smooth_time-1),2)]
    elseif position == 3
        array = array[:,:,1:end - div((smooth_time-1),2)]
    end
    return array
end

function cutborders!(array::Array{T,1},smooth_time,position) where T <: Real

    if position == 1
        array = array[1+div((smooth_time-1),2):end]
    elseif position == 2
        array = array[1+div((smooth_time-1),2):end-div((smooth_time-1),2)]
    elseif position == 3
        array = array[1:end-div((smooth_time-1),2)]
    end
    return array

end

function getsmoothdata!(U,V, W, Tv, ThetaV, RAD, Fs, smooth_x,smooth_y,smooth_time,position)
    buf3d = similar(U)
    buf2d = similar(Fs)
    filter_array!(buf3d,U,smooth_x,smooth_time,position)
    filter_array!(buf3d,V,smooth_x,smooth_time,position)
    filter_array!(buf3d,W,smooth_x,smooth_time,position)
    filter_array!(buf3d,Tv,smooth_x,smooth_time,position)
    filter_array!(buf3d,ThetaV,smooth_x,smooth_time,position)
    filter_array!(buf3d,RAD,smooth_x,smooth_time,position)
    filter_array!(buf2d,Fs,smooth_x,smooth_time,position)
    #return U, V, W, Tv, ThetaV, RAD, Fs
end

function getsmoothdata_nospace(U::Array{T,4},V::Array{T,4}, W::Array{T,4}, Tv::Array{T,4}, ThetaV::Array{T,4}, RAD::Array{T,4}, Fs::Array{T,3}, smooth_x::Int,smooth_y::Int,smooth_time::Int,position::Int) where T <: Real
    U = filter_array_nospace(U,smooth_time,position)
    V = filter_array_nospace(V,smooth_time,position)
    W = filter_array_nospace(W,smooth_time,position)
    Tv = filter_array_nospace(Tv,smooth_time,position)
    ThetaV = filter_array_nospace(ThetaV,smooth_time,position)
    RAD = filter_array_nospace(RAD,smooth_time,position)
    Fs = filter_array_nospace(Fs,smooth_time,position)
    return U, V, W, Tv, ThetaV, RAD, Fs
end


text:
module Unishox

using Unishox_jll

    export compress, decompress

    """
        compress(s::AbstractString)

    Compresses a string using the unishox library and returns it.
    """
    function compress(s::AbstractString)
        isempty(s) && return ""

        # The output should be no longer than the input
        compressed = Array{Cchar}(undef, sizeof(s))

        # The function modifies `compressed` and returns the number of bytes written
        nbytes = ccall((:unishox2_compress_simple, libunishox), Cint,
                        (Ptr{Cchar}, Cint, Ptr{Cchar}),
                        s, sizeof(s), compressed)

        # Failed compression is unishox's fault, not ours >_>
        nbytes > 0 || throw(ErrorException("Compression failed for input $s"))

        compressed = compressed[1:nbytes]

        return unsafe_string(pointer(compressed), nbytes)
    end

    """
        decompress(s::AbstractString)

    Decompresses a string previously compressed by unishox and returns it.
    """
    function decompress(s::AbstractString)
        isempty(s) && return ""

        decompressed = Array{Cchar}(undef, 3 * sizeof(s))

        nbytes = ccall((:unishox2_decompress_simple, libunishox), Cint,
                        (Ptr{Cchar}, Cint, Ptr{Cchar}),
                        s, sizeof(s), decompressed)

        if nbytes > 3 * sizeof(s)
            decompressed = Array{Cchar}(undef, nbytes + 2)
            nbytes = ccall((:unishox2_decompress_simple, libunishox), Cint,
                        (Ptr{Cchar}, Cint, Ptr{Cchar}),
                        s, sizeof(s), decompressed)
        end

        nbytes > 0 || throw(ErrorException("Decompression failed for input $s"))

        decompressed = decompressed[1:nbytes]

    return unsafe_string(pointer(decompressed), nbytes)
    end

end


text:
using Parameters

function rank_matrix(x::Array{Float64, 2}, nt_args::NamedTuple)
    @unpack n_draws, n_pars, n_chains = nt_args
    xc = zeros(Int, n_draws, n_pars, n_chains)
    for i in 1:n_pars
        ranked = Int.(ordinalrank(x[:, i]))
        for j in 1:n_chains
            start_draw = (j-1) * n_draws + 1
            last_draw = j * n_draws
            xc[:, i, j] = ranked[start_draw:last_draw]
        end
    end
    xc
end

function trankplot(nt::NamedTuple, param::Symbol;
        bins=40, n_draws=4000, n_chains=4, n_eff=0)

   if length(size(nt[param])) > 1
        params = [Symbol("$(param)[$i]") for i in 1:size(nt[param], 1)]
        n_pars = length(params)
    else
        params = [param]
        n_pars = 1
    end

    nt_args = (n_draws=n_draws, n_pars=n_pars, n_chains=n_chains)
    if length(params) > 1
        ranks = rank_matrix(Matrix(nt[param]'), nt_args)
    else
        ranks = rank_matrix(reshape(nt[param], length(nt[param]), 1), nt_args)
    end

    figs = Vector{Plots.Plot{Plots.GRBackend}}(undef, length(params))
    for i in 1:length(params)
        figs[i] = stephist(ranks[:,1,1]; bins, alpha=0.9, lab="Chain 1", leg=:bottomright)
        for i in 2:n_chains
            stephist!(ranks[:,1,i]; bins, alpha=0.9, lab="Chain $i", leg=:bottomright)
        end
        title!("Trankplot for $(params[i])")
        annotate!([(30, 0, Plots.text("n_eff = $(n_eff[i])", 10, :darkred, :left))])
    end
    figs
end

export
    rank_matrix,
    trankplot


text:
#RSDeltaSigma: Array manipulations
#-------------------------------------------------------------------------------

#==Converters to make input data uniform
===============================================================================#
"""`conv2seriesmatrix2D(x)`

Ensure we have a 2D matrix representing series data (Vector->2×N Array)
"""
conv2seriesmatrix2D(x::T) where T =
	throw(ErrorException("Cannot convert $T to a \"series data matrix\""))
conv2seriesmatrix2D(x::AbstractVector) = collect(x')
conv2seriesmatrix2D(x::Array{T, 2}) where T<:Number = x #Assume format is ok.

#==Array/Vector padding
===============================================================================#
"""`y = padl(x, n, val)`

Pad a matrix x on the left to length n with value val(0)
The empty matrix is assumed to be have 1 empty row
"""
padl(x, n::Int, val=0) =
	[ val*ones(eltype(x), max(1,size(x,1)), n-size(x,2)) x ]
padl(x::Vector, n::Int, val=0) = vcat(val*ones(eltype(x), n-length(x)), x)

"""`y = padr(x, n, val)`

Pad a matrix x on the right to length n with value val(0)
The empty matrix is assumed to be have 1 empty row
"""
padr(x, n::Int, val=0) =
	[ x val*ones(eltype(x), max(1,size(x,1)), n-size(x,2)) ]
padr(x::Vector, n::Int, val=0) = vcat(x, val*ones(eltype(x), n-length(x)))

"""`y = padt(x, n, val)`

Pad a matrix x on the top to length n with value val(0)
The empty matrix is assumed to be have 1 empty column
"""
padt(x, n::Int, val=0) =
	[ val*ones(eltype(x), n-size(x,1), max(1,size(x,2))); x ]
padt(x::Vector, n::Int, val=0) = padl(x,n,val)

"""`y = padb(x, n, val)`

Pad a matrix x on the bottom to length n with value val(0)
The empty matrix is assumed to be have 1 empty column
"""
padb(x, n::Int, val=0) =
	[ x ; val*ones(eltype(x), n-size(x,1), max(1,size(x,2))) ]
padb(x::Vector, n::Int, val=0) = padr(x,n,val)

#==Quadrature <--> real
===============================================================================#
"""`A  = mapQtoR(Z)`

Convert a quadrature matrix into its real equivalent.

Each element in `Z` is represented by a 2x2 matrix in `A`:
 - `z -> [x -y; y x]`
"""
function mapQtoR(Z)
	A = zeros(2*size(Z))
	A[1:2:end,1:2:end] = real(Z)
	A[2:2:end,2:2:end] = real(Z)
	A[1:2:end,2:2:end] = -imag(Z)
	A[2:2:end,1:2:end] = imag(Z)
	return A
end

#Last line


text:
# WATCHOUT FOR NOTHING
data_id(e::Entity, id::Int) = getfirst(x -> x.comp_id == id, e.data_ids).data_id
# All entity Assemblages go here

function assemble_sphere(position = zero(Point3f0); velocity = zero(Vec3f0),
                         color = DEFAULT_COLOR, radius = 1.0f0, specint = 0.8f0,
                         specpow = 0.8f0)
    return (Spatial(position, velocity), PolygonGeometry(Sphere(Point3f0(0.0), radius)),
            Material(specint, specpow),
            (color isa ComponentData ? color : UniformColor(color)), Shape(radius))
end

function box_coordinates()
    return [zero(Point3f0), zero(Point3f0), Vec3f0(1, 0, 0), Point3f0(1, 1, 0),
            Point3f0(0, 1, 0), zero(Point3f0), Point3f0(0, 0, 1), Point3f0(1, 0, 1),
            Point3f0(1, 1, 1), Point3f0(0, 1, 1), Point3f0(0, 0, 1), Point3f0(0, 0, 0),
            Point3f0(1, 0, 0), Point3f0(1, 0, 1), Point3f0(1, 1, 1), Point3f0(1, 1, 0),
            Point3f0(0, 1, 0), Point3f0(0, 1, 1), Point3f0(0, 1, 1)]
end

function assemble_wire_box(; velocity = zero(Vec3f0), position = zero(Point3f0),
                           color = GREY, left = Vec3f0(-10), right = Vec3f0(10),
                           linewidth = 2.0f0, miter = 0.6f0)
    coords = box_coordinates()

    return (BufferColor([color for i in 1:length(coords)]),
            Spatial(; position = position, velocity = velocity),
            LineOptions(linewidth, miter),
            LineGeometry(Point3f0.([left + ((right - left) .* c) for c in coords])))
end

function assemble_wire_axis_box(; position = zero(Point3f0), x = Vec3f0(1, 0, 0),
                                y = Vec3f0(0, 1, 0), z = Vec3f0(0, 0, 1),
                                velocity = zero(Vec3f0), color = GREY, linewidth = 2.0f0,
                                miter = 0.6f0)
    coords = (Mat3([[x...] [y...] [z...]]),) .* box_coordinates()
    return (BufferColor([color for i in 1:length(coords)]),
            Spatial(; position = position, velocity = velocity),
            LineOptions(linewidth, miter), LineGeometry(Point3f0.(coords)))
end

function assemble_box(left = Point3f0(-0.5), right = Point3f0(0.5); velocity = zero(Vec3f0),
                      color = DEFAULT_COLOR, specint = 0.8f0, specpow = 0.8f0)
    unit_direction = (right - left) / norm(right - left)

    return (Spatial((left + right) / 2, velocity),
            PolygonGeometry(GeometryBasics.HyperRectangle(Vec3f0(-unit_direction / 2),
                                                          Vec3f0(unit_direction / 2))),
            Material(specint, specpow), UniformColor(color), Shape(norm(right - left)))
end

function assemble_pyramid(; position = zero(Point3f0), velocity = zero(Vec3f0),
                          color = DEFAULT_COLOR, width = 1.0f0, height = 1.0f0,
                          specint = 0.8f0, specpow = 0.8f0, scale = 1.0f0)
    return (Spatial(position, velocity),
            PolygonGeometry(Pyramid(Point3f0(0), height, width)),
            Material(specint, specpow), UniformColor(color), Shape(scale))
end

function assemble_file_mesh(file; position = zero(Point3f0), velocity = zero(Vec3f0),
                            color = DEFAULT_COLOR, specint = 0.8f0, specpow = 0.8f0,
                            scale = 1.0f0)
    return (Spatial(position, velocity), FileGeometry(file), Material(specint, specpow),
            UniformColor(color), Shape(scale))
end

function assemble_camera3d(width_pixels, height_pixels; eyepos = -10 * Y_AXIS,
                           velocity = zero(Vec3f0), kwargs...)
    return (Spatial(eyepos, velocity),
            Camera3D(width_pixels, height_pixels; eyepos = eyepos, kwargs...))
end

function assemble_line(points; origin = zero(Point3f0), color = DEFAULT_COLOR,
                       thickness = 2.0f0, miter = 0.6f0)
    spatial = Spatial(; position = origin)
    return (spatial, UniformColor(color), LineGeometry(Point3f0.(points)),
            LineOptions(thickness, miter))
end

function assemble_arrow(origin, extremity; color = DEFAULT_COLOR, thickness = 0.2f0,
                        length_ratio = 0.4f0, radius_ratio = 2.5f0)
    spatial = Spatial(; position = origin)
    direction = extremity - origin
    scale = norm(direction)
    unit_direction = direction / scale

    geom = PolygonGeometry(Arrow(zero(Point3f0), Point3f0(0, 0, scale), thickness,
                                 length_ratio, radius_ratio))
    return (spatial, geom, Rotation(rotation(Z_AXIS, Vec3f0(unit_direction))),
            UniformColor(color), Material())
end

function assemble_axis_arrows(origin = Point3f0(0.0); axis_length = 5.0f0,
                              thickness = 0.2f0, length_ratio = 0.4f0, radius_ratio = 2.5f0)
    return (assemble_sphere(origin; radius = thickness * 3, color = RGBf0(0.6, 0.6, 0.6)),
            assemble_arrow(origin, origin + Point3f0(0, 0, axis_length);
                           thickness = thickness, length_ratio = length_ratio,
                           radius_ratio = radius_ratio, color = RGBf0(1, 0, 0)),
            assemble_arrow(origin, origin + Point3f0(0, axis_length, 0);
                           thickness = thickness, length_ratio = length_ratio,
                           radius_ratio = radius_ratio, color = RGBf0(0, 1, 0)),
            assemble_arrow(origin, origin + Point3f0(axis_length, 0, 0);
                           thickness = thickness, length_ratio = length_ratio,
                           radius_ratio = radius_ratio, color = RGBf0(0, 0, 1)))
end

function assemble_orientation_sphere(origin = Point3f0(0.0); radius = 5.0f0,
                                     thickness = 10.0f0, base_color = RGBf0(0.7, 0.7, 0.7),
                                     pieces = 60)
    in_plane_points = Vector{Point3f0}(undef, pieces + 2)

    for (i, θ) in enumerate(range(0, 2π; length = pieces))
        in_plane_points[i+1] = radius * Point3f0(cos(θ), sin(θ), 0.0)
    end
    in_plane_points[1] = radius * Point3f0(1, 0, 0.0)

    return ((assemble_line(in_plane_points; thickness = thickness, color = 0.5f0 * RED,
                           origin = origin)...,
             Rotation(Quaternions.qrotation(Z_AXIS, 0.0f0))),
            (assemble_line(in_plane_points; thickness = thickness, color = 0.5f0 * GREEN,
                           origin = origin)..., Rotation(rotation(X_AXIS, Z_AXIS))),
            (assemble_line(in_plane_points; thickness = thickness, color = 0.5f0 * BLUE,
                           origin = origin)..., Rotation(rotation(Y_AXIS, Z_AXIS))))
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/PlaneElasticity.jl/blob/master/LICENSE

using FEMBase
using PlaneElasticity

using Base.Test

@testset "PlaneElasticity.jl" begin
    # Test stiffness matrix
    K = 1.0
    K_expected = 1.0
    @test isapprox(K, K_expected)
end


text:
module BayesianLinearRegressors

using AbstractGPs
using LinearAlgebra
using Random
using Statistics

using AbstractGPs: AbstractGP, _cholesky, FiniteGP

include("bayesian_linear_regression.jl")

export logpdf, rand, mean, std, cov, BayesianLinearRegressor, marginals, posterior

end # module


text:
# Multiple variables alone a line across multiple frames, multi-process version.
# Note: this script can be made faster by updating the data without deleting the lines.
#
# Usage:
#   julia -p $ncores demo_vars_extracting_parallel.jl
#
# Hongyang Zhou, hyzhou@umich.edu

using Distributed, ParallelDataTransfer, Glob
@everywhere using Vlasiator, PyPlot, Printf

@everywhere function init_figure(x1, x2)
   fig, axs = plt.subplots(4, 1; num=myid(),
      figsize=(10, 15), sharex=true, constrained_layout=true)

   axs[end].set_xlim(x1, x2)

   axs[1].set_ylim(ρmin, ρmax)
   axs[2].set_ylim(vmin, vmax)
   axs[3].set_ylim(pmin, pmax)
   axs[4].set_ylim(bmin, bmax)
   for ax in axs
      ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
      ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
      ax.grid(true)
   end

   axs[end].set_xlabel("x [Re]"; fontsize)
   axs[1].set_ylabel("Density [amu/cc]"; fontsize)
   axs[2].set_ylabel("Velocity [km/s]"; fontsize)
   axs[3].set_ylabel("Pressure [nPa]"; fontsize)
   axs[4].set_ylabel("Magnetic field [nT]"; fontsize)
   return fig, axs
end

@everywhere function process(axs, file, cellids)
   isfile("out/"*file[end-8:end-5]*".png") && return

   println("file = $file")
   meta = load(file)

   p_extract = readvariable(meta, "vg_pressure", cellids) .* 1e9 |> vec # [nPa]
   rho_extract = readvariable(meta, "proton/vg_rho", cellids) |> vec
   v_extract = readvariable(meta, "proton/vg_v", cellids)
   vmag2_extract = sum(v_extract.^2, dims=1) |> vec
   pdyn_extract = rho_extract .* Vlasiator.mᵢ .* vmag2_extract .* 1e9 # [nPa]

   bz = readvariable(meta, "vg_b_vol", cellids)[3,:] .* 1e9 #[nT]

   imagnetopause_ = findfirst(<(0.0), bz)

   axs[1].plot(loc, rho_extract ./ 1e6, label="Proton density", color="#1f77b4")
   vl1 = axs[1].vlines(loc[imagnetopause_], ρmin, ρmax;
      colors="r", linestyle="dashed", alpha=0.5)

   axs[2].plot(loc, v_extract[1,:] ./ 1e3, label="Vx", color="#1f77b4")
   vl2 = axs[2].vlines(loc[imagnetopause_], vmin, vmax;
      colors="r", linestyle="dashed", alpha=0.5)

   axs[3].plot(loc, pdyn_extract, label="Dynamic", color="#1f77b4")
   axs[3].plot(loc, p_extract, label="Thermal", color="#ff7f0e")
   vl3 = axs[3].vlines(loc[imagnetopause_], pmin, pmax;
      colors="r", linestyle="dashed", alpha=0.5)

   axs[4].plot(loc, bz, label="Bz", color="#1f77b4")
   hl4 = axs[4].hlines(0.0, loc[1], loc[end], colors="k", linestyle="dashed", alpha=0.2)
   vl4 = axs[4].vlines(loc[imagnetopause_], bmin, bmax;
       colors="r", linestyle="dashed", alpha=0.5)

   str_title = @sprintf "t= %4.1fs" meta.time
   axs[1].set_title(str_title; fontsize)

   axs[2].legend(;loc="upper right", fontsize)
   axs[3].legend(;loc="upper left",  fontsize)
   axs[4].legend(;loc="upper right", fontsize)

   savefig("out/"*file[end-8:end-5]*".png", bbox_inches="tight")

   for ax in axs
      for line in ax.get_lines()
         line.remove()
      end
   end
   for line in (vl1, vl2, vl3, vl4, hl4)
      line.remove()
   end
   return
end

function make_jobs(files)
   for f in files
       put!(jobs, f)
   end
end

@everywhere function do_work(jobs, results)
   fig, axs = init_figure(x1, x2)
   while true
      file = take!(jobs)
      process(axs, file, cellids)
      put!(results, true)
   end
   close(fig)
end

############################################################################################
files = glob("bulk*.vlsv", ".")
nfile = length(files)

Re = Vlasiator.Re # Earth radii
x1, x2 = 8.0, 29.0
point1 = [x1, 0, 0] .* Re
point2 = [x2, 0, 0] .* Re

meta = load(files[1])
cellids, _, _ = getcellinline(meta, point1, point2)
close(meta.fid)

passobj(1, workers(), [:x1, :x2, :cellids])

@broadcast begin
   const loc = range(x1, x2, length=length(cellids))

   const ρmin, ρmax = 0.0, 10.0     # [amu/cc]
   const vmin, vmax = -640.0, 0.0   # [km/s]
   const pmin, pmax = 0.0, 1.82     # [nPa]
   const bmin, bmax = -25.0, 60.0   # [nT]

   const fontsize = 14
end

jobs    = RemoteChannel(()->Channel{String}(nfile))
results = RemoteChannel(()->Channel{Bool}(nworkers()))

println("Total number of files: $nfile")
println("Running with $(nworkers()) workers...")

@async make_jobs(files) # Feed the jobs channel with all files to process.

@sync for p in workers()
   @async remote_do(do_work, p, jobs, results)
end

let n = nfile
   t = @elapsed while n > 0 # wait for all jobs to complete
      take!(results)
      n -= 1
   end
   println("Finished in $(round(t, digits=2))s.")
end


text:
"""
Module Sampling has the following sampling methods:

 - Uniform Sampling
"""
module Sampling

"""
 uniform(x)
  x - \$ k × 2 \$ matrix with lower and upper bounds in each row
  n - Vector of number of samples in each dimension
"""
function uniform(x,n)

    dim = length(x[:,1])
    X = zeros(
    for i in 1:dim
        X =

end # uniform

export uniform

end # module Sampling

# vim:filetype=julia:foldmethod=syntax

max_stars_repo_name: devendra-ghate/surrogateModels

max_issues_repo_name: devendra-ghate/surrogateModels

max_forks_repo_name: devendra-ghate/surrogateModels


text:
import JuMP
import Ipopt
import Zygote

using JuMP: @variable, @constraint, @objective
using PartiallyObservedInverseGames.ForwardGame: IBRGameSolver, KKTGameSolver, solve_game
using PartiallyObservedInverseGames.InverseGames: InverseKKTConstraintSolver, solve_inverse_game
using PartiallyObservedInverseGames.TrajectoryVisualization: visualize_trajectory
using SparseArrays: spzeros
using Test: @test, @testset, @test_broken
using UnPack: @unpack

import Plots

unique!(push!(LOAD_PATH, joinpath(@__DIR__, "utils")))
import TestUtils
import TestDynamics

#======================================== Global parameters ========================================#

function objective_p1(x, u; weights)
    weights[:state_velocity] * sum((x[3, :] .- 1.0) .^ 2) + weights[:control_Δv] * sum(u[1, :] .^ 2)
end

function objective_gradients_p1(x, u; weights)
    T = size(x, 2)
    dJdx = 2 * weights[:state_velocity] * [zeros(2, T); x[3:3, :] .- 1.0; zeros(1, T)]
    dJdu = 2 * weights[:control_Δv] * [u[1:1, :]; zeros(1, T)]
    (; dx = dJdx, du = dJdu)
end

function objective_p2(x, u2; weights)
    weights[:state_goal] * sum(x[1:2, :] .^ 2) + weights[:control_Δθ] * sum(u2 .^ 2)
end

function objective_gradients_p2(x, u; weights)
    T = size(x, 2)
    dJdx = 2 * weights[:state_goal] * [x[1:2, :]; zeros(2, T)]
    dJdu = 2 * weights[:control_Δθ] * [zeros(1, T); u[2:2, :]]
    (; dx = dJdx, du = dJdu)
end

control_system = TestDynamics.Unicycle(0.1)
observation_model = (; σ = 0, expected_observation = identity)
x0 = [-1, 1, 0.0, 0]
T = 100

player_cost_models = (
    (;
        player_inputs = [1],
        weights = (; state_velocity = 1, control_Δv = 10),
        objective = objective_p1,
        objective_gradients = objective_gradients_p1,
        add_objective! = function (opt_model, args...; kwargs...)
            @objective(opt_model, Min, objective_p1(args...; kwargs...))
        end,
        add_objective_gradients! = function (opt_model, args...; kwargs...)
            objective_gradients_p1(args...; kwargs...)
        end,
    ),
    (;
        player_inputs = [2],
        weights = (; state_goal = 0.1, control_Δθ = 10),
        objective = objective_p2,
        objective_gradients = objective_gradients_p2,
        add_objective! = function (opt_model, args...; kwargs...)
            @objective(opt_model, Min, objective_p2(args...; kwargs...))
        end,
        add_objective_gradients! = function (opt_model, args...; kwargs...)
            objective_gradients_p2(args...; kwargs...)
        end,
    ),
)

#=============================================== Tests =============================================#

@testset "Gradient check" begin
    x = rand(4, 100)
    u = rand(2, 100)

    for (player_idx, cost_model) in enumerate(player_cost_models)
        dJdx_ad, dJdu_ad =
            Zygote.gradient((x, u) -> cost_model.objective(x, u; cost_model.weights), x, u)
        dJ = cost_model.objective_gradients(x, u; cost_model.weights)
        @test dJdx_ad == dJ.dx
        @test dJdu_ad[player_idx, :] == dJ.du[player_idx, :]
    end
end

function test_unicycle_multipliers(system, λ, x, u; player_cost_models)
    # annoying code duplication due to JuMP.expression ...
    df = let
        As = [
            [
                1 0 system.ΔT*cos(x[4, t]) -system.ΔT*x[3, t]*sin(x[4, t])
                0 1 system.ΔT*sin(x[4, t]) +system.ΔT*x[3, t]*cos(x[4, t])
                0 0 1 0
                0 0 0 1
            ] for t in 1:T
        ]

        Bs = [[
            0 0
            0 0
            1 0
            0 1
        ] for t in 1:T]

        (;
            dx = reduce((A, x) -> cat(A, x; dims = 3), As),
            du = reduce((A, x) -> cat(A, x; dims = 3), Bs),
        )
    end

    for (player_idx, cost_model) in enumerate(player_cost_models)
        @testset "λ$player_idx" begin
            dJ = cost_model.objective_gradients(x, u; cost_model.weights)

            dLdx = [
                dJ.dx[:, t] + λ[:, t - 1, player_idx] - df.dx[:, :, t]' * λ[:, t, player_idx]
                for t in 2:(T - 1)
            ]
            dLdu = [
                dJ.du[cost_model.player_inputs, t] -
                (df.du[:, cost_model.player_inputs, t]' * λ[:, t, player_idx])
                for t in 1:(T - 1)
            ]

            @test all(all(isapprox.(x, 0; atol = 1e-6)) for x in dLdx)
            @test all(all(isapprox.(x, 0; atol = 1e-6)) for x in dLdu)
        end
    end
end

@testset "Forward IBR" begin
    ibr_converged, ibr_solution, ibr_models =
        solve_game(IBRGameSolver(), control_system, player_cost_models, x0, T)
    global ibr_solution
    @test ibr_converged
end

@testset "Forward KKT Nash" begin
    kkt_converged, kkt_solution, kkt_model = solve_game(
        KKTGameSolver(),
        control_system,
        player_cost_models,
        x0,
        T;
        solver = Ipopt.Optimizer,
    )
    global kkt_solution

    @test kkt_converged

    visualize_trajectory(control_system, kkt_solution.x)

    test_unicycle_multipliers(
        control_system,
        kkt_solution.λ,
        kkt_solution.x,
        kkt_solution.u;
        player_cost_models,
    )
end

@testset "Inverse KKT Nash" begin
    converged, inverse_kkt_solution, inverse_kkt_model = solve_inverse_game(
        InverseKKTConstraintSolver(),
        ibr_solution.x;
        init = (; ibr_solution.u),
        control_system,
        observation_model,
        player_cost_models,
        max_observation_error = 0.1,
    )
    global inverse_kkt_solution

    @test converged
    for (cost_model, weights) in zip(player_cost_models, inverse_kkt_solution.player_weights)
        TestUtils.test_inverse_solution(weights, cost_model.weights)
    end
    TestUtils.test_data_fidelity(
        inverse_kkt_model,
        observation_model,
        ibr_solution.x,
        ibr_solution.x,
    )
end


text:
"""
    overlap(N::T; dims=dims::Union{Nothing,Integer}=nothing) where {T <: BipartiteNetwork}

Returns the overlap graph for a bipartite network. The `dims` keyword argument
can be `1` (default; overlap between top-level species) or `2` (overlap between
bottom-level species). See the documentation for `?overlap` for the output
format.
"""
function overlap(N::T; dims::Integer=1) where {T <: BipartiteNetwork}
    return overlap(convert(UnipartiteNetwork, N); dims=dims)
end

"""
    overlap(N::T; dims::Union{Nothing,Integer}=nothing) where {T <: UnipartiteNetwork}

Returns the overlap graph for a unipartite network. The `dims` keyword argument
can be `1` (overlap based on preys) or `2` (overlap based on predators), or
`nothing` (default; overlap based on both predators and preys). The overlap is
returned as a vector of named tuples, with elements `pair` (a tuple of species
names), and `overlap` (the number of shared interactors). The ordering within
the pair of species is unimportant, since overlap graphs are symetrical.
"""
function overlap(N::T; dims::Union{Nothing,Integer}=nothing) where {T <: UnipartiteNetwork}

    if dims !== nothing
        dims== 1 || dims == 2 || throw(ArgumentError("dims can only be nothing, 1, or 2 -- you used $(dims)"))
    end

    itype = Pair{Set{_species_type(N)},Int64}
    overlaps = Vector{itype}()

    for i in 1:(richness(N; dims=dims)-1)
        s1 = species(N; dims=dims)[i]
        s1suc = N[s1,:]
        s1pre = N[:,s1]
        isnothing(dims) && (s1set = union(s1suc, s1pre))
        dims == 1 && (s1set = s1suc)
        dims == 2 && (s1set = s1pre)
        for j in (i+1):(richness(N; dims=dims))
            s2 = species(N; dims=dims)[j]
            s2suc = N[s2,:]
            s2pre = N[:,s2]
            isnothing(dims) && (s2set = union(s2suc, s2pre))
            dims == 1 && (s2set = s2suc)
            dims == 2 && (s2set = s2pre)
            this_overlap = length(intersect(s1set, s2set))
            if this_overlap > 0
                push!(overlaps, Pair(Set((s1, s2)), this_overlap))
            end
        end
    end

    return overlaps
end

function AJS(N::T; dims::Integer=1) where {T <: BipartiteNetwork}
    return AJS(convert(UnipartiteNetwork, N); dims=dims)
end

"""
    AJS(N::T; dims::Union{Nothing,Integer}=nothing) where {T <: UnipartiteNetwork}

Additive Jaccard Similarity for pairs of species in the network. AJS
varies between 0 (no common species) to 1 (same profiles). This function can be
used to measure AJS based on only successors or predecessors, using the `dims`
argument.

Note that this function uses all *direct* preys and predators to measure the
similarity (and so does not go beyond the immediate neighbors).

#### References

Gao, P., Kupfer, J.A., 2015. Uncovering food web structure using a novel trophic
similarity measure. Ecological Informatics 30, 110–118.
https://doi.org/10.1016/j.ecoinf.2015.09.013
"""
function AJS(N::T; dims::Union{Nothing,Integer}=nothing) where {T <: UnipartiteNetwork}

    if dims !== nothing
        dims == 1 || dims == 2 || throw(ArgumentError("dims can only be nothing, 1, or 2 -- you used $(dims)"))
    end

    itype = Pair{Set{_species_type(N)},Float64}
    overlaps = Vector{itype}()

    for i in 1:(richness(N; dims=dims)-1)
        s1 = species(N; dims=dims)[i]
        s1suc = N[s1,:]
        s1pre = N[:,s1]
        isnothing(dims) && (s1set = union(s1suc, s1pre))
        dims == 1 && (s1set = s1suc)
        dims == 2 && (s1set = s1pre)
        for j in (i+1):(richness(N; dims=dims))
            s2 = species(N; dims=dims)[j]
            s2suc = N[s2,:]
            s2pre = N[:,s2]
            isnothing(dims) && (s2set = union(s2suc, s2pre))
            dims == 1 && (s2set = s2suc)
            dims == 2 && (s2set = s2pre)
            a = length(intersect(s1set, s2set))
            b = length(setdiff(s1set, s2set))
            c = length(setdiff(s2set, s1set))
            this_overlap = a/(a+b+c)
            if this_overlap > 0.0
                push!(overlaps, Pair(Set((s1, s2)), this_overlap))
            end
        end
    end

    return overlaps

end

function EAJS(N::T; dims::Integer=1) where {T <: BipartiteNetwork}
    return EAJS(convert(UnipartiteNetwork, N); dims=dims)
end

"""
    EAJS(N::T; dims::Union{Nothing,Integer}=nothing) where {T <: UnipartiteNetwork}

Extended Additive Jaccard Similarity for pairs of species in the network. AJS
varies between 0 (no common species) to 1 (same profiles). This function can be
used to measure AJS based on only successors or predecessors, using the `dims`
argument.

Note that this function counts all interactions up to a distance of 50 to define
the neighbourhood of a species. This should be more than sufficient for most
ecological networks.

#### References

Gao, P., Kupfer, J.A., 2015. Uncovering food web structure using a novel trophic
similarity measure. Ecological Informatics 30, 110–118.
https://doi.org/10.1016/j.ecoinf.2015.09.013
"""
function EAJS(N::T; dims::Union{Nothing,Integer}=nothing) where {T <: UnipartiteNetwork}

    sp = shortest_path(N)
    Y = UnipartiteNetwork(sp.>0, _species_objects(N)...)

    if dims !== nothing
        dims == 1 || dims == 2 || throw(ArgumentError("dims can only be nothing, 1, or 2 -- you used $(dims)"))
    end

    itype = Pair{Set{_species_type(N)},Float64}
    overlaps = Vector{itype}()

    for i in 1:(richness(Y; dims=dims)-1)
        s1 = species(Y; dims=dims)[i]
        s1suc = Y[s1,:]
        s1pre = Y[:,s1]
        isnothing(dims) && (s1set = union(s1suc, s1pre))
        dims == 1 && (s1set = s1suc)
        dims == 2 && (s1set = s1pre)
        for j in (i+1):(richness(Y; dims=dims))
            s2 = species(Y; dims=dims)[j]
            s2suc = Y[s2,:]
            s2pre = Y[:,s2]
            isnothing(dims) && (s2set = union(s2suc, s2pre))
            dims == 1 && (s2set = s2suc)
            dims == 2 && (s2set = s2pre)
            a = length(intersect(s1set, s2set))
            b = length(setdiff(s1set, s2set))
            c = length(setdiff(s2set, s1set))
            this_overlap = a/(a+b+c)
            if this_overlap > 0.0
                push!(overlaps, Pair(Set((s1, s2)), this_overlap))
            end
        end
    end

    return overlaps

end


text:
abstract ClassicalSpinModel{T} <: SpinModel{T}

function getBeta{T}(classicalspinmodel::ClassicalSpinModel{T})
	return one(T)/getTemperature(classicalspinmodel)
end


text:
module StringDistances

using Distances

include("distances/utils.jl")
include("distances/edit.jl")
include("distances/qgram.jl")
include("normalize.jl")

const StringDistance = Union{Hamming, Jaro, Levenshtein, DamerauLevenshtein, RatcliffObershelp, QGramDistance, Winkler, Partial, TokenSort, TokenSet, TokenMax, Normalize}
# Distances API
Distances.result_type(dist::StringDistance, s1::Type, s2::Type) = typeof(dist("", ""))
Distances.result_type(dist::StringDistance, s1, s2) = result_type(dist, typeof(s1), typeof(s2))

include("find.jl")
include("pairwise.jl")

##############################################################################
##
## Export
##
##############################################################################

export
StringDistance,
Hamming,
Levenshtein,
DamerauLevenshtein,
Jaro,
RatcliffObershelp,
QGramDistance,
QGram,
Cosine,
Jaccard,
SorensenDice,
Overlap,
MorisitaOverlap,
NMD,
QGramDict,
QGramSortedVector,
Winkler,
Partial,
TokenSort,
TokenSet,
TokenMax,
evaluate,
compare,
result_type,
qgrams,
normalize,
findnearest,
pairwise,
pairwise!
end


text:
module Concorde

using Random, LinearAlgebra
using TSPLIB

include("../deps/deps.jl")
include("dist.jl")
include("util.jl")
include("solver.jl")

export solve_tsp

end


text:
using Documenter, CommandLineEquationsOfState

makedocs(;
    modules=[CommandLineEquationsOfState],
    format=Documenter.HTML(),
    pages=[
        "Home" => "index.md",
    ],
    repo="https://github.com/singularitti/CommandLineEquationsOfState.jl/blob/{commit}{path}#L{line}",
    sitename="CommandLineEquationsOfState.jl",
    authors="Qi Zhang <singularitti@outlook.com>",
    assets=String[],
)

deploydocs(;
    repo="github.com/singularitti/CommandLineEquationsOfState.jl",
)


text:
g = EvolvingGraph()

add_edge!(g, 1, 2, 1)
add_edge!(g, 1, 3, 2)
add_edge!(g, 2, 3, 3)

l = depth_first_impl(g, 2, 1, 3, 3)
@test length(l) == 3 # number of nodes

g = IntAdjacencyList(4,3)

add_edge!(g, 1, 2, 1)
add_edge!(g, 1, 3, 2)
add_edge!(g, 2, 3, 3)
add_edge!(g, 4, 1, 1)
add_edge!(g, 4, 3, 2)
l = depth_first_impl(g, 4, 1, 3, 2)
@test length(l) == 3 # number of nodes


text:
using Printf

using Oceananigans.Operators: ℑxyᶜᶠᵃ, ℑxyᶠᶜᵃ

#####
##### Physical constants
#####

const Ω_Earth = 7.292115e-5 # [s⁻¹] https://en.wikipedia.org/wiki/Earth%27s_rotation#Angular_speed
const R_Earth = 6371.0e3    # Mean radius of the Earth [m] https://en.wikipedia.org/wiki/Earth

#####
##### Functions for non-rotating models
#####

@inline x_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)
@inline y_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)

#####
##### The 'FPlane' approximation. This is equivalent to a model with a constant
##### rotation rate around its vertical axis.
#####

"""
    FPlane{FT} <: AbstractRotation

A parameter object for constant rotation around a vertical axis.
"""
struct FPlane{FT} <: AbstractRotation
    f :: FT
end

"""
    FPlane([FT=Float64;] f=nothing, rotation_rate=Ω_Earth, latitude=nothing)

Returns a parameter object for constant rotation at the angular frequency
`f/2`, and therefore with background vorticity `f`, around a vertical axis.
If `f` is not specified, it is calculated from `rotation_rate` and
`latitude` according to the relation `f = 2*rotation_rate*sind(latitude).

By default, `rotation_rate` is assumed to be Earth's.

Also called `FPlane`, after the "f-plane" approximation for the local effect of
a planet's rotation in a planar coordinate system tangent to the planet's surface.
"""
function FPlane(FT::DataType=Float64; f=nothing, rotation_rate=Ω_Earth, latitude=nothing)

    use_f = !isnothing(f)
    use_planet_parameters = !isnothing(latitude)

    if !xor(use_f, use_planet_parameters)
        throw(ArgumentError("Either both keywords rotation_rate and latitude must be " *
                            "specified, *or* only f must be specified."))
    end

    if use_f
        return FPlane{FT}(f)
    elseif use_planet_parameters
        return FPlane{FT}(2rotation_rate*sind(latitude))
    end
end

@inline x_f_cross_U(i, j, k, grid, coriolis::FPlane, U) = - coriolis.f * ℑxyᶠᶜᵃ(i, j, k, grid, U.v)
@inline y_f_cross_U(i, j, k, grid, coriolis::FPlane, U) =   coriolis.f * ℑxyᶜᶠᵃ(i, j, k, grid, U.u)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::FPlane, U) where FT = zero(FT)

Base.show(io::IO, f_plane::FPlane{FT}) where FT =
    println(io, "FPlane{$FT}: f = ", @sprintf("%.2e", f_plane.f))

#####
##### βeta Plane
#####

"""
    BetaPlane{T} <: AbstractRotation

A parameter object for meridionally increasing Coriolis parameter (`f = f₀ + βy`).
"""
struct BetaPlane{T} <: AbstractRotation
    f₀ :: T
     β :: T
end

"""
    BetaPlane([T=Float64;] f₀=nothing, β=nothing,
                           rotation_rate=Ω_Earth, latitude=nothing, radius=R_Earth)

A parameter object for meridionally increasing Coriolis parameter (`f = f₀ + βy`).

The user may specify both `f₀` and `β`, or the three parameters `rotation_rate`,
`latitude`, and `radius` that specify the rotation rate and radius of a planet, and
the central latitude at which the `β`-plane approximation is to be made.

By default, the `rotation_rate` and planet `radius` is assumed to be Earth's.
"""
function BetaPlane(T=Float64; f₀=nothing, β=nothing,
                              rotation_rate=Ω_Earth, latitude=nothing, radius=R_Earth)

    use_f_and_β = !isnothing(f₀) && !isnothing(β)
    use_planet_parameters = !isnothing(latitude)

    if !xor(use_f_and_β, use_planet_parameters)
        throw(ArgumentError("Either both keywords f₀ and β must be specified, " *
                            "*or* all of rotation_rate, latitude, and radius."))
    end

    if use_planet_parameters
        f₀ = 2rotation_rate * sind(latitude)
         β = 2rotation_rate * cosd(latitude) / radius
     end

    return BetaPlane{T}(f₀, β)
end

@inline x_f_cross_U(i, j, k, grid, coriolis::BetaPlane, U) =
    @inbounds - (coriolis.f₀ + coriolis.β * grid.yC[j]) * ℑxyᶠᶜᵃ(i, j, k, grid, U.v)
@inline y_f_cross_U(i, j, k, grid, coriolis::BetaPlane, U) =
    @inbounds   (coriolis.f₀ + coriolis.β * grid.yF[j]) * ℑxyᶜᶠᵃ(i, j, k, grid, U.u)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::BetaPlane, U) where FT = zero(FT)

Base.show(io::IO, β_plane::BetaPlane{FT}) where FT =
    println(io, "BetaPlane{$FT}: ", @sprintf("f₀ = %.2e, β = %.2e", β_plane.f₀, β_plane.β))