text:
export Experiment

using .ArcadeLearningEnvironment
using Dates
using ReinforcementLearningCore
using .ReinforcementLearningEnvironments
using Flux
using BSON
using TensorBoardLogger
using Logging
using Statistics

clip_reward(obs) = RewardOverriddenObs(obs, clamp(get_reward(obs), -1, 1))

function atari_env_factory(name, state_size, n_frames, max_episode_steps=100_000; seed = nothing)
    WrappedEnv(
        env = AtariEnv(;
            name = string(name),
            grayscale_obs = true,
            noop_max = 30,
            frame_skip = 4,
            terminal_on_life_loss = false,
            repeat_action_probability = 0.25,
            max_num_frames_per_episode = n_frames * max_episode_steps,
            color_averaging = false,
            full_action_space = false,
            seed = seed,
        ),
        preprocessor = ComposedPreprocessor(
            ResizeImage(state_size...),  # this implementation is different from cv2.resize https://github.com/google/dopamine/blob/e7d780d7c80954b7c396d984325002d60557f7d1/dopamine/discrete_domains/atari_lib.py#L629
            StackFrames(state_size..., n_frames),
            clip_reward,
        ),
    )
end

function RLCore.Experiment(
    ::Val{:Dopamine},
    ::Val{:DQN},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "dopamine_DQN_atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = atari_env_factory(name, STATE_SIZE, N_FRAMES;)#= seed=nothing =#
    N_ACTIONS = length(get_action_space(env))
    init = seed_glorot_uniform()#= seed=341 =#

    create_model() =
        Chain(
            x -> x ./ 255,
            CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
            CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
            CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
            x -> reshape(x, :, size(x)[end]),
            Dense(11 * 11 * 64, 512, relu),
            Dense(512, N_ACTIONS),
        ) |> gpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = DQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = create_model(),
                    optimizer = RMSProp(0.00025, 0.95),
                ),  # unlike TF/PyTorch RMSProp doesn't support center
                target_approximator = NeuralNetworkApproximator(model = create_model()),
                update_freq = 4,
                γ = 0.99f0,
                update_horizon = 1,
                batch_size = 32,
                stack_size = N_FRAMES,
                min_replay_history = 20_000,
                loss_func = huber_loss,
                target_update_freq = 8_000,
            ),
            explorer = EpsilonGreedyExplorer(
                ϵ_init = 1.0,
                ϵ_stable = 0.01,
                decay_steps = 250_000,
                kind = :linear,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1_000_000,
            state_type = Float32,
            state_size = STATE_SIZE,
        ),
    )

    evaluation_result = []
    EVALUATION_FREQ = 250_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env, obs
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end] log_step_increment =
                    0
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            old_explorer = agent.policy.explorer
            agent.policy.explorer = EpsilonGreedyExplorer(0.001)  # set evaluation epsilon
            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
            s = @elapsed run(
                agent,
                atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;),#= seed=nothing =#
                StopAfterStep(125_000; is_show_progress = false),
                h,
            )
            res = (
                avg_length = mean(h[2].steps[1:end-1]),
                avg_score = mean(h[1].rewards[1:end-1]),
            )
            push!(evaluation_result, res)
            agent.policy.explorer = old_explorer
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_length = res.avg_length avg_score =
                res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_length = res.avg_length avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent; is_save_trajectory = false)  # saving trajectory will take about 27G disk space each time
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_reward_per_episode time_per_step evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    N_TRAINING_STEPS = 50_000_000
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    description = """
    This experiment uses almost the same config in [dopamine](https://github.com/google/dopamine/blob/master/dopamine/agents/dqn/configs/dqn.gin). But do notice that there are some minor differences:

    - The RMSProp in Flux do not support center option (also the epsilon is not the same).
    - The image resize method used here is provided by ImageTransformers, which is not the same with the one in cv2.
    - `max_steps_per_episode` is not set, this might affect the evaluation result slightly.

    The testing environment is $name.
    Agent and statistic info will be saved to: $(joinpath(save_dir, string(N_TRAINING_STEPS)))
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`

    To load the agent and statistic info:
    ```
    agent = RLCore.load("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", Agent)
    BSON.@load joinpath("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", "stats.bson") total_reward_per_episode time_per_step evaluation_result
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:Dopamine},
    ::Val{:Rainbow},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "Dopamine_Rainbow_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = atari_env_factory(name, STATE_SIZE, N_FRAMES;)#= seed=(135, 246) =#
    N_ACTIONS = length(get_action_space(env))
    N_ATOMS = 51
    init = seed_glorot_uniform()#= seed=341 =#

    create_model() =
        Chain(
            x -> x ./ 255,
            CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
            CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
            CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
            x -> reshape(x, :, size(x)[end]),
            Dense(11 * 11 * 64, 512, relu),
            Dense(512, N_ATOMS * N_ACTIONS),
        ) |> gpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = RainbowLearner(
                approximator = NeuralNetworkApproximator(
                    model = create_model(),
                    optimizer = ADAM(0.0000625),
                ),  # epsilon is not set here
                target_approximator = NeuralNetworkApproximator(model = create_model()),
                n_actions = N_ACTIONS,
                n_atoms = N_ATOMS,
                Vₘₐₓ = 10.0f0,
                Vₘᵢₙ = -10.0f0,
                update_freq = 4,
                γ = 0.99f0,
                update_horizon = 3,
                batch_size = 32,
                stack_size = N_FRAMES,
                min_replay_history = 20_000,
                loss_func = logitcrossentropy_unreduced,
                target_update_freq = 8_000,
                # seed=89,
            ),
            explorer = EpsilonGreedyExplorer(
                ϵ_init = 1.0,
                ϵ_stable = 0.01,
                decay_steps = 250_000,
                kind = :linear,
                # seed=97,
            ),
        ),
        trajectory = CircularCompactPSARTSATrajectory(
            capacity = 1_000_000,
            state_type = Float32,
            state_size = STATE_SIZE,
        ),
    )

    evaluation_result = []
    EVALUATION_FREQ = 250_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    steps_per_episode = StepsPerEpisode()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        steps_per_episode,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env, obs
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end] episode_length =
                    steps_per_episode.steps[end] log_step_increment = 0
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            old_explorer = agent.policy.explorer
            agent.policy.explorer = EpsilonGreedyExplorer(0.001)  # set evaluation epsilon
            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
            s = @elapsed run(
                agent,
                atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;),#= seed=nothing =#
                StopAfterStep(125_000; is_show_progress = false),
                h,
            )
            res = (
                avg_length = mean(h[2].steps[1:end-1]),
                avg_score = mean(h[1].rewards[1:end-1]),
            )
            push!(evaluation_result, res)
            agent.policy.explorer = old_explorer
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_length = res.avg_length avg_score =
                res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_length = res.avg_length avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent; is_save_trajectory = false)  # saving trajectory will take about 27G disk space each time
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_reward_per_episode time_per_step evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    N_TRAINING_STEPS = 50_000_000
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    description = """
    This experiment uses alomost the same config in [dopamine](https://github.com/google/dopamine/blob/master/dopamine/agents/rainbow/configs/rainbow.gin). But do notice that there are some minor differences:

    - The epsilon in ADAM optimizer is not changed
    - The image resize method used here is provided by ImageTransformers, which is not the same with the one in cv2.

    The testing environment is $name.
    Agent and statistic info will be saved to: `$(joinpath(save_dir, string(N_TRAINING_STEPS)))`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`

    To load the agent and statistic info:
    ```
    agent = RLCore.load("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", Agent)
    BSON.@load joinpath("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", "stats.bson") total_reward_per_episode time_per_step evaluation_result
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:Dopamine},
    ::Val{:IQN},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "Dopamine_IQN_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    MAX_STEPS_PER_EPISODE = 27_000

    env = atari_env_factory(name, STATE_SIZE, N_FRAMES;)#= seed=(135, 246) =#
    N_ACTIONS = length(get_action_space(env))
    Nₑₘ = 64

    init = seed_glorot_uniform()#= seed=341 =#

    create_model() =
        ImplicitQuantileNet(
            ψ = Chain(
                x -> x ./ 255,
                CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
                CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
                CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
                x -> reshape(x, :, size(x)[end]),
            ),
            ϕ = Dense(Nₑₘ, 11 * 11 * 64, relu),
            header = Chain(Dense(11 * 11 * 64, 512, relu), Dense(512, N_ACTIONS)),
        ) |> gpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = IQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = create_model(),
                    optimizer = ADAM(0.00005),  # epsilon is not set here
                ),
                target_approximator = NeuralNetworkApproximator(model = create_model()),
                κ = 1.0f0,
                N = 64,
                N′ = 64,
                Nₑₘ = Nₑₘ,
                K = 32,
                γ = 0.99f0,
                stack_size = 4,
                batch_size = 32,
                update_horizon = 3,
                min_replay_history = 20_000,
                update_freq = 4,
                target_update_freq = 8_000,
                default_priority = 1.0f2,
                # seed=105,
                # device_seed=237,
            ),
            explorer = EpsilonGreedyExplorer(
                ϵ_init = 1.0,
                ϵ_stable = 0.01,
                decay_steps = 250_000,
                kind = :linear,
                # seed=99,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1_000_000,
            state_type = Float32,
            state_size = STATE_SIZE,
        ),
    )

    evaluation_result = []
    EVALUATION_FREQ = 250_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    steps_per_episode = StepsPerEpisode()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        steps_per_episode,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env, obs
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end] episode_length =
                    steps_per_episode.steps[end] log_step_increment = 0
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            old_explorer = agent.policy.explorer
            agent.policy.explorer = EpsilonGreedyExplorer(0.001)  # set evaluation epsilon
            h = ComposedHook(TotalRewardPerEpisode(), StepsPerEpisode())
            s = @elapsed run(
                agent,
                atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;),#= seed=nothing =#
                StopAfterStep(125_000; is_show_progress = false),
                h,
            )
            res = (
                avg_length = mean(h[2].steps[1:end-1]),
                avg_score = mean(h[1].rewards[1:end-1]),
            )
            push!(evaluation_result, res)
            agent.policy.explorer = old_explorer
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_length = res.avg_length avg_score =
                res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_length = res.avg_length avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent; is_save_trajectory = false)  # saving trajectory will take about 27G disk space each time
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_reward_per_episode time_per_step evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    N_TRAINING_STEPS = 50_000_000
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    description = """
    This experiment uses alomost the same config in [dopamine](https://github.com/google/dopamine/blob/master/dopamine/agents/implicit_quantile/configs/implicit_quantile.gin). But do notice that there are some minor differences:

    - The epsilon in ADAM optimizer is not changed
    - The image resize method used here is provided by ImageTransformers, which is not the same with the one in cv2.

    The testing environment is $name.
    Agent and statistic info will be saved to: `$(joinpath(save_dir, string(N_TRAINING_STEPS)))`
    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`

    To load the agent and statistic info:
    ```
    agent = RLCore.load("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", Agent)
    BSON.@load joinpath("$(joinpath(save_dir, string(N_TRAINING_STEPS)))", "stats.bson") total_reward_per_episode time_per_step evaluation_result
    ```
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:A2C},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_A2C_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_ENV = 16
    UPDATE_FREQ = 5
    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = MultiThreadEnv([atari_env_factory(name, STATE_SIZE, N_FRAMES;) for i in 1:N_ENV]) #= seed = i =#
    N_ACTIONS = length(get_action_space(env[1]))

    init = seed_glorot_uniform()#= seed=341 =#

    # share model
    model = Chain(
        x -> x ./ 255,
        CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
        CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
        CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
        x -> reshape(x, :, size(x)[end]),
        Dense(11 * 11 * 64, 512, relu),
    )

    agent = Agent(
        policy = QBasedPolicy(
            learner = A2CLearner(
                approximator = ActorCritic(
                    actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
                    critic = Chain(model, Dense(512, 1; initW = init)),
                    optimizer = RMSProp(7e-4, 0.99),
                ) |> gpu,
                γ = 0.99f0,
                max_grad_norm = 0.5f0,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.25f0,
                entropy_loss_weight = 0.01f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(;)),#= seed = s=#
        ),
        trajectory = CircularCompactSARTSATrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (STATE_SIZE..., N_FRAMES, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )

    N_TRAINING_STEPS = 10_000_000
    EVALUATION_FREQ = 100_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    total_batch_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    batch_steps_per_episode = BatchStepsPerEpisode(N_ENV)
    evaluation_result = []

    hook = ComposedHook(
        total_batch_reward_per_episode,
        batch_steps_per_episode,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env, obs
            learner = agent.policy.learner
            with_logger(lg) do
                @info "training" loss = learner.loss actor_loss = learner.actor_loss critic_loss =
                    learner.critic_loss entropy_loss = learner.entropy_loss norm =
                    learner.norm log_step_increment = UPDATE_FREQ
            end
        end,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                rewards = [
                    total_batch_reward_per_episode.rewards[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(rewards) > 0
                    @info "training" rewards = mean(rewards) log_step_increment = 0
                end
                steps = [
                    batch_steps_per_episode.steps[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(steps) > 0
                    @info "training" steps = mean(steps) log_step_increment = 0
                end
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            h = TotalBatchRewardPerEpisode(N_ENV)
            s = @elapsed run(
                agent,
                MultiThreadEnv([
                    atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;) for i in 1:N_ENV
                ]),
                StopAfterStep(27_000; is_show_progress = false),
                h,
            )
            res = (avg_score = mean(Iterators.flatten(h.rewards)),)
            push!(evaluation_result, res)
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_score = res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent;)
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_batch_reward_per_episode evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    description = """
    # Play Atari($name) with A2C

    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:A2CGAE},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_A2CGAE_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_ENV = 16
    UPDATE_FREQ = 32
    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = MultiThreadEnv([atari_env_factory(name, STATE_SIZE, N_FRAMES;) for i in 1:N_ENV]) #= seed = i =#
    N_ACTIONS = length(get_action_space(env[1]))

    init = seed_glorot_uniform()#= seed=341 =#

    # share model
    model = Chain(
        x -> x ./ 255,
        CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
        CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
        CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
        x -> reshape(x, :, size(x)[end]),
        Dense(11 * 11 * 64, 512, relu),
    )

    agent = Agent(
        policy = QBasedPolicy(
            learner = A2CGAELearner(
                approximator = ActorCritic(
                    actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
                    critic = Chain(model, Dense(512, 1; initW = init)),
                    optimizer = RMSProp(7e-4, 0.99),
                ) |> gpu,
                γ = 0.99f0,
                λ = 0.95f0,
                max_grad_norm = 0.5f0,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.25f0,
                entropy_loss_weight = 0.01f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(;)),#= seed = s=#
        ),
        trajectory = CircularCompactSARTSATrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (STATE_SIZE..., N_FRAMES, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )

    N_TRAINING_STEPS = 10_000_000
    EVALUATION_FREQ = 100_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    total_batch_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    batch_steps_per_episode = BatchStepsPerEpisode(N_ENV)
    evaluation_result = []

    hook = ComposedHook(
        total_batch_reward_per_episode,
        batch_steps_per_episode,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env, obs
            learner = agent.policy.learner
            with_logger(lg) do
                @info "training" loss = learner.loss actor_loss = learner.actor_loss critic_loss =
                    learner.critic_loss entropy_loss = learner.entropy_loss norm =
                    learner.norm log_step_increment = UPDATE_FREQ
            end
        end,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                rewards = [
                    total_batch_reward_per_episode.rewards[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(rewards) > 0
                    @info "training" rewards = mean(rewards) log_step_increment = 0
                end
                steps = [
                    batch_steps_per_episode.steps[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(steps) > 0
                    @info "training" steps = mean(steps) log_step_increment = 0
                end
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            h = TotalBatchRewardPerEpisode(N_ENV)
            s = @elapsed run(
                agent,
                MultiThreadEnv([
                    atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;) for i in 1:N_ENV
                ]),
                StopAfterStep(27_000; is_show_progress = false),
                h,
            )
            res = (avg_score = mean(Iterators.flatten(h.rewards)),)
            push!(evaluation_result, res)
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_score = res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent;)
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_batch_reward_per_episode evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    description = """
    # Play Atari($name) with A2CGAE

    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    """

    Experiment(agent, env, stop_condition, hook, description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:PPO},
    ::Val{:Atari},
    name::AbstractString;
    save_dir = nothing,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyymmddHHMMSS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_PPO_Atari_$(name)_$(t)")
    end

    lg = TBLogger(joinpath(save_dir, "tb_log"), min_level = Logging.Info)

    N_ENV = 32
    UPDATE_FREQ = 64
    N_FRAMES = 4
    STATE_SIZE = (84, 84)
    env = MultiThreadEnv([atari_env_factory(name, STATE_SIZE, N_FRAMES;) for i in 1:N_ENV]) #= seed = i =#
    N_ACTIONS = length(get_action_space(env[1]))
    INIT_CLIP_RANGE = 0.1f0
    INIT_LEARNING_RATE = 1e-3

    init = seed_orthogonal()#= seed=341 =#

    # share model
    model = Chain(
        x -> x ./ 255,
        CrossCor((8, 8), N_FRAMES => 32, relu; stride = 4, pad = 2, init = init),
        CrossCor((4, 4), 32 => 64, relu; stride = 2, pad = 2, init = init),
        CrossCor((3, 3), 64 => 64, relu; stride = 1, pad = 1, init = init),
        x -> reshape(x, :, size(x)[end]),
        Dense(11 * 11 * 64, 512, relu),
    )

    agent = Agent(
        policy = QBasedPolicy(
            learner = PPOLearner(
                approximator = ActorCritic(
                    actor = Chain(model, Dense(512, N_ACTIONS; initW = init)),
                    critic = Chain(model, Dense(512, 1; initW = init)),
                    optimizer = ADAM(INIT_LEARNING_RATE),  # decrease learning rate with a hook
                ) |> gpu,
                γ = 0.99f0,
                λ = 0.98f0,
                clip_range = INIT_CLIP_RANGE,  # decrease with a hook
                max_grad_norm = 1.0f0,
                n_microbatches = 4,
                n_epochs = 4,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 1.0f0,
                entropy_loss_weight = 0.01f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(;)),#= seed = s=#
        ),
        trajectory = PPOTrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (STATE_SIZE..., N_FRAMES, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            action_log_prob_type = Float32,
            action_log_prob_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )

    N_TRAINING_STEPS = 10_000_000
    EVALUATION_FREQ = 100_000
    MAX_EPISODE_STEPS_EVAL = 27_000
    N_CHECKPOINTS = 3
    stop_condition = StopAfterStep(N_TRAINING_STEPS)

    total_batch_reward_per_episode = TotalBatchRewardPerEpisode(N_ENV)
    batch_steps_per_episode = BatchStepsPerEpisode(N_ENV)
    evaluation_result = []

    hook = ComposedHook(
        total_batch_reward_per_episode,
        batch_steps_per_episode,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env, obs
            learner = agent.policy.learner
            with_logger(lg) do
                @info "training" loss = mean(learner.loss) actor_loss =
                    mean(learner.actor_loss) critic_loss = mean(learner.critic_loss) entropy_loss =
                    mean(learner.entropy_loss) norm = mean(learner.norm) log_step_increment =
                    UPDATE_FREQ
            end
        end,
        DoEveryNStep(UPDATE_FREQ) do t, agent, env, obs
            decay = (N_TRAINING_STEPS - t) / N_TRAINING_STEPS
            agent.policy.learner.approximator.optimizer.eta =
                INIT_LEARNING_RATE * decay
            agent.policy.learner.clip_range = INIT_CLIP_RANGE * Float32(decay)
        end,
        DoEveryNStep() do t, agent, env, obs
            with_logger(lg) do
                rewards = [
                    total_batch_reward_per_episode.rewards[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(rewards) > 0
                    @info "training" rewards = mean(rewards) log_step_increment = 0
                end
                steps = [
                    batch_steps_per_episode.steps[i][end]
                    for i in 1:length(obs) if get_terminal(obs[i])
                ]
                if length(steps) > 0
                    @info "training" steps = mean(steps) log_step_increment = 0
                end
            end
        end,
        DoEveryNStep(EVALUATION_FREQ) do t, agent, env, obs
            @info "evaluating agent at $t step..."
            flush(stdout)
            Flux.testmode!(agent)
            h = TotalBatchRewardPerEpisode(N_ENV)
            s = @elapsed run(
                agent,
                MultiThreadEnv([
                    atari_env_factory(name, STATE_SIZE, N_FRAMES, MAX_EPISODE_STEPS_EVAL;) for i in 1:4
                ]),
                StopAfterStep(27_000; is_show_progress = false),
                h,
            )
            res = (avg_score = mean(Iterators.flatten(h.rewards)),)
            push!(evaluation_result, res)
            Flux.trainmode!(agent)
            @info "finished evaluating agent in $s seconds" avg_score = res.avg_score
            with_logger(lg) do
                @info "evaluating" avg_score = res.avg_score log_step_increment = 0
            end
            flush(stdout)

            RLCore.save(joinpath(save_dir, string(t)), agent;)
            BSON.@save joinpath(save_dir, string(t), "stats.bson") total_batch_reward_per_episode evaluation_result

            # only keep recent 3 checkpoints
            old_checkpoint_folder =
                joinpath(save_dir, string(t - EVALUATION_FREQ * N_CHECKPOINTS))
            if isdir(old_checkpoint_folder)
                rm(old_checkpoint_folder; force = true, recursive = true)
            end
        end,
    )

    description = """
    # Play Atari($name) with PPO

    You can also view the tensorboard logs with `tensorboard --logdir $(joinpath(save_dir, "tb_log"))`
    """

    Experiment(agent, env, stop_condition, hook, description)
end


text:
using Game2048
using Game2048: bitboard_to_array
using Colors

WIDTH = 1600
HEIGHT = 1200
BACKGROUND = colorant"grey"

board = initbboard()

CELL_WIDTH = 18*11

function drawnum(i, j, val)
    draw(Rect((CELL_WIDTH+8)j, (CELL_WIDTH+8)i, CELL_WIDTH, CELL_WIDTH), colorant"black", fill=true)
    for v in val:-1:1
        draw(Rect((CELL_WIDTH+8)j+9(11-v), (CELL_WIDTH+8)i+9(11-v), CELL_WIDTH-18(11-v), CELL_WIDTH-18(11-v)), RGB(v/11, v/11, v/11), fill=true)
    end
end

function drawboard(board)
    board_arr = bitboard_to_array(board)
    for i in 1:4, j in 1:4
        val = board_arr[i, j]
        if val != 0
            #draw(Rect(105j, 105i, 100, 100), colorant"black", fill=true)
            drawnum(i, j, val)
        end
    end
end

function on_key_down(g::Game, key)
    global board
    if key == Keys.LEFT
        new_board = move(board, left)
    elseif key == Keys.RIGHT
        new_board = move(board, right)
    elseif key == Keys.DOWN
        new_board = move(board, down)
    elseif key == Keys.UP
        new_board = move(board, up)
    end

    if new_board != board
        board = add_tile(new_board)
    end
end

function update(g::Game)

end

function draw(g::Game)
    clear()
    drawboard(board)
end

ext: jl


text:
"""
    gamma_bounds(gamma)

A bounds-generating function for [`flux_variability_analysis`](@ref) that
limits the objective value to be at least `gamma*Z₀`, as usual in COBRA
packages. Use as the `bounds` argument:
```
flux_variability_analysis(model, some_optimizer; bounds = gamma_bounds(0.9))
```
"""
gamma_bounds(gamma) = z -> (gamma * z, Inf)

"""
    (tolerance) = z -> begin

A bounds-generating function for [`flux_variability_analysis`](@ref) that
limits the objective value to a small multiple of Z₀. Use as `bounds` argument,
similarly to [`gamma_bounds`](@ref).
"""
objective_bounds(tolerance) = z -> begin
    vs = (z * tolerance, z / tolerance)
    (minimum(vs), maximum(vs))
end


text:
#using unfold
using MAT
using DataFrames
using DelimitedFiles

function import_eeglab(filename)

    file = MAT.matopen(filename)
    # open file
    EEG = read(file, "EEG")

    function parse_struct(s::Dict)
        return DataFrame(map(x->dropdims(x,dims=1),values(s)),collect(keys(s)))
    end

    evts_df = parse_struct(EEG["event"])
    chanlocs_df = parse_struct(EEG["chanlocs"])

#    epoch_df = parse_struct(EEG["epochs"])

    srate = EEG["srate"]
    if typeof(EEG["data"]) == String
        datapath = joinpath(splitdir(filename)[1],EEG["data"])
        data = Array{Float32, 3}(undef,Int(EEG["nbchan"]),Int(EEG["pnts"]),Int(EEG["trials"]))
        read!(datapath,data)
    else
        data = EEG["data"]
    end

    if (ndims(data)==3) & (size(data,3) == 1)
        data = dropdims(data,dims=3)
    end
return data,srate,evts_df,chanlocs_df,EEG
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule libcint_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("libcint")
JLLWrappers.@generate_main_file("libcint", UUID("574b78ca-bebd-517c-801d-4735c93a9686"))
end  # module libcint_jll


text:
"""
    GodunovFlux

Godunovs flux \$f(u(x/t = 0))\$, i.e. the flux of the value at zero of the solution
of the Riemann problem.
"""
struct GodunovFlux <: NumericalFlux end

"""
    LocalLaxFriedrichsFlux{MaxAbsSpeed}

The local Lax-Friedrichs flux \$\\frac{f(u_r) + f(u_l)}{2} - \\frac{\\lambda}{2} (u_r - u_l)\$.
\$\\lambda\$ is the maximal absolute value of the speed in the solution of a Riemann
problem, computed by `max_abs_speed::MaxAbsSpeed`.
"""
struct LocalLaxFriedrichsFlux{MaxAbsSpeed} <: NumericalFlux
    max_abs_speed::MaxAbsSpeed
end

LocalLaxFriedrichsFlux() = LocalLaxFriedrichsFlux(max_abs_speed)

Base.@pure function (fnum::LocalLaxFriedrichsFlux)(uₗ, uᵣ, model::AbstractBalanceLaw{1})
    λ = fnum.max_abs_speed(uₗ, uᵣ, model)

    (flux(uₗ,model) + flux(uᵣ,model))/2 - λ/2 * (uᵣ - uₗ)
end

Base.@pure function (fnum::LocalLaxFriedrichsFlux)(uₗ, uᵣ, model::AbstractBalanceLaw, direction)
    λ = fnum.max_abs_speed(uₗ, uᵣ, model, direction)

    (flux(uₗ, model, direction) + flux(uᵣ, model, direction))/2 - λ/2 * (uᵣ - uₗ)
end

"""
    max_abs_speed(uₗ, uᵣ, model::AbstractBalanceLaw)

Compute the maximal absolute value of the speed in the solution of the Riemann
problem with states `uₗ`, `uᵣ` for `model`.
"""
@inline function max_abs_speed(uₗ, uᵣ, model::AbstractBalanceLaw)
    naive_max_abs_speed(uₗ, uᵣ, model)
end

@inline function max_abs_speed(uₗ, uᵣ, model::AbstractBalanceLaw, direction)
    naive_max_abs_speed(uₗ, uᵣ, model, direction)
end

@inline function naive_max_abs_speed(uₗ, uᵣ, model::AbstractBalanceLaw)
    max(max_abs_speed(uₗ,model), max_abs_speed(uᵣ,model))
end

@inline function naive_max_abs_speed(uₗ, uᵣ, model::AbstractBalanceLaw, direction)
    max(max_abs_speed(uₗ,model,direction), max_abs_speed(uᵣ,model,direction))
end

"""
    max_abs_speed{T}(u::T, model::ScalarBalanceLaw{T,1})

Compute the maximal absolute value of speed at `u` for `model`.
"""
@inline max_abs_speed(u::T, model::ScalarBalanceLaw{T,1}) where {T} = abs(speed(u, model))

"""
    min_max_speed{T}(u::T, model::ScalarBalanceLaw{T,1})

Compute the minimal and maximal speed at `u` for `model`.
"""
@inline function min_max_speed(u::T, model::ScalarBalanceLaw{T,1}) where {T}
    λ = speed(u, model)
    λ, λ
end

"""
    HartenLaxVanLeerFlux{MinMaxSpeed}

The HLL (Harten, Lax, van Leer) flux using `min_max_speed::MinMaxSpeed` to compute
 the (approximate) minimal and maximal speed in the solution of a Riemann problem.
"""
struct HartenLaxVanLeerFlux{MinMaxSpeed} <: NumericalFlux
    min_max_speed::MinMaxSpeed
end

const HLL = HartenLaxVanLeerFlux

HartenLaxVanLeerFlux() = HartenLaxVanLeerFlux(min_max_speed)

@inline function (fnum::HartenLaxVanLeerFlux)(uₗ, uᵣ, model::AbstractBalanceLaw{1})
    λ₋, λ₊ = fnum.min_max_speed(uₗ, uᵣ, model)

    if 0 <= λ₋
        flux(uₗ, model)
    elseif 0 < λ₊
        ( λ₊*flux(uₗ, model) - λ₋*flux(uᵣ, model) + λ₋*λ₊*(uᵣ-uₗ) ) / (λ₊ - λ₋)
    else
        flux(uᵣ, model)
    end
end

@inline function (fnum::HartenLaxVanLeerFlux)(uₗ, uᵣ, model::AbstractBalanceLaw, direction)
    λ₋, λ₊ = fnum.min_max_speed(uₗ, uᵣ, model, direction)

    if 0 <= λ₋
        flux(uₗ, model, direction)
    elseif 0 < λ₊
        ( λ₊*flux(uₗ, model, direction) - λ₋*flux(uᵣ, model, direction) + λ₋*λ₊*(uᵣ-uₗ) ) / (λ₊ - λ₋)
    else
        flux(uᵣ, model, direction)
    end
end

"""
    min_max_speed(uₗ, uᵣ, model::AbstractBalanceLaw{1})

Compute the maximal absolute value of the speed in the solution of the Riemann
problem with states `uₗ`, `uᵣ` for `model`.
"""
@inline function min_max_speed(uₗ, uᵣ, model::AbstractBalanceLaw{1})
    λₗ₋, λₗ₊ = min_max_speed(uₗ, model)
    λᵣ₋, λᵣ₊ = min_max_speed(uᵣ, model)
    min(λₗ₋, λᵣ₋), max(λₗ₊, λᵣ₊)
end

@inline function min_max_speed(uₗ, uᵣ, model::AbstractBalanceLaw, direction)
    λₗ₋, λₗ₊ = min_max_speed(uₗ, model, direction)
    λᵣ₋, λᵣ₊ = min_max_speed(uᵣ, model, direction)
    min(λₗ₋, λᵣ₋), max(λₗ₊, λᵣ₊)
end

"""
    EnergyConservativeFlux

The "standard" energy conservative flux.
"""
struct EnergyConservativeFlux <: NumericalFlux end

"""
    EnergyConservativeFlux1Param

A one-parameter family of energy conservative fluxes.
"""
struct EnergyConservativeFlux1Param{T} <: NumericalFlux
    a₁::T
end

"""
    EnergyConservativeFlux2Param

A two-parameter family of energy conservative fluxes.
"""
struct EnergyConservativeFlux2Param{T} <: NumericalFlux
    a₁::T
    a₂::T
end

"""
    L2L4ConservativeFlux

A numerical flux conserving the \$L^2 \\cap L^4\$ entropy U(U) = u^2 + u^4.
"""
struct L2L4ConservativeFlux <: NumericalFlux end

"""
    L2L2sConservativeFlux

A numerical flux conserving the \$L^2 \\cap L^{2s}\$ entropy U(U) = u^2 + u^(2s).
"""
struct L2L2sConservativeFlux{s} <: NumericalFlux
    half_pow::Val{s}

    function L2L2sConservativeFlux(half_pow::Val{s}) where s
        @assert typeof(s) <: Integer && s >= 1
        new{s}(half_pow)
    end
end

"""
    jump_pow_u_r_over_jump_u(uₗ, uᵣ, ::Val{r})

Compute `(uᵣ^r - uₗ^r) / (uᵣ - uₗ)` for integer `r >= 2` via the expanded
expression `uᵣ^(r-1) + rᵣ^(r-2)*uₗ + ... + uᵣ*uₗ^(r-2) + uₗ^(r-1)`.
"""
@generated function jump_pow_u_r_over_jump_u(uₗ, uᵣ, ::Val{r}) where r
    @assert typeof(r) <: Integer && r >= 2
    ex = :( Base.FastMath.pow_fast(uₗ, $r-1) )
    ex = :( $ex + Base.FastMath.pow_fast(uᵣ, $r-1) )
    for pₗ in Base.OneTo(r-2)
        pᵣ = r-1 - pₗ
        ex = :( $ex + Base.FastMath.pow_fast(uₗ, $pₗ) * Base.FastMath.pow_fast(uᵣ, $pᵣ) )
    end
    return :(Base.@_inline_meta; $ex)
end

"""
    CentralFlux

The central flux \$\\frac{f(u_r) + f(u_l)}{2}\$.
"""
struct CentralFlux <: NumericalFlux end

Base.@pure function (fnum::CentralFlux)(uₗ, uᵣ, model)
    (flux(uₗ, model) + flux(uᵣ, model)) / 2
end

Base.@pure function (fnum::CentralFlux)(uₗ, uᵣ, model, direction)
    (flux(uₗ, model, direction) + flux(uᵣ, model, direction)) / 2
end

"""
    MorinishiFlux

The numerical flux corresponding to the splitting of Morinishi (2010),
see Gassner, Winter, Kopriva (2016).
"""
struct MorinishiFlux <: NumericalFlux end

"""
    DucrosEtAlFlux

The numerical flux corresponding to the splitting of Ducros, Laporte, Soulères,
Guinot, Moinat, Caruelle (2000), see Gassner, Winter, Kopriva (2016).
"""
struct DucrosEtAlFlux <: NumericalFlux end

"""
    KennedyGruberFlux

The numerical flux corresponding to the splitting of Kennedy and Gruber (2008),
see Gassner, Winter, Kopriva (2016).
"""
struct KennedyGruberFlux <: NumericalFlux end

"""
    PirozzoliFlux

The numerical flux corresponding to the splitting of Pirozzoli (2011),
see Gassner, Winter, Kopriva (2016).
"""
struct PirozzoliFlux <: NumericalFlux end

"""
    SuliciuFlux

The Suliciu relaxation solver.
"""
struct SuliciuFlux <: NumericalFlux end

"""
    KineticFlux

A kinetic relaxation solver.
"""
struct KineticFlux <: NumericalFlux end

"""
    ChandrashekarFluxEC

The entropy conservative flux of Chandrashekar (2013).
"""
struct ChandrashekarFluxEC <: NumericalFlux end

"""
    IsmailRoeFluxEC

The entropy conservative flux of Ismail and Roe (2009).
"""
struct IsmailRoeFluxEC <: NumericalFlux end

"""
    RanochaFluxECandKEP

The entropy conservative and kinetic energy preserving flux of Ranocha (2017).
"""
struct RanochaFluxECandKEP <: NumericalFlux end

"""
    FluxPlusDissipation{Fnum::NumericalFlux, Diss} <: NumericalFlux

A numerical flux given by the (possibly central / symmetric) numerical flux
`fnum` and the dissipation operator `diss`.
"""
struct FluxPlusDissipation{Fnum<:NumericalFlux, Diss<:DissipationOperator} <: NumericalFlux
    fnum::Fnum
    diss::Diss
end

@inline function (fnumdiss::FluxPlusDissipation)(uₗ, uᵣ, model::AbstractBalanceLaw)
    @unpack fnum, diss = fnumdiss

    fnum(uₗ, uᵣ, model) + diss(uₗ, uᵣ, model)
end

"""
    LocalLaxFriedrichsDissipation{MaxAbsSpeed}

The local Lax-Friedrichs dissipation \$- \\frac{\\lambda}{2} (u_r - u_l)\$.
\$\\lambda\$ is the approximate maximal absolute value of the speed in the solution
of a Riemann problem, computed by `max_abs_speed::MaxAbsSpeed`.
"""
struct LocalLaxFriedrichsDissipation{MaxAbsSpeed} <: DissipationOperator
    max_abs_speed::MaxAbsSpeed
end

LocalLaxFriedrichsDissipation() = LocalLaxFriedrichsDissipation(naive_max_abs_speed)

@inline function (diss::LocalLaxFriedrichsDissipation)(uₗ, uᵣ, model::AbstractBalanceLaw)
    λ = diss.max_abs_speed(uₗ, uᵣ, model)

    -λ * (uᵣ - uₗ) / 2
end

"""
    ScalarDissipation{MaxAbsSpeed}

The scalar dissipation \$- \\frac{\\lambda}{2} R \\cdot R^T \\cdot  (w_r - w_l)\$.
\$\\lambda\$ is the approximate maximal absolute value of the speed, computed by
`max_abs_speed::MaxAbsSpeed`.
"""
struct ScalarDissipation{MaxAbsSpeed} <: DissipationOperator
    max_abs_speed::MaxAbsSpeed
end

ScalarDissipation() = ScalarDissipation(naive_max_abs_speed)

"""
    MatrixDissipation

The matrix dissipation \$- \\frac{1}{2} R \\cdot |\\Lambda| \\cdot R^T \\cdot  (w_r - w_l)\$.
"""
struct MatrixDissipation <: DissipationOperator end


text:
Hello,
world.

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
module TypeParams

export @typeparams

"""
    @typeparams typedef

Automatically generate type parameters for fields of the form `fieldname::{...}`.

# Examples
```
julia> @macroexpand(
       @typeparams struct T
           a::{}
       end)
:(struct T{var"##a#1"}
      a::var"##a#1"
  end)

julia> @macroexpand(
       @typeparams struct T
           a::{<:Integer}
       end)
:(struct T{var"##a#1" <: Integer}
      a::var"##a#1"
  end)

julia> @macroexpand(
       @typeparams struct T
            a::{A}
       end)
:(struct T{A}
      a::A
  end)
```
"""
macro typeparams(typedef)
    typedef = macroexpand(__module__, typedef)
    @assert typedef isa Expr && typedef.head === :struct "@typeparams must be applied to a struct definition"
    typename = curly_typename!(typedef)
    paramnames = paramname.(typename.args[2:end])
    scan_typebody!(typedef.args[3],typename,paramnames)
    return esc(typedef)
end

"""
    curly_typename!(typedef)

Extract the type name and parameters from `typedef`.

# Examples
```jldoctest
julia> curly_typename!(:(struct Foo; end))
:(Foo{})
```
"""
function curly_typename!(typedef)
    typename_ref = Ref(typedef.args,2)
    if typename_ref[] isa Expr && typename_ref[].head === :<:
        typename_ref = Ref(typename_ref[].args,1)
    end
    if typename_ref[] isa Symbol
        typename_ref[] = Expr(:curly, typename_ref[])
    end
    return typename_ref[]
end

paramname(paramdef::Symbol) = paramdef
function paramname(paramdef::Expr)
    if paramdef.head === :<:
        @assert paramdef.args[1] isa Symbol "Invalid type parameter"
        return paramdef.args[1]
    end
    @assert false "Invalid type parameter"
end

function scan_typebody!(typebody, typename, paramnames)
    for item in typebody.args
        if !(item isa Expr); continue; end
        if item.head === :block
            scan_typebody!(item, typename)
            continue
        end
        if item.head === :(=)
            item = item.args[1]
        end
        if item.head === :(::)
            name = item.args[1]
            type = item.args[2]
            if type isa Expr && type.head === :braces
                if length(type.args) == 0
                    paramname = gensym(name)
                    paramdef = paramname
                elseif length(type.args) == 1
                    braces_content = type.args[1]
                    if braces_content isa Symbol
                        paramname = braces_content
                        paramdef = braces_content
                    elseif braces_content isa Expr && braces_content.head === :<:
                        if length(braces_content.args) == 1
                            paramname = gensym(name)
                            paramdef = :($paramname <: $(braces_content.args[1]))
                        elseif length(braces_content.args) == 2
                            paramname = braces_content.args[1]
                            paramdef = braces_content
                        else
                            @assert false "Invalid type parameter"
                        end
                    else
                        @assert false "Invalid type parameter"
                    end
                else
                    @assert false "Invalid type parameter"
                end
                item.args[2] = paramname
                if !(paramname in paramnames)
                    push!(typename.args, paramdef)
                    push!(paramnames, paramname)
                else
                    @assert paramdef isa Symbol "Type constraint on previously defined type parameter"
                end
            end
        end
    end
end

end # module


text:
using DataFrames
using XLSX

excel_df = DataFrame(XLSX.readtable("data/SalesData.xlsx", "SalesOrders")...)

using CSV
csv_df = CSV.read("data/StarWars.csv")

ss_df = CSV.read("data/Cereal.txt")

delim_df = CSV.read("data/Soccer.txt", delim = "::")

size(ss_df)

nrow(ss_df), ncol(ss_df)

describe(ss_df)

names(ss_df)

eltypes(ss_df)

csv_df[1,1]

csv_df[1]

csv_df[:Name]

csv_df.Name

csv_df[:, 1] # note that this creates a copy

csv_df[1:1, :]

csv_df[1, :] # this produces a DataFrameRow

excel_df[1:3, 5] = 1

excel_df[4:6, 5] = [4, 5, 6]

excel_df[1:2, 6:7] = DataFrame([-2 -2; -2 -2])

excel_df

passportdata = CSV.read("data/passport-index-matrix.csv", copycols = true)

for i in 1:nrow(passportdata)
    for j in 2:ncol(passportdata)
        if passportdata[i,j] == -1 || passportdata[i,j] == 3
            passportdata[i,j] = 1
        else
            passportdata[i,j] = 0
        end
    end
end

using JuMP, GLPK

# Finding number of countries
n = ncol(passportdata) - 1 # Subtract 1 for column representing country of passport

model = Model(with_optimizer(GLPK.Optimizer))
@variable(model, pass[1:n], Bin)
@constraint(model, [j = 2:n], sum(passportdata[i,j] * pass[i] for i in 1:n) >= 1)
@objective(model, Min, sum(pass))
optimize!(model)

println("Minimum number of passports needed: ", objective_value(model))

countryindex = findall(value.(pass) .== 1 )

print("Countries: ")
for i in countryindex
    print(names(passportdata)[i+1], " ")
end


text:
module IterationTools

export halt, tee, sample, stopwatch, loop

using Base.Iterators

# Halting

struct HaltingIterable{I,F}
    iter::I
    fun::F
end

Base.IteratorSize(::Type{HaltingIterable{I,F}}) where {I,F} = Base.IteratorSize(I)
Base.IteratorEltype(::Type{HaltingIterable{I,F}}) where {I,F} = Base.IteratorEltype(I)

Base.length(iter::HaltingIterable{I,F}) where {I,F} = length(iter.iter)
Base.eltype(iter::HaltingIterable{I,F}) where {I,F} = eltype(iter.iter)

function Base.iterate(iter::HaltingIterable)
    next = iterate(iter.iter)
    return dispatch(iter, next)
end

function Base.iterate(iter::HaltingIterable, (instruction, state))
    if instruction == :halt
        return nothing
    end
    next = iterate(iter.iter, state)
    return dispatch(iter, next)
end

function dispatch(iter::HaltingIterable, next)
    if next === nothing
        return nothing
    end
    return next[1], (iter.fun(next[1]) ? :halt : :continue, next[2])
end

halt(iter::I, fun::F) where {I,F} = HaltingIterable{I,F}(iter, fun)

# Side effects

struct TeeIterable{I,F}
    iter::I
    fun::F
end

Base.IteratorSize(::Type{TeeIterable{I,F}}) where {I,F} = Base.IteratorSize(I)
Base.IteratorEltype(::Type{TeeIterable{I,F}}) where {I,F} = Base.IteratorEltype(I)

Base.length(iter::TeeIterable{I,F}) where {I,F} = length(iter.iter)
Base.axes(iter::TeeIterable{I,F}) where {I,F} = axes(iter.iter)
Base.eltype(iter::TeeIterable{I,F}) where {I,F} = eltype(iter.iter)

function Base.iterate(iter::TeeIterable, args...)
    next = iterate(iter.iter, args...)
    if next !== nothing
        iter.fun(next[1])
    end
    return next
end

tee(iter::I, fun::F) where {I,F} = TeeIterable{I,F}(iter, fun)

# Sampling

struct SamplingIterable{I}
    iter::I
    period::UInt
end

Base.IteratorSize(::Type{SamplingIterable{I}}) where {I} = Base.IteratorSize(I)
Base.IteratorEltype(::Type{SamplingIterable{I}}) where {I} = Base.IteratorEltype(I)

function Base.length(iter::SamplingIterable{I}) where {I}
    remainder = length(iter.iter) % iter.period
    quotient = length(iter.iter) ÷ iter.period
    return remainder == 0 ? quotient : quotient + 1
end

Base.size(iter::SamplingIterable{I}) where {I} = (length(iter),)
Base.eltype(iter::SamplingIterable{I}) where {I} = eltype(iter.iter)

function Base.iterate(iter::SamplingIterable, state = iter.iter)
    current = iterate(state)
    if current === nothing
        return nothing
    end
    for i = 1:iter.period-1
        next = iterate(state, current[2])
        if next === nothing
            return current[1], rest(state, current[2])
        end
        current = next
    end
    return current[1], rest(state, current[2])
end

sample(iter::I, period) where {I} = SamplingIterable{I}(iter, period)

# Timing

struct StopwatchIterable{I}
    iter::I
end

Base.IteratorSize(::Type{StopwatchIterable{I}}) where {I} = Base.IteratorSize(I)
Base.IteratorEltype(::Type{StopwatchIterable{I}}) where {I} = Base.IteratorEltype(I)

Base.length(iter::StopwatchIterable{I}) where {I} = length(iter.iter)
Base.axes(iter::StopwatchIterable{I}) where {I} = axes(iter.iter)
Base.eltype(iter::StopwatchIterable{I}) where {I} = (UInt64, eltype(iter.iter))

function Base.iterate(iter::StopwatchIterable)
    t0 = time_ns()
    next = iterate(iter.iter)
    return dispatch(iter, t0, next)
end

function Base.iterate(iter::StopwatchIterable, (t0, state))
    next = iterate(iter.iter, state)
    return dispatch(iter, t0, next)
end

function dispatch(iter::StopwatchIterable, t0, next)
    if next === nothing
        return nothing
    end
    return (time_ns() - t0, next[1]), (t0, next[2])
end

stopwatch(iter::I) where {I} = StopwatchIterable{I}(iter)

# Looping

function loop(iter)
    output, state = iterate(iter)
    res = iterate(iter, state)
    while res !== nothing
        output, state = res
        res = iterate(iter, state)
    end
    return output
end

end # module


text:
#= Run this script interactively: `julia --project -i runtests.jl`
... or with arguments `julia --project runtests.jl [testfiles...]` =#

using Test, TestSetExtensions, Random
using Coverage, Logging
using Revise, GeometricAlgebra

project_root = dirname(dirname(pathof(GeometricAlgebra)))
cd(joinpath(project_root, "test"))

function test(files...)
	@eval @testset ExtendedTestSet "GeometricAlgebra" begin
		@includetests $files
	end
	nothing
end

function coverage()
	run(`julia --code-coverage --project runtests.jl`)
	score()
end

function score()
	local percentage
	with_logger(NullLogger()) do
		coverage = process_folder(joinpath(project_root, "src"))
		covered, total = get_summary(coverage)
		percentage = round(100covered/total, digits=2)
	end
	@info "Coverage: $percentage%"
end

clear_coverage_files() = run(`bash -c "rm $project_root/**/*.cov"`)

if isempty(ARGS)
	@info """Run this script in interactive mode, and call:
		 - `test("testfile1", "testfile2", ...)`
		   (without `.jl` extensions) to run tests
		 - `test()` to run all tests
		 - `coverage()` to analyse code coverage
		 - `score()` to show coverage score from last run
		 - `clear_coverage_files()` to remove `*.cov` files
		Keep the session alive; changes will be revised and successive runs will be faster.
		"""
	!isinteractive() && test()
elseif "--code-coverage" in ARGS
	coverage()
else
	test(ARGS...)
end


text:
### SaddleSystem ###

abstract type SchurSolverType end

struct SaddleSystem{T,Ns,Nc,TU,TF,TS<:SchurSolverType}
    A :: LinearMap{T}
    B₂ :: LinearMap{T}
    B₁ᵀ :: LinearMap{T}
    C :: LinearMap{T}
    A⁻¹ :: LinearMap{T}
    A⁻¹B₁ᵀf :: Vector{T}
    B₂A⁻¹r₁ :: Vector{T}
    _f_buf :: Vector{T}
    P :: LinearMap{T}
    S :: LinearMap{T}
    S⁻¹ :: LinearMap{T}
end

# Constructors

"""
    SaddleSystem

Construct a saddle-point system operator from the constituent operator blocks. The resulting object can be used
with `*` and `\\` to multiply and solve. The saddle-point problem has the form

``
\\begin{bmatrix}A & B_1^T \\\\ B_2 & C \\end{bmatrix} \\begin{pmatrix} u \\\\ f \\end{pmatrix} = \\begin{pmatrix} r_1 \\\\ r_2 \\end{pmatrix}
``

### Constructors
`SaddleSystem(A::AbstractMatrix,B₂::AbstractMatrix,B₁ᵀ::AbstractMatrix,C::AbstractMatrix[,eltype=Float64])`.
Blocks are given as matrices. Must have consistent sizes to stack appropriately. If
this is called with `SaddleSystem(A,B₂,B₁ᵀ)`, it sets `C` to zero automatically.

`SaddleSystem(A,B₂,B₁ᵀ,C,u,f[,eltype=Float64])`.
Operators `A`, `B₂`, `B₁ᵀ`, `C` are given in various forms, including matrices, functions, and function-like objects.
`u` and `f` are examples of the data types in the corresponding solution and right-hand side vectors.
Guidelines:

* The entries `A` and `B₂` must be able to act upon `u` (either by multiplication or as a function) and `B₁ᵀ` and `C` must be able to act on `f` (also, either by multiplication or as a function).
* `A` and `B₁ᵀ` should return data of type `u`, and `B₂` and `C` should return data of type `f`.
* `A` must be invertible and be outfitted with operators `\` and `ldiv!`.
* Both `u` and `f` must be subtypes of `AbstractArray`: they must be equipped with `size`
  and `vec` functions and with a constructor of the form `T(data)` where `T` is the data type of
  `u` or `f` and `data` is the wrapped data array.

If called as `SaddleSystem(A,B₂,B₁ᵀ,u,f)`, the `C` block is omitted and assumed to be zero.

If called with `SaddleSystem(A,u)`, this is equivalent to calling `SaddleSystem(A,nothing,nothing,u,[])`, then this reverts
to the unconstrained system described by operator `A`.

The list of vectors `u` and `f` in any of these constructors can be bundled together
as a [`SaddleVector`](@ref), e.g. `SaddleSystem(A,B₂,B₁ᵀ,SaddleVector(u,f))`.

An optional keyword argument `solver=` can be used to specify the type of
solution for the Schur complement system. By default, this is set to `Direct`,
and the Schur complement matrix is formed, factorized, and stored. This can be
changed to `Iterative`, in which case an iterative solver is determined
by the `lin_solve` function of [`KrylovKit.jl`](https://github.com/Jutho/KrylovKit.jl).
"""
function SaddleSystem(A::LinearMap{T},B₂::LinearMap{T},B₁ᵀ::LinearMap{T},C::LinearMap{T},
                      A⁻¹::LinearMap{T},P::LinearMap{T},TU,TF;solver::Type{TS}=Direct,kwargs...) where {T,TS<:SchurSolverType}

    ns, nc = _check_sizes(A,B₂,B₁ᵀ,C,P)

    S = C - B₂*A⁻¹*B₁ᵀ

    S⁻¹ = _schur_inverse_function(S,T,nc,solver,kwargs...)

    return SaddleSystem{T,ns,nc,TU,TF,solver}(A,B₂,B₁ᵀ,C,A⁻¹,zeros(T,ns),zeros(T,nc),zeros(T,nc),P,S,S⁻¹)
end

##### Schur complement solver functions #####

abstract type Direct <: SchurSolverType end

function _schur_inverse_function(S,T,M,::Type{Direct},kwargs...)
  Sfact = factorize(Matrix(S))
  return LinearMap{T}(x -> Sfact\x,M)
end

abstract type Iterative <: SchurSolverType end

function _schur_inverse_function(S,T,M,::Type{Iterative},kwargs...)
  return LinearMap{T}(x -> ((y,_) = linsolve(S,x); y) ,M)
end

macro createsolver(stype)
  sroutine = Symbol(lowercase(string(stype)),"!")

  return esc(quote
          export $stype
          abstract type $stype <: SchurSolverType end
          function _schur_inverse_function(S,T,M,::Type{$stype},kwargs...)
            return LinearMap{T}(x -> (y = deepcopy(x); $sroutine(y,S,x;kwargs...); return y),M)
          end
        end)
end

@createsolver CG
@createsolver BiCGStabl
@createsolver GMRES
@createsolver MINRES
@createsolver IDRS

###########

### OTHER CONSTRUCTORS

### Matrix operators
function SaddleSystem(A::AbstractMatrix{T},B₂::AbstractMatrix{T},B₁ᵀ::AbstractMatrix{T},
                      C::AbstractMatrix{T};solver::Type{TS}=Direct,filter=I,kwargs...) where {T,TS<:SchurSolverType}

    Afact = factorize(A)
    Ainv = LinearMap{T}((y,x) -> y .= Afact\x,size(A,1))

    return SaddleSystem(LinearMap{T}(A),LinearMap{T}(B₂),LinearMap{T}(B₁ᵀ),
                        LinearMap{T}(C),Ainv,linear_map(filter,zeros(T,size(C,1)),eltype=T),
                        Vector{T},Vector{T};solver=solver,kwargs...)
end

# For cases in which C is zero, no need to pass along the argument
SaddleSystem(A::AbstractMatrix{T},B₂::AbstractMatrix{T},B₁ᵀ::AbstractMatrix{T};solver::Type{TS}=Direct,filter=I,kwargs...) where {T,TS<:SchurSolverType} =
        SaddleSystem(A,B₂,B₁ᵀ,zeros(T,size(B₂,1),size(B₁ᵀ,2));solver=solver,filter=filter,kwargs...)

### Operators are functions or function-like operators
# This version should take in functions or function-like objects that act upon given
# data types u and f. Should transform them into operators that act on abstract vectors
# of the same size
# There should already be an \ operator associated with A
# NOTE: should change default value of eltype to eltype(u)
function SaddleSystem(A,B₂,B₁ᵀ,C,u::TU,f::TF;eltype=Float64,filter=I,solver::Type{TS}=Direct,kwargs...) where {TU,TF,TS<:SchurSolverType}
    return SaddleSystem(linear_map(A,u,eltype=eltype),linear_map(B₂,u,f,eltype=eltype),
                        linear_map(B₁ᵀ,f,u,eltype=eltype),
                        linear_map(C,f,eltype=eltype),
                        linear_inverse_map!(A,u,eltype=eltype),
                        linear_map(filter,f,eltype=eltype),TU,TF;solver=solver,kwargs...)
end

# No C operator provided, so set it to zero
SaddleSystem(A,B₂,B₁ᵀ,u::TU,f::TF;
            eltype=Float64,filter=I,solver::Type{TS}=Direct,kwargs...) where {TU,TF,TS<:SchurSolverType} =
            SaddleSystem(A,B₂,B₁ᵀ,C_zero(f,eltype),u,f;eltype=eltype,filter=filter,solver=solver,kwargs...)

# Unconstrained system
SaddleSystem(A,u::TU;eltype=Float64,filter=I,solver::Type{TS}=Direct,kwargs...) where {TU,TS<:SchurSolverType} = SaddleSystem(A,nothing,nothing,u,Type{eltype}[];eltype=eltype,filter=filter,solver=solver,kwargs...)

### For handling ArrayPartition arguments for the solution/rhs
SaddleSystem(A,B₂,B₁ᵀ,C,v::ArrayPartition;
            eltype=Float64,filter=I,solver::Type{TS}=Direct,kwargs...) where {TS<:SchurSolverType} =
            SaddleSystem(A,B₂,B₁ᵀ,C,v.x[1],v.x[2];eltype=eltype,filter=filter,solver=solver,kwargs...)

# No C operator
SaddleSystem(A,B₂,B₁ᵀ,v::ArrayPartition;
            eltype=Float64,filter=I,solver::Type{TS}=Direct,kwargs...) where {TS<:SchurSolverType} =
            SaddleSystem(A,B₂,B₁ᵀ,v.x[1],v.x[2];eltype=eltype,filter=filter,solver=solver,kwargs...)

# Unconstrained system
SaddleSystem(A,v::ArrayPartition;
            eltype=Float64,filter=I,solver::Type{TS}=Direct,kwargs...) where {TS<:SchurSolverType} =
            SaddleSystem(A,v.x[1];eltype=eltype,filter=filter,solver=solver,kwargs...)

### AUXILIARY ROUTINES

function Base.show(io::IO, S::SaddleSystem{T,Ns,Nc,TU,TF,TS}) where {T,Ns,Nc,TU,TF,TS<:SchurSolverType}
     println(io, "Saddle system with $Ns states and $Nc constraints and")
     println(io, "   State vector of type $TU")
     println(io, "   Constraint vector of type $TF")
     println(io, "   Elements of type $T")
     println(io, "using a $TS solver")
end

"""
    Base.size(::SaddleSystem)

Report the size of a [`SaddleSystem`](@ref).
"""
size(::SaddleSystem{T,Ns,Nc}) where {T,Ns,Nc} = (Ns+Nc,Ns+Nc)

"""
    Base.eltype(::SaddleSystem)

Report the element type of a [`SaddleSystem`](@ref).
"""
eltype(::SaddleSystem{T,Ns,Nc}) where {T,Ns,Nc} = T

C_zero(f,eltype) = zeros(eltype,length(f),length(f))

function _check_sizes(A,B₂,B₁ᵀ,C,P)
    mA, nA = size(A)
    mB1, nB1 = size(B₁ᵀ)
    mB2, nB2 = size(B₂)
    mC, nC = size(C)
    mP, nP = size(P)

    # check compatibility of sizes
    mA == nA  || error("A is not square")
    mA == mB1 || error("Incompatible number of rows in A and B₁ᵀ")
    nA == nB2 || error("Incompatible number of columns in A and B₂")
    mC == mB2 || error("Incompatible number of rows in C and B₂")
    nC == nB1 || error("Incompatible number of columns in C and B₁ᵀ")
    mP == nP == mC  || error("Filter has incompatible dimensions")

    ns = nA
    nc = nB1

    return ns, nc
end


text:
function lis(arr::Vector)
    if length(arr) == 0 return copy(arr) end
    L = Vector{typeof(arr)}(length(arr))
    L[1] = [arr[1]]

    for i in 2:length(arr)
        nextL = []
        for j in 1:i
            if arr[j] < arr[i] && length(L[j]) ≥ length(nextL)
                nextL = L[j]
            end
        end
        L[i] = vcat(nextL, arr[i])
    end

    return L[indmax(length.(L))]
end

@show lis([3, 2, 6, 4, 5, 1])
@show lis([0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15])


text:
# Following the algorithm by William and Seeger, 2001
# Cs is equivalent to X_mm and C to X_mn

function sampleindex(X::AbstractMatrix, r::Real; obsdim::Integer=defaultobs)
    0 < r <= 1 || throw(ArgumentError("Sample rate `r` must be in range (0,1]"))
    n = size(X, obsdim)
    m = ceil(Int, n*r)
    S = StatsBase.sample(1:n, m; replace=false, ordered=true)
    return S
end

function nystrom_sample(k::Kernel, X::AbstractMatrix, S::Vector{<:Integer}; obsdim::Integer=defaultobs)
    obsdim ∈ [1, 2] || throw(ArgumentError("`obsdim` should be 1 or 2 (see docs of kernelmatrix))"))
    Xₘ = obsdim == 1 ? X[S, :] : X[:, S]
    C = kernelmatrix(k, Xₘ, X; obsdim=obsdim)
    Cs = C[:, S]
    return (C, Cs)
end

function nystrom_pinv!(Cs::Matrix{T}, tol::T=eps(T)*size(Cs,1)) where {T<:Real}
    # Compute eigendecomposition of sampled component of K
    QΛQᵀ = LinearAlgebra.eigen!(LinearAlgebra.Symmetric(Cs))

    # Solve for D = Λ^(-1/2) (pseudo inverse - use tolerance from before factorization)
    D = QΛQᵀ.values
    λ_tol = maximum(D)*tol

    for i in eachindex(D)
        @inbounds D[i] = abs(D[i]) <= λ_tol ? zero(T) : one(T)/sqrt(D[i])
    end

    # Scale eigenvectors by D
    Q = QΛQᵀ.vectors
    QD = LinearAlgebra.rmul!(Q, LinearAlgebra.Diagonal(D))  # Scales column i of Q by D[i]

    # W := (QD)(QD)ᵀ = (QΛQᵀ)^(-1)  (pseudo inverse)
    W = QD*QD'

    # Symmetrize W
    return LinearAlgebra.copytri!(W, 'U')
end

@doc raw"""
    NystromFact

Type for storing a Nystrom factorization. The factorization contains two fields: `W` and
`C`, two matrices satisfying:
```math
\mathbf{K} \approx \mathbf{C}^{\intercal}\mathbf{W}\mathbf{C}
```
"""
struct NystromFact{T<:Real}
    W::Matrix{T}
    C::Matrix{T}
end

function NystromFact(W::Matrix{<:Real}, C::Matrix{<:Real})
    T = Base.promote_eltypeof(W, C)
    return NystromFact(convert(Matrix{T}, W), convert(Matrix{T}, C))
end

@doc raw"""
    nystrom(k::Kernel, X::Matrix, S::Vector; obsdim::Int=defaultobs)

Computes a factorization of Nystrom approximation of the square kernel matrix of data
matrix `X` with respect to kernel `k`. Returns a `NystromFact` struct which stores a
Nystrom factorization satisfying:
```math
\mathbf{K} \approx \mathbf{C}^{\intercal}\mathbf{W}\mathbf{C}
```
"""
function nystrom(k::Kernel, X::AbstractMatrix, S::Vector{<:Integer}; obsdim::Int=defaultobs)
    C, Cs = nystrom_sample(k, X, S; obsdim=obsdim)
    W = nystrom_pinv!(Cs)
    return NystromFact(W, C)
end

@doc raw"""
    nystrom(k::Kernel, X::Matrix, r::Real; obsdim::Int=defaultobs)

Computes a factorization of Nystrom approximation of the square kernel matrix of data
matrix `X` with respect to kernel `k` using a sample ratio of `r`.
Returns a `NystromFact` struct which stores a Nystrom factorization satisfying:
```math
\mathbf{K} \approx \mathbf{C}^{\intercal}\mathbf{W}\mathbf{C}
```
"""
function nystrom(k::Kernel, X::AbstractMatrix, r::Real; obsdim::Int=defaultobs)
    S = sampleindex(X, r; obsdim=obsdim)
    return nystrom(k, X, S; obsdim=obsdim)
end

"""
    nystrom(CᵀWC::NystromFact)

Compute the approximate kernel matrix based on the Nystrom factorization.
"""
function kernelmatrix(CᵀWC::NystromFact{<:Real})
    W = CᵀWC.W
    C = CᵀWC.C
    return C'*W*C
end


text:
usePrecaching = true
push!(LOAD_PATH, pwd())
include("$(pwd())/startup.jl")

using Distributed
import Statistics
import .TOFFunctions

function averageAllSpectra(
  filepath;
  filefilterRegexp=r"\.h5$",
  filePrecaching = true, # Precache whole next file. Only use when you have enough RAM and no SWAP has to be used.
  openWholeFile = true
  )

    allFiles = readdir(filepath)
    files = filter(s->occursin(filefilterRegexp, s), allFiles)
    nFiles = size(files,1)

    validFiles, timeSortIndices = TOFFunctions.validateHDF5Files(filepath, files)

    if !issorted(timeSortIndices)
    println("Reordering Files according to start acquisition time!")
    else
    println("File order seems fine, continuing...")
    end
    files = validFiles[timeSortIndices]
    nFiles = size(files,1)

    averagedSpectra = Array{Any,1}()

    @time for j=1:5
        println("working on $filepath")
        totalPath = joinpath(filepath, files[j])
        if j < nFiles && filePrecaching
            totalPrecachePath = joinpath(filepath, files[j+1])
        else
            totalPrecachePath = ""
        end

        totalSubSpectra = TOFFunctions.getSubSpectraCount(totalPath)
        sumspec = TOFFunctions.getSubSpectrumFromFile(totalPath,1, preloadFile=totalPrecachePath, openWholeFile = openWholeFile)
        for i=2:totalSubSpectra
            sumspec = sumspec .+ TOFFunctions.getSubSpectrumFromFile(totalPath,i, preloadFile=totalPrecachePath, openWholeFile = openWholeFile)
        end
        push!(averagedSpectra, sumspec./totalSubSpectra)
    end
    return averagedSpectra
end


text:
# demultiplex_test.jl
@testset "demultiplex test" begin
    @test reverse_complement("ATCGT") == "ACGAT"
    @test reverse_complement("ATCGTG") == "CACGAT"
    @test reverse_complement("TGCCGTGACGT") == "ACGTCACGGCA"
    @test reverse_complement("ATCG") == "CGAT"
    @test potential_mismatches("GCGT", 1) == Any["GCGT",
                                                 "CCGT",
                                                 "ACGT",
                                                 "TCGT",
                                                 "GGGT",
                                                 "GAGT",
                                                 "GTGT",
                                                 "GCCT",
                                                 "GCAT",
                                                 "GCTT",
                                                 "GCGG",
                                                 "GCGC",
                                                 "GCGA",
                                                 "CGT",
                                                 "GGT",
                                                 "GCT",
                                                 "GCG"]
end


text:
using JuMP
using Gurobi
using Suppressor
using XLSX
using DataFrames
gurobi_env = @suppress Gurobi.Env()


text:
using BlackBoxOptim
using Test

const NumTestRepetitions = 100

random_candidate(n, lo, hi) =
    BlackBoxOptim.Candidate{Float64}([clamp(randn(), lo, hi) for _ in 1:n])


text:
@sumProductRule(:node_type     => Dirichlet,
                :outbound_type => Message{Dirichlet},
                :inbound_types => (Nothing, Message{PointMass}),
                :name          => SPDirichletOutNP)

@naiveVariationalRule(:node_type     => Dirichlet,
                      :outbound_type => Message{Dirichlet},
                      :inbound_types => (Nothing, ProbabilityDistribution),
                      :name          => VBDirichletOut)

@naiveVariationalRule(:node_type     => Dirichlet,
                      :outbound_type => Message{Function},
                      :inbound_types => (ProbabilityDistribution, Nothing),
                      :name          => VBDirichletIn1)


text:
using FastRationals
if VERSION >= v"0.7.0-"
    using Test
else
    using Base.Test
end

using BenchmarkTools

macro fastest_time(this)
   :((@benchmark $this).times[1])
end

# using Int32s

numer_a, denom_a = Int32(5), Int32(77)
numer_b, denom_b = Int32(100), Int32(17)
a = Rational(numer_a, denom_a); b = Rational(numer_b, denom_b);
afast = FastRational(numer_a, denom_a); bfast = FastRational(numer_b, denom_b);

@test a + b == afast + bfast
@test a - b == afast - bfast
@test a * b == afast * bfast
@test a // b == afast // bfast

@test fld( (@fastest_time $a + $b), (@fastest_time $afast + $bfast) ) >= 1
@test fld( (@fastest_time $a * $b), (@fastest_time $afast * $bfast) ) >= 1

# using Int64s

numer_a, denom_a = Int64(5), Int64(77)
numer_b, denom_b = Int64(100), Int64(17)
a = Rational(numer_a, denom_a); b = Rational(numer_b, denom_b);
afast = FastRational(numer_a, denom_a); bfast = FastRational(numer_b, denom_b);

@test a + b == afast + bfast
@test a - b == afast - bfast
@test a * b == afast * bfast
@test a // b == afast // bfast

@test fld( (@fastest_time $a + $b), (@fastest_time $afast + $bfast) ) >= 1
@test fld( (@fastest_time $a * $b), (@fastest_time $afast * $bfast) ) >= 1

# using Int128s

numer_a, denom_a = Int128(5), Int128(77)
numer_b, denom_b = Int128(100), Int128(17)
a = Rational(numer_a, denom_a); b = Rational(numer_b, denom_b);
afast = FastRational(numer_a, denom_a); bfast = FastRational(numer_b, denom_b);

@test a + b == afast + bfast
@test a - b == afast - bfast
@test a * b == afast * bfast
@test a // b == afast // bfast

@test fld( (@fastest_time $a + $b), (@fastest_time $afast + $bfast) ) >= 1
@test fld( (@fastest_time $a * $b), (@fastest_time $afast * $bfast) ) >= 1

# using BigInt

numer_a, denom_a = BigInt(5), BigInt(77)
numer_b, denom_b = BigInt(100), BigInt(17)
a = Rational(numer_a, denom_a); b = Rational(numer_b, denom_b);
afast = FastRational(numer_a, denom_a); bfast = FastRational(numer_b, denom_b);

@test a + b == afast + bfast
@test a - b == afast - bfast
@test a * b == afast * bfast
@test a // b == afast // bfast

@test fld( (@fastest_time $a + $b), (@fastest_time $afast + $bfast) ) >= 1
@test fld( (@fastest_time $a * $b), (@fastest_time $afast * $bfast) ) >= 1


text:
using QuBase
using QuDynamics
using BenchmarkLite

type SampleProc{Alg} <: Proc end

# Quantum Harmonic Oscillator type
# The type construct takes in the parameters which are used
# for constructing the various configuration used for benchmarking.
type HarmonicOscillator
    N::Int
    lambda
    tlist
end

# Name of the solver, used in the creating the table of timing analysis
Base.string{Alg}(::SampleProc{Alg}) = lowercase(string(Alg))

# Constructing a valid configuration, this is set to `true` as we have
# set the type of the required fields in the construct of `HarmonicOscillator` itself.
Base.isvalid{Alg}(p::SampleProc{Alg}, cfg::HarmonicOscillator) = true

# This returns the `hamitlonian`, initial state`, `observable whose expectation
# value is to be calculated` taking in the `cfg` configuration of the system.
function operator_set(p::SampleProc, cfg::HarmonicOscillator)
    a = lowerop(cfg.N)
    hamiltonian = a'*a + cfg.lambda*(a + a')
    init_state = complex(statevec(1,FiniteBasis(cfg.N)))
    return hamiltonian, init_state, a+a'
end

# Start states are pre-allocated, that is they are constructed before the evaluation
# here we return the `hamiltonian`, `initial state`, `observable whose expectation
# is to be calculated` using the above function i.e., operator_set. The first two remain fixed,
# the third can be varied as required.
Base.start(p::SampleProc, cfg::HarmonicOscillator) = (operator_set(p, cfg))
Base.length(p::SampleProc, cfg::HarmonicOscillator) = cfg.N
function Base.run{Alg}(p::SampleProc{Alg}, cfg::HarmonicOscillator, s::(QuBase.QuArray, QuBase.QuArray, QuBase.QuArray))
    qprop = QuPropagator(s[1], s[2], cfg.tlist, Alg())
    for (t, psi) in qprop
        psi'*s[3]*psi
    end
end

Base.done(p::SampleProc, cfg, s) = nothing

# Creating the array of methods on which benchmarks have to be performed
procs = Proc[ SampleProc{QuExpmV}(),
              SampleProc{QuExpokit}(),
              SampleProc{QuODE45}(),
              SampleProc{QuODE78}(),
              SampleProc{QuODE23s}(),
              SampleProc{QuKrylov}(),
              SampleProc{QuCrankNicolson}(),
              SampleProc{QuEuler}()]

# Setting up various configurations.
cfgs = HarmonicOscillator[HarmonicOscillator(2, 0.3, 0.:0.1:2*pi),
        HarmonicOscillator(3, 0.3, 0.:0.1:2*pi),
        HarmonicOscillator(6, 0.3, 0.:0.1:2*pi),
        HarmonicOscillator(10, 0.3, 0.:0.1:2*pi),
        HarmonicOscillator(12, 0.3, 0.:0.1:2*pi)]

# Generating the timing analysis.
rtable = run(procs, cfgs)

# Table display from the above method.
show(rtable; unit=:usec)


text:
using BenchmarkTools, AutoARIMA

const SUITE = BenchmarkGroup()

z = cos.(1:10000) .* log.(1:10000) .* randn()

SUITE["fit"] = BenchmarkGroup()

SUITE["fit"]["ar"] = BenchmarkGroup()
for f in (yule_walker, levinson_durbin, least_squares)
    for p in (1, 10)
        SUITE["fit"]["ar"][string(f),p] = @benchmarkable $f($z, $p)
    end
end

SUITE["fit"]["ma"] = BenchmarkGroup()
for f in (innovations,)
    for q in (1, 10)
        SUITE["fit"]["ma"][string(f),q] = @benchmarkable $f($z, $q, n=20)
    end
end

SUITE["fit"]["arma"] = BenchmarkGroup()
for f in (hannan_rissanen,)
    for p in (1, 10), q in (1, 10)
        SUITE["fit"]["arma"][string(f),p,q] = @benchmarkable $f($z, $p, $q)
    end
end

SUITE["forecast"] = BenchmarkGroup()

SUITE["forecast"]["ar"] = BenchmarkGroup()
for p in (1, 10, 100)
    ar = fit(ARParams(p), z)
    SUITE["forecast"]["ar"][p] = @benchmarkable forecast($ar, $z)
end

SUITE["forecast"]["ma"] = BenchmarkGroup()
for q in (1, 10, 100)
    ma = fit(MAParams(q), z, n=200)
    SUITE["forecast"]["ma"][q] = @benchmarkable forecast($ma, $z)
end

SUITE["forecast"]["arma"] = BenchmarkGroup()
for p in (1, 10, 100), q in (1, 10, 100)
    arma = fit(ARMAParams(p, q), z)
    SUITE["forecast"]["arma"][p,q] = @benchmarkable forecast($arma, $z)
end