text:
# # Example 1: Optimization of a State-to-State Transfer in a Two-Level-System

#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`simple_state_to_state.ipynb`](@__NBVIEWER_ROOT_URL__/examples/simple_state_to_state.ipynb).
#md #
#md #     Compare this example against the [same example using the `krotov`
#md #     Python package](https://qucontrol.github.io/krotov/v1.2.1/notebooks/01_example_simple_state_to_state.html).

#md # ``\gdef\op#1{\hat{#1}}``
#md # ``\gdef\init{\text{init}}``
#md # ``\gdef\tgt{\text{tgt}}``

#nb # $
#nb # \newcommand{tr}[0]{\operatorname{tr}}
#nb # \newcommand{diag}[0]{\operatorname{diag}}
#nb # \newcommand{abs}[0]{\operatorname{abs}}
#nb # \newcommand{pop}[0]{\operatorname{pop}}
#nb # \newcommand{aux}[0]{\text{aux}}
#nb # \newcommand{opt}[0]{\text{opt}}
#nb # \newcommand{tgt}[0]{\text{tgt}}
#nb # \newcommand{init}[0]{\text{init}}
#nb # \newcommand{lab}[0]{\text{lab}}
#nb # \newcommand{rwa}[0]{\text{rwa}}
#nb # \newcommand{bra}[1]{\langle#1\vert}
#nb # \newcommand{ket}[1]{\vert#1\rangle}
#nb # \newcommand{Bra}[1]{\left\langle#1\right\vert}
#nb # \newcommand{Ket}[1]{\left\vert#1\right\rangle}
#nb # \newcommand{Braket}[2]{\left\langle #1\vphantom{#2}\mid{#2}\vphantom{#1}\right\rangle}
#nb # \newcommand{op}[1]{\hat{#1}}
#nb # \newcommand{Op}[1]{\hat{#1}}
#nb # \newcommand{dd}[0]{\,\text{d}}
#nb # \newcommand{Liouville}[0]{\mathcal{L}}
#nb # \newcommand{DynMap}[0]{\mathcal{E}}
#nb # \newcommand{identity}[0]{\mathbf{1}}
#nb # \newcommand{Norm}[1]{\lVert#1\rVert}
#nb # \newcommand{Abs}[1]{\left\vert#1\right\vert}
#nb # \newcommand{avg}[1]{\langle#1\rangle}
#nb # \newcommand{Avg}[1]{\left\langle#1\right\rangle}
#nb # \newcommand{AbsSq}[1]{\left\vert#1\right\vert^2}
#nb # \newcommand{Re}[0]{\operatorname{Re}}
#nb # \newcommand{Im}[0]{\operatorname{Im}}
#nb # $

# This first example illustrates the basic use of the `Krotov.jl` by solving a
# simple canonical optimization problem: the transfer of population in a two
# level system.

using Printf
using QuantumControl
using LinearAlgebra
using QuantumControlBase: chain_infohooks
using GRAPELinesearchAnalysis
using LineSearches
using PyPlot: matplotlib
matplotlib.use("Agg")

#jl using Test

# ## Two-level Hamiltonian

# We consider the Hamiltonian $\op{H}_{0} = - \frac{\omega}{2} \op{\sigma}_{z}$, representing
# a simple qubit with energy level splitting $\omega$ in the basis
# $\{\ket{0},\ket{1}\}$. The control field $\epsilon(t)$ is assumed to couple via
# the Hamiltonian $\op{H}_{1}(t) = \epsilon(t) \op{\sigma}_{x}$ to the qubit,
# i.e., the control field effectively drives transitions between both qubit
# states.
#
# We we will use

ϵ(t) = 0.2 * QuantumControl.Shapes.flattop(t, T = 5, t_rise = 0.3, func = :blackman);

#-
"""Two-level-system Hamiltonian."""
function hamiltonian(Ω = 1.0, ϵ = ϵ)
    σ̂_z = ComplexF64[1 0; 0 -1]
    σ̂_x = ComplexF64[0 1; 1 0]
    Ĥ₀ = -0.5 * Ω * σ̂_z
    Ĥ₁ = σ̂_x
    return (Ĥ₀, (Ĥ₁, ϵ))
end;
#-

H = hamiltonian();
#jl @test length(H) == 2

# The control field here switches on from zero at $t=0$ to it's maximum amplitude
# 0.2 within the time period 0.3 (the switch-on shape is half a [Blackman pulse](https://en.wikipedia.org/wiki/Window_function#Blackman_window)).
# It switches off again in the time period 0.3 before the
# final time $T=5$). We use a time grid with 500 time steps between 0 and $T$:

tlist = collect(range(0, 5, length = 500));

#-

function plot_control(pulse::Vector, tlist)
    fig, ax = matplotlib.pyplot.subplots(figsize = (6, 3))
    ax.plot(tlist, pulse)
    ax.set_xlabel("time")
    ax.set_ylabel("amplitude")
    return fig
end

plot_control(ϵ::T, tlist) where {T<:Function} = plot_control([ϵ(t) for t in tlist], tlist)

#!jl plot_control(H[2][2], tlist)

# ## Optimization target

# The `krotov` package requires the goal of the optimization to be described by a
# list of `Objective` instances. In this example, there is only a single
# objective: the state-to-state transfer from initial state $\ket{\Psi_{\init}} =
# \ket{0}$ to the target state $\ket{\Psi_{\tgt}} = \ket{1}$, under the dynamics
# of the Hamiltonian $\op{H}(t)$:

function ket(label)
    result = Dict("0" => Vector{ComplexF64}([1, 0]), "1" => Vector{ComplexF64}([0, 1]))
    return result[string(label)]
end;

#-
#jl @test dot(ket(0), ket(1)) ≈ 0
#-

objectives = [Objective(initial_state = ket(0), generator = H, target_state = ket(1))]

#-
#jl @test length(objectives) == 1
#-

problem = ControlProblem(
    objectives = objectives,
    tlist = tlist,
    pulse_options=Dict(),
    iter_stop = 500,
    J_T = QuantumControl.Functionals.J_T_sm,
    gradient=QuantumControl.Functionals.grad_J_T_sm!,
    check_convergence = res -> begin
        ((res.J_T < 1e-3) && (res.converged = true) && (res.message = "J_T < 10⁻³"))
    end,
);

# ## Simulate dynamics under the guess field

# Before running the optimization procedure, we first simulate the dynamics under the
# guess field $\epsilon_{0}(t)$. The following solves equation of motion for the
# defined objective, which contains the initial state $\ket{\Psi_{\init}}$ and
# the Hamiltonian $\op{H}(t)$ defining its evolution.

guess_dynamics = propagate_objective(
    objectives[1],
    problem.tlist;
    storage = true,
    observables = (Ψ -> abs.(Ψ) .^ 2,),
)

#-

function plot_population(pop0::Vector, pop1::Vector, tlist)
    fig, ax = matplotlib.pyplot.subplots(figsize = (6, 3))
    ax.plot(tlist, pop0, label = "0")
    ax.plot(tlist, pop1, label = "1")
    ax.legend()
    ax.set_xlabel("time")
    ax.set_ylabel("population")
    return fig
end

#!jl plot_population(guess_dynamics[1,:], guess_dynamics[2,:], tlist)

# ## Optimize

# In the following we optimize the guess field $\epsilon_{0}(t)$ such
# that the intended state-to-state transfer $\ket{\Psi_{\init}} \rightarrow
# \ket{\Psi_{\tgt}}$ is solved.

#jl println("")
opt_result = optimize(
        problem;
        method=:grape,
        #=show_trace=true, extended_trace=false,=#
        info_hook=chain_infohooks(
            GRAPELinesearchAnalysis.plot_linesearch(@__DIR__),
            QuantumControl.GRAPE.print_table,
        )
        #=alphaguess=LineSearches.InitialStatic(alpha=0.2),=#
        #=linesearch=LineSearches.HagerZhang(alphamax=2.0),=#
        #=linesearch=LineSearches.BackTracking(), # fails=#
        #=allow_f_increases=true,=#
);
#-
opt_result
#-
#jl display(opt_result)
#jl display(opt_result.optim_res)
#jl @test opt_result.J_T < 1e-3
#-

# We can plot the optimized field:

#-
#!jl plot_control(opt_result.optimized_controls[1], tlist)
#-

#jl using UnicodePlots
#jl println(lineplot(tlist, opt_result.optimized_controls[1]))

# ## Simulate the dynamics under the optimized field

# Having obtained the optimized control field, we can simulate the dynamics to
# verify that the optimized field indeed drives the initial state
# $\ket{\Psi_{\init}} = \ket{0}$ to the desired target state
# $\ket{\Psi_{\tgt}} = \ket{1}$.

opt_dynamics = propagate_objective(
    objectives[1],
    problem.tlist;
    controls_map = IdDict(ϵ => opt_result.optimized_controls[1]),
    storage = true,
    observables = (Ψ -> abs.(Ψ) .^ 2,),
)

#-
#!jl plot_population(opt_dynamics[1,:], opt_dynamics[2,:], tlist)
#-

#-
#jl @test opt_dynamics[2,end] > 0.99
#-


text:
using Chess, Test

include("piece.jl")
include("square.jl")
include("squareset.jl")
include("move.jl")
include("board.jl")
include("san.jl")
include("chess960.jl")
include("movegen.jl")

ext: jl


text:
using ChooseOptimizer, JuMP, Multisets

"""
`matrix_match(A,B)` returns permutation matrices `P` and `Q`
so that `P*A==B*Q` or throws an error if not possible.
"""
function matrix_match(A::AbstractMatrix, B::AbstractMatrix)
    err_msg = "Matrices do not match"

    # Check that A and B have same dimensions and same elements
    r, c = size(A)
    if (r, c) != size(B) || Multiset(collect(A)) != Multiset(collect(B))
        error(err_msg)
    end

    m = Model(get_solver())

    # Set up P and Q are permutation matrices

    @variable(m, P[1:r, 1:r], Bin)
    @variable(m, Q[1:c, 1:c], Bin)

    for i = 1:r
        @constraint(m, sum(P[i, j] for j = 1:r) == 1)
        @constraint(m, sum(P[j, i] for j = 1:r) == 1)
    end

    for i = 1:c
        @constraint(m, sum(Q[i, j] for j = 1:c) == 1)
        @constraint(m, sum(Q[j, i] for j = 1:c) == 1)
    end

    # require P*A == B*Q

    for i = 1:r
        for j = 1:c
            @constraint(
                m,
                sum(P[i, k] * A[k, j] for k = 1:r) == sum(B[i, k] * Q[k, j] for k = 1:c)
            )
        end
    end

    optimize!(m)
    status = Int(termination_status(m))
    if status != 1
        error(err_msg)
    end

    PP = Int.(value.(P))
    QQ = Int.(value.(Q))
    return PP, QQ
end


text:
int_rules_6_3_10 = @theory begin

    #= ::Subsection::Closed:: =#
    #= 6.3.10*(c+d*x)^m*(a+b*tanh)^n =#
    @apply_utils Antiderivative((~u) ^ ~(m') * (~(a') + ~(b') * tanh(~v)) ^ ~(n'), ~x) => Antiderivative(ExpandToSum(~u, ~x) ^ ~m * (~a + ~b * tanh(ExpandToSum(~v, ~x))) ^ ~n, ~x) <-- FreeQ([~a, ~b, ~m, ~n], ~x) && (LinearQ([~u, ~v], ~x) && Not(LinearMatchQ([~u, ~v], ~x)))
    @apply_utils Antiderivative((~u) ^ ~(m') * (~(a') + ~(b') * coth(~v)) ^ ~(n'), ~x) => Antiderivative(ExpandToSum(~u, ~x) ^ ~m * (~a + ~b * coth(ExpandToSum(~v, ~x))) ^ ~n, ~x) <-- FreeQ([~a, ~b, ~m, ~n], ~x) && (LinearQ([~u, ~v], ~x) && Not(LinearMatchQ([~u, ~v], ~x)))
end

max_stars_repo_path: src/rules/6 Hyperbolic functions/6.3 Hyperbolic tangent/6.3.10 (c+d x)^m (a+b tanh)^n.jl


text:
import AlphaZero.GI
using StaticArrays

const BOARD_SIDE = 3
const NUM_POSITIONS = BOARD_SIDE ^ 2

const Player = Bool
const WHITE = true
const BLACK = false

const Cell = Union{Nothing, Player}
const Board = SVector{NUM_POSITIONS, Cell}
const INITIAL_BOARD = Board(repeat([nothing], NUM_POSITIONS))

mutable struct Game <: GI.AbstractGame
  board :: Board
  curplayer :: Player
  function Game(board=INITIAL_BOARD, player=WHITE)
    new(board, player)
  end
end

GI.Board(::Type{Game}) = Board

GI.Action(::Type{Game}) = Int

#####
##### Defining winning conditions
#####

pos_of_xy((x, y)) = (y - 1) * BOARD_SIDE + (x - 1) + 1

xy_of_pos(pos) = ((pos - 1) % BOARD_SIDE + 1, (pos - 1) ÷ BOARD_SIDE + 1)

const ALIGNMENTS =
  let N = BOARD_SIDE
  let XY = [
    [[(i, j) for j in 1:N] for i in 1:N];
    [[(i, j) for i in 1:N] for j in 1:N];
    [[(i, i) for i in 1:N]];
    [[(i, N - i + 1) for i in 1:N]]]
  [map(pos_of_xy, al) for al in XY]
end end

function has_won(g::Game, player)
  any(ALIGNMENTS) do al
    all(al) do pos
      g.board[pos] == player
    end
  end
end

#####
##### Game API
#####

const ACTIONS = collect(1:NUM_POSITIONS)

GI.actions(::Type{Game}) = ACTIONS

Base.copy(g::Game) = Game(g.board, g.curplayer)

GI.actions_mask(g::Game) = map(isnothing, g.board)

GI.board(g::Game) = g.board

function GI.board_symmetric(g::Game)
  symmetric(c::Cell) = isnothing(c) ? nothing : !c
  # Inference fails when using `map`
  @SVector Cell[symmetric(g.board[i]) for i in 1:NUM_POSITIONS]
end

GI.white_playing(g::Game) = g.curplayer

function GI.white_reward(g::Game)
  isempty(GI.available_actions(g)) && return 0.
  has_won(g, WHITE) && return 1.
  has_won(g, BLACK) && return -1.
  return nothing
end

function GI.play!(g::Game, pos)
  g.board = setindex(g.board, g.curplayer, pos)
  g.curplayer = !g.curplayer
end

#####
##### Simple heuristic for minmax
#####

function alignment_value_for(g::Game, player, alignment)
  γ = 0.3
  N = 0
  for pos in alignment
    mark = g.board[pos]
    if mark == player
      N += 1
    elseif !isnothing(mark)
      return 0.
    end
  end
  return γ ^ (BOARD_SIDE - 1 - N)
end

function heuristic_value_for(g::Game, player)
  return sum(alignment_value_for(g, player, al) for al in ALIGNMENTS)
end

function GI.heuristic_value(g::Game)
  mine = heuristic_value_for(g, g.curplayer)
  yours = heuristic_value_for(g, !g.curplayer)
  return mine - yours
end

#####
##### Machine Learning API
#####

# Vectorized representation: 3x3x3 array
# Channels: free, white, black
function GI.vectorize_board(::Type{Game}, board)
  Float32[
    board[pos_of_xy((x, y))] == c
    for x in 1:BOARD_SIDE,
        y in 1:BOARD_SIDE,
        c in [nothing, WHITE, BLACK]]
end

#####
##### Symmetries
#####

function generate_dihedral_symmetries()
  N = BOARD_SIDE
  rot((x, y)) = (y, N - x + 1) # 90° rotation
  flip((x, y)) = (x, N - y + 1) # flip along vertical axis
  ap(f) = p -> pos_of_xy(f(xy_of_pos(p)))
  sym(f) = map(ap(f), collect(1:NUM_POSITIONS))
  rot2 = rot ∘ rot
  rot3 = rot2 ∘ rot
  return [
    sym(rot), sym(rot2), sym(rot3),
    sym(flip), sym(flip ∘ rot), sym(flip ∘ rot2), sym(flip ∘ rot3)]
end

const SYMMETRIES = generate_dihedral_symmetries()

function GI.symmetries(::Type{Game}, board)
  return [(board[sym], sym) for sym in SYMMETRIES]
end

#####
##### Interaction API
#####

function GI.action_string(::Type{Game}, a)
  string(Char(Int('A') + a - 1))
end

function GI.parse_action(g::Game, str)
  length(str) == 1 || (return nothing)
  x = Int(uppercase(str[1])) - Int('A')
  (0 <= x < NUM_POSITIONS) ? x + 1 : nothing
end

function read_board(::Type{Game})
  n = BOARD_SIDE
  str = reduce(*, ((readline() * "   ")[1:n] for i in 1:n))
  white = ['w', 'r', 'o']
  black = ['b', 'b', 'x']
  function cell(i)
    if (str[i] ∈ white) WHITE
    elseif (str[i] ∈ black) BLACK
    else nothing end
  end
  @SVector [cell(i) for i in 1:NUM_POSITIONS]
end

function GI.read_state(::Type{Game})
  b = read_board(Game)
  nw = count(==(WHITE), b)
  nb = count(==(BLACK), b)
  if nw == nb
    Game(b, WHITE)
  elseif nw == nb + 1
    Game(b, BLACK)
  else
    nothing
  end
end

using Crayons

player_color(p) = p == WHITE ? crayon"light_red" : crayon"light_blue"
player_name(p)  = p == WHITE ? "Red" : "Blue"
player_mark(p)  = p == WHITE ? "o" : "x"

function GI.print_state(g::Game; with_position_names=true, botmargin=true)
  pname = player_name(g.curplayer)
  pcol = player_color(g.curplayer)
  print(pcol, pname, " plays:", crayon"reset", "\n\n")
  for y in 1:BOARD_SIDE
    for x in 1:BOARD_SIDE
      pos = pos_of_xy((x, y))
      c = g.board[pos]
      if isnothing(c)
        print(" ")
      else
        print(player_color(c), player_mark(c), crayon"reset")
      end
      print(" ")
    end
    if with_position_names
      print(" | ")
      for x in 1:BOARD_SIDE
        print(GI.action_string(Game, pos_of_xy((x, y))), " ")
      end
    end
    print("\n")
  end
  botmargin && print("\n")
end


text:
using Revise
using BridgeLandmarks
const BL = BridgeLandmarks
using Distributions
using LinearAlgebra
using Random
using DelimitedFiles
using ReverseDiff
using ForwardDiff
using Zygote
using StaticArrays

n = 101
xobs0 = [PointF(2.0cos(t), sin(t))/4.0  for t in collect(0:(2pi/n):2pi)[2:end]]
θ, ψ =  π/6, 0.1
rot =  SMatrix{2,2}(cos(θ), sin(θ), -sin(θ), cos(θ))
stretch = SMatrix{2,2}(1.0 + ψ, 0.0, 0.0, 1.0 - ψ)
shift = [0.5, -1.0]
xobsT = [rot * stretch * xobs0[i] + shift  for i in 1:n]

adaptskip = 1
skip_saveITER = 50
printskip = 5#1000
ITER = 100#25_000

δmom = [0.01]

p_ahs = Pars_ahs(δmom=δmom,  db=[2.0, 1.0],stdev=.2,γinit=0.1, aprior=Pareto(1.0, 0.1), η =  n -> 0.0, dt = 0.01, σobs=0.01,
                                adaptskip=adaptskip, skip_saveITER=skip_saveITER, ρlowerbound=0.995)

p_ms = Pars_ms(δmom=δmom,  γinit=1.0/√n, aprior=Pareto(1.0, 0.1), η =  n -> 0.0, dt = 0.01, σobs=0.01,
                adaptskip=adaptskip, skip_saveITER=skip_saveITER, ρlowerbound=0.999)
pars = p_ms
model = pars.model
n = length(xobs0)
nshapes = 1

Σobs = fill([pars.σobs^2 * one(UncF) for i in 1:n],2) # noise on observations
T = 1.0; t = 0.0:pars.dt:T; tt = tc(t,T)
obs_atzero = true
fixinitmomentato0 = false
subsamples = 0:pars.skip_saveITER:ITER
obsinfo = BL.set_obsinfo(xobs0,[xobsT],Σobs, obs_atzero,fixinitmomentato0)

################################# initialise P #################################
ainit = 0.5*mean(norm.(diff(xobs0)))
cinit = pars.cinit
γinit = pars.γinit
if model == :ms
    P = MarslandShardlow(ainit, cinit, γinit, 0.0, n)
elseif model == :ahs
    nfsinit = construct_nfs(pars.db, pars.stdev, γinit)
    P = Landmarks(ainit, cinit, n, pars.db , pars.stdev, nfsinit)
end

################## prior specification with θ = (a, c, γ) ########################
priorθ = product_distribution([pars.aprior, pars.γprior])
priormom = MvNormalCanon(gramkernel(xobs0,P)/pars.κ)

xinit = State(xobs0, zeros(PointF,P.n))
mT = zeros(PointF, n)

### from lmguid.jl ###
lt = length(t)
StateW = PointF
dwiener = BL.dimwiener(P)
nshapes = obsinfo.nshapes

guidrec = [BL.GuidRecursions(t,obsinfo)  for _ ∈ 1:nshapes]  # initialise guiding terms
gramT_container = [BL.gram_matrix(obsinfo.xobsT[k], P) for k ∈ 1:nshapes]  # gram matrices at observations at time T

Paux = [BL.auxiliary(P, State(obsinfo.xobsT[k],mT), gramT_container[k]) for k ∈ 1:nshapes] # auxiliary process for each shape
Q = BL.GuidedProposal(P,Paux,t,obsinfo.xobs0,obsinfo.xobsT,guidrec,nshapes,[mT for _ ∈ 1:nshapes])
Q = BL.update_guidrec!(Q, obsinfo)   # compute backwards recursion

X = [BL.initSamplePath(t, xinit) for _ ∈ 1:nshapes]
W = [BL.initSamplePath(t,  zeros(StateW, dwiener)) for _ ∈ 1:nshapes]
for k ∈ 1:nshapes   sample!(W[k], BL.Wiener{Vector{StateW}}())  end

x = deepvec(xinit)
q, p = BL.split_state(xinit)
qᵒ = deepcopy(q); pᵒ = deepcopy(q); ∇ = deepcopy(q); ∇ᵒ = deepcopy(q)

# memory allocations, actual state at each iteration is (X, W, Q, x, ∇x) (x, ∇x are initial state and its gradient)
Xᵒ = deepcopy(X)
Wᵒ = BL.initSamplePath(t,  zeros(StateW, dwiener))
Wnew = BL.initSamplePath(t,  zeros(StateW, dwiener))
# sample guided proposal and compute loglikelihood (write into X)
At = zeros(UncF, 2P.n, 2P.n)  # container for to compute a=σσ' for auxiliary process.
@time ll, X = BL.gp!(BL.LeftRule(), X, xinit, W, Q, At; skip=sk)

# setup containers for saving objects
Xsave = typeof(zeros(length(t) * P.n * 2 * d * nshapes))[]
push!(Xsave, BL.convert_samplepath(X))
parsave = Vector{Float64}[]
push!(parsave, BL.getpars(Q))
initendstates_save = [BL.extract_initial_and_endstate(0,X[1]) ]

δ = [pars.δpos, pars.δmom, pars.δsgd_mom]
δa, δγ = pars.δa, pars.δγ
ρ = pars.ρlowerbound
adapt(i) = (i > 1.5*pars.adaptskip) & (mod(i,pars.adaptskip)==0)

x0 = X[1].yy[1]
dK = BL.gramkernel(x0.q, P)
inv_dK = inv(dK)

P = Q.target;   n = P.n
x0 = BL.deepvec2state(x)
q, p = BL.split_state(x0)
dn = d * n

u = BL.slogρ_mom!(q, Q, W, X, priormom,ll, At)
cfg = ForwardDiff.GradientConfig(u, p, ForwardDiff.Chunk{dn}()) # d*P.n is maximal

#using Profile
#Profile.clear()
#@profile ForwardDiff.gradient!(∇, u, p, cfg)
cfg = ForwardDiff.GradientConfig(u, p, ForwardDiff.Chunk{dn}())
@time ForwardDiff.gradient!(∇, u, p, cfg)
#Juno.profiler()

# ftape = ReverseDiff.GradientTape(u, p)
# compiled_ftape = ReverseDiff.compile(ftape)
# @time ReverseDiff.gradient!(∇, compiled_ftape, p)


text:
# Give me benchmarks:
# The dispatching is zero-cost like a normal struct:

using VarStructs

@var struct A end;
@var struct B end;

dispacth_varstruct(x::A) = 1

dispacth_varstruct(x::B) = 2

using BenchmarkTools
@btime dispath_varstruct(x) setup = (x = rand() < 0.5 ? A() : B())
# 0.001 ns (0 allocations: 0 bytes)


text:
# IRIS web service specification tests

using Test
import Dates
using SeisRequests

@testset "IRISWS" begin
    @testset "IRISTimeSeries" begin
        # Subtyping
        @test IRISTimeSeries <: SeisRequests.IRISRequest
        # Construction
        @test_throws ArgumentError IRISTimeSeries()
        let
            kwargs = Dict(:network => "IU",
                          :station => "ANMO",
                          :location => "--",
                          :channel => "BHZ",
                          :starttime => Dates.DateTime(2000, 1, 1),
                          :endtime => Dates.DateTime(2000, 1, 1, 0, 0, 30),
                          )
            # Missing required parameters
            for arg in keys(kwargs)
                kwargs2 = delete!(deepcopy(kwargs), arg)
                @test_throws ArgumentError IRISTimeSeries(; kwargs2...)
            end
            # Duplication of endtime with duration
            @test_throws ArgumentError IRISTimeSeries(; duration=30, kwargs...)

            req = IRISTimeSeries(; kwargs...)
            for field in fieldnames(IRISTimeSeries)
                if field in keys(kwargs)
                    @test getfield(req, field) == kwargs[field]
                elseif field == :process
                    @test isempty(getfield(req, field))
                elseif field == :format
                    # Default is for `format = "miniseed"`
                    @test getfield(req, field) == "miniseed"
                else
                    @test getfield(req, field) === missing
                end
            end

            # Replacement of missing location with '--'
            kwargs2 = delete!(deepcopy(kwargs), :location)
            @test IRISTimeSeries(; location="", kwargs2...).location == "--"

            # Processing limitations
            for taper in (-1, 2)
                @test_throws ArgumentError IRISTimeSeries(; taper=taper, kwargs...)
            end
            @test_throws ArgumentError IRISTimeSeries(; taper_type="HAMMING", kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; taper=0.1, taper_type="weird_taper_type")
            @test_throws ArgumentError IRISTimeSeries(; lpfilter=-1, kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; hpfilter=-1, kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; bpfilter=(2,1), kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; bpfilter=(-1,2), kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; bpfilter="1;2", kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; bpfilter=(1,2,3), kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; scale="auto", correct=true, kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; scale="weird_scale", kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; freqlimits=[1,2,3,4], kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; correct=true, freqlimits=[1,2,3,1], kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; correct=true, freqlimits=(-1,2,3,4), kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; correct=true, freqlimits=(1,2,3), kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; correct=true, freqlimits=(1,2,3,4),
                autolimits=(3,3), kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; correct=true, autolimits=(1,), kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; units="VEL", kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; correct=true, units="weird_units", kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; units="VEL", kwargs...)
            @test_throws ArgumentError IRISTimeSeries(; format="weird_format",
                delete!(deepcopy(kwargs), :format)...)
            @test_throws ArgumentError IRISTimeSeries(; weird_processing="val", kwargs...)

            # Preservation of processing order
            req = IRISTimeSeries(; diff=true, demean=true, int=true, bpfilter=(1,2), kwargs...)
            @test findall.(isequal.((:diff, :demean, :int, :bpfilter)),
                           Ref(collect(keys(req.process)))) |> x -> issorted(first.(x))

            # Methods
            let req = IRISTimeSeries(network="IU", station="ANMO", location="00",
                                     channel="BHZ", starttime=Dates.DateTime(2000),
                                     endtime=Dates.DateTime(2000)+Dates.Second(1),
                                     format="miniseed", int=true)
                @test SeisRequests.protocol_string(req) == "irisws"
                @test SeisRequests.service_string(req) == "timeseries"
                @test SeisRequests.request_uri(req) ==
                    "http://service.iris.edu/irisws/timeseries/1/query?" *
                    "network=IU&station=ANMO&location=00&channel=BHZ&" *
                    "starttime=2000-01-01T00:00:00&endtime=2000-01-01T00:00:01&" *
                    "format=miniseed&int=true"
                # Get request: will fail if IRISWS not working for some reason
                response = get_request(req, verbose=false)
                @test response.status in keys(SeisRequests.STATUS_CODES)
            end

            @testset "Code conversion" begin
                st = Dates.DateTime(2000)
                et = st + Dates.Hour(1)
                fmt = "miniseed"
                @test_throws ArgumentError IRISTimeSeries(code="A.B.C.D",
                    network="A", station="B", location="C", channel="D",
                    starttime=st, endtime=et, format=fmt)
                @test IRISTimeSeries(code="A.B..D", starttime=st, endtime=et, format=fmt) ==
                    IRISTimeSeries(network="A", station="B", location="", channel="D",
                        starttime=st, endtime=et, format=fmt)
            end

            # Conversion of strings to dates
            @test IRISTimeSeries(network="AN", station="XYZ", location="",
                channel="LHZ", format="miniseed",
                starttime=Dates.DateTime(2000),
                endtime=Dates.DateTime(2000, 01, 01, 01, 02, 03, 400)) ==
                IRISTimeSeries(network="AN", station="XYZ", location="",
                channel="LHZ", format="miniseed",
                starttime="2000-01-01", endtime="2000-01-01T01:02:03.4")
        end
    end
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/Mortar2D.jl/blob/master/LICENSE

using Test

@testset "Mortar2D.jl" begin
    @testset "test_calculate_normals.jl" begin include("test_calculate_normals.jl") end
    @testset "test_calculate_projections.jl" begin include("test_calculate_projections.jl") end
    @testset "test_calculate_segments.jl" begin include("test_calculate_segments.jl") end
    @testset "test_calculate_mortar_matrices.jl" begin include("test_calculate_mortar_matrices.jl") end
    @testset "test_calculate_mortar_assembly.jl" begin include("test_calculate_mortar_assembly.jl") end
end


text:
module TokamakSmallNoncanonical

    import ElectromagneticFields.AxisymmetricTokamakToroidal

    export charged_particle_3d_ode, charged_particle_3d_sode, charged_particle_3d_iode,
           hamiltonian, toroidal_momentum, ϑ

    AxisymmetricTokamakToroidal.@code() # inject magnetic field code

    const xᵢ = [1.05, 0.0, 0.0]
    const qᵢ = vcat(from_cartesian(0, xᵢ), DF̄(0, xᵢ) * [2.1E-3, 4.3E-4, 0.0])

    include("charged_particle_3d_noncanonical.jl")

end


text:
using Base.Test
import Utils: blockinv, kernelnz, transpoflip

@testset "kernelnz" begin
    M = [2 0 0; 0 0 0; 0 0 0]
    @test kernelnz(M, 1; count=-1) == [[2], [3]]
end

@testset "transpoflip" begin
    M = [1 2 3; 4 5 6; 7 8 9]
    @test transpoflip(M) == [9 6 3; 8 5 2; 7 4 1]
end

@testset "invupdate" begin
    for rep = 1 : 100
        X = rand(5,5); M = X' * X; iM = inv(M)
        n = rand(1:5); δ = 5(rand() - 0.5)
        M2 = copy(M); M2[n,n] += δ;
        if isposdef(M) && abs(δ * iM[n,n] + 1) > 1e-3
            @test Utils.invupdate(iM, n, δ) ≈ inv(M2)
        end
    end
end

@testset "invupdate Hermitian" begin
    for rep = 1 : 100
        X = rand(5,5); M = Hermitian(X' * X); iM = inv(M)
        n = rand(1:5); δ = 5(rand() - 0.5)
        M2 = copy(M); M2[n,n] += δ;
        if isposdef(M) && abs(δ * iM[n,n] + 1) > 1e-3
            @test Utils.invupdate(iM, n, δ) ≈ inv(M2)
            @test Utils.invupdate(iM, n, δ) isa Hermitian
        end
    end
end


text:
using GMM
using FactCheck

facts("Testing basic interface") do
    context("Test from vingette for the R gmm package.") do
        include("normal_dist.jl")
        cft = [3.84376, 2.06728]
        cfe = coef(step_1)
        sdt = [0.06527666675890574,0.04672629429325811]
        sde = stderr(step_1, TwoStepGMM())
        for j = 1:2
            @fact cfe[j] => roughly(cft[j])
            @fact sde[j] => roughly(sdt[j])
        end
        Je, pe = GMM.J_test(step_iid_mgr)
        ## This is the objective value, which is the J-test
        ## for other softwares
        ov     = 1.4398836656920428
        Jt, pt = (1.4378658264483137,0.23048500853597673)
        @fact Je => roughly(Jt, atol = 0.01)
        @fact pe => roughly(pt, atol = 0.01)
        @fact objval(step_iid_mgr) => roughly(ov, atol = 0.1)
    end

    context("Example 13.5 from Greene (2012) -- verified with Stata") do
        include("gamma_dist.jl")
        cf_stata = [3.358432, .1244622]
        cfe = coef(two_step)
        for j = 1:2
            @fact cfe[j] => roughly(cf_stata[j], atol=1e-3)
        end
        Je, pe = GMM.J_test(two_step)
        Jt, pt = (1.7433378725860427,0.41825292894063326)
        ## This is the stata J-test
        ov = 1.97522
        @fact Je => roughly(Jt, atol=1e-3)
        @fact pe => roughly(pt, atol=1e-3)
        @fact objval(two_step) => roughly(ov, atol = 0.1)
    end
end

facts("Test utilities") do
    context("test row_kron") do
        h = ["a" "b"; "c" "d"]
        z = ["1" "2" "3"; "4" "5" "6"]
        want = ["a1" "a2" "a3" "b1" "b2" "b3"; "c4" "c5" "c6" "d4" "d5" "d6"]
        @fact GMM.row_kron(h, z) => want

        # now test on some bigger matrices
        a = randn(400, 3)
        b = randn(400, 5)
        out = GMM.row_kron(a, b)
        @fact size(out) => (400, 15)

        rows_good = true
        for row=1:400
            rows_good &= out[row, :] == kron(a[row, :], b[row, :])
        end
        @fact rows_good => true
    end

    context("test max_args") do
        foo(x, y, z) = nothing  # standard 3 args
        bar(x, z=100) = nothing  # standard 2 args with default value
        baz = (x, y, z)-> nothing  # anonymous 3 args
        qux(a; b=100) = nothing  # standard 1 with 1 kwarg (kwarg not counted)

        @fact GMM.max_args(foo) => 3
        @fact GMM.max_args(bar) => 2
        @fact GMM.max_args(baz) => 3
        @fact GMM.max_args(qux) => 1
    end
end

FactCheck.exitstatus()


text:
using Test
using AutoEnvs
using NGSIM

function test_basics()
    # ctor
    filepath = joinpath(dirname(pathof(NGSIM)), "..", "data", "trajdata_i80_trajectories-0400-0415.txt")
    n_envs = 100
    params = Dict("trajectory_filepaths"=>[filepath], "n_envs"=>n_envs)
    env = VectorizedNGSIMEnv(params)

    # reset, step
    x = reset(env)
    a = zeros(n_envs, 2)
    nx, r, terminal, infos = step(env, a)
    @test x != nx
    nnx, r, terminal, infos = step(env, a)
    @test nx != nnx

    # obs spec
    shape, spacetype, infos = observation_space_spec(env)
    @test spacetype == "Box"
    @test in("high", keys(infos))
    @test in("low", keys(infos))

    for _ in 1:200
        _, _, terminals, _ = step(env, a)
        reset(env, terminals)
    end

end

@time test_basics()


text:
"""
Description
===========

Sort a population in decreasing order of fitness.

Arguments
=========

- `P::Array{Model}`: The population to sort.

Side Effects
============

- Mutates the order of `P`.

Returns
=======

- `nothing`.
"""
function sortpop!(P::Array{Model})
    sort!(P, lt=(a,b) -> a.fitness < b.fitness)
    return
end

"""
Description
===========

Select a member of a population through a tournament selection.

Arguments
=========

- `P::Array{Model}`: Population to select a winner from.
- `K::Integer`: Number of competitors in tournament.
- `Ps::Real`: Probability of selection. 0 <= Ps <= 1.

Returns
=======

- `w::Model`: The winning model

Notes
=====

- **Throws an error** if `!(1 <= K <= length(P))`.
"""
function select_with_tournament(P::Array{Model}, K::Integer, Ps::Real)
    if !(1 <= K <= length(P))
        error("Tournament size can not be less than 1 or greater than the population size")
    end

    competitors = shuffle(P)[1:K]
    sortpop!(competitors)

    result = rand()
    for k in 0:K-2
        if result <= Ps*(1-Ps)^k
            return competitors[1]
        else
            competitors = competitors[2:end]
        end
    end

    return competitors[1]
end


text:
export

    # curves
    CurvePt,
    Curve,
    CurveIndex,
    CurveProjection,

    CURVEINDEX_START,

    get_lerp_time,
    index_closest_to_point,
    get_curve_index,
    curveindex_end,

    # roadway
    LaneTag,
    LaneBoundary,
    Lane,
    LaneConnection,
    SpeedLimit,
    RoadSegment,
    Roadway,
    RoadProjection,
    RoadIndex,

    NULL_BOUNDARY,
    NULL_LANETAG,
    NULL_ROADINDEX,

    DEFAULT_SPEED_LIMIT,
    DEFAULT_LANE_WIDTH,

    is_in_exits,
    is_in_entrances,
    is_at_curve_end,
    is_between_segments_lo,
    is_between_segments_hi,
    is_between_segments,
    has_segment,
    has_lanetag,
    has_next,
    has_prev,
    next_lane,
    prev_lane,
    next_lane_point,
    prev_lane_point,
    connect!,
    prev,
    project_to_closest_lane,
    move_along,
    n_lanes_left,
    n_lanes_right,
    get_neighbor_lanetag_left,
    get_neighbor_lanetag_right,
    read_dxf

include("utils.jl")
include("splines.jl")
include("curves.jl")
include("roadways.jl")
include("roadway_generation.jl")


text:
# Test queries
@testset "Queries" begin
    # initialize the model, references, and other info
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @variable(m, inf, Infinite(par))
    @variable(m, pt, Point(inf, 0.5))
    @variable(m, x)
    data = TestData(par, 0, 1)
    meas = measure(x, data, name = "test")
    # test objective_sense
    @testset "JuMP.objective_sense" begin
        # test default
        @test objective_sense(m) == MOI.FEASIBILITY_SENSE
        # change sense
        m.objective_sense = MOI.MIN_SENSE
        # test new sense
        @test objective_sense(m) == MOI.MIN_SENSE
        # undo changes
        m.objective_sense = MOI.FEASIBILITY_SENSE
    end
    # test objective_function
    @testset "JuMP.objective_function" begin
        # test default
        @test isequal(objective_function(m), zero(GenericAffExpr{Float64,
                                                           GeneralVariableRef}))
        # change function
        m.objective_function = x + meas + pt
        # test new function
        @test isequal(objective_function(m), x + meas + pt)
        # undo changes
        m.objective_function = zero(GenericAffExpr{Float64, GeneralVariableRef})
    end
    # test objective_function (internal)
    @testset "JuMP.objective_function (internal)" begin
        # test normal
        @test isequal(objective_function(m, GenericAffExpr{Float64,
                                    GeneralVariableRef}), zero(GenericAffExpr{Float64,
                                                                GeneralVariableRef}))
        # test new sense
        @test_throws InexactError objective_function(m, GeneralVariableRef)
    end
    # test objective_function_type
    @testset "JuMP.objective_function_type" begin
        # test default
        @test objective_function_type(m) == GenericAffExpr{Float64,
                                                           GeneralVariableRef}
        # change function
        m.objective_function = x + meas + pt
        # test new function
        @test objective_function_type(m) == GenericAffExpr{Float64, GeneralVariableRef}
        # undo changes
        m.objective_function = zero(GenericAffExpr{Float64, GeneralVariableRef})
    end
    # test objective_has_measures
    @testset "objective_has_measures" begin
        # test default
        @test !objective_has_measures(m)
        # change function
        m.objective_has_measures = true
        # test new function
        @test objective_has_measures(m)
        # undo changes
        m.objective_has_measures = false
    end
end

# Test definition methods
@testset "Definition" begin
    # initialize the model, references, and other info
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @variable(m, inf, Infinite(par))
    @variable(m, pt, Point(inf, 0.5))
    @variable(m, x)
    data = TestData(par, 0, 1)
    meas = measure(x, data, name = "test")
    # set_objective_sense
    @testset "JuMP.set_objective_sense" begin
        # test normal
        @test isa(set_objective_sense(m, MOI.MAX_SENSE), Nothing)
        @test objective_sense(m) == MOI.MAX_SENSE
    end
    # set_objective_function (expr)
    @testset "JuMP.set_objective_function (Expr)" begin
        # test first set
        @test isa(set_objective_function(m, x + meas + pt), Nothing)
        @test isequal_canonical(objective_function(m), x + meas + pt)
        @test used_by_objective(x)
        @test used_by_objective(pt)
        @test used_by_objective(meas)
        @test !optimizer_model_ready(m)
        @test objective_has_measures(m)
        # test reset
        @test isa(set_objective_function(m, pt), Nothing)
        @test isequal(objective_function(m), pt)
        @test !used_by_objective(x)
        @test used_by_objective(pt)
        @test !used_by_objective(meas)
        @test !objective_has_measures(m)
        # test errors
        @test_throws ErrorException set_objective_function(m, inf + pt)
        @test_throws ErrorException set_objective_function(m, par + pt)
        @test_throws VariableNotOwned set_objective_function(m, @variable(InfiniteModel()))
    end
    # set_objective_function (number)
    @testset "JuMP.set_objective_function (Number)" begin
        # test normal
        @test isa(set_objective_function(m, 3), Nothing)
        @test isequal(objective_function(m), GenericAffExpr{Float64,
                                                      GeneralVariableRef}(3))
        @test !used_by_objective(x)
        @test !used_by_objective(pt)
        @test !used_by_objective(meas)
    end
    # set_objective (expr)
    @testset "JuMP.set_objective (Expr)" begin
        # test first set
        @test isa(set_objective(m, MOI.MIN_SENSE, x + meas + pt), Nothing)
        @test isequal_canonical(objective_function(m), x + meas + pt)
        @test objective_sense(m) == MOI.MIN_SENSE
        @test used_by_objective(x)
        @test used_by_objective(pt)
        @test used_by_objective(meas)
        @test !optimizer_model_ready(m)
        # test reset
        @test isa(set_objective(m, MOI.MAX_SENSE, pt), Nothing)
        @test isequal(objective_function(m), pt)
        @test objective_sense(m) == MOI.MAX_SENSE
        @test !used_by_objective(x)
        @test used_by_objective(pt)
        @test !used_by_objective(meas)
        # test errors
        @test_throws ErrorException set_objective(m, MOI.MAX_SENSE, inf + pt)
        @test_throws ErrorException set_objective(m, MOI.MAX_SENSE, par + pt)
    end
    # set_objective (number)
    @testset "JuMP.set_objective (Number)" begin
        # test normal
        @test isa(set_objective(m, MOI.MIN_SENSE, 3), Nothing)
        @test isequal(objective_function(m), GenericAffExpr{Float64,
                                                      GeneralVariableRef}(3))
        @test objective_sense(m) == MOI.MIN_SENSE
        @test !used_by_objective(x)
        @test !used_by_objective(pt)
        @test !used_by_objective(meas)
    end
    # set_objective (fallback)
    @testset "JuMP.set_objective (Fallback)" begin
        @test_throws ErrorException set_objective(m, MOI.MAX_SENSE, 1im)
    end
    # test the objective macro
    @testset "JuMP.@objective" begin
        @test isequal_canonical(@objective(m, Min, meas + x), meas + x)
        @test isequal_canonical(objective_function(m), meas + x)
        @test objective_sense(m) == MOI.MIN_SENSE
        @test isequal_canonical(@objective(m, Max, (meas + x) + (2x + 3)), meas + 3x + 3)
    end
end

# Test modification methods
@testset "Modification" begin
    # initialize the model, references, and other info
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @variable(m, inf, Infinite(par))
    @variable(m, pt, Point(inf, 0.5))
    @variable(m, x)
    data = TestData(par, 0, 1)
    meas = measure(x, data, name = "test")
    # test set_objective_coefficient
    @testset "JuMP.set_objective_coefficient" begin
        # test variable function
        set_objective_function(m, pt)
        @test isa(set_objective_coefficient(m, pt, 2), Nothing)
        @test isequal(objective_function(m), 2pt)
        @test isa(set_objective_coefficient(m, x, 2), Nothing)
        @test isequal_canonical(objective_function(m), 2pt + 2x)
        # test AffExpr
        @test isa(set_objective_coefficient(m, pt, 3), Nothing)
        @test isequal_canonical(objective_function(m), 3pt + 2x)
        @test isa(set_objective_coefficient(m, meas, 2.5), Nothing)
        @test isequal_canonical(objective_function(m), 3pt + 2x + 2.5 * meas)
        # test QuadExpr
        set_objective_function(m, pt^2 + pt)
        @test isa(set_objective_coefficient(m, pt, 2), Nothing)
        @test isequal_canonical(objective_function(m), pt^2 + 2pt)
        @test isa(set_objective_coefficient(m, x, 2), Nothing)
        @test isequal_canonical(objective_function(m), pt^2 + 2pt + 2x)
        # test number
        set_objective_function(m, 2)
        @test isa(set_objective_coefficient(m, pt, 2), Nothing)
        @test isequal_canonical(objective_function(m), 2pt + 2)
    end
end


text:
using CSV
using DataFrames
using LinearAlgebra
using RecipesBase
using StateSpaceModels
using Statistics
using Test

# Functions that are used in different tests
include("utils.jl")

# Core functionality
include("systems.jl")
include("prints.jl")

# Models
include("models/locallevel.jl")
include("models/locallineartrend.jl")
include("models/locallevelcycle.jl")
include("models/locallevelexplanatory.jl")
include("models/basicstructural.jl")
include("models/basicstructural_multivariate.jl")
include("models/sarima.jl")
include("models/linear_regression.jl")

# Visualization
include("visualization/forecast.jl")


text:
using DifferentialEquations
using LinearAlgebra
using Plots
using ForwardDiff
using ProgressMeter
using DSP

N = 4
A = zeros(N,N)

xs = 1.0:100.0
ys = 1.0:100.0
x = zeros(length(xs),length(ys))

stencil = [0.0  1.0 0.0;
           1.0 -4.0 1.0;
           0.0  1.0 0.0]

# generate_u_prime(x, dx)

# function generate1DStencil!(A, N, Δx)
#     # Generate the centered difference matrix manually
#     for i in 1:N, j in 1:N

#     if abs(i-j)<=1
#         (A[i,j]+=1)
#     end
#     if i==j
#         (A[i,j]-=3)
#     end

#     end

#     A = A/(Δx^2)
# end

function double_partial(x)
    dx = zeros(size(x))
    for i in 2:length(dx)-1
      dx[i] = 1*x[i-1] - 2*x[i] + 1*x[i+1]
    end
    dx[1] = -(2*x[1] + -1*x[2])
    dx[end] = -(-1*x[end-1] + 2*x[end])
end

function double_partial_2D(x)
    padded_du = conv2(x, stencil)
    du = padded_du[2:end-1,2:end-1]
    du
end

f(u, p, t) = gen_forcing_2D(p)

function gen_forcing_2D(positions)
    out_vec = create_2D_gaussian_matrix(xs,ys,positions)
    out_vec
end

function heat_eq(u, p, t)
    du = double_partial_2D(u) + f(u,p,t)
    du
end

function gaussian(x, mu, sig)
    return exp(-(x - mu)^2.0 / (2 * sig^2.0))
end

function gaussian_2D(loc, mu, sig)
    x, y = loc
    x₀, y₀ = mu
    σ₁, σ₂ = sig
    return gaussian(x, x₀, σ₁) * gaussian(y, y₀, σ₂)
end

function gaussian_2D_multipos(loc, positions, sig)
    x, y = loc
    σ₁, σ₂ = sig
    val = 0
    for pos in positions
        x₀, y₀ = pos
        val += gaussian(x, x₀, σ₁) * gaussian(y, y₀, σ₂)
    end
    val
end

# Loss: Given a list of positions, what is the cost / loss of the
function loss(positions)
    simulation = predict(positions)
    # display(simulation)
    final_state = simulation[end]
    # display(size(final_state)[1])
    # display(final_state)
    return norm(final_state, 2)
end

function predict(positions)
    # Simply roll this forward in time and we get our solution
    tspan = (0.0, 1.0)
    u₀ = zeros(size(x))
    u_dual = convert.(eltype(positions[1]),u₀)
    prob = ODEProblem(heat_eq, u_dual, tspan, positions)
    sol = solve(prob)
    # display("Finished ODE Solve")
    sol
end

# function create_2D_gaussian_matrix_orig(xs,ys,positions)
#     out = zeros(size(xs)[1],size(ys)[1])
#     out = convert.(eltype(positions[1]),out)
#     x_idx = 1
#     y_idx = 1

#     for pos in positions
#         x_p, y_p = pos
#         for x in xs
#             for y in ys
#                 out[x_idx, y_idx] += gaussian_2D([x,y],[x_p,y_p],[5.0,5.0])
#                 y_idx += 1
#             end
#             y_idx = 1
#             x_idx += 1
#         end
#         x_idx = 1

#     end
#     out
# end

function create_2D_gaussian_matrix(xs,ys,positions)
    xsz = length(xs)
    ysz = length(ys)
    out = zeros(xsz,ysz)
    out = convert.(eltype(positions[1]),out)

    for i=1:xsz
        for j=1:ysz
            x, y = xs[i], ys[j]
            out[i, j] += gaussian_2D_multipos([x,y],positions,[5.0,5.0])
        end
    end

    out
end

function plot_gaussian()
    positions = [[20.0,30.0],[40.0,50.0]]
    xs = 1.0:100.0
    ys = 1.0:100.0
    gm = create_2D_gaussian_matrix(xs,ys,positions)
    gm_f(x, y) = gm[x, y]
    display(plot(1:length(xs), 1:length(ys), gm_f, st=:heatmap, color=:viridis))
end

function plot_heatmap(data)
    display(plot(1:size(data,1), 1:size(data,2), data, st=:heatmap, color=:viridis))
end

function plot_gaussian_2()
    positions = [[20.0,30.0],[40.0,50.0],[66.0,3.3]]
    xs = 1.0:0.5:100.0
    ys = 1.0:0.5:100.0
    gf = (x,y) -> gaussian_2D_multipos([x,y],positions,[5.0,5.0])
    display(plot(xs, ys, gf, st=:heatmap, color=:viridis))
end

function optimize_and_save_thermal(positions, η=0.1)
    solns = []
    append!(solns, [deepcopy(positions)])
    @showprogress for idx in 1:1000
        grads = ForwardDiff.gradient(s -> loss(s), positions) # Magic
        positions .-= η*grads
        # display(positions)
        # display(loss(positions))

        if loss(positions) < 0.1
            break
        end

        if idx in [500,1000]
            append!(solns, [deepcopy(positions)])
        end
    end

    solns
end


text:
using Base.Threads
using TaskScheduling
using Base.Test

# Simple test
flag = false
submit() do
    global flag = true
end
runtasks()
@test flag == true

# Run many trivial tasks
counter = Atomic{Int}(0)
for i in 1:100
    submit() do
        global counter
        atomic_add!(counter, 1)
    end
end
runtasks()
@test counter[] == 100

# Run a long sequence of tasks
counter = Atomic{Int}(0)
function run1(i::Integer)
    global counter
    i==0 && return
    atomic_add!(counter, 1)
    submit() do
        run1(i-1)
    end
end
submit() do
    run1(100)
end
runtasks()
@test counter[] == 100

# Run a tree of tasks
counter = Atomic{Int}(0)
function run2(i::Integer)
    global counter
    i==0 && return
    atomic_add!(counter, 1)
    j = i-1
    j1 = div(j,2)
    j2 = j - j1
    submit() do
        run2(j1)
    end
    submit() do
        run2(j2)
    end
end
submit() do
    run2(100)
end
runtasks()
@test counter[] == 100

# A real-world example
mats = Matrix{Float64}[]
for i in 1:100
    sz = rand(10:100)
    A = rand(sz,sz)
    push!(mats, A)
end
svds = Vector{Any}(100)
for i in 1:100
    submit() do
        svds[i] = svd(mats[i])
    end
end
runtasks()
for i in 1:100
    @test isa(svds[i][2], Vector{Float64})
end


text:
"""

u(t+1) = A u(t) + B w(t+1)
z(t) = C u(t) + v(t+1)
"""
struct LinearStateSpaceProblem{
    isinplace,
    Atype<:AbstractArray,
    Btype<:AbstractArray,
    Ctype<:AbstractArray,
    wtype,
    Rtype, # Should be expanded later on to include distributions
    utype,
    ttype,
    otype
} <: AbstractStateSpaceProblem{isinplace}
    A::Atype # Evolution matrix
    B::Btype # Noise matrix
    C::Ctype # Observation matrix
    noise::wtype # Latent noise distribution
    R::Rtype # Observation noise matrix
    u0::utype # Initial condition
    tspan::ttype # Timespan to use
    observables::otype # Observed data to use, if any
end

function LinearStateSpaceProblem(
    A::Atype,
    B::Btype,
    C::Ctype,
    u0::utype,
    tspan::ttype;
    R = diagm(ones(size(C,1))),
    observables = nothing,
    noise = StandardGaussian(size(B, 1)),
) where {
    Atype<:AbstractArray,
    Btype<:AbstractArray,
    Ctype<:AbstractArray,
    utype,
    ttype,
}
    v = diagm(ones(size(C,1)))

    if R isa Vector
        @assert length(R) == 1
        R = hcat(R) # Convert to matrix
    end

    return LinearStateSpaceProblem{
        Val(false),
        Atype,
        Btype,
        Ctype,
        typeof(noise),
        typeof(R),
        utype,
        ttype,
        typeof(observables)
    }(
        A, # Evolution matrix
        B, # Noise matrix
        C, # Observation matrix
        noise, # Latent noise distribution
        v, # Observation noise matrix
        u0, # Initial condition
        tspan, # Timespan to use
        observables # Observed data to use, if any
    )
end

function linear_likelihood(
    prob::LinearStateSpaceProblem,
    ::LinearGaussian,
    z_t,
    ::Nothing
)
    return 0df0ads0fdsa
end

function SciMLBase.solve(
    prob::LinearStateSpaceProblem{isinplace, Atype, Btype, Ctype, wtype, Rtype, utype, ttype, otype},
    ::LinearGaussian,
    args...;
    kwargs...
) where {isinplace, Atype, Btype, Ctype, wtype, Rtype<:AbstractMatrix, utype, ttype, otype<:Nothing}
    # Preallocate values
    T = prob.tspan[2]
    A,B,C,R = prob.A, prob.B, prob.C, prob.R
    K = size(B, 1) # Rows of latent states
    L = size(C, 1) # Rows of observations

    u = Vector{utype}(undef, T+1) # Latent states
    n = Vector{utype}(undef, T+1) # Latent noise
    z = Vector{utype}(undef, T+1) # Observables generated

    u[1] = prob.u0
    z[1] = C * u[1] + R*randn(L)
    n[1] = noise(prob.noise, 1)  # This noise term is never used

    # Simulate it, homie
    for t in 2:T+1
        n[t] = noise(prob.noise, t)
        u[t] = A * u[t-1] + B * n[t] # Call latent_noise here
        z[t] = C*u[t] + R*randn(L)
    end

    return StateSpaceSolution(z, u, n, nothing)
end

function SciMLBase.solve(
    prob::LinearStateSpaceProblem{isinplace, Atype, Btype, Ctype, wtype, Rtype, utype, ttype, otype},
    ::LinearGaussian,
    args...;
    kwargs...
) where {isinplace, Atype, Btype, Ctype, wtype, Rtype<:AbstractMatrix, utype, ttype, otype}
    # Preallocate values
    T = prob.tspan[2]
    A,B,C,R = prob.A, prob.B, prob.C, prob.R
    K = size(B, 1) # Rows of latent states
    L = size(C, 1) # Rows of observations

    u = Vector{utype}(undef, T+1) # Latent states
    n = Vector{utype}(undef, T+1) # Latent noise
    z = Vector{utype}(undef, T+1) # Observables generated

    u[1] = prob.u0
    z[1] = C * u[1] + R*randn(L)
    n[1] = noise(prob.noise, 1) # This noise term is never used
    loglik = 0.0

    # Simulate it, homie
    for t in 2:T+1
        n[t] = noise(prob.noise, t)
        u[t] = A * u[t-1] + B * n[t] # Call latent_noise here
        z[t] = C*u[t] + R*randn(L)
        err = z[t] - prob.observables[t]
        loglik += logpdf(MvNormal(R), err)
    end

    return StateSpaceSolution(z, u, n, loglik)
end


text:
### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions, Random, StatsBase, DataFrames, Plots
	Random.seed!(1)
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.30"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	n = 10^6
	data = tan.(rand(n)*pi .- pi/2)
	averages = accumulate(+,data)./collect(1:n)
end

# ╔═╡ 588624fe-21c2-11eb-3d7b-f37b1c1f44ca
plot( 1:n, averages,
	c=:blue, legend=:none,
	xscale=:log10, xlims=(1,n), xlabel="n", ylabel="Running average")

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.30"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═588624fe-21c2-11eb-3d7b-f37b1c1f44ca
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


text:
@testset "fill" begin
    s = """
        @def x = 5
        {{fill x}}
        """ |> fd2html_td
    @test isapproxstr(s, "5")
    @test_throws F.HTMLFunctionError ("{{fill x y z}}" |> fd2html)
end

@testset "href" begin
    s = "{{href EQR foo}}" |> fd2html_td
    @test isapproxstr(s, "??")
end

@testset "redirect" begin
    # only one arg
    @test_throws F.HTMLFunctionError ("{{redirect foo bar}}" |> fd2html)
    # must end with html
    @test_throws F.HTMLFunctionError ("{{redirect foo}}" |> fd2html)
end

@testset "isempty" begin
    s = """
        +++
        using Dates
        d = Date(1,1,1)
        x = nothing
        +++
        {{isempty d}}A{{end}}
        {{isempty x}}B{{end}}
        """ |> fd2html_td
    @test isapproxstr(s, "
AB
")
end


text:
import Base.eps
import DSP.spectrogram, DSP.stft
export spectrogram, stft, istft, phase_vocoder

""""""
function spectrogram(audio::SampleBuf{T, 1},
                        windowsize::Int = 1024,
                        hopsize::Int = windowsize >> 2;
                        window = hanning, kwargs...) where T
    noverlap = windowsize - hopsize
    data = map(Float32, audio.data)
    DSP.spectrogram(data, windowsize, noverlap; fs = audio.samplerate.val, window = window, kwargs...)
end

""""""
function spectrogram(audio::SampleBuf{T, 2},
                        windowsize::Int = 1024,
                        hopsize::Int = windowsize >> 2;
                        window = hanning, kwargs...) where T
    noverlap = windowsize - hopsize
    data = map(Float32, audio.data)
    (mapslices(data, 1) do data
        DSP.spectrogram(data, windowsize, noverlap; fs = audio.samplerate.val, window = window, kwargs...)
    end)[:]
end

""""""
function spectrogram(audio::SampleBuf{T, N},
                           windowsize::Seconds,
                           hopsize::Seconds = windowsize / 4;
                           kwargs...) where {T, N}
    w = round(Int, ustrip(windowsize) * audio.samplerate)
    h = round(Int, ustrip(hopsize) * audio.samplerate)
    spectrogram(audio, w, h; kwargs...)
end

""""""
function stft(audio::SampleBuf{T, 1},
                 windowsize::Int = 1024,
                 hopsize::Int = windowsize >> 2;
                 window = hanning, kwargs...) where T
    noverlap = windowsize - hopsize
    data = tofloat(audio.data)
    DSP.stft(data, windowsize, noverlap; window = window, kwargs...)
end

""""""
function stft(audio::SampleBuf{T, 2},
                 windowsize::Int = 1024,
                 hopsize::Int = windowsize >> 2; kwargs...) where T
    nchannels = SampledSignals.nchannels(audio)
    noverlap = windowsize - hopsize

    stft = Array(Matrix, nchannels)
    data = tofloat(audio.data)
    for i in 1:nchannels
        stft[i] = DSP.stft(data[:, i], windowsize, noverlap; kwargs...)
    end
    cat(3, stft...)
end

""""""
function stft(audio::SampleBuf{T, N},
                    windowsize::Seconds,
                    hopsize::Seconds = windowsize / 4; kwargs...) where {T, N}
    w = round(Int, ustrip(windowsize) * audio.samplerate)
    h = round(Int, ustrip(hopsize) * audio.samplerate)
    stft(audio, w, h; kwargs...)
end

# functions to make window vector from a window function, vector or nothing.
@inline w(window::Function, size) = window(size)
@inline function w(window::AbstractVector, size)
    if size != length(window)
        error("the length of window does not match windowsize")
    end
    window
end
@inline w(window::Nothing, size) = ones(size)

"""
    istft(stft, samplerate, windowsize, hopsize; window)

Reconstruct a mono audio from its STFT.

# Arguments
* `stft::Array{Complex{T}, 2}`: the STFT matrix, usually having (1+nfft/2) rows
* `samplerate::Real`: Sample rate of the resulting audio
* `windowsize::Int`: window size used for STFT. (default: nfft)
* `hopsize::Int`: number of frames between STFT columns (default: windowsize/4)
* `nfft::Int`: the FFT size (default: 2*(size(stft,1)-1))
* `window::Union{Function, AbstractVector, Nothing}`: function or vector that has the window used for STFT (default: uniform window)
"""
function istft(stft::Array{Complex{T}, 2},
                                   samplerate::Real,
                                   windowsize::Int = 2 * (size(stft, 1) - 1),
                                   hopsize::Int = windowsize >> 2;
                                   nfft::Int = windowsize,
                                   window::Union{Function, AbstractVector, Nothing} = hanning) where {T <: AbstractFloat}

    if windowsize > nfft
      error("window size should be less than or equal to nfft")
    end

    window::Vector{T} = map(T, w(window, windowsize))

    columns = size(stft, 2)
    audio = zeros(T, windowsize + hopsize * (columns - 1))
    weights = zeros(audio)
    spectrum = zeros(T, nfft)

    nbins = size(stft, 1)
    base = Base.unsafe_convert(Ptr{Complex{T}}, stft)
    stride = nbins * sizeof(Complex{T})

    @inbounds for i in 1:columns
        irfft!(spectrum, base + stride * (i-1), nfft)

        left = (i-1) * hopsize
        for j in 1:windowsize
            audio[left + j] += window[j] * spectrum[j]
            weights[left + j] += window[j]
        end
    end

    @inbounds for i in 1:length(audio)
        if weights[i] > eps(T)
            audio[i] /= weights[i]
        end
    end

    SampleBuf{T, 1}(audio, samplerate)
end

function istft(stft::Array{Complex{T}, 2},
                                   samplerate::Hertz,
                                   windowsize::Int = 2 * (size(stft, 1) - 1),
                                   hopsize::Int = windowsize >> 2;
                                   nfft::Int = windowsize,
                                   window::Union{Function, AbstractVector, Nothing} = hanning) where {T <: AbstractFloat}
    istft(stft,
            ustrip(samplerate),
            windowsize,
            hopsize,
            nfft,
            window)
end

"""
    istft(stft, samplerate, windowsize, hopsize; window)

Reconstruct a multichannel audio from its STFT.

# Arguments
* `stft::Array{Complex{T}, 3}`: the STFT matrix, usually having (1+nfft/2) rows
* `samplerate::Real`: Sample rate of the resulting audio
* `windowsize::Int`: window size used for STFT. (default: nfft)
* `hopsize::Int`: number of frames between STFT columns (default: windowsize/4)
* `nfft::Int`: the FFT size (default: 2*(size(stft,1)-1))
* `window::Union{Function, AbstractVector, Nothing}`: function or vector that has the window used for STFT (default: uniform window)
"""
function istft(stft::Array{Complex{T}, 3},
                                   samplerate::Real,
                                   windowsize::Int = 2 * (size(stft, 1) - 1),
                                   hopsize::Int = windowsize >> 2; kwargs...) where {T <: AbstractFloat}
    nchannels = size(stft, 3)
    buffers = cell(nchannels)

    # run ISTFT for each channel
    for i = 1:nchannels
        buffers[i] = istft(stft[:, :, i], samplerate, windowsize, hopsize; kwargs...)
    end

    nsamples = size(buffers[1], 1)
    audio = Array(T, nsamples, nchannels)

    for i = 1:nchannels
        audio[:, i] = buffers[i].data
    end

    # concatenate channels
    SampleBuf{T, 2}(
        audio,
        samplerate
    )
end

function istft(stft::Array{Complex{T}, 3},
                                   samplerate::Hertz,
                                   windowsize::Int = 2 * (size(stft, 1) - 1),
                                   hopsize::Int = windowsize >> 2; kwargs...) where {T <: AbstractFloat}
    istft(stft,
            ustrip(samplerate),
            windowsize,
            hopsize;
            kwargs...)
end

"""
    phase_vocoder(stft, speed, hopsize)

Phase vocoder. Given an STFT matrix, speed it up by a factor.
"""
function phase_vocoder(stft::Array{Complex{T}, 2},
                                           speed::Real,
                                           hopsize::Int = (size(stft, 1) - 1) >> 1) where {T <: AbstractFloat}

    nbins = size(stft, 1)
    nframes = size(stft, 2)
    nfft = (nbins - 1) * 2
    timesteps = 1f0:speed:nframes
    stretched = zeros(Complex{T}, nbins, length(timesteps))
    phase_advance = collect(linspace(0f0, T(π * hopsize), nbins))
    phase_acc = angle(stft[:, 1])
    cis_phase = Array(Complex{T}, size(phase_acc))
    angle1 = Array(T, nbins)
    angle2 = Array(T, nbins)
    dphase = Array(T, nbins)
    mag = Array(T, nbins)
    twopi = T(2π)

    @inbounds for (t, step) in enumerate(timesteps)
        left = floor(Int, step)
        right = left + 1

        # weighting for linear magnitude interpolation
        alpha = step % 1f0
        if left < nframes
            for i = 1:nbins
                mag[i] = (1f0 - alpha) * abs(stft[i, left]) + alpha * abs(stft[i, right])
            end
        else
            for i = 1:nbins
                mag[i] = abs(stft[i, left])
            end
        end

        # store to output array
        cis!(cis_phase, phase_acc)
        multiply!(cis_phase, cis_phase, mag)
        for i in 1:nbins
            stretched[i, t] = cis_phase[i]
        end

        if t == size(stretched, 2)
            # nothing to do more
            break
        end

        for i in 1:nbins
            angle1[i] = angle(stft[i, left])
            angle2[i] = angle(stft[i, right])
            # compute phase advance
            dphase[i] = angle2[i] - angle1[i] - phase_advance[i]
            # wrap to -pi:pi range
            dphase[i] -= twopi * round(dphase[i] / twopi)
            # accumulate phase
            phase_acc[i] += phase_advance[i] + dphase[i]
        end
    end

    stretched
end

"""
    phase_vocoder(stft, speed, hopsize)

Phase vocoder. Given an STFT matrix, speed it up by a factor.
"""
function phase_vocoder(stft::Array{Complex{T}, 3},
                                           speed::Real,
                                           hopsize::Int = (size(stft, 1) - 1) >> 1) where {T <: AbstractFloat}
    mapslices(stft, 1:2) do stft
        phase_vocoder(stft, speed, hopsize)
    end
end


text:
# predefined adaptors for working with types from the Julia standard library

## Base

#adapt_structure(to, xs::Tuple) = Tuple(adapt(to, x) for x in xs)
#
# non-allocating version
#@generated adapt_structure(to, x::Tuple) =
#    Expr(:tuple, (:(adapt(to, x[$i])) for i in 1:fieldcount(x))...)
#
# non-allocating, non-@generated version
adapt_structure(to, xs::Tuple) = _adapt_structure(to, xs)
_adapt_structure(to, xs::Tuple{}) = ()
_adapt_structure(to, xs::Tuple) =
    tuple(adapt(to, xs[1]), _adapt_structure(to, Base.tail(xs))...)

@generated adapt_structure(to, x::NamedTuple) =
    Expr(:tuple, (:($f=adapt(to, x.$f)) for f in fieldnames(x))...)

## Array wrappers

using LinearAlgebra

# database of array wrappers, for use throughout the package
#
# LHS entries are a symbolic type with AT for the array type
#
# RHS entries consist of a closure to reconstruct the wrapper, with as arguments
# a wrapper instance and mutator function to apply to the inner array
const wrappers = (
  :(SubArray{<:Any,<:Any,AT})                     => (A,mut)->SubArray(mut(parent(A)), mut(parentindices(A))),
  :(PermutedDimsArray{<:Any,<:Any,<:Any,<:Any,AT})=> (A,mut)->PermutedDimsArray(mut(parent(A)), permutation(A)),
  :(Base.ReshapedArray{<:Any,<:Any,AT,<:Any})     => (A,mut)->Base.reshape(mut(parent(A)), size(A)),
  :(LinearAlgebra.Adjoint{<:Any,AT})              => (A,mut)->LinearAlgebra.adjoint(mut(parent(A))),
  :(LinearAlgebra.Transpose{<:Any,AT})            => (A,mut)->LinearAlgebra.transpose(mut(parent(A))),
  :(LinearAlgebra.LowerTriangular{<:Any,AT})      => (A,mut)->LinearAlgebra.LowerTriangular(mut(parent(A))),
  :(LinearAlgebra.UnitLowerTriangular{<:Any,AT})  => (A,mut)->LinearAlgebra.UnitLowerTriangular(mut(parent(A))),
  :(LinearAlgebra.UpperTriangular{<:Any,AT})      => (A,mut)->LinearAlgebra.UpperTriangular(mut(parent(A))),
  :(LinearAlgebra.UnitUpperTriangular{<:Any,AT})  => (A,mut)->LinearAlgebra.UnitUpperTriangular(mut(parent(A))),
  :(LinearAlgebra.Diagonal{<:Any,AT})             => (A,mut)->LinearAlgebra.Diagonal(mut(parent(A))),
  :(LinearAlgebra.Tridiagonal{<:Any,AT})          => (A,mut)->LinearAlgebra.Tridiagonal(mut(A.dl), mut(A.d), mut(A.du)),
)

permutation(::PermutedDimsArray{T,N,perm}) where {T,N,perm} = perm

for (W, ctor) in wrappers
    mut = :(A -> adapt(to, A))
    @eval adapt_structure(to, wrapper::$W where {AT <: Any}) = $ctor(wrapper, $mut)
end

## Broadcast

import Base.Broadcast: Broadcasted, Extruded

adapt_structure(to, bc::Broadcasted{Style}) where Style =
  Broadcasted{Style}(bc.f, map(arg->adapt(to, arg), bc.args), bc.axes)

adapt_structure(to, ex::Extruded) =
    Extruded(adapt(to, ex.x), ex.keeps, ex.defaults)


text:
abstract type AbstractShell <: AbstractREPL end

include("shells/bash.jl")
include("shells/zsh.jl")
include("shells/ksh.jl")
include("shells/fish.jl")
include("shells/cmdexe.jl")

function getshell(shell="/bin/bash")
    function findshell(sh)
        if contains(sh, "bash")
            return BASH(sh)
        elseif contains(sh, "zsh")
            return ZSH(sh)
        elseif contains(sh, "fish")
            return FISH(sh)
        elseif contains(sh, "ksh")
            return KSH(sh)
        else
            warn(logger, "SHELL $sh not supported.")
            return nothing
        end
    end

    if haskey(ENV, "SHELL")
        sh = findshell(ENV["SHELL"])

        if sh === nothing
            return findshell(shell)
        else
            return sh
        end
    elseif is_windows()
        return CmdExe()
    elseif is_unix()
        return findshell(shell)
    end
end


text:
using AbstractDifferentiation
using ForwardDiff
using GalacticNLopt
using LinearAlgebra
using Optim
using Pathfinder
using ProgressLogging
using SciMLBase
using Test

include("test_utils.jl")

@testset "build_optim_function" begin
    n = 20
    f(x) = logp_banana(x)
    ∇f(x) = ForwardDiff.gradient(f, x)
    ad_backend = AD.ForwardDiffBackend()
    x = randn(n)
    funs = [
        "user-provided gradient" => Pathfinder.build_optim_function(f, ∇f; ad_backend),
        "automatic gradient" => Pathfinder.build_optim_function(f; ad_backend),
    ]
    @testset "$name" for (name, fun) in funs
        @test fun isa SciMLBase.OptimizationFunction
        @test fun.f(x) ≈ -f(x)
        ∇fx = similar(x)
        fun.grad(∇fx, x, nothing)
        @test ∇fx ≈ -∇f(x)
        H = similar(x, n, n)
        fun.hess(H, x, nothing)
        @test H ≈ -ForwardDiff.hessian(f, x)
        Hv = similar(x)
        v = randn(n)
        fun.hv(Hv, x, v, nothing)
        @test Hv ≈ H * v
    end
end

@testset "build_optim_problem" begin
    n = 20
    f(x) = logp_banana(x)
    ∇f(x) = ForwardDiff.gradient(f, x)
    ad_backend = AD.ForwardDiffBackend()
    x0 = randn(n)
    fun = Pathfinder.build_optim_function(f; ad_backend)
    prob = Pathfinder.build_optim_problem(fun, x0)
    @test prob isa SciMLBase.OptimizationProblem
    @test prob.f === fun
    @test prob.u0 == x0
    @test prob.p === nothing
end

@testset "optimize_with_trace" begin
    n = 10
    P = inv(rand_pd_mat(Float64, n))
    μ = randn(n)
    f(x) = -dot(x - μ, P, x - μ) / 2
    ∇f(x) = ForwardDiff.gradient(f, x)

    x0 = randn(n)
    ad_backend = AD.ForwardDiffBackend()
    fun = Pathfinder.build_optim_function(f; ad_backend)
    prob = Pathfinder.build_optim_problem(fun, x0)

    optimizers = [
        Optim.BFGS(), Optim.LBFGS(), Optim.ConjugateGradient(), NLopt.Opt(:LD_LBFGS, n)
    ]
    @testset "$(typeof(optimizer))" for optimizer in optimizers
        optim_sol, optim_trace = Pathfinder.optimize_with_trace(prob, optimizer)
        @test optim_sol isa SciMLBase.OptimizationSolution
        @test optim_trace isa Pathfinder.OptimizationTrace
        @test optim_trace.points[1] ≈ x0
        @test optim_sol.prob.u0 ≈ x0
        @test optim_trace.points[end] ≈ μ
        @test optim_sol.u ≈ μ
        @test optim_trace.log_densities ≈ f.(optim_trace.points)
        @test optim_trace.gradients ≈ ∇f.(optim_trace.points) atol = 1e-4

        if !(optimizer isa NLopt.Opt)
            options = Optim.Options(; store_trace=true, extended_trace=true)
            res = Optim.optimize(
                x -> -f(x), (y, x) -> copyto!(y, -∇f(x)), x0, optimizer, options
            )
            @test Optim.iterations(res) == length(optim_trace.points) - 1
            @test Optim.x_trace(res) ≈ optim_trace.points
            @test Optim.minimizer(res) ≈ optim_trace.points[end]
        end
    end

    @testset "progress logging" begin
        logs, = Test.collect_test_logs(; min_level=ProgressLogging.ProgressLevel) do
            ProgressLogging.progress(; name="Optimizing") do progress_id
                Pathfinder.optimize_with_trace(prob, Optim.LBFGS(); progress_id)
            end
        end
        @test logs[1].kwargs[:progress] === nothing
        @test logs[1].message == "Optimizing"
        @test logs[2].kwargs[:progress] == 0.0
        @test logs[2].message == "Optimizing"
        @test logs[3].kwargs[:progress] == 0.001
        @test logs[end].kwargs[:progress] == "done"
    end
end


text:
using Hangul # YetJamo Initial Medial Final YetSyllable ComposeError
using Test

iㄱ = YetJamo(0x1100) # ㄱ U+1100
mㅏ = YetJamo(0x1161) # ㅏ U+1161
fㄱ = YetJamo(0x11A8) # ㄱ U+11A8

@test YetSyllable(iㄱ) == YetSyllable(iㄱ, nothing, nothing)
@test YetSyllable(mㅏ) == YetSyllable(nothing, mㅏ, nothing)
@test YetSyllable(fㄱ) == YetSyllable(nothing, nothing, fㄱ)

@test to_string(YetSyllable(iㄱ)) == "ᄀ"
@test to_string(YetSyllable(iㄱ, mㅏ)) == "가"
@test to_string(YetSyllable(iㄱ, fㄱ)) == "ᄀᆨ"
@test to_string(YetSyllable(iㄱ, mㅏ, fㄱ)) == "각"
@test codepoint.(collect(to_string(YetSyllable(iㄱ, mㅏ, fㄱ)))) == [0x1100, 0x1161, 0x11a8]

@test to_string(YetSyllable(mㅏ)) == "ᅡ"
@test to_string(YetSyllable(mㅏ, fㄱ)) == "ᅡᆨ"
@test to_string(YetSyllable(fㄱ)) == "ᆨ"

@test length(to_string(YetSyllable(iㄱ))) == 2
@test length(to_string(YetSyllable(mㅏ))) == 2
@test length(to_string(YetSyllable(fㄱ))) == 3

iㅿ = YetJamo(0x1140)  # ㅿ U+1140
@test iㅿ isa YetJamo{Initial}
@test to_string(YetSyllable(iㅿ)) == "ᅀ"

@test to_string(YetSyllable(iㅿ, mㅏ)) == "ᅀᅡ"
@test to_string(YetSyllable(iㅿ, mㅏ, fㄱ)) == "ᅀᅡᆨ"

@test length(to_string(YetSyllable(iㅿ))) == 2

@test Syllable(YetSyllable(iㄱ, mㅏ)) == Syllable(Jamo('ㄱ'), Jamo('ㅏ'))
@test Syllable(YetSyllable(iㄱ, mㅏ, fㄱ)) == Syllable(Jamo('ㄱ'), Jamo('ㅏ'), Jamo('ㄱ'))
@test YetSyllable(iㄱ, mㅏ) == YetSyllable(Syllable(Jamo('ㄱ'), Jamo('ㅏ')))
@test YetSyllable(iㄱ, mㅏ, fㄱ) == YetSyllable(Syllable(Jamo('ㄱ'), Jamo('ㅏ'), Jamo('ㄱ')))

@test_throws ComposeError Syllable(YetSyllable(iㅿ, mㅏ, fㄱ)) # 매핑 좀 해줘요

@test YetJamo('ᄀ') == YetJamo(0x1100) == YetJamo{Initial}(0x1100) # ㄱ U+1100
@test YetJamo('ᅡ')  == YetJamo(0x1161) == YetJamo{Medial}(0x1161)  # ㅏ U+1161
@test YetJamo('ᆨ')  == YetJamo(0x11A8) == YetJamo{Final}(0x11A8)   # ㄱ U+11A8
@test YetJamo('ᅀ') == YetJamo(0x1140) == YetJamo{Initial}(0x1140) # ㅿ U+1140

@test_throws Hangul.YetJamoError YetJamo(0)
@test_throws Hangul.YetJamoError YetJamo('a')
@test_throws Hangul.YetJamoError YetJamo('ㄱ') # 0x3131

@test YetJamo{Initial}(0x1100) == YetJamo(0x1100) == YetJamo('ᄀ') == Jamo(0x1100) == Jamo('ᄀ')


text:
using DPC, CairoMakie
#include(joinpath(@__DIR__, "utils.jl"))
include("utils.jl")

fig = Figure(resolution = (0.75textwidth, 0.75textwidth))

volleys = [
    (70.0, :i1), (120.0, :i2), (170.0, :i3),
    (315.0, :i3), (325.0, :i2), (335.0, :i1),
    (370.0, :i3), (380.0, :i2), (390.0, :i1),
    (425.0, :i3), (435.0, :i2), (445.0, :i1)
]

ax0 = fig[1, 1] = Axis(fig;
    titlealign=:left, title = "a.    Volleys of spikes from population A, B and C"
)

for (t,pop) in  volleys
    (name,y) = if pop == :i1
        "A",3
    elseif pop==:i2
        "B",2
    else
        "C",1
    end
    text!(ax0, name, position=Point2f0(t, y), align=(:center, :center), textsize=14, color=:black)
end
hideydecorations!(ax0)
hidexdecorations!(ax0, grid=false)
ylims!(ax0, 0.25, 3.75)
fig

## First setup: without inhibition

config1 = """
refractory_duration: 5.01
inputs:
- id: i11
- id: i12
- id: i13
- id: i14
- id: i15
- id: i21
- id: i22
- id: i23
- id: i24
- id: i25
- id: i31
- id: i32
- id: i33
- id: i34
- id: i35
neurons:
- id: n
  θ_syn: 5
  θ_seg: 1
  branches:
    - id: seg2
      θ_seg: 1
      θ_syn: 5
      branches:
        - id: seg1
          θ_syn: 5
synapses:
- {id: syn11, source: i11, target: seg1}
- {id: syn12, source: i12, target: seg1}
- {id: syn13, source: i13, target: seg1}
- {id: syn14, source: i14, target: seg1}
- {id: syn15, source: i15, target: seg1}
- {id: syn21, source: i21, target: seg2}
- {id: syn22, source: i22, target: seg2}
- {id: syn23, source: i23, target: seg2}
- {id: syn24, source: i24, target: seg2}
- {id: syn25, source: i25, target: seg2}
- {id: syn31, source: i31, target: n}
- {id: syn32, source: i32, target: n}
- {id: syn33, source: i33, target: n}
- {id: syn34, source: i34, target: n}
- {id: syn35, source: i35, target: n}
"""

(net1,objects1) = load_network(YAML_source=config1)

input=[ Event(:input_spikes, 0.0, t+τ-0.5, objects1[Symbol("$(pop)$(i)")]) for (t, pop) in volleys for (i,τ) in enumerate(5*rand(5))]

logger=simulate!(net1, input)

syn11 = get_trace(:syn11, logger.data)
syn12 = get_trace(:syn12, logger.data)
syn13 = get_trace(:syn13, logger.data)
syn14 = get_trace(:syn14, logger.data)
syn15 = get_trace(:syn15, logger.data)
syn21 = get_trace(:syn21, logger.data)
syn22 = get_trace(:syn22, logger.data)
syn23 = get_trace(:syn23, logger.data)
syn24 = get_trace(:syn24, logger.data)
syn25 = get_trace(:syn25, logger.data)
syn31 = get_trace(:syn31, logger.data)
syn32 = get_trace(:syn32, logger.data)
syn33 = get_trace(:syn33, logger.data)
syn34 = get_trace(:syn34, logger.data)
syn35 = get_trace(:syn35, logger.data)
seg1  = get_trace(:seg1, logger.data)
seg2  = get_trace(:seg2, logger.data)
n     = get_trace(:n, logger.data)

plateau1_starts = filter(x->(x.object == :seg1 && x.event == :plateau_starts), logger.data).t
plateau2_starts = filter(x->(x.object == :seg2 && x.event == :plateau_starts), logger.data).t
plateau1_ends = filter(x->(x.object == :seg1 && x.event == :plateau_ends), logger.data).t
plateau2_ends = filter(x->(x.object == :seg2 && x.event == :plateau_ends), logger.data).t
spike_times = filter(x->(x.object == :n && x.event == :spikes), logger.data).t

## Plot
xticks,xtickformat = make_manual_ticks([0;250;500;spike_times], ["0ms";"250ms";"500ms";["t"*['₁','₂','₃','₄','₅'][i] for i in 1:length(spike_times)]])
yticks,ytickformat = make_manual_ticks(0.5:14.5, reverse!(["$(grp)$(sub)" for grp in ["A","B","C"] for sub in ['₁','₂','₃','₄','₅']]))

ax11 = fig[2, 1] = Axis(fig;
    yticks=(2.5:5:15, ["C","B","A"]),
    yminorticks=IntervalsBetween(5, true),
    yminorticksvisible = true, yminorgridvisible = true,
    titlealign=:left, title = "b.    Without inhibition"
)

ax12 = fig[2, 2] = Axis(fig, aspect=DataAspect(), backgroundcolor=:transparent)
hidexdecorations!(ax11, grid=false)

for (i, syn) in enumerate([syn11,syn12,syn13,syn14,syn15])
    statetrace!(ax11, [0;syn.t;150], [0;syn.state;0], Dict(0=>:transparent, 1=>color_1), 9 + i, 0.9 )
end
statetrace!(ax11, [0;seg1.t;150], [DPC.voltage_low;seg1.state;DPC.voltage_low], Dict(DPC.voltage_low=>:transparent, DPC.voltage_elevated=>color_1_25, DPC.voltage_high=>color_1_50), 10, 4.9)

for (i, syn) in enumerate([syn21,syn22,syn23,syn24,syn25])
    statetrace!(ax11, [0;syn.t;150], [0;syn.state;0], Dict(0=>:transparent, 1=>color_2), 4 .+ i, 0.9 )
end
statetrace!(ax11, [0;seg2.t;150], [DPC.voltage_low;seg2.state;DPC.voltage_low], Dict(DPC.voltage_low=>:transparent, DPC.voltage_elevated=>color_2_25, DPC.voltage_high=>color_2_50), 5, 4.9)

for (i, syn) in enumerate([syn31,syn32,syn33,syn34,syn35])
    statetrace!(ax11, [0;syn.t;150], [0;syn.state;0], Dict(0=>:transparent, 1=>color_3), -1 .+ i, 0.9 )
end

arrows!(ax11, spike_times.+30, fill(2.5,length(spike_times)), fill(-20,length(spike_times)), fill(0,length(spike_times)))
text!(ax11,
    fill(" spike", length(spike_times)),
    position=Point2f0.(spike_times.+30,fill(2.5,length(spike_times))),
    align=(:left,:center),
    textsize=12
)
ylims!(ax11, (-0.5,15.5))

fig
##
pn1 = plot!(ax12, objects1[:n], ports=Dict(:n=>[:C],:seg2=>[:dummy, :B],:seg1=>[:dummy, :A]), angle_between=20/180*π, branch_width=0.2, branch_length=1.0, color=Dict(:n=>color_3, :seg1=>color_1, :seg2=>color_2))
hidedecorations!(ax12)

ports = Dict(pn1.attributes[:ports][])
arrows!(ax12, [ports[x] - Point2f0(0.4, 0) for x in [:C,:B,:A]] , fill(Point2f0(0.3,0),3), linewidth=2, arrowsize = -10)
for (label,pos)  in zip(["C","B","A"], [ports[x] - Point2f0(0.45, 0) for x in [:C,:B,:A]])
    text!(ax12, label, position=pos, align=(:right, :center), textsize=14, color=:black)
end

fig

## Second setup: with inhibition

config2 = """
refractory_duration: 5.01
inputs:
- id: i11
- id: i12
- id: i13
- id: i14
- id: i15
- id: i21
- id: i22
- id: i23
- id: i24
- id: i25
- id: i31
- id: i32
- id: i33
- id: i34
- id: i35
neurons:
- id: n
  θ_syn: 5
  θ_seg: 1
  branches:
    - id: seg2
      θ_seg: 1
      θ_syn: 5
      branches:
        - id: seg1
          θ_syn: 5
synapses:
- {id: syn11, source: i11, target: seg1}
- {id: syn12, source: i12, target: seg1}
- {id: syn13, source: i13, target: seg1}
- {id: syn14, source: i14, target: seg1}
- {id: syn15, source: i15, target: seg1}
- {id: syn21, source: i21, target: seg2}
- {id: syn22, source: i22, target: seg2}
- {id: syn23, source: i23, target: seg2}
- {id: syn24, source: i24, target: seg2}
- {id: syn25, source: i25, target: seg2}
- {id: syn31, source: i31, target: n}
- {id: syn32, source: i32, target: n}
- {id: syn33, source: i33, target: n}
- {id: syn34, source: i34, target: n}
- {id: syn35, source: i35, target: n}
- {id: syn41, source: i31, target: seg1, weight: -1, delay: 10.1}
- {id: syn42, source: i32, target: seg1, weight: -1, delay: 10.1}
- {id: syn43, source: i33, target: seg1, weight: -1, delay: 10.1}
- {id: syn44, source: i34, target: seg1, weight: -1, delay: 10.1}
- {id: syn45, source: i35, target: seg1, weight: -1, delay: 10.1}
"""

(net2,objects2) = load_network(YAML_source=config2)

volleys = [
    (70.0, :i1), (120.0, :i2), (170.0, :i3),
    (315.0, :i3), (325.0, :i2), (335.0, :i1),
    (370.0, :i3), (380.0, :i2), (390.0, :i1),
    (425.0, :i3), (435.0, :i2), (445.0, :i1)
]
input=[ Event(:input_spikes, 0.0, t+τ-0.5, objects2[Symbol("$(pop)$(i)")]) for (t, pop) in volleys for (i,τ) in enumerate(5*rand(5))]

logger=simulate!(net2, input)

syn11 = get_trace(:syn11, logger.data)
syn12 = get_trace(:syn12, logger.data)
syn13 = get_trace(:syn13, logger.data)
syn14 = get_trace(:syn14, logger.data)
syn15 = get_trace(:syn15, logger.data)
syn21 = get_trace(:syn21, logger.data)
syn22 = get_trace(:syn22, logger.data)
syn23 = get_trace(:syn23, logger.data)
syn24 = get_trace(:syn24, logger.data)
syn25 = get_trace(:syn25, logger.data)
syn31 = get_trace(:syn31, logger.data)
syn32 = get_trace(:syn32, logger.data)
syn33 = get_trace(:syn33, logger.data)
syn34 = get_trace(:syn34, logger.data)
syn35 = get_trace(:syn35, logger.data)
syn41 = get_trace(:syn41, logger.data)
syn42 = get_trace(:syn42, logger.data)
syn43 = get_trace(:syn43, logger.data)
syn44 = get_trace(:syn44, logger.data)
syn45 = get_trace(:syn45, logger.data)
syn51 = get_trace(:syn51, logger.data)
syn52 = get_trace(:syn52, logger.data)
syn53 = get_trace(:syn53, logger.data)
syn54 = get_trace(:syn54, logger.data)
syn55 = get_trace(:syn55, logger.data)
seg1  = get_trace(:seg1, logger.data)
seg2  = get_trace(:seg2, logger.data)
n     = get_trace(:n, logger.data)

plateau1_starts = filter(x->(x.object == :seg1 && x.event == :plateau_starts), logger.data).t
plateau2_starts = filter(x->(x.object == :seg2 && x.event == :plateau_starts), logger.data).t
plateau1_ends = filter(x->(x.object == :seg1 && x.event == :plateau_ends), logger.data).t
plateau2_ends = filter(x->(x.object == :seg2 && x.event == :plateau_ends), logger.data).t
spike_times = filter(x->(x.object == :n && x.event == :spikes), logger.data).t

## Plot
ax21 = fig[3, 1] = Axis(fig;
    yticks=(2.5:5:15, ["C","B","A"]),
    yminorticks=Makie.IntervalsBetween(5, true),
    yminorticksvisible = true, yminorgridvisible = true,
    titlealign=:left, title = "c.    With inhibition",
    xlabel="time [ms]"
)
ax22 = fig[3, 2] = Axis(fig, aspect=DataAspect(), backgroundcolor=:transparent)

for (i, syn) in enumerate([syn11,syn12,syn13,syn14,syn15])
    statetrace!(ax21, [0;syn.t;150], [0;syn.state;0], Dict(0=>:transparent, 1=>color_1), 9 .+ i ,  0.9 )
end
statetrace!(ax21, [0;seg1.t;150], [DPC.voltage_low;seg1.state;DPC.voltage_low], Dict(DPC.voltage_low=>:transparent, DPC.voltage_elevated=>color_1_25, DPC.voltage_high=>color_1_50), 10 , 4.9)

for (i, syn) in enumerate([syn21,syn22,syn23,syn24,syn25])
    statetrace!(ax21, [0;syn.t;150], [0;syn.state;0], Dict(0=>:transparent, 1=>color_2), 4 .+ i ,  0.9 )
end
statetrace!(ax21, [0;seg2.t;150],  [DPC.voltage_low;seg2.state;DPC.voltage_low], Dict(DPC.voltage_low=>:transparent, DPC.voltage_elevated=>color_2_25, DPC.voltage_high=>color_2_50), 5 ,  5 )

for (i, syn) in enumerate([syn31,syn32,syn33,syn34,syn35])
    statetrace!(ax21, [0;syn.t;150], [0;syn.state;0], Dict(0=>:transparent, 1=>color_3), -1 .+ i ,  0.9 )
end

arrows!(ax21, spike_times.+30, fill(2.5,length(spike_times)), fill(-20,length(spike_times)), fill(0,length(spike_times)))
text!(ax21,
    fill(" spike", length(spike_times)),
    position=Point2f0.(spike_times.+30,fill(2.5,length(spike_times))),
    align=(:left,:center),
    textsize=12
)

ylims!(ax21, (-0.5,15.5))

pn2 = plot!(ax22, objects2[:n], ports=Dict(:n=>[:C],:seg2=>[:dummy, :B],:seg1=>[:dummy, :A, :nC1]), angle_between=20/180*π, branch_width=0.2, branch_length=1.0, color=Dict(:n=>color_3, :seg1=>color_1, :seg2=>color_2))
hidedecorations!(ax22)

ports = Dict(pn2.attributes[:ports][])
arrows!(ax22, [ports[x] - Point2f0(0.4, 0) for x in [:C,:B,:A]] , fill(Point2f0(0.3,0),3), linewidth=2, arrowsize = -10)
for (label,pos)  in zip(["C","B","A"], [ports[x] - Point2f0(0.45, 0) for x in [:C,:B,:A]])
    text!(ax22, label, position=pos, align=(:right, :center), textsize=14, color=:black)
end

arrows!(ax22, [ports[:nC1] + Point2f0(0.35, 0)] , [Point2f0(-0.3,0)], linewidth=2, arrowsize = -10)
text!(ax22, "¬C", position=ports[:nC1]+Point2f0(0.4,0), align=(:left, :center), textsize=14, color=:black)

xlims!(ax12, (-1,1))
xlims!(ax22, (-1,1))
linkxaxes!(ax0,ax11)
linkaxes!(ax11, ax21)
fig

## Save
colsize!(fig.layout, 2, Relative(0.2))
rowsize!(fig.layout, 1, Relative(0.15))

save(joinpath("figures","inhibition.pdf"), fig)
save(joinpath("figures","inhibition.svg"), fig)
save(joinpath("figures","inhibition.png"), fig)
fig
################################################################################


text:
struct RandomOutlierQs <: ModelBasedQs
    occ::SVDD.OCClassifier
end

function qs_score(qs::RandomOutlierQs, x::Array{T, 2}, labels::Dict{Symbol, Array{Int, 1}})::Array{Float64, 1} where T <: Real
    prediction = classify.(predict(qs.occ, x))
    scores = rand(size(x, 2))
    scores[prediction .== :inlier] .= 0
    return scores
end


text:
"""
$(SIGNATURES)

Fixed point CPU StructArray carrier replica generation
"""
function gen_carrier_replica!(
    carrier_replica::StructArray{Complex{T}},
    carrier_frequency,
    sampling_frequency,
    start_phase,
    start_sample,
    num_samples
) where T
    c_re = carrier_replica.re; c_im = carrier_replica.im
    carrier_freq = upreferred(carrier_frequency / Hz)
    sampling_freq = upreferred(sampling_frequency / Hz)
    @avx for i in 0:num_samples - 1
        c_im[i + start_sample], c_re[i + start_sample] =
            sincos(T(2π) * (i * T(carrier_freq) / T(sampling_freq) + T(start_phase)))
    end
    carrier_replica
end

"""
$(SIGNATURES)

Updates the carrier phase.
"""
function update_carrier_phase(
    num_samples,
    carrier_frequency,
    sampling_frequency,
    start_phase,
)
    phase = num_samples * carrier_frequency / sampling_frequency + start_phase
    mod(phase + 0.5, 1) - 0.5
end

"""
$(SIGNATURES)

Calculates the current carrier frequency.
"""
function get_current_carrier_frequency(intermediate_frequency, carrier_doppler)
    carrier_doppler + intermediate_frequency
end


text:
using Aqua
using Base64
using CompatHelper
using CompatHelper: DepInfo, EntryType
using Dates
using GitForge
using GitForge: GitForge, GitHub, GitLab
using Mocking
using Pkg
using Random
using SHA
using Test
using TimeZones
using TOML
using UUIDs

Mocking.activate()

Aqua.test_all(CompatHelper; ambiguities=false)

include("patches.jl")

@testset "CompatHelper.jl" begin
    include(joinpath("utilities", "ci.jl"))
    include(joinpath("utilities", "git.jl"))
    include(joinpath("utilities", "new_versions.jl"))
    include(joinpath("utilities", "ssh.jl"))
    include(joinpath("utilities", "types.jl"))
    include(joinpath("utilities", "utilities.jl"))
    include(joinpath("utilities", "rate_limiting.jl"))
    include(joinpath("utilities", "timestamps.jl"))

    include("dependencies.jl")
    include("exceptions.jl")
    include("pull_requests.jl")
    include("main.jl")

    COMPATHELPER_RUN_INTEGRATION_TESTS = get(ENV, "COMPATHELPER_RUN_INTEGRATION_TESTS", "")
    if COMPATHELPER_RUN_INTEGRATION_TESTS == "true"
        @testset "Integration Tests" begin
            @info "Running Integration Tests"
            include(joinpath("utilities", "integration_tests.jl"))
            include("integration_tests.jl")
        end
    end
end


text:
using OrdinaryDiffEq, DiffEqJump, Test
j1 = ConstantRateJump((u,p,t)->10,(integrator)->integrator.u[1]+=1)
j2 = ConstantRateJump((u,p,t)->1u[1],(integrator)->integrator.u[1]-=1)
u0 = [0]

prob = DiscreteProblem(u0, (0., 100.), [])
jump_prob = JumpProblem(prob, Direct(), j1, j2)
sol = solve(EnsembleProblem(jump_prob), FunctionMap(), trajectories=3)
@test Array(sol[1]) !== Array(sol[2])
@test Array(sol[1]) !== Array(sol[3])
@test Array(sol[2]) !== Array(sol[3])
sol = solve(EnsembleProblem(jump_prob), SSAStepper(), trajectories=3)
@test Array(sol[1]) !== Array(sol[2])
@test Array(sol[1]) !== Array(sol[3])
@test Array(sol[2]) !== Array(sol[3])


text:
# Implement all the belief related function

const SingleOIBelief = SingleOIDistribution

mutable struct SingleOIUpdater <: Updater
    pomdp::SingleOIPOMDP
end

#for sarsop for exploration
#function POMDPs.initialstate(pomdp::SingleOIPOMDP)
#    state_space = states(pomdp)
#    states_to_add = SingleOIState[]
#    for s in state_space
#        if !s.crash
#            push!(states_to_add, s)
#        end
#    end
#    probs = ones(length(states_to_add))
#    normalize!(probs, 1)
#    return SingleOIDistribution(probs, states_to_add)
#end

function POMDPs.initialstate(pomdp::SingleOIPOMDP)
     env = pomdp.env
     v_grid = get_v_grid(pomdp)
     states = SingleOIState[]
     ego = ego_state(pomdp, pomdp.s_start, 0.)
     for lane_id in LANE_ID_LIST
         for s in get_lane_s(pomdp, lane_id)
             for v in v_grid[3:end]
                 car = car_state(pomdp, lane_id, s, v)
                 if !is_observable_fixed(ego, car, env)
                     push!(states, SingleOIState(is_crash(pomdp, ego, car), ego, car))
                 end
             end
         end
     end

     # Add also the off the grid state, weight using the probability p_birth
     n_in_grid = length(states)
     push!(states, SingleOIState(false, ego, get_off_the_grid(pomdp)))
     n_off_grid = length(states) - n_in_grid
     probs = ones(length(states))
     probs[1:n_in_grid] .= pomdp.p_birth/n_in_grid
     probs[n_in_grid + 1:end] .= (1.0 - pomdp.p_birth)/n_off_grid
     @assert n_off_grid == 1
     @assert sum(probs) ≈ 1.
     # normalize!(probs, 1)
     return SingleOIBelief(probs, states)
end

function initial_ego_state(pomdp::SingleOIPOMDP)
    return ego_state(pomdp, pomdp.s_start, 0.)
end

# Updates the belief given the current action and observation
function POMDPs.update(bu::SingleOIUpdater, bold::SingleOIBelief, a::SingleOIAction, o::SingleOIObs)
    bnew = SingleOIBelief()
    pomdp = bu.pomdp
    # initialize spaces
    pomdp_states = ordered_states(pomdp)

    # iterate through each state in belief vector
    for (i, sp) in enumerate(pomdp_states)
        # get the distributions
        od = observation(pomdp, a, sp)
        # get prob of observation o from current distribution
        probo = pdf(od, o)
        # if observation prob is 0.0, then skip rest of update b/c bnew[i] is zero
        if probo == 0.0
            continue
        end
        b_sum = 0.0 # belief for state sp
        for (j, s) in enumerate(bold.it)
            td = transition(pomdp, s, a)
            pp = pdf(td, sp)
            @inbounds b_sum += pp * bold.p[j]
        end
        if b_sum != 0.
            push!(bnew.it, sp)
            push!(bnew.p, probo*b_sum)
        end
    end

    # if norm is zero, the update was invalid - reset to uniform
    if sum(bnew.p) == 0.0
        println("Invalid update for: ", bold, " ", a, " ", o)
        throw("UpdateError")
    else
        normalize!(bnew.p, 1)
    end
    bnew
end


text:
#!/usr/bin/env julia

using PyCall
@pyimport matplotlib as mpl
mpl.rc("text", usetex=true)
mpl.use("Agg")
@pyimport matplotlib.ticker as mti
using PyPlot
using JLD2
using FileIO
using ArgParse
using Printf

o = 1
function singleplot(filename::String)
  d = load(filename*".jld2")
  x = d["x"]
  t = d["t"]
  n = d["n"]
  m = d["m"]
  println(m)
  mpl.rc("font", family="serif", size = 7)
  fig, ax = subplots(figsize = (2.5, 2.))
  col = ["black", "gray", "silver", "gray", "brown", "orange"]
  marker = [":s", ":o", ":v", ":<", ":>", ":d"]
  for i in 1:size(d["cumulants"], 2)
    if occursin("nblocks", filename)  | occursin("maxf", filename)
      str = replace(@sprintf("%.1e", t[i]), "0"=>"")
      tt = "\$"*replace(str, "e+"=>"\\times 10^{")*"}\$"
      ax[:plot](d[x], d["cumulants"][:,i], marker[i], label= "\$ t_{up} \$ = $tt", color = col[i], markersize=2.5, linewidth = 1)
    elseif occursin("nprocs", filename)
      str = replace(@sprintf("%.1e", t[i]), "0"=>"")
      tt = "\$"*replace(str, "e+"=>"\\times 10^{")*"}\$"
      y = d["cumulants"][:,i].\d["cumulants"][1,i]
      ax[:plot](d[x][1:6], y[1:6], marker[i], label= "\$ t_{up} \$ = $tt", color = col[i], markersize=2.5, linewidth = 1)
    else
      tt =  n[i]
      comptimes = d["cumulants"][:,i]./d["cumulants updat"][:,i]
      ax[:plot](d[x], comptimes, marker[i], label= "n = $tt", color = col[i], markersize=2.5, linewidth = 1)
      #i == size(d["cumulants"], 2)? ax[:plot](d[x], d["tm"], ":<", label= "theoretical", color = "black", markersize=2.5, linewidth = 1):()
    end
  end
  subplots_adjust(bottom = 0.16, top=0.92, left = 0.12, right = 0.92)
  fx = matplotlib[:ticker][:ScalarFormatter]()
  fx[:set_powerlimits]((-1, 4))
  ax[:xaxis][:set_major_formatter](fx)
  if occursin("nblocks", filename)
    subplots_adjust(left = 0.15)
    PyPlot.ylabel("computational time [s]", labelpad = 0.6)
    PyPlot.xlabel(x, labelpad = 0.6)
    ax[:xaxis][:set_major_locator](mti.MaxNLocator(integer=true))
    ax[:legend](fontsize = 5, loc = (m==5) ? 2 : 9, ncol = 1)
  elseif occursin("nprocs", filename)
    subplots_adjust(left = 0.15)
    PyPlot.ylabel("computational speedup", labelpad = 0)
    PyPlot.xlabel(x, labelpad = 0.6)
    ax[:xaxis][:set_major_locator](mti.MaxNLocator(integer=true))
    ax[:legend](fontsize = 5, loc = 2, ncol = 1)
  elseif occursin("maxf", filename)
    subplots_adjust(left = 0.18)
    PyPlot.ylabel("frequency [Hz]", labelpad = 0.)
    PyPlot.xlabel(x, labelpad = 0.6)
    ax[:legend](fontsize = 5, loc = 1, ncol = 1)
  else
    PyPlot.ylabel("computational speedup", labelpad = -0.5)
    PyPlot.xlabel("\$ t_{up} \$", labelpad = 0.6)
    subplots_adjust(bottom = 0.16, top=0.92, left = 0.14, right = 0.92)
    if m == 6
      ax[:yaxis][:set_major_locator](mti.MaxNLocator(integer=true))
    end
    ax[:legend](fontsize = 6, loc = 1, ncol = 1)
  end
  f = matplotlib[:ticker][:ScalarFormatter]()
  f[:set_powerlimits]((-3, 2))
  ax[:yaxis][:set_major_formatter](f)
  fig[:savefig]("cumulants"*filename*".eps")
end

function main(args)
  s = ArgParseSettings("description")
  @add_arg_table s begin
    "file"
    help = "the file name"
    arg_type = String
  end
  parsed_args = parse_args(s)
  filename = parsed_args["file"]
  singleplot(replace(filename, ".jld2"=>""))
end

main(ARGS)


text:
module Data

include("functions.jl")

end


text:
import Base: iterate

"""
    iterate(trigger, dt[, n=number_of_times])

Iterate from instant `dt` using trigger with a given iteration number `n`
if `n < 0` (`-1` by default), it iterates indefinitely.

# Usage
```
julia> trigger = Trigger(Dates.Time(20, 30))

julia> for dt in iterate(trigger, DateTime(2020, 1, 1), n=3)
         @show dt
       end
dt = 2020-01-01T20:30:00
dt = 2020-01-02T20:30:00
dt = 2020-01-03T20:30:00

julia> collect(iterate(trigger, DateTime(2020, 1, 1), n=3))
3-element Array{Any,1}:
 2020-01-01T20:30:00
 2020-01-02T20:30:00
 2020-01-03T20:30:00
```
"""
function iterate(trigger::AbstractTrigger, dt; n=-1)
    TriggerIterator(trigger, dt, n)
end

abstract type AbstractTriggerIterator end

struct FiniteTriggerIterator <: AbstractTriggerIterator
    trigger::AbstractTrigger
    dt
    n::Int
end

struct InfiniteTriggerIterator <: AbstractTriggerIterator
    trigger::AbstractTrigger
    dt
end

function TriggerIterator(trigger::AbstractTrigger, dt, n)
    if n < 0
        if IteratorSize(trigger) == IsInfinite()
            InfiniteTriggerIterator(trigger, dt)
        else
            FiniteTriggerIterator(trigger, dt, trigger.n)
        end
    else
        if IteratorSize(trigger) == IsInfinite()
            FiniteTriggerIterator(trigger, dt, n)
        else
            FiniteTriggerIterator(trigger, dt, min(n, trigger.n))
        end
    end
end

function iterate(itr::AbstractTriggerIterator, state=(itr.dt, 0))
    dt_now, i = state
    i >= itr.n && return nothing
    dt_next = get_next_dt_fire(itr.trigger, DateTime(0), dt_now)
    i += 1
    state = dt_next, i
    return (dt_next, state)
end

function iterate(itr::InfiniteTriggerIterator, state=(itr.dt, 0))
    dt_now, i = state
    dt_next = get_next_dt_fire(itr.trigger, DateTime(0), dt_now)
    i += 1
    state = dt_next, i
    return (dt_next, state)
end

IteratorSize(itr::InfiniteTriggerIterator) = IsInfinite()
IteratorSize(itr::FiniteTriggerIterator) = HasLength()

length(itr::FiniteTriggerIterator) = itr.n


text:
abstract type NuclearDistribution <: NonadiabaticDistribution end

abstract type PositionDistribution <: NuclearDistribution end
abstract type VelocityDistribution <: NuclearDistribution end

include("dynamical_distribution.jl")
export DynamicalDistribution

include("boltzmann_velocity.jl")
export BoltzmannVelocityDistribution

include("harmonic_wigner.jl")
export MomentumHarmonicWigner
export PositionHarmonicWigner
export VelocityHarmonicWigner

include("harmonic_ring_polymer.jl")
export HarmonicRingPolymer


text:
using Equations
using Base.Test

include("commonTests.jl")
include("DivTests.jl")
include("equationsTests.jl")
include("matchersTests.jl")
include("PowTests.jl")
include("DerTests.jl")
include("VecTests.jl")
include("plasmaTests.jl")
include("LogTests.jl")
include("TenTests.jl")
include("gravity.jl")
include("readmeTests.jl")
include("SqrtTests.jl")