text:
"""
NNMPCdriver
Driver that randomly changes lanes and speeds.

# Constructors
BafflingDriver(timestep::Float64;mlon::LaneFollowingDriver=IntelligentDriverModel(), mlat::LateralDriverModel=ProportionalLaneTracker(), mlane::LaneChangeModel=RandLaneChanger(timestep),rec::SceneRecord = SceneRecord(1, timestep))

# Fields
- `rec::SceneRecord` A record that will hold the resulting simulation results
- `mlon::LaneFollowingDriver = IntelligentDriverModel()` Longitudinal driving model
- `mlat::LateralDriverModel = ProportionalLaneTracker()` Lateral driving model
- `mlane::LaneChangeModel =RandLaneChanger` Lane change model (randomly)
"""

mutable struct NNMPCdriver <: DriverModel{AccelSteeringAngle}
    rec::SceneRecord
    n_ticks::Integer
    mlane::LaneChangeModel
    δ_max::Float64
    δ_min::Float64
    a_max::Float64
    a_min::Float64
    N_sim::Float64
    T::Float64
    ΔT::Float64
    T_obs_given::Float64
    T_obs::Float64
    T_pred::Float64
    Δa_max::Float64
    Δδ_max::Float64
    λ_v::Float64
    λ_div::Float64
    λ_δ::Float64
    λ_a::Float64
    λ_Δδ::Float64
    λ_Δa::Float64
    v_des::Float64
    a::Float64
    δ::Float64
    height::Float64
    width::Float64
    models::Dict{Int, DriverModel}
    thred_safety::Float64
    isDebugMode::Bool
    pred_model_type::AbstractString  # {"sgan", "perfect"}
    predictor::Any
    sgan_path::AbstractString
    sgan_model_path::AbstractString
    lane_change_action::LaneChangeChoice
    temp_frame_count::Integer
    veh_shape_type::AbstractString

    function NNMPCdriver(
        timestep::Float64;
        n_ticks::Integer = 5,
        mlane::LaneChangeModel= BafflingLaneChanger(timestep),
        rec::SceneRecord = SceneRecord(n_ticks, timestep),
        δ_max::Float64 = 0.6,
        δ_min::Float64 = -0.4,
        a_max::Float64 = 3.5,
        a_min::Float64 = -4.0,
        N_sim::Float64 = 200.0,
        T::Float64 = 0.2,
        ΔT::Float64 = timestep,
        T_obs::Float64 = 1.0,
        T_obs_given::Float64 = T_obs,
        T_pred::Float64 = 1.0,
        Δa_max::Float64 = 2.0,
        Δδ_max::Float64 = 0.4,
        λ_v::Float64 = 10.0,
        λ_div::Float64 = 10e2+0.0,
        λ_δ::Float64 = 1.0,
        λ_a::Float64 = 1.0,
        λ_Δδ::Float64 = 3.0,
        λ_Δa::Float64 = 1.0,
        height::Float64 = 4.0,
        width::Float64 = 1.8,
        v_des::Float64 = 30.0,
        thred_safety::Float64 = 0.3,
        isDebugMode::Bool = false,
        pred_model_type::AbstractString = "sgan",
        sgan_path::AbstractString = "/home/sbae/NNMPC.jl/python/sgan",
        sgan_model_path::AbstractString = "/home/sbae/NNMPC.jl/python/sgan/models/sgan-models/eth_8_model.pt",
        lane_change_action::LaneChangeChoice = LaneChangeChoice(DIR_RIGHT),
        veh_shape_type::AbstractString = "circle")

        retval = new()

        pushfirst!(PyVector(pyimport("sys")."path"),sgan_path)
        sganPredictor = pyimport("sgan.predictor")
        predictor = sganPredictor.Predictor(sgan_model_path)

        retval.rec = rec
        retval.mlane = mlane
        retval.δ_max = δ_max
        retval.δ_min = δ_min
        retval.a_max = a_max
        retval.a_min = a_min
        retval.N_sim = trunc(Int,N_sim)
        retval.T = T
        retval.ΔT = ΔT
        retval.T_obs = T_obs
        retval.T_obs_given = T_obs_given
        retval.T_pred = T_pred
        retval.Δa_max = Δa_max
        retval.Δδ_max = Δδ_max
        retval.λ_v = λ_v
        retval.λ_div = λ_div
        retval.λ_δ = λ_δ
        retval.λ_a = λ_a
        retval.λ_Δδ = λ_Δδ
        retval.λ_Δa = λ_Δa
        retval.v_des = v_des
        retval.a = a_min
        retval.δ = 0
        retval.height = height
        retval.width = width
        retval.models = Dict{Int64,DriverModel}()
        retval.thred_safety = thred_safety
        retval.isDebugMode = isDebugMode
        retval.predictor = predictor
        retval.pred_model_type = pred_model_type
        retval.lane_change_action = lane_change_action
        retval.temp_frame_count = 0
        retval.veh_shape_type = veh_shape_type

        retval
    end
end

get_name(::NNMPCdriver) = "NNMPCdriver"
function AutomotiveDrivingModels.set_desired_speed!(model::NNMPCdriver, v_des::Float64)
    model.v_des = v_des
    model
end

function set_other_models!(model::NNMPCdriver, models::Dict{Int, DriverModel})
    model.models = models
    model
end

"""
    compute_longitudinal_acceleration(model::NNMPCdriver, scene::Scene, roadway::Roadway, egoid::Int)

"""
function compute_longitudinal_acceleration(model::NNMPCdriver, scene::Scene, roadway::Roadway, ind_ego::Int)
    # local variables -- from IDM
    k_spd = 1.0 # proportional constant for speed tracking when in freeflow [s⁻¹]
    δ = 4.0 # acceleration exponent [-]
    T = 1.5 # desired time headway [s]
    # v_des = 29.0 # desired speed [m/s]
    s_min = 5.0 # minimum acceptable gap [m]
    a_max = 3.0 # maximum acceleration ability [m/s²]
    d_cmf = 2.0 # comfortable deceleration [m/s²] (positive)
    d_max = 9.0 # maximum deceleration [m/s²] (positive)
    ΔT = 0.2 # timestep to simulate [s]
    v_offset = 5.0 # offset from the desired speed [m/s]

    fore = get_neighbor_fore_along_lane(scene, ind_ego, roadway, VehicleTargetPointFront(), VehicleTargetPointRear(), VehicleTargetPointFront())
    v_ego = scene[ind_ego].state.v
    if fore.ind != nothing
        headway, v_oth = fore.Δs, scene[fore.ind].state.v
    else
        headway, v_oth = NaN, NaN
    end
    if !isnan(v_oth)
        Δv = v_oth - v_ego
        s_des = s_min + v_ego*T - v_ego*Δv / (2*sqrt(model.a_max*abs(model.a_min)))
        v_ratio = model.v_des > 0.0 ? (v_ego/model.v_des) : 1.0
        a = model.a_max * (1.0 - v_ratio^δ - (s_des/headway)^2)
    else
        # no lead vehicle, just drive to match desired speed
        Δv = model.v_des - v_ego
        a = Δv*k_spd # predicted accel to match target speed
    end

    if v_ego + model.ΔT * a < 0
        a = max(model.a_min, -v_ego/model.ΔT)
    end

    return a
end

"""
    observe!(model::NNMPCdriver, scene::Scene, roadway::Roadway, egoid::Int)

(1) randomly generates a sequence of control inputs (acceleration and steering angle)
(2) evaluates Social GAN to predict motions of neighboring vehicles within a range
(3) checks the safety constraint at each time step
(4) evalautes the cumulative cost over the receding time horizon T
(5) finds the optimal pair of control inputs
"""
function AutomotiveDrivingModels.observe!(model::NNMPCdriver, scene::Scene, roadway::Roadway, egoid::Int)
    # δ_max,δ_min,a_max,a_min,N_sim,T,ΔT,T_obs,T_pred,Δa_max,Δδ_max=0.6,-0.4,3.0,-3.0,1000,4,0.5,2,2,5,0.6
    ind_ego = findfirst(egoid, scene)

    update!(model.rec, scene)
    observe!(model.mlane, scene, roadway, egoid) # receive lane change action from higher level decision maker
    # global y_ref = DEFAULT_LANE_WIDTH * lane_change_action.dir
    # global y_ref = -1*DEFAULT_LANE_WIDTH

    # check if we have enough dataset to use SGAN
    isSganDataReady = model.rec.nframes*model.ΔT >= model.T_obs_given
    model.isDebugMode ? println("isSganDataReady: ", isSganDataReady) : nothing

    # if sgan data is not ready, then it just keeps the lane
    if !isSganDataReady
        @warn "Not enough records to run SGAN. IDM is running."
        model.a = compute_longitudinal_acceleration(model, scene, roadway, ind_ego)
        model.δ = 0
        return model
    end

    # ============== TEMP: test================
    # if model.lane_change_action.dir == DIR_MIDDLE
    #     if model.temp_frame_count == 20
    #         model.temp_frame_count = 0
    #
    #         left_lane_exists = convert(Float64, get(N_LANE_LEFT, model.rec, roadway, ind_ego)) > 0
    #         right_lane_exists = convert(Float64, get(N_LANE_RIGHT, model.rec, roadway, ind_ego)) > 0
    #
    #         if left_lane_exists
    #             model.lane_change_action = LaneChangeChoice(DIR_LEFT)
    #         elseif right_lane_exists
    #             model.lane_change_action = LaneChangeChoice(DIR_RIGHT)
    #         end
    #     else
    #         model.temp_frame_count += 1
    #     end
    #
    # end
    # =========================================

    # handle NaN lane offset. later, this must be addressed in the higher level
    lane_offset = get_lane_offset(model.lane_change_action, model.rec, roadway, ind_ego)
    if isnan(lane_offset)
        @warn "NaN lane offset. Fixed to keeping lane command."
        model.lane_change_action = LaneChangeChoice(DIR_MIDDLE)
        model.temp_frame_count = 0
    end

    lane_offset = get_lane_offset(model.lane_change_action, model.rec, roadway, ind_ego)
    if model.lane_change_action.dir == DIR_MIDDLE
        println("keeping lane")
    elseif model.lane_change_action.dir == DIR_LEFT
        println("turning to left lane")
    else
        println("turning to right lane")
    end

    # local variables
    N_receding = trunc(Int, round(model.T/model.ΔT))
    N_sgan_obs = trunc(Int, round(model.T_obs/model.ΔT))
    N_sgan_pred = trunc(Int, round(model.T_pred/model.ΔT))

    # generate sequences of control inputs
    # println("generates random control sequences")
    # @time begin
        a_seq, δ_seq = gen_rand_inputs(model)
    # end

    cost_min = Inf
    n_opt = -1            # index of the optimal sequence of the controls
    model.a = model.a_min # maximum braking by default
    model.δ = 0           # zero steering by default
    # initialize the vehicle object for the other vehicles
    ind_near_vehs = get_indices_near_vehs(ind_ego, scene, 15.0)
    println("number of vehicles in the range: ", length(ind_near_vehs))

    println("find optimal contorls")
    @time begin
        for n = 1 : model.N_sim
            model.isDebugMode ? println("n: ",n) : nothing

            # initialize the vehicle object
            ego_veh = scene[ind_ego] # Entity{VehicleStateBuffer,BicycleModel,Int}

            # initialize the record scene queue
            rec′ = SceneRecord(model.rec.nframes + N_receding, model.ΔT)
            for k = 1 : N_sgan_obs
                update!(rec′,model.rec[-N_sgan_obs + k])
            end

            # initialize lane change choice
            # lane_change_action′= LaneChangeChoice(lane_change_action.dir)
            lane_offset′ = lane_offset

            # randomly sample the sequence of control
            a_ = a_seq[:,rand(1:size(a_seq,2))]
            δ_ = δ_seq[:,rand(1:size(δ_seq,2))]

            # initialize the scene with the recent one
            scene′= Scene()
            for i = 1 : length(scene)
                push!(scene′,scene[i])
            end

            # evaluate the control sequences
            cost_n = 0
            for ℓ = 1 : N_receding
                model.isDebugMode ? println("ℓ: ", ℓ) : nothing
                # (1) predict other drivers' motion (SGAN) (for the next step)
                if length(ind_near_vehs) >= 1
                    if model.pred_model_type == "perfect"
                        scene′ = propagate_other_vehs(Any, model, ind_ego, ind_near_vehs, scene′, roadway, rec′)
                    elseif model.pred_model_type == "sgan"
                        scene′ = predict(model, ind_ego, ind_near_vehs, scene′, roadway, rec′)
                    else
                        error("ERROR: prediction model type error. It must be either {perfect} (benchmark) or {sgan}")
                    end
                end

                # (2) update the vehicle states with a,δ
                a,δ = a_[ℓ],δ_[ℓ]
                model.isDebugMode ? println("try a = ",a, ", δ = ", δ) : nothing
                state′ = propagate(ego_veh, AccelSteeringAngle(a,δ), roadway, model.ΔT)
                ego_veh = Entity(state′, ego_veh.def, ego_veh.id)
                scene′[ind_ego] = ego_veh
                update!(rec′,scene′)

                # (3) check the feasibility (break if violated)
                # println("check safety")
                # @time begin
                if length(ind_near_vehs) >= 1
                    if !isSafe(model, ind_ego, ind_near_vehs, scene′)
                        cost_n = Inf
                        break
                    end
                end
                # end

                # (4) compute the cumulative costs
                a_prev = ℓ == 1 ? 0 : a_[ℓ-1]
                δ_prev = ℓ == 1 ? 0 : δ_[ℓ-1]
                # if rec′[-1][ind_ego].state.posF.t * rec′[0][ind_ego].state.posF.t < 0 # lane changed
                #     lane_change_action′ = LaneChangeChoice(DIR_MIDDLE)
                # end
                lane_offset′ += ego_veh.state.v*sin(ego_veh.state.posF.ϕ)*model.ΔT
                cost_n += ( model.λ_v*(ego_veh.state.v-model.v_des)^2
                        + model.λ_div*(abs(lane_offset′)+1)^2
                        + model.λ_δ*(δ)^2
                        + model.λ_a*(a)^2
                        + model.λ_Δδ*(δ-δ_prev)^2
                        + model.λ_Δa*(a-a_prev)^2)
                        # + model.λ_div*(ego_veh.state.posG.y-y_ref)^2
            end

            if cost_n < cost_min
                # update with the optimal control sequence
                cost_min = cost_n
                model.a = a_[1]
                model.δ = δ_[1]
                n_opt = n
            end
        end
    end

    model.isDebugMode ? println("optimal controls (n = ", n_opt,") : a = ",model.a, ", δ = ", model.δ) : nothing
    if n_opt == -1
        @warn "No feasible solution found. Maximum braking applied."
    end

    # prevent the vehicle going backward
    if scene[ind_ego].state.v + model.ΔT*model.a < 0
        model.a = max(-scene[ind_ego].state.v/model.ΔT, model.a_min)
    end
    println("lane_offset: ",get_lane_offset(model.lane_change_action, model.rec, roadway, ind_ego))
    println("a: ", model.a)
    println("δ: ", model.δ)

    # check if the ego vehicle changes the lane with the next action
    state′ = propagate(scene[ind_ego], AccelSteeringAngle(model.a,model.δ), roadway, model.ΔT)
    if abs(lane_offset + state′.v*sin(state′.posF.ϕ)*model.ΔT) <= DEFAULT_LANE_WIDTH/2
        model.lane_change_action = LaneChangeChoice(DIR_MIDDLE)
    end

    model
end

"""
    get_indices_near_vehs(model::NNMPCdriver, scene::Scene; range::Float64 = 5 [m])

find indices of near vehicles around the ego vehicle
"""
function get_indices_near_vehs(ind_ego::Int, scene::Scene, range::Float64)
    ego_veh = scene[ind_ego]
    x = ego_veh.state.posG.x
    y = ego_veh.state.posG.y
    ind_near_vehs = []
    for i in 1 : length(scene)
        if i == ind_ego
            continue
        end
        xᵢ = scene[i].state.posG.x
        yᵢ = scene[i].state.posG.y
        if sqrt((x-xᵢ)^2+(y-yᵢ)^2) <= range
            append!(ind_near_vehs, i)
        end
    end
    return ind_near_vehs
end

"""
    gen_rand_inputs(model::NNMPCdriver)

returns randomly generated sequence of a and δ, both of which satisfies the
maximum rate constraints (jerk, steering rate)
"""
function gen_rand_inputs(model::NNMPCdriver)
    n_row = trunc(Int,model.T/model.ΔT)
    n_col = trunc(Int,model.N_sim)

    # generate random sequences of action pair (a, δ)
    a_seq_raw = rand(Random.MersenneTwister(),model.a_min:0.01:model.a_max, (n_row,n_col)) #0.002 sec
    δ_seq_raw = rand(Random.MersenneTwister(),model.δ_min:0.01:model.δ_max, (n_row,n_col))

    # filter out the infeasible jerks
    a_seq_check = [zeros(1,n_col);a_seq_raw]
    a_diff = abs.(reshape([diff(reshape(a_seq_check, length(a_seq_check)));0],n_row+1,n_col)[1:end-1,1:end])
    bitarr = sum(a_diff .< model.Δa_max,dims=1) .== n_row
    a_seq = []
    for i in 1:length(bitarr)
        if bitarr[i] == true
            append!(a_seq, a_seq_raw[1:end,i])
        end
    end
    a_seq = reshape(a_seq, n_row, trunc(Int,length(a_seq)/n_row))

    # filter out the infeasible steering rate
    δ_seq_check = [zeros(1,n_col);δ_seq_raw]
    δ_diff = abs.(reshape([diff(reshape(δ_seq_check, length(δ_seq_check)));0],n_row+1,n_col)[1:end-1,1:end])
    bitarr = sum(δ_diff .< model.Δδ_max,dims=1) .== n_row
    δ_seq = []
    for i in 1:length(bitarr)
        if bitarr[i] == true
            append!(δ_seq, δ_seq_raw[1:end,i])
        end
    end
    δ_seq = reshape(δ_seq, n_row, trunc(Int,length(δ_seq)/n_row))

    return a_seq, δ_seq
end

"""
    function propagate(model::NNMPCdriver, ind_near_vehs::Vector{Any}, rec::SceneRecord)
returns positions of the neighboring vehicles, based on their Driver Model. This is equivalent to the SGAN that is perfectly trained.
"""
function propagate_other_vehs(::Type{A}, model::NNMPCdriver, ind_ego::Int, ind_near_vehs::Vector{Any}, scene::Scene, roadway::Roadway, rec::SceneRecord) where {A}
    models = Dict{Int, DriverModel}()
    for i in 1 : length(model.models)
        if i == ind_ego
            models[i] = BafflingDriver(rec.timestep)
        else
            models[i] = model.models[i]
        end
    end
    actions = Array{A}(undef, length(scene))
    get_actions!(actions, scene, roadway, models)
    tick!(scene, roadway, actions, rec.timestep)

    return scene
end

"""
    predict(model::NNMPCdriver, ind_ego::Integer, ind_near_vehs::Vector{Any}, scene::Scene, roadway::Roadway, rec::SceneRecord)
predicts motions of all vehicles within a range, using Social GAN.
"""
function predict(model::NNMPCdriver, ind_ego::Integer, ind_near_vehs::Vector{Any}, scene::Scene, roadway::Roadway, rec::SceneRecord)
    # arrange input for SGAN
    # example:
    # obs_traj = [  [ 1,  1.000e+00,  8.460e+00,  3.590e+00],
    #           [ 1,  2.000e+00,  1.364e+01,  5.800e+00],
    #           [ 2,  1.000e+00,  9.570e+00,  3.790e+00],
    #           [ 2,  2.000e+00,  1.364e+01,  5.800e+00],
    #           [ 3,  1.000e+00,  1.067e+01,  3.990e+00],
    #           [ 3,  2.000e+00,  1.364e+01,  5.800e+00],
    #           [ 4,  1.000e+00,  1.173e+01,  4.320e+00],
    #           [ 4,  2.000e+00,  1.209e+01,  5.750e+00],
    #           [ 5,  1.000e+00,  1.281e+01,  4.610e+00],
    #           [ 5,  2.000e+00,  1.137e+01,  5.800e+00],
    #           [ 6,  1.000e+00,  1.281e+01,  4.610e+00],
    #           [ 6,  2.000e+00,  1.031e+01,  5.970e+00],
    #           [ 7,  1.000e+00,  1.194e+01,  6.770e+00],
    #           [ 7,  2.000e+00,  9.570e+00,  6.240e+00],
    #           [ 8,  1.000e+00,  1.103e+01,  6.840e+00],
    #           [ 8,  2.000e+00,  8.730e+00,  6.340e+00]]

    N_obs = trunc(Int, round(model.T_obs/model.ΔT))
    N_vehs = length(ind_near_vehs)

    # bulid the input matrix for sgan
    obs_traj = []
    for seq in 1 : N_obs # sequence
        for i in 1 : N_vehs + 1 # vehicles
            ind = i == 1 ? ind_ego : ind_near_vehs[i-1]
            veh = rec[0-(N_obs-1)+(seq-1)][ind]
            l = [seq, i, veh.state.posG.x, veh.state.posG.y]
            obs_traj = [obs_traj;[l]]
        end
    end

    if model.isDebugMode
        println("abs_traj:")
        for i in 1:length(obs_traj)
              println(obs_traj[i])
        end
    end

    # call SGAN and return the predicted positions
    next_pred_traj = model.predictor.predict(obs_traj)
    if model.isDebugMode
        println("pred_traj:")
        for i in 1:size(next_pred_traj,1)
            ind = i == 1 ? ind_ego : ind_near_vehs[i-1]
            if i == 1
                println("ego vehicle: ", next_pred_traj[i,:])
            else
                println("vehicle ",ind, ": ", next_pred_traj[i,:])
            end
        end
    end

    # update the scene with the predicted positions
    for i in 1 : N_vehs + 1
        ind = i == 1 ? ind_ego : ind_near_vehs[i-1]
        veh = scene[ind]

        x′ = next_pred_traj[i,1]; x_diff = x′-veh.state.posG.x
        y′ = next_pred_traj[i,2]; y_diff = y′-veh.state.posG.y
        θ′ = atan(y_diff/x_diff)
        v′ = sqrt(x_diff^2+y_diff^2)/model.ΔT

        posG = VecSE2(x′, y′, θ′)
        state′  = VehicleState(posG, roadway, v′)

        scene[ind] = Entity(state′, veh.def, veh.id)
    end

    return scene
end

"""
    isSafe(ind_ego::Int, ind_others::Array{Int}, scene::Scene)
check collision with any of adjacent vehicles
"""
function isSafe(model::NNMPCdriver, ind_ego::Int, ind_others::Vector{Any}, scene::Scene)
    ego_veh = scene[ind_ego]
    ϵ = model.thred_safety
    # println("ind_ego: ", ind_ego)
    for i in ind_others
        dist = min_distance(model, ego_veh, scene[i])
        model.isDebugMode ? println("distance to vehicle ", i, ": ", dist) : nothing
        if dist > ϵ
            continue
        else
            model.isDebugMode ? println("   COLISION DETECTED. actions discarded") : nothing
            return false
        end
    end
    return true
end

"""
    min_distance(ego::Entity{VehicleState, D, Int},
                other::Entity{VehicleState, D, Int}) where {D}

returns distance between ego vehicle and an adjacent vehicle
"""
function min_distance(model::NNMPCdriver,
                    ego::Entity{VehicleState, D, Int},
                    other::Entity{VehicleState, D, Int}) where {D}
    # ego: ego vehicle object
    # other: other vehicle object
    x,y,θ = ego.state.posG.x, ego.state.posG.y, ego.state.posG.θ
    xᵢ,yᵢ,θᵢ = other.state.posG.x, other.state.posG.y, other.state.posG.θ
    h = model.height/2
    w = model.width/2

    if model.veh_shape_type == "ellipsoid"
        prob = Model(with_optimizer(Ipopt.Optimizer, print_level=0))
        @variable(prob, x′); @variable(prob, y′)
        @variable(prob, xᵢ′); @variable(prob, yᵢ′)
        @objective(prob, MOI.MIN_SENSE, (x′-xᵢ′)^2+(y′-yᵢ′)^2)
        @constraint(prob, ((x′-x)*cos(θ) + (y′-y)*sin(θ))^2/(h^2) + ((x′-x)*sin(θ) + (y′-y)*cos(θ))^2/(w^2) <= 1.0)
        @constraint(prob, ((xᵢ′-xᵢ)*cos(θᵢ) + (yᵢ′-yᵢ)*sin(θᵢ))^2/(h^2) + ((xᵢ′-xᵢ)*sin(θᵢ) + (yᵢ′-yᵢ)*cos(θᵢ))^2/(w^2) <= 1.0)
        optimize!(prob)
        return sqrt(max(objective_value(prob), 0))
    elseif model.veh_shape_type == "circle"
        r = w
        min_dis = Inf
        for i in [-1,0,1]
            for j in [-1,0,1]
                dis = sqrt(
                          ( (x + i*r*cos(θ)) - (xᵢ+ j*r*cos(θᵢ)) )^2
                         +( (y + i*r*sin(θ)) - (yᵢ+ j*r*sin(θᵢ)) )^2
                         )     - 2*r
                min_dis = min(dis, min_dis)
            end
        end
        return max(min_dis, 0)
    else
        error("ERROR: vehicle shape type error. It must be either circle or ellipsoid")
    end
end

Base.rand(model::NNMPCdriver) = AccelSteeringAngle(model.a,model.δ)


text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã ``2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

module GrammarCounts

export rulecounts, decisioncounts

using RLESUtils, SwapBuffers, MathUtils
using DerivationTrees

typealias CountDict Dict{Symbol,Float64}
typealias DecisionCounts Dict{Symbol,Vector{Float64}}

"""
Returns a Dict where key is symbol of decision rule, value is a vector of counts for each action
Example:
see runtests.jl
"""
function decisioncounts(grammar::Grammar, r_cnts::CountDict)
  d_cnts = DecisionCounts()
  for (name, rule) in grammar.rules
    if isa(rule, DecisionRule)
      d_cnts[name] = vec_count(grammar, rule, r_cnts)
    end
  end
  d_cnts
end

#vector of counts, by decision
function vec_count(grammar::Grammar, rule::OrRule, logcnts0::CountDict)
  Float64[count(grammar, r, logcnts0) for r in rule.values]
end

function vec_count(grammar::Grammar, rule::RangeRule, logcnts0::CountDict)
  zeros(length(rule.range))::Vector{Float64}
end

#non-decisions
vec_count(grammar::Grammar, rule, logcnts0::CountDict) = count(grammar, rule, logcnts0)

#"""
#Returns a Dict where key is symbol of grammar rule and value is counts
#"""
function rulecounts(grammar::Grammar, depth::Int64)
  logcntsA = initcounts(grammar)
  logcntsB = initcounts(grammar)
  sbuf = SwapBuffer(logcntsA, logcntsB)

  for d = 1:depth
    logcnts0 = active(sbuf) #active contains current
    logcnts1 = inactive(sbuf) #inactive contains next
    for (name, rule) in grammar.rules
      logcnts1[name] = count(grammar, rule, logcnts0)
    end
    @show logcnts1[:bin_vec], logcnts1[:and], logcnts1[:or]
    swap!(sbuf)
  end
  active(sbuf)
end

function initcounts(grammar::Grammar)
  logcnts = CountDict()
  for (k, v) in grammar.rules
    logcnts[k] = init(v)
  end
  logcnts
end

init(rule::OrRule) = -Inf
init(rule::ExprRule) = -Inf
init(rule::ReferencedRule) = -Inf

#terminals
init(rule::RangeRule) = log(length(rule.range))
init(rule::Terminal) = 0.0 #log(1.0)

function count(grammar::Grammar, rule::OrRule, logcnts0::CountDict)
  logxpy(vec_count(grammar, rule, logcnts0))::Float64 #sum the counts
end

function count(grammar::Grammar, rule::ExprRule, logcnts0::CountDict)
  logcnt = 0.0
  for arg in rule.args
    if isa(arg, Rule)
      logcnt += count(grammar, arg, logcnts0)
    end
  end
  logcnt
end

function count(grammar::Grammar, rule::ReferencedRule, logcnts0::CountDict)
  logcnts0[rule.symbol] #passthrough
end

#terminals are not recursive
count(grammar::Grammar, rule::RangeRule, logcnts0::CountDict) = init(rule)
count(grammar::Grammar, rule::Terminal, logcnts0::CountDict) = init(rule)

end #module

ext: jl

max_stars_repo_name: sisl/ExprSearch

max_issues_repo_name: sisl/ExprSearch


text:
export CustomEngine

struct CustomEngine <: Engine
    mass::Float64
    thrust::Float64
    Isp::Float64
end

mass(engine::CustomEngine) = engine.mass
thrust(engine::CustomEngine) = engine.thrust
Isp(engine::CustomEngine) = engine.Isp


text:
export sudoku_print

"""
`sudoku_print(A)` prints the Sudoku board `A` in an attractive manner.
"""
function sudoku_print(A::Matrix)
    A = replace(i -> i == 0 ? "." : string(i), Matrix{Any}(A))
    nn = size(A)[1]
    n = Int(sqrt(nn))

    digitsLength = length(string(nn))
    A = lpad.(A, digitsLength, " ") # for working properly with n > 4

    header = "+" * prod(["-"^(digitsLength * n + n + 1) * "+" for _ ∈ 1:n])
    println(header)

    # print every line
    for i ∈ 1:nn
        to_print = ""
        for j ∈ 1:nn
            if j == 1
                to_print *= "| " * A[i, j] * " "
            else
                to_print *= A[i, j] * " "
            end
            if j % n == 0
                to_print *= "| "
            end
        end
        println(to_print)
        if i % n == 0
            println(header)
        end
    end
    nothing
end


text:
module LotkaVolterra3dPlots

    using LaTeXStrings
    using Measures: mm
    using RecipesBase

    import GeometricProblems.Diagnostics: compute_invariant_error, compute_momentum_error

    # export compute_energy_error, compute_momentum_error

    compute_energy_error(t,q,params) = compute_invariant_error(t,q, (t,q) -> hamiltonian(t,q,params))
    # compute_momentum_error(t,q,p) = compute_momentum_error(t,q,p,ϑ)

    @userplot PlotLotkaVolterra3d
    @recipe function f(p::PlotLotkaVolterra3d; latex=true)
        # if length(p.args) != 2 || !(typeof(p.args[1][1]) <: Solution) || !(typeof(p.args[2]) <: NamedTuple)
        #     error("Lotka-Volterra plots should be given a solution. Got: $(typeof(p.args))")
        # end
        sol = p.args[1]
        params = p.args[2]

        H, ΔH = compute_energy_error(sol.t, sol.q, params);

        size   := (800,1200)
        legend := :none
        guidefontsize := 12
        tickfontsize  := 8
        left_margin := 12mm
        right_margin := 8mm

        # traces
        layout := (4,1)

        if latex
            ylabels = (L"x_1", L"x_2", L"x_3")
        else
            ylabels = ("x₁", "x₂", "x₃")
        end

        for i in 1:3
            @series begin
                subplot := i
                yguide := ylabels[i]
                xlims  := (sol.t[0], Inf)
                xaxis := false
                sol.t, sol.q[i,:]
            end
        end

        @series begin
            subplot := 4
            if latex
                xguide := L"t"
                yguide := L"[H(t) - H(0)] / H(0)"
            else
                xguide := "t"
                yguide := "[H(t) - H(0)] / H(0)"
            end
            xlims  := (sol.t[0], Inf)
            yformatter := :scientific
            sol.t, ΔH
        end
    end

end


text:
using Test
using StaticArrays: SVector
using QuantumLattices.Essentials.SpinPackage
using QuantumLattices.Essentials.Spatials: PID, Point, Bond
using QuantumLattices.Essentials.DegreesOfFreedom: Table, OID, isHermitian, IDFConfig, Operators, oidtype, otype, script, iid, latexformat
using QuantumLattices.Essentials.Terms: Couplings, @subscript, statistics, abbr
using QuantumLattices.Interfaces: dims, inds, expand, matrix, permute, rank
using QuantumLattices.Prerequisites: Float
using QuantumLattices.Mathematics.Combinatorics: Permutations
using QuantumLattices.Mathematics.AlgebraOverFields: ID
using QuantumLattices.Mathematics.VectorSpaces: IsMultiIndexable, MultiIndexOrderStyle

@testset "SID" begin
    @test SID|>fieldnames == (:orbital, :spin, :tag)
    @test SID(orbital=1, spin=0.5, tag='z')' == SID(orbital=1, spin=0.5, tag='z')
    @test SID(orbital=1, spin=0.5, tag='x')' == SID(orbital=1, spin=0.5, tag='x')
    @test SID(orbital=1, spin=0.5, tag='y')' == SID(orbital=1, spin=0.5, tag='y')
    @test SID(orbital=1, spin=0.5, tag='+')' == SID(orbital=1, spin=0.5, tag='-')
    @test SID(orbital=1, spin=0.5, tag='-')' == SID(orbital=1, spin=0.5, tag='+')
end

@testset "matrix" begin
    @test isapprox(matrix(SID(1, 0.5, 'z')), [[-0.5, 0.0] [0.0, 0.5]])
    @test isapprox(matrix(SID(1, 0.5, 'x')), [[0.0, 0.5] [0.5, 0.0]])
    @test isapprox(matrix(SID(1, 0.5, 'y')), [[0.0, -0.5im] [0.5im, 0.0]])
    @test isapprox(matrix(SID(1, 0.5, '+')), [[0.0, 1.0] [0.0, 0.0]])
    @test isapprox(matrix(SID(1, 0.5, '-')), [[0.0, 0.0] [1.0, 0.0]])

    @test isapprox(matrix(SID(1, 1.0, 'z')), [[-1.0, 0.0, 0.0] [0.0, 0.0, 0.0] [0.0, 0.0, 1.0]])
    @test isapprox(matrix(SID(1, 1.0, 'x')), [[0.0, √2/2, 0.0] [√2/2, 0.0, √2/2] [0.0, √2/2, 0.0]])
    @test isapprox(matrix(SID(1, 1.0, 'y')), [[0.0, -√2im/2, 0.0] [√2im/2, 0.0, -√2im/2] [0.0, √2im/2, 0.0]])
    @test isapprox(matrix(SID(1, 1.0, '+')), [[0.0, √2, 0.0] [0.0, 0.0, √2] [0.0, 0.0, 0.0]])
    @test isapprox(matrix(SID(1, 1.0, '-')), [[0.0, 0.0, 0.0] [√2, 0.0, 0.0] [0.0, √2, 0.0]])
end

@testset "Spin" begin
    spin = Spin(atom=1, norbital=2, spin=1.0)
    @test IsMultiIndexable(Spin) == IsMultiIndexable(true)
    @test MultiIndexOrderStyle(Spin) == MultiIndexOrderStyle('C')
    @test dims(spin) == (2, 5)
    @test inds(SID(1, 1.0, 'z'), spin) == (1, 3)
    @test SID((1, 1), spin) == SID(1, 1.0, 'x')
    @test spin|>collect == [    SID(1, 1.0, 'x'), SID(1, 1.0, 'y'), SID(1, 1.0, 'z'), SID(1, 1.0, '+'), SID(1, 1.0, '-'),
                                SID(2, 1.0, 'x'), SID(2, 1.0, 'y'), SID(2, 1.0, 'z'), SID(2, 1.0, '+'), SID(2, 1.0, '-')
    ]
end

@testset "SIndex" begin
    @test SIndex|>fieldnames == (:scope, :site, :orbital, :spin, :tag)
    @test union(PID{Char}, SID) == SIndex{Char}
end

@testset "script" begin
    @test script(OID(SIndex('S', 1, 2, 0.5, 'z')), Val(:site)) == 1
    @test script(OID(SIndex('S', 1, 2, 0.5, 'z')), Val(:orbital)) == 2
    @test script(OID(SIndex('S', 1, 2, 0.5, 'z')), Val(:spin)) == 0.5
    @test script(OID(SIndex('S', 1, 2, 0.5, 'z')), Val(:tag)) == 'z'
end

@testset "oidtype" begin
    @test oidtype(SID, Point{2, PID{Int}}, Nothing, Val(true)) == OID{SIndex{Int}, SVector{2, Float}, SVector{2, Float}, Nothing}
    @test oidtype(SID, Point{2, PID{Int}}, Table, Val(true)) == OID{SIndex{Int}, SVector{2, Float}, SVector{2, Float}, Int}
    @test oidtype(SID, Point{2, PID{Int}}, Nothing, Val(false)) == OID{SIndex{Int}, Nothing, Nothing, Nothing}
    @test oidtype(SID, Point{2, PID{Int}}, Table, Val(false)) == OID{SIndex{Int}, Nothing, Nothing, Int}
end

@testset "SOperator" begin
    @test latexformat(SOperator) == soptdefaultlatex
    opt = SOperator(1.0, (SIndex('a', 1, 1, 0.5, '+'), SIndex('a', 1, 1, 0.5, '-')))
    @test opt|>statistics == opt|>typeof|>statistics == 'B'
    @test opt' == SOperator(1.0, (SIndex('a', 1, 1, 0.5, '+'), SIndex('a', 1, 1, 0.5, '-')))
    @test isHermitian(opt)
    @test repr(opt) == "S^{+}_{1, 1}S^{-}_{1, 1}"
end

@testset "permute" begin
    soptrep(opt::SOperator) = opt.value * prod([matrix(opt.id[i].index|>iid) for i = 1:rank(opt)])
    for S in (0.5, 1.0, 1.5)
        oids = [OID(SIndex('S', 1, 2, S, tag), seq=i) for (i, tag) in enumerate(('x', 'y', 'z', '+', '-'))]
        table = Dict(oid.index=>oid.seq for oid in oids)
        for (id1, id2) in Permutations{2}(oids)
            left = soptrep(SOperator(1, ID(id1, id2)))
            right = sum([soptrep(opt) for opt in permute(SOperator, id1, id2, table)])
            @test isapprox(left, right)
        end
    end
    id1 = OID(SIndex('S', 1, 2, 0.5, 'z'))
    id2 = OID(SIndex('S', 2, 2, 0.5, 'z'))
    @test permute(SOperator, id1, id2, nothing) == (SOperator(1, ID(id2, id1)),)
end

@testset "SCID" begin
    @test SCID(center=1, atom=1, orbital=1, tag='z')|>typeof|>fieldnames == (:center, :atom, :orbital, :tag, :subscript)
end

@testset "SpinCoupling" begin
    @test SpinCoupling{2}(1.0, tags=('+', '-'))|>string == "SpinCoupling(value=1.0, tags=(+, -))"
    @test SpinCoupling{2}(1.0, atoms=(1, 1), tags=('z', 'z'))|>string == "SpinCoupling(value=1.0, atoms=(1, 1), tags=(z, z))"
    @test SpinCoupling{2}(1.0, atoms=(1, 1), orbitals=(1, 2), tags=('-', '+'))|>string == "SpinCoupling(value=1.0, atoms=(1, 1), orbitals=(1, 2), tags=(-, +))"
    @test SpinCoupling{2}(2.0, tags=('x', 'y'))|>repr == "2.0 SxSy"

    sc1 = SpinCoupling{2}(1.5, tags=('+', '-'), atoms=(1, 2), orbitals=(@subscript (x, y) with x > y), centers=(1, 2))
    sc2 = SpinCoupling{2}(2.0, tags=('+', '-'), atoms=(1, 2), orbitals=(@subscript (x, y) with x < y), centers=(1, 2))
    @test sc1|>repr == "1.5 S+S- sl(1, 2)⊗ob(x, y)@(1, 2) with $(sc1.id[1].subscript) && $(sc1.id[2].subscript)"
    @test sc2|>repr == "2.0 S+S- sl(1, 2)⊗ob(x, y)@(1, 2) with $(sc2.id[1].subscript) && $(sc2.id[2].subscript)"

    sc = sc1 * sc2
    @test sc|>repr == "3.0 S+S-S+S- sl(1, 2, 1, 2)⊗ob(x, y, x, y)@(1, 2, 1, 2) with $(sc1.id[1].subscript) && $(sc1.id[2].subscript) && $(sc2.id[1].subscript) && $(sc2.id[2].subscript)"

    ex = expand(SpinCoupling{2}(2.0, tags=('+', '-'), atoms=(1, 1)), PID(1, 1), Spin(atom=2, norbital=2, spin=1.0))
    @test collect(ex) == []

    ex = expand(SpinCoupling{2}(2.0, tags=('+', '-'), atoms=(1, 2), orbitals=(1, 2)), (PID(1, 1), PID(1, 2)), (Spin(atom=1, norbital=2, spin=1.0), Spin(atom=2, norbital=2, spin=1.0)))
    @test IsMultiIndexable(typeof(ex)) == IsMultiIndexable(true)
    @test MultiIndexOrderStyle(typeof(ex)) == MultiIndexOrderStyle('C')
    @test dims(ex) == (1,)
    @test collect(ex) == [(2.0, (SIndex(1, 1, 1, 1.0, '+'), SIndex(1, 2, 2, 1.0, '-')))]

    ex = expand(SpinCoupling{4}(2.0, tags=('+', '-', '+', '-'), orbitals=(@subscript (α, α, β, β) with α < β)), PID(1, 1), Spin(norbital=3, spin=1.0))
    @test dims(ex) == (3,)
    @test collect(ex) == [  (2.0, (SIndex(1, 1, 1, 1.0, '+'), SIndex(1, 1, 1, 1.0, '-'), SIndex(1, 1, 2, 1.0, '+'), SIndex(1, 1, 2, 1.0, '-'))),
                            (2.0, (SIndex(1, 1, 1, 1.0, '+'), SIndex(1, 1, 1, 1.0, '-'), SIndex(1, 1, 3, 1.0, '+'), SIndex(1, 1, 3, 1.0, '-'))),
                            (2.0, (SIndex(1, 1, 2, 1.0, '+'), SIndex(1, 1, 2, 1.0, '-'), SIndex(1, 1, 3, 1.0, '+'), SIndex(1, 1, 3, 1.0, '-')))
    ]
end

@testset "Heisenberg" begin
    @test Heisenberg(orbitals=(1, 2)) == Couplings( SpinCoupling{2}(1//1, tags=('z', 'z'), orbitals=(1, 2)),
                                                    SpinCoupling{2}(1//2, tags=('+', '-'), orbitals=(1, 2)),
                                                    SpinCoupling{2}(1//2, tags=('-', '+'), orbitals=(1, 2))
    )
    @test Heisenberg("xyz") == Couplings(   SpinCoupling{2}(1, tags=('x', 'x')),
                                            SpinCoupling{2}(1, tags=('y', 'y')),
                                            SpinCoupling{2}(1, tags=('z', 'z'))
    )
end

@testset "Ising" begin
    @test Ising('x', atoms=(1, 2)) == Couplings(SpinCoupling{2}(1, tags=('x', 'x'), atoms=(1, 2)))
    @test Ising('y', atoms=(1, 2)) == Couplings(SpinCoupling{2}(1, tags=('y', 'y'), atoms=(1, 2)))
    @test Ising('z', atoms=(1, 2)) == Couplings(SpinCoupling{2}(1, tags=('z', 'z'), atoms=(1, 2)))
end

@testset "Gamma" begin
    @test Gamma('x', orbitals=(1, 1)) == SpinCoupling{2}(1, tags=('y', 'z'), orbitals=(1, 1))+SpinCoupling{2}(1, tags=('z', 'y'), orbitals=(1, 1))
    @test Gamma('y', atoms=(1, 2)) == SpinCoupling{2}(1, tags=('z', 'x'), atoms=(1, 2))+SpinCoupling{2}(1, tags=('x', 'z'), atoms=(1, 2))
    @test Gamma('z') == SpinCoupling{2}(1, tags=('x', 'y'))+SpinCoupling{2}(1, tags=('y', 'x'))
end

@testset "DM" begin
    @test DM('x', orbitals=(1, 1)) == SpinCoupling{2}(1, tags=('y', 'z'), orbitals=(1, 1))-SpinCoupling{2}(1, tags=('z', 'y'), orbitals=(1, 1))
    @test DM('y', atoms=(1, 2)) == SpinCoupling{2}(1, tags=('z', 'x'), atoms=(1, 2))-SpinCoupling{2}(1, tags=('x', 'z'), atoms=(1, 2))
    @test DM('z') == SpinCoupling{2}(1, tags=('x', 'y'))-SpinCoupling{2}(1, tags=('y', 'x'))
end

@testset "Sᵅ" begin
    @test Sˣ(atom=1, orbital=1) == Couplings(SpinCoupling{1}(1, tags=('x',), atoms=(1,), orbitals=(1,)))
    @test Sʸ(atom=1) == Couplings(SpinCoupling{1}(1, tags=('y',), atoms=(1,)))
    @test Sᶻ(orbital=1) == Couplings(SpinCoupling{1}(1, tags=('z',), orbitals=(1,)))
end

@testset "spincoupling" begin
    sc = sc"1.0 S+S- sl(1, 1)⊗ob(α, β)@(1, 2) with α<β"
    ob = sc.subscripts[1].identifier
    @test repr(sc) == "1.0 S+S- sl(1, 1)⊗ob(α, β)@(1, 2) with $ob && $ob"

    sc = sc"1.0 S+S- sl(1, 1)⊗ob(α, β)@(1, 2)"
    ob = sc.subscripts[1].identifier
    @test repr(sc) == "1.0 S+S- sl(1, 1)⊗ob(α, β)@(1, 2) with $ob && $ob"

    sc = sc"1.0 S+S- sl(1, 1)⊗ob(α, β)"
    ob = sc.subscripts[1].identifier
    @test repr(sc) == "1.0 S+S- sl(1, 1)⊗ob(α, β) with $ob && $ob"

    sc = sc"1.0 S+S- sl(1, 1)⊗ob(1, 2)@(1, 2)"
    @test repr(sc) == "1.0 S+S- sl(1, 1)⊗ob(1, 2)@(1, 2)"

    sc = sc"1.0 S+S- sl(1, 1)@(1, 2)"
    @test repr(sc) == "1.0 S+S- sl(1, 1)@(1, 2)"

    sc = sc"1.0 S+S- @(1, 2)"
    @test repr(sc) == "1.0 S+S- @(1, 2)"

    sc = sc"1.0 S+S-"
    @test repr(sc) == "1.0 S+S-"
end

@testset "spincouplings" begin
    @test heisenberg"sl(1, 1)⊗ob(1, 3)@(1, 2)" == Heisenberg(centers=(1, 2), atoms=(1, 1), orbitals=(1, 3))
    @test heisenberg"sl(1, 1)⊗ob(1, 3)" == Heisenberg(atoms=(1, 1), orbitals=(1, 3))
    @test heisenberg"@(1, 2)" == Heisenberg(centers=(1, 2))
    @test heisenberg"ob(1, 3)" == Heisenberg(orbitals=(1, 3))
    @test heisenberg"sl(1, 1)" == Heisenberg(atoms=(1, 1))
    @test heisenberg"" == heisenberg"+-z" == Heisenberg()
    @test heisenberg"xyz" == Heisenberg("xyz")

    @test ising"x" == Ising('x') && ising"x sl(1, 1)⊗ob(1, 3)@(1, 2)" == Ising('x', centers=(1, 2), atoms=(1, 1), orbitals=(1, 3))
    @test ising"y" == Ising('y') && ising"y sl(1, 1)⊗ob(1, 3)@(1, 2)" == Ising('y', centers=(1, 2), atoms=(1, 1), orbitals=(1, 3))
    @test ising"z" == Ising('z') && ising"z sl(1, 1)⊗ob(1, 3)@(1, 2)" == Ising('z', centers=(1, 2), atoms=(1, 1), orbitals=(1, 3))

    @test gamma"x" == Gamma('x') && gamma"x sl(1, 1)⊗ob(1, 3)@(1, 2)" == Gamma('x', centers=(1, 2), atoms=(1, 1), orbitals=(1, 3))
    @test gamma"y" == Gamma('y') && gamma"y sl(1, 1)⊗ob(1, 3)@(1, 2)" == Gamma('y', centers=(1, 2), atoms=(1, 1), orbitals=(1, 3))
    @test gamma"z" == Gamma('z') && gamma"z sl(1, 1)⊗ob(1, 3)@(1, 2)" == Gamma('z', centers=(1, 2), atoms=(1, 1), orbitals=(1, 3))

    @test dm"x" == DM('x') && dm"x sl(1, 1)⊗ob(1, 3)@(1, 2)" == DM('x', centers=(1, 2), atoms=(1, 1), orbitals=(1, 3))
    @test dm"y" == DM('y') && dm"y sl(1, 1)⊗ob(1, 3)@(1, 2)" == DM('y', centers=(1, 2), atoms=(1, 1), orbitals=(1, 3))
    @test dm"z" == DM('z') && dm"z sl(1, 1)⊗ob(1, 3)@(1, 2)" == DM('z', centers=(1, 2), atoms=(1, 1), orbitals=(1, 3))

    @test sˣ"" == Sˣ() && sˣ"sl(1)⊗ob(2)" == Sˣ(atom=1, orbital=2)
    @test sʸ"" == Sʸ() && sʸ"sl(1)⊗ob(2)" == Sʸ(atom=1, orbital=2)
    @test sᶻ"" == Sᶻ() && sᶻ"sl(1)⊗ob(2)" == Sᶻ(atom=1, orbital=2)
end

@testset "SpinTerm" begin
    term = SpinTerm{1}(:h, 1.5, 0, couplings=Sᶻ())
    @test term|>abbr == :sp
    @test otype(term|>typeof, OID{SIndex{Int}, SVector{2, Float}, SVector{2, Float}, Nothing}) == SOperator{Float, ID{OID{SIndex{Int}, SVector{2, Float}, SVector{2, Float}, Nothing}, 1}}
    @test otype(term|>typeof, OID{SIndex{Int}, SVector{2, Float}, SVector{2, Float}, Int}) == SOperator{Float, ID{OID{SIndex{Int}, SVector{2, Float}, SVector{2, Float}, Int}, 1}}

    point = Point(PID('a', 1), (0.5, 0.5), (0.0, 0.0))
    config = IDFConfig{Spin}(pid->Spin(atom=pid.site%2, norbital=2, spin=0.5), [point.pid])
    table = Table(config, usualspinindextotuple)
    term = SpinTerm{1}(:h, 1.5, 0, couplings=Sᶻ())
    operators = Operators(  SOperator(1.5, ID(OID(SIndex('a', 1, 1, 0.5, 'z'), [0.5, 0.5], [0.0, 0.0], 1))),
                            SOperator(1.5, ID(OID(SIndex('a', 1, 2, 0.5, 'z'), [0.5, 0.5], [0.0, 0.0], 2)))
    )
    @test expand(term, point, config, table) == operators

    bond = Bond(1, Point(PID('a', 1), (0.0, 0.0), (0.0, 0.0)), Point(PID('b', 1), (0.5, 0.5), (0.0, 0.0)))
    config = IDFConfig{Spin}(pid->Spin(atom=pid.site%2, norbital=2, spin=0.5), [bond.spoint.pid, bond.epoint.pid])
    table = Table(config, usualspinindextotuple)
    term = SpinTerm{2}(:J, 1.5, 1, couplings=Heisenberg())
    operators = Operators(  SOperator(1.50, ID(OID(SIndex('b', 1, 2, 0.5, 'z'), [0.5, 0.5], [0.0, 0.0], 4), OID(SIndex('a', 1, 2, 0.5, 'z'), [0.0, 0.0], [0.0, 0.0], 2))),
                            SOperator(0.75, ID(OID(SIndex('b', 1, 2, 0.5, '-'), [0.5, 0.5], [0.0, 0.0], 4), OID(SIndex('a', 1, 2, 0.5, '+'), [0.0, 0.0], [0.0, 0.0], 2))),
                            SOperator(0.75, ID(OID(SIndex('b', 1, 1, 0.5, '-'), [0.5, 0.5], [0.0, 0.0], 3), OID(SIndex('a', 1, 1, 0.5, '+'), [0.0, 0.0], [0.0, 0.0], 1))),
                            SOperator(0.75, ID(OID(SIndex('b', 1, 1, 0.5, '+'), [0.5, 0.5], [0.0, 0.0], 3), OID(SIndex('a', 1, 1, 0.5, '-'), [0.0, 0.0], [0.0, 0.0], 1))),
                            SOperator(1.50, ID(OID(SIndex('b', 1, 1, 0.5, 'z'), [0.5, 0.5], [0.0, 0.0], 3), OID(SIndex('a', 1, 1, 0.5, 'z'), [0.0, 0.0], [0.0, 0.0], 1))),
                            SOperator(0.75, ID(OID(SIndex('b', 1, 2, 0.5, '+'), [0.5, 0.5], [0.0, 0.0], 4), OID(SIndex('a', 1, 2, 0.5, '-'), [0.0, 0.0], [0.0, 0.0], 2)))
    )
    @test expand(term, bond, config, table) == operators
end


text:
s1 = "12.125"; v1 = 12.125
s2 = "-0.125"; v2 = -0.125
s3 = "5"     ; v3 = 5

@testset "parse $T" for T in (Double16, Double32, Double64)

@test T(s1) === T(v1)
@test T(s2) === T(v2)
@test T(s3) === T(v3)

end

@testset "T_str" begin

@test df64"12.125" === Double64(s1)
@test df32"-0.125" === Double32(s2)
@test df16"5" === Double16(s3)

end


text:
cd("$(homedir())/Documents/GitHub")

# Push Documents/GitHub on the search path
push!(LOAD_PATH, "$(homedir())/Documents/GitHub")

using FileIO, ExcelFiles, DataFrames
#using ExcelReaders
#using MAT
using NEPA

#mat = matread("C:/Users/u0093191/Documents/MATLAB/imputeLCM.mat")
#println(mat)

X = DataFrame(load("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!F4:H195"))
W = DataFrame(load("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!J4:L195"))
Y = DataFrame(load("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!E4:E195"))

#X = readxl("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!F4:H195")
#W = readxl("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!J4:L195")
#Y = readxl("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","Data!E4:E195")

# Convert DataFrame or DataArray to Array so that the  function WACM doesn't crash
X = convert(Array{Float64},X)
W = convert(Array{Float64},W)
Y = convert(Array{Float64},Y)

# Input-oriented efficiencies computed using Maple. We use these to check our computed results
#MapleEff = readxl("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","MapleEfficiencies!B2:O193")
MapleEff = DataFrame(load("C:/Users/u0093191/Documents/MATLAB/WACM/Atkinson-RSanalyis-11May2011.xls","MapleEfficiencies!B2:O193"))
MapleEff = convert(Array{Float64},MapleEff)

#eff = WACM(W,X,Y)
#Time execution using @time
#@time eff = WACM(w,x,y)
#@code_warntype WACM(w,x,y)
#eff = WACM(w,x,y)
#print(eff)

#Input- or output-oriented?
input = true

D = FDH_VRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,8])))
D = FDH_CRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,9])))
D = FDH_NIRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,10])))
D = FDH_NDRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,11])))

#M = convert(Int,round(3*size(D,1)/4))
#thetaM = orderm(D,M,100)
#println(thetaM - D())

# This should be equal to the FDH_VRS efficiency!
D = DDF{FreeDisposal,VRS}(X,Y,X,zeros(size(Y)))
println(maximum(abs((1 - geteff(D())) - MapleEff[:,8])))
#println(hcat((1 - geteff(D())), MapleEff[:,8]))

# Hyperbolic equals square root of input oriented FDH under CRS
GDR = Hyperbolic{FreeDisposal,CRS}(X,Y)
D = FDH_CRS(X,Y,input)
println(maximum(abs(geteff(GDR()) - sqrt(geteff(D())))))

D = DEA_VRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,1])))
D = DEA_CRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,2])))
D = DEA_NIRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,3])))
D = DEA_NDRS(X,Y,input)
println(maximum(abs(geteff(D()) - MapleEff[:,4])))
#@time theta = D() #4.389059 seconds
#@time theta = D() #0.123012 seconds

D = DDF{Convex,VRS}(X,Y,X,zeros(size(Y)))
println(maximum(abs((1 - geteff(D())) - MapleEff[:,1])))

# Test indexing of DDF object
println(D[1])
println(X[1,:],Y[1,:],X[1,:],zeros(1,size(Y,2)))

# Do order-m efficiency
#M = convert(Int,round(0.9*size(D,1)))
#thetaM = orderm(D,M,100)
#println(thetaM - geteff(D()))

#data = readdlm("./NEPA/data/GriffellTatjéLovell.txt")
#X0 = data[:,1]
#Y0 = data[:,2]
#X1 = data[:,3]
#Y1 = data[:,4]

#LTFP = Luenberger{Convex,VRS}(X,Y,X,Y,X,Y,X,Y)
#println(LTFP())
#println(TEI(LTFP))
#println(TC(LTFP))

# HMBTFP = HMB{Convex,VRS}(X,Y,X,Y)
# println(HMBTFP())
# println(TEI(HMBTFP))
# println(TC(HMBTFP))
# println(SEC(HMBTFP))

max_stars_repo_name: kepiej/NEPA


text:
# ---
# title: 89. Gray Code
# id: problem89
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Backtracking
# link: <https://leetcode.com/problems/gray-code/description/>
# hidden: true
# ---
#
# The gray code is a binary numeral system where two successive values differ in
# only one bit.
#
# Given a non-negative integer _n_ representing the total number of bits in the
# code, print the sequence of gray code. A gray code sequence must begin with 0.
#
# **Example 1:**
#
#
#
#     Input:  2
#     Output:  [0,1,3,2]
#     Explanation:
#     00 - 0
#     01 - 1
#     11 - 3
#     10 - 2
#
#     For a given  _n_ , a gray code sequence may not be uniquely defined.
#     For example, [0,2,3,1] is also a valid gray code sequence.
#
#     00 - 0
#     10 - 2
#     11 - 3
#     01 - 1
#
#
# **Example 2:**
#
#
#
#     Input:  0
#     Output:  [0]
#     Explanation: We define the gray code sequence to begin with 0.
#                  A gray code sequence of _n_ has size = 2n, which for _n_ = 0 the size is 20 = 1.
#                  Therefore, for _n_ = 0 the gray code sequence is [0].
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
using YAActL, Printf

struct Prt <: Message        # define a message
    txt::String
end

# define two behaviors accepting a msg::Message as their last argument
function pr(msg::Prt)
    print(@sprintf("%s\n", msg.txt))
    become(pr, "Next") # change behavior
end
pr(info, msg::Prt) = print(@sprintf("%s: %s\n", info, msg.txt))

# a behavior for doing arithmetic
function calc(op::F, v::U, msg::Request) where {F<:Function,U<:Number}
    send!(msg.from, Response(op(v,msg.x)))
end

# start an actor with the first behavior and save the returned link
myactor = Actor(pr)

send!(myactor, Prt("My first actor"))  # send a message to it

send!(myactor, Prt("Something else"))  # send again a message

become!(myactor, pr, "New behavior")   # change the behavior to another one

send!(myactor, Prt("bla bla bla"))     # and send again a message

become!(myactor, calc, +, 10);         # now become a adding machine

send!(myactor, Request(5, USR));       # send a request to add 5

take!(USR.chn)                             # take the result


text:
@inline _reddims(y) = ((i for i=1:ndims(y)-2)..., ndims(y))

"""
Initializes an identity convolutional filter of the given size.
Useful for Net2DeeperNet.
"""
function id_filter(w1, w2, cx, cy)
    w = zeros(w1, w2, cx, cy)
    center_1::Int = ceil(w1/2)
    center_2::Int = ceil(w2/2)
    for i in 1:cy
        tmp = zeros(w1, w2, cy)
        tmp[center_1, center_2, i] = 1
        w[:, :, :, i] = tmp
    end
    return w
end

"""
    deeper_conv(layers, layer_index, dtrn=nothing)

Net2WiderDeeperNet method for Conv layers.
Creates a new convolutional layer with identity filters that fit with the previous layer.
Function is preserved during the deepening.
If layer uses batchnorm, training minibatches are needed for initialization.

`layers` should be the array of layers of your network
`layer_index` should be the index of the layer that you want to use for deepening
If the layer is of type ConvBN, `dtrn` should be the training minibatches that your network was trained with.
If batch normalization is not used, `dtrn` can be left as `nothing`
"""
function deeper_conv(layers, layer_index, dtrn=nothing)
    prev_layer = layers[layer_index]
    if !(prev_layer isa Conv) && !(prev_layer isa ConvBN)
        print("Layer is not convolutional!")
    end

    prev_w = prev_layer.w
    deeper_w = param(size(prev_w, 1), size(prev_w, 2), size(prev_w, 4), size(prev_w, 4); init=id_filter, atype=atype())
    deeper_b = param0(1, 1, size(prev_w, 4), 1; atype=atype())
    if prev_layer isa Conv
        deeper_layer = Conv(deeper_w, deeper_b, prev_layer.f, Int(floor(size(prev_w, 1)/2)), prev_layer.stride)
    else
        deeper_layer = ConvBN(deeper_w, deeper_b, prev_layer.f, Int(floor(size(prev_w, 1)/2)), prev_layer.stride, Param(convert(atype(), bnparams(size(prev_w, 4)))), bnmoments())
        # Forward computation to set the new bn moments and params correctly
        for (x, y) in dtrn
            h = x
            for i in 1:layer_index
                h = layers[i](h)
            end
            ah = convert(Array, h)
            dims = _reddims(ah)
            _lazy_init!(deeper_layer.bn_moments, ah)

            mu = mean(ah, dims=dims)
            # sigma2 = var(h; corrected=false, mean=mu, dims=dims)
            sigma2 = Statistics._var(ah, false, convert(Array, mu), dims)

            _update_moments!(deeper_layer.bn_moments, mu, sigma2)
        end

        deeper_layer.bn_moments.mean = convert(atype(), deeper_layer.bn_moments.mean)
        deeper_layer.bn_moments.var = convert(atype(), deeper_layer.bn_moments.var)
        eps = 1e-5
        gamma = sqrt.(deeper_layer.bn_moments.var .+ eps)
        beta = deeper_layer.bn_moments.mean

        deeper_layer.bn_params = vcat(reshape(gamma, size(gamma, 3)), reshape(beta, size(beta, 3)))
    end
    insert!(layers, layer_index+1, deeper_layer)
    return deeper_layer
end

"""
    deeper_inception(layers, layer_index, dtrn, deepening_factor=2)

Net2WiderDeeperNet method for Inception modules.
Works with both InceptionA and InceptionB.
Creates new conv layers for each conv that is not 1x1 in the module.
Given deepening factor determines how many layers per layer are created.
Function is preserved during the deepening.

`layers` should be the array of layers of your network
`layer_index` should be the index of the layer that you want to use for deepening
The layer at 'layer_index' can be of type `InceptionA` or `InceptionB`
`dtrn` should be the training minibatches that your network was trained with.
`deepening_factor` determines the amount of deepening
"""
function deeper_inception(layers, layer_index, dtrn, deepening_factor=2)
    inc_layer = layers[layer_index]
    if !(inc_layer isa InceptionA) && !(inc_layer isa InceptionB)
        print("Layer is not an Inception module!")
        return
    end

    if inc_layer isa InceptionA
        deeper_layer = InceptionADeeper(
            inc_layer.c1_before_3, inc_layer.c1_before_d3, inc_layer.c1_after_pool,
            inc_layer.c1_alone,
            [inc_layer.c3],
            [inc_layer.cd3_1, inc_layer.cd3_2],
            inc_layer.pool_mode
        )

        c3_layers = layers[1:layer_index-1]
        push!(c3_layers, inc_layer.c1_before_3, inc_layer.c3)
        new_conv = deeper_conv(c3_layers, layer_index+1, dtrn)
        for i in 1:deepening_factor
            push!(deeper_layer.c3s, deepcopy(new_conv))
        end

        cd3_layers = layers[1:layer_index-1]
        push!(cd3_layers, inc_layer.c1_before_d3, inc_layer.cd3_1, inc_layer.cd3_2)
        new_conv = deeper_conv(cd3_layers, layer_index+2, dtrn)
        for i in 1:deepening_factor
            push!(deeper_layer.cd3s, deepcopy(new_conv))
        end

        layers[layer_index] = deeper_layer
    else
        deeper_layer = InceptionBDeeper(
            inc_layer.c1_before_3, inc_layer.c1_before_d3,
            [inc_layer.c3],
            [inc_layer.cd3_1, inc_layer.cd3_2]
        )

        c3_layers = layers[1:layer_index-1]
        push!(c3_layers, inc_layer.c1_before_3, inc_layer.c3)
        new_conv = deeper_conv(c3_layers, layer_index+1, dtrn)
        new_conv.stride = 1
        for i in 1:deepening_factor
            push!(deeper_layer.c3s, deepcopy(new_conv))
        end

        cd3_layers = layers[1:layer_index-1]
        push!(cd3_layers, inc_layer.c1_before_d3, inc_layer.cd3_1, inc_layer.cd3_2)
        new_conv = deeper_conv(cd3_layers, layer_index+2, dtrn)
        new_conv.stride = 1
        for i in 1:deepening_factor
            push!(deeper_layer.cd3s, deepcopy(new_conv))
        end

        layers[layer_index] = deeper_layer
    end
end

#-------------------------------------------------------------------------------
# ---------------------------- TESTS BEGIN -------------------------------------
#-------------------------------------------------------------------------------

"""
Test for checking if deeper functions for Conv layers deepen properly and are function preserving
"""
function test_deeper_conv(with_bn=true)
    (xtrn, ytrn), (xtst, ytst) = load_data()

    dtrn = minibatch(xtrn, ytrn, 50, xtype=atype())
    dtst = minibatch(xtst, ytst, 50, xtype=atype())

    chosen_layer = 4

    cnn_model = create_cnn_model(3, 10, with_bn)
    file_name = with_bn ? "cnn.jld2" : "cnn_no_bn.jld2"
    cnn_results, cnn_model = train_results(dtrn, dtst, file_name, cnn_model, 5, false)
    old_layer_count = length(cnn_model.layers)

    cnn_deeper = deepcopy(cnn_model)
    deeper_conv(cnn_deeper.layers, chosen_layer, dtrn)

    @assert (length(cnn_model.layers) == old_layer_count) "Old model is modified"
    @assert (length(cnn_deeper.layers) == old_layer_count + 1) "New model is not deepened correctly"

    sames = 0
    total = 0
    for (x, y) in dtst
        y_olds = x
        y_news = x

        for i in 1:chosen_layer
            if sum(y_news .- y_olds) != 0
                println("Old layers are changed! i = ", i)
            end
            y_olds = cnn_model.layers[i](y_olds)
            y_news = cnn_deeper.layers[i](y_news)
        end

        y_news = cnn_deeper.layers[chosen_layer+1](y_news)
        differents = abs.(y_olds .- y_news) .< 0.01

        sames_mb = Int(sum(differents))

        sames += sames_mb
        total += length(y_olds)
    end

    @show sames/total
    @assert sames/total == 1 "Function is not preserved"
    println("deeper conv test passed")
end

"""
Test for checking if deeper functions for Inception modules deepen properly and are function preserving
"""
function test_deeper_inception()
    (xtrn, ytrn), (xtst, ytst) = load_data()

    dtrn = minibatch(xtrn, ytrn, 50, xtype=atype())
    dtst = minibatch(xtst, ytst, 50, xtype=atype())

    chosen_layer_a = 3
    chosen_layer_b = 5

    model = create_inception_bn_sm_model(3, 10)

    results, model = train_results(dtrn, dtst, "inception_sm.jld2", model, 5, false, false)

    deeper = deepcopy(model)
    deeper_inception(deeper.layers, chosen_layer_a, dtrn)
    deeper_inception(deeper.layers, chosen_layer_b, dtrn)

    @assert (model.layers[chosen_layer_a] isa InceptionA) "Old model is modified"
    @assert (model.layers[chosen_layer_b] isa InceptionB) "Old model is modified"
    @assert (deeper.layers[chosen_layer_a] isa InceptionADeeper) "New model is not deepened correctly"
    @assert (deeper.layers[chosen_layer_b] isa InceptionBDeeper) "New model is not deepened correctly"

    sames = 0
    total = 0
    for (x, y) in dtst
        y_olds = x
        y_news = x

        for i in 1:chosen_layer_b
            y_olds = model.layers[i](y_olds)
            y_news = deeper.layers[i](y_news)
        end

        differents = abs.(y_olds .- y_news) .< 0.01

        sames_mb = Int(sum(differents))

        sames += sames_mb
        total += length(y_olds)
    end

    @show sames/total
    @assert sames/total == 1 "Function is not preserved"
    println("deeper inception test passed")
end


text:
## ================================================================================================
## FlatVariables - used for packing variables for optimization
## ================================================================================================

struct FlatVariables{T<:Real}
  X::Vector{T}
  idx::Dict{Symbol, UnitRange{Int}}
end

function FlatVariables(fg::AbstractDFG, varIds::Vector{Symbol})

  index = 1
  idx = Dict{Symbol, UnitRange{Int}}()
  for vid = varIds
    v = getVariable(fg, vid)
    dims = getDimension(v)
    idx[vid] = index:(index+dims-1)
    index += dims
  end
  return FlatVariables(Vector{Float64}(undef, index-1), idx)
end

function Base.setindex!(flatVar::FlatVariables{T}, val::Vector{T}, vId::Symbol) where T<:Real
  if length(val) == length(flatVar.idx[vId])
    flatVar.X[flatVar.idx[vId]] .= val
  else
    error("array could not be broadcast to match destination")
  end
end

function Base.getindex(flatVar::FlatVariables{T}, vId::Symbol) where T<:Real
  return flatVar.X[flatVar.idx[vId]]
end

## ================================================================================================
## Parametric Factors
## ================================================================================================

"""
    $SIGNATURES

Returns the parametric measurement for a factor as a tuple (measurement, inverse covariance) for parametric inference (assuming Gaussian).
Defaults to find the parametric measurement at field `Zij`, `Z`, or `z`.

Notes
- Users should overload this method should their factor not default to `.Z<:ParametricType` or `.z<:ParametricType`
- First design choice was to restrict this function to returning coordinates
  - See https://github.com/JuliaRobotics/RoME.jl/issues/465
  - Pay attention to which tangent space point is used for converting points on a manifold to coordinates,
    - Originally written just for Lie Groups to support legacy, but future needs may well alter the design.
- Original design driven by parametric solve and dead reckon tethering.

See also: [`accumulateFactorMeans`](@ref), [`solveFactorParameteric`](@ref)
"""
function getMeasurementParametric end

function getMeasurementParametric(Z)
  error("$(typeof(Z)) is not supported, please use non-parametric or open an issue if it should be")
end

function getMeasurementParametric(Z::Normal)
  meas = mean(Z)
  iσ = 1/std(Z)^2
  return [meas], reshape([iσ],1,1)
end

function getMeasurementParametric(Z::MvNormal)
  meas = mean(Z)
  iΣ = invcov(Z)
  return meas, iΣ
end

# the point `p` on the manifold is the mean
function getMeasurementParametric(s::ManifoldPrior)
  meas = s.p
  iΣ = invcov(s.Z)
  return meas, iΣ
end

function getMeasurementParametric(s::AbstractFactor)
  if hasfield(typeof(s), :Zij)
    Z = s.Zij
    @info "getMeasurementParametric falls back to using field `.Zij` by default. Extend it for more complex factors." maxlog=1
  elseif hasfield(typeof(s), :Z)
    Z = s.Z
    @info "getMeasurementParametric falls back to using field `.Z` by default. Extend it for more complex factors." maxlog=1
  elseif hasfield(typeof(s), :z)
    Z = s.z
    @info "getMeasurementParametric falls back to using field `.z` by default. Extend it for more complex factors." maxlog=1
  else
    error("getMeasurementParametric(::$(typeof(s))) not defined, please add it, or use non-parametric, or open an issue for help.")
  end

  return getMeasurementParametric(Z)
end

## ================================================================================================
## Parametric solve with Mahalanobis distance - CalcFactor
## ================================================================================================

getFactorMechanics(f::AbstractFactor) = f
getFactorMechanics(f::Mixture) = f.mechanics

function CalcFactorMahalanobis(fct::DFGFactor)
  cf = getFactorType(fct)
  varOrder = getVariableOrder(fct)

  _meas, _iΣ = getMeasurementParametric(cf)
  meas = typeof(_meas) <: Tuple ? _meas : (_meas,)
  iΣ = typeof(_iΣ) <: Tuple ? _iΣ : (_iΣ,)

  calcf = CalcFactor(getFactorMechanics(cf), _getFMdThread(fct), 0, 0, (), [])

  multihypo = getSolverData(fct).multihypo
  nullhypo = getSolverData(fct).nullhypo

  # FIXME, type instability, use dispatch instead of if-else
  if length(multihypo) > 0
    special = MaxMultihypo(multihypo)
  elseif nullhypo > 0
    special = MaxNullhypo(nullhypo)
  elseif cf isa Mixture
    special = MaxMixture(cf.diversity.p, Ref(0))
  else
    special = nothing
  end

  return CalcFactorMahalanobis(calcf, varOrder, meas, iΣ, special)
end

# This is where the actual parametric calculation happens, CalcFactor equivalent for parametric
function (cfp::CalcFactorMahalanobis{CalcFactor{T, U, V, W}})(variables...) where {T<:AbstractFactor, U, V, W}
  # call the user function (be careful to call the new CalcFactor version only!!!)
  res = cfp.calcfactor!(cfp.meas[1], variables...)
  # 1/2*log(1/(  sqrt(det(Σ)*(2pi)^k) ))  ## k = dim(μ)
  return res' * cfp.iΣ[1] * res
end

function (cfp::CalcFactorMahalanobis{CalcFactor{T, U, V, W}})(variables...) where {T<:Union{ManifoldFactor, ManifoldPrior} , U, V, W}
  # call the user function (be careful to call the new CalcFactor version only!!!)
  M = cfp.calcfactor!.factor.M
  X = cfp.calcfactor!(cfp.meas[1], variables...)
  # 1/2*log(1/(  sqrt(det(Σ)*(2pi)^k) ))  ## k = dim(μ)
  # return mahalanobus_distance2(M, X, cfp.iΣ[1])

  #TODO do something about basis?
  # Xc = get_coordinates(M, variables[1], X, DefaultOrthogonalBasis())
  # Xc = get_coordinates(M, variables[1], X, DefaultOrthonormalBasis())
  Xc = vee(M, variables[1], X)
  return X, Xc' * cfp.iΣ[1] * Xc
end

function calcFactorMahalanobisDict(fg)
  calcFactors = Dict{Symbol, CalcFactorMahalanobis}()
  for fct in getFactors(fg)
    calcFactors[fct.label] = CalcFactorMahalanobis(fct)
  end
  return calcFactors
end

# new experimental Manifold ProductRepr version
function _totalCost(M, cfvec::Vector{<:CalcFactorMahalanobis}, labeldict, points)
  #
  obj = 0

  for cfp in cfvec

      varOrder = getindex.(Ref(labeldict), cfp.varOrder)

      factor_point_params = [points[M, i] for i in varOrder]

      # call the user function
      _,retval = cfp(factor_point_params...)

      # 1/2*log(1/(  sqrt(det(Σ)*(2pi)^k) ))  ## k = dim(μ)
      obj += 1/2*retval
  end

  return obj
end

function _totalCost(cfdict::Dict{Symbol, CalcFactorMahalanobis},
                    flatvar,
                    X )
  #
  obj = 0
  for (fid, cfp) in cfdict

    varOrder = cfp.varOrder

    Xparams = [view(X, flatvar.idx[varId]) for varId in varOrder]

    # call the user function
    retval = cfp(Xparams...)

    # 1/2*log(1/(  sqrt(det(Σ)*(2pi)^k) ))  ## k = dim(μ)
    obj += 1/2*retval
  end

  return obj
end

# DEPRECATED slower version
function _totalCost(fg::AbstractDFG,
                    flatvar,
                    X )
  #
  obj = 0
  for fct in getFactors(fg)

    varOrder = getVariableOrder(fct)

    Xparams = [view(X, flatvar.idx[varId]) for varId in varOrder]

    cfp = CalcFactorMahalanobis(fct)
    retval = cfp(Xparams...)

    # 1/2*log(1/(  sqrt(det(Σ)*(2pi)^k) ))  ## k = dim(μ)
    obj += 1/2*retval
  end

  return obj
end

export solveGraphParametric

"""
    $SIGNATURES

Batch solve a Gaussian factor graph using Optim.jl. Parameters can be passed directly to optim.
Notes:
  - Only :Euclid and :Circular manifolds are currently supported, own manifold are supported with `algorithmkwargs` (code may need updating though)
"""
function solveGraphParametric(fg::AbstractDFG;
                              useCalcFactor::Bool=true, #TODO dev param will be removed
                              solvekey::Symbol=:parametric,
                              autodiff = :forward,
                              algorithm=Optim.BFGS,
                              algorithmkwargs=(), # add manifold to overwrite computed one
                              options = Optim.Options(allow_f_increases=true,
                                                      time_limit = 100,
                                                      # show_trace = true,
                                                      # show_every = 1,
                                                      ))

  #Other options
  # options = Optim.Options(time_limit = 100,
  #                     iterations = 1000,
  #                     show_trace = true,
  #                     show_every = 1,
  #                     allow_f_increases=true,
  #                     g_tol = 1e-6,
  #                     )
  varIds = listVariables(fg)

  #TODO mabye remove sorting, just for convenience
  sort!(varIds, lt=natural_lt)

  flatvar = FlatVariables(fg, varIds)

  for vId in varIds
    flatvar[vId] = getVariableSolverData(fg, vId, solvekey).val[1][:]
  end

  initValues = flatvar.X

  mc_mani = IIF.MixedCircular(fg, varIds)
  alg = algorithm(;manifold=mc_mani, algorithmkwargs...)
  # alg = algorithm(; algorithmkwargs...)

  if useCalcFactor
    cfd = calcFactorMahalanobisDict(fg)
    tdtotalCost = Optim.TwiceDifferentiable((x)->_totalCost(cfd, flatvar, x), initValues, autodiff = autodiff)
  else
    tdtotalCost = Optim.TwiceDifferentiable((x)->_totalCost(fg, flatvar, x), initValues, autodiff = autodiff)
  end

  result = Optim.optimize(tdtotalCost, initValues, alg, options)
  rv = Optim.minimizer(result)

  H = Optim.hessian!(tdtotalCost, rv)

  Σ = pinv(H)

  d = Dict{Symbol,NamedTuple{(:val, :cov),Tuple{Vector{Float64},Matrix{Float64}}}}()

  for key in varIds
    r = flatvar.idx[key]
    push!(d,key=>(val=rv[r],cov=Σ[r,r]))
  end

  return d, result, flatvar.idx, Σ
end

#TODO maybe consolidate with solveGraphParametric
#TODO WIP
```
    $SIGNATURES
Solve for frontal values only with values in seprarators fixed
```
function solveConditionalsParametric(fg::AbstractDFG,
                                    frontals::Vector{Symbol};
                                    solvekey::Symbol=:parametric,
                                    autodiff = :forward,
                                    algorithm=Optim.BFGS,
                                    algorithmkwargs=(), # add manifold to overwrite computed one
                                    options = Optim.Options(allow_f_increases=true,
                                                            time_limit = 100,
                                                            # show_trace = true,
                                                            # show_every = 1,
                                                            ))

  varIds = listVariables(fg)

  #TODO mabye remove sorting, just for convenience
  sort!(varIds, lt=natural_lt)
  separators = setdiff(varIds, frontals)

  varIds = [frontals; separators]

  flatvar = FlatVariables(fg, varIds)

  for vId in varIds
    flatvar[vId] = getVariableSolverData(fg, vId, solvekey).val[1][:]
  end
  initValues = flatvar.X

  frontalsLength = sum(map(v->getDimension(getVariable(fg, v)), frontals))

  # build variables for frontals and seperators
  # fX = view(initValues, 1:frontalsLength)
  fX = initValues[1:frontalsLength]
  # sX = view(initValues, (frontalsLength+1):length(initValues))
  sX = initValues[frontalsLength+1:end]

  mc_mani = MixedCircular(fg, varIds)
  alg = algorithm(;manifold=mc_mani, algorithmkwargs...)
  # alg = algorithm(; algorithmkwargs...)

  tdtotalCost = Optim.TwiceDifferentiable((x)->_totalCost(fg, flatvar,x), initValues, autodiff = autodiff)

  result = Optim.optimize((x)->_totalCost(fg, flatvar, [x;sX]), fX, alg, options)
  # result = optimize(x->totalCost([x;sX]), fX, alg, options)

  rv = Optim.minimizer(result)

  H = Optim.hessian!(tdtotalCost, [rv; sX])

  Σ = pinv(H)

  d = Dict{Symbol,NamedTuple{(:val, :cov),Tuple{Vector{Float64},Matrix{Float64}}}}()

  for key in frontals
    r = flatvar.idx[key]
    push!(d,key=>(val=rv[r],cov=Σ[r,r]))
  end

  return d, result, flatvar.idx, Σ
end

## ================================================================================================
## MixedCircular Manifold for Optim.jl
## ================================================================================================

"""
    MixedCircular
Mixed Circular Manifold. Simple manifold for circular and cartesian mixed for use in optim

DevNotes
- Consolidate around `ManifoldsBase.AbstractManifold` instead, with possible wrapper-type solution for `Optim.Manifold`
"""
struct MixedCircular <: Optim.Manifold
  isCircular::BitArray
end

function MixedCircular(fg::AbstractDFG, varIds::Vector{Symbol})
  circMask = Bool[]
  for k = varIds
    append!(circMask, convert(Tuple, getManifold(getVariableType(fg, k))) .== :Circular)
  end
  MixedCircular(circMask)
end

# https://github.com/JuliaNLSolvers/Optim.jl/blob/e439de4c997a727f3f724ae76da54b9cc08456b2/src/Manifolds.jl#L3
# retract!(m, x): map x back to a point on the manifold m
function Optim.retract!(c::MixedCircular, x)
  for (i,v) = enumerate(x)
    c.isCircular[i] && (x[i] = rem2pi(v, RoundNearest))
  end
  return x
end
# https://github.com/JuliaNLSolvers/Optim.jl/blob/e439de4c997a727f3f724ae76da54b9cc08456b2/src/Manifolds.jl#L2
# project_tangent!(m, g, x): project g on the tangent space to m at x
Optim.project_tangent!(S::MixedCircular,g,x) = g

## ================================================================================================
## Manifolds.jl Consolidation
## TODO: Still to be completed and tested.
## ================================================================================================
# struct ManifoldsVector <: Optim.Manifold
#   manis::Vector{Manifold}
# end

# Base.getindex(mv::ManifoldsVector, inds...) = getindex(mv.mani, inds...)
# Base.setindex!(mv, X, inds...) =  setindex!(mv.mani, X, inds...)

# function ManifoldsVector(fg::AbstractDFG, varIds::Vector{Symbol})
#   manis = Bool[]
#   for k = varIds
#     push!(manis, getVariableType(fg, k) |> getManifold)
#   end
#   ManifoldsVector(manis)
# end

# function Optim.retract!(manis::ManifoldsVector, x)
#   for (i,M) = enumerate(manis)
#     x[i] = project(M, x[i])
#   end
#   return x
# end
# function Optim.project_tangent!(manis::ManifoldsVector, G, x)
#   for (i, M) = enumerate(manis)
#     G[i] = project(M, x[i], G)
#   end
#   return G
# end

## ================================================================================================
## UNDER DEVELOPMENT Parametric solveTree utils
## ================================================================================================

"""
    $SIGNATURES
Get the indexes for labels in FlatVariables
"""
function collectIdx(varinds, labels)
  idx = Int[]
  for lbl in labels
    append!(idx, varinds[lbl])
  end
  return idx
end

"""
    $SIGNATURES
Calculate the marginal distribution for a clique over subsetVarIds.
"""
function calculateMarginalCliqueLikelihood(vardict, Σ, varindxs, subsetVarIds)

  μₘ = Float64[]
  for lbl in subsetVarIds
    append!(μₘ, vardict[lbl].val)
  end

  Aidx = collectIdx(varindxs, subsetVarIds)
  Σₘ = Σ[Aidx, Aidx]

  return createMvNormal(μₘ, Σₘ)
end

"""
    $SIGNATURES

"""
function calculateCoBeliefMessage(soldict, Σ, flatvars, separators, frontals)
  Aidx = IIF.collectIdx(flatvars,separators)
  Cidx = IIF.collectIdx(flatvars,frontals)

  #marginalize separators
  A = Σ[Aidx, Aidx]
  #marginalize frontals
  C = Σ[Cidx, Cidx]
  # cross
  B = Σ[Aidx, Cidx]

  Σₘ = deepcopy(A)
  if length(separators) == 0

    return (varlbl=Symbol[], μ=Float64[], Σ=Matrix{Float64}(undef,0,0))

  elseif length(separators) == 1

    # create messages
    return (varlbl = deepcopy(separators), μ = soldict[separators[1]].val, Σ = A)

  elseif length(separators) == 2
    A = Σₘ[1, 1]
    C = Σₘ[2, 2]
    B = Σₘ[1, 2]

    #calculate covariance between separators
    ΣA_B = A - B*inv(C)*B'
    # create messages
    m2lbl = deepcopy(separators)
    m2cov = isa(ΣA_B, Matrix) ? ΣA_B : fill(ΣA_B,1,1)
    m2val = soldict[m2lbl[2]].val - soldict[m2lbl[1]].val
    return (varlbl = m2lbl, μ = m2val, Σ = m2cov)

  else
    error("Messages with more than 2 seperators are not supported yet")
  end
end

## ================================================================================================
## Parametric utils
## ================================================================================================

## SANDBOX of usefull development functions to be cleaned up

"""
    $SIGNATURES
Add parametric solver to fg, batch solve using [`solveGraphParametric`](@ref) and update fg.
"""
function solveGraphParametric!(fg::AbstractDFG; init::Bool=true, kwargs...)
  if !(:parametric in fg.solverParams.algorithms)
    addParametricSolver!(fg; init=init)
  elseif init
    initParametricFrom!(fg)
  end

  vardict, result, varIds, Σ = solveGraphParametric(fg; kwargs...)

  updateParametricSolution!(fg, vardict)

  return vardict, result, varIds, Σ
end

"""
    $SIGNATURES
Initialize the parametric solver data from a different solution in `fromkey`.
"""
function initParametricFrom!(fg::AbstractDFG, fromkey::Symbol = :default; parkey::Symbol = :parametric, onepoint=false)
  if onepoint
    for v in getVariables(fg)
      fromvnd = getSolverData(v, fromkey)
      dims = getDimension(v)
      getSolverData(v, parkey).val[1:dims] .= fromvnd.val[1:dims]#zeros(dims)*1e-5
      getSolverData(v, parkey).bw[1:dims,1:dims] .= LinearAlgebra.I(dims)
    end
  else
    for var in getVariables(fg)
        fromvnd = getSolverData(var, fromkey)
        ptr_ = fromvnd.val
        @cast ptr[i,j] := ptr_[j][i]
        if fromvnd.dims == 1
          nf = fit(Normal, ptr)
          getSolverData(var, parkey).val[1][1] = nf.μ
          getSolverData(var, parkey).bw[1,1] = nf.σ
          # m = var.estimateDict[:default].mean
        else
          #FIXME circular will not work correctly with MvNormal
          nf = fit(MvNormal, ptr)
          getSolverData(var, parkey).val[1][1:fromvnd.dims] .= mean(nf)[1:fromvnd.dims]
          getSolverData(var, parkey).bw = cov(nf)
        end
    end
  end
end

"""
    $SIGNATURES
Add the parametric solveKey to all the variables in fg if it doesn't exists.
"""
function addParametricSolver!(fg; init=true)
  if !(:parametric in fg.solverParams.algorithms)
      push!(fg.solverParams.algorithms, :parametric)
      foreach(v->IIF.setDefaultNodeDataParametric!(v, getVariableType(v), initialized=false), getVariables(fg))
      if init
        initParametricFrom!(fg)
      end
  else
      error("parametric solvekey already exists")
  end
  nothing
end

function updateVariablesFromParametricSolution!(fg::AbstractDFG, vardict)
  for (v,val) in vardict
    vnd = getVariableSolverData(fg, v, :parametric)
    vnd.val .= val.val
    if size(vnd.bw) != size(val.cov)
      vnd.bw = val.cov
    else
      vnd.bw .= val.cov
    end
  end
end

"""
    $SIGNATURES
Update the fg from solution in vardict and add MeanMaxPPE (all just mean). Usefull for plotting
"""
function updateParametricSolution!(sfg, vardict)

  for (v,val) in vardict
      vnd = getSolverData(getVariable(sfg, v), :parametric)
      # fill in the variable node data value
      vnd.val .= val.val
      #calculate and fill in covariance
      vnd.bw = val.cov
      #fill in ppe as mean
      ppe = MeanMaxPPE(:parametric, val.val, val.val, val.val)
      getPPEDict(getVariable(sfg, v))[:parametric] = ppe
  end
end

function createMvNormal(val,cov)
    #TODO do something better for properly formed covariance, but for now just a hack...FIXME
    if all(diag(cov) .> 0.001) && isapprox(cov, transpose(cov), rtol=1e-4)
        return MvNormal(val,Symmetric(cov))
    else
        @error("Covariance matrix error", cov)
        # return nothing # FIXME, blanking nothing during #459 consolidation
        return MvNormal(val, ones(length(val)))
    end
end

function createMvNormal(v::DFGVariable, key=:parametric)
    if key == :parametric
        vnd = getSolverData(v, :parametric)
        dims = vnd.dims
        return createMvNormal(vnd.val[1:dims,1], vnd.bw[1:dims, 1:dims])
    else
        @warn "Trying MvNormal Fit, replace with PPE fits in future"
        return fit(MvNormal,getSolverData(v, key).val)
    end
end

## ================================================================================================
## Experimental specialized dispatch for Mixture
## ================================================================================================
# To sort out how to dispatch on specialized functions.
# related to #931 and #1069

struct MaxMixture
  p::Vector{Float64}
  # the chosen component to be used for the optimization
  choice::Base.RefValue{Int}
end

function getMeasurementParametric(s::Mixture{N,F,S,T}) where {N,F,S,T}
  meas = map(c->getMeasurementParametric(c)[1], values(s.components))
  iΣ = map(c->getMeasurementParametric(c)[2], values(s.components))
  return meas, iΣ
end

function _calcFactorMahalanobis(cfp, meas, iΣ, variables...)
  res = cfp.calcfactor!(meas, variables...)
  r = res' * iΣ * res
  return r
end

# DEV NOTE: function with other options including select once and use
# function (cfp::CalcFactorMahalanobis{<:CalcFactor, MaxMixture})(variables...)
#   if cfp.specialAlg.choice[] == 0
#     #calculate all mixture options
#     r = [_calcFactorMahalanobis(cfp, cfp.meas[i], cfp.iΣ[i], variables...) for i = 1:length(cfp.meas)]

#     p = cfp.specialAlg.p

#     k = size(cfp.iΣ[1], 2)
#     # α = 1 ./ sqrt.(2pi .* k .* det.(inv.(cfp.iΣ)))
#     α = sqrt.(det.(cfp.iΣ) ./ ((2pi)^k))

#     # mm, at = findmax(α .* p .* exp.(-0.5 .* r))
#     # mm = sum(α .* p .* exp.(-0.5 .* r) )

#     mm, at = findmin( 0.5 .* r .- log.(α .* p))
#     # mm = -log(sum(α .* p .* exp.(-0.5 .* r) ))
#     # return mm + maximum(log.(α .* p))

#     cfp.specialAlg.choice[] = at

#     return r[at]

#   else
#     at = cfp.specialAlg.choice[]
#     return _calcFactorMahalanobis(cfp, cfp.meas[at], cfp.iΣ[at], variables...)
#   end

# end

function (cfp::CalcFactorMahalanobis{<:CalcFactor, MaxMixture})(variables...)

  r = [_calcFactorMahalanobis(cfp, cfp.meas[i], cfp.iΣ[i], variables...) for i = 1:length(cfp.meas)]

  p = cfp.specialAlg.p

  k = size(cfp.iΣ[1], 2)
  # α = 1 ./ sqrt.(2pi .* k .* det.(inv.(cfp.iΣ)))
  α = sqrt.(det.(cfp.iΣ) ./ ((2pi)^k))

  mm, at = findmin(r .- log.(α .* p))
  # mm = -log(sum(α .* p .* exp.(-0.5 .* r) ))
  return mm + maximum(log.(α .* p))

end

## ================================================================================================
## Experimental specialised dispatch for multihypo and nullhypo
## ================================================================================================
#TODO better dispatch

struct MaxMultihypo
  multihypo::Vector{Float64}
end
struct MaxNullhypo
  nullhypo::Float64
end

function (cfp::CalcFactorMahalanobis{<:CalcFactor, MaxMultihypo})(X1, L1, L2)
  mh = cfp.specialAlg.multihypo
  @assert length(mh) == 3 "multihypo $mh  not supported with parametric, length should be 3"
  @assert mh[1] == 0 "multihypo $mh  not supported with parametric, first should be 0"

  #calculate both multihypo options
  r1 = cfp(X1, L1)
  r2 = cfp(X1, L2)
  r = [r1, r2]

  # hacky multihypo to start of with
  mm, at = findmin(r .* (1 .- mh[2:end]))
  nat = at == 1 ? 1 : 2
  k = length(X1)*one(r1) * 1e-3
  return r[at] + r[nat]*k

end

function (cfp::CalcFactorMahalanobis{<:CalcFactor, MaxNullhypo})(X1, X2)
  nh = cfp.specialAlg.nullhypo
  @assert nh > 0 "nullhypo $nh not as expected"

  #calculate factor residual
  res = cfp.calcfactor!(cfp.meas[1], X1, X2)
  r1 =  res' * cfp.iΣ * res

  # compare to uniform nullhypo
  r2 = length(res)*one(r1)
  r = [r1,r2]
  mm, at = findmin(r .* [nh, (1-nh)])

  residual = at == 1 ? r1 : r1*1e-3

  return residual

  # rand residual option
  # idx = rand(Categorical([(1-nh), nh]))
  # nh == 0.05 && cfp.varOrder==[:x1,:l1] && println("$idx -> $(r1.value), $r2")
  # return r[idx]

end


text:
# ---
# title: 373. Find K Pairs with Smallest Sums
# id: problem373
# author: Indigo
# date: 2021-06-14
# difficulty: Medium
# categories: Heap
# link: <https://leetcode.com/problems/find-k-pairs-with-smallest-sums/description/>
# hidden: true
# ---
#
# You are given two integer arrays **nums1** and **nums2** sorted in ascending
# order and an integer **k**.
#
# Define a pair **(u,v)** which consists of one element from the first array and
# one element from the second array.
#
# Find the k pairs **(u 1,v1),(u2,v2) ...(uk,vk)** with the smallest sums.
#
# **Example 1:**
#
#
#
#     Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
#     Output: [[1,2],[1,4],[1,6]]
#     Explanation: The first 3 pairs are returned from the sequence:
#                  [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
#
# **Example 2:**
#
#
#
#     Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
#     Output: [1,1],[1,1]
#     Explanation: The first 2 pairs are returned from the sequence:
#                  [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
#
# **Example 3:**
#
#
#
#     Input: nums1 = [1,2], nums2 = [3], k = 3
#     Output: [1,3],[2,3]
#     Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]
#
#
#
## @lc code=start
using LeetCode
using DataStructures
function k_smallest_pairs(nums1::Vector{Int}, nums2::Vector{Int}, k::Int)
    hp = Tuple{Int, Int}[]
    odr = Base.Order.By(x -> -x[1] - x[2])
    for n1 in nums1, n2 in nums2
        heappush!(hp, (n1, n2), odr)
        length(hp) > k && heappop!(hp, odr)
    end
    sort!(hp; by = x -> x[1] + x[2])
end
## @lc code=end


text:
# Temperature and pressure coupling

export
    NoCoupling,
    apply_coupling!,
    AndersenThermostat,
    RescaleThermostat,
    FrictionThermostat,
    maxwell_boltzmann,
    temperature

"""
    NoCoupling()

Placeholder coupler that does nothing.
"""
struct NoCoupling <: AbstractCoupler end

"""
    apply_coupling!(system, simulator, coupling)

Apply a coupler to modify a simulation.
Custom couplers should implement this function.
"""
function apply_coupling!(sys::System, simulator, ::NoCoupling)
    return sys
end

"""
    AndersenThermostat(temperature, coupling_const)

Rescale random velocities according to the Andersen thermostat.
"""
struct AndersenThermostat{T, C} <: AbstractCoupler
    temperature::T
    coupling_const::C
end

function apply_coupling!(sys::System{D}, sim, thermostat::AndersenThermostat) where D
    for i in 1:length(sys)
        if rand() < (sim.dt / thermostat.coupling_const)
            mass = sys.atoms[i].mass
            sys.velocities[i] = velocity(mass, thermostat.temperature; dims=D)
        end
    end
    return sys
end

struct RescaleThermostat{T} <: AbstractCoupler
    temperature::T
end

function apply_coupling!(sys::System, sim, thermostat::RescaleThermostat)
    sys.velocities *= sqrt(thermostat.temperature / temperature(s))
    return sys
end

struct FrictionThermostat{T} <: AbstractCoupler
    friction_const::T
end

function apply_coupling!(sys::System, sim, thermostat::FrictionThermostat)
    sys.velocities *= thermostat.friction_const
    return sys
end

"""
    velocity(mass, temperature; dims=3)

Generate a random velocity from the Maxwell-Boltzmann distribution.
"""
function AtomsBase.velocity(mass, temp; dims::Integer=3)
    return SVector([maxwell_boltzmann(mass, temp) for i in 1:dims]...)
end

"""
    maxwell_boltzmann(mass, temperature)

Generate a random speed along one dimension from the Maxwell-Boltzmann distribution.
"""
function maxwell_boltzmann(mass, temp)
    T = typeof(convert(AbstractFloat, ustrip(temp)))
    k = unit(temp) == NoUnits ? one(T) : uconvert(u"u * nm^2 * ps^-2 * K^-1", T(Unitful.k))
    σ = sqrt(k * temp / mass)
    return rand(Normal(zero(T), T(ustrip(σ)))) * unit(σ)
end

"""
    temperature(system)

Calculate the temperature of a system from the kinetic energy of the atoms.
"""
function temperature(s::System{D, S, false}) where {D, S}
    ke = sum([a.mass * dot(s.velocities[i], s.velocities[i]) for (i, a) in enumerate(s.atoms)]) / 2
    df = 3 * length(s) - 3
    T = typeof(ustrip(ke))
    k = unit(ke) == NoUnits ? one(T) : uconvert(u"K^-1" * unit(ke), T(Unitful.k))
    return 2 * ke / (df * k)
end

function temperature(s::System{D, S, true}) where {D, S}
    ke = sum(mass.(s.atoms) .* sum.(abs2, s.velocities)) / 2
    df = 3 * length(s) - 3
    T = typeof(ustrip(ke))
    k = unit(ke) == NoUnits ? one(T) : uconvert(u"K^-1" * unit(ke), T(Unitful.k))
    return 2 * ke / (df * k)
end