module Abstractsmearing_module
    using LinearAlgebra
    import ..Wilsonloops_module:Wilson_loop_set,make_staples,Wilson_loop_set,
            make_cloverloops,Tensor_derivative_set, make_loops
    import ..AbstractGaugefields_module:AbstractGaugefields,Abstractfields,initialize_TA_Gaugefields,add_force!,
                        exp_aF_U!,clear_U!,add_U!,evaluate_wilson_loops!,exptU!,
                        Traceless_antihermitian_add!,set_wing_U!,Traceless_antihermitian,evaluate_gaugelinks!,
                        construct_Λmatrix_forSTOUT!,Traceless_antihermitian!,shift_U
    import Wilsonloop:Wilsonline,DwDU,make_loopforactions,make_Cμ,derive_U,derive_Udag,get_leftlinks,get_rightlinks
    import ..Verboseprint_mpi:Verbose_print,println_verbose_level1,println_verbose_level2,println_verbose_level3
    #import ..GaugeAction_module:GaugeAction

    abstract type Abstractsmearing end

    struct Nosmearing <: Abstractsmearing
    end

    abstract type CovLayer{Dim} end

    struct CovNeuralnet{Dim} <: Abstractsmearing
        #numlayers::Int64
        layers::Array{CovLayer{Dim},1}
    end

    function get_numlayers(c::CovNeuralnet)
        return length(c.layers)
    end

    function CovNeuralnet(;Dim=4)
        layers = CovLayer{Dim}[]
        return CovNeuralnet{Dim}(layers)
    end

    function Base.push!(c::CovNeuralnet,layer::CovLayer)
        #c.numlayers += 1
        push!(c.layers,layer)
    end

    function Base.getindex(c::CovNeuralnet,i)
        return c.layers[i]
    end

    function Base.show(c::CovNeuralnet)
        numlayers = get_numlayers(c)
        println("num. of layers: ",numlayers)
        for i=1:numlayers
            if i==1
                string = "st"
            elseif i==2
                string = "nd"
            elseif i==3
                string = "rd"
            else
                string = "th"
            end
            println("- $i-$string layer: $(get_name(c[i]))")
            show(c[i])
        end
    end

    function get_name(s::CovLayer)
        error("layer $s has no name")
    end

    include("./stout.jl")
    #include("./gradientflow.jl")

    function set_parameters(s::CovNeuralnet,i,v)
        set_parameters(s[i],v)
    end

    function construct_smearing(smearingparameters,loops_list,L,coefficients,numlayers)
        if smearingparameters == "nothing"
            smearing = Nosmearing()
        elseif smearingparameters == "stout"
            @assert loops_list != nothing "loops should be put if you want to use smearing schemes"
            loops = make_loops(loops_list,L)

            @assert coefficients != nothing "coefficients should be put if you want to use smearing schemes"
            println("stout smearing will be used")
            if numlayers == 1
                input_coefficients = [coefficients]
            else
                input_coefficients = coefficients
            end

            smearing = Stoutsmearing(loops,input_coefficients)
        elseif smearingparameters == "covnet_stout"
            if numlayers == 1
                input_coefficients = [coefficients]
            else
                input_coefficients = coefficients
            end
            println("covnet verion of the stout smearing will be used")
            smearing = CovNeuralnet_STOUT(loops_list,input_coefficients,L;Dim=length(L))
        else
            error("smearing = $smearing is not supported")
        end
        return smearing
    end

    function calc_smearedU(Uin::Array{T,1},smearing;calcdSdU = false,temps = nothing) where T<: AbstractGaugefields
        if smearing != nothing && typeof(smearing) != Nosmearing
            #println(smearing)
            #println(typeof(smearing))
            Uout_multi = apply_smearing_U(Uin,smearing)
            U = Uout_multi[end]

            #=
            if typeof(smearing) <: SmearingParam_single
                Uout_multi = nothing
                U = apply_smearing_U(Uin,smearing)
            elseif typeof(smearing) <: SmearingParam_multi
                Uout_multi = apply_smearing_U(Uin,smearing)
                U = Uout_multi[end]
            else
                error("something is wrong in calc_smearingU")
            end
            =#
            set_wing_U!(U)  #we want to remove this.
            if calcdSdU
                dSdU = [temps[end-3],temps[end-2],temps[end-1],temps[end]]
            else
                dSdU = nothing
            end
        else
            dSdU = nothing
            Uout_multi = nothing
            U = Uin
        end
        return U,Uout_multi,dSdU
    end

    function apply_smearing_U(Uin::Array{T,1},smearing) where T<: Abstractfields
        error("apply_smearing_U is not implemented in type $(typeof(Uin)) and smearing type $(typeof(smearing))")
    end

    function apply_smearing_U(Uin::Array{T,1},smearing::CovNeuralnet{Dim}) where {T<: Abstractfields,Dim}
        temp1  = similar(Uin[1])
        temp2  = similar(Uin[1])
        temp3  = similar(Uin[1])
        temp4  = similar(Uin[1])
        F0 = initialize_TA_Gaugefields(Uin[1])
        numlayers = get_numlayers(smearing)
        Uout_multi = Array{typeof(Uin),1}(undef,numlayers)
        for i=1:numlayers
            Uout_multi[i] = similar(Uin)
        end
        apply_neuralnet!(Uout_multi,smearing,Uin,[temp1,temp2,temp3,temp4],[F0])
        return Uout_multi

        error("apply_smearing_U is not implemented in type $(typeof(Uin)) and smearing type $(typeof(smearing))")
    end

    function back_prop(δL,net::CovNeuralnet{Dim},Uout_multi,Uin) where Dim
        temps = similar(Uout_multi[1])
        temps_F1 = initialize_TA_Gaugefields(temps[1])
        tempf = [temps_F1]

        layer = net.layers[get_numlayers(net)]
        δ_prev = similar(δL)
        δ_current = deepcopy(δL)
        set_wing_U!(δ_current)

        for i=get_numlayers(net):-1:2
            layer = net.layers[i]
            layer_pullback!(δ_prev,δ_current,layer,Uout_multi[i-1],temps,tempf)
            δ_current,δ_prev = δ_prev,δ_current
            #set_wing_U!(δ_current)
        end
        layer = net.layers[1]
        layer_pullback!(δ_prev,δ_current,layer,Uin,temps,tempf)
        δ_current,δ_prev = δ_prev,δ_current

        return δ_current
    end

    function get_parameter_derivatives(δL,net::CovNeuralnet{Dim},Uout_multi,Uin) where Dim
        temps = similar(Uout_multi[1])
        δs = get_δ_from_back_prop(δL,net,Uout_multi,Uin)
        numlayer = get_numlayers(net)
        dSdp = Array{Vector{Float64},1}(undef,numlayer)

        for i=1:numlayer
            layer = net.layers[i]
            U_current = Uout_multi[i]
            dSdp[i] = parameter_derivatives(δs[i],layer,U_current,temps)
        end
        return dSdp
    end

    function get_δ_from_back_prop(δL,net::CovNeuralnet{Dim},Uout_multi,Uin) where Dim
        temps = similar(Uout_multi[1])
        temps_F1 = initialize_TA_Gaugefields(temps[1])
        tempf = [temps_F1]

        layer = net.layers[get_numlayers(net)]
        δ_prev = similar(δL)
        δ_current = deepcopy(δL)
        set_wing_U!(δ_current)

        δs = Array{typeof(δL),1}(undef,get_numlayers(net))

        for i=get_numlayers(net):-1:2
            layer_pullback!(δ_prev,δ_current,layer,Uout_multi[i-1],temps,tempf)
            δs[i] = deepcopy(δ_prev)
            δ_current,δ_prev = δ_prev,δ_current
            #set_wing_U!(δ_current)
        end
        layer = net.layers[1]
        layer_pullback!(δ_prev,δ_current,layer,Uin,temps,tempf)
        δs[1] = deepcopy(δ_prev)
        δ_current,δ_prev = δ_prev,δ_current

        return δs
    end

    function apply_neuralnet!(Uout_multi,net::CovNeuralnet{Dim},Uin,temps,temps_F) where {Dim}
        layer = net.layers[1]
        println_verbose_level3(Uin[1],"apply_neuralnet! on 1st layer")
        apply_layer!(Uout_multi[1],layer,Uin,temps,temps_F)
        #set_wing_U!(Uout_multi[1])
        for i=2:get_numlayers(net)
            println_verbose_level3(Uin[1],"apply_neuralnet! on $i -th layer")
            layer = net.layers[i]
            apply_layer!(Uout_multi[i],layer,Uout_multi[i-1],temps,temps_F)
            #set_wing_U!(Uout_multi[i])
        end
    end

    function apply_layer!(Uout,layer::T,Uin,temps,temps_F) where T <: CovLayer
        error("apply_layer!(Uout,layer,Uin,temps,temps_F) is not implemented with type $(typeof(layer)) of layer.")
    end

    """
    layer_pullback!(δ_prev,δ_next,layer::T,Uprev,temps,tempf)
    This is a function for a back propagation
        δ_next,Uprev -> δ_prev
    """
    function layer_pullback!(δ_prev,δ_next,layer::T,Uprev,temps,tempf) where T <: CovLayer
        error("layer_pullback!(δ_prev,δ_next,layer::T,Uprev,temps,tempf) is not implemented with type $(typeof(layer)) of layer.")
    end

    function parameter_derivatives(δ_current,layer::T,U_current,temps) where T <: CovLayer
        error("parameter_derivatives(δ_current,layer::T,U_current,temps) is not implemented with type $(typeof(layer)) of layer.")
    end

    function apply_smearing_U(U::Array{T,1},smearing::Stoutsmearing) where T<: Abstractfields
        numlayer = length(smearing.ρs)
        Uout_multi = Array{typeof(U),1}(undef,numlayer)
        for i=1:numlayer
            Uout_multi[i] = similar(U)
        end
        #println("smearing.ρs ", smearing.ρs)
        #println("type U ",typeof(Uout_multi))
        #Uout = similar(U)
        calc_stout_multi!(Uout_multi,U,smearing.ρs,smearing.staples_for_stout)
        #Uout_multi = calc_stout_multi!(U,smearing.ρs,smearing.staples_for_stout)
        return Uout_multi
    end

    #=
    function calc_stout_multi(Uin::Array{<: AbstractGaugefields{NC,Dim},1},ρs::Array{Array{T,1},1},staples)  where {NC,Dim,T <: Number}
        numlayer = length(ρs)
        #println("numlayer = ",numlayer,"\t",ρs)
        Utype = eltype(Uin)
        Uout_multi = Array{Array{Utype,1}}(undef,numlayer)
        for i=1:numlayer
            Uout_multi[i] = similar(Uin)
        end
        calc_stout_multi!(Uout_multi,Uin,ρs,staples)

        return Uout_multi
    end
    =#

    function calc_stout_multi!(Uout_multi::Vector{<: Vector{<: AbstractGaugefields{NC,Dim}}},Uin::Array{<: AbstractGaugefields{NC,Dim},1},ρs::Array{Array{T,1},1},staples)  where {NC,Dim,T <: Number}
        numlayer = length(ρs)
        Utmp = similar(Uin)
        #Uout_multi = Array{Array{GaugeFields{SU{NC}},1}}(undef,numlayer)
        U = deepcopy(Uin)
        for i = 1:numlayer
            if i != numlayer
                apply_stout_smearing!(Utmp,U,ρs[i],staples)
                Uout_multi[i] = deepcopy(Utmp)
                Utmp,U = U,Utmp
            else
                apply_stout_smearing!(Uout_multi[i],U,ρs[i],staples)
            end
        end
    end

    function calc_stout_multi!(Uout::Array{<: AbstractGaugefields{NC,Dim},1},Uin::Array{<: AbstractGaugefields{NC,Dim},1},ρs::Array{Array{T,1},1},staples)  where {NC,Dim,T <: Number}
        numlayer = length(ρs)
        Utmp = similar(Uin)
        #Uout_multi = Array{Array{GaugeFields{SU{NC}},1}}(undef,numlayer)
        U = deepcopy(Uin)
        for i = 1:numlayer
            if i != numlayer
                apply_stout_smearing!(Utmp,U,ρs[i],staples)
                Utmp,U = U,Utmp
            else
                apply_stout_smearing!(Uout,U,ρs[i],staples)
            end
        end

    end

    function apply_stout_smearing!(Uout::Array{<: AbstractGaugefields{NC,Dim},1},
        U::Array{<: AbstractGaugefields{NC,Dim},1},ρs::Array{T,1},staples) where {NC,Dim,T <: Number}
        @assert Uout != U "input U and output U should not be same!"
        V  = similar(U[1])
        temp1  = similar(U[1])
        temp2  = similar(U[1])
        temp3  = similar(U[1])
        F0 = initialize_TA_Gaugefields(U[1])

        num = length(ρs)

        for μ=1:Dim
            clear_U!(V)
            for i=1:num
                loops = staples[i][μ]
                evaluate_wilson_loops!(temp3,loops,U,[temp1,temp2])
                add_U!(V,ρs[i],temp3)
            end
            mul!(temp1,V,U[μ]') #U U*V
            clear_U!(F0)
            Traceless_antihermitian_add!(F0,1,temp1)

            exptU!(temp3,1,F0,[temp1,temp2])

            mul!(Uout[μ],temp3,U[μ])
        end
        set_wing_U!(Uout)

        #error("ee")

        #error("not implemented")
    end

    #=

    abstract type SmearingParam end

    struct SmearingParam_nosmearing <: SmearingParam
    end

    abstract type SmearingParam_single <: SmearingParam
    end

    abstract type SmearingParam_multi <: SmearingParam
    end

    mutable struct SmearingParam_stout <: SmearingParam_single
        staples_for_stout::Array{Array{Wilson_loop_set,1},1}
        tensor_derivative::Array{Tensor_derivative_set,1}
        staples_for_stout_dag::Array{Array{Wilson_loop_set,1},1}
        tensor_derivative_dag::Array{Tensor_derivative_set,1}
        ρs::Array{Float64,1}
        #ρs::Array{Float64,1}
    end

    mutable struct SmearingParam_stout_multi <: SmearingParam_multi
        staples_for_stout::Array{Array{Wilson_loop_set,1},1}
        tensor_derivative::Array{Tensor_derivative_set,1}
        staples_for_stout_dag::Array{Array{Wilson_loop_set,1},1}
        tensor_derivative_dag::Array{Tensor_derivative_set,1}
        ρs::Array{Array{Float64,1},1}
        #ρs::Array{Float64,1}
    end

    const Nosmearing = SmearingParam_nosmearing
    const Stout = SmearingParam_stout

    =#
end


"""
The module `PDB` defines types and methods to work with protein structures inside Julia.
It is useful to link structural and sequential information, and
needed for measure the predictive performance at protein contact prediction of mutual information scores.

**Features**

- Read and parse PDF and PDBML files
- Calculate distance and contacts between atoms or residues
- Determine interaction between residues

```julia
using MIToS.PDB
```
"""
module PDB

using RecipesBase           # Plots for PDB Residues
using AutoHashEquals
using StaticArrays
using DataStructures
using PairwiseListMatrices
using NamedArrays
using LightXML
using LinearAlgebra
using Statistics            # mean
using MIToS.Utils
using Formatting

export  # PDBResidues
        PDBResidueIdentifier,
        Coordinates,
        PDBAtom,
        PDBResidue,
        squared_distance,
        distance,
        contact,
        isresidue,
        isatom,
        residues,
        @residues,
        residuesdict,
        @residuesdict,
        atoms,
        @atoms,
        findheavy,
        findatoms,
        findCB,
        selectbestoccupancy,
        bestoccupancy,
        residuepairsmatrix,
        proximitymean,
        # AtomsData
        covalentradius,
        vanderwaalsradius,
        check_atoms_for_interactions,
        # Interaction
        ishydrophobic,
        isaromatic,
        iscationic,
        isanionic,
        ishbonddonor,
        ishbondacceptor,
        hydrogenbond,
        vanderwaals,
        vanderwaalsclash,
        covalent,
        disulphide,
        aromaticsulphur,
        pication,
        aromatic,
        ionic,
        hydrophobic,
        stridehydrogenbond,
        chimerahydrogenbond,
        # PDBParser
        PDBFile,
        # PDBMLParser
        PDBML,
        downloadpdb,
        PDBMLHeader,
        downloadpdbheader,
        getpdbdescription,
        # Kabsch
        kabsch,
        center!,
        rmsd,
        getCA,
        CAmatrix,
        coordinatesmatrix,
        change_coordinates,
        centeredcoordinates,
        centeredresidues,
        superimpose,
        mean_coordinates,
        rmsf,
        # MIToS.Utils
        All,
        # Imported from Base (and exported for docs)
        any,
        parse,
        print,
        angle

include("PDBResidues.jl")
include("AtomsData.jl")
include("Interaction.jl")
include("PDBParser.jl")
include("PDBMLParser.jl")
include("Kabsch.jl")
include("Plots.jl")

end


using
    Oceananigans,
    Oceananigans.Operators,
    Oceananigans.Grids,
    Oceananigans.Solvers,
    Oceananigans.Diagnostics,
    Oceananigans.OutputWriters,
    Oceananigans.TurbulenceClosures,
    Test,
    Random,
    JLD2,
    Printf,
    Statistics,
    OffsetArrays,
    FFTW

@hascuda using CuArrays

using Statistics: mean
using LinearAlgebra: norm
using GPUifyLoops: @launch, @loop
using NCDatasets: Dataset

using Oceananigans: architecture, device, launch_config, datatuples, with_tracers,
                    Face, Cell, interiorparent, location, TracerFields, fill_halo_regions!

import Oceananigans: interior, datatuple

using Oceananigans.Solvers: PoissonSolver, PPN, PNN, solve_poisson_3d!

using Oceananigans.Diagnostics: run_diagnostic, velocity_div!

using Oceananigans.TimeSteppers: _compute_w_from_continuity!

using Oceananigans.TurbulenceClosures

using Oceananigans.TurbulenceClosures: ∂x_caa, ∂x_faa, ∂x²_caa, ∂x²_faa,
                                       ∂y_aca, ∂y_afa, ∂y²_aca, ∂y²_afa,
                                       ∂z_aac, ∂z_aaf, ∂z²_aac, ∂z²_aaf,
                                       ▶x_caa, ▶x_faa, ▶y_aca, ▶y_afa,
                                       ▶z_aac, ▶z_aaf

using Oceananigans.AbstractOperations

using Oceananigans.AbstractOperations: Computation, compute!

const seed = 420  # Random seed to use for all pseudorandom number generators.

datatuple(A) = NamedTuple{propertynames(A)}(Array(data(a)) for a in A)

function get_output_tuple(output, iter, tuplename)
    file = jldopen(output.filepath, "r")
    output_tuple = file["timeseries/$tuplename/$iter"]
    close(file)
    return output_tuple
end

float_types = (Float32, Float64)

archs = (CPU(),)
@hascuda archs = (CPU(), GPU())

closures = (
            :ConstantIsotropicDiffusivity,
            :ConstantAnisotropicDiffusivity,
            :AnisotropicBiharmonicDiffusivity,
            :SmagorinskyLilly,
            :BlasiusSmagorinsky,
            :RozemaAnisotropicMinimumDissipation,
            :VerstappenAnisotropicMinimumDissipation
           )

@testset "Oceananigans" begin
    include("test_grids.jl")
    include("test_fields.jl")
    include("test_halo_regions.jl")
    include("test_operators.jl")
    include("test_poisson_solvers.jl")
    include("test_coriolis.jl")
    include("test_surface_waves.jl")
    include("test_buoyancy.jl")
    include("test_models.jl")
    include("test_time_stepping.jl")
    include("test_boundary_conditions.jl")
    include("test_forcings.jl")
    include("test_turbulence_closures.jl")
    include("test_dynamics.jl")
    include("test_diagnostics.jl")
    include("test_output_writers.jl")
    include("test_regression.jl")
    include("test_examples.jl")
    include("test_abstract_operations.jl")
    include("test_verification.jl")
end


using Distributed, Distributions, Serialization

test_e_pth = "/bench/PhD/NGS_binaries/BMI/test_e"

@info "Assembling worker pool..."

#DISTRIBUTED CLUSTERs CONSTANTS
remote_machine = "10.0.0.3"
no_local_processes = 1

@info "Spawning local cluster workers..."
worker_pool=addprocs(no_local_processes, topology=:master_worker)

@info "Loading worker libraries everywhere..."
@everywhere using BioMotifInference,Random
@everywhere Random.seed!(myid()*100000)

#JOB CONSTANTS
models_to_permute=5000
func_limit=30
clamp=.02
funcvec=full_perm_funcvec
push!(funcvec, BioMotifInference.permute_source)
push!(funcvec, BioMotifInference.permute_source)
args=[Vector{Tuple{Symbol,Any}}() for i in 1:length(funcvec)]
args[end-1]=[(:weight_shift_freq,0.),(:length_change_freq,1.),(:length_perm_range,1:1)]
args[end]=[(:weight_shift_freq,.1),(:length_change_freq,0.),(:weight_shift_dist,Uniform(.00001,.01))]

instruct = Permute_Instruct(funcvec, ones(length(funcvec))./length(funcvec), models_to_permute, func_limit, clamp; args=args)

display_rotation=[true,10,1,[[:tuning_disp,:lh_disp,:src_disp],[:conv_plot,:liwi_disp,:ens_disp]]]

@info "Beginning nested sampling for test ensemble..."
test_e=deserialize(test_e_pth*"/ens")
test_logZ = converge_ensemble!(test_e, instruct, worker_pool, .001, backup=(true,1),tuning_disp=true,lh_disp=true,src_disp=true, disp_rotate_inst=display_rotation)

rmprocs(worker_pool)

@info "Job done!"


import .ArcStd: State, expand_gold, expand_pred, toheads

type DepParser
end

const num_feats = 1 << 26
const weights = [fill(0.0,num_feats) for i=1:3]

function tohash(data::Vector{Int})
    hash = 1
    for x in data
        hash += x
        hash += (hash << 10)
        hash $= (hash >> 6)
    end
    hash += (hash << 3)
    hash $= (hash >> 11)
    hash += (hash << 15)
    abs(hash) % num_feats + 1
end

function scorefun(s::State, acts::Vector{Int})
    feats = Int[]
    tokens = s.tokens
    s1 = tokens[s.top]
    s2 = s.left == nothing ? nulltoken : tokens[s.left.top]
    b1 = s.right <= length(tokens) ? tokens[s.right] : nulltoken
    b2 = s.right+1 <= length(tokens) ? tokens[s.right+1] : nulltoken
    xs = Any[
    Int[1, s1.formid],
    Int[2, s1.catid],
    Int[3, s1.formid, s2.catid],
    Int[4, s2.formid],
    Int[5, s2.catid],
    Int[6, s2.formid, s2.catid],
    Int[7, b1.formid],
    Int[8, b1.catid],
    Int[9, b1.formid, b1.catid],
    Int[10, s1.formid, s2.formid],
    Int[11, s1.catid, s2.catid],
    Int[12, s1.catid, b1.catid],
    Int[13, s1.formid, s1.catid, s2.catid],
    Int[14, s1.catid, s2.formid, s2.catid],
    Int[15, s1.formid, s2.formid, s2.catid],
    Int[16, s1.formid, s1.catid, s2.formid],
    Int[17, s1.formid, s1.catid, s2.formid, s2.catid],
    Int[18, s1.catid, b1.catid, b2.catid],
    Int[19, s2.catid, s1.catid, b1.catid],
    Int[20, s1.formid, b1.catid, b2.catid],
    Int[21, s2.catid, s1.formid, b1.catid],
    ]
    for x in xs
        push!(feats, tohash(x))
    end
    s.feats = feats

    scores = Float64[]
    for a in acts
        s = mapreduce(f -> weights[a][f], +, feats)
        push!(scores, s)
    end
    scores
end

function train_gold(s::State)
    feats::Vector{Int} = s.prev.feats
    for f in feats
        weights[s.prevact][f] += 1.0
    end
end

function train_pred(s::State)
    feats::Vector{Int} = s.prev.feats
    for f in feats
        weights[s.prevact][f] -= 1.0
    end
end

function decode(p::DepParser, tokens::Vector{Token})
end

function train(p::DepParser, traindata, testdata)
    for epoch = 1:20
        println("epoch: $(epoch)")
        loss = 0.0
        for i in randperm(length(traindata))
            #for x in traindata
            x = traindata[i]
            sy = State(x, scorefun)
            sz = State(x, scorefun)
            y = beamsearch(sy, 1, expand_gold)[end][1]
            z = beamsearch(sz, 1, expand_pred)[end][1]
            loss += z.score - y.score
            max_violation!(y, z, train_gold, train_pred)
        end
        println("loss: $(loss)")
        golds, preds = Int[], Int[]
        for x in testdata
            s = State(x, scorefun)
            z = beamsearch(s, 1, expand_pred)[end][1]
            append!(golds, map(t -> t.headid, x))
            append!(preds, toheads(z))
        end
        acc = accuracy(golds, preds)
        println("test acc: $(acc)")
        println("")
    end
end


export DenseConvDims

"""
    DenseConvDims

Concrete subclass of `ConvDims` for a normal, dense, conv2d/conv3d.
"""
struct DenseConvDims{N, K, S, P, D} <: ConvDims{N}
    input_size::NTuple{N, Int}

    kernel_size::NTuple{K, Int}
    channels_in::Int
    channels_out::Int
    groupcount::Int

    stride::NTuple{S, Int}
    padding::NTuple{P, Int}
    dilation::NTuple{D, Int}
    flipkernel::Bool
end

function DenseConvDims(
    x_size::NTuple{M}, w_size::NTuple{M};
    stride = 1, padding = 0, dilation = 1, groups = 1,
    flipkernel::Bool = false,
) where {M}
    sstride, ppadding, ddilation = check_spdf(
        x_size, w_size, stride, padding, dilation)

    # Ensure channels are equal
    if x_size[end - 1] != w_size[end - 1] * groups
        xs = x_size[end - 1]
        ws = w_size[end - 1]
        throw(DimensionMismatch("Input channels must match! ($xs vs. $ws)"))
    end

    # Ensure groups are valid
    if x_size[end - 1] % w_size[end - 1] != 0 || w_size[end] % groups != 0
        throw(DimensionMismatch(
            "Group count should be divisble by input and output channels ($groups vs. $(w_size[end-1:end]))"))
    end

    DenseConvDims(
        x_size[1:(end - 2)],
        w_size[1:(end - 2)], x_size[end - 1], w_size[end], groups,
        sstride, ppadding, ddilation, flipkernel)
end

function DenseConvDims(x::AbstractArray, w::AbstractArray; kwargs...)
    if ndims(x) != ndims(w)
        throw(DimensionMismatch(
            "Rank of x and w must match! ($(ndims(x)) vs. $(ndims(w)))"))
    end
    return DenseConvDims(size(x), size(w); kwargs...)
end

# Useful for constructing a new DenseConvDims that has only a few elements different
# from the original progenitor object that it inherits shapes from.
@inline DenseConvDims(
    c::C; I=input_size(c), K=kernel_size(c),
    C_in=channels_in(c), C_out=channels_out(c), S=stride(c),
    P=padding(c), D=dilation(c), F=flipkernel(c), G=groupcount(c),
) where C <: ConvDims = DenseConvDims(
    I,
    K, C_in, C_out, G,
    S, P, D, F)

@inline groupcount(c::DenseConvDims) = c.groupcount
@inline channels_in(c::DenseConvDims) = c.channels_in
@inline channels_out(c::DenseConvDims) = c.channels_out

@inline input_size(c::DenseConvDims) = c.input_size
@inline kernel_size(c::DenseConvDims) = c.kernel_size

@inline stride(c::DenseConvDims) = c.stride
@inline padding(c::DenseConvDims) = c.padding
@inline dilation(c::DenseConvDims) = c.dilation
@inline flipkernel(c::DenseConvDims) = c.flipkernel

function check_dims(x::NTuple{M}, w::NTuple{M}, y::NTuple{M}, cdims::DenseConvDims) where {M}
    # First, check that channel counts are all correct:
    @assert x[M-1] * groupcount(cdims) == channels_in(cdims) DimensionMismatch("Data input channel count ($(x[M-1]) vs. $(channels_in(cdims)))")
    @assert y[M-1] == channels_out(cdims) ÷ groupcount(cdims)  DimensionMismatch("Data output channel count ($(y[M-1]) vs. $(channels_out(cdims)))")
    @assert w[M-1] * groupcount(cdims) == channels_in(cdims) DimensionMismatch("Kernel input channel count ($(w[M-1]) vs. $(channels_in(cdims)))")
    @assert w[M] * groupcount(cdims) == channels_out(cdims) DimensionMismatch("Kernel output channel count ($(w[M]) vs. $(channels_out(cdims)))")

    # Next, check that the spatial dimensions match up
    @assert x[1:M-2] == input_size(cdims) DimensionMismatch("Data input spatial size ($(x[1:M-2]) vs. $(input_size(cdims)))")
    @assert y[1:M-2] == output_size(cdims) DimensionMismatch("Data output spatial size ($(y[1:M-2]) vs. $(output_size(cdims)))")
    @assert w[1:M-2] == kernel_size(cdims) DimensionMismatch("Kernel spatial size ($(w[1:M-2]) vs. $(kernel_size(cdims)))")

    # Check the groups match
    @assert channels_in(cdims) % groupcount(cdims) == 0 DimensionMismatch("Groups ($(groupcount(cdims))) should be divisble by input channels $(channels_in(cdims))")

    # Finally, check that the batch size matches
    @assert x[M] == y[M] DimensionMismatch("Batch size ($(x[M]) vs. $(y[M]))")
end


function mcmc(setup_mcmc::MCMCSetup, schedule::MCMCSchedule, setups::Vector{T}) where T <: ModelSetup
    workspaces = [create_workspace(setup) for setup in setups]
    wss, lls = [ws[1] for ws in workspaces], [ws[2] for ws in workspaces]
    θ = workspaces[1][3]
    ws_mcmc = create_workspace(setup_mcmc, schedule, θ)
    # adpt = adaptation_object(setup, ws) # adaptive auxiliary law currently not supported

    aux = nothing
    for step in schedule
        step.save && [save_imputed!(ws) for ws in wss]
        for i in step.idx
            if step.param_updt || typeof(ws_mcmc.updates[i]) <: Imputation
                wss = [next(ws, ws_mcmc.updates[i]) for ws in wss]
                lls, acc, θ = update!(ws_mcmc.updates[i], wss, θ, lls, step, aux)
                aux = aux_params(ws_mcmc.updates[i], aux)
                #TODO: make it nicer, currently we need a little dance here
                dance = typeof(ws_mcmc.updates[i]) <: Imputation
                dance && begin
                    update!(ws_mcmc, acc[1], θ, i)
                    for k in 2:length(setups)
                        register_accpt!(ws_mcmc.updates[i], acc[k])
                    end
                end
                !dance && update!(ws_mcmc, acc, θ, i)
                step.verbose && print("\n")
            end
        end
        step.verbose && print("-----------------------------------------------",
                              "------\n")
        step.readjust && readjust!(ws_mcmc, step.iter)
        step.fuse && fuse!(ws_mcmc, schedule)
        # ll = adaptation!(ws, adpt, step.iter, ll) adaptation currently not supported
    end
    [display_summary(ws, ws_mcmc) for ws in wss]
    wss, ws_mcmc
end

function update!(updt::Imputation{NoBlocking}, wss::Vector{<:Workspace{OS}}, θ,
                 lls, step, ::Any, headstart=false) where OS
    ρ, K = updt.ρs[1], length(wss)
    accpts = repeat([false], K)

    for k in 1:K
        ws = wss[k]
        # sample proposal starting point
        zᵒ, yᵒ = proposal_start_pt(ws, ws.P[1], ρ)

        # sample proposal path
        m = length(ws.WWᵒ)
        success = sample_segments!(1:m, ws, yᵒ, ρ, Val{headstart}())
        llᵒ = success ? (logpdf(ws.x0_prior, yᵒ) +
                         path_log_likhd(OS(), ws.XXᵒ, ws.P, 1:m, ws.fpt) +
                         lobslikelihood(ws.P[1], yᵒ)) : -Inf

        print_info(step, value(lls[k]), value(llᵒ), "impute")

        if accept_sample(llᵒ-lls[k], step.verbose)
            swap!(ws.XX, ws.XXᵒ, ws.WW, ws.WWᵒ, 1:m)
            set!(ws.z, zᵒ)
            lls[k] = llᵒ
            accpts[k] = true
        end
    end
    lls, accpts, θ
end

function update!(pu::ParamUpdate{MetropolisHastingsUpdt},
                 wss::Vector{<:Workspace{OS}}, θ, lls, step,
                 blocking::NoBlocking) where OS
    K = length(wss)

    #NOTE for now, let's sample parameter corresponding to a single path, this
    # will need to be changed for Mixed effect models
    θᵒ = rand(pu.t_kernel, θ, pu.updt_coord)               # sample new parameter
    (logpdf(pu.priors, θᵒ) === -Inf) && (return lls, false, θ)

    llᵒs = copy(lls)
    zᵒs = [copy(ws.z.val) for ws in wss]
    for k in 1:K
        ws = wss[k]
        WW, Pᵒ, P, XXᵒ, XX, fpt = ws.WW, ws.Pᵒ, ws.P, ws.XXᵒ, ws.XX, ws.fpt
        m = length(WW)
        update_laws!(Pᵒ, θᵒ)
        pu.aux.recompute_ODEs && solve_back_rec!(blocking, pu.aux.solver, Pᵒ) # compute (H, Hν, c)

        # find white noise which for a given θᵒ gives a correct starting point
        y = XX[1].yy[1]
        zᵒs[k] = inv_start_pt(y, ws.x0_prior, Pᵒ[1])

        success = find_path_from_wiener!(XXᵒ, y, WW, Pᵒ, 1:m)

        llᵒs[k] = success ? (logpdf(ws.x0_prior, y) +
                             path_log_likhd(OS(), XXᵒ, Pᵒ, 1:m, fpt) +
                             lobslikelihood(Pᵒ[1], y)) : -Inf
    end
    ll, llᵒ = sum(lls), sum(llᵒs)
    print_info(step, value(ll), value(llᵒ))

    llr = ( llᵒ - ll + prior_kernel_contrib(pu, θ, θᵒ))

    # Accept / reject
    if accept_sample(llr, step.verbose)
        for k in 1:K
            ws = wss[k]
            WW, Pᵒ, P, XXᵒ, XX, fpt = ws.WW, ws.Pᵒ, ws.P, ws.XXᵒ, ws.XX, ws.fpt
            m = length(WW)
            swap!(XX, XXᵒ, P, Pᵒ, 1:m)
            set!(ws.z, zᵒs[k])
        end
        return llᵒs, true, θᵒ
    else
        return lls, false, θ
    end
end

function update!(pu::ParamUpdate{ConjugateUpdt}, wss::Vector{<:Workspace{OS}},
                 θ, lls, step, blocking::NoBlocking) where OS
    K = length(wss)

    θᵒ = conjugate_draw(θ, [ws.XX for ws in wss], wss[1].P[1].Target, pu.priors[1], pu.updt_coord)

    total_ll_old = sum(lls)
    for k in 1:K
        ws = wss[k]
        WW, P, XX, fpt = ws.WW, ws.P, ws.XX, ws.fpt
        m = length(WW)

        update_laws!(P, θᵒ)
        pu.aux.recompute_ODEs && solve_back_rec!(blocking, pu.aux.solver, P) # compute (H, Hν, c)

        for i in 1:m    # compute wiener path WW that generates XX
            inv_solve!(EulerMaruyamaBounded(), XX[i], WW[i], P[i])
        end
        # compute white noise that generates starting point
        y = XX[1].yy[1]
        z = inv_start_pt(y, ws.x0_prior, P[1])

        llᵒ = logpdf(ws.x0_prior, y)
        llᵒ += path_log_likhd(OS(), XX, P, 1:m, fpt; skipFPT=true)
        llᵒ += lobslikelihood(P[1], y)
        lls[k] = llᵒ
        set!(ws.z, z)
    end
    print_info(step, sum(total_ll_old), sum(lls))
    return lls, true, θᵒ
end

function conjugate_draw(θ, XX::Vector{<:Vector}, PT, prior, updtIdx)
    μ = mustart(updtIdx)
    𝓦 = μ*μ'
    ϑ = SVector(thetaex(updtIdx, θ))
    for X in XX
        μ, 𝓦 = _conjugate_draw(ϑ, μ, 𝓦, X, PT, updtIdx)
    end
    Σ = inv(𝓦 + inv(Matrix(prior.Σ)))
    Σ = (Σ + Σ')/2 # eliminates numerical inconsistencies
    μ_post = Σ * (μ + Vector(prior.Σ\prior.μ))
    ϑ = rand(Gaussian(μ_post, Σ))
    move_to_proper_place(ϑ, θ, updtIdx)     # align so that dimensions agree
end


export DQNLearner

mutable struct DQNLearner{Tq,Tt,Tf,R<:AbstractRNG} <: Any
    approximator::Tq
    target_approximator::Tt
    loss_func::Tf
    min_replay_history::Int
    update_freq::Int
    update_step::Int
    target_update_freq::Int
    sampler::NStepBatchSampler
    rng::R
    # for logging
    loss::Float32
    is_enable_double_DQN::Bool
end

"""
    DQNLearner(;kwargs...)

See paper: [Human-level control through deep reinforcement learning](https://www.nature.com/articles/nature14236)

# Keywords

- `approximator`::[`AbstractApproximator`](@ref): used to get Q-values of a state.
- `target_approximator`::[`AbstractApproximator`](@ref): similar to `approximator`, but used to estimate the target (the next state).
- `loss_func`: the loss function.
- `γ::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=100`: the frequency of syncing `target_approximator`.
- `stack_size::Union{Int, Nothing}=4`: use the recent `stack_size` frames to form a stacked state.
- `traces = SARTS`: set to `SLARTSL` if you are to apply to an environment of `FULL_ACTION_SET`.
- `rng = Random.GLOBAL_RNG`
- `is_enable_double_DQN = Bool`: enable double dqn, enabled by default.
"""
function DQNLearner(;
    approximator::Tq,
    target_approximator::Tt,
    loss_func::Tf,
    stack_size::Union{Int,Nothing} = nothing,
    γ::Float32 = 0.99f0,
    batch_size::Int = 32,
    update_horizon::Int = 1,
    min_replay_history::Int = 32,
    update_freq::Int = 1,
    target_update_freq::Int = 100,
    traces = SARTS,
    update_step = 0,
    rng = Random.GLOBAL_RNG,
    is_enable_double_DQN::Bool = true,
) where {Tq,Tt,Tf}
    copyto!(approximator, target_approximator)
    sampler = NStepBatchSampler{traces}(;
        γ = γ,
        n = update_horizon,
        stack_size = stack_size,
        batch_size = batch_size,
    )
    DQNLearner(
        approximator,
        target_approximator,
        loss_func,
        min_replay_history,
        update_freq,
        update_step,
        target_update_freq,
        sampler,
        rng,
        0.0f0,
        is_enable_double_DQN,
    )
end

Functors.functor(x::DQNLearner) = (Q = x.approximator, Qₜ = x.target_approximator),
y -> begin
    x = @set x.approximator = y.Q
    x = @set x.target_approximator = y.Qₜ
    x
end

function (learner::DQNLearner)(env)
    env |>
    state |>
    x ->
        Flux.unsqueeze(x, ndims(x) + 1) |>
        x ->
            send_to_device(device(learner), x) |>
            learner.approximator |>
            vec |>
            send_to_host
end

function RLBase.update!(learner::DQNLearner, batch::NamedTuple)
    Q = learner.approximator
    Qₜ = learner.target_approximator
    γ = learner.sampler.γ
    loss_func = learner.loss_func
    n = learner.sampler.n
    batch_size = learner.sampler.batch_size
    is_enable_double_DQN = learner.is_enable_double_DQN
    D = device(Q)

    s, a, r, t, s′ = (send_to_device(D, batch[x]) for x in SARTS)
    a = CartesianIndex.(a, 1:batch_size)

    if is_enable_double_DQN
        q_values = Q(s′)
    else
        q_values = Qₜ(s′)
    end

    if haskey(batch, :next_legal_actions_mask)
        l′ = send_to_device(D, batch[:next_legal_actions_mask])
        q_values .+= ifelse.(l′, 0.0f0, typemin(Float32))
    end

    if is_enable_double_DQN
        selected_actions = dropdims(argmax(q_values, dims = 1), dims = 1)
        q′ = Qₜ(s′)[selected_actions]
    else
        q′ = dropdims(maximum(q_values; dims = 1), dims = 1)
    end

    G = r .+ γ^n .* (1 .- t) .* q′

    gs = gradient(params(Q)) do
        q = Q(s)[a]
        loss = loss_func(G, q)
        ignore() do
            learner.loss = loss
        end
        loss
    end

    update!(Q, gs)
end


###################################################
## network.jl
## the network basics
###################################################

"""
`Node`, represents one network node. `x` and `y` are projections of `lon` and `lat`
- `lon` and `lat` do not have to be defined
"""
struct Node
    x::Float64
    y::Float64
    lon::Float32
    lat::Float32
end
Node(x::Float64,y::Float64) = Node(x,y,0.,0.)

"""
`Road`, represents a road in the network. roadType from 1 to 8
"""
struct Road
    orig::Int
    dest::Int
    # Distance in meters
    distance::Float64
    roadType::Int
end

"""
`Network`, represents a routing network
"""
mutable struct Network
    "Graph of network"
    graph::DiGraph
    "Nodes information"
    nodes::Vector{Node}
    "Roads information"
    roads::Dict{Tuple{Int,Int},Road}
    "Center coordinates for ENU projection (only defined if used)"
    projcenter::Tuple{Float32,Float32}

    function Network(graph::DiGraph, nodes::Vector{Node}, roads::Dict{Tuple{Int,Int},Road})
        n = new()
        n.graph = graph; n.nodes = nodes; n.roads = roads;
        return n
    end
end

function Network(nodes::Vector{Node}, roads::Dict{Tuple{Int,Int},Road})
    g = DiGraph(length(nodes))
    for (s,d) in keys(roads)
        add_edge!(g,s,d)
    end
    return Network(g,nodes,roads)
end

function Base.show(io::IO, n::Network)
    println(io, "Network with $(nNodes(n)) nodes and $(nRoads(n)) edges")
end

"""
`nRoads`: returns number of roads in network (in both directions)
"""
nRoads(n::Network) = ne(n.graph)
"""
`nNodes`: returns number of nodes in network
"""
nNodes(n::Network) = nv(n.graph)


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

# this makes work e.g. TropicalWeight{Float32}(x::TropicalWeight{Float64})
for W in [:TropicalWeight=>:AbstractFloat,
          :ProbabilityWeight=>:AbstractFloat,
          :NLogWeight=>:AbstractFloat,
          :LogWeight=>:AbstractFloat,
         ]
  @eval $(W[1]){T}(x::$(W[1]){TT}) where {T<:$(W[2]),TT<:$(W[2])} = $(W[1]){T}(T(x.x))
  @eval $(W[1]){T}(x::$(W[1]){T}) where {T<:$(W[2])} = x
end
for W in [:LeftStringWeight,:RightStringWeight,:BoolWeight,:ProductWeight]
  @eval $(W)(x::$(W)) = x
end

#parse
for W in [:TropicalWeight,
          :ProbabilityWeight,
          :NLogWeight,
          :LogWeight,
         ]
  @eval parse(::Type{$(W){T}},str) where T = $(W)(parse(T,str))
end
parse(::Type{S},str) where {S <: BoolWeight} = S(parse(Bool,str))


@testset "games simulation" begin

MAX_ACTIONS_PER_GAME = 1000

SPIEL_GAMES_LIST = registered_games()

SPIEL_LOADABLE_GAMES_LIST = [
    g for g in SPIEL_GAMES_LIST if default_loadable(g)
]

@test length(SPIEL_LOADABLE_GAMES_LIST) >= 38

SPIEL_SIMULTANEOUS_GAMES_LIST = [
    g for g in SPIEL_LOADABLE_GAMES_LIST
    if dynamics(g) == SIMULTANEOUS
]

@test length(SPIEL_SIMULTANEOUS_GAMES_LIST) >= 14

SPIEL_MULTIPLAYER_GAMES_LIST = [
    (g, p)
    for g in SPIEL_LOADABLE_GAMES_LIST
    for p in max(min_num_players(g), 2) : min(max_num_players(g), 6)
    if (max_num_players(g) > 2) &&
        (max_num_players(g) > min_num_players(g)) &&
        (short_name(g) != "tiny_hanabi") &&  # default payoff only works for 2p
        (short_name(g) != "universal_poker")
]

@test length(SPIEL_MULTIPLAYER_GAMES_LIST) >= 35

function apply_action_test_clone(state, action)
    state_copy = copy(state)
    @test string(state) == string(state_copy)
    @test history(state) == history(state_copy)

    apply_action(state, action)
    apply_action(state_copy, action)

    @test string(state) == string(state_copy)
    @test history(state) == history(state_copy)
end

function serialize_deserialize(game, state)
    ser_str = serialize_game_and_state(game, state)
    new_game, new_state = deserialize_game_and_state(ser_str)
    @test string(game) == string(new_game)
    @test string(state) == string(new_state)
end

function simulate_game(game)
    @info "simulating game $(short_name(get_type(game)))"
    min_u, max_u = min_utility(game), max_utility(game)
    @test min_u < max_u

    state = new_initial_state(game)
    total_actions = 0

    next_serialize_check = 1

    while !is_terminal(state) && (total_actions <= MAX_ACTIONS_PER_GAME)
        total_actions += 1

        # Serialize/Deserialize is costly. Only do it every power of 2 actions.
        if total_actions >= next_serialize_check
            serialize_deserialize(game, state)
            next_serialize_check *= 2
        end

        # The state can be three different types: chance node,
        # simultaneous node, or decision node
        if is_chance_node(state)
            # Chance node: sample an outcome
            outcomes = chance_outcomes(state)
            @test length(outcomes) > 0
            action_list, prob_list = zip(outcomes...)
            action = action_list[sample(weights(collect(prob_list)))]
            apply_action(state, action)
        elseif is_simultaneous_node(state)
            chosen_actions = [
                rand(legal_actions(state, pid-1))
                for pid in 1:num_players(game)
            ]  # in julia, index starts with 1
            # Apply the joint action and test cloning states.
            apply_action_test_clone(state, chosen_actions)
        else
            # Decision node: sample action for the single current player
            action = rand(legal_actions(state, current_player(state)))
            # Apply action and test state cloning.
            apply_action_test_clone(state, action)
        end
    end

    @test total_actions > 0

    if is_terminal(state)
        # Check there are no legal actions.
        @test length(legal_actions(state)) == 0
        for player in 0:(num_players(game)-1)
            @test length(legal_actions(state, player)) == 0
        end

        utilities = returns(state)

        for u in utilities
            @test u >= min_utility(game)
            @test u <= max_utility(game)
        end

        @info "Simulation of game $game" total_actions utilities
    else
        @info "Simulation of game $game terminated after maximum number of actions $MAX_ACTIONS_PER_GAME"
    end
end

for game_info in SPIEL_LOADABLE_GAMES_LIST
    game = load_game(short_name(game_info))
    @test num_players(game) >= min_num_players(game_info)
    @test num_players(game) <= max_num_players(game_info)
    simulate_game(game)
end

for game_info in SPIEL_SIMULTANEOUS_GAMES_LIST
    converted_game = load_game_as_turn_based(short_name(game_info))
    simulate_game(converted_game)
end

for (game_info, n) in SPIEL_MULTIPLAYER_GAMES_LIST
    game = load_game(short_name(game_info); players=GameParameter(n))
    simulate_game(game)
end

simulate_game(load_game("breakthrough(rows=6,columns=6)"))
simulate_game(load_game("pig(players=2,winscore=15)"))

end


"""

```
contrib(s::Selection, atom_contributions::Array{Float64}, selection)
```

Extract the contribution of a given atom type selection from the
solute or solvent atomic contributions to the MDDF.

`s` here is the solute or solvent selection (type `ComplexMixtures.Selection`)
`atom_contributions` is the `R.solute_atom` or `R.solvent_atom` arrays of the `Result` structure,
and the last argument is the selection of atoms from the solute to be considered, given as a list
of indexes, list of atom names, vector of `PDBTools.Atom`s, or a `PDBTools.Residue`.

"""
function contrib(s::Selection, atom_contributions::Array{Float64}, indexes::Vector{Int})
    nbins = size(atom_contributions, 1)
    c = zeros(nbins)
    # If the selection is a single molecule, the indexes are anything
    if s.nmols == 1
        for it in indexes
            ind = findfirst(isequal(it),s.index)
            if isnothing(ind)
                error("Index $it of input list not found in selection indexes list.")
            end
            c += atom_contributions[:, ind]
        end
    # If more than one molecule, the index must correspond to an atom within one molecule
    else
        for it in indexes
            if it > s.natomspermol
                error(
                    "The index list contains atoms with indexes greater than the number of atoms of the molecule.",
                )
            end
            c += atom_contributions[:, it]
        end
    end
    return c
end

#
# If a list of atom names is provided
#
function contrib(s::Selection, atom_contributions::Array{Float64}, names::Vector{String})
    indexes = Vector{Int}(undef, 0)
    for name in names
        index = findall(isequal(name), s.names)
        if length(index) == 0
            error(" Atom in input list is not part of solvent (or solute): $name")
        end
        append!(indexes, index)
    end
    return contrib(s, atom_contributions, indexes)
end

#
# If a list of atoms of PDBTools.Atom is provided
#
function contrib(
    s::Selection,
    atom_contributions::Array{Float64},
    atoms::Vector{PDBTools.Atom};
    warning = true,
)
    (warning && s.nmols > 1) && warning_nmols_types()
    indexes = [atom.index for atom in atoms]
    # Check which types of atoms belong to this selection
    selected_types = which_types(s, indexes, warning = warning)
    return contrib(s, atom_contributions, selected_types)
end

#
# If a residue of type PDBTools.Residue is provided
#
function contrib(
    s::Selection,
    atom_contributions::Array{Float64},
    residue::Residue;
    warning = true,
)
    (warning && s.nmols > 1) && warning_nmols_types()
    indexes = collect(residue.range)
    # Check which types of atoms belong to this selection
    selected_types = which_types(s, indexes, warning = warning)
    return contrib(s, atom_contributions, selected_types)
end

function warning_nmols_types()
    println(
        "WARNING: there is more than one molecule in this selection. " *
        "Contributions are summed over all atoms of the same type.",
    )
end


"""
Start FlowAtlas using:
```julia
julia> FlowAtlas.run()
```
"""
module FlowAtlas
include_dependency("../Project.toml")

using JSServe: Session, @js_str, file_server, response_404, Routes, App, DOM, Observable, on
using JSServe, HTTP, JSON, Images, FileIO, ImageIO
using GigaScatter, ColorSchemes

using Serialization: serialize,deserialize
using FlowWorkspace: inpolygon
using Base: NamedTuple

using FlowWorkspace, StaticArrays, DataFrames, MetaGraphs, OrderedCollections
using GigaSOM, StatsBase, Impute

include("lib/colors.jl")
include("lib/selection.jl")
include("lib/embed.jl")

include("lib/tile.jl")
include("lib/gate.jl")
include("lib/count.jl")

include("lib/sys.jl")
include("lib/map.jl")
include("lib/sidebar.jl")

############################################## javascript dependencies
const ol = JSServe.Dependency( :ol, [# OpenLayers
    joinpath(@__DIR__,"assets/ol/ol.js"), joinpath(@__DIR__,"assets/ol/ol.css")
])

const d3 = JSServe.Dependency( :d3, [ # data-driven documents
    joinpath(@__DIR__, "assets/d3/d3.v6.min.js"),
])

const FileSaver = JSServe.Dependency( :FileSaver, [ # saving files
    joinpath(@__DIR__, "assets/FileSaver.min.js")
])

const style = JSServe.Dependency( :style, [ # custom styling todo(@gszep) remove remote dep
    "//cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.2/css/all.min.css", joinpath(@__DIR__, "assets/style.css"),
])

const extensions = JSServe.Dependency( :extensions, map(  extension -> joinpath(@__DIR__,extension),
    ["assets/ol/sidebar.css","assets/ol/colorbar.js","assets/Sortable.js","assets/violins.js","assets/boxplots.js","assets/utils.js"]
))

function run( path::String; files::String=joinpath(dirname(path),"*.fcs"), transform::Function=x->asinh(x/250),
        port::Int = 3141, url::String = "http://localhost:$port", cols::Symbol=:union, drop::Union{Vector{String},Nothing}=String[],
        nlevels::Int=10, nbins::Int=50, p::Real=0.1, q::Real=5e-4, channelScheme=reverse(ColorSchemes.matter), labelScheme=ColorSchemes.seaborn_colorblind,
        perplexity::Number=300, maxIter::Integer=10000, hold::Bool=true )

    indexTransform(x::AbstractVector{<:Union{Number,Missing}}) = toIndex(x; nlevels=nlevels)
    channelRange(x::AbstractVector{<:Union{Number,Missing}}) = range(quantile(skipmissing(x),(p,1-p))...,length=nlevels)

    @info "Loading FCS files..."
    data, labels, groups, gating = FlowWorkspace.load( path; files=files, transform=transform, cols=cols)

    keep = map( graph -> map( idx -> MetaGraphs.get_prop(graph,idx,:name), filter(idx -> MetaGraphs.outdegree(graph,idx) == 0, 1:MetaGraphs.nv(graph))), values(gating) )
    select!( labels, isnothing(drop) ? union(keep...,["Unlabelled"]) : Not(drop) )

    names = (
        channels = Base.names(data), populations = Base.names(labels),
        conditions = filter(name -> ~occursin(r"\d", name), Base.names(groups)),
        groups = filter(name -> occursin(r"\d", name), Base.names(groups))
    )

    ###############################################################################
    ########################################################### calculate embedding
    som_path, _ = splitext(path)
    _, embedding = embed(data, path = som_path*".som", perplexity = perplexity, maxIter = maxIter)
    embedding = (
        coordinates = map(SVector, embedding[2,:], -embedding[1,:]),
        array = embedding, extrema = extrema(embedding,dims=2)
    )

    ###############################################################################
    ######################################################### colormaps and filters
    colors = (
        labels = (
            names = OrderedDict([ name=>'#'*hex(get(labelScheme,x)) for (name,x) ∈ zip(Base.names(labels),range(0,1,length=size(labels,2))) ]),
            rows = channelview(fill(parse(RGBA{Float64},"#663F46"), size(data,1))),

        ),
        channels = (
            range =  Dict([ col => channelRange(data[!,col]) for col ∈ Base.names(data) ]),
            rows = combine(data, [ col => indexTransform => col for col ∈ Base.names(data) ]),
            colors = channelview([map(x->RGBA(get(channelScheme, x)),range(0,1,length=nlevels)); RGBA(0,0,0,0)]),
            hex = map(x->'#'*hex(x),colorview(RGB,view(channelview(map(x->RGBA(get(channelScheme, x)),range(0,1,length=nlevels))),1:3,:)))

        )
    )

    flattened_data = vcat(map(col->data[!,col],Base.names(data))...)
    channelMin,channelMax = quantile(skipmissing(flattened_data),(q,1-q))
    for (name,color) ∈ colors.labels.names colors!( Dict(["name"=>name,"color"=>color]),labels,groups,colors) end

    ##################################################### initalise filter settings
    selections = (
        codes = select(hcat(labels,groups), AsTable(:) => ByRow(encode ∘ values) => "encoding")[:,:encoding],
        names = OrderedDict([ name => true for name ∈ [Base.names(labels);Base.names(groups)] ]),
        rows = fill(true, size(data, 1))
    )

    ########################################################################## body
    app = App() do session::Session
        Map = DOM.div(id = "map", class = "sidebar-map")

        JSServe.onload(session, Map, olMap( embedding.extrema, colors; port=port))
        return DOM.div(id = "application", extensions, style, DOM.title("FlowAtlas.jl"), Map,

            HTML("""<script src="$(JSServe.Asset(joinpath(@__DIR__,"assets/ol/sidebar.js")))" ></script>"""),
            sidebar(session, names, colors; port=port)
        )
    end

    ###############################################################################
    ############################################# open app as locally hosted server
    try
        global server = JSServe.Server(app, "127.0.0.1", port;
            verbose = true, routes = Routes( "/" => app,
                r"/assetserver/" * r"[\da-f]"^40 * r"-.*" => file_server,

                r"/\d+/\d+/\d+.png" => x->tile(x,embedding,selections,colors,names),
                r"/colors" => x->colors!(x,labels,groups,colors), r"/selection" => x->selection!(x,selections,names),

                r"/gate" => x->gate(x,data,embedding,selections;channelRange=range(channelMin,channelMax,length=nbins)),
                r"/count" => x->count(x,selections),

                r"/favicon.ico" => context -> HTTP.Response(500),
                r".*" => context -> response_404() )
        )

        browser(url)
        printstyled("[FlowAtlas] ", color = :blue)

        println("Running at $url ~ Happy exploring! ")
        hold && wait()

    catch exception
        if exception isa InterruptException
            close(server)

            printstyled("[FlowAtlas] ", color = :blue)
            println("Closing... have a lovely day! 🎈")

        else
            rethrow()
        end
    end
end

@info """\n
    Welcome to FlowAtlas 🎈
    Start FlowAtlas server using:
        julia> FlowAtlas.run("workspace.wsp"; files="*.fcs")
\n"""
end


# all other glasses should follow the format below, new glasses must be added to OTHER_GLASSES and OTHER_GLASS_NAMES where the index in the array matches the numeric part of the GlassID

module CARGILLE
using ..GlassCat: Glass, GlassID, OTHER
using StaticArrays: SVector

"""    CARGILLE.OG0608
```
ID:                      OTHER:1
RI @ 587nm:              1.457518
Abbe Number:             57.18978
ΔPgF:                    0.008
TCE (÷1e-6):             800.0
Density:                 0.878g/m³
Valid wavelengths:       0.32μm to 1.55μm
Reference Temp:          25°C
```
"""
const OG0608 = Glass(GlassID(OTHER, 1), -2, 1.4451400, 0.0043176, -1.80659e-5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.32, 1.55, -0.0009083144750540808, 0.0, 0.0, 0.0, 0.0, 0.0, 25.0, 0.008, -1.0, -1.0, 800.0, -1.0, 0, -1.0, [SVector(0.32, 0.03, 10.0), SVector(0.365, 0.16, 100.0), SVector(0.4047, 0.40, 100.0), SVector(0.480, 0.71, 100.0), SVector(0.4861, 0.72, 100.0), SVector(0.5461, 0.80, 100.0), SVector(0.5893, 0.90, 100.0), SVector(0.6328, 0.92, 100.0), SVector(0.6439, 0.95, 100.0), SVector(0.6563, 0.96, 100.0), SVector(0.6943, 0.99, 100.0), SVector(0.840, 0.99, 100.0), SVector(0.10648, 0.74, 100.0), SVector(0.1300, 0.39, 100.0), SVector(0.1550, 0.16, 100.0)], 1.457518, -1.0, -1.0, 0, 57.18978, 0, 0.878, -1)

"""    CARGILLE.OG0607
```
ID:                      OTHER:2
RI @ 587nm:              1.457587
Abbe Number:             57.19833
ΔPgF:                    0.008
TCE (÷1e-6):             700.0
Density:                 0.878g/m³
Valid wavelengths:       0.32μm to 1.55μm
Reference Temp:          25°C
```
"""
const OG0607 = Glass(GlassID(OTHER, 2), -2, 1.44503, 0.0044096, -2.85878e-5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.32, 1.55, -0.0009083144750540808, 0.0, 0.0, 0.0, 0.0, 0.0, 25.0, 0.008, -1.0, -1.0, 700.0, -1.0, 0, -1.0, [SVector(0.32, 0.15, 10.0), SVector(0.365, 0.12, 100.0), SVector(0.4047, 0.42, 100.0), SVector(0.480, 0.78, 100.0), SVector(0.4861, 0.79, 100.0), SVector(0.5461, 0.86, 100.0), SVector(0.5893, 0.90, 100.0), SVector(0.6328, 0.92, 100.0), SVector(0.6439, 0.90, 100.0), SVector(0.6563, 0.92, 100.0), SVector(0.6943, 0.98, 100.0), SVector(0.840, 0.99, 100.0), SVector(0.10648, 0.61, 100.0), SVector(0.1300, 0.39, 100.0), SVector(0.1550, 0.11, 100.0)], 1.457587, -1.0, -1.0, 0, 57.19833, 0, 0.878, -1)

"""    CARGILLE.OG081160
```
ID:                      OTHER:3
RI @ 587nm:              1.515549
Abbe Number:             36.82493
ΔPgF:                    0.014
TCE (÷1e-6):             700.0
Density:                 1.11g/m³
Valid wavelengths:       0.32μm to 1.55μm
Reference Temp:          25°C
```
"""
const OG081160 = Glass(GlassID(OTHER, 3), -2, 1.49614, 0.00692199, -8.07052e-5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.32, 1.55, -0.000885983052189022, 0.0, 0.0, 0.0, 0.0, 0.0, 25.0, 0.014, -1.0, -1.0, 700.0, -1.0, 0, -1.0, [SVector(0.32, 0.04, 100.0), SVector(0.365, 0.13, 100.0), SVector(0.4047, 0.26, 100.0), SVector(0.480, 0.48, 100.0), SVector(0.4861, 0.49, 100.0), SVector(0.5461, 0.60, 100.0), SVector(0.5893, 0.68, 100.0), SVector(0.6328, 0.71, 100.0), SVector(0.6439, 0.73, 100.0), SVector(0.6563, 0.74, 100.0), SVector(0.6943, 0.76, 100.0), SVector(0.840, 0.83, 100.0), SVector(0.10648, 0.86, 100.0), SVector(0.1300, 0.89, 100.0), SVector(0.1550, 0.90, 100.0)], 1.515549, -1.0, -1.0, 0, 36.82493, 0, 1.11, -1)

end

const OTHER_GLASSES = [CARGILLE.OG0607, CARGILLE.OG0608, CARGILLE.OG081160]
const OTHER_GLASS_NAMES = ["CARGILLE.OG0607", "CARGILLE.OG0608", "CARGILLE.OG081160"]


#=
 = This module returns a dictionary of symbols and alphanumeric values, as given by a single
 = ASCII line from an International Maritime Meteorological Archive file (IMMA Release 2.5;
 = following http://rda.ucar.edu/datasets/ds540.0/docs/R2.5-imma_short.pdf).  This format
 = consists of a core entry and up to six attachments, some of which are still in the process
 = of being standardized.  Note that the dictionary returns only defined values (that is, the
 = original string if a scale parameter is zero, a scaled float if the scale is positive, and
 = a float converted from a different base if the scale is negative) and for the time being,
 = attachments 5 and 6 are not parsed (at least until 5 is formalized) - RD March 2016.
 =#

module ICOADSDict                                       # first define the values contained in the core
export imma                                             # and attachments by symbol, location, and scale

const symb0 = [:YR :MO :DY  :HR :LAT :LON :IM :ATTC :TI :LI :DS :VS :NID :II :ID :C1 :DI :D :WI  :W :VI :VV :WW :W1 :SLP :A :PPP :IT :AT :WBTI :WBT :DPTI :DPT :SI :SST :N :NH :CL :HI :H :CM :CH :WD :WP :WH :SD :SP :SH]'
const leng0 = [  4   2   2    4    5    6   2     1   1   1   1   1    2   2   9   2   1  3   1   3   1   2   2   1    5  1    3   1   4     1    4     1    4   2    4  1   1   1   1  1   1   1   2   2   2   2   2   2]'
const scal0 = [  1   1   1 0.01 0.01 0.01   1     1   1   1   1   1    1   1   0   0   1  1   1 0.1   1   1   1   1  0.1  1  0.1   1 0.1     1  0.1     1  0.1   1  0.1  1   1   0   1  0   0   0   1   1   1   1   1   1]'

const symb1 = [:ATTI1 :ATTL1 :BSI :B10 :B1 :DCK :SID :PT :DUPS :DUPC :TC :PB :WX :SX :C2 :SQZ :SQA :AQZ :AQA :UQZ :UQA :VQZ :VQA :PQZ :PQA :DQZ :DQA :ND :SF :AF :UF :VF :PF :RF :ZNC :WNC :BNC :XNC :YNC :PNC :ANC :GNC :DNC :SNC :CNC :ENC :FNC :TNC :QCE :LZ :QCZ]'
const leng1 = [     2      2    1    3   2    3    3   2     2     1   1   1   1   1   2    1    1    1    1    1    1    1    1    1    1    1    1   1   1   1   1   1   1   1    1    1    1    1    1    1    1    1    1    1    1    1    1    1    2   1    2]'
const scal1 = [     1      1    0    1   1    1    1   1     1     1   1   1   1   1   1  -36  -36  -36  -36  -36  -36  -36  -36  -36  -36  -36  -36   1 -36 -36 -36 -36 -36 -36  -36  -36  -36  -36  -36  -36  -36  -36  -36  -36  -36  -36  -36  -36    1   1    1]'

const symb2 = [:ATTI2 :ATTL2 :OS :OP :FM :IX :W2 :SGN :SGT :SGH :WMI :SD2 :SP2 :SH2 :IS :ES :RS :IC1 :IC2 :IC3 :IC4 :IC5 :IR :RRR :TR :QCI :QI1 :QI2 :QI3 :QI4 :QI5 :QI6 :QI7 :QI8 :QI9 :QI10 :QI11 :QI12 :QI13 :QI14 :QI15 :QI16 :QI17 :QI18 :QI19 :QI20 :QI21 :HDG :COG :SOG :SLL :SLHH :RWD :RWS]'
const leng2 = [     2      2   1   1   2   1   1    1    1    2    1    2    2    2   1   2   1    1    1    1    1    1   1    3   1    1    1    1    1    1    1    1    1    1    1     1     1     1     1     1     1     1     1     1     1     1     1    3    3    2    2     3    3    3]'
const scal2 = [     1      1   1   1   1   1   1    1    0    1    1    1    1    1   1   1   1    0    0    0    0    0   1    1   1    1    1    1    1    1    1    1    1    1    1     1     1     1     1     1     1     1     1     1     1     1     1    1    1    1    1     1    1    1]'

const symb3 = [:ATTI3 :ATTL3 :CCCC :BUID :BMP :BSWU :SWU :BSWV :SWV :BSAT :BSRH :SRH :SIX :BSST :MST :MSH :BY :BM :BD :BH :BFL]'
const leng3 = [     2      2     4     6    5     4    4     4    4     4     3    3    1     4    1    3   4   2   2   2    2]'
const scal3 = [     1      1     0     0  0.1   0.1  0.1   0.1  0.1   0.1     1    1    1   0.1    1    1   1   1   1   1    1]'

const symb4 = [:ATTI4 :ATTL4 :C1M :OPM :KOV :COR :TOB :TOT :EOT :LOT :TOH :EOH :SIM :LOV :DOS :HOP :HOT :HOB :HOA :SMF :SME :SMV]'
const leng4 = [     2      2    2    2    2    2    3    3    2    2    1    2    3    3    2    3    3    3    3    5    5    2]'
const scal4 = [     1      1    0    1    0    0    0    0    0    0    0    0    0    1    1    1    1    1    1    1    1    1]'

const symb5 = [:ATTI5 :ATTL5 :WFI :WF :XWI :XW :XDI :XD :SLPI :TAI :TA :XNI :XN :OTHERSTUFF5]'
const leng5 = [     2      2    1   2    1   3    1   2     1    1   4    1   2           71]'
const scal5 = [     1      1    0   1    0   1    0   0     0    0   1    0   0            0]'

const symb6 = [:ATTI6 :ATTL6 :ATTE :OTHERSTUFF6]'
const leng6 = [     2      2     1          999]'
const scal6 = [     1      1     0            0]'

function imma(line::AbstractString)
  vals = Dict{Symbol, Any}()                                                  # initialize a "vals" dictionary
  linlen = length(line)

  subsec = 0                                                                  # then loop through the core and
  substa = 1                                                                  # all available attachments until
  while substa < linlen                                                       # we're past the end of the line
    if substa > 1                                                             # (sta/end are section start/end)
      subsec = parse(Int, line[substa:substa+1])
    end

    symb = eval(symbol("symb", subsec))                                       # use the relevant section arrays
    leng = eval(symbol("leng", subsec))
    scal = eval(symbol("scal", subsec))
    sumlen = sum(leng)
    subend = substa + sumlen - 1

    if subsec < 5                                                             # ignore historical/supplemental
      if subend <= linlen                                                     # isolate this section and pad if
        sublin =      line[substa:subend]                                     # at the end of a truncated line
      else
        sublin = rpad(line[substa:   end], sumlen, ' ')
      end
#     println(subsec, " :", sublin, ": ")

      index = [0; cumsum(leng)] + 1                                           # add any valid entries from each
      for (a, sym) in enumerate(symb)                                         # section and return the dictionary
        tmpval = sublin[index[a]:index[a]+leng[a]-1]
        if tmpval != " "^leng[a]
          if scal[a] > 0
            vals[sym] = float(tmpval) * scal[a]
          elseif scal[a] == 0
            vals[sym] = strip(tmpval)
          else
            vals[sym] = float(parse(Int, tmpval, 36))
          end
#         println(sym, " ", vals[sym])
        end
      end
    end
    substa = subend + 1
  end
  return vals
end

end                                                                           # end ICOADSDict module


using SeaPearl
using SeaPearlExtras
using ReinforcementLearning
const RL = ReinforcementLearning
using Flux
using GeometricFlux
using BSON: @save, @load
using Random
using Statistics

include("rewards.jl")
include("features.jl")

# -------------------
# Generator
# -------------------
nqueens_generator = SeaPearl.NQueensGenerator(25)
#model = model_queens(4)
SR = SeaPearl.DefaultStateRepresentation{SeaPearl.DefaultFeaturization, SeaPearl.DefaultTrajectoryState}
#gplot(SR.cplayergraph)

# -------------------
# Internal variables
# -------------------
numInFeatures = SeaPearl.feature_length(SR)

# -------------------
# Experience variables
# -------------------
nbEpisodes = 100
evalFreq = 300
nbInstances = 1
nbRandomHeuristics = 0

# -------------------
# Agent definition
# -------------------
include("agents.jl")

# -------------------
# Value Heuristic definition
# -------------------
learnedHeuristic = SeaPearl.LearnedHeuristic{SR, InspectReward, SeaPearl.FixedOutput}(agent)

# Basic value-selection heuristic
selectMin(x::SeaPearl.IntVar; cpmodel=nothing) = SeaPearl.minimum(x.domain)
heuristic_min = SeaPearl.BasicHeuristic(selectMin)

function select_random_value(x::SeaPearl.IntVar; cpmodel=nothing)
    selected_number = rand(1:length(x.domain))
    i = 1
    for value in x.domain
        if i == selected_number
            return value
        end
        i += 1
    end
    @assert false "This should not happen"
end

randomHeuristics = []
for i in 1:nbRandomHeuristics
    push!(randomHeuristics, SeaPearl.BasicHeuristic(select_random_value))
end

valueSelectionArray = [learnedHeuristic, heuristic_min]
append!(valueSelectionArray, randomHeuristics)
# -------------------
# Variable Heuristic definition
# -------------------
variableSelection = SeaPearl.MinDomainVariableSelection{false}()

# -------------------
# -------------------
# Core function
# -------------------
# -------------------

function trytrain(nbEpisodes::Int)

    metricsArray, eval_metricsArray = SeaPearl.train!(
        valueSelectionArray=valueSelectionArray,
        generator=nqueens_generator,
        nbEpisodes=nbEpisodes,
        strategy=SeaPearl.DFSearch(),
        variableHeuristic=variableSelection,
        out_solver=false,
        verbose = false,
        evaluator=SeaPearl.SameInstancesEvaluator(valueSelectionArray,nqueens_generator; evalFreq = evalFreq, nbInstances = nbInstances),
        restartPerInstances = 1
    )

    #saving model weights

    return metricsArray, eval_metricsArray
end

# -------------------
# -------------------

metricsArray, eval_metricsArray = trytrain(nbEpisodes)
nothing


export BackwardPlanner, BackwardGreedyPlanner, BackwardAStarPlanner

"Heuristic-guided best-first backward search."
@kwdef mutable struct BackwardPlanner <: Planner
    heuristic::Heuristic = GoalCountHeuristic(:backward)
    g_mult::Float64 = 1.0 # Path cost multiplier
    h_mult::Float64 = 1.0 # Heuristic multiplier
    max_nodes::Int = typemax(Int)
    save_search::Bool = false # Flag to save search info
end

"Backward greedy search, with cycle checking."
BackwardGreedyPlanner(heuristic::Heuristic; kwargs...) =
    BackwardPlanner(;heuristic=heuristic, g_mult=0, kwargs...)

"Backward A* search."
BackwardAStarPlanner(heuristic::Heuristic; kwargs...) =
    BackwardPlanner(;heuristic=heuristic, kwargs...)

function solve(planner::BackwardPlanner,
               domain::Domain, state::State, spec::Specification)
    @unpack h_mult, heuristic, save_search = planner
    # Precompute heuristic information
    precompute!(heuristic, domain, state, spec)
    # Convert to backward search goal specification
    spec = BackwardSearchGoal(spec, state)
    state = State(get_goal_terms(spec), PDDL.get_types(state))
    # Initialize search tree and priority queue
    node_id = hash(state)
    search_tree = Dict{UInt,PathNode}(node_id => PathNode(node_id, state, 0))
    est_cost = h_mult * heuristic(domain, state, spec)
    queue = PriorityQueue{UInt,Float64}(node_id => est_cost)
    # Run the search
    status, node_id = search!(planner, domain, spec, search_tree, queue)
    # Reconstruct plan and return solution
    if status != :failure
        plan, traj = reconstruct(node_id, search_tree)
        reverse!(plan); reverse!(traj)
        if save_search
            return PathSearchSolution(status, plan, traj, search_tree, queue)
        else
            return PathSearchSolution(status, plan, traj)
        end
    elseif save_search
        return PathSearchSolution(status, [], [], search_tree, queue)
    else
        return NullSolution()
    end
end

function search!(planner::BackwardPlanner,
                 domain::Domain, spec::BackwardSearchGoal,
                 search_tree::Dict{UInt,PathNode}, queue::PriorityQueue)
    count = 1
    while length(queue) > 0
        # Get state with lowest estimated cost to start state
        node_id = dequeue!(queue)
        node = search_tree[node_id]
        # Return status and current state if search terminates
        if is_goal(spec, domain, node.state)
            return :success, node_id # Start state reached
        elseif count >= planner.max_nodes
            return :max_nodes, node_id # Node budget reached
        end
        count += 1
        # Expand current node
        expand!(planner, node, search_tree, queue, domain, spec)
    end
    return :failure, nothing
end

function expand!(planner::BackwardPlanner, node::PathNode,
                 search_tree::Dict{UInt,PathNode}, queue::PriorityQueue,
                 domain::Domain, spec::BackwardSearchGoal)
    @unpack g_mult, h_mult, heuristic = planner
    state = node.state
    # Iterate over relevant actions
    actions = relevant(state, domain)
    for act in actions
        # Regress (reverse-execute) the action
        next_state = regress(act, state, domain; check=false)
        # Add constraints to regression state
        add_constraints!(spec, state)
        next_id = hash(next_state)
        # Compute path cost
        act_cost = get_cost(spec, domain, state, act, next_state)
        path_cost = node.path_cost + act_cost
        # Update path costs if new path is shorter
        next_node = get!(search_tree, next_id,
                         PathNode(next_id, next_state, Inf))
        cost_diff = next_node.path_cost - path_cost
        if cost_diff > 0
            next_node.parent_id = node.id
            next_node.parent_action = act
            next_node.path_cost = path_cost
            # Update estimated cost from next state to start
            if !(next_id in keys(queue))
                g_val = g_mult * path_cost
                h_val = h_mult * heuristic(domain, next_state, spec)
                enqueue!(queue, next_id, g_val + h_val)
            else
                queue[next_id] -= cost_diff
            end
        end
    end
end


# This file is a part of EncodedArrays.jl, licensed under the MIT License (MIT).

using EncodedArrays
using Test

using BitOperations

@testset "varlen_diff_codec" begin
    a = zigzagdec(UInt64(0x4c6b5f94758))
    data = [a, a + 3, a - 2]
    codec = VarlenDiffArrayCodec()

    encoded = Vector{UInt8}()
    EncodedArrays.encode_data!(encoded, codec, data)
    @test encoded == [0xd8, 0x8e, 0xe5, 0xaf, 0xeb, 0x98, 0x01, 0x06, 0x09]

    data_dec = view(similar(data), :) # Use view to test that decoder doesn't resize
    @test EncodedArrays.decode_data!(data_dec, codec, encoded) === data_dec
    @test data_dec == data

    data_dec = Vector{Int64}()
    @test EncodedArrays.decode_data!(data_dec, codec, encoded) === data_dec
    @test data_dec == data

    data_dec = Vector{Int64}(undef, 2)
    @test EncodedArrays.decode_data!(data_dec, codec, encoded) === data_dec
    @test data_dec == data

    data_dec = Vector{UInt64}()
    @test EncodedArrays.decode_data!(data_dec, codec, encoded) === data_dec
    @test data_dec == data

    b = unsigned(a)
    data = [b, b + 3, b - 2]

    encoded = Vector{UInt8}()
    EncodedArrays.encode_data!(encoded, codec, data)
    @test encoded == [0xd8, 0x8e, 0xe5, 0xaf, 0xeb, 0x98, 0x01, 0x06, 0x09]

    data_dec = view(similar(data), :) # Use view to test that decoder doesn't resize
    @test EncodedArrays.decode_data!(data_dec, codec, encoded) === data_dec
    @test data_dec == data

    data_dec = Vector{Int64}()
    @test EncodedArrays.decode_data!(data_dec, codec, encoded) === data_dec
    @test data_dec == data
end # testset


# test compute fft

T = Float32
fs = 100.  # sampling rate
N = 2^14   # number of points
Nwin = 10
freqmin = 10. # minimum frequency
freqmax = 20. # maximum frequency
corners = 4   # number of corners
zerophase = true # do zerophase filtering
starttime = d2u(DateTime(Date(now()))) # starttime for SeisChanel / SeisData

# test process raw
@testset "process raw" begin

    S = SeisData(2)
    S.x[1] = rand(T,N)
    S.id[1] = "EO.SYT.02.LLE"
    S.fs[1] = fs
    S.t[1] = [1 starttime * 1e6;N 0]
    S.x[2] = rand(T,N)
    S.id[2] = "MZ.KTFD.08.ENZ"
    S.fs[2] = fs
    S.t[2] = [1 (starttime + 0.004) * 1e6;N 0]

    # test without downsampling
    Snew = process_raw(S,fs)
    @test size(Snew.x[1]) == size(S.x[1])
    @test eltype(Snew.x[1]) == eltype(S.x[1])
    @test size(Snew.x[2]) == size(S.x[2])
    @test eltype(Snew.x[2]) == eltype(S.x[2])
    @test Snew.fs == S.fs
    @test Snew.t[1] == S.t[1]
    @test Snew.t[2][1,2] == S.t[2][1,2] - 0.004 * 1e6 # test phase shift

    # test with integer sampling rate
    SInt = process_raw(S,Int(fs))
    @test size(SInt.x[1]) == size(S.x[1])
    @test eltype(SInt.x[1]) == eltype(S.x[1])
    @test size(SInt.x[2]) == size(S.x[2])
    @test eltype(SInt.x[2]) == eltype(S.x[2])
    @test SInt.fs == S.fs
    @test SInt.t[1] == S.t[1]
    @test SInt.t[2][1,2] == S.t[2][1,2] - 0.004 * 1e6 # test phase shift

    # test with downsampling
    Snew = process_raw(S,20.)
    @test size(Snew.x[1]) != size(S.x[1])
    @test eltype(Snew.x[1]) == eltype(S.x[1])
    @test size(Snew.x[2]) != size(S.x[2])
    @test eltype(Snew.x[2]) == eltype(S.x[2])
    @test Snew.fs == [20.,20.]
    @test Snew.t[1][2,1] == length(Snew.x[1])
    @test Snew.t[2][2,1] == length(Snew.x[2])
    @test Snew.t[1] != S.t[1]
    @test Snew.t[2][1,2] == S.t[2][1,2] - 0.004 * 1e6 # test phase shift

    # test without ϕshift
    Sphi = process_raw(S,20.,ϕshift=false)
    @test Sphi.x[1] == Snew.x[1]
    @test Sphi.x[2] != Snew.x[2]
    @test Sphi.t[2] == Snew.t[2]

    # test in-place
    process_raw!(S,20.)
    @test Snew == S

    # test with SeisChannel
    C = SeisChannel()
    C.x = rand(T,N)
    C.fs = fs
    C.t = [1 (starttime + 0.004) * 1e6;N 0]

    # test no downsample
    Cnew = process_raw(C,fs)
    @test size(Cnew.x) == size(C.x)
    @test eltype(Cnew.x) == eltype(C.x)
    @test Cnew.fs == C.fs
    @test Cnew.t[1,2] == C.t[1,2] - 0.004 * 1e6 # test phase shift

    # test with int sampling rate
    CInt = process_raw(C,Int(fs))
    @test size(CInt.x) == size(C.x)
    @test eltype(CInt.x) == eltype(C.x)
    @test CInt.fs == C.fs
    @test CInt.t[1,2] == C.t[1,2] - 0.004 * 1e6 # test phase shift

    # test downsample
    Cnew = process_raw(C,20.)
    @test size(Cnew.x) != size(C.x)
    @test eltype(Cnew.x) == eltype(C.x)
    @test Cnew.fs == 20.
    @test Cnew.t[1,2] == C.t[1,2] - 0.004 * 1e6 # test phase shift

    # test in-place
    process_raw!(C,20.)
    @test Cnew == C
end

# test FFTData creation with rfft
@testset "rfft" begin

    cc_len = 10.24
    cc_step = 5.12
    C = SeisChannel()
    C.x = rand(T,N)
    C.fs = fs
    C.id = "MZ.KTFD.08.ENZ"
    C.t = [1 (starttime + 0.004) * 1e6;N 0]
    R = RawData(C,cc_len,cc_step)
    R.gain = 1.234234e6
    F = rfft(R)
    @test F.name == C.id == R.name
    @test F.id == R.id
    @test F.loc == R.loc
    @test F.fs == R.fs
    @test F.gain == R.gain
    @test F.freqmin == R.freqmin
    @test F.freqmax == F.freqmax
    @test F.cc_len == R.cc_len == cc_len
    @test F.cc_step == R.cc_step == cc_step
    @test F.whitened == R.whitened
    @test F.time_norm == R.time_norm
    @test F.resp == R.resp
    @test F.t == R.t
    @test length(F.t) == size(F.fft,2)
end

@testset "phase" begin

    # test with all zero row
    A = rand(T,N,Nwin)
    A[:,1] .= T(0)
    p = phase(A)
    @test all(.!SeisNoise.isweird.(p)) # test no NaNs

    # test phase
    A = rand(T,N,Nwin)
    p = phase(A)
    @test size(A) == size(p)
    @test all(p[N ÷ 2 + 2 + isodd(N),1] .== T(0))

    # test phase RawData
    R = RawData()
    R.x = A
    R.fs = fs
    F = phase(R)
    @test isa(F,FFTData)
    @test all(.!SeisNoise.isweird.(F.fft))
    @test F.fft == p
end

@testset "mute" begin

    # generate random signal with zero mean
    A = rand(T,N,Nwin)
    A .-= mean(A,dims=1)
    A[1000:1200,:] .*= 1000
    M = mute(A)
    @test size(M) == size(A)
    @test eltype(M) == eltype(A)

    # test begin , middle, end changes
    @test all(M[1:500,:] .== A[1:500,:])
    @test median(M[1000:1200,:]) == T(0)
    @test all(M[end-500:end,:] .== A[end-500:end,:])

    # test in-place
    mute!(A)
    @test A == M

    # test RawData
    R = RawData()
    A = rand(T,N,Nwin)
    A .-= mean(A,dims=1)
    A[1000:1200,:] .*= 1000
    R.x = A
    M = mute(R)
    @test isa(M,RawData)
    @test size(M.x) == size(A)
    @test eltype(M.x) == eltype(A)

    # test in-place
    mute!(R)
    @test R.x == M.x
end

@testset "clip/clamp" begin
    Nspikes = 100
    A = rand(T,N,Nwin) .- 0.5
    ind = rand(1:N,Nspikes)
    A[ind,:] .= randn(Nspikes,Nwin) .* 100000
    A[:,1] .= T(0)

    # test clipping
    B = clip(A,3)
    @test size(B) == size(A)
    @test eltype(B) == eltype(A)
    @test B != A
    @test all(maximum(abs.(B),dims=1) .<= maximum(abs.(A),dims=1))

    # test in-place
    clip!(A,3)
    @test A == B

    # test 1D array
    A = rand(T,N) .- 0.5
    ind = rand(1:N,Nspikes)
    A[ind] .= randn(Nspikes) .* 100000
    B = clip(A,3)
    @test size(B) == size(A)
    @test eltype(B) == eltype(A)
    @test all(maximum(abs.(B),dims=1) .<= maximum(abs.(A),dims=1))

    # test RawData
    Nspikes = 100
    A = rand(T,N,Nwin) .- 0.5
    ind = rand(1:N,Nspikes)
    A[ind,:] .= randn(Nspikes,Nwin) .* 100000
    A[:,1] .= T(0)
    R = RawData()
    R.x = A
    Rclip = clip(R,3)
    @test isa(Rclip,RawData)
    @test size(Rclip.x) == size(R.x)
    @test eltype(Rclip.x) == eltype(R.x)
    @test Rclip.x != R.x
    @test all(maximum(abs.(Rclip.x),dims=1) .<= maximum(abs.(R.x),dims=1))

    # test in-place
    clip!(R,3)
    @test R.x == Rclip.x

    # test clamping w RawData
    A = rand(T,N,Nwin) .- 0.5
    ind = rand(1:N,Nspikes)
    A[ind,:] .= randn(Nspikes,Nwin) .* 100000
    A[:,1] .= T(0)
    R = RawData()
    R.x = deepcopy(A)
    Rclamp = clamp(R,π)
    @test isa(Rclamp,RawData)
    @test size(Rclamp.x) == size(R.x)
    @test eltype(Rclamp.x) == eltype(R.x)
    @test all(abs.(Rclamp.x) .<= π)
    @test Rclamp.x != R.x

    # test in place
    clamp!(R,π)
    @test R.x == Rclamp.x

    # test clamping with two values
    R.x = deepcopy(A)
    Rclamp = clamp(R,-0.75,0.5)
    @test isa(Rclamp,RawData)
    @test size(Rclamp.x) == size(R.x)
    @test eltype(Rclamp.x) == eltype(R.x)
    @test all(Rclamp.x .>= -0.75)
    @test all(Rclamp.x .<= 0.5)

    # test switching hi and lo
    @test_throws ArgumentError clamp(R,0.5,-0.75)

    # test in-place
    clamp!(R,-0.75,0.5)
    @test R.x == Rclamp.x
end

@testset "remove_amp" begin
    # create array with std == 1
    A = rand(T,N,Nwin) .- 0.5
    A ./= std(A,dims=1)

    # add amplitude spikes
    Nspikes = 100
    cols = unique(rand(1:Nwin,Nwin÷2))
    ind = rand(1:N,Nspikes)
    R = RawData()
    R.x = deepcopy(A)
    R.x[ind,cols] .= randn(Nspikes,length(cols)) .* 100000
    R.t = collect(1.:10.)
    Ramp = remove_amp(R)
    @test isa(Ramp,RawData)
    @test size(Ramp.x,2) == Nwin - length(cols)
    @test length(Ramp.t) == Nwin - length(cols)

    # test in-place
    remove_amp!(R)
    @test R.x == Ramp.x

    # add amplitude spikes to all columns
    R.x = deepcopy(A)
    R.x[ind,:] .= randn(Nspikes) .* 100000
    @test_throws ErrorException remove_amp(R)

    # test with zeros
    cols = unique(rand(1:Nwin,Nwin÷2))
    R.x = deepcopy(A)
    R.x[:,cols] .= T(0)
    R.t = collect(1.:10.)
    Ramp = remove_amp(R)
    @test isa(Ramp,RawData)
    @test size(Ramp.x,2) == Nwin - length(cols)
    @test length(Ramp.t) == Nwin - length(cols)

    # test with all zeros
    R.x .= zero(T)
    @test_throws ErrorException remove_amp(R)
end

@testset "one-bit" begin
    A = rand(T,N,Nwin) .- T(0.5)
    R = RawData()
    R.x = A

    # one-bit RawData
    R1 = onebit(R)
    @test isa(R1,RawData)
    @test maximum(R1.x) == T(1)
    @test size(R1.x) == size(A)
    @test eltype(R1.x) == eltype(A)
    @test minimum(R1.x) == T(-1)
    @test maximum(R1.x) == T(1)

    # test in-place
    onebit!(R)
    @test R.x == R1.x

end

@testset "nonzero" begin
    A = rand(T,N,Nwin)
    cols = unique(rand(1:Nwin,Nwin÷2))
    A[:,cols] .= T(0)
    ind = SeisNoise.nonzero(A)
    # test no intersect between cols and ind
    @test setdiff(cols,ind) == cols
    # test ind are Int
    @test all(isa.(ind,Int))
end


const default_attrs = (disabled = false, style = "", class = "", changes = 0)
const default_input_attrs = merge(default_attrs, (value = "", type = "text"))

function input(;
               oninput = js"function () {data.value = this.value;}",
               onchange = js"function () {data.changes = data.changes + 1;}",
               selector = "input.input",
               kwargs...)

    datavals = merge(default_input_attrs, values(kwargs))
    data = map(to_observable, datavals)

    attrs = Dict(key => js"data[$key]" for key in keys(data))
    attrs[:oninput] = oninput
    attrs[:onchange] = onchange
    template = js"{view: () => m($selector, $attrs)}"
    m = MithrilComponent(template, data)
    return MithrilWidget{:input}(m, data.value)
end

_parse(::Type{S}, x) where{S} = parse(S, x)
function _parse(::Type{Dates.Time}, x)
    segments = split(x, ':')
    length(segments) >= 2 && all(!isempty, segments) || return nothing
    h, m = parse.(Int, segments)
    Dates.Time(h, m)
end

function _string(x::Dates.Time)
    h = Dates.hour(x)
    m = Dates.minute(x)
    string(lpad(h, 2, "0"), ":", lpad(m, 2, "0"))
end
_string(x) = string(x)
_string(x::Colorant) = "#"*hex(x)
_string(::Nothing) = ""

for (func, typ, type) in [(:timepicker, :(Dates.Time), "time"), (:datepicker, :(Dates.Date), "date"), (:colorpicker, Colorant, "color")]
    @eval function $func(; value = nothing, type = "", kwargs...)
        (value isa AbstractObservable) || (value = Observable{Union{$typ, Nothing}}(value))
        f = _string
        g = t -> _parse($typ, t)
        string_value = Observable{String}(f(value[]))
        ObservablePair(value, string_value, f = f, g = g)
        m = input(; value = string_value, type = $type, kwargs...)
        MithrilWidget{$(Expr(:quote, func))}(component(m), value)
    end
end

function spinbox(; value = Observable{Union{Int, Nothing}}(nothing), type = "", kwargs...)
    if !isa(value, AbstractObservable)
        T = value === nothing ? Union{Int, Nothing} : Union{typeof(value), Nothing}
        value = Observable{T}(value)
    end
    jsvalue = Observable{Any}(value[])
    m = input(; value = jsvalue, type = "number", kwargs...)

    excluded = Function[]

    push!(excluded, on(value) do val
        str = jsvalue[]
        if !isa(str, AbstractString) || tryparse(Float64, str) != val
            jsvalue[notify = !in(excluded)] = val
        end
    end)

    push!(excluded, on(jsvalue) do val
        if val isa AbstractString
            res = tryparse(Float64, val)
            res !== nothing && (value[notify = !in(excluded)] = res)
        end
    end)

    MithrilWidget{:spinbox}(component(m), value)
end

const default_button_attrs = merge(default_attrs, (changes = 0,))

function button(children...; kwargs...)
    datavals = merge(default_button_attrs, values(kwargs))
    data = map(to_observable, datavals)

    attrs = Dict(key => js"data[$key]" for key in keys(data))
    attrs[:onclick] = js"function () {data.changes = data.changes + 1;}"

    template = node("button.button", children..., attributes = attrs)
    MithrilWidget{:button}(MithrilComponent(template, data), data.changes)
end


module CUDD

# cudd_path = joinpath(dirname(@__FILE__), "..", "deps", "cudd-3.0.0", "cudd", ".libs")
const depfile = joinpath(@__DIR__, "..", "deps", "deps.jl")

if isfile(depfile)
    include(depfile)
else
    error("CUDD not properly installed. Please run Pkg.build(\"CUDD\")")
end

using CEnum

# Files automatically generated by Clang.jl
include(joinpath(@__DIR__, "gen", "libcudd_common.jl"))
include(joinpath(@__DIR__, "gen", "libcudd_api.jl"))

# Missing C macros omitted by Clang.jl

Cudd_Not(node) =
    convert(Ptr{Nothing}, xor(convert(Int,node), 1))

Cudd_IsComplement(node) =
    isone(convert(Int,node) & 1)

Cudd_Regular(node) =
    convert(Ptr{Nothing}, convert(Int,node) & ~1)

# export everything
foreach(names(@__MODULE__, all=true)) do s
    if startswith(string(s), "Cudd_")
        @eval export $s
    end
end

# More user-friendly interface for ADD
include("ADD.jl")
include("ADD_apply.jl")
include("ADD_find.jl")
include("ADD_abs.jl")
include("ADD_inverse.jl")
include("ADD_ITE.jl")
include("ADD_negate.jl")
include("ADD_walsh.jl")

end # module


module AutoGrad_gpu

import AutoGrad: addto!, addtoindex!, matches, ungetindex, Sparse, zeroslike
import LinearAlgebra: axpy!
import Base: +, -, copyto!
import CUDA, Knet, AutoGrad
using CUDA: CuArray, CuPtr, functional
using Knet.KnetArrays: DevArray, KnetArray, Cptr, cuallocator
using Knet.LibKnet8: @knet8
using Knet.CuArrays: cuarrays
using AutoGrad: AutoGrad, Sparse, recording, Result, Node, Tape, Value, Arg, value, set_gc_function
using Base.Broadcast: Broadcasted

include("addto.jl")
include("convert.jl")
include("getindex.jl")
include("sparse.jl")

include("priorityqueue.jl") # using DataStructures: DataStructures, PriorityQueue, dequeue!, dequeue_pair! # peek
include("gcnode.jl")
include("gcnode_kptr.jl")

end


module TestMOIWrapper

using Test
using MathOptInterface
import Cbc

const MOI = MathOptInterface

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

function test_SolverName()
    @test MOI.get(Cbc.Optimizer(), MOI.SolverName()) ==
          "COIN Branch-and-Cut (Cbc)"
end

function test_supports_incremental_interface()
    @test !MOI.supports_incremental_interface(Cbc.Optimizer())
    return
end

function test_runtests()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            Cbc.Optimizer(),
        ),
        Float64,
    )
    MOI.set(model, MOI.Silent(), true)
    MOI.Test.runtests(
        model,
        MOI.Test.Config(
            exclude = Any[
                MOI.ConstraintDual,
                MOI.DualObjectiveValue,
                MOI.ConstraintBasisStatus,
                MOI.VariableBasisStatus,
            ],
        ),
        exclude = [
            # TODO(odow): bug in Cbc.jl
            "test_model_copy_to_UnsupportedAttribute",
            "test_model_ModelFilter_AbstractConstraintAttribute",
            # TODO(odow): bug in MOI
            "test_model_LowerBoundAlreadySet",
            "test_model_UpperBoundAlreadySet",
            # TODO(odow): upstream bug in Cbc
            "_Indicator_",
            "_SOS1_",
            "_SOS2_",
            # Can't prove infeasible.
            "test_conic_NormInfinityCone_INFEASIBLE",
            "test_conic_NormOneCone_INFEASIBLE",
            "test_solve_TerminationStatus_DUAL_INFEASIBLE",
        ],
    )
    return
end

function test_params()
    # Note: we generate a non-trivial problem to ensure that Cbc struggles to
    # find a solution at the root node.
    knapsack_model = MOI.Utilities.Model{Float64}()
    N = 100
    x = MOI.add_variables(knapsack_model, N)
    MOI.add_constraint.(knapsack_model, x, MOI.ZeroOne())
    MOI.add_constraint(
        knapsack_model,
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1 + sin(i) for i in 1:N], x),
            0.0,
        ),
        MOI.LessThan(10.0),
    )
    MOI.set(
        knapsack_model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([cos(i) for i in 1:N], x),
            0.0,
        ),
    )
    model = Cbc.Optimizer(
        maxSol = 1,
        presolve = "off",
        cuts = "off",
        heur = "off",
        logLevel = 0,
    )
    MOI.copy_to(model, knapsack_model)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.SOLUTION_LIMIT
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    MOI.empty!(model)
    MOI.copy_to(model, knapsack_model)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.SOLUTION_LIMIT
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    return
end

function test_PrimalStatus()
    model = MOI.Utilities.Model{Float64}()
    x = MOI.add_variable(model)
    MOI.add_constraint(model, x, MOI.GreaterThan(1.0))
    MOI.add_constraint(model, x, MOI.LessThan(0.0))
    cbc = Cbc.Optimizer()
    MOI.copy_to(cbc, model)
    MOI.optimize!(cbc)
    MOI.get(cbc, MOI.PrimalStatus()) == MOI.NO_SOLUTION
    return
end

end

TestMOIWrapper.runtests()


# Copyright (c) 2022 Code Komali
#
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT

# module Vehicles

# export plot_vehicles!

using Agents
using InteractiveDynamics
using GLMakie
using Observables

include("Environment.jl")
include("Parameters.jl")
include("DrawEnvironment.jl")

import .Parameters as P

mutable struct VehicleAgent <: AbstractAgent
    id::Int
    pos::NTuple{2,Float64}
    vel::NTuple{2,Float64}
end

function initialize()
    space2d = ContinuousSpace((P.EXTENT_WIDTH, P.EXTENT_HEIGHT))
    properties = Dict()
    properties[:tick] = 0
    properties[:steptext] = Observable("Step: " * string(properties[:tick]))
    intersectingRoads = TwoWayIntersectingRoads(2000, 0, 4000, 2000, 4000, 0)
    properties[:env] = intersectingRoads
    properties[:spawn_rate] = 1400
    model = ABM(VehicleAgent, space2d, scheduler=Schedulers.randomly; properties=properties)
    add_vehicle!((4000.0, 1924.0), model, (-0.1, 0.0))
    return model
end

function add_vehicle!(spawn_pos, model, initial_vel=(P.VEHICLE_INITIAL_SPEED, 0.0))
    add_agent!(
        spawn_pos,
        model,
        initial_vel
    )
end

# Generalize
function vehicle_poly()
    hw = P.VEHICLE_WIDTH / 2
    pt = Point2f[(0, -hw), (P.VEHICLE_LENGTH, -hw), (P.VEHICLE_LENGTH, hw), (0, hw)]
    return Polygon(pt)
end

function vehicle_marker(v::VehicleAgent)
    φ = atan(v.vel[2], v.vel[1])
    rotate2D(vehicle_poly(), φ)
end

function vehicle_step!(agent, model)
    move_agent!(agent, model)
end

function plot_environment!(model)
    custom_setup_environment!(model.env)
    drawRoad!(model.env)
end

function custom_setup_environment!(twir::TwoWayIntersectingRoads)
    # Make the signals on the edge inactive
    twir.leftRoadSeg.R2Lroad.signal.active = false
    twir.rightRoadSeg.L2Rroad.signal.active = false
    twir.topRoadSeg.B2Troad.signal.active = false
    twir.bottomRoadSeg.T2Broad.signal.active = false
    # Alternate L2R and T2B road signals
    setSignalState!(twir.leftRoadSeg.L2Rroad.signal, :green)
    setSignalState!(twir.rightRoadSeg.R2Lroad.signal, :green)
    setSignalState!(twir.topRoadSeg.T2Broad.signal, :red)
    setSignalState!(twir.bottomRoadSeg.B2Troad.signal, :red)
    foreach(draw_cross_path!, cross_paths(twir))
end

function cross_paths(twir)
    cps = []
    xpos = top_boundary(twir.leftRoadSeg)[2][1] - P.CROSS_PATH_WIDTH - P.CROSS_PATH_MARGIN
    push!(cps, HcrossPath(twir.leftRoadSeg, xpos))
    xpos = top_boundary(twir.rightRoadSeg)[1][1] + P.CROSS_PATH_WIDTH + P.CROSS_PATH_MARGIN
    push!(cps, HcrossPath(twir.rightRoadSeg, xpos))
    ypos = right_boundary(twir.topRoadSeg)[2][2] + P.CROSS_PATH_WIDTH + P.CROSS_PATH_MARGIN
    push!(cps, VcrossPath(twir.topRoadSeg, ypos))
    ypos = right_boundary(twir.bottomRoadSeg)[1][2] - P.CROSS_PATH_WIDTH - P.CROSS_PATH_MARGIN
    push!(cps, VcrossPath(twir.bottomRoadSeg, ypos))
    return cps
end

function mark_spawn_positions(twir::TwoWayIntersectingRoads)
    for sp in spawnPos(twir)
        spawnPt = Point2f(sp.pos...)
        scatter!(spawnPt, color=:green, markersize=P.SIGNAL_MS)
    end
end

function model_step!(model)
    model.tick += 1
    model.steptext[] = "Step: " * string(model.tick)
    if (model.tick % model.spawn_rate == 0)
        rnd_spawn_pos = rand(spawnPos(model.env))
        @show rnd_spawn_pos
        spawn_vel = rnd_spawn_pos.orient .* P.VEHICLE_INITIAL_SPEED
        @show spawn_vel
        add_vehicle!(rnd_spawn_pos.pos, model, spawn_vel)
    end
    draw_signal!(model.env)
    @show model.tick
end

function plot_vehicles!()
    axiskwargs = (title=P.PLOT_TITLE, titlealign=P.PLOT_TITLE_ALIGN)
    model = initialize()
    params = Dict(
        :spawn_rate => P.VSR_MIN:P.VSR_INC:P.VSR_MAX,
    )
    fig, _ = abmplot(
        model;
        agent_step! = vehicle_step!,
        model_step! = model_step!,
        am=vehicle_marker,
        params=params,
        ac=:green,
        axis=axiskwargs
    )
    text!(model.steptext, position = (1000, 1000), textsize = 10)
    plot_environment!(model)
    return fig
end

#intersectingRoads = TwoWayIntersectingRoads(2000,0,4000,2000,4000,0)

#a macro for easing development only. will be removed later
macro pv()
    quote
        plot_vehicles!()
    end
end

# end


#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*--------------------------------------------------------------------
#* Author: Gerson Balbueno Bicca
#* $Id$
#*--------------------------------------------------------------------
type Baffles_Main
	Baffles_Main()=begin
		new(
			DanaInteger(Dict{Symbol,Any}(
				:Brief=>"Baffle Cut",
				:Default=>25,
				:Lower=>25
			)),
			DanaReal(Dict{Symbol,Any}(
				:Brief=>"Number of Baffles",
				:Symbol=>"N_{baffles}",
				:Lower=>1
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Inlet Baffle Spacing",
				:Lower=>1e-8,
				:Symbol=>"L_{si}",
				:DisplayUnit=>"mm"
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Central Baffle Spacing",
				:Lower=>1e-8,
				:Symbol=>"L_s",
				:DisplayUnit=>"mm"
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Outlet Baffle Spacing",
				:Lower=>1e-8,
				:Symbol=>"L_{so}",
				:DisplayUnit=>"mm"
			)),
			[:BaffleCut,:NumberOfBaffles,],
			[:Inlet_Spacing,:Central_Spacing,:Outlet_Spacing,]
		)
	end
	BaffleCut::DanaInteger
	NumberOfBaffles::DanaReal
	Inlet_Spacing::length
	Central_Spacing::length
	Outlet_Spacing::length
	parameters::Array{Symbol,1}
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export Baffles_Main
function atributes(in::Baffles_Main,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=false
	fields[:Brief]="Main variables in the Baffle section of a shell and tube heat exchanger."
	drive!(fields,_)
	return fields
end
Baffles_Main(_::Dict{Symbol,Any})=begin
	newModel=Baffles_Main()
	newModel.attributes=atributes(newModel,_)
	newModel
end