export RainbowLearner, update!

using Flux, Zygote
using StatsBase

"""
    RainbowLearner(;kwargs...)

See paper: [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://arxiv.org/abs/1710.02298)

# Keywords

- `approximator`::[`AbstractQApproximator`](@ref): used to get Q-values of a state.
- `target_approximator`::[`AbstractQApproximator`](@ref): similar to `approximator`, but used to estimate the target (the next state).
- `loss_fun`: the loss function.
- `Vₘₐₓ::Float32`: the maximum value of distribution.
- `Vₘᵢₙ::Float32`: the minimum value of distribution.
- `n_actions::Int`: number of possible actions.
- `γ::Float32=0.99f0`: discount rate.
- `batch_size::Int=32`
- `update_horizon::Int=1`: length of update ('n' in n-step update).
- `min_replay_history::Int=32`: number of transitions that should be experienced before updating the `approximator`.
- `update_freq::Int=4`: the frequency of updating the `approximator`.
- `target_update_freq::Int=500`: the frequency of syncing `target_approximator`.
- `stack_size::Union{Int, Nothing}=4`: use the recent `stack_size` frames to form a stacked state.
- `default_priority::Float64=100.`: the default priority for newly added transitions.
- `n_atoms::Int=51`: the number of buckets of the value function distribution.
- `stack_size::Union{Int, Nothing}=4`: use the recent `stack_size` frames to form a stacked state.
- `default_priority::Float64=100.`: the default priority for newly added transitions.

!!! warning
    The `Zygote_gpu` version is slow due to that the `argmax(A, dims=1)` falls back to the CPU version in CuArrays.
"""
Base.@kwdef mutable struct RainbowLearner{Tq<:AbstractQApproximator, Tt<:AbstractQApproximator,Tf,Ts,Tss<:Union{Int, Nothing}} <: AbstractLearner
    approximator::Tq
    target_approximator::Tq
    loss_fun::Tf
    Vₘₐₓ::Float32
    Vₘᵢₙ::Float32
    n_actions::Int
    n_atoms::Int = 51
    support::Ts = collect(range(Float32(-Vₘₐₓ), Float32(Vₘₐₓ), length = n_atoms))
    stack_size::Tss = 4
    delta_z::Float32 = Float32(support[2] - support[1])
    γ::Float32 = 0.99
    batch_size::Int = 32
    update_horizon::Int = 1
    min_replay_history::Int = 32
    update_freq::Int = 1
    target_update_freq::Int = 500
    update_step::Int = 0
    default_priority::Float64 = 100.0

    function RainbowLearner(approximator::Tq, target_approximator::Tt, loss_fun::Tf, Vₘₐₓ::Float32, Vₘᵢₙ::Float32, n_actions::Int, n_atoms::Int, support::Ts, stack_size::Tss, args...) where {Tq,Tt, Tf, Ts, Tss}
        copyto!(approximator, target_approximator)  # force sync
        support = to_device(approximator, support)
        new{Tq,Tt,Tf, typeof(support), Tss}(approximator, target_approximator, loss_fun, Vₘₐₓ, Vₘᵢₙ, n_actions, n_atoms, support, stack_size, args...)
    end
end

function (learner::RainbowLearner)(obs::Observation)
    state = get_state(obs)
    logits = batch_estimate(learner.approximator, to_device(learner.approximator, state))
    q = learner.support .* softmax(reshape(logits, :, learner.n_actions))
    # probs = vec(sum(q, dims=1)) .+ legal_action
    vec(sum(q, dims = 1)) |> to_host
end

function update!(learner::RainbowLearner, batch)
    learner.update_step += 1
    learner.update_step % learner.update_freq == 0 || return nothing

    Q, Qₜ, γ, loss_fun, n_atoms, n_actions, support, delta_z, update_horizon, batch_size = learner.approximator,
        learner.target_approximator,
        learner.γ,
        learner.loss_fun,
        learner.n_atoms,
        learner.n_actions,
        learner.support,
        learner.delta_z,
        learner.update_horizon,
        learner.batch_size

    states, rewards, terminals, next_states = map(x->to_device(Q, x), (batch.states, batch.rewards, batch.terminals, batch.next_states))
    actions = CartesianIndex.(batch.actions, 1:batch_size)

    target_support = reshape(rewards, 1, :) .+
                     (reshape(support, :, 1) *
                      reshape((γ^update_horizon) .* (1 .- terminals), 1, :))

    updated_priorities = Vector{Float32}()

    gs = gradient(Q) do
        logits = reshape(batch_estimate(Q, states), n_atoms, n_actions, :)
        select_logits = logits[:, actions]
        next_logits = batch_estimate(Qₜ, next_states)
        next_probs = reshape(softmax(reshape(next_logits, n_atoms, :)), n_atoms, n_actions, :)
        next_q = reshape(sum(support .* next_probs, dims = 1), n_actions, :)
        # next_q_argmax = argmax(cpu(next_q .+ next_legal_actions), dims=1)
        next_q_argmax = argmax(next_q, dims = 1)
        next_prob_select = reshape(next_probs[:, next_q_argmax], n_atoms, :)

        target_distribution = project_distribution(
            target_support,
            next_prob_select,
            support,
            delta_z,
            learner.Vₘᵢₙ,
            learner.Vₘₐₓ,
        )

        batch_losses = loss_fun(select_logits, Zygote.dropgrad(target_distribution))
        updated_priorities = vec(clamp.(sqrt.(batch_losses .+ 1f-10), 1.f0, 1.f2))

        target_priorities = 1.0f0 ./ sqrt.(updated_priorities .+ 1f-10)
        target_priorities = target_priorities ./ maximum(target_priorities)

        mean(Zygote.dropgrad(target_priorities) .* batch_losses)
    end

    update!(Q, gs)

    if learner.update_step % learner.target_update_freq == 0
        copyto!(Qₜ, Q)
    end

    updated_priorities |> to_host
end

function project_distribution(supports, weights, target_support, delta_z, vmin, vmax)
    batch_size, n_atoms = size(supports, 2), length(target_support)
    clampped_support = clamp.(supports, vmin, vmax)
    tiled_support = reshape(repeat(clampped_support; outer=(n_atoms, 1)), n_atoms, n_atoms, batch_size)

    projection = clamp.(
        1 .- abs.(tiled_support .- reshape(target_support, 1, :)) ./ delta_z,
        0,
        1,
    ) .* reshape(weights, n_atoms, 1, batch_size)
    reshape(sum(projection, dims = 1), n_atoms, batch_size)
end

function extract_transitions(buffer::CircularTurnBuffer{PRTSA}, learner::RainbowLearner)
    if length(buffer) > learner.min_replay_history
        inds, consecutive_batch = sample(buffer, learner.batch_size, learner.update_horizon, learner.stack_size)
        inds, extract_SARTS(consecutive_batch, learner.γ)
    else
        nothing
    end
end


using CartesianGeneticProgramming
using Cambrian
using ArcadeLearningEnvironment
using ArgParse
import Cambrian.mutate
import Random

```
Playing Atari games using classic CGP on RAM values

If an individual is provided using --ind, an evaluation loop with rendering will
be performed. Otherwise, an evolution is launched. By default this uses seed=0
for each evaluation for a deterministic environment, but this can be removed for
a stochastic and more realistic result.
```

s = ArgParseSettings()
@add_arg_table s begin
    "--cfg"
    help = "configuration script"
    default = "cfg/atari_ram.yaml"
    "--game"
    help = "game rom name"
    default = "centipede"
    "--seed"
    help = "random seed for evolution"
    arg_type = Int
    default = 0
    "--ind"
    help = "individual for evaluation"
    arg_type = String
    default = ""
end
args = parse_args(ARGS, s)

function play_atari(ind::CGPInd, rom_name::String; seed=0, max_frames=18000, render=false)
    ale = ALE_new()
    setInt(ale, "random_seed", Cint(seed))
    if render
        setBool(ale, "display_screen", true)
        setBool(ale, "sound", true)
    end
    loadROM(ale, rom_name)
    actions = getMinimalActionSet(ale)
    reward = 0.0
    frames = 0
    while ~game_over(ale)
        ram = getRAM(ale) ./ typemax(UInt8)
        output = process(ind, ram)
        action = actions[argmax(output)]
        reward += act(ale, action)
        frames += 1
        if frames > max_frames
            break
        end
    end
    ALE_del(ale)
    [reward]
end

ale = ALE_new()
loadROM(ale, args["game"])
n_in = length(getRAM(ale))
n_out = length(getMinimalActionSet(ale))
ALE_del(ale)

cfg = get_config(args["cfg"]; game=args["game"], n_in=n_in, n_out=n_out)
Random.seed!(args["seed"])

if length(args["ind"]) > 0
    ind = CGPInd(cfg, read(args["ind"], String))
    reward = play_atari(ind, args["game"]; seed=args["seed"], render=true)
    println(reward)
else
    mutate(i::CGPInd) = goldman_mutate(cfg, i)
    fit(i::CGPInd) = play_atari(i, cfg.game)
    e = CGPEvolution(cfg, fit)
    run!(e)
end


#=##############################################################################
# DESCRIPTION
    Vehicle with vortex-lattice method (VLM) and blade-element momentum (BEM)
    models assuming quasi-steady wakes. This means that at every time step the
    wake of VLM models are represented as rigid, semi-infinite filaments, and
    wakes of BEM models are obtained from a conservation of momentum
    assumptions. In this approach only VLM-on-VLM and VLM-on-BEM interactions
    are accounted for, meaning that the wake of each BEM rotor has no effect
    on other rotors and wings.

# AUTHORSHIP
  * Author    : Eduardo J. Alvarez
  * Email     : Edo.AlvarezR@gmail.com
  * Created   : Apr 2020
  * License   : MIT
=###############################################################################

################################################################################
# QUASI-STEADY VLM VEHICLE TYPE
################################################################################
"""
    `QVLMVehicle(system; optargs...)`

Type handling all geometries and subsystems that define a flight vehicle made
out of VLM (Wing, WingSystem, Rotor) components.

NOTE: For the solver to work correctly, all components in `wake_system` (if any)
need to be also components of `vlm_system`.

NOTE: It is recommended that `wake_system` doesn't include any Rotor object.
Otherwise, blades will generate a wake going straight out of every blade
trailing edge pointing oposite to the direction of rotation instead of
generating a streamtube.

# ARGUMENTS
* `system::vlm.WingSystem`:        System of all FLOWVLM objects. This system
                                    is considered as the entire vehicle. Not all
                                    components in this system will be solved,
                                    but they will all be rotated and translated
                                    during maneuver.
# OPTIONAL ARGUMENTS
* `tilting_systems::Tuple(vlm.WingSystem, ...)`:   Tuple of all FLOWVLM
                                    tilting objects, where `tilting_systems[i]`
                                    contains the i-th FLOWVLM system of lifting
                                    surfaces and rotors that tilt together.
* `rotors_systems::Tuple(Array{vlm.Rotor,1}, ...)`:   Tuple of groups of Rotors
                                    that share a common RPM.
* `vlm_system::vlm.WingSystem`:    System of all FLOWVLM objects to be solved
                                    through the VLM solver.
* `wake_system::vlm.WingSystem`:   System of all FLOWVLM objects that will
                                    have a rigid, semi-infinite VLM wake.
* `grids::Array{gt.GridTypes, 1}`: Array of grids that will be translated and
                                    rotated along with `system`.
"""
struct QVLMVehicle{N, M, R} <: AbstractVLMVehicle{N, M, R}

    # Required inputs
    system::vlm.WingSystem

    # Optional inputs
    tilting_systems::NTuple{N, vlm.WingSystem}
    rotor_systems::NTuple{M, Array{vlm.Rotor, 1}}
    vlm_system::vlm.WingSystem
    wake_system::vlm.WingSystem
    grids::Array{gt.GridTypes, 1}

    # Internal properties
    V::Array{R, 1}                          # Current vehicle velocity
    W::Array{R, 1}                          # Current vehicle angular velocity
    prev_data::Array{Any, 1}                # Information about previous step
    grid_O::Array{Array{R, 1}, 1}           # Origin of every grid

    QVLMVehicle{N, M, R}(
                    system;
                    tilting_systems=NTuple{0, vlm.WingSystem}(),
                    rotor_systems=NTuple{0, Array{vlm.Rotor, 1}}(),
                    vlm_system=vlm.WingSystem(),
                    wake_system=vlm.WingSystem(),
                    grids=Array{gt.GridTypes, 1}(),
                    V=zeros(3), W=zeros(3),
                    prev_data=[deepcopy(vlm_system), deepcopy(wake_system),
                                                    deepcopy(rotor_systems)],
                    grid_O=Array{Array{Float64, 1}, 1}(),
                ) where {N, M, R} = new(
                    system,
                    tilting_systems,
                    rotor_systems,
                    vlm_system,
                    wake_system,
                    grids,
                    V, W,
                    prev_data,
                    grid_O,
                )
end

# Implicit N and M constructor
QVLMVehicle(system::vlm.WingSystem;
        V::Array{R, 1}=zeros(3), W::Array{R, 1}=zeros(3),
        tilting_systems::NTuple{N, vlm.WingSystem}=NTuple{0, vlm.WingSystem}(),
        rotor_systems::NTuple{M, Array{vlm.Rotor, 1}}=NTuple{0, Array{vlm.Rotor, 1}}(),
        grids=Array{gt.GridTypes, 1}(),
        optargs...
        ) where {N, M, R} = QVLMVehicle{N, M, R}( system;
                                V=V, W=W,
                                tilting_systems=tilting_systems,
                                rotor_systems=rotor_systems,
                                grids=Array{gt.GridTypes, 1}(grids),
                                grid_O=[zeros(R, 3) for i in 1:length(grids)],
                                optargs...)

##### FUNCTIONS  ###############################################################
function shed_wake(self::QVLMVehicle, Vinf::Function, pfield::vpm.ParticleField,
                                                dt::Real, nt::Int; optargs...)
    nothing
end

function generate_static_particle_fun(pfield, self::QVLMVehicle, sigma;
                                                                     optargs...)

    function static_particles_function(args...)
        return nothing
    end

    return static_particles_function
end

function save_vtk(self::QVLMVehicle, filename; path=nothing, num=nothing,
                                                save_wopwopin=false, optargs...)

    # Filter optional arguments for wake
    optargs1 = [(sym, val) for (sym, val) in optargs
                         if !(sym in [:infinite_vortex, :save_horseshoes,
                                      :only_horseshoes, :only_infinite_vortex])]

    # Save wake
    strn = vlm.save(self.wake_system, filename; suff="_wake",
                                infinite_vortex=true, save_horseshoes=true,
                                only_horseshoes=true, only_infinite_vortex=true,
                                path=path, num=num, optargs1...)

    # Filter optional arguments of VLM
    optargs2 = [(sym, val) for (sym, val) in optargs if sym!=:infinite_vortex]

    # Save VLM
    strn *= save_vtk_base(self, filename; path=path, num=num,
                                    save_wopwopin=save_wopwopin,
                                    infinite_vortex=false, optargs2...)
    return strn
end
##### INTERNAL FUNCTIONS  ######################################################

##### END OF VEHICLE ###########################################################


#=##############################################################################
# DESCRIPTION
    Plotting and outputting-processing functions for monitoring vehicle and
    component performance. The functions in this module generate
    monitor-functions that can be passed to the simulation engine through the
    optional arguments `extra_runtime_function`.

# AUTHORSHIP
  * Author    : Eduardo J. Alvarez
  * Email     : Edo.AlvarezR@gmail.com
  * Created   : Apr 2020
  * License   : MIT
=###############################################################################

"""
    `generate_monitor_rotors( rotors::Array{vlm.Rotor, 1}, J_ref::Real,
rho_ref::Real, RPM_ref::Real, nsteps_sim::Int)`

Generate monitor of load distribution and performance of a rotor system.
`J_ref` and `rho_ref` are the reference advance ratio and air density used for
calculating propulsive efficiency and coefficients. `RPM_ref` is the reference
RPM used to estimate the age of the wake. `nsteps_sim` is the number of time
steps by the end of the simulation (used for generating the color gradient).
"""
function generate_monitor_rotors( rotors::Array{vlm.Rotor, 1},
                                    J_ref::Real, rho_ref::Real, RPM_ref::Real,
                                    nsteps_sim::Int;
                                    t_scale=1.0,                    # Time scaling factor
                                    t_lbl="Simulation time (s)",    # Time-axis label
                                    # OUTPUT OPTIONS
                                    save_path=nothing,
                                    run_name="rotor",
                                    figname="monitor_rotor",
                                    disp_conv=true,
                                    conv_suff="_convergence.csv",
                                    nsteps_savefig=10,
                                    colors="rgbcmy"^100,
                                    stls="o^*.px"^100, )

    fcalls = 0                  # Number of function calls

    # Name of convergence file
    if save_path!=nothing
        fname = joinpath(save_path, run_name*conv_suff)
    end

    # Call figure
    if disp_conv
        fig = figure(figname, figsize=(7*3,5*2))
        axs = fig.subplots(2, 3)
        axs = [axs[1], axs[3], axs[5], axs[2], axs[4], axs[6]]
    end

    # Function for run_vpm! to call on each iteration
    function extra_runtime_function(sim::Simulation{V, M, R},
                                    PFIELD::vpm.ParticleField,
                                    T::Real, DT::Real
                                   ) where{V<:AbstractVLMVehicle, M, R}

        # rotors = vcat(sim.vehicle.rotor_systems...)
        angle = T*360*RPM_ref/60
        t_scaled = T*t_scale

        if fcalls==0
            # Format subplots
            if disp_conv
                ax = axs[1]
                ax.title.set_text("Circulation Distribution")
                ax.set_xlabel("Element index")
                ax.set_ylabel(L"Circulation $\Gamma$ (m$^2$/s)")
                ax.grid(true, color="0.8", linestyle="--")
                ax = axs[2]
                ax.title.set_text("Plane-of-rotation Normal Force")
                ax.set_xlabel("Element index")
                ax.set_ylabel(L"Normal Force $N_p$ (N/m)")
                ax.grid(true, color="0.8", linestyle="--")
                ax = axs[3]
                ax.title.set_text("Plane-of-rotation Tangential Force")
                ax.set_xlabel("Element index")
                ax.set_ylabel(L"Tangential Force $T_p$ (N/m)")
                ax.grid(true, color="0.8", linestyle="--")
                ax = axs[4]
                ax.set_xlabel(t_lbl)
                ax.set_ylabel(L"Thrust $C_T$")
                ax.grid(true, color="0.8", linestyle="--")
                ax = axs[5]
                ax.set_xlabel(t_lbl)
                ax.set_ylabel(L"Torque $C_Q$")
                ax.grid(true, color="0.8", linestyle="--")
                ax = axs[6]
                ax.set_xlabel(t_lbl)
                ax.set_ylabel(L"Propulsive efficiency $\eta$")
                ax.grid(true, color="0.8", linestyle="--")
            end

            # Convergence file header
            if save_path!=nothing
                f = open(fname, "w")
                print(f, "ref age (deg),T,DT")
                for (i, rotor) in enumerate(rotors)
                    print(f, ",RPM_$i,CT_$i,CQ_$i,eta_$i")
                end
                print(f, "\n")
                close(f)
            end
        end

        # Write rotor position and time on convergence file
        if save_path!=nothing
            f = open(fname, "a")
            print(f, angle, ",", T, ",", DT)
        end

        # Plot circulation and loads distributions
        if disp_conv

            cratio = PFIELD.nt/nsteps_sim
            cratio = cratio > 1 ? 1 : cratio
            clr = fcalls==0 && false ? (0,0,0) : (1-cratio, 0, cratio)
            stl = fcalls==0 && false ? "o" : "-"
            alpha = fcalls==0 && false ? 1 : 0.5

            # Circulation distribution
            this_sol = []
            for rotor in rotors
                this_sol = vcat(this_sol, [vlm.get_blade(rotor, j).sol["Gamma"] for j in 1:rotor.B]...)
            end
            axs[1].plot(1:size(this_sol,1), this_sol, stl, alpha=alpha, color=clr)

            # Np distribution
            this_sol = []
            for rotor in rotors
                this_sol = vcat(this_sol, rotor.sol["Np"]["field_data"]...)
            end
            axs[2].plot(1:size(this_sol,1), this_sol, stl, alpha=alpha, color=clr)

            # Tp distribution
            this_sol = []
            for rotor in rotors
                this_sol = vcat(this_sol, rotor.sol["Tp"]["field_data"]...)
            end
            axs[3].plot(1:size(this_sol,1), this_sol, stl, alpha=alpha, color=clr)
        end

        # Plot performance parameters
        for (i,rotor) in enumerate(rotors)
            CT, CQ = vlm.calc_thrust_torque_coeffs(rotor, rho_ref)
            eta = J_ref*CT/(2*pi*CQ)

            if disp_conv
                axs[4].plot([t_scaled], [CT], "$(stls[i])", alpha=alpha, color=clr)
                axs[5].plot([t_scaled], [CQ], "$(stls[i])", alpha=alpha, color=clr)
                axs[6].plot([t_scaled], [eta], "$(stls[i])", alpha=alpha, color=clr)
            end

            if save_path!=nothing
                print(f, ",", rotor.RPM, ",", CT, ",", CQ, ",", eta)
            end
        end

        if disp_conv
            # Save figure
            if fcalls%nsteps_savefig==0 && fcalls!=0 && save_path!=nothing
                fig.savefig(joinpath(save_path, run_name*"_convergence.png"),
                                                    transparent=false, dpi=300)
            end
        end

        # Close convergence file
        if save_path!=nothing
            print(f, "\n")
            close(f)
        end

        fcalls += 1

        return false
    end

    return extra_runtime_function
end

"""
    `generate_monitor_wing(wing, Vinf::Function, b_ref::Real, ar_ref::Real,
rho_ref::Real, qinf_ref::Real, nsteps_sim::Int)`

Generate monitor of load distribution and performance of a wing.
"""
function generate_monitor_wing(wing, Vinf::Function, b_ref::Real, ar_ref::Real,
                                rho_ref::Real, qinf_ref::Real, nsteps_sim::Int;
                                lencrit_f=0.5,      # Factor for critical length to ignore horseshoe forces
                                L_dir=[0,0,1],      # Direction of lift component
                                D_dir=[-1,0,0],     # Direction of drag component
                                # OUTPUT OPTIONS
                                out_Lwing=nothing,
                                out_Dwing=nothing,
                                out_CLwing=nothing,
                                out_CDwing=nothing,
                                save_path=nothing,
                                run_name="wing",
                                figname="monitor_wing",
                                disp_plot=true,
                                title_lbl="",
                                CL_lbl=L"Lift Coefficient $C_L$",
                                CD_lbl=L"Drag Coefficient $C_D$",
                                y2b_i=2,
                                y2b_ref=nothing,
                                ClCL_ref=nothing, CdCD_ref=nothing,
                                CL_ref=nothing, CD_ref=nothing,
                                ref_lbl="Reference", ref_stl="ok",
                                conv_suff="_convergence.csv",
                                figsize_factor=5/6,
                                nsteps_plot=10,
                                nsteps_savefig=10)

    fcalls = 0                  # Number of function calls

    # y2b = 2*wing._ym/b_ref          # Initial span position
    y2b = 2*[vlm.getControlPoint(wing, i)[y2b_i] for i in 1:vlm.get_m(wing)]/b_ref
    prev_wing = nothing

    # Critical length to ignore horseshoe forces
    meanchord = b_ref/ar_ref
    lencrit = lencrit_f*meanchord/vlm.get_m(wing)

    # Name of convergence file
    if save_path!=nothing
        fname = joinpath(save_path, run_name*conv_suff)
    end

    if disp_plot
        fig1 = figure(figname, figsize=[7*2, 5*2]*figsize_factor)
        fig1.suptitle(title_lbl)
        axs1 = fig1.subplots(2, 2)
        axs1 = [axs1[1], axs1[3], axs1[2], axs1[4]]
        ax = axs1[1]
        # xlim([-1,1])
        ax.set_xlabel(L"$\frac{2y}{b}$")
        ax.set_ylabel(L"$\frac{Cl}{CL}$")
        ax.title.set_text("Spanwise lift distribution")

        ax = axs1[2]
        # xlim([-1,1])
        ax.set_xlabel(L"$\frac{2y}{b}$")
        ax.set_ylabel(L"$\frac{Cd}{CD}$")
        ax.title.set_text("Spanwise drag distribution")

        ax = axs1[3]
        ax.set_xlabel("Simulation time (s)")
        ax.set_ylabel(CL_lbl)

        ax = axs1[4]
        ax.set_xlabel("Simulation time (s)")
        ax.set_ylabel(CD_lbl)

        fig2 = figure(figname*"_2", figsize=[7*2, 5*1]*figsize_factor)
        axs2 = fig2.subplots(1, 2)
        ax = axs2[1]
        ax.set_xlabel(L"$\frac{2y}{b}$")
        ax.set_ylabel(L"Circulation $\Gamma$")
        ax = axs2[2]
        ax.set_xlabel(L"$\frac{2y}{b}$")
        ax.set_ylabel(L"Effective velocity $V_\infty$")
    end

    function extra_runtime_function(sim, PFIELD, T, DT)

        aux = PFIELD.nt/nsteps_sim
        clr = (1-aux, 0, aux)

        if fcalls==0
            # Convergence file header
            if save_path!=nothing
                f = open(fname, "w")
                print(f, "T,CL,CD\n")
                close(f)
            end
        end

        if PFIELD.nt>2

            # Force at each VLM element
            Ftot = calc_aerodynamicforce(wing, prev_wing, PFIELD, Vinf, DT,
                                                            rho_ref; t=PFIELD.t,
                                                            lencrit=lencrit)
            L, D, S = decompose(Ftot, L_dir, D_dir)
            vlm._addsolution(wing, "L", L)
            vlm._addsolution(wing, "D", D)
            vlm._addsolution(wing, "S", S)

            # Force per unit span at each VLM element
            ftot = calc_aerodynamicforce(wing, prev_wing, PFIELD, Vinf, DT,
                                        rho_ref; t=PFIELD.t, per_unit_span=true,
                                        lencrit=lencrit)
            l, d, s = decompose(ftot, L_dir, D_dir)
            vlm._addsolution(wing, "l", l)
            vlm._addsolution(wing, "d", d)
            vlm._addsolution(wing, "s", s)

            # Lift of the wing
            Lwing = norm(sum(L))
            CLwing = Lwing/(qinf_ref*b_ref^2/ar_ref)
            ClCL = norm.(l) / (Lwing/b_ref)

            # Drag of the wing
            Dwing = norm(sum(D))
            CDwing = Dwing/(qinf_ref*b_ref^2/ar_ref)
            CdCD = [sign(dot(this_d, [1,0,0])) for this_d in d].*norm.(d) / (Dwing/b_ref) # Preserves the sign of drag

            vlm._addsolution(wing, "Cl/CL", ClCL)
            vlm._addsolution(wing, "Cd/CD", CdCD)

            if out_Lwing!=nothing; push!(out_Lwing, Lwing); end;
            if out_Dwing!=nothing; push!(out_Dwing, Dwing); end;
            if out_CLwing!=nothing; push!(out_CLwing, CLwing); end;
            if out_CDwing!=nothing; push!(out_CDwing, CDwing); end;

            if PFIELD.nt%nsteps_plot==0 && disp_plot

                ax = axs1[1]
                if y2b_ref!=nothing && ClCL_ref!=nothing
                    ax.plot(y2b_ref, ClCL_ref, ref_stl, label=ref_lbl)
                end
                ax.plot(y2b, ClCL, "-", alpha=0.5, color=clr)

                ax = axs1[2]
                if y2b_ref!=nothing && CdCD_ref!=nothing
                    ax.plot(y2b_ref, CdCD_ref, ref_stl, label=ref_lbl)
                end
                ax.plot(y2b, CdCD, "-", alpha=0.5, color=clr)

                ax = axs1[3]
                if CL_ref!=nothing
                    ax.plot([0, T], CL_ref*ones(2), ":k", label=ref_lbl)
                end
                ax.plot([T], [CLwing], "o", alpha=0.5, color=clr)

                ax = axs1[4]
                if CD_ref!=nothing
                    ax.plot([0, T], CD_ref*ones(2), ":k", label=ref_lbl)
                end
                ax.plot([T], [CDwing], "o", alpha=0.5, color=clr)

                ax = axs2[1]
                ax.plot(y2b, wing.sol["Gamma"], "-", alpha=0.5, color=clr)
                ax = axs2[2]
                if "Vkin" in keys(wing.sol)
                    ax.plot(y2b, norm.(wing.sol["Vkin"]), "-", alpha=0.5, color=[clr[1], 1, clr[3]])
                end
                if "Vvpm" in keys(wing.sol)
                    ax.plot(y2b, norm.(wing.sol["Vvpm"]), "-", alpha=0.5, color=clr)
                elseif "Vind" in keys(wing.sol)
                    ax.plot(y2b, norm.(wing.sol["Vind"]), "-", alpha=0.5, color=clr)
                end
                ax.plot(y2b, [norm(Vinf(vlm.getControlPoint(wing, i), T)) for i in 1:vlm.get_m(wing)],
                                                            "-k", alpha=0.5)

                if save_path!=nothing
                    # Save figures
                    if PFIELD.nt%nsteps_savefig==0
                        fig1.savefig(joinpath(save_path, run_name*"_convergence.png"),
                                                    transparent=false, dpi=300)

                        fig2.savefig(joinpath(save_path, run_name*"_convergence2.png"),
                                                    transparent=false, dpi=300)
                    end
                end
            end

            if save_path != nothing
                # Write rotor position and time on convergence file
                f = open(fname, "a")
                print(f, T, ",", CLwing, ",", CDwing, "\n")
                close(f)
            end
        end

        prev_wing = deepcopy(wing)
        fcalls += 1
        return false
    end
end

function generate_monitor_statevariables(; figname="monitor_statevariables",
                                           save_path=nothing,
                                           run_name="",
                                           nsteps_savefig=10)

    fig = figure(figname, figsize=[7*2, 5*1])
    axs = fig.subplots(1, 3)
    ax = axs[1]
    ax.set_xlabel("Simulation time")
    ax.set_ylabel("Velocity")
    Vlbls = [L"V_x", L"V_y", L"V_z"]
    ax = axs[2]
    ax.set_xlabel("Simulation time")
    ax.set_ylabel(L"Angular velocity ($^\circ/t$)")
    Wlbls = [L"\Omega_x", L"\Omega_y", L"\Omega_z"]
    ax = axs[3]
    ax.set_xlabel("Simulation time")
    ax.set_ylabel(L"$O$ position")
    Olbls = [L"O_x", L"O_y", L"O_z"]

    function extra_runtime_function(sim, PFIELD, T, DT)
        for j in 1:3
            axs[1].plot(sim.t, sim.vehicle.V[j], ".", label=Vlbls[j], alpha=0.8,
                                                            color=clrs[j])
            axs[2].plot(sim.t, sim.vehicle.W[j], ".", label=Wlbls[j], alpha=0.8,
                                                            color=clrs[j])
            axs[3].plot(sim.t, sim.vehicle.system.O[j], ".", label=Olbls[j], alpha=0.8,
                                                            color=clrs[j])
        end
        if sim.nt==0
            for j in 1:2
                axs[j].legend(loc="best", frameon=false)
            end
        end

        if save_path!=nothing
            # Save figures
            if PFIELD.nt%nsteps_savefig==0
                fig.savefig(joinpath(save_path, run_name*"statevariables.png"),
                                                    transparent=false, dpi=300)
            end
        end

        return false
    end

    return extra_runtime_function
end

function generate_monitor_enstrophy(; save_path=nothing,
                                     figname="monitor_enstrophy", run_name="",
                                     nsteps_savefig=10,
                                     nsteps_plot=10)

    fig = figure(figname, figsize=[7*1, 5*1])
    ax = fig.gca()
    ax.set_xlabel("Simulation time")
    ax.set_ylabel(L"Enstrophy ($\mathrm{m}^6/\mathrm{s}^2$)")

    enstrophy = []
    ts = []

    function extra_runtime_function(sim, PFIELD, T, DT)
        vpm.monitor_enstrophy(PFIELD, T, DT; save_path=save_path,
                                               run_name=run_name, out=enstrophy)

        push!(ts, T)

        if PFIELD.nt%nsteps_plot == 0
            ax.plot(ts, enstrophy[end-length(ts)+1:end], ".k")
            ts = []
        end

        if save_path!=nothing
            if PFIELD.nt%nsteps_savefig==0
                fig.savefig(joinpath(save_path, run_name*"enstrophy.png"),
                                                    transparent=false, dpi=300)
            end
        end

        return false
    end

    return extra_runtime_function
end


type MHRemoteSegments <: AbstractRemoteSegments
    remote::Vector{RemoteRef}
    numproposalspersegment::Int
    numsegments::Int

    prop2collected::Dict{Int,Tuple{Int,Int}}
    collectedsamples::Vector{AbstractSample}
    MHRemoteSegments(r::Vector{RemoteRef},npps::Int,ns::Int) = new(r,npps,ns,Dict{Int,Tuple{Int,Int}}(),Vector{AbstractSample}(ns))
end

@inline _numjobsegments(policy_::MHRuntimePolicy,nproposals::Int) = min(nproposals,_numjobsegments(traittype(policy_.jobsegments)))
@inline _processnumbers(policy_::MHRuntimePolicy,njobsegments::Int) = take(cycle(_processnumbers(traittype(policy_.jobsegments))),njobsegments)

function _remotesegments(policy_::MHRuntimePolicy,model_::AbstractModel,sampler_::AbstractSampler,nproposals::Int)
    njobsegments = _numjobsegments(policy_,nproposals)
    nproposalspersegment = _numproposalspersegment(nproposals,njobsegments)
    procnumbers = collect(_processnumbers(policy_,njobsegments))
    r = RemoteRef[remotecall(i,segment,policy_,model_,sampler_,nproposalspersegment) for i in procnumbers]
    MHRemoteSegments(r,nproposalspersegment,njobsegments)
end

_prop2seg(segments_::MHRemoteSegments,j::Int) = ind2sub((segments_.numproposalspersegment,segments_.numsegments),j)

function _insegmentindex(segments_::MHRemoteSegments,sampleindex_::AbstractVector)
    boolindex_ = [fill(false,segments_.numproposalspersegment) for k=1:segments_.numsegments]
    ntp = numtotalproposals(segments_)
    for i=1:length(sampleindex_)
        @inbounds i1 = sampleindex_[i]-1
        i1==ntp?continue:nothing #if this the final
        si = div(i1,segments_.numproposalspersegment) + 1
        pi = mod(i1,segments_.numproposalspersegment) + 1
        @inbounds boolindex_[si][pi] = true
    end
    map(find,boolindex_)
end

function _insegmentindex2collected!(segments_::MHRemoteSegments,insegmentindex::Array{Array{Int,1},1})
    empty!(segments_.prop2collected)
    d = (segments_.numproposalspersegment,segments_.numsegments)
    for i=1:segments_.numsegments
        for j=1:length(insegmentindex[i])
            segments_.prop2collected[sub2ind(d,insegmentindex[i][j],i)] = (j,i)
        end
    end
    segments_.prop2collected
end

#start the jobs on each process
function iterate!(segments_::MHRemoteSegments,indicatorstate::AbstractSamplerState)
    a = Array{RemoteRef}(segments_.numsegments)
    @sync begin
        for j=1:segments_.numsegments
            a[j] = remotecall(segments_.remote[j].where,iterate!,segments_.remote[j],indicatorstate)
        end
    end #@sync means wait for all processes to finish
    map(fetch,a)
end

#call prepare in order to copy over the new indicator state
function prepare!(segments_::MHRemoteSegments,indicatorstate::AbstractSamplerState,j::Int)
    p,s = _prop2seg(segments_,j)
    r = segments_.remote[s]
    remotecall_fetch(r.where,prepare!,r,indicatorstate,p)
end

function retrievesamples!(segments_::MHRemoteSegments,sampleindex_::AbstractVector)
    insegmentindex = _insegmentindex(segments_,sampleindex_)
    segments_.prop2collected = _insegmentindex2collected!(segments_,insegmentindex)
    @sync begin
        map!((r,i)->(~isempty(i)?remotecall_fetch(r.where,getsamples,r,i):samples(:base,0,0,Float64,Float64)),segments_.collectedsamples,segments_.remote,insegmentindex)
    end
    segments_.collectedsamples
end

function getsamples(segments_::MHRemoteSegments,j::Int)
    if !haskey(segments_.prop2collected,j)
        p,s = _prop2seg(segments_,j)
        r = segments_.remote[s]
        return remotecall_fetch(r.where,getsamples,r,p)
    else
        p,s = segments_.prop2collected[j]
        return copy(segments_.collectedsamples[s],p)
    end
end

function store!(segments_::MHRemoteSegments,chain_::AbstractChain,j::Int)
    p,s = segments_.prop2collected[j]
    store!(chain_,segments_.collectedsamples[s],p)
end

function tune!(segments_::MHRemoteSegments,tvals...)
    @sync begin
        for i=1:segments_.numsegments
            r = segments_.remote[i]
            remotecall_wait(r.where,tune!,r,tvals...)
        end
    end
end

function show(io::IO,r::MHRemoteSegments)
    println(io,"RemoteSegments with $(r.numsegments) segment",r.numsegments==1?"":"s"," and $(r.numproposalspersegment) proposal",r.numproposalspersegment==1?"":"s"," per segment.")
    println(io,"Additional fields: :remote, :collectedsamples, :prop2collected")
    nothing
end


push!(LOAD_PATH,"../src/")

using Documenter, DocumenterTools
using EAGO, IntervalArithmetic, MathOptInterface, McCormick, JuMP

import McCormick: final_cut, mid3v, precond_and_contract!, AbstractMCCallback, populate_affine!,
                  dline_seg, correct_exp!, cut, mid_grad, preconditioner_storage, newton, secant,
                  MCCallback, contract!, affine_exp!

import EAGO: ExtensionType, Evaluator, variable_dbbt!, set_current_node!,
             VariableInfo, Log, aggressive_filtering!,
             bool_indx_diff, trivial_filtering!, SIPResult, SIPProblem,
             GlobalOptimizer, InputProblem, ParsedProblem, is_integer_feasible_relaxed,
             local_problem_status, default_upper_heurestic, label_branch_variables!,
             label_fixed_variables!, AbstractDirectedGraph, AbstractCache,
             AbstractCacheAttribute, initialize!, f_init!, r_init!, fprop!, rprop!,
             Variable, Subexpression, Expression, Constant, Parameter

import EAGO.Script: dag_flattening!, register_substitution!, Template_Graph,
                    Template_Node, scrub, scrub!, flatten_expression!

const MOI = MathOptInterface

@info "Making documentation..."
makedocs(modules = [EAGO, McCormick],
         doctest = false,
         format = Documenter.HTML(
                prettyurls = get(ENV, "CI", nothing) == "true",
                canonical = "https://PSORLab.github.io/EAGO.jl/stable/",
                collapselevel = 1,
         ),
         authors = "Matthew E. Wilhelm",
         #repo = "https://github.com/PSORLab/EAGO.jl/blob/{commit}{path}#L{line}",
         sitename = "EAGO.jl: Easy Advanced Global Optimization",
         pages = Any["Introduction" => "index.md",
                     "McCormick Operator Library" => Any["McCormick/overview.md",
                                                         "McCormick/usage.md",
                                                         "McCormick/operators.md",
                                                         "McCormick/type.md",
                                                         "McCormick/implicit.md"
                                                        ],
                     "Global Optimizer" => Any["Optimizer/starting.md",
                                               "Optimizer/optimizer.md",
                                               "Optimizer/bnb_back.md",
                                               "Optimizer/relax_back.md",
                                               "Optimizer/domain_reduction.md",
                                               "Optimizer/high_performance.md",
                                               "Optimizer/udf_utilities.md"
                                               ],
                     "Semi-Infinite Programming" => "SemiInfinite/semiinfinite.md",
                     "Contributing to EAGO"      => Any["Dev/contributing.md",
                                                        "Dev/future.md"
                                                        ],
                     "References"  => "ref.md",
                     "Citing EAGO" => "cite.md"]
)

@info "Deploying documentation..."
deploydocs(repo = "github.com/PSORLab/EAGO.jl.git",
           push_preview  = true)


abstract type MAPPlot <: ESDLPlot end
import Colors: Color
using YAXArrayBase: getattributes
import YAXArrays.Cubes.Axes: get_step, abshalf, getAxis

mutable struct MAPPlotRGB <: MAPPlot
  xaxis
  yaxis
  rgbaxis
  dmin
  dmax
  c_1
  c_2
  c_3
  misscol::Color
  oceancol::Color
  cType
  overlay
end
plotAxVars(p::MAPPlotRGB)=[
  FixedAx(p.xaxis,"X Axis",true,false,1),
  FixedAx(p.yaxis,"Y Axis",true,false,2),
  FixedAx(p.rgbaxis,"RGB Axis", true, true,-1),
  FixedVar(p.rgbaxis,p.c_1,channel_names(p.cType)[1],true),
  FixedVar(p.rgbaxis,p.c_2,channel_names(p.cType)[2],true),
  FixedVar(p.rgbaxis,p.c_3,channel_names(p.cType)[3],true)
  ]
function plotCall(p::MAPPlotRGB, d, ixaxis, iyaxis, irgbaxis, c1,c2,c3,otherinds...)

  axlist = caxes(d)
  inds1 = ntuple(  i->(i==irgbaxis) ? axVal2Index(axlist[i],c1,fuzzy=true)  : in(i,(ixaxis,iyaxis)) ? (:) : axVal2Index(axlist[i],otherinds[i],fuzzy=true), length(otherinds))
  inds2 = ntuple(  i->(i==irgbaxis) ? axVal2Index(axlist[i],c2,fuzzy=true)  : in(i,(ixaxis,iyaxis)) ? (:) : axVal2Index(axlist[i],otherinds[i],fuzzy=true), length(otherinds))
  inds3 = ntuple(  i->(i==irgbaxis) ? axVal2Index(axlist[i],c3,fuzzy=true)  : in(i,(ixaxis,iyaxis)) ? (:) : axVal2Index(axlist[i],otherinds[i],fuzzy=true), length(otherinds))

  a1 = d[inds1...]
  a2 = d[inds2...]
  a3 = d[inds3...]

  if p.dmin==p.dmax
    mir,mar,mig,mag,mib,mab=(getMinMax(a1)...,getMinMax(a2)...,getMinMax(a3)...)
  else
    mir,mar = dmin[1],dmax[1]
    mig,mag = dmin[2],dmax[2]
    mib,mab = dmin[3],dmax[3]
  end

  rgbar=getRGBAR(p.cType,a1,a2,a3,(mir,mig,mib),(mar,mag,mab),p.misscol,p.oceancol)
  pngbuf=IOBuffer()
  show(pngbuf,"image/png",transpose(reshape(rgbar,size(rgbar,1),size(rgbar,2))))
  themap=compose(context(0,0,1,1),bitmap("image/png",pngbuf.data,0,0,1,1))
end

abstract type MAPPlotMapped <: MAPPlot end

plotAxVars(p::MAPPlotMapped)=[FixedAx(p.xaxis,"X Axis",true,false,1),FixedAx(p.yaxis,"Y Axis",true,false,2)]

mutable struct MAPPlotCategory <: MAPPlotMapped
  colorm
  colorm2
  oceancol
  misscol
  xaxis
  yaxis
  im_only
  overlay
end

toMatrix(a::Array)=reshape(a,size(a,1),size(a,2))
function plotCall(p::MAPPlotCategory,d, ixaxis, iyaxis,otherinds...)

  axlist = caxes(d)
  inds = ntuple(i->in(i,(ixaxis,iyaxis)) ? (:) : axVal2Index(axlist[i],otherinds[i],fuzzy=true), length(otherinds))

  a = d[inds...]

  if p.im_only
    _makeMaprgb(a,0.0,0.0,(p.colorm,p.colorm2),p.oceancol,p.misscol,:right,true,false,[])[5]
  else
    _makeMap(a,0.0,0.0,(p.colorm,p.colorm2),p.oceancol,p.misscol,:right,true,false,[],p.overlay)
  end
end

import Shapefile
#Here comes some Shapefile stuff
#PesudoMercator
#y2lat(y) = (2 * atan(exp( y/6378137)) - pi/2) * 180 / pi
#x2lon(x) = (x/6378137)*180/pi
x2lon(x) = x
y2lat(x) = x
getXY(p::Shapefile.Polygon) = map(i->x2lon(i.x),p.points),map(i->y2lat(i.y),p.points)
getarea(ii) = abs((ii.MBR.right-ii.MBR.left)*(ii.MBR.top-ii.MBR.bottom))
allpoints(rect) = ((rect.left, rect.bottom),(rect.right,rect.bottom), (rect.right,rect.top),(rect.left,rect.top))
pointinrect(p,r) = (r.left < p[1] < r.right) && (r.bottom < p[2] < r.top)
function getXY2(p::Shapefile.Polygon)
    points = map(i->(x2lon(i.x),y2lat(i.y)),p.points)
    map(1:length(p.parts)) do ipart
        i0 = p.parts[ipart]+1
        iend = ipart == length(p.parts) ? length(p.points) : p.parts[ipart+1]
        points[i0:iend]
    end
end
function getshapeplot(shapepath,userbb)
    handle = open(shapepath, "r") do io
        read(io, Shapefile.Handle)
    end
    p = handle.shapes;
    bbs = map(i->(left = x2lon(i.MBR.left), right = x2lon(i.MBR.right), top = y2lat(i.MBR.top), bottom=y2lat(i.MBR.bottom)),p)
    allp = allpoints(userbb)
    filter!(p) do i
        mbr = (left = x2lon(i.MBR.left), right = x2lon(i.MBR.right), top = y2lat(i.MBR.top), bottom=y2lat(i.MBR.bottom))
        let usb = userbb, alp = allp
            any(i->pointinrect(i,usb),allpoints(mbr)) || any(i->pointinrect(i,mbr),allp)
        end
    end
    xys = Vector{Tuple{Float64,Float64}}[]
    foreach(p) do ip
      append!(xys,getXY2(ip))
    end
    polBB = Compose.UnitBox(userbb.left,userbb.top,userbb.right-userbb.left, userbb.bottom-userbb.top)
    earthland = Compose.compose(context(units = polBB),Compose.line(xys))
end

mutable struct MAPPlotContin <: MAPPlotMapped
  colorm
  dmin
  dmax
  symmetric
  oceancol
  misscol
  xaxis
  yaxis
  tickspos
  im_only
  overlay
end

function plotCall(p::MAPPlotContin, d, ixaxis, iyaxis, otherinds...)

  axlist = caxes(d)
  inds = ntuple(i->in(i,(ixaxis,iyaxis)) ? (:) : axVal2Index(axlist[i],otherinds[i],fuzzy=true), length(otherinds))
  a = d[inds...]

  if p.dmin==p.dmax
    mi,ma=getMinMax(a)
  else
    mi=p.dmin
    ma=p.dmax
  end

  if p.im_only
    _makeMaprgb(a,mi,ma,p.colorm,p.oceancol,p.misscol,:bottom,false,p.symmetric,p.tickspos)[5]
  else
    _makeMap(a,mi,ma,p.colorm,p.oceancol,p.misscol,:bottom,false,p.symmetric,p.tickspos,p.overlay)
  end
end

function interpretoverlay(overlay,xaxis,yaxis,axlist)
  if overlay !== nothing
    xax,yax = getAxis(xaxis,axlist),getAxis(yaxis,axlist)
    xstep,ystep = get_step(xax.values),get_step(yax.values)
    userbb = (left = first(xax.values)-abshalf(xstep), right = last(xax.values)+abshalf(xstep), top = first(yax.values)-abshalf(ystep), bottom = last(yax.values)+abshalf(ystep))
    overlay = getshapeplot(overlay,userbb)
  end
  overlay
end

"""
    plotMAP(cube; dmin=datamin, dmax=datamax, colorm=colormap("oranges"), oceancol=colorant"darkblue", misscol=colorant"gray", kwargs...)

Map plotting tool for cube objects, can be called on any type of cube data

### Keyword arguments

* `dmin, dmax` Minimum and maximum value to be used for color transformation
* `colorm` colormap to be used. Find a list of colormaps in the [Colors.jl](https://github.com/JuliaGraphics/Colors.jl) package
* `oceancol` color to fill the ocean with, defaults to `colorant"darkblue"`
* `misscol` color to represent missing values, defaults to `colorant"gray"`
* `symmetric` make the color scale symmetric around zero
* `xaxis` which axis should be used for x axis, defaults to `LonAxis`
* `yaxis` which axis should be used for y axis, defaults to `LatAxis`
* `dim=value` can set other dimensions to certain values, for example `var="air_temperature_2m"` will fix the variable for the resulting plot
If a dimension is neither longitude or latitude and is not fixed through an additional keyword, a slider or dropdown menu will appear to select the axis value.

If the properties field of `cube` contains a "labels" field with a dictionary mapping field values to
the name of the class represented.
"""
function plotMAP(cube;xaxis="Lon", yaxis="Lat", dmin=zero(eltype(cube)),dmax=zero(eltype(cube)),
  colorm=:inferno,oceancol=colorant"darkblue",misscol=colorant"gray",symmetric=false, tickspos=[],im_only=false,
  overlay=nothing, kwargs...)

  isa(colorm,Symbol) && (colorm=get(namedcolms,colorm,namedcolms[:inferno]))
  dmin,dmax=typed_dminmax(eltype(cube),dmin,dmax)
  axlist=caxes(cube)

  props=getattributes(cube)

  overlay = interpretoverlay(overlay, xaxis,yaxis,axlist)

  if haskey(props,"labels")
    labels = props["labels"]
    _colorm  = distinguishable_colors(length(labels)+2,[misscol,oceancol])[3:end]
    colorm   = Dict(k=>_colorm[i] for (i,k) in enumerate(keys(labels)))
    colorm2  = Dict(k=>_colorm[i] for (i,k) in enumerate(values(labels)))
    plotGeneric(MAPPlotCategory(colorm,colorm2,oceancol,misscol,xaxis,yaxis,im_only,overlay),cube;kwargs...)
  else
    plotGeneric(MAPPlotContin(colorm,dmin,dmax,symmetric,oceancol,misscol,xaxis,yaxis,tickspos,im_only,overlay),cube;kwargs...)
  end
end

"""
    plotMAPRGB(cube; dmin=datamin, dmax=datamax, colorm=colormap("oranges"), oceancol=colorant"darkblue", misscol=colorant"gray", kwargs...)

Map plotting tool for colored plots that use up to 3 variables as input into the several color channels.
Several color representations from the `Colortypes.jl` package are supported, so that besides RGB (XYZ)-plots
one can create HSL, HSI, HSV or Lab and Luv plots.

### Keyword arguments

* `dmin, dmax` Minimum and maximum value to be used for color transformation, can be either a single value or a tuple, when min/max values are given for each channel
* `rgbaxis` which axis should be used to select RGB channels from
* `oceancol` color to fill the ocean with, defaults to `colorant"darkblue"`
* `misscol` color to represent missing values, defaults to `colorant"gray"`
* `labels` given a list of labels this will create a plot with a non-continouous color scale where integer cube values [1..N] are mapped to the given labels.
* `cType` ColorType to use for the color representation. Can be one of `RGB`, `XYZ`, `Lab`, `Luv`, `HSV`, `HSI`, `HSL`
* `dim=value` can set other dimensions to certain values, for example `var="air_temperature_2m"` will fix the variable for the resulting plot
* `c1`, `c2`, `c3` values on the first, second, third colour channel

If a dimension is neither longitude or latitude and is not fixed through an additional keyword, a slider or dropdown menu will appear to select the axis value.
"""
function plotMAPRGB(cube;dmin=zero(eltype(cube)),dmax=zero(eltype(cube)),
  rgbaxis="Var",oceancol=colorant"darkblue",misscol=colorant"gray",symmetric=false,
  c1 = nothing, c2=nothing, c3=nothing, cType=XYZ, xaxis="Lon",yaxis="Lat",
  overlay=nothing, kwargs...) where T

  dmin,dmax = typed_dminmax2(eltype(cube),dmin,dmax)
  axlist    = caxes(cube)

  overlay = interpretoverlay(overlay,xaxis,yaxis,axlist)

  irgb = findAxis(rgbaxis,axlist)
  if length(axlist[irgb])==3 && c1==nothing && c2==nothing && c3==nothing
    c1=CartesianIndex((1,))
    c2=CartesianIndex((2,))
    c3=CartesianIndex((3,))
  end
  return plotGeneric(MAPPlotRGB(xaxis,yaxis,rgbaxis,dmin,dmax,c1,c2,c3,misscol,oceancol,cType,overlay),cube;kwargs...)
end

@noinline getRGBAR(a,colorm,mi,ma,misscol,oceancol)=RGB{U8}[val2col(a[i,j],colorm,mi,ma,misscol,oceancol) for j=1:size(a,2),i=1:size(a,1)]
@noinline getRGBAR(cType,ar,ag,ab,mi,ma,misscol,oceancol)=map((ar,ag,ab)->RGB(val2col(cType,ar,ag,ab,mi,ma,misscol,oceancol)),ar,ag,ab)
@noinline getRGBAR(a,colorm::Dict,mi,ma,misscol,oceancol)=RGB{U8}[val2col(a[i,j],colorm,misscol,oceancol) for j=1:size(a,2),i=1:size(a,1)]

using ColorTypes
channel_max(d::DataType)="Colortype $d not yet added"
const RGBlike=Union{Type{XYZ},Type{RGB},Type{xyY}}
const HSVlike=Union{Type{HSV},Type{HSI},Type{HSL}}
const Lablike=Union{Type{Lab},Type{Luv}}
channel_min(::RGBlike)=(0.0,0.0,0.0)
channel_max(::RGBlike)=(1.0,1.0,1.0)
channel_min(::Lablike)=(0.0,-170.0,-100.0)
channel_max(::Lablike)=(100.0,100.0,150.0)
channel_min(::HSVlike)=(0.0,0.0,0.0)
channel_max(::HSVlike)=(360.0,1.0,1.0)
channel_names(::Type{RGB})=("R","G","B")
channel_names(::Type{XYZ})=("X","Y","Z")
channel_names(::Type{xyY})=("x","y","Y")
channel_names(::Type{HSV})=("H","S","V")
channel_names(::Type{HSI})=("H","S","I")
channel_names(::Type{HSL})=("H","S","L")
channel_names(::Type{Lab})=("L","a","b")
channel_names(::Type{Luv})=("L","u","v")

import Showoff.showoff
function getlegend(xmin,xmax,colm,legheight,tickspos)
  xoffs=0.05
  xl=1-2xoffs
  if isempty(tickspos)
      tlabs,smin,smax=optimize_ticks(Float64(xmin),Float64(xmax),extend_ticks=false,k_min=4)
  else
      tlabs = tickspos
  end
  tpos=[(tlabs[i]-xmin)/(xmax-xmin) for i=1:length(tlabs)]
  r=rectangle([(i-1)/length(colm) for i in 1:length(colm)],[0],[1/(length(colm)-1)],[1])
  f=fill([colm[div((i-1)*length(colm),length(colm))+1] for i=1:length(colm)])
  bar=compose(context(xoffs,0.35,xl,0.55),r,f,stroke(nothing),svgattribute("shape-rendering","crispEdges"))
  tlabels=compose(context(xoffs,0,xl,0.2),text(tpos,[1],showoff(tlabs),[HCenter()],[VBottom()]))
  dlines=compose(context(xoffs,0.25,xl,0.1),line([[(tpx,0.1),(tpx,0.9)] for tpx in tpos]),stroke(colorant"black"))
  compose(context(0,1Measures.h-legheight,1,legheight),bar,tlabels,dlines)
end

function getlegend(colm,width)
  texth1=Compose.max_text_extents(Compose.default_font_family, Compose.default_font_size, first(keys(colm)))[2]
  texth=texth1*1.05*length(colm)
  yoffs=(Measures.h-texth)/2
  yl=texth
  ncol=length(colm)
  tpos=[(i-0.5)/ncol for i=1:ncol]
  r=Compose.circle([0.5],[(i-0.5)/ncol for i in 1:ncol],[max(1/(ncol-1),0.5)])
  f=fill(collect(values(colm)))
  bar=compose(context(0.9,yoffs,0.1,yl),r,f,stroke(nothing),svgattribute("shape-rendering","crispEdges"))
  tlabels=compose(context(0,yoffs,0.85,yl),Compose.text([1],tpos,collect(keys(colm)),[HRight()],[VCenter()]))
  compose(context(1Measures.w-width,0,width,1),bar,tlabels)
end

function getMinMax(x::AbstractArray{<:Union{T,Missing}};symmetric=false) where T
  mi=typemax(T)
  ma=typemin(T)
  for ix in x
    if !ismissing(ix)
      if ix<mi mi=ix end
      if ix>ma ma=ix end
    end
  end
  if mi==typemax(T) || ma==typemin(T)
    mi,ma=(zero(T),one(T))
  elseif mi==ma
    mi,ma=(mi,mi+1)
  end
  if symmetric
    m=max(abs(mi),abs(ma))
    mi=-m
    ma=m
  end
  mi,ma
end

function val2col(x,colorm,mi,ma,misscol,oceancol)
  N=length(colorm)
  #println(x)
  if !ismissing(x) && x<typemax(x)
    i=ceil(Int,min(N,max(1.0,(x-mi)/(ma-mi)*N)))
    return colorm[i]
  #elseif (m & OCEAN)==OCEAN
  #  return oceancol
  else
    return misscol
  end
end

function val2col(cType,xr,xg,xb,mi,ma,misscol,oceancol)
  mi1,mi2,mi3=channel_min(cType)
  ma1,ma2,ma3=channel_max(cType)
  if !any(ismissing,(xr,xg,xb))
    return cType((xr-mi[1])/(ma[1]-mi[1])*(ma1-mi1)+mi1,(xg-mi[2])/(ma[2]-mi[2])*(ma2-mi2)+mi2,(xb-mi[3])/(ma[3]-mi[3])*(ma3-mi3)+mi3)
  #elseif (mr & OCEAN)==OCEAN
  #  return oceancol
  else
    return misscol
  end
end

function val2col(x,colorm::Dict,misscol,oceancol)
  if !ismissing(x) && !isnan(x)
    return get(colorm,x,misscol)
  else
    return misscol
  end
end

function _makeMaprgb(a::AbstractArray{T},mi,ma,colorm,oceancol,misscol,legPos,iscategorical,symmetric,tickspos) where T
  if iscategorical
    @assert isa(colorm, Tuple)
    colorm,colorm2=colorm
  else
    mi==ma && ((mi,ma)=getMinMax(a,symmetric=symmetric))
    colorm2=nothing
  end
  colorm, colorm2, mi,ma,getRGBAR(a,colorm,convert(T,mi),convert(T,ma),misscol,oceancol)
end
function _makeMap(a::AbstractArray{T},mi,ma,colorm,oceancol,misscol,legPos,iscategorical,symmetric,tickspos,overlay) where T
  if !iscategorical
    mi==ma && ((mi,ma)=getMinMax(a,symmetric=symmetric))
  end
  colorm, colorm2, mi,ma,rgbar = _makeMaprgb(a,mi,ma,colorm,oceancol,misscol,legPos,iscategorical,symmetric,tickspos)
  pngbuf=IOBuffer()
  show(pngbuf,"image/png",rgbar)
  if overlay !== nothing
    ratio = size(rgbar,1)/size(rgbar,2)
    if ratio<1
      ww,hh = Measures.w, Measures.w*ratio
    else
      ww,hh = Measures.h/ratio, Measures.h
    end
    backim = compose(context(0,0,ww,hh),bitmap("image/png",pngbuf.data,0,0,1Measures.w,1Measures.h))
    earthland = compose(context(0,0,ww,hh),overlay,Compose.fill(RGBA(0,0,0,0)),Compose.stroke("white"),Compose.linewidth(0.2))
    rgbar = compose(context(),earthland, backim)
  else
    rgbar = compose(context(0,0,1,1),bitmap("image/png",pngbuf.data,0,0,1,1))
  end
  legheight=legPos==:bottom ? max(0.1*Measures.h,1.6Measures.cm) : 0Measures.h
  legwidth =legPos==:right  ? max(0.2*Measures.w,3.2Measures.cm) : 0Measures.w
  themap=compose(context(0,0,1Measures.w-legwidth,1Measures.h-legheight),rgbar)
  theleg=iscategorical ? getlegend(colorm2,legwidth) : getlegend(mi,ma,colorm,legheight,tickspos)
  compose(context(),themap,theleg)
end


#include("../src/CGP.jl")
export p, funcs, MyInt, Ones, goallist2, goallist3
#using Base.Test

#using Main.CGP

#const MyInt = UInt16   # Also set in CGP.jl

numinputs = 2
numoutputs = 2
nodearity = 2
numinteriors = 5
numlevelsback = 8
const context = construct_contexts(numinputs)[numinputs]

goal1 = MyInt[ 0xb, 0xe ]
goal2 = MyInt[ 0x6, 0x9 ]
goal3 = MyInt[ 0xa, 0xe ]
goallist2 = [goal1, goal2,goal3]
goallist3 = [[0x50], [ 0x05], [ 0xf5], [ 0xa0], [ 0xe0], [ 0x0e], [ 0xf1], [ 0x80], [ 0x08], [ 0xfa]]

p = Parameters(numinputs, numoutputs, nodearity, numinteriors, numlevelsback)
const Ones = Main.CGP.construct_ones(numinputs)[numinputs]

#const funcs = default_funcs()
println("SetParams: ")
print_parameters(p)


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

export
    sssr, atom_sssrmem, bond_sssrmem,
    atom_sssrsizes, bond_sssrsizes,
    isringatom, isringbond,
    atom_sssrcount, bond_sssrcount,
    scaffoldmem, componentmem,
    atomsymbol, charge, multiplicity, bondorder, nodedegree,
    apparentvalence, valence, lonepair, heavyatomcount,
    explicithcount, implicithcount, hcount,
    connectivity, pielectron, ishdonor, ishacceptor, stdweight,
    isrotatable, isaromatic, isaromaticbond,
    molweight,
    hacceptorcount,
    hdonorcount,
    wclogp,
    rotatablecount

# Molecular graph topology

sssr = mincyclenodes
atom_sssrmem = node_cyclemem
bond_sssrmem = edge_cyclemem
scaffoldmem = two_edge_membership
componentmem = connected_membership

function atom_sssrsizes(mol::UndirectedGraph)
    vec = [Set{Int}() for i in 1:nodecount(mol)]
    rings = sssr(mol)
    for (i, cs) in enumerate(atom_sssrmem(mol))
        for c in cs
            push!(vec[i], length(rings[c]))
        end
    end
    return vec
end

atom_sssrsizes(view::SubgraphView) = atom_sssrsizes(view.graph)

function bond_sssrsizes(mol::UndirectedGraph)
    vec = [Set{Int}() for i in 1:edgecount(mol)]
    rings = sssr(mol)
    for (i, cs) in enumerate(bond_sssrmem(mol))
        for c in cs
            push!(vec[i], length(rings[c]))
        end
    end
    return vec
end

bond_sssrsizes(view::SubgraphView) = bond_sssrsizes(view.graph)

# TODO: waiting for fix #28992
# @cache atom_isringmem(mol) = .!isempty.(node_cyclemem(mol))
isringatom(mol::UndirectedGraph) = [!i for i in isempty.(atom_sssrmem(mol))]
isringatom(view::SubgraphView) = isringatom(view.graph)

# @cache bond_isringmem(mol) = .!isempty.(node_cyclemem(mol))
isringbond(mol::UndirectedGraph) = [!i for i in isempty.(bond_sssrmem(mol))]
isringbond(view::SubgraphView) = isringbond(view.graph)

atom_sssrcount(mol::UndirectedGraph) = length.(atom_sssrmem(mol))
atom_sssrcount(view::SubgraphView) = atom_sssrcount(view.graph)

bond_sssrcount(mol::UndirectedGraph) = length.(bond_sssrmem(mol))
bond_sssrcount(view::SubgraphView) = bond_sssrcount(view.graph)

# Elemental properties

"""
    atomsymbol(mol::GraphMol) -> Vector{Int}

Return the size `n` vector of atom symbols within the molecule that have `n`
atoms.
"""
@cache atomsymbol(mol::GraphMol) = getproperty.(nodeattrs(mol), :symbol)
atomsymbol(view::SubgraphView) = atomsymbol(view.graph)

"""
    charge(mol::GraphMol) -> Vector{Int}

Return the size `n` vector of atom charges within the molecule that have `n`
atoms.
"""
@cache charge(mol::GraphMol) = getproperty.(nodeattrs(mol), :charge)
charge(view::SubgraphView) = charge(view.graph)

"""
    multiplicity(mol::GraphMol) -> Vector{Int}

Return the size `n` vector of the atom multiplicity within the molecule that
have `n` atoms (1: non-radical, 2: radical, 3: biradical).
"""
@cache multiplicity(mol::GraphMol) = getproperty.(nodeattrs(mol), :multiplicity)
multiplicity(view::SubgraphView) = multiplicity(view.graph)

"""
    bondorder(mol::GraphMol) -> Vector{Int}

Return the size `n` vector of the bond order within the molecule that have `n`
bonds.
"""
@cache bondorder(mol::GraphMol) = getproperty.(edgeattrs(mol), :order)
bondorder(view::SubgraphView) = bondorder(view.graph)

"""
    nodedegree(mol::GraphMol) -> Vector{Int}

Return the size `n` vector of the node degree within the molecular graph that
have `n` atom nodes.
"""
@cache nodedegree(mol::GraphMol) = [degree(mol, n) for n in 1:nodecount(mol)]
nodedegree(view::SubgraphView) = nodedegree(view.graph)

"""
    apparentvalence(mol::GraphMol) -> Vector{Int}

Return the size `n` vector of the apparent valence (the sum of incident bond
order) within the molecule that have `n` atoms.
"""
@cache function apparentvalence(mol::GraphMol)
    vec = zeros(Int, nodecount(mol))
    bondorder_ = bondorder(mol)
    for i in 1:nodecount(mol)
        for inc in incidences(mol, i)
            vec[i] += bondorder_[inc]
        end
    end
    return vec
end

apparentvalence(view::SubgraphView) = apparentvalence(view.graph)

"""
    valence(mol::GraphMol) -> Vector{Union{Int,Nothing}}

Return the size `n` vector of the intrinsic valence (with considering implicit
hydrogens) within the molecule that have `n` atoms.

Note that implicit hydrogens are available for only organic atoms. The intrinsic
valence value of inorganic atoms would be `nothing`.
"""
@cache function valence(mol::GraphMol)
    defs = Dict(
        :H => 1, :B => 3, :C => 4, :N => 3, :O => 2, :F => 1,
        :Si => 4, :P => 3, :S => 2, :Cl => 1,
        :As => 3, :Se => 2, :Br => 1, :I => 1
    )
    vec = Union{Int,Nothing}[]
    for (sym, chg) in zip(atomsymbol(mol), charge(mol))
        num = get(defs, sym, nothing)
        v = num === nothing ? nothing : num + chg
        push!(vec, v)
    end
    return vec
end

valence(view::SubgraphView) = valence(view.graph)

"""
    lonepair(mol::GraphMol) -> Vector{Union{Int,Nothing}}

Return the size `n` vector of the number of lone pairs within the molecule that
have `n` atoms.

Note that implicit hydrogens are available for only organic atoms. The lonepair
value of inorganic atoms would be `nothing`.
"""
@cache function lonepair(mol::GraphMol)
    defs = Dict(
        :H => 0, :B => -1, :C => 0, :N => 1, :O => 2, :F => 3,
        :Si => 0, :P => 1, :S => 2, :Cl => 3,
        :As => 1, :Se => 2, :Br => 3, :I => 3
    )
    vec = Union{Int,Nothing}[]
    for (sym, chg) in zip(atomsymbol(mol), charge(mol))
        num = get(defs, sym, nothing)
        v = num === nothing ? nothing : num - chg
        push!(vec, v)
    end
    return vec
end

lonepair(view::SubgraphView) = lonepair(view.graph)

"""
    heavyatomcount(mol::GraphMol) -> Vector{Int}

Return the size `n` vector of the number of adjacent non-hydrogen atoms within
the molecule that have `n` atoms.
"""
@cache function heavyatomcount(mol::GraphMol)
    vec = zeros(Int, nodecount(mol))
    atomsymbol_ = atomsymbol(mol)
    for i in 1:nodecount(mol)
        for adj in adjacencies(mol, i)
            if atomsymbol_[adj] != :H
                vec[i] += 1
            end
        end
    end
    return vec
end

heavyatomcount(view::SubgraphView) = heavyatomcount(view.graph)

@cache function implicithcount(mol::GraphMol)
    hcnt = (v, av) -> v === nothing ? 0 : max(0, v - av)
    return hcnt.(valence(mol), apparentvalence(mol))
end

implicithcount(view::SubgraphView) = implicithcount(view.graph)

@cache explicithcount(mol::GraphMol) = nodedegree(mol) - heavyatomcount(mol)
explicithcount(view::SubgraphView) = explicithcount(view.graph)

@cache hcount(mol::GraphMol) = explicithcount(mol) + implicithcount(mol)
hcount(view::SubgraphView) = hcount(view.graph)

@cache connectivity(mol::GraphMol) = nodedegree(mol) + implicithcount(mol)
connectivity(view::SubgraphView) = connectivity(view.graph)

@cache pielectron(mol::GraphMol) = apparentvalence(mol) - nodedegree(mol)
pielectron(view::SubgraphView) = pielectron(view.graph)

@cache function ishdonor(mol::GraphMol)
    dc = (sym, h) -> sym in (:N, :O) && h > 0
    return dc.(atomsymbol(mol), hcount(mol))
end

ishdonor(view::SubgraphView) = ishdonor(view.graph)

@cache function ishacceptor(mol::GraphMol)
    ac = (sym, lp) -> lp === nothing ? false : sym in (:N, :O, :F) && lp > 0
    return ac.(atomsymbol(mol), lonepair(mol))
end

ishacceptor(view::SubgraphView) = ishacceptor(view.graph)

function stdweight(mol::GraphMol)
    weight = (atom, imh) -> atomweight(atom) + H_WEIGHT * imh
    return weight.(nodeattrs(mol), implicithcount(mol))
end

stdweight(view::SubgraphView) = stdweight(view.graph)

# Rotatable bonds

"""
    isrotatable(mol::GraphMol)

Return whether the bonds are rotatable or not.
"""
@cache function isrotatable(mol::GraphMol)
    nodedegree_ = nodedegree(mol)
    isringbond_ = isringbond(mol)
    bondorder_ = bondorder(mol)
    vec = Bool[]
    for (i, (u, v)) in enumerate(edgesiter(mol))
        rot = (!isringbond_[i] && bondorder_[i] == 1
            && nodedegree_[u] != 1 && nodedegree_[v] != 1)
        push!(vec, rot)
    end
    return vec
end

isrotatable(view::SubgraphView) = isrotatable(view.graph)

# Aromatic rings

@cache function isaromaticring(mol::GraphMol)
    vec = falses(circuitrank(mol))
    for (i, ring) in enumerate(sssr(mol))
        if satisfyhuckel(mol, ring)
            vec[i] = true
        elseif nodeattrtype(mol) === SmilesAtom # SMILES aromatic atom
            sub = nodesubgraph(mol, Set(ring))
            if all(nodeattr(mol, n).isaromatic for n in nodeset(sub))
                vec[i] = true
            end
        end
    end
    return vec
end

"""
    isaromatic(mol::GraphMol)

Return the vector whether the atom belongs to an aromatic ring.

Note that aromaticity described here means simplified binary descriptor
(aromatic or not) based on classical Huckel's rule. This is intended for use in
some kind of pharmaceutical researches. Non-classical aromaticity such as
Moebius aromaticity is not considered.
"""
@cache function isaromatic(mol::GraphMol)
    aromatic = falses(nodecount(mol))
    for ring in sssr(mol)[isaromaticring(mol)]
        sub = nodesubgraph(mol, Set(ring))
        for n in nodeset(sub)
            aromatic[n] = true
        end
    end
    return aromatic
end

isaromatic(view::SubgraphView) = isaromatic(view.graph)

@cache function isaromaticbond(mol::GraphMol)
    aromaticbond = falses(edgecount(mol))
    for ring in sssr(mol)[isaromaticring(mol)]
        sub = nodesubgraph(mol, Set(ring))
        for e in edgeset(sub)
            aromaticbond[e] = true
        end
    end
    return aromaticbond
end

isaromaticbond(view::SubgraphView) = isaromaticbond(view.graph)

function satisfyhuckel(mol::GraphMol, ring)
    cnt = 0
    atomsymbol_ = atomsymbol(mol)
    nodedegree_ = nodedegree(mol)
    pielectron_ = pielectron(mol)
    lonepair_ = lonepair(mol)
    carbonylO = findall(
        (atomsymbol_ .== :O) .* (nodedegree_ .== 1) .* (pielectron_ .== 1))
    carbonylC = Int[]
    for o in carbonylO
        c = iterate(adjacencies(mol, o))[1]
        if atomsymbol_[c] == :C
            push!(carbonylC, c)
        end
    end
    for r in ring
        if r in carbonylC
            continue
        elseif pielectron_[r] == 1
            cnt += 1
        elseif lonepair_[r] === nothing
            return false
        elseif lonepair_[r] > 0
            cnt += 2
        elseif lonepair_[r] < 0
            continue
        else
            return false
        end
    end
    return cnt % 4 == 2
end

# Molecular properties

"""
    molweight(mol::GraphMol; digits=2) -> Float64

Return standard molecular weight.
"""
molweight(mol::GraphMol; digits=2
    ) = round(reduce(+, stdweight(mol); init=0), digits=digits)

"""
    hacceptorcount(mol::GraphMol) -> Int

Return the number of hydrogen bond acceptors (N, O and F).
"""
hacceptorcount(mol::GraphMol) = reduce(+, ishacceptor(mol); init=0)

"""
    hdonorcount(mol::GraphMol) -> Int

Return the number of hydrogen bond donors (O and N attached to hydrogens).
"""
hdonorcount(mol::GraphMol) = reduce(+, ishdonor(mol); init=0)

"""
    wclogp(mol::GraphMol) -> Float64

Return predicted logP value calculated by using Wildman and Crippen method.

# Reference

1. Wildman, S. A. and Crippen, G. M. (1999). Prediction of Physicochemical
   Parameters by Atomic Contributions. Journal of Chemical Information and
   Modeling, 39(5), 868–873. https://doi.org/10.1021/ci990307l
"""
wclogp(mol::GraphMol; digits=2
    ) = round(reduce(+, wclogpcontrib(mol); init=0), digits=digits)

"""
    rotatablecount(mol::GraphMol) -> Int

Return the number of rotatable bonds.
"""
rotatablecount(mol::GraphMol) = reduce(+, isrotatable(mol); init=0)


using Images
using Plots
include("funcs.jl")

let

    cd("../../erickson/output_files/d4_p4_spinup")
    @printf "found directory...\n"
    slip_file = collect(eachline(open("slip.dat", "r")))
    stress_data = collect(eachline(open("stress.dat","r")))

    nn = Int64(get_line(slip_file,1)[1])
    y = get_line(slip_file, 2)[3:end]
    slip_data = @view slip_file[3:end , :]

    img1 = load(string("V_png/velocities.", lpad(0,4,"0"), ".png"))
    plt1 = plot(img1, yaxis=false, xaxis=false, ticks=false)
    img2 = load(string("U_png/displacements.", lpad(0,4,"0"), ".png"))
    plt2 = plot(img2,  yaxis=false, xaxis=false, ticks=false)
    plt3 = plot(legend=false, yflip = true, ylabel="Depth(Km)", xlabel="Slip(m)")
    plt4 = plot(legend=false, yflip = true, ylabel="Depth(Km)", xlabel="Stress(MPa)")

    @printf "intialized plotting...\n"
    l = @layout [
        a{.7w} b{.3w}
        c{.7w} d{.3w}
    ]

    offset = zeros(nn)
    inter_flag = 1
    break_count = 0
    inter_count = 0
    cycle_count = 1
    final_index = size(slip_data)[1]

    for i = 1:final_index
        @printf "\r%f" i/final_index
        stress_index = i - break_count

        if inter_flag == 1
            if slip_data[i] == "BREAK"
                inter_flag = 0
                break_count += 1
            else

                slip = get_line(slip_data, i)[3:end]
                stress = get_line(stress_data, stress_index)[3:end]
                plt3 = plot!(plt3, slip .- offset, y, linecolor=:blue, linewidth=.1, title=string("Cycle # ", cycle_count))
                plt4 = plot(stress, y,
                            legend=false,
                            linecolor=:black,
                            flip = true,
                            ylabel="Depth(Km)",
                            xlabel="Stress(MPa)")
                plot(plt1, plt3, plt2, plt4, layout=l)
                #gui()
                inter_count += 1
            end
        else

            if slip_data[i] == "BREAK"
                inter_flag = 1
                break_count += 1
                cycle_count += 1
                if break_count == 12
                    offset .= get_line(slip_data, i+1)[3:end]
                    plt3 = plot(legend=false, yflip = true, ylabel="Depth(Km)", xlabel="Slip(m)")
                end
            else

                slip = get_line(slip_data, i)
                V_max = slip[2]

                if V_max > log(.1)

                    particle_index = i - break_count - inter_count
                    img1 = load(string("V_png/velocities.", lpad(particle_index,4,"0"), ".png"))
                    plt1 = plot(img1, yaxis=false, xaxis=false, ticks=false)
                    img2 = load(string("U_png/displacements.", lpad(particle_index,4,"0"), ".png"))
                    plt2 = plot(img2, yaxis=false, xaxis=false, ticks=false)
                    slip = slip[3:end]
                    stress = get_line(stress_data, stress_index)[3:end]
                    plt3 = plot!(plt3, slip .- offset, y, linecolor=:red, linewidth=.1)
                    plt4 = plot(stress, y,
                                legend=false,
                                linecolor=:black,
                                flip = true,
                                ylabel="Depth(Km)",
                                xlabel="Stress(MPa)")

                    plot(plt1, plt3, plt2, plt4, layout=l)
                    #gui()
                end

            end
        end
    end

    gif(anim, "ruptures.gif", fps=16)
    cd("/home/tharvey2/Basin_Simulations/utils")
end


function generate_scenario(scenario, ego_v, hit_point)

    ego_y = 0.0
    ped_x = 100.0
    ped_y = 0.0

    if scenario == "CPCN"
        scenario_id = 1
        ped_v = 5/3.6
        obstacles = [ConvexPolygon([VecE2(ped_x-VehicleDef().length-1, ped_y-2),
                VecE2(ped_x-VehicleDef().length-1, ped_y-2-VehicleDef().width),
                VecE2(ped_x-1, ped_y-2-VehicleDef().width),
                VecE2(ped_x-1, ped_y-2)],4)]

    elseif scenario == "CPAN25"
        scenario_id = 2
        ped_v = 5 /3.6
        obstacles = []

    elseif scenario == "CPAN75"
        scenario_id = 3
        ped_v = 5 /3.6
        hit_point = hit_point + 50
        obstacles = []

    elseif scenario == "CPFA"
        scenario_id = 4
        ped_v = 8 / 3.6
        obstacles = []

    elseif scenario == "FP"
        scenario_id = 9
        ped_v = 5 /3.6
        obstacles = []
    else
        scenario_id = -1
        ped_v = 5 / 3.6
        obstacles = []
    end

    # fix values
    ped_right_side = true
    ped_theta = 1.5707
    ped_t_collision = 3.0

    ped_y_end = VehicleDef().width * hit_point / 100 - VehicleDef().width/2
    if (ped_right_side)
        ped_y_start = ped_y_end - ped_t_collision * ped_v
    else
        ped_y_start = ped_y_end + ped_t_collision * ped_v
    end
    ped_y = ped_y_start
    ego_x = ped_x - ego_v * ped_t_collision - VehicleDef().length/2;
    #println("Scenario: ", scenario, " v_ego=", ego_v*3.6, "km/h v_ped=", ped_v*3.6, "km/h HP=", hit_point, " ped_y: ", ped_y)
    return (ego_x, ego_y, ego_v, ped_x, ped_y, ped_v, ped_theta, obstacles, scenario_id)
end

function evaluate_scenario(ego_x, ego_y, ego_v, ped_x, ped_y, ped_v, ped_theta, obstacles)

    timestep = 0.05

    params = CrosswalkParams()
    if (length(obstacles) > 0)
        params.obstacles = obstacles
        params.obstacles_visible = true
    else
        params.obstacles_visible = false
    end
    params.roadway_length = 400.0

    env = CrosswalkEnv(params);

    ego_id = EGO_ID
    ped_id = 2
    ped2_id = 3
    ped3_id = 4

    # Car definition
    ego_initial_state = VehicleState(VecSE2(ego_x, ego_y, 0.), env.roadway.segments[1].lanes[1], env.roadway, ego_v)
    ego = Vehicle(ego_initial_state, VehicleDef(), ego_id)

    # Pedestrian definition using our new Vehicle definition
    ped_initial_state = VehicleState(VecSE2(ped_x, ped_y, ped_theta), env.crosswalk, env.roadway, ped_v)
    ped = Vehicle(ped_initial_state, AutomotivePOMDPs.PEDESTRIAN_DEF, ped_id)

    ped2 = Vehicle(VehicleState(VecSE2(90., 3., -1.57), env.crosswalk, env.roadway, 0.), AutomotivePOMDPs.PEDESTRIAN_DEF, ped2_id)
    ped3 = Vehicle(VehicleState(VecSE2(103., 10., -1.57), env.crosswalk, env.roadway, 1.), AutomotivePOMDPs.PEDESTRIAN_DEF, ped3_id)

    scene = Scene()
    push!(scene, ego)
    push!(scene, ped)
    push!(scene, ped2)
    push!(scene, ped3)

    pos_noise = 0.0
    vel_noise = 0.0
    false_positive_rate = 0.0
    false_negative_rate = 0.0
    rng = MersenneTwister(1);
    sensor = AutomotiveSensors.GaussianSensor(AutomotiveSensors.LinearNoise(10, pos_noise, 0.0),
                     AutomotiveSensors.LinearNoise(10, vel_noise, 0.0), false_positive_rate, false_negative_rate, rng)

    # define a model for each entities present in the scene
    models = Dict{Int, DriverModel}()

    models[ego_id] = EmergencyBrakingSystem.EmergencySystem(a=LatLonAccel(0.0, 0.0),
        env=env,
        sensor=sensor,
        obstacles=env.obstacles,
        SAFETY_DISTANCE_LON=1.0,
        AX_MAX=-10.0,
        THRESHOLD_COLLISION_RATE = 0.5,
        THRESHOLD_TIME_TO_REACT = 0.99,
        timestep=timestep)

    models[ped_id] = ConstantPedestrian(v_desired=ped_v, dawdling_amp=0.0) # dumb model
    models[ped2_id] = ConstantPedestrian(v_desired=0.0, dawdling_amp=0.05) # dumb model
    models[ped3_id] = ConstantPedestrian(v_desired=1.0, dawdling_amp=0.05) # dumb model

    nticks = 100
    sensor_observations = [Vehicle[]]
    ego_vehicle = Vehicle[]
    ego_a = Float64[]
    collision = Bool[]

    collision_rate = Float64[]  # not used
    ttc = Float64[]             # not used
    risk = Float64[]
    emergency_brake_request = Bool[]      # not used
    prediction_obstacle = Vector{Array{Float64}}()  # not used

    obs_callback = (EmergencyBrakingSystem.ObservationCallback(sensor_observations, ego_vehicle, ego_a, collision, collision_rate, ttc, risk, emergency_brake_request, prediction_obstacle),)

    scenes = simulate(scene, env.roadway, models, nticks, timestep, callbacks=obs_callback)

    return (scenes, timestep, env, sensor, sensor_observations, ego_vehicle, ego_a, collision, collision_rate, ttc, risk, emergency_brake_request, prediction_obstacle)

end

function evaluateScenarioMetric(ego_vehicle, emergency_brake_request, ego_a, collision_vector, ped_x)

    emergency_brake_intervention = false
    collision = false
    dv_collision = 0.

    v = []
    a = []
    a_jerk = 0.
    a_last = 0.
    for i=1:length(ego_vehicle)
        # do not consider acceleration part, only until the pedestrian is reached
        if (ego_vehicle[i].state.v > 0 && ego_vehicle[i].state.posG.x < ped_x - 3 )
        #if (ego_vehicle[i].state.v > 0)
            push!(v,ego_vehicle[i].state.v)
        end

        if (ego_a[i] != 0. && ego_a[i] < 0.)
            push!(a,ego_a[i])
        end

        if ( a_last != ego_a[i] )
            a_jerk = a_jerk + abs(ego_a[i] - a_last)
        end
        a_last = ego_a[i]

        if length(emergency_brake_request) > 0 && emergency_brake_request[i] == 1
            emergency_brake_intervention = true
        end

    end
    v_mean = mean(v)
    if ( length(a) > 0)
        a_mean = mean(a)
        a_min = minimum(a)
    else
        a_mean = 0.
        a_min = 0.
    end

    if (collision_vector[end] == true && ego_vehicle[end].state.posG.x < ped_x - 2.)   # ignore collisions when the pedestrian runs into the side of the car (ego vehicle too slow)
        collision = true
        dv_collision = ego_vehicle[end].state.v
    end

    return (collision, emergency_brake_intervention, dv_collision, v_mean, a_mean, a_jerk, a_min)
end

function evaluateScenariosMetric(results)

    sum_collisions = 0
    sum_eb = 0

    dv = []
    v_mean = []
    a_mean = []
    a_jerk = []
    a_min = []

    for i=1:length(results)
        #println(results[i])
        if ( results[i][4] == 1.0 )
            sum_collisions += 1
        end
        if ( results[i][5] == 1.0 )
            sum_eb += 1
        end

        push!(dv,results[i][6])
        push!(v_mean,results[i][7])
        push!(a_mean,results[i][8])
        push!(a_jerk,results[i][9])
        push!(a_min,results[i][10])

    end
    dv = mean(dv)
    v_mean = mean(v_mean)
    a_mean = mean(a_mean)
    a_jerk = mean(a_jerk)
    a_min = mean(a_min)

    return (sum_collisions, sum_eb, dv, v_mean, a_mean, a_jerk, a_min)
end


"""
Container type for the environment parameters
"""
@with_kw mutable struct UrbanParams
    # geometry
    x_min::Float64 = -30.0
    x_max::Float64 = 30.0
    y_min::Float64 = -30.0
    inter_x::Float64 = 0.
    inter_y::Float64 = 0.
    inter_width::Float64 = 6.0
    inter_height::Float64 = 1.0
    lane_width::Float64 = 3.0
    nlanes_main::Int64 = 2
    nlanes::Int64 = 1
    # crosswalk
    crosswalk_pos::Vector{VecSE2} = [VecSE2(6, 0., pi/2), VecSE2(-6, 0., pi/2), VecSE2(0., -8., 0.)]
    crosswalk_length::Vector{Float64}  = [20.0, 20., 20.0]
    crosswalk_width::Vector{Float64} = [4.0, 4.0, 3.1]

    stop_line::Float64 = 19. # in m along the ego car lane
    # cars
    speed_limit::Float64 = 10.0 # in m/s
    car_rate::Float64 = 0.3 # probability of a car appearing every time steps
    # peds
    ped_rate::Float64 = 0.3
    ped_max_speed::Float64 = 2.0

end

"""
Urban Environment with intersections and crosswalk
"""
mutable struct UrbanEnv <: OccludedEnv
    roadway::Roadway
    ped_roadway::Roadway
    crosswalks::Vector{Lane{Float64}}
    obstacles::Vector{ConvexPolygon}
    params::UrbanParams
    priorities::Dict{Tuple{LaneTag, LaneTag}, Bool}
    directions::Dict{Symbol, Tuple{LaneTag, LaneTag}}
end

function UrbanEnv(;params=UrbanParams())
    intersection_params = TInterParams(params.x_min, params.x_max, params.y_min, params.inter_x,
                                       params.inter_y, params.inter_width, params.inter_height,
                                       params. lane_width, params.nlanes_main, params.nlanes,
                                       params. stop_line, params.speed_limit, params.car_rate)
    roadway = gen_T_roadway(intersection_params)
    ped_roadway = Roadway()
    crosswalks = Lane[]
    cw_id = length(roadway.segments)+1
    for i =1:length(params.crosswalk_pos)
        cw_tag = LaneTag(cw_id, 1)

        crosswalk = Lane(cw_tag, gen_straight_curve(VecE2(params.crosswalk_pos[i]) - polar(params.crosswalk_length[i]/2, params.crosswalk_pos[i].θ),
                                                              VecE2(params.crosswalk_pos[i]) + polar(params.crosswalk_length[i]/2, params.crosswalk_pos[i].θ),
                                                               2), width = params.crosswalk_width[i])
        cw_segment = RoadSegment(length(roadway.segments)+1, [crosswalk])
        push!(crosswalks, crosswalk)
        push!(roadway.segments, cw_segment)
        push!(ped_roadway.segments, cw_segment)
        cw_id += 1
    end
    obstacles = Vector{ConvexPolygon}[]
    priorities, directions = priority_map(roadway)
    return UrbanEnv(roadway, ped_roadway, crosswalks, obstacles, params, priorities, directions)
end

function car_roadway(env::UrbanEnv)
    params = env.params
    intersection_params = TInterParams(params.x_min, params.x_max, params.y_min, params.inter_x,
                                    params.inter_y, params.inter_width, params.inter_height,
                                    params. lane_width, params.nlanes_main, params.nlanes,
                                    params. stop_line, params.speed_limit, params.car_rate)
    car_roadway = gen_T_roadway(intersection_params)
    return car_roadway
end

# return (turn, right)::Tuple{Bool, Bool}
function turn_map(env::UrbanEnv, route::Tuple{LaneTag, LaneTag})
    if route ∈ [(LaneTag(3,1), LaneTag(5,1)), (LaneTag(6,1), LaneTag(4,1))]
        return (true, true)
    elseif route ∈ [(LaneTag(6,1), LaneTag(2,2)),  (LaneTag(6,1), LaneTag(2,1)), (LaneTag(1,1), LaneTag(5,1)), (LaneTag(1,2), LaneTag(5,1))]
        return (true, false)
    else
        return (false, false)
    end
end


function MOI.get(model::Optimizer, ::MOI.TerminationStatus)
    solution_info = model.inner.solution_info
    if solution_info === nothing
        return MOI.OPTIMIZE_NOT_CALLED
    end
    solver_status = solution_info.solver_status
    model_status = solution_info.model_status

    if solver_status == NORMAL_COMPLETION
        if model_status == OPTIMAL
            return MOI.OPTIMAL
        elseif model_status == INFEASIBLE
            return MOI.INFEASIBLE
        elseif model_status == UNBOUNDED
            return MOI.DUAL_INFEASIBLE
        elseif model_status == INTERMEDIATE_FEASIBLE
            return LOCALLY_SOLVED
        elseif model_status == UNKNOWN
            return MOI.OTHER_ERROR
        end
    elseif solver_status == INSUFFICIENT_MEMORY_FOR_NODES
        return MOI.MEMORY_LIMIT
    elseif solver_status == ITERATION_LIMIT
        return MOI.ITERATION_LIMIT
    elseif solver_status == TIME_LIMIT
        return MOI.TIME_LIMIT
    elseif solver_status == NUMERICAL_SENSITIVITY
        return MOI.NUMERICAL_ERROR
    elseif solver_status == INSUFFICIENT_MEMORY_FOR_SETUP
        return MOI.MEMORY_LIMIT
    elseif solver_status == RESERVED
        return MOI.OTHER_ERROR
    elseif solver_status == TERMINATED_BY_BARON
        return MOI.OTHER_ERROR
    elseif solver_status == SYNTAX_ERROR
        return MOI.INVALID_MODEL
    elseif solver_status == LICENSING_ERROR
        return MOI.OTHER_ERROR
    elseif solver_status == USER_HEURISTIC_TERMINATION
        return MOI.OTHER_LIMIT
    end

    error("Unrecognized Baron status: $solver_status, $model_status")
end

function MOI.get(model::Optimizer, ::MOI.ResultCount)
    return (model.inner.solution_info.feasible_point === nothing) ? 0 : 1
end

function MOI.get(model::Optimizer, ::MOI.PrimalStatus)
    solution_info = model.inner.solution_info
    if solution_info === nothing || solution_info.feasible_point === nothing
        return MOI.NO_SOLUTION
    else
        return solution_info.model_status == UNBOUNDED ? MOI.INFEASIBILITY_CERTIFICATE : MOI.FEASIBLE_POINT
    end
end

MOI.get(model::Optimizer, ::MOI.ObjectiveValue) = model.inner.solution_info.objective_value

function MOI.get(model::Optimizer, ::MOI.VariablePrimal, vi::VI)
    solution_info = model.inner.solution_info
    if solution_info === nothing || solution_info.feasible_point === nothing
        error("VariablePrimal not available.")
    end
    check_variable_indices(model, vi)
    return solution_info.feasible_point[vi.value]
end

function MOI.get(model::Optimizer, ::MOI.ObjectiveBound)
    solution_info = model.inner.solution_info
    return solution_info.dual_bound
end

function MOI.get(model::Optimizer, ::MOI.SolveTime)
    solution_info = model.inner.solution_info
    return solution_info.wall_time
end

# TODO: desirable?
function MOI.get(model::MOIU.CachingOptimizer{BARON.Optimizer}, attr::MOI.ConstraintPrimal, ci::MOI.ConstraintIndex)
    return MOIU.get_fallback(model, attr, ci)
end

# TODO: MOI getter for solvetime


using Test, YaoBase, YaoBlocks, BitBasis, YaoArrayRegister, Random
using StatsBase: mean

@testset "measure ghz" begin
    # GHZ state
    st = normalize!(arrayreg(bit"0000") + arrayreg(bit"1111"))
    Random.seed!(1234)

    # measure it at 1, 2
    # should collapse to 0000 or 1111 since entangled
    g = Measure(4; locs = (1, 2))
    @test occupied_locs(g) == (1, 2)
    st |> g
    @test (copy(st) |> g) isa ArrayReg

    @test g.results[1] == 0 ? st.state[end] == 0 : st.state[1] == 0
    g = Measure(4; locs = (1, 2), resetto = 2)
    @test g.postprocess isa ResetTo{BitStr64{2}}

    m = Measure(4)
    @test occupied_locs(m) == (1, 2, 3, 4)
    @test m.postprocess isa NoPostProcess
end

@testset "resetto" begin
    Random.seed!(1234)

    st = rand_state(5; nbatch = 3)
    g = Measure(5; locs = (1, 2), resetto = bit"00011")
    @test num_measured(g) == 2
    st |> g
    for k = 1:32
        if !(st.state[k] ≈ 0.0)
            @test all(BitStr64{5}(k - 1)[1:2] .== 1)
        end
    end
    @test Measure(5; locs = (1, 2), resetto = 0b0011).postprocess isa ResetTo{BitStr64{2}}
end

@testset "error handling" begin
    @test_throws ErrorException Measure(
        5;
        locs = (1, 2),
        resetto = bit"00011",
        remove = true,
    )
    @test_throws ErrorException mat(Measure(5; locs = (1, 2), resetto = bit"00011"))
end

@testset "op-measures" begin
    Random.seed!(1234)

    reg = rand_state(6, nbatch = 10)
    op = repeat(3, X)

    # measure!
    reg2 = reg |> copy
    res = measure!(op, reg2, 2:4)
    res2 = measure!(op, reg2, 2:4)
    @test size(res) == (10,)
    @test res2 == res

    # measure_resetto! and measure_remove! for operators are no-longer supported due to its ill property.
    reg = repeat(ArrayReg(ComplexF64[1, -1] / sqrt(2.0)), 10)
    @test measure!(X, reg) |> mean ≈ -1
    reg = repeat(ArrayReg(ComplexF64[1.0, 0]), 1000)
    @test abs(measure!(X, reg) |> mean) < 0.1
end

@testset "op-measures" begin
    reg = rand_state(8; nbatch = 32)
    op = repeat(5, X, 1:5)

    # measure!
    reg2 = reg |> copy
    res = measure!(op, reg2, 2:6)
    res2 = measure!(op, reg2, 2:6)
    @test size(res) == (32,)
    @test res2 == res

    # measure
    reg2 = reg |> copy
    res = measure(op, reg2, 2:6; nshots = 100)
    @test size(res) == (100, 32)
    @test reg ≈ reg2

    reg = repeat(ArrayReg([1, -1 + 0im] / sqrt(2.0)), 10)
    @test measure!(X, reg) |> mean ≈ -1
    reg = repeat(ArrayReg([1.0, 0 + 0im]), 1000)
    @test abs(measure!(X, reg) |> mean) < 0.1

    m = Measure(5)
    @test chmeasureoperator(m, X) == Measure(5, operator = X)
end

@testset "measure an operator correctly" begin
    op = kron(Z, Z)
    reg = ArrayReg(ComplexF64[1/sqrt(2), 0, 0, 1/sqrt(2)])
    res = measure!(op, reg)
    @test res == 1
    @test reg ≈ ArrayReg(ComplexF64[1/sqrt(2), 0, 0, 1/sqrt(2)])
    reg = uniform_state(2)
    res = measure!(op, reg)
    @test count(!iszero, reg.state) == 2

    # batched
    reg =
        BatchedArrayReg(reshape(ComplexF64[1/sqrt(2), 0, 0, 1/sqrt(2), 0.5, 0.5, 0.5, 0.5], 4, 2))
    res = measure!(op, reg)
    @test length(reg) == 2 && res[1] == 1
    @test reg.state[:, 1] ≈ ComplexF64[1/sqrt(2), 0, 0, 1/sqrt(2)]
    @test count(!iszero, reg.state) == 4
    @test isnormalized(reg)

    # with virtual dimension
    reg =
        ArrayReg(
            reshape(ComplexF64[1/sqrt(2), 0, 0, 1/sqrt(2), 0.5, 0.5, 0.5, 0.5], 4, 2),
        ) / sqrt(2)
    res = measure!(op, reg)
    @test reg isa ArrayReg
    c = count(!iszero, reg.state)
    @test (c == 4 && res ≈ 1) || (c == 2 && res == -1)
    @test isnormalized(reg)

    # measure zero space
    reg = ArrayReg(ComplexF64[1.0])
    res = measure!(matblock(fill(3.0 + 0im, 1, 1)), reg)
    @test res ≈ 3.0
    @test reg == zero_state(0)
end


using Adapt
using Statistics
using KernelAbstractions: @kernel, @index, Event
using Oceananigans.Grids

"""
    struct ComputedField{X, Y, Z, A, G, O} <: AbstractField{X, Y, Z, A, G}

Type representing a field computed from an operand.
"""
struct ComputedField{X, Y, Z, S, A, G, O} <: AbstractField{X, Y, Z, A, G}
       data :: A
       grid :: G
    operand :: O
     status :: S

    function ComputedField{X, Y, Z}(data, grid, operand; recompute_safely=true) where {X, Y, Z}
        validate_field_data(X, Y, Z, data, grid)

        # Use FieldStatus if we want to avoid always recomputing
        status = recompute_safely ? nothing : FieldStatus(0.0)

        return new{X, Y, Z, typeof(status), typeof(data),
                   typeof(grid), typeof(operand)}(data, grid, operand, status)
    end

    function ComputedField{X, Y, Z}(data, grid, operand, status) where {X, Y, Z}
        validate_field_data(X, Y, Z, data, grid)

        return new{X, Y, Z, typeof(status), typeof(data),
                   typeof(grid), typeof(operand)}(data, grid, operand, status)
    end
end

"""
    ComputedField(operand; data=nothing)

Returns a field whose data is `computed` from `operand`.
If the keyword argument `data` is not provided, memory is allocated to store
the result. The `arch`itecture of `data` is inferred from `operand`.
"""
function ComputedField(operand; data=nothing, recompute_safely=true)

    loc = location(operand)
    arch = architecture(operand)
    grid = operand.grid

    if isnothing(data)
        data = new_data(arch, grid, loc)
        recompute_safely = false
    end

    return ComputedField{loc[1], loc[2], loc[3]}(data, grid, operand; recompute_safely=recompute_safely)
end

"""
    compute!(comp::ComputedField)

Compute `comp.operand` and store the result in `comp.data`.
"""
function compute!(comp::ComputedField{X, Y, Z}, time=nothing) where {X, Y, Z}
    compute_at!(comp.operand, time) # ensures any 'dependencies' of the computation are computed first

    arch = architecture(comp.data)

    workgroup, worksize = work_layout(comp.grid,
                                      :xyz,
                                      include_right_boundaries=true,
                                      location=(X, Y, Z))

    compute_kernel! = _compute!(device(arch), workgroup, worksize)

    event = compute_kernel!(comp.data, comp.operand; dependencies=Event(device(arch)))

    wait(device(arch), event)

    return nothing
end

compute_at!(field::ComputedField{X, Y, Z, <:FieldStatus}, time) where {X, Y, Z} =
    conditional_compute!(field, time)

"""Compute an `operand` and store in `data`."""
@kernel function _compute!(data, operand)
    i, j, k = @index(Global, NTuple)
    @inbounds data[i, j, k] = operand[i, j, k]
end

#####
##### Adapt
#####

Adapt.adapt_structure(to, computed_field::ComputedField) = Adapt.adapt(to, computed_field.data)


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

using JuliaFEM
using JuliaFEM.Testing

@testset "get nodal values" begin
    el1 = Element(Seg2, [1, 2])
    el2 = Element(Seg2, [2, 3])
    X = Dict{Int64, Vector{Float64}}(
        1 => [0.0],
        2 => [1.0],
        3 => [2.0])
    T = Dict{Int64, Vector{Float64}}(
        1 => [0.0],
        2 => [1.0],
        3 => [0.0])
    P = Problem(Heat, "foo", 1)
    push!(P, el1, el2)
    update!(P, "geometry", X)
    update!(P, "temperature", T)
    @test isnan(P("temperature", [-0.1]))
    @test isapprox(P("temperature", [0.0]), [0.0])
    @test isapprox(P("temperature", [0.5]), [0.5])
    @test isapprox(P("temperature", [1.0]), [1.0])
    @test isapprox(P("temperature", [1.5]), [0.5])
    @test isapprox(P("temperature", [2.0]), [0.0])
    @test isnan(P("temperature", [ 2.1]))
end

@testset "interpolate from set of elements" begin
    el1 = Element(Seg2, [1, 2])
    el2 = Element(Seg2, [2, 3])
    X = Dict{Int64, Vector{Float64}}(
        1 => [0.0],
        2 => [1.0],
        3 => [2.0])
    T = Dict{Int64, Vector{Float64}}(
        1 => [0.0],
        2 => [1.0],
        3 => [0.0])
    P = Problem(Heat, "foo", 1)
    push!(P, el1, el2)
    update!(P, "geometry", X)
    update!(P, "temperature", T)
    @test isnan(P("temperature", [-0.1]))
    @test isapprox(P("temperature", [0.0]), [0.0])
    @test isapprox(P("temperature", [0.5]), [0.5])
    @test isapprox(P("temperature", [1.0]), [1.0])
    @test isapprox(P("temperature", [1.5]), [0.5])
    @test isapprox(P("temperature", [2.0]), [0.0])
    @test isnan(P("temperature", [ 2.1]))
end


@testset "RCall-based plot Test" begin
  # testing for absence of errors, not for correctness

  #net = readTopology("(((A,(B)#H1:::0.9),(C,#H1:::0.1)),D);")
  net = readTopology("(((Ag,(#H1:7.159::0.056,((Ak,(E:0.08,#H2:0.0::0.004):0.023):0.078,(M:0.0)#H2:::0.996):2.49):2.214):0.026,(((((Az:0.002,Ag2:0.023):2.11,As:2.027):1.697)#H1:0.0::0.944,Ap):0.187,Ar):0.723):5.943,(P,20):1.863,165);");
  @test_logs plot(net,:RCall);
  @test_logs (:warn, "At least one non-missing edge length: plotting any missing length as 1.0") plot(net,:RCall, useEdgeLength=true);
  @test_logs plot(net,:RCall, showTipLabel=false);
  @test_logs plot(net,:RCall, showNodeNumber=true, showIntNodeLabel=true);
  @test_logs plot(net,:RCall, tipOffset=1, showGamma=true);
  @test_logs plot(net,:RCall, showEdgeLength=true, showEdgeNumber=true);
  @test_logs plot(net,:RCall, edgeColor="tomato4",minorHybridEdgeColor="skyblue",
          majorHybridEdgeColor="tan");
  dat = DataFrame(node=[-5,-10,-1],bs=["90","95","100"],edge=[11,22,26]);
  @test_logs (:warn, "Some node numbers in the nodeLabel data frame are not found in the network:\n -1") plot(net,:RCall, nodeLabel=dat);
  @test_logs plot(net,:RCall, edgeLabel=dat[!,[:edge,:bs]]);

  @test_logs plot(net,:RCall, style=:majortree, arrowlen=0.1);
  @test_logs (:warn, "Style bogus is unknown. Defaulted to :fulltree.") plot(net,:RCall, style=:bogus);

  # coverage for the nomajorchild
  net2 = readTopology("((((B)#H1)#H2,((D,C,#H2)S1,(#H1:::.8,A)S2)S3)S4);")
  @test_logs plot(net2, :R, style=:majortree);

  # plot based on RCall and ape:
  tre = readTopology("(((((((1,2),3),4),5),(6,7)),(8,9)),10);");
  # fixit: plot(tre, :ape)
end


#Script para migrar información de las series de tiempo conjeturadas a la tabla CTD
using  CSV, DataFrames, Dates

abreviaturas_entidades = ["AGU", "BCN", "BCS", "CAM", "CHP", "CHH", "CMX", "COA", "COL", "DUR", "GUA", "GRO", "HID", "JAL", "MIC", "MOR", "MEX", "NAY", "NLE", "OAX", "PUE", "QUE", "ROO", "SLP", "SIN", "SON", "TAB", "TAM", "TLA", "VER", "YUC", "ZAC"]

#Cálculo de los números de tipos de casos por estado a partir de las series de tiempo conjeturadas para construir la tabla resumen.
function tipo_casos(datos, abreviatura_estado)

    #El número total de casos sospechosos en el estado:
    sospechosos = datos[!, Symbol(abreviatura_estado*"_S")][end]
    #El número total de casos confirmados en el estado:
    positivos = datos[!, Symbol(abreviatura_estado*"_P")][end]
    #El número total de casos confirmados fallecidos en el estado:
    fallecidos = datos[!, Symbol(abreviatura_estado*"_PF")][end]

    return [sospechosos, missing, missing, positivos, missing, fallecidos]
end

#Actualiza la tabla cumulativa del día a partir de la fecha y los datos de las series de tiempo conjeturadas.
function fila_actualización()

    if Time(now()) < Time(19)
        fecha = today() - Day(1)
    else
        fecha = today()
    end

    fecha_string = Dates.format(fecha, "yyyymmdd")

    #Cargamos las series de tiempo
    datos = CSV.read("Open_data/Time_series/$(fecha_string[1:6])/Mexico_COVID19_$fecha_string.csv", header = 1, DataFrame)

    #Calcula los casos de cada estado
    casos_estados = map(estado -> tipo_casos(datos, estado), abreviaturas_entidades)

    #Calcula los casos en el país
    acumulado_país = sum(casos_estados)

    #Calcula el número total de muestras analizadas.
    total_tests = acumulado_país[4] + missing
    #Genera un vector con el número de casos positivos importados, locales, total de casos; positivos reportados, sospechosos reportados, negativos reportados, número total de pruebas, número de casos recuperados y número de casos fallecidos (para mantener el orden de la tabla original):
    reporte_país = [acumulado_país[1:4]..., missing, missing, missing, missing, total_tests, acumulado_país[5:6]...]

    #Genera la fila de datos del día correspondiente:
    fila = string.(vcat(casos_estados..., reporte_país))

    #Agrega la fecha a la primer columna:
    pushfirst!(fila, string(fecha))

    #Agrega comas para el csv y lo junta todo en un string
    fila_csv = prod(fila[1:(end - 1)].*",")*fila[end]

    #Elimina los missing (para que quede el espacio vacío)
    fila_csv = replace(fila_csv, r"missing" => "")

    #Agrega el resultado al csv acumulativo:
    open("Mexico_COVID19_CTD.csv", "a") do io

        write(io, fila_csv*"\n")
    end

    return "Done"
end

fila_actualización()