valid_moves(piece::Bishop) = _valid_moves_diagonal(piece)

valid_moves(piece::Queen) = vcat(
  _valid_moves_straight(piece), _valid_moves_diagonal(piece)
)

function valid_moves(piece::Pawn)
  is_white = piece.player.is_white
  other_player = filter(tmp_player -> tmp_player != piece.player, piece.player.game.players)[1]

  if is_white
    @assert piece.row > 1
    direction = +1
  else
    @assert piece.row < 8
    direction = -1
  end

  @assert 1 <= piece.row+direction <= 8

  board = piece.player.game.board
  moves = []

  work_i, work_j = piece.row+direction, piece.col
  if isnothing(piece.player.game.board[work_i,work_j])
    push!(moves, [work_i, work_j])

    is_init = (
      ( is_white && piece.row == 2 ) ||
      ( !is_white && piece.row == 7 )
    )

    if is_init
      work_i, work_j = piece.row+2*direction, piece.col
      isnothing(piece.player.game.board[work_i,work_j]) && push!(moves, [work_i, work_j])
    end
  end

  work_i = piece.row+direction
  for work_j in piece.col .+ [-1,+1]
    ( 1 <= work_j <= 8 ) || continue

    anti_pawn = other_player.anti_pawn
    work_piece = board[work_i,work_j]

    can_move = !isnothing(anti_pawn) && anti_pawn.row == work_i && anti_pawn.col == work_j
    can_move |= !isnothing(work_piece) && work_piece.player != piece.player

    can_move && push!(moves, [work_i, work_j])
  end

  moves
end

function valid_moves(piece::Knight)
  board = piece.player.game.board
  moves = []

  for work_i in piece.row .+ [+1,-1]
    ( 1 <= work_i <= 8 ) || continue
    for work_j in piece.col .+ [+2,-2]
      ( 1 <= work_j <= 8 ) || continue
      is_valid = isnothing(piece.player.game.board[work_i,work_j])
      is_valid || ( is_valid = ( board[work_i,work_j].player != piece.player ) )
      is_valid && push!(moves, [work_i, work_j])
    end
  end

  for work_i in piece.row .+ [+2,-2]
    ( 1 <= work_i <= 8 ) || continue
    for work_j in piece.col .+ [+1,-1]
      ( 1 <= work_j <= 8 ) || continue
      is_valid = isnothing(piece.player.game.board[work_i,work_j])
      is_valid || ( is_valid = ( board[work_i,work_j].player != piece.player ) )
      is_valid && push!(moves, [work_i, work_j])
    end
  end

  moves
end

function valid_moves(piece::Rook)
  moves = _valid_moves_straight(piece)
  ( piece.player.can_castle_left || piece.player.can_castle_right ) || return moves
  piece.player.has_castled && return moves

  if piece.player.is_white
    work_i = 1
  else
    work_i = 8
  end

  king = piece.player.pieces[findfirst(work_piece -> isa(work_piece,King), piece.player.pieces)]

  ( piece.row == work_i ) || return moves
  ( king.row == work_i ) || return moves

  @assert king.col == 5

  if piece.col < king.col
    j_list = piece.col+1:king.col-1
  else
    j_list = king.col+1:piece.col-1
  end

  for work_j in j_list
    isnothing(piece.player.game.board[work_i,work_j]) || return moves
  end

  push!(moves, [king.row,king.col])

  moves
end

function valid_moves(piece::King)
  board = piece.player.game.board
  moves = []

  for work_i in piece.row .+ (-1:+1)
    ( 1 <= work_i <= 8 ) || continue
    for work_j in piece.col .+ (-1:+1)
      ( 1 <= work_j <= 8 ) || continue

      if !isnothing(piece.player.game.board[work_i,work_j]) && board[work_i,work_j].player == piece.player
        continue
      end

      push!(moves, [work_i, work_j])
    end
  end

  moves
end

function _valid_moves_diagonal(piece::Piece)
  board = piece.player.game.board
  moves = []

  for i_dir = [-1,+1]
    for j_dir = [-1,+1]
      work_i = piece.row + i_dir
      work_j = piece.col + j_dir

      while 1 <= work_i <= 8 && 1 <= work_j <= 8
        if isnothing(piece.player.game.board[work_i,work_j])
          push!(moves, [work_i, work_j])
        elseif board[work_i,work_j].player == piece.player
          break
        else
          push!(moves, [work_i, work_j])
          break
        end

        work_i += i_dir
        work_j += j_dir
      end
    end
  end

  moves
end

function _valid_moves_straight(piece::Piece)
  board = piece.player.game.board
  moves = []

  work_j = piece.col

  for work_i in piece.row-1:-1:1
    push!(moves, [work_i, work_j])
    isnothing(board[work_i,work_j]) && continue
    ( board[work_i,work_j].player == piece.player ) || break
    pop!(moves) ; break
  end

  for work_i in piece.row+1:+1:8
    push!(moves, [work_i, work_j])
    isnothing(board[work_i,work_j]) && continue
    ( board[work_i,work_j].player == piece.player ) || break
    pop!(moves) ; break
  end

  work_i = piece.row

  for work_j in piece.col-1:-1:1
    push!(moves, [work_i, work_j])
    isnothing(board[work_i,work_j]) && continue
    ( board[work_i,work_j].player == piece.player ) || break
    pop!(moves) ; break
  end

  for work_j in piece.col+1:+1:8
    push!(moves, [work_i, work_j])
    isnothing(board[work_i,work_j]) && continue
    ( board[work_i,work_j].player == piece.player ) || break
    pop!(moves) ; break
  end

  moves
end


"""
    ScalarBatchAccumulator(net, state, batch_size)

A ScalarBatchAccumulator is used to evaluate the contribution to
local observables ⟨σ|Ô|σ'⟩ψ(σ')/ψ(σ) , but by computing the neural
network ψ(σ) in batches of size `batch_size`. This is essential
to extrract a speedup when using GPUs.

This is an internal implementation detail of NeuralQuantum, and should
not be relied upon.

Once constructed, a ScalarBatchAccumulator is supposed to be used as follows:
- if `isfull(sba) == true` you should not push new elements to it (an error
will be throw otherwise)
- data is pushed as `sba(⟨σ|Ô|σ'⟩, σ', ψ(σ))`.
The configuration should be passed as a vector (if ket) or as a tuple
of two vectors (if density matrix).
"""
mutable struct AccumulatorLogGradPsi{N,A,B,C} <: AbstractMachineGradAccumulator
    bnet::N         # A batched version of the cached neural network

    in_buf::A       # the matrix of Nsites x batchsz used as input
    out_buf::B      # The row vector of outputs
    ∇out_buf::C

    buf_n::Int      # Counter for elements in buffer
    batch_sz::Int   # batch size
end

function AccumulatorLogGradPsi(net::NeuralNetwork, v, batch_sz)
    #bnet     = cached(net, batch_sz)
    bnet     = net

    w        = trainable_first(net)
    RT       = real(eltype(w))
    in_buf   = deepcopy(v)
    out_buf  = out_similar(net)
    ∇out_buf = grad_cache(net, batch_sz)

    return AccumulatorLogGradPsi(
        bnet,
        in_buf, out_buf, ∇out_buf,
        0, batch_sz)
end

Base.length(a::AccumulatorLogGradPsi) = a.batch_sz
Base.count(a::AccumulatorLogGradPsi)  = a.buf_n
isfull(a::AccumulatorLogGradPsi) = count(a) == length(a)

"""
    init!(c::ScalarBatchAccumulator)

Resets the internal counter of the accumulator, deleting
all previously accumulated (but not computed) values.
"""
init!(c::AccumulatorLogGradPsi) = c.buf_n = 0
function init!(c::AccumulatorLogGradPsi, σ)
    init!(c)

    # Reset the state so that later is faster to apply the changes
    state_copy!(c.in_buf, σ)
    return nothing
end

function (c::AccumulatorLogGradPsi)(v::Union{AState, ADoubleState})
    @assert !isfull(c) "Pushed data to a full accumulator."

    # Increase the step in our internal buffer
    # this should be guaranteed to always be < max_capacity
    c.buf_n = c.buf_n + 1

    buf_i = NeuralQuantum.unsafe_get_batch(c.in_buf, c.buf_n)
    state_copy!(buf_i, v)

    return nothing
end

function (c::AccumulatorLogGradPsi)(cngs)
    @assert !isfull(c) "Pushed data to a full accumulator."
    c.buf_n = c.buf_n + 1

    σp = unsafe_get_batch(c.in_buf, c.buf_n)
    @inbounds apply!(σp, cngs)
end

function (c::AccumulatorLogGradPsi)(cngs_l, cngs_r)
    @assert !isfull(c) "Pushed data to a full accumulator."
    c.buf_n = c.buf_n + 1

    σp = unsafe_get_batch(c.in_buf, c.buf_n)
    @inbounds apply!(σp, cngs_l, cngs_r)
end

"""
    process_accumulator!(c)

Processes all states stored in the accumulator, by computing their
relative local contribution.

It is safe to call this even if the accumulator is not full. In this
case all data beyond the count should be disregarded as it was
not initialized.

The output will be returned. You should not assume ownership of
the output, as it is preallocated and will be used for further
computations of the accumulator.
"""
function process_accumulator!(c::AccumulatorLogGradPsi)
    out  = c.out_buf
    ∇out = c.∇out_buf

    # Compute the batch of logψ neural networks
    logψ_and_∇logψ!(∇out, out, c.bnet, c.in_buf)

    # Reset the counter of the batch accumulator
    init!(c)

    return ∇out, out
end


using .NumericalFluxes: numerical_volume_flux_first_order!

# {{{ FIXME: remove this after we've figure out how to pass through to kernel
const _ξ1x1, _ξ2x1, _ξ3x1 = Grids._ξ1x1, Grids._ξ2x1, Grids._ξ3x1
const _ξ1x2, _ξ2x2, _ξ3x2 = Grids._ξ1x2, Grids._ξ2x2, Grids._ξ3x2
const _ξ1x3, _ξ2x3, _ξ3x3 = Grids._ξ1x3, Grids._ξ2x3, Grids._ξ3x3
const _M = Grids._M

const _n1, _n2, _n3 = Grids._n1, Grids._n2, Grids._n3
const _sM, _vMI = Grids._sM, Grids._vMI
# }}}

@doc """
    esdg_volume_tendency!(
        balance_law::BalanceLaw,
        ::Val{dim},
        ::Val{polyorder},
        tendency,
        state_prognostic,
        state_auxiliary,
        vgeo,
        D,
        α,
        β,
    ) where {dim, polyorder}

Computes the reference element horizontal tendency using a two-point flux
approximation of all derivatives.
"""
@kernel function esdg_volume_tendency!(
    balance_law::BalanceLaw,
    ::Val{dir},
    ::Val{info},
    volume_numerical_flux_first_order,
    tendency,
    state_prognostic,
    state_auxiliary,
    vgeo,
    D,
    α,
    β,
    add_source = false,
) where {dir, info}
    @uniform begin
        dim = info.dim

        FT = eltype(state_prognostic)
        num_state = number_states(balance_law, Prognostic(), FT)
        num_aux = number_states(balance_law, Auxiliary(), FT)

        local_H = MArray{Tuple{3, num_state}, FT}(undef)

        state_2 = MArray{Tuple{num_state}, FT}(undef)
        aux_2 = MArray{Tuple{num_aux}, FT}(undef)

        local_source = MArray{Tuple{num_state}, FT}(undef)

        @inbounds Nq1 = info.Nq[1]
        @inbounds Nq2 = info.Nq[2]
        Nq3 = info.Nqk

        if dir == 1
            _ξdx1, _ξdx2, _ξdx3 = _ξ1x1, _ξ1x2, _ξ1x3
            Nqd = Nq1
        elseif dir == 2
            _ξdx1, _ξdx2, _ξdx3 = _ξ2x1, _ξ2x2, _ξ2x3
            Nqd = Nq2
        elseif dir == 3
            _ξdx1, _ξdx2, _ξdx3 = _ξ3x1, _ξ3x2, _ξ3x3
            Nqd = Nq3
        end
    end

    e = @index(Group, Linear)
    i, j, k = @index(Local, NTuple)

    state_1 = @private FT (num_state,)
    aux_1 = @private FT (num_aux,)
    local_tendency = @private FT (num_state,)
    local_MI = @private FT (1,)
    shared_G = @localmem FT (Nq1 * Nq2 * Nq3, 3)
    shared_D = @localmem FT (Nqd, Nqd)

    @inbounds begin
        ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))

        # generalization for different polynomial orders
        @unroll for l in 1:Nqd
            if dir == 1
                id = i
                ild = l + Nq1 * ((j - 1) + Nq2 * (k - 1))
            elseif dir == 2
                id = j
                ild = i + Nq1 * ((l - 1) + Nq2 * (k - 1))
            elseif dir == 3
                id = k
                ild = i + Nq1 * ((j - 1) + Nq2 * (l - 1))
            end
            shared_D[id, l] = D[id, l]
        end

        M = vgeo[ijk, _M, e]
        shared_G[ijk, 1] = M * vgeo[ijk, _ξdx1, e]
        shared_G[ijk, 2] = M * vgeo[ijk, _ξdx2, e]
        if dim == 3
            shared_G[ijk, 3] = M * vgeo[ijk, _ξdx3, e]
        end

        # Build ode scaling into mass matrix (so it doesn't show up later)
        local_MI[1] = α / M

        # Get the volume tendency (scaling by β done below)
        @unroll for s in 1:num_state
            local_tendency[s] = β != 0 ? tendency[ijk, s, e] : -zero(FT)
        end

        # Scale β into the volume tendency
        @unroll for s in 1:num_state
            local_tendency[s] *= β
        end

        # Compute the volume tendency
        # ∑_{i,j}^{d} ( G_ij (Q_i ∘ H_j) - (H_j ∘ Q_i^T) G_ij)
        #  =
        # ( G_11 (Q_1 ∘ H_1) - (H_1 ∘ Q_1^T) G_11) 1 +
        # ( G_12 (Q_1 ∘ H_2) - (H_2 ∘ Q_1^T) G_12) 1 +
        # ( G_13 (Q_1 ∘ H_3) - (H_3 ∘ Q_1^T) G_13) 1 +
        # ( G_21 (Q_2 ∘ H_1) - (H_1 ∘ Q_2^T) G_21) 1 +
        # ( G_22 (Q_2 ∘ H_2) - (H_2 ∘ Q_2^T) G_22) 1 +
        # ( G_23 (Q_2 ∘ H_3) - (H_3 ∘ Q_2^T) G_23) 1 +
        # ( G_31 (Q_3 ∘ H_1) - (H_1 ∘ Q_3^T) G_31) 1 +
        # ( G_32 (Q_3 ∘ H_2) - (H_2 ∘ Q_3^T) G_32) 1 +
        # ( G_33 (Q_3 ∘ H_3) - (H_3 ∘ Q_3^T) G_33) 1
        @unroll for s in 1:num_state
            state_1[s] = state_prognostic[ijk, s, e]
        end
        @unroll for s in 1:num_aux
            aux_1[s] = state_auxiliary[ijk, s, e]
        end

        if add_source
            fill!(local_source, -zero(eltype(local_source)))
            source!(
                balance_law,
                Vars{vars_state(balance_law, Prognostic(), FT)}(local_source),
                Vars{vars_state(balance_law, Prognostic(), FT)}(state_1),
                Vars{vars_state(balance_law, Auxiliary(), FT)}(aux_1),
            )

            @unroll for s in 1:num_state
                local_tendency[s] += α * local_source[s]
            end
        end

        @synchronize

        ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))

        # Note: unrolling this loop makes things slower
        @views for l in 1:Nqd
            if dir == 1
                id = i
                ild = l + Nq1 * ((j - 1) + Nq2 * (k - 1))
            elseif dir == 2
                id = j
                ild = i + Nq1 * ((l - 1) + Nq2 * (k - 1))
            elseif dir == 3
                id = k
                ild = i + Nq1 * ((j - 1) + Nq2 * (l - 1))
            end

            # Compute derivatives wrt ξd
            # ( G_31 (Q_3 ∘ H_1) - (H_1 ∘ Q_3^T) G_31) 1 +
            # ( G_32 (Q_3 ∘ H_2) - (H_2 ∘ Q_3^T) G_32) 1 +
            # ( G_33 (Q_3 ∘ H_3) - (H_3 ∘ Q_3^T) G_33) 1 +
            @unroll for s in 1:num_state
                state_2[s] = state_prognostic[ild, s, e]
            end
            @unroll for s in 1:num_aux
                aux_2[s] = state_auxiliary[ild, s, e]
            end
            fill!(local_H, -zero(FT))
            numerical_volume_flux_first_order!(
                volume_numerical_flux_first_order,
                balance_law,
                local_H,
                state_1[:],
                aux_1[:],
                state_2,
                aux_2,
            )
            @unroll for s in 1:num_state
                # G_21 (Q_2 ∘ H_1) 1 +
                # G_22 (Q_2 ∘ H_2) 1 +
                # G_23 (Q_2 ∘ H_3) 1
                local_tendency[s] -=
                    local_MI[1] *
                    shared_D[id, l] *
                    (
                        shared_G[ijk, 1] * local_H[1, s] +
                        shared_G[ijk, 2] * local_H[2, s] +
                        (
                            dim == 3 ? shared_G[ijk, 3] * local_H[3, s] :
                            -zero(FT)
                        )
                    )
                #  (H_1 ∘ Q_2^T) G_21 1 +
                #  (H_2 ∘ Q_2^T) G_22 1 +
                #  (H_3 ∘ Q_2^T) G_23 1
                local_tendency[s] +=
                    local_MI[1] *
                    (
                        local_H[1, s] * shared_G[ild, 1] +
                        local_H[2, s] * shared_G[ild, 2] +
                        (
                            dim == 3 ? local_H[3, s] * shared_G[ild, 3] :
                            -zero(FT)
                        )
                    ) *
                    shared_D[l, id]
            end
        end

        @unroll for s in 1:num_state
            tendency[ijk, s, e] = local_tendency[s]
        end
    end
end


"""
    NeighborLongitudinalResult

A structure to retrieve information about a neihbor in the longitudinal direction i.e. rear and front neighbors on the same lane.
If the neighbor index is equal to `nothing` it means there is no neighbor.

# Fields
- `ind::Union{Nothing, Int64}` index of the neighbor in the scene
- `Δs::Float64` positive distance along the lane between vehicles positions
"""
struct NeighborLongitudinalResult
    ind::Union{Nothing, Int64} # index in scene of the neighbor
    Δs::Float64 # positive distance along lane between vehicles' positions
end

"""
    VehicleTargetPoint

Defines a point on an entity that is used to measure distances.
The following target points are supported and are subtypes of VehicleTargetPoint:
- `VehicleTargetPointFront`
- `VehicleTargetPointCenter`
- `VehicleTargetPointRear`

The method `targetpoint_delta(::VehicleTargetPoint, ::Entity)` can be used to
compute the delta in the longitudinal direction to add when considering a specific target point.
"""
abstract type VehicleTargetPoint end

"""
    VehicleTargetPointFront <: VehicleTargetPoint
Set the target point at the front (and center) of the entity.
"""
struct VehicleTargetPointFront <: VehicleTargetPoint end

targetpoint_delta(::VehicleTargetPointFront, veh::Entity{S, D, I}) where {S,D<:AbstractAgentDefinition, I} = length(veh.def)/2*cos(posf(veh.state).ϕ)

"""
    VehicleTargetPointCenter <: VehicleTargetPoint
Set the target point at the center of the entity.
"""
struct VehicleTargetPointCenter <: VehicleTargetPoint end

targetpoint_delta(::VehicleTargetPointCenter, veh::Entity{S, D, I}) where {S,D<:AbstractAgentDefinition, I} = 0.0

"""
    VehicleTargetPointFront <: VehicleTargetPoint
Set the target point at the front (and center) of the entity.
"""
struct VehicleTargetPointRear <: VehicleTargetPoint end

targetpoint_delta(::VehicleTargetPointRear, veh::Entity{S, D, I}) where {S,D<:AbstractAgentDefinition, I} = -length(veh.def)/2*cos(posf(veh.state).ϕ)

const VEHICLE_TARGET_POINT_CENTER = VehicleTargetPointCenter()

"""
    find_neighbor(scene::Scene, roadway::Roawday, ego::Entity; kwargs...)

Search through lanes and road segments to find a neighbor of `ego` in the `scene`.

# Arguments

- `scene::Scene` the scene containing all the entities
- `roadway::Roadway` the road topology on which entities are driving
- `ego::Entity` the entity that we want to compute the neighbor of.

# Keyword arguments
- `lane::Union{Nothing, Lane}` the lane on which to search the neighbor, if different from the ego vehicle current lane, it uses the projection of the ego vehicle on the given lane as a reference point. If nothing, returns nothing.
- `rear::Bool = false` set to true to search for rear neighbor, search forward neighbor by default
- `max_distance::Float64 = 250.0` stop searching after this distance is reached, if the neighbor is further than max_distance, returns nothing
- `targetpoint_ego::VehicleTargetPoint` the point on the ego vehicle used for distance calculation, see `VehicleTargetPoint` for more info
- `targetpoint_neighbor::VehicleTargetPoint` the point on the neighbor vehicle used for distance calculation, see `VehicleTargetPoint` for more info
- `ids_to_ignore::Union{Nothing, Set{I}} = nothing` a list of entity ids to ignore for the search, ego is always ignored.
"""
function find_neighbor(scene::Scene, roadway::Roadway, ego::Entity{S,D,I};
                       lane::Union{Nothing, Lane} = get_lane(roadway, ego),
                       rear::Bool=false,
                       max_distance::Float64=250.0,
                       targetpoint_ego::VehicleTargetPoint = VehicleTargetPointCenter(),
                       targetpoint_neighbor::VehicleTargetPoint = VehicleTargetPointCenter(),
                       ids_to_ignore::Union{Nothing, Set{I}} = nothing) where {S,D,I}

    if lane == nothing
        return NeighborLongitudinalResult(nothing, max_distance)
    elseif get_lane(roadway, ego).tag == lane.tag
        tag_start = lane.tag
    else  # project ego on desired lane
        roadproj = proj(posg(ego.state), lane, roadway)
        tag_start = roadproj.tag
    end
    s_base = posf(ego.state).s + targetpoint_delta(targetpoint_ego, ego)

    tag_target = tag_start
    best_ind = nothing
    best_dist = max_distance

    dist_searched = 0.0
    while dist_searched < max_distance

        curr_lane = roadway[tag_target]

        for (i, veh) in enumerate(scene)

            if veh.id == ego.id
                continue
            end

            if ids_to_ignore !== nothing && veh.id ∈ ids_to_ignore
                continue
            end

            # check if veh is on thislane
            s_adjust = NaN

            if get_lane(roadway, veh).tag == curr_lane.tag
                s_adjust = 0.0

            elseif is_between_segments_hi(posf(veh.state).roadind.ind, curr_lane.curve) &&
                   is_in_entrances(roadway[tag_target], posf(veh.state).roadind.tag)

                distance_between_lanes = norm(VecE2(roadway[tag_target].curve[1].pos - roadway[posf(veh.state).roadind.tag].curve[end].pos))
                s_adjust = -(roadway[posf(veh.state).roadind.tag].curve[end].s + distance_between_lanes)

            elseif is_between_segments_lo(posf(veh.state).roadind.ind) &&
                    is_in_exits(roadway[tag_target], posf(veh.state).roadind.tag)

                distance_between_lanes = norm(VecE2(roadway[tag_target].curve[end].pos - roadway[posf(veh.state).roadind.tag].curve[1].pos))
                s_adjust = roadway[tag_target].curve[end].s + distance_between_lanes
            end

            if !isnan(s_adjust)
                s_valid = posf(veh.state).s + targetpoint_delta(targetpoint_neighbor, veh) + s_adjust
                if rear
                    dist_valid = s_base - s_valid + dist_searched
                else
                    dist_valid = s_valid - s_base + dist_searched
                end

                if dist_valid ≥ 0.0
                    s_primary = posf(veh.state).s + targetpoint_delta(targetpoint_neighbor, veh) + s_adjust
                    if rear
                        dist=  s_base - s_primary + dist_searched
                    else
                        dist = s_primary - s_base + dist_searched
                    end
                    if dist < best_dist
                        best_dist = dist
                        best_ind = i
                    end
                end
            end
        end

        # neighbor has been found exit
        if best_ind != nothing
            break
        end

        # no next lane and no neighbor found exit
        if !has_next(curr_lane) ||
            (tag_target == tag_start && dist_searched != 0.0) # exit after visiting this lane a 2nd time
            break
        end

        # go to the connected lane.
        if rear
            dist_searched += s_base
            s_base = curr_lane.curve[end].s + norm(VecE2(lane.curve[end].pos - prev_lane_point(curr_lane, roadway).pos))
            tag_target = prev_lane(lane, roadway).tag
        else
            dist_searched += (curr_lane.curve[end].s - s_base)
            s_base = -norm(VecE2(curr_lane.curve[end].pos - next_lane_point(curr_lane, roadway).pos)) # negative distance between lanes
            tag_target = next_lane(curr_lane, roadway).tag
        end
    end

    return NeighborLongitudinalResult(best_ind, best_dist)
end

"""
    FrenetRelativePosition

Contains information about the projection of a point on a lane. See `get_frenet_relative_position`.

# Fields
- `origin::RoadIndex` original roadindex used for the projection, contains the target lane ID.
- `target::RoadIndex` roadindex reached after projection
- `Δs::Float64` longitudinal distance to the original roadindex
- `t::Float64` lateral distance to the original roadindex in the frame of the target lane
- `ϕ::Float64` angle with the original roadindex in the frame of the target lane
"""
struct FrenetRelativePosition
    origin::RoadIndex
    target::RoadIndex
    Δs::Float64
    t::Float64
    ϕ::Float64
end

"""
    get_frenet_relative_position(veh_fore::Entity, veh_rear::Entity, roadway::Roadway)
return the Frenet relative position between the two vehicles. It projects the position of the first vehicle onto the lane of the second vehicle.
The result is stored as a `FrenetRelativePosition`.

Lower level:
    get_frenet_relative_position(posG::VecSE2{Float64}, roadind::RoadIndex, roadway::Roadway;
        max_distance_fore::Float64 = 250.0, # max distance to search forward [m]
        max_distance_rear::Float64 = 250.0, # max distance to search backward [m]
        improvement_threshold::Float64 = 1e-4,
    )

Project the given point to the same lane as the given RoadIndex.
This will return the projection of the point, along with the Δs along the lane from the RoadIndex.
The returned type is a `FrenetRelativePosition` object.
"""
function get_frenet_relative_position(posG::VecSE2{Float64}, roadind::RoadIndex, roadway::Roadway;
    max_distance_fore::Float64 = 250.0, # max distance to search forward [m]
    max_distance_rear::Float64 = 250.0, # max distance to search backward [m]
    improvement_threshold::Float64 = 1e-4,
    )

    # project to current lane first
    tag_start = roadind.tag
    lane_start = roadway[tag_start]
    curveproj_start = proj(posG, lane_start, roadway, move_along_curves=false).curveproj
    curvept_start = lane_start[curveproj_start.ind, roadway]
    s_base = lane_start[roadind.ind, roadway].s
    s_proj = curvept_start.s
    Δs = s_proj - s_base
    sq_dist_to_curve = normsquared(VecE2(posG - curvept_start.pos))
    retval = FrenetRelativePosition(roadind,
                RoadIndex(curveproj_start.ind, tag_start), Δs, curveproj_start.t, curveproj_start.ϕ)

    # search downstream
    if has_next(lane_start)
        dist_searched = lane_start.curve[end].s - s_base
        s_base = -norm(VecE2(lane_start.curve[end].pos - next_lane_point(lane_start, roadway).pos)) # negative distance between lanes
        tag_target = next_lane(lane_start, roadway).tag
        while dist_searched < max_distance_fore

            lane = roadway[tag_target]
            curveproj = proj(posG, lane, roadway, move_along_curves=false).curveproj
            sq_dist_to_curve2 = normsquared(VecE2(posG - lane[curveproj.ind, roadway].pos))

            if sq_dist_to_curve2 < sq_dist_to_curve - improvement_threshold

                sq_dist_to_curve = sq_dist_to_curve2
                s_proj = lane[curveproj.ind, roadway].s
                Δs = s_proj - s_base + dist_searched
                retval = FrenetRelativePosition(
                            roadind, RoadIndex(curveproj.ind, tag_target),
                            Δs, curveproj.t, curveproj.ϕ)
            end

            if !has_next(lane)
                break
            end

            dist_searched += (lane.curve[end].s - s_base)
            s_base = -norm(VecE2(lane.curve[end].pos - next_lane_point(lane, roadway).pos)) # negative distance between lanes
            tag_target = next_lane(lane, roadway).tag

            if tag_target == tag_start
                break
            end
        end
    end

    # search upstream
    if has_prev(lane_start)
        dist_searched = s_base
        tag_target = prev_lane(lane_start, roadway).tag
        s_base = roadway[tag_target].curve[end].s + norm(VecE2(lane_start.curve[1].pos - prev_lane_point(lane_start, roadway).pos)) # end of the lane

        while dist_searched < max_distance_fore

            lane = roadway[tag_target]
            curveproj = proj(posG, lane, roadway, move_along_curves=false).curveproj
            sq_dist_to_curve2 = normsquared(VecE2(posG - lane[curveproj.ind, roadway].pos))

            if sq_dist_to_curve2 < sq_dist_to_curve - improvement_threshold

                sq_dist_to_curve = sq_dist_to_curve2
                s_proj = lane[curveproj.ind, roadway].s
                dist = s_base - s_proj  + dist_searched
                Δs = -dist
                retval = FrenetRelativePosition(
                            roadind, RoadIndex(curveproj.ind, tag_target),
                            Δs, curveproj.t, curveproj.ϕ)
            end

            if !has_prev(lane)
                break
            end

            dist_searched += s_base
            s_base = lane.curve[end].s + norm(VecE2(lane.curve[1].pos - prev_lane_point(lane, roadway).pos)) # length of this lane plus crossover
            tag_target = prev_lane(lane, roadway).tag

            if tag_target == tag_start
                break
            end
        end
    end

    retval
end
get_frenet_relative_position(veh_fore::Entity{S, D, I}, veh_rear::Entity{S, D, I}, roadway::Roadway) where {S,D, I} = get_frenet_relative_position(posg(veh_fore.state), posf(veh_rear.state).roadind, roadway)

"""
    dist_to_front_neighbor(roadway::Roadway, scene::Scene, veh::Entity)
Feature function to extract the longitudinal distance to the front neighbor (in the Frenet frame).
Returns `missing` if there are no front neighbor.
"""
function dist_to_front_neighbor(roadway::Roadway, scene::Scene, veh::Entity)
    neighbor = find_neighbor(scene, roadway, veh)
    if neighbor.ind === nothing
        return missing
    else
        return neighbor.Δs
    end
end

"""
    front_neighbor_speed(roadway::Roadway, scene::Scene, veh::Entity)
Feature function to extract the velocity of the front neighbor.
Returns `missing` if there are no front neighbor.
"""
function front_neighbor_speed(roadway::Roadway, scene::Scene, veh::Entity)
    neighbor = find_neighbor(scene, roadway, veh)
    if neighbor.ind === nothing
        return missing
    else
        return vel(scene[neighbor.ind])
    end
end

"""
    time_to_collision(roadway::Roadway, scene::Scene, veh::Entity)
Feature function to extract the time to collision with the front neighbor.
Returns `missing` if there are no front neighbor.
"""
function time_to_collision(roadway::Roadway, scene::Scene, veh::Entity)
    neighbor = find_neighbor(scene, roadway, veh)
    if neighbor.ind === nothing
        return missing
    else
        len_ego = length(veh.def)
        len_oth = length(scene[neighbor.ind].def)
        Δs = neighbor.Δs - len_ego/2 - len_oth/2
        Δv = vel(scene[neighbor.ind]) - vel(veh)
        return -Δs / Δv
    end
end


# Demands that have no flexibility.

# Import packages.
using Dates, TimeSeries

# Timeseries of possibly mobile demands.

"""
Demands that vary with time.

# Type parameters
- `T <: TimeType`: timestamp

# Example
```
example = TimeArray(
    [Time(0), Time(8), Time(9), Time(17), Time(18)],
    [     0.,     10.,      0.,      11.,       0.]
) :: TemporalDemand{Time}
```
"""
const TemporalDemand{T <: TimeType} = TimeArray{Float64,1,T,Array{Float64,1}}

"""
Locations where demand moves.

# Type parameters
- `T <: TimeType`: timestamp
- `L`            : network location

# Example
```
example = TimeArray(
    [Time(0)   , Time(8)   , Time(9)       , Time(17) , Time(18)  ],
    ["Home #23", "Road #14", "Workplace #3", "Road #9", "Home #23"]
) :: MobileDemand{Time,String}
```
"""
const MobileDemand{T <: TimeType, L} = TimeArray{L,1,T,Array{L,1}}

"""
Demands that move.

# Type parameters
- `T <: TimeType`: timestamp
- `L`            : network location

# Example
```
example = TimeArray(
    [Time(0)         , Time(8)          , Time(9)             , Time(17)        , Time(18)        ],
    [("Home #23", 0.), ("Road #14", 10.), ("Workplace #3", 0.), ("Road #9", 11.), ("Home #23", 0.)]
) :: LocatedDemand{Time,String}
```
"""
const LocatedDemand{T <: TimeType, L} = TimeArray{Tuple{L,Float64},1,T,Array{Tuple{L,Float64},1}}

# Abstract inflexbile demands.

abstract type InflexibleDemand{T,L} <: Demand{T,L} end

"""
Time-varying demands and their locations.

This must be implemented by subtypes of `InflexibleDemand`.

# Arguments
- `demand :: InflexibleDemand{T,L}`: the demand
"""
function demands(demand :: InflexibleDemand{T,L}) :: LocatedDemand{T,L} where L where T <: TimeType
end

"""
The "envelope" of minimum and maximum allowable demands at each time pont, with a location for the demand.

# Arguments
- `demand :: InflexibleDemand{T,L}`: the demand
"""
function envelope(demand :: InflexibleDemand{T,L}) :: LocatedEnvelope{T,L} where L where T <: TimeType
    simple = demands(demand)
    TimeArray(
        timestamp(simple),
        map(x -> (x[1], x[2], x[2]), values(simple))
    )
end


# Test typing in a typed gripper domain
@testset "typing" begin

path = joinpath(dirname(pathof(PDDL)), "..", "test", "typing")

domain = load_domain(joinpath(path, "domain.pddl"))
@test domain.name == Symbol("gripper-typed")
@test convert(Term, domain.predicates[:free]) == pddl"(free ?g)"
@test domain.predicates[:carry].argtypes == (:ball, :gripper)
@test :gripper in PDDL.get_types(domain)

problem = load_problem(joinpath(path, "problem.pddl"))
@test problem.name == Symbol("gripper-problem")
@test problem.objects == @pddl("rooma", "roomb", "ball1", "ball2", "left", "right")
@test problem.objtypes[Const(:ball1)] == :ball

state = initstate(domain, problem)
implementations = [
    "concrete interpreter" => (domain, state),
    "abstract interpreter" => abstracted(domain, state),
    "ground interpreter" => (ground(domain, state), state),
    "concrete compiler" => compiled(domain, state),
    "abstract compiler" => compiled(abstracted(domain), state)
]

@testset "typing ($name)" for (name, (domain, _)) in implementations

    # Test forward execution of plans
    state = initstate(domain, problem)
    state = execute(domain, state, pddl"(pick ball1 rooma left)")
    @test domain[state => pddl"(carry ball1 left)"] ≃ true
    state = execute(domain, state, pddl"(move rooma roomb)")
    @test domain[state => pddl"(robbyat roomb)"] ≃ true
    state = execute(domain, state, pddl"(drop ball1 roomb left)")
    @test domain[state => pddl"(at ball1 roomb)"] ≃ true

    @test satisfy(domain, state, problem.goal) ≃ true

    # Test action availability
    state = initstate(domain, problem)
    @test Set{Term}(available(domain, state)) == Set{Term}(@pddl(
        "(pick ball1 rooma right)", "(pick ball1 rooma left)",
        "(pick ball2 rooma right)", "(pick ball2 rooma left)",
        "(move rooma roomb)", "(move rooma rooma)"
    ))

end

# Test backward regression of plans
state = goalstate(domain, problem)
state = regress(domain, state, pddl"(drop ball1 roomb left)")
@test domain[state => pddl"(carry ball1 left)"] == true
state = regress(domain, state, pddl"(move rooma roomb)")
@test domain[state => pddl"(robbyat rooma)"] == true
state = regress(domain, state, pddl"(pick ball1 rooma left)")
@test domain[state => pddl"(at ball1 rooma)"] == true
@test issubset(state, initstate(domain, problem))

# Test action relevance
state = goalstate(domain, problem)
@test Set{Term}(relevant(domain, state)) == Set{Term}(@pddl(
    "(drop ball1 roomb left)", "(drop ball1 roomb right)"
))

end # typing


include("filtering.jl");
##
list = ["PreVehicle",
    "Altanserin",
    "SB242084",
    "Way_100135",
    "Methysergide",
    "Citalopram",
    "SB242084_opt",
    "Saline"]
s = filter(r->r.Treatment in list &&
    r.Trial < 61 &&
    r.Trial_Travel_to < 10 &&
    r.MouseID != "pc7",
    streaks)
s[s.Treatment .== "PreVehicle",:Treatment] .= "Control"
s[s.Treatment .== "Saline",:Treatment] = [o ? "Optogenetic" : "Control" for o in s[s.Treatment .== "Saline",:Stim]]
s[s.Treatment .== "SB242084_opt",:Phase] .=  "SB242084_opt"
s[s.Treatment .== "SB242084_opt",:Treatment] = [o ? "SB242084_opt" : "Control" for o in s[s.Treatment .== "SB242084_opt",:Stim]]
##Webers Law

release = filter(r -> r.Treatment in ["Citalopram","Optogenetic"],s)
WebersLaw(release,:AfterLast,:Treatment)
savefig(joinpath(figs_loc,"WebersLaw/ReleaseManipulation.pdf"))

selective = filter(r -> r.Treatment in ["SB242084","Altanserin"],s)
WebersLaw(selective,:AfterLast,:Treatment)
savefig(joinpath(figs_loc,"WebersLaw/SelectiveManipulation.pdf"))

both = filter(r -> r.Treatment in ["Citalopram","Optogenetic","SB242084","Altanserin"],s)
WebersLaw(both,:AfterLast,:Treatment)
savefig(joinpath(figs_loc,"WebersLaw/EffectiveManipulation.pdf"))

manipulations = filter(r -> r.Treatment != "Control",s)
WebersLaw(manipulations,:AfterLast,:Treatment)
savefig(joinpath(figs_loc,"WebersLaw/AllManipulation.pdf"))

WebersLaw(s,:AfterLast,:Treatment)
savefig(joinpath(figs_loc,"WebersLaw/WithControl.pdf"))
##
Drug_colors!(s)

WebersLaw(s,:AfterLast,:Treatment)

################ 2c opto comparison across days ##############
optopharma = filter(r ->
    r.Phase in ["Optogenetic", "SB242084_opt"], s)
gd = groupby(optopharma,[:MouseID,:Phase,:Stim])
df1 = combine([:Num_pokes,:AfterLast,:Trial_Travel_to] => (n,a,t) ->
    (Num_pokes = mean(n), AfterLast = mean(a), Travel = mean(t))
    ,gd)
df2 = select(df1, Not([:Travel,:Num_pokes]))
df3 = combine(groupby(df2,:Stim)) do dd
    subdf = unstack(dd,:Phase,:AfterLast)
    # current_drug = Symbol(subdf[1,:Phase])
    # rename!(subdf, current_drug => :Drug)
end
open_html_table(df3)
df4 = combine(groupby(df3,[:Stim])) do dd
    wilcoxon(dd,:SB242084_opt, :Optogenetic; f = x -> mean(skipmissing(x)))
end

plot_wilcoxon(df4)
############ Find the exact day before SB manipulation as control ###############
using Dates
stim_drug_days = [Date(2016,07,22),
    Date(2016,07,23),
    Date(2016,07,26),
    Date(2016,07,28),
    Date(2016,07,30),
    Date(2016,08,01)]
optopharma = filter(r ->
    r.Phase in ["Optogenetic", "SB242084_opt"], s)
open_html_table(optopharma)
union(filter(r -> r.Treatment == "SB242084_opt",optopharma).Day)
gd = groupby(optopharma,[:MouseID,:Phase,:Stim])
df1 = combine([:Num_pokes,:AfterLast,:Trial_Travel_to] => (n,a,t) ->
    (Num_pokes = mean(n), AfterLast = mean(a), Travel = mean(t))
    ,gd)
open_html_table(df1)
stim = filter(r -> r.Phase == "Optogenetic" && r.Stim, df1).AfterLast
nostim = filter(r -> r.Phase == "Optogenetic" && !r.Stim, df1).AfterLast
sb_stim = filter(r -> r.Phase == "SB242084_opt" && r.Stim, df1).AfterLast
sb_nostim = filter(r -> r.Phase == "SB242084_opt" && !r.Stim, df1).AfterLast

SignedRankTest(nostim,sb_nostim)


@testset "simple workflow" begin
    env = StateOverriddenEnv(CartPoleEnv{Float32}(), deepcopy)
    policy = RandomPolicy(action_space(env))
    N_EPISODE = 10_000
    hook = TotalRewardPerEpisode()
    run(policy, env, StopAfterEpisode(N_EPISODE), hook)

    @test isapprox(sum(hook[]) / N_EPISODE, 21; atol = 2)
end


##
using DelimitedFiles, Logging, Pkg, Test
cd(@__DIR__())
pkg"activate ./Intcode"
using Intcode
logfile = "intcode.log"
function do_with_logger(f::Function, level = Logging.Info, logfile=logfile)
    open(logfile, "w+") do io
        with_logger(ConsoleLogger(io, level)) do
            f(io)
        end
        flush(io)
    end
end

##
try
    pkg"test Intcode"
catch e
    @error e
end

##
init_data = readdlm("input.txt", ',', Int) |> vec;

##
do_with_logger() do io
    ic = IntcodeMachine(init_data)
    run_intcode!(ic)
    res = fetch(ic)
    isnothing(res) && @error "TEST program failed"
    @info something(res)
    !isnothing(io) && flush(io)
end

##
include(joinpath("Intcode", "test", "test_cmp_jmp.jl"))

##
do_with_logger(test_cmp, Logging.Debug)

##
do_with_logger(test_jmp, Logging.Debug)

##
do_with_logger(test_larger, Logging.Debug)


using ADatasets
using Test
using ADatasets: sampleindices, subsampletrainvalidation

@testset "test functions for construction of datasets" begin
	@test length(sampleindices([1,2,3,4,5], 0.2)) == 1
	@test length(sampleindices([1,2,3,4,5], 1.2)) == 5
	@test length(sampleindices([1,2,3,4,5], 10)) == 5
	@test length(sampleindices([1,2,3,4,5], 0)) == 0
	@test length(sampleindices([1,2,3,4,5], 1)) == 1
end

@testset "test functions for construction of datasets" begin
	x = (randn(2,10), fill(0,10))
	xtrn, xval = subsampletrainvalidation(x, 0.1,0.8)
	labels = xtrn[2]
	@test length(xtrn[2]) == 8
	@test size(xtrn[1],2) == 8
	@test length(xval[2]) == 2
	@test size(xval[1],2) == 2

	x = (randn(2,20), vcat(fill(0,10),fill(1,10)))
	xtrn, xval = subsampletrainvalidation(x, 0.1,0.8)
	@test length(xtrn[2]) == 9
	@test size(xtrn[1],2) == 9
	@test sum(xtrn[2] .== 1) == 1
	@test length(xval[2]) == 11
	@test size(xval[1],2) == 11
	@test sum(xval[2] .== 1) == 9

	@test length(sampleindices([1,2,3,4,5], 1.2)) == 5
	@test length(sampleindices([1,2,3,4,5], 10)) == 5
	@test length(sampleindices([1,2,3,4,5], 0)) == 0
	@test length(sampleindices([1,2,3,4,5], 1)) == 1
end

@testset "degenerative case" begin
	x = ones(3)
	for y in [[1, 0, 1], [0, 1, 1], [1, 1, 0]]
		fpr, tpr = ADatasets.roccurve(x, y)
		@test all(fpr .== [0.0, 1.0]) && all(tpr .== [0.0, 1.0])
		@test ADatasets.auc(fpr,tpr) == 0.5
	end
end

# Compare the these functions to SciKit Learn functions
using PyCall
@pyimport sklearn.metrics as sm

function compareskauc(labels, ascores)
	pyfpr, pytpr, _ = sm.roc_curve(labels, ascores, drop_intermediate = false)
	pyauc = sm.auc(pyfpr, pytpr)

	fpr, tpr = ADatasets.roccurve(ascores, labels)
	auc = ADatasets.auc(fpr, tpr)

	@test pyauc ≈ auc
end

@testset "Scikit learn auc comparison" begin
	# rand dataset
	for i in 1:10
		counts = rand(1:1000, 2)
		labels = vcat(zeros(counts[1]), ones(counts[2]))
		ascores = rand(size(labels))
		compareskauc(labels, ascores)
	end
end


"""
    couple_models(aero::QuasiSteady, stru::LiftingLineSection, flap::SimpleFlap,
        ctrl::LiftingLineSectionControl)

Create an aerostructural model using a quasi-steady aerodynamics model, a
lifting line aerodynamic model, and a linear, steady-state control surface model.
The existence of this coupling allows the [`QuasiSteady`](@ref) and [`SimpleFlap`](@ref)
to be used with [`LiftingLine`](@ref) and [`LiftingLineFlaps`](@ref). This model
introduces the freestream air density ``\\rho`` as an additional parameter.
"""
function couple_models(aero::QuasiSteady, stru::LiftingLineSection, flap::SimpleFlap,
    ctrl::LiftingLineSectionControl)

    return (aero, stru, flap, ctrl)
end

# --- Traits --- #

# steady
function number_of_additional_parameters(::Type{QuasiSteady{0}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return 1
end

function coupling_inplaceness(::Type{QuasiSteady{0}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return OutOfPlace()
end

function coupling_rate_jacobian_type(::Type{QuasiSteady{0}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Zeros()
end

function coupling_state_jacobian_type(::Type{QuasiSteady{0}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Nonlinear()
end

function coupling_parameter_jacobian_type(::Type{QuasiSteady{0}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Nonlinear()
end

function coupling_time_gradient_type(::Type{QuasiSteady{0}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Zeros()
end

# quasisteady without acceleration terms
function number_of_additional_parameters(::Type{QuasiSteady{1}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return 1
end

function coupling_inplaceness(::Type{QuasiSteady{1}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return OutOfPlace()
end

function coupling_rate_jacobian_type(::Type{QuasiSteady{1}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Zeros()
end

function coupling_state_jacobian_type(::Type{QuasiSteady{1}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Nonlinear()
end

function coupling_parameter_jacobian_type(::Type{QuasiSteady{1}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Nonlinear()
end

function coupling_time_gradient_type(::Type{QuasiSteady{1}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Zeros()
end

# quasisteady with acceleration terms
function number_of_additional_parameters(::Type{QuasiSteady{2}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})
    return 1
end

function coupling_inplaceness(::Type{QuasiSteady{2}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return OutOfPlace()
end

function coupling_rate_jacobian_type(::Type{QuasiSteady{2}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Linear()
end

function coupling_state_jacobian_type(::Type{QuasiSteady{2}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Nonlinear()
end

function coupling_parameter_jacobian_type(::Type{QuasiSteady{2}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Nonlinear()
end

function coupling_time_gradient_type(::Type{QuasiSteady{2}}, ::Type{LiftingLineSection},
    ::Type{SimpleFlap}, ::Type{LiftingLineSectionControl})

    return Zeros()
end

# --- Methods --- #

# steady
function get_coupling_inputs(aero::QuasiSteady{0}, stru::LiftingLineSection,
    flap::SimpleFlap, ctrl::LiftingLineSectionControl, dx, x, p, t)
    # extract rate variables
    dvx, dvy, dvz, dωx, dωy, dωz, dδ = dx
    # extract state variables
    vx, vy, vz, ωx, ωy, ωz, δ = x
    # extract parameters
    a, b, a0, α0, cd0, cm0, cnδ, caδ, cmδ, ρ = p
    # local freestream velocity components
    u, v, ω = liftingline_velocities(vx, vz, ωy)
    # calculate loads
    Na, Aa, Ma = quasisteady0_loads(a, b, ρ, a0, α0, cd0, cm0, u, v)
    # loads due to flap deflections
    Nf, Af, Mf = simpleflap_loads(b, u, ρ, cnδ, caδ, cmδ, δ)
    # loads per unit span
    f = SVector(Aa+Af, 0, Na+Nf)
    m = SVector(0, Ma+Mf, 0)
    # return inputs
    return SVector(f..., m...)
end

# quasisteady without acceleration terms
function get_coupling_inputs(aero::QuasiSteady{1}, stru::LiftingLineSection,
    flap::SimpleFlap, ctrl::LiftingLineSectionControl, dx, x, p, t)
    # extract rate variables
    dvx, dvy, dvz, dωx, dωy, dωz, dδ = dx
    # extract state variables
    vx, vy, vz, ωx, ωy, ωz, δ = x
    # extract parameters
    a, b, a0, α0, cd0, cm0, cnδ, caδ, cmδ, ρ = p
    # local freestream velocity components
    u, v, ω = liftingline_velocities(vx, vz, ωy)
    # calculate aerodynamic loads
    Na, Aa, Ma = quasisteady1_loads(a, b, ρ, a0, α0, cd0, cm0, u, v, ω)
    # loads due to flap deflections
    Nf, Af, Mf = simpleflap_loads(b, u, ρ, cnδ, caδ, cmδ, δ)
    # loads per unit span
    f = SVector(Aa+Af, 0, Na+Nf)
    m = SVector(0, Ma+Mf, 0)
    # return inputs
    return SVector(f..., m...)
end

# quasisteady with acceleration terms
function get_coupling_inputs(aero::QuasiSteady{2}, stru::LiftingLineSection,
    flap::SimpleFlap, ctrl::LiftingLineSectionControl, dx, x, p, t)
    # extract rate variables
    dvx, dvy, dvz, dωx, dωy, dωz, dδ = dx
    # extract state variables
    vx, vy, vz, ωx, ωy, ωz, δ = x
    # extract parameters
    a, b, a0, α0, cd0, cm0, cnδ, caδ, cmδ, ρ = p
    # local freestream velocity components
    u, v, ω = liftingline_velocities(vx, vz, ωy)
    udot, vdot, ωdot = liftingline_accelerations(dvx, dvz, dωy)
    # calculate aerodynamic loads
    Na, Aa, Ma = quasisteady2_loads(a, b, ρ, a0, α0, cd0, cm0, u, v, ω, vdot, ωdot)
    # loads due to flap deflections
    Nf, Af, Mf = simpleflap_loads(b, u, ρ, cnδ, caδ, cmδ, δ)
    # loads per unit span
    f = SVector(Aa+Af, 0, Na+Nf)
    m = SVector(0, Ma+Mf, 0)
    # return inputs
    return SVector(f..., m...)
end

# --- Performance Overloads --- #

# TODO: State rate, state, and parameter jacobians

# --- Convenience Methods --- #

function set_additional_parameters!(padd, aero::QuasiSteady, stru::LiftingLineSection,
    flap::SimpleFlap, ctrl::LiftingLineSectionControl; rho)

    padd[1] = rho

    return padd
end

function separate_additional_parameters(aero::QuasiSteady, stru::LiftingLineSection,
    flap::SimpleFlap, ctrl::LiftingLineSectionControl, padd)

    return (rho = padd[1],)
end


function event_is_enabled(obj::EnablingPredicate, arg0::SpacecraftState, arg1::EventDetector, arg2::jdouble)
    return jcall(obj, "eventIsEnabled", jboolean, (SpacecraftState, EventDetector, jdouble), arg0, arg1, arg2)
end


# Base case to test the distribution network model.
#
# Runs a single-period problem without storage and flexible loads.
# Uses the original CIGRE network with the following changes:
# - added generator cost: all have equal prices, HV grid exchanges cost twice;
# - a fixed 1.0 tap ratio is assigned to the transformer of branch 17;
# - added some candidate branches.

## Import packages and choose a solver

import PowerModels; const _PM = PowerModels
import FlexPlan; const _FP = FlexPlan
import Cbc
optimizer = _FP.optimizer_with_attributes(Cbc.Optimizer, "logLevel"=>0)

## Settings

bus_1_voltage_fixed        = false
gen_7_power_reactive_fixed = false
branch_1_impedance_zeroed  = false
branch_12_rating_zeroed    = false
branch_16_tap_ratio_fixed  = false
load_bus_1_big             = false
load_bus_12_big            = false
load_bus_13_big            = false
load_bus_14_big            = false

result_file = "test/data/output_files/test_dist_nw_model_result.txt"

## Load and preprocess data

data_file = "test/data/cigre_mv_eu/cigre_mv_eu_unit_test.m"
data = _FP.parse_file(data_file)
_FP.add_dimension!(data, :hour, 1)
_FP.add_dimension!(data, :year, 1)
data = _FP.make_multinetwork(data)

## Apply changes for testing purpose

branch = data["nw"]["1"]["branch"]
bus    = data["nw"]["1"]["bus"]
gen    = data["nw"]["1"]["gen"]
load   = data["nw"]["1"]["load"]
if bus_1_voltage_fixed
    bus["1"]["vmin"] = 1.0
    bus["1"]["vmax"] = 1.0
end
if gen_7_power_reactive_fixed
    gen["7"]["qmin"] = 0.0
    gen["7"]["qmax"] = 0.0
end
if branch_1_impedance_zeroed
    branch["1"]["br_r"] = 0.0
    branch["1"]["br_x"] = 0.0
end
if branch_12_rating_zeroed
    branch["12"]["rate_a"] = 0.0
    branch["12"]["rate_b"] = 0.0
    branch["12"]["rate_c"] = 0.0
end
if branch_16_tap_ratio_fixed
    branch["16"]["tm_min"] = 0.95
    branch["16"]["tap"]    = 0.95
    branch["16"]["tm_max"] = 0.95
end
if load_bus_1_big
    load["1"]["pd"] += 10.0
end
if load_bus_12_big
    load["11"]["pd"] += 10.0
end
if load_bus_13_big
    load["12"]["pd"] += 8.0
end
if load_bus_14_big
    load["13"]["pd"] += 1.0
end

## Solve problem

result = _FP.strg_tnep(data, _FP.BFARadPowerModel, optimizer)
# Two-step alternative
#pm = _PM.instantiate_model(data, _FP.BFARadPowerModel, _FP.post_strg_tnep; ref_extensions=[_FP.ref_add_gen!, _FP.ref_add_storage!, _FP.ref_add_ne_storage!, _PM.ref_add_on_off_va_bounds!, _FP.ref_add_ne_branch_allbranches!, _FP.ref_add_frb_branch!, _FP.ref_add_oltc_branch!])
#result = _PM.optimize_model!(pm; optimizer, solution_processors=[_PM.sol_data_model!])
@assert result["termination_status"] ∈ (_PM.OPTIMAL, _PM.LOCALLY_SOLVED) "$(result["optimizer"]) termination status: $(result["termination_status"])"

## Write result

_PM.print_summary(result["solution"]["nw"]["1"])
mkpath(dirname(result_file))
open(result_file, "w") do io
    _PM.summary(io, result["solution"]["nw"]["1"])
end

## Perform other unit tests

gen = result["solution"]["nw"]["1"]["gen"]

pg = sum(g["pg"] for g in values(gen))
pd = sum(l["pd"] for l in values(load))
@assert pg ≈ pd "generated active power ($pg) does not match demand ($pd)"

qg = sum(g["qg"] for g in values(gen))
qd = sum(l["qd"] for l in values(load))
@assert qg ≈ qd "generated reactive power ($qg) does not match demand ($qd)"


using .Turing, Random, ExprTools, Distributions, Test
using DynamicPPL: DynamicPPL, vsym, vinds, @varname, VarInfo, VarName, build_model_info, namedtuple, build_model_info

dir = splitdir(splitdir(pathof(DynamicPPL))[1])[1]
include(dir*"/test/test_utils/AllUtils.jl")

Random.seed!(129)

priors = 0 # See "new grammar" test.

macro custom(expr)
    (Meta.isexpr(expr, :call, 3) && expr.args[1] === :~) ||
        error("incorrect macro usage")
    quote
        $(esc(expr.args[2])) = 0.0
    end
end

@testset "build_model_info" begin
    @testset "single arg" begin
        expr = :(model(obs) = begin
            obs
        end)

        result = build_model_info(expr)
        @test result[:args] == [:obs]
        @test result[:arg_syms] == [:obs]
        @test result[:name] == :model
        result
    end

    @testset "default arg missing" begin
        expr = :(model(x = missing) = begin
            x
        end)

        result = build_model_info(expr)

        @test result[:args] == [:($(Expr(:kw, :x, :missing)))]
        @test result[:arg_syms] == [:x]
        args_nt = result[:args_nt]
        @test args_nt.args[2] ==  :(NamedTuple{(:x,), Tuple{Core.Typeof(x)}})
        @test args_nt.args[3] ==  :((x,))
        defaults_nt = result[:defaults_nt]
        @test defaults_nt.args[2] ==  :(NamedTuple{(:x,), Tuple{Core.Typeof(missing)}})
        @test defaults_nt.args[3] ==  :((missing,))
    end

    @testset "default arg with type annotation" begin
        expr1 = :(model(x::Int = 2) = begin
            x
        end)

        result1 = build_model_info(expr1)

        @test result1[:args] == [:($(Expr(:kw, :(x::Int), 2)))]
        @test result1[:arg_syms] == [:x]
        args_nt = result1[:args_nt]
        @test args_nt.args[2] ==  :(NamedTuple{(:x,), Tuple{Core.Typeof(x)}})
        @test args_nt.args[3] ==  :((x,))
        defaults_nt = result1[:defaults_nt]
        @test defaults_nt.args[2] ==  :(NamedTuple{(:x,), Tuple{Core.Typeof(2)}})
        @test defaults_nt.args[3] ==  :((2,))

        expr2 = :(model(x::Vector{Float64} = [1,2,3]) = begin
            x
        end)

        result2 = build_model_info(expr2)
        @test result2[:arg_syms] == [:x]
    end

    @testset "default arg type" begin
        expr1 = :(model(::Type{T}=Float64) where {T <: Float64} = begin
            T
        end)

        result1 = build_model_info(expr1)
        @test result1[:args] == [:($(Expr(:kw, :(T::Type{<:Float64}), :Float64)))]
        @test result1[:arg_syms] == [:T]

        expr2 = :(model(::Type{T}=Float64) where {T} = begin
            T
        end)

        result2 = build_model_info(expr2)
        @test result2[:args] == [:($(Expr(:kw, :(T::Type{<:Any}), :Float64)))]
        @test result2[:arg_syms] == [:T]

        #TODO support t::Type{T}
        expr3 = :(model(t::Type{T}=Float64) where {T <: Float64} = begin
            T
        end)

        @test_throws ArgumentError build_model_info(expr3)

        #Invalid expression
        expr4 = :(model(::Type{T}=Float64) where {X} = begin
            T
        end)
        @test_throws ArgumentError build_model_info(expr4)
    end
end

@testset "compiler.jl" begin
    @testset "assume" begin
        @model test_assume() = begin
            x ~ Bernoulli(1)
            y ~ Bernoulli(x / 2)
            x, y
        end

        smc = SMC()
        pg = PG(10)

        res1 = sample(test_assume(), smc, 1000)
        res2 = sample(test_assume(), pg, 1000)

        check_numerical(res1, [:y], [0.5], atol=0.1)
        check_numerical(res2, [:y], [0.5], atol=0.1)

        # Check that all xs are 1.
        @test all(isone, res1[:x].value)
        @test all(isone, res2[:x].value)
    end
    @testset "beta binomial" begin
        prior = Beta(2,2)
        obs = [0,1,0,1,1,1,1,1,1,1]
        exact = Beta(prior.α + sum(obs), prior.β + length(obs) - sum(obs))
        meanp = exact.α / (exact.α + exact.β)

        @model testbb(obs) = begin
            p ~ Beta(2,2)
            x ~ Bernoulli(p)
            for i = 1:length(obs)
                obs[i] ~ Bernoulli(p)
            end
            p, x
        end

        smc = SMC()
        pg = PG(10)
        gibbs = Gibbs(HMC(0.2, 3, :p), PG(10, :x))

        chn_s = sample(testbb(obs), smc, 1000)
        chn_p = sample(testbb(obs), pg, 2000)
        chn_g = sample(testbb(obs), gibbs, 1500)

        check_numerical(chn_s, [:p], [meanp], atol=0.05)
        check_numerical(chn_p, [:x], [meanp], atol=0.1)
        check_numerical(chn_g, [:x], [meanp], atol=0.1)
    end
    @testset "forbid global" begin
        xs = [1.5 2.0]
        # xx = 1

        @model fggibbstest(xs) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))
            # xx ~ Normal(m, sqrt(s)) # this is illegal

            for i = 1:length(xs)
                xs[i] ~ Normal(m, sqrt(s))
                # for xx in xs
                # xx ~ Normal(m, sqrt(s))
            end
            s, m
        end

        gibbs = Gibbs(PG(10, :s), HMC(0.4, 8, :m))
        chain = sample(fggibbstest(xs), gibbs, 2);
    end
    @testset "model macro" begin
        @model testmodel_comp(x, y) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x ~ Normal(m, sqrt(s))
            y ~ Normal(m, sqrt(s))

            return x, y
        end
        testmodel_comp(1.0, 1.2)

        # check if drawing from the prior works
        @model testmodel0(x = missing) = begin
            x ~ Normal()
            return x
        end
        f0_mm = testmodel0()
        @test mean(f0_mm() for _ in 1:1000) ≈ 0. atol=0.1

        # Test #544
        @model testmodel0(x = missing) = begin
            if x === missing
                x = Vector{Float64}(undef, 2)
            end
            x[1] ~ Normal()
            x[2] ~ Normal()
            return x
        end
        f0_mm = testmodel0()
        @test all(x -> isapprox(x, 0; atol = 0.1), mean(f0_mm() for _ in 1:1000))

        @model testmodel01(x = missing) = begin
            x ~ Bernoulli(0.5)
            return x
        end
        f01_mm = testmodel01()
        @test mean(f01_mm() for _ in 1:1000) ≈ 0.5 atol=0.1

        # test if we get the correct return values
        @model testmodel1(x1, x2) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x1 ~ Normal(m, sqrt(s))
            x2 ~ Normal(m, sqrt(s))

            return x1, x2
        end
        f1_mm = testmodel1(1., 10.)
        @test f1_mm() == (1, 10)
        f1_mm = testmodel1(x1=1., x2=10.)
        @test f1_mm() == (1, 10)

        @info "Testing the compiler's ability to catch bad models..."

        # Test for assertions in observe statements.
        @model brokentestmodel_observe1(x1, x2) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x1 ~ Normal(m, sqrt(s))
            x2 ~ x1 + 2

            return x1, x2
        end

        btest = brokentestmodel_observe1(1., 2.)
        @test_throws ArgumentError btest()

        @model brokentestmodel_observe2(x) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x = Vector{Float64}(undef, 2)
            x ~ [Normal(m, sqrt(s)), 2.0]

            return x
        end

        btest = brokentestmodel_observe2([1., 2.])
        @test_throws ArgumentError btest()

        # Test for assertions in assume statements.
        @model brokentestmodel_assume1() = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x1 ~ Normal(m, sqrt(s))
            x2 ~ x1 + 2

            return x1, x2
        end

        btest = brokentestmodel_assume1()
        @test_throws ArgumentError btest()

        @model brokentestmodel_assume2() = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x = Vector{Float64}(undef, 2)
            x ~ [Normal(m, sqrt(s)), 2.0]

            return x
        end

        btest = brokentestmodel_assume2()
        @test_throws ArgumentError btest()

        # Test missing input arguments
        @model testmodel(x) = begin
            x ~ Bernoulli(0.5)
            return x
        end
        @test_throws UndefKeywordError testmodel()

        # Test missing initialization for vector observation turned parameter
        @model testmodel(x) = begin
            x[1] ~ Bernoulli(0.5)
            return x
        end
        @test_throws MethodError testmodel(missing)()

        # Test use of internal names
        @model testmodel(x) = begin
            x[1] ~  Bernoulli(0.5)
            global varinfo_ = _varinfo
            global sampler_ = _sampler
            global model_ = _model
            global context_ = _context
            global lp = getlogp(_varinfo)
            return x
        end
        model = testmodel([1.0])
        varinfo = DynamicPPL.VarInfo(model)
        model(varinfo)
        @test getlogp(varinfo) == lp
        @test varinfo_ isa AbstractVarInfo
        @test model_ === model
        @test sampler_ === SampleFromPrior()
        @test context_ === DefaultContext()

        # disable warnings
        @model testmodel(x) = begin
            x[1] ~  Bernoulli(0.5)
            global varinfo_ = _varinfo
            global sampler_ = _sampler
            global model_ = _model
            global context_ = _context
            global lp = getlogp(_varinfo)
            return x
        end false
        lpold = lp
        model = testmodel([1.0])
        varinfo = DynamicPPL.VarInfo(model)
        model(varinfo)
        @test getlogp(varinfo) == lp == lpold

        # test DPPL#61
        @model testmodel(z) = begin
            m ~ Normal()
            z[1:end] ~ MvNormal(fill(m, length(z)), 1.0)
            return m
        end
        model = testmodel(rand(10))
        @test all(z -> isapprox(z, 0; atol = 0.2), mean(model() for _ in 1:1000))
    end
    @testset "nested model" begin
        function makemodel(p)
            @model testmodel(x) = begin
                x[1] ~ Bernoulli(p)
                global lp = getlogp(_varinfo)
                return x
            end
            return testmodel
        end
        model = makemodel(0.5)([1.0])
        varinfo = DynamicPPL.VarInfo(model)
        model(varinfo)
        @test getlogp(varinfo) == lp
    end
    @testset "new grammar" begin
        x = Float64[1 2]

        @model gauss(x) = begin
            priors = TArray{Float64}(2)
            priors[1] ~ InverseGamma(2,3)         # s
            priors[2] ~ Normal(0, sqrt(priors[1])) # m
            for i in 1:length(x)
                x[i] ~ Normal(priors[2], sqrt(priors[1]))
            end
            priors
        end

        chain = sample(gauss(x), PG(10), 10)
        chain = sample(gauss(x), SMC(), 10)

        @model gauss2(x, ::Type{TV}=Vector{Float64}) where {TV} = begin
            priors = TV(undef, 2)
            priors[1] ~ InverseGamma(2,3)         # s
            priors[2] ~ Normal(0, sqrt(priors[1])) # m
            for i in 1:length(x)
                x[i] ~ Normal(priors[2], sqrt(priors[1]))
            end
            priors
        end

        chain = sample(gauss2(x), PG(10), 10)
        chain = sample(gauss2(x=x, TV=Vector{Float64}), PG(10), 10)
        chain = sample(gauss2(x), SMC(), 10)
        chain = sample(gauss2(x=x, TV=Vector{Float64}), SMC(), 10)
    end
    @testset "new interface" begin
        obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1]

        @model newinterface(obs) = begin
          p ~ Beta(2,2)
          for i = 1:length(obs)
            obs[i] ~ Bernoulli(p)
          end
          p
        end

        chain = sample(
            newinterface(obs),
            HMC{Turing.ForwardDiffAD{2}}(0.75, 3, :p, :x),
            100)
    end
    @testset "no return" begin
        @model noreturn(x) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0, sqrt(s))
            for i in 1:length(x)
                x[i] ~ Normal(m, sqrt(s))
            end
        end

        chain = sample(noreturn([1.5 2.0]), HMC(0.15, 6), 1000)
        check_numerical(chain, [:s, :m], [49/24, 7/6])
    end
    @testset "observe" begin
        @model test() = begin
          z ~ Normal(0,1)
          x ~ Bernoulli(1)
          1 ~ Bernoulli(x / 2)
          0 ~ Bernoulli(x / 2)
          x
        end

        is  = IS()
        smc = SMC()
        pg  = PG(10)

        res_is = sample(test(), is, 10000)
        res_smc = sample(test(), smc, 1000)
        res_pg = sample(test(), pg, 100)

        @test all(isone, res_is[:x].value)
        @test res_is.logevidence ≈ 2 * log(0.5)

        @test all(isone, res_smc[:x].value)
        @test res_smc.logevidence ≈ 2 * log(0.5)

        @test all(isone, res_pg[:x].value)
    end
    @testset "sample" begin
        alg = Gibbs(HMC(0.2, 3, :m), PG(10, :s))
        chn = sample(gdemo_default, alg, 1000);
    end
    @testset "vectorization @." begin
        @model vdemo1(x) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0, sqrt(s))
            @. x ~ Normal(m, sqrt(s))
            return s, m
        end

        alg = HMC(0.01, 5)
        x = randn(100)
        res = sample(vdemo1(x), alg, 250)

        @model vdemo1b(x) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0, sqrt(s))
            @. x ~ Normal(m, $(sqrt(s)))
            return s, m
        end

        res = sample(vdemo1b(x), alg, 250)

        D = 2
        @model vdemo2(x) = begin
            μ ~ MvNormal(zeros(D), ones(D))
            @. x ~ $(MvNormal(μ, ones(D)))
        end

        alg = HMC(0.01, 5)
        res = sample(vdemo2(randn(D, 100)), alg, 250)

        # Vector assumptions
        N = 10
        setchunksize(N)
        alg = HMC(0.2, 4)

        @model vdemo3() = begin
            x = Vector{Real}(undef, N)
            for i = 1:N
                x[i] ~ Normal(0, sqrt(4))
            end
        end

        t_loop = @elapsed res = sample(vdemo3(), alg, 1000)

        # Test for vectorize UnivariateDistribution
        @model vdemo4() = begin
          x = Vector{Real}(undef, N)
          @. x ~ Normal(0, 2)
        end

        t_vec = @elapsed res = sample(vdemo4(), alg, 1000)

        @model vdemo5() = begin
            x ~ MvNormal(zeros(N), 2 * ones(N))
        end

        t_mv = @elapsed res = sample(vdemo5(), alg, 1000)

        println("Time for")
        println("  Loop : ", t_loop)
        println("  Vec  : ", t_vec)
        println("  Mv   : ", t_mv)

        # Transformed test
        @model vdemo6() = begin
            x = Vector{Real}(undef, N)
            @. x ~ InverseGamma(2, 3)
        end

        sample(vdemo6(), alg, 1000)

        N = 3
        @model vdemo7() = begin
            x = Array{Real}(undef, N, N)
            @. x ~ [InverseGamma(2, 3) for i in 1:N]
        end

        sample(vdemo7(), alg, 1000)
    end

    # Notation is ugly since `x .~ Normal(μ, σ)` cannot be parsed in Julia 1.0
    @testset "vectorization .~" begin
        @model vdemo1(x) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0, sqrt(s))
            (.~)(x, Normal(m, sqrt(s)))
            return s, m
        end

        alg = HMC(0.01, 5)
        x = randn(100)
        res = sample(vdemo1(x), alg, 250)

        D = 2
        @model vdemo2(x) = begin
            μ ~ MvNormal(zeros(D), ones(D))
            (.~)(x, MvNormal(μ, ones(D)))
        end

        alg = HMC(0.01, 5)
        res = sample(vdemo2(randn(D,100)), alg, 250)

        # Vector assumptions
        N = 10
        setchunksize(N)
        alg = HMC(0.2, 4)

        @model vdemo3() = begin
            x = Vector{Real}(undef, N)
            for i = 1:N
                x[i] ~ Normal(0, sqrt(4))
            end
        end

        t_loop = @elapsed res = sample(vdemo3(), alg, 1000)

        # Test for vectorize UnivariateDistribution
        @model vdemo4() = begin
            x = Vector{Real}(undef, N)
            (.~)(x, Normal(0, 2))
        end

        t_vec = @elapsed res = sample(vdemo4(), alg, 1000)

        @model vdemo5() = begin
            x ~ MvNormal(zeros(N), 2 * ones(N))
        end

        t_mv = @elapsed res = sample(vdemo5(), alg, 1000)

        println("Time for")
        println("  Loop : ", t_loop)
        println("  Vec  : ", t_vec)
        println("  Mv   : ", t_mv)

        # Transformed test
        @model vdemo6() = begin
            x = Vector{Real}(undef, N)
            (.~)(x, InverseGamma(2, 3))
        end

        sample(vdemo6(), alg, 1000)

        @model vdemo7() = begin
            x = Array{Real}(undef, N, N)
            (.~)(x, [InverseGamma(2, 3) for i in 1:N])
        end

        sample(vdemo7(), alg, 1000)
    end

    @testset "Type parameters" begin
        N = 10
        setchunksize(N)
        alg = HMC(0.01, 5)
        x = randn(1000)
        @model vdemo1(::Type{T}=Float64) where {T} = begin
            x = Vector{T}(undef, N)
            for i = 1:N
                x[i] ~ Normal(0, sqrt(4))
            end
        end

        t_loop = @elapsed res = sample(vdemo1(), alg, 250)
        t_loop = @elapsed res = sample(vdemo1(Float64), alg, 250)
        t_loop = @elapsed res = sample(vdemo1(T=Float64), alg, 250)

        @model vdemo2(::Type{T}=Float64) where {T <: Real} = begin
            x = Vector{T}(undef, N)
            @. x ~ Normal(0, 2)
        end

        t_vec = @elapsed res = sample(vdemo2(), alg, 250)
        t_vec = @elapsed res = sample(vdemo2(Float64), alg, 250)
        t_vec = @elapsed res = sample(vdemo2(T=Float64), alg, 250)

        @model vdemo3(::Type{TV}=Vector{Float64}) where {TV <: AbstractVector} = begin
            x = TV(undef, N)
            @. x ~ InverseGamma(2, 3)
        end

        sample(vdemo3(), alg, 250)
        sample(vdemo3(Vector{Float64}), alg, 250)
        sample(vdemo3(TV=Vector{Float64}), alg, 250)
    end
    @testset "var name splitting" begin
        var_expr = :(x)
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(())

        var_expr = :(x[1,1][2,3])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(((1, 1), (2, 3)))

        var_expr = :(x[:,1][2,:])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(((:, 1), (2, :)))

        var_expr = :(x[2:3,1][2,1:2])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(((2:3, 1), (2, 1:2)))

        var_expr = :(x[2:3,2:3][[1,2],[1,2]])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(((2:3, 2:3), ([1, 2], [1, 2])))

        var_expr = :(x[end])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :((($lastindex(x),),))

        var_expr = :(x[1, end])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(((1, $lastindex(x, 2)),))
    end
    @testset "user-defined variable name" begin
        @model f1() = begin
            x ~ NamedDist(Normal(), :y)
        end
        @model f2() = begin
            x ~ NamedDist(Normal(), @varname(y[2][:,1]))
        end
        @model f3() = begin
            x ~ NamedDist(Normal(), @varname(y[1]))
        end
        vi1 = VarInfo(f1())
        vi2 = VarInfo(f2())
        vi3 = VarInfo(f3())
        @test haskey(vi1.metadata, :y)
        @test vi1.metadata.y.vns[1] == VarName(:y)
        @test haskey(vi2.metadata, :y)
        @test vi2.metadata.y.vns[1] == VarName(:y, ((2,), (Colon(), 1)))
        @test haskey(vi3.metadata, :y)
        @test vi3.metadata.y.vns[1] == VarName(:y, ((1,),))
    end
    @testset "custom tilde" begin
        @model demo() = begin
            $(@custom m ~ Normal())
            return m
        end
        model = demo()
        @test all(iszero(model()) for _ in 1:1000)
    end
    @testset "docstring" begin
        "This is a test"
        @model function demo(x)
            m ~ Normal()
            x ~ Normal(m, 1)
        end

        s = @doc(demo)
        @test string(s) == "This is a test\n"

        # Verify that adding docstring didn't completely break execution of model
        m = demo(0.)
        @test m() isa Float64
    end
end


"""DNN Feature decoding - feature prediction script"""

using Statistics

using BrainDecoder
using FastL2LiR
using HDF5
using Printf
using Glob
using Debugger

# Settings ###################################################################

# Brain data
brain_dir = "/home/kiss/data/fmri_shared/datasets/Deeprecon/fmriprep"
subjects = Dict(
    "AM" => "AM_ImageNetTest_volume_native.h5",
    # "ES" => "ES_ImageNetTest_volume_native.h5",
    # "FC" => "FC_ImageNetTest_volume_native.h5",
    # "JK" => "JK_ImageNetTest_volume_native.h5",
    # "JP" => "JP_ImageNetTest_volume_native.h5",
    # "KS" => "KS_ImageNetTest_volume_native.h5",
    # "TH" => "TH_ImageNetTest_volume_native.h5",
)

label_key = "stimulus_name"

rois = Dict(
    "VC"  => "ROI_VC",
    # "LVC" => "ROI_LVC",
    # "HVC" => "ROI_HVC",
    # "V1"  => "ROI_V1",
    # "V2"  => "ROI_V2",
    # "V3"  => "ROI_V3",
    # "V4"  => "ROI_hV4",
    # "LOC" => "ROI_LOC",
    # "FFA" => "ROI_FFA",
    # "PPA" => "ROI_PPA",
)

# DNN features
features_dir = "/home/nu/data/contents_shared/ImageNetTest/derivatives/features/default"
network = "caffe/bvlc_alexnet"
layers = [
    "conv1", #"relu1", "pool1",
    "conv2", #"relu2", "pool2",
    "conv3", #"relu3",
    "conv4", #"relu4",
    "conv5", #"relu5", "pool5",
    "fc6", #"relu6",
    "fc7", #"relu7",
    "fc8",
]

chunk_axis = 0
# Note that Julia's indexing is one-based.

# Feature decoders
models_dir_root = "./results/feature_decoders/deeprecon_fmriprep_rep5_500voxel_allunits_fastl2lir_alpha100"

# Results directory
#results_dir_root = "/home/nu/data/contents_shared/ImageNetTraining/derivatives/feature_decoders/deeprecon_fmriprep_rep5_500voxel_allunits_fastl2lir_alpha100"
results_dir_root = "./results/decoded_features/deeprecon_fmriprep_rep5_500voxel_allunits_fastl2lir_alpha100"

# Functions ##################################################################

function predict_features(model_dir, x; x_mean=[], x_norm=[], y_mean=[], y_norm=[], chunk_axis=0)

    # Normalize X
    x = (x .- x_mean) ./ x_norm

    model = load_model(model_dir, chunk_axis=chunk_axis)

    @show size(model.W)
    @show size(model.b)

    y = predict(model, x)

    # Denormalize Y
    y = y .* y_norm .+ y_mean

    return y
end

# Main #######################################################################

# Load data ------------------------------------------------------------------
features = Features(joinpath(features_dir, network))

# Analysis loop --------------------------------------------------------------
println("----------------------------------------")
println("Analysis loop")

for lay in layers

    # Load DNN features
    println("Loading features: " * lay)
    y = features.get_features(lay)  # TODO: speed-up
    y_labels = features.labels

    y = Float32.(y)

    @show size(y)
    @show size(y_labels)

    for sub in keys(subjects)

        # Load fMRI data
        brain_data_file = joinpath(brain_dir, subjects[sub])
        println("Loading " * brain_data_file)
        brain_data = BData(brain_data_file)

        for roi in keys(rois)
            println("Layer:   " * lay)
            println("Subject: " * sub)
            println("ROI:     " * roi)

            # Feature prediction
            save_dir = joinpath(results_dir_root, "decoded_features", network, lay, sub, roi)
            if !isdir(save_dir)
                mkpath(save_dir)
            end
            # TODO: check existing files

            model_dir = joinpath(models_dir_root, network, lay, sub, roi, "model")

            x = brain_data.select(rois[roi])
            x_labels = brain_data.get_labels(label_key)

            x = Float32.(x)

            @show size(x)
            @show size(x_labels)

            # Average fMRI samples
            x_labels_unique = unique(x_labels)
            x_ave = cat([mean(x[findall(x -> x == xl, x_labels), :], dims=1)
                         for xl in unique(x_labels)]..., dims=1)

            @show size(x_ave)
            @show size(x_labels_unique)

            # Preprocessing parameters
            x_mean = load_array(joinpath(model_dir, "x_mean.mat"), "x_mean")
            x_norm = load_array(joinpath(model_dir, "x_norm.mat"), "x_norm")
            y_mean = load_array(joinpath(model_dir, "y_mean.mat"), "y_mean")
            y_norm = load_array(joinpath(model_dir, "y_norm.mat"), "y_norm")

            # Get sample index of Y
            y_index = [findall(x -> x == xl, y_labels)[1] for xl in x_labels_unique]

            # Call training routine
            @time y_pred = predict_features(
                model_dir, x_ave,
                x_mean=x_mean, x_norm=x_norm,
                y_mean=y_mean, y_norm=y_norm,
                chunk_axis=chunk_axis
            )

            @show size(y_pred)

            # Save decoded features
            y_idx = axes(y_pred)
            for (i, xl) in enumerate(x_labels_unique)
                save_file = joinpath(save_dir, xl * ".mat")
                save_array(save_file, "feat", y_pred[[i], y_idx[2:end]...])
            end
        end
    end
end

println("All done")


using ESGIWeatherLogistics2021: ESGIWeatherLogistics2021, Data
using Plots
using Dates
using DataFrames

const obs = Data.get_obs()
const obsdaily = ESGIWeatherLogistics2021.aggregate_daily_obs(obs)
const prices = Data.get_prices()

function plot_year(obsdaily, year; area, crop)
    obsidx = (obsdaily.area .== area) .& (Dates.year.(obsdaily.date) .== year)
    pricesidx = (prices.crop .== crop) .& (Dates.year.(prices.date) .== year)
    plot(obsdaily.date[obsidx], obsdaily.temperature_mean[obsidx])
    plot!(obsdaily.date[obsidx], obsdaily.temperature_minimum[obsidx])
    plot!(obsdaily.date[obsidx], obsdaily.temperature_maximum[obsidx])
    plot!(prices.date[pricesidx], 10 .* prices.price[pricesidx])
end


abstract type AbstractArrayStream{S,T,N,I} <: AbstractArray{T,N} end

# IO interactions
stream(s::AbstractArrayStream) = error("stream is not defined for $(typeof(s))")
Base.seek(s::AbstractArrayeStream, n::Integer) = seek(stream(s), n)
Base.position(s::AbstractArrayStream) = position(stream(s))
Base.skip(s::AbstractArrayStream, n::Integer) = skip(stream(s), n)
Base.close(s::AbstractArrayStream) = close(stream(s))
Base.eof(s::AbstractArrayStream) = eof(stream(s))
Base.isopen(s::AbstractArrayStream) = isopen(stream(s))

# Field accessors
ownstream(s::AbstractArrayStream) = getfield(s, :ownstream)
streamindices(s::AbstractArrayStream) = getfield(s, :streamindices)
needswap(s::AbstractArrayStream) = getfield(s :needswap)

Base.size(s::AbstractArrayStream{S,T,N}) where {S,T,N} = (S.parameters...,)
Base.size(s::AbstractArrayStream{S,T,N}, i::Int) where {S,T,N} = S.parameters[i]
Base.ndims(s::AbstractArrayStream{S,T,N}) where {S,T,N} = N
Base.length(s::AbstractArrayStream{S,T,N}) where {S,T,N} = prod(S.parameters)
Base.eltype(s::AbstractArrayStream{S,T,N}) where {S,T,N} = T

Base.IndexStyle(::Type{AbstractArrayStream}) = IndexLinear

Base.firstindex(s::AbstractArrayStream) = streamindices(s)[1]
Base.lastindex(s::AbstractArrayStream) = streamindices(s)[end]

"""
    getindex(AbstractArrayStream, I)

While an `ArrayStream` may be indexed, this will only result in performance gains
if the indices represent a significantly smaller portion of the original `ArrayStream`.
This is because an indexed ArrayStream has to perform a `seek(ArrayStream, index)`
step at every index before reading. This is in contrast to reading in a single
chunk of memory, which is considerably faster.
"""
Base.getindex(s::AbstractArrayStream{S,T,N,IndexLinear}, i::Int) where {S,T,N} =
    ArrayStreams(stream(s), streamindices(s)[i], ownstream(s), needswap(s))

function Base.getindex(s::AbstractArrayStream{S,T,N,L}, I::Vararg{Int, N}) where {S,T,N,L}
    si = s.streamindices[LinearIndices(size(s))[I...]]
    ArrayStreams{Tuple{1},T,1,1}(stream(s), s.streamindices[LinearIndices(size(s))[I...]], ownstream(s), needswap(s))
end

# leverage SubArrays which determine whether fast indexing is possible
function Base.getindex(s::AbstractArrayStream{S,T,N,IndexLinear}, I...) where {S,T,N}
    li = view(LinearIndices(s), I...)
    # fast linear indexing, so just need first and last
    if typeof(li).parameters[5]
        ArrayStream{Tuple{size(li)...},T,ndims(li)}(stream(s), streamindices(s)[first(li):last(li)], ownstream(s), needswap(s))
    else
        ArrayStream{Tuple{size(li)...},T,ndims(li)}(stream(s), streamindices(s)[li...], ownstream(s), needswap(s))
    end
end

Base.setindex!(s::AbstractArrayStream, I) = error("setindex is not defined for $(typeof(s)). Use write to change stream content.")

function Base.show(io::IO, s::AbstractArrayStream)
    println(summary(s))
    println("ownstream: ", ownstream(s))
    println("needswap: ", needswap(s))
    println("streamindices: \n", streamindices(s))
end


q1 = nominal_configuration(quadruped)
v1 = zeros(quadruped.nq)

h = 0.01
T = 100

s = Simulator(quadruped, T, h=h)
simulate!(s, q1, v1)

vis = Visualizer()
open(vis)
visualize!(vis, s)


# This code is part of QuantumCircuits.
#
# (C) Copyright Rafał Pracht 2022.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"Quantum machine learning"
module QML

include("Optimization.jl")
include("CircuitOptimization.jl")

using QuantumCircuits.QCircuits.Gates: ParameterT
using QuantumCircuits.QCircuits.Circuit: getparameters, getRandParameters
using QuantumCircuits.QCircuits.QBase: setparameters!, bindparameters!
using QuantumCircuits.QML.CircuitOptimization

export ParameterT, getparameters, getRandParameters, setparameters!,
       bindparameters!, findparam

end # module


# calculates colsums of a large matrix stored in HDFS
# julia -p 3
# include("test_colsum.jl")

using Elly
using HadoopBlocks

const INP = "hdfs://tan@localhost:9000/colsuminp.csv"

if myid() == 1
    function gendata(M, N)
        open(HDFSFile(INP), "w") do f
            for m in 1:M
                write(f, join(map(string, rand(N)), ","), "\n")
            end
        end
    end
    gendata(10^5, 5)
end

# split script here and uncomment below lines to run data generation and computation separately
#using Elly
#using HadoopBlocks
#const INP = "hdfs://tan@localhost:9000/colsuminp.csv"

function findrow(r::HdfsBlockReader, iter_status)
    rec = HadoopBlocks.find_rec(r, iter_status, Vector)
    #HadoopBlocks.logmsg("findrow found rec:$rec")
    rec
end

function maprow(rec)
    #HadoopBlocks.logmsg("maprow called with rec:$rec")
    [tuple([parse(Float64, x) for x in rec]...)]
end

function collectrow(results, rec)
    #HadoopBlocks.logmsg("collectrow called with results:$results rec:$rec")
    isempty(rec) && (return results)
    (results == nothing) && (return rec)
    tuple([results[x]+rec[x] for x in 1:length(results)]...)
end

function reducerow(reduced, results...)
    #HadoopBlocks.logmsg("reducerow called with reduced:$reduced results:$results")
    (nothing == reduced) && (reduced = zeros(Float64, length(results[1])))
    for res in results
        (nothing == res) && continue
        #HadoopBlocks.logmsg("reducerow res:$res")
        #HadoopBlocks.logmsg("reducerow res:$([res...])")
        for x in 1:length(res)
            reduced[x] += res[x]
        end
    end
    #HadoopBlocks.logmsg("reducerow returning reduced:$reduced")
    reduced
end

function wait_results(j_mon)
    loopstatus = true
    while(loopstatus)
        sleep(5)
        jstatus,jstatusinfo = status(j_mon,true)
        ((jstatus == "error") || (jstatus == "complete")) && (loopstatus = false)
        (jstatus == "running") && println("$(j_mon): $(jstatusinfo)% complete...")
    end
    wait(j_mon)
    println("time taken (total time, wait time, run time): $(times(j_mon))")
    println("")
end

j = dmapreduce(MRHdfsFileInput([INP], findrow), maprow, collectrow, reducerow)
wait_results(j)
println(results(j))


module SpikeNetOpt

    using PyCall
    using OrderedCollections
    using LinearAlgebra
    using SpikeSynchrony
    using SpikingNeuralNetworks
    SNN = SpikingNeuralNetworks
    SNN.@load_units
    using Evolutionary#, Test, Random
    import DataStructures
    using JLD
    #using Reexport
    using UnicodePlots
    using Statistics
    using JLD
    using Distributed
    using SharedArrays
    using Plots
    using Evolutionary
    using Distributions
    using LightGraphs
    using Metaheuristics
    using ClearStacktrace
    using SharedArrays
    #using WaspNet
    #using SpikingNN
    using SignalAnalysis
    #using Reexport
    #using Spikes
    # Export Function names and struct names found in file utils
    #=
    export get_trains
    export get_vm
    export sim_net_darsnack

    function get_trains end
    function get_vm end
    function checkmodel end
    function make_net_from_graph_structure end
    function make_net end
    function spike_train_difference end
    function get_spikes end
    function get_trains end
    function sim_net_darsnack end
    =#
    include("utils.jl")
    #include("current_search.jl")
    include("sdo_network.jl")

    #include("sciunit.jl")

end


fname = :ARCqKdense
c = Combi(hessian_dense,PDataK,solve_modelKARC,preprocessKARC,decreaseKARC,TparamsKARC())
include("Template.jl")


"""
    getTRMMlsm(
        geo :: GeoRegion = GeoRegion("GLB");
	    sroot :: AbstractString = homedir()
    ) -> nothing

Downloads the TRMM-TMPA Land-Sea Mask for a GeoRegion specified by `geo` and saves it into `sroot`

Arguments
=========
- `geo` : a `GeoRegion` (see [GeoRegions.jl](https://github.com/JuliaClimate/GeoRegions.jl)) that sets the geographic bounds of the data array in lon-lat

Keyword Arguments
=================
- `sroot` : the directory to save the TRMM-TMPA Land-Sea Mask, default is the home directoy called by `homedir()`
"""
function getTRMMlsm(
	geo   :: GeoRegion = GeoRegion("GLB");
    sroot :: AbstractString = homedir(),
)

    @info "$(modulelog()) - Retrieving Land-Sea Mask for the TRMM-TMPA dataset"

	if geo.regID == "GLB"
		@info "$(modulelog()) - Global dataset request has been detected, switching to the TRMM LandSea Mask GeoRegion"
		addNPDGeoRegions(); geo = GeoRegion("TRMMLSM")
	else
		isinGeoRegion(geo,GeoRegion("TRMMLSM"))
	end

    lon,lat = trmmlonlat(full=true); nlon = length(lon); nlat = length(lat)
	ginfo = RegionGrid(geo,lon,lat)

	@info "$(modulelog()) - Preallocating temporary array for extraction of TRMM-TMPA Land-Sea Mask data for the $(geo.name) GeoRegion from the original global gridded dataset"
	glon = ginfo.glon; nglon = length(glon); iglon = ginfo.ilon
	glat = ginfo.glat; nglat = length(glat); iglat = ginfo.ilat
	tmp  = zeros(Float32,nlat,nlon)
	var  = zeros(Float32,nglon,nglat)

	@info "$(modulelog()) - Retrieving the original TRMM-TMPA Land-Sea Mask data from NASA's EOSDIS OPeNDAP servers"
	hroot = "https://gpm1.gesdisc.eosdis.nasa.gov/opendap/AUXILIARY"
    npdnc = "TRMM_TMPA_LandSeaMask.2/TRMM_TMPA_LandSeaMask.2.nc4"
    npdds = NCDataset(joinpath(hroot,npdnc))
	NCDatasets.load!(npdds["landseamask"].var,tmp,:,:)
	close(npdds)

	@debug "$(modulelog()) - Extraction of data from temporary array for the $(geo.name) GeoRegion"

	extractregionlsm!(var,tmp,ginfo)

	saveTRMMlsm(var,geo,ginfo,sroot)

end

function saveTRMMlsm(
	lsm   :: AbstractArray{<:Real,2},
	geo   :: GeoRegion,
	ginfo :: RegionGrid,
	sroot :: AbstractString
)

	@info "$(modulelog()) - Saving TRMM-TMPA Land-Sea Mask data in the $(geo.name) GeoRegion"

	if !isdir(sroot); mkpath(sroot) end
	fnc = joinpath(sroot,"trmmlsm-$(geo.regID).nc")
	if isfile(fnc)
		@info "$(modulelog()) - Overwrite stale NetCDF file $(fnc) ..."
		rm(fnc);
	end

	@info "$(modulelog()) - Creating NetCDF file $(fnc) ..."
	ds = NCDataset(fnc,"c",attrib = Dict(
		"Conventions" => "CF-1.4",
		"NCO"         => "4.5.1",
		"Authors"     => "Bill Olson, Dave Bolvin, George Huffman",
		"Description" => "This TMPA land sea mask contains values ranging from 0% to 100% with 0% representing all land and 100% representing all ocean",
		"Title"       => "Land/Sea static mask relevant to TMPA precipitation 0.25x0.25 degree",
		"Version"     => "2",
		"History"     => "This land sea mask originated from the NOAA group at SSEC in the 1980s. It was originally produced at 1/6 deg resolution, and then regridded for the purposes of GPCP, TMPA, and IMERG precipitation products. NASA code 610.2, Global Change Data Center, restructured this TMPA land sea mask to match the TMPA grid, and converted the file to CF-compliant netCDF4. Version 2 was created in May, 2019 to resolve detected inaccuracies in coastal regions.",
		"DOI"         => "10.5067/K1VR7BN7YDWN",
	))

	ds.dim["longitude"] = length(ginfo.glon)
	ds.dim["latitude"]  = length(ginfo.glat)

	nclon = defVar(ds,"longitude",Float32,("longitude",),attrib = Dict(
	    "units"     => "degrees_east",
	    "long_name" => "longitude",
	))

	nclat = defVar(ds,"latitude",Float32,("latitude",),attrib = Dict(
	    "units"     => "degrees_north",
	    "long_name" => "latitude",
	))

	nclsm = defVar(ds,"lsm",Float32,("longitude","latitude",),attrib = Dict(
	    "units"         => "0-1",
	    "long_name"     => "land_sea_mask",
		"full_name"     => "Land-Sea Mask",
	))

	nclon[:] = ginfo.glon
	nclat[:] = ginfo.glat
	nclsm[:] = 1 .- lsm/100

	close(ds)

	@info "$(modulelog()) - TRMM-TMPA Land-Sea Mask data in the $(geo.name) GeoRegion has been saved into $(fnc)"

end


module Sentiment

using ZipFile
using ..Data: deps

function load()
  isfile(deps("sentiment.zip")) || return
  info("Downloading sentiment treebank dataset")
  download("https://cache.julialang.org/https://nlp.stanford.edu/sentiment/trainDevTestTrees_PTB.zip",
           deps("sentiment.zip"))
end

getfile(r, name) = r.files[findfirst(x -> x.name == name, r.files)]

function getfile(name)
  r = ZipFile.Reader(deps("sentiment.zip"))
  text = readstring(getfile(r, "trees/$name"))
  close(r)
  return text
end

using ..Data: Tree

totree_(n, w) = Tree{Any}((parse(Int, n), w))
totree_(n, a, b) = Tree{Any}((parse(Int, n), nothing), totree(a), totree(b))
totree(t::Expr) = totree_(t.args...)

function parsetree(s)
  s = replace(s, r"\$", s -> "\\\$")
  s = replace(s, r"[^\s\(\)]+", s -> "\"$s\"")
  s = replace(s, " ", ", ")
  return totree(parse(s))
end

function gettrees(name)
  load()
  ss = split(getfile("$name.txt"), '\n', keep = false)
  return parsetree.(ss)
end

train() = gettrees("train")
test() = gettrees("test")
dev() = gettrees("dev")

end


@testset "capture" begin
    @testset "basic" begin
        @system SCapture(Controller) begin
            a => 1 ~ track
            b(a) => a + 1 ~ capture
            c(a) => a + 1 ~ accumulate
        end
        s = instance(SCapture)
        @test s.b' == 0 && s.c' == 0
        update!(s)
        @test s.b' == 2 && s.c' == 2
        update!(s)
        @test s.b' == 2 && s.c' == 4
    end

    @testset "time" begin
        @system SCaptureTime(Controller) begin
            t(x=context.clock.time) => 2x ~ track(u"hr")
            a => 1 ~ track
            b(a) => a + 1 ~ capture(time=t)
            c(a) => a + 1 ~ accumulate(time=t)
        end
        s = instance(SCaptureTime)
        @test s.b' == 0 && s.c' == 0
        update!(s)
        @test s.b' == 4 && s.c' == 4
        update!(s)
        @test s.b' == 4 && s.c' == 8
    end

    @testset "unit hour" begin
        @system SCaptureUnitHour(Controller) begin
            a => 1 ~ capture(u"hr")
        end
        s = instance(SCaptureUnitHour)
        @test iszero(s.a')
        update!(s)
        @test s.a' == 1u"hr"
        update!(s)
        @test s.a' == 1u"hr"
    end

    @testset "unit day" begin
        @system SCaptureUnitDay(Controller) begin
            a => 1 ~ capture(u"d")
        end
        s = instance(SCaptureUnitDay)
        @test iszero(s.a')
        update!(s)
        @test s.a' == 1u"hr"
        update!(s)
        @test s.a' == 1u"hr"
    end

    @testset "when" begin
        @system SCaptureWhen(Controller) begin
            t(context.clock.tick) ~ track::int
            f ~ preserve(parameter)
            w(t, f) => t <= f ~ flag
            a => 1 ~ capture
            b => 1 ~ capture(when=w)
            c => 1 ~ capture(when=!w)
        end
        n = 5
        s = instance(SCaptureWhen; config=:0 => :f => n)
        simulate!(s, stop=n)
        @test s.a' == 1
        @test s.b' == 1
        @test s.c' == 0
        simulate!(s, stop=n)
        @test s.a' == 1
        @test s.b' == 0
        @test s.c' == 1
    end
end


# Apportion total learning to different human capital types

abstract type AbstractLearnSwitches end
abstract type AbstractLearningSet <: ModelObject end
abstract type AbstractLearning end

## ----------------  All colleges only teach college h

struct LearnCollegeOnlySet <: AbstractLearningSet
    objId :: ObjectId
end
struct LearnCollegeOnlySwitches <: AbstractLearnSwitches end
struct LearnCollegeOnly <: AbstractLearning end

ModelParams.has_pvector(l :: LearnCollegeOnlySet) = false;

init_learning_set(objId :: ObjectId, switches :: LearnCollegeOnlySwitches) =
    LearnCollegeOnlySet(objId);

init_learning(l :: LearnCollegeOnlySet, isTwoYear :: Bool) =
    LearnCollegeOnly();

frac_h(l :: LearnCollegeOnly) = 1.0;

# Map learning and endowments into human capital stocks at work start
h_stocks(l :: LearnCollegeOnly, dh, h0, hHs) =
    (dh .+ h0, hHs .+ zeros(size(dh)));

StructLH.describe(switches :: LearnCollegeOnlySwitches) =
    ["Learning"  "Colleges only teach college human capital"];

## ---------  Two and four year colleges

Base.@kwdef mutable struct LearnTwoFourSwitches <: AbstractLearnSwitches
    # Fraction learned that becomes hCollege for both types of colleges
    fracHc2 :: Double = 0.2
    fracHc4 :: Double = 0.8
    calFracHc2 :: Bool = false
    calFracHc4 :: Bool = false
    # 4y colleges produce more college h than 2y colleges
    isIncreasing :: Bool = true
end

mutable struct LearnTwoFourSet <: AbstractLearningSet
    objId :: ObjectId
    switches :: LearnTwoFourSwitches
    pvec :: ParamVector
    fracHc2 :: Double
    fracHc4 :: Double
end

struct LearnTwoFour <: AbstractLearning
    fracHc :: Double
end

function StructLH.describe(l :: LearnTwoFourSet)
    frac2y = round(frac_h(l, true); digits = 2);
    frac4y = round(frac_h(l, false); digits = 2);
    return [
        "Learning in college"  "2y/4y produce different combinations of h";
        "Fraction college h 2y"  "$frac2y";
        "Fraction college h 4y"  "$frac4y"
    ]
end

function StructLH.describe(switches :: LearnTwoFourSwitches)
    if switches.calFracHc2
        h2Str = "calibrated";
    else
        h2Str = "$(switches.fracHc2)"
    end
    if switches.calFracHc4
        h4Str = "calibrated";
    else
        h4Str = "$(switches.fracHc4)"
    end
    return ["Learning"  "Two and four year colleges teach different h mixes.";
    "Frac college h"  "Two yr: $h2Str   Four yr: $h4Str"];
end

# Displays parameters in levels, not as intercept and increments.
function ModelParams.param_table(l :: LearnTwoFourSet, isCalibrated :: Bool)
    if l.switches.isIncreasing
        if (l.switches.calFracHc2 == isCalibrated) ||
            (l.switches.calFracHc4 == isCalibrated)
            symbols = chain_strings(param_symbols(l), ", ");
            values = [frac_h(l, true), frac_h(l, false)];
            # This is where we get the description and symbol from
            p = l.pvec[1];
            pt = ModelParams.ParamTable(1);
            ModelParams.set_row!(pt, 1, string(p.name), symbols,
                "Fraction college h produced by quality",
                format_vector("", values, 2));
        else
            pt = nothing;
        end
    else
        # The standard reporting is easily readable.
        pt = ModelParams.report_params(l.pvec, isCalibrated);
    end
    return pt
end

param_symbols(l :: LearnTwoFourSet) = [p.symbol  for p in l.pvec];

function init_learning_set(objId :: ObjectId, switches :: LearnTwoFourSwitches)
    fracHc2 = switches.fracHc2;
    pFrac2 = Param(:fracHc2, ldescription(:fracHcTwo), lsymbol(:fracHcTwo),
        fracHc2, fracHc2, 0.0, 1.0, switches.calFracHc2);

    fracHc4 = switches.fracHc4;
    descrStr = ldescription(:fracHcFour);
    if switches.isIncreasing
        descrStr = descrStr * " (increment)";
    end
    pFrac4 = Param(:fracHc4, descrStr, lsymbol(:fracHcFour),
        fracHc4, fracHc4, 0.0, 1.0, switches.calFracHc4);

    pvec = ParamVector(objId, [pFrac2, pFrac4]);
    return LearnTwoFourSet(objId, switches, pvec, fracHc2, fracHc4);
end

init_learning(l :: LearnTwoFourSet, isTwoYear :: Bool) =
    LearnTwoFour(frac_h(l, isTwoYear));

# Fraction of learning that goes to college h
frac_h(switches :: LearnTwoFourSwitches, isTwoYear :: Bool) =
    frac_h(switches.isIncreasing, switches.fracHc2, switches.fracHc4, isTwoYear);
frac_h(l :: LearnTwoFourSet, isTwoYear :: Bool) =
    frac_h(l.switches.isIncreasing, l.fracHc2, l.fracHc4, isTwoYear);
frac_h(l :: LearnTwoFour) = l.fracHc;

function frac_h(isIncreasing, fracHc2, fracHc4, isTwoYear)
    if isTwoYear
        fracH = fracHc2;
    elseif isIncreasing
        fracH = fracHc2 + fracHc4 * (1.0 - fracHc2);
    else
        fracH = fracHc4;
    end
    return fracH
end

# Map learning and endowments into human capital stocks at work start
h_stocks(l :: LearnTwoFour, dh, h0, hHs) =
    (h0 .+ frac_h(l) .* dh, hHs .+ (1.0 - frac_h(l)) .* dh);

## ---------  Generic

make_test_learning() = LearnCollegeOnly();

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