text:
struct A1 <: Layout
    a::Bool
    b::Cint
    c::Float64
    d::Cstring
end

@testset "field access" begin
    a = fill(UInt64(0), 3)
    p = pointer(a)
    cs = CStruct{A1}(p)
    @test cs.a == 0
    @test cs.b == 0
    @test cs.c == 0
    @test cs.d == ""
    v1 = true
    v2 = 0x12345678
    v3 = 47.11
    v4 = "hallo"
    cs.a = v1
    cs.b = v2
    cs.c = v3
    cs.d = v4
    @test cs.a == v1
    @test cs.b == v2
    @test cs.c == v3
    @test cs.d == v4
end

@testset "index access" begin
    a = fill(UInt64(0), 100)
    p = pointer(a)
    cv = CVector{Int}(p, 3)
    @test length(cv) == 3
    cv[1:3] .= (1, 2, 3)
    @test cv[2] == 2
    @test cv[[1,3]] == [1, 3]
end

struct A2 <: Layout
    a::Ptr{A2}
end

@testset "self-referencing" begin
    a = fill(UInt8(0), 100)
    p = pointer(a)
    cs = CStruct{A2}(p)
    @test cs.a === nothing
    io = IOBuffer()
    show(io, cs)
    @test String(take!(io)) == "CStruct{A2}(nothing)"
    cs.a = cs
    @test cs.a === cs
    show(io, cs)
    @test String(take!(io)) == "CStruct{A2}(CStruct{A2}(#= circular reference @-1 =#))"
end

struct A3 <: Layout
    len::Int
    vec::LVarVector{Float64, (x) -> x.len}
end

@testset "variable vector at end of struct" begin
    a = fill(Int(0), 1024)
    p = pointer(a)
    LEN = 25
    cs = CStruct{A3}(p)
    cs.len = LEN
    @test cs.vec isa CVector{Float64}
    @test length(cs.vec) == cs.len == LEN
end

struct A4 <: Layout
    len::Int
    vec::Ptr{LVarVector{Float64, (x) -> x.len}}
end

@testset "pointer to variable vector" begin
    a = fill(Int(0), 1024)
    p = pointer(a)
    a[2] = p + 32
    LEN = 25
    cs = CStruct{A4}(p)
    cs.len = LEN
    @test cs.vec isa CVector{Float64}
    @test length(cs.vec) == cs.len == LEN
end

struct B <: Layout
    a::Int
    b::LVarVector{Float64, :a}
end

@testset "variable length vector in struct" begin
    @test_throws ArgumentError Cserialize(B, ())
    cs = CStruct{B}(Cserialize(B, (a=0, b=Float64[])))
    @test length(cs.b) == cs.a == 0
    cs = CStruct{B}(Cserialize(B, (a=2, b=[1.0; 2; 3])))
    @test length(cs.b) == cs.a == 2
    cs.a = 3
    @test length(cs.b) == cs.a == 3
end

struct I1
    a::Int8
    b::Int16
    I1() = new(1, 1)
end

mutable struct M1
    a::Int8
    b::Float64
    M1() = new(2, 2)
end

struct I2
    a::Int16
    b::M1
    I2() = new(3, M1())
end

mutable struct M2
    a::Int16
    b::M1
    M2() = new(4, M1())
end

@testset "construct function $T" for (T, a1, a2) in ((I1, 12, 12), (M1, 12, 12), (I2, 12, M1()), (M2, 12, M1()))
    @test construct(T, a1, a2) isa T
end


text:
@testset "Variable" begin
    var = Variable{Float64}(2, "dog")

    @test var.id == 2
    @test var.name == "dog"

    # Generate shape function and Variable values for a dummy element
    qr = QuadratureRule{2,RefCube}(:legendre, 2)
    func_space = Lagrange{2, RefCube, 1}()
    fe_values = FECellValues(qr, func_space)
    n_qp = length(points(qr))

    # Nodal coordinates
    coords = Vec{2, Float64}[Vec{2}((-1, -1)),
                             Vec{2}((1, -1)),
                             Vec{2}((1, 1)),
                             Vec{2}((-1, 1))]

    reinit!(fe_values, coords)

    dof_values = [0., 1., 1., 0.]

    MOOSE.reinit!(var, fe_values, [1,2,3,4], dof_values)

    @test var.dofs == [1,2,3,4]
    @test var.n_dofs == 4
    @test var.n_qp == 4

    for grad in var.grad
        @test abs(grad[1] - 0.5) < 1e-9
        @test abs(grad[2] - 0.0) < 1e-9
    end

    # Test reinit at a node
    MOOSE.reinit!(var, 3, 2.7)

    @test var.nodal_dof == 3
    @test var.nodal_value == 2.7
end

@testset "FDVariable" begin
    var = Variable{Dual{4, Float64}}(2, "dog")

    @test var.id == 2
    @test var.name == "dog"

    # Generate shape function and Variable values for a dummy element
    qr = QuadratureRule{2,RefCube}(:legendre, 2)
    func_space = Lagrange{2, RefCube, 1}()
    fe_values = FECellValues(qr, func_space)
    n_qp = length(points(qr))

    # Nodal coordinates
    coords = Vec{2, Float64}[Vec{2}((-1, -1)),
                             Vec{2}((1, -1)),
                             Vec{2}((1, 1)),
                             Vec{2}((-1, 1))]

    reinit!(fe_values, coords)

    dof_values = [MOOSE.dualVariable(0.,1,4),
                  MOOSE.dualVariable(1.,2,4),
                  MOOSE.dualVariable(1.,3,4),
                  MOOSE.dualVariable(0.,4,4)]

    MOOSE.reinit!(var, fe_values, [1,2,3,4], dof_values)

    @test partials(var.value[1])[1] == var.phi[1][1]
    @test partials(var.grad[1][1])[1] == var.grad_phi[1][1][1]

    @test partials(var.value[2])[1] == var.phi[1][2]
    @test partials(var.grad[2][1])[1] == var.grad_phi[2][1][1]

    @test var.dofs == [1,2,3,4]
    @test var.n_dofs == 4
    @test var.n_qp == 4

    for grad in var.grad
        @test abs(grad[1] - 0.5) < 1e-9
        @test abs(grad[2] - 0.0) < 1e-9
    end

    # Test reinit at a node
    MOOSE.reinit!(var, 3, 2.7)

    @test var.nodal_dof == 3
    @test var.nodal_value == 2.7
end

ext: jl


text:
using PyPlot;

function plotFunc_timecourse(Sim::Module)
    rc("figure",figsize = (12,15));
    rc("font",size = 10);
    rc("axes",linewidth = 1);
    rc("lines",linewidth = 2);

    subplots_adjust(wspace=0.4, hspace=0.5);

    subplot(4,2,1)
    plot(Sim.t,Sim.ERK_act[:,1]./589.5,"-",label="EGF=0 nM");
    plot(Sim.t,Sim.ERK_act[:,4]./589.5,"--",label="EGF=0.5 nM");
    plot(Sim.t,Sim.ERK_act[:,7]./589.5,":",label="EGF=10 nM");
    title("EGF increasing; HRG=0.5 nM");
    xlim(0,1900);
    xticks([0,500,1000,1500]);
    xlabel("Time (s)");
    ylim(0,1.5);
    yticks([0,0.5,1,1.5]);
    ylabel("Normalized ERK*");
    legend(loc="upper right",fontsize=8);

    subplot(4,2,2)
    plot(Sim.t,Sim.ERK_act[:,2]./589.5,"-",label="EGF=0 nM");
    plot(Sim.t,Sim.ERK_act[:,5]./589.5,"--",label="EGF=0.5 nM");
    plot(Sim.t,Sim.ERK_act[:,8]./589.5,":",label="EGF=10 nM");
    title("EGF increasing; HRG=10 nM");
    xlim(0,1900);
    xticks([0,500,1000,1500]);
    xlabel("Time (s)");
    ylim(0,1.5);
    yticks([0,0.5,1,1.5]);
    ylabel("Normalized ERK*");
    legend(loc="upper right",fontsize=8);

    subplot(4,2,3)
    plot(Sim.t,Sim.ERK_act[:,3]./589.5,"-",label="HRG=0 nM");
    plot(Sim.t,Sim.ERK_act[:,4]./589.5,"--",label="HRG=0.5 nM");
    plot(Sim.t,Sim.ERK_act[:,5]./589.5,":",label="HRG=10 nM");
    title("EGF=0.5 nM; HRG increasing");
    xlim(0,1900);
    xticks([0,500,1000,1500]);
    xlabel("Time (s)");
    ylim(0,1.5);
    yticks([0,0.5,1,1.5]);
    ylabel("Normalized ERK*");
    legend(loc="upper right",fontsize=8);

    subplot(4,2,4)
    plot(Sim.t,Sim.ERK_act[:,6]./589.5,"-",label="HRG=0 nM");
    plot(Sim.t,Sim.ERK_act[:,7]./589.5,"--",label="HRG=0.5 nM");
    plot(Sim.t,Sim.ERK_act[:,8]./589.5,":",label="HRG=10 nM");
    title("EGF=10 nM; HRG increasing");
    xlim(0,1900);
    xticks([0,500,1000,1500]);
    xlabel("Time (s)");
    ylim(0,1.5);
    yticks([0,0.5,1,1.5]);
    ylabel("Normalized ERK*");
    legend(loc="upper right",fontsize=8);

    subplot(4,2,5)
    plot(Sim.t,Sim.Akt_act[:,1]./18.8,"-",label="EGF=0 nM");
    plot(Sim.t,Sim.Akt_act[:,4]./18.8,"--",label="EGF=0.5 nM");
    plot(Sim.t,Sim.Akt_act[:,7]./18.8,":",label="EGF=10 nM");
    title("EGF increasing; HRG=0.5 nM");
    xlim(0,1900);
    xticks([0,500,1000,1500]);
    xlabel("Time (s)");
    ylim(0,1.5);
    yticks([0,0.5,1,1.5]);
    ylabel("Normalized Akt*");
    legend(loc="upper right",fontsize=8);

    subplot(4,2,6)
    plot(Sim.t,Sim.Akt_act[:,2]./18.8,"-",label="EGF=0 nM");
    plot(Sim.t,Sim.Akt_act[:,5]./18.8,"--",label="EGF=0.5 nM");
    plot(Sim.t,Sim.Akt_act[:,8]./18.8,":",label="EGF=10 nM");
    title("EGF increasing; HRG=10 nM");
    xlim(0,1900);
    xticks([0,500,1000,1500]);
    xlabel("Time (s)");
    ylim(0,1.5);
    yticks([0,0.5,1,1.5]);
    ylabel("Normalized Akt*");
    legend(loc="upper right",fontsize=8);

    subplot(4,2,7)
    plot(Sim.t,Sim.Akt_act[:,3]./18.8,"-",label="HRG=0 nM");
    plot(Sim.t,Sim.Akt_act[:,4]./18.8,"--",label="HRG=0.5 nM");
    plot(Sim.t,Sim.Akt_act[:,5]./18.8,":",label="HRG=10 nM");
    title("EGF=0.5 nM; HRG increasing");
    xlim(0,1900);
    xticks([0,500,1000,1500]);
    xlabel("Time (s)");
    ylim(0,1.5);
    yticks([0,0.5,1,1.5]);
    ylabel("Normalized Akt*");
    legend(loc="upper right",fontsize=8);

    subplot(4,2,8)
    plot(Sim.t,Sim.Akt_act[:,6]./18.8,"-",label="HRG=0 nM");
    plot(Sim.t,Sim.Akt_act[:,7]./18.8,"--",label="HRG=0.5 nM");
    plot(Sim.t,Sim.Akt_act[:,8]./18.8,":",label="HRG=10 nM");
    title("EGF=10 nM; HRG increasing");
    xlim(0,1900);
    xticks([0,500,1000,1500]);
    xlabel("Time (s)");
    ylim(0,1.5);
    yticks([0,0.5,1,1.5]);
    ylabel("Normalized Akt*");
    legend(loc="upper right",fontsize=8);

    show()

end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct NetworkProfile <: SwaggerModel
    networkInterfaces::Any # spec type: Union{ Nothing, Vector{NetworkInterfaceReference} } # spec name: networkInterfaces

    function NetworkProfile(;networkInterfaces=nothing)
        o = new()
        validate_property(NetworkProfile, Symbol("networkInterfaces"), networkInterfaces)
        setfield!(o, Symbol("networkInterfaces"), networkInterfaces)
        o
    end
end # type NetworkProfile

const _property_map_NetworkProfile = Dict{Symbol,Symbol}(Symbol("networkInterfaces")=>Symbol("networkInterfaces"))
const _property_types_NetworkProfile = Dict{Symbol,String}(Symbol("networkInterfaces")=>"Vector{NetworkInterfaceReference}")
Base.propertynames(::Type{ NetworkProfile }) = collect(keys(_property_map_NetworkProfile))
Swagger.property_type(::Type{ NetworkProfile }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_NetworkProfile[name]))}
Swagger.field_name(::Type{ NetworkProfile }, property_name::Symbol) =  _property_map_NetworkProfile[property_name]

function check_required(o::NetworkProfile)
    true
end

function validate_property(::Type{ NetworkProfile }, name::Symbol, val)
end


text:
# this represents a circular buffer, where new items are pushed to the back of the list,
# overwriting values in a circular fashion
# iterating occurs from the "start" until the "end", and will return "length(buffer)" nodes

type CircularBuffer{T} <: AbstractVector{T}
  capacity::Int
  startIdx::Int
  buffer::Vector{T}

  CircularBuffer{T}(::Type{T}, capacity::Int) = new(capacity, 1, T[])
end

CircularBuffer{T}(::Type{T}, capacity::Int) = CircularBuffer{T}(T, capacity)

function CircularBuffer{T}(::Type{T}, capacity::Int, fillval::T)
  cb = CircularBuffer(T, capacity)
  for i in 1:capacity
    push!(cb, copy(fillval))
  end
  cb
end

function CircularBuffer{T}(::Type{T}, capacity::Int, v::Vector{T})
  cb = CircularBuffer(T, capacity)
  for x in v
    push!(cb, copy(x))
  end
  cb
end

function Base.print{T}(io::IO, cb::CircularBuffer{T})
  print(io, "{")
  print(io, join(collect(cb), ", "))
  print(io, "}")
end

function Base.show{T}(io::IO, cb::CircularBuffer{T})
  print(io, "CircularBuffer(cap=$(cb.capacity) start=$(cb.startIdx) n=$(length(cb)))")
  print(io, cb)
end

# cap = 8
# n = 6
# start = 6

# i = 1
# x = 6

# i = 4
# x = 1

# 4 5 6 _ _ 1 2 3

function bufferIndex(cb::CircularBuffer, i::Int)
  if i < 1 || i > length(cb)
    error("CircularBuffer out of range. cb=$cb i=$i")
  end
  (i + cb.startIdx - 2) % cb.capacity + 1
end

Base.getindex(cb::CircularBuffer, i::Int) = cb.buffer[bufferIndex(cb, i)]
# Base.getindex{T}(cb::CircularBuffer{T}, itr::IntIterable) = T[cb[i] for i in itr]
function Base.setindex!{T}(cb::CircularBuffer, data::T, i::Int)
  cb.buffer[bufferIndex(cb, i)] = data
  nothing
end

# Base.endof(cb::CircularBuffer) = length(cb)

# Note: iterating like "for x in X; (do something); end" translates to:
# state = start(X)
# while !done(X, state)
#   (x, state) = next(X, state)
#   (do something)
# end

# # iterate with the syntax "for x in list"
# Base.start{T}(cb::CircularBuffer{T}) = 1
# Base.done{T}(cb::CircularBuffer{T}, state::Int) = state > length(cb)
# Base.next{T}(cb::CircularBuffer{T}, state::Int) = (cb.buffer[bufferIndex(cb, state)], state+1)

# add to back and return the node
function Base.push!{T}(cb::CircularBuffer{T}, data::T)

  if length(cb) == cb.capacity
    # if the buffer is full, increment the startIdx, and add to the end
    cb.startIdx = (cb.startIdx == cb.capacity ? 1 : cb.startIdx + 1)
    cb[length(cb)] = data

  else

    # add to the end
    push!(cb.buffer, data)

  end
end

function Base.append!{T}(cb::CircularBuffer{T}, datavec::AbstractVector{T})
  for d in datavec
    push!(cb, d)
  end
end

Base.length(cb::CircularBuffer) = length(cb.buffer)
# Base.isempty(cb::CircularBuffer) = length(cb) == 0
Base.size(cb::CircularBuffer) = (length(cb),)

capacity(cb::CircularBuffer) = cb.capacity
isfull(cb::CircularBuffer) = length(cb) == cb.capacity
toarray{T}(cb::CircularBuffer{T}) = T[copy(x) for x in cb]

# Base.first(cb::CircularBuffer) = cb[1]
# Base.last(cb::CircularBuffer) = cb[length(cb)]


text:
###
### Tensor
###

struct Tensor{T,N,ipos,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
    parent::A
end

Base.size(t::Tensor) = size(t.parent)
Base.firstindex(t::Tensor, i...) = firstindex(t.parent, i...)
Base.lastindex(t::Tensor, i...) = lastindex(t.parent, i...)
Base.parent(t::Tensor) = t.parent

tto_indices(t, inds) = expand_tbeginends(t, to_indices(t, (), inds))
Base.@propagate_inbounds function Base.getindex(t::Tensor, inds...)
    t.parent[tto_indices(t, inds)...]
end
Base.@propagate_inbounds function Base.setindex!(t::Tensor, x, inds...)
    t.parent[tto_indices(t, inds)...] = x
end
Base.checkbounds(t::Tensor, inds...) = checkbounds(t.parent, tto_indices(t, inds)...)

###
### IndexPos
###

struct IndexPos{N}
    # TODO: Should this be UInt64 instead?
    x::UInt
end

index_pos(::Tensor{<:Any,<:Any,ipos}) where {ipos} = ipos
function (ipos::IndexPos{N})(a::AbstractArray{T,N}) where {T,N}
    Tensor{T,N,ipos,typeof(a)}(a)
end
(::IndexPos{0})(x::Number) = x

position(ipos::IndexPos, i) = (ipos.x >> (i-1)) % Bool
function deleteat(ipos::IndexPos{N}, i::Int) where {N}
    mask = (UInt(1) << i) - UInt(1)
    IndexPos{N-1}((ipos.x & (mask >> 1)) | ((ipos.x & ~mask) >> 1))
end
function Base.show(io::IO, ipos::IndexPos{N}) where {N}
    print(io, "T\"")
    for i in 1:N
        print(io, position(ipos, i) ? ''' : ',')
    end
    print(io, '\"')
    nothing
end

###
### T"..."
###

macro T_str(s)
    i = 0
    x = UInt(0)
    for c in s
        isspace(c) && continue
        x |= if c === ''' || c === '^'
            UInt(1) << i
        elseif c === ',' || c === '_' || c === '.'
            UInt(0)
        else
            throw(ArgumentError("Index positions can only be ''' and ','"))
        end
        i += 1
    end
    :(IndexPos{$i}($x))
end

###
### T".."[a b; c d]
###

for (typed_f, f) in [
    (:getindex, :vect), (:typed_hcat, :hcat), (:typed_vcat, :vcat), (:typed_hvcat, :hvcat),
]
    @eval Base.$typed_f(ipos::IndexPos, args...) = ipos(Base.$f(args...))
end
Base.getindex(::IndexPos{0}, x) = T""(fill(x))


text:
module NeuralQuantum

# Using statements
using Reexport

using QuantumOptics
using LightGraphs

using Zygote: gradient, forward
using Random: AbstractRNG, MersenneTwister, GLOBAL_RNG
using LinearAlgebra, SparseArrays, Strided
using NNlib

include("IterativeSolvers/minresqlp.jl")
using .MinresQlp

# Quantum Lattices, used to construct lattice hamiltonians (custom package)
@reexport using QuantumLattices

# Optimisers, that will be split in a separate package at some point
include("Optimisers/Optimisers.jl")
using .Optimisers
import .Optimisers: update, update!
export Optimisers

# Imports
import Base: length, UInt, eltype, copy, deepcopy, iterate
import Random: rand!
import QuantumOptics: basis

# Abstract Types
abstract type NeuralNetwork end

abstract type State end
abstract type FiniteBasisState <: State end

abstract type AbstractProblem end
abstract type AbstractSteadyStateProblem <: AbstractProblem end
abstract type HermitianMatrixProblem <: AbstractSteadyStateProblem end
abstract type LRhoSquaredProblem <: AbstractSteadyStateProblem end
abstract type OpenTimeEvolutionProblem <: AbstractSteadyStateProblem end
abstract type OperatorEstimationProblem <: AbstractProblem end

abstract type Sampler end

# Type describing the parallel backend used by a solver.
abstract type ParallelType end
struct NotParallel <: ParallelType end
struct ParallelThreaded <: ParallelType end
export NotParallel, ParallelThreaded

# Universal defines
const STD_REAL_PREC =  Float32

# Base elements
include("base_states.jl")
include("base_derivatives.jl")
include("base_networks.jl")
include("base_cached_networks.jl")
include("base_lookup.jl")
include("base_batched_networks.jl")
include("treelike.jl") # from flux
include("tuple_logic.jl")

# Useful
include("utils/math.jl")

# Basic states for uniform systems
include("States/StateChanges.jl")
include("States/NAryState.jl")
include("States/DoubleState.jl")
include("States/PurifiedState.jl")
include("States/DiagonalStateWrapper.jl")
export local_index
include("States/ModifiedState.jl")
include("States/LUState.jl")
export ModifiedState

# Linear Operators
import Base: +
include("Operators/BaseOperators.jl")
include("Operators/OpConnection.jl")
include("Operators/OpConnectionIndex.jl")
include("Operators/KLocalOperator.jl")
include("Operators/KLocalOperatorSum.jl")
include("Operators/GraphConversion.jl")
export OpConnection
export KLocalOperator, KLocalOperatorSum, KLocalOperatorRow, operators
export row_valdiff, row_valdiff_index, col_valdiff, sites, conn_type
export duplicate

# Neural Networks
include("Networks/utils.jl")
include("Networks/RBMSplit.jl")
include("Networks/NDM.jl")
include("Networks/NDMComplex.jl")
include("Networks/NDMSymm.jl")

# LT
include("Networks/RBMSplitLT.jl")
include("Networks/NDMLT.jl")

# Batch
include("Networks/RBMSplitBatched.jl")

#
include("Networks/ClosedSystems/RBM.jl")

# Problems
export LdagLSparseOpProblem, LRhoSparseSuperopProblem, LdagLProblem, LdagLFullProblem, LdagLSparseSuperopProblem, LdagLSparseSuperopProblemlem, LRhoSparseOpProblem
include("Problems/SteadyStateLindblad/LdagLSparseOpProblem.jl")
include("Problems/SteadyStateLindblad/LdagLSparseSuperopProblem.jl")
include("Problems/SteadyStateLindblad/LRhoKLocalOpProblem.jl")
include("Problems/SteadyStateLindblad/LRhoSparseOpProblem.jl")
include("Problems/SteadyStateLindblad/LRhoSparseSuperopProblem.jl")
const LdagLFullProblem = LRhoSparseSuperopProblem
const LdagLProblem = LdagLSparseOpProblem
const LdagLSparseSuperopProblemlem = LRhoSparseOpProblem

include("Problems/SteadyStateLindblad/build_SteadyStateProblem.jl")

# Hamiltonian problems
include("Problems/Hamiltonian/HamiltonianGSEnergyProblem.jl")
include("Problems/Hamiltonian/build_GroundStateProblem.jl")

# Observables problem
include("Problems/ObservablesProblem.jl")

# gen state
export state, state_lut
include("generate_state.jl")

# Algorithms
abstract type Algorithm end
abstract type EvaluatedAlgorithm end
abstract type EvaluationSamplingCache end
include("Algorithms/base_algorithms.jl")
export EvaluatedNetwork, evaluation_post_sampling!, precondition!, SamplingCache

# SR
include("Algorithms/SR/SR.jl")
include("Algorithms/SR/SampledSRCache.jl")
include("Algorithms/SR/SampledSRCache_L.jl")
include("Algorithms/SR/SR_eval.jl")
# Gradient
include("Algorithms/Gradient/Gradient.jl")
include("Algorithms/Gradient/SampledGradientCache.jl")
include("Algorithms/Gradient/SampledGradientCache_L.jl")
include("Algorithms/Gradient/Gradient_eval.jl")
# Observables
include("Algorithms/Observables/Obs.jl")
include("Algorithms/Observables/Obs_eval.jl")
include("Algorithms/Observables/Obs_ket_eval.jl")

# Sampling
include("Samplers/base_samplers.jl")
include("Samplers/base_samplers_parallel.jl")
export cache, init_sampler!, done, samplenext!
export get_sampler, sampler_list, multithread

include("Samplers/Exact.jl")
include("Samplers/FullSum.jl")
include("Samplers/MCMCSampler.jl")
include("Samplers/MCMCRules/Metropolis.jl")
include("Samplers/MCMCRules/Nagy.jl")

include("base_diffeval.jl")

# other
include("utils/densitymatrix.jl")
include("utils/expectation_values.jl")
include("utils/translational_symm.jl")
#include("utils/logging.jl")
include("utils/loading.jl")

# interface
include("IterativeInterface/BaseIterativeSampler.jl")
include("IterativeInterface/IterativeSampler.jl")
include("IterativeInterface/MTIterativeSampler.jl")
export sample!

end # module


text:
include("./imports.jl")

@testset ExtendedTestSet "runstep, runepoch" begin
    learner = testlearner(coeff = 3)
    struct TestPhase <: Phase end
    state = runstep(learner, TestPhase()) do handle, state
        state.xs, state.ys = learner.data.training[1]
        state.loss = sum(abs.(state.xs - state.ys))
    end
    @test state.xs isa Vector
    @test state.loss isa Number
    @test_nowarn runepoch(learner, TestPhase()) do handle
        throw(FluxTraining.CancelEpochException("test"))
    end
end

@testset ExtendedTestSet "TrainingPhase" begin
    learner = testlearner(coeff = 3)
    runepoch(learner, TrainingPhase()) do _
        state = step!(learner, TrainingPhase(), learner.data.training[1])
        @test state.xs isa Vector
        @test state.ys isa Vector
        @test state.ŷs isa Vector
        @test state.grads isa Zygote.Grads
        @test state.loss isa Number
    end
end

@testset ExtendedTestSet "ValidationPhase" begin
    learner = testlearner(coeff = 3)
    runepoch(learner, ValidationPhase()) do _
        state = step!(learner, ValidationPhase(), learner.data.validation[1])
        @test state.xs isa Vector
        @test state.ys isa Vector
        @test state.ŷs isa Vector
        @test state.loss isa Number
    end
end

@testset ExtendedTestSet "Basic `Learner` convergence" begin
    learner = testlearner(coeff = 3)
    fit!(learner, 5)
    @test learner.model.coeff[1] ≈ 3 atol = 0.1
end


text:
module Grids
using ..Topologies
import ..Metrics, ..Elements
import ..BrickMesh

using MPI
using LinearAlgebra
using KernelAbstractions

export DiscontinuousSpectralElementGrid, AbstractGrid
export dofs_per_element, arraytype, dimensionality, polynomialorders
export referencepoints, min_node_distance, get_z
export EveryDirection, HorizontalDirection, VerticalDirection, Direction

abstract type Direction end
struct EveryDirection <: Direction end
struct HorizontalDirection <: Direction end
struct VerticalDirection <: Direction end
Base.in(::T, ::S) where {T <: Direction, S <: Direction} = T == S

abstract type AbstractGrid{
    FloatType,
    dim,
    polynomialorder,
    numberofDOFs,
    DeviceArray,
} end

dofs_per_element(::AbstractGrid{T, D, N, Np}) where {T, D, N, Np} = Np

polynomialorders(::AbstractGrid{T, dim, N}) where {T, dim, N} = N

dimensionality(::AbstractGrid{T, dim}) where {T, dim} = dim

Base.eltype(::AbstractGrid{T}) where {T} = T

arraytype(::AbstractGrid{T, D, N, Np, DA}) where {T, D, N, Np, DA} = DA

"""
    referencepoints(::AbstractGrid)

Returns the points on the reference element.
"""
referencepoints(::AbstractGrid) = error("needs to be implemented")

"""
    min_node_distance(::AbstractGrid, direction::Direction=EveryDirection() )

Returns an approximation of the minimum node distance in physical space.
"""
function min_node_distance(
    ::AbstractGrid,
    direction::Direction = EveryDirection(),
)
    error("needs to be implemented")
end

# {{{
const _nvgeo = 16
const _ξ1x1,
_ξ2x1,
_ξ3x1,
_ξ1x2,
_ξ2x2,
_ξ3x2,
_ξ1x3,
_ξ2x3,
_ξ3x3,
_M,
_MI,
_MH,
_x1,
_x2,
_x3,
_JcV = 1:_nvgeo
const vgeoid = (
    ξ1x1id = _ξ1x1,
    ξ2x1id = _ξ2x1,
    ξ3x1id = _ξ3x1,
    ξ1x2id = _ξ1x2,
    ξ2x2id = _ξ2x2,
    ξ3x2id = _ξ3x2,
    ξ1x3id = _ξ1x3,
    ξ2x3id = _ξ2x3,
    ξ3x3id = _ξ3x3,
    Mid = _M,
    MIid = _MI,
    MHid = _MH,
    x1id = _x1,
    x2id = _x2,
    x3id = _x3,
    JcVid = _JcV,
)
# JcV is the vertical line integral Jacobian
# The MH terms are for integrating over a plane.
const _nsgeo = 5
const _n1, _n2, _n3, _sM, _vMI = 1:_nsgeo
const sgeoid = (n1id = _n1, n2id = _n2, n3id = _n3, sMid = _sM, vMIid = _vMI)
# }}}

"""
    DiscontinuousSpectralElementGrid(topology; FloatType, DeviceArray,
                                     polynomialorder,
                                     meshwarp = (x...)->identity(x))

Generate a discontinuous spectral element (tensor product,
Legendre-Gauss-Lobatto) grid/mesh from a `topology`, where the order of the
elements is given by `polynomialorder`. `DeviceArray` gives the array type used
to store the data (`CuArray` or `Array`), and the coordinate points will be of
`FloatType`.

The polynomial order can be different in each direction (specified as a
`NTuple`). If only a single integer is specified, then each dimension will use
the same order. If the topology dimension is 3 and the `polynomialorder` has
dimension 2, then the first value will be used for horizontal and the second for
the vertical.

The optional `meshwarp` function allows the coordinate points to be warped after
the mesh is created; the mesh degrees of freedom are orginally assigned using a
trilinear blend of the element corner locations.
"""
struct DiscontinuousSpectralElementGrid{
    T,
    dim,
    N,
    Np,
    DA,
    DAT1,
    DAT2,
    DAT3,
    DAT4,
    DAI1,
    DAI2,
    DAI3,
    TOP,
    TVTK,
} <: AbstractGrid{T, dim, N, Np, DA}
    "mesh topology"
    topology::TOP

    "volume metric terms"
    vgeo::DAT3

    "surface metric terms"
    sgeo::DAT4

    "element to boundary condition map"
    elemtobndy::DAI2

    "volume DOF to element minus side map"
    vmap⁻::DAI3

    "volume DOF to element plus side map"
    vmap⁺::DAI3

    "list of DOFs that need to be received (in neighbors order)"
    vmaprecv::DAI1

    "list of DOFs that need to be sent (in neighbors order)"
    vmapsend::DAI1

    "An array of ranges in `vmaprecv` to receive from each neighbor"
    nabrtovmaprecv::Any

    "An array of ranges in `vmapsend` to send to each neighbor"
    nabrtovmapsend::Any

    "Array of real elements that do not have a ghost element as a neighbor"
    interiorelems::Any

    "Array of real elements that have at least one ghost element as a neighbor"
    exteriorelems::Any

    "Array indicating if a degree of freedom (real or ghost) is active"
    activedofs::Any

    "1-D lgl weights on the device (one for each dimension)"
    ω::DAT1

    "1-D derivative operator on the device (one for each dimension)"
    D::DAT2

    "1-D indefinite integral operator on the device (one for each dimension)"
    Imat::DAT2

    """
    tuple of (x1, x2, x3) to use for vtk output (Needed for the `N = 0` case) in
    other cases these match `vgeo` values
    """
    x_vtk::TVTK

    # Constructor for a tuple of polynomial orders
    function DiscontinuousSpectralElementGrid(
        topology::AbstractTopology{dim};
        polynomialorder,
        FloatType,
        DeviceArray,
        meshwarp::Function = (x...) -> identity(x),
    ) where {dim}

        if polynomialorder isa Integer
            polynomialorder = ntuple(j -> polynomialorder, dim)
        elseif polynomialorder isa NTuple{2} && dim == 3
            polynomialorder =
                (polynomialorder[1], polynomialorder[1], polynomialorder[2])
        end

        @assert dim == length(polynomialorder)

        N = polynomialorder

        (vmap⁻, vmap⁺) = mappings(
            N,
            topology.elemtoelem,
            topology.elemtoface,
            topology.elemtoordr,
        )

        (vmaprecv, nabrtovmaprecv) = commmapping(
            N,
            topology.ghostelems,
            topology.ghostfaces,
            topology.nabrtorecv,
        )
        (vmapsend, nabrtovmapsend) = commmapping(
            N,
            topology.sendelems,
            topology.sendfaces,
            topology.nabrtosend,
        )

        Np = prod(N .+ 1)

        # Create element operators for each polynomial order
        ξω = ntuple(
            j ->
                N[j] == 0 ? Elements.glpoints(FloatType, N[j]) :
                Elements.lglpoints(FloatType, N[j]),
            dim,
        )
        ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)

        Imat = ntuple(
            j -> indefinite_integral_interpolation_matrix(ξ[j], ω[j]),
            dim,
        )
        D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)

        (vgeo, sgeo, x_vtk) =
            computegeometry(topology.elemtocoord, D, ξ, ω, meshwarp)

        @assert Np == size(vgeo, 1)

        activedofs = zeros(Bool, Np * length(topology.elems))
        activedofs[1:(Np * length(topology.realelems))] .= true
        activedofs[vmaprecv] .= true

        # Create arrays on the device
        vgeo = DeviceArray(vgeo)
        sgeo = DeviceArray(sgeo)
        elemtobndy = DeviceArray(topology.elemtobndy)
        vmap⁻ = DeviceArray(vmap⁻)
        vmap⁺ = DeviceArray(vmap⁺)
        vmapsend = DeviceArray(vmapsend)
        vmaprecv = DeviceArray(vmaprecv)
        activedofs = DeviceArray(activedofs)
        ω = DeviceArray.(ω)
        D = DeviceArray.(D)
        Imat = DeviceArray.(Imat)

        # FIXME: There has got to be a better way!
        DAT1 = typeof(ω)
        DAT2 = typeof(D)
        DAT3 = typeof(vgeo)
        DAT4 = typeof(sgeo)
        DAI1 = typeof(vmapsend)
        DAI2 = typeof(elemtobndy)
        DAI3 = typeof(vmap⁻)
        TOP = typeof(topology)
        TVTK = typeof(x_vtk)

        new{
            FloatType,
            dim,
            N,
            Np,
            DeviceArray,
            DAT1,
            DAT2,
            DAT3,
            DAT4,
            DAI1,
            DAI2,
            DAI3,
            TOP,
            TVTK,
        }(
            topology,
            vgeo,
            sgeo,
            elemtobndy,
            vmap⁻,
            vmap⁺,
            vmaprecv,
            vmapsend,
            nabrtovmaprecv,
            nabrtovmapsend,
            DeviceArray(topology.interiorelems),
            DeviceArray(topology.exteriorelems),
            activedofs,
            ω,
            D,
            Imat,
            x_vtk,
        )
    end
end

"""
    referencepoints(::DiscontinuousSpectralElementGrid)

Returns the 1D interpolation points used for the reference element.
"""
function referencepoints(
    ::DiscontinuousSpectralElementGrid{FT, dim, N},
) where {FT, dim, N}
    ξω = ntuple(
        j ->
            N[j] == 0 ? Elements.glpoints(FT, N[j]) :
            Elements.lglpoints(FT, N[j]),
        dim,
    )
    ξ, _ = ntuple(j -> map(x -> x[j], ξω), 2)
    return ξ
end

"""
    min_node_distance(::DiscontinuousSpectralElementGrid,
                      direction::Direction=EveryDirection()))

Returns an approximation of the minimum node distance in physical space along
the reference coordinate directions.  The direction controls which reference
directions are considered.
"""
function min_node_distance(
    grid::DiscontinuousSpectralElementGrid{T, dim, N},
    direction::Direction = EveryDirection(),
) where {T, dim, N}
    topology = grid.topology
    nrealelem = length(topology.realelems)

    if nrealelem > 0
        Nq = N .+ 1
        Np = prod(Nq)
        device = grid.vgeo isa Array ? CPU() : CUDADevice()
        min_neighbor_distance = similar(grid.vgeo, Np, nrealelem)
        event = Event(device)
        event = kernel_min_neighbor_distance!(device, min(Np, 1024))(
            Val(N),
            Val(dim),
            direction,
            min_neighbor_distance,
            grid.vgeo,
            topology.realelems;
            ndrange = (Np * nrealelem),
            dependencies = (event,),
        )
        wait(device, event)
        locmin = minimum(min_neighbor_distance)
    else
        locmin = typemax(T)
    end

    MPI.Allreduce(locmin, min, topology.mpicomm)
end

"""
    get_z(grid; z_scale = 1, rm_dupes = false)

Get the Gauss-Lobatto points along the Z-coordinate.

 - `grid`: DG grid
 - `z_scale`: multiplies `z-coordinate`
 - `rm_dupes`: removes duplicate Gauss-Lobatto points
"""
function get_z(
    grid::DiscontinuousSpectralElementGrid{T, dim, N};
    z_scale = 1,
    rm_dupes = false,
) where {T, dim, N}
    # Assumes same polynomial orders in all horizontal directions
    @assert dim < 3 || N[1] == N[2]
    Nhoriz = N[1]
    Nvert = N[end]
    Nph = (Nhoriz + 1)^2
    Np = Nph * (Nvert + 1)
    if rm_dupes
        ijk_range = (1:Nph:(Np - Nph))
        vgeo = Array(grid.vgeo)
        z = reshape(vgeo[ijk_range, _x3, :], :)
        z = [z..., vgeo[Np, _x3, end]]
        return z * z_scale
    else
        ijk_range = (1:Nph:Np)
        z = Array(reshape(grid.vgeo[ijk_range, _x3, :], :))
        return z * z_scale
    end
    return reshape(grid.vgeo[(1:Nph:Np), _x3, :], :) * z_scale
end

function Base.getproperty(G::DiscontinuousSpectralElementGrid, s::Symbol)
    if s ∈ keys(vgeoid)
        vgeoid[s]
    elseif s ∈ keys(sgeoid)
        sgeoid[s]
    else
        getfield(G, s)
    end
end

function Base.propertynames(G::DiscontinuousSpectralElementGrid)
    (
        fieldnames(DiscontinuousSpectralElementGrid)...,
        keys(vgeoid)...,
        keys(sgeoid)...,
    )
end

# {{{ mappings
"""
    mappings(N, elemtoelem, elemtoface, elemtoordr)

This function takes in a tuple of polynomial orders `N` and parts of a topology
(as returned from `connectmesh`) and returns index mappings for the element
surface flux computation. The returned `Tuple` contains:

 - `vmap⁻` an array of linear indices into the volume degrees of freedom where
   `vmap⁻[:,f,e]` are the degrees of freedom indices for face `f` of element
    `e`.

 - `vmap⁺` an array of linear indices into the volume degrees of freedom where
   `vmap⁺[:,f,e]` are the degrees of freedom indices for the face neighboring
   face `f` of element `e`.
"""
function mappings(N, elemtoelem, elemtoface, elemtoordr)
    nfaces, nelem = size(elemtoelem)

    d = div(nfaces, 2)
    Nq = N .+ 1
    # number of points in the element
    Np = prod(Nq)

    # Compute the maximum number of points on a face
    Nfp = div.(Np, Nq)

    # linear index for each direction, e.g., (i, j, k) -> n
    p = reshape(1:Np, ntuple(j -> Nq[j], d))

    # fmask[f] -> returns an array of all degrees of freedom on face f
    fmask = if d == 1
        (
            p[1:1],    # Face 1
            p[Nq[1]:Nq[1]], # Face 2
        )
    elseif d == 2
        (
            p[1, :][:],     # Face 1
            p[Nq[1], :][:], # Face 2
            p[:, 1][:],     # Face 3
            p[:, Nq[2]][:], # Face 4
        )
    elseif d == 3
        (
            p[1, :, :][:],     # Face 1
            p[Nq[1], :, :][:], # Face 2
            p[:, 1, :][:],     # Face 3
            p[:, Nq[2], :][:], # Face 4
            p[:, :, 1][:],     # Face 5
            p[:, :, Nq[3]][:], # Face 6
        )
    else
        error("unknown dimensionality")
    end

    # Create a map from Cartesian face dof number to linear face dof numbering
    # inds[face][i, j] -> n
    inds = ntuple(
        f -> dropdims(
            LinearIndices(ntuple(j -> j == cld(f, 2) ? 1 : Nq[j], d));
            dims = cld(f, 2),
        ),
        nfaces,
    )

    # Use the largest possible storage
    vmap⁻ = fill!(similar(elemtoelem, maximum(Nfp), nfaces, nelem), 0)
    vmap⁺ = fill!(similar(elemtoelem, maximum(Nfp), nfaces, nelem), 0)

    for e1 in 1:nelem, f1 in 1:nfaces
        e2 = elemtoelem[f1, e1]
        f2 = elemtoface[f1, e1]
        o2 = elemtoordr[f1, e1]
        d1, d2 = cld(f1, 2), cld(f2, 2)

        # Check to make sure the dof grid is conforming
        @assert Nfp[d1] == Nfp[d2]

        # Always pull out minus side without any flips / rotations
        vmap⁻[1:Nfp[d1], f1, e1] .= Np * (e1 - 1) .+ fmask[f1][1:Nfp[d1]][:]

        # Orientation codes defined in BrickMesh.jl (arbitrary numbers in 3D)
        if o2 == 1 # Neighbor oriented same as minus
            vmap⁺[1:Nfp[d1], f1, e1] .= Np * (e2 - 1) .+ fmask[f2][1:Nfp[d1]][:]
        elseif d == 3 && o2 == 3 # Neighbor fliped in first index
            vmap⁺[1:Nfp[d1], f1, e1] =
                Np * (e2 - 1) .+ fmask[f2][inds[f2][end:-1:1, :]][:]
        else
            error("Orientation '$o2' with dim '$d' not supported yet")
        end
    end

    (vmap⁻, vmap⁺)
end
# }}}

"""
   commmapping(N, commelems, commfaces, nabrtocomm)

This function takes in a tuple of polynomial orders `N` and parts of a mesh (as
returned from `connectmesh` such as `sendelems`, `sendfaces`, and `nabrtosend`)
and returns index mappings for the element surface flux parallel communcation.
The returned `Tuple` contains:

 - `vmapC` an array of linear indices into the volume degrees of freedom to be
   communicated.

 - `nabrtovmapC` a range in `vmapC` to communicate with each neighbor.
"""
function commmapping(N, commelems, commfaces, nabrtocomm)
    nface, nelem = size(commfaces)

    @assert nelem == length(commelems)

    d = div(nface, 2)
    Nq = N .+ 1
    Np = prod(Nq)

    vmapC = similar(commelems, nelem * Np)
    nabrtovmapC = similar(nabrtocomm)

    i = 1
    e = 1
    for neighbor in 1:length(nabrtocomm)
        rbegin = i
        for ne in nabrtocomm[neighbor]
            ce = commelems[ne]

            # Whole element sending
            # for n = 1:Np
            #   vmapC[i] = (ce-1)*Np + n
            #   i += 1
            # end

            CI = CartesianIndices(ntuple(j -> 1:Nq[j], d))
            for (ci, li) in zip(CI, LinearIndices(CI))
                addpoint = false
                for j in 1:d
                    addpoint |=
                        (commfaces[2 * (j - 1) + 1, e] && ci[j] == 1) ||
                        (commfaces[2 * (j - 1) + 2, e] && ci[j] == Nq[j])
                end

                if addpoint
                    vmapC[i] = (ce - 1) * Np + li
                    i += 1
                end
            end

            e += 1
        end
        rend = i - 1

        nabrtovmapC[neighbor] = rbegin:rend
    end

    resize!(vmapC, i - 1)

    (vmapC, nabrtovmapC)
end

# Compute geometry
function computegeometry_fvm(elemtocoord, D, ξ, ω, meshwarp)
    FT = eltype(D[1])
    dim = length(D)
    nface = 2dim
    nelem = size(elemtocoord, 3)

    Nq = ntuple(j -> size(D[j], 1), dim)
    Np = prod(Nq)
    Nfp = div.(Np, Nq)

    Nq_N1 = max.(Nq, 2)
    Np_N1 = prod(Nq_N1)
    Nfp_N1 = div.(Np_N1, Nq_N1)

    # First we compute the geometry with all the N = 0 dimension set to N = 1
    # so that we can later compute the geometry for the case N = 0, as the
    # average of two N = 1 cases
    ξ1, ω1 = Elements.lglpoints(FT, 1)
    D1 = Elements.spectralderivative(ξ1)
    D_N1 = ntuple(j -> Nq[j] == 1 ? D1 : D[j], dim)
    ξ_N1 = ntuple(j -> Nq[j] == 1 ? ξ1 : ξ[j], dim)
    ω_N1 = ntuple(j -> Nq[j] == 1 ? ω1 : ω[j], dim)
    (vgeo_N1, sgeo_N1, x_vtk) =
        computegeometry(elemtocoord, D_N1, ξ_N1, ω_N1, meshwarp)

    # Sort out the vgeo terms
    @views begin
        vgeo_N1_flds =
            ntuple(fld -> reshape(vgeo_N1[:, fld, :], Nq_N1..., nelem), _nvgeo)
    end

    # Allocate the storage for N = 0 volume metrics
    vgeo = zeros(FT, Np, _nvgeo, nelem)

    # Counter to make sure we got all the vgeo terms
    num_vgeo_handled = 0

    X = ntuple(j -> (@view vgeo[:, _x1 + j - 1, :]), dim)
    Metrics.creategrid!(X..., elemtocoord, ξ...)
    x1 = @view vgeo[:, _x1, :]
    x2 = @view vgeo[:, _x2, :]
    x3 = @view vgeo[:, _x3, :]
    @inbounds for j in 1:length(x1)
        (x1[j], x2[j], x3[j]) = meshwarp(x1[j], x2[j], x3[j])
    end

    num_vgeo_handled += 3

    @views begin
        # _M should be a sum
        vgeo[:, _M, :][:] .= sum(vgeo_N1_flds[_M], dims = findall(Nq .== 1))[:]
        num_vgeo_handled += 1

        # need to recompute _MI
        vgeo[:, _MI, :] = 1 ./ vgeo[:, _M, :]
        num_vgeo_handled += 1

        # coordinates should just be averages
        avg_den = 2 .^ sum(Nq .== 1)
        for fld in (_JcV,)
            vgeo[:, fld, :] =
                sum(vgeo_N1_flds[fld], dims = findall(Nq .== 1))[:] / avg_den
            num_vgeo_handled += 1
        end

        # For the metrics it is J * ξixk we approximate so multiply and divide the
        # mass matrix (which has the Jacobian determinant and the proper averaging
        # due to the quadrature weights)
        M_N1 = vgeo_N1_flds[_M]
        MI = vgeo[:, _MI, :]
        for fld in
            (_ξ1x1, _ξ2x1, _ξ3x1, _ξ1x2, _ξ2x2, _ξ3x2, _ξ1x3, _ξ2x3, _ξ3x3)
            vgeo[:, fld, :] =
                sum(M_N1 .* vgeo_N1_flds[fld], dims = findall(Nq .== 1))[:] .*
                MI[:]
            num_vgeo_handled += 1
        end

        # compute MH and JvC
        horizontal_metrics(vgeo, Nq, ω)
        num_vgeo_handled += 1

        # Make sure we handled all the vgeo terms
        @assert _nvgeo == num_vgeo_handled
    end

    # Sort out the sgeo terms
    @views begin
        sgeo = zeros(FT, _nsgeo, maximum(Nfp), nface, nelem)

        # for the volume inverse mass matrix
        MI = vgeo[:, _MI, :]
        p = reshape(1:Np, Nq)
        if dim == 1
            fmask = (p[1:1], p[Nq[1]:Nq[1]])
        elseif dim == 2
            fmask = (p[1, :][:], p[Nq[1], :][:], p[:, 1][:], p[:, Nq[2]][:])
        elseif dim == 3
            fmask = (
                p[1, :, :][:],
                p[Nq[1], :, :][:],
                p[:, 1, :][:],
                p[:, Nq[2], :][:],
                p[:, :, 1][:],
                p[:, :, Nq[3]][:],
            )
        end
        for d in 1:dim
            for f in (2d - 1):(2d)
                # number of points matches means that we keep all the data
                # (N = 0 is not on the face)
                if Nfp[d] == Nfp_N1[d]
                    sgeo[:, 1:Nfp[d], f, :] .= sgeo_N1[:, 1:Nfp[d], f, :]

                    # Volume inverse mass will be wrong so reset it
                    sgeo[_vMI, 1:Nfp[d], f, :] .= MI[fmask[f], :]
                else
                    # Counter to make sure we got all the sgeo terms
                    num_sgeo_handled = 0

                    # sum to get sM
                    Nq_f = (Nq[1:(d - 1)]..., Nq[(d + 1):dim]...)
                    Nq_f_N1 = (Nq_N1[1:(d - 1)]..., Nq_N1[(d + 1):dim]...)
                    sM_N1 = reshape(
                        sgeo_N1[_sM, 1:Nfp_N1[d], f, :],
                        Nq_f_N1...,
                        nelem,
                    )
                    sgeo[_sM, 1:Nfp[d], f, :][:] .=
                        sum(sM_N1, dims = findall(Nq_f .== 1))[:]
                    num_sgeo_handled += 1

                    # Normals (like metrics in the volume) need to be computed
                    # scaled by surface Jacobian which we can do with the
                    # surface mass matrices
                    sM = sgeo[_sM, 1:Nfp[d], f, :]
                    for fld in (_n1, _n2, _n3)
                        fld_N1 = reshape(
                            sgeo_N1[fld, 1:Nfp_N1[d], f, :],
                            Nq_f_N1...,
                            nelem,
                        )
                        sgeo[fld, 1:Nfp[d], f, :][:] .=
                            sum(sM_N1 .* fld_N1, dims = findall(Nq_f .== 1))[:] ./
                            sM[:]
                        num_sgeo_handled += 1
                    end

                    # set the volume inverse mass matrix
                    sgeo[_vMI, 1:Nfp[d], f, :] .= MI[fmask[f], :]
                    num_sgeo_handled += 1

                    # Make sure we handled all the vgeo terms
                    @assert _nsgeo == num_sgeo_handled
                end
            end
        end
    end

    (vgeo, sgeo, x_vtk)
end

function computegeometry(elemtocoord, D, ξ, ω, meshwarp)
    dim = length(D)
    nface = 2dim
    nelem = size(elemtocoord, 3)

    Nq = ntuple(j -> size(D[j], 1), dim)

    # Compute metric terms for FVM
    if any(Nq .== 1)
        return computegeometry_fvm(elemtocoord, D, ξ, ω, meshwarp)
    end

    Np = prod(Nq)
    Nfp = div.(Np, Nq)

    FT = eltype(D[1])

    vgeo = zeros(FT, Np, _nvgeo, nelem)
    sgeo = zeros(FT, _nsgeo, maximum(Nfp), nface, nelem)

    (
        #! format: off
        ξ1x1, ξ2x1, ξ3x1, ξ1x2, ξ2x2, ξ3x2, ξ1x3, ξ2x3, ξ3x3,
        MJ, MJI, MHJH,
        x1, x2, x3,
        JcV,
       #! format: on
    ) = ntuple(j -> (@view vgeo[:, j, :]), _nvgeo)
    J = similar(x1)
    (n1, n2, n3, sMJ, vMJI) = ntuple(j -> (@view sgeo[j, :, :, :]), _nsgeo)
    sJ = similar(sMJ)

    X = ntuple(j -> (@view vgeo[:, _x1 + j - 1, :]), dim)
    Metrics.creategrid!(X..., elemtocoord, ξ...)

    @inbounds for j in 1:length(x1)
        (x1[j], x2[j], x3[j]) = meshwarp(x1[j], x2[j], x3[j])
    end

    # Compute the metric terms
    p = reshape(1:Np, Nq)
    if dim == 1
        Metrics.computemetric!(x1, J, JcV, ξ1x1, sJ, n1, D...)
        fmask = (p[1:1], p[Nq[1]:Nq[1]])
    elseif dim == 2
        Metrics.computemetric!(
            #! format: off
            x1, x2,
            J, JcV,
            ξ1x1, ξ2x1, ξ1x2, ξ2x2,
            sJ,
            n1, n2,
            D...,
            #! format: on
        )
        fmask = (p[1, :][:], p[Nq[1], :][:], p[:, 1][:], p[:, Nq[2]][:])
    elseif dim == 3
        Metrics.computemetric!(
            #! format: off
            x1, x2, x3,
            J, JcV,
            ξ1x1, ξ2x1, ξ3x1, ξ1x2, ξ2x2, ξ3x2, ξ1x3, ξ2x3, ξ3x3,
            sJ,
            n1, n2, n3,
            D...,
            #! format: on
        )
        fmask = (
            p[1, :, :][:],
            p[Nq[1], :, :][:],
            p[:, 1, :][:],
            p[:, Nq[2], :][:],
            p[:, :, 1][:],
            p[:, :, Nq[3]][:],
        )
    end

    # since `ξ1` is the fastest dimension and `ξdim` the slowest the tensor
    # product order is reversed
    M = kron(1, reverse(ω)...)
    MJ .= M .* J
    MJI .= 1 ./ MJ
    for d in 1:dim
        for f in (2d - 1):(2d)
            vMJI[1:Nfp[d], f, :] .= MJI[fmask[f], :]
        end
    end

    sM = fill!(similar(sJ, maximum(Nfp), nface), NaN)
    for d in 1:dim
        for f in (2d - 1):(2d)
            ωf = ntuple(j -> ω[mod1(d + j, dim)], dim - 1)
            # Because of the `mod1` this face is already flipped
            if !(dim == 3 && d == 2)
                ωf = reverse(ωf)
            end
            sM[1:Nfp[d], f] .= dim > 1 ? kron(1, ωf...) : one(FT)
        end
    end
    sMJ .= sM .* sJ

    # compute MH and JvC
    horizontal_metrics(vgeo, Nq, ω)

    # This is mainly done to support FVM plotting when N=0 (since we need cell
    # edge values)
    x_vtk = (vgeo[:, _x1, :], vgeo[:, _x2, :], vgeo[:, _x3, :])

    (vgeo, sgeo, x_vtk)
end

function horizontal_metrics(vgeo, Nq, ω)
    dim = length(Nq)

    MH = dim == 1 ? 1 : kron(ones(1, Nq[dim]), reverse(ω[1:(dim - 1)])...)[:]
    M = kron(1, reverse(ω)...)[:]

    (
        #! format: off
        ξ1x1, ξ2x1, ξ3x1, ξ1x2, ξ2x2, ξ3x2, ξ1x3, ξ2x3, ξ3x3,
        MJ, MJI, MHJH,
        x1, x2, x3,
        JcV,
       #! format: on
    ) = ntuple(j -> (@view vgeo[:, j, :]), _nvgeo)
    J = MJ ./ M[:]

    # Compute |r'(ξ3)| for vertical line integrals
    if dim == 1
        MHJH .= 1
    elseif dim == 2
        map!(MHJH, J, ξ2x1, ξ2x2) do J, ξ2x1, ξ2x2
            hypot(J * ξ2x1, J * ξ2x2)
        end
        MHJH .= MH .* MHJH

    elseif dim == 3
        map!(MHJH, J, ξ3x1, ξ3x2, ξ3x3) do J, ξ3x1, ξ3x2, ξ3x3
            hypot(J * ξ3x1, J * ξ3x2, J * ξ3x3)
        end
        MHJH .= MH .* MHJH
    else
        error("dim $dim not implemented")
    end
end

"""
    indefinite_integral_interpolation_matrix(r, ω)

Given a set of integration points `r` and integration weights `ω` this computes
a matrix that will compute the indefinite integral of the (interpolant) of a
function and evaluate the indefinite integral at the points `r`.

Namely, let
```math
    q(ξ) = ∫_{ξ_{0}}^{ξ} f(ξ') dξ'
```
then we have that
```
I∫ * f.(r) = q.(r)
```
where `I∫` is the integration and interpolation matrix defined by this function.

!!! note

    The integration is done using the provided quadrature weight, so if these
    cannot integrate `f(ξ)` exactly, `f` is first interpolated and then
    integrated using quadrature. Namely, we have that:
    ```math
        q(ξ) = ∫_{ξ_{0}}^{ξ} I(f(ξ')) dξ'
    ```
    where `I` is the interpolation operator.

"""
function indefinite_integral_interpolation_matrix(r, ω)
    Nq = length(r)

    I∫ = similar(r, Nq, Nq)
    # first value is zero
    I∫[1, :] .= Nq == 1 ? ω[1] : 0

    # barycentric weights for interpolation
    wbary = Elements.baryweights(r)

    # Compute the interpolant of the indefinite integral
    for n in 2:Nq
        # grid from first dof to current point
        rdst = (1 .- r) / 2 * r[1] + (1 .+ r) / 2 * r[n]
        # interpolation matrix
        In = Elements.interpolationmatrix(r, rdst, wbary)
        # scaling from LGL to current of the interval
        Δ = (r[n] - r[1]) / 2
        # row of the matrix we have computed
        I∫[n, :] .= (Δ * ω' * In)[:]
    end
    I∫
end
# }}}

using KernelAbstractions.Extras: @unroll

using StaticArrays

const _x1 = Grids._x1
const _x2 = Grids._x2
const _x3 = Grids._x3

@doc """
    kernel_min_neighbor_distance!(::Val{N}, ::Val{dim}, direction,
                             min_neighbor_distance, vgeo, topology.realelems)

Computational kernel: Computes the minimum physical distance between node
neighbors within an element.

The `direction` in the reference element controls which nodes are considered
neighbors.
""" kernel_min_neighbor_distance!
@kernel function kernel_min_neighbor_distance!(
    ::Val{N},
    ::Val{dim},
    direction,
    min_neighbor_distance,
    vgeo,
    elems,
) where {N, dim}

    @uniform begin
        FT = eltype(min_neighbor_distance)
        Nq = N .+ 1
        Np = prod(Nq)

        if direction isa EveryDirection
            mininξ = (true, true, true)
        elseif direction isa HorizontalDirection
            mininξ = (true, dim == 2 ? false : true, false)
        elseif direction isa VerticalDirection
            mininξ = (false, dim == 2 ? true : false, dim == 2 ? false : true)
        end

        @inbounds begin
            Nq1 = Nq[1]
            Nq2 = Nq[2]
            Nqk = dim == 2 ? 1 : Nq[end]
            mininξ1 = mininξ[1]
            mininξ2 = mininξ[2]
            mininξ3 = mininξ[3]
        end
    end

    I = @index(Global, Linear)
    e = (I - 1) ÷ Np + 1
    ijk = (I - 1) % Np + 1

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

    md = typemax(FT)

    x = SVector(vgeo[ijk, _x1, e], vgeo[ijk, _x2, e], vgeo[ijk, _x3, e])

    if mininξ1
        @unroll for î in (i - 1, i + 1)
            if 1 ≤ î ≤ Nq1
                îjk = î + Nq1 * (j - 1) + Nq1 * Nq2 * (k - 1)
                x̂ = SVector(
                    vgeo[îjk, _x1, e],
                    vgeo[îjk, _x2, e],
                    vgeo[îjk, _x3, e],
                )
                md = min(md, norm(x - x̂))
            end
        end
    end

    if mininξ2
        @unroll for ĵ in (j - 1, j + 1)
            if 1 ≤ ĵ ≤ Nq2
                iĵk = i + Nq1 * (ĵ - 1) + Nq1 * Nq2 * (k - 1)
                x̂ = SVector(
                    vgeo[iĵk, _x1, e],
                    vgeo[iĵk, _x2, e],
                    vgeo[iĵk, _x3, e],
                )
                md = min(md, norm(x - x̂))
            end
        end
    end

    if mininξ3
        @unroll for k̂ in (k - 1, k + 1)
            if 1 ≤ k̂ ≤ Nqk
                ijk̂ = i + Nq1 * (j - 1) + Nq1 * Nq2 * (k̂ - 1)
                x̂ = SVector(
                    vgeo[ijk̂, _x1, e],
                    vgeo[ijk̂, _x2, e],
                    vgeo[ijk̂, _x3, e],
                )
                md = min(md, norm(x - x̂))
            end
        end
    end

    min_neighbor_distance[ijk, e] = md
end

end # module

max_stars_repo_name: jm-c/CLIMA

max_issues_repo_name: jm-c/CLIMA

max_forks_repo_name: jm-c/CLIMA


text:
using DSP

"""
    dbv(x) = 20 × log₁₀(|x|)

The dB equivalent of the voltage `x`
"""
function dbv(x)
    return amp2db.(abs.(x))
end

max_stars_repo_name: Mahmoud-Kharsa/DeltaSigma

max_issues_repo_name: Mahmoud-Kharsa/DeltaSigma

max_forks_repo_name: Mahmoud-Kharsa/DeltaSigma


text:
# Test index sets and vector gather/scatter

@testset "IS{$ST}" begin
  let i = IS(ST, 10:-2:1), j = sort(i)
    @test !issorted(i)
    @test issorted(j)
    @test i == j
    @test i ∪ j == i
    @test Vector{Int}(i) == [10:-2:1;]
    @test Vector{Int}(j) == [2:2:10;]
    @test length(i) == 5
    @test lengthlocal(i) == 5
    @test minimum(i) == 2
    @test maximum(i) == 10
    @test Set(setdiff(i, IS(ST, 1:4))) == Set(6:2:10)
    @test Set(union(i, IS(ST, [1,2,3,4]))) == Set(2:2:10) ∪ Set(1:4)
  end

  let x = Vec(ST[1,17,24,2]), y = Vec(ST, 2)
    @test scatter!(x, 2:3, y, 1:2) == [17,24]
  end

  let x = Vec(ST[1,17,24,2]), y = Vec(ST, 2)
    VS = VecScatter(x, IS(ST, 2:3, comm=comm(x)), y, IS(ST, 1:2, comm=comm(y)))
    @test scatter!(copy(VS),x,y) == [17,24]
  end

  # Commented out this test. It fails with PETSc 3.13.4
  # let
  #   idx = 2:4
  #   is = IS(ST, idx)
  #   bs = 2
  #   set_blocksize(is, bs)
  #   @test get_blocksize(is) == bs
  # end

  let
    idx = 2:5
    bs = 3
    is = ISBlock(ST, 3, idx)
    @test get_blocksize(is) == bs
    @test length(is) == length(idx)*bs
    @test Vector{PetscInt}(is) == [4:15;]
  end

  let
     is = ISLocalToGlobalMapping(ST,[4,5,6])
     petscview(is)
  end

end


text:
#=
Here an optimization oracle is a linear optimization oracle for the set S, i.e., it solves
    z^\ast(c) :=    min c^Tw
                    s.t. w \in S
The solution of the above problem returned by the oracle is referred to as w^\ast(\cdot) in
the paper

A specific implementation of an optimization oracle should take a single input argument c,
and should return the pair (z^\ast(c), w^\ast(c))

This file contains the code for the shortest path oracle, which relies on JuMP
=#

using Gurobi, JuMP, Clp, LightGraphs, SparseArrays, LinearAlgebra

"""
    sp_flow_jump_setup(sources, destinations, start_node, end_node; solver=:Gurobi, gurobiEnv = Gurobi.Env())

Sets up a shortest path oracle using JuMP. Sets up constraints once then changes only the
objective coefficients dynamically. Takes as input a directed graph via a list of edges in the format of
(sources[i], destinations[i]) for i = 1,...,num_edges. Nodes are assumed to be integers in the range
1,...,num_nodes.

# Arguments
- `sources`:  list of sources for each edge, length of this vector is number of edges
- `destinations`:  list of destinations of each edge, length of this vector is number of edges
- `start_node`:  start node for the shortest path problem
- `end_node`:  end node for the shortest path problem
- `solver`:  either `:Gurobi` or `:Clp`
- `gurobiEnv`:  optionally pass in an existing Gurobi enviornemnet
"""
function sp_flow_jump_setup(sources, destinations, start_node, end_node; solver=:Gurobi, gurobiEnv = Gurobi.Env(), small_coefficient_tolerance = 0.01)
    nodes = unique(union(sources, destinations))
    n_nodes = length(nodes)
    n_edges = length(sources)

    # Need to hard code incidence matrix!
    # A_mat = zeros(n_nodes, n_edges)
    # for i = 1:n_edges
    #     A_mat[sources[i], i] = -1
    #     A_mat[destinations[i], i] = 1
    # end

    # Hard code sparse node-edge incidence matrix!
    I_vec = [sources; destinations]
    J_vec = [collect(1:n_edges); collect(1:n_edges)]
    V_vec = [-ones(n_edges); ones(n_edges)]
    A_mat = sparse(I_vec, J_vec, V_vec)

    # Set up RHS
    b_vec = sparse(zeros(n_nodes))
    b_vec[start_node] = -1
    b_vec[end_node] = 1

    # Remove sparse
    A_mat = round.(Int, Matrix(A_mat))
    b_vec = round.(Int, Vector(b_vec))

    # Set up JuMP model
    if solver == :Gurobi
        mod = Model(with_optimizer(Gurobi.Optimizer, gurobiEnv))
    elseif solver == :Clp
        mod = Model(with_optimizer(Clp.Optimizer))
    else
        error("Not a valid solver, either :Clp or :Gurobi")
    end

    @variable(mod, 0 <= w[1:(n_edges)] <= 1) # bound constraints
    @constraint(mod, A_mat*w .== b_vec)

    # Set up local oracle function
    function local_sp_oracle_jump(c::Vector{Float64})
        d = length(c)
        for i = 1:d
            if abs(c[i]) < small_coefficient_tolerance
                c[i] = 0
            end
        end

        @objective(mod, Min, dot(c, w))
        optimize!(mod)
        z_ast = objective_value(mod)
        w_ast = value.(w)

        return (z_ast, w_ast)
    end

    return c -> local_sp_oracle_jump(c)
end

"""
    convert_grid_to_list(dim1, dim2)

Utility function to convert a dim1 x dim2 grid graph into a (sources[i], detinations[i])
format.
"""
function convert_grid_to_list(dim1, dim2)
    g = Grid([dim1, dim2])
    sources = Int64[]
    destinations = Int64[]

    for e in edges(g)
        push!(sources, src(e))
        push!(destinations, dst(e))

        # Make it directed by commenting out the below
        #push!(sources, dst(e))
        #push!(destinations, src(e))
    end

    return sources, destinations
end

max_stars_repo_name: paulgrigas/SmartPredictThenOptimize

max_issues_repo_name: paulgrigas/SmartPredictThenOptimize

max_forks_repo_name: paulgrigas/SmartPredictThenOptimize


text:
# Rational approximations to generalized hypergeometric functions
# using Weniger's sequence transformation

# ₘFₙ(α;β;z)
# γ ∉ ℕ
function wenigerpFq(α::AbstractVector{T1}, β::AbstractVector{T2}, z::T3) where {T1, T2, T3}
    T = promote_type(T1, T2, T3)
    if norm(z) < eps(real(T)) || norm(prod(α)) < eps(real(T))
        return one(T)
    end
    γ = T(3)/2
    ζ = inv(z)
    p = length(α)
    q = length(β)
    r = max(p, q)+3
    ρ = max(p, q)+1
    C = zeros(T, r)
    C[1] = one(T)
    Cρ = zeros(T, ρ+2)
    Cρ[ρ+2] = one(T)
    for s in ρ:-1:0
        Cρ[s+1] = -(s+1)*Cρ[s+2]/(ρ+1-s)
    end
    C1 = zeros(T, ρ+1)
    C2 = zeros(T, ρ+2)
    C2[ρ+2] = one(T)/pochhammer(γ-ρ-2, ρ+2)
    C3 = zeros(T, ρ+2)
    C3[ρ+1] = one(T)/pochhammer(γ-ρ-1, ρ+2)
    C3[ρ+2] = one(T)/pochhammer(γ-ρ, ρ+2)
    P = zeros(T, p+2)
    t = one(T)
    for j in 1:p
        t *= α[j]+1
    end
    P[1] = γ*t
    Q = zeros(T, q+2)
    t = one(T)
    for j in 1:q
        t *= β[j]+1
    end
    Q[1] = 2t
    N = zeros(T, r+1)
    ΔN = zeros(T, r+1)
    ΔNold = zeros(T, r+1)
    D = zeros(T, r+1)
    ΔD = zeros(T, r+1)
    ΔDold = zeros(T, r+1)
    R = zeros(T, r+1)
    N[r+1] = prod(β)*ζ/prod(α)/(γ-1)
    ΔN[r] = N[r+1]/pochhammer(γ-ρ-1, ρ)
    D[r+1] = prod(β)*ζ/prod(α)/(γ-1)
    ΔD[r] = D[r+1]/pochhammer(γ-ρ-1, ρ)
    R[r+1] = N[r+1]/D[r+1]
    k = 0
    while (abs(R[r+1]-R[r]) > 10*abs(R[r+1])*eps(real(T)) && k < 10_000) || k < r
        for j in 1:r
            N[j] = N[j+1]
            D[j] = D[j+1]
            R[j] = R[j+1]
        end
        t1 = zero(T)
        for j in 0:min(k, q+1)
            t1 += C[j+1]*Q[j+1]*ΔN[r-j]
        end
        if k ≤ ρ
            for j in 0:k
                t2 = one(T)
                for i in 1:q
                    t2 *= β[i]+j+1
                end
                t2 *= j+2
                t1 += C[j+1]*(-one(T))^(k-j)*pochhammer(j+γ, k-ρ-1)*t2
            end
        end
        t2 = zero(T)
        s2 = zero(T)
        for s in max(0,ρ+1-k):ρ
            s2 += Cρ[s+1]*C1[s+1]*(N[r-ρ+s]+(γ+k-ρ+s-2)*N[r-ρ+s-1])
        end
        s2 += (γ+k-1)*N[r]/pochhammer(γ+2k-ρ-1, ρ+2)
        t2 += P[1]*s2
        s2 = zero(T)
        for s in max(0,ρ+1-k):ρ+1
            s2 += Cρ[s+1]*C2[s+1]*(γ+2k-2ρ+2s-3)*N[r-ρ+s-1]
        end
        ΔNold[r+1] = s2
        for j in 1:min(k, p+1)
            t2 += C[j+1]*P[j+1]*(ΔNold[r+2-j] + ΔNold[r+1-j])
        end
        N[r+1] = (ζ*t1-t2)/(P[1]/pochhammer(γ+2k-ρ-1, ρ+2))
        s1 = zero(T)
        for s in max(0,ρ-k):ρ+1
            s1 += Cρ[s+1]*C3[s+1]*(γ+2k-2ρ+2s-1)*N[r-ρ+s]
        end
        ΔN[r+1] = s1
        t1 = zero(T)
        for j in 0:min(k, q+1)
            t1 += C[j+1]*Q[j+1]*ΔD[r-j]
        end
        t2 = zero(T)
        s2 = zero(T)
        for s in max(0,ρ+1-k):ρ
            s2 += Cρ[s+1]*C1[s+1]*(D[r-ρ+s]+(γ+k-ρ+s-2)*D[r-ρ+s-1])
        end
        s2 += (γ+k-1)*D[r]/pochhammer(γ+2k-ρ-1, ρ+2)
        t2 += P[1]*s2
        s2 = zero(T)
        for s in max(0,ρ+1-k):ρ+1
            s2 += Cρ[s+1]*C2[s+1]*(γ+2k-2ρ+2s-3)*D[r-ρ+s-1]
        end
        ΔDold[r+1] = s2
        for j in 1:min(k, p+1)
            t2 += C[j+1]*P[j+1]*(ΔDold[r+2-j] + ΔDold[r+1-j])
        end
        D[r+1] = (ζ*t1-t2)/(P[1]/pochhammer(γ+2k-ρ-1, ρ+2))
        s1 = zero(T)
        for s in max(0,ρ-k):ρ+1
            s1 += Cρ[s+1]*C3[s+1]*(γ+2k-2ρ+2s-1)*D[r-ρ+s]
        end
        ΔD[r+1] = s1
        R[r+1] = N[r+1]/D[r+1]
        k += 1
        for j in min(k, p+1):-1:0
            ΔNold[r-j] = (γ+2k-ρ-j-4)*ΔNold[r-j+1] + (k-j-1)*ΔNold[r-j]
        end
        for j in min(k, q+1):-1:0
            ΔN[r-j] = (γ+2k-ρ-j-2)*ΔN[r-j+1] + (k-j)*ΔN[r-j]
        end
        for j in min(k, p+1):-1:0
            ΔDold[r-j] = (γ+2k-ρ-j-4)*ΔDold[r-j+1] + (k-j-1)*ΔDold[r-j]
        end
        for j in min(k, q+1):-1:0
            ΔD[r-j] = (γ+2k-ρ-j-2)*ΔD[r-j+1] + (k-j)*ΔD[r-j]
        end
        for j in min(k, ρ):-1:1
            C[j+1] += C[j]
        end
        if k ≤ ρ+1
            for s in max(0, ρ+1-k):ρ
                C1[s+1] = pochhammer(k-ρ+s, ρ+1-s)/pochhammer(γ+2k-2ρ+s-2, ρ+2)
            end
            for s in max(0, ρ+1-k):ρ+1
                C2[s+1] = pochhammer(k-ρ+s, ρ+1-s)/pochhammer(γ+2k-2ρ+s-3, ρ+2)
            end
        else
            for s in 0:ρ
                C1[s+1] *= k/(k-one(T)-ρ+s)*(γ+2k-2ρ+s-4)/(γ+2k-ρ+s-2)*(γ+2k-2ρ+s-3)/(γ+2k-ρ+s-1)
            end
            for s in 0:ρ+1
                C2[s+1] *= k/(k-one(T)-ρ+s)*(γ+2k-2ρ+s-5)/(γ+2k-ρ+s-3)*(γ+2k-2ρ+s-4)/(γ+2k-ρ+s-2)
            end
        end
        if k ≤ ρ
            for s in max(0, ρ-k):ρ+1
                C3[s+1] = pochhammer(k-ρ+s+1, ρ+1-s)/pochhammer(γ+2k-2ρ+s-1, ρ+2)
            end
        else
            for s in max(0, ρ-k):ρ+1
                C3[s+1] *= (k+one(T))/(k-ρ+s)*(γ+2k-2ρ+s-3)/(γ+2k-ρ+s-1)*(γ+2k-2ρ+s-2)/(γ+2k-ρ+s)
            end
        end
        t = one(T)
        for j in 1:p
            t *= α[j]+k+1
        end
        t *= k+γ
        for j in 2:p+2
            s = t - P[j-1]
            P[j-1] = t
            t = s
        end
        P[p+2] = t
        t = one(T)
        for j in 1:q
            t *= β[j]+k+1
        end
        t *= k+2
        for j in 2:q+2
            s = t - Q[j-1]
            Q[j-1] = t
            t = s
        end
        Q[q+2] = t
    end
    return isfinite(R[r+1]) ? R[r+1] : R[r]
end


text:
import Base: convert, promote_rule, length, ==

export X90, X, X90m, Y90, Y, Y90m, U90, Uθ, Z90, Z, Z90m, Id, ⊗, MEAS, AC, DiAC, ZX90, CNOT

# NOTE A suggested type hierarchy:
# abstract AbstractBlock
#   - abstract AbstractPulse
#       * immutable Pulse
#       * immutable ZPulse
#   - type PulseBlock
@compat abstract type AbstractPulse end

immutable Pulse <: AbstractPulse
	label::String
	channel::Channel
	length::Float64
	amp::Float64
	phase::Float64
	frequency::Float64
	shape_params::Dict{Symbol, Any}
	hash::UInt
end

function Pulse(
		label::String,
		chan::Channel,
		length::Real=0.0,
		amp::Real=0.0,
		phase::Real=0.0,
		frequency::Real=0.0,
		shape_params=Dict{Symbol,Any}())

	# override the default channel parameters with any passed in
	shape_params = merge(chan.shape_params, shape_params)

	# amplitude and phase will be applied later in translators
	delete!(shape_params, :amp)
	delete!(shape_params, :phase)

	# detemine the necessary parameters
	# WARNING! this is obviously subject to change in Julia Base
	ml = methods(shape_params[:shape_function])
	m = collect(ml)[1]
	kwargs = Base.kwarg_decl(m, typeof(ml.mt.kwsorter))

	filter!((k,v) -> k == :shape_function || k in kwargs, shape_params)

	# precompute pulse hash we'll evaluate it for each pulse we compile
	pulse_hash = hash((label, chan, length, amp, phase, frequency, shape_params))
	Pulse(label, chan, length, amp, phase, frequency, shape_params, pulse_hash)
end

==(a::Pulse, b::Pulse) = a.hash == b.hash
hash(p::Pulse, h::UInt) = hash(p.hash, h)

show(io::IO, p::Pulse) = print(io, "$(p.label)($(p.channel.label))")

immutable ZPulse <: AbstractPulse
	label::String
	channel::Channel
	angle::Float64
end

show(io::IO, z::ZPulse) = print(io, "$(z.label)($(z.channel.label), $(z.angle))")

# loop through and eval to  create some basic 90/180 pulses
# quote the pi2Amp/piAmp symbols to interpolate the symbol
for (func, label, amp, phase) in [
	(:X90,  "X90",  :(:pi2Amp), 0),
	(:X,    "X",    :(:piAmp),  0),
	(:X90m, "X90m", :(:pi2Amp), 0.5),
	(:Y90,  "Y90",  :(:pi2Amp), 0.25),
	(:Y,    "Y",    :(:piAmp),  0.25),
	(:Y90m, "Y90m", :(:pi2Amp), 0.75)
	]
	@eval $func(q) = Pulse($label, q, q.shape_params[:length], q.shape_params[$amp], $phase, 0)
end

U90(q::Qubit, phase::Float64=0.0) = Pulse("U90", q, q.shape_params[:length], q.shape_params[:pi2Amp], phase, 0)
U180(q::Qubit, phase::Float64=0.0) = Pulse("U180", q, q.shape_params[:length], q.shape_params[:piAmp], phase, 0)
Uθ(q::Union{Qubit, Edge}, length, amp, phase) = Pulse("Uθ", q, length, amp, phase)
Uθ(q::Union{Qubit, Edge}, length, amp, phase, freq, shape_params) = Pulse("Uθ", q, length, amp, phase, freq, shape_params)

Z(q::Union{Qubit, Edge}, angle=0.5) = ZPulse("Z", q, angle)
Z90(q::Union{Qubit, Edge}) = ZPulse("Z90", q, 0.25)
Z90m(q::Union{Qubit, Edge}) = ZPulse("Z90m", q, 0.75)
length(z::ZPulse) = 0

function Id(c::Channel, length)
	# TODO: inject constant pulse shape
	Pulse("Id", c, length, 0.0)
end
Id(c::Channel) = Id(c, c.shape_params[:length])

"""
	AC(q::Qubit, num; sampling_rate=1.2e9)

Atomic Clifford `num` on a single qubit `q`. AC pulses are enumerated 0:24.
"""
function AC(q::Qubit, num; sampling_rate=1.2e9)
	if 0 < num < 10
		pulses = [
			Id,
			X90,
			X,
			X90m,
			Y90,
			Y,
			Y90m,
			Z90,
			Z,
			Z90m,
		]
		return pulses[num](q)
	elseif num == 11
		return U180(q, 1/8);
	elseif num == 12
		return U180(q, -1/8)
	elseif num <= 24
		# figure out the approximate nutation frequency calibration from the X180
		# and the sampling_rate.  This isn't ideal as it breaks the split between
		# the pulse and the hardware representation
		Xp = X(q)
		xpulse = waveform(Xp, sampling_rate)
		nut_freq = 0.5/sum(xpulse)*sampling_rate

		rot_angle = [fill(0.5, 4); fill(1/3, 8)]

		# rotation axis polar angle in portions of circle
		# comes in three flavours: x+z or y+z, xy+z and xy-z
		xz = 1/8
		xyzₚ = acos(1/√3) / 2π
		xyzₘ = (π - acos(1/√3) ) / 2π
		Θ = [xz, xz, xz, xz, xyzₚ, xyzₘ, xyzₚ, xyzₘ, xyzₘ, xyzₚ, xyzₚ, xyzₘ]
		# rotation axis azimuthal angle in portions of circle
		ϕ = [0, 1/2, 1/4, -1/4, 1/4, 5/8, -1/8, 3/8, 1/8, 5/8, 3/8, -1/8]

		#TODO: reduce code duplication with pulseshapes.jl
		if q.shape_params[:length] > 0
			# start from a gaussian shaped pulse
			gauss_pulse = PulseShapes.gaussian(pulse_length=q.shape_params[:length], sampling_rate=sampling_rate)
			# scale to achieve to the desired rotation
			cal_scale = (rot_angle/2/pi)*sampling_rate/sum(gauss_pulse)
			# calculate the phase ramp steps to achieve the desired Z component to the rotation axis
			phase_steps = -2π*cos(Θ)*cal_scale*gauss_pulse/sampling_rate
			# Calculate Z DRAG correction to phase steps
			# β is a conversion between XY drag scaling and Z drag scaling
			β = q.shape_params[:drag_scaling]/sampling_rate
			instantaneous_detuning = β * (2π*cal_scale*sin(Θ)*gauss_pulse).^2
			phase_steps += instantaneous_detuning*(1/sampling_rate)
			frame_change = sum(phase_steps)
		elseif abs(Θ) <1e-10
			# Otherwise assume we have a zero-length Z rotation
			frame_change = -rot_angle
		end

		return PulseBlock(Dict(q => [Pulse("AC", q, q.shape_params[:length], 1.0,
		             ϕ[num-12], q.frequency,
		             Dict(:shape_function => getfield(QGL.PulseShapes, :arb_axis_drag),
		                  :nut_freq => nut_freq,
		                  :rot_angle => rot_angle[num-12],
		                  :Θ => Θ[num-12])), Z(q, frame_change)]))
	else
		error("Invalid single qubit Atomic Clifford number")
	end

end

"""
	DiAC(q, num)

	Return the single qubit Clifford `num` in "diatomic" form: Z(α) - X90 - Z(β) - X90 - Z(γ)
"""
function DiAC(q::Qubit, num)
	angles = 0.5 * [
		[ 0.0,  1.0,  1.0],
		[ 0.5, -0.5,  0.5],
		[ 0.0,  0.0,  0.0],
		[ 0.5,  0.5,  0.5],
		[ 0.0, -0.5,  1.0],
		[ 0.0,  0.0,  1.0],
		[ 0.0,  0.5,  1.0],
		[ 0.0,  1.0, -0.5],
		[ 0.0,  1.0,  0.0],
		[ 0.0,  1.0,  0.5],
		[ 0.0,  0.0,  0.5],
		[ 0.0,  0.0, -0.5],
		[ 1.0, -0.5,  1.0],
		[ 1.0,  0.5,  1.0],
		[ 0.5, -0.5, -0.5],
		[ 0.5,  0.5, -0.5],
		[ 0.5, -0.5,  1.0],
		[ 1.0, -0.5, -0.5],
		[ 0.0,  0.5, -0.5],
		[-0.5, -0.5,  1.0],
		[ 1.0,  0.5, -0.5],
		[ 0.5,  0.5,  1.0],
		[ 0.0, -0.5, -0.5],
		[-0.5,  0.5,  1.0]
		]
	return PulseBlock(Dict(q => [Z(q, angles[num][1]), X90(q), Z(q, angles[num][2]), X90(q), Z(q, angles[num][3])]))
end

type PulseBlock
	pulses::Dict{Channel, Vector{AbstractPulse}}
end

convert(::Type{PulseBlock}, p::AbstractPulse) = PulseBlock(Dict(p.channel => [p]))
PulseBlock(p::AbstractPulse) = convert(PulseBlock, p)
PulseBlock{T<:Channel}(chans::Set{T}) = PulseBlock(Dict(chan => AbstractPulse[] for chan in chans))
PulseBlock{T<:Channel}(chans::Vector{T}) = PulseBlock(Dict(chan => AbstractPulse[] for chan in chans))

promote_rule{T<:AbstractPulse}(::Type{T}, ::Type{PulseBlock}) = PulseBlock

⊗(x::AbstractPulse, y::AbstractPulse) = PulseBlock(x) ⊗ PulseBlock(y)
⊗(x::PulseBlock, y::PulseBlock) = PulseBlock(merge(x.pulses, y.pulses))
⊗(x::AbstractPulse, y::PulseBlock) = ⊗(promote(x,y)...)
⊗(x::PulseBlock, y::AbstractPulse) = ⊗(promote(x,y)...)

channels(pb::PulseBlock) = keys(pb.pulses)

function show(io::IO, pb::PulseBlock)
	strs = []
	for ps = values(pb.pulses)
		if length(ps) > 0
			push!(strs, "(" * join([string(p) for p in ps], ", ") * ")")
		end
	end
	str = join(strs, "⊗")
	print(io, str)
end

length(p::Pulse) = p.length
length(pb::PulseBlock) = maximum(sum(length(p) for p in ps) for ps in values(pb.pulses))

"""
	waveform(p::Pulse, sampling_rate)

Render a waveform for a pulse at a given sampling rate.
"""
function waveform(p::Pulse, sampling_rate)
	params = copy(p.shape_params)
	shape_function = pop!(params, :shape_function)
	return shape_function(;pulse_length=p.length, sampling_rate=sampling_rate, params...)
end

function MEAS(q::Qubit)
	m = measurement_channel(q.label)
	meas_pulse = Pulse("MEAS", m, m.shape_params[:length], m.shape_params[:amp], 0.0, m.shape_params[:autodyne_freq])
	pb = PulseBlock(meas_pulse)
	if m.trigger_channel != ""
		t = Marker(m.trigger_channel)
		trig_pulse = Pulse("TRIG", t, t.shape_params[:length], 1.0)
		pb = pb ⊗ trig_pulse
	end
	return pb
end

"""
	flat_top_gaussian(chan; pi_shift=false)

Helper function that returns an Array of pulses to implement a flat-topped gaussian pulse.
"""
function flat_top_gaussian(chan; pi_shift=false)
	pulse_phase = chan.shape_params[:phase]/2π + 0.5*pi_shift
	pulse_amp = chan.shape_params[:amp]
	return [
		Uθ(chan, chan.shape_params[:riseFall], pulse_amp, pulse_phase, 0.0,
			Dict(:shape_function => getfield(QGL.PulseShapes, :half_gaussian), :direction => QGL.PulseShapes.HALF_GAUSSIAN_RISE)),
		Uθ(chan, chan.shape_params[:length], pulse_amp, pulse_phase, 0.0, Dict(:shape_function => getfield(QGL.PulseShapes, :constant))),
		Uθ(chan, chan.shape_params[:riseFall], pulse_amp, pulse_phase, 0.0,
			Dict(:shape_function => getfield(QGL.PulseShapes, :half_gaussian), :direction => QGL.PulseShapes.HALF_GAUSSIAN_FALL))
	]
end

"""
	ZX90(qc::Qubit, qt::Qubit)

Implements an echoed ZX90 "cross-resonance" gate.
"""
function ZX90(qc::Qubit, qt::Qubit)
	CRchan = Edge(qc,qt)
	flat_top_length = CRchan.shape_params[:length] + 2*CRchan.shape_params[:riseFall]
  return PulseBlock( Dict(
		CRchan => [flat_top_gaussian(CRchan); Id(CRchan, qc.shape_params[:length]); flat_top_gaussian(CRchan; pi_shift=true); Id(CRchan, qc.shape_params[:length])],
		qc => [Id(qc, flat_top_length), X(qc), Id(qc, flat_top_length), X(qc)]
	))
end

function CNOT(qc::Qubit, qt::Qubit)
	zx90 = ZX90(qc, qt)
	CRchan = Edge(qc,qt)
	x90m = X90m(qt)
	zx90.pulses[qt] = [Id(qt, length(zx90)), x90m]
	push!(zx90.pulses[qc], Z90m(qc), Id(qc, length(x90m)))
	push!(zx90.pulses[CRchan], Id(CRchan, length(x90m)))
	return zx90
end

# function CNOT(qc::Qubit, qt::Qubit)
# 	CRchan = Edge(qc,qt)
# 	flat_top_length = CRchan.shape_params[:length] + 2*CRchan.shape_params[:riseFall]
#   return QGL.PulseBlock( Dict(
# 		CRchan => [QGL.flat_top_gaussian(CRchan); Id(CRchan, qc.shape_params[:length]); QGL.flat_top_gaussian(CRchan; pi_shift=true); Id(CRchan, qc.shape_params[:length] + qt.shape_params[:length])],
# 		qc => [Id(qc, flat_top_length), X(qc), Id(qc, flat_top_length), X(qc), Z90m(qc), Id(qc, qt.shape_params[:length])],
# 		qt => [Id(qt, 2*flat_top_length+2*qc.shape_params[:length]), X90m(qt)]
# 	))
# end


text:
# Set up
import TrajectoryOptimization._solve

function dubinscar_benchmarks()
    group = BenchmarkGroup()
    stats = BenchmarkGroup()
    group["parallel park"], stats["parallel park"] = parallelpark_benchmark()
    return group, stats
end

function parallelpark_benchmark()
    group = BenchmarkGroup()
    stat_group = BenchmarkGroup()

    stats = Dict("iterations"=>0, "setup_time"=>0.0, "c_max"=>0.0, "major iterations"=>0.0, "cost"=>0.0)
    stats = BenchmarkGroup(["stats"],stats)

    # Params
    disable_logging(Logging.Warn)
    N = 501
    integration = :rk3
    model, obj = Dynamics.dubinscar
    n,m = model.n, model.m

    # Objective
    x0 = [0.0;0.0;0.]
    xf = [0.0;1.0;0.]
    tf =  3.
    Qf = 100.0*Diagonal(I,n)
    Q = (1e-3)*Diagonal(I,n)
    R = (1e-2)*Diagonal(I,m)
    obj = LQRObjective(Q, R, Qf, tf, x0, xf)

    # Unconstrained
    solver_uncon = Solver(model, obj, N=N, integration=integration)
    U0 = ones(m,N)

    stat_group["unconstrained"] = deepcopy(stats)
    group["unconstrained"] = @benchmarkable _solve($solver_uncon, $U0, bmark_stats=$stat_group["unconstrained"])

    # Constrained
    x_min = [-0.25; -0.001; -Inf]
    x_max = [0.25; 1.001; Inf]
    obj_con = TrajectoryOptimization.ConstrainedObjective(obj,x_min=x_min,x_max=x_max)

    solver_con = Solver(model, obj_con, N=N, integration=integration)
    stat_group["constrained"] = deepcopy(stats)
    group["constrained"] = @benchmarkable _solve($solver_con, $U0, bmark_stats=$stat_group["constrained"])

    return group, stat_group
end


text:
@testset "kfolds" begin
    @test collect(kfolds(1:15, k=5)) == collect(kfolds(1:15, 5))
    @test kfolds(15, k=5) == kfolds(15, 5)

    @testset "int" begin
        itrain, ival = kfolds(15, k=5)
        @test size(itrain) == (5,)
        @test size(ival) == (5,)
        @test all(i -> size(i) == (12,), itrain)
        @test all(i -> size(i) == (3,), ival)
        tot = [[it;iv] for (it, iv) in zip(itrain, ival)]
        @test all(x -> sort(x) == 1:15, tot)
    end

    @testset "data" begin
        for (xtr, xval) in kfolds(1:15, k=5)
            @test size(xtr) == (12,)
            @test size(xval) == (3,)
            @test sort([xtr; xval]) == 1:15
        end

        for ((xtr,ytr), (xval, yval)) in kfolds((1:15, 11:25), k=5)
            @test size(xtr) == (12,)
            @test size(xval) == (3,)
            @test size(ytr) == (12,)
            @test size(yval) == (3,)
            @test sort([xtr; xval]) == 1:15
            @test sort([ytr; yval]) == 11:25
        end
    end
end

@testset "leavepout" begin
    @test leavepout(15) == kfolds(15, k=15)
    @test collect(leavepout(1:15)) == collect(kfolds(1:15, 15))
    @test collect(leavepout(1:15, p=5)) == collect(kfolds(1:15, 3))
end


text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

struct PowerBasisCurve{P, S, N, M} <: Spline{P, S, N, M}
    controlpolygon::Array{S, 2} #not exactly the right name, should be coefficients instead, but consistent with the other splines
    #data stored like this:
    # a0x a1x ... aM+1x
    # a0y a1y ... aM+1x
    # and so on for however many spatial dimensions

    function PowerBasisCurve{P, S, N, M}(
        coefficients::Array{S, 2},
    ) where {P, S, N, M}
        @assert (N, M + 1) == size(coefficients)

        return new{P, S, N, M}(copy(coefficients))
    end
end

function coefficients(
    curve::PowerBasisCurve{P, S, N, M},
    spatialindex,
) where {P, S, N, M}
    return curve.controlpolygon[spatialindex, :]
end

function PowerBasisCurve{P, S, N, M}(
    curve::BezierCurve{P, S, N, M},
) where {P, S, N, M}
    controlpolygon = curve.controlpolygon
    coefficients = Array{S, 2}(undef, N, M + 1) # array stores coefficients for all N spatial dimensions

    #this is probably not the most efficient way to convert from Bernstein to Power basis. Optimize later if necessary.
    for k in 0:M
        for i in k:M
            @. coefficients[:, i + 1] +=
                (-1)^(i - k) *
                binomial(M, i) *
                binomial(i, k) *
                controlpolygon[k + 1]
        end
    end
    return PowerBasisCurve{P, S, N, M}(coefficients)
end

function beziertopowerbasis(k, n)
    coefficient = 0
    for i in k:n
        @. coefficients[:, i + 1] +=
            (-1)^(i - k) * binomial(n, i) * binomial(i, k)
    end
end

function point(curve::PowerBasisCurve, u::T) where {T <: Number}
    _, n = size(curve.controlpolygon)
    sum = curve.controlpolygon[:, n]

    for i in (n - 1):-1:1
        @. sum = u * sum + curve.controlpolygon[:, i]
    end

    return sum
end


text:
module test_bernoulli
using VoronoiFVM
using Test

function test()

    function B_Big(x)
	bx=BigFloat(x)
	Float64(bx/(exp(bx)-one(bx)))
    end
    smallrange=-1:1.00001e-5:1
    largerange=-100:1.00001e-3:100
    maxerror(X,b,f)=maximum(abs.(b.(X).-f.(X)))
    @test maxerror(smallrange,B_Big,fbernoulli)<1.0e-14
    @test maxerror(largerange,B_Big,fbernoulli)<1.0e-14

    @test maxerror(smallrange,B_Big,(x)->fbernoulli_pm(x)[1])<1.0e-14
    @test maxerror(largerange,B_Big,(x)->fbernoulli_pm(x)[1])<1.0e-14

    @test maxerror(smallrange,x->B_Big(-x),(x)->fbernoulli_pm(x)[2])<1.0e-14
    @test maxerror(largerange,x->B_Big(-x),(x)->fbernoulli_pm(x)[2])<1.0e-14

    true
end

end


text:
module TestTSClassifier

using Test
using TSML

function test_realdatatsclassifier()
  tscl=TSClassifier(Dict())
  @test_throws ArgumentError fit!(tscl)
  trdirname = joinpath(dirname(pathof(TSML)),"../data/realdatatsclassification/training")
  tstdirname = joinpath(dirname(pathof(TSML)),"../data/realdatatsclassification/testing")
  modeldirname = joinpath(dirname(pathof(TSML)),"../data/realdatatsclassification/model")

  tscl = TSClassifier(Dict(:trdirectory=>trdirname,
			   :tstdirectory=>tstdirname,
			   :modeldirectory=>modeldirname,
			   :feature_range => 5:20,
			   :num_trees=>100))
  fit!(tscl)
  dfresults = transform!(tscl)
  m=fit(tscl)
  dfresults1 = transform(m)
  @test (dfresults .== dfresults1) |> Matrix |> sum > 0

  apredict = dfresults.predtype
  fnames = dfresults.fname
  myregex = r"(?<dtype>[A-Z _ - a-z]+)(?<number>\d*).(?<ext>\w+)"
  mtypes=map(fnames) do fname
    mymatch=match(myregex,fname)
    mymatch[:dtype]
  end
  # misclassified one data
  @test sum(mtypes .== apredict) > 0
end
@testset "TSClassifier" begin
  test_realdatatsclassifier()
end

end


text:
include("startup.jl")

using CairoMakie
using GLM
using MixedModels
using AlgebraOfGraphics
using Statistics
using MultipleTesting
using CategoricalArrays

lq, uq = quantile(skipmissing(get(species, :cogScore)), [0.25, 0.75])
for smp in samplenames(species)
    cs = get(species, smp, :cogScore)
    qrt = ismissing(cs) ? missing :
          cs <= lq      ? "lower" :
          cs >= uq      ? "upper" :
                          "middle"

    set!(species, smp, :cogQuartile, qrt)
end

specdf = metadata(species) |> DataFrame

##

prevspecies = prevalence(species) |> vec .> 0.1
totest = map(x-> !ismissing(x) && x >=18, get(species, :ageMonths)) .&
         map(x-> !ismissing(x), get(species, :cogScore))

@info count(totest)

testcomm = species[prevspecies, totest]
lm_results = DataFrame()

##

for spc in features(testcomm)
    @info name(spc)
    ab = vec(abundances(testcomm[spc, :]))
    over0 = ab .> 0

    df = DataFrame(
        age = get(testcomm, :ageMonths)[over0],
        cogScore = get(testcomm, :cogScore)[over0],
        cogQuartile = categorical(get(testcomm, :cogQuartile)[over0], levels=["lower", "middle", "upper"]),
        bug = log.(collect(ab[over0]))
    )

    mod = lm(@formula(bug ~ age + cogScore), df)
    ct = DataFrame(coeftable(mod))
    ct.species .= name(spc)
    ct.testvar .= "cogScore"
    append!(lm_results, ct)

    mod = lm(@formula(bug ~ age + levelcode(cogQuartile)), df)
    ct = DataFrame(coeftable(mod))
    ct.species .= name(spc)
    ct.testvar .= "cogQuartile"
    append!(lm_results, ct)

end

select!(lm_results, Cols(:species, :testvar, :))
rename!(lm_results, "Pr(>|t|)"=>"pvalue")

##

adjust = MultipleTesting.adjust

@chain lm_results begin
    groupby(:testvar)
    transform!(:pvalue => (col-> adjust(collect(col), BenjaminiHochberg())) => :qvalue)
end

sort!(lm_results, :qvalue)

CSV.write("data/lm_results/species_cogscore.csv", lm_results)
##

cl = ["lower"=> :seagreen, "middle"=> :lightgray, "upper"=> :dodgerblue]

dt = data(specdf[completecases(specdf, [:ageMonths, :cogScore]), :]) * mapping(:ageMonths => "Age (months)", :cogScore=> "Cognitive function score", color = :cogQuartile)

out = draw(dt * visual(Scatter, strokewidth=0.5), palettes=(;color=cl))

save("figures/cogScore_quartiles.png", out)

##

bugs = [
    "s__Bifidobacterium_catenulatum",
    "s__Coprococcus_eutactus", # butyrate producer
    "s__Streptococcus_thermophilus",
    "s__Eubacterium_eligens",
    "s__Ruminococcus_gnavus",
    "s__Veillonella_parvula"
]

bugdf = DataFrame()
for bug in bugs
    ab = vec(abundances(testcomm[bug, :]))
    over0 = ab .> 0
    df = DataFrame(
        age = get(testcomm, :ageMonths)[over0],
        cogScore = get(testcomm, :cogScore)[over0],
        cogQuartile = categorical(get(testcomm, :cogQuartile)[over0], levels=["lower", "middle", "upper"]),
        bug = log.(collect(ab[over0])),
        regbug = collect(ab[over0])
    )

    df.bugname .= replace(bug, r"s__([A-Z])(\w+)_(\w+)"=> s"\1. \3")
    append!(bugdf, df)
end

##

fig = Figure(resolution=(800,1200))
draw!(fig[1,1],
      data(bugdf) *
      mapping(:cogQuartile=>"Quartile",
              :bug=> "Log(relative abundance)";
              color=:cogQuartile, row=:bugname
              ) *
      (visual(BoxPlot) + visual(Scatter, strokewidth=0.5)),
      palettes= (; color=cl),
      axis = (; titlevisible=false)
)

draw!(fig[1,2],
      data(bugdf) *
      mapping(:cogScore=>"Cognitive Function Score",
              :bug=> "Log(relative abundance)";
              color=:cogQuartile, row=:bugname
              );
      palettes= (; color=cl),
      axis=(; ylabel="", yticklabelsvisible=false)
)

save("figures/lms.png", fig)
fig

##

##

fig = Figure(resolution=(900,600))
# draw!(fig[1,1],
#       data(subset(bugdf, :bugname => ByRow(b-> b in replace.(bugs[1:3], r"s__([A-Z])(\w+)_(\w+)"=> s"\1. \3")))) *
#       mapping(:cogQuartile=>"Quartile",
#               :bug=> "Log(relative abundance)";
#               color=:cogQuartile, row=:bugname
#               ) *
#       (visual(BoxPlot) + visual(Scatter, strokewidth=0.5)),
#       palettes= (; color=cl),
#       axis = (; titlevisible=false, xticklabelsize=14)
# )

draw!(fig[1,1],
      data(subset(bugdf, :bugname => ByRow(b-> b in replace.(bugs[1:3], r"s__([A-Z])(\w+)_(\w+)"=> s"\1. \3")))) *
      (mapping(:cogScore=>"Cognitive Function Score",
              :bug=> "Log(relative abundance)";
              color=:cogQuartile, row=:bugname
              ) +
       mapping(:cogScore=>"Cognitive Function Score",
              :bug=> "Log(relative abundance)";
              row=:bugname
              ) * linear());
      palettes= (; color=cl),
      axis=(; xticklabelsize=14)
)

# draw!(fig[1,3],
#       data(subset(bugdf, :bugname => ByRow(b-> b in replace.(bugs[4:6], r"s__([A-Z])(\w+)_(\w+)"=> s"\1. \3")))) *
#       mapping(:cogQuartile=>"Quartile",
#               :bug=> "Log(relative abundance)";
#               color=:cogQuartile, row=:bugname
#               ) *
#       (visual(BoxPlot) + visual(Scatter, strokewidth=0.5)),
#       palettes= (; color=cl),
#       axis = (; titlevisible=false, ylabelvisible = false, yticklabelsvisible=false, xticklabelsize=14)
# )

draw!(fig[1,2],
      data(subset(bugdf, :bugname => ByRow(b-> b in replace.(bugs[4:6], r"s__([A-Z])(\w+)_(\w+)"=> s"\1. \3")))) *
      (mapping(:cogScore=>"Cognitive Function Score",
              :bug=> "Log(relative abundance)";
              color=:cogQuartile, row=:bugname
              ) +
       mapping(:cogScore=>"Cognitive Function Score",
              :bug=> "Log(relative abundance)";
              row=:bugname
              ) * linear());
      palettes= (; color=cl),
)

save("figures/lms-split-2.png", fig)
fig

##

################
# Longitudinal #
################

sort!(tps, [:subject, :timepoint])
grps = groupby(tps, :subject)
fecals = Set(zip(get(species, :subject), get(species, :timepoint)))

##

fig = Figure()
ax = Axis(fig[1,1], xlabel="timepoint", ylabel="cogScore")

for grp in grps
    any(!ismissing, grp.cogScore) || continue
    keep = subset(grp, :cogScore => ByRow(!ismissing))
    scatter!(ax, keep.timepoint, keep.cogScore)
    lines!(ax, keep.timepoint, keep.cogScore)
end

fig

##

diffs = Float64[]

for key in keys(grps)
    grp = grps[key]
    scores = collect(skipmissing(grp.cogScore))
    n = length(scores)
    n > 1 || continue
    d = collect(Iterators.map(x-> abs(x[1] - x[2]), Iterators.product(scores, scores)))
    for i in 1:n, j in 1:n
        i < j || continue
        push!(diffs, d[i] - d[j])
    end
end

##

hist(abs.(diffs), axis=(; xlabel="abs(score difference)", ylabel="count"))

##

tpsfecal = leftjoin(tps, DataFrame(subject=get(species, :subject), timepoint=get(species, :timepoint), sample=samplenames(species)), on=[:subject, :timepoint])
sort!(tpsfecal, [:subject, :timepoint])
grps = groupby(tpsfecal, :subject)

##

youngfecal = Dict()
oldfecal = Dict()

for grp in grps
    subject = first(grp.subject)
    issorted(grp.timepoint) || error()

    firstyoung = findfirst(((a, f),) -> !ismissing(f) && !ismissing(a) && a < 12, collect(zip(grp.ageMonths, grp.sample)))
    firstold   = findfirst(((a, f),) -> !ismissing(f) && !ismissing(a) && a > 18, collect(zip(grp.ageMonths, grp.sample)))

    if !isnothing(firstyoung)
        laterscores = map(((i, s),) -> i > firstyoung && !ismissing(s), enumerate(grp.cogScore))
        if sum(laterscores) > 0
            lscr = grp[laterscores, :cogScore]
            ages = grp[laterscores, :ageMonths]
            youngfecal[subject] = (;fecal = grp[firstyoung, :sample], cogscores = collect(skipmissing(lscr[.!ismissing.(lscr)])), ages = collect(skipmissing(ages[.!ismissing.(lscr)])))
        end
    end

    if !isnothing(firstold)
        laterscores = map(((i, s),) -> i > firstold && !ismissing(s), enumerate(grp.cogScore))
        if sum(laterscores) > 0
            lscr = grp[laterscores, :cogScore]
            ages = grp[laterscores, :ageMonths]
            oldfecal[subject] = (;fecal = grp[firstold, :sample], cogscores = collect(skipmissing(lscr[.!ismissing.(lscr)])), ages = collect(skipmissing(ages[.!ismissing.(lscr)])))
        end
    end
end

##

yfdf = DataFrame(sample = [youngfecal[k][:fecal] for k in keys(youngfecal)])
yfdf.age = [get(species, s, :ageMonths) for s in yfdf.sample]
yfdf.hhs = [get(species, s, :mother_HHS_Education) for s in yfdf.sample]
yfdf.firstscore = [first(youngfecal[k][:cogscores]) for k in keys(youngfecal)]
yfdf.meanscore = [mean(youngfecal[k][:cogscores]) for k in keys(youngfecal)]
yfdf.allscores = [youngfecal[k][:cogscores] for k in keys(youngfecal)]
yfdf.allages = [youngfecal[k][:ages] for k in keys(youngfecal)]

##

youngsamples = map(a-> !ismissing(a) && a <= 12, get(species, :ageMonths))
youngprev = vec((prevalence(species[:, youngsamples]) .> 0.1))

youngcomm = species[youngprev, youngsamples]

for sp in features(youngcomm)
    yfdf[!, Microbiome.name(sp)] = collect(vec(abundances(species[sp, yfdf.sample])))
end

CSV.write("data/lm_results/under12m_cogScores_taxa.csv", yfdf)

##

lm_results = DataFrame()

for spc in featurenames(youngcomm)
    @info spc
    df = yfdf[:, ["sample", "age", "firstscore", "meanscore", spc]]
    subset!(df, spc => ByRow(!=(0)))
    size(df, 1) > 5 || continue

    df.bug = log.(df[!, spc])

    mod = lm(@formula(bug ~ age + firstscore), df)
    ct = DataFrame(coeftable(mod))
    ct.species .= spc
    ct.testvar .= "firstscore"
    append!(lm_results, ct)

    mod = lm(@formula(bug ~ age + meanscore), df)
    ct = DataFrame(coeftable(mod))
    ct.species .= spc
    ct.testvar .= "meanscore"
    append!(lm_results, ct)

end

select!(lm_results, Cols(:species, :testvar, :))
rename!(lm_results, "Pr(>|t|)"=>"pvalue")

##

@chain lm_results begin
    groupby(:testvar)
    transform!(:pvalue => (col-> adjust(collect(col), BenjaminiHochberg())) => :qvalue)
    @rsubset!(:Name in ("firstscore", "meanscore"))
end
lm_results

sort!(lm_results, :qvalue)
CSV.write("data/lm_results/young_species_predict_cogscore.csv", lm_results)

##

##

ofdf = DataFrame(sample = [oldfecal[k][:fecal] for k in keys(oldfecal)])
ofdf.age = [get(species, s, :ageMonths) for s in ofdf.sample]
ofdf.hhs = [get(species, s, :mother_HHS_Education) for s in ofdf.sample]
ofdf.firstscore = [first(oldfecal[k][:cogscores]) for k in keys(oldfecal)]
ofdf.meanscore = [mean(oldfecal[k][:cogscores]) for k in keys(oldfecal)]
ofdf.allscores = [oldfecal[k][:cogscores] for k in keys(oldfecal)]
ofdf.allages = [oldfecal[k][:ages] for k in keys(oldfecal)]

##

oldsamples = map(a-> !ismissing(a) && a > 12, get(species, :ageMonths))
oldprev = vec((prevalence(species[:, oldsamples]) .> 0.1))

oldcomm = species[oldprev, oldsamples]

for sp in features(oldcomm)
    ofdf[!, Microbiome.name(sp)] = collect(vec(abundances(species[sp, ofdf.sample])))
end

CSV.write("data/lm_results/over12m_cogScores_taxa.csv", ofdf)

##

lm_results = DataFrame()

for spc in featurenames(oldcomm)
    @info spc
    df = ofdf[:, ["sample", "age", "firstscore", "meanscore", spc]]
    subset!(df, spc => ByRow(!=(0)))
    size(df, 1) > 5 || continue

    df.bug = log.(df[!, spc])

    mod = lm(@formula(bug ~ age + firstscore), df)
    ct = DataFrame(coeftable(mod))
    ct.species .= spc
    ct.testvar .= "firstscore"
    append!(lm_results, ct)

    mod = lm(@formula(bug ~ age + meanscore), df)
    ct = DataFrame(coeftable(mod))
    ct.species .= spc
    ct.testvar .= "meanscore"
    append!(lm_results, ct)

end

select!(lm_results, Cols(:species, :testvar, :))
rename!(lm_results, "Pr(>|t|)"=>"pvalue")

##

@chain lm_results begin
    groupby(:testvar)
    transform!(:pvalue => (col-> adjust(collect(col), BenjaminiHochberg())) => :qvalue)
    @rsubset!(:Name in ("firstscore", "meanscore"))
end
lm_results

sort!(lm_results, :qvalue)
CSV.write("data/lm_results/old_species_predict_cogscore.csv", lm_results)

##

using MLJ

yfdf.quartile = categorical([x < lq ? "lower" : x > uq ? "upper" : "middle" for x in yfdf.firstscore], levels=["lower", "middle", "upper"])

y, X = unpack(yfdf, ==(:quartile), (x-> !in(x, (:sample, :meanscore, :firstscore, :age))))

Tree = @load DecisionTreeClassifier pkg=DecisionTree
tree = Tree()

fprs = Float64[]
tprs = Float64[]

for _ in 1:100
    res = MLJ.evaluate(tree, X, y,
        resampling=Holdout(; fraction_train=0.9, shuffle=true),
                measures=[Accuracy(), MulticlassTruePositiveRate(), MulticlassFalsePositiveRate()],
                verbosity=2
    )

    (ac, tpr, fpr) = res.measurement
    push!(fprs, fpr)
    push!(tprs, tpr)
end

srt = sortperm(fprs)
scatter(fprs, tprs, axis=(; limits=(0,1,0,1)))

##

tree3 = Tree(max_depth=3)
tree5 = Tree(max_depth=5)
tree8 = Tree(max_depth=8)
tree10 = Tree(max_depth=10)

mach3 = machine(tree3, X, y)
mach5 = machine(tree5, X, y)
mach8 = machine(tree8, X, y)
mach10 = machine(tree10, X, y)
fit!.([mach3,mach5,mach8,mach10])

yhat3 = MLJ.predict(mach3)
yhat5 = MLJ.predict(mach5)
yhat8 = MLJ.predict(mach8)
yhat10 = MLJ.predict(mach10)

fprs3, tprs3, ts3 = roc_curve(yhat3, y)
fprs5, tprs5, ts5 = roc_curve(yhat5, y)
fprs8, tprs8, ts8 = roc_curve(yhat8, y)
fprs10, tprs10, ts10 = roc_curve(yhat10, y)

##
fig = Figure()
ax = Axis(fig[1,1], ylabel="true positive rate (sensitivity)", xlabel="false positive rate (specificity)")
noinf = lines!(ax, [0, 1], [0, 1], linestyle=:dash, color=:black)

sc3 = scatter!(ax, fprs3, tprs3)
ln3 = lines!(ax, fprs3, tprs3; color=sc3.attributes[:color][])
sc5 = scatter!(ax, fprs5, tprs5)
ln5 = lines!(ax, fprs5, tprs5; color=sc5.attributes[:color][])
sc8 = scatter!(ax, fprs8, tprs8)
ln8 = lines!(ax, fprs8, tprs8; color=sc8.attributes[:color][])
sc10 = scatter!(ax, fprs10, tprs10)
ln10 = lines!(ax, fprs10, tprs10; color=sc10.attributes[:color][])

Legend(fig[1,2], [[sc3, ln3], [sc5, ln5], [sc8, ln8], [sc10, ln10], noinf],
                 ["Depth = 3", "Depth = 5", "Depth = 8", "Depth = 10", "no information"])

fig

max_stars_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis

max_issues_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis

max_forks_repo_name: Klepac-Ceraj-Lab/ResonanceAnalysis


text:
@testset "KLdivergences" begin

end


text:
ENV["MPLBACKEND"]="agg" # no GUI
using SeisIO, SeisIO.RandSeis, PlotSeis
using Random: randn
using Dates: now
using Printf: @sprintf
path = Base.source_dir()
cd(dirname(pathof(PlotSeis))*"/../test")
test_start = now()
printstyled(stdout, string(test_start, ": tests begin, source_dir = ", path, "/\n"), color=:light_green, bold=true)

printstyled("Generating data\n", color=:light_green, bold=true)
n = (4,5,10,15,20,23,25)
A = Array{SeisData,1}(undef, 12)
for i = 1:7
  A[i] = randSeisData(n[i])
end

S = SeisData(randSeisChannel(s=true))
# timespan crosses a year boundary
S.x[1] = randn(200000)
S.fs[1] = 50.0
S.t[1] = [1 -100000; 200000 0]
A[8] = deepcopy(S)

# timespan crosses a day boundary
S.t[1] = [1 86300000000; 200000 0]
A[9] = deepcopy(S)

# timespan crosses an hour boundary
S.t[1] = [1 1; 200000 0]
A[10] = deepcopy(S)

# timespan crosses a minute boundary
S.x[1] = rand(120000)
S.fs[1] = 1000.0
S.t[1] = [1 1; 120000 0]
A[11] = deepcopy(S)

# timespan under a minute
S.x[1] = randn(2000)
S.t[1] = [1 0; 2000 0]
S.fs[1] = 4000.0
A[12] = deepcopy(S)

printstyled("Testing plots\n", color=:light_green, bold=true)
for S in A
  plotseis(S)
  uptimes(S)
  uptimes(S, summed=true)
  for i=1:3
    close()
  end
  if length(S.x[1]) ≥ 10000 && S.fs[1] > 0.0
    logspec(S, 1, fmin = S.fs[1] > 100.0 ? 1.0 : 0.01)
    close()
  end
end
plotseis(A[9], fmt = "%H:%M:%S")

test_end = now()
δt = 0.001*(test_end-test_start).value
mm = round(Int, div(δt, 60))
ss = rem(δt, 60)
printstyled(string(test_end, ": tests end, elapsed time (mm:ss.μμμ) = ",
                          @sprintf("%02i", mm), ":",
                          @sprintf("%06.3f", ss), "\n"), color=:light_green, bold=true)