text:
using ClimateMachine.ODESolvers
using ClimateMachine.LinearSolvers

const slow_mrrk_methods =
    ((LSRK54CarpenterKennedy, 4), (LSRK144NiegemannDiehlBusch, 4))
const fast_mrrk_methods = (
    (LSRK54CarpenterKennedy, 4),
    (LSRK144NiegemannDiehlBusch, 4),
    (SSPRK33ShuOsher, 3),
    (SSPRK34SpiteriRuuth, 3),
)
const explicit_methods = (
    (LSRK54CarpenterKennedy, 4),
    (LSRK144NiegemannDiehlBusch, 4),
    (SSPRK33ShuOsher, 3),
    (SSPRK34SpiteriRuuth, 3),
    (LSRKEulerMethod, 1),
)

const imex_methods = (
    (ARK2GiraldoKellyConstantinescu, 2),
    (ARK437L2SA1KennedyCarpenter, 4),
    (ARK548L2SA2KennedyCarpenter, 5),
)

const mis_methods =
    ((MIS2, 2), (MIS3C, 2), (MIS4, 3), (MIS4a, 3), (TVDMISA, 2), (TVDMISB, 2))

const mrigark_erk_methods = ((MRIGARKERK33aSandu, 3), (MRIGARKERK45aSandu, 4))

const mrigark_irk_methods = (
    (MRIGARKESDIRK24LSA, 2),
    (MRIGARKESDIRK23LSA, 2),
    (MRIGARKIRK21aSandu, 2),
    (MRIGARKESDIRK34aSandu, 3),
    (MRIGARKESDIRK46aSandu, 4),
)

const fast_mrigark_methods =
    ((LSRK54CarpenterKennedy, 4), (LSRK144NiegemannDiehlBusch, 4))

struct DivideLinearSolver <: AbstractLinearSolver end
function LinearSolvers.prefactorize(
    linearoperator!,
    ::DivideLinearSolver,
    args...,
)
    linearoperator!
end
function LinearSolvers.linearsolve!(
    linearoperator!,
    ::DivideLinearSolver,
    Qtt,
    Qhat,
    args...,
)
    @. Qhat = 1 / Qhat
    linearoperator!(Qtt, Qhat, args...)
    @. Qtt = 1 / Qtt
end


text:
const kb_eV = 8.6173303e-5 #Boltzmann const in eV/K
const h_eV = 4.135667662e-15 #Planck const in eV*s
const hbar_eV = 6.582119514e-16 #Reduced Planck const in eV*s
const qe = 1.6021766208e-19 #Electron charge in C
const mu_0 = 1.2566370614359173e-06 #Vacuum permeability H/m

ext: jl


text:
module GateOps

export GateOps
export pauli_x, pauli_y, pauli_z, hadamard, u
export r_x, r_y, r_z, r_phase, swap
export c_pauli_x, c_pauli_y, c_pauli_z, c_u
export c_r_x, c_r_y, c_r_z, c_r_phase

using LinearAlgebra
using Memoize

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
#                       Redirect I/O streams
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
#(rd, wr) = redirect_stdout();

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
#                       Lightweight gate calls
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

"Abstract Gate type"
abstract type AGate end
"Abstract Gate Label, for unique gates"
abstract type AGateLabel <: AGate end
"Abstract Gate Call, for tracking Gate labels applied to specific qubits"
abstract type AGateCall <: AGate end

"Parametric Gate label. Tracks the gate symbol (:x,:y,:z, etc), and arbitrary parameters"
struct GateLabelP{IType<:Integer, NType<:Number} <: AGateLabel
    label::Symbol
    params::Union{Nothing, Dict{String, Union{IType, NType, Bool}}}
    GateLabelP{IType,NType}(gate_label) where {IType<:Integer, NType<:Number} = new(gate_label, nothing)
    GateLabelP{IType,NType}(gate_label, params) where {IType<:Integer, NType<:Number} = new(gate_label, params)
end

"Parametric single qubit Gate call. Has a GateLabelP, target qubit index and register label"
struct GateCall1P{IType<:Integer, NType<:Number} <: AGateCall
    gate_label::GateLabelP{IType, NType}
    target::IType
    reg::Union{String, Nothing}
    GateCall1P{IType,NType}(gate_label, target, reg) where {IType<:Integer, NType<:Number} = new(gate_label, target, reg)
    GateCall1P{IType,NType}(gate_label, target) where {IType<:Integer, NType<:Number} = new(gate_label, target, nothing)
end

"""
Parametric two qubit Gate call.
Has a GateLabelP, control and target qubit indices, register label, and base gate (assumes controlled U).
"""
struct GateCall2P{IType<:Integer, NType<:Number} <: AGateCall
    gate_label::GateLabelP{IType, NType}
    ctrl::IType
    target::IType
    base_gate::Union{GateCall1P{IType, NType}, GateCall2P{IType, NType}, Nothing}
    #base_gate::Union{GateLabelP{IType, NType}, Nothing}
    reg::Union{String, Nothing}
    GateCall2P{IType,NType}(gate_label, ctrl, target, base_gate, reg) where {IType<:Integer, NType<:Number} =
    new(gate_label, ctrl, target, base_gate, reg)

    GateCall2P{IType,NType}(gate_label, ctrl, target, base_gate) where {IType<:Integer, NType<:Number} =
    new(gate_label, ctrl, target, base_gate, nothing)

    GateCall2P{IType,NType}(gate_label, ctrl, target) where {IType<:Integer, NType<:Number} =
    new(gate_label, ctrl, target, nothing, nothing)
end

"""
Parametric n-qubit Gate call.
Has a GateLabelP, control vector indices, target qubit index, register label, and base gate.
"""
struct GateCallNP{IType<:Integer, NType<:Number} <: AGateCall
    gate_label::GateLabelP{IType, NType}
    ctrl::Vector{IType}
    target::IType
    base_gate::Union{GateCall1P{IType, NType}, GateCall2P{IType, NType}, GateCallNP{IType, NType}, Nothing}
    reg::Union{String,Nothing}
    GateCallNP{IType,NType}(gate_label, ctrl, target, base_gate, reg) where {IType<:Integer, NType<:Number} = new(gate_label, ctrl, target, base_gate{IType,NType}, reg)
    GateCallNP{IType,NType}(gate_label, ctrl, target, base_gate) where {IType<:Integer, NType<:Number} = new(gate_label, ctrl, target, base_gate{IType,NType}, "")
end

# Useful aliases
GateLabel = GateLabelP{Int64, Float64}
GateCall1 = GateCall1P{Int64, Float64}
GateCall2 = GateCall2P{Int64, Float64}
GateCallN = GateCallNP{Int64, Float64}

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
#                           Gate implementations
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

"Wrap of Matrix to fit type system hierarchy"
struct Gate <: AGateCall
    mat::Matrix{<:Number}
    Gate(mat) = new(mat)
end

"Placeholder struct for Euler-angle gate"
struct GateEulerP{NType<:Number} <: AGateCall
    θ::NType
    ϕ::NType
    λ::NType
    phase::NType
    GateEulerP{NType}(θ, ϕ, λ) where {NType<:Number} = new(θ, ϕ, λ, 0.0)
    GateEulerP{NType}(θ, ϕ, λ, phase) where {NType<:Number} = new(θ, ϕ, λ, phase)
    function GateEulerP{NType}(params::Dict{String,NType}) where {NType<:Number}
        θ, ϕ, λ, phase = params["θ"], params["ϕ"], params["λ"], params["phase"]
        return GateEulerP{NType}(θ, ϕ, λ, phase)
    end
end
GateEuler = GateEulerP{Float64}

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
#                            1Q Gate calls ops
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

"""
    pauli_x(q_target::Int, register::Union{String, Nothing}=nothing)

Generate a single qubit Pauli-x GateCall (GateCall1P) applied to the target qubit (on given register, if provided)

# Examples
```julia-repl
julia> GateOps.pauli_x(0, "q")
QuantExQASM.GateOps.GateCall1P{Int64,Float64}(QuantExQASM.GateOps.GateLabelP{Int64,Float64}(:x, nothing), 0, "q")
```
"""
function pauli_x(q_target::Int, register::Union{String, Nothing}=nothing)
    if register == nothing
        return GateCall1(GateLabel(:x), q_target)
    else
        return GateCall1(GateLabel(:x), q_target, register)
    end
end

"""
    pauli_y(q_target::Int, register::Union{String, Nothing}=nothing)

Generate a single qubit Pauli-y GateCall (GateCall1P) applied to the target qubit (on given register, if provided)

# Examples
```julia-repl
julia> GateOps.pauli_y(0, "q")
QuantExQASM.GateOps.GateCall1P{Int64,Float64}(QuantExQASM.GateOps.GateLabelP{Int64,Float64}(:y, nothing), 0, "q")
```
"""
function pauli_y(q_target::Int, register::Union{String, Nothing}=nothing)
    if register == nothing
        return GateCall1(GateLabel(:y), q_target)
    else
        return GateCall1(GateLabel(:y), q_target, register)
    end
end

"""
    pauli_z(q_target::Int, register::Union{String, Nothing}=nothing)

Generate a single qubit Pauli-z GateCall (GateCall1P) applied to the target qubit (on given register, if provided)

# Examples
```julia-repl
julia> GateOps.pauli_z(0, "q")
QuantExQASM.GateOps.GateCall1P{Int64,Float64}(QuantExQASM.GateOps.GateLabelP{Int64,Float64}(:z, nothing), 0, "q")
```
"""
function pauli_z(q_target::Int, register::Union{String, Nothing}=nothing)
    if register == nothing
        return GateCall1(GateLabel(:z), q_target)
    else
        return GateCall1(GateLabel(:z), q_target, register)
    end
end

"""
    hadamard(q_target::Int, register::Union{String, Nothing}=nothing)

Generate a single qubit Hadamard GateCall (GateCall1P) applied to the target qubit (on given register, if provided)

# Examples
```julia-repl
julia> GateOps.hadamard(0, "q")
GateOps.GateCall1P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(:h, nothing), 0, "q")
```
"""
function hadamard(q_target::Int, register::Union{String, Nothing}=nothing)
    if register == nothing
        return GateCall1(GateLabel(:h), q_target)
    else
        return GateCall1(GateLabel(:h), q_target, register)
    end
end

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

struct u_gate_schema{NType<:Number}
    schema::Union{Tuple{NType,NType,NType,NType}, Tuple{NType, Bool}}
    type::String
    u_gate_schema{NType}(schema) where {NType<:Number} = length(schema) == 2 ? new(schema, "r") : new(schema, "e")
end

"""
    u(label::GateLabel, q_target::Int, register::Union{String, Nothing}=nothing)

Generate a single qubit arbitrary unitary GateCall (GateCall1P) applied to the target qubit (on given register, if provided)

# Examples
```julia-repl
julia> GateOps.u(GateOps.GateLabel(:mygate), 1, "qr")
GateOps.GateCall1P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(:mygate, nothing), 1, "qr")
```
"""
function u(label::GateLabel, q_target::Int, register::Union{String, Nothing}=nothing)
    return GateCall1( label, q_target, register)
end
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

"""
    r_x(q_target::Int, theta::Real, register::Union{String, Nothing}=nothing)

Generate a single qubit R_x(θ) GateCall (GateCall1P) applied to the target qubit (on given register, if provided)

# Examples
```julia-repl
julia> GateOps.r_x(3,pi/2)
GateOps.GateCall1P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(Symbol("r_x_angle=1.5707963267948966"), Dict{String,Union{Bool, Float64, Int64}}("angle"=>1.5708)), 3, nothing)
```
"""
function r_x(q_target::Int, theta::Number, register::Union{String, Nothing}=nothing)
    if register == nothing
        return GateCall1(GateLabel(Symbol("r_x" * "_" * "angle=" * string(theta)), Dict("angle"=>theta)), q_target )
    else
        return GateCall1(GateLabel(Symbol("r_x" * "_" * "angle=" * string(theta)), Dict("angle"=>theta)), q_target, register )
    end
end

"""
    r_y(q_target::Int, theta::Real, register::Union{String, Nothing}=nothing)

Generate a single qubit R_y(θ) GateCall (GateCall1P) applied to the target qubit (on given register, if provided)

# Examples
```julia-repl
julia> GateOps.r_y(3,pi/2)
GateOps.GateCall1P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(Symbol("r_y_angle=1.5707963267948966"), Dict{String,Union{Bool, Float64, Int64}}("angle"=>1.5708)), 3, nothing)
```
"""
function r_y(q_target::Int, theta::Number, register::Union{String, Nothing}=nothing)
    if register == nothing
        return GateCall1(GateLabel(Symbol("r_y" * "_" * "angle=" * string(theta)), Dict("angle"=>theta)), q_target )
    else
        return GateCall1(GateLabel(Symbol("r_y" * "_" * "angle=" * string(theta)), Dict("angle"=>theta)), q_target, register )
    end
end

"""
    r_z(q_target::Int, theta::Real, register::Union{String, Nothing}=nothing)

Generate a single qubit R_z(θ) GateCall (GateCall1P) applied to the target qubit (on given register, if provided)

# Examples
```julia-repl
julia> GateOps.r_z(3,pi/2)
GateOps.GateCall1P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(Symbol("r_z_angle=1.5707963267948966"), Dict{String,Union{Bool, Float64, Int64}}("angle"=>1.5708)), 3, nothing)
```
"""
function r_z(q_target::Int, theta::Number, register::Union{String, Nothing}=nothing)
    if register == nothing
        return GateCall1(GateLabel(Symbol("r_z" * "_" * "angle=" * string(theta)), Dict("angle"=>theta)), q_target )
    else
        return GateCall1(GateLabel(Symbol("r_z" * "_" * "angle=" * string(theta)), Dict("angle"=>theta)), q_target, register )
    end
end

"""
    r_phase(q_target::Int, theta::Real, register::Union{String, Nothing}=nothing)

Generate a single qubit phase shift GateCall ( diag(1, exp(1im*theta)) , GateCall1P) applied to the target qubit (on given register, if provided)

# Examples
```julia-repl
julia> GateOps.r_phase(3,pi/2)
GateOps.GateCall1P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(Symbol("r_phase_angle=1.5707963267948966"), Dict{String,Union{Bool, Float64, Int64}}("angle"=>1.5708)), 3, nothing)
```
"""
function r_phase(q_target::Int, theta::Number, register::Union{String, Nothing}=nothing)
    if register == nothing
        return GateCall1(GateLabel(Symbol("r_phase" * "_" * "angle=" * string(theta)), Dict("angle"=>theta)), q_target )
    else
        return GateCall1(GateLabel(Symbol("r_phase" * "_" * "angle=" * string(theta)), Dict("angle"=>theta)), q_target, register )
    end
end

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
#                            2Q Gate calls ops
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

function swap(q_target::Int, q_ctrl::Int, register::Union{String, Nothing}=nothing)
    return GateCall2( GateLabel(:swap), q_target, q_ctrl)
end

"""
    c_pauli_x(q_target::Int, q_ctrl::Int, register::Union{String, Nothing}=nothing)

Generate a controlled Pauli-x (two qubit) GateCall (GateCall2P), controlled on the index `q_ctrl` applied to the target `q_target` (on given register, if provided).

# Examples
```julia-repl
julia> GateOps.c_pauli_x(0, 1, "qreg")
GateOps.GateCall2P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(:c_x, nothing), 0, 1, GateOps.GateCall1P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(:x, nothing), 0, "qreg"), "qreg")
```
"""
function c_pauli_x(q_target::Int, q_ctrl::Int, register::Union{String, Nothing}=nothing)
    return GateCall2( GateLabel(:c_x), q_target, q_ctrl, pauli_x(q_target,register), register)
end

"""
    c_pauli_y(q_target::Int, q_ctrl::Int, register::Union{String, Nothing}=nothing)

Generate a controlled Pauli-y (two qubit) GateCall (GateCall2P), controlled on the index `q_ctrl` applied to the target `q_target` (on given register, if provided).

# Examples
```julia-repl
julia> GateOps.c_pauli_y(0, 1, "qreg")
GateOps.GateCall2P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(:c_y, nothing), 0, 1, GateOps.GateCall1P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(:y, nothing), 0, "qreg"), "qreg")
```
"""
function c_pauli_y(q_target::Int, q_ctrl::Int, register::Union{String, Nothing}=nothing)
    return GateCall2( GateLabel(:c_y), q_target, q_ctrl, pauli_y(q_target,register), register)
end

"""
    c_pauli_z(q_target::Int, q_ctrl::Int, register::Union{String, Nothing}=nothing)

Generate a controlled Pauli-x (two qubit) GateCall (GateCall2P), controlled on the index `q_ctrl` applied to the target `q_target` (on given register, if provided).

# Examples
```julia-repl
julia> GateOps.c_pauli_z(0, 1, "qreg")
GateOps.GateCall2P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(:c_z, nothing), 0, 1, GateOps.GateCall1P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(:z, nothing), 0, "qreg"), "qreg")
```
"""
function c_pauli_z(q_target::Int, q_ctrl::Int, register::Union{String, Nothing}=nothing)
    return GateCall2( GateLabel(:c_z), q_target, q_ctrl, pauli_z(q_target,register), register)
end

function c_u(label::String, q_target::Int, q_ctrl::Int, params::Dict, register::Union{String, Nothing}=nothing)
    return GateCall2( GateLabel( Symbol(label), params), q_target, q_ctrl, register)
end

"""
    c_u(label::GateLabel, q_target::Int, q_ctrl::Int, gc::GateCall1, register::Union{String, Nothing}=nothing)

Generate a controlled unitary (two qubit) GateCall (GateCall2P), controlled on the index `q_ctrl` applied to the target `q_target` (on given register, if provided).

# Examples
```julia-repl
julia> GateOps.c_u(GateOps.GateLabel(:myCU), 0, 1, GateOps.pauli_x(0), "q")
GateOps.GateCall2P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(:myCU, nothing), 0, 1, GateOps.GateCall1P{Int64,Float64}(GateOps.GateLabelP{Int64,Float64}(:x, nothing), 0, nothing), "q")
```
"""
function c_u(label::GateLabel, q_target::Int, q_ctrl::Int, gc::GateCall1, register::Union{String, Nothing}=nothing)
    return GateCall2( label, q_target, q_ctrl, gc, register)
end

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

"""
    c_r_x(q_target::Int, q_ctrl::Int, theta::Real, register::Union{String, Nothing}=nothing)

Generate a controlled rotation about x (exp(iθσ_x/2)) GateCall (GateCall2P), controlled on the index `q_ctrl` applied to the target `q_target` (on given register, if provided).

# Examples
```julia-repl
julia> QuantExQASM.GateOps.c_r_x( 0, 1, pi/2, "q")
QuantExQASM.GateOps.GateCall2P{Int64,Float64}(QuantExQASM.GateOps.GateLabelP{Int64,Float64}(Symbol("c_r_x_angle=1.5707963267948966"), Dict{String,Union{Bool, Float64, Int64}}("angle" => 1.5707963267948966)), 0, 1, nothing, nothing)
```
"""
function c_r_x(q_target::Int, q_ctrl::Int, theta::Real, register::Union{String, Nothing}=nothing)
    return GateCall2( GateLabel( Symbol("c_r_x" * "_" * "angle=" * string(theta)), Dict("angle"=>theta) ), q_target, q_ctrl )
end

"""
    c_r_y(q_target::Int, q_ctrl::Int, theta::Real, register::Union{String, Nothing}=nothing)

Generate a controlled rotation about y (exp(iθσ_y/2)) GateCall (GateCall2P), controlled on the index `q_ctrl` applied to the target `q_target` (on given register, if provided).

# Examples
```julia-repl
julia> QuantExQASM.GateOps.c_r_y( 0, 1, pi/3, "q")
QuantExQASM.GateOps.GateCall2P{Int64,Float64}(QuantExQASM.GateOps.GateLabelP{Int64,Float64}(Symbol("c_r_y_angle=1.0471975511965976"), Dict{String,Union{Bool, Float64, Int64}}("angle" => 1.0471975511965976)), 0, 1, nothing, nothing)
```
"""
function c_r_y(q_target::Int, q_ctrl::Int, theta::Real, register::Union{String, Nothing}=nothing)
    return GateCall2( GateLabel( Symbol("c_r_y" * "_" * "angle=" * string(theta)), Dict("angle"=>theta) ), q_target, q_ctrl )
end

"""
    c_r_z(q_target::Int, q_ctrl::Int, theta::Real, register::Union{String, Nothing}=nothing)

Generate a controlled rotation about z (exp(iθσ_z/2)) GateCall (GateCall2P), controlled on the index `q_ctrl` applied to the target `q_target` (on given register, if provided).

# Examples
```julia-repl
julia> QuantExQASM.GateOps.c_r_z( 0, 1, pi/4, "q")
QuantExQASM.GateOps.GateCall2P{Int64,Float64}(QuantExQASM.GateOps.GateLabelP{Int64,Float64}(Symbol("c_r_z_angle=0.7853981633974483"), Dict{String,Union{Bool, Float64, Int64}}("angle" => 0.7853981633974483)), 0, 1, nothing, nothing)
```
"""
function c_r_z(q_target::Int, q_ctrl::Int, theta::Real, register::Union{String, Nothing}=nothing)
    return GateCall2( GateLabel( Symbol("c_r_z" * "_" * "angle=" * string(theta)), Dict("angle"=>theta) ), q_target, q_ctrl )
end

"""
    c_r_phase(q_target::Int, q_ctrl::Int, theta::Real, register::Union{String, Nothing}=nothing)

Generate a controlled phase rotation about (controlled [1 0; 0 exp(iθ)] GateCall (GateCall2P), controlled on the index `q_ctrl` applied to the target `q_target` (on given register, if provided).

# Examples
```julia-repl
julia> QuantExQASM.GateOps.c_r_phase( 0, 1, pi/7, "q")
QuantExQASM.GateOps.GateCall2P{Int64,Float64}(QuantExQASM.GateOps.GateLabelP{Int64,Float64}(Symbol("c_r_phase_angle=0.4487989505128276"), Dict{String,Union{Bool, Float64, Int64}}("angle" => 0.4487989505128276)), 0, 1, nothing, nothing)
```
"""
function c_r_phase(q_target::Int, q_ctrl::Int, theta::Real, register::Union{String, Nothing}=nothing)
    return GateCall2( GateLabel( Symbol("c_r_phase" * "_" * "angle=" * string(theta)), Dict("angle"=>theta) ), q_target, q_ctrl )
end

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
#                           Utility functions
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

function create_label_r(params::Dict)
    label = params["label"]
    sqrt_depth = params["depth"]
    adj = params["adj"]
    return "$(label)_$(sqrt_depth)_$(adj)"
end
function create_label_e(params::Dict)
    label = params["label"]
    θ = params["θ"]
    ϕ = params["ϕ"]
    λ = params["λ"]
    phase = params["phase"]
    return "$(label)_$(θ)_$(λ)_$(phase)"
end

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

@memoize function Base.:sqrt(g::Gate)
    return Gate(sqrt(g.mat))
end

@memoize function Base.:adjoint(gate::Gate)
    return Gate(adjoint(gate.mat))
end

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

"""Finds rotation angles (a,b,c,d) in the decomposition u=exp(id)*Rz(c).Ry(b).Rz(a).
Direct port to Julia from qiskit: https://qiskit.org/documentation/_modules/qiskit/extensions/quantum_initializer/squ.html
"""
@memoize function mat_to_euler(unitary::Matrix{<:Number})

    u00 = unitary[1, 1]
    u01 = unitary[1, 2]
    u10 = unitary[2, 1]
    u11 = unitary[2, 2]

    tol = 1e-8 # Tolerance defines proximity to 0; assume 0 if less than tol

    # Handle special case if the entry (0,0) of the unitary is equal to zero
    if abs(u00) < tol
        # Note that u10 can't be zero, since u is unitary (and u00 == 0)
        gamma = angle(-u01 / u10)
        delta = angle(u01 * exp(-1im * gamma / 2))
        return 0., pi, -gamma, delta
    end
    # Handle special case if the entry (0,1) of the unitary is equal to zero
    if abs(u01) < tol
        # Note that u11 can't be zero, since u is unitary (and u01 == 0)
        gamma = angle(u00 / u11)
        delta = angle(u00 * exp(-1im * gamma / 2))
        return 0., 0., -gamma, delta
    end
    beta = 2 * acos(abs(u00))
    if sin(beta / 2) - cos(beta / 2) > 0
        gamma = angle(-u00 / u10)
        alpha = angle(u00 / u01)
    else
        gamma = -angle(-u10 / u00)
        alpha = -angle(u01 / u00)
    end
    delta = angle(u00 * exp(-1im * (alpha + gamma) / 2))
    # The decomposition works with another convention for the rotation gates
    # (the one using negative angles).
    # Therefore, we have to take the inverse of the angles at the end.
    return -alpha, -beta, -gamma, delta
end

"""
theta, phi, lam, phase - 0.5 * (phi + lam)

- Z(phi) Y(theta) Z(lambda)
- e^{i gamma}{2})} U_3(theta,phi,lambda)
"""
@memoize function zyz_to_u3(a,b,c,d)

    # Convert exp(id)*Rz(c).Ry(b).Rz(a) angles to u3
    return (b, c, a, d -0.5*(a+c) )
end

"""
theta, phi, lam, phase - 0.5 * (phi + lam)

- Z(phi) Y(theta) Z(lambda)
- e^{i gamma}{2})} U_3(theta,phi,lambda)
"""
@memoize function u3_to_zyz(a,b,c,d)
    # Convert exp(id)*Rz(c).Ry(b).Rz(a) angles to u3
    return (b, c, a, d -0.5*(a+c) )
end

@memoize function euler_to_mat(a::Number,b::Number,c::Number,d::Number)
    """Creates unitary from angles defined by mat_to_euler function as u=exp(id)*Rz(c).Ry(b).Rz(a).
    """
    u00 = exp(-1im*(c+a)/2)*cos(b/2)
    u01 = -exp(1im*(a-c)/2)*sin(b/2)
    u10 = exp(1im*(c-a)/2)*sin(b/2)
    u11 = exp(1im*(a+c)/2)*cos(b/2)
    return [u00 u01; u10 u11].*exp(1im*d)
end

@memoize function u3_to_mat(θ::Number,ϕ::Number,λ::Number,ph::Number)
    """Creates unitary from angles defined by mat_to_euler function as u=exp(id)*Rz(c).Ry(b).Rz(a).
    """
    u00 = cos(θ/2)
    u01 = -exp(1im*λ)*sin(θ/2)
    u10 = exp(1im*ϕ)*sin(θ/2)
    u11 = exp(1im*(ϕ+λ)/2)*cos(θ/2)
    return [u00 u01; u10 u11].*exp(1im*ph)
end

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

function Base.:String(gc::AGateCall)
    return String(gc.gate_label);
end

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

"""
    create_gates_nonparam(gate_call::AGateCall)

    Create QASM syntax for gates cached in g_map_lg dictionary.
    Single and controlled versions generated.
    Maps X^{1/2^n} gates to H Z^{1/2^n} H

"""
function create_gates_nonparam(gate_call::AGateCall)
    if isdefined(gate_call, :ctrl)
        ct = "ctrl,tgt";
        u_pre = "c"
    else
        ct = "tgt";
        u_pre = ""
    end
    label = gate_call.gate_label
    gate = g_map_lg[label]

    # Problematic if z-gates not populated before x-gates; init to sufficient depth to prevent issues
    if String(label.gate_label)[end] == 'x'
        gate_label_sub = Symbol(String(label.gate_label)[1:end-1] * "z")
        gate = g_map[gate_label_sub]
    end

    euler_vals = GateOps.mat_to_euler(gate)
    u3_vals = GateOps.zyz_to_u3(euler_vals...)

    # QASM and X^{1/2^n} gates do not work well, so map to H Z^{1/2^n} H
    if String(label.gate_label)[end] == 'x'
        g = """gate $(u_pre)$(label.gate_label) tgt{\nh tgt;\n$(u_pre)u3($(u3_vals[1]),$(u3_vals[2]),$(u3_vals[3])) tgt;\nh tgt;\n}"""
    elseif String(label.gate_label)[end] == 'z'
        g = """gate $(u_pre)$(label.gate_label) tgt{\n$(u_pre)u3($(u3_vals[1]),$(u3_vals[2]),$(u3_vals[3])) tgt;\n}"""
    else
        error("Unsupported gate: Pauli-X and Pauli-Z currently only supported")
    end

    return g
end

"""
    create_gates_nonparam(gate_label::Symbol, num_qubits::Int

    Create QASM syntax for gates cached in g_map_lg dictionary.
    Single and controlled versions generated.
    Maps X^{1/2^n} gates to H Z^{1/2^n} H

"""
function create_gates_nonparam(gate_label::Symbol, num_qubits::Int)
    ct=""
    gate = g_map[gate_label]

    if String(gate_label)[end] == 'x'
        gate_label_sub = Symbol(String(gate_label)[1:end-1] * "z")
        gate = g_map[gate_label_sub]
    end

    euler_vals = GateOps.mat_to_euler(gate)
    u3_vals = GateOps.zyz_to_u3(euler_vals...)

    # QASM and X^{1/2^n} gates do not work well, so map to H Z^{1/2^n} H
    if num_qubits == 1 && String(gate_label)[end] == 'x'
        ct = "tgt"
        g = """gate $(gate_label) tgt{\nh tgt;\nu3($(u3_vals[1]),$(u3_vals[2]),$(u3_vals[3])) tgt;\nh tgt;\n}"""
    elseif num_qubits == 2 && String(gate_label)[end] == 'x'
        ct = "ctrl,tgt"
        g = """gate c$(gate_label) $(ct){\nh tgt;\ncu3($(u3_vals[1]),$(u3_vals[2]),$(u3_vals[3])) $(ct);\nh tgt;\n}"""
    elseif num_qubits == 1 && String(gate_label)[end] == 'z'
        ct = "tgt"
        g = """gate $(gate_label) tgt{\nu3($(u3_vals[1]),$(u3_vals[2]),$(u3_vals[3])) tgt;\n}"""
    else
        ct = "ctrl,tgt"
        g = """gate c$(gate_label) $(ct){\ncu3($(u3_vals[1]),$(u3_vals[2]),$(u3_vals[3])) $(ct);\n}"""
    end

    return g
end

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

# QASM interop functions
function mat_to_u3(unitary::Matrix{<:Number})
    # Direct port from Qiskit https://qiskit.org/documentation/_modules/qiskit/quantum_info/synthesis/one_qubit_decompose.html

    coeff = 1/sqrt(det(unitary))
    phase = -angle(coeff)
    su_mat = coeff * unitary  # U in SU(2)

    theta = 2 * atan(abs(su_mat[2, 1]), abs(su_mat[1, 1]))
    phiplambda = 2 * angle(su_mat[2, 2])
    phimlambda = 2 * angle(su_mat[2, 1])
    phi = (phiplambda + phimlambda) / 2.0
    lam = (phiplambda - phimlambda) / 2.0

    return theta, phi, lam, phase - 0.5 * (phi + lam)
end

# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #

end


text:
@namedargs struct Tuners
  y_imputed::Dict{Tuple{Int64, Int64, Int64}, MCMC.TuningParam{Float64}}
  Z::MCMC.TuningParam{Float64}
  v::Vector{MCMC.TuningParam{Float64}}
end

function Tuners(y, K; z_prob_flip=nothing)
  I = length(y)
  J = ncol(y[1])
  N = nrow.(y)

  y_tuner = Dict{Tuple{Int, Int, Int}, MCMC.TuningParam{Float64}}()
  for i in 1:I
    for n in 1:N[i]
      for j in 1:J
        if isnan(y[i][n, j])
          y_tuner[i, n, j] = MCMC.TuningParam(1.0)
        end
      end
    end
  end

  if z_prob_flip == nothing
    z_prob_flip = 1. / (J * K)
  end

  Z_tuner = MCMC.TuningParam(z_prob_flip)
  v_tuner = [MCMC.TuningParam(1.0) for k in 1:K]

  return Tuners(y_imputed=y_tuner, Z=Z_tuner, v=v_tuner)
end

max_stars_repo_name: luiarthur/CytofResearch

max_issues_repo_name: luiarthur/CytofResearch

max_forks_repo_name: luiarthur/CytofResearch


text:
export Thunk, delayed, delayedmap

const ID_COUNTER = Threads.Atomic{Int}(1)
next_id() = Threads.atomic_add!(ID_COUNTER, 1)

"""
    Thunk

Wraps a callable object to be run with Dagger. A `Thunk` is typically
created through a call to `delayed` or its macro equivalent `@par`.

## Constructors
```julia
delayed(f; kwargs...)(args...)
@par [option=value]... f(args...)
```

## Examples
```julia
julia> t = delayed(sin)(π)  # creates a Thunk to be computed later
Thunk(sin, (π,))

julia> collect(t)  # computes the result and returns it to the current process
1.2246467991473532e-16
```

## Arguments
- `f`: The function to be called upon execution of the `Thunk`.
- `args`: The arguments to be passed to the `Thunk`.
- `kwargs`: The properties describing unique behavior of this `Thunk`. Details
for each property are described in the next section.
- `option=value`: The same as passing `kwargs` to `delayed`.

## Public Properties
- `meta::Bool=false`: If `true`, instead of fetching cached arguments from
`Chunk`s and passing the raw arguments to `f`, instead pass the `Chunk`. Useful
for doing manual fetching or manipulation of `Chunk` references. Non-`Chunk`
arguments are still passed as-is.
- `processor::Processor=OSProc()` - The processor associated with `f`. Useful if
`f` is a callable struct that exists on a given processor and should be
transferred appropriately.
- `scope::Dagger.AbstractScope=AnyScope()` - The scope associated with `f`.
Useful if `f` is a function or callable struct that may only be transferred to,
and executed within, the specified scope.

## Options
- `options`: A `Sch.ThunkOptions` struct providing the options for the `Thunk`.
If omitted, options can also be specified by passing key-value pairs as
`kwargs`.
"""
mutable struct Thunk
    f::Any # usually a Function, but could be any callable
    inputs::Vector{Any} # TODO: Use `ImmutableArray` in 1.8
    id::Int
    get_result::Bool # whether the worker should send the result or only the metadata
    meta::Bool
    persist::Bool # don't `free!` result after computing
    cache::Bool   # release the result giving the worker an opportunity to
                  # cache it
    cache_ref::Any
    affinity::Union{Nothing, Pair{OSProc, Int}}
    eager_ref::Union{DRef,Nothing}
    options::Any # stores scheduler-specific options
    propagates::Tuple # which options we'll propagate
    function Thunk(f, xs...;
                   id::Int=next_id(),
                   get_result::Bool=false,
                   meta::Bool=false,
                   persist::Bool=false,
                   cache::Bool=false,
                   cache_ref=nothing,
                   affinity=nothing,
                   eager_ref=nothing,
                   processor=nothing,
                   scope=nothing,
                   options=nothing,
                   propagates=(),
                   kwargs...
                  )
        if !isa(f, Chunk) && (!isnothing(processor) || !isnothing(scope))
            f = tochunk(f,
                        something(processor, OSProc()),
                        something(scope, AnyScope()))
        end
        xs = Any[xs...]
        if options !== nothing
            @assert isempty(kwargs)
            new(f, xs, id, get_result, meta, persist, cache, cache_ref,
                affinity, eager_ref, options, propagates)
        else
            new(f, xs, id, get_result, meta, persist, cache, cache_ref,
                affinity, eager_ref, Sch.ThunkOptions(;kwargs...), propagates)
        end
    end
end

function affinity(t::Thunk)
    if t.affinity !== nothing
        return t.affinity
    end

    if t.cache && t.cache_ref !== nothing
        aff_vec = affinity(t.cache_ref)
    else
        aff = Dict{OSProc,Int}()
        for inp in inputs(t)
           #if haskey(state.cache, inp)
           #    as = affinity(state.cache[inp])
           #    for a in as
           #        proc, sz = a
           #        aff[proc] = get(aff, proc, 0) + sz
           #    end
           #else
                if isa(inp, Union{Chunk, Thunk})
                    # TODO if inp is a FileRef, affinity[1] will always be OSProc(1)
                    proc, sz = affinity(inp)
                    aff[proc] = get(aff, proc, 0) + sz
                end
           #end
        end
        aff_vec = collect(aff)
    end
    aff_vec
   #if length(aff) > 1
   #    return sort!(aff_vec, by=last,rev=true)
   #else
   #    return aff_vec
   #end
end

"""
    delayed(f; kwargs...)(args...)

Creates a [`Thunk`](@ref) object which can be executed later, which will call
`f` with `args`. `kwargs` controls various properties of the resulting `Thunk`.
"""
function delayed(f; kwargs...)
    (args...) -> Thunk(f, args...; kwargs...)
end

delayedmap(f, xs...) = map(delayed(f), xs...)

"A future holding the result of a `Thunk`."
struct ThunkFuture
    future::Future
end
ThunkFuture(x::Integer) = ThunkFuture(Future(x))
ThunkFuture() = ThunkFuture(Future())
Base.isready(t::ThunkFuture) = isready(t.future)
Base.wait(t::ThunkFuture) = Dagger.Sch.thunk_yield() do
    wait(t.future)
end
function Base.fetch(t::ThunkFuture; proc=OSProc(), raw=false)
    error, value = Dagger.Sch.thunk_yield() do
        if raw
            fetch(t.future)
        else
            move(proc, fetch(t.future))
        end
    end
    if error
        throw(value)
    end
    value
end
Base.put!(t::ThunkFuture, x; error=false) = put!(t.future, (error, x))

"A weak reference to a `Thunk`."
struct WeakThunk
    x::WeakRef
    WeakThunk(t::Thunk) = new(WeakRef(t))
end
istask(::WeakThunk) = true
unwrap_weak(t::WeakThunk) = t.x.value
unwrap_weak(t) = t
function unwrap_weak_checked(t::WeakThunk)
    t = unwrap_weak(t)
    @assert t !== nothing
    t
end
unwrap_weak_checked(t) = t
Base.show(io::IO, t::WeakThunk) = (print(io, "~"); Base.show(io, t.x.value))
Base.convert(::Type{WeakThunk}, t::Thunk) = WeakThunk(t)

struct ThunkFailedException{E<:Exception} <: Exception
    thunk::WeakThunk
    origin::WeakThunk
    ex::E
end
ThunkFailedException(thunk, origin, ex::E) where E =
    ThunkFailedException{E}(convert(WeakThunk, thunk), convert(WeakThunk, origin), ex)
function Base.showerror(io::IO, ex::ThunkFailedException)
    t = unwrap_weak(ex.thunk)
    o = unwrap_weak(ex.origin)
    t_str = t !== nothing ? "$t" : "?"
    o_str = o !== nothing ? "$o" : "?"
    t_id = t !== nothing ? t.id : '?'
    o_id = o !== nothing ? o.id : '?'
    println(io, "ThunkFailedException ($t failure",
                (o !== nothing && t != o) ? " due to a failure in $o)" : ")",
                ":")
    Base.showerror(io, ex.ex)
end

"""
    EagerThunk

Returned from `spawn`/`@spawn` calls. Represents a task that is in the
scheduler, potentially ready to execute, executing, or finished executing. May
be `fetch`'d or `wait`'d on at any time.
"""
mutable struct EagerThunk
    uid::UInt
    future::ThunkFuture
    finalizer_ref::DRef
    thunk_ref::DRef
end
Base.isready(t::EagerThunk) = isready(t.future)
Base.wait(t::EagerThunk) = wait(t.future)
function Base.fetch(t::EagerThunk; raw=false)
    if raw
        fetch(t.future; raw=true)
    else
        move(OSProc(), fetch(t.future))
    end
end
function Base.show(io::IO, t::EagerThunk)
    print(io, "EagerThunk ($(isready(t) ? "finished" : "running"))")
end

"When finalized, cleans-up the associated `EagerThunk`."
mutable struct EagerThunkFinalizer
    uid::UInt
    function EagerThunkFinalizer(uid)
        x = new(uid)
        finalizer(Sch.eager_cleanup, x)
        x
    end
end

const EAGER_ID_COUNTER = Threads.Atomic{UInt64}(1)
eager_next_id() = Threads.atomic_add!(EAGER_ID_COUNTER, one(UInt64))

function _spawn(f, args...; options, kwargs...)
    Dagger.Sch.init_eager()
    uid = eager_next_id()
    future = ThunkFuture()
    finalizer_ref = poolset(EagerThunkFinalizer(uid))
    added_future = Future()
    propagates = keys(options)
    put!(Dagger.Sch.EAGER_THUNK_CHAN, (added_future, future, uid, finalizer_ref, f, (args...,), (;propagates, options..., kwargs...,)))
    thunk_ref = fetch(added_future)
    return (uid, future, finalizer_ref, thunk_ref)
end
"""
    spawn(f, args...; kwargs...) -> EagerThunk

Spawns a task with `f` as the function and `args` as the arguments, returning
an `EagerThunk`. Uses a scheduler running in the background to execute code.

Note that `kwargs` are passed to the `Thunk` constructor, and are documented in
its docstring.
"""
function spawn(f, args...; processor=nothing, scope=nothing, kwargs...)
    options = get_options()
    if !isnothing(processor) || !isnothing(scope)
        f = tochunk(f,
                    something(processor, get_options(:processor, OSProc())),
                    something(scope, get_options(:scope, AnyScope())))
    end
    uid, future, finalizer_ref, thunk_ref = if myid() == 1
        _spawn(f, args...; options, kwargs...)
    else
        remotecall_fetch(_spawn, 1, f, args...; options, kwargs...)
    end
    return EagerThunk(uid, future, finalizer_ref, thunk_ref)
end

"""
    @par [opts] f(args...) -> Thunk

Convenience macro to call `Dagger.delayed` on `f` with arguments `args`.
May also be called with a series of assignments like so:

```julia
x = @par begin
    a = f(1,2)
    b = g(a,3)
    h(a,b)
end
```

`x` will hold the Thunk representing `h(a,b)`; additionally, `a` and `b`
will be defined in the same local scope and will be equally accessible
for later calls.

Options to the `Thunk` can be set as `opts` with namedtuple syntax, e.g.
`single=1`. Multiple options may be provided, and will be applied to all
generated thunks.
"""
macro par(exs...)
    opts = exs[1:end-1]
    ex = exs[end]
    _par(ex; lazy=true, opts=opts)
end

"""
    @spawn [opts] f(args...) -> Thunk

Convenience macro like `Dagger.@par`, but eagerly executed from the moment it's
called (equivalent to `spawn`).

See the docs for `@par` for more information and usage examples.
"""
macro spawn(exs...)
    opts = exs[1:end-1]
    ex = exs[end]
    _par(ex; lazy=false, opts=opts)
end

function _par(ex::Expr; lazy=true, recur=true, opts=())
    if ex.head == :call && recur
        f = ex.args[1]
        args = ex.args[2:end]
        opts = esc.(opts)
        if lazy
            return :(Dagger.delayed($(esc(f)); $(opts...))($(_par.(args; lazy=lazy, recur=false)...)))
        else
            return quote
                let args = ($(_par.(args; lazy=lazy, recur=false)...),)
                    $spawn($(esc(f)), args...; $(opts...))
                end
            end
        end
    else
        return Expr(ex.head, _par.(ex.args, lazy=lazy, recur=recur, opts=opts)...)
    end
end
_par(ex::Symbol; kwargs...) = esc(ex)
_par(ex; kwargs...) = ex

persist!(t::Thunk) = (t.persist=true; t)
cache_result!(t::Thunk) = (t.cache=true; t)

# @generated function compose{N}(f, g, t::NTuple{N})
#     if N <= 4
#       ( :(()->f(g())),
#         :((a)->f(g(a))),
#         :((a,b)->f(g(a,b))),
#         :((a,b,c)->f(g(a,b,c))),
#         :((a,b,c,d)->f(g(a,b,c,d))), )[N+1]
#     else
#         :((xs...) -> f(g(xs...)))
#     end
# end

# function Thunk(f::Function, t::Tuple{Thunk})
#     g = compose(f, t[1].f, t[1].inputs)
#     Thunk(g, t[1].inputs)
# end

# this gives a ~30x speedup in hashing
Base.hash(x::Thunk, h::UInt) = hash(x.id, hash(h, 0x7ad3bac49089a05f % UInt))
Base.isequal(x::Thunk, y::Thunk) = x.id==y.id

function Base.show(io::IO, z::Thunk)
    lvl = get(io, :lazy_level, 2)
    f = if z.f isa Chunk
        Tf = z.f.chunktype
        if isdefined(Tf, :instance)
            Tf.instance
        else
            "instance of $Tf"
        end
    else
        z.f
    end
    print(io, "Thunk[$(z.id)]($f, ")
    if lvl > 0
        show(IOContext(io, :lazy_level => lvl-1), z.inputs)
    else
        print(io, "...")
    end
    print(io, ")")
end

Base.summary(z::Thunk) = repr(z)

inputs(x::Thunk) = x.inputs
inputs(x) = ()

istask(x::Thunk) = true
istask(x) = false


text:
module ColorVectorSpaceTests

using LinearAlgebra, Statistics
using ColorVectorSpace, Colors, FixedPointNumbers, StatsBase

using Test

const var = Statistics.var

macro test_colortype_approx_eq(a, b)
    :(test_colortype_approx_eq($(esc(a)), $(esc(b)), $(string(a)), $(string(b))))
end

n8sum(x,y) = Float64(N0f8(x)) + Float64(N0f8(y))

function test_colortype_approx_eq(a::Colorant, b::Colorant, astr, bstr)
    @test typeof(a) == typeof(b)
    n = length(fieldnames(typeof(a)))
    for i = 1:n
        @test getfield(a, i) ≈ getfield(b,i)
    end
end

@testset "Colortypes" begin

    @testset "nan" begin
        function make_checked_nan(::Type{T}) where T
            x = nan(T)
            isa(x, T) && isnan(x)
        end
        for S in (Float32, Float64)
            @test make_checked_nan(S)
            @test make_checked_nan(Gray{S})
            @test make_checked_nan(AGray{S})
            @test make_checked_nan(GrayA{S})
            @test make_checked_nan(RGB{S})
            @test make_checked_nan(ARGB{S})
            @test make_checked_nan(ARGB{S})
        end
    end

    @testset "Arithmetic with Gray" begin
        cf = Gray{Float32}(0.1)
        @test +cf == cf
        @test -cf == Gray(-0.1f0)
        ccmp = Gray{Float32}(0.2)
        @test 2*cf == ccmp
        @test cf*2 == ccmp
        @test ccmp/2 == cf
        @test 2.0f0*cf == ccmp
        @test_colortype_approx_eq cf*cf Gray{Float32}(0.01)
        @test_colortype_approx_eq cf^2 Gray{Float32}(0.01)
        @test_colortype_approx_eq cf^3.0f0 Gray{Float32}(0.001)
        @test eltype(2.0*cf) == Float64
        @test abs2(ccmp) == 0.2f0^2
        @test norm(cf) == 0.1f0
        @test sum(abs2, ccmp) == 0.2f0^2
        cu = Gray{N0f8}(0.1)
        @test 2*cu == Gray(2*cu.val)
        @test 2.0f0*cu == Gray(2.0f0*cu.val)
        f = N0f8(0.5)
        @test (f*cu).val ≈ f*cu.val
        @test cf/2.0f0 == Gray{Float32}(0.05)
        @test cu/2 == Gray(cu.val/2)
        @test cu/0.5f0 == Gray(cu.val/0.5f0)
        @test cf+cf == ccmp
        @test isfinite(cf)
        @test isfinite(Gray(true))
        @test !isinf(cf)
        @test !isnan(cf)
        @test !isfinite(Gray(NaN))
        @test !isinf(Gray(NaN))
        @test isnan(Gray(NaN))
        @test !isfinite(Gray(Inf))
        @test Gray(Inf) |> isinf
        @test !isnan(Gray(Inf))
        @test abs(Gray(0.1)) ≈ 0.1
        @test eps(Gray{N0f8}) == Gray(eps(N0f8))  # #282
        @test atan(Gray(0.1), Gray(0.2)) == atan(0.1, 0.2)

        acu = Gray{N0f8}[cu]
        acf = Gray{Float32}[cf]
        @test +acu === acu
        @test @inferred(acu./trues(1)) == acu
        @test typeof(acu./trues(1)) == Vector{typeof(cu/true)}
        @test @inferred(ones(Int, 1)./acu) == [1/cu]
        @test typeof(ones(Int, 1)./acu) == Vector{typeof(1/cu)}
        @test @inferred(acu./acu) == [1]
        @test typeof(acu./acu) == Vector{typeof(cu/cu)}
        @test typeof(acu+acf) == Vector{Gray{Float32}}
        @test typeof(acu-acf) == Vector{Gray{Float32}}
        @test typeof(acu.+acf) == Vector{Gray{Float32}}
        @test typeof(acu.-acf) == Vector{Gray{Float32}}
        @test typeof(acu.+cf) == Vector{Gray{Float32}}
        @test typeof(acu.-cf) == Vector{Gray{Float32}}
        @test typeof(2*acf) == Vector{Gray{Float32}}
        @test typeof(2 .* acf) == Vector{Gray{Float32}}
        @test typeof(0x02*acu) == Vector{Gray{Float32}}
        @test typeof(acu/2) == Vector{Gray{typeof(N0f8(0.5)/2)}}
        @test typeof(acf.^2) == Vector{Gray{Float32}}
        @test (acu./Gray{N0f8}(0.5))[1] == gray(acu[1])/N0f8(0.5)
        @test (acf./Gray{Float32}(2))[1] ≈ 0.05f0
        @test (acu/2)[1] == Gray(gray(acu[1])/2)
        @test (acf/2)[1] ≈ Gray{Float32}(0.05f0)
        @test sum(abs2, [cf, ccmp]) ≈ 0.05f0

        @test gray(0.8) == 0.8

        a = Gray{N0f8}[0.8,0.7]
        @test sum(a) == Gray(n8sum(0.8,0.7))
        @test abs( var(a) - (a[1]-a[2])^2 / 2 ) <= 0.001
        @test isapprox(a, a)
        @test real(Gray{Float32}) <: Real
        @test zero(ColorTypes.Gray) == 0
        @test oneunit(ColorTypes.Gray) == 1
        a = Gray{N0f8}[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
        @test StatsBase.histrange(a,10) == 0.1f0:0.1f0:1f0

        @test typeof(float(Gray{N0f16}(0.5))) <: AbstractFloat
        @test quantile( Gray{N0f16}[0.0,0.5,1.0], 0.1) ≈ 0.10000152590218968
        @test middle(Gray(0.2)) === Gray(0.2)
        @test middle(Gray(0.2), Gray(0.4)) === Gray((0.2+0.4)/2)

        # issue #56
        @test Gray24(0.8)*N0f8(0.5) === Gray{N0f8}(0.4)
        @test Gray24(0.8)*0.5 === Gray(0.4)
        @test Gray24(0.8)/2   === Gray(0.5f0*N0f8(0.8))
        @test Gray24(0.8)/2.0 === Gray(0.4)
    end

    @testset "Comparisons with Gray" begin
        g1 = Gray{N0f8}(0.2)
        g2 = Gray{N0f8}(0.3)
        @test isless(g1, g2)
        @test !(isless(g2, g1))
        @test g1 < g2
        @test !(g2 < g1)
        @test isless(g1, 0.5)
        @test !(isless(0.5, g1))
        @test g1 < 0.5
        @test !(0.5 < g1)
        @test (@inferred(max(g1, g2)) ) == g2
        @test max(g1, 0.1) == 0.2
        @test (@inferred(min(g1, g2)) ) == g1
        @test min(g1, 0.1) == 0.1
        a = Gray{Float64}(0.9999999999999999)
        b = Gray{Float64}(1.0)

        @test isapprox(a, b)
        a = Gray{Float64}(0.99)
        @test !(isapprox(a, b, rtol = 0.01))
        @test isapprox(a, b, rtol = 0.1)
    end

    @testset "Unary operations with Gray" begin
        for g in (Gray(0.4), Gray{N0f8}(0.4))
            @test @inferred(zero(g)) === typeof(g)(0)
            @test @inferred(oneunit(g)) === typeof(g)(1)
            for op in ColorVectorSpace.unaryOps
                try
                    v = @eval $op(gray(g))  # if this fails, don't bother
                    @show op
                    @test op(g) == v
                catch
                end
            end
        end
        u = N0f8(0.4)
        @test ~Gray(u) == Gray(~u)
        @test -Gray(u) == Gray(-u)
    end

    @testset "Arithmetic with GrayA" begin
        p1 = GrayA{Float32}(Gray(0.8), 0.2)
        @test @inferred(zero(p1)) === GrayA{Float32}(0,0)
        @test @inferred(oneunit(p1)) === GrayA{Float32}(1,1)
        @test +p1 == p1
        @test -p1 == GrayA(-0.8f0, -0.2f0)
        p2 = GrayA{Float32}(Gray(0.6), 0.3)
        @test_colortype_approx_eq p1+p2 GrayA{Float32}(Gray(1.4),0.5)
        @test_colortype_approx_eq (p1+p2)/2 GrayA{Float32}(Gray(0.7),0.25)
        @test_colortype_approx_eq 0.4f0*p1+0.6f0*p2 GrayA{Float32}(Gray(0.68),0.26)
        @test_colortype_approx_eq ([p1]+[p2])[1] GrayA{Float32}(Gray(1.4),0.5)
        @test_colortype_approx_eq ([p1].+[p2])[1] GrayA{Float32}(Gray(1.4),0.5)
        @test_colortype_approx_eq ([p1].+p2)[1] GrayA{Float32}(Gray(1.4),0.5)
        @test_colortype_approx_eq ([p1]-[p2])[1] GrayA{Float32}(Gray(0.2),-0.1)
        @test_colortype_approx_eq ([p1].-[p2])[1] GrayA{Float32}(Gray(0.2),-0.1)
        @test_colortype_approx_eq ([p1].-p2)[1] GrayA{Float32}(Gray(0.2),-0.1)
        @test_colortype_approx_eq ([p1]/2)[1] GrayA{Float32}(Gray(0.4),0.1)
        @test_colortype_approx_eq (0.4f0*[p1]+0.6f0*[p2])[1] GrayA{Float32}(Gray(0.68),0.26)

        a = GrayA{N0f8}[GrayA(0.8,0.7), GrayA(0.5,0.2)]
        @test sum(a) == GrayA(n8sum(0.8,0.5), n8sum(0.7,0.2))
        @test isapprox(a, a)
        a = AGray{Float64}(1.0, 0.9999999999999999)
        b = AGray{Float64}(1.0, 1.0)

        @test a ≈ b
        a = AGray{Float64}(1.0, 0.99)
        @test !isapprox(a, b, rtol = 0.01)
        @test isapprox(a, b, rtol = 0.1)

        # issue #56
        @test AGray32(0.8,0.2)*N0f8(0.5) === AGray{N0f8}(0.4,0.1)
        @test AGray32(0.8,0.2)*0.5 === AGray(0.4,0.1)
        @test AGray32(0.8,0.2)/2   === AGray(0.5f0*N0f8(0.8),0.5f0*N0f8(0.2))
        @test AGray32(0.8,0.2)/2.0 === AGray(0.4,0.1)
    end

    @testset "Arithemtic with RGB" begin
        cf = RGB{Float32}(0.1,0.2,0.3)
        @test @inferred(zero(cf)) === RGB{Float32}(0,0,0)
        @test @inferred(oneunit(cf)) === RGB{Float32}(1,1,1)
        @test +cf == cf
        @test -cf == RGB(-0.1f0, -0.2f0, -0.3f0)
        ccmp = RGB{Float32}(0.2,0.4,0.6)
        @test 2*cf == ccmp
        @test cf*2 == ccmp
        @test ccmp/2 == cf
        @test 2.0f0*cf == ccmp
        @test eltype(2.0*cf) == Float64
        cu = RGB{N0f8}(0.1,0.2,0.3)
        @test_colortype_approx_eq 2*cu RGB(2*cu.r, 2*cu.g, 2*cu.b)
        @test_colortype_approx_eq 2.0f0*cu RGB(2.0f0*cu.r, 2.0f0*cu.g, 2.0f0*cu.b)
        f = N0f8(0.5)
        @test (f*cu).r ≈ f*cu.r
        @test cf/2.0f0 == RGB{Float32}(0.05,0.1,0.15)
        @test cu/2 ≈ RGB(cu.r/2,cu.g/2,cu.b/2)
        @test cu/0.5f0 == RGB(cu.r/0.5f0, cu.g/0.5f0, cu.b/0.5f0)
        @test cf+cf == ccmp
        @test cu * 1//2 == mapc(x->Float64(Rational(x)/2), cu)
        @test_colortype_approx_eq (cf.*[0.8f0])[1] RGB{Float32}(0.8*0.1,0.8*0.2,0.8*0.3)
        @test_colortype_approx_eq ([0.8f0].*cf)[1] RGB{Float32}(0.8*0.1,0.8*0.2,0.8*0.3)
        @test isfinite(cf)
        @test !isinf(cf)
        @test !isnan(cf)
        @test !isfinite(RGB(NaN, 1, 0.5))
        @test !isinf(RGB(NaN, 1, 0.5))
        @test isnan(RGB(NaN, 1, 0.5))
        @test !isfinite(RGB(1, Inf, 0.5))
        @test isinf(RGB(1, Inf, 0.5))
        @test !isnan(RGB(1, Inf, 0.5))
        @test abs(RGB(0.1,0.2,0.3)) ≈ 0.6
        @test sum(abs2, RGB(0.1,0.2,0.3)) ≈ 0.14
        @test norm(RGB(0.1,0.2,0.3)) ≈ sqrt(0.14)

        acu = RGB{N0f8}[cu]
        acf = RGB{Float32}[cf]
        @test typeof(acu+acf) == Vector{RGB{Float32}}
        @test typeof(acu-acf) == Vector{RGB{Float32}}
        @test typeof(acu.+acf) == Vector{RGB{Float32}}
        @test typeof(acu.-acf) == Vector{RGB{Float32}}
        @test typeof(acu.+cf) == Vector{RGB{Float32}}
        @test typeof(acu.-cf) == Vector{RGB{Float32}}
        @test typeof(2*acf) == Vector{RGB{Float32}}
        @test typeof(convert(UInt8, 2)*acu) == Vector{RGB{Float32}}
        @test typeof(acu/2) == Vector{RGB{typeof(N0f8(0.5)/2)}}
        rcu = rand(RGB{N0f8}, 3, 5)
        @test @inferred(rcu./trues(3, 5)) == rcu
        @test typeof(rcu./trues(3, 5)) == Matrix{typeof(cu/true)}

        a = RGB{N0f8}[RGB(1,0,0), RGB(1,0.8,0)]
        @test sum(a) == RGB(2.0,0.8,0)
        @test isapprox(a, a)
        a = RGB{Float64}(1.0, 1.0, 0.9999999999999999)
        b = RGB{Float64}(1.0, 1.0, 1.0)

        @test isapprox(a, b)
        a = RGB{Float64}(1.0, 1.0, 0.99)
        @test !(isapprox(a, b, rtol = 0.01))
        @test isapprox(a, b, rtol = 0.1)
        # issue #56
        @test RGB24(1,0,0)*N0f8(0.5) === RGB{N0f8}(0.5,0,0)
        @test RGB24(1,0,0)*0.5 === RGB(0.5,0,0)
        @test RGB24(1,0,0)/2   === RGB(0.5f0,0,0)
        @test RGB24(1,0,0)/2.0 === RGB(0.5,0,0)
    end

    @testset "Arithemtic with RGBA" begin
        cf = RGBA{Float32}(0.1,0.2,0.3,0.4)
        @test @inferred(zero(cf)) === RGBA{Float32}(0,0,0,0)
        @test @inferred(oneunit(cf)) === RGBA{Float32}(1,1,1,1)
        @test +cf == cf
        @test -cf == RGBA(-0.1f0, -0.2f0, -0.3f0, -0.4f0)
        ccmp = RGBA{Float32}(0.2,0.4,0.6,0.8)
        @test 2*cf == ccmp
        @test cf*2 == ccmp
        @test ccmp/2 == cf
        @test 2.0f0*cf == ccmp
        @test eltype(2.0*cf) == Float64
        cu = RGBA{N0f8}(0.1,0.2,0.3,0.4)
        @test_colortype_approx_eq 2*cu RGBA(2*cu.r, 2*cu.g, 2*cu.b, 2*cu.alpha)
        @test_colortype_approx_eq 2.0f0*cu RGBA(2.0f0*cu.r, 2.0f0*cu.g, 2.0f0*cu.b, 2.0f0*cu.alpha)
        f = N0f8(0.5)
        @test (f*cu).r ≈ f*cu.r
        @test cf/2.0f0 == RGBA{Float32}(0.05,0.1,0.15,0.2)
        @test cu/2 == RGBA(cu.r/2,cu.g/2,cu.b/2,cu.alpha/2)
        @test cu/0.5f0 == RGBA(cu.r/0.5f0, cu.g/0.5f0, cu.b/0.5f0, cu.alpha/0.5f0)
        @test cf+cf == ccmp
        @test_colortype_approx_eq (cf.*[0.8f0])[1] RGBA{Float32}(0.8*0.1,0.8*0.2,0.8*0.3,0.8*0.4)
        @test_colortype_approx_eq ([0.8f0].*cf)[1] RGBA{Float32}(0.8*0.1,0.8*0.2,0.8*0.3,0.8*0.4)
        @test isfinite(cf)
        @test !isinf(cf)
        @test !isnan(cf)
        @test isnan(RGBA(NaN, 1, 0.5, 0.8))
        @test !isinf(RGBA(NaN, 1, 0.5))
        @test isnan(RGBA(NaN, 1, 0.5))
        @test !isfinite(RGBA(1, Inf, 0.5))
        @test RGBA(1, Inf, 0.5) |> isinf
        @test !isnan(RGBA(1, Inf, 0.5))
        @test !isfinite(RGBA(0.2, 1, 0.5, NaN))
        @test !isinf(RGBA(0.2, 1, 0.5, NaN))
        @test isnan(RGBA(0.2, 1, 0.5, NaN))
        @test !isfinite(RGBA(0.2, 1, 0.5, Inf))
        @test RGBA(0.2, 1, 0.5, Inf) |> isinf
        @test !isnan(RGBA(0.2, 1, 0.5, Inf))
        @test abs(RGBA(0.1,0.2,0.3,0.2)) ≈ 0.8

        acu = RGBA{N0f8}[cu]
        acf = RGBA{Float32}[cf]
        @test typeof(acu+acf) == Vector{RGBA{Float32}}
        @test typeof(acu-acf) == Vector{RGBA{Float32}}
        @test typeof(acu.+acf) == Vector{RGBA{Float32}}
        @test typeof(acu.-acf) == Vector{RGBA{Float32}}
        @test typeof(acu.+cf) == Vector{RGBA{Float32}}
        @test typeof(acu.-cf) == Vector{RGBA{Float32}}
        @test typeof(2*acf) == Vector{RGBA{Float32}}
        @test typeof(convert(UInt8, 2)*acu) == Vector{RGBA{Float32}}
        @test typeof(acu/2) == Vector{RGBA{typeof(N0f8(0.5)/2)}}

        a = RGBA{N0f8}[RGBA(1,0,0,0.8), RGBA(0.7,0.8,0,0.9)]
        @test sum(a) == RGBA(n8sum(1,0.7),0.8,0,n8sum(0.8,0.9))
        @test isapprox(a, a)
        a = ARGB{Float64}(1.0, 1.0, 1.0, 0.9999999999999999)
        b = ARGB{Float64}(1.0, 1.0, 1.0, 1.0)

        @test isapprox(a, b)
        a = ARGB{Float64}(1.0, 1.0, 1.0, 0.99)
        @test !(isapprox(a, b, rtol = 0.01))
        @test isapprox(a, b, rtol = 0.1)
        # issue #56
        @test ARGB32(1,0,0,0.8)*N0f8(0.5) === ARGB{N0f8}(0.5,0,0,0.4)
        @test ARGB32(1,0,0,0.8)*0.5 === ARGB(0.5,0,0,0.4)
        @test ARGB32(1,0,0,0.8)/2   === ARGB(0.5f0,0,0,0.5f0*N0f8(0.8))
        @test ARGB32(1,0,0,0.8)/2.0 === ARGB(0.5,0,0,0.4)
    end

    @testset "Mixed-type arithmetic" begin
        @test RGB(1,0,0) + Gray(0.2f0) == RGB{Float32}(1.2,0.2,0.2)
        @test RGB(1,0,0) - Gray(0.2f0) == RGB{Float32}(0.8,-0.2,-0.2)
    end

    @testset "dotc" begin
        @test dotc(0.2, 0.2) == 0.2^2
        @test dotc(0.2, 0.3f0) == 0.2*0.3f0
        @test dotc(N0f8(0.2), N0f8(0.3)) == Float32(N0f8(0.2))*Float32(N0f8(0.3))
        @test dotc(Gray{N0f8}(0.2), Gray24(0.3)) == Float32(N0f8(0.2))*Float32(N0f8(0.3))
        xc, yc = RGB(0.2,0.2,0.2), RGB{N0f8}(0.3,0.3,0.3)
        @test isapprox(dotc(xc, yc) , dotc(convert(Gray, xc), convert(Gray, yc)), atol=1e-6)
        @test dotc(RGB(1,0,0), RGB(0,1,1)) == 0
    end

    @testset "typemin/max" begin
        for T in (Normed{UInt8,8}, Normed{UInt8,6}, Normed{UInt16,16}, Normed{UInt16,14}, Float32, Float64)
            @test typemin(Gray{T}) === Gray{T}(typemin(T))
            @test typemax(Gray{T}) === Gray{T}(typemax(T))
            @test typemin(Gray{T}(0.5)) === Gray{T}(typemin(T))
            @test typemax(Gray{T}(0.5)) === Gray{T}(typemax(T))
            A = maximum(Gray{T}.([1 0 0; 0 1 0]); dims=1)  # see PR#44 discussion
            @test isa(A, Matrix{Gray{T}})
            @test size(A) == (1,3)
        end
    end

    @testset "Colors issue #326" begin
        A = rand(RGB{N0f8}, 2, 2)
        @test @inferred mean(A) == mean(map(c->mapc(FixedPointNumbers.Treduce, c), A))
    end
end

end


text:
"""
Module for working with [GPX file format](https://en.wikipedia.org/wiki/GPS_Exchange_Format).

Be aware that this module don't fully follow 1.1 version of this format.
"""
module GPX

using LightXML
using Dates
using TimeZones

import Base: push!, iterate, length, getindex

export GPXAuthor, GPXMetadata, GPXPoint,
    GPXTrackSegment, GPXDocument, GPXRoute
export read_gpx_file, parse_gpx_string
export new_track, new_track_segment, last_segment

# not exported but very useful: GPXTracks

"""
Constant defining GPX Format version
"""
const GPX_VERSION = "1.1"

"""
Constant dictionary defining default namespaces
"""
const GPX_NS = Dict(
    "" => "http://www.topografix.com/GPX/1/1"
)
# http://www.topografix.com/GPX/1/1/gpx.xsd

"""
Constant defining default GPX file creator name
"""
const GPX_CREATOR = "GPX.jl"

"""
    GPXAuthor(name; email="", link="")

The person or organization who created the GPX file.
"""
struct GPXAuthor
    name::String
    email::String
	link::String
end
function GPXAuthor(name; email="", link="")
    GPXAuthor(name, email, link)
end

"""
    GPXMetadata(; name="", author=GPXAuthor(""), time=now(tz"UTC"))

Information about the GPX file, author, and copyright restrictions goes in the metadata section.
Providing rich, meaningful information about your GPX files allows others to search for and use your GPS data.

# Parameters
- `name`: The name of the GPX file.
- `desc`: A description of the contents of the GPX file.
- `author`: The person or organization who created the GPX file.
- `copyright`: Copyright and license information governing use of the file.
- `link`: URLs associated with the location described in the file.
- `time`: The creation date of the file.
- `keywords`: Keywords associated with the file. Search engines or databases can use this information to classify the data.
- `bounds`: Minimum and maximum coordinates which describe the extent of the coordinates in the file.
- `extensions`:
"""
struct GPXMetadata
    name::String
    desc::String
    author::GPXAuthor
    copyright::String
    link::String
    time::ZonedDateTime
    keywords::String
    bounds::String
end
GPXMetadata(; name="", desc="", author=GPXAuthor(""), copyright="", link="", time=now(tz"UTC"), keywords="", bounds="") = GPXMetadata(name, desc, author, copyright, link, time, keywords, bounds)

"""
    GPXPoint(lat::Float64, lon::Float64, ele::Float64, time::ZonedDateTime, desc::String)

A geographic point with optional elevation and time. Available for use by other schemas.

A Track Point holds the coordinates, elevation, timestamp, and metadata for a single point in a track.

# Parameters
- `lat`: The latitude of the point. This is always in decimal degrees, and always in WGS84 datum.
- `lon`: The longitude of the point. This is always in decimal degrees, and always in WGS84 datum.
- `ele`: The elevation (in meters) of the point.
- `time`: The time that the point was recorded.
- `desc`: A text description of the element.
"""
struct GPXPoint
    lat::Float64
    lon::Float64
    ele::Float64
    time::ZonedDateTime
    desc::String
end

"""
    GPXTrackSegment()

A Track Segment holds a list of Track Points which are logically connected in order. To represent a single GPS track where GPS reception was lost, or the GPS receiver was turned off, start a new Track Segment for each continuous span of track data
"""
struct GPXTrackSegment
    points::Vector{GPXPoint}
    GPXTrackSegment() = new(GPXPoint[])
end

"""
    GPXTrack()

A Track holds a list of Track Segments.
"""
struct GPXTrack
    segments::Vector{GPXTrackSegment}
    GPXTrack() = new(GPXTrackSegment[])
end

"""
    GPXTracks()

A list of tracks.
"""
struct GPXTracks
    collection::Vector{GPXTrack}
    GPXTracks() = new(Vector{GPXTrack}[])
end
push!(tracks::GPXTracks, track::GPXTrack) = push!(tracks.collection, track)
iterate(tracks::GPXTracks) = iterate(tracks.collection)
iterate(tracks::GPXTracks, i) = iterate(tracks.collection, i)
length(tracks::GPXTracks) = length(tracks.collection)
getindex(tracks::GPXTracks, i) = getindex(tracks.collection, i)

"""

ToDo
"""
struct GPXRoute
end

"""
    GPXDocument(metadata::GPXMetadata; tracks=GPXTracks(), namespaces=GPX_NS, version=GPX_VERSION, creator=GPX_CREATOR)

GPX documents contain a metadata header, followed by waypoints, routes, and tracks.
You can add your own elements to the extensions section of the GPX document.
"""
struct GPXDocument
    namespaces::Dict{String, String}

    version::String
    creator::String
    metadata::GPXMetadata

	waypoints::Vector{GPXPoint}
	routes::Vector{GPXRoute}
    tracks::GPXTracks

    function GPXDocument(metadata::GPXMetadata; tracks=GPXTracks(), namespaces=GPX_NS, version=GPX_VERSION, creator=GPX_CREATOR)
        waypoints = GPXPoint[]
        routes = GPXRoute[]
        new(namespaces, version, creator, metadata, waypoints, routes, tracks)
    end
end

"""
    new_track(gpx::GPXDocument) -> GPXTrack

Get a new track from a `GPXDocument`.
"""
function new_track(gpx::GPXDocument)
    return new_track(gpx.tracks)
end

"""
    new_track(tracks::GPXTracks) -> GPXTrack

Get a new track from a collection of tracks (ie `GPXTracks`).
"""
function new_track(tracks::GPXTracks)
    track = GPXTrack()
    push!(tracks, track)
    return track
end

"""
    new_track_segment(track::GPXTrack) -> GPXTrackSegment

Get a new track segment from a `GPXTrack`.
"""
function new_track_segment(track::GPXTrack)
    track_segment = GPXTrackSegment()
    push!(track.segments, track_segment)
    return track_segment
end

"""
    push!(track_segment::GPXTrackSegment, point::GPXPoint)

Insert one `GPXPoint` at the end of a `GPXTrackSegment`.
"""
function push!(track_segment::GPXTrackSegment, point::GPXPoint)
    push!(track_segment.points, point)
end

"""
    last_segment(gpx) -> GPXTrackSegment

"""
function last_segment(gpx::GPXDocument)
    track = isempty(gpx.tracks) ? new_track(gpx) : gpx.tracks[end]
    segment = isempty(track.segments) ? new_track_segment(track) : track.segments[end]
    return segment
end

"""
    XMLDocument(gpx::GPXDocument)

Create a `LightXML.XMLDocument` from a `GPXDocument`.

`XMLDocument` can then be saved to file or printed.
"""
function LightXML.XMLDocument(gpx::GPXDocument)
    xdoc = LightXML.XMLDocument()

    xroot = create_root(xdoc, "gpx")
    for (ns_key, ns_url) in gpx.namespaces
        key = ns_key == "" ? "xmlns" : "xmlns:$ns_key"
        set_attribute(xroot, key, ns_url)
    end
    set_attribute(xroot, "version", gpx.version)
    set_attribute(xroot, "creator", gpx.creator)

    x_metadata = new_child(xroot, "metadata")
    x_metadata_name = new_child(x_metadata, "name")
    add_text(x_metadata_name, gpx.metadata.name)
    x_metadata_author = new_child(x_metadata, "author")
    x_metadata_author_name = new_child(x_metadata_author, "name")
    add_text(x_metadata_author_name, gpx.metadata.author.name)

    for track in gpx.tracks
        x_trk = new_child(xroot, "trk")
        for segment in track.segments
            x_trkseg = new_child(x_trk, "trkseg")
            for point in segment.points
                x_trkpt = new_child(x_trkseg, "trkpt")
                lat, lon = point.lat, point.lon
                set_attribute(x_trkpt, "lat", "$lat")
                set_attribute(x_trkpt, "lon", "$lon")
                x_ele = new_child(x_trkpt, "ele")
                ele = point.ele
                add_text(x_ele, "$ele")
                x_time = new_child(x_trkpt, "time")
                dt = point.time
                s_dt_utc = "$(dt.utc_datetime)Z"
                add_text(x_time, s_dt_utc)
                desc = point.desc
                if desc != ""
                    x_desc = new_child(x_trkpt, "desc")
                    add_text(x_desc, desc)
                end
            end
        end
    end

    return xdoc
end

"""
    read_gpx_file(fname) -> GPXDocument

Read GPX file from filename `fname`.
"""
function read_gpx_file(fname)
    xdoc = parse_file(fname)
    return _parse_gpx(xdoc)
end

"""
    parse_gpx_string(s) -> GPXDocument

Parse GPX data from String `s`.
"""
function parse_gpx_string(s)
    xdoc = parse_string(s)
    return _parse_gpx(xdoc)
end

"""
    _parse_gpx(xdoc::XMLDocument) -> GPXDocument

Parse `XMLDocument` and return a `GPXDocument`.
"""
function _parse_gpx(xdoc::XMLDocument)
    gpxs = root(xdoc)

    metadata = GPXMetadata(
        name="title",
        author=GPXAuthor("author"),
        time=ZonedDateTime("2019-01-01T00:00:00.000+00:00"),
    )

    tracks = GPXTracks()

    for x_gpx in child_elements(gpxs)
        if name(x_gpx) == "trk"
            track = new_track(tracks)
            for x_track in child_elements(x_gpx)
                if name(x_track) == "trkseg"
                    track_segment = new_track_segment(track)
                    for x_track_segment in child_elements(x_track)
                        if name(x_track_segment) == "trkpt"
                            d_att = attributes_dict(x_track_segment)
                            dt = ZonedDateTime(0, tz"UTC")
                            lat = parse(Float64, d_att["lat"])
                            lon = parse(Float64, d_att["lon"])
                            ele = 0.0
                            desc = ""
                            for x_track_point in child_elements(x_track_segment)
                                if name(x_track_point) == "time"
                                    s = content(x_track_point)
                                    # Required for optional millisecond field
                                    # see https://github.com/JuliaTime/TimeZones.jl/issues/83
                                    # Check for period ('.') as indication of
                                    # datetime format
                                    dt = if occursin('.', s)
                                        parse(ZonedDateTime, s, dateformat"yyyy-mm-ddTHH:MM:SS.ssszzz")
                                    else
                                        parse(ZonedDateTime, s, dateformat"yyyy-mm-ddTHH:MM:SSzzz")
                                    end
                                elseif name(x_track_point) == "ele"
                                    s = content(x_track_point)
                                    ele = parse(Float64, s)
                                elseif name(x_track_point) == "desc"
                                    desc = content(x_track_point)
                                    # desc = ""  # for debug
                                end
                            end
                            point = GPXPoint(lat, lon, ele, dt, desc)
                            push!(track_segment, point)
                        end
                    end
                end
            end
        end
    end

    free(xdoc)

    gpx = GPXDocument(metadata, tracks=tracks)

    return gpx
end

end # module


text:
module SlurmHyperopt

abstract type AbstractHyperparameterSampler end

include("slurm_file.jl")
include("hyperopt.jl")

export SlurmParams, SlurmHyperoptimizer, RandomSampler, HyperoptResults
export merge_results!, get_params, get_results, save_result

end


text:
@reexport module _GMP

import Base.GMP: ClongMax, CulongMax, CdoubleMax
import Brainstorm._Math: add!, sub!, mul!

export
    set!,
    fld!, div!, mod!, rem!,
    gcd!, lcm!, and!, or!, xor!,
    neg!, com!,
    lsh!, rsh!,
    pow!, isqrt!

function set!(x::BigInt, y::BigInt)
    ccall((:__gmpz_set, :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}), &x, &y)
    return x
end
function set!(x::BigInt, y::ClongMax)
    ccall((:__gmpz_set_si, :libgmp), Void, (Ptr{BigInt}, Clong), &x, y)
    return x
end
function set!(x::BigInt, y::CulongMax)
    ccall((:__gmpz_set_ui, :libgmp), Void, (Ptr{BigInt}, Culong), &x, y)
    return x
end
set!(x::BigInt, y::Integer) = set!(x, convert(BigInt, y))

# Binary ops
for (fJ, fC) in ((:add!, :add), (:sub!,:sub), (:mul!, :mul),
                 (:fld!, :fdiv_q), (:div!, :tdiv_q), (:mod!, :fdiv_r), (:rem!, :tdiv_r),
                 (:gcd!, :gcd), (:lcm!, :lcm),
                 (:and!, :and), (:or!, :ior), (:xor!, :xor))
    @eval begin
        function ($fJ)(x::BigInt, y::BigInt)
            ccall(($(string(:__gmpz_,fC)), :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}, Ptr{BigInt}), &x, &x, &y)
            return x
        end
    end
end

# unary ops
for (fJ, fC) in ((:neg!, :neg), (:com!, :com))
    @eval begin
        function ($fJ)(x::BigInt)
            ccall(($(string(:__gmpz_,fC)), :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}), &x, &x)
            return x
        end
    end
end

function lsh!(x::BigInt, c::CulongMax)
    c < 0 && throw(DomainError())
    c == 0 && return x
    ccall((:__gmpz_mul_2exp, :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}, Culong), &x, &x, c)
    return x
end
lsh!(x::BigInt, c::Integer) = lsh!(x, convert(Culong, c))

function rsh!(x::BigInt, c::CulongMax)
    c < 0 && throw(DomainError())
    c == 0 && return x
    ccall((:__gmpz_fdiv_q_2exp, :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}, Culong), &x, &x, c)
    return x
end
rsh!(x::BigInt, c::Integer) = rsh!(x, convert(Culong, c))

function isqrt!(x::BigInt)
    ccall((:__gmpz_sqrt, :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}), &x, &x)
    return x
end

function pow!(x::BigInt, y::CulongMax)
    ccall((:__gmpz_pow_ui, :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}, Culong), &x, &x, y)
    return x
end
pow!(x::BigInt, c::Integer) = pow!(x, convert(Culong, c))

# Basic arithmetic without promotion
function add!(x::BigInt, c::CulongMax)
    ccall((:__gmpz_add_ui, :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}, Culong), &x, &x, c)
    return x
end
add!(c::CulongMax, x::BigInt) = add!(x, c)

function sub!(x::BigInt, c::CulongMax)
    ccall((:__gmpz_sub_ui, :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}, Culong), &x, &x, c)
    return x
end
function sub!(c::CulongMax, x::BigInt)
    ccall((:__gmpz_ui_sub, :libgmp), Void, (Ptr{BigInt}, Culong, Ptr{BigInt}), &x, c, &x)
    return x
end

add!(x::BigInt, c::ClongMax) = c < 0 ? sub!(x, -(c % Culong)) : add!(x, convert(Culong, c))
add!(c::ClongMax, x::BigInt) = c < 0 ? sub!(x, -(c % Culong)) : add!(x, convert(Culong, c))
sub!(x::BigInt, c::ClongMax) = c < 0 ? add!(x, -(c % Culong)) : sub!(x, convert(Culong, c))
sub!(c::ClongMax, x::BigInt) = c < 0 ? neg!(add!(x, -(c % Culong))) : sub!(convert(Culong, c), x)

add!(x::BigInt, c::Integer) = add!(x, convert(BigInt, c))
add!(c::Integer, x::BigInt) = add!(convert(BigInt, c), x)
sub!(x::BigInt, c::Integer) = sub!(x, convert(BigInt, c))
sub!(c::Integer, x::BigInt) = sub!(convert(BigInt, c), x)

function mul!(x::BigInt, c::CulongMax)
    ccall((:__gmpz_mul_ui, :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}, Culong), &x, &x, c)
    return x
end
mul!(c::CulongMax, x::BigInt) = mul!(x, c)

function mul!(x::BigInt, c::ClongMax)
    ccall((:__gmpz_mul_si, :libgmp), Void, (Ptr{BigInt}, Ptr{BigInt}, Clong), &x, &x, c)
    return x
end
mul!(c::ClongMax, x::BigInt) = mul!(x, c)

mul!(x::BigInt, c::Integer) = mul!(x, convert(BigInt, c))
mul!(c::Integer, x::BigInt) = mul!(convert(BigInt, c), x)

end


text:
"""
Main module for `MinimaxEstimation.jl`

"""
module MinimaxEstimation
    using LinearAlgebra
    using JuMP

    export KalmanFilter, MinimaxMMAE, BayesianMMAE
    export update!, predict
    export Tsyn

    include("types.jl")
    include("constructors.jl")
    include("methods.jl")
end


text:
module BLAS

using Dates
using DataFrames
using Distributions

# default problem is set up for tx and rx being DataFrames
# with only a time::DateTime column needed in each

Base.@kwdef struct Problem
  tx::DataFrame
  rx::DataFrame
  σₜ::Float64 = 10.0
  time::Function = x -> datetime2unix(DateTime(x.time))
  delay::Function = (tx, rx) -> 0.0
  passoc::Function = (tx, rx) -> 1.0
  pfalse::Function = rx -> 0.1
end

struct State
  score::Float64
  backlink::Union{State,Missing}
  assoc::Union{Pair{Int},Missing}
  i::Int
  j::Int
  μₜ::Float64
  σₜ::Float64
end

Base.show(io::IO, s::State) = print(io, "<i=$(s.i), j=$(s.j), μₜ=$(s.μₜ), $(s.assoc), $(s.score)>")

function Base.show(io::IO, ss::AbstractArray{State})
  print(io, "[\n")
  for s ∈ ss
    print(io, "  <i=$(s.i), j=$(s.j), μₜ=$(s.μₜ), $(s.assoc), $(s.score)>\n")
  end
  print(io, "]\n")
end

score(state) = state.score

function isduplicate(state, 𝒮)
  for s ∈ 𝒮
    s === state && continue
    (s.i != state.i || s.j != state.j || s.assoc !== state.assoc) && continue
    state.score > s.score && continue
    state.score < s.score && return true
    hash(state) < hash(s) && return true
  end
  false
end

function solve(P::Problem; nhypothesis=30)
  𝒮 = [State(0.0, missing, missing, 0, 0, 0.0, P.σₜ)]
  for (j, rx) ∈ enumerate(eachrow(P.rx))
    𝒮⁺ = State[]
    for state ∈ 𝒮
      # TODO: track σₜ better, adjusting it based on the time difference for each entry
      𝒟ₜ = Normal(state.μₜ, state.σₜ)
      pfalse = P.pfalse(rx)
      p = pfalse * pdf(𝒟ₜ, state.μₜ)
      push!(𝒮⁺, State(state.score + log10(p), state, missing, state.i, j, state.μₜ, state.σₜ))
      for i ∈ state.i+1:size(P.tx,1)
        tx = P.tx[i,:]
        Δt = P.time(rx) - P.time(tx) - P.delay(tx, rx)
        Δt < -3 * state.σₜ && break
        p = (1 - pfalse) * pdf(𝒟ₜ, Δt) * P.passoc(tx, rx)
        push!(𝒮⁺, State(state.score + log10(p), state, i => j, i, j, Δt, state.σₜ))
      end
    end
    sort!(𝒮⁺; by=score, rev=true)
    if isinf(𝒮⁺[1].score)
      @warn "Ran out of possibilities for RX[$j]!"
      @show 𝒮⁺
      break
    end
    filter!(s -> s.score ≥ 𝒮⁺[1].score - 1, 𝒮⁺)
    filter!(s -> !isduplicate(s, 𝒮⁺), 𝒮⁺)
    length(𝒮⁺) > nhypothesis && (𝒮⁺ = 𝒮⁺[1:nhypothesis])
    𝒮 = 𝒮⁺
  end
  assoc = DataFrame(txid=Int[], rxid=[])
  state = 𝒮[1]
  while state !== missing
    state.assoc === missing || push!(assoc, (state.assoc...,))
    state = state.backlink
  end
  sort!(assoc, :txid)
  assoc.txtime = P.tx[assoc.txid,:time]
  assoc.rxtime = P.rx[assoc.rxid,:time]
  assoc.Δt = [P.time(P.rx[assoc.rxid[i],:]) - P.time(P.tx[assoc.txid[i],:]) - P.delay(P.tx[assoc.txid[i],:], P.rx[assoc.rxid[i],:]) for i ∈ 1:size(assoc, 1)]
  assoc
end

end # module


text:
using CUDA
# using Flux

items = Int(1e6)
δ = rand(Float32, items, 1)
δ² = rand(Float32, items, 1)
𝑤 = rand(Float32, items)
pred = rand(Float32, items, 1)
target = rand(Float32, items)

δ_gpu = CuArray(δ)
δ²_gpu = CuArray(δ²)
𝑤_gpu = CuArray(𝑤)
pred_gpu = CuArray(pred)
target_gpu = CuArray(target)

function update_grads_gpu_linear_1!(pred::AbstractMatrix{T}, target::AbstractVector{T}, δ::AbstractMatrix{T}, δ²::AbstractMatrix{T}, 𝑤::AbstractVector{T}) where {T <: AbstractFloat}
    @. δ = 2f0 * (pred - target) * 𝑤
    @. δ² = 2f0 * 𝑤
    return
end

function kernel_linear_δ!(δ::CuDeviceMatrix{T}, p::CuDeviceMatrix{T}, t::CuDeviceVector{T}, 𝑤::CuDeviceVector{T}) where {T<:AbstractFloat}
    i = threadIdx().x + (blockIdx().x - 1) * blockDim().x
    if i <= length(t)
        @inbounds δ[i] = 2 * (p[i] - t[i]) * 𝑤[i]
    end
    return
end

function kernel_linear_δ²!(δ::CuDeviceMatrix{T}, p::CuDeviceMatrix{T}, t::CuDeviceVector{T}, 𝑤::CuDeviceVector{T}) where {T<:AbstractFloat}
    i = threadIdx().x + (blockIdx().x - 1) * blockDim().x
    if i <= length(t)
        @inbounds δ[i] = 2 * 𝑤[i]
    end
    return
end

# base approach - block built along the cols first, the rows (limit collisions)
function grad_linear!(δ::CuMatrix{T}, δ²::CuMatrix{T}, p::CuMatrix{T}, t::CuVector{T}, 𝑤::CuVector{T}; MAX_THREADS=1024) where {T<:AbstractFloat}
    thread_i = min(MAX_THREADS, length(t))
    threads = (thread_i)
    blocks = ceil.(Int, (length(t)) ./ threads)
    @cuda blocks=blocks threads=threads kernel_linear_δ!(δ, p, t, 𝑤)
    @cuda blocks=blocks threads=threads kernel_linear_δ²!(δ², p, t, 𝑤)
    return
end

CUDA.@time update_grads_gpu_linear_1!(pred_gpu, target_gpu, δ_gpu, δ²_gpu, 𝑤_gpu)
CUDA.@time grad_linear!(δ_gpu, δ²_gpu, pred_gpu, target_gpu, 𝑤_gpu, MAX_THREADS=1024)

#################################################
# Gaussian
#################################################
items = Int(1e6)
δ = zeros(Float32, items, 1)
δ² = zeros(Float32, items, 1)
𝑤 = rand(Float32, items)
pred = rand(Float32, items, 1)
target = rand(Float32, items)

δ_gpu = CuArray(δ)
δ²_gpu = CuArray(δ²)
𝑤_gpu = CuArray(𝑤)
pred_gpu = CuArray(pred)
target_gpu = CuArray(target)

function kernel_gauss_δ!(δ::CuDeviceMatrix{T}, p::CuDeviceMatrix{T}, t::CuDeviceVector{T}, 𝑤::CuDeviceVector{T}) where {T<:AbstractFloat}
    i = threadIdx().x + (blockIdx().x - 1) * blockDim().x
    if i <= length(t)
        δ[i,1] = (p[i,1] - t[i]) / max(Cfloat(1e-8), exp(2f0 * p[i,2])) * 𝑤[i]
        δ[i,2] = (1f0 - (p[i,1] - t[i])^2f0 / max(Cfloat(1e-8), exp(2f0 * p[i,2]))) * 𝑤[i]
    end
    return
end

function kernel_gauss_δ²!(δ::CuDeviceMatrix{T}, p::CuDeviceMatrix{T}, t::CuDeviceVector{T}, 𝑤::CuDeviceVector{T}) where {T<:AbstractFloat}
    i = threadIdx().x + (blockIdx().x - 1) * blockDim().x
    if i <= length(t)
        δ[i,1] = 𝑤[i] / max(Cfloat(1e-8), exp(2 * p[i,2]))
        δ[i,2] = 2 * 𝑤[i] / max(Cfloat(1e-8), exp(2 * pred[i,2])) * (p[i,1] - target[i])^2
    end
end

# base approach - block built along the cols first, the rows (limit collisions)
function grad_gaussian!(δ::CuMatrix{T}, δ²::CuMatrix{T}, p::CuMatrix{T}, t::CuVector{T}, 𝑤::CuVector{T}; MAX_THREADS=1024) where {T<:AbstractFloat}
    thread_i = min(MAX_THREADS, length(t))
    threads = (thread_i)
    blocks = ceil.(Int, (length(t)) ./ threads)
    @cuda blocks=blocks threads=threads kernel_linear_δ!(δ, p, t, 𝑤)
    @cuda blocks=blocks threads=threads kernel_linear_δ²!(δ², p, t, 𝑤)
    return
end

CUDA.@time grad_gaussian!(δ_gpu, δ²_gpu, pred_gpu, target_gpu, 𝑤_gpu, MAX_THREADS=1024)


text:
import AdversarialPrediction: define, constraint, generate_constraints_on_p!
using LinearAlgebra

# sample metrics
@metric Prec
function define(::Type{Prec}, C::ConfusionMatrix)
    return C.tp / C.pp
end

@metric F1_score
function define(::Type{F1_score}, C::ConfusionMatrix)
    return (2 * C.tp) / (C.ap + C.pp)
end

@metric Kappa
function define(::Type{Kappa}, C::ConfusionMatrix)
    num = (C.tp + C.tn) / C.all - (C.ap * C.pp + C.an * C.pn) / C.all^2
    den = 1 - (C.ap * C.pp + C.an * C.pn) / C.all^2
    return num / den
end

@metric Prec_F1_score
function define(::Type{Prec_F1_score}, C::ConfusionMatrix)
    return C.tp / C.pp
end
function constraint(::Type{Prec_F1_score}, C::ConfusionMatrix)
    return (2 * C.tp) / (C.ap + C.pp) >= 0.6
end

@testset "compute metric" begin

    # hardcoded functions
    prec_fn(yhat, y) = (dot(yhat, y) / sum(yhat))
    f1_fn(yhat, y) = (2 * dot(yhat, y) / (sum(yhat) + sum(y)))

    function kappa_fn(yhat, y)
        n = length(y)
        num = ((dot(yhat, y) + dot(1 .- yhat, 1 .- y)) / n -
            (sum(y) * sum(yhat) + sum(1 .- y) * sum(1 .- yhat)) / n^2 )
        den = 1 - (sum(y) * sum(yhat) + sum(1 .- y) * sum(1 .- yhat)) / n^2
        return num / den
    end

    y = rand(0:1, 100)
    yhat = rand(0:1, 100)

    prec = Prec()
    @test compute_metric(prec, yhat, y) ≈ prec_fn(yhat, y)

    f1 = F1_score()
    @test compute_metric(f1, yhat, y) ≈ f1_fn(yhat, y)

    kappa = Kappa()
    @test compute_metric(kappa, yhat, y) ≈ kappa_fn(yhat, y)
end

@testset "constraints metric" begin

    # hardcoded functions
    function f1_fn(y)
        n = length(y)
        iq = sum(y .== 1)

        ks = 1:n
        B = zeros(n,n)
        c = 0.0

        for ip = 1:n
            B[:,ip] = ( 2 * y ) / (ip + iq)
        end

        if iq == 0
            B .+= (-1 ./ ks')
            c += 1.0
        end

        return B, c
    end

    prec_f1 = Prec_F1_score()
    special_case_positive!(prec_f1)
    cs_special_case_positive!(prec_f1, true)

    y = rand(0:1, 20)
    generate_constraints_on_p!(prec_f1, y)

    B, c = f1_fn(y)
    @test B ≈ prec_f1.data.B_list[1]
    @test isapprox(c, prec_f1.data.c_list[1], atol = 1e-6)

    # if y == zeros
    y = zeros(20)
    generate_constraints_on_p!(prec_f1, y)

    B, c = f1_fn(y)
    @test B ≈ prec_f1.data.B_list[1]
    @test isapprox(c, prec_f1.data.c_list[1], atol = 1e-6)

    # if y == ones
    y = ones(20)
    generate_constraints_on_p!(prec_f1, y)

    B, c = f1_fn(y)
    @test B ≈ prec_f1.data.B_list[1]
    @test isapprox(c, prec_f1.data.c_list[1], atol = 1e-6)

end


text:
struct PlutoBarycenter <: abstractBarycenter end

naif_id(::Type{PlutoBarycenter}) = 9
parent(::Type{PlutoBarycenter}) = SSB


text:
using ITensors,
      ITensorsGPU,
      LinearAlgebra, # For tr()
      Combinatorics, # For permutations()
      Random,
      CUDA,
      Test

@testset "cuITensor $T Contractions" for T ∈ (Float64,ComplexF64)
  mi,mj,mk,ml,ma = 2,3,4,5,6,7
  i = Index(mi,"i")
  j = Index(mj,"j")
  k = Index(mk,"k")
  l = Index(ml,"l")
  a = Index(ma,"a")
  @testset "Test contract cuITensors" begin
      A = cuITensor(randomITensor(T))
      B = cuITensor(randomITensor(T))
      Ai = cuITensor(randomITensor(T,i))
      Bi = cuITensor(randomITensor(T,i))
      Aj = cuITensor(randomITensor(T,j))
      Aij = cuITensor(randomITensor(T,i,j))
      Aji = cuITensor(randomITensor(T,j,i))
      Bij = cuITensor(randomITensor(T,i,j))
      Aik = cuITensor(randomITensor(T,i,k))
      Ajk = cuITensor(randomITensor(T,j,k))
      Ajl = cuITensor(randomITensor(T,j,l))
      Akl = cuITensor(randomITensor(T,k,l))
      Aijk = cuITensor(randomITensor(T,i,j,k))
      Ajkl = cuITensor(randomITensor(T,j,k,l))
      Aikl = cuITensor(randomITensor(T,i,k,l))
      Akla = cuITensor(randomITensor(T,k,l,a))
      Aijkl = cuITensor(randomITensor(T,i,j,k,l))
    @testset "Test contract cuITensor (Scalar*Scalar -> Scalar)" begin
      C = A*B
      @test scalar(C)≈scalar(A)*scalar(B)
      C = cuITensor(T(2.0))*cuITensor(T(2.0))
      @test scalar(C)≈T(4.0)
    end
    @testset "Test contract cuITensor (Scalar*Vector -> Vector)" begin
      C = A*Ai
      @test cpu(C)≈scalar(A)*cpu(Ai)
      C = cuITensor(T(2.0))*Ai
      @test cpu(C)≈T(2.0)*cpu(Ai)
      C = Ai*cuITensor(T(2.0))
      @test cpu(C)≈T(2.0)*cpu(Ai)
    end
    @testset "Test contract cuITensor (Vector*Scalar -> Vector)" begin
      C = Aj*A
      @test cpu(C)≈scalar(A)*cpu(Aj)
    end
    @testset "Test contract cuITensors (Vectorᵀ*Vector -> Scalar)" begin
      C = Ai*Bi
      Ccpu = cpu(Ai) * cpu(Bi)
      @test scalar(Ccpu)≈scalar(C)
    end
    @testset "Test contract cuITensors (Vector*Vectorᵀ -> Matrix)" begin
      C = Ai*Aj
      Ccpu = cpu(Ai)*cpu(Aj)
      @test Ccpu≈cpu(permute(C,i,j))
    end
    @testset "Test contract cuITensors (Matrix*Scalar -> Matrix)" begin
      Aij = permute(Aij,i,j)
      C = Aij*A
      @test cpu(permute(C,i,j))≈scalar(A)*cpu(Aij)
    end
    @testset "Test contract cuITensors (Matrixᵀ*Vector -> Vector)" begin
      cAij = permute(copy(Aij),j,i)
      Ccpu = cpu(Aij)*cpu(Aj)
      C = cAij*Aj
      @test Ccpu ≈ cpu(C)
    end
    @testset "Test contract cuITensors (Matrix*Vector -> Vector)" begin
      cpAij = permute(copy(Aij),i,j)
      Ccpu = cpu(cpAij)*cpu(Aj)
      C   = copy(cpAij)*copy(Aj)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Vector*Matrix -> Vector)" begin
      Aij = permute(Aij,i,j)
      C = Ai*Aij
      Ccpu = cpu(Ai) * cpu(Aij)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Vector*Matrixᵀ -> Vector)" begin
      C = Ai*Aji
      Ccpu = cpu(Ai) * cpu(Aji)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Matrix*Matrix -> Scalar)" begin
      Aij = permute(Aij,i,j)
      Bij = permute(Bij,i,j)
      C = Aij*Bij
      Ccpu = cpu(Aij) * cpu(Bij)
      @test scalar(Ccpu)≈scalar(C)
    end
    @testset "Test contract cuITensors (Matrix*Matrix -> Matrix)" begin
      Aij = permute(Aij,i,j)
      Ajk = permute(Ajk,j,k)
      C = Aij*Ajk
      Ccpu = cpu(Aij)*cpu(Ajk)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Matrixᵀ*Matrix -> Matrix)" begin
      Aij = permute(Aij,j,i)
      Ajk = permute(Ajk,j,k)
      C = Aij*Ajk
      Ccpu = cpu(Aij) * cpu(Ajk)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Matrix*Matrixᵀ -> Matrix)" begin
      Aij = permute(Aij,i,j)
      Ajk = permute(Ajk,k,j)
      C = Aij*Ajk
      Ccpu = cpu(Aij) * cpu(Ajk)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Matrixᵀ*Matrixᵀ -> Matrix)" begin
      Aij = permute(Aij,j,i)
      Ajk = permute(Ajk,k,j)
      C = Aij*Ajk
      Ccpu = cpu(Aij) * cpu(Ajk)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (3-Tensor*Scalar -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      C = Aijk*A
      @test cpu(permute(C,i,j,k))≈scalar(A)*cpu(Aijk)
    end
    @testset "Test contract cuITensors (3-Tensor*Vector -> Matrix)" begin
      cAijk = permute(copy(Aijk),i,j,k)
      C = cAijk*Ai
      Ccpu = cpu(cAijk) * cpu(Ai)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (Vector*3-Tensor -> Matrix)" begin
      Aijk = permute(Aijk,i,j,k)
      C = Aj*Aijk
      Ccpu = cpu(Aj)*cpu(Aijk)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (3-Tensor*Matrix -> Vector)" begin
      Aijk = permute(Aijk,i,j,k)
      Aik = permute(Aik,i,k)
      C = Aijk*Aik
      Ccpu = cpu(Aijk)*cpu(Aik)
      @test Ccpu≈cpu(C)
    end
    @testset "Test contract cuITensors (3-Tensor*Matrix -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Ajl = permute(Ajl,j,l)
      C = Aijk*Ajl
      Ccpu = cpu(Aijk)*cpu(Ajl)
      @test Ccpu≈cpu(permute(C,i,k,l))
    end
    @testset "Test contract cuITensors (Matrix*3-Tensor -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Akl  = permute(Akl,k,l)
      C = Akl*Aijk
      Ccpu = cpu(Aijk)*cpu(Akl)
      @test Ccpu≈cpu(permute(C,l,i,j))
    end
    @testset "Test contract cuITensors (3-Tensor*3-Tensor -> 3-Tensor)" begin
      Aijk = permute(Aijk,i,j,k)
      Ajkl = permute(Ajkl,j,k,l)
      C = Aijk*Ajkl
      Ccpu = cpu(Aijk)*cpu(Ajkl)
      @test Ccpu≈cpu(permute(C,i,l))
    end
    @testset "Test contract cuITensors (3-Tensor*3-Tensor -> 3-Tensor)" begin
      for inds_ijk ∈ permutations([i,j,k]), inds_jkl ∈ permutations([j,k,l])
        Aijk = permute(Aijk,inds_ijk...)
        Ajkl = permute(Ajkl,inds_jkl...)
        C = Ajkl*Aijk
        Ccpu = cpu(Ajkl) * cpu(Aijk)
        @test Ccpu≈cpu(C)
      end
    end
    @testset "Test contract cuITensors (4-Tensor*3-Tensor -> 3-Tensor)" begin
      for inds_ijkl ∈ permutations([i,j,k,l]), inds_kla ∈ permutations([k,l,a])
        Aijkl = permute(Aijkl,inds_ijkl...)
        Akla = permute(Akla,inds_kla...)
        C = Akla*Aijkl
        Ccpu = cpu(Akla) * cpu(Aijkl)
        @test Ccpu≈cpu(C)
      end
    end
    @testset "Test contract cuITensors (4-Tensor*3-Tensor -> 1-Tensor)" begin
      for inds_ijkl ∈ permutations([i,j,k,l]), inds_jkl ∈ permutations([j,k,l])
        Aijkl = permute(Aijkl,inds_ijkl...)
        Ajkl = permute(Ajkl,inds_jkl...)
        C = Ajkl*Aijkl
        Ccpu = cpu(Ajkl)*cpu(Aijkl)
        @test Ccpu≈cpu(C)
      end
    end
    @testset "Test supersized contract cuITensors (14-Tensor*14-Tensor -> 14-Tensor)" begin
        a_only_inds = [Index(2) for ii in 1:7]
        b_only_inds = [Index(2) for ii in 1:7]
        shared_inds = [Index(2) for ii in 1:7]
        A = randomITensor(IndexSet(vcat(a_only_inds, shared_inds)))
        B = randomITensor(IndexSet(vcat(b_only_inds, shared_inds)))
        cA = cuITensor(A)
        cB = cuITensor(B)
        inds_a = vcat(a_only_inds, shared_inds)
        inds_b = vcat(b_only_inds, shared_inds)
        cA_ = permute(cA,inds_a...)
        cB_ = permute(cB,inds_b...)
        C = cA_*cB_
        Ccpu = cpu(cA_)*cpu(cB_)
        @test Ccpu≈cpu(C)
        for shuffles in 1:1 # too many permutations to test all
            inds_a = shuffle(vcat(a_only_inds, shared_inds))
            inds_b = shuffle(vcat(b_only_inds, shared_inds))
            cA_ = permute(cA,inds_a...)
            cB_ = permute(cB,inds_b...)
            C = cA_*cB_
            Ccpu = cpu(cA_)*cpu(cB_)
            @test Ccpu≈cpu(C)
        end
    end
  end # End contraction testset
end


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    BoundedSearcher(searcher, nmax)

A method for searching at most `nmax` neighbors using `searcher`.
"""
struct BoundedSearcher{S<:AbstractNeighborSearcher} <: AbstractBoundedNeighborSearcher
  searcher::S
  nmax::Int
end

object(searcher::BoundedSearcher) = object(searcher.searcher)

maxneighbors(searcher::BoundedSearcher) = searcher.nmax

function search!(neighbors, xₒ::AbstractVector,
                 searcher::BoundedSearcher; mask=nothing)
  inds = search(xₒ, searcher.searcher)
  nmax = searcher.nmax

  if isnothing(mask)
    nneigh = min(length(inds), nmax)
    @inbounds for i in 1:nneigh
      neighbors[i] = inds[i]
    end
  else
    nneigh = 0
    @inbounds for ind in inds
      if mask[ind]
        nneigh += 1
        neighbors[nneigh] = ind
      end
      nneigh == nmax && break
    end
  end

  nneigh
end


text:
include("update.jl")

function randomwalk(nsteps::Integer, prob::Real=0.5)
    result = zeros(Int, nsteps+1)
    x = 0
    result[1] = x
    for t in 1:nsteps
        x = update(x,prob)
        result[t+1] = x
    end
    return result
end

function main(nsteps, prob)
    result = randomwalk(nsteps,prob)

    for (t,x) in enumerate(result)
        println("$(t-1) $x")
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    nsteps = length(ARGS)>0 ? parse(Int, ARGS[1]) : 1000
    prob = length(ARGS)>1 ? parse(Float64, ARGS[1]) : 0.5

    main(nsteps, prob)
end

max_stars_repo_name: JuliaTagBot/JuliaBook-Samples

max_issues_repo_name: JuliaTagBot/JuliaBook-Samples

max_forks_repo_name: JuliaTagBot/JuliaBook-Samples


text:
using StaticArrays

export ForwardPredictProblem, forward_predict!, self_estimate, self_z_estimate

struct ForwardPredictProblem{N,Hf,X,Z,U,Dy,Ctrl}
    world_dynamics::Dy
    π::Ctrl
    δx::MMatrix{Hf,N,X}
    z_obs::SizedMatrix{Hf,N,Optional{Z}}

    function ForwardPredictProblem(
        world_dynamics::Dy, π::Ctrl; X, Z, Hf,
    ) where {N,U,Dy <: WorldDynamics{N},Ctrl <: CentralControl{U}}
        @assert isbitstype(X) "X = $X is not of bits type"
        @assert isbitstype(U) "U = $U is not of bits type"
        new{N,Hf,X,Z,U,Dy,Ctrl}(
            world_dynamics, π,
            MMatrix{Hf,N,X}(fill(zero(X), Hf, N)),
            SizedMatrix{Hf,N,Optional{Z}}(fill(missing, Hf, N)))
    end
end

"""
Forward-predict the ideal fleet trajectory.
Returns `(ũ[t in t0: t0+H-1], x̃[t in t0+1: t0+H], z̃[t in t0+1: t0+H])`.
Note that `init_s` will be modified from `t=t0` to `t=t0+ΔT`. (It will then be deep copied)

# Argument Details
- `prob.δx`: state disturbance ∈ [t0, t0+H-1]
- `prob.z_obs`: observation ∈ [t0, t0+H-1]
- `init_xz`: state `x` and observation `z` at time t0
- `init_s`: central control state at time t0, will be modified to t0+1.
"""
function forward_predict!(
    prob::ForwardPredictProblem{N,Hf,X,Z,U},
    init_xz::Each{Tuple{X,Z}},
    init_s::S,
    τ0::𝕋,
    ΔT::𝕋,
) where {X,Z,U,S,N,Hf}
    x_traj = MMatrix{Hf,N,X}(undef)
    z_traj = SizedMatrix{Hf,N,Z}(undef)
    u_traj = MMatrix{Hf,N,U}(undef)

    xs = MVector{N,X}(first.(init_xz))
    zs = SizedVector{N,Z}(last.(init_xz))
    x_dy = prob.world_dynamics.dynamics
    z_dy = prob.world_dynamics.obs_dynamics

    s_c = init_s

    for t in 1:Hf
        us = control_all(prob.π, s_c, xs, zs, τ0 + t - 1, SOneTo(N))
        us = limit_control.(x_dy, us, xs, τ0 + t - 1)
        @inbounds for i in 1:N
            xs[i] = sys_forward(x_dy[i], xs[i], us[i], N) + prob.δx[t, i]
            z = prob.z_obs[t, i]
            zs[i] = (z === missing) ? obs_forward(z_dy[i], xs[i], zs[i], N) : z
        end
        x_traj[t,:] = xs
        z_traj[t,:] = zs
        u_traj[t,:] = us
        if t == ΔT
            s_c = deepcopy(s_c)
        end
    end
    u_traj, x_traj, z_traj
end

"""
Estimate the self state history using an old state and the actuation history.
Returns `x_history`.
"""
function self_estimate(
    self_dynamics::SysDynamics,
    s0::Tuple{X,𝕋},
    u_history,
)::Vector{X} where {X}
    x, τ = s0
    x_history = Vector{X}(undef, length(u_history))
    for (t, u) in enumerate(u_history)
        x = sys_forward(self_dynamics, x, u, τ + t - 1)
        x_history[t] = x
    end
    x_history
end

function self_estimate(
    self_dynamics::SysDynamics,
    s0::Timed{X},
    u_history,
)::Vector{Timed{X}} where {X}
    x, τ = s0.value, s0.time
    x_history = Vector{Timed{X}}(undef, length(u_history))
    for (t, u) in enumerate(u_history)
        u::Timed
        x = sys_forward(self_dynamics, x, u[τ + t - 1], τ + t - 1)
        x_history[t] = Timed(τ + t, x)
    end
    x_history
end

function self_z_estimate(
    z_dynamics::ObsDynamics,
    z0::Timed{Z},
    x_history,
)::Vector{Timed{Z}} where {Z}
    z, τ = z0.value, z0.time
    z_history = Vector{Timed{Z}}(undef, length(x_history))
    for (t, x) in enumerate(x_history)
        x::Timed
        z = obs_forward(z_dynamics, x[τ + t - 1], z, τ + t - 1)
        z_history[t] = Timed(τ + t, z)
    end
    z_history
end