export Experiment

using Dates
using ReinforcementLearningBase
using ReinforcementLearningCore
using .ReinforcementLearningEnvironments
using Flux
using BSON
using TensorBoardLogger
using Logging
using Random

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:BasicDQN},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_BasicDQN_CartPole_$(t)")
    end

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

    env = CartPoleEnv(; T = Float32, rng = rng)
    ns, na = length(get_state(env)), length(get_actions(env))
    agent = Agent(
        policy = QBasedPolicy(
            learner = BasicDQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, na; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                batch_size = 32,
                min_replay_history = 100,
                loss_func = huber_loss,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                rng = rng,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(10000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end]
            end
        end,
        DoEveryNStep(10000) do t, agent, env
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses three dense layers to approximate the Q value.
    The testing environment is CartPoleEnv.

    Agent and statistic info will be saved to: `$save_dir`
    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("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

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

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:DQN},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_DQN_CartPole_$(t)")
    end

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

    env = CartPoleEnv(; T = Float32, rng = rng)
    ns, na = length(get_state(env)), length(get_actions(env))

    agent = Agent(
        policy = QBasedPolicy(
            learner = DQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, na; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                target_approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, na; initW = glorot_uniform(rng)),
                    ) |> cpu,
                ),
                loss_func = huber_loss,
                stack_size = nothing,
                batch_size = 32,
                update_horizon = 1,
                min_replay_history = 100,
                update_freq = 1,
                target_update_freq = 100,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                rng = rng,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(10_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env
            if agent.policy.learner.update_step % agent.policy.learner.update_freq == 0
                with_logger(lg) do
                    @info "training" loss = agent.policy.learner.loss
                end
            end
        end,
        DoEveryNStep(10000) do t, agent, env
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses the `DQNLearner` method with three dense layers to approximate the Q value.
    The testing environment is CartPoleEnv.

    Agent and statistic info will be saved to: `$save_dir`
    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("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

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

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:PrioritizedDQN},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_PrioritizedDQN_CartPole_$(t)")
    end

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

    env = CartPoleEnv(; T = Float32, rng = rng)
    ns, na = length(get_state(env)), length(get_actions(env))

    agent = Agent(
        policy = QBasedPolicy(
            learner = PrioritizedDQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, na; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                target_approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, na; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                loss_func = huber_loss_unreduced,
                stack_size = nothing,
                batch_size = 32,
                update_horizon = 1,
                min_replay_history = 100,
                update_freq = 1,
                target_update_freq = 100,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                rng = rng,
            ),
        ),
        trajectory = CircularCompactPSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(10_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env
            if agent.policy.learner.update_step % agent.policy.learner.update_freq == 0
                with_logger(lg) do
                    @info "training" loss = agent.policy.learner.loss
                end
            end
        end,
        DoEveryNStep(10000) do t, agent, env
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses the `PrioritizedDQNLearner` method with three dense layers to approximate the Q value.
    The testing environment is CartPoleEnv.

    Agent and statistic info will be saved to: `$save_dir`
    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("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

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

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:Rainbow},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_Rainbow_CartPole_$(t)")
    end

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

    env = CartPoleEnv(; T = Float32, rng = rng)
    ns, na = length(get_state(env)), length(get_actions(env))

    n_atoms = 51
    agent = Agent(
        policy = QBasedPolicy(
            learner = RainbowLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, na * n_atoms; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(0.0005),
                ),
                target_approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, 128, relu; initW = glorot_uniform(rng)),
                        Dense(128, na * n_atoms; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(0.0005),
                ),
                n_actions = na,
                n_atoms = n_atoms,
                Vₘₐₓ = 200.0f0,
                Vₘᵢₙ = 0.0f0,
                update_freq = 1,
                γ = 0.99f0,
                update_horizon = 1,
                batch_size = 32,
                stack_size = nothing,
                min_replay_history = 100,
                loss_func = logitcrossentropy_unreduced,
                target_update_freq = 100,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                rng = rng,
            ),
        ),
        trajectory = CircularCompactPSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(10_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env
            if agent.policy.learner.update_step % agent.policy.learner.update_freq == 0
                with_logger(lg) do
                    @info "training" loss = agent.policy.learner.loss
                end
            end
        end,
        DoEveryNStep(10000) do t, agent, env
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses the `RainbowLearner` method with three dense layers to approximate the distributed Q value.
    The testing environment is CartPoleEnv.

    Agent and statistic info will be saved to: `$save_dir`
    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("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

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

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:IQN},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_IQN_CartPole_$(t)")
    end

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

    env = CartPoleEnv(; T = Float32, rng = rng)
    ns, na = length(get_state(env)), length(get_actions(env))

    init = glorot_uniform(rng)
    Nₑₘ = 16
    n_hidden = 64
    κ = 1.0f0

    nn_creator() =
        ImplicitQuantileNet(
            ψ = Dense(ns, n_hidden, relu; initW = init),
            ϕ = Dense(Nₑₘ, n_hidden, relu; initW = init),
            header = Dense(n_hidden, na; initW = init),
        ) |> cpu

    agent = Agent(
        policy = QBasedPolicy(
            learner = IQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = nn_creator(),
                    optimizer = ADAM(0.001),
                ),
                target_approximator = NeuralNetworkApproximator(model = nn_creator()),
                κ = κ,
                N = 8,
                N′ = 8,
                Nₑₘ = Nₑₘ,
                K = 32,
                γ = 0.99f0,
                stack_size = nothing,
                batch_size = 32,
                update_horizon = 1,
                min_replay_history = 100,
                update_freq = 1,
                target_update_freq = 100,
                default_priority = 1.0f2,
                rng = rng,
                device_rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                rng = rng,
            ),
        ),
        trajectory = CircularCompactPSARTSATrajectory(
            capacity = 1000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(10_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env
            if agent.policy.learner.update_step % agent.policy.learner.update_freq == 0
                with_logger(lg) do
                    @info "training" loss = agent.policy.learner.loss
                end
            end
        end,
        DoEveryNStep(10000) do t, agent, env
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses the `IQNLearner` method with a `ImplicitQuantileNet`.
    The testing environment is CartPoleEnv.

    Agent and statistic info will be saved to: `$save_dir`
    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("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

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

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:A2C},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    rng = MersenneTwister(seed)
    N_ENV = 16
    UPDATE_FREQ = 10
    env = MultiThreadEnv([
        CartPoleEnv(; T = Float32, rng = MersenneTwister(hash(seed + i))) for i in 1:N_ENV
    ])
    ns, na = length(get_state(env[1])), length(get_actions(env[1]))
    RLBase.reset!(env, is_force = true)
    agent = Agent(
        policy = QBasedPolicy(
            learner = A2CLearner(
                approximator = ActorCritic(
                    actor = Chain(
                        Dense(ns, 256, relu; initW = glorot_uniform(rng)),
                        Dense(256, na; initW = glorot_uniform(rng)),
                    ),
                    critic = Chain(
                        Dense(ns, 256, relu; initW = glorot_uniform(rng)),
                        Dense(256, 1; initW = glorot_uniform(rng)),
                    ),
                    optimizer = ADAM(1e-3),
                ) |> cpu,
                γ = 0.99f0,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.5f0,
                entropy_loss_weight = 0.001f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer()),#= seed = nothing =#
        ),
        trajectory = CircularCompactSARTSATrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (ns, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )

    hook = TotalBatchRewardPerEpisode(N_ENV)
    stop_condition = StopAfterStep(100000)
    Experiment(agent, env, stop_condition, hook, "# A2C with CartPole")
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:A2CGAE},
    ::Val{:CartPole},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    rng = MersenneTwister(seed)
    N_ENV = 16
    UPDATE_FREQ = 10
    env = MultiThreadEnv([
        CartPoleEnv(; T = Float32, rng = MersenneTwister(hash(seed + i))) for i in 1:N_ENV
    ])
    ns, na = length(get_state(env[1])), length(get_actions(env[1]))
    RLBase.reset!(env, is_force = true)
    agent = Agent(
        policy = QBasedPolicy(
            learner = A2CGAELearner(
                approximator = ActorCritic(
                    actor = NeuralNetworkApproximator(
                        model = Chain(
                            Dense(ns, 256, relu; initW = glorot_uniform(rng)),
                            Dense(256, na; initW = glorot_uniform(rng)),
                        ),
                        optimizer = ADAM(1e-3),
                    ),
                    critic = NeuralNetworkApproximator(
                        model = Chain(
                            Dense(ns, 256, relu; initW = glorot_uniform(rng)),
                            Dense(256, 1; initW = glorot_uniform(rng)),
                        ),
                        optimizer = ADAM(1e-3),
                    ),
                ) |> cpu,
                γ = 0.99f0,
                λ = 0.97f0,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.5f0,
                entropy_loss_weight = 0.001f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(;)), #= seed = nothing =#
        ),
        trajectory = CircularCompactSARTSATrajectory(;
            capacity = UPDATE_FREQ,
            state_type = Float32,
            state_size = (ns, N_ENV),
            action_type = Int,
            action_size = (N_ENV,),
            reward_type = Float32,
            reward_size = (N_ENV,),
            terminal_type = Bool,
            terminal_size = (N_ENV,),
        ),
    )
    Experiment(
        agent,
        env,
        StopAfterStep(100000),
        TotalBatchRewardPerEpisode(N_ENV),
        "# A2CGAE with CartPole",
    )
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:DDPG},
    ::Val{:Pendulum},
    ::Nothing;
    seed = 123,
)
    rng = MersenneTwister(seed)
    inner_env = PendulumEnv(T = Float32, rng = rng)
    action_space = get_actions(inner_env)
    low = action_space.low
    high = action_space.high
    ns = length(get_state(inner_env))

    env = inner_env |> ActionTransformedEnv(x -> low + (x + 1) * 0.5 * (high - low))
    init = glorot_uniform(rng)

    create_actor() = Chain(
        Dense(ns, 30, relu; initW = init),
        Dense(30, 30, relu; initW = init),
        Dense(30, 1, tanh; initW = init),
    )

    create_critic() = Chain(
        Dense(ns + 1, 30, relu; initW = init),
        Dense(30, 30, relu; initW = init),
        Dense(30, 1; initW = init),
    )

    agent = Agent(
        policy = DDPGPolicy(
            behavior_actor = NeuralNetworkApproximator(
                model = create_actor(),
                optimizer = ADAM(),
            ),
            behavior_critic = NeuralNetworkApproximator(
                model = create_critic(),
                optimizer = ADAM(),
            ),
            target_actor = NeuralNetworkApproximator(
                model = create_actor(),
                optimizer = ADAM(),
            ),
            target_critic = NeuralNetworkApproximator(
                model = create_critic(),
                optimizer = ADAM(),
            ),
            γ = 0.99f0,
            ρ = 0.995f0,
            batch_size = 64,
            start_steps = 1000,
            start_policy = RandomPolicy(ContinuousSpace(-1.0, 1.0); rng = rng),
            update_after = 1000,
            update_every = 1,
            act_limit = 1.0,
            act_noise = 0.1,
            rng = rng,
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 10000,
            state_type = Float32,
            state_size = (ns,),
            action_type = Float32,
        ),
    )

    description = """
    # Play Pendulum with DDPG
    """

    Experiment(agent, env, StopAfterStep(10000), TotalRewardPerEpisode(), description)
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:PPO},
    ::Val{:CartPole},
    ::Nothing;
    seed = 123,
)
    rng = MersenneTwister(seed)
    N_ENV = 8
    UPDATE_FREQ = 16
    env = MultiThreadEnv([
        CartPoleEnv(; T = Float32, rng = MersenneTwister(hash(seed + i))) for i in 1:N_ENV
    ])
    ns, na = length(get_state(env[1])), length(get_actions(env[1]))
    RLBase.reset!(env, is_force = true)
    agent = Agent(
        policy = QBasedPolicy(
            learner = PPOLearner(
                approximator = ActorCritic(
                    actor = NeuralNetworkApproximator(
                        model = Chain(
                            Dense(ns, 256, relu; initW = glorot_uniform(rng)),
                            Dense(256, na; initW = glorot_uniform(rng)),
                        ),
                        optimizer = ADAM(1e-3),
                    ),
                    critic = NeuralNetworkApproximator(
                        model = Chain(
                            Dense(ns, 256, relu; initW = glorot_uniform(rng)),
                            Dense(256, 1; initW = glorot_uniform(rng)),
                        ),
                        optimizer = ADAM(1e-3),
                    ),
                ) |> cpu,
                γ = 0.99f0,
                λ = 0.95f0,
                clip_range = 0.1f0,
                max_grad_norm = 0.5f0,
                n_epochs = 4,
                n_microbatches = 4,
                actor_loss_weight = 1.0f0,
                critic_loss_weight = 0.5f0,
                entropy_loss_weight = 0.001f0,
            ),
            explorer = BatchExplorer(GumbelSoftmaxExplorer(; rng = rng)),
        ),
        trajectory = PPOTrajectory(;
            capacity = 32,
            state_type = Float32,
            state_size = (ns, 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,),
        ),
    )
    Experiment(
        agent,
        env,
        StopAfterStep(100000),
        TotalBatchRewardPerEpisode(N_ENV),
        "# PPO with CartPole",
    )
end

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:BasicDQN},
    ::Val{:MountainCar},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    rng = MersenneTwister(seed)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_BasicDQN_MountainCar_$(t)")
    end

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

    env = MountainCarEnv(; T = Float32, max_steps = 5000, rng = rng)
    ns, na = length(get_state(env)), length(get_actions(env))
    agent = Agent(
        policy = QBasedPolicy(
            learner = BasicDQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 64, relu; initW = glorot_uniform(rng)),
                        Dense(64, 64, relu; initW = glorot_uniform(rng)),
                        Dense(64, na; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                batch_size = 32,
                min_replay_history = 100,
                loss_func = huber_loss,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                rng = rng,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 50000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(40000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env
            with_logger(lg) do
                @info "training" loss = agent.policy.learner.loss
            end
        end,
        DoEveryNEpisode() do t, agent, env
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end]
            end
        end,
        DoEveryNStep(10000) do t, agent, env
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses three dense layers to approximate the Q value.
    The testing environment is MountainCarEnv.

    Agent and statistic info will be saved to: `$save_dir`
    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("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

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

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:DQN},
    ::Val{:MountainCar},
    ::Nothing;
    save_dir = nothing,
    seed = 123,
)
    if isnothing(save_dir)
        t = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
        save_dir = joinpath(pwd(), "checkpoints", "JuliaRL_DQN_MountainCar_$(t)")
    end

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

    env = MountainCarEnv(; T = Float32, max_steps = 5000, rng = rng)
    ns, na = length(get_state(env)), length(get_actions(env))

    agent = Agent(
        policy = QBasedPolicy(
            learner = DQNLearner(
                approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 64, relu; initW = glorot_uniform(rng)),
                        Dense(64, 64, relu; initW = glorot_uniform(rng)),
                        Dense(64, na; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                target_approximator = NeuralNetworkApproximator(
                    model = Chain(
                        Dense(ns, 64, relu; initW = glorot_uniform(rng)),
                        Dense(64, 64, relu; initW = glorot_uniform(rng)),
                        Dense(64, na; initW = glorot_uniform(rng)),
                    ) |> cpu,
                    optimizer = ADAM(),
                ),
                loss_func = huber_loss,
                stack_size = nothing,
                batch_size = 32,
                update_horizon = 1,
                min_replay_history = 100,
                update_freq = 1,
                target_update_freq = 100,
                rng = rng,
            ),
            explorer = EpsilonGreedyExplorer(
                kind = :exp,
                ϵ_stable = 0.01,
                decay_steps = 500,
                rng = rng,
            ),
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 50000,
            state_type = Float32,
            state_size = (ns,),
        ),
    )

    stop_condition = StopAfterStep(40_000)

    total_reward_per_episode = TotalRewardPerEpisode()
    time_per_step = TimePerStep()
    hook = ComposedHook(
        total_reward_per_episode,
        time_per_step,
        DoEveryNStep() do t, agent, env
            if agent.policy.learner.update_step % agent.policy.learner.update_freq == 0
                with_logger(lg) do
                    @info "training" loss = agent.policy.learner.loss
                end
            end
        end,
        DoEveryNEpisode() do t, agent, env
            with_logger(lg) do
                @info "training" reward = total_reward_per_episode.rewards[end]
            end
        end,
        DoEveryNStep(10000) do t, agent, env
            RLCore.save(save_dir, agent)
            BSON.@save joinpath(save_dir, "stats.bson") total_reward_per_episode time_per_step
        end,
    )

    description = """
    This experiment uses the `DQNLearner` method with three dense layers to approximate the Q value.
    The testing environment is MountainCarEnv.

    Agent and statistic info will be saved to: `$save_dir`
    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("$save_dir", Agent)
    BSON.@load joinpath("$save_dir", "stats.bson") total_reward_per_episode time_per_step
    ```
    """

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

function RLCore.Experiment(
    ::Val{:JuliaRL},
    ::Val{:SAC},
    ::Val{:Pendulum},
    ::Nothing;
    seed = 123,
)
    rng = MersenneTwister(seed)
    inner_env = PendulumEnv(T = Float32, rng = rng)
    action_space = get_actions(inner_env)
    low = action_space.low
    high = action_space.high
    ns = length(get_state(inner_env))

    env = inner_env |> ActionTransformedEnv(x -> low + (x + 1) * 0.5 * (high - low))
    init = glorot_uniform(rng)

    create_policy_net() = NeuralNetworkApproximator(
        model = SACPolicyNetwork(
            Chain(Dense(ns, 30, relu), Dense(30, 30, relu)),
            Chain(Dense(30, 1, initW = init)),
            Chain(Dense(
                30,
                1,
                x -> min(max(x, typeof(x)(-20)), typeof(x)(2)),
                initW = init,
            )),
        ),
        optimizer = ADAM(0.003),
    )

    create_q_net() = NeuralNetworkApproximator(
        model = Chain(
            Dense(ns + 1, 30, relu; initW = init),
            Dense(30, 30, relu; initW = init),
            Dense(30, 1; initW = init),
        ),
        optimizer = ADAM(0.003),
    )

    agent = Agent(
        policy = SACPolicy(
            policy = create_policy_net(),
            qnetwork1 = create_q_net(),
            qnetwork2 = create_q_net(),
            target_qnetwork1 = create_q_net(),
            target_qnetwork2 = create_q_net(),
            γ = 0.99f0,
            ρ = 0.995f0,
            α = 0.2f0,
            batch_size = 64,
            start_steps = 1000,
            start_policy = RandomPolicy(ContinuousSpace(-1.0, 1.0); rng = rng),
            update_after = 1000,
            update_every = 1,
            rng = rng,
        ),
        trajectory = CircularCompactSARTSATrajectory(
            capacity = 10000,
            state_type = Float32,
            state_size = (ns,),
            action_type = Float32,
        ),
    )

    description = """
    # Play Pendulum with SAC
    """

    Experiment(agent, env, StopAfterStep(10000), TotalRewardPerEpisode(), description)
end


# used for constructing appropraite dict from YAML object.
function construct_type_map(t::Symbol, constructor::YAML.Constructor,
                            node::YAML.MappingNode)
    mapping = _symbol_dict(YAML.construct_mapping(constructor, node))
    mapping[:tag] = t
    mapping
end

function construct_type_map(t::Symbol, constructor::YAML.Constructor,
                            node::YAML.SequenceNode)
    YAML.construct_sequence(constructor, node)
end

const yaml_types = let
    pairs = [("!Cartesian", :Cartesian),
            ("!Smolyak", :Smolyak),
            ("!Random", :Random),
             ("!Normal", :Normal),
             ("!MarkovChain", :MarkovChain),
             ("!Product", :Product),
             ("!PoissonProcess", :PoissonProcess),
             ("!DeathProcess", :DeathProcess),
             ("!AgingProcess", :AgingProcess),
             ("!VAR1", :VAR1)]
    Dict{AbstractString,Function}([(t, (c, n) -> construct_type_map(s, c, n))
                           for (t, s) in pairs])
end

mutable struct SModel{ID} <: ASModel{ID}
    data::Dict{Symbol,Any}
end

function SModel(url::AbstractString)
    if match(r"(http|https):.*", url) != nothing
        res = HTTP.request("GET", url)
        txt = (String(res.body))
        data = _symbol_dict(YAML.load(txt, yaml_types))
    else
        data = _symbol_dict(YAML.load_file(url, yaml_types))
    end
    id = gensym(basename(url))
    return SModel{id}(_symbol_dict(data))
end

function get_symbols(model::ASModel)
    syms = model.data[:symbols]
    symbols = OrderedDict{Symbol,Vector{Symbol}}()
    for k in keys(syms)
        symbols[Symbol(k)] = [Symbol(e) for e in syms[k]]
    end
    return symbols
end

function get_variables(model::ASModel)
    symbols = get_symbols(model)
    vars = cat(values(symbols)...; dims=1)
    dynvars = setdiff(vars, symbols[:parameters] )
end

function get_name(model::ASModel)
    get(model.data, :name, "modeldoesnotwork")
end

### TODO: look for duplicates
eq_to_expr(eq::Symbol) = eq
function eq_to_expr(eq::Expr)
   if !(eq.head==:call && eq.args[1]==:(==))
      eq
   else
      rhs = eq.args[3]
      lhs = eq.args[2]
      :($rhs-$lhs)
   end
end

function get_equations(model::ASModel)

    eqs = model.data[:equations]
    recipe = RECIPES[:dtcc]
    _symbols = get_symbols(model)

    # prep equations: parse to Expr
    _eqs = OrderedDict{Symbol,Vector{Expr}}()
    # for k in Dolo.keys(recipe[:specs])
    for k in Dolo.keys(eqs)
      if !(k in Dolo.keys(recipe[:specs]))
           ii=0
           dist = zeros(length(Dolo.keys(recipe[:specs])))
           for kk in Dolo.keys(recipe[:specs])
             ii+=1
             dist[ii]=compare(Hamming(), string(kk), string(k))
           end
           kk = collect(keys(Dolo.keys(recipe[:specs]).dict))[findmax(dist)[2]]
           msg = string("You are using a wrong name for a sub-section title.`$(k)` is not accepted. ",
                        "Maybe you meant `$(kk)`?")
           error(msg)
      end
      (k in [:arbitrage,]) && continue
      these_eq = get(eqs, (k), [])
      # verify that we have at least 1 equation if section is required
      if !get(recipe[:specs][k], :optional, false)
          length(these_eq) == 0 && error("equation section $k required")
      end
      # finally pass in the expressions
      _eqs[k] = Expr[_to_expr(eq) for eq in these_eq]
    end
    # end
    # handle the arbitrage, arbitrage_exp, controls_lb, and controls_ub
    if haskey(recipe[:specs], :arbitrage) && (:arbitrage in keys(eqs))
        c_lb, c_ub, arb = _handle_arbitrage(eqs[:arbitrage],
                                            _symbols[:controls])
        _eqs[:arbitrage] = arb
        _eqs[:controls_lb] = c_lb
        _eqs[:controls_ub] = c_ub
    end

    defs = get_definitions(model)
    dynvars = cat(get_variables(model), [k[1] for k in keys(defs)]...; dims=1)

    for eqtype in keys(_eqs)
        _eqs[eqtype] = [sanitize(eq,dynvars) for eq in _eqs[eqtype]]
    end

    return _eqs
end

function get_infos(model::ASModel)
    get(model.data, :infos, Dict())
end

function get_options(model::ASModel; options=Dict())
    opts = get(model.data, :options, Dict())
    return rmerge(opts, options)
end

function get_definitions(model::ASModel)::OrderedDict{Tuple{Symbol,Int},SymExpr}
    # parse defs so values are Expr
    defs = get(model.data,:definitions, Dict())
    dynvars = cat(get_variables(model), keys(defs)...; dims=1)
    _defs = OrderedDict{Tuple{Symbol,Int},SymExpr}(
        [ (k,0) => sanitize(_to_expr(v), dynvars) for (k, v) in defs]
    )
    return _defs
end

function get_calibration(model::ASModel)
    calib = get(model.data, :calibration, Dict())
    # prep calib: parse to Expr, Symbol, or Number
    _calib  = OrderedDict{Symbol,Union{Expr,Symbol,Number}}()
    # add calibration for definitions
    for k in keys(calib)
        _calib[Symbol(k)] = _expr_or_number(calib[k])
    end
    # so far _calib is a symbolic calibration
    calibration = solve_triangular_system(_calib)
    symbols = get_symbols(model)
    return ModelCalibration(calibration, symbols)
end

function get_domain(model::ASModel)
    domain = deepcopy(get(model.data, :domain, Dict()))
    calib = get_calibration(model)
    # TODO deal with empty dict and make robust construction
    states = get_symbols(model)[:states]
    min = [eval_with(calib, domain[(k)][1]) for k in states]
    max = [eval_with(calib, domain[(k)][2]) for k in states]
    return Domain(states, min, max)
end

function get_grid(model::ASModel; options=Dict())
    domain = get_domain(model)
    d = length(domain.states)
    options = get_options(model; options=Dict(:grid=>options))
    grid_dict = options[:grid]
    if :type in keys(grid_dict)
        grid_dict[:tag] = grid_dict[:type]
    end
    if grid_dict[:tag] == :Cartesian
        orders = get(grid_dict, :orders, [20 for i=1:d])
        grid = CartesianGrid{d}(domain.min, domain.max, orders)
        if length(orders)!=length(model.calibration[:states])
            msg = string("Check the dimension of the matrix given in the yaml file, section: options-grid-orders. ",
                         "Expected to be of dimension $([length(model.calibration[:states])])")
            error(msg)
        end
    elseif grid_dict[:tag] == :Smolyak
        mu = get(grid_dict, :mu, 3)
        grid = SmolyakGrid{d}(domain.min, domain.max, mu)
    elseif grid_dict[:tag] == :Random
        n = get(grid_dict, :N, 200)
        grid = RandomGrid{d}(domain.min, domain.max, n)
    else
        error("Unknown grid type.")
    end
    return grid
end

function get_exogenous_old(model::ASModel)
    exo_dict = get(model.data,:exogenous,Dict{Symbol,Any}())
    if length(exo_dict)==0
        exo_dict = get(model.data[:options], :exogenous, Dict{Symbol,Any}())
    end
    exo = _symbol_dict(exo_dict)
    calib = get_calibration(model)
    exogenous = _build_exogenous_entry(exo, calib)
    return exogenous
end

#

function get_exogenous(model::AModel)
    rdata = model.rdata
    exo_dict = model.rdata[:exogenous]
    cond = !(exo_dict.tag=="tag:yaml.org,2002:map")
    if cond
        # old style exogenous block
        return get_exogenous_old(model)
    else
        syms = cat( [[Symbol(strip(e)) for e in split(k, ",")] for k in keys(exo_dict)]..., dims=1)
        expected = model.symbols[:exogenous]
        if (syms != expected)
            msg = string("In 'exogenous' section, shocks must be declared in the same order as shock symbol. Found: ", syms, ". Expected: ", expected, ".")
            throw(ErrorException(msg))
        end
        calibration = model.calibration.flat
        processes = []
        for k in keys(exo_dict)
            v = exo_dict[k]
            p = Dolang.eval_node(v, calibration, minilang, FromGreek())
            push!(processes, p)
        end
        return ProductProcess(processes...)
    end
end

function set_calibration!(model::ASModel, key::Symbol, value::Union{Real,Expr, Symbol})
    model.data[:calibration][key] = value
end

mutable struct Model{ID} <: AModel{ID}

    data::Dict{Symbol,Any}
    rdata::YAML.Node
    name::String           # weakly immutable
    filename::String
    symbols::OrderedDict{Symbol,Array{Symbol,1}}        # weakly immutable
    equations      # weakly immutable
    definitions    # weakly immutable
    # functions      # weakly immutable
    factories::Dict
    calibration::ModelCalibration
    exogenous
    domain
    grid
    options

    function Model{ID}(data::Dict{Symbol,Any}, rdata::YAML.Node; print_code::Bool=false, filename="<string>") where ID

        model = new{ID}(data, rdata)
        model.name = get_name(model)
        model.filename = filename
        model.symbols = get_symbols(model)
        model.equations = get_equations(model)
        model.definitions = get_definitions(model)
        model.factories = Dict()
        model.calibration = get_calibration(model)
        model.exogenous = get_exogenous(model)
        model.domain = get_domain(model)
        model.options = get_options(model)
        model.grid = get_grid(model)

        factories = Dict{Symbol,Dolang.FlatFunctionFactory}()
        for (eqtype,equations) in model.equations
            spec = Dolo.RECIPES[:dtcc][:specs][eqtype]
            spec_eqs = spec[:eqs]
            symbols = model.symbols
            arguments = OrderedDict{Symbol,Union{Vector{Symbol},Vector{Tuple{Symbol,Int}}}}([
                Symbol(s)=>[(e,t) for e in symbols[Symbol(sg)]]
                for (sg,t,s) in spec_eqs
            ])
            arguments[:p] = [e[1] for e in arguments[:p]]
            tdefs = Dolo.get_definitions(model)
            if :target in keys(spec)
                sg,t,s = spec[:target]
                targets = [(e,t) for e in symbols[Symbol(sg)]]
                lhs = [eq.args[2] for eq in equations]
                rhs = [eq.args[3] for eq in equations]
                expressions = OrderedDict(zip(targets,rhs))
                @assert Dolang.stringify.(lhs) == Dolang.stringify.(targets)
            else
                expressions = OrderedDict( [Symbol("out_",i)=>eq_to_expr(eq) for (i,eq) in enumerate(equations)])
            end
            ff = Dolang.FlatFunctionFactory(expressions, arguments, tdefs; funname=eqtype )
            factories[eqtype] = ff
            code = Dolang.gen_generated_gufun(ff; dispatch=typeof(model))
            print_code && println(code)
            Core.eval(Dolo, code)
        end

        model.factories = factories

        # TEMP: until we have a better method in Dolang
        # Create definitions
        vars = cat(model.symbols[:exogenous], model.symbols[:states], model.symbols[:controls]; dims=1)
        args = OrderedDict(
            :past => [(v,-1) for v in vars],
            :present => [(v,0) for v in vars],
            :future => [(v,1) for v in vars],
            :params => model.symbols[:parameters]
        )
        fff = Dolang.FlatFunctionFactory(model.definitions, args, typeof(model.definitions)())
        code = Dolang.gen_generated_gufun(fff;dispatch=typeof(model), funname=:evaluate_definitions)
        Core.eval(Dolo,code)

        return model
    end

end

_numeric_mod_type(::Model{ID}) where {ID} = Model{ID}

function Base.show(io::IO, model::Model)
    print(io, "Model")
end

#### import functions

function Model(url::AbstractString; print_code=false)
    # it looks like it would be cool to use the super constructor ;-)
    if match(r"(http|https):.*", url) != nothing
        res = HTTP.request("GET", url)
        txt = String(res.body)
    else
        txt = read(open(url), String)
    end
    data = _symbol_dict(load(txt, yaml_types))
    rdata = Dolang.yaml_node_from_string(txt)
    id = gensym()
    fname = basename(url)
    return Model{id}(data, rdata; print_code=print_code, filename=fname)
end

function set_calibration!(model::Model, key::Symbol, value::Union{Real,Expr, Symbol})
    model.data[:calibration][key] = value
    model.calibration = get_calibration(model)
    model.exogenous = get_exogenous(model)
    model.domain = get_domain(model)
    # model.options = get_options(model) # we don't substitute calib here
    model.grid = get_grid(model)
    model.calibration
end

function set_calibration!(model::Model, values::AbstractDict{Symbol,T}) where T
    for (key,value) in values
        model.data[:calibration][key] = value
    end
    model.calibration = get_calibration(model)
    model.exogenous = get_exogenous(model)
    model.domain = get_domain(model)
    # model.options = get_options(model) # we don't substitute calib here
    model.grid = get_grid(model)
    model.calibration
end

function set_calibration!(model::Model; kwargs...)
    set_calibration!(model, Dict(kwargs))
end

"""
Imports the model from a yaml file specified by the `url` input
parameter, and returns the corresponding `Dolo` model object.
"""
function yaml_import(url; print_code::Bool=false)
    Model(url; print_code=print_code)
end

function features(model::Dolo.ASModel)
    features = Dict{Symbol, Bool}()
    features[:one_state] = (length(model.symbols[:states])==1)
    features[:one_control] = (length(model.symbols[:controls])==1)
    features[:one_dimensional] = features[:one_state] && features[:one_state]
    rhs_transitions = [eq.args[3] for eq in model.equations[:transition]]
    it = Dolang.IncidenceTable( model.equations[:transition] )
    features[:nonstochastic_transitions] = (length( intersect(it.by_date[0], model.symbols[:exogenous] ) )==0)
    it = Dolang.IncidenceTable( [model.equations[:controls_lb]; model.equations[:controls_lb]] )
    features[:bounds_are_constant] = (length(it.by_var) == 0)
    return features
end


using Blpapi
using CSV

# Verbindungsaufbau
con = createConn()

t = floor((now() - Dates.Month(0)), Dates.Day(1))
startDate = floor((t - Dates.Day(1)), Dates.Minute(30))

    # Instrumen deklaration
inst = instrument("RXU8 Comdty", "TRADE", startDate, t)

    #request data
BarFrame(con, inst, 30)

z = bundDaten[[:Dates, :Volume]]

sort(z, cols = [:Volume])[(end-19):end,:]

i
TS(con, inst)

bdh(connection.session, "IBM US Equity", "TRADE", startDate, t)

bdh(connection.session, ["IBM US Equity"], ["PX_LAST"], "20180101", "20180331")
bdh(connection.session, [inst.ticker], "PX_Last", "20180101", "20180331")

barsize = 30

# Pulling all from BBG
Response = bar(connection.session, inst.ticker, inst.fields, string(inst.startDate), string(inst.endDate); interval=barsize)

#sorting into DataFrame

df = DataFrame()
df[:Dates] = Response.datetimes
df[:Open] = zeros(length(Response.datetimes))
df[:High] = zeros(length(Response.datetimes))
df[:Low] = zeros(length(Response.datetimes))
df[:Close] = zeros(length(Response.datetimes))
df[:Volume] = zeros(length(Response.datetimes))

for i in 1:length(Response.datetimes)
df[i,:Open] = Response[Response.datetimes[i]].openVar
df[i,:High] = Response[Response.datetimes[i]].highVar
df[i,:Low] = Response[Response.datetimes[i]].lowVar
df[i,:Close] = Response[Response.datetimes[i]].closeVar
df[i,:Volume] = Response[Response.datetimes[i]].volumeVar
end

using TimeZones

df

d = df[:,1]

zd = ZonedDateTime.(d, tz"UTC")

finalDT = astimezone.(zd, tz"Europe/Zurich")

df[:,1] = finalDT

df

Date(finalDT[1], dateformat"y-m-d")

con = createConn()
"PX_High", "PX_Low", "PX_Last"
# For historical end of day, mind the date format
inst = instrument("RXU8 Comdty", ["PX_Open", "PX_High", "PX_Low", "PX_Last"], "20180501", "20180625")

#function TS(con = connection, inst = ibm)

#endData = Array()

for i in (1:length(inst.fields))

    response = bdh(con.session, [inst.ticker], [inst.fields[i]], inst.startDate, inst.endDate; periodicitySelection = "DAILY")
    data = response[inst.ticker, inst.fields[i]]

    dformat = Dates.DateFormat("y-m-d")
    dates = similar(data[:,1], Date)
    prices = similar(data[:,2], Float64)

    for j in (1:length(data[:,1]))

        dates[j] = Date(data[j,1], dformat)
        prices[j] = data[j,2]

    end

    k = sortperm(dates)
    dates = dates[k]
    prices = prices[k]

    #if i > 1
    #    data2 = similar(data)

    data = TimeArray(dates, prices, [inst.fields[i]])

    # if i == 1
    #     dataFinal = similar(1:length(data), Float64)
    #     dataFinal = data
    # end
    #
    # if i > 1
    #     dataFinal = merge(dataFinal, data)
    # end
    #

end

connection = createConn()
bund = instrument("RXU8 Comdty", ["PX_Open", "PX_High", "PX_Low", "PX_Last"], "20180601", "20180627")

function TS(con = connection, inst = ibm, field = "PX_Last")

  response = bdh(con.session, [inst.ticker], [field], inst.startDate, inst.endDate; periodicitySelection = "DAILY")
  data = response[inst.ticker, field]

  dformat = Dates.DateFormat("y-m-d")
  dates = similar(data[:,1], Date)
  prices = similar(data[:,2], Float64)

  for i in (1:length(data[:,1]))
    dates[i] = Date(data[i,1], dformat)
    prices[i] = data[i,2]
  end

  k = sortperm(dates)
  dates = dates[k]
  prices = prices[k]

  data = TimeArray(dates, prices, [field])

end

TS(connection, bund, bund.fields[1])

data = TS(connection, bund)

colnames(data)

Dates.firstdayofweek(Date(now()-Dates.Day(5))) + Dates.Day(1)
Dates.firstdayofweek(Date(now()-Dates.Day(5)))

data[Dates.firstdayofweek(Date(now()-Dates.Day(5)))]
data[Dates.firstdayofweek(Date(now()-Dates.Day(5))) + Dates.Day(1)]

data.timestamp

#end
for

    #refresh end time to now

    #call BarFrame

    #write result to csv./ or call this function from excel

function barUpdate(contract, con = con)

    t = now()
    startDate = t - Dates.Month(2)

    inst = instrument(contract.ticker, contract.fieldnames, startDate, t)

    BarFrame(con, inst, 30)

end

quantile(barUpdate()[(end-300):end,:Volume],0.95)
log(bund2[:High]./bund2[:Low]) ./ log(bund2[:Close])

(bund2[end,:High] - bund2[end,:Low]) / maximum(bund2[:High] - bund2[:Low])

for i in keys(instContracts)
       contract = instContracts[i]
       barUpdate(contract)
end

function volIndicatorXL(commodity, days_back)

    inst = bbgTS.contracts[commodity]
    inst.startDate = now() - Dates.Month(2)
    inst.endDate = now()
    df = BarFrame(bbgTS.conInternal, inst, bbgTS.n)
    df = Array{Any}(df[Date(df[:,:Dates]) .== Date(now())-Dates.Day(days_back), [:Dates, :Volume, :RotFact]])
    df[:,1] = toXLDate.(df[:,1])
    return df

end


using .TransformVariables: AbstractTransform, transform, inverse

setting(xf::AbstractTransform) = _setting(xf)
getting(xf::AbstractTransform) = _getting(xf)

XFBijection{INV}(xf::T) where {INV, T} = XFBijection{INV, T}(xf)
XFBijection(xf) = XFBijection{false}(xf)

Base.inv(bijection::XFBijection{INV}) where INV = XFBijection{!INV}(bijection.xf)

Bijection(xf::AbstractTransform) = XFBijection(xf)
tofield(b::XFBijection{false}, x) = transform(b.xf, x)
fromfield(b::XFBijection{false}, y) = inverse(b.xf, y)
tofield(b::XFBijection{true}, x) = inverse(b.xf, x)
fromfield(b::XFBijection{true}, y) = transform(b.xf, y)

Base.show(io::IO, lens::BijectionLens{<:XFBijection}) =
    print_apply(io, typeof(lens), _getfields(lens.bijection))


using Test
using SafeTestsets

@testset "shallow_water" begin
  @safetestset "entropy_conservation_1d" begin include("entropy_conservation_1d.jl") end
end


module ZeroProblems

using ..Lenny: AbstractContinuationProblem

#--- Forward definitions

"""
	name(x)

Return the name of a variable or zero problem as a Symbol.
"""
function name end

"""
	dependencies(x)

Return the variable dependencies of a zero problem as a Tuple.
"""
function dependencies end

#--- Variables that zero problems depend on

"""
	Var

A placeholder for a continuation variable. As a minimum it comprises a name and
length (in units of the underlying numerical type). Optionally it can include a
parent variable and offset into the parent variable (negative offsets represent
the offset from the end).

# Example

```
coll = Var(:collocation, 20)  # an array of 20 Float64 (or other numerical type)
x0 = Var(:x0, 1, coll, 0)  # a single Float64 at the start of the collocation array
x1 = Var(:x1, 1, coll, -1)  # a single Float64 at the end of the collocation array
```
"""
mutable struct Var
	name::Symbol
	len::Int
	parent::Union{Var, Nothing}
	offset::Int
end

Var(name::Symbol, len::Integer) = Var(name, len, nothing, 0)

name(u::Var) = u.name
Base.length(u::Var) = u.len

#--- ZeroSubproblem

# For simple subproblems the ZeroSubproblem type can be used to define a
# subproblem and its associated variable dependencies.

# For more complicated subproblems (or subproblems that require access to the
# full problem structure) you should inherit from AbstractZeroSubproblem.

abstract type AbstractZeroSubproblem end

name(ϕ::AbstractZeroSubproblem) = ϕ.name
dependencies(ϕ::AbstractZeroSubproblem) = ϕ.deps

struct ZeroSubproblem{F} <: AbstractZeroSubproblem
	name::Symbol
	deps::Vector{Var}
	f!::F
end

@inline (z::ZeroSubproblem)(res, prob::AbstractContinuationProblem, u...) = z.f!(res, u...)
@inline (z::ZeroSubproblem)(res, u...) = z.f!(res, u...)

#--- ZeroProblem - the full problem structure

struct ZeroProblem{D, U, Φ}
	u::U
	ui::Vector{UnitRange{Int}}
	ϕ::Φ
	ϕi::Vector{UnitRange{Int}}
	ϕdeps::Vector{Tuple{Vararg{Int, N} where N}}
end

ZeroProblem() = ZeroProblem{Nothing, Vector{Var}, Vector{Any}}(Vector{Var}(), Vector{UnitRange{Int}}(), Vector{Any}(), Vector{UnitRange{Int}}())

function Base.push!(zp::ZeroProblem{Nothing}, u::Var)
	if !(u in zp.u)
		if u.parent === nothing
			last = maximum(maximum.(zp.ui))
			push!(zp.u, u)
			push!(zp.ui, last + 1:last + length(u))
		else
			idx = findfirst(isequal(u.parent), zp.u)
			start = (u.offset < 0) ? (zp.ui[idx][end] + u.offset + 1) : (zp.ui[idx][1] + u.offset)
			push!(zp.u, u)
			push!(zp.ui, start:start + length(u) - 1)
		end
	end
	return zp
end

function Base.push!(zp::ZeroProblem{Nothing}, subprob::AbstractZeroSubproblem)
	if subprob in zp.ϕ
		throw(ArgumentError("Subproblem is already part of the zero problem"))
	end
	depidx = Vector{Int}()
	for dep in dependencies(subprob)
		push!(zp, dep)
		push!(depidx, findfirst(isequal(dep), zp.u))
	end
	last = maximum(maximum.(zp.ϕi))
	push!(zp.ϕ, ϕ)
	push!(zp.ϕi, last + 1:last + length(ϕ))
	push!(zp.ϕdeps, (depidx...,))
	return zp
end

function specialize(zp::ZeroProblem)
	u = (zp.u...,)
	ui = zp.ui
	ϕ = (zp.ϕ...,)
	ϕi = zp.ϕi
	ϕdeps = zp.ϕdeps
	return ZeroProblem{(ϕdeps...,), typeof(u), typeof(ϕ)}(u, ui, ϕ, ϕi, ϕdeps)
end

function residual!(res, zp::ZeroProblem{Nothing}, prob::AbstractContinuationProblem, u)
	# TODO: implement anyway?
	throw(ArgumentError("Specialize the zero problem before calling residual!"))
end

@generated function residual!(res, zp::ZeroProblem{D, U, Φ}, prob::AbstractContinuationProblem, u) where {D, U <: Tuple, Φ <: Tuple}
	body = quote
		# Construct views into u for each variable
		uv = ($((:(view(u, zp.ui[$i])) for i in eachindex(U.parameters))...),)
		# Construct views into res for each subproblem
		resv = ($((:(view(res, zp.ϕi[$i])) for i in eachindex(Φ.parameters))...),)
	end
	# Call each of the subproblems
	for i in eachindex(D)
        if length(D[i]) == 0
            # No dependencies means pass everything
            push!(body.args, :(zp.ϕ[$i](resv[$i], prob, u)))
        else
            push!(body.args, :(zp.ϕ[$i](resv[$i], prob, $((:(uv[$(D[i][j])]) for j in eachindex(D[i]))...))))
        end
	end
	# Return res
    push!(body.args, :res)
    @show body
    body
end

end


module FastBEAST

include("tree/tree.jl")

include("aca.jl")
include("skeletons.jl")
include("utils.jl")
include("hmatrix.jl")
include("beast.jl")

export BoundingBox
export getboxframe
export getchildbox
export whichchildbox

export BoxTreeNode
export create_tree
export BoxTreeOptions
export KMeansTreeOptions
export KMeansTreeNode

export aca_compression

export FullMatrixView, FullMatrixView2, MatrixView, LowRankMatrixView

export HMatrix
export buildhmatrix
export adjoint
export estimate_norm
export estimate_reldifference
export nnz
export compressionrate

export hassemble
end


### Generators are functions used for generating various specific models as instances of GenericModel

## likelihood_model represents a likelihood L(Vector{Parameter} | Vector{Data}, Vector{Hyperparameter})

function likelihood_model(vs::VariableVector; isdirected::Bool=true, isindexed::Bool=true)
  m = GenericModel(vs, Dependence[], isdirected=isdirected, isindexed=isindexed)

  pindex = find(v::Variable -> isa(v, Parameter), m.vertices)
  for i in pindex
    for j in setdiff(1:length(vs), pindex)
      add_edge!(m, m.vertices[j], m.vertices[i])
    end
  end

  return m
end

likelihood_model(v::Variable, isindexed::Bool=true) = GenericModel(Variable[v], Dependence[], isindexed=isindexed)


#!/usr/bin/env bash
    #=
    exec julia --project="$(realpath $(dirname $0))/" "${BASH_SOURCE[0]}" "$@" -e "include(popfirst!(ARGS))" \
    "${BASH_SOURCE[0]}" "$@"
    =#

#=
Adapted from https://github.com/Simon-Hohberg/Viola-Jones/
=#

println("\033[1;34m===>\033[0;38m\033[1;38m\tLoading required libraries (it will take a moment to precompile if it is your first time doing this)...\033[0;38m")

include(joinpath(dirname(@__DIR__), "src", "FaceDetection.jl"))

using .FaceDetection
const FD = FaceDetection
using Printf: @printf
using Serialization: serialize

println("...done\n")

function main(;
    smart_choose_feats::Bool=false,
	scale::Bool=false,
	scale_to::Tuple=(200, 200)
)
	include("constants.jl")
	include("main_data.jl")

	min_size_img = (19, 19) # default for our test dataset
	if smart_choose_feats
		# For performance reasons restricting feature size
		notify_user("Selecting best feature width and height...")

		max_feature_width, max_feature_height, min_feature_height, min_feature_width, min_size_img = determine_feature_size(pos_training_path, neg_training_path; scale = scale, scale_to = scale_to)

		println("...done.  Maximum feature width selected is $max_feature_width pixels; minimum feature width is $min_feature_width; maximum feature height is $max_feature_height pixels; minimum feature height is $min_feature_height.\n")
	else
		min_feature_height = 8
		max_feature_height = 10
		min_feature_width = 8
		max_feature_width = 10
	end

	# classifiers are haar like features
	votes = FD.get_feature_votes(pos_training_path, neg_training_path, num_classifiers, min_feature_height, max_feature_height, min_feature_width, max_feature_width; scale = scale, scale_to = scale_to)

	# write classifiers to file
	img_size = scale ? scale_to : min_size_img
	data_file = joinpath(@__DIR__, "data", "feature_votes_$(img_size)")
	serialize(data_file, votes)
end

@time main(smart_choose_feats=true, scale=true, scale_to=(20, 20))


# mbe_demo.jl
#
# This script contains the work towards a "reduced effort" control law for an
# underactuated vehicle. In includes the initial simulation, the Monte-Carlo
# runs, and TODO.
#
# To run this file, start Julia, navigate to the directory containing this
# file, and just include it in your main workspace, like so:
#
# ```
# $ julia
# julia> cd(Pkg.dir("Eratosthenes") * "/examples/5_mbe")
# julia> include("mbe_demo.jl")
# ```

# Include our custom models.
include("models.jl")

# We work inside a module because it's more flexible than the main workspace.
module MBEDemo

#############
# Resources #
#############

# Use/import the modules we'll need.
using Eratosthenes # The simulation engine (setup, simulate, mc)
using ..MBEModels  # The MBEModels module included above (and "above" this module)
import HDF5        # For loading the log file
using Plots        # For plotting our results
using LinearAlgebra

# Create a directory for the log file.
if !isdir("out"); mkdir("out"); end

##############
# Simulation #
##############

# Specify the name of the YAML file that defines our scenario.
yaml_file = joinpath(@__DIR__, "underactuated.yaml")

# Run the scenario defined by our YAML file. The returned scenario will
# have all of the final states.
scenario = simulate(yaml_file, @__MODULE__)

#########
# Plots #
#########

# Read some data from the logs. This do block takes care of
# error-handling, so we can just read whatever we want and return it.
t, q_BI, ω_BI_B = HDF5.h5open(scenario.sim.log_file, "r") do logs
    (read(logs, "/smallsat/body/state/time"),
     read(logs, "/smallsat/body/state/data/q_BI"),
     read(logs, "/smallsat/body/state/data/ω_BI_B"))
end

# Scalars are logged as 1-by-n. Convert to just n for the plot.
t = dropdims(t, dims=1)

# Choose a friendly plotting package.
pyplot()
# plotlyjs()

# Define each plot that we'll need.
display(plot(t, transpose(q_BI),
             label  = ["q1" "q2" "q3" "q4"],
             xlabel = "Time (s)",
             ylabel = "Attitude Quaternion"))
display(plot(t, 180/π * transpose(ω_BI_B),
             label  = ["ω1" "ω2" "ω3"],
             xlabel = "Time (s)",
             ylabel = "Rotation Rate (deg/s)"))

# The plot looks fine, so we can now try some Monte-Carlo runs.
# See mbe_demo_mc.jl.

end # module


abstract type AbstractLShaped end

StochasticPrograms.num_subproblems(lshaped::AbstractLShaped) = lshaped.num_subproblems
num_cuts(lshaped::AbstractLShaped) = lshaped.data.num_cuts
num_iterations(lshaped::AbstractLShaped) = lshaped.data.iterations
tolerance(lshaped::AbstractLShaped) = lshaped.parameters.τ

# Initialization #
# ======================================================================== #
function initialize!(lshaped::AbstractLShaped)
    # Initialize progress meter
    lshaped.progress.thresh = lshaped.parameters.τ
    # Initialize subproblems
    initialize_subproblems!(lshaped, scenarioproblems(lshaped.structure, 2))
    # Prepare the master optimization problem
    prepare_master!(lshaped)
    # Initialize regularization policy
    initialize_regularization!(lshaped)
    # Finish initialization
    finish_initilization!(lshaped)
    return nothing
end
# ======================================================================== #

# Functions #
# ======================================================================== #
function prepare_master!(lshaped::AbstractLShaped)
    # Check sense first
    sense = MOI.get(lshaped.master, MOI.ObjectiveSense())
    if sense == MOI.FEASIBILITY_SENSE
        lshaped.data.no_objective = true
        # Use min-sense during L-shaped procedure
        MOI.set(lshaped.master, MOI.ObjectiveSense(), MOI.MIN_SENSE)
        F = AffineDecisionFunction{Float64}
        MOI.set(lshaped.master, MOI.ObjectiveFunction{F}(), zero(F))
    else
        # Cache the objective function
        F = MOI.get(lshaped.master, MOI.ObjectiveFunctionType())
        lshaped.data.master_objective = MOI.get(lshaped.master, MOI.ObjectiveFunction{F}())
    end
    # Initialize the required number of master variables. Use
    # MOI.VariableIndex(0) as an undef value
    append!(lshaped.master_variables, fill(MOI.VariableIndex(0), num_thetas(lshaped)))
end

function add_master_variable!(lshaped::AbstractLShaped, index::Integer)
    F = MOI.get(lshaped.master, MOI.ObjectiveFunctionType())
    master_variable = MOI.add_variable(lshaped.master)
    # Set name
    MOI.set(lshaped.master, MOI.VariableName(), master_variable,
            add_subscript("θ", index))
    # Get sense
    sense = MOI.get(lshaped.master, MOI.ObjectiveSense())
    coeff = sense == MOI.MIN_SENSE ? 1.0 : -1.0
    # Add to objective
    MOI.modify(lshaped.master, MOI.ObjectiveFunction{F}(), MOI.ScalarCoefficientChange(master_variable, coeff))
    lshaped.master_variables[index] = master_variable
    return nothing
end

function remove_cut_constraints!(lshaped::AbstractLShaped)
    # Decrease count
    lshaped.data.num_cuts -= length(lshaped.cut_constraints)
    # Remove cuts
    for ci in lshaped.cut_constraints
        if !iszero(ci.value)
            MOI.delete(lshaped.master, ci)
        end
    end
    empty!(lshaped.cut_constraints)
    return nothing
end

function restore_master!(lshaped::AbstractLShaped)
    # Remove cut constraints
    remove_cut_constraints!(lshaped)
    # Remove master variables
    for var in lshaped.master_variables
        if !iszero(var.value)
            MOI.delete(lshaped.master, var)
        end
    end
    empty!(lshaped.master_variables)
    # Remove any regularization terms
    restore_regularized_master!(lshaped)
    if lshaped.data.no_objective
        # Re-set FEASIBILITY_SENSE
        MOI.set(lshaped.master, MOI.ObjectiveSense(), MOI.FEASIBILITY_SENSE)
    else
        # Re-add original objective
        @unpack master_objective = lshaped.data
        F = typeof(master_objective)
        MOI.set(lshaped.master, MOI.ObjectiveFunction{F}(), master_objective)
    end
    return nothing
end

function filter_cuts!(lshaped::AbstractLShaped, list)
    # Nothing to do if constraints do not match
    return nothing
end

function filter_cuts!(lshaped::AbstractLShaped, list::Vector{<:CutConstraint})
    # Filter out the cut constraints
    for cut in optimizer.lshaped.cut_constraints
        i = something(findfirst(isequal(c), list), 0)
        if !iszero(i)
            MOI.deleteat!(list, i)
        end
    end
    return nothing
end

function active_model_objectives(lshaped::AbstractLShaped)
    return map(lshaped.master_variables) do var
        var.value != 0
    end
end

function update_solution!(lshaped::AbstractLShaped)
    ncols = num_decisions(lshaped.structure)
    nb = num_thetas(lshaped)
    lshaped.x .= MOI.get.(lshaped.master, MOI.VariablePrimal(), all_decisions(lshaped.decisions))
    θs = map(lshaped.master_variables) do vi
        if vi.value == 0
            -1e10
        else
            MOI.get(lshaped.master, MOI.VariablePrimal(), vi)
        end
    end
    set_model_objectives(lshaped, θs)
    return nothing
end

function decision(lshaped::AbstractLShaped, index::MOI.VariableIndex)
    i = something(findfirst(i -> i == index, all_decisions(lshaped.decisions)), 0)
    if iszero(i)
        throw(MOI.InvalidIndex(index))
    end
    return decision(lshaped)[i]
end

function evaluate_first_stage(lshaped::AbstractLShaped, x::AbstractVector)
    model = lshaped.master
    # Get objective
    @unpack master_objective = lshaped.data
    # Evaluate objective
    decisions = all_decisions(lshaped.decisions)
    obj_val = MOIU.eval_variables(master_objective) do vi
        if vi in decisions
            # Only evaluate decision
            x[vi.value]
        else
            0.0
        end
    end
    # Return value
    return obj_val
end

function current_objective_value(lshaped::AbstractLShaped)
    # Get sense
    sense = MOI.get(lshaped.master, MOI.ObjectiveSense())
    correction = sense == MOI.MIN_SENSE ? 1.0 : -1.0
    # Return sense-corrected value
    return evaluate_first_stage(lshaped, current_decision(lshaped)) +
        correction * sum(subobjectives(lshaped))
end

function calculate_estimate(lshaped::AbstractLShaped)
    # Get sense
    sense = MOI.get(lshaped.master, MOI.ObjectiveSense())
    correction = sense == MOI.MIN_SENSE ? 1.0 : -1.0
    # Return sense-corrected value
    return evaluate_first_stage(lshaped, lshaped.x) +
        correction * sum(model_objectives(lshaped))
end

function log!(lshaped::AbstractLShaped; optimal = false, status = nothing)
    @unpack Q, θ, iterations = lshaped.data
    @unpack keep, offset, indent = lshaped.parameters
    # Early termination log
    if status != nothing && lshaped.parameters.log
        lshaped.progress.thresh = Inf
        lshaped.progress.printed = true
        val = if status == MOI.INFEASIBLE
            Inf
        elseif status == MOI.DUAL_INFEASIBLE
            -Inf
        else
            0.0
        end
        ProgressMeter.update!(lshaped.progress, val,
                              showvalues = [
                                  ("$(indentstr(indent))Objective", val),
                                  ("$(indentstr(indent))Early termination", status),
                                  ("$(indentstr(indent))Number of cuts", num_cuts(lshaped)),
                                  ("$(indentstr(indent))Iterations", iterations)
                              ], keep = keep, offset = offset)
        return nothing
    end
    # Value update
    push!(lshaped.Q_history, Q)
    push!(lshaped.θ_history, θ)
    lshaped.data.iterations += 1
    # Log
    log_regularization!(lshaped)
    if lshaped.parameters.log
        current_gap = optimal ? 0.0 : gap(lshaped)
        ProgressMeter.update!(lshaped.progress, current_gap,
                              showvalues = [
                                  ("$(indentstr(indent))Objective", objective_value(lshaped)),
                                  ("$(indentstr(indent))Gap", current_gap),
                                  ("$(indentstr(indent))Number of cuts", num_cuts(lshaped)),
                                  ("$(indentstr(indent))Iterations", iterations)
                              ], keep = keep, offset = offset)
    end
    return nothing
end

function log!(lshaped::AbstractLShaped, t::Integer; optimal = false, status = nothing)
    @unpack Q,θ,iterations = lshaped.data
    @unpack keep, offset, indent = lshaped.parameters
    lshaped.Q_history[t] = Q
    lshaped.θ_history[t] = θ
    if status != nothing && lshaped.parameters.log
        lshaped.progress.thresh = Inf
        lshaped.progress.printed = true
        val = if status == MOI.INFEASIBLE
            Inf
        elseif status == MOI.DUAL_INFEASIBLE
            -Inf
        else
            0.0
        end
        ProgressMeter.update!(lshaped.progress, val,
                              showvalues = [
                                  ("$(indentstr(indent))Objective", val),
                                  ("$(indentstr(indent))Early termination", status),
                                  ("$(indentstr(indent))Number of cuts", num_cuts(lshaped)),
                                  ("$(indentstr(indent))Iterations", iterations)
                              ], keep = keep, offset = offset)
        return nothing
    end
    log_regularization!(lshaped,t)
    if lshaped.parameters.log
        current_gap = optimal ? 0.0 : gap(lshaped)
        ProgressMeter.update!(lshaped.progress, current_gap,
                              showvalues = [
                                  ("$(indentstr(indent))Objective", objective_value(lshaped)),
                                  ("$(indentstr(indent))Gap", current_gap),
                                  ("$(indentstr(indent))Number of cuts", num_cuts(lshaped)),
                                  ("$(indentstr(indent))Iterations", iterations)
                              ], keep = keep, offset = offset)
    end
    return nothing
end

function indentstr(n::Integer)
    return repeat(" ", n)
end

function check_optimality(lshaped::AbstractLShaped)
    @unpack τ = lshaped.parameters
    @unpack θ = lshaped.data
    return θ > -Inf && gap(lshaped) <= τ
end
# ======================================================================== #

# Cut functions #
# ======================================================================== #
active(lshaped::AbstractLShaped, hyperplane::AbstractHyperPlane) = active(hyperplane, decision(lshaped), tolerance(lshaped))
active(lshaped::AbstractLShaped, cut::HyperPlane{OptimalityCut}) = optimal(cut, decision(lshaped), model_objectives(lshaped)[cut.id], tolerance(lshaped))
active(lshaped::AbstractLShaped, cut::AggregatedOptimalityCut) = optimal(cut, decision(lshaped), sum(model_objectives(lshaped)[cut.ids]), tolerance(lshaped))
satisfied(lshaped::AbstractLShaped, hyperplane::AbstractHyperPlane) = satisfied(hyperplane, decision(lshaped), tolerance(lshaped))
satisfied(lshaped::AbstractLShaped, cut::HyperPlane{OptimalityCut}) = satisfied(cut, decision(lshaped), model_objectives(lshaped)[cut.id], tolerance(lshaped))
satisfied(lshaped::AbstractLShaped, cut::AggregatedOptimalityCut) = satisfied(cut, decision(lshaped), sum(model_objectives(lshaped)[cut.ids]), tolerance(lshaped))
violated(lshaped::AbstractLShaped, hyperplane::AbstractHyperPlane) = !satisfied(lshaped, hyperplane)
gap(lshaped::AbstractLShaped, hyperplane::AbstractHyperPlane) = gap(hyperplane, decision(lshaped))
gap(lshaped::AbstractLShaped, cut::HyperPlane{OptimalityCut}) = gap(cut, decision(lshaped), model_objectives(lshaped)[cut.id])
gap(lshaped::AbstractLShaped, cut::AggregatedOptimalityCut) = gap(cut, decision(lshaped), sum(model_objectives(lshaped)[cut.ids]))

function add_cut!(lshaped::AbstractLShaped, cut::AbstractHyperPlane; consider_consolidation = true, check = true)
    added = add_cut!(lshaped, cut, model_objectives(lshaped), subobjectives(lshaped), check = check)
    if consider_consolidation
        added && add_cut!(lshaped, lshaped.consolidation, cut)
    end
    return added
end

function add_cut!(lshaped::AbstractLShaped, cut::HyperPlane{OptimalityCut}, θs::AbstractVector, subobjectives::AbstractVector, Q::AbstractFloat; check = true)
    if lshaped.master_variables[cut.id].value == 0
        # Add master variable if cut is encountered for the first time
        add_master_variable!(lshaped, cut.id)
    end
    # Get the model value
    θ = θs[cut.id]
    @unpack cut_scaling = lshaped.parameters
    # Update objective
    subobjectives[cut.id] = Q
    # Check if cut gives new information
    if check && θ > -Inf && (θ + sqrt(eps()) >= Q || θ + sqrt(eps()) >= cut(lshaped.x))
        # Optimal with respect to this subproblem
        return false
    end
    # Add optimality cut
    process_cut!(lshaped, cut)
    f, set = moi_constraint(cut, lshaped.master_variables, cut_scaling)
    push!(lshaped.cut_constraints, MOI.add_constraint(lshaped.master, f, set))
    lshaped.data.num_cuts += 1
    if lshaped.parameters.debug
        push!(lshaped.cuts, cut)
    end
    return true
end
function add_cut!(lshaped::AbstractLShaped, cut::AggregatedOptimalityCut, θs::AbstractVector, subobjectives::AbstractVector, Q::AbstractFloat; check = true)
    for id in cut.ids
        if lshaped.master_variables[id].value == 0
            # Add master variable if cut is encountered for the first time
            add_master_variable!(lshaped, id)
        end
    end
    θs = θs[cut.ids]
    θ = sum(θs)
    @unpack cut_scaling = lshaped.parameters
    # Update objective
    subobjectives[cut.ids] .= Q/length(cut.ids)
    # Check if cut gives new information
    if check && θ > -Inf && (θ + sqrt(eps()) >= Q || θ + sqrt(eps()) >= cut(lshaped.x))
        # Optimal with respect to these subproblems
        return false
    end
    # Add optimality cut
    process_cut!(lshaped, cut)
    f, set = moi_constraint(cut, lshaped.master_variables, cut_scaling)
    push!(lshaped.cut_constraints, MOI.add_constraint(lshaped.master, f, set))
    lshaped.data.num_cuts += 1
    if lshaped.parameters.debug
        push!(lshaped.cuts, cut)
    end
    return true
end
add_cut!(lshaped::AbstractLShaped, cut::AnyOptimalityCut, θs::AbstractVector, subobjectives::AbstractVector, x::AbstractVector; check = true) = add_cut!(lshaped, cut, θs, subobjectives, cut(x); check = check)
add_cut!(lshaped::AbstractLShaped, cut::AnyOptimalityCut, θs::AbstractVector, subobjectives::AbstractVector; check = true) = add_cut!(lshaped, cut, θs, subobjectives, current_decision(lshaped); check = check)

function add_cut!(lshaped::AbstractLShaped, cut::HyperPlane{FeasibilityCut}, ::AbstractVector, subobjectives::AbstractVector, Q::AbstractFloat; check = true)
    # Ensure that there is no false convergence
    subobjectives[cut.id] = Q
    # Add feasibility cut
    process_cut!(lshaped, cut)
    f, set = moi_constraint(cut, lshaped.master_variables, 1.0)
    push!(lshaped.cut_constraints, MOI.add_constraint(lshaped.master, f, set))
    lshaped.data.num_cuts += 1
    if lshaped.parameters.debug
        push!(lshaped.feasibility.cuts, cut)
    end
    return true
end
add_cut!(lshaped::AbstractLShaped, cut::HyperPlane{FeasibilityCut}, θs::AbstractVector, subobjectives::AbstractVector; check = true) = add_cut!(lshaped, cut, θs, subobjectives, Inf)

function add_cut!(lshaped::AbstractLShaped, cut::HyperPlane{Infeasible}, ::AbstractVector, subobjectives::AbstractVector; check = true)
    subobjectives[cut.id] = Inf
    return true
end

function add_cut!(lshaped::AbstractLShaped, cut::HyperPlane{Unbounded}, ::AbstractVector, subobjectives::AbstractVector; check = true)
    subobjectives[cut.id] = -Inf
    return true
end


function pointcoords(geom::AbstractGeometry)
    @assert geotype(geom) == :Point
    [tuple(coordinates(geom)...)]
end

function multipointcoords(geom::AbstractGeometry)
    @assert geotype(geom) == :MultiPoint
    coords = coordinates(geom)
    first.(coords), last.(coords)
end

function linestringcoords(geom::AbstractGeometry)
    @assert geotype(geom) == :LineString
    coords = coordinates(geom)
    first.(coords), last.(coords)
end

function multilinestringcoords(geom::AbstractGeometry)
    @assert geotype(geom) == :MultiLineString
    x, y = Float64[], Float64[]
    for line in coordinates(geom)
        append!(x, first.(line)); push!(x, NaN)
        append!(y, last.(line)); push!(y, NaN)
    end
    x, y
end

function polygoncoords(geom::AbstractGeometry)
    @assert geotype(geom) == :Polygon
    ring = first(coordinates(geom)) # currently doesn't plot holes
    first.(ring), last.(ring)
end

function multipolygoncoords(geom::AbstractGeometry)
    @assert geotype(geom) == :MultiPolygon
    x, y = Float64[], Float64[]
    for poly in coordinates(geom)
        ring = first(coordinates(poly)) # currently doesn't plot holes
        append!(x, first.(ring)); push!(x, NaN)
        append!(y, last.(ring)); push!(y, NaN)
    end
    x, y
end

shapecoords(geom::AbstractPoint) = pointcoords(geom)
shapecoords(geom::AbstractMultiPoint) = multipointcoords(geom)
shapecoords(geom::AbstractLineString) = linestringcoords(geom)
shapecoords(geom::AbstractMultiLineString) = multilinestringcoords(geom)
shapecoords(geom::AbstractPolygon) = polygoncoords(geom)
shapecoords(geom::AbstractMultiPolygon) = multipolygoncoords(geom)

function shapecoords(geom::AbstractGeometry)
    gtype = geotype(geom)
    if gtype == :Point
        return pointcoords(geom)
    elseif gtype == :MultiPoint
        return multipointcoords(geom)
    elseif gtype == :LineString
        return linestringcoords(geom)
    elseif gtype == :MultiLineString
        return multilinestringcoords(geom)
    elseif gtype == :Polygon
        return polygoncoords(geom)
    elseif gtype == :MultiPolygon
        return multipolygoncoords(geom)
    else
        warn("unknown geometry type: $gtype")
    end
end

RecipesBase.@recipe f(geom::AbstractPoint) = (
    aspect_ratio --> 1;
    seriestype --> :scatter;
    legend --> :false;
    shapecoords(geom)
)

RecipesBase.@recipe f(geom::AbstractMultiPoint) = (
    aspect_ratio --> 1;
    seriestype --> :scatter;
    legend --> :false;
    shapecoords(geom)
)

RecipesBase.@recipe f(geom::AbstractLineString) = (
    aspect_ratio --> 1;
    seriestype --> :path;
    legend --> :false;
    shapecoords(geom)
)

RecipesBase.@recipe f(geom::AbstractMultiLineString) = (
    aspect_ratio --> 1;
    seriestype --> :path;
    legend --> :false;
    shapecoords(geom)
)

RecipesBase.@recipe f(geom::AbstractPolygon) = (
    aspect_ratio --> 1;
    seriestype --> :shape;
    legend --> :false;
    shapecoords(geom)
)

RecipesBase.@recipe f(geom::AbstractMultiPolygon) = (
    aspect_ratio --> 1;
    seriestype --> :shape;
    legend --> :false;
    shapecoords(geom)
)

RecipesBase.@recipe function f(geom::AbstractGeometry)
    aspect_ratio --> 1
    legend --> :false
    gtype = geotype(geom)
    if gtype == :Point || gtype == :MultiPoint
        seriestype := :scatter
    elseif gtype == :LineString || gtype == :MultiLineString
        seriestype := :path
    elseif gtype == :Polygon || gtype == :MultiPolygon
        seriestype := :shape
    else
        @warn("unknown geometry type: $gtype")
    end
    shapecoords(geom)
end

RecipesBase.@recipe function f(geom::Vector{<:Union{Missing, AbstractGeometry}})
    aspect_ratio --> 1
    legend --> :false
    for g in skipmissing(geom)
        @series begin
            gtype = geotype(g)
            if gtype == :Point || gtype == :MultiPoint
                seriestype := :scatter
            elseif gtype == :LineString || gtype == :MultiLineString
                seriestype := :path
            elseif gtype == :Polygon || gtype == :MultiPolygon
                seriestype := :shape
            else
                @warn("unknown geometry type: $gtype")
            end
            shapecoords(g)
        end
    end
end

RecipesBase.@recipe f(feature::AbstractFeature) = geometry(feature)
RecipesBase.@recipe f(features::Vector{<:AbstractFeature}) = geometry.(features)
RecipesBase.@recipe f(collection::AbstractFeatureCollection) = features(collection)
RecipesBase.@recipe f(collection::AbstractGeometryCollection) = geometries(collection)


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
    AbstractProgressiveHedgingAttribute

Abstract supertype for attribute objects specific to the progressive-hedging algorithm.
"""
abstract type AbstractProgressiveHedgingAttribute <: AbstractStructuredOptimizerAttribute end
"""
    PrimalTolerance

An optimizer attribute for specifying the primal tolerance in the progressive-hedging algorithm.
"""
struct PrimalTolerance <: AbstractProgressiveHedgingAttribute end
"""
    DualTolerance

An optimizer attribute for specifying the dual tolerance in the progressive-hedging algorithm.
"""
struct DualTolerance <: AbstractProgressiveHedgingAttribute end
"""
    Regularizer

An optimizer attribute for specifying a penalization procedure to be used in the progressive-hedging algorithm. Options are:

- [`Fixed`](@ref):  Fixed penalty (default) ?Fixed for parameter descriptions.
- [`Adaptive`](@ref): Adaptive penalty update ?Adaptive for parameter descriptions.
"""
struct Penalizer <: AbstractProgressiveHedgingAttribute end
"""
    PenaltyTerm

An optimizer attribute for specifying what proximal term should be used in subproblemsof the progressive-hedging algorithm. Options are:

- [`Quadratic`](@ref) (default)
- [`Linearized`](@ref)
- [`InfNorm`](@ref)
- [`ManhattanNorm`](@ref)
"""
struct PenaltyTerm <: AbstractProgressiveHedgingAttribute end
"""
    PenalizationParameter

Abstract supertype for penalization-specific attributes.
"""
abstract type PenalizationParameter <: AbstractProgressiveHedgingAttribute end


#####
##### apply flux boundary conditions to tracer tendency `Gc`
#####

# a positive west flux is associated with an *increase* in `Gc` near the west boundary.
# same for south and bottom.
@kernel function apply_west_bcs_kernel!(Gc, grid, west_bc, iter, ΔT)
    j, k = @index(Global, NTuple)
    jj = j + grid.Hy
    kk = k + grid.Hz
    @inbounds Gc[1+grid.Hx, jj, kk] += getbc(west_bc, j, k, iter) * ΔT * Ax(1+grid.Hx, jj, kk, grid) / volume(1+grid.Hx, jj, kk, grid)
end

@kernel function apply_south_bcs_kernel!(Gc, grid, south_bc, iter, ΔT)
    i, k = @index(Global, NTuple)
    ii = i + grid.Hx
    kk = k + grid.Hz
    @inbounds Gc[ii, 1+grid.Hy, kk] += getbc(south_bc, i, k, iter) * ΔT * Ay(ii, 1+grid.Hy, kk, grid) / volume(ii, 1+grid.Hy, kk, grid)
end

@kernel function apply_bottom_bcs_kernel!(Gc, grid, bottom_bc, iter, ΔT)
    i, j = @index(Global, NTuple)
    ii = i + grid.Hx
    jj = j + grid.Hy
    @inbounds Gc[ii, jj, grid.Nz+grid.Hz] += getbc(bottom_bc, i, j, iter) * ΔT * Az(ii, jj, grid.Nz+grid.Hz, grid) / volume(ii, jj, grid.Nz+grid.Hz, grid)
end

# a positive east flux is associated with an *decrease* in `Gc` near the east boundary.
# same for north and top.
@kernel function apply_east_bcs_kernel!(Gc, grid, east_bc, iter, ΔT)
    j, k = @index(Global, NTuple)
    jj = j + grid.Hy
    kk = k + grid.Hz
    @inbounds Gc[grid.Nx+grid.Hx, jj, kk] -= getbc(east_bc, j, k, iter) * ΔT * Ax(grid.Nx+grid.Hx, jj, kk, grid) / volume(grid.Nx+grid.Hx, jj, kk, grid)
end

@kernel function apply_north_bcs_kernel!(Gc, grid, north_bc, iter, ΔT)
    i, k = @index(Global, NTuple)
    ii = i + grid.Hx
    kk = k + grid.Hz
    @inbounds Gc[ii, grid.Ny+grid.Hy, kk] -= getbc(north_bc, i, k, iter) * ΔT * Ay(ii, grid.Ny+grid.Hy, kk, grid) / volume(ii, grid.Ny+grid.Hy, kk, grid)
end

@kernel function apply_top_bcs_kernel!(Gc, grid, top_bc, iter, ΔT)
    i, j = @index(Global, NTuple)
    ii = i + grid.Hx
    jj = j + grid.Hy
    @inbounds Gc[ii, jj, 1+grid.Hz] -= getbc(top_bc, i, j, iter) * ΔT * Az(ii, jj, 1+grid.Hz, grid) / volume(ii, jj, 1+grid.Hz, grid)
end

##### apply flux boundary conditions in x direction (west and east)
function apply_west_bcs!(Gc, grid, west_bc, iter, ΔT, arch, dep)
    kernel! = apply_west_bcs_kernel!(device(arch), (16,16), (grid.Ny, grid.Nz))
    event = kernel!(Gc, grid, west_bc, iter, ΔT, dependencies=dep)
    wait(device(arch), event)
    return nothing
end
function apply_east_bcs!(Gc, grid, east_bc, iter, ΔT, arch, dep)
    kernel! = apply_east_bcs_kernel!(device(arch), (16,16), (grid.Ny, grid.Nz))
    event = kernel!(Gc, grid, east_bc, iter, ΔT, dependencies=dep)
    wait(device(arch), event)
    return nothing
end

##### apply flux boundary conditions in y direction (south and north)
function apply_south_bcs!(Gc, grid, south_bc, iter, ΔT, arch, dep)
    kernel! = apply_south_bcs_kernel!(device(arch), (16,16), (grid.Nx, grid.Nz))
    event = kernel!(Gc, grid, south_bc, iter, ΔT, dependencies=dep)
    wait(device(arch), event)
    return nothing
end
function apply_north_bcs!(Gc, grid, north_bc, iter, ΔT, arch, dep)
    kernel! = apply_north_bcs_kernel!(device(arch), (16,16), (grid.Nx, grid.Nz))
    event = kernel!(Gc, grid, north_bc, iter, ΔT, dependencies=dep)
    wait(device(arch), event)
    return nothing
end

##### apply flux boundary conditions in z direction (bottom and top)
function apply_bottom_bcs!(Gc, grid, bottom_bc, iter, ΔT, arch, dep)
    kernel! = apply_bottom_bcs_kernel!(device(arch), (16,16), (grid.Nx, grid.Ny))
    event = kernel!(Gc, grid, bottom_bc, iter, ΔT, dependencies=dep)
    wait(device(arch), event)
    return nothing
end
function apply_top_bcs!(Gc, grid, top_bc, iter, ΔT, arch, dep)
    kernel! = apply_top_bcs_kernel!(device(arch), (16,16), (grid.Nx, grid.Ny))
    event = kernel!(Gc, grid, top_bc, iter, ΔT, dependencies=dep)
    wait(device(arch), event)
    return nothing
end

##### apply flux boundary conditions in all three directions, six faces.
function apply_bcs!(Gcs, nuts, grid, iter, ΔT, arch)
    barrier = Event(device(arch))
    events = []
    for name in nut_names
        west_event   =   apply_west_bcs!(Gcs[name].data, grid, nuts[name].bc.west,   iter, ΔT, arch, barrier)
        east_event   =   apply_east_bcs!(Gcs[name].data, grid, nuts[name].bc.east,   iter, ΔT, arch, barrier)
        south_event  =  apply_south_bcs!(Gcs[name].data, grid, nuts[name].bc.south,  iter, ΔT, arch, barrier)
        north_event  =  apply_north_bcs!(Gcs[name].data, grid, nuts[name].bc.north,  iter, ΔT, arch, barrier)
        bottom_event = apply_bottom_bcs!(Gcs[name].data, grid, nuts[name].bc.bottom, iter, ΔT, arch, barrier)
        top_event    =    apply_top_bcs!(Gcs[name].data, grid, nuts[name].bc.top,    iter, ΔT, arch, barrier)
        push!(events, west_event, east_event, south_event, north_event, bottom_event, top_event)
    end

    # filter out ::Nothing
    events = filter(e -> typeof(e) <: Event, events)

    wait(device(arch), MultiEvent(Tuple(events)))
    return nothing
end

# Avoid some computation / memory accesses when no flux boundary conditions will be applied
  apply_west_bcs!(Gc, grid, ::Nothing, iter, ΔT, arch, dep) = nothing
  apply_east_bcs!(Gc, grid, ::Nothing, iter, ΔT, arch, dep) = nothing
 apply_north_bcs!(Gc, grid, ::Nothing, iter, ΔT, arch, dep) = nothing
 apply_south_bcs!(Gc, grid, ::Nothing, iter, ΔT, arch, dep) = nothing
apply_bottom_bcs!(Gc, grid, ::Nothing, iter, ΔT, arch, dep) = nothing
   apply_top_bcs!(Gc, grid, ::Nothing, iter, ΔT, arch, dep) = nothing