text:
"""
A generic, standalone implementation of Monte Carlo Tree Search.
It can be used on any game that implements `GameInterface`
and with any external oracle.

Both a synchronous and an asynchronous version are implemented, which
share most of their code. When browsing the sources for the first time,
we recommend that you study the sychronous version first.
"""
module MCTS

using DataStructures: Stack
using Distributions: Categorical, Dirichlet

using ..Util: @printing_errors, @unimplemented
import ..GI, ..GameType

#####
##### Interface for External Oracles
#####

"""
    MCTS.Oracle{Game}

Abstract base type for an oracle. Oracles must implement
[`MCTS.evaluate`](@ref) and [`MCTS.evaluate_batch`](@ref).
"""
abstract type Oracle{Game} end

"""
    MCTS.evaluate(oracle::Oracle, board)

Evaluate a single board position (assuming white is playing).

Return a pair `(P, V)` where:

  - `P` is a probability vector on `GI.available_actions(Game(board))`
  - `V` is a scalar estimating the value or win probability for white.
"""
function evaluate(oracle::Oracle, board)
  @unimplemented
end

"""
    MCTS.evaluate_batch(oracle::Oracle, boards)

Evaluate a batch of board positions.

Expect a vector of boards and return a vector of `(P, V)` pairs.

A default implementation is provided that calls [`MCTS.evaluate`](@ref)
sequentially on each position.
"""
function evaluate_batch(oracle::Oracle, boards)
  return [evaluate(oracle, b) for b in boards]
end

GameType(::Oracle{Game}) where Game = Game

#####
##### Some Example Oracles
#####

"""
    MCTS.RolloutOracle{Game} <: MCTS.Oracle{Game}

This oracle estimates the value of a position by simulating a random game
from it (a rollout). Moreover, it puts a uniform prior on available actions.
Therefore, it can be used to implement the "vanilla" MCTS algorithm.
"""
struct RolloutOracle{Game} <: Oracle{Game} end

function rollout!(state, board)
  while true
    reward = GI.white_reward(state)
    isnothing(reward) || (return reward)
    action = rand(GI.available_actions(state))
    GI.play!(state, action)
   end
end

function evaluate(::RolloutOracle{Game}, board) where Game
  state = Game(board)
  n = length(GI.available_actions(state))
  P = ones(n) ./ n
  V = rollout!(state, board)
  return P, V
end

struct RandomOracle{Game} <: Oracle{Game} end

function evaluate(::RandomOracle{Game}, board) where Game
  s = Game(board)
  n = length(GI.available_actions(s))
  P = ones(n) ./ n
  V = 0.
  return P, V
end

#####
##### Workers and Queries
#####

# Workers can send three types of messages to the inference server:
# - Done: there is no more work available for the worker
# - Query: the worker requests the server to evaluate a position
# - NoQuery: the worker finished an MCTS simulation without sending
#   an evaluation query (see `worker_yield!` for why this matters)

abstract type Message{B} end
struct Done{B} <: Message{B} end
struct NoQuery{B} <: Message{B} end
struct Query{B} <: Message{B} board :: B end
const AnyMessage{B} = Union{Done{B}, NoQuery{B}, Query{B}}
const EvaluationResult{R} = Tuple{Vector{R}, R}

mutable struct Worker{Board}
  id    :: Int # useful for debugging purposes
  stack :: Stack{Tuple{Board, Bool, Int}} # board, white_playing, action_number
  send  :: Channel{AnyMessage{Board}}
  recv  :: Channel{EvaluationResult{Float32}}
  queried :: Bool # the worker queried the server during the current simulation

  function Worker{B}(id) where B
    stack = Stack{Tuple{B, Bool, Int}}()
    send = Channel{AnyMessage{B}}(1)
    recv = Channel{EvaluationResult{Float32}}(1)
    new{B}(id, stack, send, recv, false)
  end
end

#####
##### Board Statistics
#####

struct ActionStats
  P :: Float32
  W :: Float64
  N :: Int
  nworkers :: UInt16 # Number of workers currently exploring this branch
end

struct BoardInfo
  stats :: Vector{ActionStats}
  Vest  :: Float32
end

Ntot(b::BoardInfo) = sum(s.N for s in b.stats)

#####
##### MCTS Environment
#####

"""
    MCTS.Env{Game}(oracle; <keyword args>) where Game

Create and initialize an MCTS environment with a given `oracle`.

## Keyword Arguments

  - `nworkers=1`: numbers of asynchronous workers (see below)
  - `fill_batches=false`: if true, a constant batch size is enforced for
     evaluation requests, by completing batches with dummy entries if necessary
  - `cpuct=1.`: exploration constant in the UCT formula
  - `noise_ϵ=0., noise_α=1.`: parameters for the dirichlet exploration noise
     (see below)

## Asynchronous MCTS

  - If `nworkers == 1`, MCTS is run in a synchronous fashion and the oracle is
    invoked through [`MCTS.evaluate`](@ref).

  - If `nworkers > 1`, `nworkers` asynchronous workers are spawned,
    along with an additional task to serve board evaluation requests.
    Such requests are processed by batches of
    size `nworkers` using [`MCTS.evaluate_batch`](@ref).

## Dirichlet Noise

A naive way to ensure exploration during training is to adopt an ϵ-greedy
policy, playing a random move at every turn instead of using the policy
prescribed by [`MCTS.policy`](@ref) with probability ϵ.
The problem with this naive strategy is that it may lead the player to make
terrible moves at critical moments, thereby biasing the policy evaluation
mechanism.

A superior alternative is to add a random bias to the neural prior for the root
node during MCTS exploration: instead of considering the policy ``p`` output
by the neural network in the UCT formula, one uses ``(1-ϵ)p + ϵη`` where ``η``
is drawn once per call to [`MCTS.explore!`](@ref) from a Dirichlet distribution
of parameter ``α``.
"""
mutable struct Env{Game, Board, Oracle}
  # Store (nonterminal) state statistics assuming player one is to play
  tree :: Dict{Board, BoardInfo}
  # External oracle to evaluate positions
  oracle :: Oracle
  # Workers
  workers :: Vector{Worker{Board}}
  global_lock :: ReentrantLock
  remaining :: Int # counts the number of remaining simulations to do
  # Parameters
  fill_batches :: Bool
  cpuct :: Float64
  noise_ϵ :: Float64
  noise_α :: Float64
  # Performance statistics
  total_time :: Float64
  inference_time :: Float64
  total_simulations :: Int64
  total_nodes_traversed :: Int64

  function Env{G}(oracle;
      nworkers=1, fill_batches=false,
      cpuct=1., noise_ϵ=0., noise_α=1.) where G
    B = GI.Board(G)
    tree = Dict{B, BoardInfo}()
    total_time = 0.
    inference_time = 0.
    total_simulations = 0
    total_nodes_traversed = 0
    lock = ReentrantLock()
    remaining = 0
    workers = [Worker{B}(i) for i in 1:nworkers]
    new{G, B, typeof(oracle)}(
      tree, oracle, workers, lock, remaining, fill_batches,
      cpuct, noise_ϵ, noise_α,
      total_time, inference_time, total_simulations, total_nodes_traversed)
  end
end

asynchronous(env::Env) = length(env.workers) > 1

GameType(::Env{Game}) where Game = Game

Done(::Env{G,B}) where {G,B} = Done{B}()
NoQuery(::Env{G,B}) where {G,B} = NoQuery{B}()

#####
##### Access and initialize board information
#####

function init_board_info(P, V)
  stats = [ActionStats(p, 0, 0, 0) for p in P]
  return BoardInfo(stats, V)
end

# Returns statistics for the current player, true if new node
# Synchronous version
function board_info_sync(env, worker, board)
  if haskey(env.tree, board)
    return (env.tree[board], false)
  else
    (P, V), time = @timed evaluate(env.oracle, board)
    env.inference_time += time
    info = init_board_info(P, V)
    env.tree[board] = info
    return (info, true)
  end
end

# Equivalent of `board_info_sync` for asynchronous MCTS
function board_info_async(env, worker, board)
  if haskey(env.tree, board)
    return (env.tree[board], false)
  else
    # Send a request to the inference server
    put!(worker.send, Query(board))
    unlock(env.global_lock)
    P, V = take!(worker.recv)
    lock(env.global_lock)
    worker.queried = true
    # Another worker may have sent the same request and initialized
    # the node before. Therefore, we have to test membership again.
    if !haskey(env.tree, board)
      info = init_board_info(P, V)
      env.tree[board] = info
      return (info, true)
    else
      # The inference result is ignored and we proceed as if
      # the node was already in the tree.
      return (env.tree[board], false)
    end
  end
end

function board_info(env, worker, board)
  asynchronous(env) ?
    board_info_async(env, worker, board) :
    board_info_sync(env, worker, board)
end

#####
##### Main algorithm
#####

function uct_scores(info::BoardInfo, cpuct, ϵ, η)
  @assert iszero(ϵ) || length(η) == length(info.stats)
  sqrtNtot = sqrt(Ntot(info))
  return map(enumerate(info.stats)) do (i, a)
    Q = (a.W - a.nworkers) / max(a.N, 1)
    P = iszero(ϵ) ? a.P : (1-ϵ) * a.P + ϵ * η[i]
    Q + cpuct * P * sqrtNtot / (a.N + 1)
  end
end

function push_board_action!(env, worker, (b, wp, aid))
  push!(worker.stack, (b, wp, aid))
  stats = env.tree[b].stats
  astats = stats[aid]
  stats[aid] = ActionStats(
    astats.P, astats.W, astats.N + 1, astats.nworkers + 1)
end

function select!(env, worker, state, η)
  state = copy(state)
  env.total_simulations += 1
  isroot = true
  while true
    wr = GI.white_reward(state)
    isnothing(wr) || (return wr)
    wp = GI.white_playing(state)
    board = GI.canonical_board(state)
    actions = GI.available_actions(state)
    let (info, new_node) = board_info(env, worker, board)
      new_node && (return info.Vest)
      ϵ = isroot ? env.noise_ϵ : 0.
      scores = uct_scores(info, env.cpuct, ϵ, η)
      best_action_id = argmax(scores)
      best_action = actions[best_action_id]
      push_board_action!(env, worker, (board, wp, best_action_id))
      GI.play!(state, best_action)
      env.total_nodes_traversed += 1
      isroot = false
    end
  end
end

function backprop!(env, worker, white_reward)
  while !isempty(worker.stack)
    board, white_playing, action_id = pop!(worker.stack)
    reward = white_playing ?
      white_reward :
      GI.symmetric_reward(white_reward)
    stats = env.tree[board].stats
    astats = stats[action_id]
    stats[action_id] = ActionStats(
      astats.P, astats.W + reward, astats.N, astats.nworkers - 1)
  end
end

# It is important to guarantee that a worker sends one request to the
# inference server per simulation. Otherwise, a worker could block all
# the others if it repeatedly doesn't need to make evaluation queries.
function worker_yield!(env::Env, worker::Worker)
  if asynchronous(env) && !worker.queried
    put!(worker.send, NoQuery(env))
    unlock(env.global_lock)
    take!(worker.recv)
    lock(env.global_lock)
  end
end

function worker_explore!(env::Env, worker::Worker, state, η)
  @assert isempty(worker.stack)
  worker.queried = false
  white_reward = select!(env, worker, state, η)
  worker_yield!(env, worker)
  backprop!(env, worker, white_reward)
  @assert isempty(worker.stack)
end

function inference_server(env::Env{G, B, A}) where {G, B, A}
  to_watch = env.workers
  while true
    # Updating the list of workers to watch
    requests = [take!(w.send) for w in to_watch]
    done = [isa(r, Done) for r in requests]
    active = .~ done
    any(active) || break
    to_watch = to_watch[active]
    requests = requests[active]
    # Gathering queries
    batch = [q.board for q in requests if isa(q, Query)]
    if isempty(batch)
      answers, time = EvaluationResult{Float32}[], 0.
    else
      if env.fill_batches
        nmissing = length(env.workers) - length(batch)
        if nmissing > 0
          append!(batch, [batch[1] for i in 1:nmissing])
        end
        @assert length(batch) == length(env.workers)
      end
      answers, time = @timed evaluate_batch(env.oracle, batch)
    end
    dummy_answer = (Float32[], 0f0)
    env.inference_time += time
    for (i, q) in enumerate(requests)
      if isa(q, Query)
        put!(to_watch[i].recv, popfirst!(answers))
      else
        put!(to_watch[i].recv, dummy_answer)
      end
    end
  end
end

function dirichlet_noise(state, α)
  actions = GI.available_actions(state)
  n = length(actions)
  return rand(Dirichlet(n, α))
end

function explore_sync!(env::Env, state, nsims)
  η = dirichlet_noise(state, env.noise_α)
  elapsed = @elapsed for i in 1:nsims
    worker_explore!(env, env.workers[1], state, η)
  end
  env.total_time += elapsed
end

function explore_async!(env::Env, state, nsims)
  env.remaining = nsims
  η = dirichlet_noise(state, env.noise_α)
  elapsed = @elapsed begin
    @sync begin
      @async @printing_errors inference_server(env)
      for w in env.workers
        @async @printing_errors begin
          lock(env.global_lock)
          while env.remaining > 0
            env.remaining -= 1
            worker_explore!(env, w, state, η)
          end
          put!(w.send, Done(env))
          unlock(env.global_lock)
        end
      end
    end
  end
  env.total_time += elapsed
  return
end

"""
    MCTS.explore!(env, state, nsims)

Run `nsims` MCTS simulations from `state`.
"""
function explore!(env::Env, state, nsims)
  asynchronous(env) ?
    explore_async!(env, state, nsims) :
    explore_sync!(env, state, nsims)
end

"""
    MCTS.policy(env, state; τ=1.)

Return the recommended stochastic policy on `state`,
with temperature parameter equal to `τ`. If `τ` is zero, all the weight
goes to the action with the highest visits count.

A call to this function must always be preceded by
a call to [`MCTS.explore!`](@ref).
"""
function policy(env::Env, state; τ=1.0)
  actions = GI.available_actions(state)
  board = GI.canonical_board(state)
  info =
    try env.tree[board]
    catch e
      if isa(e, KeyError)
        error("MCTS.explore! must be called before MCTS.policy")
      else
        rethrow(e)
      end
    end
  if iszero(τ)
    best = argmax([a.N for a in info.stats])
    π = zeros(length(actions))
    π[best] = 1.0
    return actions, π
  else
    τinv = 1 / τ
    Ntot = sum(a.N for a in info.stats)
    π = [(a.N / Ntot) ^ τinv for a in info.stats]
    π ./= sum(π)
  end
  return actions, π
end

"""
    MCTS.reset!(env)

Empty the MCTS tree.
"""
function reset!(env)
  empty!(env.tree)
  GC.gc(true)
end

#####
##### Profiling Utilities
#####

"""
    MCTS.inference_time_ratio(env)

Return the ratio of time spent by [`MCTS.explore!`](@ref)
on position evaluation (through functions [`MCTS.evaluate`](@ref) or
[`MCTS.evaluate_batch`](@ref)) since the environment's creation.
"""
function inference_time_ratio(env)
  T = env.total_time
  iszero(T) ? 0. : env.inference_time / T
end

"""
    MCTS.average_exploration_depth(env)

Return the average number of nodes that are traversed during an
MCTS simulation, not counting the root.
"""
function average_exploration_depth(env)
  return env.total_nodes_traversed / env.total_simulations
end

"""
    MCTS.memory_footprint_per_node(env)

Return an estimate of the memory footprint of a single node
of the MCTS tree (in bytes).
"""
function memory_footprint_per_node(env::Env{G}) where G
  # The hashtable is at most twice the number of stored elements
  # For every element, a board and a pointer are stored
  size_key = 2 * (GI.board_memsize(G) + sizeof(Int))
  dummy_stats = BoardInfo([
    ActionStats(0, 0, 0, 0) for i in 1:GI.num_actions(G)], 0)
  size_stats = Base.summarysize(dummy_stats)
  return size_key + size_stats
end

"""
    MCTS.approximate_memory_footprint(env)

Return an estimate of the memory footprint of the MCTS tree (in bytes).
"""
function approximate_memory_footprint(env::Env)
  return memory_footprint_per_node(env) * length(env.tree)
end

# Possibly very slow for large trees
memory_footprint(env::Env) = Base.summarysize(env.tree)

end


text:
"""
    varima_test(Y::Array{Float64,2}, d::Int64, p::Int64, q::Int64, benchmark_data::Tuple)

Run a series of tests to check whether the varima functions in uc_models.jl work.
"""
function varima_test(Y::JMatrix{Float64}, d::Int64, p::Int64, q::Int64, benchmark_data::Tuple)

    # Benchmark data
    benchmark_X0, benchmark_P0, benchmark_B, benchmark_R, benchmark_C, benchmark_V, benchmark_fc = benchmark_data;

    # Tests on VARIMASettings
    varima_settings = VARIMASettings(Y, d, p, q);
    @test varima_settings.d == d;
    @test varima_settings.p == p;
    @test varima_settings.q == q;
    @test varima_settings.n == size(Y,1);
    @test varima_settings.nr == size(Y,1)*max(p, q+1);
    @test varima_settings.np == size(Y,1)*p;
    @test varima_settings.nq == size(Y,1)*q;
    @test varima_settings.nnp == size(Y,1)^2*p;
    @test varima_settings.nnq == size(Y,1)^2*q;

    # Estimate parameters
    varima_out = varima(varima_settings, NelderMead(), Optim.Options(iterations=10000, f_tol=1e-2, x_tol=1e-2, g_tol=1e-2, show_trace=true, show_every=500));

    # Test on the parameters
    @test round.(varima_out.B, digits=10) == benchmark_B;
    @test round.(varima_out.R, digits=10) == benchmark_R;
    @test round.(varima_out.C, digits=10) == benchmark_C;
    @test round.(varima_out.V, digits=10) == benchmark_V;
    @test round.(varima_out.X0, digits=10) == benchmark_X0;
    @test round.(varima_out.P0, digits=10) == benchmark_P0;

    # 12-step ahead forecast
    fc = forecast(varima_out, 12, varima_settings);
    @test round.(fc, digits=10) == benchmark_fc;
end

@testset "arma" begin

    # Load data
    Y = permutedims(read_test_input("./input/arma/data"));

    # Settings for ARMA(1,1)
    d = 0;
    p = 1;
    q = 1;

    # Correct estimates: initial conditions
    benchmark_X0 = read_test_input("./input/arma/benchmark_X0");
    benchmark_P0 = read_test_input("./input/arma/benchmark_P0");

    # Correct estimates: observation equation
    benchmark_B = read_test_input("./input/arma/benchmark_B");
    benchmark_R = read_test_input("./input/arma/benchmark_R");

    # Correct estimates: transition equation
    benchmark_C = read_test_input("./input/arma/benchmark_C");
    benchmark_V = read_test_input("./input/arma/benchmark_V");

    # Correct estimates: forecast
    benchmark_fc = read_test_input("./input/arma/benchmark_fc");

    # Benchmark data
    benchmark_data = (benchmark_X0, benchmark_P0, benchmark_B, benchmark_R, benchmark_C, benchmark_V, benchmark_fc);

    # Run tests
    varima_test(Y, d, p, q, benchmark_data);
end

@testset "arima" begin

    # Load data
    Y = permutedims(read_test_input("./input/arima/data"));

    # Settings for arima(1,1)
    d = 1;
    p = 1;
    q = 1;

    # Correct estimates: initial conditions
    benchmark_X0 = read_test_input("./input/arima/benchmark_X0");
    benchmark_P0 = read_test_input("./input/arima/benchmark_P0");

    # Correct estimates: observation equation
    benchmark_B = read_test_input("./input/arima/benchmark_B");
    benchmark_R = read_test_input("./input/arima/benchmark_R");

    # Correct estimates: transition equation
    benchmark_C = read_test_input("./input/arima/benchmark_C");
    benchmark_V = read_test_input("./input/arima/benchmark_V");

    # Correct estimates: forecast
    benchmark_fc = read_test_input("./input/arima/benchmark_fc");

    # Benchmark data
    benchmark_data = (benchmark_X0, benchmark_P0, benchmark_B, benchmark_R, benchmark_C, benchmark_V, benchmark_fc);

    # Run tests
    varima_test(Y, d, p, q, benchmark_data);
end

@testset "varma" begin

    # Load data
    Y = read_test_input("./input/varma/data");

    # Settings for ARMA(1,1)
    d = 0;
    p = 1;
    q = 1;

    # Correct estimates: initial conditions
    benchmark_X0 = read_test_input("./input/varma/benchmark_X0");
    benchmark_P0 = read_test_input("./input/varma/benchmark_P0");

    # Correct estimates: observation equation
    benchmark_B = read_test_input("./input/varma/benchmark_B");
    benchmark_R = read_test_input("./input/varma/benchmark_R");

    # Correct estimates: transition equation
    benchmark_C = read_test_input("./input/varma/benchmark_C");
    benchmark_V = read_test_input("./input/varma/benchmark_V");

    # Correct estimates: forecast
    benchmark_fc = read_test_input("./input/varma/benchmark_fc");

    # Benchmark data
    benchmark_data = (benchmark_X0, benchmark_P0, benchmark_B, benchmark_R, benchmark_C, benchmark_V, benchmark_fc);

    # Run tests
    varima_test(Y, d, p, q, benchmark_data);
end

@testset "varima" begin

    # Load data
    Y = read_test_input("./input/varima/data");

    # Settings for ARMA(1,1)
    d = 1;
    p = 1;
    q = 1;

    # Correct estimates: initial conditions
    benchmark_X0 = read_test_input("./input/varima/benchmark_X0");
    benchmark_P0 = read_test_input("./input/varima/benchmark_P0");

    # Correct estimates: observation equation
    benchmark_B = read_test_input("./input/varima/benchmark_B");
    benchmark_R = read_test_input("./input/varima/benchmark_R");

    # Correct estimates: transition equation
    benchmark_C = read_test_input("./input/varima/benchmark_C");
    benchmark_V = read_test_input("./input/varima/benchmark_V");

    # Correct estimates: forecast
    benchmark_fc = read_test_input("./input/varima/benchmark_fc");

    # Benchmark data
    benchmark_data = (benchmark_X0, benchmark_P0, benchmark_B, benchmark_R, benchmark_C, benchmark_V, benchmark_fc);

    # Run tests
    varima_test(Y, d, p, q, benchmark_data);
end

ext: jl


text:
import TransferEntropy: TransferEntropyEstimator, EmbeddingTE, transferentropy
export predictive_asymmetry

function causalbalance(lags, tes)
    posinds = findall(lags .> 0)
    neginds = findall(lags .< 0)
    sum(tes[posinds]) - sum(tes[neginds])
end

function verified_prediction_lags(lag::Int)
    # If only one η is provided (either positive or negative), just make a prediction
    # for that lag.
    ηs = [-abs(lag), abs(lag)]
end

function verified_prediction_lags(lags)
    ηs = sort(lags[lags .!= 0])

    if length(ηs) % 2 != 0
        throw(ArgumentError("Need an even number of lags (as many negative as positive prediction lags)"))
    end

    ηs_neg = ηs[ηs .< 0]
    ηs_pos = ηs[ηs .> 0]

    if length(ηs_neg) != length(ηs_pos)
        throw(ArgumentError("There must be as many negative as positive prediction lags. Got $ηs_neg and $ηs_pos"))
    end

    ηs_neg_sort = sort(abs.(ηs_neg))
    ηs_pos_sort = sort(ηs_pos)
    if !(all(ηs_neg_sort .== ηs_pos_sort))
        throw(ArgumentError("Negative amd positive prediction lags must be symmetric. Got $ηs_neg and $ηs_pos"))
    end

    return ηs
end

"""
## General interface

    predictive_asymmetry(s, t, [c],
        estimator::TransferEntropyEstimator, ηs;
        d𝒯 = 1, dT = 1, dS = 1, τT = -1, τS = -1,
        [dC = 1, τC = -1,],
        normalize::Bool = false, f::Real = 1.0) → Vector{Float64}

Compute the predictive asymmetry[^Haaga2020] 𝔸(`s` → `t`) for source time series `s` and
target time series `t` over prediction lags `ηs`, using the given `estimator` and embedding
parameters `d𝒯`, `dT`, `dS`, `τT`, `τS`.

If a conditional time series `c` is provided, compute 𝔸(`s` → `t` | `c`). Then, `dC` and
`τC` controls the embedding dimension and embedding lag for the conditional variable.

## Returns

Returns a vector containing the predictive asymmetry for each value of `ηs`.

## Normalization (hypothesis test)

If `normalize == true` (the default), then compute the normalized predictive asymmetry 𝒜.

In this case, for each ``\\eta`` in `ηs`, compute 𝒜(η) by normalizing 𝔸(η) to some fraction `f` of the
mean transfer entropy over prediction lags ``-\\eta, ..., \\eta`` (exluding lag 0).

Haaga et al. (2020)[^Haaga2020] uses a normalization with `f=1.0` as a built-in hypothesis test,
avoiding more computationally costly surrogate testing.

## Estimators

Any [estimator](@ref) that works for [`transferentropy`](@ref) will also work with
`predictive_asymmetry`. It is recommended to use either the rectangular
binning-based methods or the symbolic estimators for the fastest computations.

## Examples

```julia
using CausalityTools

# Some example time series
x, y = rand(100), rand(100)

# 𝔸(x → y) over prediction lags 1:5
𝔸reg  = predictive_asymmetry(x, y, VisitationFrequency(RectangularBinning(5)), 1:5)
```

[^Haaga2020]: Haaga, Kristian Agasøster, David Diego, Jo Brendryen, and Bjarte Hannisdal. "A simple test for causality in complex systems." arXiv preprint arXiv:2005.01860 (2020).
"""
function predictive_asymmetry end

function check_ηs(ηs)
    all(ηs .> 0) || throw(ArgumentError("all ηs must be >= 1, got $(ηs)"))
    issorted(ηs) || throw(ArgumentError("ηs must be provided in increasing order, got $(ηs)"))
end

function predictive_asymmetry(source, target, estimator, ηs;
        normalize = false, f::Real = 1.0,
        d𝒯 = 1, dT = 1, dS = 1, τT = -1, τS = -1)

    check_ηs(ηs)
    Nη = length(ηs)

    te_fws = zeros(Nη)
    te_bws = zeros(Nη)
    𝔸s = zeros(Nη)

    for (i, η) in enumerate(ηs)
        te_fws[i] = transferentropy(source, target, estimator, d𝒯 = d𝒯, dT = dT, dS = dS, τT = τT, τS = τS, η𝒯 = η)
        te_bws[i] = transferentropy(source, target, estimator, d𝒯 = d𝒯, dT = dT, dS = dS, τT = τT, τS = τS, η𝒯 = -η)

        if normalize
            𝔸ᵢ = (sum(te_fws[1:i]) - sum(te_bws[1:i])) / η
            avg_te = (sum(te_fws[1:i]) + sum(te_bws[1:i])) / (2*η)
            𝔸s[i] = 𝔸ᵢ / (f*avg_te)
        else
            𝔸ᵢ = (sum(te_fws[1:i]) - sum(te_bws[1:i])) / η
            𝔸s[i] = 𝔸ᵢ
        end
    end

    return 𝔸s
end

function predictive_asymmetry(source, target, cond, estimator, ηs;
        normalize = false, f::Real = 1.0,
        d𝒯 = 1, dT = 1, dS = 1, dC = 1, τT = -1, τS = -1, τC = -1)

    check_ηs(ηs)
    Nη = length(ηs)

    te_fws = zeros(Nη)
    te_bws = zeros(Nη)
    𝔸s = zeros(Nη)

    for (i, η) in enumerate(ηs)
        te_fws[i] = transferentropy(source, target, cond, estimator, d𝒯 = d𝒯, dT = dT, dS = dS, τT = τT, τS = τS, dC = dC, τC = τC, η𝒯 = η)
        te_bws[i] = transferentropy(source, target, cond, estimator, d𝒯 = d𝒯, dT = dT, dS = dS, τT = τT, τS = τS, dC = dC, τC = τC, η𝒯 = -η)
        if normalize
            𝔸ᵢ = (sum(te_fws[1:i]) - sum(te_bws[1:i])) / η
            avg_te = (sum(te_fws[1:i]) + sum(te_bws[1:i])) / (2*η)
            𝔸s[i] = 𝔸ᵢ / (f*avg_te)
        else
            𝔸ᵢ = (sum(te_fws[1:i]) - sum(te_bws[1:i])) / η
            𝔸s[i] = 𝔸ᵢ
        end
    end

    return 𝔸s
end


text:
export MultirateSolverType

"""
# Description
    MultirateSolverType(;
        splitting_type = SlowFastSplitting(),
        fast_model = AtmosAcousticGravityLinearModel,
        implicit_solver = ManyColumnLU,
        implicit_solver_adjustable = false,
        slow_method = LSRK54CarpenterKennedy,
        fast_method = LSRK54CarpenterKennedy,
        timestep_ratio = 100,
    )

This solver type constructs an ODE solver using a standard multirate
Runge-Kutta implementation. This solver computes solutions to ODEs with
the partitioned form:

```math
    \\dot{Q} = f_fast(Q, t) + f_slow(Q, t)
```

where the right-hand-side functions `f_fast` and `f_slow` denote
fast and slow dynamics respectively, depending on the state `Q`.

# Arguments
- `splitting_type` (DiscreteSplittingType): The type of discrete
    splitting to apply to the right-hand side.
    Default: `SlowFastSplitting()`
- `fast_model` (Type): The model describing fast dynamics.
    Default: `AtmosAcousticGravityLinearModel`
- `implicit_solver` (Type): An implicit solver for inverting the
    implicit system of equations (if using `HEVISplitting()`).
    Default: `ManyColumnLU`
- `implicit_solver_adjustable` (Bool): A flag identifying whether
    or not the `implicit_solver` can be updated as the time-step
    size changes. This is particularly important when using
    an implicit solver within a multirate scheme.
    Default: `false`
- `slow_method` (Function): Function defining the particular explicit
    Runge-Kutta method to be used for the slow processes.
    Default: `LSRK54CarpenterKennedy`
- `fast_method` (Function): Function defining the fast solver.
    Depending on the choice of `splitting_type`, this can be
    an explicit Runge Kutta method or a 1-D IMEX (additive Runge-Kutta)
    method.
    Default: `LSRK54CarpenterKennedy`
- `timestep_ratio` (Int): Integer denoting the ratio between the slow
    and fast time-step sizes.
    Default: `100`

### References
    @article{SchlegelKnothArnoldWolke2012,
        title={Implementation of multirate time integration methods for air
            pollution modelling},
        author={Schlegel, M and Knoth, O and Arnold, M and Wolke, R},
        journal={Geoscientific Model Development},
        volume={5},
        number={6},
        pages={1395--1405},
        year={2012},
        publisher={Copernicus GmbH}
    }
"""
struct MultirateSolverType{DS} <: AbstractSolverType
    # The type of discrete splitting to apply to the right-hand side
    splitting_type::DS
    # The model describing fast dynamics
    fast_model::Type
    # Choice of implicit solver
    implicit_solver::Type
    # Can the implicit solver be updated with changing dt?
    implicit_solver_adjustable::Bool
    # RK method for evaluating the slow processes
    slow_method::Function
    # RK method for evaluating the fast processes
    fast_method::Function
    # The ratio between slow and fast time-step sizes
    timestep_ratio::Int

    function MultirateSolverType(;
        splitting_type = SlowFastSplitting(),
        fast_model = AtmosAcousticGravityLinearModel,
        implicit_solver = ManyColumnLU,
        implicit_solver_adjustable = false,
        slow_method = LSRK54CarpenterKennedy,
        fast_method = LSRK54CarpenterKennedy,
        timestep_ratio = 100,
    )

        DS = typeof(splitting_type)

        return new{DS}(
            splitting_type,
            fast_model,
            implicit_solver,
            implicit_solver_adjustable,
            slow_method,
            fast_method,
            timestep_ratio,
        )
    end
end

"""
    getdtmodel(ode_solver::MultirateSolverType, bl)

A function which returns a model representing the dynamics
with the most restrictive time-stepping requirements.
"""
function getdtmodel(ode_solver::MultirateSolverType, bl)
    # Most restrictive dynamics are part of the fast model
    return ode_solver.fast_model(bl)
end

"""
# Description
    function solversetup(
        ode_solver::MultirateSolverType{SlowFastSplitting},
        dg,
        Q,
        dt,
        t0,
        diffusion_direction,
    )

Creates an ODE solver for the partition slow-fast ODE
using a multirate method with explicit time-integration.
The splitting of the fast (acoustic and gravity waves)
dynamics is done in _all_ spatial directions. Examples
of a similar implementations include the following
references.

### References
    @article{SchlegelKnothArnoldWolke2012,
        title={Implementation of multirate time integration methods for air
            pollution modelling},
        author={Schlegel, M and Knoth, O and Arnold, M and Wolke, R},
        journal={Geoscientific Model Development},
        volume={5},
        number={6},
        pages={1395--1405},
        year={2012},
        publisher={Copernicus GmbH}
    }

    @article{Schlegel_2009,
        doi = {10.1016/j.cam.2008.08.009},
        year = {2009},
        month = {apr},
        publisher = {Elsevier {BV}},
        volume = {226},
        number = {2},
        pages = {345--357},
        author = {Martin Schlegel and Oswald Knoth and Martin Arnold and Ralf Wolke},
        title = {Multirate Runge-Kutta schemes for advection equations},
        journal = {Journal of Computational and Applied Mathematics}
    }
"""
function solversetup(
    ode_solver::MultirateSolverType{SlowFastSplitting},
    dg,
    Q,
    dt,
    t0,
    diffusion_direction,
)

    # Extract fast model and define a DG model
    # for the fast processes (acoustic/gravity waves
    # in all spatial directions)
    fast_model = ode_solver.fast_model(dg.balance_law)
    fast_dg = DGModel(
        fast_model,
        dg.grid,
        dg.numerical_flux_first_order,
        dg.numerical_flux_second_order,
        dg.numerical_flux_gradient,
        state_auxiliary = dg.state_auxiliary,
        state_gradient_flux = dg.state_gradient_flux,
        states_higher_order = dg.states_higher_order,
        direction = EveryDirection(),
    )

    # Using the RemainderModel, we subtract away the
    # fast processes and define a DG model for the
    # slower processes (advection and diffusion)
    slow_model = RemainderModel(dg.balance_law, (fast_model,))
    slow_dg = DGModel(
        slow_model,
        dg.grid,
        dg.numerical_flux_first_order,
        dg.numerical_flux_second_order,
        dg.numerical_flux_gradient,
        state_auxiliary = dg.state_auxiliary,
        state_gradient_flux = dg.state_gradient_flux,
        states_higher_order = dg.states_higher_order,
        # Ensure diffusion direction is passed to the correct
        # DG model
        diffusion_direction = diffusion_direction,
    )

    slow_solver = ode_solver.slow_method(slow_dg, Q; dt = dt)
    fast_dt = dt / ode_solver.timestep_ratio
    fast_solver = ode_solver.fast_method(fast_dg, Q; dt = fast_dt)

    solver = MultirateRungeKutta((slow_solver, fast_solver), t0 = t0)

    return solver
end

"""
# Description
    solversetup(
        ode_solver::MultirateSolverType{HEVISplitting},
        dg,
        Q,
        dt,
        t0,
        diffusion_direction,
    )

Creates an ODE solver for the partition slow-fast ODE
using a multirate method with HEVI time-integration.
The splitting of the fast (acoustic and gravity waves)
dynamics is performed by splitting the fast model
into horizontal and vertical directions. All horizontal
acoustic waves are treated explicitly, while the 1-D
vertical problem is treated implicitly. The HEVI-splitting
of the acoustic waves is handled using an IMEX additive
Runge-Kutta method in the fast (inner-most) solver.

Examples of similar multirate-HEVI approaches include
the following references.

### References
    @article{doms2011description,
        title={A Description of the Nonhydrostatic Regional COSMO model.
            Part I: Dynamics and Numerics},
        author={Doms, G{\"u}nther and Baldauf, M},
        journal={Deutscher Wetterdienst, Offenbach},
        year={2011}
    }

    @article{Tomita_2008,
        doi = {10.1137/070692273},
        year = {2008},
        month = {jan},
        publisher = {Society for Industrial and Applied Mathematics ({SIAM})},
        volume = {30},
        number = {6},
        pages = {2755--2776},
        author = {Hirofumi Tomita and Koji Goto and Masaki Satoh},
        title = {A New Approach to Atmospheric General Circulation Model:
            Global Cloud Resolving Model {NICAM} and its Computational Performance},
        journal = {{SIAM} Journal on Scientific Computing}
    }

# Comments:
Currently, the only HEVI-type splitting ClimateMachine can currently
do only involves splitting the acoustic processes; it is not currently
possible to perform more fine-grained separation of tendencies
(for example, including vertical advection or diffusion in the 1-D implicit problem)
"""
function solversetup(
    ode_solver::MultirateSolverType{HEVISplitting},
    dg,
    Q,
    dt,
    t0,
    diffusion_direction,
)

    # Extract fast model and define a DG model
    # for the fast processes
    fast_model = ode_solver.fast_model(dg.balance_law)

    # Full DG model for the acoustic waves in all directions
    acoustic_dg_full = DGModel(
        fast_model,
        dg.grid,
        dg.numerical_flux_first_order,
        dg.numerical_flux_second_order,
        dg.numerical_flux_gradient,
        state_auxiliary = dg.state_auxiliary,
        state_gradient_flux = dg.state_gradient_flux,
        states_higher_order = dg.states_higher_order,
        direction = EveryDirection(),
    )

    # DG model for the vertical acoustic waves only
    acoustic_dg_vert = DGModel(
        fast_model,
        dg.grid,
        dg.numerical_flux_first_order,
        dg.numerical_flux_second_order,
        dg.numerical_flux_gradient,
        state_auxiliary = dg.state_auxiliary,
        state_gradient_flux = dg.state_gradient_flux,
        states_higher_order = dg.states_higher_order,
        direction = VerticalDirection(),
    )

    # Compute fast time-step size from target ratio
    fast_dt = dt / ode_solver.timestep_ratio

    # Fast solver for the acoustic/gravity waves using
    # a HEVI-type splitting and a 1-D IMEX method
    fast_solver = ode_solver.fast_method(
        acoustic_dg_full,
        acoustic_dg_vert,
        LinearBackwardEulerSolver(
            ode_solver.implicit_solver();
            isadjustable = ode_solver.implicit_solver_adjustable,
        ),
        Q;
        dt = fast_dt,
        t0 = t0,
        # NOTE: This needs to be `false` since the ARK method will
        # evaluate the explicit part using the RemainderModel
        # (Difference between acoustic_dg_full and acoustic_dg_vert)
        split_explicit_implicit = false,
        # NOTE: Do we want to support more variants?
        variant = LowStorageVariant(),
    )

    # Finally, we subtract away the
    # fast processes and define a DG model for the
    # slower processes (advection and diffusion)
    slow_model = RemainderModel(dg.balance_law, (fast_model,))
    slow_dg = DGModel(
        slow_model,
        dg.grid,
        dg.numerical_flux_first_order,
        dg.numerical_flux_second_order,
        dg.numerical_flux_gradient,
        state_auxiliary = dg.state_auxiliary,
        state_gradient_flux = dg.state_gradient_flux,
        states_higher_order = dg.states_higher_order,
        # Ensure diffusion direction is passed to the correct
        # DG model
        diffusion_direction = diffusion_direction,
    )

    slow_solver = ode_solver.slow_method(slow_dg, Q; dt = dt, t0 = t0)

    solver = MultirateRungeKutta((slow_solver, fast_solver), t0 = t0)

    return solver
end

max_stars_repo_name: leios/CLIMA

max_issues_repo_name: leios/CLIMA

max_forks_repo_name: leios/CLIMA


text:
module Utils

using Test

function random_sleep(spawn::Bool = true)
    if rand() < 0.1
        # no sleep
    elseif spawn && rand(Bool)
        nspins = rand(0:10000)
        for _ in 1:nspins
            GC.safepoint()
            ccall(:jl_cpu_pause, Cvoid, ())
        end
    else
        sleep(rand() / 1_000_000)
    end
end

function trywait(task::Task, seconds::Real = 0.1)
    istaskdone(task) && wait(task)
    delta = 0.01
    for _ in 0.0:delta:seconds
        sleep(delta)
        if istaskdone(task)
            wait(task)
            break
        end
    end
    return istaskdone(task)
end

macro test_error(expr)
    @gensym err tmp
    quote
        local $err = nothing
        $Test.@test try
            $expr
            false
        catch $tmp
            $err = $tmp
            true
        end
        $err
    end |> esc
end

const ⊏ = occursin

end  # module


text:
"""
    function random_mate(nsire, ndam, nsib)
---
Random mate `nsire` with `ndam`, each dam will have `nsib`.
Assuming the first `nsire` are sires.
The rest are dams.
This is for generation one.
And can be replaced by call `random_mate(shuffle(1:nid), nsire, ndam, nsib)`.
"""
function random_mate(nsire, ndam, nsib)
    r  = Int(ceil(ndam/nsire))
    pa = repeat(repeat(shuffle(1:nsire), r), inner=nsib)
    ma = repeat(1:ndam, inner = nsib) .+ nsire
    [pa[1:length(ma)] ma]
end

"""
    function random_mate(ID, nsire, ndam, nsib)
---
Random mate `ID`.
Its first `nsire` are sires, last `ndam` are dams.

**Note**: `nsire + ndam` must be of `length(ID)`.
"""
function random_mate(ID, nsire, ndam, nsib)
    r  = Int(ceil(ndam/nsire))
    pa = repeat(repeat(shuffle(ID[1:nsire]), r), inner = nsib)
    ma = repeat(ID[(nsire+1):(ndam+nsire)], inner = nsib)
    [pa[1:length(ma)] ma]
end

"""
    function haldane(pos; M = 1e10)
---
Return recommbination rate column given bp distances.

Haldane function:
- Recombination rate: r = .5(1 - exp(-2d))
- Distance in Morgan: d = -.5log(1 - 2r)
"""
function haldane(pos; M = 1e8)
    chr = bp = 0
    nl = size(pos)[1]
    r = zeros(Float64, nl)
    @inbounds for i in 1:nl
        if pos[i, 1] == chr
            d = (pos[i, 2] - bp) / M # distance in cM
            r[i] = .5(1 - exp(-2d))
        else
            chr = pos[i, 1]
            r[i] = 0.5
        end
        bp = pos[i, 2]
    end
    r
end

"""
    function gdrop(base, ped)
---
Drop genotypes of a subset or all of `snp`, through `ped` into `osnp`.
The last will be returned.

## Note:
If to use `osnp = BitArray(undef, nsnp, 2noff)` to save memory,
the speed is greatly reduced.

## ToDo:
Needs optimization.
May be parallel into 10 tasks.
Least urgent.
"""
function gdrop(snp, ped, r)
    noff = size(ped)[1]
    nsnp = size(snp)[1]
    osnp = zeros(Bool, nsnp, 2noff)
    i = 1
    sample = zeros(nsnp)
    @inbounds for prt in eachrow(ped)
        @inbounds for p in prt
            hap = view(snp, :, 2p-1:2p)
            off = view(osnp, :, i)
            k   = 1
            rand!(sample)
            @inbounds for j in 1:nsnp
                sample[j] < r[j] && (k = 3 - k)
                off[j] = hap[j, k]
            end
            i += 1
        end
    end
    osnp
end


text:
# Periodic mean function

"""
    MeanPeriodic <: Mean

Periodic mean function
```math
m(x) = a'cos(2πx/p) + b'sin(2πx/p)
```
with polynomial coefficients ``βᵢⱼ`` of shape ``d × D`` where ``d`` is the dimension of
observations and ``D`` is the degree of the polynomial.
"""
mutable struct MeanPeriodic <: Mean
    "Cosine and sine coefficients"
    a::Vector{Float64}
    b::Vector{Float64}
    "Period"
    p::Vector{Float64}
    "Priors for mean parameters"
    priors::Array

    """
        MeanPeriodic(a::Vector{Float64}, b::Vector{Float64}, p::Vector{Float64})

    Create `MeanPeriodic` with cosine and sine coefficients `a` and `b`
    and period `p`.
    """
    MeanPeriodic(a::Vector{Float64}, b::Vector{Float64}, lp::Vector{Float64}) = new(a, b, exp.(lp), [])
end
MeanPeriodic(a::Real, b::Real, lp::Real) = MeanPeriodic([a], [b], [lp])

function mean(mPer::MeanPeriodic, x::AbstractVector)
    return dot(mPer.a, cos.(2π .* x ./ mPer.p)) +
           dot(mPer.b, sin.(2π .* x ./ mPer.p))
end

get_params(mPer::MeanPeriodic) = [mPer.a; mPer.b; log.(mPer.p)]
get_param_names(mPer::MeanPeriodic) = [get_param_names(mPer.a, :a);
                                       get_param_names(mPer.b, :b);
                                       get_param_names(mPer.p, :lp)]
num_params(mPer::MeanPeriodic) = length(mPer.a)+length(mPer.b)+length(mPer.p)

function set_params!(mPer::MeanPeriodic, hyp::AbstractVector)
    length(hyp) == num_params(mPer) || throw(ArgumentError("MeanPeriodic mean function has $(num_params(mPer)) parameters"))
    dim = length(mPer.a)
    copyto!(mPer.a, hyp[1:dim])
    copyto!(mPer.b, hyp[dim+1:2dim])
    copyto!(mPer.p, exp.(hyp[2dim+1:3dim]))
end

function grad_mean(mPer::MeanPeriodic, x::AbstractVector)
    dim = length(x)
    costerm = cos.(2π .* x ./ mPer.p)
    sinterm = sin.(2π .* x ./ mPer.p)
    dM_da = costerm
    dM_db = sinterm
    dM_dlp =  (mPer.a .* sinterm .- mPer.b .* costerm) .*
                       (2π .* x ./ mPer.p)
    return [dM_da; dM_db; dM_dlp]
end


text:
struct InfiniteDepthOrderBook{Tprice, Tvolume} <: AbstractOrderBook{Tprice, Tvolume}
    lowest_ask::Tprice
    highest_bid::Tprice
    volume::Tvolume

    function InfiniteDepthOrderBook(lowest_ask::Tprice, highest_bid::Tprice; volume::Tvolume=null_volume, unchecked=false) where {Tprice, Tvolume}
        unchecked && return new{Tprice, Tvolume}(lowest_ask, highest_bid, volume)
        highest_bid > lowest_ask && throw(OrderBookException("lowest_ask ($lowest_ask) must be greater than or equal to highest_bid ($highest_bid)"))
        new{Tprice, Tvolume}(lowest_ask, highest_bid, volume)
    end

end

function OrderBook(lowest_ask, highest_bid; volume=null_volume)
    InfiniteDepthOrderBook(lowest_ask, highest_bid; volume=volume)
end

function OrderBook(price; volume=null_volume)
    InfiniteDepthOrderBook(price, price; volume=volume)
end

function bid(ob::InfiniteDepthOrderBook, volume::Tvolume; raise=false) where {Tvolume}
    ob.highest_bid, zero(Tvolume)
end

function ask(ob::InfiniteDepthOrderBook, volume::Tvolume; raise=false) where {Tvolume}
    ob.lowest_ask, zero(Tvolume)
end


text:
# This file is part of InvertedFiles.jl

const EmptyWeightsPostingList = Float32[]

"""
    struct PostingList

A paired list of identifiers and weights
"""
struct PostingList
    I::Vector{UInt32}
    W::Vector{Float32}
    val::Float64  # useful for search time and saving global data
end

"""
    PostingList(I)

Creates a posting lists with an empty array of weights
"""
PostingList(I) = PostingList(I, EmptyWeightsPostingList, 1.0)

@inline Base.length(plist::PostingList) = length(plist.I)
@inline Intersections._get_key(plist::PostingList, i) = @inbounds plist.I[i]


text:
import Base: (:)

(:)(start::Epoch{S}, stop::Epoch{S}) where {S} = (:)(start, 1.0days, stop)

function (:)(start::Epoch{S}, step::Period, stop::Epoch{S}) where S
    step = seconds(step)
    step = start < stop ? step : -step
    StepRangeLen(start, step, floor(Int, value(stop-start)/value(step))+1)
end

Base.step(r::StepRangeLen{T}) where {T<:Epoch} = r.step


text:
import SpecialFunctions: gamma, lgamma, erf, erfc, beta

# Implicitly defined:
#binary: mod
#unary: exp, abs, log1p, sin, cos, tan, sinh, cosh, tanh, asin, acos, atan, asinh, acosh, atanh, erf, erfc, gamma, relu, round, floor, ceil, round, sign.

# Use some fast operators from https://github.com/JuliaLang/julia/blob/81597635c4ad1e8c2e1c5753fda4ec0e7397543f/base/fastmath.jl
# Define allowed operators. Any julia operator can also be used.
function plus(x::T, y::T)::T where {T<:Real}
	x + y #Do not change the name of this operator.
end
function sub(x::T, y::T)::T where {T<:Real}
	x - y #Do not change the name of this operator.
end
function mult(x::T, y::T)::T where {T<:Real}
	x * y #Do not change the name of this operator.
end
function square(x::T)::T where {T<:Real}
	x * x
end
function cube(x::T)::T where {T<:Real}
	x ^ 3
end
function pow(x::T, y::T)::T where {T<:Real}
	abs(x)^y
end
function div(x::T, y::T)::T where {T<:Real}
	x / y
end
function logm(x::T)::T where {T<:Real}
    log(abs(x) + convert(T, 1//100000000))
end
function logm2(x::T)::T where {T<:Real}
    log2(abs(x) + convert(T, 1//100000000))
end
function logm10(x::T)::T where {T<:Real}
    log10(abs(x) + convert(T, 1//100000000))
end

# Generics:
square(x) = x * x
cube(x) = x * x * x
plus(x, y) = x + y
sub(x, y) = x - y
mult(x, y) = x * y
pow(x, y) = abs(x)^y
div(x, y) = x / y
logm(x) = log(abs(x) + 1//100000000)
logm2(x) = log2(abs(x) + 1//100000000)
logm10(x) = log10(abs(x) + 1//100000000)

function sqrtm(x::T)::T where {T}
	sqrt(abs(x))
end
function neg(x::T)::T where {T}
	- x
end

function greater(x::T, y::T)::T where {T}
    if x > y
        return convert(T, 1)
    end
    return convert(T, 0)
end

function relu(x::T)::T where {T}
    if x > convert(T, 0)
        return x
    end
    return convert(T, 0)
end

function logical_or(x::T, y::T)::T where {T}
    if x > convert(T, 0) || y > convert(T, 0)
        return convert(T, 1)
    end
    return convert(T, 0)
end

# (Just use multiplication normally)
function logical_and(x::T, y::T)::T where {T}
    if x > convert(T, 0) && y > convert(T, 0)
        return convert(T, 1)
    end
    return convert(T, 0)
end


text:
@init global hp = safe_pyimport("healpy")

function θϕ_to_xy((θ,ϕ))
    r = 2cos(θ/2)
    x =  r*cos(ϕ)
    y = -r*sin(ϕ)
    x, y
end

function xy_to_θϕ((x,y))
    r = sqrt(x^2+y^2)
    θ = 2*acos(r/2)
    ϕ = -atan(y,x)
    θ, ϕ
end

function healpix_to_flat(healpix_map::Vector{T}, ::Type{P}; rot=(0,0,0)) where {Nside, θpix, T, P<:Flat{Nside,θpix}}

    Nside_sphere = hp.npix2nside(length(healpix_map))
    @unpack Δx = fieldinfo(P)

    # compute projection coordinate mapping
    xs = ys = Δx*((-Nside÷2:Nside÷2-1) .+ 0.5)
    xys = tuple.(xs,ys')[:]
    θϕs = xy_to_θϕ.(xys)
    (θs, ϕs) = first.(θϕs), last.(θϕs)

    # rotate the pole to the equator to match Healpy's azeqview convention, in
    # addition to applying the user rotation
    R = hp.Rotator((0,90,0), eulertype="ZYX") * hp.Rotator(rot, eulertype="ZYX")
    (θs′, ϕs′) = eachrow(R.get_inverse()(θs, ϕs))

    # interpolate map
    FlatMap{P,T}(reshape(hp.get_interp_val(healpix_map, θs′, ϕs′), Nside, Nside))

end

function flat_to_healpix(f::FlatMap, Nside_sphere; rot=(0,0,0))

    @unpack Δx, Nside = fieldinfo(f)

    (θs, ϕs) = hp.pix2ang(Nside_sphere, 0:(12*Nside_sphere^2-1))

    # rotate the pole to the equator to match Healpy's azeqview convention, in
    # addition to applying the user rotation
    R = hp.Rotator((0,90,0), eulertype="ZYX") * hp.Rotator(rot, eulertype="ZYX")
    (θs′, ϕs′) = eachrow(R(θs, ϕs))

    # compute projection coordinate mapping
    xys = θϕ_to_xy.(tuple.(θs′, ϕs′))
    xs = first.(xys) ./ Δx .+ Nside÷2 .+ 0.5
    ys = last.(xys) ./ Δx .+ Nside÷2 .+ 0.5

    # interpolate map
    @ondemand(Images.bilinear_interpolation).(Ref(f[:Ix]), xs, ys)

end


text:
@testset "convenience" begin
    x = OrderedDict{String, Any}("x" => 1, "y" => 2)
    buf = bson_write(UInt8[], x)
    @test bson_read(buf) == x
    io = IOBuffer()
    bson_write(io, x)
    seekstart(io)
    @test bson_read(io) == x
    path, io = mktemp()
    bson_write(path, x)
    @test bson_read(path) == x
end


text:
_is_incremental(::Type{Clp.Optimizer}) = false

function set_default_config!(ext::ExtensionType, d::GlobalOptimizer, m::Clp.Optimizer, local_solver::Bool)
    MOI.set(m, MOI.RawOptimizerAttribute("PrimalTolerance"),       _absolute_tol(d)*1E-2)
    MOI.set(m, MOI.RawOptimizerAttribute("DualTolerance"),         _absolute_tol(d)*1E-2)
    MOI.set(m, MOI.RawOptimizerAttribute("DualObjectiveLimit"),    1e308)
    MOI.set(m, MOI.RawOptimizerAttribute("MaximumIterations"),     2147483647)
    MOI.set(m, MOI.RawOptimizerAttribute("PresolveType"),          0)
    MOI.set(m, MOI.RawOptimizerAttribute("SolveType"),             5)
    MOI.set(m, MOI.RawOptimizerAttribute("InfeasibleReturn"),      1)
    MOI.set(m, MOI.RawOptimizerAttribute("Scaling"),               3)
    MOI.set(m, MOI.RawOptimizerAttribute("Perturbation"),          100)
    return
end

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
using Polytopes
using LightGraphs
using Test

@testset "Polytopes.jl" begin
    _edges = [[1, 2], [2, 3], [3, 1], [2, 3], [3, 4], [4, 2]]
    _faces = [[1, 2, 3], [4, 5, 6]]
    p = Polytope(_edges, _faces)
    @test rank(p) == 2
    @test nverts(p) == 4
    @test faces(p, 0).graph == SimpleDiGraph(nverts(p))
    @test faces(p, 1) == facets(p) == Polytope([_edges])
    @test faces(p, 2) === p
    @test p.graph == SimpleDiGraph([
        Edge(5, 1),
        Edge(5, 2),
        Edge(6, 2),
        Edge(6, 3),
        Edge(7, 3),
        Edge(7, 1),
        Edge(8, 2),
        Edge(8, 3),
        Edge(9, 3),
        Edge(9, 4),
        Edge(10, 4),
        Edge(10, 2),
        Edge(11, 5),
        Edge(11, 6),
        Edge(11, 7),
        Edge(12, 8),
        Edge(12, 9),
        Edge(12, 10),
    ])
    @test string(p) == "2-polytope with 4 vertices, 6 edges, 2 faces"

    gp = GeometricPolytope(p, map(x -> Vertex(randn(3)), 1:nverts(p)))
    @test embeddim(gp) == 3
    @test gp.polytope == p
    @test string(gp) == "3-dimensional geometric 2-polytope with 4 vertices, 6 edges, 2 faces"
end


text:
@testset "Pendulum" begin
    # get pendulum environment and simulate
    timestep=0.1
    env = Dojo.get_environment("pendulum",
        timestep=timestep,
        gravity=-10.0);
    reset(env);
    Dojo.initialize_pendulum!(env.mechanism,
        angle=0.25 * π)
    u1 = rand(Dojo.input_dimension(env.mechanism))
    z1 = Dojo.get_maximal_state(env.mechanism)

    storage = Dojo.simulate!(env.mechanism, 1.0,
        record=true)
    @test norm(z1 - Dojo.get_maximal_state(storage, 1)) < 1.0e-6
    zTs = Dojo.get_maximal_state(env.mechanism)
    @test norm(zTs - Dojo.get_maximal_state(storage, Dojo.length(storage))) < 1.0e-6
    zT = Dojo.get_next_state(env.mechanism)
    Dojo.initialize_pendulum!(env.mechanism,
        angle=0.25 * π)
    storage = Dojo.simulate!(env.mechanism, 1.0 + timestep,
        record=true)
    @test norm(zT - Dojo.get_maximal_state(storage, Dojo.length(storage))) < 1.0e-6

    # maximal gradients
    Dojo.step!(env.mechanism, z1, u1)
    Fx1, Fu1 = Dojo.get_maximal_gradients(env.mechanism)
    Fx2, Fu2 = Dojo.get_maximal_gradients!(env.mechanism, z1, u1)
    @test norm(Fx1 - Fx2, Inf) < 1.0e-6
    @test norm(Fu1 - Fu2, Inf) < 1.0e-6

    timestep=0.1
    env = Dojo.get_environment("halfcheetah",
        timestep=timestep,
        gravity=-10.0);
    reset(env);

    # get body
    @test Dojo.get_body(env.mechanism, :origin).name == :origin

    # get contact
    contact_name = env.mechanism.contacts[1].name
    @test Dojo.get_contact(env.mechanism, contact_name).name == contact_name

    # velocity indices
    @test Dojo.velocity_index(env.mechanism) == [4, 5, 6, 8, 10, 12, 14, 16, 18]

    # set state
    z = Dojo.get_maximal_state(env.mechanism)
    env2 = Dojo.get_environment(:halfcheetah,
        timestep=timestep,
        gravity=-10.0);
    reset(env);
    Dojo.set_maximal_state!(env2.mechanism, zeros(Dojo.maximal_dimension(env2.mechanism)))
    z2 = Dojo.get_maximal_state(env2.mechanism)
    @test norm(z2) < 1.0e-8
    Dojo.set_maximal_state!(env2.mechanism, z)
    @test norm(z - Dojo.get_maximal_state(env2.mechanism)) < 1.0e-8

    # change atlas floating base
    mech = get_mechanism(:atlas)
    j = get_joint(mech, :floating_base)
    base_origin = get_body(mech, j.child_id).name

    mech = set_floating_base(mech, :head)
    j = get_joint(mech, :floating_base)
    base_new = get_body(mech, j.child_id).name

    @test base_origin != base_new
    @test base_new == :head
end

#TODO: get and set methods


text:
let
    roadway = get_test_roadway()

    trajdatas = [get_test_trajdata(roadway), get_test_trajdata(roadway)]
    segments = [TrajdataSegment(1, 2, 1, 2), TrajdataSegment(1, 2, 1, 2)]
    evaldata = EvaluationData(trajdatas, segments)

    evaldata2 = create_evaldata_with_random_subsegments_of_equal_length(evaldata, nsegs=2, nframes=1)
    @test length(evaldata2.segments) == 2
    @test evaldata2.trajdatas === evaldata.trajdatas
    for seg in evaldata2.segments
        @test nframes(seg) == 1
    end

    segments = [TrajdataSegment(1, 2, 1, 2), TrajdataSegment(3, 4, 10, 20)]
    evaldata = EvaluationData(trajdatas, segments)
    evaldata2 = create_evaldata_with_random_subsegments_of_equal_length(evaldata, foldset_match([1,2], 1), nsegs=2, nframes=1)
    @test length(evaldata2.segments) == 2
    @test evaldata2.trajdatas === evaldata.trajdatas
    @test evaldata2.segments[1].trajdata_index == evaldata.segments[1].trajdata_index
    @test evaldata2.segments[1].egoid == evaldata.segments[1].egoid
    @test evaldata2.segments[2].trajdata_index == evaldata.segments[1].trajdata_index
    @test evaldata2.segments[2].egoid == evaldata.segments[1].egoid

    metrics = TraceMetricExtractor[RootWeightedSquareError(SPEED, 2.0)]
    metric_df = allocate_metrics_dataframe(metrics, 10)
    @test ncol(metric_df) == 3
    @test nrow(metric_df) == 10
    @test names(metric_df) == [:time, :logl, Symbol(metrics[1])]

    rec = pull_record(segments[1], evaldata)
    @test rec[0][1].def == get(trajdatas[1], 1, 2).def
    @test rec[0][1].state.posG == get(trajdatas[1], 1, 2).state.posG
end


text:
include("autopar.jl")
#input the Original Image `Img_O`, Image with noise  `Img_N`, sets of parameters `Par`
#use Iterative method to get Estimated Image `Img_E`
#output the Estimated Image `Img_E` in the form of Array{Float32,3} H x W x 3 (3 = RGB Channels)

include("blockMatching.jl") #SLOWWWWWWW
include("patEstimation.jl") #SLOWWWW because of ↓↓↓↓↓↓
include("WNNM.jl")          #CORE and its SLOWWWWWWWW

function denoising(Img_O::Array{Float32,3}, Img_N::Array{Float32,3}, Par::PAR)
    #use Iterative method to get Estimated Image
    Img_E = Img_N

    (H,W,D) = size(Img_N)

    #the total number of patches
    TotalPatNum = (H - Par.patsize + 1) * (W - Par.patsize + 1)

    (Neighbor_arr, Num_arr, Self_arr) = neighborIndex(H, W, Par)
    NL_mat              =   zeros(Int32,   Par.patnum[], length(Num_arr))

    #PSNR reflect the quality of `Img_E`
    PSNR_arr            =   zeros(Float64, 1, Par.Iter)

    for iter in 1 : Par.Iter

        #Iterative method
        Img_E =	Img_E + Par.delta * (Img_N - Img_E)

        #image to patch `CurPat` and estimate local noise variance `Sigma_arr`
        (CurPat, Sigma_arr)	= im2patch(Img_E, Img_N, Par)

        if (mod(iter - 1, Par.Innerloop) == 0)

            #Lower noise level, less None-local patches
            Par.patnum[] = Par.patnum[] - 10

            #Caculate Non-local similar patches for each
            NL_mat  =  blockMatching(CurPat, Par, Neighbor_arr, Num_arr, Self_arr)

            #First Iteration use the input noise parameter
            if(iter == 1)
                Sigma_arr = ones(Float32, size(Sigma_arr)) .* Par.nSig
            end
        end
        #Estimate all the patches
        (EPat, WPat) = patEstimation(NL_mat, Self_arr, Sigma_arr, permutedims(CurPat, (2,3,1)), Par)

        #get Estimated image from patches
        Img_E = patch2Im(EPat, WPat, Par.patsize, H, W, D)

        #record PSNR in each Iteration
        PSNR_arr[iter]  = psnr(Img_O, Img_E)
        println("in Iteration:", iter, " PSNR = ", PSNR_arr[iter])

        #saturation situation
        if (iter > 1) && (PSNR_arr[iter] - PSNR_arr[iter-1] < 0.01)
            break
        end

    end
    Img_E
end

function neighborIndex(H, W, Par::PAR)
    SW          =   Par.SearchWin
    s           =   Par.step
    patsr       =   H - Par.patsize + 1
    patsc       =   W - Par.patsize + 1

    #Generate the GridIndex of patches
    r_Grid	    =   Vector(1 : s : patsr)::Array{Int64,1}
    r_Grid	    =   [r_Grid; (last(r_Grid) + 1):patsr]
    c_Grid	    =   Vector(1 : s : patsc)::Array{Int64,1}
    c_Grid	    =   [c_Grid; (last(c_Grid) + 1):patsc]

    #index of patches ( col-first↓ )
    IdxV        =   Vector(1 : patsr * patsc):: Array{Int64,1}
    Idx         =   reshape(IdxV, patsr, patsc):: Array{Int64,2}

    r_Gridn     =   length(r_Grid)
    c_Gridm     =   length(c_Grid)

    Neighbor_arr    =   zeros(Int32, (2 * SW + 1)^2, r_Gridn * c_Gridm)
    Num_arr         =   zeros(Int32,  r_Gridn * c_Gridm)
    SelfIndex_arr   =   zeros(Int32,  r_Gridn * c_Gridm)

    for i in 1 : r_Gridn
        for j in 1 : c_Gridm
            #offsrt r and c of pats
            OffsetR     =   r_Grid[i] :: Int64
            OffsetC     =   c_Grid[j] :: Int64

            #Offset1 is the index of pats(the left-up conor of grid)
            Offsetpats   	=  (OffsetC - 1) * patsr + OffsetR :: Int64

            #Offset2 is the index of grid
            Offsetgrid   	=  (j - 1) * r_Gridn + i :: Int64

            #make sure the windows is in the valid region
            top         =   max(OffsetR - SW, 1)
            button      =   min(OffsetR + SW, patsr)
            left        =   max(OffsetC - SW, 1)
            right       =   min(OffsetC + SW, patsc)

            NL_IdxV     =   Idx[top:button, left:right] :: Array{Int64,2}
            NL_Idx     =   NL_IdxV[:]                   :: Array{Int64,1}

            Num_arr[Offsetgrid]  =  length(NL_Idx)

            Neighbor_arr[1:Num_arr[Offsetgrid], Offsetgrid]  =  NL_Idx[1:Num_arr[Offsetgrid]]
            SelfIndex_arr[Offsetgrid] = Offsetpats
        end
    end
    (Neighbor_arr, Num_arr, SelfIndex_arr)
end

function im2patch(Img_E::Array{Float32,3}, Img_N::Array{Float32,3}, Par::PAR)

    (H, W, D) = size(Img_E)
    N           =   H - Par.patsize + 1
    M           =   W - Par.patsize + 1
    TotalPatNum =   N * M

    Y           =   zeros(Float32, D, Par.patsize^2, TotalPatNum)       #Current Patches
    N_Y         =   zeros(Float32, D, Par.patsize^2, TotalPatNum)       #Noisy   Patches

    k = 0
    for i in 1 : Par.patsize
        for j in 1 : Par.patsize
            k += 1

            n = H - Par.patsize + i
            m = W - Par.patsize + j

            E_patch      =  Img_E[i:n, j:m, :]::Array{Float32,3}
            N_patch      =  Img_N[i:n, j:m, :]::Array{Float32,3}

            Y[:,k,:]     .=  reshape(E_patch, (N*M, D))'
            N_Y[:,k,:]   .=  reshape(N_patch, (N*M, D))'

        end
    end

    #Estimated Local Noise Level
    SigmaArr = Par.lamada .* sqrt.(abs.(dropdims(mean((N_Y .- Y).^2, dims=2), dims=2)' .- Par.nSig^2 ))

    (Y, SigmaArr)
end

function patch2Im(ImPat::Array{Float32,3}, WPat::Array{Int32,3}, PatSize, H, W, D)
    N       =   H - PatSize + 1
    M       =   W - PatSize + 1
    TempOffsetR  =   Vector(1:N)
    TempOffsetC  =   Vector(1:M)

    Img_E  	=  zeros(Float32,H, W, D)
    Img_W 	=  zeros(Float32,H, W, D)

    k = 0
    for i in 1 : PatSize
        for j in 1 : PatSize
            k += 1

            Img_E[TempOffsetR .+ (i-1), TempOffsetC .+ (j-1), :]  += reshape( ImPat[k,:,:], (N, M, D))
            Img_W[TempOffsetR .+ (i-1), TempOffsetC .+ (j-1), :]  += reshape( WPat[k,:,:],  (N, M, D))

        end
    end
    Img_E   =  Float32.(Img_E ./ (Img_W .+ eps()))
end


text:
module IterativeClosestPoint

using NearestNeighbors
using Distances
using CoordinateTransformations
using Rotations
using StaticArrays
using Statistics
using LinearAlgebra

export kabsch
export iterate_kabsch

include("kabsch.jl")
include("correspondence.jl")
include("iterate.jl")
include("goicp/uncertaintyregion.jl")

include("tmalign/align.jl")
include("tmalign/tmscore.jl")
include("tmalign/needlemanwunsch.jl")

end


text:
# The map_subscripts function is stolen from Symbolics.jl
const IndexMap = Dict{Char,Char}(
    '-' => '₋',
    '0' => '₀',
    '1' => '₁',
    '2' => '₂',
    '3' => '₃',
    '4' => '₄',
    '5' => '₅',
    '6' => '₆',
    '7' => '₇',
    '8' => '₈',
    '9' => '₉')

function map_subscripts(indices)
    str = string(indices)
    join(IndexMap[c] for c in str)
end

# Define a type hierarchy to describe a variable declaration. This is mainly for convenient pattern matching later.
abstract type VarDecl end

struct SymDecl <: VarDecl
    sym :: Symbol
end

struct NamedDecl <: VarDecl
    name :: String
    rest :: VarDecl
end

struct FunctionDecl <: VarDecl
    rest :: VarDecl
end

struct TensorDecl <: VarDecl
    ranges :: Vector{AbstractRange}
    rest :: VarDecl
end

struct AssumptionsDecl <: VarDecl
    assumptions :: Vector{Symbol}
    rest :: VarDecl
end

# Transform a Decl struct in an Expression that calls SymPy to declare the corresponding symbol
function gendecl(x::VarDecl)
    asstokw(a) = Expr(:kw, esc(a), true)
    val = :($(ctor(x))($(name(x, missing)), $(map(asstokw, assumptions(x))...)))
    :($(esc(sym(x))) = $(genreshape(val, x)))
end

# Transform an expression in a Decl struct
function parsedecl(expr)
    # @syms x
    if isa(expr, Symbol)
        return SymDecl(expr)

    # @syms x::assumptions, where assumption = assumptionkw | (assumptionkw...)
    elseif isa(expr, Expr) && expr.head == :(::)
        symexpr, assumptions = expr.args
        assumptions = isa(assumptions, Symbol) ? [assumptions] : assumptions.args
        return AssumptionsDecl(assumptions, parsedecl(symexpr))

    # @syms x=>"name"
    elseif isa(expr, Expr) && expr.head == :call && expr.args[1] == :(=>)
        length(expr.args) == 3 || parseerror()
        isa(expr.args[3], String) || parseerror()

        expr, strname = expr.args[2:end]
        return NamedDecl(strname, parsedecl(expr))

    # @syms x()
    elseif isa(expr, Expr) && expr.head == :call && expr.args[1] != :(=>)
        length(expr.args) == 1 || parseerror()
        return FunctionDecl(parsedecl(expr.args[1]))

    # @syms x[1:5, 3:9]
    elseif isa(expr, Expr) && expr.head == :ref
        length(expr.args) > 1 || parseerror()
        ranges = map(parserange, expr.args[2:end])
        return TensorDecl(ranges, parsedecl(expr.args[1]))
    else
        parseerror()
    end
end

function parserange(expr)
    range = eval(expr)
    isa(range, AbstractRange) || parseerror()
    range
end

sym(x::SymDecl) = x.sym
sym(x::NamedDecl) = sym(x.rest)
sym(x::FunctionDecl) = sym(x.rest)
sym(x::TensorDecl) = sym(x.rest)
sym(x::AssumptionsDecl) = sym(x.rest)

ctor(::SymDecl) = :symbols
ctor(x::NamedDecl) = ctor(x.rest)
ctor(::FunctionDecl) = :SymFunction
ctor(x::TensorDecl) = ctor(x.rest)
ctor(x::AssumptionsDecl) = ctor(x.rest)

assumptions(::SymDecl) = []
assumptions(x::NamedDecl) = assumptions(x.rest)
assumptions(x::FunctionDecl) = assumptions(x.rest)
assumptions(x::TensorDecl) = assumptions(x.rest)
assumptions(x::AssumptionsDecl) = x.assumptions

# Reshape is not used by most nodes, but TensorNodes require the output to be given
# the shape matching the specification. For instance if @syms x[1:3, 2:6], we should
# have size(x) = (3, 5)
genreshape(expr, ::SymDecl) = expr
genreshape(expr, x::NamedDecl) = genreshape(expr, x.rest)
genreshape(expr, x::FunctionDecl) = genreshape(expr, x.rest)
genreshape(expr, x::TensorDecl) = let
    shape = tuple(length.(x.ranges)...)
    :(reshape(collect($(expr)), $(shape)))
end
genreshape(expr, x::AssumptionsDecl) = genreshape(expr, x.rest)

# To find out the name, we need to traverse in both directions to make sure that each node can get
# information from parents and children about possible name.
# This is done because the expr tree will always look like NamedDecl -> ... -> TensorDecl -> ... -> SymDecl
# and the TensorDecl node will need to know if it should create names base on a NamedDecl parent or
# based on the SymDecl leaf.
name(x::SymDecl, parentname) = coalesce(parentname, String(x.sym))
name(x::NamedDecl, parentname) = coalesce(name(x.rest, x.name), x.name)
name(x::FunctionDecl, parentname) = name(x.rest, parentname)
name(x::AssumptionsDecl, parentname) = name(x.rest, parentname)
name(x::TensorDecl, parentname) = let
    basename = name(x.rest, parentname)
    # we need to double reverse the indices to make sure that we traverse them in the natural order
    namestensor = map(Iterators.product(x.ranges...)) do ind
        sub = join(map(map_subscripts, ind), "_")
        string(basename, sub)
    end
    join(namestensor[:], ", ")
end

function parseerror()
    error("Incorrect @syms syntax. Try `@syms x::(real,positive)=>\"x₀\" y() z::complex n::integer` for instance.")
end


text:
cpu(x) = x
gpu(x) = x

cpu(x::CuArray) = adapt(Array, x)
gpu(x::AbstractArray) = cu(x)

function cpu(structure::Union{Network,DelayNetwork,NeuronType,LearningRule})
    return typeof(structure)(
        [cpu(getfield(structure, property)) for property in propertynames(structure)]...
    )
end

function gpu(structure::Union{Network,DelayNetwork,NeuronType,LearningRule})
    return typeof(structure)(
        [gpu(getfield(structure, property)) for property in propertynames(structure)]...
    )
end


text:
using Lazy: @_
using DataStructures
using MacroTools
import Base: run, display
export addtrial, addbreak, addpractice, moment, await_response, record, timeout,
  when, looping, @addtrials

function findkwd(kwds,sym,default)
  for (k,v) in kwds
    if k == sym
      return v
    end
  end

  default
end

const null_record = []
function write_csv_line(exp::Experiment{NullWindow},header,kwds)
  push!(null_record,Dict{Symbol,Any}(map(c -> c => findkwd(kwds,c,""),header)))
end

function write_csv_line(exp::Experiment{SDLWindow},header,kwds)
  if !isnull(info(exp).file)
    open(get(info(exp).file),"a") do stream
      @_ (string(findkwd(kwds,c,"")) for c in header) begin
        join(_,",")
        println(stream,_)
      end
    end
  end
end

function record_helper(exp::Experiment,kwds,header)
  columns = map(x -> x[1],kwds)

  if !isempty(columns) && !all(map(c -> c ∈ header,columns))
    missing = collect(Iterators.filter(c -> c ∉ header,columns))

    error("Unexpected column $(length(missing) > 1 ? "s" : "")"*
          "$(join(missing,", "," and ")). "*
          "Make sure you specify all columns you plan to use "*
          "during experiment initialization.")
  end

  kwds = reverse(kwds) # this ensures that if the user overwrites a value
                       # it will be honored

  write_csv_line(exp,header,kwds)
end

function record_header(exp::Experiment)
  extra_keys = [:weber_version,:start_date,:start_time,:offset,:trial,:time]
  info_keys = map(x->x[1],info(exp).values)

  reserved_keys = Set([extra_keys...;info_keys...])
  reserved = filter(x -> x ∈ reserved_keys,info(exp).header)
  if length(reserved) == 1
    error("The column name \"$(reserved[1])\" is reserved. Please use "*
          " a different name.")
  elseif length(reserved) > 1
    error("The column names "*
          join(map(x -> "\""*x*"\"",reserved),", "," and ")*
          " are reserved. Please use different names.")
  end

  columns = [extra_keys...,info_keys...,:code,info(exp).header...]
  if !isnull(info(exp).file)
    open(x -> println(x,join(columns,",")),get(info(exp).file),"w")
  end
end

function record(exp::ExtendedExperiment,code;kwds...)
  record(next(exp),code;kwds...)
end
function record{T <: BaseExperiment}(exp::T,code;kwds...)
  extra = [:weber_version => Weber.version,
           :start_date => Dates.format(info(exp).start,"yyyy-mm-dd"),
           :start_time => Dates.format(info(exp).start,"HH:MM:SS"),
           :offset => data(exp).offset,
           :trial => data(exp).trial,
           :time => data(exp).last_time]

  info_keys = map(x->x[1],info(exp).values)
  extra_keys = map(x->x[1],extra)
  record_helper(exp,tuple(extra...,info(exp).values...,:code => code,kwds...),
                [extra_keys...,info_keys...,:code,info(exp).header...])
end

"""
    record(code;keys...)

Record a row to the experiment data file using a given `code`.

Each event has a code which identifies it as being a particular type of
experiment event. This is normally a string. Each keyword argument is the value
of a column (with the same name). By convention when you record something with
the same code you should specify the same set of columns.

All calls to record also result in many additional values being written to the
data file. The start time and date of the experiment, the trial and offset
number, the version of Weber, and the time at which the last moment started are
all stored.  Additional information can be added during creation of the
experiment (see [`Experiment`](@ref)).

Each call to record writes a new row to the data file used for the experiment, so
there should be no loss of data if the program is terminated prematurely for
some reason.

!!! note "Automatically Recorded Codes"

    There are several codes that are automatically recorded by Weber.
    They include:
    1. **trial_start** - recorded at the start of moments added by `addtrial`
    2. **practice_start** - recorded at the start of moments added by `addpractice`
    3. **break_start** - recorded at the start of moments added by `addbreak`
    4. **high_latency** - recorded whenever a high latency warning is triggered.
       The "value" column is set to the error between the actual and the desired
       timing of a moment, in seconds.
    5. **paused** - recorded when user hits 'escape' and the experiment is
       paused.
    6. **unpaused** - recorded when the user ends the pause, continuuing the
       experiment.
    7. **terminated** - recorded when the user manually terminates the
       experiment (via 'escape')
    8. **closed** - recorded just before the experiment window closes
"""
function record(code;kwds...)
  record(get_experiment(),code;kwds...)
end

function addmoment(q::ExpandingMoment,m::AbstractMoment)
  m = flag_expanding(m)
  if !isempty(q.data) && can_continue_sequence(m) && sequenceable(top(q.data))
    m = sequence(pop!(q.data),m)
  end
  push!(q.data,m)
end

function addmoment(q::MomentQueue,m::AbstractMoment)
  if !isempty(q) && can_continue_sequence(m) && sequenceable(back(q))
    m = sequence(pop!(q),m)
  end
  enqueue!(q,m)
end

function addmoment(q::Vector{AbstractMoment},m::AbstractMoment)
  if !isempty(q) && can_continue_sequence(m) && sequenceable(last(q))
    m = sequence(pop!(q),m)
  end
  push!(q,m)
end

addmoment(e::Experiment,m) = addmoment(data(e).moments,m)
addmoment(q::Array{MomentQueue},m::AbstractMoment) = addmoment(first(q),m)
function addmoment(q::Union{ExpandingMoment,MomentQueue,Array{MomentQueue}},watcher::Function)
  for t in concrete_events
    precompile(watcher,(t,))
  end
  addmoment(q,moment(() -> data(get_experiment()).trial_watcher = watcher))
end
function addmoment(q,ms)
  function handle_error()
    if !(typeof(ms) <: AbstractMoment || typeof(ms) <: Function)
      error("Expected some kind of moment, but got a value of type",
            " $(typeof(ms)) instead.")
    else
      error("Cannot add moment to an object of type $(typeof(q))")
    end
  end

  try
    first(ms)
  catch e
    if e isa MethodError
      handle_error()
    else
      rethrow(e)
    end
  end

  for m in ms
    # some types iterate over themselves (e.g. numbers);
    # check for this to avoid infinite recursion
    if m == ms
      handle_error()
    end
    addmoment(q,m)
  end
  q
end

const addtrial_block = Stack(ExpandingMoment)
function addmoments(exp,moments)
  if isempty(addtrial_block)
    foreach(m -> addmoment(exp,m),moments)
  else
    block = top(addtrial_block)
    foreach(m -> addmoment(block,m),moments)
  end
end

"""
    @addtrials expr...

Marks a let block, a for loop, or an if expression as dependent on experiment
[run-time](@ref setup_time) state, leaving the offset counter unincremented within that block.  The
immediately proceeding loop or conditional logic will be run during
experiment run-time rather than setup-time.

Refer to the [`Advanced Topics`](advanced.md) of the manual section for
more details.
"""
macro addtrials(expr)
  if isexpr(expr,:let)
    quote
      trial_block(() -> true) do
        $(esc(expr))
      end
    end
  elseif isexpr(expr,:if)
    cond,ifbody,elsebody = @match expr begin
      if cond_
        ifbody_
      end => (cond,ifbody,nothing)

      if cond_
        ifbody_
      else
        elsebody_
      end => (cond,ifbody,elsebody)
    end

    elsebody = @match elsebody begin
      ifelse_if => :(@addtrials($ifelse))
      begin
        ifelse_if
      end => :(@addtrials($ifelse))
      other_ => elsebody
    end

    if elsebody == nothing
      quote
        trial_block(() -> $(esc(cond))) do
          $(esc(ifbody))
        end
      end
    else
      quote
        let ifpassed = false
          trial_block(() -> ifpassed = $(esc(cond))) do
            $(esc(ifbody))
          end
          trial_block(() -> !ifpassed) do
            $(esc(elsebody))
          end
        end
      end
    end
  elseif isexpr(expr,:while)
    cond,body = @match expr begin
      while cond_
        body_
      end => (cond,body)
    end

    quote
      trial_block(loop=true,() -> $(esc(cond))) do
        $(esc(body))
      end
    end
  else
    error("@addtrials expects a `let`, `if` or `while` expression.")
  end
end

function trial_block(body,condition;keys...)
  trial_block(get_experiment(),body,condition;keys...)
end

function trial_block(exp::Experiment,body::Function,condition::Function;loop=false)
  moment = ExpandingMoment(condition,Stack(AbstractMoment),loop,true)
  push!(addtrial_block,moment)
  body()
  pop!(addtrial_block)

  addmoments(exp,[moment])
end

"""
    addtrial(moments...)

Adds a trial to the experiment, consisting of the specified moments.

Each trial records a "trial_start" code, and increments a counter tracking the
number of trials, and (normally) an offset counter. These two numbers are
reported on every line of the resulting data file (see [`record`](@ref)). They can be
retrieved using [`Weber.trial()`](@ref) and [`Weber.offset()`](@ref).
"""

function addtrial(moments...)
  addtrial_helper(get_experiment(),"trial_start",moments)
end

addtrial(exp::ExtendedExperiment,moments...) = addtrial(next(exp),moments...)
function addtrial{T <: BaseExperiment}(exp::T,moments...)
  addtrial_helper(exp,"trial_start",moments)
end

function addtrial_helper(exp::Experiment,start_code,moments)
  start_trial = offset_start_moment(start_code == "trial_start") do
    record(start_code)
  end

  addmoments(exp,[start_trial,moments])
end

"""
    addpractice(moments...)

Identical to [`addtrial`](@ref), except that it does not incriment the trial count,
and records a "practice_start" instead of "trial_start" code.
"""
function addpractice(moments...)
  addtrial_helper(get_experiment(),"practice_start",moments)
end

addpractice(exp::ExtendedExperiment,moments...) = addpractice(next(exp),moments...)
function addpractice{T <: BaseExperiment}(exp::T,moments...)
  addtrial_helper(exp,"practice_start",moments...)
end

"""
    addbreak(moments...)

Identical to [`addpractice`](@ref), but records "break_start" instead of "practice_start".
"""
function addbreak(moments...)
  addbreak(get_experiment(),moments...)
end

addbreak(exp::ExtendedExperiment,moments...) = addbreak(next(exp),moments...)
function addbreak{T <: BaseExperiment}(exp::T,moments...)
  addtrial_helper(exp,"break_start",moments)
end

"""
    moment([delta_t],[fn],args...;keys...)

Create a moment that occurs `delta_t` (default 0) seconds after the onset of
the previous moment, running the specified function.

The function `fn` is passed the arguments specified in `args` and `keys`.
"""
function moment(delta_t::Number=0.0s,fn::Function=()->nothing,args...;keys...)
  precompile(fn,map(typeof,args))
  delta_t = ustrip(inseconds(delta_t))
  TimedMoment(delta_t,() -> fn(args...;keys...),stacktrace()[2:end])
end

function moment(fn::Function,args...;keys...)
  moment(0.0s,fn,args...;keys...)
end

const PlayFunction = typeof(play)
function moment(delta_t::Number,::PlayFunction,x;channel=0)
  PlayMoment(ustrip(inseconds(delta_t)),playable(x),channel,stacktrace()[2:end])
end
function moment(delta_t::Number,::PlayFunction,fn::Function;channel=0)
  PlayFunctionMoment(ustrip(inseconds(delta_t)),fn,channel,stacktrace()[2:end])
end

# const StreamFunction = typeof(TimedSound.stream)
# function moment(delta_t::Number,::StreamFunction,itr,channel::Int)
#   StreamMoment(ustrip(inseconds(delta_t)),itr,channel,stacktrace()[2:end])
# end

const DisplayFunction = typeof(display)
function moment(delta_t::Number,::DisplayFunction,x;keys...)
  DisplayMoment(ustrip(inseconds(delta_t)),visual(x;keys...),stacktrace()[2:end])
end
function moment(delta_t::Number,::DisplayFunction,fn::Function;keys...)
  DisplayFunctionMoment(ustrip(inseconds(delta_t)),fn,keys,stacktrace()[2:end])
end

"""
    moment(moments...)
    moment(moments::Array)

Create a single, compound moment by concatentating several moments togethor.
"""
moment(moments...) = moment(collect(moments))
moment(moments::Array) = CompoundMoment(addmoment(Array{AbstractMoment,1}(),moments))
function moment(moments)
  try
    start(moments)
  catch
    throw(MethodError(moment,typeof(moments)))
  end
  moment(collect(Any,moments))
end

function offset_start_moment(fn::Function=()->nothing,count_trials=false)
  precompile(fn,(Float64,))
  OffsetStartMoment(fn,count_trials,false,stacktrace()[2:end])
end

"""
    await_response(isresponse;[atleast=0.0])

This moment starts when the `isresponse` function evaluates to true.

The `isresponse` function will be called anytime an event occurs. It should
take one parameter (the event that just occured).

If the response is provided before `atleast` seconds, the moment does not start
until `atleast` seconds have passed.
"""
function await_response(fn::Function;atleast=0.0)
  for t in concrete_events
    precompile(fn,(t,))
  end

  ResponseMoment(fn,() -> nothing,0,atleast,stacktrace()[2:end])
end

"""
    timeout(fn,isresponse,timeout,[atleast=0s])

This moment starts when either `isresponse` evaluates to true or
timeout time (in seconds) passes.

The `isresponse` function will be called anytime an event occurs. It should
take one parameter (the event that just occured).

If the moment times out, the function `fn` (with no arguments) will be called.

If the response is provided before `atleast` seconds, the moment does not begin
until `atleast` seconds (`fn` will not be called).
"""
function timeout(fn::Function,isresponse::Function,timeout;atleast=0.0s)
  precompile(fn,())
  for t in concrete_events
    precompile(isresponse,(t,))
  end

  ResponseMoment(isresponse,fn,ustrip(inseconds(timeout)),
                 ustrip(inseconds(atleast)),stacktrace()[2:end])
end

flag_expanding(m::AbstractMoment) = m
function flag_expanding(m::OffsetStartMoment)
  OffsetStartMoment(m.run,m.count_trials,true,m.trace)
end
function flag_expanding(m::ExpandingMoment)
  if m.update_offset
    ExpandingMoment(m.condition,m.data,m.repeat,false)
  else
    m
  end
end

"""
    looping(when=fn,moments...)

This moment will begin at the *start* of the previous moment, and repeats the
listed moments (possibly in nested iterable objects) until the `when` function
(which takes no arguments) evaluates to false.
"""
function looping(moments...;when=() -> error("infinite loop!"))
  Weber.when(when,moments...;loop=true)
end

"""
    when(condition,moments...)

This moment will begin at the *start* of the previous moment, and presents the
following moments (possibly in nested iterable objects) if the `condition`
function (which takes no arguments) evaluates to true.
"""
function when(condition::Function,moments...;loop=false,update_offset=false)
  precompile(condition,())
  e = ExpandingMoment(condition,Stack(AbstractMoment),loop,update_offset)
  foreach(m -> addmoment(e,m),moments)
  e
end

"""
    Weber.prepare!(m,[onset_s])

If there is anything the moment needs to do before it occurs, it is done during
`prepare!`. Prepare can be used to set up precise timing even when hardware
latency is high, if that latency can be predicted, and accounted for. A moment's
prepare! method is called just before the first non-zero pause between moments
that occurs before this moment: in the simplest case, when this moment has a
non-zero value for [`delta_t`](@ref), preapre! will occur `delta_t` seconds
before this moment. However, if several moments with no pause occur, prepare!
will occur before all of those moments as well.

Prepare accepts an optional second argument used to indicate the time, in
seconds from the start of the experiemnt when this moment will begin (as a
Float64).  This argument may be Inf, indicating that it is not possible to
predict when the moment will occur at this point, because the timing depends on
some stateful information (e.g. a participant's response). It is accetable in
this case to throw an error, explaining that this kind of moment must be able to
know precisely when it occurs to be prepared.

!!! note

    This method is part of the private interface for moments. It
    should not be called directly, but implemented as part of an extension.
    You need only extend the method taking a single arugment unless you
    intend to use this information during prepartion.
"""
prepare!(m::AbstractMoment,onset_s::Float64) = prepare!(m)
prepare!(m::AbstractMoment) = nothing
function prepare!(ms::MomentSequence,onset_s::Float64)
  for m in ms.data prepare!(m,onset_s) end
end

function prepare!(m::DisplayFunctionMoment)
  m.visual = Nullable(visual(m.fn();m.keys...))
end

function prepare!(m::PlayMoment,onset_s::Float64)
  if !isinf(onset_s)
    TimedSound.play_(m.sound,onset_s,m.channel)
  else
    m.prepared = true
  end
end

function prepare!(m::PlayFunctionMoment,onset_s::Float64)
  if !isinf(onset_s)
    TimedSound.play_(playable(m.fn()),onset_s,m.channel)
  else
    m.prepared = Nullable(playable(m.fn()))
  end
end

run(exp,q,m::TimedMoment) = m.run()
run(exp,q,m::OffsetStartMoment) = m.run()
run(exp,q,m::MomentSequence) = foreach(x -> run(exp,q,x),m.data)

run(exp,q,m::DisplayMoment) = display(win(exp),m.visual)
run(exp,q,m::DisplayFunctionMoment) = display(win(exp),get(m.visual))

function run(exp,q,m::PlayMoment)
  if m.prepared
    m.prepared = false
    TimedSound.play_(m.sound,0.0,m.channel)
  end
end

function run(exp,q,m::PlayFunctionMoment)
  if !isnull(m.prepared)
    TimedSound.play_(get(m.prepared),0.0,m.channel)
    m.prepared = Nullable()
  end
end

"""
    handle(exp,queue,moment,to_handle)

Internal method to handle the given moment object in a manner specific to its
type.

The function `handle` is only called when the appropriate time has been
reached for the next moment to be presented (according to [`delta_t`](@ref)) or
when an event occurs.

The `to_handle` object is either a `Float64`, indicating the current experiment
time, or it is an `ExpEvent` indicating the event that just occured. As an
example, a timed moment, will run when it recieves any `Float64` value, but
nothing occurs when passed an event.

The queue is a `MomentQueue` object, which has the same interface as the
`Dequeue` object (from the `DataStructures` package) but it is also
iterable. Upon calling handle, `top(queue) == moment`.

Handle should return a boolean indicating whether the event was "handled" or
not. If unhandled, the moment should remain on top of the queue. If returning
true, handle should *normally* remove the top moment from the queue. Exceptions
exist (for instance, to allow for loops), but one does not normally need to
implement custom moments that have such behavior.

!!! note

    This method is part of the private interface for moments. It
    should not be called directly, but implemented as part of an extension.
    It is called during the course of running an experiment.

"""
function handle(exp::Experiment,q::MomentQueue,
                moment::AbstractTimedMoment,time::Float64)
  run(exp,q,moment)
  q.last = time
  dequeue!(q)
  true
end

# function handle(exp::Experiment,q::MomentQueue,
#                 moment::StreamMoment,time::Float64)
#   stream(moment.itr,moment.channel)
#   dequeue!(q)
#   true
# end

function handle(exp::Experiment,q::MomentQueue,
                moment::AbstractTimedMoment,event::ExpEvent)
  false
end

function handle(exp::Experiment,q::MomentQueue,
                moment::ResponseMoment,time::Float64)
  moment.timeout()
  q.last = time
  dequeue!(q)
  true
end

function handle(exp::Experiment,q::MomentQueue,m::ResponseMoment,event::ExpEvent)
  if m.respond(event)
    dequeue!(q)
    unshift!(q,ResponseMomentMin(max(0.0,m.minimum_delta_t),m.trace))
    true
  end
  false
end

function handle(exp::Experiment,q::MomentQueue,
                moment::ResponseMomentMin,time::Float64)
  q.last = time
  dequeue!(q)
  true
end

function handle(exp::Experiment,q::MomentQueue,
                moment::ResponseMomentMin,evt::ExpEvent)
  false
end

function handle(exp::Experiment,q::MomentQueue,moments::CompoundMoment,x)
  compq = MomentQueue(moments.data,q.last)
  push!(data(exp).moments,compq)
  prepare!(compq,q.last)
  dequeue!(q)
  true
end

function handle(exp::Experiment,q::MomentQueue,m::ExpandingMoment,x)
  if m.condition()
    if !m.repeat
      dequeue!(q)
    end

    for x in m.data
      unshift!(q,x)
    end
    unshift!(q,expanding_stub)
  else
    dequeue!(q)
  end
  true
end

function handle(exp::Experiment,q::MomentQueue,m::ExpandingMomentStub,x)
  dequeue!(q)
  true
end

is_moment_skipped(exp,moment::AbstractMoment) = data(exp).offset < data(exp).skip_offsets
function is_moment_skipped(exp,moment::OffsetStartMoment)
  if !moment.expanding
    data(exp).offset += 1
  end
  if moment.count_trials
    data(exp).trial += 1
  end
  data(exp).offset < data(exp).skip_offsets
end
function is_moment_skipped(exp,moment::ExpandingMoment)
  if moment.update_offset
    data(exp).offset += 1
    # each expanding moment only ever incriments the offset
    # counter once, event if it creates a loop.
    moment.update_offset = false
  end
  data(exp).offset < data(exp).skip_offsets
end


text:
oyoy = "Salut"

include("../F.jl")


text:
function solve_boat(
    t_end;
    x0 = 0.0,
    z0 = 0.0,
    θ0 = 0.0,
    u0 = 0.0,
    w0 = 0.0,
    θ′0 = 0.0,
    settings = Boat_settings(),
)
    # "decode" settings
    p = Boat_ode_params(θ0, settings)

    # Compute center of mass location to get consistent initial conditions
    M = system_lhs(Boat_angles(p.angle_ranges, p.timing, 0.0))
    z_CF = M[1, 3] / m
    x_CF = -M[2, 3] / m
    y0 = [x0, z0, θ0, u0 - θ′0 * z_CF, w0 + θ′0 * x_CF, θ′0]

    tspan = (0.0, t_end)
    prob = ODEProblem(boatode!, y0, tspan, p)

    solve(prob, abstol = 1e-9, reltol = 1e-7, saveat = 0.01)
end