module HydrostaticBoussinesq

export HydrostaticBoussinesqModel, HydrostaticBoussinesqProblem, OceanDGModel

using StaticArrays
using LinearAlgebra: I, dot, Diagonal
using ..VariableTemplates
using ..MPIStateArrays
using ..DGmethods: init_ode_state
using ..PlanetParameters: grav
using ..Mesh.Filters: CutoffFilter, apply!, ExponentialFilter
using ..Mesh.Grids: polynomialorder, VerticalDirection

using ..DGmethods.NumericalFluxes: Rusanov, CentralGradPenalty,
                                   CentralNumericalFluxDiffusive,
                                   CentralNumericalFluxNonDiffusive

import ..DGmethods.NumericalFluxes: update_penalty!, numerical_flux_diffusive!,
                                    NumericalFluxNonDiffusive

import ..DGmethods: BalanceLaw, vars_aux, vars_state, vars_gradient,
                    vars_diffusive, vars_integrals, flux_nondiffusive!,
                    flux_diffusive!, source!, wavespeed,
                    boundary_state!, update_aux!, update_aux_diffusive!,
                    gradvariables!, init_aux!, init_state!,
                    LocalGeometry, indefinite_stack_integral!,
                    reverse_indefinite_stack_integral!, integrate_aux!,
                    DGModel, nodal_update_aux!, diffusive!,
                    copy_stack_field_down!, create_state

×(a::SVector, b::SVector) = StaticArrays.cross(a, b)
∘(a::SVector, b::SVector) = StaticArrays.dot(a, b)

abstract type OceanBoundaryCondition end
struct CoastlineFreeSlip             <: OceanBoundaryCondition end
struct CoastlineNoSlip               <: OceanBoundaryCondition end
struct OceanFloorFreeSlip            <: OceanBoundaryCondition end
struct OceanFloorNoSlip              <: OceanBoundaryCondition end
struct OceanSurfaceNoStressNoForcing <: OceanBoundaryCondition end
struct OceanSurfaceStressNoForcing   <: OceanBoundaryCondition end
struct OceanSurfaceNoStressForcing   <: OceanBoundaryCondition end
struct OceanSurfaceStressForcing     <: OceanBoundaryCondition end

abstract type HydrostaticBoussinesqProblem end

struct HydrostaticBoussinesqModel{P,T} <: BalanceLaw
  problem::P
  c₁::T
  c₂::T
  c₃::T
  αᵀ::T
  νʰ::T
  νᶻ::T
  κʰ::T
  κᶻ::T
end

HBModel   = HydrostaticBoussinesqModel
HBProblem = HydrostaticBoussinesqProblem

function OceanDGModel(bl::HBModel, grid, numfluxnondiff, numfluxdiff,
                      gradnumflux; kwargs...)
  vert_filter = CutoffFilter(grid, polynomialorder(grid)-1)
  exp_filter  = ExponentialFilter(grid, 1, 8)

  modeldata = (vert_filter = vert_filter, exp_filter=exp_filter)

  return DGModel(bl, grid, numfluxnondiff, numfluxdiff, gradnumflux;
                 kwargs..., modeldata=modeldata)
end

# If this order is changed check the filter usage!
function vars_state(m::HBModel, T)
  @vars begin
    u::SVector{2, T}
    η::T # real a 2-D variable TODO: should be 2D
    θ::T
  end
end

# If this order is changed check update_aux!
function vars_aux(m::HBModel, T)
  @vars begin
    w::T
    pkin_reverse::T # ∫(-αᵀ θ) # TODO: remove me after better integral interface
    w_reverse::T               # TODO: remove me after better integral interface
    pkin::T         # ∫(-αᵀ θ)
    wz0::T          # w at z=0
    θʳ::T           # SST given    # TODO: Should be 2D
    f::T            # coriolis
    τ::T            # wind stress  # TODO: Should be 2D
    # κ::SMatrix{3, 3, T, 9} # diffusivity tensor (for convective adjustment)
    κᶻ::T
  end
end

function vars_gradient(m::HBModel, T)
  @vars begin
    u::SVector{2, T}
    θ::T
  end
end

function vars_diffusive(m::HBModel, T)
  @vars begin
    ∇u::SMatrix{3, 2, T, 6}
    ∇θ::SVector{3, T}
  end
end

function vars_integrals(m::HBModel, T)
  @vars begin
    ∇hu::T
    αᵀθ::T
  end
end

@inline function flux_nondiffusive!(m::HBModel, F::Grad, Q::Vars,
                                    A::Vars, t::Real)
  @inbounds begin
    u = Q.u # Horizontal components of velocity
    η = Q.η
    θ = Q.θ
    w = A.w   # vertical velocity
    pkin = A.pkin
    v = @SVector [u[1], u[2], w]
    Ih = @SMatrix [ 1 -0;
                   -0  1;
                   -0 -0]

    # ∇ • (u θ)
    F.θ += v * θ

    # ∇h • (g η)
    F.u += grav * η * Ih

    # ∇h • (- ∫(αᵀ θ))
    F.u += grav * pkin * Ih

    # ∇h • (v ⊗ u)
    # F.u += v * u'

  end

  return nothing
end

@inline wavespeed(m::HBModel, n⁻, _...) = abs(SVector(m.c₁, m.c₂, m.c₃)' * n⁻)

# We want not have jump penalties on η (since not a flux variable)
function update_penalty!(::Rusanov, ::HBModel, n⁻, λ, ΔQ::Vars,
                         Q⁻, A⁻, Q⁺, A⁺, t)
  ΔQ.η = -0

  #=
  θ⁻ = Q⁻.θ
  u⁻ = Q⁻.u
  w⁻ = A⁻.w
  @inbounds v⁻ = @SVector [u⁻[1], u⁻[2], w⁻]
  n̂_v⁻ = n⁻∘v⁻

  θ⁺ = Q⁺.θ
  u⁺ = Q⁺.u
  w⁺ = A⁺.w
  @inbounds v⁺ = @SVector [u⁺[1], u⁺[2], w⁺]
  n̂_v⁺ = n⁻∘v⁺

  # max velocity
  # n̂∘v = (abs(n̂∘v⁺) > abs(n̂∘v⁻) ? n̂∘v⁺ : n̂∘v⁻

  # average velocity
  n̂_v = (n̂_v⁻ + n̂_v⁺) / 2

  ΔQ.θ = ((n̂_v > 0) ? 1 : -1) * (n̂_v⁻ * θ⁻ - n̂_v⁺ * θ⁺)
  # ΔQ.θ = abs(n̂_v⁻) * θ⁻ - abs(n̂_v⁺) * θ⁺
  =#

  return nothing
end

@inline function flux_diffusive!(m::HBModel, F::Grad, Q::Vars, D::Vars,
                                 A::Vars, t::Real)
  ν = Diagonal(@SVector [m.νʰ, m.νʰ, m.νᶻ])
  F.u -= ν * D.∇u

  κ = Diagonal(@SVector [m.κʰ, m.κʰ, A.κᶻ])
  F.θ -= κ * D.∇θ

  return nothing
end

@inline function gradvariables!(m::HBModel, G::Vars, Q::Vars, A, t)
  G.u = Q.u
  G.θ = Q.θ

  return nothing
end

@inline function diffusive!(m::HBModel, D::Vars, G::Grad, Q::Vars,
                            A::Vars, t)
  D.∇u = G.u
  D.∇θ = G.θ

  return nothing
end

@inline function source!(m::HBModel{P}, source::Vars, Q::Vars, A::Vars,
                         t::Real) where P
  @inbounds begin
    u = Q.u # Horizontal components of velocity
    f = A.f
    wz0 = A.wz0

    # f × u
    source.u -= @SVector [-f * u[2], f * u[1]]

    source.η += wz0
  end

  return nothing
end

@inline function integrate_aux!(m::HBModel, integrand::Vars, Q::Vars, A::Vars)
  αᵀ = m.αᵀ
  integrand.αᵀθ = -αᵀ * Q.θ
  integrand.∇hu = A.w # borrow the w value from A...

  return nothing
end

function update_aux!(dg::DGModel, m::HBModel, Q::MPIStateArray, t::Real)
  MD = dg.modeldata

  # required to ensure that after integration velocity field is divergence free
  vert_filter = MD.vert_filter
  # Q[1] = u[1] = u, Q[2] = u[2] = v
  apply!(Q, (1, 2), dg.grid, vert_filter, VerticalDirection())

  exp_filter = MD.exp_filter
  # Q[4] = θ
  apply!(Q, (4,), dg.grid, exp_filter, VerticalDirection())

  return true
end

function update_aux_diffusive!(dg::DGModel, m::HBModel, Q::MPIStateArray, t::Real)
  A  = dg.auxstate

  # store ∇ʰu as integrand for w
  # update vertical diffusivity for convective adjustment
  function f!(::HBModel, Q, A, D, t)
    @inbounds begin
      A.w = -(D.∇u[1,1] + D.∇u[2,2])

      # κʰ = m.κʰ
      # A.κ = @SMatrix [κʰ -0 -0; -0 κʰ -0; -0 -0 κᶻ]
      D.∇θ[3] < 0 ? A.κᶻ = 1000 * m.κᶻ : A.κᶻ = m.κᶻ
    end

    return nothing
  end
  nodal_update_aux!(f!, dg, m, Q, t; diffusive=true)

  # compute integrals for w and pkin
  indefinite_stack_integral!(dg, m, Q, A, t) # bottom -> top
  reverse_indefinite_stack_integral!(dg, m, A, t) # top -> bottom

  # project w(z=0) down the stack
  # Need to be consistent with vars_aux
  # A[1] = w, A[5] = wz0
  copy_stack_field_down!(dg, m, A, 1, 5)

  return true
end

function ocean_init_aux! end
function init_aux!(m::HBModel, A::Vars, geom::LocalGeometry)
  return ocean_init_aux!(m, m.problem, A, geom)
end

function ocean_init_state! end
function init_state!(m::HBModel, Q::Vars, A::Vars, coords, t)
  return ocean_init_state!(m.problem, Q, A, coords, t)
end

@inline function boundary_state!(nf, m::HBModel, Q⁺::Vars, A⁺::Vars, n⁻,
                                 Q⁻::Vars, A⁻::Vars, bctype, t, _...)
  return ocean_boundary_state!(m, bctype, nf, Q⁺, A⁺, n⁻, Q⁻, A⁻, t)
end
@inline function boundary_state!(nf, m::HBModel,
                                 Q⁺::Vars, D⁺::Vars, A⁺::Vars,
                                 n⁻,
                                 Q⁻::Vars, D⁻::Vars, A⁻::Vars,
                                 bctype, t, _...)
  return ocean_boundary_state!(m, bctype, nf, Q⁺, D⁺, A⁺, n⁻, Q⁻, D⁻, A⁻, t)
end

@inline function ocean_boundary_state!(::HBModel, ::CoastlineFreeSlip,
                                       ::Union{Rusanov, CentralGradPenalty},
                                       Q⁺, A⁺, n⁻, Q⁻, A⁻, t)
  return nothing
end

@inline function ocean_boundary_state!(::HBModel, ::CoastlineFreeSlip,
                                       ::CentralNumericalFluxDiffusive, Q⁺,
                                       D⁺, A⁺, n⁻, Q⁻, D⁻, A⁻, t)
  D⁺.∇u = -D⁻.∇u

  D⁺.∇θ = -D⁻.∇θ

  return nothing
end

@inline function ocean_boundary_state!(::HBModel, ::CoastlineNoSlip,
                                       ::Union{Rusanov, CentralGradPenalty},
                                       Q⁺, A⁺, n⁻, Q⁻, A⁻, t)
  Q⁺.u = -Q⁻.u

  return nothing
end

@inline function ocean_boundary_state!(::HBModel, ::CoastlineNoSlip,
                                       ::CentralNumericalFluxDiffusive, Q⁺,
                                       D⁺, A⁺, n⁻, Q⁻, D⁻, A⁻, t)
  Q⁺.u = -Q⁻.u

  D⁺.∇θ = -D⁻.∇θ

  return nothing
end

@inline function ocean_boundary_state!(m::HBModel, ::OceanFloorFreeSlip,
                                       ::Union{Rusanov, CentralGradPenalty},
                                       Q⁺, A⁺, n⁻, Q⁻, A⁻, t)
  A⁺.w = -A⁻.w

  return nothing
end

@inline function ocean_boundary_state!(m::HBModel, ::OceanFloorFreeSlip,
                                       ::CentralNumericalFluxDiffusive, Q⁺,
                                       D⁺, A⁺, n⁻, Q⁻, D⁻, A⁻, t)
  A⁺.w = -A⁻.w
  D⁺.∇u = -D⁻.∇u

  D⁺.∇θ = -D⁻.∇θ

  return nothing
end

@inline function ocean_boundary_state!(m::HBModel, ::OceanFloorNoSlip,
                                       ::Union{Rusanov, CentralGradPenalty},
                                       Q⁺, A⁺, n⁻, Q⁻, A⁻, t)
  Q⁺.u = -Q⁻.u
  A⁺.w = -A⁻.w

  return nothing
end

@inline function ocean_boundary_state!(m::HBModel, ::OceanFloorNoSlip,
                                       ::CentralNumericalFluxDiffusive, Q⁺,
                                       D⁺, A⁺, n⁻, Q⁻, D⁻, A⁻, t)

  Q⁺.u = -Q⁻.u
  A⁺.w = -A⁻.w

  D⁺.∇θ = -D⁻.∇θ

  return nothing
end

@inline function ocean_boundary_state!(m::HBModel, ::Union{
                                       OceanSurfaceNoStressNoForcing,
                                       OceanSurfaceStressNoForcing,
                                       OceanSurfaceNoStressForcing,
                                       OceanSurfaceStressForcing},
                                       ::Union{Rusanov, CentralGradPenalty},
                                       Q⁺, A⁺, n⁻, Q⁻, A⁻, t)
  return nothing
end

@inline function ocean_boundary_state!(m::HBModel,
                                       ::OceanSurfaceNoStressNoForcing,
                                       ::CentralNumericalFluxDiffusive,
                                       Q⁺, D⁺, A⁺, n⁻, Q⁻, D⁻, A⁻, t)
  D⁺.∇u = -D⁻.∇u

  D⁺.∇θ = -D⁻.∇θ

  return nothing
end

@inline function ocean_boundary_state!(m::HBModel,
                                       ::OceanSurfaceStressNoForcing,
                                       ::CentralNumericalFluxDiffusive,
                                       Q⁺, D⁺, A⁺, n⁻, Q⁻, D⁻, A⁻, t)
  τ = A⁻.τ
  D⁺.∇u = -D⁻.∇u + 2 * @SMatrix [ -0 -0;
                                  -0 -0;
                                  τ / 1000 -0]

  D⁺.∇θ = -D⁻.∇θ

  return nothing
end

@inline function ocean_boundary_state!(m::HBModel,
                                       ::OceanSurfaceNoStressForcing,
                                       ::CentralNumericalFluxDiffusive,
                                       Q⁺, D⁺, A⁺, n⁻, Q⁻, D⁻, A⁻, t)
  D⁺.∇u = -D⁻.∇u

  θ  = Q⁻.θ
  θʳ = A⁻.θʳ
  λʳ = m.problem.λʳ
  D⁺.∇θ = -D⁻.∇θ + 2 * @SVector [-0, -0, λʳ * (θʳ - θ)]

  return nothing
end

@inline function ocean_boundary_state!(m::HBModel,
                                       ::OceanSurfaceStressForcing,
                                       ::CentralNumericalFluxDiffusive,
                                       Q⁺, D⁺, A⁺, n⁻, Q⁻, D⁻, A⁻, t)
  τ = A⁻.τ
  D⁺.∇u = -D⁻.∇u + 2 * @SMatrix [ -0 -0;
                                  -0 -0;
                                  τ / 1000 -0]

  θ  = Q⁻.θ
  θʳ = A⁻.θʳ
  λʳ = m.problem.λʳ
  D⁺.∇θ = -D⁻.∇θ + 2 * @SVector [-0, -0, λʳ * (θʳ - θ)]

  return nothing
end

end


module CuNFFT

using Reexport
@reexport using AbstractNFFTs
using NFFT
using LinearAlgebra
using AbstractFFTs
using CUDA

if CUDA.functional()
  using CUDA.CUSPARSE
  include("implementation.jl")
end

end


using Parameters, AutomotiveDrivingModels, POMDPs

include("occluded_crosswalk_env.jl")
include("helpers.jl")
include("pomdp_types.jl")
include("constant_pedestrian.jl")
include("generative_model.jl")

rng = MersenneTwister(1)
pomdp = OCPOMDP()

s0 = initialstate(pomdp, rng)
ego = s0[findfirst(EGO_ID,s0)]
println("Initial scene with $(length(s0)) cars")
println("Initial ego state: $ego")


export moveispseudo,
    pseudocaptures, pseudochecks, pseudoislegal, pseudomoves, pseudoquiets, pseudoevasions

function is_attacked_after_move(b::Board, s::Square, m::Move)::Bool
    f = from(m)
    t = to(m)
    us = sidetomove(b)
    theirs = pieces(b, -us)

    # Remove captured piece, if any:
    if pieceon(b, t) != EMPTY
        theirs -= t
    elseif moveisep(b, m)
        theirs -= t - (us == WHITE ? DELTA_N : DELTA_S)
    end

    # Test for attacks by non-sliding pieces:
    if !isempty(pawnattacks(us, s) ∩ pawns(b) ∩ theirs)
        return true
    end
    if !isempty(knightattacks(s) ∩ knights(b) ∩ theirs)
        return true
    end
    if !isempty(kingattacks(s) ∩ kings(b) ∩ theirs)
        return true
    end

    # Test for attacks by sliding pieces:
    ours = pieces(b, us) - f + t
    occ = ours ∪ theirs

    if !isempty(bishopattacks(occ, s) ∩ bishoplike(b) ∩ theirs)
        return true
    end
    if !isempty(rookattacks(occ, s) ∩ rooklike(b) ∩ theirs)
        return true
    end

    # Square not attacked!
    return false
end

"""
    moveispseudo(b::Board, m::Move)

Test whether the move `m` is pseudo-legal.

This means that the move `m` is legal except for possibly leaving the king in
check.

# Examples
```julia-repl
julia> b = @startboard e4 c5 Nf3 d6 Bb5;

julia> moveispseudo(b, Move(SQ_G8, SQ_F6))
true

julia> moveispseudo(b, Move(SQ_B8, SQ_C6))
true

julia> moveispseudo(b, Move(SQ_G8, SQ_G6))
false
```
"""
function moveispseudo(b::Board, m::Move)::Bool
    f = from(m)
    t = to(m)
    us = sidetomove(b)
    them = -us

    # The piece on the source square must belong to the current side to move:
    pcolor(pieceon(b, f)) == us || return false

    # The piece on the destination square cannot belong to the current side to
    # move. TODO: Chess960 castling
    pcolor(pieceon(b, t)) ≠ us || return false

    p = pieceon(b, f)
    Δ = t - f

    if ptype(p) == PAWN
        # If the destination square is on the 1st or 8th rank, the move must be
        # a promotion:
        if t ∈ SS_RANK_1 ∪ SS_RANK_8 && !ispromotion(m)
            return false
        end

        # Is this a normal pawn push?
        push = us == WHITE ? DELTA_N : DELTA_S
        if Δ == push
            return pieceon(b, t) == EMPTY
        end

        # Is this a double pawn push?
        if Δ == 2 * push && f ∈ SS_RANK_2 ∪ SS_RANK_7
            return pieceon(b, t) == EMPTY && pieceon(b, f + push) == EMPTY
        end

        # Is this a capture?
        if t ∈ pawnattacks(us, f)
            return pcolor(pieceon(b, t)) == them || t == epsquare(b)
        end

        # Neither a pawn push nor a pawn capture -- cannot be a legal move:
        return false
    else # Non-pawn move
        # Non-pawn  moves cannot be promotions:
        ispromotion(m) && return false

        pt = ptype(p)

        if pt == KNIGHT
            return t ∈ knightattacks(f)
        elseif pt == BISHOP
            return t ∈ bishopattacks(b, f)
        elseif pt == ROOK
            return t ∈ rookattacks(b, f)
        elseif pt == QUEEN
            return t ∈ queenattacks(b, f)
        elseif pt == KING
            if t ∈ kingattacks(f)
                return true
            elseif Δ == 2 * DELTA_E
                # Kingside castling
                return cancastlekingside(b, us) &&
                       !ischeck(b) &&
                       pieceon(b, f + DELTA_E) == EMPTY &&
                       pieceon(b, t) == EMPTY
            elseif Δ == 2 * DELTA_W
                # Queenside castling
                return cancastlequeenside(b, us) &&
                       !ischeck(b) &&
                       pieceon(b, f + DELTA_W) == EMPTY &&
                       pieceon(b, t) == EMPTY &&
                       pieceon(b, f + 3 * DELTA_W) == EMPTY
            else
                return false
            end
        else
            return false
        end
    end
end

"""
    pseudoislegal(b::Board, m::Move)::Bool

Tests whether the pseudo-legal move `m` is legal.

# Examples
```julia-repl
julia> b = @startboard e4 c5 Nf3 d6 Bb5;

julia> pseudoislegal(b, Move(SQ_C8, SQ_D7))
true

julia> pseudoislegal(b, Move(SQ_G8, SQ_F6))
false
```
"""
function pseudoislegal(b::Board, m::Move)::Bool
    f = from(m)
    t = to(m)
    us = sidetomove(b)
    them = -us
    ksq = kingsquare(b, us)
    check = ischeck(b)

    # If we're not in check, any non-king, non-en-passant move of a non-pinned
    # piece is legal:
    if !check && f ∉ pinned(b) && f ≠ ksq && !moveisep(b, m)
        return true
    end

    # Castling moves are legal if we're not in check and both the destination
    # square and the square we're passing over are not under attack by the
    # opponent:
    if moveisshortcastle(b, m)
        return !check && !isattacked(b, f + DELTA_E, them) && !isattacked(b, t, them)
    end
    if moveislongcastle(b, m)
        return !check && !isattacked(b, f + DELTA_W, them) && !isattacked(b, t, them)
    end

    # For king moves, make sure the king is not under attack on the new square:
    if f == ksq
        return !is_attacked_after_move(b, t, m)
    end

    # When in check or when making an en passant capture, check in a slow way
    # that the move does not leave the king in check:
    if check || moveisep(b, m)
        return !is_attacked_after_move(b, ksq, m)
    end

    # If we are here, we are not in check, and we're moving a pinned piece. The
    # move is legal if it's moving along the ray towards/away from the friendly
    # king.
    return t ∈ squaresbetween(ksq, f) || f ∈ squaresbetween(ksq, t)
end

function addmove!(list::MoveList, f::Square, t::Square)
    push!(list, Move(f, t))
end

function addpromotions!(list::MoveList, f::Square, t::Square)
    for p ∈ [QUEEN, ROOK, BISHOP, KNIGHT]
        push!(list, Move(f, t, p))
    end
end

function addpawnmoves!(list::MoveList, target::SquareSet, Δ::SquareDelta)
    for t ∈ target
        addmove!(list, t - Δ, t)
    end
end

function addpawnpromotions!(list::MoveList, target::SquareSet, Δ::SquareDelta)
    for t ∈ target
        addpromotions!(list, t - Δ, t)
    end
end

function knight_pseudomoves(b::Board, list::MoveList, src::SquareSet, target::SquareSet)
    for f ∈ src
        for t ∈ knightattacks(f) ∩ target
            addmove!(list, f, t)
        end
    end
end

function bishop_pseudomoves(b::Board, list::MoveList, src::SquareSet, target::SquareSet)
    for f ∈ src
        for t ∈ bishopattacks(b, f) ∩ target
            addmove!(list, f, t)
        end
    end
end

function rook_pseudomoves(b::Board, list::MoveList, src::SquareSet, target::SquareSet)
    for f ∈ src
        for t ∈ rookattacks(b, f) ∩ target
            addmove!(list, f, t)
        end
    end
end

function queen_pseudomoves(b::Board, list::MoveList, src::SquareSet, target::SquareSet)
    for f ∈ src
        for t ∈ queenattacks(b, f) ∩ target
            addmove!(list, f, t)
        end
    end
end

function king_pseudomoves(b::Board, list::MoveList, src::SquareSet, target::SquareSet)
    for f ∈ src
        for t ∈ kingattacks(f) ∩ target
            addmove!(list, f, t)
        end
    end
end

function pawn_pseudocaptures(b::Board, list::MoveList)
    us = sidetomove(b)
    them = -us
    ps = pawns(b, us)

    if us == WHITE
        # Promotions (captures and non-captures)
        source1 = ps ∩ SS_RANK_7
        addpawnpromotions!(list, pawnshift_ne(source1) ∩ pieces(b, them), DELTA_NE)
        addpawnpromotions!(list, pawnshift_nw(source1) ∩ pieces(b, them), DELTA_NW)
        addpawnpromotions!(list, pawnshift_n(source1) ∩ emptysquares(b), DELTA_N)

        # Non-promotion captures
        source2 = ps ∩ -SS_RANK_7
        addpawnmoves!(list, pawnshift_ne(source2) ∩ pieces(b, them), DELTA_NE)
        addpawnmoves!(list, pawnshift_nw(source2) ∩ pieces(b, them), DELTA_NW)
    else
        # Promotions (captures and non-captures)
        source1 = ps ∩ SS_RANK_2
        addpawnpromotions!(list, pawnshift_se(source1) ∩ pieces(b, them), DELTA_SE)
        addpawnpromotions!(list, pawnshift_sw(source1) ∩ pieces(b, them), DELTA_SW)
        addpawnpromotions!(list, pawnshift_s(source1) ∩ emptysquares(b), DELTA_S)

        # Non-promotion captures
        source2 = ps ∩ -SS_RANK_2
        addpawnmoves!(list, pawnshift_se(source2) ∩ pieces(b, them), DELTA_SE)
        addpawnmoves!(list, pawnshift_sw(source2) ∩ pieces(b, them), DELTA_SW)
    end
end

function pawn_pseudopushes(b::Board, list::MoveList, target::SquareSet)
    us = sidetomove(b)
    ps = pawns(b, us)
    if us == WHITE
        source1 = ps ∩ -SS_RANK_7
        source2 = ps ∩ SS_RANK_2
        target1 = pawnshift_n(source1) ∩ emptysquares(b) ∩ target
        target2 =
            pawnshift_n(pawnshift_n(source2) ∩ emptysquares(b)) ∩ emptysquares(b) ∩ target
        addpawnmoves!(list, target1, DELTA_N)
        addpawnmoves!(list, target2, 2 * DELTA_N)
    else
        source1 = ps ∩ -SS_RANK_2
        source2 = ps ∩ SS_RANK_7
        target1 = pawnshift_s(source1) ∩ emptysquares(b) ∩ target
        target2 =
            pawnshift_s(pawnshift_s(source2) ∩ emptysquares(b)) ∩ emptysquares(b) ∩ target
        addpawnmoves!(list, target1, DELTA_S)
        addpawnmoves!(list, target2, 2 * DELTA_S)
    end
end

function pawn_promotion_pseudopushes(b::Board, list::MoveList, target::SquareSet)
    us = sidetomove(b)
    ps = pawns(b, us)
    if us == WHITE
        source = ps ∩ SS_RANK_7
        target = pawnshift_n(source) ∩ emptysquares(b) ∩ target
        addpawnpromotions!(list, target, DELTA_N)
    else
        source = ps ∩ SS_RANK_2
        target = pawnshift_s(source) ∩ emptysquares(b) ∩ target
        addpawnpromotions!(list, target, DELTA_N)
    end
end

function pseudo_ep_captures(b::Board, list::MoveList)
    us = sidetomove(b)
    if epsquare(b) != SQ_NONE
        for f ∈ pawnattacks(-us, epsquare(b)) ∩ pawns(b, us)
            addmove!(list, f, epsquare(b))
        end
    end
end

function pseudo_castles(b::Board, list::MoveList)
    us = sidetomove(b)
    ksq = kingsquare(b, us)
    if cancastlekingside(b, us)
        if isempty(squaresbetween(ksq, ksq + 3 * DELTA_E) ∩ occupiedsquares(b))
            addmove!(list, ksq, ksq + 2 * DELTA_E)
        end
    end
    if cancastlequeenside(b, us)
        if isempty(squaresbetween(ksq, ksq + 4 * DELTA_W) ∩ occupiedsquares(b))
            addmove!(list, ksq, ksq + 2 * DELTA_W)
        end
    end
end

function non_pawn_moves_to_target(b::Board, list::MoveList, target::SquareSet)
    us = sidetomove(b)
    knight_pseudomoves(b, list, knights(b, us) ∩ -pinned(b), target)
    bishop_pseudomoves(b, list, bishops(b, us), target)
    rook_pseudomoves(b, list, rooks(b, us), target)
    queen_pseudomoves(b, list, queens(b, us), target)
    king_pseudomoves(b, list, kings(b, us), target)
end

function pawn_discovered_checks(b::Board, list::MoveList, src::SquareSet)
    us = sidetomove(b)
    source = src ∩ pawns(b, us) ∩ -filesquares(file(kingsquare(b, -us)))
    push = us == WHITE ? DELTA_N : DELTA_S
    for f ∈ source
        if f + push ∉ (SS_RANK_1 ∪ SS_RANK_8) && pieceon(b, f + push) == EMPTY
            addmove!(list, f, f + push)
            if f ∈ SS_RANK_2 ∪ SS_RANK_7 && pieceon(b, f + 2 * push) == EMPTY
                addmove!(list, f, f + 2 * push)
            end
        end
    end
end

function pawn_plain_checks(b::Board, list::MoveList, src::SquareSet)
    us = sidetomove(b)
    ksq = kingsquare(b, -us)
    source = src ∩ pawns(b, us) ∩ adjacentfilesquares(ksq)
    push = us == WHITE ? DELTA_N : DELTA_S

    for f ∈ source
        if f + push ∉ (SS_RANK_1 ∪ SS_RANK_8) && pieceon(b, f + push) == EMPTY
            if ksq ∈ pawnattacks(us, f + push)
                addmove!(list, f, f + push)
            end
            if f ∈ (SS_RANK_2 ∪ SS_RANK_7) && pieceon(b, f + 2 * push) == EMPTY
                if ksq ∈ pawnattacks(us, f + 2 * push)
                    addmove!(list, f, f + 2 * push)
                end
            end
        end
    end
end

function discovered_checks(b::Board, list::MoveList, source::SquareSet)
    us = sidetomove(b)
    them = -us
    ksq = kingsquare(b, them)
    target = emptysquares(b)

    knight_pseudomoves(b, list, knights(b, us) ∩ source, target)
    bishop_pseudomoves(b, list, bishops(b, us) ∩ source, target)
    rook_pseudomoves(b, list, rooks(b, us) ∩ source, target)
    pawn_discovered_checks(b, list, source)
    for f ∈ kings(b, us) ∩ source
        for t ∈ kingattacks(f) ∩ target
            if t ∉ squaresbetween(f, ksq) && f ∉ squaresbetween(t, ksq)
                addmove!(list, f, t)
            end
        end
    end
end

function plain_checks(b::Board, list::MoveList, source::SquareSet)
    us = sidetomove(b)
    them = -us
    ksq = kingsquare(b, them)
    target = emptysquares(b)
    b_check_sqs = bishopattacks(b, ksq) ∩ target
    r_check_sqs = rookattacks(b, ksq) ∩ target
    q_check_sqs = b_check_sqs ∪ r_check_sqs

    knight_pseudomoves(b, list, knights(b, us) ∩ source, knightattacks(ksq) ∩ target)
    bishop_pseudomoves(b, list, bishops(b, us) ∩ source, b_check_sqs)
    rook_pseudomoves(b, list, rooks(b, us) ∩ source, r_check_sqs)
    queen_pseudomoves(b, list, queens(b, us) ∩ source, q_check_sqs)
    pawn_plain_checks(b, list, source)
end

"""
    pseudocaptures(b::Board, list::MoveList)

Generates all pseudo-legal captures and promotions.
"""
function pseudocaptures(b::Board, list::MoveList)
    target = pieces(b, -sidetomove(b))
    pawn_pseudocaptures(b, list)
    non_pawn_moves_to_target(b, list, target)
    pseudo_ep_captures(b, list)
end

"""
    pseudoquiets(b::Board, list::MoveList)

Generates all pseudo-legal quiet moves.

Captures and promotions will not be included.
"""
function pseudoquiets(b::Board, list::MoveList)
    target = emptysquares(b)
    pawn_pseudopushes(b, list, target)
    non_pawn_moves_to_target(b, list, target)
    pseudo_castles(b, list)
end

"""
    pseudomoves(b::Board, list::MoveList)

Generates all pseudo-legal moves.
"""
function pseudomoves(b::Board, list::MoveList)
    pseudocaptures(b, list)
    pseudoquiets(b, list)
end

"""
    pseudoevasions(b::Board, list::MoveList)

Generates pseudo-legal moves when in check.

Most obviously illegal moves are excluded.
"""
function pseudoevasions(b::Board, list::MoveList)
    us = sidetomove(b)

    # Generate all pseudo-legal moves for the king.
    king_pseudomoves(b, list, kings(b, us), -pieces(b, us))

    # Other moves are only possible if this is not a double check:
    if issingleton(b.checkers)
        ksq = kingsquare(b, us)
        chsq = first(b.checkers)
        target = squaresbetween(ksq, chsq) + chsq

        # Blocking pawn pushes:
        pawn_pseudopushes(b, list, target)
        pawn_promotion_pseudopushes(b, list, target)

        # Pawn captures of checking piece:
        for f ∈ pawnattacks(-us, chsq) ∩ pawns(b, us) ∩ -pinned(b)
            if chsq ∈ (SS_RANK_8 ∪ SS_RANK_1)
                addpromotions!(list, f, chsq)
            else
                addmove!(list, f, chsq)
            end
        end
        if chsq == epsquare(b) - (us == WHITE ? DELTA_N : DELTA_S)
            pseudo_ep_captures(b, list)
        end

        # Remaining pieces: Generate moves to target for non-pinned pieces.
        knight_pseudomoves(b, list, knights(b, us) ∩ -pinned(b), target)
        bishop_pseudomoves(b, list, bishops(b, us) ∩ -pinned(b), target)
        rook_pseudomoves(b, list, rooks(b, us) ∩ -pinned(b), target)
        queen_pseudomoves(b, list, queens(b, us) ∩ -pinned(b), target)
    end
end

"""
    pseudochecks(b::Board, list::MoveList)

Generates all pseudo-legal non-capturing checking moves.
"""
function pseudochecks(b::Board, list::MoveList)
    dc = discovered_check_candidates(b)
    discovered_checks(b, list, dc)
    plain_checks(b, list, pieces(b, sidetomove(b)) ∩ -dc)
end


using .JSort
import YAML

function yamlvalue(yaml::Dict{String, Any}, key::String, default)
    if key ∉ keys(yaml)
        default
    else
        param = yaml[key]
    end
end

struct Parameters
    root::String
    readpath::String
    savepath::String
    #datafiles::Vector{String}
    isotope::String
    gates::Dict{Symbol, Gate}
    calibrator::Dict{Symbol, AbstractCalibrator}
    percent::Float64
    badchannels::Set{Int}
    qkinz::Vector{Poly{Float32}}
end

function Parameters(path::AbstractString)
    parameters::Dict{String,Any} = YAML.load_file(path)
    readkeys = Set{String}()
    function get(key::String, default)
        push!(readkeys, key)
        yamlvalue(parameters, key, default)
    end
    root     = splitdir(path)[1]
    savepath = get("save path", root)
    if startswith(savepath, "~")
        savepath = expanduser(savepath)
    else
        savepath = if savepath == root savepath else joinpath(root, savepath) end
    end

    readpath = get("read path", savepath)
    if startswith(readpath, "~")
        readpath = expanduser(readpath)
    else
        readpath = if readpath == savepath readpath else joinpath(root, readpath) end
    end
    isotope  = get("isotope", "")

    # Different calibration coefficients
    calibrationpaths = Dict{Symbol, Union{String, Vector{String}}}(
        :e  => get("coefficients e",     "coefficients_e.csv"),
        :Δe => get("coefficients de",    "coefficients_de.csv"),
        :γ  => get("coefficients gamma", "coefficients_g.csv"),
#        :γi => get("coefficients gamma inverse", "coefficients_g_i.csv"),
        :t  => get("coefficients time",  "coefficients_t.csv"),
        :tc => get("coefficients time correction", "coefficients_t_c.csv")
    )
    # Try to read from the coefficient paths, defaulting to identity polynomials
    calibrator = Dict{Symbol, AbstractCalibrator}(
        :e  => make(ParticleCalibrator, calibrationpaths[:e],  :e),
        :Δe => make(ParticleCalibrator, calibrationpaths[:Δe], :Δe),
        :γ  => make(GammaCalibrator,    calibrationpaths[:γ],  :eᵧ),
#        :γi => make(GammaCalibrator,    calibrationpaths[:γi], :eᵧ),
        :t  => make(TimeCalibrator,     calibrationpaths[:t],  :t),
        :tc => make(TimeCorrector,      calibrationpaths[:tc], :t)
    )

    gates = getgates(get("gates", nothing))

    # Read the datafile paths
    # datafiles = get("data files", nothing)
    # if isnothing(datafiles)
    #     @warn "No data files will be processed!"
    # else
    #     map!(x -> joinpath(readpath, x), datafiles, datafiles)
    # end

    percent = get("percent to read", 100)

    # Qkinz
    qkinz = get("qkinz", nothing) |> abspath |> parseqkinz

    # Bad channels
    badchannels = Set{Int}(i for i in get("channels to ignore", Set{Int}()))

    atleast1d(x::AbstractArray) = x
    atleast1d(x) = [x]
    for key in keys(parameters)
        if key ∉ readkeys
            @warn "The parameter $key was supplied but not used."
        end
        if occursin("coefficient", key) || occursin("file", key)
            paths = get(key, "") |> atleast1d
            for path in paths
              if !isfile(path)
                  @warn "Could not open $key"
              end
            end
        end
    end

    Parameters(path, readpath, savepath, isotope,
               gates, calibrator,
               percent, badchannels, qkinz)
end

function make(::Type{T}, path::AbstractString, var)::T where T<:AbstractCalibrator
    if isfile(path)
        read(path, T, var=var)
    else
        T(var=var)
    end
end

function make(::Type{T}, paths::AbstractArray{<:AbstractString}, var)::T where T<:AbstractCalibrator
    base = make(T, paths[1], var)
    if length(paths) > 1
      for path in paths[2:end]
          calib = make(T, path, var)
          combine!(base, calib)
      end
    end
    base
end

### PARSE GATES ###
getgates(::Nothing) = Dict(:ex => Gate(), :t => Gate())
function getgates(unparsedgates::AbstractDict)
    # Gates on the excitation energy
    ukeys = lowercase.(keys(unparsedgates))
    ekey =  ukeys ∩ ["ex", "excitation", "ex_int"]
    promptkey = ukeys ∩ ["prompt"]
    bgkey = ukeys ∩ ["bg", "background"]

    gates = Dict{Symbol,Gate}(:ex => Gate(), :background => Gate(),
                              :promt => Gate())

    if length(ekey) == 1
        for gate in unparsedgates[ekey[1]]
            gates[:ex] = gates[:ex] ∩ parsegate(gate)
        end
    end

    if length(promptkey) == 1
        gates[:prompt] = unparsedgates[promptkey[1]] |> parsegate
    end

    if length(bgkey) == 1
        gates[:background] = unparsedgates[bgkey[1]] |> parsegate
    end

    if length(ukeys) > 0 && length(ekey ∪ promptkey ∪ bgkey) == 0
        throw(IOError("Gates $ukeys not supported."))
    end
    gates
end

function parsegate(gate::AbstractDict)::Gate
    ukeys = gate |> keys .|> lowercase
    low = ukeys ∩ ["low", "l"]
    high = ukeys ∩ ["high", "h"]
    wrong = setdiff(ukeys, low ∪ high)
    if length(wrong) > 0
        throw(ArgumentError("Unsupported gate bounds: $wrong"))
    end

    if length(low) > 0 && length(high) > 0
        return Gate(low=gate[low[1]], high=gate[high[1]])
    elseif length(low) > 0
        return Gate(low=gate[low[1]])
    else
        return Gate(high=gate[high[1]])
    end
end

function parsegate(gate::AbstractArray)::Gate
    if length(gate) ≠ 2
        throw(ArgumentError("Prove an upper and lower limit for gate"))
    end
    return Gate(low=gate[1], high=gate[2])
end

## Parse Qkinz
function parseqkinz(path::AbstractString)::Vector{Poly}
    dir = dirname(path)
    pattern = Regex(basename(path))
    polys = Tuple{Int, Poly{Float64}}[]
    j(x) = joinpath(dir, x)
    for fname in readdir(dir)
        m = match(pattern, fname)
        if m ≢ nothing
            f = parse(Int, m[:f])
            push!(polys, (f, JSort.BetheBloch.Banana(j(fname)).exfromede))
        end
    end
    (f, poly) = polys |> sort |> x -> zip(x...)
    @assert length(f) == 8 "Missing Qkinz strips. Got $f"
    [p for p in poly]
end


module Aheui

export 아희
export 스택, 큐, 통로, 순방향, 역방향
export 없음, 끝냄
export 덧셈, 곱셈, 뺄셈, 나눗셈, 나머지
export 집어넣기, 뽑기, 중복, 바꿔치기
export 선택, 이동, 비교, 조건
export 가자, 종성획수

abstract type 저장공간 end
abstract type 커서타입 end

mutable struct 진행방향
  어느쪽::Function
end

mutable struct 스택 <: 저장공간
  공간::Vector
  출력::Vector
  방향::진행방향
end
struct 큐 <: 저장공간
  공간::Vector
  출력::Vector
  방향::진행방향
end
struct 통로 <: 저장공간
  공간::Vector
  출력::Vector
  방향::진행방향
end

순방향 = 진행방향(+)
역방향 = 진행방향(-)

없음(저장::저장공간, _) = nothing
끝냄(저장::저장공간, _) = isempty(저장.공간) ? 0 : 뽑(저장)
덧셈(저장::저장공간, _) = 집어넣기(저장, +(뽑(저장), 뽑(저장)))
곱셈(저장::저장공간, _) = 집어넣기(저장, *(뽑(저장), 뽑(저장)))
뺄셈(저장::저장공간, _) = 집어넣기(저장, -뽑(저장)+뽑(저장))
나눗셈(저장::저장공간, _) = 집어넣기(저장, \(뽑(저장), 뽑(저장)))
나머지(저장::저장공간, _) = 집어넣기(저장, ((a,b)->%(b,a))(뽑(저장), 뽑(저장)))

뽑(저장::저장공간) = pop!(저장.공간)

function 뽑기(저장::저장공간, 받침::Symbol)
  if isempty(저장.공간)
    nothing
  else
    값 = 뽑(저장)
    출력(저장, :ㅎ==받침 ? Char(값) : 값)
  end
end
뽑기(저장::저장공간, ::Union{Int,Float64}) = isempty(저장.공간) ? nothing : 출력(저장, 뽑(저장))
뽑기(저장::저장공간) = 뽑기(저장, 0)

function 출력(저장::저장공간, 값)
  push!(저장.출력, 값)
  값
end

집어넣기(저장::저장공간, ::Symbol) = nothing
집어넣기(저장::저장공간, 값::Union{Int,Float64}) = push!(저장.공간, 값)

중복(저장::저장공간, _) = 집어넣기(저장, last(저장.공간))

바꿔치기(저장::저장공간, _) = 저장.공간[end], 저장.공간[end-1] = 저장.공간[end-1], 저장.공간[end]
선택(공간들::Dict, 공간키) = 공간들[공간키]
function 이동(저장::저장공간, 공간들::Dict, 공간키)
  공간 = 선택(공간들, 공간키)
  집어넣기(공간, 뽑기(저장))
  저장
end
function 비교(저장::저장공간)
  값 = 뽑기(저장) >= 뽑기(저장) ? 1 : 0
  집어넣기(저장, 값)
end
function 조건(저장::저장공간)
  저장.방향 = 0 == 뽑기(저장) ? 역방향 : 순방향
  저장
end
가자(f, a...) = f(a...)

struct 그냥내비둠 <: 커서타입
end
struct 칸이동 <: 커서타입
  얼만큼::Int
  어느쪽::Function
end

mutable struct 커서움직임
  이전위치
  위치
  방향::Function
end

내비둠 = 그냥내비둠()
한칸 = 칸이동(1, identity)
두칸 = 칸이동(2, identity)
import Base: +, -
+(칸::칸이동) = 칸이동(칸.얼만큼, +)
-(칸::칸이동) = 칸이동(칸.얼만큼, -)

function 커서이동(상하::칸이동, 좌우::그냥내비둠)
  function 어디로(저장::저장공간, 커서::커서움직임)
    행값,렬값 = 커서.위치
    커서.이전위치 = [행값 렬값]
    커서.위치 = [상하.어느쪽(행값, 상하.얼만큼) 렬값]
    저장.방향.어느쪽 = 상하.어느쪽
    커서
  end
  어디로
end
function 커서이동(상하::그냥내비둠, 좌우::칸이동)
  function 어디로(저장::저장공간, 커서::커서움직임)
    행값,렬값 = 커서.위치
    커서.이전위치 = [행값 렬값]
    커서.위치 = [행값 좌우.어느쪽(렬값, 좌우.얼만큼)]
    저장.방향.어느쪽 = 좌우.어느쪽
    커서
  end
  어디로
end
기능없음(저장::저장공간, 커서::커서움직임) = 커서

const 닿소리표 = [없음 없음 나눗셈 덧셈 곱셈 나머지 뽑기 집어넣기 중복 선택 이동 없음 비교 없음 조건 없음 뺄셈 바꿔치기 끝냄]
const 홀소리표 = [
  커서이동(내비둠, +한칸)
  기능없음
  커서이동(내비둠, +두칸)
  기능없음
  커서이동(내비둠, -한칸)
  기능없음
  커서이동(내비둠, -두칸)
  기능없음
  커서이동(-한칸, 내비둠)
  기능없음
  기능없음
  기능없음
  커서이동(-두칸, 내비둠)
  커서이동(+한칸, 내비둠)
  기능없음
  기능없음
  기능없음
  커서이동(+두칸, 내비둠)
  기능없음 # ㅡ  FIXME: 방향 바꾸기 처리 추가
  기능없음 # ㅢ
  기능없음 # ㅣ
]
const 종성획수표 = [0 2 4 4 2 5 5 3 5 7 9 9 7 9 9 8 4 4 6 2 4 :ㅇ 3 4 3 4 4 :ㅎ]
const 유니코드_가 = 0xAC00
const 유니코드_히흫 = 0xD7A3
const 초성오프셋 = 21 * 28
const 중성오프셋 = 28

function 기능(저장::저장공간, 소리::Char, 커서::커서움직임)
  값 = Int(소리)
  if 값 >= 유니코드_가 && 값 <= 유니코드_히흫
    값 = 값 - 유니코드_가
    초성인덱스 = trunc(Int, 값 / 초성오프셋) + 1
    값 = 값 % 초성오프셋
    중성인덱스 = trunc(Int, 값 / 중성오프셋) + 1
    종성인덱스 = trunc(Int, 값 % 중성오프셋) + 1
    닿소리표[초성인덱스], 홀소리표[중성인덱스], 종성획수표[종성인덱스]
  else
  end
end

function 아희(입력::String)
  여러줄 = split(입력, "\n")
  행,렬 = length(여러줄), maximum(map(length, 여러줄))
  격자 = Array{Char}(undef, 행, 렬)
  for (행값,줄) in enumerate(여러줄)
    for (렬값,글자) in enumerate(줄)
      격자[행값,렬값] = 글자
    end
  end
  아희(격자)
end

function 아희(입력::Array{Char,2})
  isempty(입력) && return ""
  공간들 = Dict(:큐=>큐([], [], 순방향), :스택=>스택([], [], 순방향), :통로=>통로([], [], 순방향))
  저장 = 가자(선택, 공간들, :스택)
  커서 = 커서움직임([0 0], [1 1], 기능없음)
  while true
    소리 = 입력[커서.위치...]
    초,중,종 = 기능(저장, 소리, 커서)
    if 끝냄==초
      break
    else
      가자(초, 저장, 종)
      커서 = 중(저장, 커서)
    end
  end
  join(저장.출력)
end

end


#///////////////////////////////////////
#// File Name: crowd_nav_controller.jl
#// Author: Haruki Nishimura (hnishimura@stanford.edu)
#// Date Created: 2021/04/02
#// Description: Controller based on CrowdNav RL Policies
#///////////////////////////////////////

struct CrowdNavControlParameter <: Parameter
    model_dir::String # Directory of the trained policy
    env_config::String # environment config file name
    policy_config::String # policy config file name
    policy_name::String # policy name
    goal_pos::Vector{Float64} # [x, y] goal position
    tcalc::Float64 # Allocated control computation time
    dtr::Float64 # Replanning time interval
    target_speed::Float64 # target speed for robot
end

mutable struct CrowdNavController
    sim_param::SimulationParameter
    cnt_param::CrowdNavControlParameter

    rl_robot::PyObject

    tcalc_actual::Union{Nothing, Float64}
    u_value::Union{Nothing, Vector{Float64}}

    control_update_task::Union{Nothing, Task}

    tcalc_actual_tmp::Union{Nothing, Float64}
    u_value_tmp::Union{Nothing, Vector{Float64}}
end

function CrowdNavController(sim_param::SimulationParameter,
                            cnt_param::CrowdNavControlParameter)
    rl_robot = py"configure_rl_robot"(cnt_param.model_dir, cnt_param.env_config,
                                      cnt_param.policy_config, cnt_param.policy_name);
    @assert rl_robot.visible == false "RL Robot has to be invisible."
    @assert rl_robot.time_step == sim_param.dto "RL Robot has to use the same timestep as sim_param.dto."
    rl_robot.v_pref = cnt_param.target_speed;
    rl_robot.gx, rl_robot.gy = cnt_param.goal_pos[1], cnt_param.goal_pos[2];
    rl_robot.theta = 0.0
    return CrowdNavController(sim_param, cnt_param, rl_robot, nothing, nothing,
                              nothing, nothing, nothing);
end

function get_action!(controller::CrowdNavController,
                     e_init::RobotState,
                     ado_state_dict::Dict{T, Vector{Float64}} where T <: Union{PyObject, String})
    # update rl_robot state
    controller.rl_robot.set_position(copy(get_position(e_init)));
    controller.rl_robot.set_velocity(copy(get_velocity(e_init)));

    # ado_state_dict has to have velocity information
    ado_state_array = [];
    radius::Float64 = controller.rl_robot.radius
    if length(ado_state_dict) < 2 # crowdnav seems to need >1 ado agent
        vec = controller.cnt_param.goal_pos - get_position(e_init);
        action = vec./norm(vec).*controller.rl_robot.v_pref;
        return action
    else
        for key in keys(ado_state_dict)
            @assert length(ado_state_dict[key]) >= 4 "ado_state_dict has to have velocity information."
            px, py = ado_state_dict[key][1], ado_state_dict[key][2];
            vx, vy = ado_state_dict[key][3], ado_state_dict[key][4];
            observable_state = py"ObservableState"(px, py, vx, vy, radius);
            push!(ado_state_array, observable_state)
        end
        action = controller.rl_robot.act(ado_state_array);
        return action
    end
end

@inline function crowdnav_control_update!(controller::CrowdNavController,
                                          e_init::RobotState,
                                          ado_state_dict::Dict{T, Vector{Float64}} where T <: Union{PyObject, String})
    tcalc_actual =
        @elapsed u_vel = get_action!(controller, e_init, ado_state_dict);
    tcalc_actual +=
        @elapsed u_acc = (u_vel .- get_velocity(e_init)) ./ controller.sim_param.dto;
        #@elapsed u_acc = 2 * (u_vel .- get_velocity(e_init)) ./ controller.sim_param.dto;

    if tcalc_actual >= controller.cnt_param.tcalc
        # tcalc_actual has exceeded allowable computation time
        time = @sprintf "Time %.2f" round(to_sec(e_init.t), digits=5)
        @warn "$(time) [sec]: CrowdNav computation took $(round(tcalc_actual, digits=3)) [sec], which exceeds the maximum computation time allowed."
    end

    return tcalc_actual, u_acc
end

function control!(controller::CrowdNavController,
                  current_time::Time;
                  log::Union{Nothing, Vector{Tuple{Time, String}}}=nothing)

    if !isnothing(controller.control_update_task) &&
        istaskdone(controller.control_update_task)
        if !isnothing(log)
            msg = "New CrowdNav control is available to the controller."
            push!(log, (current_time, msg))
        end
        controller.tcalc_actual = controller.tcalc_actual_tmp;
        controller.u_value = copy(controller.u_value_tmp);
        controller.control_update_task = nothing;
    end
    u = copy(controller.u_value);
    if !isnothing(log)
        msg = "Control: $(u) is applied to the system."
        push!(log, (current_time, msg))
    end
    return u
end

function schedule_control_update!(controller::CrowdNavController,
                                  e_init::RobotState,
                                  ado_state_dict::Dict{T, Vector{Float64}} where T <: Union{PyObject, String};
                                  log::Union{Nothing, Vector{Tuple{Time, String}}}=nothing)
    if !isnothing(log)
        msg = "New CrowdNav control computation is scheduled."
        push!(log, (e_init.t, msg))
    end

    controller.control_update_task = @task begin
        controller.tcalc_actual_tmp, controller.u_value_tmp =
            crowdnav_control_update!(controller, e_init, ado_state_dict);
    end
    schedule(controller.control_update_task);
end


#
include("objects.jl")
import Base.show

mutable struct Collector
    health
    player  # who owns this unit
    pos::WorldPos
    Collector(player) = new(max_health(Collector), player, WorldPos(0,0))
    Collector(player, pos) = new(max_health(Collector), player, pos)
end
mutable struct Fighter
    health
    player  # who owns this unit
    pos::WorldPos
    attackTargetUnit  # who currently attacking
    Fighter(player) = new(max_health(Fighter), player, WorldPos(0,0), nothing)
    Fighter(player, pos) = new(max_health(Fighter), player, pos, nothing)
end

Base.show(io::IO, u::Collector) = print(io, "Collector($(u.health), $(Ptr{PlayerUnits}(pointer_from_objref(u.player))), $(u.pos)")
health_percent(u) = health(u) / max_health(typeof(u))

health(x::Collector) = x.health
set_health!(x::Collector, h) = (x.health = h)

health(x::Fighter) = x.health
set_health!(x::Fighter, h) = (x.health = h)

# Unit attributes
@noinline max_health(::Type{Collector}) = (global kMaxHealth_Collector; kMaxHealth_Collector)
@noinline max_health(::Type{Fighter}) = (global kMaxHealth_Collector; kMaxHealth_Fighter)

@noinline attack_damage(::Collector) = (global kMaxHealth_Collector; kAttackDamage_Collector)
@noinline attack_damage(::Fighter) = (global kMaxHealth_Collector; kAttackDamage_Fighter)

@noinline build_time(::Type{Collector}) = (global kMaxHealth_Collector; kBuildTime_Collector)
@noinline build_time(::Type{Fighter}) = (global kMaxHealth_Collector; kBuildTime_Fighter)
@noinline build_cost(::Type{Collector}) = (global kMaxHealth_Collector; kBuildCost_Collector)
@noinline build_cost(::Type{Fighter}) = (global kMaxHealth_Collector; kBuildCost_Fighter)

unit_types = (Collector, Fighter)
UnitTypes = Union{unit_types...}

# ---------------------
# PLAYER
# --------------------
mutable struct PlayerUnits
    units::Array
    collectors::Array{Collector}
    fighters::Array{Fighter}
end
PlayerUnits() = PlayerUnits([],[],[])
Base.show(io::IO, pu::PlayerUnits) = print(io,
                     "PlayerUnits(\n"
                    *"            units[$(length(pu.units))]\n"
                    *"            collectors[$(length(pu.collectors))]\n"
                    *"            fighters[$(length(pu.fighters))]\n"
                    *"           )")

function add_unit_helper(p::PlayerUnits, u)
    @assert u.player == p
    push!(p.units, u)
    return u
end
function add_unit!(p::PlayerUnits, u)
    return add_unit_helper(p,u)
end
function add_unit!(p::PlayerUnits, u::Collector)
    add_unit_helper(p,u)
    push!(p.collectors, u)
    return u
end
function add_unit!(p::PlayerUnits, u::Fighter)
    add_unit_helper(p,u)
    push!(p.fighters, u)
    return u
end
add_collector!(p::PlayerUnits) = add_unit!(p, Collector(p))
add_fighter!(p::PlayerUnits) = add_unit!(p, Fighter(p))

remove_unit!(p, unit) = filter!(u->u≠unit, p.units)
function remove_unit!(p, unit::Collector)
    filter!(u->u≠unit, p.collectors)
    filter!(u->u≠unit, p.units)
end
function remove_unit!(p, unit::Fighter)
    filter!(u->u≠unit, p.fighters)
    filter!(u->u≠unit, p.units)
end

function clear_units!(p)
    empty!(p.units)
    empty!(p.collectors)
    empty!(p.fighters)
end
function clear_collectors!(p)
    filter!(u->!isa(u, Collector), p.units)
    empty!(p.collectors)
end
function clear_fighters!(p)
    filter!(u->!isa(u, Fighter), p.units)
    empty!(p.fighters)
end

owning_player(x)::PlayerUnits = x.player

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

function attack!(target, attacker)
    set_health!(target, health(target) - attack_damage(attacker))
    if (health(target) <= 0)
        destroy_unit!(target)
    end
    return health(target)
end

function destroy_unit!(target)
    remove_unit!(owning_player(target), target)
end

# Tests
using Test
collectors = PlayerUnits()
add_collector!(collectors)
destroy_unit!(collectors.units[1])
@test 0 == length(collectors.units)
w = add_collector!(collectors)
w = add_collector!(collectors)
@test 2 == length(collectors.units)
attack!(w, Fighter(nothing))
attack!(w, Fighter(nothing))
@test 1 == length(collectors.units)


#heatmap plots for grouped and ungrouped genotype and read depth viz

"""
    genotype_heatmap2_new_legend(input::Array{Any,2},title::AbstractString,chrom_label_info,sample_names,chr_pos_tuple_list_rev,y_axis_label_option,x_axis_label_option)
generate heatmap of genotype data.
"""
function genotype_heatmap2_new_legend(input,title,chrom_label_info,sample_names,chr_pos_tuple_list_rev,y_axis_label_option,x_axis_label_option) #chr_pos_tuple_list_rev is rev because heatmap in plotly mirrors list for some reason.

            chr_pos_tuple_indices,chr_pos_tuple_list,sample_name_indices,sample_names,chrom_labels,chrom_label_indices,font_size = process_plot_inputs(chrom_label_info,sample_names,chr_pos_tuple_list_rev)
            title_no_underscores=replace(title, "_"=>' ')

             hover_text_array=generate_hover_text_array(chr_pos_tuple_list,sample_names,input,"GT")

             position_1,position_2,position_3,position_4,position_5,position_6,position_7,position_8,position_9,position_10,text_point_increment,legend_xaxis_anchor,legend_xaxis_anchor_2,x_axis_text_anchor,extra_right_margin_space,extra_right_margin_space_dp = generate_legend_increments_ungrouped(input)

               trace=heatmap(
                   z = input, x = 1:size(input, 2),y = 1:size(input, 1),

                   hoverinfo="text",
                   text=hover_text_array,

                   zauto=false,zmax=3,zmin=0,

                   transpose=true,
                   colorscale =    [[0, "rgb(255,255,255)"],
                                   [0.33, "rgb(51,106,145)"],
                                   [0.66, "rgb(65,165,137)"],
                                   [1, "rgb(251,231,65)"]],

                   gridcolor = "#E2E2E2",
                   showscale = false, #we define a custom legend using shapes and a text scatter trace
                   colorbar = attr(tickvals = [0, 1, 2, 3],
                   ticktext = ["No Call (0)", "Homozygous Reference", "Heterozygous Variant", "Homozygous Variant"])
                   );

                   #we define a custom legend using shapes and a text scatter trace #x_axis_text_anchor+extra_right_margin_space
                   trace2 = scatter(
                   ;x=[x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor+extra_right_margin_space], y=y=[position_1+text_point_increment,position_3+text_point_increment,position_5+text_point_increment,position_7+text_point_increment,position_9+text_point_increment],
                     mode="text", name="legend label",
                     textposition="middle right",
                     text=["No Call", "Homozygous Reference", "Heterozygous Variant", "Homozygous Variant",""],
                     marker_size=300, textfont_family="Raleway, sans-serif")

                   shapes = [
                   rect(x0=legend_xaxis_anchor, y0=position_7, x1=legend_xaxis_anchor_2, y1=position_8, opacity=1.0, fillcolor="rgb(251,231,65)", line_color="black"),
                   rect(x0=legend_xaxis_anchor, y0=position_5, x1=legend_xaxis_anchor_2, y1=position_6, opacity=1.0, fillcolor="rgb(65,165,137)", line_color="black"),
                   rect(x0=legend_xaxis_anchor, y0= position_3, x1=legend_xaxis_anchor_2, y1=position_4, opacity=1.0, fillcolor="rgb(51,106,145)", line_color="black"),
                   rect(x0=legend_xaxis_anchor, y0= position_1, x1=legend_xaxis_anchor_2, y1=position_2, opacity=1.0, fillcolor="rgb(255,255,255)", line_color="black")]

                   data = [trace, trace2]

                if y_axis_label_option == "chromosomes"

                   layout = Layout(
                                   shapes=shapes,

                                   title = "$title_no_underscores",
                                   xaxis=attr(title="Sample ID", showgrid=false, zeroline=false, tickvals=sample_name_indices,
                                   ticktext=sample_names, tickangle=45,showticklabels=x_axis_label_option),
                                   yaxis=attr(title="Genomic Location", zeroline=false, tickvals=chrom_label_indices,
                                   ticktext=chrom_labels,tickfont_size=font_size,hovermode=true,automargin=true,showgrid=false)
                   )

                      plot(data,layout)

               elseif y_axis_label_option == "positions"

                   layout = Layout(shapes=shapes,
                                   title = "$title_no_underscores",
                                   xaxis=attr(title="Sample ID", showgrid=false, zeroline=false, tickvals=sample_name_indices,
                                   ticktext=sample_names, tickangle=45, showticklabels=x_axis_label_option),
                                   yaxis=attr(title="Genomic Location", zeroline=false, tickvals=chr_pos_tuple_indices,
                                   ticktext=chr_pos_tuple_list,tickfont_size=font_size,hovermode=true,automargin=true,showgrid=false)
                   )

                   plot(data,layout)

               elseif y_axis_label_option == "hover_positions"

                   layout = Layout(shapes=shapes,
                                   title = "$title_no_underscores",
                                   xaxis=attr(title="Sample ID", showgrid=false, zeroline=false, tickvals=sample_name_indices,
                                   ticktext=sample_names, tickangle=45,showticklabels=x_axis_label_option),
                                   yaxis=attr(title="Genomic Location", zeroline=false, tickvals=chr_pos_tuple_indices,
                                   ticktext=chr_pos_tuple_list,tickfont_size=font_size,hovermode=true,automargin=true,showticklabels=false)
                   )

                   plot(data,layout)

                else
                    println("--y_axis_labels is not a valid option. Choose positions or chromosomes")
                end
end

"""
   genotype_heatmap_with_groups(input::Array{Int64,2},title::String,chrom_label_info::Tuple{Array{String,1},Array{Int64,1},String},group_label_pack::Array{Any,1},id_list,chr_pos_tuple_list_rev,y_axis_label_option,trait_label_array,x_axis_label_option,number_rows)
generate heatmap of genotype data.
"""
function genotype_heatmap_with_groups(input::Array{Int64,2},title::String,chrom_label_info::Tuple{Array{String,1},Array{Int64,1},String},group_label_pack::Array{Any,1},id_list,chr_pos_tuple_list_rev,y_axis_label_option,trait_label_array,x_axis_label_option,number_rows)

    sample_name_indices,id_list,chrom_labels,chrom_label_indices,font_size,group_dividing_line,group1_label,group2_label,chr_pos_tuple_indices,chr_pos_tuple_list,font_size = process_plot_inputs_for_grouped_data(chrom_label_info::Tuple{Array{String,1},Array{Int64,1},String},group_label_pack::Array{Any,1},id_list,chr_pos_tuple_list_rev,trait_label_array)

    position_1,position_2,position_3,position_4,position_5,position_6,position_7,position_8,position_9,text_point_increment,trait_positions,legend_xaxis_anchor,legend_xaxis_anchor_2,x_axis_text_anchor,extra_right_margin_space = generate_legend_increments_grouped(input)

    position_1_trait_1,position_2_trait_1,position_1_trait_2,position_2_trait_2 = trait_positions

    title_no_underscores=replace(title, "_"=>' ')

    hover_text_array=generate_hover_text_array_grouped(chr_pos_tuple_list,id_list,input,"GT",number_rows)

                  trace = heatmap(
                       z = input, x=1:size(input, 2),y=1:size(input, 1),

                       hoverinfo="text",
                       text=hover_text_array,

                       zauto=false,zmax=3,zmin=-2,

                       transpose=true,

                                    colorscale = [
                                                  [0, "rgb(208, 211, 212)"], #light grey
                                                  [0.2, "rgb(151, 154, 154)"], #dark grey
                                                  [0.4, "rgb(255,255,255)"], #white
                                                  #[0.4, "rgb(56,25,90)"], #dark blue
                                                  [0.6, "rgb(51,106,145)"], #blue
                                                  [0.8, "rgb(65,165,137)"], #green
                                                  [1, "rgb(251,231,65)"] #yellow
                                                  ],

                      gridcolor = "#E2E2E2",
                      showscale = false, #we define a custom legend using shapes and a text scatter trace
                      colorbar = attr(tickvals = [-2, -1, 0, 1, 2, 3],
                      ticktext = ["Trait 2", "Trait 1", "No Call", "Homozygous Reference", "Heterozygous Variant", "Homozygous Variant"])
                      );

                    #we define a custom legend using shapes and a text scatter trace

                    trace2 = scatter(
                    ;x=[x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor+extra_right_margin_space,x_axis_text_anchor,x_axis_text_anchor], y=y=[position_1+text_point_increment,position_2+text_point_increment,position_3+text_point_increment,position_4+text_point_increment,position_5+text_point_increment,position_7+text_point_increment,position_8+text_point_increment],
                      mode="text", name="legend label",
                      textposition="middle right",
                      text=["No Call", "Homozygous Reference", "Heterozygous Variant", "Homozygous Variant"," ", "Trait 2","Trait 1"],
                      marker_size=300, textfont_family="Raleway, sans-serif")

                    shapes = [
                    rect(x0=legend_xaxis_anchor, y0=position_4, x1=legend_xaxis_anchor_2, y1=position_5, opacity=1.0, fillcolor="rgb(251,231,65)", line_color="black"),
                    rect(x0=legend_xaxis_anchor, y0=position_3, x1=legend_xaxis_anchor_2, y1=position_4, opacity=1.0, fillcolor="rgb(65,165,137)", line_color="black"),
                    rect(x0=legend_xaxis_anchor, y0= position_2, x1=legend_xaxis_anchor_2, y1=position_3, opacity=1.0, fillcolor="rgb(51,106,145)", line_color="black"),
                    rect(x0=legend_xaxis_anchor, y0= position_1, x1=legend_xaxis_anchor_2, y1=position_2, opacity=1.0, fillcolor="rgb(255,255,255)", line_color="black"),
                    rect(x0=legend_xaxis_anchor, y0= position_7, x1=legend_xaxis_anchor_2, y1=position_8, opacity=1.0, fillcolor="rgb(208, 211, 212)", line_color="black"),
                    rect(x0=legend_xaxis_anchor, y0= position_8, x1=legend_xaxis_anchor_2, y1=position_9, opacity=1.0, fillcolor="rgb(151, 154, 154)", line_color="black")
                    ]

                     data = [trace, trace2]

      if y_axis_label_option == "chromosomes"

          layout = Layout(
                          title = "$title_no_underscores",

                          xaxis=attr(
                          title="Sample ID (Grouped by $(group1_label) | $(group2_label))",
                          showgrid=false,
                          zeroline=false,
                          tickvals=sample_name_indices,
                          ticktext=id_list,
                          tickangle=45,
                          showticklabels=x_axis_label_option),

                          yaxis=attr(
                          title="Genomic Location",
                          zeroline=false,
                          tickvals=chrom_label_indices,
                          ticktext=chrom_labels,
                          tickfont_size=font_size,
                          hovermode=true,
                          automargin=true,showgrid=false),

                          shapes=shapes,
                          yaxis_range = [1:size(input,1)],
                          xaxis_range = [1:size(input,2)]
                        )

                        plot(data,layout)

      elseif y_axis_label_option == "positions"

          layout = Layout(
                          title = "$title_no_underscores",

                          xaxis=attr(title="Sample ID (Grouped by $(group1_label) | $(group2_label))", showgrid=false, zeroline=false, tickvals=sample_name_indices,
                          ticktext=id_list, tickangle=45,showticklabels=x_axis_label_option),

                          yaxis=attr(title="Genomic Location", zeroline=false, tickvals=chr_pos_tuple_indices,
                          ticktext=chr_pos_tuple_list,tickfont_size=font_size,hovermode=true,automargin=true,showgrid=false),

                          shapes=shapes, yaxis_range = [1:size(input,1)], xaxis_range = [1:size(input,2)] #line bright white
                        )

                        plot(data,layout)

      elseif y_axis_label_option == "hover_positions"

        layout = Layout(
                        title = "$title_no_underscores",
                        xaxis=attr(title="Sample ID (Grouped by $(group1_label) | $(group2_label))", showgrid=false, zeroline=false, tickvals=sample_name_indices,
                        ticktext=id_list, tickangle=45,showticklabels=x_axis_label_option),

                        yaxis=attr(title="Genomic Location", zeroline=false, tickvals=chr_pos_tuple_indices,
                        ticktext=chr_pos_tuple_list,tickfont_size=font_size,hovermode=true,automargin=true,showticklabels=false),

                        shapes=shapes, yaxis_range = [1:size(input,1)], xaxis_range = [1:size(input,2)]
        )

                        plot(data,layout)

      else
          println("--y_axis_labels is not a valid option. Choose positions or chromosomes")
      end
  end

"""
    dp_heatmap2(input::Array{Int64,2},title::String,chrom_label_info::Tuple{Array{String,1},Array{Int64,1},String}, sample_names,chr_pos_tuple_list_rev,y_axis_label_option,x_axis_label_option)
generate heatmap of read depth data.
"""
function dp_heatmap2(input, title, chrom_label_info, sample_names,chr_pos_tuple_list_rev,y_axis_label_option,x_axis_label_option)

    chr_pos_tuple_indices,chr_pos_tuple_list,sample_name_indices,sample_names,chrom_labels,chrom_label_indices,font_size = process_plot_inputs(chrom_label_info,sample_names,chr_pos_tuple_list_rev)

    title_no_underscores=replace(title, "_"=>' ')

    hover_text_array=generate_hover_text_array(chr_pos_tuple_list,sample_names,input,"DP")

    position_1,position_2,position_3,position_4,position_5,position_6,position_7,position_8,position_9,position_10,text_point_increment,legend_xaxis_anchor,legend_xaxis_anchor_2,x_axis_text_anchor,extra_right_margin_space,extra_right_margin_space_dp = generate_legend_increments_ungrouped(input)

    trace=heatmap(
        z = input, x=1:size(input, 2),y=1:size(input, 1),

       hoverinfo="text",
       text=hover_text_array,

        zauto=false,zmax=100,zmin=-20,

        transpose=true,
        colorscale = [
                     [0, "rgb(255,255,255)"],
                     [0.2, "rgb(153,231,255)"],
                     [0.3, "rgb(79,146,255)"],
                     [0.41, "rgb(43,124,255)"],
                     [1, "rgb(0,64,168)"]
                     ],

        colorbar = attr(tickvals = [-20,0,20,40,60,80,99],
        title="Read Depth",
        ticktext = ["No Call","0","20","40","60","80","100+"]),

        gridcolor = "#E2E2E2",
        showscale = false, #we define a custom legend using shapes and a text scatter trace
        );

        #we define a custom legend using shapes and a text scatter trace
        trace2 = scatter(
        ;x=[x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor+extra_right_margin_space_dp],
        y=[position_1+text_point_increment,position_2+text_point_increment,position_3+text_point_increment,position_4+text_point_increment,position_5+text_point_increment,position_6+text_point_increment],
          mode="text", name="legend label",
          textposition="middle right",
          text=["0", "10", "30", "70","100+",""],
          marker_size=300, textfont_family="Raleway, sans-serif")

        shapes = [
        rect(x0=legend_xaxis_anchor, y0=position_5, x1=legend_xaxis_anchor_2, y1=position_6, opacity=1.0, fillcolor="rgb(0,64,168)", line_color="black"),
        rect(x0=legend_xaxis_anchor, y0=position_4, x1=legend_xaxis_anchor_2, y1=position_5, opacity=1.0, fillcolor="rgb(43,124,255)", line_color="black"),
        rect(x0=legend_xaxis_anchor, y0=position_3, x1=legend_xaxis_anchor_2, y1=position_4, opacity=1.0, fillcolor="rgb(79,146,255)", line_color="black"),
        rect(x0=legend_xaxis_anchor, y0= position_2, x1=legend_xaxis_anchor_2, y1=position_3, opacity=1.0, fillcolor="rgb(153,231,255)", line_color="black"),
        rect(x0=legend_xaxis_anchor, y0= position_1, x1=legend_xaxis_anchor_2, y1=position_2, opacity=1.0, fillcolor="rgb(255,255,255)", line_color="black")
        ]

        if y_axis_label_option == "chromosomes"

           layout = Layout(
                           shapes=shapes,
                           title = "$title_no_underscores",
                           xaxis=attr(title="Sample ID", showgrid=false, zeroline=false, tickvals=sample_name_indices,
                           ticktext=sample_names, tickangle=45,showticklabels=x_axis_label_option),
                           yaxis=attr(title="Genomic Location", zeroline=false, tickvals=chrom_label_indices,
                           ticktext=chrom_labels,tickfont_size=font_size,hovermode=true,automargin=true,showgrid=false)
           )

          data = [trace, trace2]
              plot(data,layout)

       elseif y_axis_label_option == "positions"

           layout = Layout(
                           shapes=shapes,
                           title = "$title_no_underscores",
                           xaxis=attr(title="Sample ID", showgrid=false, zeroline=false, tickvals=sample_name_indices,
                           ticktext=sample_names, tickangle=45,showticklabels=x_axis_label_option),
                           yaxis=attr(title="Genomic Location", zeroline=false, tickvals=chr_pos_tuple_indices,
                           ticktext=chr_pos_tuple_list,tickfont_size=font_size,hovermode=true,automargin=true,showgrid=false)
           )

      data = [trace, trace2]
           plot(data,layout)

       elseif y_axis_label_option == "hover_positions"

           layout = Layout(
                           shapes=shapes,
                           title = "$title_no_underscores",
                           xaxis=attr(title="Sample ID", showgrid=false, zeroline=false, tickvals=sample_name_indices,
                           ticktext=sample_names, tickangle=45,showticklabels=x_axis_label_option),
                           yaxis=attr(title="Genomic Location", zeroline=false, tickvals=chr_pos_tuple_indices,
                           ticktext=chr_pos_tuple_list,tickfont_size=font_size,hovermode=true,automargin=true,showticklabels=false)
           )
          data = [trace, trace2]
               plot(data,layout)

        else
            println("--y_axis_labels is not a valid option. Choose positions or chromosomes")
        end
end

"""
    dp_heatmap2_with_groups(input::Array{Int64,2},title::String,chrom_label_info::Tuple{Array{String,1},Array{Int64,1},String},group_label_pack::Array{Any,1},id_list,chr_pos_tuple_list_rev,y_axis_label_option,trait_label_array,x_axis_label_option,number_rows)
generate heatmap of read depth data with grouped samples.
"""
function dp_heatmap2_with_groups(input::Array{Int64,2},title::String,chrom_label_info::Tuple{Array{String,1},Array{Int64,1},String},group_label_pack::Array{Any,1},id_list,chr_pos_tuple_list_rev,y_axis_label_option,trait_label_array,x_axis_label_option,number_rows)

    sample_name_indices,id_list,chrom_labels,chrom_label_indices,font_size,group_dividing_line,group1_label,group2_label,chr_pos_tuple_indices,chr_pos_tuple_list,font_size = process_plot_inputs_for_grouped_data(chrom_label_info::Tuple{Array{String,1},Array{Int64,1},String},group_label_pack::Array{Any,1},id_list,chr_pos_tuple_list_rev,trait_label_array)

    title_no_underscores=replace(title, "_"=>' ')

    hover_text_array=generate_hover_text_array_grouped(chr_pos_tuple_list,id_list,input,"DP",number_rows)

    position_1,position_2,position_3,position_4,position_5,position_6,position_7,position_8,position_9,text_point_increment,trait_positions,legend_xaxis_anchor,legend_xaxis_anchor_2,x_axis_text_anchor,extra_right_margin_space = generate_legend_increments_grouped(input)

    position_1_trait_1,position_2_trait_1,position_1_trait_2,position_2_trait_2 = trait_positions

    trace=heatmap(

        z = input, x=1:size(input, 2),y=1:size(input, 1),

        hoverinfo="text",
        text=hover_text_array,

        zauto=false,zmax=100,zmin=-60,

        transpose=true,

                     colorscale = [
                                      [0, "rgb(208, 211, 212)"], #light grey
                                      [0.125, "rgb(151, 154, 154)"], #dark grey
                                      [0.25, "rgb(255, 255, 255)"],
                                      [0.4, "rgb(153,231,255)"],
                                      [0.5, "rgb(79,146,255)"],
                                      [0.5625, "rgb(43,124,255)"],
                                      [1, "rgb(0,64,168)"]
                                  ],

                     colorbar = attr(tickvals = [-60, -40, -20, 0, 20, 40, 60, 80, 99],
                     title="Depth / Trait",
                     ticktext = ["Trait 2","Trait 1","No Call","0","20","40","60","80","100+"]),
                     gridcolor = "#E2E2E2",
                     showscale = false, #we define a custom legend using shapes and a text scatter trace
                     );

         #we define a custom legend using shapes and a text scatter trace
         trace2 = scatter(
         ;x=[x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor,x_axis_text_anchor+extra_right_margin_space,x_axis_text_anchor,x_axis_text_anchor],
         y=[position_1+text_point_increment,position_2+text_point_increment,position_3+text_point_increment,position_4+text_point_increment,position_5+text_point_increment,position_6 + text_point_increment, position_7+text_point_increment,position_8+text_point_increment],
           mode="text", name="legend label",
           textposition="middle right",
           text=["No Call", "0", "10", "20","100+","","Trait 2","Trait 1"],
           marker_size=300, textfont_family="Raleway, sans-serif")

         data = [trace, trace2]

         shapes = [
         rect(x0=legend_xaxis_anchor, y0= position_7, x1=legend_xaxis_anchor_2, y1=position_8, opacity=1.0, fillcolor="rgb(208, 211, 212)", line_color="black"),
         rect(x0=legend_xaxis_anchor, y0= position_8, x1=legend_xaxis_anchor_2, y1=position_9, opacity=1.0, fillcolor="rgb(151, 154, 154)", line_color="black"),
         rect(x0=legend_xaxis_anchor, y0=position_5, x1=legend_xaxis_anchor_2, y1=position_6, opacity=1.0, fillcolor="rgb(37,114,242)", line_color="black"),
         rect(x0=legend_xaxis_anchor, y0=position_4, x1=legend_xaxis_anchor_2, y1=position_5, opacity=1.0, fillcolor="rgb(67,138,254)", line_color="black"),
         rect(x0=legend_xaxis_anchor, y0=position_3, x1=legend_xaxis_anchor_2, y1=position_4, opacity=1.0, fillcolor="rgb(110,182,255)", line_color="black"),
         rect(x0=legend_xaxis_anchor, y0= position_2, x1=legend_xaxis_anchor_2, y1=position_3, opacity=1.0, fillcolor="rgb(153,231,255)", line_color="black"),
         rect(x0=legend_xaxis_anchor, y0= position_1, x1=legend_xaxis_anchor_2, y1=position_2, opacity=1.0, fillcolor="rgb(255,255,255)", line_color="black")
         ]

if y_axis_label_option == "chromosomes"

    layout = Layout(
                    shapes=shapes,
                    title = "$title_no_underscores",

                    xaxis=attr(title="Sample ID (Grouped by $(group1_label) | $(group2_label))",
                    showgrid=false, zeroline=false, tickvals=sample_name_indices,
                    ticktext=id_list, tickangle=45,showticklabels=x_axis_label_option),

                    yaxis=attr(title="Genomic Location", zeroline=false,
                    tickvals=chrom_label_indices,ticktext=chrom_labels,
                    tickfont_size=font_size,hovermode=true,automargin=true,showgrid=false),

                    yaxis_range = [1:size(input,1)],
                    xaxis_range = [1:size(input,2)]
                  )

      plot(data,layout)

elseif y_axis_label_option == "positions"

    layout = Layout(
                    shapes=shapes,
                    title = "$title_no_underscores",

                    xaxis=attr(title="Sample ID (Grouped by $(group1_label) | $(group2_label))",
                    showgrid=false, zeroline=false, tickvals=sample_name_indices,
                    ticktext=id_list, tickangle=45,showticklabels=x_axis_label_option),

                    yaxis=attr(title="Genomic Location", zeroline=false,
                    tickvals=chr_pos_tuple_indices,
                    ticktext=chr_pos_tuple_list,tickfont_size=font_size,
                    hovermode=true,automargin=true,showgrid=false),

                    yaxis_range = [1:size(input,1)],
                    xaxis_range = [1:size(input,2)]
                  )

      plot(data,layout)

elseif y_axis_label_option == "hover_positions"

  layout = Layout(
                  shapes=shapes,
                  title = "$title_no_underscores",
                  xaxis=attr(title="Sample ID (Grouped by $(group1_label) | $(group2_label))", showgrid=false, zeroline=false, tickvals=sample_name_indices,
                  ticktext=id_list, tickangle=45,showticklabels=x_axis_label_option),

                  yaxis=attr(title="Genomic Location", zeroline=false, tickvals=chr_pos_tuple_indices,
                  ticktext=chr_pos_tuple_list,tickfont_size=font_size,hovermode=true,automargin=true,showticklabels=false),

                  yaxis_range = [1:size(input,1)], xaxis_range = [1:size(input,2)] #line bright white
  )
  plot(data,layout)

else
    println("--y_axis_labels is not a valid option. Choose positions or chromosomes")
end

end

#Scatter plots for average read depth viz

"""
    avg_sample_dp_scatter(sample_avg_list::Array{Float64,1},sample_names,x_axis_label_option)
generate line chart of average read depths of each sample.
"""
function avg_sample_dp_scatter(sample_avg_list::Array{Float64,1},sample_names,x_axis_label_option)

    sample_name_indices = collect(1:1:size(sample_names,2))

    trace = scatter(;x=1:size(sample_avg_list,1), y=sample_avg_list,mode="markers")

    layout = Layout(title="Average Sample Read Depth",
    xaxis=attr(title="Samples",ticktext=sample_names,tickvals=sample_name_indices,
    tickangle=45,tickfont_size=5,showticklabels=x_axis_label_option),
    yaxis=attr(title="Average Read Depth"),showticklabels=false)

    plot(trace,layout)
end

"""
    avg_variant_dp_line_chart(variant_avg_list::Array{Float64,1},chr_pos_tuple_list,y_axis_label_option,chrom_label_info)
generate line chart of average read depths of each variant.
"""
function avg_variant_dp_line_chart(variant_avg_list::Array{Float64,1},chr_pos_tuple_list,y_axis_label_option,chrom_label_info)

    chr_pos_tuple_indices = collect(1:1:size(chr_pos_tuple_list,1))
    chrom_labels = chrom_label_info[1]
    returnXY_column1!(chrom_labels)
    chrom_label_indices = chrom_label_info[2]
    font_size = chrom_label_info[3]

    trace = scatter(;x=1:size(variant_avg_list,1), y=variant_avg_list,mode="markers")

    if y_axis_label_option == "positions"
        layout = Layout(title="Average Variant Read Depth",xaxis=attr(title="Variant Positions",
        ticktext=chr_pos_tuple_list,
        tickvals=chr_pos_tuple_indices,
        tickangle=45,showticklabels=true),yaxis=attr(title="Average Read Depth"))
    elseif y_axis_label_option == "chromosomes"
        layout = Layout(title="Average Variant Read Depth",xaxis=attr(title="Variant Positions",
        tickvals=chrom_label_indices,
        ticktext=chrom_labels,
        showticklabels=true),yaxis=attr(title="Average Read Depth"))
    elseif y_axis_label_option == "hover_positions"
        layout = Layout(title="Average Variant Read Depth",xaxis=attr(title="Variant Positions",
        ticktext=chr_pos_tuple_list,
        tickvals=chr_pos_tuple_indices,
        tickangle=45,showticklabels=false,hovermode=true),yaxis=attr(title="Average Read Depth"))
    else
        println("--y_axis_labels is not a valid option. Choose positions or chromosomes")
    end

    plot(trace,layout)

end

"""
    process_plot_inputs(chrom_label_info,sample_names,chr_pos_tuple_list_rev)
Prepares input for heatmap plot function for both genotype and read depth plots without --group_samples flag.
"""
function process_plot_inputs(chrom_label_info,sample_names,chr_pos_tuple_list_rev)

    chr_pos_tuple_indices = collect(1:1:size(chr_pos_tuple_list_rev,1))

    #chr_pos_tuple_list=Array{Tuple{Any,Int64}}(undef,0) #use this if tuple to string doesnt work #version
    chr_pos_tuple_list=Array{String}(undef,0)

    for i in chr_pos_tuple_list_rev
        push!(chr_pos_tuple_list, i)
    end

    sample_name_indices = collect(1:1:size(sample_names,2))

    chrom_labels = chrom_label_info[1]
    returnXY_column1!(chrom_labels)
    chrom_label_indices = chrom_label_info[2]
    font_size = chrom_label_info[3]

    return chr_pos_tuple_indices,chr_pos_tuple_list,sample_name_indices,sample_names,chrom_labels,chrom_label_indices,font_size
end

"""
    process_plot_inputs_for_grouped_data(chrom_label_info::Tuple{Array{String,1},Array{Int64,1},String},group_label_pack::Array{Any,1},id_list,chr_pos_tuple_list_rev,trait_label_array)
Prepares input for heatmap plot function for both genotype and read depth plots with --group_samples flag.
"""
function process_plot_inputs_for_grouped_data(chrom_label_info::Tuple{Array{String,1},Array{Int64,1},String},group_label_pack::Array{Any,1},id_list,chr_pos_tuple_list_rev,trait_label_array)

    sample_name_indices = collect(1:1:size(id_list,2))

    chrom_labels = chrom_label_info[1]
    returnXY_column1!(chrom_labels)
    chrom_label_indices = chrom_label_info[2]
    font_size = chrom_label_info[3]

    group_dividing_line=group_label_pack[3]
    group1_label=group_label_pack[4]
    group2_label=group_label_pack[5]

    chr_pos_tuple_list=Array{String}(undef,0)

    for i in chr_pos_tuple_list_rev
        push!(chr_pos_tuple_list, i)
    end

    y_labels=vcat(chr_pos_tuple_list,trait_label_array)

    y_labels_indices = collect(1:1:size(y_labels,1))

    return sample_name_indices,id_list,chrom_labels,chrom_label_indices,font_size,group_dividing_line,group1_label,group2_label,y_labels_indices,y_labels,font_size

end

"""
     generate_legend_increments_ungrouped(input)
Dynamically generates positons for shapes that build categorical colorscale.
"""
function generate_legend_increments_ungrouped(input)

    number_rows = size(input,1)
    increment = number_rows/30

    number_samples = size(input,2)
    distance_from_plot_edge = number_samples/40
    width_of_color_box = number_samples/40
    legend_xaxis_anchor = number_samples + distance_from_plot_edge
    legend_xaxis_anchor_2 = legend_xaxis_anchor + width_of_color_box
    x_axis_text_anchor = legend_xaxis_anchor_2 + number_samples/40
    extra_right_margin_space=number_samples/4
    extra_right_margin_space_dp = number_samples/10

    position_1 = number_rows/2
    position_2 = position_1 + increment
    position_3 = position_2 + increment
    position_4 = position_3 + increment
    position_5 = position_4 + increment
    position_6 = position_5 + increment
    position_7 = position_6 + increment
    position_8 = position_7 + increment
    position_9 = position_8 + increment
    position_10 = position_9 + increment
    text_point_increment = increment/2
    #trait_positions = position_1_trait_1,position_2_trait_1,position_1_trait_2,position_1_trait_2

    return position_1,position_2,position_3,position_4,position_5,position_6,position_7,position_8,position_9,position_10,text_point_increment,legend_xaxis_anchor,legend_xaxis_anchor_2,x_axis_text_anchor,extra_right_margin_space,extra_right_margin_space_dp
end

"""
     generate_legend_increments_grouped(input)
Dynamically generates positons for shapes that build categorical colorscale including two color boxes for traits 1 and 2
"""
function generate_legend_increments_grouped(input)

    number_rows = size(input,1)
    increment = number_rows/30

    number_samples = size(input,2)
    distance_from_plot_edge = number_samples/40
    width_of_color_box = number_samples/40
    legend_xaxis_anchor = number_samples + distance_from_plot_edge
    legend_xaxis_anchor_2 = legend_xaxis_anchor + width_of_color_box
    x_axis_text_anchor = legend_xaxis_anchor_2 + number_samples/40
    extra_right_margin_space = number_samples/3

    position_1 = number_rows/3
    position_2 = position_1 + increment
    position_3 = position_2 + increment
    position_4 = position_3 + increment
    position_5 = position_4 + increment
    position_6 = position_5 + increment
    position_7 = position_6 + increment
    position_8 = position_7 + increment
    position_9 = position_8 + increment
    #position_10 = position_9 + increment
    text_point_increment = increment/2
    position_1_trait_1 = position_9 + increment
    position_2_trait_1 = position_1_trait_1 + increment
    position_1_trait_2 = position_2_trait_1 + increment
    position_2_trait_2 = position_1_trait_2 + increment
    trait_positions = position_1_trait_1,position_2_trait_1,position_1_trait_2,position_2_trait_2

    return position_1,position_2,position_3,position_4,position_5,position_6,position_7,position_8,position_9,text_point_increment,trait_positions,legend_xaxis_anchor,legend_xaxis_anchor_2,x_axis_text_anchor,extra_right_margin_space
end


module MiniCheetahRobot

using RigidBodyDynamics
using RigidBodyDynamics.Contact

packagepath() = joinpath(@__DIR__, "..", "deps")
urdfpath() = joinpath(packagepath(), "cheetah_description", "mini_cheetah.urdf")

function default_contact_model()
    SoftContactModel(hunt_crossley_hertz(k = 500e3), ViscoelasticCoulombModel(0.8, 20e3, 100.))
end

function mechanism(::Type{T} = Float64;
                   floating = true,
                   contactmodel = default_contact_model(),
                   remove_fixed_tree_joints = false,
                   add_flat_ground = false) where {T}
    mechanism = RigidBodyDynamics.parse_urdf(urdfpath(); scalar_type = T,
                floating = floating, remove_fixed_tree_joints = remove_fixed_tree_joints)

    if contactmodel != nothing
        for leg in (:FR, :FL, :RR, :RL)
            foot = findbody(mechanism, "$(string(leg))_foot")
            frame = default_frame(foot)
            radius = 0.015

            # bottom of foot -- in foot frame, positive x
            add_contact_point!(foot, ContactPoint(Point3D(frame, radius, 0, 0), contactmodel))
        end
    end

    if add_flat_ground
        frame = root_frame(mechanism)
        ground = HalfSpace3D(Point3D(frame, 0., 0., 0.), FreeVector3D(frame, 0., 0., 1.))
        add_environment_primitive!(mechanism, ground)
    end

    remove_fixed_tree_joints && remove_fixed_tree_joints!(mechanism)

    return mechanism
end

function setnominal!(cheetah_state::MechanismState)
    mechanism = cheetah_state.mechanism
    zero!(cheetah_state)

    # Give the legs their natural bend
    hip_pitch = -0.8
    knee_pitch = -1.5
    for leg in (:FR, :FL, :RR, :RL)
        hip = findjoint(mechanism, "$(string(leg))_thigh_joint")
        knee = findjoint(mechanism, "$(string(leg))_calf_joint")
        set_configuration!(cheetah_state, hip, hip_pitch)
        set_configuration!(cheetah_state, knee, knee_pitch)
    end

    # Lift the floating base joint 0.30 off the ground
    floating_joint = first(out_joints(root_body(mechanism), mechanism))
    set_configuration!(cheetah_state, floating_joint, [1; 0; 0; 0; 0; 0; 0.30])
    return cheetah_state
end

function __init__()
    if !isfile(urdfpath())
        error("Could not find $(urdfpath()). Please run `import Pkg; Pkg.build(\"MiniCheetahRobot\")`.")
    end
end

end # module


module MadNLPGPU

import CUDA: CUBLAS, CUSOLVER, CuVector, CuMatrix, toolkit_version, R_64F, has_cuda
import MadNLP:
    @kwdef, Logger, @debug, @warn, @error,
    AbstractOptions, AbstractLinearSolver, set_options!, MadNLPLapackCPU,
    SymbolicException,FactorizationException,SolveException,InertiaException,
    introduce, factorize!, solve!, improve!, is_inertia, inertia, tril_to_full!

if has_cuda()
    include("lapackgpu.jl")
    export MadNLPLapackGPU
end

end # module


#=
# Mandelbrot Set
In the last tutorial we got a brief overview of the architecture of a GPU. In this tutorial We will elborate a bit more on Grid configuration and compute the mandelbrot set on the GPU.

## Grid configuration

When we launch a CUDA kernel we use the syntax `@cuda blocks=x threads=y f(a, b, c...)` where `f` is the function and `a, b, c...` are its arguments.

The `blocks` and `threads` options to `@cuda` are used to specify the **grid** configuration of the kernel which is being launched.

`blocks` specifies the dimension and the size of the grid of blocks. This can be one, two or three dimensional. The reason for having multiple dimensions is to make it easier to express algorithms which index over two or three dimensional spaces.

`threads` specifies the *cooperative thread array* (CTA). Threads in the same CTA have access to better coordination and communication utilities. CTA's can also two or three dimensional for convenient indexing.

There are restrictions related to the grid given below.
* Maximum x-dimension of a grid of thread blocks : $2^{31} - 1$
* Maximum y-, or z-dimension of a grid of thread blocks : $65535 (2^{16} - 1)$
* Maximum x- or y-dimension of a CTA: $1024$
* Maximum z-dimension of a CTA: $64$
* Maximum number of threads per block: $1024$

Now let's go back to our SAXPY example and verify the flexibility in choosing grid configurations.
=#

using CUDA, BenchmarkTools

function gpu_axpy!(A, X, Y)
    ## set tid to thread rank
    tid = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    tid > length(X) && return
    Y[tid] = A*X[tid] + Y[tid]
    return
end

N = 2^20
gpu_v1 = CUDA.rand(N)
gpu_v2 = CUDA.rand(N)
gpu_v3 = copy(gpu_v2)

α = 0.48
gpu_v2 .+= α * gpu_v1

function verify_grid(args, result, numthreads, numblocks = cld(N, numthreads))
    u = copy(args[3])

    @cuda threads=numthreads blocks=numblocks gpu_axpy!(args...)
    println("Explicit kernel launch with threads=$numthreads and blocks=$numblocks is correct: "
    ,result == args[3])

    args[3] = u
    return
end

args = [α, gpu_v1, gpu_v3]

verify_grid(args, gpu_v2, 1024)
verify_grid(args, gpu_v2, 1)
verify_grid(args, gpu_v2, 33)

#=
## Occupancy

The above exercise shows the flexibility in deciding grid configuration. However this raises an important question of how the configuration affects performance. The best way to determine is to actually measure what works best in a given scenario however that may prove to be cumbersome for most workflows. Expecially while developing an application it may not be worth our time to find the optimal configuration.

While each SM(streaming multiprocessor) might be executing 100's of threads from the perspective of the programmer, the GPU Hardware deals with `warps`. Each warp is a set of fixed number of threads(32 on NVIDIA hardware). Scheduling and issuing instructions is done at a per warp basis rather than a per thread basis. There is atleast one *warp scheduler* inside each SM whose job it is to keep the SM as busy as possible. Switching between warps(context switch) is extremely fast and essential to hide latencies such as global memory access. Whenever a warp stalls another warp is immediately switched to.

Coming back to out original question of how to determine the optimal grid configuration. One possible solution is to use the occupancy heurestic. $\mbox{occupancy} = \frac{\mbox{active warps}}{\mbox{maximum number of active warps}}$.
Since each SM has a finite amount of resources. As the number of resources per thread increases, fewer of them can be concurrently executed. Occupancy can limited by register usage, shared memory and block size.

The `launch_configuration` function analyses the kernel's resource usage and suggests a configuration that maximises occupancy.
=#

@show kernel_args = cudaconvert((α, gpu_v1, gpu_v3)) # Convert to GPU friendly types
@show kernel_tt = Tuple{Core.Typeof.(kernel_args)...}
kernel = cufunction(gpu_axpy!, kernel_tt)
kernel_config = launch_configuration(kernel.fun, shmem = 0)

#=
For our example the configurator returns 20 blocks and 1024 threads. The occupancy API does not understand what our kernel is doing, it can only see the input types and the function definition. It's our job to figure out if the suggested configuration will work or not. It's best to keep the suggested block size in account while deciding the launch config.

For this example it's perhaps best to set the block size to 1024 and determine the grid size based on that.

# Mandelbrot Set

A popular example of mathematical visualization is the Mandelbrot set. Mathematically it is defined as the set of [complex numbers](https://simple.wikipedia.org/wiki/Complex_number) $c$ such $f_c(z) = z^2 + c$ is bounded.

In other words:
* $Z_0 = 0$
* $Z_{n + 1} = {Z_n}^2 + c$
* $c$ is in the mandelbrot set if the value of $Z_{n}$ is bounded.

It can be mathematically shown $|Z_n| \leq 2.0 $ for bounded points.
=#

using Images

img = rand(Bool, 10, 10)
Gray.(img)

#-

dims = (3000, 3000)

mset = Array{Bool, 2}(undef, dims)

function mandelbrot_cpu(mset::AbstractArray, dims, iterations)
    origin = CartesianIndex(div.(dims, (2, 2), RoundUp))
    for ind in CartesianIndices(mset)
        ## Compute coordinates for true canvas
        coordinates = Tuple(ind - origin) ./ 1000.
        c = ComplexF32(coordinates[1]im + coordinates[2])
        mset[ind] = mandelbrot(c, iterations)
    end
end

function mandelbrot(c, iterations)
    z = ComplexF32(0, 0)
    for i in 1:iterations
        z = z^2 + c
        abs(z) > 2.0 && return false
    end
    return true
end

#-

mandelbrot_cpu(mset, dims, 32)
Gray.(mset)

# This black and white image is no fun. To add color let's map a color to the number iterations it took $z$ to become greater than two.

mset_color = Array{UInt8}(undef, dims)
function mandelbrot(c, iterations)
    z = ComplexF32(0, 0)
    for i in 1:iterations
        z = z^2 + c
        abs2(z) > 4.0 && return i % UInt8
    end
    return zero(UInt8)
end
mandelbrot_cpu(mset_color, dims, 32)

#-

cmap = colormap("RdBu", 32 + 1)

#-

map(x -> cmap[x + 1], mset_color)

#=
Our task is to move this computation to the GPU. The tricky part with moving to the GPU is that the idexing gets tricky. We can use 1-D indexing then figure out inside the kernel what our 2-D index is or use 2-D index from the get go.
=#

mset_gpu = CuArray{UInt8}(undef, dims)
function mandelbrot_gpu(mset::AbstractArray, dims, iterations)
    ind = CartesianIndex((blockIdx().y - 1)*blockDim().y + threadIdx().y,
                        (blockIdx().x - 1)*blockDim().x + threadIdx().x)
    ## Check if index is valid, if not then exit
    !(ind in CartesianIndices(dims)) && return
    origin = CartesianIndex(div.(dims, (2, 2), RoundUp))

    ## Scale the 3000x3000 image to -1.5 to 1.5
    coordinates = Tuple(ind - origin) ./ 1000.
    c = ComplexF32(coordinates[1]im + coordinates[2]) # x + yi
    mset[ind] = mandelbrot(c, iterations)
    return
end

#-

blkdim = (16, 16)
@cuda blocks=cld.(dims, blkdim) threads=blkdim mandelbrot_gpu(mset_gpu, dims, 32)

#-

## copy back to host and display the same image
map(x -> cmap[x + 1], Array(mset_gpu))

#=
## Thread Divergence

We mentioned in the last tutorial that threads in a warp execute the same instruction. This was not the entire picture as you can guess from the mandelbrot set example. Inside the `mandelbrot` inner loop we have consecutive threads exiting at different points during iteration.

```julia
    for i in 1:iterations
        z = z^2 + c
        abs2(z) > 4.0 && return i % UInt8
    end
```

When threads of the same warp are following different execution paths we call it thread divergence. Even when 1 thread out of 32 follows takes a different branch there is thread divergence.

For example

```julia
if threadIdx().x % 2 == 0
    # Do something
else
    # Do something else
end
```

In the above example threads with an even index will follow the first path and the odd indexed ones will follow the second path. Inside the GPU when a branching condition is evaluated a 32-bit mask is generated for that warp (1-bit for each lane). All lanes whose corresponding mask bit is true will be active and the remaining lanes will be idle. When execution reaches a convergence point, the mask is inverted so that all lanes which were idle become active and vice versa. This kind of an IF-ELSE branch has 50% efficiency. Even if a single thread had diverged (say condition was `threadIdx().x % 32 == 0`) we would still be at 50% efficiency. However, nesting IF-ELSE will further reduce efficiency.

In our mandelbrot example while we definitely had a lot of thread divergence it was still beneficial because threads in a warp represented physically close pixels which diverge less.

Also, note that thread divergence refers to intra-warp divergence rather than inter-warp divergence which does not matter to performance. Because doing work efficiently at the warp level is extremely important for performance we will consider it while writing algorithms. In addition there are a number of functions that work at the warp-level such as `sync_warp()`and `shfl_up_sync()`. We will explore these in the `reduction` and `prefix scan` tutorials.
=#