### Specify simulation parameters ###
@with_kw mutable struct Simulation_Parameters
    TTC_car_freq::Int64
    max_load_gov_restr::Float64 # 1.0 means there is no restrictions, because TTC car can take 1.0*max_load.
    when_to_run_people::Float64
    when_to_run_wagons::Float64
	sim_run_time_in_secs::Int
	N_agents::Int
	agents_speed_in_sec_per_m::Float64 # 0.8 means 0.8s per 1m of distance. The average, healthy human walks at a pace of 100m in 80 sec. Here we assume 100m=80s, so 1m = 0.8s.
    max_load_subway::Int64
    max_load_streetcar::Int64
	p0::Float64
end

### Agent definition ###
@with_kw mutable struct Commuter
    id::Int
    pos::Int64
    work_loc::Int64
    home_loc::Int64
    direction::Symbol = :work_to_home # It will be changed to home --> work in a first step

    infected::Bool = false
    infected_in::Symbol = :not_infected

    path::Dict{Int,Int} = Dict{Int,Int}()
    path_distances::Dict{Int,Float64} = Dict{Int,Float64}()
    paths_home_work::Dict{Int, Dict} = Dict{Int, Dict}()
    distances_home_work::Dict{Int,Dict} = Dict{Int,Dict}()
    paths_work_home::Dict{Int, Dict} = Dict{Int, Dict}()
    distances_work_home::Dict{Int,Dict} = Dict{Int,Dict}()
    no_of_trips::Int = 0

    wagon_id::Int64 = -999
	met_infected_in_wagon::Set{Int} = Set{Int}()

end

Commuter(id, pos, home_loc, work_loc) = Commuter(id = id, pos = pos, home_loc = home_loc, work_loc = work_loc)

### Articial agents definition ###
@with_kw mutable struct Artificial_Commuter_Stats
    id::Int = 999
	enter_every_sec::Int = 60
end

### Wagon definition ###
@with_kw mutable struct Wagon
    id::Int
    pos::Int64
    line::Symbol
    type::Symbol

    path::OrderedDict{Int,Int} = OrderedDict{Int,Int}() #seqence of nodes
    path_distances::OrderedDict{Int,Float64} = OrderedDict{Int,Float64}() # distance from each node in the sequence

    passengers::Set{Commuter} = Set{Commuter}()

    max_load::Union{Int, Float64}
    run_every::Int64 #sec

    time_start::Float64 = 0.0
    next_wagon_requested::Bool = false

    max_passengers_ever::Int64 = 0
end

### Simulation definiton ###
# EventTime represents time when event occurs. the second value in tuple is for sort ordering when 2 or more events are scheduled for the same point in time.
const EventTime = Tuple{Float64,Float64}

@with_kw mutable struct Simulation
    m::OpenStreetMapX.MapData

    agents::Vector{Commuter}
    wagons::Vector{Wagon}
    artificial_agent_stats::Artificial_Commuter_Stats

    recent_wagon::Int = -9999
    nodes_agents::Vector{Set{Int}}
    nodes_wagons::Vector{Set{Int}}

	infected_agents_count_current::Int
	infected_agents_wagon_current::Int
	infected_agents_street_current::Int

    infected_agents_count::Vector{Int}
    infected_agents_wagon::Vector{Int}
    infected_agents_street::Vector{Int}

    events = PriorityQueue{Union{Commuter,Wagon,Artificial_Commuter_Stats}, EventTime}()

    p0::Float64 # = 0.01

    max_load_subway::Int # = 15
    max_load_streetcar::Int # = 15
    max_load_gov_restr::Float64 # 1.0 means there is no restrictions, because TTC car can take 1.0*max_load.

    max_passengers_per_TTC_car::Dict{Tuple,Int64} = Dict{Tuple,Int64}() # [string(wagon.id,"|") * wagon.line * string("|",t[1])] = wagon.max_passengers_ever

end

function Simulation(m::OpenStreetMapX.MapData
					, orig_map_nodes_num::Int
					, routes_path::OrderedDict{Symbol, OrderedDict{Int64,Int64}}
					, routes_distances::OrderedDict{Symbol, OrderedDict{Int64,Float64}}
					, routes_types_both_dir::Dict{Symbol,Symbol}
					, N_agents::Int
					, sim_run_time_in_secs::Int
					, TTC_freq::Int64
                    , max_load_subway::Int64
                    , max_load_streetcar::Int64
					, max_load_gov_restr::Float64
					, p0::Float64)

	println("TTC_freq: ",TTC_freq)
    println("max_load_subway: ",max_load_subway)
	println("max_load_streetcar: ",max_load_streetcar)
	println("max_load_gov_restr: ",max_load_gov_restr)
	println("N_agents: ",N_agents)
	println("sim_run_time_in_secs: ",sim_run_time_in_secs)
	println("p0: ", p0)

    vv = size(m.g)[1] # number of vertices

    ### Agents creation ###
    home_loc = rand(1:orig_map_nodes_num,N_agents)
    work_loc = rand(1:orig_map_nodes_num,N_agents)
    agents = Commuter.(1:N_agents, home_loc, home_loc, work_loc)
    for ag in agents
       while ag.home_loc == ag.work_loc
            ag.work_loc = rand(1:orig_map_nodes_num)
       end
    end
    # Routes from home to work and inverse
    for agent in agents
        # Home to work
        ys = yen_k_shortest_paths(m.g, agent.home_loc, agent.work_loc, m.w, 3, maxdist=100_000)
        for (index, nodes) in enumerate(ys.paths)
            agent.paths_home_work[index] = Dict()
            agent.distances_home_work[index] = Dict()
            setindex!.(Ref(agent.paths_home_work[index]), nodes[2:end], nodes[1:end-1])
            for i in 1:(length(nodes)-1)
                agent.distances_home_work[index][nodes[i]] = m.w[nodes[i], nodes[i+1]]
            end
        end

        # Work to home
        ys = yen_k_shortest_paths(m.g, agent.work_loc, agent.home_loc, m.w, 3, maxdist=100_000)
        for (index, nodes) in enumerate(ys.paths)
            agent.paths_work_home[index] = Dict()
            agent.distances_work_home[index] = Dict()
            setindex!.(Ref(agent.paths_work_home[index]), nodes[2:end], nodes[1:end-1])
            for i in 1:(length(nodes)-1)
                agent.distances_work_home[index][nodes[i]] = m.w[nodes[i], nodes[i+1]]
            end
        end
        # Current path: home --> work
        # !!! We should not assign it here, because it will negatively influence timing.
        # It will be 1-step delayed comparing to the reality.
        # That is why it is mentioned here, but commented!
        # agent.path = deepcopy(agent.path_home_work)
        # agent.path_distances = deepcopy(agent.path_distances_home_work)

    end

    ### Wagons creation ###
    wagons = Wagon[]
    wagon_id = 1
    for r in collect(keys(routes_path))
        push!(wagons, Wagon(id = wagon_id
                ,pos = collect(keys(routes_path[r]))[1]
                ,line = r
                ,type = routes_types_both_dir[r]
                ,path = routes_path[r]
                ,path_distances = routes_distances[r]
                ,run_every = TTC_freq
                ,max_load = routes_types_both_dir[r] == :subway ? max_load_subway : max_load_streetcar
                ))
        wagon_id += 1
    end
    nodes_agents = [Set{Int}() for _ in 1:vv]
    nodes_wagons = [Set{Int}() for _ in 1:vv]
	# Add agents and wagons current positions.
    for agent in agents
        push!(nodes_agents[agent.pos], agent.id)
    end

    for wagon in wagons
        push!(nodes_wagons[wagon.pos], wagon.id)
    end
	# Identify patient zero and ensure that his/her trip is long enough, so we are sure that the only infected agent is not locked at home.
	target_n_of_patients_zero = N_agents*0.01
	n_of_patients_zero = 0
	while n_of_patients_zero < target_n_of_patients_zero
		patient_zero = rand(agents)
		if patient_zero.infected == false
			patient_zero.infected = true
			n_of_patients_zero += 1
		end
	end
    println("Number of patients zero: ", n_of_patients_zero)
    Simulation(agents = agents, wagons = wagons, recent_wagon = length(collect(keys(routes_path))), m = m, artificial_agent_stats = Artificial_Commuter_Stats(),
        nodes_agents = nodes_agents, nodes_wagons = nodes_wagons,
		infected_agents_count_current = n_of_patients_zero, infected_agents_wagon_current = 0, infected_agents_street_current = 0,
        infected_agents_count = fill(n_of_patients_zero, sim_run_time_in_secs + 1), infected_agents_wagon = zeros(Int, sim_run_time_in_secs + 1),
        infected_agents_street = zeros(Int, sim_run_time_in_secs + 1), max_load_subway = max_load_subway,
        max_load_streetcar = max_load_streetcar, max_load_gov_restr = max_load_gov_restr, p0 = p0)
end

println("Simulations objects were created")


immutable Protocol{N,C,U}
    insertread::Int
    insertpos::UnitRange{Int}
    cellpos::NTuple{N,UnitRange{Int}}
    umipos::NTuple{N,UnitRange{Int}}
    readnames::NTuple{N,String}
    readlengths::NTuple{N,Int}
end

# const config_path = joinpath(Pkg.dir("FASTQDemultiplexer"),"config","protocols")
const config_path = joinpath(@__DIR__,"..","config","protocols")
const protocol_config_paths = Dict(:marsseq => joinpath(config_path,"marsseq.yml"),
                                   :dropseq => joinpath(config_path,"dropseq.yml"),
                                   :x10 => joinpath(config_path,"10x.yml"))

function Protocol(prot::Symbol)

    allowed_protocols = keys(protocol_config_paths)

    if prot in allowed_protocols
        return Protocol(protocol_config_paths[prot])
    else
        error("Unknown protocol: $prot, use one of $allowed_protocol")
    end
end

function Protocol(yamlfile::String)
    config = YAML.load_file(yamlfile)

    nreads = length(config["reads"])

    cellpos = Array{UnitRange{Int}}(nreads)
    umipos  = Array{UnitRange{Int}}(nreads)
    readnames = Array{String}(nreads)
    readlengths = Array{Int}(nreads)
    insertpos = 1:0
    insertread = 0
    insertdefined = false

    # parse the reads
    for (i,read) in enumerate(config["reads"])

        if haskey(read,"cell")
            # TODO add safer parsing?
            cellpos[i] = eval(parse(read["cell"]))
        else
            cellpos[i] = 1:0
        end

        if haskey(read,"umi")
            umipos[i] = eval(parse(read["umi"]))
        else
            umipos[i] = 1:0
        end

        if haskey(read,"insert")
            if insertdefined
                error("The field insert can be defined in only one read")
            end
            insertpos = eval(parse(read["insert"]))
            insertread = i
            insertdefined = true
        end

        if haskey(read,"name")
            readnames[i]=read["name"]
        else
            error("Name of the read is required")
        end

        if haskey(read,"length")
            readlengths[i]=read["length"]
        else
            error("Length of the read is required")
        end

    end

    C = sum(map(length,cellpos))
    U = sum(map(length,umipos))

    return Protocol{nreads,C,U}(insertread,
                                insertpos,
                                (cellpos...),
                                (umipos...),
                                (readnames...),
                                (readlengths...))

end

cellidtoname(id::UInt,::Protocol{N,C,U}) where {N,C,U} = idtoname(id,C)

umiidtoname(id::UInt,::Protocol{N,C,U}) where {N,C,U} = idtoname(id,U)


import MPI

export distribute, MPI_rank, MPI_size

if !MPI.Initialized()
    MPI.Init()
end

MPI_rank = MPI.Comm_rank(MPI.COMM_WORLD)
MPI_size = MPI.Comm_size(MPI.COMM_WORLD)
MPI_COMM_WORLD = MPI.COMM_WORLD

mpirank(comm::MPI.Comm) = MPI.Comm_rank(comm)
mpisize(comm::MPI.Comm) = MPI.Comm_size(comm)

function distribute(size::Int, comm_size::Int, rank::Int)
    """
    Compute the first element and size for a given rank

    Parameters
    ----------
    size : Int
        Length of array
    comm_size : Int
        Number of MPI  processes
    rank: Int
        MPI rank

    Returns
    -------
    size : Int
        Size
    offset : Int
        Offsets for the MPI process
    """
    if comm_size > size
        if rank + 1 <= size
            return rank + 1, 1
        else
            return 1, 0
        end
    end

    base = div(size, comm_size)
    leftover = size % comm_size
    if rank+1 <= leftover
        size = base+1
        start = rank * (base+1) + 1
    else
        size = base
        start = (base+1) * leftover + (rank-leftover) * base + 1
    end
    return start, size
end

function Allreduce(data, mpi_sum, comm)
    if comm === nothing
        return data
    end
    return MPI.Allreduce(data, mpi_sum, comm)
end


export translate_resp!, translate_resp, remove_resp!, remove_resp

# =====================================================================

@doc """
    translate_resp!(S, resp_new[, chans=CC, wl=γ])
    translate_resp(S, resp_new[, chans=CC, wl=γ])

Translate the instrument response of seismic data channels `CC` in `S` to
`resp_new`. Replaces field `:resp` with `resp_new` for all affected channels.

    remove_resp!(S, chans=CC, wl=γ])
    remove_resp(S, chans=CC, wl=γ])

Remove (flatten to DC) the instrument response of seismic data channels `cha`
in `S`. Replaces field `:resp` with the appropriate (all-pass) response.

    translate_resp!(Ch, resp_new[, wl=γ])
    translate_resp(Ch, resp_new[, wl=γ])

Translate the instrument response of seismic data in SeisChannel object `Ch` to
`resp_new`. Replaces field `:resp` with `resp_new`.

    remove_resp!(Ch, wl=γ])
    remove_resp(Ch, wl=γ])

Remove (flatten to DC) the instrument response of seismic data in `Ch`.
Replaces field `:resp` with the appropriate (all-pass) response.

Keywords:
* **chans=CC** restricts response translation to channel(s) `CC`. `chans` can be
an Integer, UnitRange, or Array{Int64,1}. By default, all seismic data channels
have their responses translated to `resp_new`.
* **wl=γ** sets the waterlevel to γ (default: `γ` = eps(Float32) ≈ ~1f-7)

### Interaction with the :resp field

`translate_resp` and `remove_resp` only work on a channel `i` that satisfies
`S.resp[i] <: PZResp, PZResp64, MultiStageResp`. In the last case,
`S.resp[i].stage[1]` must be a PZResp or PZResp64, only the first stage of
the response is changed, and the stage gain is ignored; instead, the
sensitivity `S.resp[i].stage[1].a0` is used.

### Poles and zeros should be rad/s

Always check when loading from an unsupported data format. Responses read
from station XML are corrected to rad/s automatically (most use rad/s);
responses read from a SACPZ or SEED RESP file already use rad/s.
""" translate_resp!
function translate_resp!(S::GphysData,
                         resp_new::Union{PZResp, PZResp64};
                         chans::Union{Integer, UnitRange, Array{Int64,1}}=Int64[],
                         wl::Float32=eps(Float32))

  # first ensure that there is something to do
  chans = mkchans(chans, S.n)
  @inbounds for i in chans
    if S.resp[i] != resp_new
      break
    end
    if i == last(chans)
      @info(string(timestamp(), ": nothing done (no valid responses to translate)."))
      return nothing
    end
  end

  # remove channels with inappropriate response types
  k = Int64[]
  for (n,i) in enumerate(chans)
    if (typeof(S.resp[i]) <: Union{PZResp, PZResp64, MultiStageResp}) == false
      push!(k,n)
    end
  end
  deleteat!(chans, k)

  # initialize complex "work" vectors to the largest size we need
  Nx      = nx_max(S, chans)
  N2      = nextpow(2, Nx)
  Xw      = Array{Complex{Float32},1}(undef, N2)
  ff_old  = Array{Complex{Float32},1}(undef, N2)
  ff_new  = Array{Complex{Float32},1}(undef, N2)
  f       = Array{Float32,1}(undef, N2)

  GRPS = get_unique(S, ["fs", "resp", "units"], chans)
  for grp in GRPS

    # get fs, resp
    j = grp[1]
    RT = typeof(getindex(getfield(S, :resp), j))

    # no translating instrument responses unless we're dealing with seismometers
    codes = inst_codes(S)
    kill = falses(length(grp))
    for i = 1:length(grp)
      if codes[i] in seis_inst_codes
        continue
      else
        kill[i] = true
      end
    end
    deleteat!(grp, kill)
    isempty(grp) && continue
    j = grp[1]
    uu = lowercase(S.units[j])

    # onward
    resp_old = RT == MultiStageResp ? deepcopy(S.resp[j].stage[1]) : deepcopy(S.resp[j])
    fs = Float32(S.fs[j])

    if fs > 0.0f0 && resp_old != resp_new

      # we'll need this for logging
      resp_str = string("translate_resp!, wl = ", wl, ", resp_old = ", typeof(resp_old),
        "(a0=", resp_old.a0, ", f0=", resp_old.f0, ", p=", resp_old.p, ", z=", resp_old.z, ")")

      # for accelerometers, working *in* acceleration units, we check: are there any poles below the nyquist?
      if uu == "m/s2"
        P = resp_old.p
        k = trues(length(P))
        for (i,p) in enumerate(P)
          if abs(p)/pi ≤ fs # equivalent to abs(p)/2pi ≤ fn/2 which is the true condition
            k[i] = false
          end
        end
        deleteat!(P, k)

        # Most accelerometers have two complex zeros at the origin, but many XML resp files don't show it.
        if isempty(resp_old.z)
          resp_old.z = map(eltype(resp_old.z), [0.0 + 0.0im, 0.0 - 0.0im])
        end
      end

      # get views to segments from each target channel
      (L,X) = get_views(S, grp)

      # initialize Nx, N2, xre
      Nx = first(L)
      N2 = nextpow(2, Nx)
      xfl, xre = update_resp_vecs!(Xw, f, ff_old, ff_new, N2)
      update_resp!(f, ff_old, ff_new, N2, fs, resp_old, resp_new, wl)

      j = 0
      while j < length(L)
        j = j + 1
        if L[j] != Nx
          Nx = L[j]
          N2 = nextpow(2, Nx)
          xfl, xre = update_resp_vecs!(Xw, f, ff_old, ff_new, N2)
          update_resp!(f, ff_old, ff_new, N2, fs, resp_old, resp_new, wl)
        end

        # copy X[j] to Xw and compute FFT
        fill!(Xw, zero(Complex{Float32}))
        copyto!(Xw, 1, X[j], 1, Nx)
        fft!(Xw)
        broadcast!(*, Xw, Xw, ff_new)
        ifft!(Xw)
        copyto!(X[j], 1, xre, 1, Nx)
      end

      # post-processing: set resp and log to :notes
      for k in grp
        if RT == MultiStageResp
          S.resp[k].stage[1] = deepcopy(resp_new)
        else
          setindex!(S.resp, deepcopy(resp_new), k)
        end
        note!(S, k, resp_str)
      end
    end
  end
 return nothing
end

@doc (@doc translate_resp!)
function translate_resp( S::GphysData,
                    resp_new::Union{PZResp, PZResp64};
                    chans::Union{Integer, UnitRange, Array{Int64,1}}=Int64[],
                    wl::Float32=eps(Float32))

  U = deepcopy(S)
  translate_resp!(U, resp_new, chans=chans, wl=wl)
  return U
end

function translate_resp!(C::GphysChannel,
                    resp_new::Union{PZResp, PZResp64};
                    wl::Float32=eps(Float32))

  # first ensure that there is something to do
  uu = lowercase(C.units)
  fs = Float32(C.fs)
  if any([C.resp == resp_new,
          uu in ("m/s", "m/s2", "m") == false,
          fs ≤ 0.0f0,
          inst_code(C) in seis_inst_codes == false])
    @info(string(timestamp(), ": nothing done (no valid responses to translate)."))
    return nothing
  end

  # initialize complex "work" vectors to the largest size we need
  Nx      = nx_max(C)
  N2      = nextpow(2, Nx)
  Xw      = zeros(Complex{Float32}, N2)
  ff_old  = Array{Complex{Float32},1}(undef, N2)
  ff_new  = Array{Complex{Float32},1}(undef, N2)
  f       = Array{Float32,1}(undef, N2)

  # for accelerometers, working *in* acceleration units, we check: are there any poles below the nyquist?
  if uu == "m/s2"
    P = C.resp.p
    k = trues(length(P))
    for (i,p) in enumerate(P)
      if abs(p)/pi ≤ fs # equivalent to abs(p)/2pi ≤ fn/2 which is the true condition
        k[i] = false
      end
    end
    deleteat!(P, k)

    # Most accelerometers have two complex zeros at the origin, but many XML resp files don't show it.
    if isempty(C.resp.z)
      C.resp.z = map(eltype(C.resp.z), [0.0 + 0.0im, 0.0 - 0.0im])
    end
  end

  # Get views
  if size(C.t,1) == 2
    xfl, xre = update_resp_vecs!(Xw, f, ff_old, ff_new, N2)
    update_resp!(f, ff_old, ff_new, N2, fs, C.resp, resp_new, wl)

    # copy X[j] to Xw and compute FFT
    copyto!(Xw, 1, C.x, 1, Nx)
    fft!(Xw)
    broadcast!(*, Xw, Xw, ff_new)
    ifft!(Xw)
    copyto!(C.x, 1, xre, 1, Nx)
  else
    (L,X) = get_views(C)
    xfl = reinterpret(Float32, Xw)
    xre = view(xfl, 1:2:2*N2-1)
    update_resp!(f, ff_old, ff_new, N2, fs, C.resp, resp_new, wl)

    j = 0
    while j < length(L)
      j = j + 1

      if L[j] != Nx
        Nx = L[j]
        N2 = nextpow(2, Nx)
        xfl, xre = update_resp_vecs!(Xw, f, ff_old, ff_new, N2)
        update_resp!(f, ff_old, ff_new, N2, fs, C.resp, resp_new, wl)
      end

      if j > 1
        fill!(Xw, zero(Complex{Float32}))
      end

      copyto!(Xw, 1, X[j], 1, Nx)
      fft!(Xw)
      broadcast!(*, Xw, Xw, ff_new)
      ifft!(Xw)
      copyto!(X[j], 1, xre, 1, Nx)
    end
  end
  C.resp = deepcopy(resp_new)
  return nothing
end

function translate_resp(C::GphysChannel,
                        resp_new::Union{PZResp, PZResp64};
                        wl::Float32=eps(Float32))

  U = deepcopy(C)
  translate_resp!(U, resp_new, wl=wl)
  return U
end

@doc (@doc translate_resp!)
remove_resp!(S::GphysData;
             chans::Union{Integer, UnitRange, Array{Int64,1}}=Int64[],
             wl::Float32=eps(Float32)) = translate_resp!(S, flat_resp, chans=chans, wl=wl)

@doc (@doc translate_resp!)
remove_resp(S::GphysData;
            chans::Union{Integer, UnitRange, Array{Int64,1}}=Int64[],
            wl::Float32=eps(Float32)) = translate_resp(S, flat_resp, chans=chans, wl=wl)

remove_resp!(Ch::GphysChannel;
             wl::Float32=eps(Float32)) = translate_resp!(Ch, flat_resp, wl=wl)

remove_resp(Ch::GphysChannel;
            wl::Float32=eps(Float32)) = translate_resp(Ch, flat_resp, wl=wl)


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

# ToDo: Replace with custom UnshapedUvDist:
unshaped(d::UnivariateDistribution) = Distributions.Product(Fill(d, 1))


# Array info
for (m, f) in ((:Base, :size), (:Base, :axes), (:Base, :firstindex), (:Base, :lastindex))
    @eval begin
        @inline $m.$f(A::AbstractDimArray, dims::AllDims) = $m.$f(A, dimnum(A, dims))
    end
end

# Reducing methods

# With a function arg version
for (m, f) in ((:Base, :sum), (:Base, :prod), (:Base, :maximum), (:Base, :minimum),
                     (:Base, :extrema), (:Statistics, :mean))
    _f = Symbol('_', f)
    @eval begin
        # Base methods
        $m.$f(A::AbstractDimArray; dims=:, kw...) = $_f(A::AbstractDimArray, dims; kw...)
        $m.$f(f, A::AbstractDimArray; dims=:, kw...) =
            $_f(f, A::AbstractDimArray, dims; kw...)
        # Local dispatch methods
        # - Return a reduced DimArray
        $_f(A::AbstractDimArray, dims; kw...) =
            rebuild(A, $m.$f(parent(A); dims=dimnum(A, _astuple(dims)), kw...), reducedims(A, dims))
        $_f(f, A::AbstractDimArray, dims; kw...) =
            rebuild(A, $m.$f(f, parent(A); dims=dimnum(A, _astuple(dims)), kw...), reducedims(A, dims))
        # - Return a scalar
        $_f(A::AbstractDimArray, dims::Colon; kw...) = $m.$f(parent(A); kw...)
        $_f(f, A::AbstractDimArray, dims::Colon; kw...) =
            $m.$f(f, parent(A); dims=dims, kw...)
    end
end
# With no function arg version
for (m, f) in ((:Statistics, :std), (:Statistics, :var))
    _f = Symbol('_', f)
    @eval begin
        # Base methods
        $m.$f(A::AbstractDimArray; corrected::Bool=true, mean=nothing, dims=:) =
            $_f(A, corrected, mean, dims)
        # Local dispatch methods - Returns a reduced array
        $_f(A::AbstractDimArray, corrected, mean, dims) =
            rebuild(A, $m.$f(parent(A); corrected=corrected, mean=mean, dims=dimnum(A, _astuple(dims))), reducedims(A, dims))
        # - Returns a scalar
        $_f(A::AbstractDimArray, corrected, mean, dims::Colon) =
            $m.$f(parent(A); corrected=corrected, mean=mean, dims=:)
    end
end
for (m, f) in ((:Statistics, :median), (:Base, :any), (:Base, :all))
    _f = Symbol('_', f)
    @eval begin
        # Base methods
        $m.$f(A::AbstractDimArray; dims=:) = $_f(A, dims)
        # Local dispatch methods - Returns a reduced array
        $_f(A::AbstractDimArray, dims) =
            rebuild(A, $m.$f(parent(A); dims=dimnum(A, _astuple(dims))), reducedims(A, dims))
        # - Returns a scalar
        $_f(A::AbstractDimArray, dims::Colon) = $m.$f(parent(A); dims=:)
    end
end

# These are not exported but it makes a lot of things easier using them
function Base._mapreduce_dim(f, op, nt::NamedTuple{(),<:Tuple}, A::AbstractDimArray, dims)
    rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, _astuple(dims))), reducedims(A, dims))
end
function Base._mapreduce_dim(f, op, nt::NamedTuple{(),<:Tuple}, A::AbstractDimArray, dims::Colon)
    Base._mapreduce_dim(f, op, nt, parent(A), dims)
end
function Base._mapreduce_dim(f, op, nt, A::AbstractDimArray, dims)
    rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, dims)), reducedims(A, dims))
end
@static if VERSION >= v"1.6"
    function Base._mapreduce_dim(f, op, nt::Base._InitialValue, A::AbstractDimArray, dims)
        rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, dims)), reducedims(A, dims))
    end
    function Base._mapreduce_dim(f, op, nt::Base._InitialValue, A::AbstractDimArray, dims::Colon)
        Base._mapreduce_dim(f, op, nt, parent(A), dims)
    end
end

# TODO: Unfortunately Base/accumulate.jl kw methods all force dims to be Integer.
# accumulate wont work unless that is relaxed, or we copy half of the file here.
# Base._accumulate!(op, B, A, dims::AllDims, init::Union{Nothing, Some}) =
    # Base._accumulate!(op, B, A, dimnum(A, dims), init)

# Dimension dropping

function Base.dropdims(A::AbstractDimArray; dims)
    dims = DD.dims(A, dims)
    data = Base.dropdims(parent(A); dims=dimnum(A, dims))
    rebuildsliced(A, data, _dropinds(A, dims))
end

@inline _dropinds(A, dims::DimTuple) = dims2indices(A, map(d -> basetypeof(d)(1), dims))
@inline _dropinds(A, dim::Dimension) = dims2indices(A, basetypeof(dim)(1))

# Function application

@inline Base.map(f, A::AbstractDimArray) = rebuild(A, map(f, parent(A)))

function Base.mapslices(f, A::AbstractDimArray; dims=1, kw...)
    dimnums = dimnum(A, _astuple(dims))
    _data = mapslices(f, parent(A); dims=dimnums, kw...)
    rebuild(A, _data, reducedims(A, DimensionalData.dims(A, dimnums)))
end

# This is copied from base as we can't efficiently wrap this function
# through the kw with a rebuild in the generator. Doing it this way
# also makes it faster to use a dim than an integer.
function Base.eachslice(A::AbstractDimArray; dims=1, kw...)
    if dims isa Tuple && length(dims) != 1
        throw(ArgumentError("only single dimensions are supported"))
    end
    dim = first(dimnum(A, dims))
    dim <= ndims(A) || throw(DimensionMismatch("A doesn't have $dim dimensions"))
    idx1, idx2 = ntuple(d->(:), dim-1), ntuple(d->(:), ndims(A)-dim)
    return (view(A, idx1..., i, idx2...) for i in axes(A, dim))
end

# Duplicated dims

for fname in (:cor, :cov)
    @eval function Statistics.$fname(A::AbstractDimArray{<:Any,2}; dims=1, kw...)
        newdata = Statistics.$fname(parent(A); dims=dimnum(A, dims), kw...)
        removed_idx = dimnum(A, dims)
        newrefdims = $dims(A)[removed_idx]
        newdim = $dims(A)[3 - removed_idx]
        rebuild(A, newdata, (newdim, newdim), (newrefdims,))
    end
end

# Rotations

struct _Rot90 end
struct _Rot180 end
struct _Rot270 end
struct _Rot360 end

Base.rotl90(A::AbstractDimMatrix) = rebuild(A, rotl90(parent(A)), _rot(_Rot90(), dims(A)))
Base.rotl90(A::AbstractDimMatrix, k::Integer) =
    rebuild(A, rotl90(parent(A), k), _rot(_rottype(k), dims(A)))

Base.rotr90(A::AbstractDimMatrix) = rebuild(A, rotr90(parent(A)), _rot(_Rot270(), dims(A)))
Base.rotr90(A::AbstractDimMatrix, k::Integer) =
    rebuild(A, rotr90(parent(A), k), _rot(_rottype(-k), dims(A)))

Base.rot180(A::AbstractDimMatrix) = rebuild(A, rot180(parent(A)), _rot(_Rot180(), dims(A)))

# Not type stable - but we have to lose type stability somewhere when
# dims are being swapped, by an Int value, so it may as well be here
function _rottype(k)
    k = mod(k, 4)
    if k == 1
        _Rot90()
    elseif k == 2
        _Rot180()
    elseif k == 3
        _Rot270()
    else
        _Rot360()
    end
end

_rot(::_Rot90, (dima, dimb)) = (flip(Relation, dimb), dima)
_rot(::_Rot180, dims) = map(d -> flip(Relation, d), dims)
_rot(::_Rot270, (dima, dimb)) = (dimb, flip(Relation, dima))
_rot(::_Rot360, dims) = dims

# Dimension reordering

for (pkg, fname) in [(:Base, :permutedims), (:Base, :adjoint),
                     (:Base, :transpose), (:LinearAlgebra, :Transpose)]
    @eval begin
        @inline $pkg.$fname(A::AbstractDimArray{<:Any,2}) =
            rebuild(A, $pkg.$fname(parent(A)), reverse(dims(A)))
        @inline $pkg.$fname(A::AbstractDimArray{<:Any,1}) =
            rebuild(A, $pkg.$fname(parent(A)), (AnonDim(Base.OneTo(1)), dims(A)...))
    end
end
for fname in [:permutedims, :PermutedDimsArray]
    @eval begin
        @inline function Base.$fname(A::AbstractDimArray, perm)
            rebuild(A, $fname(parent(A), dimnum(A, Tuple(perm))), sortdims(dims(A), Tuple(perm)))
        end
    end
end

# Concatenation
function Base._cat(catdims::Union{Int,Base.Dims}, Xin::AbstractDimArray...)
    Base._cat(dims(first(Xin), catdims), Xin...)
end
function Base._cat(catdims::AllDims, Xin::AbstractDimArray...)
    A1 = first(Xin)
    comparedims(Xin...)
    if all(hasdim(A1, catdims))
        # Concatenate an existing dim
        dnum = dimnum(A1, catdims)
        # cat the catdim, ignore others
        newdims = Tuple(_catifcatdim(catdims, ds) for ds in zip(map(dims, Xin)...))
    else
        # Concatenate a new dim
        add_dims = if (catdims isa Tuple)
            Tuple(d for d in catdims if !hasdim(A1, d))
        else
            (catdims,)
        end
        dnum = ndims(A1) + length(add_dims)
        newdims = (dims(A1)..., add_dims...)
    end
    newA = Base._cat(dnum, map(data, Xin)...)
    rebuild(A1, newA, formatdims(newA, newdims))
end

function _catifcatdim(catdims::Tuple, dims)
    anydimsmatch = any(map(d -> basetypeof(d) <: basetypeof(dims[1]), catdims))
    anydimsmatch ? vcat(dims...) : dims[1]
end
function _catifcatdim(catdim, dims)
    basetypeof(catdim) <: basetypeof(dims[1]) ? vcat(dims...) : dims[1]
end

function Base.vcat(dims::Dimension...)
    newmode = _vcat_modes(mode(dims)...)
    rebuild(dims[1], _vcat_index(newmode, index(dims)...), newmode)
end

# IndexModes may need adjustment for `cat`
_vcat_modes(modes::IndexMode...) = first(modes)
function _vcat_modes(modes::AbstractSampled...)
    _vcat_modes(sampling(first(modes)), span(first(modes)), modes...)
end
function _vcat_modes(::Any, ::Regular, modes...)
    _step = step(first(modes))
    map(modes) do mode
        step(span(mode)) == _step || error("Step sizes $(step(span(mode))) and $_step do not match ")
    end
    first(modes)
end
function _vcat_modes(::Intervals, ::Irregular, modes...)
    allbounds = map(bounds ∘ span, modes)
    newbounds = minimum(map(first, allbounds)), maximum(map(last, allbounds))
    rebuild(modes[1]; span=Irregular(newbounds))
end
_vcat_modes(::Points, ::Irregular, modes...) =
    rebuild(first(modes); span=Irregular(nothing, nothing))

# Index vcat depends on mode: NoIndex is always Colon()
_vcat_index(mode::NoIndex, A...) = OneTo(sum(map(length, A)))
# TODO: handle vcat OffsetArrays?
# Otherwise just vcat. TODO: handle order breaking vcat?
_vcat_index(mode::IndexMode, A...) = vcat(A...)

Base.inv(A::AbstractDimArray{T,2}) where T =
    rebuild(A, inv(parent(A)), reverse(map(d -> flip(IndexOrder, d), dims(A))))

# Index breaking

# TODO: change the index and traits of the reduced dimension and return a DimArray.
Base.unique(A::AbstractDimArray; dims::Union{DimOrDimType,Colon}=:) = _unique(A, dims)
Base.unique(A::AbstractDimArray{<:Any,1}) = unique(parent(A))

_unique(A::AbstractDimArray, dims::DimOrDimType) =
    unique(parent(A); dims=dimnum(A, dims))
_unique(A::AbstractDimArray, dims::Colon) = unique(parent(A); dims=:)

Base.diff(A::AbstractDimVector; dims=1) = _diff(A, dimnum(A, dims))
Base.diff(A::AbstractDimArray; dims) = _diff(A, dimnum(A, dims))

@inline function _diff(A::AbstractDimArray{<:Any,N}, dims::Integer) where {N}
    r = axes(A)
    # Copied from Base.diff
    r0 = ntuple(i -> i == dims ? UnitRange(1, last(r[i]) - 1) : UnitRange(r[i]), N)
    rebuildsliced(A, diff(parent(A); dims=dimnum(A, dims)), r0)
end


export Conv
using Knet.Ops21: conv, zeropad, symmetric
import CUDA

using CUDA.CUDNN:
    cudnnConvolutionForward,
    cudnnConvolutionForward!,
    cudnnConvolutionBackwardFilter,
    cudnnConvolutionBackwardData,
    cudnnGetConvolutionNdForwardOutputDim,
    cudnnSetConvolutionMathType,
    cudnnSetConvolutionReorderType,
    cudnnSetConvolutionGroupCount,
    cudnnFindConvolutionForwardAlgorithmEx,
        cudnnConvolutionFwdAlgoPerf_t,
    cudnnFindConvolutionBackwardFilterAlgorithmEx,
        cudnnConvolutionBwdFilterAlgoPerf_t,
    cudnnFindConvolutionBackwardDataAlgorithmEx,
        cudnnConvolutionBwdDataAlgoPerf_t,
    cudnnConvolutionDescriptor,
        cudnnConvolutionDescriptor_t,
        cudnnCreateConvolutionDescriptor,
        cudnnSetConvolutionNdDescriptor,
        cudnnDestroyConvolutionDescriptor,
    cudnnConvolutionMode_t,
        CUDNN_CONVOLUTION,       # 0
        CUDNN_CROSS_CORRELATION, # 1
    cudnnActivationMode_t,
        CUDNN_ACTIVATION_SIGMOID,      # 0
        CUDNN_ACTIVATION_RELU,         # 1
        CUDNN_ACTIVATION_TANH,         # 2
        CUDNN_ACTIVATION_CLIPPED_RELU, # 3
        CUDNN_ACTIVATION_ELU,          # 4
        CUDNN_ACTIVATION_IDENTITY,     # 5
    cudnnNanPropagation_t,
        CUDNN_NOT_PROPAGATE_NAN, # 0
        CUDNN_PROPAGATE_NAN,     # 1
    cudnnMathType_t,
        CUDNN_DEFAULT_MATH,                    # 0
        CUDNN_TENSOR_OP_MATH,                  # 1
        CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION, # 2
        CUDNN_FMA_MATH,                        # 3
    cudnnReorderType_t,
        CUDNN_DEFAULT_REORDER, # 0
        CUDNN_NO_REORDER,      # 1
    cudnnConvolutionFwdAlgo_t,
        CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM,         # 0
        CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM, # 1
        CUDNN_CONVOLUTION_FWD_ALGO_GEMM,                  # 2
        CUDNN_CONVOLUTION_FWD_ALGO_DIRECT,                # 3
        CUDNN_CONVOLUTION_FWD_ALGO_FFT,                   # 4
        CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING,            # 5
        CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD,              # 6
        CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED,     # 7
        CUDNN_CONVOLUTION_FWD_ALGO_COUNT,                 # 8
    cudnnConvolutionBwdFilterAlgo_t,
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_0,                 # 0, /* non-deterministic */
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_1,                 # 1,
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT,               # 2,
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_3,                 # 3, /* non-deterministic */
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD,          # 4, /* not implemented */
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD_NONFUSED, # 5,
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT_TILING,        # 6,
        CUDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT,             # 7
    cudnnConvolutionBwdDataAlgo_t,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_0,                 # 0, /* non-deterministic */
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_1,                 # 1,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_FFT,               # 2,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_FFT_TILING,        # 3,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD,          # 4,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD_NONFUSED, # 5,
        CUDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT,             # 6
    cudnnTensorFormat_t,
        CUDNN_TENSOR_NCHW,        # 0, /* row major (wStride = 1, hStride = w) */
        CUDNN_TENSOR_NHWC,        # 1, /* feature maps interleaved ( cStride = 1 )*/
        CUDNN_TENSOR_NCHW_VECT_C, # 2, /* each image point is vector of element of C, vector length in data type */
    cudnnDataType,
    convdims,
    math_mode,
    handle

"""
    c = Conv(wdims...; winit, binit, kwargs...)
    c = Conv(w; bias, kwargs...)
    y = c(x, [z=nothing])

Return a layer that can perform convolution and optionally scaling, bias/residual addition,
normalization and/or activation. The constructor can take the convolution weight tensor `w`
or its dimensions `wdims`. If `wdims` is used a weight tensor will be initialized using the
`winit` function which defaults to uniform random in `±√(6/(fanin+fanout))` and a bias will
be initialized using the `binit` function if provided. No bias will be used unless either
`bias` or `binit` is specified.  The layer can be called with a single input `x`, or two
inputs `x,z` where `z` has the same size as the output `y`. The computed result is:

    y = activation(normalization(alpha * conv(w,x) + beta * z .+ bias))

For tensor sizes and keyword arguments with their defaults see `@doc Knet.Ops21.conv`.
"""
mutable struct Conv
    w  # No type here, we do not want to restrict the type of array
    wdims::Dims
    winit

    bias  # Can be nothing
    bdims::Dims
    binit

    # y; dw; dbias; dx; dz; These are overwritten if multiple calls to same layer.

    padding::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}},Tuple{Vararg{NTuple{2,<:Integer}}}}
    stride::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}}}
    dilation::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}}}
    groups::Integer
    crosscorrelation::Bool
    channelmajor::Bool
    activation
    normalization
    convDesc::Union{Nothing,cudnnConvolutionDescriptor}

    alpha::Real
    beta::Real
end

# If the user provides size, we can't immediately initialize w because we do not know the
# array type until we see the first input
function Conv(
    wdims::Integer...;
    w = nothing,
    bias = nothing,
    crosscorrelation::Bool = true,
    channelmajor::Bool = false, # CUDNN_TENSOR_NHWC if true.
    winit = 𝑼(√(6/(fanin(wdims; channelmajor)+fanout(wdims; channelmajor)))),
    binit = nothing,
    activation = nothing,
    normalization = nothing,
    padding::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}},Tuple{Vararg{NTuple{2,<:Integer}}}} = 0,
    stride::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}}} = 1,
    dilation::Union{Integer,Vector{<:Integer},Tuple{Vararg{<:Integer}}} = 1,
    groups::Integer = 1,
    alpha::Real = 1,
    beta::Real = 0,
)
    wdims = Int.(wdims)
    dw = Ref{Any}(nothing)
    dbias = Ref{Any}(nothing)
    ndims = length(wdims)
    cdim = channelmajor ? 1 : ndims-1
    bdims = ntuple(i->(i===cdim ? wdims[end] : 1), ndims)
    convDesc = nothing          # To be initialized at first call
    Conv(w, wdims, winit, bias, bdims, binit, padding, stride, dilation, groups, crosscorrelation, channelmajor, activation, normalization, convDesc, alpha, beta)
end

Conv(w; o...) = Conv(size(w)...; w, o...)

# Some of the initialization can only be done at the first call when we know the type of x
# TODO: are we sure about this? should we use atype() instead?
function initconv(c::Conv, x, z, padding)
    issimilar(u,v,s=size(v))=(typeof(value(u)) === typeof(value(v)) && size(u) === s)
    if c.w === nothing
        c.w = Param(oftype(value(x), c.winit(c.wdims...)))
    end
    @assert issimilar(c.w, x, c.wdims)  "w=$(summary(c.w)) x=$(summary(x)) wdims=$(c.wdims)"
    @assert (c.channelmajor ? size(x,1) === size(c.w,1)*c.groups : size(x)[end-1] === size(c.w)[end-1]*c.groups)  "x=$(size(x)) w=$(size(c.w)) grp=$(c.groups)"
    if c.bias === nothing && c.binit !== nothing
        c.bias = Param(oftype(value(x), c.binit(c.bdims...)))
    end
    if c.bias === nothing && ((c.activation === relu && c.normalization ∈ (nothing, identity)) || (z !== nothing && c.beta != 0))
        # will call cudnnConvolutionBiasActivationForward, must have bias
        c.bias = fill!(similar(c.w, c.bdims), 0)
    end
    @assert c.bias === nothing || issimilar(c.bias, c.w, c.bdims)
    if CUDA.functional() && (c.convDesc === nothing || c.convDesc.ptr === C_NULL)
        mode = c.crosscorrelation ? CUDNN_CROSS_CORRELATION : CUDNN_CONVOLUTION
        format = c.channelmajor ? CUDNN_TENSOR_NHWC : CUDNN_TENSOR_NCHW
        reorderType, mathType = CUDNN_DEFAULT_REORDER, math_mode()
        c.convDesc = cudnnConvolutionDescriptor(convdims(padding,size(x),format), convdims(c.stride,size(x),format), convdims(c.dilation,size(x),format), mode, cudnnDataType(eltype(x)), mathType, reorderType, Cint(c.groups))
    end
end

function (c::Conv)(x, z=nothing)
    if symmetric(c.padding)
        padding = c.padding
    else
        x = zeropad(x, c.padding)
        padding = 0
    end
    initconv(c, x, z, padding)
    conv(c.w, x; z, c.bias, c.activation, c.normalization, c.alpha, c.beta, c.channelmajor,
         c.convDesc, c.crosscorrelation, c.dilation, c.groups, padding, c.stride)
end


##  Copyright 2022 Myrl Marmarelis
##
##  Licensed under the Apache License, Version 2.0 (the "License");
##  you may not use this file except in compliance with the License.
##  You may obtain a copy of the License at
##
##    http://www.apache.org/licenses/LICENSE-2.0
##
##  Unless required by applicable law or agreed to in writing, software
##  distributed under the License is distributed on an "AS IS" BASIS,
##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
##  See the License for the specific language governing permissions and
##  limitations under the License.

using Flux, Zygote
using Random
using ProgressMeter
using StatsFuns

const Optimizer = Flux.Optimise.AbstractOptimiser

function ensemble_log_likelihood(predictors::Vector{P},
    parametrization::PredictionParam{T}, instance::Instance{In,Out,T},
    )::T where {In,Out,POut,T, P <: Predictor{In,POut,T}}
  ensemble_size = length(predictors)
  # `logsumexp` needs an intermediary array... (is reduce(logaddexp, ...) as stable?
  #   yes because logsumexp is implemented that way basically)
  log_likelihoods = Iterators.map(predictors) do predictor
    log_likelihood(predictor, parametrization, instance)
  end
  (reduce(logaddexp, log_likelihoods) - log(ensemble_size)) / Out
end

# INTERESTING THOUGHT, BUT UNFITTING
#  ensemble_size = length(predictors)
#  log_likelihoods = Iterators.map(zip(predictor_weights, predictors)) do (weights, predictor)
#    mean(zip(weights, instances)) do (weight, instance)
#      weight * log_likelihood(predictor, parametrization, instance)
#    end
#  end
#  (logsumexp(log_likelihoods) - log(ensemble_size)) / Out

using Base.Threads

# is "parameterization" more correct?
function train_directly!(predictors::Vector{P},
    instances::Vector{Instance{In,Out,T}}, parametrization::PredictionParam{T};
    weights::WeightScheme{T}, n_iterations::Int, optimizer::Optimizer, batch_size::Int,
    constraints::NTuple{C,Constraint{T}}=(), constraint_weight::T=T(0), n_constraint_draws::Int=0,
    validation_instances::Vector{Instance{In,Out,T}}=Instance{In,Out,T}[], verbose::Bool=false) where
      { In, Out, POut, C, T, P <: Predictor{In,POut,T} }
  @assert POut == n_predictor_outputs(parametrization, Out)
  @assert length(predictors) > 0
  sample_size = length(instances)
  validation_size = length(validation_instances)
  n_predictors = length(predictors)
  sample_indices = collect(1:sample_size) # filter for nonzero weights HERE
  predictor_weights = [
    draw_weights(weights, sample_size) for p in 1:n_predictors ]
  n_batches = div(sample_size, batch_size)
  n_constraints = length(constraints) # C
  CIs = [ constraint_instance_type(constraint, predictors[1], parametrization)
    for constraint in constraints ] # broadcasting doesn't seem to work here
  constraint_instances = [ Matrix{CI}(undef, n_constraint_draws, n_batches) for CI in CIs ]
  scores = zeros(T, n_iterations)
  validations = zeros(T, n_iterations)
  progress = Progress(n_iterations, enabled=verbose)
  for iteration in 1:n_iterations
    shuffle!(sample_indices)
    batch_indices = [
      @view sample_indices[(b-1)*batch_size+1 : b*batch_size]
      for b in 1:n_batches ]
    for (constraint_sample, constraint) in zip(constraint_instances, constraints)
      for draw in 1:n_constraint_draws, batch in 1:n_batches
        constraint_sample[draw, batch] = draw_instance(
          constraint, predictors[1], parametrization, instances)
      end
    end
    @threads for (weights, predictor) in # `collect` is needed for `@threads`
        zip(predictor_weights, predictors) |> collect
      Flux.train!(get_params(predictor),
          enumerate(batch_indices), optimizer) do batch, indices
        batch_weights = @view weights[indices]
        batch_instances = @view instances[indices]
        fit = mean(zip(batch_weights, batch_instances)) do (weight, instance)
          weight * log_likelihood(predictor, parametrization, instance)
        end
        penalty = n_constraints > 0 ? (
          mean(zip(constraint_instances, constraints)) do (constraint_sample, constraint)
            mean(@view constraint_sample[:, batch]) do constraint_instance
              apply_operator(
                constraint, predictor, parametrization, constraint_instance)
            end
          end) : T(0)
        -fit + constraint_weight*penalty
      end
    end
    score = mean(instances) do instance
      ensemble_log_likelihood(
        predictors, parametrization, instance)
    end
    validation = ( validation_size == 0 ? T(NaN) :
      mean(validation_instances) do instance
        ensemble_log_likelihood(
          predictors, parametrization, instance)
      end )
    scores[iteration] = score
    validations[iteration] = validation
    report = () -> [
      (:score, Float16(score)),
      (:validation, Float16(validation)) ]
    next!(progress, showvalues=report)
  end
  (; scores, validations, predictor_weights,
     predictors, instances, validation_instances ) # these ones are just returned back as passed in
end


module AtlasSim

using LinearAlgebra
using RigidBodySim
using HumanoidLCMSim
using RigidBodyDynamics
using RigidBodyDynamics.Contact

import AtlasRobot

using LCMCore: LCM, publish
using DiffEqBase: CallbackSet, solve, init
using OrdinaryDiffEq: Tsit5
using MechanismGeometries: URDFVisuals
using Blink: Window

function atlasrobotinfo(mechanism::Mechanism)
    actuators = [Actuator(string(j), JointID(j)) for j in tree_joints(mechanism) if num_positions(j) == 1]
    feet = Dict(
        Sides.left => findbody(mechanism, "l_foot"),
        Sides.right => findbody(mechanism, "r_foot"))
    hands = Dict(
        Sides.left => findbody(mechanism, "l_hand"),
        Sides.right => findbody(mechanism, "r_hand"))
    HumanoidRobotInfo(mechanism, feet, hands, actuators)
end

function send_init_messages(state::MechanismState, receiver::LCMControlReceiver)
    HumanoidLCMSim.publish_robot_state(receiver, 0.0, state)
    utime = HumanoidLCMSim.utime_t()
    lcm = LCM()
    @info "Publishing START_MIT_STAND."
    publish(lcm, "START_MIT_STAND", utime)
    sleep(2)
    nothing
end

function make_callback(state::MechanismState, headless::Bool, max_rate; controls::Bool = true)
    callback = CallbackSet()
    if max_rate < Inf
        callback = CallbackSet(callback, RealtimeRateLimiter(max_rate = max_rate))
    end
    if headless
        if controls
            controls = SimulationControls()
            open(controls, Window())
            callback = CallbackSet(callback, CallbackSet(controls))
        end
    else
        visuals = URDFVisuals(AtlasRobot.urdfpath(); package_path = [AtlasRobot.packagepath()])
        gui = GUI(state.mechanism, visuals)
        open(gui)
        copyto!(gui.visualizer, state)
        callback = CallbackSet(callback, CallbackSet(gui))
    end
    callback
end

function simulate(dynamics::Dynamics, state0, final_time, callback)
    problem = ODEProblem(dynamics, state0, (0., final_time), callback = callback)
    integrator = init(problem, Tsit5(); abs_tol = 1e-8, dt = 1e-6, dtmin = 0.0)
    @info "Starting simulation."
    walltime = @elapsed solve!(integrator)
    integrator.sol, walltime
end

"""
    run(; controlΔt::Float64 = 1 / 300; headless = false)

Start a simulation of the Atlas robot, controlled using a `HumanoidLCMSim.LCMControlReceiver` running
at a rate of `1 / controlΔt` Hertz.

# Examples

Default usage:

```julia
using HumanoidLCMSim; AtlasSim.run()
```
"""
function run(; final_time = Inf, controlΔt::Float64 = 1 / 300, headless = false, max_rate = Inf, controls = true)
    BLAS.set_num_threads(max(floor(Int, Sys.CPU_THREADS / 2  - 1), 1)) # leave some cores for other processes
    mechanism = AtlasRobot.mechanism(add_flat_ground=true)
    info = atlasrobotinfo(mechanism)
    state0 = MechanismState(mechanism)
    AtlasRobot.setnominal!(state0)
    receiver = LCMControlReceiver(info;
        robot_state_channel="EST_ROBOT_STATE",
        robot_command_channel="ATLAS_COMMAND")
    pcontroller = PeriodicController(controlΔt, receiver)
    control! = let pcontroller = pcontroller
        function (τ, t, state)
            pcontroller(τ, t, state)
            # add some damping. TODO: do in a more appropriate place
            v = velocity(state)
            τ .-= 1.0 .* v
            τ
        end
    end
    send_init_messages(state0, receiver)
    callback = CallbackSet(make_callback(state0, headless, max_rate; controls = controls),
        PeriodicCallback(pcontroller))
    sol, walltime = simulate(Dynamics(mechanism, control!), state0, final_time, callback)
    simtime = sol.t[end]
    println("Simulated $simtime s in $walltime s ($(simtime / walltime) × realtime).")
    sol
end

end # module


#md# It is the code to draw this repo's logo
using WordCloud
Luxor = WordCloud.Render.Luxor

function cloudshape(height, args...; backgroundcolor=(0, 0, 0, 0))
    height = ceil(height)
    r = height / 2
    d = Luxor.Drawing((2height, height)..., :svg)
    Luxor.origin()
    Luxor.background(parsecolor(backgroundcolor))
    Luxor.setcolor(parsecolor((0.22, 0.596, 0.149)))
    Luxor.pie(0, 0, r, 0, 2π, :fill)
    Luxor.setcolor(parsecolor((0.584, 0.345, 0.698)))
    Luxor.pie(r, r, r, -π, 0, :fill)
    Luxor.setcolor(parsecolor((0.796, 0.235, 0.20)))
    Luxor.Luxor.pie(-r, r, r, -π, 0, :fill)
    Luxor.finish()
    d
end
wc = wordcloud(
    repeat(string.('a':'z'), 5),
    repeat([1], 26 * 5),
    mask=cloudshape(500),
    transparent=(0, 0, 0, 0),
    colors=1,
    angles=-90:0,
    font="JuliaMono Black",
    density=0.7,
    )
generate!(wc, 2000, retry=1)
recolor!(wc, style=:main)
println("results are saved to logo.svg")
paint(wc, "logo.svg", background=false)
wc
#eval# runexample(:logo)
#md# ![](logo.svg)


function test_progress_messenger(messenger; LES=false, model_kwargs...)
    model = create_model(; model_kwargs...)
    test_progress_messenger(model, messenger)
end

function test_progress_messenger(model, messenger; LES=false)
    simulation = Simulation(model; Δt=1e-2,
                            stop_iteration=10,
                            iteration_interval=1,
                            progress=messenger,
                            )

    run!(simulation)

    return nothing
end


# ------------------------------------------------------------------
# Copyright (c) 2015, Júlio Hoffimann Mendes <juliohm@stanford.edu>
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

module GeoStats

using Reexport
using Random
using Distances
using Statistics
using LinearAlgebra
using Distributed
using StaticArrays
using RecipesBase

# export project modules
@reexport using GeoStatsBase
@reexport using GeoStatsDevTools
@reexport using KrigingEstimators

# extend base module
import GeoStatsBase: solve, solve_single, preprocess

# solvers installed by default
include("solvers/kriging.jl")
include("solvers/sgsim.jl")
include("solvers/cookiecutter.jl")

# solver comparisons
include("comparisons/cross_validation.jl")
include("comparisons/visual_comparison.jl")

export
  # solvers
  Kriging,
  SeqGaussSim,
  CookieCutter,

  # solver comparisons
  VisualComparison,
  CrossValidation

end # module


"""
< DEFINITION of DATA TYPES for JuliaPBF >
"""
module StructPBF
import Base.+, Base.*, Base.-, Base.isless, Base.abs2 # To add generic functions
export
    AnalysisDataStruct,
    SimulationDataStruct

    include("BasicStruct.jl")
    include("AnalysisDataStruct.jl")
    include("SimulationDataStruct.jl")

    """시뮬레이션 데이터"""
    struct AnalysisDataStruct
        timeStep            ::TimeStep
        endTime             ::Float64
        gravity             ::Vec2
        analysisBox         ::Box2
        nMaxParticles       ::Int64
        maxParticleSize     ::Float64
        kernelRadius        ::Float64
        tensileInstability  ::TensileInstability
        a_initialBox        ::Vector{InitialBox}
        a_phases            ::Vector{PhaseData}
    end

    struct SimulationDataStruct
        particles           ::Vector{ParticleData} # For AoS
        grids               ::GridData
        phases              ::Vector{PhaseData}
        boundaries          ::Vector{BoundaryData}
        alivedParticles     ::Vector{Int64} # Fluids / Solids
    end

end


import Pkg

# Julia Packages
using PGFPlots
using Profile
using StatProfilerHTML

# Allow for easy updating
using Revise

# Satellite Dynamics Packages
using SatelliteDynamics

# Load SatelliteTasking
using SatelliteTasking

# Add Map Plots
using PGFMapPlots

##
# Simulation Definition
##

# Set simulation time window
t_start = Epoch(2019, 1, 1, 0, 0, 0, tsys="UTC"); # Start of time span
t_end = Epoch(2019, 1, 2, 0, 0, 0, tsys="UTC"); # End of simulation time span

# Initialize Planning Problem
problem = SatPlanningProblem(t_start=t_start, t_end=t_end);

# Add constraints to problem
push!(problem.constraints, constraint_agility_single_axis)

# Create Spacecraft
line1 = "1     1U          19  1.00000000  .00000000  00000-0  00000-0 0    05"
line2 = "2     1  90.0000   0.0000 0010000   0.0000   0.0000 15.24162312    00"
sc1 = Spacecraft(id=1, tle=TLE(line1, line2), slew_rate=1.0);
push!(problem.spacecraft, sc1);

# Compute Orital Period
T = orbit_period(state(problem.spacecraft[1].tle, problem.spacecraft[1].tle.epoch)[1]);

println("Loadded $(length(problem.spacecraft)) spacecraft")

##
# Load Groundstation Data
##

stations = load_stations("./data/groundstations.json");
num_stations = length(stations);

# Load Image Data - [100, 200, 500, 1000, 2000, 5000, 10000, 12000, 15000]
num_requests = 1000
requests = load_requests("./data/landsat_test_$num_requests.json", id_offset=num_stations);
@assert num_requests == length(requests);

# Set Request Parameters
collect_duration = 30.0
set_collect_duration(requests, collect_duration)

# Add Locations To Problem
# add_locations(problem, stations)
add_locations(problem, requests)

##
# Compute Opportunities for collects and contacts
##

# Remove all existing opportunities
clear_opportunities(problem)

# Compute problem accesses
@time compute_access(problem, orbit_fraction=0.75)

println("Found $(length(problem.contacts)) passes in planning horizon")
println("Found $(length(problem.collects)) collects in planning horizon")
num_accessible_requests = sum([1 for req in problem.requests if length(problem.lt_loc_opps[req.id]) > 0])
pct_accessible_requests = floor(num_accessible_requests/length(problem.requests) * 10000)/100
println("Found $num_accessible_requests out of $(length(problem.requests)) requests have accesses. $pct_accessible_requests%")

##
# Baseline Solve (Single-Threaded)
##

# @time baseline_plan, baseline_reward = satellite_plan_baseline(problem, allow_repeats=false)

# # Check feasibility, and compute actual reward
# baseline_feasible, baseline_realized_reward, baseline_nur, baseline_ndr, baseline_nr, baseline_nc = analyze_plan(problem, baseline_plan);

# println("Computed plan is feasible: $(uppercase(string(baseline_feasible)))")
# println("Computed optimal schedule. Computed reward: $baseline_reward - Realized reward: $baseline_realized_reward.")
# println("Requests Satisfied: $baseline_nur/$(length(problem.requests)), Duplicate Collects: $baseline_ndr")
# println("Ground Contacts Taken: $baseline_nc")

##
# Graph Solve (Single-Threaded)
##

# @time graph_plan, graph_reward = satellite_plan_graph(problem, allow_repeats=false)

# # Check feasibility, and compute actual reward
# graph_feasible, graph_realized_reward, graph_nur, graph_ndr, graph_nr, graph_nc = analyze_plan(problem, graph_plan);

# println("Computed plan is feasible: $(uppercase(string(graph_feasible)))")
# println("Computed optimal schedule. Computed reward: $graph_reward - Realized reward: $graph_realized_reward.")
# println("Requests Satisfied: $graph_nur/$(length(problem.requests)), Duplicate Collects: $graph_ndr")
# println("Ground Contacts Taken: $graph_nc")

##
# Mixed-Integer Linear Programing Solve (Single-Threaded)
##

# @time milp_plan, milp_reward = satellite_plan_milp(problem, allow_repeats=false)

# # Check feasibility, and compute actual reward
# milp_feasible, milp_realized_reward, milp_nur, milp_ndr, milp_nr, milp_nc = analyze_plan(problem, milp_plan);

# println("")
# println("Computed plan is feasible: $(uppercase(string(milp_feasible)))")
# println("Computed optimal schedule. Computed reward: $milp_reward - Realized reward: $milp_realized_reward.")
# println("Requests Satisfied: $milp_nur/$(length(problem.requests)), Duplicate Collects: $milp_ndr")
# println("Ground Contacts Taken: $milp_nc")

##
# Markov Decision Process (No Resources)
##

# # Adjust solve parameters
# problem.solve_discount = 1.0 # Typical values 0.999 - 0.9999
# problem.solve_depth = 5
# problem.solve_breadth = 3
# problem.solve_horizon = T

# @time mdp_fs_plan, mdp_fs_reward = satellite_plan_mdp_fs(problem)

# # Check feasibility, and compute actual reward
# mdp_fs_feasible, mdp_fs_realized_reward, mdp_fs_nur, mdp_fs_ndr, mdp_fs_nr, mdp_fs_nc = analyze_plan(problem, mdp_fs_plan);

# println("Computed plan is feasible: $(uppercase(string(mdp_fs_feasible)))")
# println("Computed optimal schedule. Computed reward: $mdp_fs_reward - Realized reward: $mdp_fs_realized_reward.")
# println("Requests Satisfied: $mdp_fs_nur/$(length(problem.requests)), Duplicate Collects: $mdp_fs_ndr")
# println("Ground Contacts Taken: $mdp_fs_nc")

# # Profile MDP Planning code
# @profilehtml satellite_plan_mdp_fs(problem);

##
# Markov Decision Process - Monte Carlo Tree Search (No Resources)
##

# Adjust solve parameters
problem.solve_allow_repeats = false
problem.mdp_reward_scarcity = false
problem.solve_discount = 1.0 # Typical values 0.99 - 0.9999
problem.solve_depth = 10
problem.solve_breadth = 0
problem.solve_horizon = T
problem.mcts_sim_iterations = 30
problem.mcts_c = 3.0

mdp_mcts_plan, mdp_mcts_reward = satellite_plan_mdp_mcts(problem, parallel=false)

# Check feasibility, and compute actual reward
mdp_mcts_feasible, mdp_mcts_realized_reward, mdp_mcts_nur, mdp_mcts_ndr, mdp_mcts_nr, mdp_mcts_nc = analyze_plan(problem, mdp_mcts_plan);

println("Computed plan is feasible: $(uppercase(string(mdp_mcts_feasible)))")
println("Computed optimal schedule. Computed reward: $mdp_mcts_reward - Realized reward: $mdp_mcts_realized_reward.")
println("Requests Satisfied: $mdp_mcts_nur/$(length(problem.requests)), Duplicate Collects: $mdp_mcts_ndr")
println("Ground Contacts Taken: $mdp_mcts_nc")

# Profile MCTS
@profilehtml satellite_plan_mdp_mcts(problem);


### Functions for plotting stellar cluster results
###

#### aux. function to plot text in a subplot
## plot array of text in a box
function show_text(posx,posy, text, ywidth = 1.)
    dct= Dict("color"=> "black", "fontsize"=> 11)
    dy = ywidth / length(text)
    for t in text
        PyPlot.plt.text(posx,posy,t, dct)
        posy += dy
    end
end

### plot the stats of the DBSCAN parameters from the MCMC
function plot_dbscan_mcmc(plotdir, voname, mc::GaiaClustering.mc , showplot = true)
    PyPlot.plt[:figure](figsize=(10.0,12.0))

    nbins = 50
    PyPlot.plt.subplot(3, 2, 1  )
    h = PyPlot.plt.hist(mc.eps,nbins, alpha=0.75)
    PyPlot.plt.xlabel("ϵ")
    PyPlot.plt.grid(true)

    nbins = 20
    PyPlot.plt.subplot(3, 2, 2 )
    h = PyPlot.plt.hist(mc.mne,nbins, alpha=0.75)
    PyPlot.plt.xlabel("min_neighbor")
    PyPlot.plt.grid(true)

    nbins = 20
    PyPlot.plt.subplot(3, 2, 3 )
    h = PyPlot.plt.hist(mc.mcl,nbins, alpha=0.75)
    PyPlot.plt.xlabel("min_cluster")
    PyPlot.plt.grid(true)

    ### text

    PyPlot.plt.subplot(3, 2, 4)
    PyPlot.plt.axis("off")
    ## text to display
        text =[]
        push!(text, "VOT: "*voname)
        vtext  = mean(mc.eps)
        v2text = std(mc.eps)
        txt = "ϵ : $vtext +/- $v2text "
        push!(text,txt)
        vtext  = mean(mc.mne)
        v2text = std(mc.mne)
        txt = "min_neigh : $vtext +/- $v2text "
        push!(text,txt)
        vtext  = mean(mc.mcl)
        v2text = std(mc.mcl)
        txt = "min_clus : $vtext +/- $v2text "
        push!(text,txt)
        vtext  = mean(mc.qn)
        v2text = std(mc.qn)
        txt = "Qn : $vtext +/- $v2text "
        push!(text,txt)
        vtext  = mean(mc.qc)
        v2text = std(mc.qc)
        txt = "Qc : $vtext +/- $v2text "
        push!(text,txt)
        show_text(-0.01,0.0, text)

    #######
    nbins = 50
    PyPlot.plt.subplot(3, 2, 6 )
    PyPlot.plt.axis("on")
    PyPlot.plt.axis("on")
    h = PyPlot.plt.hist(mc.qn,nbins, color = "g", alpha=0.75)
    PyPlot.plt.xlabel("Qn")
    PyPlot.plt.grid(true)

    nbins = 50
    PyPlot.plt.subplot(3, 2, 5 )
    h = PyPlot.plt.hist(mc.qc,nbins, color="g", alpha=0.75)
    PyPlot.plt.grid(true)

    PyPlot.plt.xlabel("Qc")
    figname = plotdir*"/"*voname*".mcmc.png"
    PyPlot.plt.savefig(figname)
    if showplot PyPlot.plt.show() end
end

### plot the stats of the DBSCAN parameters from the MCMC
function plot_dbscanfull_mcmc(plotdir, voname, mc::mcfull , showplot = true)
    patch= pyimport("matplotlib.patches")
    PyPlot.plt.figure(figsize=(12.0,13.0))

    nbins = 50
    PyPlot.plt.subplot(3, 3, 1  )
    h = PyPlot.plt.hist(mc.eps,nbins, alpha=0.75)
    PyPlot.plt.xlabel("ϵ")
    PyPlot.plt.grid(true)

    nbins = 20
    PyPlot.plt.subplot(3, 3, 2 )
    h = PyPlot.plt.hist(mc.mne,nbins, alpha=0.75)
    PyPlot.plt.xlabel("min_neighbor")
    PyPlot.plt.grid(true)

    nbins = 20
    PyPlot.plt.subplot(3, 3, 3 )
    h = PyPlot.plt.hist(mc.mcl,nbins, alpha=0.75)
    PyPlot.plt.xlabel("min_cluster")
    PyPlot.plt.grid(true)

    nbins = 50
    PyPlot.plt.subplot(3, 3, 4 )
    h = PyPlot.plt.hist(mc.w3d,nbins, alpha=0.75)
    PyPlot.plt.xlabel("W3d")
    PyPlot.plt.grid(true)

    nbins = 50
    PyPlot.plt.subplot(3, 3, 5 )
    h = PyPlot.plt.hist(mc.wvel,nbins, alpha=0.75)
    PyPlot.plt.xlabel("Wvel")
    PyPlot.plt.grid(true)

    nbins = 50
    PyPlot.plt.subplot(3, 3, 6 )
    h = PyPlot.plt.hist(mc.whrd,nbins, alpha=0.75)
    PyPlot.plt.xlabel("Whrd")
    PyPlot.plt.grid(true)

    #######
    nbins = 50
    PyPlot.plt.subplot(3, 3, 7 )
    PyPlot.plt.axis("on")
    PyPlot.plt.axis("on")
    h = PyPlot.plt.hist(mc.qn,nbins, color = "g", alpha=0.75)
    PyPlot.plt.xlabel("Qn")
    PyPlot.plt.grid(true)

    nbins = 50
    PyPlot.plt.subplot(3, 3, 8 )
    h = PyPlot.plt.hist(mc.qc,nbins, color="g", alpha=0.75)
    PyPlot.plt.xlabel("Qc")
    PyPlot.plt.grid(true)

        ### text

    axt= PyPlot.plt.subplot(3, 3, 9)
    PyPlot.plt.axis("off")
    ## text to display
        dct= Dict( :color=> "black", :fontsize=> 10)
        text =[]
        push!(text, "votable: "*voname)
        vtext  =  fmt("3.3f", median(mc.eps))
        v2text =  fmt("3.3f", std(mc.eps))
        txt = "ϵ : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.mne))
        v2text =  fmt("3.3f", std(mc.mne))
        txt = "min_neigh : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.mcl))
        v2text =  fmt("3.3f", std(mc.mcl))
        txt = "min_clus : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.w3d))
        v2text =  fmt("3.3f", std(mc.w3d))
        txt = "W3d : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.wvel))
        v2text =  fmt("3.3f", std(mc.wvel))
        txt = "Wvel : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.whrd))
        v2text =  fmt("3.3f", std(mc.whrd))
        txt = "Whrd : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.qn))
        v2text =  fmt("3.3f", std(mc.qn))
        txt = "Qn : $vtext +/- $v2text "
        push!(text,txt)
        vtext  =  fmt("3.3f", median(mc.qc))
        v2text =  fmt("3.3f", std(mc.qc))
        txt = "Qc : $vtext +/- $v2text "
        push!(text,txt)
        show_text(-0.01,0.0, text)

        rec= patch.Rectangle((-0.07, -0.05), 1.2 , 1.05, color="mediumaquamarine",
            alpha= 0.4, clip_on=false)
        axt.add_artist(rec)

    PyPlot.plt.xlabel("Qc")
    figname = plotdir*"/"*voname*".mcmc.png"
    PyPlot.plt.savefig(figname)
    if showplot PyPlot.plt.show() end
end

function plot_cluster(plotdir, voname, indx, sc::GaiaClustering.SCproperties, df::GaiaClustering.Df, showplot = true , cmap = "gist_stern")

    println("### Cluster plot is centered in Y,Z...")

    PyPlot.plt.figure(figsize=(13.0,12.0))
    PyPlot.plt.subplot(3, 3, 1 , xlim = [-20,20] , ylim = [-20,20])

    xx = df.data[2,indx] .- mean(df.data[2,indx])
    yy = df.data[3,indx] .- mean(df.data[3,indx])

    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 2 , ylim = [-20,20])
    xx = df.data[1,indx]
    yy = df.data[3,indx] .- mean(df.data[3,indx])
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("X (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 4 , xlim = [-20,20])
    xx = df.data[2,indx] .- mean(df.data[2,indx])
    yy = df.data[1,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("X (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 3 )
    xx = df.data[1,indx]
    yy = df.raw[13,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("X(pc)")
    PyPlot.plt.ylabel("Vrad (km/s)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 5)
    PyPlot.plt.axis("off")
    ## text to display
    text =[]
    v = sc.nstars ; txt = "N stars   : $v" ; push!(text,txt)
    v = fmt("3.1f",sc.distance) ; txt = "Distance  : $v (pc)" ; push!(text,txt)
    v = fmt("3.3f",sc.l) ; txt = "l         : $v (degree)" ; push!(text,txt)
    v = fmt("3.3f",sc.b) ; txt = "b         : $v (degree)" ; push!(text,txt)
    v = fmt("3.2f", sc.vl) ; txt = "Vl      : $v (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vb) ; txt = "Vb      : $v (km/s)"; push!(text,txt)
    v = fmt("3.2f",sc.vrad) ; txt = "Vradial   : $v (km/s)"; push!(text,txt)
    v = fmt("3.2f",sc.xdisp) ; txt = "X disp.   : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.ydisp) ; txt = "Y disp.   : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.zdisp) ; txt = "Z disp.   : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.vldisp) ; txt = "Vl disp. : $v (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vbdisp) ; txt = "Vb disp.: $v (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vraddisp) ; txt = "Vradial disp.: $v (km/s)" ; push!(text,txt)
    show_text(-0.01,0.0, text , 1.0)

    PyPlot.plt.subplot(3, 3, 7 )
    PyPlot.plt.axis("on")
    xx = df.data[7,indx]
    yy = -df.data[6,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("G-Rp")
    PyPlot.plt.ylabel("G")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 8 )
    xx = df.data[4,indx]
    yy = df.data[5,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Vl (km/s)")
    PyPlot.plt.ylabel("Vb (km/s)")
    PyPlot.plt.grid(true)

    figname = plotdir*"/"*voname*".cluster.png"
    PyPlot.plt.savefig(figname)
    if showplot PyPlot.plt.show() end
end

function plot_cluster2(plotdir, voname, indx, sc::GaiaClustering.SCproperties2, df::GaiaClustering.Df,
    showplot = true , extra= 0, cmap = "gist_stern")
    patch= pyimport("matplotlib.patches")

    println("### Cluster plot is centered in Y,Z...")
    PyPlot.plt.rcParams["font.size"]= 25

    PyPlot.plt.figure(figsize=(13.0,12.0))
    PyPlot.plt.subplot(3, 3, 1 , xlim = [-20,20] , ylim = [-20,20])

    xx = df.data[2,indx] .- mean(df.data[2,indx])
    yy = df.data[3,indx] .- mean(df.data[3,indx])

    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 2 , ylim = [-20,20])
    xx = df.data[1,indx]
    yy = df.data[3,indx] .- mean(df.data[3,indx])
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("X (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 4 , xlim = [-20,20])
    xx = df.data[2,indx] .- mean(df.data[2,indx])
    yy = df.data[1,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("X (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 3 )
    xx = df.data[1,indx]
    yy = df.raw[13,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("X(pc)")
    PyPlot.plt.ylabel("Vrad (km/s)")
    PyPlot.plt.grid(true)

    axt= PyPlot.plt.subplot(3, 3, 5)
    PyPlot.plt.axis("off")

    ## text to display
    dct= Dict("color"=> "black", "fontsize"=> 10)
    text =[]
    v = sc.nstars ; txt = "N stars : $v" ; push!(text,txt)
    v = fmt("3.1f",sc.distance) ; txt = "Distance : $v (pc)" ; push!(text,txt)
    v1 = fmt("3.3f",sc.l) ; v2 = fmt("3.3f",sc.b) ;
    txt = "l , b : $v1  ,  $v2  (degree)" ; push!(text,txt)
    v1 = fmt("3.3f",sc.ra) ; v2 = fmt("3.3f",sc.dec) ;
    txt = "RA , Dec : $v1  ,  $v2  (degree)" ; push!(text,txt)

    v1 = fmt("3.3f",sc.vl) ; v2 = fmt("3.3f",sc.vb) ;
    txt = "vl , vb : $v1  ,  $v2  (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vrad) ; txt  = "Vradial : $v (km/s)"; push!(text,txt)
    v = fmt("3.2f",sc.xdisp) ; txt = "X disp. : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.ydisp) ; txt = "Y disp. : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.zdisp) ; txt = "Z disp. : $v (pc)" ; push!(text,txt)
    v = fmt("3.2f",sc.vldisp) ; txt = "Vl disp. : $v (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vbdisp) ; txt = "Vb disp. : $v (km/s)" ; push!(text,txt)
    v = fmt("3.2f",sc.vraddisp) ; txt = "Vradial disp. : $v (km/s)" ; push!(text,txt)
    show_text(-0.01,-0.1, text , 1.1 )

    if extra != 0
        text =[]
        v1= "$(extra.votname[1])"
        txt = "Votable : $v1" ; push!(text,txt)
        v1= "$(extra.cycle[1])"
        txt = "Cycle : $v1" ; push!(text,txt)
        v1= fmt("3.3f",extra.qc[1])
        txt = "Qc : $v1" ; push!(text,txt)
        v1= fmt("3.3f",extra.score_cycle[1])
        txt = "Score : $v1" ; push!(text,txt)
        v1= fmt("3.3f",sc.offdeg)
        txt = "Offset : $v1 (degree)" ; push!(text,txt)
        v1= fmt("3.3f",sc.edgratg)
        txt = "Edge ratio(g) : $v1 " ; push!(text,txt)
        v1= fmt("3.3f",sc.edgratm)
        txt = "Edge ratio(m) : $v1 " ; push!(text,txt)
        v = @sprintf("XG : %6.1f (pc)", extra.xg[1]) ; push!(text,v)
        v = @sprintf("YG : %6.1f (pc)", extra.yg[1]) ; push!(text,v)
        v = @sprintf("ZG : %6.1f (pc)", extra.zg[1]) ; push!(text,v)
        show_text(1.2,-0.1, text , 0.95 )

        rec= patch.Rectangle((-0.07, -0.15), 2.2, 1.15, color="salmon", alpha= 0.4, clip_on=false)
        axt.add_artist(rec)
    end

    PyPlot.plt.subplot(3, 3, 7 )
    PyPlot.plt.axis("on")
    xx = df.data[7,indx]
    yy = -df.data[6,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("G-Rp")
    PyPlot.plt.ylabel("G")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 8 )
    xx = df.data[4,indx]
    yy = df.data[5,indx]
    PyPlot.plt.scatter(xx, yy , s = 1.0 )
    PyPlot.plt.xlabel("Vl (km/s)")
    PyPlot.plt.ylabel("Vb (km/s)")
    PyPlot.plt.grid(true)

    figname = plotdir*"/"*voname*".cluster.png"
    PyPlot.plt.savefig(figname)
    if showplot PyPlot.plt.show() end
end

## Raw data of the field with Principal Component plots of the extracted cluster
function plot_rawdata(plotdir, voname, indx, sc::GaiaClustering.SCproperties2, df::GaiaClustering.Df, pc, ijump=100,
    showplot = true , extra= 0, cmap = "gist_stern")
    patch= pyimport("matplotlib.patches")

    nsize= size(df.data)
    nstar= nsize[2]
    iter= 1:ijump:nstar

    PyPlot.plt.rcParams["font.size"]= 25

    PyPlot.plt.figure(figsize=(13.0,12.0))
    PyPlot.plt.subplot(3, 3, 1)

    xx = df.data[2,1iter]
    yy = df.data[3,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[2,indx]
    yy = df.data[3,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 2 )
    xx = df.data[1,iter]
    yy = df.data[3,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[1,indx]
    yy = df.data[3,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("X (pc)")
    PyPlot.plt.ylabel("Z (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 4 )
    xx = df.data[2,iter]
    yy = df.data[1,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[2,indx]
    yy = df.data[1,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("Y (pc)")
    PyPlot.plt.ylabel("X (pc)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 3 )
    xx = df.data[1,iter]
    yy = df.raw[13,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[1,indx]
    yy = df.raw[13,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("X(pc)")
    PyPlot.plt.ylabel("Vrad (km/s)")
    PyPlot.plt.grid(true)

    ## text to display
    axt= PyPlot.plt.subplot(3, 3, 5)
    PyPlot.plt.axis("off")
    ## text to display

    if extra != 0
        text =[]
        v1= "$(extra.votname[1])"
        txt = "Votable : $v1" ; push!(text,txt)
        v = nstar ; txt = "N stars in the field  : $v" ; push!(text,txt)
        v1= "$(extra.cycle[1])"
        txt = "Cycle : $v1" ; push!(text,txt)
        v1= @sprintf("%3.3f",extra.qc[1])
        txt = "Qc : $v1" ; push!(text,txt)
        v1= @sprintf("%3.3f",extra.score_cycle[1])
        txt = "Score : $v1" ; push!(text,txt)
        v1= @sprintf("%3.3f",sc.offdeg)
        txt = "Offset : $v1 (degree)" ; push!(text,txt)
        v1= @sprintf("%3.3f",sc.edgratg)
        txt = "Edge ratio(g) : $v1 " ; push!(text,txt)
        v1= @sprintf("%3.3f",sc.edgratm)
        txt = "Edge ratio(m) : $v1 " ; push!(text,txt)
        txt= @sprintf("PC1: %3.1f , PC2: %3.1f , PC3: %3.1f", extra.pc1[1], extra.pc2[1], extra.pc3[1])
        push!(text,txt)

        show_text(-0.01,-0.1, text , 1.1 )
        rec= patch.Rectangle((-0.07, -0.15), 1.2, 1.15, color="skyblue", alpha= 0.4, clip_on=false)
        axt.add_artist(rec)
    end

    PyPlot.plt.subplot(3, 3, 7 )
    PyPlot.plt.axis("on")
    xx = df.data[7,iter]
    yy = -df.data[6,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[7,indx]
    yy = -df.data[6,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("G-Rp")
    PyPlot.plt.ylabel("G")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 8 )
    xx = df.data[4,iter]
    yy = df.data[5,iter]
    PyPlot.plt.scatter(xx, yy , s = 0.1 )
    xx = df.data[4,indx]
    yy = df.data[5,indx]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r", alpha=0.5 )
    PyPlot.plt.xlabel("Vl (km/s)")
    PyPlot.plt.ylabel("Vb (km/s)")
    PyPlot.plt.grid(true)

    PyPlot.plt.subplot(3, 3, 9 )
    xx = pc[1,:]
    yy = pc[2,:]
    PyPlot.plt.scatter(xx, yy , s = 1, c="r" )
    PyPlot.plt.xlabel("PC1")
    PyPlot.plt.ylabel("PC2")
    PyPlot.plt.grid(true)

    figname = plotdir*"/"*voname*".raw.png"
    PyPlot.plt.savefig(figname)

    if showplot PyPlot.plt.show() end
end


"""Implementation of RRT-connect

ref:
- Kuffner, J. J., & LaValle, S. M. (2000).
  RRT-connect: An efficient approach to single-query path planning.
  In ICRA.
"""
module LibRRT_connect
export RRT_connect

import Printf: @sprintf, @printf
import ...MRMP: AbsState, Node, now, elapsed_sec, gen_uniform_sampling, get_mid_status, dist

"""
    RRT_connect(
        config_init::Vector{State},
        config_goal::Vector{State},
        connect::Function,
        collide::Function,
        check_goal::Function;
        steering_depth::Int64 = 4,
        epsilon::Union{Float64,Nothing} = 0.2,
        TIME_LIMIT::Union{Nothing,Real} = nothing,
        VERBOSE::Int64 = 0,
    )::Tuple{
        Union{Nothing,Vector{Vector{Node{State}}}},  # solution
        Vector{Vector{Node{State}}},  # roadmaps
    } where {State<:AbsState}

implementation of RRT-connect
"""
function RRT_connect(
    config_init::Vector{State},
    config_goal::Vector{State},
    connect::Function,
    collide::Function,
    check_goal::Function;
    steering_depth::Int64 = 4,
    epsilon::Union{Float64,Nothing} = 0.2,
    TIME_LIMIT::Union{Nothing,Real} = nothing,
    VERBOSE::Int64 = 0,
)::Tuple{
    Union{Nothing,Vector{Vector{Node{State}}}},  # solution
    Vector{Vector{Node{State}}},  # roadmaps
} where {State<:AbsState}

    # utilities for timeout
    t_s = now()
    elapsed() = elapsed_sec(t_s)
    timeover() = TIME_LIMIT != nothing && elapsed() > TIME_LIMIT

    # number of agents
    N = length(config_init)

    # define sampler
    _sampler = gen_uniform_sampling(config_init[1])
    sampler() = map(i -> _sampler(), 1:N)

    # re-define connect function
    connect_C(C_from::Vector{State}, C_to::Vector{State})::Bool = begin
        # check distance
        !(isnothing(epsilon) || dist(C_from, C_to) <= epsilon) && return false

        # check collision with obstacles
        any(i -> !connect(C_from[i], C_to[i], i), 1:N) && return false

        # check inter-agent collisions
        for i = 1:N, j = i+1:N
            collide(C_from[i], C_to[i], C_from[j], C_to[j], i, j) && return false
        end

        return true
    end

    connect_C(C::Vector{State})::Bool = all(i -> connect(C[i], i), 1:N)

    # store all samples
    V1 = [config_init]
    V2 = [config_goal]
    # parents index
    P1 = [-1]
    P2 = [-1]

    V = V1
    P = P1

    # special case
    if connect_C(config_init, config_goal)
        return (
            [
                map(q -> Node{State}(q, 1, []), config_init),
                map(q -> Node{State}(q, 2, []), config_goal),
            ],
            get_roadmaps(V1, P1, V2, P2),
        )
    end

    # main loop
    iter = 0
    from_start = false
    while !timeover()
        iter += 1
        from_start = !from_start

        VERBOSE > 1 &&
            iter % 100 == 0 &&
            @printf("\r\t%6.4f sec, iteration: %02d", elapsed(), iter)

        C_new = extend!(connect_C, sampler(), V1, P1, from_start, steering_depth)[end]
        if extend!(connect_C, C_new, V2, P2, !from_start, steering_depth)[1] == :reached
            # solved
            VERBOSE > 1 && @info @sprintf("\r\t%6.4f sec, iteration: %02d", elapsed(), iter)
            VERBOSE > 0 && @info @sprintf("\t%6.4f sec: found solution\n", elapsed())
            return (
                from_start ? (backtrack(V1, P1, V2, P2), get_roadmaps(V1, P1, V2, P2)) :
                (backtrack(V2, P2, V1, P1), get_roadmaps(V2, P2, V1, P1))
            )
        end

        # swap tree
        V1, V2 = V2, V1
        P1, P2 = P2, P1
    end

    VERBOSE > 0 && @info @sprintf("\t%6.4f sec: failed to find solution\n", elapsed())

    return (
        nothing,
        from_start ? get_roadmaps(V2, P2, V1, P1) : get_roadmaps(V1, P1, V2, P2),
    )
end

# define expand function
function extend!(
    connect::Function,
    C_rand::Vector{State},  # sampled point
    V::Vector{Vector{State}},  # vertices
    P::Vector{Int64},  # parent indexes
    from_start::Bool,  # true -> tree rooted C_init, false -> tree rooted C_goal
    steering_depth::Int64,
)::Tuple{Symbol,Vector{State}} where {State<:AbsState}

    # find nearest existing vertex
    ind_near = (
        from_start ? findmin(C -> dist(C, C_rand), V)[end] :
        findmin(C -> dist(C_rand, C), V)[end]
    )
    C_near = V[ind_near]
    flg = :reached

    C_l = C_near  # safe sample
    C_h = C_rand  # probably unsafe sample -> eventually safe

    # steering, binary search
    if from_start ? !connect(C_near, C_h) : !(connect(C_h) && connect(C_h, C_near))
        flg = :advanced
        for _ = 1:steering_depth
            C = map(e -> get_mid_status(e...), from_start ? zip(C_l, C_h) : zip(C_h, C_l))
            if from_start ? connect(C_near, C) : (connect(C) && connect(C, C_near))
                C_l = C
            else
                C_h = C
            end
        end
        C_h = C_l
    end
    push!(V, C_h)
    push!(P, ind_near)

    return (flg, C_h)
end

function backtrack(
    V1::Vector{Vector{State}},  # tree rooted config_init
    P1::Vector{Int64},
    V2::Vector{Vector{State}},  # tree rooted config_goal
    P2::Vector{Int64},
)::Vector{Vector{Node{State}}} where {State<:AbsState}
    """create solution"""

    solution = Vector{Vector{Node{State}}}()
    ind = length(V1)
    while ind != -1
        pushfirst!(solution, map(q -> Node{State}(q, ind, []), V1[ind]))
        ind = P1[ind]
    end
    solution = solution[1:end-1]
    ind = length(V2)
    while ind != -1
        push!(solution, map(q -> Node{State}(q, ind, []), V2[ind]))
        ind = P2[ind]
    end
    solution
end

function get_roadmaps(
    V1::Vector{Vector{State}},  # tree rooted config_init
    P1::Vector{Int64},
    V2::Vector{Vector{State}},  # tree rooted config_goal
    P2::Vector{Int64},
)::Vector{Vector{Node{State}}} where {State<:AbsState}
    """create roadmaps for visualization"""

    roadmaps = map(i -> Vector{Node{State}}(), 1:length(V1[1]))
    for (k, C) in enumerate(V1)
        for (i, q) in enumerate(C)
            push!(roadmaps[i], Node{State}(q, k, []))
            P1[k] > 0 && push!(roadmaps[i][P1[k]].neighbors, k)
        end
    end
    K = length(V1)
    for (k, C) in enumerate(V2)
        for (i, q) in enumerate(C)
            push!(roadmaps[i], Node{State}(q, k, P2[k] > 0 ? [P2[k] + K] : []))
        end
    end
    roadmaps
end

end


const time_names_table = ["ns", "µs", "ms", "s", "m", "h", "d", "y"]
const time_step_table =     [1000, 1000, 1000, 60,  60,  24, 365]

using Formatting
using Logging
using LightGraphs, SimpleWeightedGraphs

function format_seconds_readable(seconds, digits=2)::String
    fspec = FormatSpec(".$(digits)f")
    if 0.1 <= seconds <= 120.0 || seconds == 0.0
        return "$(fmt(fspec, seconds))s"
    end

    if seconds > 60.0
        # Up
        minutes = seconds / 60.0
        if minutes > 120.0
            hours = minutes / 60.0
            return "$(fmt(fspec, hours))h"
        end
        return "$(fmt(fspec, minutes))m"
    else
        # Down
        milliseconds = seconds * 1000.0
        if milliseconds < 0.1
            microseconds = milliseconds * 1000.0
            if microseconds < 0.1
                nanoseconds = milliseconds * 1000.0
                return "$(fmt(fspec, nanoseconds))ns"
            end
            return "$(fmt(fspec, microseconds))µs"
        end
        return "$(fmt(fspec, milliseconds))ms"
    end
    return string(seconds)
end

function format_ratio_readable(amount::Union{Int, Float64}, max::Union{Int, Float64} = 1, digits=2)::String
    if amount == 0
        if max == 0
            return "100%"
        else
            return "0%"
        end
    end
    fspec = FormatSpec(".$(digits)f")
    perc::Union{Float64, Rational{Int}} = 0.0
    if typeof(amount) == Float64
        perc = amount/max
    else
        max::Int
        perc = amount//max
    end
    perc *= 100
    if 1//5 <= perc || perc == 0
        return "$(fmt(fspec, perc))%"
    end

    perc *= 100

    if 1//5 <= perc
        return "$(fmt(fspec, perc))‰"
    end

    perc *= 100

    if 1//5 <= perc
        return "$(fmt(fspec, perc))‱"
    end

    perc *= 100

    if 1//5 <= perc
        return "$(fmt(fspec, perc))pcm"
    end

    perc *= 100

    return "$(fmt(fspec, perc))ppm"

end

function print_weighted_graph(graph::SimpleWeightedGraph, level=nothing, numberofvertices=nv(graph), numberofedges=ne(graph))
    if level === nothing
        println(numberofvertices)
    else
        @logmsg level numberofvertices
    end
    if level === nothing
        println(numberofedges)
    else
        @logmsg level numberofedges
    end
    line::Int = 0
    for (i, e) in enumerate(edges(graph))
        msg::String = "$line $(src(e) - 1) $(dst(e) - 1) $(Int(round(weight(e))))"
        if level === nothing
            println(msg)
        else
            @logmsg level msg
        end
        line += 1
    end
end

function print_graph(graph::AbstractGraph, level=nothing, weight=nothing, numberofvertices=nv(graph), numberofedges=ne(graph))
    if level === nothing
        println(numberofvertices)
    else
        @logmsg level numberofvertices
    end
    if level === nothing
        println(numberofedges)
    else
        @logmsg level numberofedges
    end
    if weight !== nothing
        if level === nothing
            println(weight)
        else
            @logmsg level weight
        end
    end
    line::Int = 0
    for (i, e) in enumerate(edges(graph))
        msg::String = "$line $(src(e) - 1) $(dst(e) - 1)"
        if level === nothing
            println(msg)
        else
            @logmsg level msg
        end
        line += 1
    end
end


# use only activations without "kinks" for numerical gradient checking
# see e.g. https://stackoverflow.com/questions/40623512/how-to-check-relu-gradient
const ACTIVATIONS = [identity, σ, swish, softplus, logcosh, mish, tanhshrink, lisht]

@testset "testing simple matrix model" begin
    layerbuilder(k) = Flux.Dense(k, 2, NNlib.relu)
    x = ArrayNode(randn(Float32, 4, 5))
    m = reflectinmodel(x, layerbuilder)
    @test size(m(x).data) == (2, 5)
    @test m isa ArrayModel
    @test eltype(m(x).data) == Float32
    @inferred m(x)
end

@testset "testing simple aggregation model" begin
    layerbuilder(k) = Flux.Dense(k, 2, NNlib.relu)
    x = BagNode(ArrayNode(randn(Float32, 4, 4)), [1:2, 3:4])
    m = reflectinmodel(x, layerbuilder)
    @test size(m(x).data) == (2, 2)
    @test m isa BagModel
    @test eltype(m(x).data) == Float32
    @inferred m(x)
end

@testset "testing simple tuple models" begin
    layerbuilder(k) = Flux.Dense(k, 2, NNlib.relu)
    x = ProductNode((a=ArrayNode(randn(Float32, 3, 4)), b=ArrayNode(randn(Float32, 4, 4))))
    m = reflectinmodel(x, layerbuilder)
    @test eltype(m(x).data) == Float32
    @test size(m(x).data) == (2, 4)
    @test m isa ProductModel
    @test m.ms[:a] isa ArrayModel
    @test m.ms[:b] isa ArrayModel
    @inferred m(x)

    x = ProductNode((BagNode(ArrayNode(randn(Float32, 3, 4)), [1:2, 3:4]),
                     BagNode(ArrayNode(randn(Float32, 4, 4)), [1:1, 2:4])))
    m = reflectinmodel(x, layerbuilder)
    @test size(m(x).data) == (2, 2)
    @test m isa ProductModel
    @test m.ms[1] isa BagModel
    @test m.ms[1].im isa ArrayModel
    @test m.ms[1].bm isa ArrayModel
    @test m.ms[2] isa BagModel
    @test m.ms[2].im isa ArrayModel
    @test m.ms[2].bm isa ArrayModel
    @inferred m(x)
end

@testset "testing nested bag model" begin
    bn = BagNode(ArrayNode(randn(Float32, 2, 8)), [1:1, 2:2, 3:6, 7:8])
    x = BagNode(bn, [1:2, 3:4])
    m = reflectinmodel(x, d -> Flux.Dense(d, 2))
    @test size(m(x).data) == (2, 2)
    @test m isa BagModel
    @test m.im isa BagModel
    @test m.im.im isa ArrayModel
    @test m.im.bm isa ArrayModel
    @test m.bm isa ArrayModel
    @test eltype(m(x).data) == Float32
    @inferred m(x)

    a = BagNode(BagNode(ArrayNode(randn(2,2)),[1:2]),[1:1])
    b = BagNode(missing,[0:-1])
    c = BagNode(a.data[1:0], [0:-1])
    m = reflectinmodel(a)
    abc = catobs(a, b, c)
    bca = catobs(b, c, a)
    ma = m(a).data
    mb = m(b).data
    mc = m(c).data
    mabc = m(abc).data
    mbca = m(bca).data
    @test mb ≈ mc
    @test mabc[:,1] ≈ ma
    @test mabc[:,2] ≈ mb
    @test mabc[:,3] ≈ mc
    @test mbca[:,1] ≈ mb
    @test mbca[:,2] ≈ mc
    @test mbca[:,3] ≈ ma
    @inferred m(a)
    @inferred m(b)
    @inferred m(c)
end

@testset "reflectinmodel for missing + all_imputing" begin
    f1 = x -> reflectinmodel(x)
    f2 = x -> reflectinmodel(x, d -> Flux.Chain(Dense(d, 10), Dense(10, 10)))
    f3 = x -> reflectinmodel(x; all_imputing=true)
    _get_first(x::ArrayModel{<:Flux.Chain}) = x.m[1]
    _get_first(x) = x.m

    x1 = maybehot(2, 1:3) |> ArrayNode
    x2 = maybehot(missing, 1:3) |> ArrayNode
    for f in [f1, f2, f3], x in [x1, x2]
        @test _get_first(f(x)) isa PostImputingDense
        @inferred f(x)(x)
    end

    x1 = maybehotbatch([1, 2, 3], 1:3) |> ArrayNode
    x2 = maybehotbatch([1, 2, missing], 1:3) |> ArrayNode
    x3 = maybehotbatch(fill(missing, 3), 1:3) |> ArrayNode
    for f in [f1, f2, f3], x in [x1, x2, x3]
        @test _get_first(f(x)) isa PostImputingDense
        @inferred f(x)(x)
    end

    x1 = NGramMatrix(["a", "b", "c"]) |> ArrayNode
    x2 = NGramMatrix(["a", missing, "c"]) |> ArrayNode
    x3 = NGramMatrix(fill(missing, 3)) |> ArrayNode
    for f in [f1, f2]
        @test _get_first(f(x1)) isa Flux.Dense
        @test _get_first(f(x2)) isa PostImputingDense
        @test _get_first(f(x3)) isa PostImputingDense
        for x in [x1, x2, x3] @inferred f(x)(x) end
    end
    for x in [x1, x2, x3]
        @test _get_first(f3(x)) isa PostImputingDense
        @inferred f3(x)(x)
    end

    x1 = rand([1, 2], 3, 3) |> ArrayNode
    x2 = rand([1, 2, missing], 3, 3) |> ArrayNode
    x3 = fill(missing, 3, 3) |> ArrayNode
    for f in [f1, f2]
        @test _get_first(f(x1)) isa Flux.Dense
        @test _get_first(f(x2)) isa PreImputingDense
        @test _get_first(f(x3)) isa PreImputingDense
        for x in [x1, x2, x3] @inferred f(x)(x) end
    end
    for x in [x1, x2, x3]
        @test _get_first(f3(x)) isa PreImputingDense
        @inferred f3(x)(x)
    end
end

# pn.m should be identity for any product node pn with a single key
@testset "single key dictionary reflect in model" begin
    layerbuilder(k) = Flux.Dense(k, 2, NNlib.relu)
    fsm = Dict("" => layerbuilder)

    x1 = (ArrayNode(randn(Float32, 3, 4)),) |> ProductNode
    x2 = (a = ArrayNode(randn(Float32, 3, 4)),) |> ProductNode
    x3 = (a = ArrayNode(randn(Float32, 3, 4)), b = ArrayNode(randn(Float32, 3, 4))) |> ProductNode

    m1 = reflectinmodel(x1, layerbuilder; single_key_identity=false)
    m1_ski = reflectinmodel(x1, layerbuilder; single_key_identity=true)
    m1_ski_fsm = reflectinmodel(x1, layerbuilder; fsm=fsm, single_key_identity=true)
    m2 = reflectinmodel(x2, layerbuilder; single_key_identity=false)
    m2_ski = reflectinmodel(x2, layerbuilder; single_key_identity=true)
    m2_ski_fsm = reflectinmodel(x2, layerbuilder; fsm=fsm, single_key_identity=true)
    m3 = reflectinmodel(x3, layerbuilder; single_key_identity=false)
    m3_ski = reflectinmodel(x3, layerbuilder; single_key_identity=true)
    m3_ski_fsm = reflectinmodel(x3, layerbuilder; fsm=fsm, single_key_identity=true)

    for m in [m1, m1_ski, m1_ski_fsm]
        @test eltype(m(x1).data) == Float32
        @test size(m(x1).data) == (2, 4)
        @inferred m(x1)
        @test m isa ProductModel
        @test m.ms[1] isa ArrayModel
    end

    for m in [m2, m2_ski, m2_ski_fsm]
        @test eltype(m(x2).data) == Float32
        @test size(m(x2).data) == (2, 4)
        @inferred m(x2)
        @test m isa ProductModel
        @test m.ms[1] isa ArrayModel
    end

    for m in [m3, m3_ski, m3_ski_fsm]
        @test eltype(m(x3).data) == Float32
        @test size(m(x3).data) == (2, 4)
        @inferred m(x3)
        @test m isa ProductModel
        @test m.ms[1] isa ArrayModel
        @test m.ms[2] isa ArrayModel
    end

    @test m1.m isa ArrayModel{<:Dense}
    @test m2.m isa ArrayModel{<:Dense}
    @test m3.m isa ArrayModel{<:Dense}
    @test m1_ski.m isa IdentityModel
    @test m2_ski.m isa IdentityModel
    @test m3_ski.m isa ArrayModel{<:Dense}
    # fsm overrides ski
    @test m1_ski_fsm.m isa ArrayModel{<:Dense}
    @test m2_ski_fsm.m isa ArrayModel{<:Dense}
    @test m3_ski_fsm.m isa ArrayModel{<:Dense}
end

# array model for matrices with one row should implement identity
@testset "single scalar as identity" begin
    layerbuilder(k) = Flux.Dense(k, 2, NNlib.relu)
    x1 = ArrayNode(randn(Float32, 1, 3))
    x2 = BagNode(ArrayNode(randn(Float32, 1, 5)), [1:2, 3:5])
    x3 = (a = ArrayNode(randn(Float32, 1, 4)), b = ArrayNode(randn(Float32, 2, 4))) |> ProductNode

    m1 = reflectinmodel(x1, layerbuilder; single_scalar_identity=false)
    m1_sci = reflectinmodel(x1, layerbuilder; single_scalar_identity=true)
    m2 = reflectinmodel(x2, layerbuilder; single_scalar_identity=false)
    m2_sci = reflectinmodel(x2, layerbuilder; single_scalar_identity=true)
    m3 = reflectinmodel(x3, layerbuilder; single_scalar_identity=false)
    m3_sci = reflectinmodel(x3, layerbuilder; single_scalar_identity=true)

    @test size(m1(x1).data) == (2, 3)
    @test eltype(m1(x1).data) == Float32
    @test m1 isa ArrayModel
    @test size(m1_sci(x1).data) == (1, 3)
    @test eltype(m1_sci(x1).data) == Float32
    @test m1_sci isa ArrayModel
    @inferred m1(x1)

    @test size(m2(x2).data) == (2, 2)
    @test eltype(m2(x2).data) == Float32
    @test m2 isa BagModel
    @test m2.im isa ArrayModel{<:Dense}
    @test size(m2_sci(x2).data) == (2, 2)
    @test eltype(m2_sci(x2).data) == Float32
    @test m2_sci isa BagModel
    @test m2_sci.im isa IdentityModel
    @inferred m2(x2)

    @test size(m3(x3).data) == (2, 4)
    @test eltype(m3(x3).data) == Float32
    @test m3 isa ProductModel
    @test m3.ms[1] isa ArrayModel{<:Dense}
    @test m3.ms[2] isa ArrayModel{<:Dense}
    @test size(m3_sci(x3).data) == (2, 4)
    @test eltype(m3_sci(x3).data) == Float32
    @test m3_sci isa ProductModel
    @test m3_sci.ms[1] isa IdentityModel
    @test m3_sci.ms[2] isa ArrayModel{<:Dense}
    @inferred m3(x3)
end

@testset "model aggregation grad check w.r.t. inputs" begin
    for (bags1, bags2, bags3) in BAGS3
        layerbuilder(k) = Dense(k, 2, rand(ACTIVATIONS)) |> f64
        abuilder(d) = BagCount(all_aggregations(Float64, d))
        x = randn(4, 4)
        y = randn(3, 4)
        z = randn(2, 8)

        for ds in [x -> ArrayNode(x),
                   x -> BagNode(ArrayNode(x), bags1)]
            m = reflectinmodel(ds(x), layerbuilder, abuilder)
            @inferred m(ds(x))
            @test gradtest(x -> m(ds(x)).data, x)
        end

        for ds in [(x, y) -> ProductNode((ArrayNode(x), ArrayNode(y))),
                   (x, y) -> ProductNode((a=BagNode(ArrayNode(x), bags1), b=BagNode(ArrayNode(y), bags2)))]
            m = reflectinmodel(ds(x, y), layerbuilder, abuilder)
            @inferred m(ds(x, y))
            @test gradtest((x, y) -> m(ds(x, y)).data, x, y)
        end

        ds = z -> BagNode(BagNode(ArrayNode(z), bags3), bags1)
        m = reflectinmodel(ds(z), layerbuilder, abuilder)
        @inferred m(ds(z))
        @test gradtest(z -> m(ds(z)).data, z)
    end
end

@testset "model aggregation grad check w.r.t. inputs weighted" begin
    for (bags1, bags2, bags3) in BAGS3
        layerbuilder(k) = Dense(k, 2, rand(ACTIVATIONS)) |> f64
        abuilder(d) = BagCount(all_aggregations(Float64, d))
        x = randn(4, 4)
        y = randn(3, 4)
        z = randn(2, 8)
        w1 = abs.(randn(4)) .+ 0.1
        w2 = abs.(randn(4)) .+ 0.1
        w3 = abs.(randn(8)) .+ 0.1

        ds = x -> WeightedBagNode(ArrayNode(x), bags1, w1)
        m = reflectinmodel(ds(x), layerbuilder, abuilder)
        @inferred m(ds(x))
        @test gradtest(x -> m(ds(x)).data, x)

        ds = (x, y) -> ProductNode((WeightedBagNode(ArrayNode(x), bags1, w1),
                                    WeightedBagNode(ArrayNode(y), bags2, w2)))
        m = reflectinmodel(ds(x, y), layerbuilder, abuilder)
        @inferred m(ds(x, y))
        @test gradtest((x, y) -> m(ds(x, y)).data, x, y)

        ds = z -> WeightedBagNode(WeightedBagNode(ArrayNode(z), bags3, w3), bags1, w1)
        m = reflectinmodel(ds(z), layerbuilder, abuilder)
        @inferred m(ds(z))
        @test gradtest(z -> m(ds(z)).data, z)
    end
end

@testset "model aggregation grad check w.r.t. params" begin
    for (bags1, bags2, bags3) in BAGS3
        layerbuilder(k) = Dense(k, 2, rand(ACTIVATIONS)) |> f64
        abuilder(d) = BagCount(all_aggregations(Float64, d))
        x = randn(4, 4)
        y = randn(3, 4)
        z = randn(2, 8)

        for ds in [
                   ArrayNode(x),
                   BagNode(ArrayNode(x), bags1),
                   ProductNode((ArrayNode(x), ArrayNode(y))),
                   ProductNode((a=BagNode(ArrayNode(y), bags1), b=BagNode(ArrayNode(x), bags2))),
                   BagNode(BagNode(ArrayNode(z), bags3), bags1)
                  ]
            m = reflectinmodel(ds, layerbuilder, abuilder)
            @inferred m(ds)
            @test gradtest(() -> m(ds).data, Flux.params(m))
        end
    end
end

@testset "model aggregation grad check w.r.t. params weighted" begin
    for (bags1, bags2, bags3) in BAGS3
        layerbuilder(k) = Dense(k, 2) |> f64
        abuilder(d) = BagCount(all_aggregations(Float64, d))
        x = randn(4, 4)
        y = randn(3, 4)
        z = randn(2, 8)
        w1 = abs.(randn(4)) .+ 0.1
        w2 = abs.(randn(4)) .+ 0.1
        w3 = abs.(randn(8)) .+ 0.1

        for ds in [
                   WeightedBagNode(ArrayNode(x), bags1, w1),
                   ProductNode((WeightedBagNode(ArrayNode(x), bags1, w1),
                                WeightedBagNode(ArrayNode(y), bags2, w2))) ,
                   WeightedBagNode(WeightedBagNode(ArrayNode(z), bags3, w3), bags1, w1)
                  ]
            m = reflectinmodel(ds, layerbuilder, abuilder)
            @inferred m(ds)
            @test gradtest(() -> m(ds).data, Flux.params(m))
        end
    end
end


using Test
using NQCBase
using StaticArrays

atoms = Atoms{Float64}([:C, :H])
@test atoms isa Atoms{2, Float64}
@test length(atoms) == 2
@test atoms.numbers == [6, 1]
@test atoms.types == [:C, :H]
@test atoms.masses isa SVector

@test range(atoms) == 1:2

@test atoms[1] isa Atoms{1,Float64}
@test atoms[1:2] isa Atoms{2,Float64}

@test Atoms(2000) isa Atoms{1, Float64}