text:
# tests for strategy functions.

# using Base.Test

# push!(LOAD_PATH, ".")
# include("includes.jl")

function get_test_state()
    uav_state_own = UAVState(100., -10., 0., 0., 1., 0., 0., 0.)
    uav_state_int = UAVState(0., 10., 0., 0., 1., 0., 0., 0.)
    s = State([uav_state_own; uav_state_int])
    s.responding_states = [1,1]
    return s
end

function test_modify_utilities()
    uav = get_real_uav()
    state = get_test_state()
    # second signal doesn't matter here
    state.signals = [:dont_turn_left, :dont_turn_right]
    if uav.is_policy_joint # joint policy
        utilities = zeros(size(uav.joint_actions, 2))
        # intruder signalling to take same sense
        mod_utilities = modify_utilities(uav, state, utilities)
        exp_utilities = zeros(length(uav.joint_actions[1, :]))
        for aidx in 1:size(uav.joint_actions, 2)
            if sign(uav.joint_actions[1, aidx]) == 1
                exp_utilities[aidx] = -Inf
            end
        end
    else
        utilities = zeros(length(uav.actions))
        # intruder signalling to take same sense
        mod_utilities = modify_utilities(uav, state, utilities)
        exp_utilities = zeros(length(uav.actions))
        for aidx in 1:length(uav.actions)
            if sign(uav.actions[aidx]) == 1
                exp_utilities[aidx] = -Inf
            end
        end
    end
    @test mod_utilities == exp_utilities
end

function test_select_action_and_signals_best()

    environment = build_environment()
    environment.uavs[1].action_selection_method = "best"
    environment.uavs[2].action_selection_method = "best"

    no_prev_signals = [:no_signal, :no_signal]

    uav = environment.uavs[1]
    if uav.is_policy_joint # joint policy
        utilities = zeros(size(uav.joint_actions, 2))
        utilities[1] = 1
        sense = :same_sense
        action, signals = select_action_and_signals(uav, utilities, sense, no_prev_signals)
        @test signals[1] == signals[2]
        @test action == -deg2rad(12)

        sense = :different_sense
        utilities = zeros(size(uav.joint_actions, 2))
        utilities[3] = 1
        action, signals = select_action_and_signals(uav, utilities, sense, no_prev_signals)
        @test signals[1] != signals[2]
        @test action == 0.0

        uav = environment.uavs[2]
        m = MersenneTwister(1)
        utilities = rand(m, size(uav.joint_actions, 2))
        sense = :neither_sense
        action, signals = select_action_and_signals(uav, utilities, sense, no_prev_signals)
        best_action_idx = indmax(utilities)
        best_action = uav.joint_actions[uav.uav_id, best_action_idx]
        @test action == best_action
        @test signals == [:no_signal, :no_signal]

        uav = environment.uavs[2]
        utilities = zeros(size(uav.joint_actions, 2))
        utilities[21] = 10
        sense = :same_sense
        action, signals = select_action_and_signals(uav, utilities, sense, no_prev_signals)
        @test action == deg2rad(6)
        @test signals[1] == signals[2]
        uav = environment.uavs[1]
        action, signals = select_action_and_signals(uav, utilities, sense, no_prev_signals)
        @test action == 0.0
        @test signals[1] == signals[2]

        utilities = collect(1.:36.)
        sense = :same_sense
        action, signals = select_action_and_signals(uav, utilities, sense, no_prev_signals)
        @test action == COC_ACTION
        @test signals[1] == signals[2]
        @test signals[1] == :no_signal
    else
        utilities = zeros(length(uav.actions))
        utilities[1] = 1
        sense = :same_sense
        action, signals = select_action_and_signals(uav, utilities, sense, no_prev_signals)
        @test signals[1] == signals[2]
        @test action == -deg2rad(12)

        sense = :different_sense
        utilities = zeros(length(uav.actions))
        utilities[3] = 1
        action, signals = select_action_and_signals(uav, utilities, sense, no_prev_signals)
        @test signals[1] != signals[2]
        @test action == 0.0

        uav = environment.uavs[2]
        m = MersenneTwister(1)
        utilities = rand(m, length(uav.actions))
        sense = :neither_sense
        action, signals = select_action_and_signals(uav, utilities, sense, no_prev_signals)
        best_action_idx = indmax(utilities)
        best_action = uav.actions[best_action_idx]
        @test action == best_action
        @test signals == [:no_signal, :no_signal]

        utilities = collect(1.:6.)
        sense = :same_sense
        action, signals = select_action_and_signals(uav, utilities, sense, no_prev_signals)
        @test action == COC_ACTION
        @test signals[1] == signals[2]
        @test signals[1] == :no_signal
    end
end

function test_select_action_and_signals_worst()
    environment = build_environment()
    environment.uavs[1].action_selection_method = "worst"
    environment.uavs[2].action_selection_method = "worst"
    uav = environment.uavs[1]

    if uav.is_policy_joint # joint policy
        utilities = ones(size(uav.joint_actions, 2)) * -1
        utilities[[6, 12, 18, 24, 30, 36]] = 1
        sense = :same_sense
        action, signals = select_action_and_signals(uav, utilities, sense, [:no_signal, :no_signal])
        @test action == COC_ACTION
        @test signals[1] == :no_signal
        @test signals[1] == signals[2]

        uav = environment.uavs[2]
        utilities[30:end] = 1
        action, signals = select_action_and_signals(uav, utilities, sense, [:no_signal, :no_signal])
        @test action == COC_ACTION
        @test signals[1] == :no_signal
        @test signals[1] == signals[2]

        uav = environment.uavs[1]
        utilities = ones(size(uav.joint_actions, 2)) * -1
        utilities[[1, 7, 13, 19, 25, 31]] = 1
        utilities[1] = -100
        action, signals = select_action_and_signals(uav, utilities, sense, [:no_signal, :no_signal])
        @test action != -deg2rad(12)

        utilities = ones(size(uav.joint_actions, 2))
        utilities[32:end] = -1
        action, signals = select_action_and_signals(uav, utilities, sense, [:no_signal, :no_signal])
        @test action == -deg2rad(12)
        @test signals[1] == :dont_turn_left
        @test signals[2] == :dont_turn_left

        uav = environment.uavs[1]
        utilities = ones(size(uav.joint_actions, 2))
        state = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        state.signals = [:dont_turn_left, :dont_turn_left]
        utilities = modify_utilities(uav, state, utilities)
        action, signals = select_action_and_signals(uav, utilities, sense, state.signals)
        @test action == 0.0
        @test signals[1] == :dont_turn_left
        @test signals[2] == :dont_turn_left

        uav = environment.uavs[1]
        utilities = ones(size(uav.joint_actions, 2))
        state = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        state.signals = [:dont_turn_right, :dont_turn_right]
        utilities = modify_utilities(uav, state, utilities)
        utilities[[3, 9, 15, 21, 27, 33]] = -1
        utilities[[4, 10, 16, 22, 28, 34]] = -1
        utilities[[6, 12, 18, 24, 30, 36]] = -1
        action, signals = select_action_and_signals(uav, utilities, sense, state.signals)
        @test action == deg2rad(12)
        @test signals[1] == :dont_turn_right
        @test signals[2] == :dont_turn_right

        uav = environment.uavs[1]
        utilities = ones(size(uav.joint_actions, 2))
        sense = :different_sense
        state = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        state.signals = [:dont_turn_right, :dont_turn_right]
        utilities = modify_utilities(uav, state, utilities)
        utilities[[3, 9, 15, 21, 27, 33]] = -1
        utilities[[4, 10, 16, 22, 28, 34]] = -1
        utilities[[6, 12, 18, 24, 30, 36]] = -1
        action, signals = select_action_and_signals(uav, utilities, sense, state.signals)
        @test action == deg2rad(12)
        @test signals[1] == :dont_turn_right
        @test signals[2] == :dont_turn_left
    end
end

function test_select_action_and_signals_average()
    environment = build_environment()
    environment.uavs[1].action_selection_method = "average"
    environment.uavs[2].action_selection_method = "average"
    sense = :same_sense
    uav = environment.uavs[1]
    if uav.is_policy_joint # joint policy
        utilities = ones(size(uav.joint_actions, 2))
        state = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        state.signals = [:dont_turn_right, :dont_turn_right]
        utilities = modify_utilities(uav, state, utilities)
        utilities[[3, 9, 15, 21, 27, 33]] = -1
        utilities[[4, 10, 16, 22, 28, 34]] = -.5
        utilities[[6, 12, 18, 24, 30, 36]] = -1
        utilities[17] = -10000
        action, signals = select_action_and_signals(uav, utilities, sense, state.signals)
        @test action == deg2rad(6)
        @test signals[1] == :dont_turn_right
        @test signals[2] == :dont_turn_right
    end
end

# non joint action tests
function test_select_equal_non_joint()
    environmentB = build_environment()
    environmentB.uavs[1].action_selection_method = "best"
    environmentB.uavs[2].action_selection_method = "best"

    environmentW = build_environment()
    environmentW.uavs[1].action_selection_method = "worst"
    environmentW.uavs[2].action_selection_method = "worst"

    environmentA = build_environment()
    environmentA.uavs[1].action_selection_method = "average"
    environmentA.uavs[2].action_selection_method = "average"

    no_prev_signals = [:no_signal, :no_signal]

    uavB = environmentB.uavs[1]
    uavW = environmentW.uavs[1]
    uavA = environmentA.uavs[1]

    # only a test for non-joint policies
    if !uavB.is_policy_joint # non joint policy
        println("i am in the correct test")
        utilitiesB = zeros(length(uavB.actions))
        utilitiesB[1] = 1
        senseB = :same_sense
        actionB, signalsB = select_action_and_signals(uavB, utilitiesB, senseB, no_prev_signals)
        @test signalsB[1] == signalsB[2]
        @test actionB == -deg2rad(12)

        utilitiesW = zeros(length(uavW.actions))
        utilitiesW[1] = 1
        senseW = :same_sense
        actionW, signalsW = select_action_and_signals(uavW, utilitiesW, senseW, no_prev_signals)
        @test signalsW[1] == signalsW[2]
        @test actionW == -deg2rad(12)

        utilitiesA = zeros(length(uavA.actions))
        utilitiesA[1] = 1
        senseA = :same_sense
        actionA, signalsA = select_action_and_signals(uavA, utilitiesA, senseA, no_prev_signals)
        @test signalsA[1] == signalsA[2]
        @test actionA == -deg2rad(12)

        @test actionB == actionW == actionA
        @test signalsB == signalsW == signalsA

        senseB = :different_sense
        utilitiesB = zeros(length(uavB.actions))
        utilitiesB[3] = 1
        actionB, signalsB = select_action_and_signals(uavB, utilitiesB, senseB, no_prev_signals)
        @test signalsB[1] != signalsB[2]
        @test actionB == 0.0

        senseW = :different_sense
        utilitiesW = zeros(length(uavW.actions))
        utilitiesW[3] = 1
        actionW, signalsW = select_action_and_signals(uavW, utilitiesW, senseW, no_prev_signals)
        @test signalsW[1] != signalsW[2]
        @test actionW == 0.0

        senseA = :different_sense
        utilitiesA = zeros(length(uavA.actions))
        utilitiesA[3] = 1
        actionA, signalsA = select_action_and_signals(uavA, utilitiesA, senseA, no_prev_signals)
        @test signalsA[1] != signalsA[2]
        @test actionA == 0.0

        @test actionB == actionW == actionA
        @test signalsB == signalsW == signalsA

        uavB = environmentB.uavs[2]
        m = MersenneTwister(1)
        utilitiesB = rand(m, length(uavB.actions))
        senseB = :neither_sense
        actionB, signalsB = select_action_and_signals(uavB, utilitiesB, senseB, no_prev_signals)
        best_action_idxB = indmax(utilitiesB)
        best_actionB = uavB.actions[best_action_idxB]
        @test actionB == best_actionB
        @test signalsB == [:no_signal, :no_signal]

        uavW = environmentW.uavs[2]
        m = MersenneTwister(1)
        utilitiesW = rand(m, length(uavW.actions))
        senseW = :neither_sense
        actionW, signalsW = select_action_and_signals(uavW, utilitiesW, senseW, no_prev_signals)
        best_action_idxW = indmax(utilitiesW)
        best_actionW = uavW.actions[best_action_idxW]
        @test actionW == best_actionW
        @test signalsW == [:no_signal, :no_signal]

        uavA = environmentA.uavs[2]
        m = MersenneTwister(1)
        utilitiesA = rand(m, length(uavA.actions))
        senseA = :neither_sense
        actionA, signalsA = select_action_and_signals(uavA, utilitiesA, senseA, no_prev_signals)
        best_action_idxA = indmax(utilitiesA)
        best_actionA = uavA.actions[best_action_idxA]
        @test actionA == best_actionA
        @test signalsA == [:no_signal, :no_signal]

        @test actionB == actionW == actionA
        @test signalsB == signalsW == signalsA

        utilitiesB = collect(1.:6.)
        senseB = :same_sense
        actionB, signalsB = select_action_and_signals(uavB, utilitiesB, senseB, no_prev_signals)
        @test actionB == COC_ACTION
        @test signalsB[1] == signalsB[2]
        @test signalsB[1] == :no_signal

        utilitiesW = collect(1.:6.)
        senseW = :same_sense
        actionW, signalsW = select_action_and_signals(uavW, utilitiesW, senseW, no_prev_signals)
        @test actionW == COC_ACTION
        @test signalsW[1] == signalsW[2]
        @test signalsW[1] == :no_signal

        utilitiesA = collect(1.:6.)
        senseA = :same_sense
        actionA, signalsA = select_action_and_signals(uavA, utilitiesA, senseA, no_prev_signals)
        @test actionA == COC_ACTION
        @test signalsA[1] == signalsA[2]
        @test signalsA[1] == :no_signal

        @test actionB == actionW == actionA
        @test signalsB == signalsW == signalsA
    end

    uavB = environmentB.uavs[1]
    uavW = environmentW.uavs[1]
    uavA = environmentA.uavs[1]

    if !uavB.is_policy_joint # non joint policy
        println("i am in the correct test")

        utilitiesB = ones(length(uavB.actions)) * -1
        utilitiesB[6] = 1
        senseB = :same_sense
        actionB, signalsB = select_action_and_signals(uavB, utilitiesB, senseB, [:no_signal, :no_signal])
        @test actionB == COC_ACTION
        @test signalsB[1] == :no_signal
        @test signalsB[1] == signalsB[2]

        utilitiesW = ones(length(uavW.actions)) * -1
        utilitiesW[6] = 1
        senseW = :same_sense
        actionW, signalsW = select_action_and_signals(uavW, utilitiesW, senseW, [:no_signal, :no_signal])
        @test actionW == COC_ACTION
        @test signalsW[1] == :no_signal
        @test signalsW[1] == signalsW[2]

        utilitiesA = ones(length(uavA.actions)) * -1
        utilitiesA[6] = 1
        senseA = :same_sense
        actionA, signalsB = select_action_and_signals(uavA, utilitiesA, senseA, [:no_signal, :no_signal])
        @test actionA == COC_ACTION
        @test signalsA[1] == :no_signal
        @test signalsA[1] == signalsA[2]

        @test actionB == actionW == actionA
        @test signalsB == signalsW == signalsA

        uavB = environmentB.uavs[2]
        utilitiesB[end] = 1
        actionB, signalsB = select_action_and_signals(uavB, utilitiesB, senseB, [:no_signal, :no_signal])
        @test actionB == COC_ACTION
        @test signalsB[1] == :no_signal
        @test signalsB[1] == signalsB[2]

        uavW = environmentW.uavs[2]
        utilitiesW[end] = 1
        actionW, signalsW = select_action_and_signals(uavW, utilitiesW, senseW, [:no_signal, :no_signal])
        @test actionW == COC_ACTION
        @test signalsW[1] == :no_signal
        @test signalsW[1] == signalsW[2]

        uavA = environmentA.uavs[2]
        utilitiesA[end] = 1
        actionA, signalsA = select_action_and_signals(uavA, utilitiesA, senseA, [:no_signal, :no_signal])
        @test actionA == COC_ACTION
        @test signalsA[1] == :no_signal
        @test signalsA[1] == signalsA[2]

        @test actionB == actionW == actionA
        @test signalsB == signalsW == signalsA

        uavB = environmentB.uavs[1]
        utilitiesB = ones(length(uavB.actions)) * -1
        utilitiesB[1] = -100
        actionB, signalsB = select_action_and_signals(uavB, utilitiesB, senseB, [:no_signal, :no_signal])
        @test actionB != -deg2rad(12)

        uavW = environmentW.uavs[1]
        utilitiesW = ones(length(uavW.actions)) * -1
        utilitiesW[1] = -100
        actionW, signalsW = select_action_and_signals(uavW, utilitiesW, senseW, [:no_signal, :no_signal])
        @test actionW != -deg2rad(12)

        uavA = environmentA.uavs[1]
        utilitiesA = ones(length(uavA.actions)) * -1
        utilitiesA[1] = -100
        actionA, signalsA = select_action_and_signals(uavA, utilitiesA, senseA, [:no_signal, :no_signal])
        @test actionA != -deg2rad(12)

        @test actionB == actionW == actionA
        @test signalsB == signalsW == signalsA

        uavB = environmentB.uavs[1]
        utilitiesB = ones(length(uavB.actions))
        stateB = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        stateB.signals = [:dont_turn_left, :dont_turn_left]
        utilitiesB = modify_utilities(uavB, stateB, utilitiesB)
        actionB, signalsB = select_action_and_signals(uavB, utilitiesB, senseB, stateB.signals)
        @test actionB == 0.0
        @test signalsB[1] == :dont_turn_left
        @test signalsB[2] == :dont_turn_left

        uavW = environmentW.uavs[1]
        utilitiesW = ones(length(uavW.actions))
        stateW = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        stateW.signals = [:dont_turn_left, :dont_turn_left]
        utilitiesW = modify_utilities(uavW, stateW, utilitiesW)
        actionW, signalsW = select_action_and_signals(uavW, utilitiesW, senseW, stateW.signals)
        @test actionW == 0.0
        @test signalsW[1] == :dont_turn_left
        @test signalsW[2] == :dont_turn_left

        uavA = environmentA.uavs[1]
        utilitiesA = ones(length(uavA.actions))
        stateA = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        stateA.signals = [:dont_turn_left, :dont_turn_left]
        utilitiesA = modify_utilities(uavA, stateA, utilitiesA)
        actionA, signalsA = select_action_and_signals(uavA, utilitiesA, senseA, stateA.signals)
        @test actionA == 0.0
        @test signalsA[1] == :dont_turn_left
        @test signalsA[2] == :dont_turn_left

        @test actionB == actionW == actionA
        @test signalsB == signalsW == signalsA

        uavB = environmentB.uavs[1]
        utilitiesB = ones(length(uavB.actions))
        stateB = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        stateB.signals = [:dont_turn_right, :dont_turn_right]
        utilitiesB = modify_utilities(uavB, stateB, utilitiesB)
        utilitiesB[3] = -1
        utilitiesB[4] = -1
        utilitiesB[6] = -1
        actionB, signalsB = select_action_and_signals(uavB, utilitiesB, senseB, stateB.signals)
        @test actionB == deg2rad(12)
        @test signalsB[1] == :dont_turn_right
        @test signalsB[2] == :dont_turn_right

        uavW = environmentW.uavs[1]
        utilitiesW = ones(length(uavW.actions))
        stateW = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        stateW.signals = [:dont_turn_right, :dont_turn_right]
        utilitiesW = modify_utilities(uavW, stateW, utilitiesW)
        utilitiesW[3] = -1
        utilitiesW[4] = -1
        utilitiesW[6] = -1
        actionW, signalsW = select_action_and_signals(uavW, utilitiesW, senseW, stateW.signals)
        @test actionW == deg2rad(12)
        @test signalsW[1] == :dont_turn_right
        @test signalsW[2] == :dont_turn_right

        uavA = environmentA.uavs[1]
        utilitiesA = ones(length(uavA.actions))
        stateA = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        stateA.signals = [:dont_turn_right, :dont_turn_right]
        utilitiesA = modify_utilities(uavA, stateA, utilitiesA)
        utilitiesA[3] = -1
        utilitiesA[4] = -1
        utilitiesA[6] = -1
        actionA, signalsA = select_action_and_signals(uavA, utilitiesA, senseA, stateA.signals)
        @test actionA == deg2rad(12)
        @test signalsA[1] == :dont_turn_right
        @test signalsA[2] == :dont_turn_right

        @test actionB == actionW == actionA
        @test signalsB == signalsW == signalsA
    end

    uavB = environmentB.uavs[1]
    uavW = environmentW.uavs[1]
    uavA = environmentA.uavs[1]
    sense = :same_sense

    if !uavB.is_policy_joint # non joint policy
        println("i am in the correct test")

        utilitiesB = ones(length(uavB.actions))
        stateB = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        stateB.signals = [:dont_turn_right, :dont_turn_right]
        utilitiesB = modify_utilities(uavB, stateB, utilitiesB)
        utilitiesB[3] = -1
        utilitiesB[4] = -.5
        utilitiesB[6] = -1
        utilitiesB[5] = -10000
        actionB, signalsB = select_action_and_signals(uavB, utilitiesB, senseB, stateB.signals)
        @test actionB == deg2rad(6)
        @test signalsB[1] == :dont_turn_right
        @test signalsB[2] == :dont_turn_right

        utilitiesW = ones(length(uavW.actions))
        stateW = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        stateW.signals = [:dont_turn_right, :dont_turn_right]
        utilitiesW = modify_utilities(uavW, stateW, utilitiesW)
        utilitiesW[3] = -1
        utilitiesW[4] = -.5
        utilitiesW[6] = -1
        utilitiesW[5] = -10000
        actionW, signalsW = select_action_and_signals(uavW, utilitiesW, senseW, stateW.signals)
        @test actionW == deg2rad(6)
        @test signalsW[1] == :dont_turn_right
        @test signalsW[2] == :dont_turn_right

        utilitiesA = ones(length(uavA.actions))
        stateA = State([UAVState(ones(Float64, 8)...), UAVState(ones(Float64, 8)...)])
        stateA.signals = [:dont_turn_right, :dont_turn_right]
        utilitiesA = modify_utilities(uavA, stateA, utilitiesA)
        utilitiesA[3] = -1
        utilitiesA[4] = -.5
        utilitiesA[6] = -1
        utilitiesA[5] = -10000
        actionA, signalsA = select_action_and_signals(uavA, utilitiesA, senseA, stateA.signals)
        @test actionA == deg2rad(6)
        @test signalsA[1] == :dont_turn_right
        @test signalsA[2] == :dont_turn_right

        @test actionB == actionW == actionA
        @test signalsB == signalsW == signalsA
    end
end
tic()
test_modify_utilities()
test_select_action_and_signals_best()
test_select_action_and_signals_worst()
test_select_action_and_signals_average()
test_select_equal_non_joint()
toc()


text:
module Colors

import Term:
    NAMED_COLORS,
    nospaces,
    COLORS_16b,
    remove_brackets,
    ANSICode,
    int,
    CODES,
    CODES_16BIT_COLORS

export NamedColor, BitColor, RGBColor, get_color

# ----------------------------- types definition ----------------------------- #
"""
    AbstractColor

Abstract color type.
"""
abstract type AbstractColor end

Base.show(io::IO, color::AbstractColor) = print(io, "$(typeof(color))")

struct NamedColor <: AbstractColor
    color::String
end

struct BitColor <: AbstractColor
    color::String
end

struct RGBColor <: AbstractColor
    r::Int
    g::Int
    b::Int
end

function RGBColor(s)
    to_number(x) = '.' ∈ x ? parse(Float64, x) : parse(Int, x)
    r, g, b = to_number.(match(RGB_REGEX, s).captures)
    if r < 1 || g < 1 || b < 1
        r *= 255
        g *= 255
        b *= 255
    end
    return RGBColor(int(r), int(g), int(b))
end

"""
    ANSICode(color; bg::Bool=false)

Create ANSI tags for colors.
"""
function ANSICode(color::NamedColor; bg::Bool = false)
    Δ = bg ? 40 : 30
    v = CODES[color.color]
    return ANSICode("\e[$(Δ + v)m", "\e[$(Δ+9)m")
end

function ANSICode(color::BitColor; bg::Bool = false)
    Δ = bg ? 48 : 38
    v = CODES_16BIT_COLORS[color.color]
    return ANSICode("\e[$Δ;5;$(v)m", "\e[$(Δ+1)m")
end

function ANSICode(color::RGBColor; bg::Bool = false)
    Δ = bg ? 48 : 38
    rgb = "$(color.r);$(color.g);$(color.b)"
    return ANSICode("\e[$Δ;2;$(rgb)m", "\e[$(Δ+1)m")
end

# --------------------------------- is color? -------------------------------- #

RGB_REGEX = r"\(\s*([\d\.]{1,3})\s*,\s*([\d\.]{1,3})\s*,\s*([\d\.]{1,3})\s*\)"
HEX_REGEX = r"#(?:[0-9a-fA-F]{3}){1,2}$"

"""
    is_named_color(string::String)::Bool

Check if a string represents a named color.
"""
is_named_color(string)::Bool = string ∈ NAMED_COLORS || string ∈ COLORS_16b

"""
    is_rgb_color(string::String)::Bool

Check if a string represents a RGB color.
"""
is_rgb_color(string)::Bool = !occursin("on_", string) && occursin(RGB_REGEX, string)

"""
    is_hex_color(string::String)::Bool

Check if a string represents a hex color.
"""
is_hex_color(string)::Bool = !occursin("on_", string) && occursin(HEX_REGEX, string)

"""
    is_color(string::String)::Bool

Check if a string represents color information, of any type.
"""
is_color(string)::Bool =
    is_named_color(string) || is_rgb_color(string) || is_hex_color(string)

"""
    is_background(string::String)::Bool

Check if a string represents background color information, of any type.
"""
function is_background(string)::Bool
    stripped = nospaces(string)
    length(stripped) < 3 && return false
    return stripped[1:3] == "on_" && is_color(stripped[4:end])
end

# --------------------------------- get color -------------------------------- #
"""
    hex2rgb(hex::String)

Converts a string hex color code to a RGB color
"""
function hex2rgb(hex)::RGBColor
    to_int(h) = parse(Int, h; base = 16)
    r, g, b = [to_int(hex[i:(i + 1)]) for i in (2, 4, 6)]
    return RGBColor(r, g, b)
end

"""
    get_color(string::String; bg=false)::AbstractColor

Extract a color type from a string with color information.
"""
function get_color(string; bg = false)::AbstractColor
    bg && (string = nospaces(string)[4:end])

    if is_named_color(string)
        if string ∈ COLORS_16b
            idx = findfirst((c) -> c == string, COLORS_16b)[1]
            return BitColor(COLORS_16b[idx])
        else
            idx = findfirst((c) -> c == string, NAMED_COLORS)[1]
            return NamedColor(NAMED_COLORS[idx])
        end
    elseif is_rgb_color(string)
        return RGBColor(string)
    else
        # convert hex to rgb
        return hex2rgb(string)
    end
end

end

ext: jl


text:
using Base.Test
using RNGTest

include("test/tests.jl")

max_stars_repo_name: josejuanmijares/RNGTest

max_issues_repo_name: josejuanmijares/RNGTest

max_forks_repo_name: josejuanmijares/RNGTest


text:
###############################################################################
#
#   GF.jl : Julia finite fields
#
###############################################################################

export GF

###############################################################################
#
#   Type and parent object methods
#
###############################################################################

parent_type(::Type{gfelem{T}}) where T <: Integer = GFField{T}

elem_type(::Type{GFField{T}}) where T <: Integer = gfelem{T}

doc"""
    base_ring(a::GFField)
> Returns `Union{}` as this field is not dependent on another field.
"""
base_ring(a::GFField) = Union{}

doc"""
    base_ring(a::gfelem)
> Returns `Union{}` as this field is not dependent on another field.
"""
base_ring(a::gfelem) = Union{}

doc"""
    parent(a::gfelem)
> Returns the parent of the given finite field element.
"""
parent(a::gfelem) = a.parent

isexact_type(::Type{gfelem{T}}) where T <: Integer = true

isdomain_type(::Type{gfelem{T}}) where T <: Integer = true

function check_parent(a::gfelem, b::gfelem)
   a.parent != b.parent && error("Operations on distinct finite fields not supported")
end

###############################################################################
#
#   Basic manipulation
#
###############################################################################

function Base.hash(a::gfelem, h::UInt)
   b = 0xe08f2b4ea1cd2a13%UInt
   return xor(xor(hash(a.d), h), b)
end

doc"""
    zero{T <: Integer}(a::GFField{T})
> Return the additive identity, zero, in the given finite field.
"""
function zero(R::GFField{T}) where T <: Integer
   return gfelem{T}(T(0), R)
end

doc"""
    one{T <: Integer}(a::GFField{T})
> Return the additive identity, zero, in the given finite field.
"""
function one(R::GFField{T}) where T <: Integer
      return gfelem{T}(T(1), R)
end

doc"""
    iszero{T <: Integer}(a::gfelem{T})
> Returns true if the given element of the finite field is zero.
"""
iszero(a::gfelem{T}) where T <: Integer = a.d == 0

doc"""
    isone{T <: Integer}(a::gfelem{T})
> Returns true if the given element of the finite field is one.
"""
isone(a::gfelem{T}) where T <: Integer = a.d == 1

doc"""
    isunit(a::gfelem)
> Return `true` if the given finite field element is invertible, i.e. nonzero,
> otherwise return `false`.
"""
isunit(a::gfelem) = a.d != 0

doc"""
    characteristic(R::GFField)
> Return the characteristic of the given finite field.
"""
function characteristic(R::GFField)
   return R.p
end

doc"""
    order(R::GFField)
> Return the order, i.e. the number of element in, the given finite field.
"""
function order(R::GFField)
   return R.p
end

doc"""
    degree(R::GFField)
> Return the degree of the given finite field.
"""
function degree(R::GFField)
   return 1
end

function deepcopy_internal(a::gfelem{T}, dict::ObjectIdDict) where T <: Integer
   R = parent(a)
   return gfelem{T}(deepcopy(a.d), R)
end

###############################################################################
#
#   Canonicalisation
#
###############################################################################

canonical_unit(x::gfelem) = x

###############################################################################
#
#   AbstractString I/O
#
###############################################################################

function show(io::IO, x::gfelem)
   print(io, x.d)
end

function show(io::IO, R::GFField)
   print(io, "Finite field F_", R.p)
end

needs_parentheses(x::gfelem) = false

isnegative(x::gfelem) = false

show_minus_one(::Type{gfelem{T}}) where T <: Integer = true

###############################################################################
#
#   Unary operations
#
###############################################################################

function -(x::gfelem{T}) where T <: Integer
   if x.d == 0
      return deepcopy(x)
   else
      R = parent(x)
      return gfelem{T}(R.p - x.d, R)
   end
end

###############################################################################
#
#   Binary operations
#
###############################################################################

function +(x::gfelem{T}, y::gfelem{T}) where T <: Integer
   check_parent(x, y)
   R = parent(x)
   p = characteristic(R)::T
   d = x.d + y.d - p
   if d < 0
      return gfelem{T}(d + p, R)
   else
      return gfelem{T}(d, R)
   end
end

function -(x::gfelem{T}, y::gfelem{T}) where T <: Integer
   check_parent(x, y)
   R = parent(x)
   p = characteristic(R)::T
   d = x.d - y.d
   if d < 0
      return gfelem{T}(d + p, R)
   else
      return gfelem{T}(d, R)
   end
end

function *(x::gfelem{T}, y::gfelem{T}) where T <: Integer
   check_parent(x, y)
   R = parent(x)
   return R(widen(x.d)*widen(y.d))
end

###############################################################################
#
#   Ad hoc binary operators
#
###############################################################################

function *(x::Integer, y::gfelem{T}) where T <: Integer
   R = parent(y)
   return R(widen(x)*widen(y.d))
end

*(x::gfelem{T}, y::Integer) where T <: Integer = y*x

###############################################################################
#
#   Powering
#
###############################################################################

function ^(x::gfelem{T}, y::Integer) where T <: Integer
   R = parent(x)
   p = R.p::T
   if x.d == 0
      return deepcopy(x)
   end
   if y < 0
      x = inv(x)
      y = -y
   end
   if y >= p - 1
      y1 = y%(p - 1)
      y = convert(T, y1)::T
   else
      y = convert(T, y)::T
   end
   if y == 0
      return one(R)
   elseif y == 1
      return x
   end
   bit = T(1) << (ndigits(y, 2) - 1)
   z = x
   bit >>= 1
   while bit != 0
      z = z*z
      if (bit & y) != 0
         z *= x
      end
      bit >>= 1
   end
   return z
end

###############################################################################
#
#   Comparison
#
###############################################################################

function ==(x::gfelem{T}, y::gfelem{T}) where T <: Integer
   check_parent(x, y)
   return x.d == y.d
end

###############################################################################
#
#   Inversion
#
###############################################################################

function inv(x::gfelem{T}) where T <: Integer
   x == 0 && throw(DivideError())
   R = parent(x)
   p = R.p::T
   g, s, t = gcdx(x.d, p)
   g != 1 && error("Characteristic not prime in ", R)
   return R(s)
end

###############################################################################
#
#   Exact division
#
###############################################################################

function divexact(x::gfelem{T}, y::gfelem{T}) where T <: Integer
   check_parent(x, y)
   return x*inv(y)
end

divides(a::gfelem{T}, b::gfelem{T}) where T <: Integer = true, divexact(a, b)

###############################################################################
#
#   Unsafe functions
#
###############################################################################

function zero!(z::gfelem{T}) where T <: Integer
   R = parent(z)
   d = zero!(z.d)
   return gfelem{T}(d, R)
end

function mul!(z::gfelem{T}, x::gfelem{T}, y::gfelem{T}) where T <: Integer
   return x*y
end

function mul!(z::gfelem{BigInt}, x::gfelem{BigInt}, y::gfelem{BigInt})
   R = parent(x)
   p = R.p::BigInt
   d = mul!(z.d, x.d, y.d)
   if d >= p
      return gfelem{BigInt}(d%p, R)
   else
      return gfelem{BigInt}(d, R)
   end
end

function addeq!(z::gfelem{T}, x::gfelem{T}) where T <: Integer
   R = parent(x)
   p = R.p::T
   d = addeq!(z.d, x.d)
   if d < p
      return gfelem{T}(d, R)
   else
      return gfelem{T}(d - p, R)
   end
end

function add!(z::gfelem{T}, x::gfelem{T}, y::gfelem{T}) where T <: Integer
   R = parent(x)
   p = R.p::T
   d = add!(z.d, x.d, y.d)
   if d < p
      return gfelem{T}(d, R)
   else
      return gfelem{T}(d - p, R)
   end
end

###############################################################################
#
#   Random functions
#
###############################################################################

function rand(R::GFField{T}) where T <: Integer
   p = R.p::T
   d = rand(0:p - 1)
   return gfelem{T}(d, R)
end

###############################################################################
#
#   Promotions
#
###############################################################################

promote_rule(::Type{gfelem{S}}, ::Type{T}) where {S <: Integer, T <: Integer} = gfelem{S}

###############################################################################
#
#   Parent object call overload
#
###############################################################################

function (R::GFField{T})() where T <: Integer
   return gfelem{T}(T(0), R)
end

function (R::GFField{T})(a::Integer) where T <: Integer
   p = R.p::T
   d = convert(T, a%p)::T
   if d < 0
      d += p
   end
   return gfelem{T}(d, R)
end

function (R::GFField{T})(a::gfelem{T}) where T <: Integer
   parent(a) != R && error("Coercion between finite fields not implemented")
   return a
end

###############################################################################
#
#   GF(p) constructor
#
###############################################################################

function GF(p::T) where T <: Integer
   p <= 0 && error("Characteristic is not prime in GF(p)")
   return GFField{T}(p)
end


text:
"""
Example script how to manually create Bayesian cache from already computed results.
Intended to run interactively as user input is needed to rewrite parts of it.
Workflow:
    0. Install `scikit-optimize`
        (from now on the environment that this is installed into has to be the one to activate before each run script)
    1. pick a model and dataset
    2. Define parameter ranges by `create_space` function.
    3. read recursively all results from corresponding dataset folder
    4. run the main forloop to create the cache
        uses built-in `register_run!` with `ignored_hyperparams` as input to indicate which parameters are not optimizable
        patching of individual results is possible using `result_postprocess!`
    5. postprocess cache by filter based on extracted values - define `cache_postprocess!`
    6. sample `init_n` entries from cache with fixed seed
    7. check if the conversions to_skopt and from_skopt work
    8. check if the bayesian optimizer can fit the values provided
    9. after this initial test, run it for all datasets

Bayesian cache structure
- !ordered! dictionary indexed by `ophash` (hash of the named tuple containing only optimizable parameters)
- values - named tuples with following fields: parameters, runs, phashes

ophash:
    parameters: 		named tuple containing optimizable parameters
    runs: 				dictionary indexed by (seed, anomaly_class) containing computed metric for each run and scores
    phashes:            array of hashes of full parameters (this is used to find the original files during evaluation)

example with `ocsvm`
0x1234:
    (
        parameters = (gamma = 1.94199, nu = 0.34065, kernel = "rbf"),
        runs = Dict(
            (1,-1) => [-0.92567],
            (2,-1) => [-0.92832],
            (3,-1) => [-0.94627],
            (4,-1) => [-0.93186],
            (5,-1) => [-0.95032]),
        phashes = [0x123456, 0x123465, 0x123478, 0x123487, 0x123489]
    )
"""

using PyCall
using DrWatson
using Random
using BSON
using GenerativeAD
using ValueHistories
using LinearAlgebra
using OrderedCollections
using Suppressor

# If run outside our main repository - RECOMMENDED
DrWatson.projectdir() = "/home/skvarvit/generativead/GenerativeAD.jl"

### choose model, metric and dataset type
metric = :val_auc
modelname = "fill in modelname"
prefix = "experiments/tabular"
target_prefix = "experiments_bayes/tabular"

### how many samples of fit results to take into account (50 by default)
init_n = 50

folder = datadir("$(prefix)/$(modelname)")
datasets = readdir(folder)

#######################################################################################
################ THIS PART IS TO BE PROVIDED FOR EACH MODEL SEPARATELY ################
"""
    create_space()

Creates a named tuple of python `skotp.space` classes which corresponds
to the hyperparameter ranges in each script.
Some naming conventions apply, see `?GenerativeAD.skopt_parse`.
Copy this from/to the corresponds run script.
"""
function create_space()
    pyReal = pyimport("skopt.space")["Real"]
    pyInt = pyimport("skopt.space")["Integer"]
    pyCat = pyimport("skopt.space")["Categorical"]

    (;
    gamma 		= pyReal(1e-4, 1e2, prior="log-uniform",		name="gamma"),
    kernel 		= pyCat(categories=["poly", "rbf", "sigmoid"],	name="kernel"),
    nu 			= pyReal(0.01, 0.99, 							name="nu")
    )
end

"""
    file_filter!(files)

This filter is applied to an array of all files from a given dataset.
In this first example we throw out all result files, that were trained on data seed>2 and
we also usually want to filter out model files.
(image datasets still have some samples computed on seed>2)
"""
file_filter(files) = filter(x -> occursin("seed=1", x) && !occursin("model", x), files)
file_filter(files) = filter(x -> !occursin("model", x), files) # default - behavior

### Before storing the parameter entry these fields are filtered.
### By default we don't want to optimize `init_seed` or parameters related to anomaly score.
### this has to be copied into `register_run!` in the corresponding run script.
ignored_hyperparams = Set([:init_seed, :percentile])

"""
    cache_postprocess!(cache)

This function is applied to the resulting cache.
In this example we filter out samples with only one computed seed and samples
that have regularization turned off as it is not possible to model continuous parameter from [1f-6,1f-2] + 0.0f0
"""
function cache_postprocess!(cache)
    # filter out those that have only one seed
    filter!(c -> length(c[2][:runs]) > 1, cache)

    # filter out runs with no regularization (throws domain errors from python)
    filter!(c -> c[2][:parameters][:wreg] > 0.0, cache)

    # filter out runs on the old implementation
    filter!(c -> :act_loc in keys(c[2][:parameters]), cache)
end
cache_postprocess!(cache) = filter!(c -> length(c[2][:runs]) > 1, cache) # default filter

################                                                       ################
#######################################################################################

#######################################################################################
################            		MAIN LOOP 				           ################
#######################################################################################
# dataset = datasets[1]             # dry dry run
# for dataset in [datasets[1]]      # dry run
for dataset in datasets           # hot run

    @info "Processing result of $modelname on $dataset."
    dataset_folder = joinpath(folder, dataset)

    # list recursively all files
    files = GenerativeAD.Evaluation.collect_files_th(dataset_folder);
    @info "Collected all $(length(files)) files from $(dataset_folder) folder."

    # file filter (for example when we don't want results from seed>2 for image datasets)
    files_filtered = file_filter(files)
    @info "Applied filter: currently left $(length(files_filtered)) files."

    @info "Running cache builder."
    # load files from each seed and call register_run! in the same way as after training
    cache = OrderedDict{UInt64, Any}()
    for f in files_filtered
        r = BSON.load(f)

        try
            @suppress begin
                GenerativeAD.register_run!(
                    cache, r;
                    metric=:val_auc,
                    flip_sign=true,
                    ignore=ignored_hyperparams) # ignore hyperparameters
            end
        catch e
            @warn("Register run on $(f) has failed due to $(e)")
        end
    end
    @info "Cache creation completed: $(length(cache)) entries"
    if length(cache) > 140
        @warn "There may be too many entries in the cache. Check `ignored_hyperparams`."
    end

    cache_postprocess!(cache)
    @info "Cache postprocess completed: $(length(cache)) entries"

    n = length(cache)
    Random.seed!(7);
    mask = Set(randperm(n)[1:min(init_n, n)])
    cache = OrderedDict(k => v for (i, (k,v)) in enumerate(cache) if i in mask)
    Random.seed!();
    @info "Sampling completed: $(length(cache)) entries"

    ### The following code should not fail and create a valid set of hyperparameters.
    space = create_space()
    x0 = [v[:parameters] for v in values(cache)];
    y0 = [GenerativeAD.objective_value(v[:runs]) for v in values(cache)]

    @info "Testing conversion from named tuples to skopt."
    x0s = [(GenerativeAD.to_skopt(space, x)..., ) for x in x0]
    @info "Testing conversion back from skopt to named tuples."
    x0t = [GenerativeAD.from_skopt(space, x) for x in x0s]

    @info "Testing BayesianHyperOpt fit using the tell method."
    try
        opt = GenerativeAD.BayesianHyperOpt(collect(space))
        GenerativeAD.tell(opt, x0s, y0)

        x1s = GenerativeAD.ask(opt)
        @info "Tested sampling of new parameters using ask: $x1s"

        x1 = GenerativeAD.from_skopt(space, x1s)
        @info "Tested conversion of sampled parameters from skopt to named tuples: $x1"
    catch e
        # helpful for debugging problems with definition of space
        # all hyperparameter dimensions must be contained in their bounds
        @warn "Failed during dry run of optimization due to $e"
        @info "This may be due to some points out of bounds showing debug:"
        @info("", [(x, [py"$(x[i]) in $(p)" for (i, p) in enumerate(space)]) for x in x0s])
        break
    end

    # if all is well save the cache
    # by default this saves it back to the root directory
    DrWatson.projectdir() = pwd()
    target_folder = datadir("$(target_prefix)/$(modelname)/$(dataset)")
    @info "Saving cache to $(target_folder)"
    GenerativeAD.save_bayes_cache(target_folder, cache)
    @info "________________________________________________________________________________"
end


text:
using JSON3
using HTTP

include("../src/parser.jl")

uri = "https://raw.githubusercontent.com/beerjson/beerjson/master/json/beer.json"
@assert isvalid(HTTP.URI(uri)) "Invalid URI"
json_string = String(HTTP.get(uri).body)
json_schema = JSON3.read(json_string)

folder_for_generated_types = "../examples/beerjson_parsed_schemas/"

generate_types(json_schema;
    gen_folder=folder_for_generated_types,
    module_name = "BeerJSON")


text:
using Knet, Random
if !isdefined(Main,:atype)
    global atype = gpu() < 0 ? Array{Float32} : KnetArray{Float32}
end
abstract type Model;end;
struct ResNet <: Model; w; end
function (M::ResNet)(m,imgurl::String,avgimg;stage=3)
    img = imgdata(imgurl, avgimg)
    return M(w,m,atype(img);stage=stage);
end
function ResNet(atype::Type;stage=3)
    w,m,meta = ResNetLib.resnet101init(;trained=true,stage=stage)
    global avgimg = meta["normalization"]["averageImage"]
    global descriptions = meta["classes"]["description"]
    return w,m,meta,avgimg
end
ResNet() = ResNet(nothing);

Par(x)     = Param(atype(x))
init(o...) = Par(xavier(Float32,o...))
bias(o...) = Par(zeros(Float32,o...))
#elu(x)     = relu.(x) + (exp.(min.(0,x)) .- 1.0f0)
#softmax(x,dims) = exp.(logp(x;dims=dims))

struct Linear <: Model; w; b; end
(m::Linear)(x) = m.w * x .+ m.b
Linear(input::Int,output::Int;winit=init,binit=bias) = Linear(winit(output,input),binit(output,1))
Linear() = Linear(nothing,nothing)

struct Embed <: Model; w ; end
(m::Embed)(x)  = m.w[:,x]
Embed(input,embed;winit=rand) = Embed(Par(winit(Float32,embed,input)))
Embed() = Embed(nothing)

struct Conv4D <: Model; w; b; end
(m::Conv4D)(x) = conv4(m.w,x;padding=1,stride=1) .+ m.b
Conv4D(h,w,c,o;winit=init,binit=bias) = Conv4D(winit(h,w,c,o),binit(1,1,o,1))

struct CNN <: Model
    layer1::Conv4D
    layer2::Conv4D
end

function (m::CNN)(x;train=false)
    if train;x=dropout(x,0.18);end
    x1 = elu.(m.layer1(x))
    if train;x1=dropout(x1,0.18);end
    x2 = elu.(m.layer2(x1))
    h,w,c,b = size(x2)
    permutedims(reshape(x2,h*w,c,b),(2,3,1))
end

CNN(h,w,c,d;winit=init,binit=bias) = CNN(Conv4D(h,w,c,d),Conv4D(h,w,d,d))

struct mRNN  <: Model
    rnn
end

function (m::mRNN)(x;batchSizes=[1],train=false)
    B = first(batchSizes)
    if last(batchSizes)!=B
        y,hyout,_,_ = rnnforw(m.rnn,m.rnn.w,x;batchSizes=batchSizes,hy=true,cy=false)
    else
        x           = reshape(x,size(x,1),B,div(size(x,2),B))
        y,hyout,_,_ = rnnforw(m.rnn,m.rnn.w,x;hy=true,cy=false)
    end
    return y,hyout
end
mRNN(input::Int,hidden::Int;o...) = mRNN(RNN(input, hidden;o...))

struct QUnit  <: Model
    embed::Embed
    rnn::mRNN
    linear::Linear
end
function (m::QUnit)(x;batchSizes=[1],train=false)
    xe = m.embed(x)
    if train; xe=dropout(xe,0.15); end;
    y,hyout = m.rnn(xe;batchSizes=batchSizes)
    q            = vcat(hyout[:,:,1],hyout[:,:,2])
    if train; q=dropout(q,0.08); end;
    B = batchSizes[1]
    if ndims(y) == 2
        indices      = bs2ind(batchSizes)
        lngths       = length.(indices)
        Tmax         = maximum(lngths)
        td,B         = size(q)
        d            = div(td,2)
        cw           = Any[];
        for i=1:length(indices)
            y1 = y[:,indices[i]]
            df = Tmax-lngths[i]
            if df > 0
                cpad = zeros(Float32,2d*df) # zeros(Float32,2d,df)
                kpad = atype(cpad)
                ypad = reshape(cat1d(y1,kpad),2d,Tmax) # hcat(y1,kpad)
                push!(cw,ypad)
            else
                push!(cw,y1)
            end
        end
        cws_2d =  reshape(vcat(cw...),2d,B*Tmax)
    else
        d      = div(size(y,1),2)
        Tmax   = size(y,3)
        cws_2d = reshape(y,2d,B*Tmax)
    end
    cws_3d =  reshape(m.linear(cws_2d),(d,B,Tmax))
    return q,cws_3d;
end
QUnit(vocab::Int,embed::Int,hidden::Int;bidir=true) = QUnit(Embed(vocab,embed),
                                                            mRNN(embed,hidden;bidirectional=bidir),
                                                            Linear(2hidden,hidden))

function bs2ind(batchSizes)
    B = batchSizes[1]
    indices = Any[]
    for i=1:B
        ind = i.+cumsum(filter(x->(x>=i),batchSizes)[1:end-1])
        push!(indices,append!(Int[i],ind))
    end
    return indices
end

struct Control  <: Model
    cq::Linear
    att::Linear
end
function (m::Control)(c,q,cws,pad;train=false,tap=nothing)
      d,B,T = size(cws)
      cqi   = reshape(m.cq(vcat(c,q)),(d,B,1))
      cvis  = reshape(cqi .* cws,(d,B*T))
      cvis_2d = reshape(m.att(cvis),(B,T)) #eq c2.1.2
      if pad != nothing
          cvi = reshape(softmax(cvis_2d .- pad,dims=2),(1,B,T)) #eq c2.2
      else
          cvi = reshape(softmax(cvis_2d,dims=2),(1,B,T)) #eq c2.2
      end
      tap!=nothing && get!(tap,"w_attn_$(tap["cnt"])",Array(reshape(cvi,B,T)))
      cnew = reshape(sum(cvi.*cws;dims=3),(d,B))
end
Control(d::Int) = Control(Linear(2d,d),Linear(d,1))

struct Read  <: Model
    me::Linear
    Kbe::Linear
    Kbe2::Linear
    Ime
    att::Linear
end

function (m::Read)(mp,ci,cws,KBhw′,KBhw′′;train=false,tap=nothing)
    d,B,N = size(KBhw′); BN = B*N
    mi_3d = reshape(m.me(mp),(d,B,1))
    ImKB  = reshape(mi_3d .* KBhw′,(d,BN)) # eq r1.2
    ImKB′ = reshape(elu.(m.Ime*ImKB .+ KBhw′′),(d,B,N)) #eq r2
    ci_3d = reshape(ci,(d,B,1))
    IcmKB_pre = elu.(reshape(ci_3d .* ImKB′,(d,BN))) #eq r3.1.1
    if train; IcmKB_pre = dropout(IcmKB_pre,0.15); end;
    IcmKB = reshape(m.att(IcmKB_pre),(B,N)) #eq r3.1.2
    mvi = reshape(softmax(IcmKB,dims=2),(1,B,N)) #eq r3.2
    tap!=nothing && get!(tap,"KB_attn_$(tap["cnt"])",Array(reshape(mvi,B,N)))
    mnew = reshape(sum(mvi.*KBhw′;dims=3),(d,B)) #eq r3.3
end
Read(d::Int) = Read(Linear(d,d),Linear(d,d),Linear(d,d),init(d,d),Linear(d,1))

struct Write  <: Model
    me::Linear
    cproj::Linear
    att::Linear
    mpp
    gating::Linear
end

function (m::Write)(m_new,mi₋1,mj,ci,cj;train=false,selfattn=true,gating=true,tap=nothing)
    d,B        = size(m_new)
    T          = length(mj)
    mi         = m.me(vcat(m_new,mi₋1))
    !selfattn && return mi
    ciproj     = m.cproj(ci)
    ci_3d      = reshape(ciproj,d,B,1)
    cj_3d      = reshape(cat1d(cj...),(d,B,T)) #reshape(hcat(cj...),(d,B,T)) #
    sap        = reshape(ci_3d.*cj_3d,(d,B*T)) #eq w2.1.1
    sa         = reshape(m.att(sap),(B,T)) #eq w2.1.2
    sa′        = reshape(softmax(sa,dims=2),(1,B,T)) #eq w2.1.3
    mj_3d      = reshape(cat1d(mj...),(d,B,T)) #reshape(hcat(mj...),(d,B,T)) #
    mi_sa      = reshape(sum(sa′ .* mj_3d;dims=3),(d,B))
    mi′′       = m.mpp*mi_sa .+ mi #eq w2.3
    !gating && return mi′′
    σci′       = sigm.(m.gating(ci))  #eq w3.1
    mi′′′      = (σci′ .* mi₋1) .+  ((1 .- σci′) .* mi′′) #eq w3.2
end

function Write(d::Int;selfattn=true,gating=true)
    if selfattn
        if gating
            Write(Linear(2d,d),Linear(d,d),Linear(d,1),init(d,d),Linear(d,1))
        else
            Write(Linear(2d,d),Linear(d,d),Linear(d,1),init(d,d),Linear())
        end
    else
        Write(Linear(2d,d),Linear(),Linear(),nothing,Linear())
    end
end

struct MAC <: Model
    control::Control
    read::Read
    write::Write
end
function (m::MAC)(qi,cws,mi,mj,ci,cj,KBhw′,KBhw′′,pad;train=false,selfattn=true,gating=true,tap=nothing)
    cnew = m.control(ci,qi,cws,pad;train=train,tap=tap)
    ri   = m.read(mi,ci,cws,KBhw′,KBhw′′;train=train,tap=tap)
    mnew = m.write(ri,mi,mj,ci,cj;train=train,selfattn=selfattn,gating=gating)
    return cnew,mnew
end
MAC(d::Int;selfattn=false,gating=false) = MAC(Control(d),Read(d),Write(d))

struct Output <: Model
    qe::Linear
    l1::Linear
    l2::Linear
end

function (m::Output)(q,mp;train=false)
  qe = elu.(m.qe(q))
  x  = elu.(m.l1(cat(qe,mp;dims=1)))
  m.l2(x)
end
Output(d::Int) = Output(Linear(2d,d),Linear(2d,d),Linear(d,28))

struct MACNetwork <: Model
    resnet::ResNet
    cnn::CNN
    qunit::QUnit
    qindex::Linear
    mac::MAC
    output::Output
    c0
    m0
end

function (M::MACNetwork)(qs,batchSizes,xS,xB,xP;answers=nothing,p=12,selfattn=false,gating=false,tap=nothing,allsteps=false)
    train         = answers!=nothing
    #STEM Processing
    KBhw          = M.cnn(xS;train=train)
    #Read Unit Precalculations
    d,B,N         = size(KBhw)
    KBhw_2d       = reshape(KBhw,(d,B*N))
    if train; KBhw_2d = dropout(KBhw_2d,0.15); end;
    KBhw′_pre     = M.mac.read.Kbe(KBhw_2d) # look if it is necessary
    KBhw′′        = M.mac.read.Kbe2(KBhw′_pre)
    KBhw′         = reshape(KBhw′_pre,(d,B,N))

    #Question Unit
    q,cws         = M.qunit(qs;batchSizes=batchSizes,train=train)
    qi_c          = M.qindex(q)
    #Memory Initialization
    ci            = M.c0*xB
    mi            = M.m0*xB

    if selfattn
        cj=[ci]; mj=[mi]
    else
        cj=nothing; mj=nothing
    end

    for i=1:p
        qi        = qi_c[(i-1)*d+1:i*d,:]
        if train; ci = dropout(ci,0.15); mi = dropout(mi,0.15); end
        ci,mi = M.mac(qi,cws,mi,mj,ci,cj,KBhw′,KBhw′′,xP;train=train,selfattn=selfattn,gating=gating,tap=tap)
        if selfattn; push!(cj,ci); push!(mj,mi); end
        tap!=nothing && (tap["cnt"]+=1)
    end

    y = M.output(q,mi;train=train)

    if answers==nothing
        predmat = convert(Array{Float32},y)
        tap!=nothing && get!(tap,"y",predmat)
        predictions = mapslices(argmax,predmat,dims=1)[1,:]
        if allsteps
            outputs = []
            for i=1:p-1
                yi = M.output(q,mj[i];train=train)
                yi = convert(Array{Float32},yi)
                push!(outputs,mapslices(argmax,yi,dims=1)[1,:])
            end
            push!(outputs,predictions)
            return outputs
        end
        return predictions
    else
        return nll(y,answers)
    end
end

function MACNetwork(o::Dict)
           MACNetwork(ResNet(),
                      CNN(3,3,1024,o[:d]),
                      QUnit(o[:vocab_size],o[:embed_size],o[:d]),
                      Linear(2*o[:d],o[:p]*o[:d]),
                      MAC(o[:d];selfattn=o[:selfattn],gating=o[:gating]),
                      Output(o[:d]),
                      init(o[:d],1), Par(randn(Float32,o[:d],1)))
end

function setoptim!(m::MACNetwork,o)
    for param in Knet.params(m)
        param.opt = Adam(;lr=o[:lr])
    end
end

function benchmark(M::MACNetwork,feats,o;N=10)
    getter(id) = view(feats,:,:,:,id)
    B=32;L=25
    @time for i=1:N
        ids  = randperm(128)[1:B]
        xB   = atype(ones(Float32,1,B))
        xS   = atype(batcher(map(getter,ids)))
        xQ   = [rand(1:84) for i=1:B*L]
        answers = [rand(1:28) for i=1:B]
        batchSizes = [B for i=1:L]
        xP   = nothing
        y    = @diff M(xQ,batchSizes,xS,xB,xP;answers=answers,p=o[:p],selfattn=o[:selfattn],gating=o[:gating])
    end
end
function benchmark(feats,o;N=30)
    M     = MACNetwork(o);
    benchmark(M,feats,o;N=N)
end

const feats_L = 200704;
const feats_H = 14;
const feats_C = 1024;
function batcher1(feats,args)
    B = length(args)
    totlen = feats_L*B
    result = similar(Array{Float32}, totlen)
    starts = (0:B-1) .*feats_L .+ 1; ends = starts .+ feats_L .- 1;
    for i=1:B
        result[starts[i]:ends[i]] = view(feats,:,:,:,args[i])
    end
    return reshape(result,feats_H,feats_H,feats_C,B)
end

function inplace_batcher(result,data,args)
     B = length(args)
     totlen = feats_L*B
     starts = (0:B-1) .* feats_L .+ 1; ends = starts .+ feats_L .- 1;
     for i=1:B
          result[starts[i]:ends[i]] = view(data,:,:,:,args[i])
     end
     return reshape(result,feats_H,feats_H,feats_C,B)
end

max_stars_repo_path: Courses/World of machine learning/visualize/src/model.jl

max_stars_repo_name: fercarozzi/JuliaAcademyMaterials

max_issues_repo_name: fercarozzi/JuliaAcademyMaterials

max_forks_repo_name: fercarozzi/JuliaAcademyMaterials


text:
"""
$(SIGNATURES)

Construct the matrix of a random 2D Ising lattice ``∑ᵢⱼJᵢⱼZᵢZⱼ``, where ``Jᵢⱼ``s are uniformly distributed between [-1, 1). Generate sparse matrix when `sp` is set to true.
"""
function random_ising(num_qubits::Integer; sp=false)
    J = Vector{Float64}()
    idx = Vector{Vector{Int64}}()
    for i ∈ 1:num_qubits
        for j ∈ i+1:num_qubits
            push!(J, 2*rand()-1)
            push!(idx, [i, j])
        end
    end
    two_local_term(J, idx, num_qubits, sp=sp)
end

"""
$(SIGNATURES)

Construct the matrix of the alternating sectors chain model described in (https://www.nature.com/articles/s41467-018-05239-9#Sec1). Generate sparse matrix when `sp` is set to true.
"""
function alt_sec_chain(w1, w2, n, num_qubits; sp=false)
    J = Vector{Float64}()
    idx = Vector{Vector{Int64}}()
    for i in 1:num_qubits-1
        push!(J, isodd(ceil(Int, i/n)) ? w1 : w2)
        push!(idx, [i, i+1])
    end
    two_local_term(J, idx, num_qubits, sp=sp)
end


text:
module TransientFETools

using Test

using Gridap.Helpers

export ∂t

import GridapODEs.ODETools: ∂t
import GridapODEs.ODETools: time_derivative

export TransientTrialFESpace
export TransientMultiFieldFESpace
export test_transient_trial_fe_space
import Gridap.Fields: evaluate
import Gridap.Fields: evaluate!
import Gridap.MultiField: MultiFieldFESpace
using Gridap.FESpaces: FESpace
using Gridap.FESpaces: SingleFieldFESpace
using Gridap.FESpaces: TrialFESpace
using Gridap.FESpaces: ZeroMeanFESpace
using Gridap.FESpaces: get_free_values
using Gridap.FESpaces: get_dirichlet_values
using Gridap.FESpaces: TrialFESpace!
using Gridap.FESpaces: HomogeneousTrialFESpace
using Gridap.FESpaces: jacobian

import Gridap.Geometry: Triangulation
import Gridap.CellData: Measure
using Gridap.FESpaces: ∫

export TransientFEOperator
export TransientAffineFEOperator
export TransientConstantFEOperator
using Gridap.FESpaces: Assembler
using Gridap.FESpaces: SparseMatrixAssembler
import GridapODEs.ODETools: allocate_cache
import GridapODEs.ODETools: update_cache!
import GridapODEs.ODETools: ODEOperator
import GridapODEs.ODETools: AffineODEOperator
import GridapODEs.ODETools: ConstantODEOperator
import GridapODEs.ODETools: allocate_residual
import GridapODEs.ODETools: allocate_jacobian
import GridapODEs.ODETools: residual!
import GridapODEs.ODETools: jacobian!
import GridapODEs.ODETools: jacobian_t!
import GridapODEs.ODETools: jacobian_and_jacobian_t!
import GridapODEs.ODETools: OperatorType
using GridapODEs.ODETools: Nonlinear
using GridapODEs.ODETools: Affine
using GridapODEs.ODETools: Constant
import Gridap.FESpaces: get_algebraic_operator
import Gridap.FESpaces: assemble_vector!
import Gridap.FESpaces: assemble_matrix_add!
import Gridap.FESpaces: allocate_vector
import Gridap.FESpaces: allocate_matrix
using Gridap.FESpaces: get_cell_shapefuns
using Gridap.FESpaces: get_cell_shapefuns_trial
using Gridap.FESpaces: collect_cell_vector
using Gridap.FESpaces: collect_cell_matrix
using Gridap.FESpaces: return_type
import Gridap.FESpaces: SparseMatrixAssembler
import Gridap.FESpaces: get_trial
import Gridap.FESpaces: get_test
using GridapODEs.ODETools: test_ode_operator
export test_transient_fe_operator

export TransientFESolver
import Gridap.FESpaces: FESolver
import GridapODEs.ODETools: ODESolver
import Gridap: solve
import Gridap.Algebra: solve!
import GridapODEs.ODETools: solve_step!
export test_transient_fe_solver

export TransientFEFunction
import Gridap.FESpaces: FEFunction
import Gridap.FESpaces: EvaluationFunction

export TransientFESolution
import Gridap: solve
import GridapODEs.ODETools: ODESolution
import GridapODEs.ODETools: GenericODESolution
import Base: iterate
export test_transient_fe_solution

include("TransientFESpaces.jl")

include("TransientFEOperators.jl")

include("ODEOperatorInterfaces.jl")

include("TransientFESolvers.jl")

include("TransientFESolutions.jl")

export FETerm
function FETerm(args...)
  Helpers.@unreachable """\n
  Function FETerm has been removed. The API for specifying the weak form has changed significantly.
  See the gridap/Tutorials repo for some examples of how to use the new API.
  This error message will be deleted in future versions.
  """
end

end #module


text:
# module SQLite_Octo.SQL

function table_from(adapter::QueryBuilders.AbstractAdapter, tablename::String)
    so = SQL.prepare(adapter, QueryBuilders.FromClause(tablename))
    fields = tuple(map(Symbol, so.schema.header)...)
    Octo.Adapters.Base.build_table(adapter, tablename, nothing, fields)
end


text:
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
module COM
"""function to initialize a container"""
init_container() = []
"""function to read a message from a pipe, assume it reads from the argument for simplicity"""
read_message(x) = x
# function to check if instruction is authorized (omitted for simplicity)

"""
function doing something based on the message received
compute! function will be defined elsewhere
"""
process!(compute!, container, msg) = compute!(container, msg)
end

module SCI
fun1(x) = x
fun2(x) = x^2

function compute!(container, msg)
    if msg == :fun1
        push!(container, fun1(rand()))
    elseif msg ==:fun2
        push!(container, fun2(rand()))
    end
    return container
end
end

@show compute! = SCI.compute!
@show container = COM.init_container()
@show msg = COM.read_message(:fun1)
@show COM.process!(compute!, container, msg)
@show msg = COM.read_message(:fun2)
@show COM.process!(compute!, container, msg);

# %%

max_stars_repo_path: 0012/On revision of BambOoxX's MWE.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_path: 0012/On revision of BambOoxX's MWE.jl

max_issues_repo_name: genkuroki/public

max_forks_repo_path: 0012/On revision of BambOoxX's MWE.jl

max_forks_repo_name: genkuroki/public


text:
@testset "SpanningTree" begin
    @testset "Symmetric difference" begin
        a = [Edge(1, 2), Edge(1, 3)]
        b = [Edge(1, 3), Edge(1, 4)]
        @test Kombinator._solution_symmetric_difference_size(a, b) == 2
        res_a, res_b = Kombinator._solution_symmetric_difference(a, b)
        @test res_a == [Edge(1, 2)] # Elements that are in a but not in b
        @test res_b == [Edge(1, 4)] # Elements that are in b but not in a
    end

    @testset "Maximum spanning tree" begin
        @testset "Interface" begin
            graph = complete_graph(5)
            rewards = Dict(Edge(1, 2) => 1.0, Edge(1, 3) => 0.5, Edge(2, 3) => 3.0)

            @testset "Global interface" begin
                i = SpanningTreeInstance(graph, rewards)
                @test objective(i) == Maximise()
            end

            @testset "Copying" begin
                i = SpanningTreeInstance(graph, rewards)
                i2 = copy(i)
                @test i.graph == i2.graph
                @test i.rewards == i2.rewards
                @test i.objective == i2.objective
            end

            @testset "Reward accessor" begin
                i = SpanningTreeInstance(graph, rewards)

                @test reward(i, Edge(1, 2)) == reward(i, Edge(2, 1))
            end

            @testset "Approximation" begin
                i = SpanningTreeInstance(graph, rewards)

                @test approximation_ratio(i, PrimAlgorithm()) == 1.0
                @test approximation_term(i, PrimAlgorithm()) == 0.0

                @test approximation_ratio(i, DynamicProgramming()) == 1.0
                @test approximation_term(i, DynamicProgramming()) == 0.0

                @test approximation_ratio(i, DefaultLinearFormulation(Cbc.Optimizer)) == 1.0
                @test approximation_term(i, DefaultLinearFormulation(Cbc.Optimizer)) == 0.0
            end

            @testset "Make solution" begin
                i = SpanningTreeInstance(graph, rewards)
                d = Dict(Edge(1, 2) => 0.8, Edge(1, 3) => 0.2, Edge(2, 3) => 1.0)
                s = make_solution(i, d)

                @test s.instance == i
                @test length(s.variables) == 2
                @test Edge(1, 2) ∈ s.variables
                @test Edge(1, 3) ∉ s.variables
                @test Edge(2, 3) ∈ s.variables

                @test value(s) == 4.0
            end

            @testset "Solution helpers" begin
                i = SpanningTreeInstance(graph, rewards)

                # Infeasible solution.
                s = SpanningTreeSolution(i, Edge{Int}[])
                @test value(s) == -Inf

                # Valid solution.
                d = Dict(Edge(1, 2) => 0.8, Edge(1, 3) => 0.2, Edge(2, 3) => 1.0)
                s = make_solution(i, d)
                @test value(s) == 4

                # Valid solution with reversed edges.
                s = SpanningTreeSolution(i, [Edge(2, 1), Edge(3, 2)])
                @test value(s) == 4
            end

            @testset "Value of a raw tree" begin
                i = SpanningTreeInstance(graph, rewards)

                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_value(
                    i,
                    Edge{Int}[],
                ) == 0
                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_value(
                    i,
                    [Edge(1, 2)],
                ) == 1.0
                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_value(
                    i,
                    [Edge(1, 2), Edge(1, 3)],
                ) == 1.5
            end
        end

        @testset "Basic" begin
            graph = complete_graph(5)
            rewards = Dict(
                Edge(1, 2) => 121.0,
                Edge(1, 3) => 10.0,
                Edge(1, 4) => 10.0,
                Edge(1, 5) => 10.0,
                Edge(2, 3) => 121.0,
                Edge(2, 4) => 10.0,
                Edge(2, 5) => 10.0,
                Edge(3, 4) => 121.0,
                Edge(3, 5) => 10.0,
                Edge(4, 5) => 121.0,
            )

            i = SpanningTreeInstance(graph, rewards)
            p = solve(i, PrimAlgorithm())
            g = solve(i, GreedyAlgorithm()) # Same as Prim.
            d = solve(i, DynamicProgramming())
            l = solve(i, DefaultLinearFormulation(Cbc.Optimizer))

            for s in [p, g, d, l]
                @test s.instance == i
                @test length(s.variables) == 4
                @test Edge(1, 2) in s.variables
                @test Edge(2, 3) in s.variables
                @test Edge(3, 4) in s.variables
                @test Edge(4, 5) in s.variables
            end
        end

        @testset "Conformity" begin
            graph = complete_graph(5)
            rewards = Dict(
                Edge(2, 5) => 0.0,
                Edge(3, 5) => 0.0,
                Edge(4, 5) => 1.0,
                Edge(1, 2) => 0.0,
                Edge(2, 3) => 0.0,
                Edge(1, 4) => 0.0,
                Edge(2, 4) => 0.0,
                Edge(1, 5) => 0.0,
                Edge(1, 3) => 0.0,
                Edge(3, 4) => 0.0,
            )

            i = SpanningTreeInstance(graph, rewards)
            p = solve(i, PrimAlgorithm())
            g = solve(i, GreedyAlgorithm()) # Same as Prim.
            d = solve(i, DynamicProgramming())
            l = solve(i, DefaultLinearFormulation(Cbc.Optimizer))

            for s in [p, g, d, l]
                @test s.instance == i
                @test length(s.variables) == 4 # Five nodes in the graph.
                @test length(unique(s.variables)) == 4 # Only unique edges.
                @test Edge(4, 5) in s.variables || Edge(5, 4) in s.variables # Only edge with nonzero cost.
            end
        end
    end

    @testset "Budgeted maximum spanning tree" begin
        @testset "Interface" begin
            graph = complete_graph(3)
            rewards = Dict(Edge(1, 2) => 1.0, Edge(1, 3) => 0.5, Edge(2, 3) => 3.0)
            weights = Dict(Edge(1, 2) => 0, Edge(1, 3) => 2, Edge(2, 3) => 0)
            i = MinimumBudget(SpanningTreeInstance(graph, rewards), weights, 0)

            @testset "Global interface" begin
                i = MinimumBudget(SpanningTreeInstance(graph, rewards), weights, 0)
                @test objective(i) == Maximise()
            end

            @testset "Copying" begin
                i = MinimumBudget(SpanningTreeInstance(graph, rewards), weights, 0)
                i2 = copy(i)
                @test i.weights == i2.weights
                @test i.min_budget == i2.min_budget
                @test i.compute_all_values == i2.compute_all_values
                @test i.instance.graph == i2.instance.graph
                @test i.instance.rewards == i2.instance.rewards
                @test i.instance.objective == i2.instance.objective
            end

            @testset "Reward accessor" begin
                i = MinimumBudget(SpanningTreeInstance(graph, rewards), weights, 0)

                @test reward(i, Edge(1, 2)) == reward(i, Edge(2, 1))
            end

            @testset "Approximation" begin
                i = MinimumBudget(SpanningTreeInstance(graph, rewards), weights, 0)

                @test approximation_ratio(i, DynamicProgramming()) == 1.0
                @test approximation_term(i, DynamicProgramming()) == 0.0

                @test approximation_ratio(i, DefaultLinearFormulation(Cbc.Optimizer)) == 1.0
                @test approximation_term(i, DefaultLinearFormulation(Cbc.Optimizer)) == 0.0

                @test isnan(approximation_ratio(i, LagrangianAlgorithm()))
                @test isnan(approximation_term(i, LagrangianAlgorithm()))

                @test isnan(approximation_ratio(i, LagrangianRefinementAlgorithm()))
                @test approximation_term(i, LagrangianRefinementAlgorithm()) == 3.0 # Depends on the instance.

                @test approximation_ratio(i, IteratedLagrangianRefinementAlgorithm()) == 0.5
                @test isnan(approximation_term(i, IteratedLagrangianRefinementAlgorithm()))
            end

            @testset "Make solution" begin
                i = MinimumBudget(SpanningTreeInstance(graph, rewards), weights, 0)
                d = Dict(Edge(1, 2) => 0.8, Edge(1, 3) => 0.2, Edge(2, 3) => 1.0)
                s = make_solution(i, d)

                @test s.instance == i
                @test length(s.variables) == 2
                @test Edge(1, 2) ∈ s.variables
                @test Edge(1, 3) ∉ s.variables
                @test Edge(2, 3) ∈ s.variables

                @test value(s) == 4.0
            end

            @testset "Solution helpers" begin
                i = MinimumBudget(SpanningTreeInstance(graph, rewards), weights, 0)

                # Infeasible solution.
                s = SimpleBudgetedSpanningTreeSolution(i, Edge{Int}[])
                @test value(s) == -Inf

                # Valid solution.
                d = Dict(Edge(1, 2) => 0.8, Edge(1, 3) => 0.2, Edge(2, 3) => 1.0)
                s = make_solution(i, d)
                @test value(s) == 4

                # Valid solution with reversed edges.
                s = SimpleBudgetedSpanningTreeSolution(i, [Edge(2, 1), Edge(3, 2)])
                @test value(s) == 4
            end

            @testset "Value of a tree" begin
                i = MinimumBudget(SpanningTreeInstance(graph, rewards), weights, 0)

                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_value(
                    i,
                    Edge{Int}[],
                ) == 0
                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_value(
                    i,
                    [Edge(1, 2)],
                ) == 1.0
                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_value(
                    i,
                    [Edge(1, 2), Edge(1, 3)],
                ) == 1.5

                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_weight(
                    i,
                    Edge{Int}[],
                ) == 0
                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_weight(
                    i,
                    [Edge(1, 2)],
                ) == 0
                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_weight(
                    i,
                    [Edge(1, 2), Edge(1, 3)],
                ) == 2
            end

            @testset "Approximation: Lagrangian refinement" begin
                graph = complete_graph(3)
                rewards = Dict(
                    Edge(1, 2) => 1.0,
                    Edge(1, 3) => 0.5,
                    Edge(2, 3) => 3.0,
                )
                weights =
                    Dict(Edge(1, 2) => 0, Edge(1, 3) => 2, Edge(2, 3) => 0)
                i = MinimumBudget(
                    SpanningTreeInstance(graph, rewards),
                    weights,
                    0,
                )

                @test_throws ErrorException solve(
                    i,
                    LagrangianRefinementAlgorithm(),
                    ζ⁻=1.0,
                )
                @test_throws ErrorException solve(
                    i,
                    LagrangianRefinementAlgorithm(),
                    ζ⁻=2.0,
                )
                @test_throws ErrorException solve(
                    i,
                    LagrangianRefinementAlgorithm(),
                    ζ⁺=1.0,
                )
                @test_throws ErrorException solve(
                    i,
                    LagrangianRefinementAlgorithm(),
                    ζ⁺=0.2,
                )
            end
        end

        @testset "Basic" begin
            graph = complete_graph(3)
            rewards =
                Dict(Edge(1, 2) => 1.0, Edge(1, 3) => 0.5, Edge(2, 3) => 3.0)
            weights = Dict(Edge(1, 2) => 0, Edge(1, 3) => 2, Edge(2, 3) => 0)
            # weights = Dict(Edge(1, 2) => 1, Edge(1, 3) => 2, Edge(2, 3) => 1)

            ε = 0.0001
            budget = 1

            i = MinimumBudget(
                SpanningTreeInstance(graph, rewards),
                weights,
                budget,
            )

            @test dimension(SpanningTreeInstance(graph, rewards)) == 3
            @test dimension(i) == 3

            @testset "Lagrangian relaxation" begin
                lagrangian = solve(i, LagrangianAlgorithm(), ε=ε)
                @test lagrangian.λ ≈ 0.25 atol = ε
                @test lagrangian.value ≈ 3.75 atol = ε
                @test length(lagrangian.variables) == 2

                # Two solutions have this Lagrangian cost: a feasible one and an infeasible one.
                if Edge(1, 3) in lagrangian.variables # Strictly feasible solution.
                    @test Edge(1, 3) in lagrangian.variables
                    @test Edge(2, 3) in lagrangian.variables
                    @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_weight(
                        i,
                        lagrangian.variables,
                    ) > budget
                else # Infeasible solution.
                    @test Edge(1, 2) in lagrangian.variables
                    @test Edge(2, 3) in lagrangian.variables
                    @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_weight(
                        i,
                        lagrangian.variables,
                    ) < budget
                end
            end

            @testset "Lagrangian refinement (additive approximation)" begin
                sol = solve(i, LagrangianRefinementAlgorithm())
                @test sol !== nothing
                @test sol.instance == i
                s = sol.variables
                @test length(s) == 2
                @test Edge(1, 3) in s # Only important edge in this instance: the only one to have a non-zero weight.
                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_weight(
                    i,
                    s,
                ) >= budget
            end

            @testset "Iterated Lagrangian refinement (multiplicative approximation algorithm)" begin
                sol = solve(i, IteratedLagrangianRefinementAlgorithm())
                @test sol !== nothing
                @test sol.instance == i
                s = sol.variables
                @test length(s) == 2
                @test Edge(1, 3) in s # Only important edge in this instance: the only one to have a non-zero weight.
                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_weight(
                    i,
                    s,
                ) >= budget
            end

            @testset "Dynamic programming" begin
                # Mostly like the multiplicate approximation algorithm.
                sol = solve(i, DynamicProgramming())
                @test sol !== nothing
                @test sol.instance == i
                s = sol.variables
                @test length(s) == 2
                @test Edge(1, 3) in s # Only important edge in this instance: the only one to have a non-zero weight.
                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_weight(
                    i,
                    s,
                ) >= budget
            end

            @testset "Linear programming" begin
                # Mostly like the multiplicate approximation algorithm.
                sol = solve(i, DefaultLinearFormulation(Cbc.Optimizer))
                @test sol !== nothing
                @test sol.instance == i
                s = sol.variables
                @test length(s) == 2
                @test Edge(1, 3) in s # Only important edge in this instance: the only one to have a non-zero weight.
                @test Kombinator.SpanningTree._budgeted_spanning_tree_compute_weight(
                    i,
                    s,
                ) >= budget
            end
        end

        @testset "Conformity" begin
            graph = complete_graph(2)
            r = Dict(Edge(1, 2) => 0.0)
            w = Dict(Edge(1, 2) => 5)

            @testset "Feasible budget" begin
                i = MinimumBudget(SpanningTreeInstance(graph, r), w, 0)

                @testset "Lagrangian refinement" begin
                    s = solve(i, LagrangianRefinementAlgorithm())
                    @test s !== nothing
                    @test s.variables == [Edge(1, 2)]
                end

                @testset "Iterated Lagrangian refinement" begin
                    s = solve(i, IteratedLagrangianRefinementAlgorithm())
                    @test s !== nothing
                    @test s.variables == [Edge(1, 2)]
                end

                @testset "Dynamic programming" begin
                    s = solve(i, DynamicProgramming())
                    @test s !== nothing
                    @test s.variables == [Edge(1, 2)]
                end
            end

            @testset "Infeasible budget" begin
                i = MinimumBudget(SpanningTreeInstance(graph, r), w, 20)

                @testset "Lagrangian refinement" begin
                    s = solve(i, LagrangianRefinementAlgorithm())
                    @test s !== nothing
                    @test s.variables == Edge{Int}[]
                end

                @testset "Iterated Lagrangian refinement" begin
                    s = solve(i, IteratedLagrangianRefinementAlgorithm())
                    @test s !== nothing
                    @test s.variables == Edge{Int}[]
                end

                @testset "Dynamic programming" begin
                    s = solve(i, DynamicProgramming())
                    @test s !== nothing
                    @test s.variables == Edge{Int}[]
                end
            end
        end
    end
end


text:
using JuliaInterpreter
using JuliaInterpreter: Frame, FrameData, enter_call_expr,
  do_assignment!, @lookup, pc_expr, isassign, getlhs, caller, traverse
using Debugger: execute_command, DebuggerState, print_locals, active_frame, print_frame

struct InterpreterError <: Exception
  err
  state

  InterpreterError(ierr::InterpreterError, s) = new(ierr.err, s)
  InterpreterError(err, s) = new(err, s)
end

frame(state) = active_frame(state)
pc(frame) = frame.pc
isdone(state) = state.frame == nothing

function expr(state)
  fr = frame(state)
  expr = pc_expr(fr, pc(fr))
  isassign(fr) || return expr
  Expr(:(=), getlhs(pc(fr)), expr)
end

step!(state) =
  execute_command(state, Val{:se}(), "se")

stepin!(state) =
  execute_command(state, Val{:s}(), "s")

lookup(frame, var) = @lookup(frame, var)
lookup(frame, x::QuoteNode) = x.value

function callargs(state)
  ex = expr(state)
  isexpr(ex, :(=)) && (ex = ex.args[2])
  isexpr(ex, :call) || return
  args = lookup.(Ref(frame(state)), ex.args)
  args[1] == Core._apply && (args = [args[2], Iterators.flatten(args[3:end])...])
  return args
end

primitive_(ctx, state, a...) = primitive(ctx, a...)

function provide_result!(state, x)
  do_assignment!(frame(state), expr(state).args[1], x)
end

function inc_pc!(state)
  fr = frame(state)
  fr.pc += 1
end

unwrap(x) = x
unwrap(x::QuoteNode) = x.value
unwrap(x::Expr) = isexpr(x,:copyast) ? unwrap(x.args[1]) : x

function runall(ctx, state)
  while !isdone(state)
    try
      if (ex = callargs(state)) ≠ nothing
        hook(ctx, ex...)
        if isprimitive(ctx, ex...)
          result = primitive_(ctx, state, ex...)
          isexpr(expr(state), :(=)) && provide_result!(state, result)
          inc_pc!(state)
        else
          stepin!(state)
        end
      else
        step!(state)
      end
    catch err
      throw(InterpreterError(err, state))
    end
  end
  return unwrap(state.overall_result)
end

function overdub(ctx, f, args...)
  frame = enter_call_expr(:($f($(args...))))
  frame == nothing && return f(args...)
  runall(ctx, DebuggerState(frame=frame))
end

macro overdub(ctx, ex)
  :(overdub($(esc(ctx)), () -> $(esc(ex))))
end

function Base.showerror(io::IOContext, ierr::InterpreterError)
  println("InterpreterError: ")
  showerror(io, ierr.err)
  println(io, "\nStacktrace of evaluated expression:")
  l = 0
  traverse(frame(ierr.state)) do fr
    print_frame(io, l, fr)
    l += 1
    return caller(fr)
  end
end


text:
prepare_call!(vars, expr) = expr
prepare_call!(vars, s::Symbol) = esc(s)

function prepare_call!(vars, e::Expr)
    e.head == :$ || return Expr(e.head, map(x->prepare_call!(vars, x), e.args)...)

    var = gensym()
    push!(vars, :($var = $(prepare_call!(vars, e.args[1]))))
    var
end

prepare_call(e) = let v=[]
    e2 = prepare_call!(v, e)
    v, e2
end

function count_ops(funcall)
    v, e = prepare_call(funcall)
    quote
        let
            ctx = CounterCtx(metadata=Counter())
            $(v...)
            Cassette.overdub(ctx, ()->begin
                             $e
                             end)
            ctx.metadata
        end
    end
end

macro count_ops(funcall)
    count_ops(funcall)
end

# Helper accessor (that can be overriden to fake benchmarking times in tests)
times(t::BenchmarkTools.Trial) = t.times

macro gflops(funcall)
    benchmark = quote
        $BenchmarkTools.@benchmark $funcall
    end

    quote
        let
            b = $(esc(benchmark))
            ns = minimum(times(b))

            cnt = flop($(count_ops(funcall)))
            gflops = cnt / ns
            peakfraction = 1e9 * gflops / peakflops()
            memory = $BenchmarkTools.prettymemory(b.memory)
            @printf("  %.2f GFlops,  %.2f%% peak  (%.2e flop, %.2e s, %d alloc: %s)\n",
                    gflops, peakfraction*100, cnt, ns*1e-9,
                    b.allocs, memory)
            gflops
        end
    end
end


text:
using Test
using RobotDynamics
using StaticArrays
using ForwardDiff
using LinearAlgebra
using FiniteDiff
using Random
const RD = RobotDynamics

include("cartpole_model.jl")
include("random_linear.jl")
include("quadrotor.jl")

##
@testset "Basic Dynamics" begin
    @testset "Cartpole" begin
        include("cartpole_test.jl")
    end
    @testset "Integration" begin
        include("jacobian_test.jl")
        include("integration_tests.jl")
    end
    @testset "KnotPoints" begin
        include("knotpoints.jl")
        include("trajectories.jl")
    end
end

@testset "Lie State" begin
    include("liestate.jl")
    include("liemodel.jl")
end

@testset "Linear Systems" begin
    include("linmodel_type.jl")
    include("test_random_linear.jl")
    include("linear_quad.jl")
    # include("linear_tests.jl")
    include("linearization.jl")
end

@testset "Rigid Bodies" begin
    @testset "RBState" begin
        include("rbstate.jl")
    end
    @testset "Dynamics" begin
        include("rigidbody.jl")
    end
end
@testset "Jacobians" begin
    include("rigid_body_jacobians.jl")
end

# only test plotting on a desktop since it takes a while to compile Plots on CI...
if !haskey(ENV, "CI")
    @testset "Plotting" begin
        include("plotting.jl")
    end
end


text:
using Plots; pyplot()
using Reproduce
using FileIO
using Statistics
using ProgressMeter
using JLD2
using LaTeXStrings
using StatsPlots

include("./plot_utils.jl")
GPU = GeneralPlotUtils
LU = LabelUtils
include("./plotting.jl")

function load_data(data_home)
    experiment_folders = [data_home]
    ic = ItemCollection(joinpath(experiment_folders[1], "data"))
    return ic
end

function load_best(ic)
    algo_divisor_keys = ["behaviour_learner", "demon_learner", "demon_opt", "demon_update"]
    sweep_params = ["alpha_init", "eta", "alpha_init","batch_size"]
    # data_key = :oned_tmaze_dmu_error
    data_key = :oned_tmaze_old_error

    algo_specs_full = GPU.split_algo(ic, algo_divisor_keys)

    all_algos_ics = [search(ic,algo_spec) for algo_spec in algo_specs_full]
    @show length.(all_algos_ics)

    valid_algos_ind = findall(a -> length(a) != 0, all_algos_ics)

    algo_specs = algo_specs_full[valid_algos_ind]
    algo_ics = all_algos_ics[valid_algos_ind]
    best_per_algo_ics = []
    for (i,algo_ic) in enumerate(algo_ics)
        push!(best_per_algo_ics, GPU.get_best_final_perf(algo_ic,sweep_params, data_key, 0.1))
    end
    @show length.(best_per_algo_ics)
    return best_per_algo_ics
end

SR_ER_ic_all = load_data("../data/ER/OneDTMaze_RR_ER/")
SR_ER_ic_all = search(SR_ER_ic_all, Dict("batch_size" => 8))
SR_ER_ic = load_best(SR_ER_ic_all)

Q_ER_ic_all = load_data("../data/ER/OneDTMaze_RR_ER_2/")
Q_ER_ic_all = search(Q_ER_ic_all, Dict("batch_size" => 8))
Q_ER_ic = load_best(Q_ER_ic_all)

SR_ic_all = load_data("../data/OneDTMaze_RR/")
SR_ic_all = search(SR_ic_all, Dict("demon_learner" => "SR"))
SR_ic = load_best(SR_ic_all)

Q_ic_all = load_data("../data/OneDTMaze_RR/")
Q_ic_all = search(Q_ic_all, Dict("demon_learner" => "Q"))
Q_ic = load_best(Q_ic_all)

LSTD_ic_all = load_data("../data/OneDTMaze_RR/")
LSTD_ic_all = search(LSTD_ic_all, Dict("demon_learner" => "LSTD"))
LSTD_ic = load_best(LSTD_ic_all)

@show length(SR_ER_ic)
@show length(Q_ER_ic)
@show length(SR_ic)
@show length(Q_ic)
@show length(LSTD_ic)

all_ic = [
    SR_ER_ic[1],
    Q_ER_ic[1],
    SR_ic[1],
    Q_ic[1],
    LSTD_ic[1],
]

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
struct Embed; w; end

function Embed(vocabsize::Int, embedsize::Int)
    Embed(param(embedsize,vocabsize))
end

function (l::Embed)(x)
    l.w[:,x]
end

max_stars_repo_name: denizyuret/research

max_issues_repo_name: denizyuret/research

max_forks_repo_name: denizyuret/research


text:
"""
    lz4_compress(input::Union{Vector{UInt8},Base.CodeUnits{UInt8}}, acceleration::Integer=0)

Compresses `input` using LZ4_compress_fast. Returns a Vector{UInt8} of the compressed data.
"""
function lz4_compress(
    input::Union{Vector{UInt8},Base.CodeUnits{UInt8}},
    acceleration::Integer=0
)
    bound = LZ4_compressBound(length(input))
    bound == 0 && throw(LZ4Exception(
        "lz4_compress",
        "Input size $(length(input)) greater than LZ4_MAX_INPUT_VALUE. Cannot encode as a single block."
    ))

    out_buffer = Vector{UInt8}(undef, bound)
    out_size = LZ4_compress_fast(pointer(input), pointer(out_buffer), length(input), bound, acceleration)
    resize!(out_buffer, out_size)
end

"""
    lz4_hc_compress(input::Union{Vector{UInt8},Base.CodeUnits{UInt8}}, acceleration::Integer=$LZ4HC_CLEVEL_DEFAULT)

Compresses `input` using LZ4_compress_HC. Returns a Vector{UInt8} of the compressed data.
"""
function lz4_hc_compress(
    input::Union{Vector{UInt8},Base.CodeUnits{UInt8}},
    compressionlevel::Integer=LZ4HC_CLEVEL_DEFAULT
)
    bound = LZ4_compressBound(length(input))
    bound == 0 && throw(LZ4Exception(
        "lz4_hc_compress",
        "Input size $(length(input)) greater than LZ4_MAX_INPUT_VALUE. Cannot encode as a single block."
    ))

    out_buffer = Vector{UInt8}(undef, bound)
    out_size = LZ4_compress_HC(pointer(input), pointer(out_buffer), length(input), bound, compressionlevel)
    resize!(out_buffer, out_size)
end

"""
    lz4_decompress(input::Union{Vector{UInt8},Base.CodeUnits{UInt8}}, expected_size::Integer=input.size * 2)

Decompresses `input` using LZ4_decompress_safe.
`expected_size` must be equal to or larger than the expected decompressed size of the input or decompression will fail.
Returns a Vector{UInt8} of the decompressed data.
"""
function lz4_decompress(
    input::Union{Vector{UInt8},Base.CodeUnits{UInt8}},
    expected_size::Integer=length(input) * 2
)
    out_buffer = Vector{UInt8}(undef, expected_size)
    out_size = LZ4_decompress_safe(pointer(input), pointer(out_buffer), length(input), expected_size)
    resize!(out_buffer, out_size)
end


text:
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.4
#   kernelspec:
#     display_name: Julia 1.3.1
#     language: julia
#     name: julia-1.3
# ---

# # Example 5
#
# A `ClimGrid` struct from the `ClimateTools.jl` package, including metadata read from file, is written to a netcdf file via `NCTiles.jl`.
#
# _Note: `ClimateTools.jl` requires that requires specific dimension names as follows_
#
# - latitude: lat, latitude, rlat, y, yc
# - longitude: lon, longitude, rlon, x, xc
# - time: time

# +
using ClimateTools,NCDatasets,NCTiles

# File Paths
inputs = "input/"
NCTiles.get_testcases_if_needed(inputs)

outputs = "output/"
if ~ispath(outputs); mkpath(outputs); end

savedir = joinpath(outputs,"ex5")
if ~ispath(savedir); mkpath(savedir); end

#original climate model output file
file_in="tas_day_MIROC5_piControl_r1i1p1_20000101-20091231.nc"
field_name = "tas"

if ~isfile(joinpath(inputs,file_in))
    run(`wget http://esgf-data1.diasjp.net/thredds/fileServer/esg_dataroot/cmip5/output1/MIROC/MIROC5/piControl/day/atmos/day/r1i1p1/v20161012/tas/tas_day_MIROC5_piControl_r1i1p1_20000101-20091231.nc`)
    run(`mv tas_day_MIROC5_piControl_r1i1p1_20000101-20091231.nc ../inputs/`)
end
# -

"""
        climgridtoncvar(C::ClimGrid,N::String)

Creates an NCvar struct from a ClimGrid object. NCvar struct can then be written
to a NetCDF file using write().

Ex: C = load(fname,fldname)
writefld = climgridtoncvar(C,fldname)
write(writefld,"myfile.nc")
"""
function climgridtoncvar(C::ClimGrid,N::String)
        x, y, timevec = ClimateTools.getdims(C) # may need to check number of dims first
        timevec = NCDatasets.timeencode(timevec, C.timeattrib["units"], get(C.timeattrib,"calendar","standard"))

        dims = [NCvar(C.dimension_dict["lon"],C.lonunits,size(C.data)[1],x,Dict("long_name" => "longitude"),NCDatasets),
                NCvar(C.dimension_dict["lat"],C.latunits,size(C.data)[2],y,Dict("long_name" => "latitude"),NCDatasets),
                NCvar("time",C.timeattrib["units"],Inf,timevec,Dict(("long_name" => "tim","standard_name" => "time")),NCDatasets)
                ]

        return NCvar(N,C.dataunits,dims,C.data.data,C.varattribs,NCDatasets)
end

# Whole Data Set
C = load(joinpath(inputs,file_in),field_name)
writefld = climgridtoncvar(C,field_name)
write(writefld,joinpath(savedir,"ex5_whole.nc"))

# Extracted Sub-Set
poly_reg = [[NaN -65 -80 -80 -65 -65];[NaN 42 42 52 52 42]]
E = load(joinpath(inputs,file_in),field_name, poly=poly_reg)
writefld = climgridtoncvar(E,field_name)
write(writefld,joinpath(savedir,"ex5_extract.nc"),
    globalattribs=E.globalattribs)


text:
using InteractBulma, Observables, Colors, DataStructures
import Widgets: scope
@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

@testset "options" begin
    a = dropdown(["a", "b", "c"])
    @test observe(a)[] == "a"
    a = dropdown(OrderedDict("a" => 1, "b" => 2, "c" => 3))
    @test observe(a)[] == 1
    a = dropdown(OrderedDict("a" => 1, "b" => 2, "c" => 3), value = 3)
    @test observe(a)[] == 3
end

@testset "ijulia" begin
    @test !InteractBase.isijulia()
end


text:
using PowerSystems
using NLsolve
const PSY = PowerSystems

############### Data Network ########################
include(joinpath(dirname(@__FILE__), "dynamic_test_data.jl"))
include(joinpath(dirname(@__FILE__), "data_utils.jl"))
############### Data Network ########################
threebus_file_dir = joinpath(dirname(@__FILE__), "ThreeBusNetwork.raw")
threebus_sys = System(threebus_file_dir, runchecks = false)
add_source_to_ref(threebus_sys)
#Reduce generator output
for g in get_components(Generator, threebus_sys)
    g.active_power = 0.75
end
res = solve_powerflow!(threebus_sys)

function dyn_gen_five_mass_shaft_order(generator)
    return PSY.DynamicGenerator(
        name = get_name(generator),
        ω_ref = 1.0, # ω_ref,
        machine = machine_oneDoneQ(), #machine
        shaft = shaft_fivemass(), #shaft
        avr = avr_type1(), #avr
        prime_mover = tg_none(),
        pss = pss_none(),
    ) #pss
end

function dyn_gen_first_order(generator)
    return PSY.DynamicGenerator(
        name = get_name(generator),
        ω_ref = 1.0, # ω_ref,
        machine = machine_oneDoneQ(), #machine
        shaft = shaft_damping(), #shaft
        avr = avr_type1(), #avr
        prime_mover = tg_none(), #tg
        pss = pss_none(),
    ) #pss
end

for g in get_components(Generator, threebus_sys)
    if get_number(get_bus(g)) == 103
        case_gen = dyn_gen_five_mass_shaft_order(g)
        add_component!(threebus_sys, case_gen, g)
    elseif get_number(get_bus(g)) == 102
        case_inv = dyn_gen_first_order(g)
        add_component!(threebus_sys, case_inv, g)
    end
end
#Compute Y_bus after fault
fault_branches = deepcopy(collect(get_components(Branch, threebus_sys))[2:end])
Ybus_fault = PSY.Ybus(fault_branches, get_components(Bus, threebus_sys))[:, :]


text:
function vandermonde_1d(N, r)
    V = zeros(length(r), N+1)
    for j in 1:N+1
        @. V[:, j] = jacobi(r, 0, 0, j-1)
    end
    return V
end

" Initialize the gradient of the modal basis (i) at (r) at order N. "
function ∂vandermonde_1d(N, r)
    DVr = zeros(length(r), N+1)
    for i in 0:N
        @. DVr[:, i+1] = ∂jacobi(r, 0, 0, i)
    end
    return DVr
end

" Initialize the (r) differentiation matrices on the interval, evaluated at (r) at order N. "
function Dmatrix1D(N, r, V)
    Vr = ∂vandermonde_1d(N, r)
    Dr = Vr / V
    return Dr
end