text:
using PoincareInvariants.SecondPoincareInvariants.ChebyshevImplementation.PaduaTransforms

module ChebyshevTestUtils
    using LinearAlgebra: UpperTriangular

    export T0, T1, T2, T3, T6, T11
    export evalT, remove_lower_right

    T0(x) = evalpoly(x, (1,))
    T1(x) = evalpoly(x, (0, 1))
    T2(x) = evalpoly(x, (-1, 0, 2))
    T3(x) = evalpoly(x, (0, -3, 0, 4))
    T6(x) = evalpoly(x, (-1, 0, 18, 0, -48, 0, 32))
    T11(x) = evalpoly(x, (0, -11, 0, 220, 0, -1232, 0, 2816, 0, -2816, 0, 1024))

    function evalT(x::T, n::Integer) where T
        if n == 0
            return one(T)
        elseif n == 1
            return x
        elseif n > 1
            a = one(T)
            b = x
            for _ in 2:n
                b, a = 2x * b - a, b
            end
            return b
        else
            throw(ArgumentError("n must be greater than 0"))
        end
    end

    function evalT(x::T, y::T, coeffs::AbstractMatrix) where T
        axes(coeffs, 1) == axes(coeffs, 2) || error()
        s = zero(T)
        fi = firstindex(coeffs, 1)

        for i in axes(coeffs, 1)
            for j in axes(coeffs, 2)
                s += coeffs[i, j] * evalT(y, i - fi) * evalT(x, j - fi)
            end
        end

        s
    end

    remove_lower_right(A) = reverse(Matrix(UpperTriangular(A)); dims=2)
end  # ChebyshevTestUtils

@safetestset "ChebyshevTestUtils" begin
    using ..ChebyshevTestUtils

    @test evalT(0.225, 0) == 1.0
    @test evalT(0.225, 1) == 0.225
    @test evalT(0.225, 6) ≈ T6(0.225) atol=5eps()
    @test evalT(0.225, 11) == let x = 0.225
        1024x^11 - 2816x^9 + 2816x^7 - 1232x^5 + 220x^3 - 11x
    end
end

@safetestset "getpaduanum, getdegree and nextpaduanum" begin
    using ..PaduaTransforms

    for n in [1:10..., 50, 123, 511, 10_000]
        paduanum = getpaduanum(n)
        @test paduanum == (n + 1) * (n + 2) ÷ 2
        @test getdegree(paduanum) == n

        @test_throws ArgumentError getdegree(paduanum + 1)
        @test_throws ArgumentError getdegree(paduanum - 1)

        @test nextpaduanum(paduanum - 1) == paduanum
        @test nextpaduanum(paduanum) == paduanum
        @test nextpaduanum(paduanum + 1) == getpaduanum(n + 1)
    end
end

@safetestset "paduapoint, ispadua and getpaduapoints" begin
    using ..PaduaTransforms: ispadua, paduapoint
    using ..PaduaTransforms

    @test paduapoint(Float64, 1, 3, 5)[1] ≈ cos(π * 1 / 5) atol=3eps()
    @test paduapoint(Float64, 1, 3, 5)[2] ≈ cos(π * 3 / 6) atol=3eps()

    @test paduapoint(Float64, 23, 42, 100)[1] ≈ cos(π * 23 / 100) atol=3eps()
    @test paduapoint(Float64, 23, 42, 100)[2] ≈ cos(π * 42 / 101) atol=3eps()

    @test eltype(paduapoint(Float32, 1, 2, 4)) == Float32
    @test eltype(paduapoint(Float64, 5, 6, 7)) == Float64

    # collect tuple to vector, so ≈ works
    @test [collect(paduapoint(Float32, x, y, 1)) for y in 0:1+1, x in 0:0] ≈ [
        [1,  1],
        [1,  0],
        [1, -1]
    ]

    @test ispadua(0, 0) == true
    @test ispadua(1, 0) == false
    @test ispadua(0, 1) == false
    @test ispadua(5, 4) == false
    @test ispadua(4, 4) == true
    @test ispadua(2, 0) == true

    @test getpaduapoints(1) ≈ [
         1  1;
         1 -1;
        -1  0
    ]

    @test getpaduapoints(4) ≈ [
        cospi(0/4) cospi(0/5);
        cospi(0/4) cospi(2/5);
        cospi(0/4) cospi(4/5);
        cospi(1/4) cospi(1/5);
        cospi(1/4) cospi(3/5);
        cospi(1/4) cospi(5/5);
        cospi(2/4) cospi(0/5);
        cospi(2/4) cospi(2/5);
        cospi(2/4) cospi(4/5);
        cospi(3/4) cospi(1/5);
        cospi(3/4) cospi(3/5);
        cospi(3/4) cospi(5/5);
        cospi(4/4) cospi(0/5);
        cospi(4/4) cospi(2/5);
        cospi(4/4) cospi(4/5)
    ]

    torowvec(t) = [t[1] t[2]]
    for n in [21, 22, 23, 100]
        pnts = getpaduapoints(n)

        @test size(pnts) == (getpaduanum(n), 2)

        testpnts = reduce(vcat, [
            paduapoint(Float64, x, y, n) |> torowvec
            for y in 0:n+1, x in 0:n if ispadua(x, y)
        ])

        @test pnts == testpnts
    end

    dopoints = getpaduapoints(5) do x, y
        x, y, x * y, x + y
    end

    @test size(dopoints) == (getpaduanum(5), 4)

    @test dopoints[:, 1] ≈ getpaduapoints(5)[:, 1]
    @test dopoints[:, 2] ≈ getpaduapoints(5)[:, 2]
    @test dopoints[:, 3] ≈ getpaduapoints(5)[:, 1] .* getpaduapoints(5)[:, 2]
    @test dopoints[:, 4] ≈ getpaduapoints(5)[:, 1] .+ getpaduapoints(5)[:, 2]
end

@safetestset "weight! and invweight!" begin
    using ..PaduaTransforms:
        weight!, invweight!

    @test weight!(ones(4+2, 4+1), 4) == [
        0.025  0.05  0.05  0.05  0.025;
         0.05   0.1   0.1   0.1   0.05;
         0.05   0.1   0.1   0.1   0.05;
         0.05   0.1   0.1   0.1   0.05;
         0.05   0.1   0.1   0.1   0.05;
        0.025  0.05  0.05  0.05  0.025
    ]

    @test invweight!(ones(4+2, 4+1)) == [
          1    0.5    0.5    0.5    1;
        0.5   0.25   0.25   0.25  0.5;
        0.5   0.25   0.25   0.25  0.5;
        0.5   0.25   0.25   0.25  0.5;
        0.5   0.25   0.25   0.25  0.5;
          1    0.5    0.5    0.5    1
    ]
end

@safetestset "tovalsmat!" begin
    using ..PaduaTransforms

    @test PaduaTransforms.tovalsmat!(ones(3 + 2, 3 + 1), 1:getpaduanum(3), 3) == [
        1 0 6 0 ;
        0 4 0 9 ;
        2 0 7 0 ;
        0 5 0 10;
        3 0 8 0
    ]

    @test PaduaTransforms.tovalsmat!(ones(2 + 2, 2 + 1), 1:getpaduanum(2), 2) == [
        1 0 5;
        0 3 0;
        2 0 6;
        0 4 0;
    ]
end

@safetestset "fromcoeffsmat!" begin
    using ..PaduaTransforms

    mat = [(x, y) for y in 0:2+1, x in 0:2]
    to = similar(mat, getpaduanum(2))

    PaduaTransforms.fromcoeffsmat!(to, mat, 2, Val(true))
    @test to == [(0, 0), (1, 0), (0, 1), (2, 0), (1, 1), (0, 2)]

    PaduaTransforms.fromcoeffsmat!(to, mat, 2, Val(false))
    @test to == [(0, 0), (0, 1), (1, 0), (0, 2), (1, 1), (2, 0)]

    @test PaduaTransforms.fromcoeffsmat!(zeros(4, 4), reshape(1:20, 5, 4), 3) == [
        1  6  11 16;
        2  7  12  0;
        3  8   0  0;
        4  0   0  0
    ]
end

@safetestset "tocoeffsmat!" begin
    using ..PaduaTransforms

    mat = PaduaTransforms.tocoeffsmat!(zeros(5, 4), reshape(1:16, 4, 4))

    @test mat == [
        1.0  5.0   9.0  13.0
        2.0  6.0  10.0  14.0
        3.0  7.0  11.0  15.0
        4.0  8.0  12.0  16.0
        0.0  0.0   0.0   0.0
    ]
end

@safetestset "fromvalsmat!" begin
    using ..PaduaTransforms

    mat1 = reshape(1:20, 5, 4)
    out1 = zeros(10)

    PaduaTransforms.fromvalsmat!(out1, mat1, 3)

    @test out1 == [1, 3, 5,  7, 9,  11, 13, 15,  17, 19]

    mat2 = reshape(1:12, 4, 3)
    out2 = zeros(6)

    PaduaTransforms.fromvalsmat!(out2, mat2, 2)

    @test out2 == [1, 3,  6, 8,  9, 11]
end

@safetestset "1D paduatransform!" begin
    using ..PaduaTransforms
    using ..ChebyshevTestUtils

    @testset "Low degree tests" begin
        let n = 3; vals = getpaduapoints(n) do x, y
                1.5 * T0(x) * T1(y) + 2 * T3(x) * T0(y) + 100 * T2(x) * T1(y)
            end

            plan = PaduaTransformPlan{Float64}(n)

            out = paduatransform!(zeros(n+1, n+1), plan, vals)

            @test maximum(abs, out .- [
                  0 0   0 2
                1.5 0 100 0
                  0 0   0 0
                  0 0   0 0
            ]) / eps() < 100
        end

        let n = 6; vals = getpaduapoints(n) do x, y
                15 * T0(x) * T0(y) + 2.5 * T0(x) * T6(y) + 0.1 * T3(x) * T3(y)
            end

            plan = PaduaTransformPlan{Float64}(n)

            out = paduatransform!(zeros(n+1, n+1), plan, vals)

            @test maximum(abs, out .- [
                 15 0 0   0 0 0 0
                  0 0 0   0 0 0 0
                  0 0 0   0 0 0 0
                  0 0 0 0.1 0 0 0
                  0 0 0   0 0 0 0
                  0 0 0   0 0 0 0
                2.5 0 0   0 0 0 0
            ]) / eps() < 100
        end
    end

    @testset "Degree $n tests" for n in [11, 20, 29, 40, 51]
        plan = PaduaTransformPlan{Float64}(n)

        testcoeffs = remove_lower_right(PaduaTransforms.rand(n+1, n+1))

        vals = getpaduapoints(n) do x, y
            evalT(x, y, testcoeffs)
        end

        out = paduatransform!(zeros(n+1, n+1), plan, vals)

        @test maximum(abs, out .- testcoeffs) / eps() < 100
    end
end

@safetestset "1D invpaduatransform!" begin
    using ..PaduaTransforms
    using ..ChebyshevTestUtils

    @testset "Low degree tests" begin
        let n = 3; testvals = getpaduapoints(n) do x, y
                1.5 * T0(x) * T1(y) + 2 * T3(x) * T0(y) + 100 * T2(x) * T1(y)
            end

            invplan = InvPaduaTransformPlan{Float64}(n)

            coeffs = [
                  0 0   0 2;
                1.5 0 100 0;
                  0 0   0 0;
                  0 0   0 0
            ]

            vals = Vector{Float64}(undef, getpaduanum(n))
            invpaduatransform!(vals, invplan, coeffs)

            @test maximum(abs, vals .- testvals) / eps() < 100
        end

        let n = 6; testvals = getpaduapoints(n) do x, y
                15 * T0(x) * T0(y) + 2.5 * T0(x) * T6(y) + 0.1 * T3(x) * T3(y)
            end

            invplan = InvPaduaTransformPlan{Float64}(n)

            coeffs = [
                 15 0 0   0 0 0 0
                  0 0 0   0 0 0 0
                  0 0 0   0 0 0 0
                  0 0 0 0.1 0 0 0
                  0 0 0   0 0 0 0
                  0 0 0   0 0 0 0
                2.5 0 0   0 0 0 0
            ]

            vals = Vector{Float64}(undef, getpaduanum(n))
            invpaduatransform!(vals, invplan, coeffs)

            @test maximum(abs, vals .- testvals) / eps() < 100
        end
    end

    @testset "Degree $n tests" for n in [11, 20, 29, 40, 51]
        invplan = InvPaduaTransformPlan{Float64}(n)

        coeffs = remove_lower_right(rand(n+1, n+1))

        testvals = getpaduapoints(n) do x, y
            evalT(x, y, coeffs)
        end

        vals = Vector{Float64}(undef, getpaduanum(n))
        invpaduatransform!(vals, invplan, coeffs)

        # These errors get pretty damn big
        # max abs error is about 14_000 for degree 50
        @test maximum(abs, vals .- testvals) / eps() < 100 * 5^(n / 10)
    end
end

@safetestset "ND paduatransform! and invpaduatransform!" begin
    using ..PaduaTransforms

    n = 500
    D = 6
    plan = PaduaTransformPlan{Float64}(n)
    invplan = InvPaduaTransformPlan{Float64}(n)

    ndtestvals = rand(getpaduanum(n), D)
    ndcoeffs = zeros(n+1, n+1, D)
    paduatransform!(ndcoeffs, plan, ndtestvals)

    ndvals = Matrix{Float64}(undef, getpaduanum(n), D)
    invpaduatransform!(ndvals, invplan, ndcoeffs)

    @test maximum(abs, ndvals .- ndtestvals) / eps() < 100

    for i in 1:D
        coeffs = zeros(n+1, n+1)
        @test ndcoeffs[:, :, i] == paduatransform!(coeffs, plan, ndtestvals[:, i])
    end

    for i in 1:D
        vals = Vector{Float64}(undef, getpaduanum(n))
        @test ndvals[:, i] == invpaduatransform!(vals, invplan, ndcoeffs[:, :, i])
    end
end


text:
using Revise

using DataFrames, Statistics, Distributions, StatsBase,
  LinearAlgebra, BenchmarkTools, Random, Optim, Zygote, UnPack, Random

MFloat64 = Union{Missing, Float64}

Random.seed!(11)

# a simple AR(1) model
struct AR1{TV<:AbstractVector, TM<:AbstractMatrix}
  ϕ::TV
  μ::TV
  Σ::TM
  ΣcholL::TM
end

#primary constructor, performs an integrity check on Σ integrity checks
function AR1(;ϕ::TV, μ::TV,  Σ::TM, ΣcholL::TM=cholesky(Σ).L |> Matrix) where {TV,TM}

  @assert (Σ .≈ ΣcholL * ΣcholL') |> all

  return AR1(ϕ, μ, Σ, ΣcholL)
end

centerar1(ar1::AR1) = AR1(ar1.ϕ, 0.0, ar1.Σ, ar1.ΣcholL)

function AR1(Xy::AbstractMatrix{T}) where T

  X̃y = @view Xy[2:end, :]
  LX̃y = @view Xy[1:(end-1), :]

  #compute the persistence as a regression
  #the special case is for a constant value in all terms, which would give ϕ=0.0, σ²=0.0, μ=c
  ϕ = ((x̃,Lx̃)->ifelse(maximum(x̃)-minimum(x̃) + 1.0 ≈ 1.0, 0.0, cov(x̃,Lx̃)/var(Lx̃))
    ).(eachcol(X̃y), eachcol(LX̃y))

  μ = mean(Xy, dims=1) |> vec

  ε = vcat((Xy[1,:] .- μ)', X̃y .- LX̃y .* ϕ' .- μ')
  Σ = cov(ε)
  ΣcholL = semicovcholl(ε)

  #println("μ: $μ, ϕ: $ϕ, b: $b, sum(Xy): $(sum(Xy)), Xy[1:10, :]: $(Xy[1:10, :])")
  #=global sigma = Σ
  display(Σ)
  println("b=$b
    ϕ: $ϕ")=#

  @assert length(μ) == length(ϕ) == size(Xy,2) == size(Σ,1)
  @assert (Σ .≈ ΣcholL * ΣcholL') |> all
  @assert (ϕ .≈ ((x̃,Lx̃)-> ifelse(maximum(x̃) - minimum(x̃) + 1.0 ≈ 1.0, 0.0,
    cov(x̃ .- mean(x̃), Lx̃ .- mean(Lx̃))/var(Lx̃ .- mean(Lx̃)))
    ).(eachcol(X̃y), eachcol(LX̃y))) |> all

  return AR1(;ϕ,μ,Σ,ΣcholL)
end

function simulatear1!(ar1::AR1, Xyt::AbstractMatrix;
    ) where TV<:AbstractVector
  @unpack ϕ, Σ, μ = ar1
  K::Int = length(ϕ)

  #first simulate the centered distribution
  #we will do this all in place

  #first generate the shocks
  #can't plug the covar matrix into MVnormal directly as its
  Xyt .= ar1.ΣcholL * rand!(MvNormal(K,1.0), Xyt)

  Xyt[:,1] .*= 1 ./ (1 .- ϕ.^2) .^ 0.5 # the first entry is simulated from the unconditional distribution

  #use a profiled breakpoint- large matrices its better to vectorize, small and not
  if K ≤ 3
    for t ∈ 2:size(Xyt,2)
      for r ∈ 1:size(Xyt,1)
        Xyt[r,t] += ϕ[r] * Xyt[r,t-1]
      end
    end
  else
    for t ∈ 2:size(Xyt,2)
      Xyt[:,t] .+= ϕ .* Xyt[:,t-1]
    end
  end
  Xyt .+= μ
  return nothing
end
#non-inplace version for testing
function simulatear1(ar1, N)
  res = Matrix{Float64}(undef, length(ar1.ϕ), N, )
  simulatear1!(ar1, res)
  return res'
end

function compareoncear1xy!(ar1::AR1, Xyt::AbstractMatrix,
  comparefunc::Tcomparefunc) where Tcomparefunc

  simulatear1!(ar1, Xyt)
  comparefunc(Xyt')
end

#this runs the simulations- basically a parametric bootstrap under the null
#statfunc is computed for each pass, while aggfunc is computed on the statistics
#at the end of the pass
function comparear1undernull(X::Tx,y::Ty,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-5,
    ) where {Tx<:AbstractMatrix{Float64}, Ty <: AbstractVector{Float64}, Tstatfunc, Taggfunc}

  Xy = hcat(X,y)
  K = size(X,2)
  ar1 = AR1(Xy)

  results = Dict(s => missings(Float64, Nsimulations, K) for s ∈ keys(statfunc(Xy)))
  @assert (k->ismissing.(results[k]) |> all).(keys(results)) |> all

  #this assumes stationary results are meaningless
  if checkstationarity && ((ar1.ϕ .≥ 1.0-stationaritytol) |> any)
    #display(ar1.Σ)
    #throw("ar1.Σ is not stationary!: eigen(ar1.Σ).values: $(eigen(ar1.Σ).values)")
    return (; Dict(s=>aggfunc(results[s], s) for s ∈ keys(results))...)
  end

  Xyts = [deepcopy(Xy' |> Matrix{Float64}) for t ∈ 1:Threads.nthreads()]

  #draw a sample and compute the aggregate function, sotring the reuslts in an array
  Threads.@threads for n ∈ 1:Nsimulations
  #WARNING: threading crashes it for some unknown reason
  #for n ∈ 1:Nsimulations

    t::Int = Threads.threadid()
    #res = compareoncear1xy!(xar1, xs[1], yar1, ys[1], statfunc)
    res = compareoncear1xy!(ar1, Xyts[t], statfunc;)

    #this structure allows us to compute multiple statistics
    #println(res)
    for s ∈ keys(res)

      @assert ismissing.(results[s][n, :]) |> all

      results[s][n,:] .= res[s] |> vec

    end
  end

  return (; Dict(s=>aggfunc(results[s], s) for s ∈ keys(results))...)
end

comparear1undernull(X::Tx,y::Ty,args...; kwargs...
    ) where {Tx <: AbstractMatrix, Ty<:AbstractVector} = comparear1undernull(
    X|>Matrix{Float64},y|>Vector{Float64},args...; kwargs...)

#this is just a wrapper for the 2-vector case
function compareundernull(x::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc, comparefunc::Tcomparefunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-6
    ) where {Tx<:AbstractVector, Tstatfunc, Taggfunc, Tcomparefunc}

  return comparefunc(reshape(x,length(x),1) |> Matrix{Float64}, y, statfunc, aggfunc;
    Nsimulations, checkstationarity, stationaritytol)
end

#wrapper for the wrapper
comparear1undernull(x::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4
    ) where {Tx<:AbstractVector, Tstatfunc, Taggfunc} = compareundernull(
      x,y,statfunc,aggfunc, comparear1undernull;
      Nsimulations, checkstationarity, stationaritytol)

#=function comparearima11undernull(X::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4,
    dXts = (size(X,2) > 1 ?
      [similar(X[2:end, :]' |> Matrix{Float64}) for t ∈ 1:Threads.nthreads()] :
      [similar(X[2:end, :] |> vec |> Vector{Float64}) for t ∈ 1:Threads.nthreads()]),
    dys = [similar(y[2:end]) for t ∈ 1:Threads.nthreads()],
    ) where {Tx<:AbstractMatrix, Tstatfunc, Taggfunc}

  #difference everything
  dX = reduce(hcat, (xi->xi[2:end] .- xi[1:(end-1)]).(eachcol(X)))

  #handle the intercept as a special case
  for (dxi, xi) ∈ zip(eachcol(dX), eachcol(X))
    if maximum(xi) - minimum(xi) + 1.0 ≈ 1.0
      dxi .= mean(xi)
    end
  end

  dx = dX |> eachcol |> collect
  dy = y[2:end] .- y[1:(end-1)]

  K = length(dx)

  dxar1 = (AR1.(dx))# .|> centerar1
  dyar1 = AR1(dy)# |> centerar1

  results = Dict(s => missings(Float64, Nsimulations, K) for s ∈ keys(statfunc(dX,dy)))
  @assert (k->ismissing.(results[k]) |> all).(keys(results)) |> all

  #this assumes stationary results are meaningless
  if checkstationarity && max((dxar1i -> abs(dxar1i.ϕ)).(dxar1)..., abs(dyar1.ϕ)) ≥ 1.0-stationaritytol
    return (; Dict(s=>aggfunc(results[s], s) for s ∈ keys(results))...)
  end

  ϕdx::Vector{Float64} = dxar1 .|> dxar1i->dxar1i.ϕ
  σdx::Vector{Float64} = dxar1 .|> dxar1i->dxar1i.σ²^0.5
  μdx::Vector{Float64} = dxar1 .|> dxar1i->dxar1i.μ

  #draw a sample and compute the aggregate function, sotring the reuslts in an array
  Threads.@threads for n ∈ 1:Nsimulations
  #for n ∈ 1:Nsimulations

    t::Int = Threads.threadid()
    #res = compareoncear1xy!(xar1, xs[1], yar1, ys[1], statfunc)
    res = compareoncear1xy!(dXts[t], dyar1, dys[t], statfunc; ϕx=ϕdx, σx=σdx, μx=μdx)

    #this structure allows us to compute multiple statistics
    for s ∈ keys(res)

      @assert ismissing.(results[s][n, :]) |> all
      results[s][n,:] .= res[s]
    end
  end

  return (; Dict(s=>aggfunc(results[s], s) for s ∈ keys(results))...)
end=#

#this runs the simulations- basically a parametric bootstrap under the null
#statfunc is computed for each pass, while aggfunc is computed on the statistics
#at the end of the pass
function comparearima11undernull(X::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-5,
    ) where {Tx<:AbstractMatrix, Tstatfunc, Taggfunc}

  Xy = hcat(X,y)
  dXy = Xy[2:end,:] .- Xy[1:(end-1),:]
  K = size(X,2)

  #handle the intercept as a special case
  for (dxi, xi) ∈ zip(eachcol(dXy), eachcol(Xy))
    if maximum(xi) - minimum(xi) + 1.0 ≈ 1.0
      dxi .= mean(xi)
    end
  end

  ar1 = AR1(dXy)

  results = Dict(s => missings(Float64, Nsimulations, K) for s ∈ keys(statfunc(dXy)))
  @assert (k->ismissing.(results[k]) |> all).(keys(results)) |> all

  #this assumes stationary results are meaningless
  if checkstationarity &&  ((ar1.ϕ .≥ 1.0-stationaritytol) |> any)
    return (; Dict(s=>aggfunc(results[s], s) for s ∈ keys(results))...)
  end

  dXyts = [deepcopy(dXy' |> Matrix{Float64}) for t ∈ 1:Threads.nthreads()]

  #draw a sample and compute the aggregate function, sotring the reuslts in an array
  Threads.@threads for n ∈ 1:Nsimulations
  #WARNING: threading crashes it for some unknown reason
  #for n ∈ 1:Nsimulations

    t::Int = Threads.threadid()
    #res = compareoncear1xy!(xar1, xs[1], yar1, ys[1], statfunc)
    res = compareoncear1xy!(ar1, dXyts[t], statfunc;)

    #this structure allows us to compute multiple statistics
    #println(res)
    for s ∈ keys(res)

      @assert ismissing.(results[s][n, :]) |> all
      results[s][n,:] .= res[s] |> vec
    end
  end

  return (; Dict(s=>aggfunc(results[s], s) for s ∈ keys(results))...)
end

comparearima11undernull(X::Tx,y::Ty,args...; kwargs...
    ) where {Tx <: AbstractMatrix, Ty<:AbstractVector} = comparearima11undernull(
    X|>Matrix{Float64},y|>Vector{Float64},args...; kwargs...)

comparearima11undernull(x::Tx,y::AbstractVector,
    statfunc::Tstatfunc, aggfunc::Taggfunc; Nsimulations::Int,
    checkstationarity::Bool = true, stationaritytol::Float64=10^-4
    ) where {Tx<:AbstractVector, Tstatfunc, Taggfunc} = compareundernull(
      x,y,statfunc,aggfunc, comparearima11undernull;
      Nsimulations, checkstationarity, stationaritytol)

nonparametricnullp(v::AbstractVector{T}, x) where T<:Real = (
  x ≥ 0 ? sum(v .> x)/length(v) : sum(v .< x)/length(v))
nonparametricnullp(v::AbstractVector, x) = nonparametricnullp(v |> skipmissing |> collect, x)
nonparametricnullp(M::AbstractMatrix, b) = (
  (v,bᵢ)->nonparametricnullp(v |> skipmissing |> collect, bᵢ)).(eachcol(M),b)

#this creates a cholesky decomposition for the covariance of a data matrix
#under the special case where one of the columns is all constant
#the purpose is to use this to generate sensible random numbers
function semicovcholl(X::TM) where TM
  K = size(X,2)

  #determine the columns that are the same
  defcols = mapslices(c->!(maximum(c)-minimum(c) + 1.0 ≈ 1.0), X, dims=1) |> vec

  #compute the full covariance matrix
  defΣ = cov(@view X[:, defcols])
  defL = cholesky(defΣ).L

  ΣcholL = zeros(eltype(TM), K, K)

  ΣcholL[defcols, defcols] .= defL

  return ΣcholL
end

function testsemicovcholl()
  Xdat = rand(Normal(),10^4,5) .+ (rand(Normal(),10^4) .* [1,2,3,4,5]') .+ collect(1:5)'
  X = deepcopy(Xdat)
  X[:,4] .= 1
  Σcheck = cov(X)
  ΣcholLcheck = deepcopy(Σcheck)
  defcols = [1,2,3,5]
  ΣcholLcheck[defcols,defcols] .= cholesky(cov(X[:,defcols])).L

  #check the computational integrity
  ΣcholL = semicovcholl(X)
  @assert (Σcheck .≈ ΣcholL * ΣcholL') |> all
  @assert (ΣcholLcheck .≈ ΣcholL) |> all

  #now check what we plan to do with it- random number generation
  X[:,4] .= Xdat[:,4]
  Σcheck = cov(X)
  ΣcholL = semicovcholl(X)
  @assert (Σcheck .≈ ΣcholL * ΣcholL') |> all

  μs = mean(X,dims=1) |> vec
  simcheck = rand(MvNormal(μs, Σcheck), 10^6)'
  sim = (ΣcholL * rand(Normal(0,1),5, 10^6).+ μs)'

  atol = 0.02
  @assert ((μ, μcheck)->isapprox(μ, μcheck; atol)
    ).(mean(sim, dims=1), mean(simcheck, dims=1)) |> all "
      mean(sim, dims=1) = $(mean(sim, dims=1))
      mean(simcheck, dims=1) = $(mean(simcheck, dims=1))"

  if !all(((σij, σcheckij)->isapprox(σij, σcheckij; atol=atol^0.5)).(cov(sim), cov(simcheck)))
    @info "cov(sim)"
    display(cov(sim))
    @info "cov(simcheck)"
    display(cov(simcheck))
    throw("failed simulation check on Σ")
  end
  @info "passed semicovcholl check"

end

function testar1(;N,K)
  testsemicovcholl()

  ϕ = [[0.8/k for k ∈ 1:K]; 0.8; ]

  #simulate a covariance matrix
  ND = max(N, min(N^2, 10^6))
  D = rand(Normal(0,1), ND,K+1) .+ [1 ./ collect(1:K);1.0;]' .* rand(Normal(0,1), ND)
  #Σtrue = reduce(vcat, 4 .* [[1/(k+r) for k ∈ 1:(K+1)]' for r ∈ 1:(K+1)])
  Σtrue  = cov(D)
    #println(Σtrue)

  μtrue = [1.0 for k ∈ 1:(K+1)]
  ε = (rand(MvNormal(Σtrue), N) .+ μtrue)'
  Xy = missings(Float64, N,K+1)
  Xy[1,:] .= ε[1]

  for t ∈ 2:N
    Xy[t,:] .= Xy[t-1, :] .* ϕ .+ ε[t, :]
  end
  global Xy = Xy

  #assign the components
  X = Xy[:,1:(end-1)]
  y = Xy[:,end]
  @info "ϕ=$ϕ, b=$(mean(ε, dims=1)), μ=$(mean(Xy, dims=1))"

  atol = 1/N^0.5
  actar1=AR1(;ϕ=[ϕ[1]], μ=[2.5], Σ=reshape([2.0], 1,1))
  xsim = simulatear1(actar1, 10^6)
  @assert isapprox(mean(xsim), 2.5; atol)  "mean(xsim): $(mean(xsim))"
  @assert isapprox(var(xsim), 2/(1-ϕ[1]^2), atol=atol^0.5) "var(xsim): $(var(xsim)), expected=$(2/(1-ϕ[1]^2))"
  println("size of xsim: $(size(xsim))")
  ar1sim = AR1(xsim)
  @info "ar1sim: $(ar1sim)"
  @assert isapprox(ar1sim.ϕ[1], ϕ[1]; atol)
  @assert isapprox(ar1sim.μ[1], 2.5; atol)
  @assert isapprox(ar1sim.Σ[1], 2; atol)

  Nsimulations=10^4
  dmeantest = 0.2
  #recall the compare func will pass the name of the relevant stat func when it is called
  aggfunc(v, s) = nonparametricnullp(v, Dict(:cor=>0.5, :dmean=>dmeantest, :beta=>1.0)[s])
  function statfunc(Xy)
    Xs = view(Xy,:, 1:(size(Xy,2)-1))
    ys = view(Xy, :, size(Xy,2))

    res = Dict(:beta=>cholesky!(Symmetric(Xs'*Xs))\(Xs'*ys),
      :cor=>cor(Xs, ys),
      :dmean=>mean(Xs, dims=1) .- mean(ys))

    return res
  end

  @info "std(Xy, dims=1): $(std(Xy, dims=1))"
  @info "expected std(Xy): $((2 ./ (1 .- ϕ.^2)).^0.5)"

  #derive the distribution for the difference in means of two AR(1) dists
  R = [0.8^abs(i-j) for i ∈ 1:N, j ∈ 1:N]
  sidmean = (1/N^2*(ones(N)' * ( R .* 2/(1-ϕ[1]^2)) * ones(N)))^0.5
  ρdmean = Σtrue[1,end] ./ (Σtrue[1,1]*Σtrue[end,end])^0.5
  #because its a difference, use a negative sign for the t-stat
  Sdmean = (2*sidmean^2-2*ρdmean*sidmean^2)^0.5
  Epdmean = cdf(Normal(), - dmeantest/Sdmean)

  #println("X[1:5,1]: $(X[1:5,1])")
  x=X[:,1]
  res = comparear1undernull(x, y, statfunc, aggfunc; Nsimulations)
  @info "compareundernull result: $res (Epdmean: $Epdmean)"

  Xtest = X#[:,1]
  @btime comparear1undernull($Xtest, $y, $statfunc, $aggfunc; Nsimulations=$Nsimulations)
  #@btime comparear1undernull($x, $y, $statfunc, $aggfunc; Nsimulations=$Nsimulations)
  #throw("simulation complete")

  X = hcat(X[:,1:2], ones(N)) |> Matrix{Float64}
  function mstatfunc(Xy)
    Xs = view(Xy,:, 1:(size(Xy,2)-1))
    ys = view(Xy, :, size(Xy,2))

    try
      return Dict(
        :beta=>cholesky!(Symmetric(Xs'*Xs))\(Xs'*ys),
        :cor=>cor(Xs,ys))
    catch err
      global X = X
      global y = y
      throw("$err:
      iter: $iter
      Xy[1:5,:]: $(Xy[1:5, :])")
    end
  end
  maggfunc(M,s) = (v->nonparametricnullp(v, Dict(:cor=>0.5, :beta=>0.3)[s])).(eachcol(M))

  res = comparear1undernull(X, y, mstatfunc, maggfunc; Nsimulations)
  @info "compareundernull matrix result: $res"

  ix = cumsum(X[:,1])
  iy = cumsum(y)
  res = comparear1undernull(ix, iy, statfunc, aggfunc; Nsimulations)
  @info "integrated compareundernull (mispecified) dmean: $res"
  res = comparearima11undernull(ix, iy, statfunc, aggfunc; Nsimulations)
  @info "integrated compareundernull dmean: $res"
  res = comparearima11undernull(X, y, mstatfunc, maggfunc; Nsimulations)
  @info "compareundernull matrix result: $res"

end

#@time testar1(N=1000, K=3)

ext: jl

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
function incident_field(D, f, n, gridL, cellL)
    ω = 2π * f
    k = n * ω

    function efield(x, y)
        r = √(x^2 + y^2 + D^2)
        ℯ ^ (k * r * im) / (4π * r)
    end

    grid = range(-gridL/2 + 1/2, gridL/2 - 1/2, length = gridL) .* cellL
    [efield(x, y) for x in grid, y in grid]
end

"""
    greens(D, f, ϵ, μ, gridL, cellL)

Returns the Green's function for
the near to far transformation.

# Arguments
- `D`: distance between near and far plane
- `f`: frequency
- `ϵ`: relative permittivity
- `μ`: relative permeability
- `gridL`: side length of E-field grid (in cells)
- `cellL`: length of each cell of grid
"""
function greens(D, f, ϵ, μ, gridL, cellL)
    ω = 2π * f
    n = √(ϵ*μ)
    k = n * ω

    function efield(x, y)
        r = √(x^2 + y^2 + D^2)
        D * (-1 + k * r * im) * ℯ ^ (k * r * im) / (4 * π * r^3)
    end

    gridout = range(-gridL, gridL - 1, length = gridL * 2) .* cellL
    g = [efield(x, y) * -μ / ϵ for x in gridout, y in gridout]

    g
end


text:
# Sequence rules over the [`SimData`](@ref) object,
# calling [`maprule!`](@ref) for each individual `Rule`.
function sequencerules!(simdata::AbstractSimData)
    newsimdata = sequencerules!(simdata, rules(simdata))
    # We run masking here to mask out cells that are `false` in the
    # `mask` array, if it exists. Not all rules run masking, so it is
    # applied here so that the final grid is always masked.
    _maybemask!(grids(newsimdata))
    newsimdata
end
function sequencerules!(simdata::AbstractSimData, rules::Tuple)
    rule = rules[1]
    rest = tail(rules)
    # Run the first rules
    newsimdata = maprule!(simdata, rule)
    # Run the rest of the rules recursively
    sequencerules!(newsimdata, rest)
end
sequencerules!(simdata::AbstractSimData, rules::Tuple{}) = simdata


text:
module EquivariantOperators

include("operators.jl")
include("plotutils.jl")
export Grid, get,put!
export xcor,conv
export Op,Radfunc,remake!,nae
export vfplot
#
end # module


text:
# This file is part of Kpax3. License is MIT.

function test_loss_binder()
  P = [1.000 0.571 0.458 0.458 0.442 0.357;
       0.571 1.000 0.458 0.458 0.442 0.357;
       0.458 0.458 1.000 0.571 0.357 0.442;
       0.458 0.458 0.571 1.000 0.357 0.442;
       0.442 0.442 0.357 0.357 1.000 0.426;
       0.357 0.357 0.442 0.442 0.426 1.000]

  # k = 1
  R = Int[1; 1; 1; 1; 1; 1]
  A = Float64[1 1 1 1 1 1;
              1 1 1 1 1 1;
              1 1 1 1 1 1;
              1 1 1 1 1 1;
              1 1 1 1 1 1;
              1 1 1 1 1 1]

  losscorrect = sum(abs.(A - P)) / 2
  lossfunction = Kpax3.loss_binder(R, P)

  @test isapprox(lossfunction, losscorrect, atol=ε)

  # k = 2
  R = Int[1; 1; 2; 2; 1; 2]
  A = Float64[1 1 0 0 1 0;
              1 1 0 0 1 0;
              0 0 1 1 0 1;
              0 0 1 1 0 1;
              1 1 0 0 1 0;
              0 0 1 1 0 1]

  losscorrect = sum(abs.(A - P)) / 2
  lossfunction = Kpax3.loss_binder(R, P)

  @test isapprox(lossfunction, losscorrect, atol=ε)

  R = Int[2; 1; 2; 1; 2; 1]
  A = Float64[1 0 1 0 1 0;
              0 1 0 1 0 1;
              1 0 1 0 1 0;
              0 1 0 1 0 1;
              1 0 1 0 1 0;
              0 1 0 1 0 1]

  losscorrect = sum(abs.(A - P)) / 2
  lossfunction = Kpax3.loss_binder(R, P)

  @test isapprox(lossfunction, losscorrect, atol=ε)

  # k = 3
  R = Int[1; 1; 2; 2; 3; 2]
  A = Float64[1 1 0 0 0 0;
              1 1 0 0 0 0;
              0 0 1 1 0 1;
              0 0 1 1 0 1;
              0 0 0 0 1 0;
              0 0 1 1 0 1]

  losscorrect = sum(abs.(A - P)) / 2
  lossfunction = Kpax3.loss_binder(R, P)

  @test isapprox(lossfunction, losscorrect, atol=ε)

  R = Int[1; 2; 3; 1; 2; 3]
  A = Float64[1 0 0 1 0 0;
              0 1 0 0 1 0;
              0 0 1 0 0 1;
              1 0 0 1 0 0;
              0 1 0 0 1 0;
              0 0 1 0 0 1]

  losscorrect = sum(abs.(A - P)) / 2
  lossfunction = Kpax3.loss_binder(R, P)

  @test isapprox(lossfunction, losscorrect, atol=ε)

  # k = 4
  R = Int[1; 1; 2; 2; 3; 4]
  A = Float64[1 1 0 0 0 0;
              1 1 0 0 0 0;
              0 0 1 1 0 0;
              0 0 1 1 0 0;
              0 0 0 0 1 0;
              0 0 0 0 0 1]

  losscorrect = sum(abs.(A - P)) / 2
  lossfunction = Kpax3.loss_binder(R, P)

  @test isapprox(lossfunction, losscorrect, atol=ε)

  R = Int[1; 2; 3; 4; 3; 2]
  A = Float64[1 0 0 0 0 0;
              0 1 0 0 0 1;
              0 0 1 0 1 0;
              0 0 0 1 0 0;
              0 0 1 0 1 0;
              0 1 0 0 0 1]

  losscorrect = sum(abs.(A - P)) / 2
  lossfunction = Kpax3.loss_binder(R, P)

  @test isapprox(lossfunction, losscorrect, atol=ε)

  # k = 5
  R = Int[1; 2; 3; 3; 4; 5]
  A = Float64[1 0 0 0 0 0;
              0 1 0 0 0 0;
              0 0 1 1 0 0;
              0 0 1 1 0 0;
              0 0 0 0 1 0;
              0 0 0 0 0 1]

  losscorrect = sum(abs.(A - P)) / 2
  lossfunction = Kpax3.loss_binder(R, P)

  @test isapprox(lossfunction, losscorrect, atol=ε)

  R = Int[1; 2; 3; 4; 1; 5]
  A = Float64[1 0 0 0 1 0;
              0 1 0 0 0 0;
              0 0 1 0 0 0;
              0 0 0 1 0 0;
              1 0 0 0 1 0;
              0 0 0 0 0 1]

  losscorrect = sum(abs.(A - P)) / 2
  lossfunction = Kpax3.loss_binder(R, P)

  @test isapprox(lossfunction, losscorrect, atol=ε)

  # k = 6
  R = Int[1; 2; 3; 4; 5; 6]
  A = Float64[1 0 0 0 0 0;
              0 1 0 0 0 0;
              0 0 1 0 0 0;
              0 0 0 1 0 0;
              0 0 0 0 1 0;
              0 0 0 0 0 1]

  losscorrect = sum(abs.(A - P)) / 2
  lossfunction = Kpax3.loss_binder(R, P)

  @test isapprox(lossfunction, losscorrect, atol=ε)

  nothing
end

test_loss_binder()


text:
function threshold!(A::AbstractArray, ϵ)
    for i in eachindex(A)
        if abs(A[i]) < ϵ A[i] = 0 end
    end
    A
end

using FastTransforms, LinearAlgebra, Random, Test

# The colatitudinal grid (mod π):
N = 10
θ = (0.5:N-0.5)/N

# The longitudinal grid (mod π):
M = 2*N-1
φ = (0:M-1)*2/M

x = [cospi(φ)*sinpi(θ) for θ in θ, φ in φ]
y = [sinpi(φ)*sinpi(θ) for θ in θ, φ in φ]
z = [cospi(θ) for θ in θ, φ in φ]

P = plan_sph2fourier(Float64, N)
PA = plan_sph_analysis(Float64, N, M)
J = FastTransforms.plan_sph_isometry(Float64, N)

f = (x, y, z) -> x^2+y^4+x^2*y*z^3-x*y*z^2

F = f.(x, y, z)
V = PA*F
U = threshold!(P\V, 100eps())
FastTransforms.execute_sph_yz_axis_exchange!(J, U)
FR = f.(x, -z, -y)
VR = PA*FR
UR = threshold!(P\VR, 100eps())
@test U ≈ UR
norm(U-UR)

α, β, γ = 0.123, 0.456, 0.789

# Isometry built up from ZYZR
A = [cos(α) -sin(α) 0; sin(α) cos(α) 0; 0 0 1]
B = [cos(β) 0 -sin(β); 0 1 0; sin(β) 0 cos(β)]
C = [cos(γ) -sin(γ) 0; sin(γ) cos(γ) 0; 0 0 1]
R = diagm([1, 1, 1.0])
Q = A*B*C*R

# Transform the sampling grid. Note that `Q` is transposed here.
u = Q[1,1]*x + Q[2,1]*y + Q[3,1]*z
v = Q[1,2]*x + Q[2,2]*y + Q[3,2]*z
w = Q[1,3]*x + Q[2,3]*y + Q[3,3]*z

F = f.(x, y, z)
V = PA*F
U = threshold!(P\V, 100eps())
FastTransforms.execute_sph_rotation!(J, α, β, γ, U)
FR = f.(u, v, w)
VR = PA*FR
UR = threshold!(P\VR, 100eps())
@test U ≈ UR
norm(U-UR)

F = f.(x, y, z)
V = PA*F
U = threshold!(P\V, 100eps())
FastTransforms.execute_sph_polar_reflection!(U)
FR = f.(x, y, -z)
VR = PA*FR
UR = threshold!(P\VR, 100eps())
@test U ≈ UR
norm(U-UR)

# Isometry built up from planar reflection
W = [0.123, 0.456, 0.789]
H = w -> I - 2/(w'w)*w*w'
Q = H(W)

# Transform the sampling grid. Note that `Q` is transposed here.
u = Q[1,1]*x + Q[2,1]*y + Q[3,1]*z
v = Q[1,2]*x + Q[2,2]*y + Q[3,2]*z
w = Q[1,3]*x + Q[2,3]*y + Q[3,3]*z

F = f.(x, y, z)
V = PA*F
U = threshold!(P\V, 100eps())
FastTransforms.execute_sph_reflection!(J, W, U)
FR = f.(u, v, w)
VR = PA*FR
UR = threshold!(P\VR, 100eps())
@test U ≈ UR
norm(U-UR)

F = f.(x, y, z)
V = PA*F
U = threshold!(P\V, 100eps())
FastTransforms.execute_sph_reflection!(J, (W[1], W[2], W[3]), U)
FR = f.(u, v, w)
VR = PA*FR
UR = threshold!(P\VR, 100eps())
@test U ≈ UR
norm(U-UR)

# Random orthogonal transformation
Random.seed!(0)
Q = qr(rand(3, 3)).Q

# Transform the sampling grid, note that `Q` is transposed here.
u = Q[1,1]*x + Q[2,1]*y + Q[3,1]*z
v = Q[1,2]*x + Q[2,2]*y + Q[3,2]*z
w = Q[1,3]*x + Q[2,3]*y + Q[3,3]*z

F = f.(x, y, z)
V = PA*F
U = threshold!(P\V, 100eps())
FastTransforms.execute_sph_orthogonal_transformation!(J, Q, U)
FR = f.(u, v, w)
VR = PA*FR
UR = threshold!(P\VR, 100eps())
@test U ≈ UR
norm(U-UR)


text:
import Limbdark

nu = 20
u=ones(nu)/nu

r=0.1
b=0.5

t = Limbdark.transit_init(r,b,u,false)

using PyPlot

plot(t.g_n)
xlabel(L"n")
ylabel(L"$g_n$")

savefig("g_n_vs_n.pdf",bbox_inches="tight")


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

using Markdown
using InteractiveUtils

# ╔═╡ 9608fdc8-807e-4be2-91f8-950c0ceec2dd
begin
	using Test
	using PlutoUI
end

# ╔═╡ 45d72214-919c-11eb-2e69-23c0786baacb
md"""
### Subtyping

Src: [Julia 1.0 Cookbook -  Bogumił Kamiński , Przemysław Szufel - 2018](https://www.packtpub.com/product/julia-1-0-programming-cookbook/9781788998369)

**Context**
A Gaussian integer is a complex number whose real and imaginary parts are both integers.
"""

# ╔═╡ 82503139-3e37-4bf5-99f6-daf686f591e8
md"""
Let us start with defining a custom `Point` type, and two outer constructors for this type:
"""

# ╔═╡ 25b36455-0e5a-4318-9bd3-0440ac0b0080
begin

	struct Point{T <: Integer, S<:AbstractString}
		pos::Complex{T}
		label::S
	end

	Point(x::T, y::T, label::S) where {T <: Integer, S <: AbstractString} =
		Point{T, S}(Complex(x, y), label)

	Point(x, y, label) = Point(promote(Integer.((x, y))...)..., label)

end

# ╔═╡ c07b43bd-4aff-4e0a-9a34-12bc8a49e8ca
begin
	p₁ = Point(1, 0, "1")
	p₂ = Point(1, 0, SubString("1", 1))
	p₃ = Point(true, false, "1")         ## will be promoted
	p₄ = Point(2, 0, "2")

	ary = Point[p₁, p₂, p₃, p₄]
	(p₁, p₂, p₃, p₄)
end

# ╔═╡ ced242ca-2352-4b3f-bf32-a095705756ef
md"""
let us write a method that takes an array of Point and returns a new Point type that is the sum of passed Point instances with an empty label:
"""

# ╔═╡ f886d114-114e-4f66-8d5c-9edd5e1fdfe2
function sum_points_v₁(ary::Vector{Point})::Point
	Point(sum(getfield.(ary, :pos)), "")
end

# ╔═╡ 81840539-f363-414d-a263-227dcb10abb0
function sum_points_v₂(ary::Vector{Point})::Point
	Point(sum(p.pos for p in ary), "")
end

# ╔═╡ e4d9bae1-1a63-45c6-97a2-9d9fa54f3b54
sum_points_v₁(ary)  ## slower than  sum_points_v₂

# ╔═╡ 20bdff6e-c631-4de0-8631-9a9de2f82edb
md"""
Expecting an **error** while invoking `sum_points_v₂([p₁, p₂])`, but why?
"""

# ╔═╡ d8e2b4db-ec24-411a-9a62-d7f2a6b4d078
## with Reactive cell - the test will pass and then fail because of added
## definition to fix the problem later

# @test_throws MethodError sum_points_v₂([p₁, p₂])

# ╔═╡ c2745a46-6258-4b92-b7bb-2a770c9273ec
md"""
our `Vector{Point{Int64, S}` (see cell below) does not match the function signature expectation, which is `Vector{Point}`.
"""

# ╔═╡ a70143ea-4920-413f-a6c1-515eef6c4f78
typeof([p₁, p₂])

# ╔═╡ 268af4e1-eb7a-4c0c-959c-c434c8e094b4
md"""
vs using `Point[...]`, which is fine:
"""

# ╔═╡ d74742e6-1aed-47fa-8c29-0e31c4459e3f
md"""
OK, so let us modified our function signature as follows. Note the use of `AbstractVector`.
"""

# ╔═╡ 6fa1df97-16e4-4b45-814c-cba39ebf9efa
function sum_points_v₂(ary::AbstractVector{<: Point})::Point
	Point(sum(p.pos for p in ary), "")
end

# ╔═╡ bc1b65b1-5063-4567-8928-f9847f0382b4
sum_points_v₂(ary)

# ╔═╡ e9a32b85-e697-4a78-8cbb-0d242e86d8fe
sum_points_v₂(Point[p₁, p₂])

# ╔═╡ 224a9eaa-7062-481f-8af3-098ba1988c5d
md"""
And now we can see this is working (cf. explanation in later cells):
"""

# ╔═╡ e432f0a0-45de-4f23-aab7-09a30b1b7bdc
sum_points_v₁(ary)

# ╔═╡ d7f7ab10-f69b-4143-962c-72ea4b764e40
md"""
##### Remarks:
When storing parametric types in collections, the standard constructors try to find the narrowest possible type that can hold the passed data. This means that if we can expect to encounter heterogeneous data types in the collection, then we *should manually specify an appropriate type*.
For instance, the first example shown below fails, while the second works:
"""

# ╔═╡ 424a9736-1c65-440a-9079-5fd63c74fd50
@test_throws MethodError push!([p₁], p₂)

# ╔═╡ 810473f2-d8c2-44e5-811b-185a445694a9
@test push!(Point[p₁], p₂) == Point[Point{Int64, String}(1 + 0im, "1"),
	Point{Int64, SubString{String}}(1 + 0im, "1")]

# ╔═╡ 4b8455fd-b54c-4eee-8e95-612b3132a0bc
md"""
##### Methods

Investigating `Julia` subtyping algorithm in method dispatch.$(html"<br />")
Let us define another function `foo` which takes a single Point:
"""

# ╔═╡ 87e8df24-dc9d-42bd-a38f-9d5f16533e75
begin
	foo(p::Point) = "Generic Definition."

	foo(p::Point{Int, <: AbstractString}) = "Default Int Passed."

	foo(p::Point{<: Integer, String}) = "Default String Passed."
end

# ╔═╡ 639755ac-0d13-4087-b632-223ced8d334c
md"""
The last test shows that it is necessary to define a specific method for the intersection of types: `foo(p::Point{Int, String}) = "Most Specific Method."`:
"""

# ╔═╡ 6bcfbc7a-fa47-4471-bedc-071bbd58decc
foo(p::Point{Int, String}) = "Most Specific Method."

# ╔═╡ cfd030ff-ddc5-46d8-b478-56a87d2f8f11
with_terminal() do
	@show methods(foo)
end

# ╔═╡ af97433f-e26d-4178-97dd-1f580b9c66d3
begin
	@test foo(Point(true, true, s"12")) == "Generic Definition."

	@test foo(Point(1, 1, s"12")) == "Default Int Passed."

	@test foo(Point(true, true, "12")) == "Default String Passed."

	##
	## OK: with reactive cells - this will fail one the specific foo method is added!
	##
	## @test_throws MethodError foo(Point(1, 1, "12")) # == "Default String Passed"
	##
	## => foo(::Main.workspace36.Point{Int64, String}) is ambiguous. Candidates:
	## Possible fix, define
	##
	## foo(::Main.workspace36.Point{Int64, String})
end

# ╔═╡ cdc06666-2646-46cc-9a45-890ae5c88abe
with_terminal() do
	@show methods(foo)
end

# ╔═╡ 42a5ca4c-cd2c-4ced-a59c-638d7429c2c7
@test foo(Point(1, 1, "12")) == "Most Specific Method."

# ╔═╡ 03b637ec-6537-4019-9332-a9881a4ad54c
md"""
##### Remarks:

In `Julia`, most parametric types are invariant.
**This means that subtyping of type parameters does not translate to subtyping of parametric types**.

Here is an example - consider the following:
"""

# ╔═╡ 3f337631-6d8a-4618-9e9f-a0f0a1eb5438
begin
	@test Int <: Integer                                  ## => true
	@test !(Point{Int, String} <: Point{Integer, String}) ## NOTE the ! ≡ negation
end

# ╔═╡ 5fc62545-fc1f-4385-9a79-bccdaa665953
md"""
In order to allow for subtyping of types in function signatures, the <: operator is  required:
"""

# ╔═╡ 37a08ce3-7b85-4540-a080-ec1f47c3bf6a
begin
	@test Point{Int, String} <: Point{<:Integer, String} ## NOTE: <:Integer

	# equivalent:
	@test Point{Int, String} <: Point{T, String} where {T <: Integer}
end

# ╔═╡ 122211d1-4ba1-440f-a632-b4b90e47302a
md"""
If a parametric type has more than one parameter, we can partially restrict them to a subset. Consider the following examples:
"""

# ╔═╡ 428c43cb-2c9b-4c1f-9509-1bd93af3d0bb
Point{Int}

# ╔═╡ b3056d5a-7790-4b94-9c77-854087a01cb2
Point{<: Signed, String}

# ╔═╡ b0162138-03e1-4dcc-b76c-073e6bea17d9
Point{Int}{String}, Point{Int, String}
## equivalent

# ╔═╡ 61a7e87a-9dc1-4519-ad78-51659c383496
md"""
We can see that type invariance is important in the definition of the `sum_points_v₂` function, as it only accepts subtypes of `AbstractVector{Point}` and `[p1, p2]` creates `Vector{Point{Int}}`, which is *NOT its subtype*.

However, the following statement works: `sum_points_v₂(Point[p1, p2])` because we have explicitly set the type of the vector holding p1 and p2. This is a common pitfall in Julia.

Fortunately, it is simple to solve this using the `<:` operator, as shown above with the defintion of `sum_points_v₂(ary::AbstractVector{<: Point})` above...
"""

# ╔═╡ a29183fa-99eb-4bc9-bffd-4bf8fe1c1cb1
md"""
The preceding type invariance rules are NOT applicable to **tuples that are covariant**. This means that tuple subtyping behaves as if the `<:` operator were added to every variable in the tuple definition.

For instance, consider this example (it has been made complex on purpose to show that subtyping of every entry of Tuple type is subtyped):

"""

# ╔═╡ 43655522-24a7-4497-a305-4cb1a53e8791
@test Tuple{Point{Int, String}, Point{Bool, SubString{String}}} <: Tuple{Point{Int}, Point}   ## true

# ╔═╡ 8a1e851c-8e86-4e56-86be-089667044cd4
md"""
The consequence is that with tuples, we do not specify `<:` to get the subtyping behavior we discussed earlier:
"""

# ╔═╡ f40cdaeb-8474-4bae-a46f-44e2199e3712
sum_point_tuple(v::Tuple{Vararg{Point}}) = Point(sum(p.pos for p in v), "")

# ╔═╡ 4c940780-fbd3-44dd-af56-2559809caa73
@test sum_point_tuple((p₁, p₂, p₃)) == Point{Int64,String}(3 + 0im, "")

# ╔═╡ Cell order:
# ╟─45d72214-919c-11eb-2e69-23c0786baacb
# ╠═9608fdc8-807e-4be2-91f8-950c0ceec2dd
# ╟─82503139-3e37-4bf5-99f6-daf686f591e8
# ╠═25b36455-0e5a-4318-9bd3-0440ac0b0080
# ╠═c07b43bd-4aff-4e0a-9a34-12bc8a49e8ca
# ╟─ced242ca-2352-4b3f-bf32-a095705756ef
# ╠═f886d114-114e-4f66-8d5c-9edd5e1fdfe2
# ╠═81840539-f363-414d-a263-227dcb10abb0
# ╠═bc1b65b1-5063-4567-8928-f9847f0382b4
# ╠═e4d9bae1-1a63-45c6-97a2-9d9fa54f3b54
# ╟─20bdff6e-c631-4de0-8631-9a9de2f82edb
# ╠═d8e2b4db-ec24-411a-9a62-d7f2a6b4d078
# ╟─c2745a46-6258-4b92-b7bb-2a770c9273ec
# ╠═a70143ea-4920-413f-a6c1-515eef6c4f78
# ╟─268af4e1-eb7a-4c0c-959c-c434c8e094b4
# ╠═e9a32b85-e697-4a78-8cbb-0d242e86d8fe
# ╟─d74742e6-1aed-47fa-8c29-0e31c4459e3f
# ╠═6fa1df97-16e4-4b45-814c-cba39ebf9efa
# ╟─224a9eaa-7062-481f-8af3-098ba1988c5d
# ╠═e432f0a0-45de-4f23-aab7-09a30b1b7bdc
# ╟─d7f7ab10-f69b-4143-962c-72ea4b764e40
# ╠═424a9736-1c65-440a-9079-5fd63c74fd50
# ╠═810473f2-d8c2-44e5-811b-185a445694a9
# ╟─4b8455fd-b54c-4eee-8e95-612b3132a0bc
# ╠═87e8df24-dc9d-42bd-a38f-9d5f16533e75
# ╠═cfd030ff-ddc5-46d8-b478-56a87d2f8f11
# ╠═af97433f-e26d-4178-97dd-1f580b9c66d3
# ╟─639755ac-0d13-4087-b632-223ced8d334c
# ╠═6bcfbc7a-fa47-4471-bedc-071bbd58decc
# ╠═cdc06666-2646-46cc-9a45-890ae5c88abe
# ╠═42a5ca4c-cd2c-4ced-a59c-638d7429c2c7
# ╟─03b637ec-6537-4019-9332-a9881a4ad54c
# ╠═3f337631-6d8a-4618-9e9f-a0f0a1eb5438
# ╟─5fc62545-fc1f-4385-9a79-bccdaa665953
# ╠═37a08ce3-7b85-4540-a080-ec1f47c3bf6a
# ╟─122211d1-4ba1-440f-a632-b4b90e47302a
# ╠═428c43cb-2c9b-4c1f-9509-1bd93af3d0bb
# ╠═b3056d5a-7790-4b94-9c77-854087a01cb2
# ╠═b0162138-03e1-4dcc-b76c-073e6bea17d9
# ╟─61a7e87a-9dc1-4519-ad78-51659c383496
# ╟─a29183fa-99eb-4bc9-bffd-4bf8fe1c1cb1
# ╠═43655522-24a7-4497-a305-4cb1a53e8791
# ╟─8a1e851c-8e86-4e56-86be-089667044cd4
# ╠═f40cdaeb-8474-4bae-a46f-44e2199e3712
# ╠═4c940780-fbd3-44dd-af56-2559809caa73

max_stars_repo_name: pascal-p/julia-notebooks

max_issues_repo_name: pascal-p/julia-notebooks

max_forks_repo_name: pascal-p/julia-notebooks


text:
get_aggregation_topology_accessor(::Type{Area}) = get_area
get_aggregation_topology_accessor(::Type{LoadZone}) = get_load_zone

set_load_zone!(bus::Bus, load_zone::LoadZone) = bus.load_zone = load_zone
set_area!(bus::Bus, area::Area) = bus.area = area

"""
Remove the aggregation topology in a Bus
"""
_remove_aggregration_topology!(bus::Bus, ::LoadZone) = bus.load_zone = nothing
_remove_aggregration_topology!(bus::Bus, ::Area) = bus.area = nothing

"""
Calculate the admittance of AC branches
"""
get_series_susceptance(b::ACBranch) = 1 / get_x(b)
get_series_susceptance(b::TapTransformer) = 1 / (get_x(b) * get_tap(b))
function get_series_susceptance(b::PhaseShiftingTransformer)
    y = 1 / (get_r(b) + get_x(b) * 1im)
    y_a = y / (get_tap(b) * exp(get_α(b) * 1im * (π / 180)))
    return 1 / imag(y_a)
end

get_series_admittance(b::ACBranch) = 1 / (get_r(b) + get_x(b) * 1im)

"""
Return the max active power for a device from get_active_power_limits.max
"""
function get_max_active_power(d::T) where {T <: Device}
    return get_active_power_limits(d).max
end

"""
Return the max reactive power for a device from get_reactive_power_limits.max
"""
function get_max_reactive_power(d::T) where {T <: Device}
    return get_reactive_power_limits(d).max
end

"""
Return the max reactive power for the Renewable Generation calculated as the rating * power_factor if
reactive_power_limits is nothing
"""
function get_max_reactive_power(d::RenewableDispatch)
    reactive_power_limits = get_reactive_power_limits(d)
    if isnothing(reactive_power_limits)
        return get_rating(d) * sin(acos(get_power_factor(d)))
    end
    return reactive_power_limits.max
end


text:
import .MCMCChains: AbstractChains, ChainDataFrame, sections
using .DataFrames

export from_mcmcchains

const turing_key_map = Dict(
    "acceptance_rate" => "mean_tree_accept",
    "hamiltonian_energy" => "energy",
    "hamiltonian_energy_error" => "energy_error",
    "is_adapt" => "tune",
    "max_hamiltonian_energy_error" => "max_energy_error",
    "n_steps" => "tree_size",
    "numerical_error" => "diverging",
    "tree_depth" => "depth",
)
const stan_key_map = Dict(
    "accept_stat__" => "accept_stat",
    "divergent__" => "diverging",
    "energy__" => "energy",
    "lp__" => "lp",
    "n_leapfrog__" => "n_leapfrog",
    "stepsize__" => "stepsize",
    "treedepth__" => "treedepth",
)
const stats_key_map = merge(turing_key_map, stan_key_map)

"""
    topandas(df::DataFrames.DataFrame) -> Pandas.DataFrame
    topandas(df::MCMCChains.ChainDataFrame) -> Pandas.DataFrame

Convert `df` into a Pandas format, maintaining column order and replacing
`missing` with `NaN`.
"""
function topandas(df::DataFrames.DataFrame)
    cols = replacemissing.(eachcol(df))
    colnames = names(df)
    df = DataFrames.DataFrame(cols, colnames)
    pdf = Pandas.DataFrame(df)
    return pdf[colnames]
end

topandas(df::ChainDataFrame) = topandas(df.df)

"""
    reshape_values(x::AbstractArray) -> AbstractArray

Convert from `MCMCChains` variable values with dimensions
`(ndraw, size..., nchain)` to ArviZ's expected `(nchain, ndraw, size...)`.
"""
reshape_values(x::AbstractArray{T,N}) where {T,N} = permutedims(x, (N, 1, 2:(N-1)...))

headtail(x) = x[1], x[2:end]

function split_locname(name)
    name = replace(name, r"[\[,]" => '.')
    name = replace(name, ']' => "")
    name, loc = headtail(split(name, '.'))
    length(loc) == 0 && return name, ()
    loc = tryparse.(Int, loc)
    Nothing <: eltype(loc) && return name, ()
    return name, tuple(loc...)
end

function varnames_locs_dict(loc_names)
    vars_to_locs = Dict()
    for loc_name in loc_names
        var_name, loc = split_locname(loc_name)
        if var_name ∉ keys(vars_to_locs)
            vars_to_locs[var_name] = ([loc_name], [loc])
        else
            push!(vars_to_locs[var_name][1], loc_name)
            push!(vars_to_locs[var_name][2], loc)
        end
    end
    return vars_to_locs
end

function attributes_dict(chns::AbstractChains)
    info = chns.info
    :hashedsummary in propertynames(info) || return info
    chndfs = info.hashedsummary.x[2]
    names = tuple(snakecase.(getproperty.(chndfs, :name))...)
    dfs = tuple(topandas.(chndfs)...)
    info = delete(info, :hashedsummary)
    attrs = merge(info, (mcmcchains_summary = Dict(zip(names, dfs)),))
    return Dict{String,Any}((string(k), v) for (k, v) in pairs(attrs))
end

attributes_dict(::Nothing) = Dict()

function section_dict(chns::AbstractChains, section)
    ndraws, _, nchains = size(chns)
    loc_names = string.(getfield(chns.name_map, section))
    vars_to_locs = varnames_locs_dict(loc_names)
    vars_to_arrays = Dict{String,Array}()
    for (var_name, names_locs) in vars_to_locs
        loc_names, locs = names_locs
        max_loc = maximum(hcat([[loc...] for loc in locs]...); dims = 2)
        ndim = length(max_loc)
        sizes = tuple(max_loc...)

        oldarr = reshape_values(replacemissing(Array(chns.value[:, loc_names, :])))
        if ndim == 0
            arr = dropdims(oldarr; dims = 3)
        else
            arr = Array{Union{typeof(NaN),eltype(oldarr)}}(undef, nchains, ndraws, sizes...)
            fill!(arr, NaN)
            for i in eachindex(locs)
                arr[:, :, locs[i]...] = oldarr[:, :, i]
            end
        end
        vars_to_arrays[var_name] = arr
    end
    return vars_to_arrays
end

function chains_to_dict(
    chns::AbstractChains;
    ignore = String[],
    section = :parameters,
    rekey_fun = identity,
)
    section in sections(chns) || return Dict()
    chns_dict = section_dict(chns, section)
    removekeys!(chns_dict, ignore)
    return rekey_fun(chns_dict)
end

chains_to_dict(::Nothing; kwargs...) = nothing

"""
    convert_to_dataset(chns::AbstractChains; library = MCMCChains, kwargs...) -> Dataset

Convert the chains `obj` to a [`Dataset`](@ref). `library` is the library that
created the chains. Remaining `kwargs` are forwarded to
[`dict_to_dataset`](@ref).
"""
function convert_to_dataset(chns::AbstractChains; library = MCMCChains, kwargs...)
    chns_dict = chains_to_dict(chns)
    return dict_to_dataset(chns_dict; library = library, kwargs...)
end

"""
    convert_to_inference_data(obj::AbstractChains; group = :posterior, kwargs...) -> InferenceData

Convert the chains `obj` to an [`InferenceData`](@ref) with the specified `group`.
Remaining `kwargs` are forwarded to [`from_mcmcchains`](@ref).
"""
function convert_to_inference_data(chns::AbstractChains; group = :posterior, kwargs...)
    group = Symbol(group)
    group == :posterior && return from_mcmcchains(chns; kwargs...)
    return from_mcmcchains(; group => chns)
end

"""
    from_mcmcchains(posterior::AbstractChains; kwargs...) -> InferenceData
    from_mcmcchains(; kwargs...) -> InferenceData

Convert data in an `MCMCChains.AbstractChains` format into an
[`InferenceData`](@ref).

Any keyword argument below without an an explicitly annotated type above is
allowed, so long as it can be passed to [`convert_to_inference_data`](@ref).

# Arguments
- `posterior::AbstractChains`: Draws from the posterior

# Keywords
- `posterior_predictive=nothing`: Draws from the posterior predictive distribution
     or name(s) of predictive variables in `posterior`
- `prior::AbstractChains=nothing`: Draws from the prior
- `prior_predictive=nothing`: Draws from the prior predictive distribution
     or name(s) of predictive variables in `prior`
- `observed_data=nothing`: Data actually observed or name(s) of observed data
     variables in `posterior`
- `constant_data=nothing`: Data that are constant, not observed (e.g. metadata),
     or name(s) of variables in posterior
- `log_likelihood::String=nothing`: Name of variable in `posterior` with log likelihoods
- `library=MCMCChains`: Name of library that generated the chains
- `coords::Dict{String,Vector}=nothing`: Map from named dimension to named
    indices
- `dims::Dict{String,Vector{String}}=nothing`: Map from variable name to names
    of its dimensions

# Returns
- `InferenceData`: The data with groups corresponding to the provided data
"""
function from_mcmcchains(
    posterior = nothing;
    prior = nothing,
    posterior_predictive = nothing,
    prior_predictive = nothing,
    observed_data = nothing,
    constant_data = nothing,
    log_likelihood = nothing,
    library = MCMCChains,
    kwargs...,
)
    attrs = attributes_dict(posterior)
    attrs = merge(attrs, Dict("inference_library" => string(library)))
    kwargs = convert(Dict, merge((; attrs = attrs, dims = nothing), kwargs))
    post_groups = (
        posterior_predictive = posterior_predictive,
        observed_data = observed_data,
        constant_data = constant_data,
    )

    group_dicts = Dict{Symbol,Dict}()
    group_datasets = Dict{Symbol,Dataset}()
    rekey_fun = d -> rekey(d, stats_key_map)

    # Convert chains to dicts
    post_dict = chains_to_dict(posterior)
    stats_dict = chains_to_dict(posterior; section = :internals, rekey_fun = rekey_fun)
    stats_dict = enforce_stat_types(stats_dict)

    prior_dict = chains_to_dict(prior)
    prior_stats_dict = chains_to_dict(prior; section = :internals, rekey_fun = rekey_fun)
    prior_stats_dict = enforce_stat_types(prior_stats_dict)

    # Remove variables from posterior/prior or convert to dataset
    for (group, data) in pairs(post_groups)
        if typeof(data) <: Union{String,Vector{String}}
            group_dicts[group] = popsubdict!(post_dict, data)
        elseif data !== nothing
            group_datasets[group] = convert_to_dataset(data; kwargs...)
        end
    end
    if typeof(prior_predictive) <: Union{String,Vector{String}}
        group_dicts[:prior_predictive] = popsubdict!(prior_dict, prior_predictive)
    elseif prior_predictive !== nothing
        group_datasets[:prior_predictive] = convert_to_dataset(prior_predictive; kwargs...)
    end

    # Handle log likelihood as a special case
    if log_likelihood !== nothing
        log_like_dict = popsubdict!(post_dict, log_likelihood)
        log_like_dict = Dict("log_likelihood" => log_like_dict[log_likelihood])
        if stats_dict === nothing
            stats_dict = log_like_dict
        else
            stats_dict = merge(stats_dict, log_like_dict)
        end
        dims = kwargs[:dims]
        if dims !== nothing && log_likelihood in keys(dims)
            dims["log_likelihood"] = dims[log_likelihood]
        end
    end

    idata1 = _from_dict(
        post_dict;
        sample_stats = stats_dict,
        prior = prior_dict,
        sample_stats_prior = prior_stats_dict,
        group_dicts...,
        kwargs...,
    )
    idata2 = InferenceData(; group_datasets...)
    isempty(idata1._groups) && return idata2
    isempty(idata2._groups) || concat!(idata1, idata2)
    return idata1
end

"""
    from_cmdstan(posterior::AbstractChains; kwargs...) -> InferenceData

Call [`from_mcmcchains`](@ref) on output of `CmdStan`.
"""
from_cmdstan(data::AbstractChains; kwargs...) =
    from_mcmcchains(data; library = "CmdStan", kwargs...)


text:
struct DummyError <: Exception end

function _setup_transaction_pool(size, file)
    ispath(file) && rm(file)

    pool = Lib.create(file, 2^24)

    # The root is a pointer to an array that we will dynamically allocate.
    roottype = PersistentOID{Int64}
    rootoid = Lib.root(pool, sizeof(roottype), roottype)

    # Do a transaction to allocate the array.
    transaction(pool) do
        oid = Lib.tx_zalloc(size, Int64)

        Lib.add_range(rootoid, 0, sizeof(roottype))
        rootptr = Lib.direct(rootoid)
        unsafe_store!(rootptr, oid)
    end

    return pool
end

@testset "Testing Transactions" begin
    using Persistence.Lib
    using Persistence.Transaction
    # The best way to test transactions is probably through the "transaction" function.
    #
    # The general approach is to have a root that is pointing to an array some where, and
    # try various ways of screwing up the transaction and making sure everything is in a
    # consistent state.

    testfile = "test.pool"
    roottype = PersistentOID{Int64}

    # Set up some ground truth arrays.
    testlength = 100000
    A = collect(1:testlength)
    B = collect(testlength:-1:1)

    @testset "Testing Normal Transactions" begin
        pool = _setup_transaction_pool(sizeof(A), testfile)

        transaction(pool) do
            # Get the persistent oid to the array
            rootoid = Lib.root(pool, sizeof(roottype), roottype)
            rootptr = Lib.direct(rootoid)

            oid = unsafe_load(rootptr)

            # Mark the memory for the array for transaction and copy A
            Lib.add_range(oid, 0, sizeof(A))
            ptr = Lib.direct(oid)
            unsafe_copyto!(ptr, pointer(A), length(A))
        end

        Lib.close(pool)
        pool = Lib.open(testfile)

        rootoid = Lib.root(pool, sizeof(roottype), roottype)
        rootptr = Lib.direct(rootoid)

        # Load the persistent pointer from the root.
        # Convert the persistent pointer to a normal pointer and load the stored array
        oid = unsafe_load(rootptr)
        ptr = Lib.direct(oid)

        dest = zeros(Int64, testlength)
        unsafe_copyto!(pointer(dest), ptr, testlength)

        @test dest == A

        # Cleanup
        Lib.close(pool)
    end

    @testset "Testing Aborted Transaction" begin
        # Manually call "abort" during the loading of A
        pool = _setup_transaction_pool(sizeof(A), testfile)

        f() = transaction(pool) do
            # Get the persistent oid to the array
            rootoid = Lib.root(pool, sizeof(roottype), roottype)
            rootptr = Lib.direct(rootoid)

            oid = unsafe_load(rootptr)

            # Mark the memory for the array for transaction and copy A
            Lib.add_range(oid, 0, sizeof(A))
            ptr = Lib.direct(oid)
            unsafe_copyto!(ptr, pointer(A), length(A))

            # Now, everything is copied directly, call "abort"
            Lib.tx_abort(10)
        end
        @test_throws TransactionAborted f()

        Lib.close(pool)
        pool = Lib.open(testfile)

        rootoid = Lib.root(pool, sizeof(roottype), roottype)
        rootptr = Lib.direct(rootoid)

        ## Load the persistent pointer from the root.
        ## Convert the persistent pointer to a normal pointer and load the stored array
        oid = unsafe_load(rootptr)
        ptr = Lib.direct(oid)

        dest = zeros(Int64, testlength)
        unsafe_copyto!(pointer(dest), ptr, testlength)

        @test dest == zeros(Int64, testlength)

        ## Cleanup
        Lib.close(pool)
    end

    @testset "Testing Errored Transactions" begin
        # Manually call "abort" during the loading of A
        pool = _setup_transaction_pool(sizeof(A), testfile)

        f() = transaction(pool) do
            # Get the persistent oid to the array
            rootoid = Lib.root(pool, sizeof(roottype), roottype)
            rootptr = Lib.direct(rootoid)

            oid = unsafe_load(rootptr)

            # Mark the memory for the array for transaction and copy A
            Lib.add_range(oid, 0, sizeof(A))
            ptr = Lib.direct(oid)
            unsafe_copyto!(ptr, pointer(A), length(A))

            # Throw dummy exception
            throw(DummyError())
        end
        @test_throws DummyError f()

        Lib.close(pool)
        pool = Lib.open(testfile)

        rootoid = Lib.root(pool, sizeof(roottype), roottype)
        rootptr = Lib.direct(rootoid)

        ## Load the persistent pointer from the root.
        ## Convert the persistent pointer to a normal pointer and load the stored array
        oid = unsafe_load(rootptr)
        ptr = Lib.direct(oid)

        dest = zeros(Int64, testlength)
        unsafe_copyto!(pointer(dest), ptr, testlength)

        @test dest == zeros(Int64, testlength)

        ## Cleanup
        Lib.close(pool)
    end
end


text:
function Base.dropdims(a::IndicesArray; dims)
    d = to_dims(a, dims)
    return rebuild(a, dropdims(parent(a); dims=d), drop_axes(a, dims=d))
end

function Base.reshape(a::IndicesArray, dims)
    d = to_dims(a, dims)
    return rebuild(a, reshape(parent(a), d), reshape_axes(a, d))
end

Base.selectdim(a::IndicesArray, d, i) = selectdim(a, to_dims(nda, d), i)

function Base.permutedims(a::IndicesArray, p)
    return rebuild(a, permutedims(parent(a), p), permute_axes(a, p))
end
function Base.permutedims(a::IndicesArray, p...)
    return rebuild(a, permutedims(parent(a), p...), permute_axes(a, p...))
end

for f in (
    :(Base.transpose),
    :(Base.adjoint),
    :(Base.permutedims),
    :(LinearAlgebra.pinv))

    # Vector
    @eval function $f(a::IndicesVector)
        return rebuild(a, $f(parent(a)), (axes(a, 1), axes(a, 1)))
    end

    # Vector Double Transpose
    if f !== :(Base.permutedims)
        @eval begin
            function $f(a::Union{IndicesAdjoint,IndicesTranspose})
                return rebuild(a, $f(parent(a)), (axes(a, 2),))
            end
        end
    end

    # Matrix
    @eval function $f(a::IndicesMatrix)
        return rebuild(a, $f(parent(a)), (axes(a, 2), axes(a, 1)))
    end
end

Base.inv(a::IndicesMatrix) = rebuild(a, inv(parent(a)), inv_axes(a))

function Base.eachslice(a::IndicesArray; dims, kwargs...)
    slices = eachslice(parent(a); dims=d, kwargs...)
    return Base.Generator(slices) do slice
        return _maybe_indarray(slice, reduce_axes(a, d))
    end
end


text:
using PlotUtils
using Test
using Statistics: mean
using Dates

# TODO: real tests

# ----------------------
# colors

const C = RGBA{Float64}

@testset "colors" begin

    @test plot_color(nothing) == C(0,0,0,0)
    @test plot_color(false) == C(0,0,0,0)
    @test_throws ErrorException plot_color(true)

    @test plot_color(:red) == parse(C, :red)
    @test plot_color("red") == parse(C, "red")
    @test_throws ErrorException plot_color("notacolor")

    @test plot_color(colorant"red") == C(1,0,0,1)

    grad = cgrad()
    @test typeof(grad) == ColorGradient
    @test plot_color(grad) === grad

    grad = cgrad([:red, "blue"])
    @test grad.colors == C[colorant"red", colorant"blue"]
    @test grad.values == collect(range(0,stop=1,length=2))

    grad = cgrad([:red, "blue"], alpha = 0.5)
    @test grad.colors == C[C(1,0,0,0.5), C(0,0,1,0.5)]
    @test grad.values == collect(range(0,stop=1,length=2))

    grad = cgrad([:red,:blue], [0,0.1,1])
    @test grad.colors == C[C(1,0,0), C(0.5,0,0.5), C(0,0,1)]
    @test grad.values == [0,0.1,1]

    cs = plot_color(rand(10))
    @test typeof(cs) == Vector{C}
    @test length(cs) == 10

    cs = plot_color(rand(4,4))
    @test typeof(cs) == Matrix{C}
    @test length(cs) == 16
    @test size(cs) == (4,4)

    cs = plot_color(rand(10), 0.5)
    @test typeof(cs) == Vector{C}
    @test length(cs) == 10
    for c in cs
        @test alpha(c) == 0.5
    end

    cs = plot_color(rand(4,4), 0.5)
    @test typeof(cs) == Matrix{C}
    @test length(cs) == 16
    @test size(cs) == (4,4)
    for c in cs
        @test alpha(c) == 0.5
    end
end

# ----------------------
# gradients

@testset "gradients" begin
    grad = cgrad(:inferno)
    @test length(grad.colors) == 30
    @test length(grad.values) == 30
    @test RGB(grad.colors[1]) == RGB(0.001462, 0.000466, 0.013866)
    @test RGB(grad.colors[end]) == RGB(0.988362, 0.998364, 0.644924)
end

# ----------------------
# ticks

# Copied from Plots.is_uniformly_spaced to avoid dependency on recent version
# on Plots which is not used on Travis.
function is_uniformly_spaced(v; tol=1e-6)
  dv = diff(v)
  maximum(dv) - minimum(dv) < tol * mean(abs.(dv))
end

@testset "ticks" begin
    @test optimize_ticks(-1,2) == ([-1.0,0.0,1.0,2.0],-1.0,2.0)
    dt1, dt2 = Dates.value(DateTime(2000)), Dates.value(DateTime(2100))
    @test optimize_datetime_ticks(dt1, dt2) == (
        [63113990400000, 63902908800000, 64691827200000, 65480745600000],
        ["2001-01-01", "2026-01-01", "2051-01-01", "2076-01-01"]
    )

    @testset "small range $x, $(i)ϵ" for x in exp10.(-12:12), i in -5:5
        y = x + i*eps(x)
        x,y = minmax(x,y)
        ticks = PlotUtils.optimize_ticks(x, y)[1]
        @test issorted(ticks)
        @test all(x .<= ticks .<= y)
        # Fails:
        # @test allunique(ticks)
    end

    @testset "digits $((10^n)-1)*10^$i" for n in 1:9, i in -9:9
        y0 = 10^n
        x0 = y0-1
        x, y = (x0,y0) .* 10.0^i
        ticks = optimize_ticks(x, y)[1]
        @test length(ticks) >= 2
        @test issorted(ticks)
        @test all(x .<= ticks .<= y)
        @test is_uniformly_spaced(ticks)
    end
end


text:
# -----------------------------------------------------------
#
# Using boundary terms and integration (summation) by parts
# to construct more accurate derivative matrices for time
# evolution using Lax-Friedrichs-type fluxes
#
# All derivatives are computed here in pos basis first
# Then conjugated into hier basis
#
# -----------------------------------------------------------

# Efficiency criticality: MED
# Most computations time is taken up by functions called from
# other script files

# Accuracy criticality: LOW
# The accuracy is limited only by functions in other scripts
# and ODE.jl

# Construct the time evolution matrix for the wave equation
# given a Kaplacian and the initial position and velocity f0, v0
function wave_data(laplac::A, f0coeffs::AbstractArray{T, 1},
        v0coeffs::AbstractArray{T, 1}) where {A <: AbstractArray, T <: Real}

    len = length(f0coeffs)
    rows = rowvals(laplac)
    vals = nonzeros(laplac)

    I = Int[]
    J = Int[]
    V = T[]
    for col = 1:len
        for i in nzrange(laplac, col)
            row = rows[i]
            val = vals[i]
            push!(I,row+len)
            push!(J,col)
            push!(V,val)
        end
    end
    for i = 1:len
        push!(I, i)
        push!(J, i+len)
        push!(V, 1.0)
    end
    # println("len = ", len) # to make debugging quicker
    # println("max(rows) = ", maximum(rows))
    RHS = sparse(I, J, V, 2*len, 2*len, +)
    y0 = Array{T}([i<=len ? f0coeffs[i] : v0coeffs[i-len] for i in 1:2*len])
    return RHS, y0
end

# Evolves the wave equation from time t0 to t1
# given coefficients for an initial position profile of f0
# and an initial velocity profile of v0
function wave_evolve(D::Int, k::Int, n::Int,
                              f0coeffs::Array{T, 1}, v0coeffs::Array{T, 1},
                              time0::Real, time1::Real;
                              order="45", scheme="sparse", kwargs...) where T <: Real

    laplac   = laplacian_matrix(D, k, n; scheme=scheme)

    RHS, y0 = wave_data(laplac, f0coeffs, v0coeffs)
    if order == "45"
        soln = ode45((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    elseif order == "78"
        soln = ode78((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    else
        throw(ArgumentError(:order))
    end
    return soln
end

# Same function but for anisotropic grids
function wave_evolve(D::Int, k::Int, n::Vector{Int},
                    f0coeffs::Array{T, 1}, v0coeffs::Array{T, 1},
                    time0::Real, time1::Real;
                    order="45", scheme="full", kwargs...) where T <: Real

    laplac   = laplacian_matrix(D, k, n; scheme=scheme)

    RHS, y0 = wave_data(laplac, f0coeffs, v0coeffs)
    if order == "45"
        soln = ode45((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    elseif order == "78"
        soln = ode78((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    else
        throw(ArgumentError(:order))
    end
    return soln
end

# Evolves the wave equation from time t0 to t1
# given initial position profile f0 and velocity profile v0
function wave_evolve(D::Int, k::Int, n::Int,
                              f0::Function, v0::Function,
                              time0::Real, time1::Real;
                              order="45", scheme="sparse", kwargs...)

    f0coeffs = vcoeffs_DG(D, k, n, f0; scheme=scheme)
    v0coeffs = vcoeffs_DG(D, k, n, v0; scheme=scheme)
    return wave_evolve(D, k, n, f0coeffs, v0coeffs, time0, time1;
                        order=order, scheme=scheme, kwargs...)
end

# wave_evolve for anisotropic grids
function wave_evolve(D::Int, k::Int, n::Vector{Int},
                                f0::Function, v0::Function,
                                time0::Real, time1::Real;
                                order="45", scheme="full", kwargs...)

    f0coeffs = vcoeffs_DG(D, k, n, f0; scheme=scheme)
    v0coeffs = vcoeffs_DG(D, k, n, v0; scheme=scheme)
    return wave_evolve(D, k, n, f0coeffs, v0coeffs, time0, time1;
                         order=order, scheme=scheme, kwargs...)
end

# The reason we have a wave evolution in 1D is to test
# that the standard "position" and multiresolution "heirarchical"
# bases agree for wave PDE evolution
function wave_evolve_1D(k::Int, max_level::Int,
                            f0::Function, v0::Function,
                            time0::Real, time1::Real; basis="hier", order = "45",
                            kwargs...)
    if basis == "pos"
        f0coeffs = pos_vcoeffs_DG(k, max_level, f0)
        v0coeffs = pos_vcoeffs_DG(k, max_level, v0)
    elseif basis == "hier"
        f0coeffs = vcoeffs_DG(1, k, max_level, f0)
        v0coeffs = vcoeffs_DG(1, k, max_level, v0)
    elseif basis == "nodal"
        throw(MethodError("Nodal basis wave evolution not yet implemented"))
    elseif basis == "point"
        m2p = modal2points_1D(k, max_level)
        f0coeffs = m2p*vcoeffs_DG(1, k, max_level, f0)
        v0coeffs = m2p*vcoeffs_DG(1, k, max_level, v0)
    else
        throw(ArgumentError(:basis))
    end

    D_op = periodic_DLF_matrix(k, max_level; basis=basis)
    laplac= *(D_op,D_op)

    RHS, y0 = wave_data(laplac, f0coeffs, v0coeffs)
    if order == "45"
        soln = ode45((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    elseif order == "78"
        soln = ode78((t,x)->*(RHS,x), y0, [time0,time1]; kwargs...)
    else
        throw(ArgumentError(:order))
    end
    return soln
end

# Evolve the Vlassov-Poisson (AKA collisionless Boltzmann) equation
# for a mass distribution in D dimensions
# -> in a 2*D-dimensional phase space
# D, k, n (good)
# f0_coeffs_modal (start with anything- like a gaussian) (good)
# F_point (Depends on the potential I pick. Use Isothermal Truncated)
# v_point (its a constant and simple vector for any potential) (good)
# time0, time1, order, scheme (good)
function vlasov_evolve(D::Int, k::Int, n::Int,
            m2n::SparseMatrixCSC{T, Int}, n2p::SparseMatrixCSC{T, Int},
            p2n::SparseMatrixCSC{T, Int}, n2m::SparseMatrixCSC{T, Int},
            f0_modal::Array{T,1}, F_point::Array{Array{T,1}, 1},
            time0::T, time1::T, nout::Int=2;
            order="45", scheme="sparse", kwargs...) where T <: Real

    # The grad matrix- using the same derivative operator as we did
    # for the wave equation.
    # We may need to apply filtering intermittently in the evolution
    Ds = grad_matrix(2*D, k, n; scheme=scheme)

    println("Writing operators...")
    flush(stdout)
    jldopen("vlasov.h5", "w") do f
        f["dimensions"] = D
        f["order"] = k
        f["levels"] = n
        function w(n::AbstractString, A::SparseMatrixCSC)
            f["$n.m"] = A.m
            f["$n.n"] = A.n
            f["$n.colptr"] = A.colptr
            f["$n.rowval"] = A.rowval
            f["$n.nzval"] = A.nzval
        end
        w("m2n", m2n)
        w("n2p", n2p)
        w("p2n", p2n)
        w("n2m", n2m)
        for (d, ds) in enumerate(Ds)
            w("Ds[$d]", ds)
        end
    end

    # Coeffs for the constant 1 in D-dim space
    one_1D = get_one_modal(1, k, n)
    # Coeffs for velocity in D-dim space
    v_modal_1D = get_xi_modal(1, 1, k, n)
    # Coeffs for velocity in 2*D-dim phase space - the tensor product of the above
    v_modal = [tensor_construct(2*D, k, n, [j-D == i ? v_modal_1D : one_1D for j in 1:2*D])
              for i in 1:D]
    v_point = [n2p * (m2n * v) for v in v_modal]

    function steprule(t::T, f_modal::Array{T, 1})
        println("    evaluating RHS at t=$t...")
        flush(stdout)
        # We want this in a modal basis for differentiation to work fast
        # Matrix multiplication is the bottle neck here.
        dfdxs_modal = [Ds[d] * f_modal for d in 1:D]
        dfdps_modal = [Ds[d] * f_modal for d in (D+1):(2*D)]

        dfdxs_point = [n2p * (m2n * x) for x in dfdxs_modal]
        dfdps_point = [n2p * (m2n * x) for x in dfdps_modal]

        # dH/dp * df/dx
        contrib1 = sum(v_point[d] .* dfdxs_point[d] for d in 1:D)
        # dH/dx * df/dp
        contrib2 = sum(F_point[d] .* dfdps_point[d] for d in 1:D)

        # df/dt = - (dH/dp * df/dx - dH/dx * df/dp)
        return n2m * (p2n * (- contrib1 + contrib2))
    end

    println("State vector size: $(length(f0_modal))")
    println("Operator sizes:")
    for d in 1:D
        println("    Ds[$d]: $(nnz(Ds[d]))")
    end
    println("    m2n:   $(nnz(m2n))")
    println("    n2p:   $(nnz(n2p))")
    println("    p2n:   $(nnz(p2n))")
    println("    n2m:   $(nnz(n2m))")
    println("Beginning PDE evolution...")
    flush(stdout)

    if order == "45"
        solver = ode45
    elseif order == "78"
        solver = ode78
    else
        throw(ArgumentError(:order))
    end
    soln = solver(steprule, f0_modal, range(time0, time1, length=nout); kwargs...)

    println("Writing solution...")
    flush(stdout)
    jldopen("vlasov.h5", "r+") do f
        f["times"] = soln[1]
        for (i,sol) in enumerate(soln[2])
            ifmt = @sprintf "%06d" i
            f["f_modal.$ifmt"] = sol
        end
    end

    println("Done.")
    flush(stdout)
end

function norm_squared(coeffs::AbstractArray{T},; basis = "hier") where T <: Real

    if basis=="hier" || basis=="pos"
        return sum(i^2 for i in coeffs)
    else
        throw(MethodError("Nodal and point basis norms not yet implemented"))
    end
end

function energy_func_1D(k, level, soln::Tuple{Array{T, 1}, Array{Array{T, 1}, 1}};
    basis = "hier") where T <: Real

    len            = length(soln[1])
    num_coeffs    = Int(round(length(soln[2][1])/2))
    times        = copy(soln[1])
    energies    = Array{T}(undef, len)
    D_op        = periodic_DLF_matrix(k, level; basis=basis)

    for i in 1:len
        ux   = D_op *  (@view soln[2][i][1:num_coeffs])
        udot = @view soln[2][i][num_coeffs+1:end]
        energies[i] = norm_squared(ux) + norm_squared(udot)

    end
    return (times, energies)
end

function energy_func(D::Int, k::Int, n::Int, soln::Tuple{Array{T, 1}, Array{Array{T, 1}, 1}};
        scheme = "sparse") where T <: Real

    len         = length(soln[1])
    num_coeffs    = div(length(soln[2][1]), 2)
    times         = copy(soln[1])
    energies    = Array{T}(undef, len)
    D_ops        = grad_matrix(D, k, n; scheme=scheme)

    for i in 1:len
        ux   = [D_ops[j] * (@view soln[2][i][1:num_coeffs]) for j in 1:D]
        udot = @view soln[2][i][num_coeffs+1:end]
        energies[i] = sum(norm_squared(ux[d]) for d in 1:D) + norm_squared(udot)
    end
    return (times, energies)
end


text:
function MCTS.RandomBandit(mdp::InvertedPendulum; kwargs...)
    RandomBandit(;
        enable_action_pw=true,
        check_repeat_action=true,
        kwargs...)
end
function MCTS.DPWBandit(mdp::InvertedPendulum; kwargs...)
    DPWBandit(;
        enable_action_pw=true,
        check_repeat_action=true,
        exploration_constant=1.0,
        k_action=10.0,
        alpha_action=0.5,
        kwargs...)
end
function MCTS.CBTSBandit(mdp::InvertedPendulum; kwargs...)
    CBTSBandit(;
        enable_action_pw=true,
        check_repeat_action=true,
        exploration_constant=1.0,
        A_max=20,
        n_proposes=100,
        log_length_scale=0.0,
        log_signal_sigma=0.0,
        log_obs_noise=-1.0,
        action_dims=2,
        n_sig=2.0,  #number of standard deviations for GP-UCB
        kwargs...)
end
function MCTS.ModularSolver(mdp::InvertedPendulum, b::ModularBandit, seed=0;
                            kwargs...)
    rng = MersenneTwister(seed)
    ModularSolver(;
        n_iterations=100,
        bandit=b,
        depth=10,
        k_state=10.0,
        alpha_state=0.5,
        keep_tree=true,
        check_repeat_state=true,
        rng=rng,
        init_Q=0.0,
        init_N=0,
        next_action=RandomActionGenerator(rng),
        estimate_value=RolloutEstimator(RandomSolver(rng)),
        track_best_path=true,
        kwargs...
       )
end
Base.string(::InvertedPendulum) = "Inverted Pendulum"


text:
function GalileoAlmanac(arg0::jint, arg1::jint, arg10::jdouble, arg11::jdouble, arg12::jdouble, arg13::jint, arg14::jint, arg15::jint, arg16::AbsoluteDate, arg2::jdouble, arg3::jdouble, arg4::jdouble, arg5::jdouble, arg6::jint, arg7::jdouble, arg8::jdouble, arg9::jdouble)
    return GalileoAlmanac((jint, jint, jdouble, jdouble, jdouble, jint, jint, jint, AbsoluteDate, jdouble, jdouble, jdouble, jdouble, jint, jdouble, jdouble, jdouble), arg0, arg1, arg10, arg11, arg12, arg13, arg14, arg15, arg16, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9)
end

function GalileoAlmanac(arg0::jint, arg1::jint, arg10::jdouble, arg11::jdouble, arg12::jdouble, arg13::jint, arg14::jint, arg15::jint, arg2::jdouble, arg3::jdouble, arg4::jdouble, arg5::jdouble, arg6::jint, arg7::jdouble, arg8::jdouble, arg9::jdouble)
    return GalileoAlmanac((jint, jint, jdouble, jdouble, jdouble, jint, jint, jint, jdouble, jdouble, jdouble, jdouble, jint, jdouble, jdouble, jdouble), arg0, arg1, arg10, arg11, arg12, arg13, arg14, arg15, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9)
end

function get_af0(obj::GalileoAlmanac)
    return jcall(obj, "getAf0", jdouble, ())
end

function get_af1(obj::GalileoAlmanac)
    return jcall(obj, "getAf1", jdouble, ())
end

function get_bgd(obj::GalileoOrbitalElements)
    return jcall(obj, "getBGD", jdouble, ())
end

function get_bgde1e5a(obj::GalileoOrbitalElements)
    return jcall(obj, "getBGDE1E5a", jdouble, ())
end

function get_bgde5b_e1(obj::GalileoOrbitalElements)
    return jcall(obj, "getBGDE5bE1", jdouble, ())
end

function get_cic(obj::GalileoAlmanac)
    return jcall(obj, "getCic", jdouble, ())
end

function get_cis(obj::GalileoAlmanac)
    return jcall(obj, "getCis", jdouble, ())
end

function get_crc(obj::GalileoAlmanac)
    return jcall(obj, "getCrc", jdouble, ())
end

function get_crs(obj::GalileoAlmanac)
    return jcall(obj, "getCrs", jdouble, ())
end

function get_cuc(obj::GalileoAlmanac)
    return jcall(obj, "getCuc", jdouble, ())
end

function get_cus(obj::GalileoAlmanac)
    return jcall(obj, "getCus", jdouble, ())
end

function get_date(obj::GalileoAlmanac)
    return jcall(obj, "getDate", AbsoluteDate, ())
end

function get_e(obj::GalileoAlmanac)
    return jcall(obj, "getE", jdouble, ())
end

function get_health_e1(obj::GalileoAlmanac)
    return jcall(obj, "getHealthE1", jint, ())
end

function get_health_e5a(obj::GalileoAlmanac)
    return jcall(obj, "getHealthE5a", jint, ())
end

function get_health_e5b(obj::GalileoAlmanac)
    return jcall(obj, "getHealthE5b", jint, ())
end

function get_i0(obj::GalileoAlmanac)
    return jcall(obj, "getI0", jdouble, ())
end

function get_i_dot(obj::GalileoAlmanac)
    return jcall(obj, "getIDot", jdouble, ())
end

function get_iod(obj::GalileoAlmanac)
    return jcall(obj, "getIOD", jint, ())
end

function get_iod_nav(obj::GalileoOrbitalElements)
    return jcall(obj, "getIODNav", jint, ())
end

function get_m0(obj::GalileoAlmanac)
    return jcall(obj, "getM0", jdouble, ())
end

function get_mean_motion(obj::GalileoAlmanac)
    return jcall(obj, "getMeanMotion", jdouble, ())
end

function get_omega0(obj::GalileoAlmanac)
    return jcall(obj, "getOmega0", jdouble, ())
end

function get_omega_dot(obj::GalileoAlmanac)
    return jcall(obj, "getOmegaDot", jdouble, ())
end

function get_pa(obj::GalileoAlmanac)
    return jcall(obj, "getPa", jdouble, ())
end

function get_prn(obj::GalileoAlmanac)
    return jcall(obj, "getPRN", jint, ())
end

function get_sma(obj::GalileoAlmanac)
    return jcall(obj, "getSma", jdouble, ())
end

function get_time(obj::GalileoAlmanac)
    return jcall(obj, "getTime", jdouble, ())
end

function get_week(obj::GalileoAlmanac)
    return jcall(obj, "getWeek", jint, ())
end


text:
@with_kw mutable struct TMazeState
    x::Int64 = 1 # position in corridor
    g::Symbol = :north# goal north or south
    term::Bool = false
end

==(s1::TMazeState, s2::TMazeState) = s1.x == s2.x && s1.g == s2.g
hash(s::TMazeState, h::UInt64 = zero(UInt64)) = hash(s.x, hash(s.g, h))
function Base.copy!(s1::TMazeState, s2::TMazeState)
    s1.x = s2.x
    s1.g = s2.g
    s1.term = s2.term
    return s1
end

@with_kw struct TMaze <: POMDP{TMazeState, Int64, Int64}
    n::Int64 = 10 # corridor length
    discount::Float64 = 0.99 # discount factor
end

# state space is length of corr + 3 cells at the end
#                   |G|
# |S| | | | | | | | | |
#                   | |
# depending on where the goal is
function states(maze::TMaze)
    space = TMazeState[]
    for x in 1:(maze.n + 1), g in [:north, :south]
        push!(space, TMazeState(x, g, false))
    end
    push!(space, TMazeState(1,:none,true)) # terminal
    return space
end
# 4 actions: go North, East, South, West (1, 2, 3, 4)
actions(maze::TMaze) = 1:4
actionindex(maze::TMaze, i::Int) = i
# 5 observations: 2 for goal (left or right) + 2 for in corridor or at intersection + 1 term
observations(maze::TMaze) = 1:5
obsindex(maze::TMaze, i::Int) = i

# transition distribution (actions are deterministic)
mutable struct TMazeStateDistribution
    current_state::TMazeState # deterministic
    reset::Bool
    reset_states::Vector{TMazeState}
    reset_probs::Vector{Float64}
end
function create_transition_distribution(::TMaze)
    rs = [TMazeState(1,:north,false), TMazeState(1,:south,false)]
    rp = [0.5, 0.5]
    TMazeStateDistribution(TMazeState(), false, rs, rp)
end
support(d::TMazeStateDistribution) = d.reset ? (return [d.current_state]) : (return zip(d.reset_states, d.reset_probs))

function pdf(d::TMazeStateDistribution, s::TMazeState)
    if d.reset
        in(s, d.reset_states) ? (return 0.5) : (return 0.0)
    else
        s == d.current_state ? (return 1.0) : (return 0.0)
    end
end
function rand(rng::AbstractRNG, d::TMazeStateDistribution)
    s = TMazeState()
    if d.reset
        rand(rng) < 0.5 ? (copy!(s, d.reset_states[1])) : (copy!(s, d.reset_states[2]))
        return s
    else
        copy!(s, d.current_state)
        return s
    end
end
#rand(rng::AbstractRNG, d::TMazeStateDistribution)

struct TMazeInit
    states::Vector{TMazeState}
    probs::Vector{Float64}
end
support(d::TMazeInit) = d.states
function initialstate_distribution(maze::TMaze)
    s = states(maze)
    ns = length(s)
    p = zeros(ns) .+ 1.0 / (ns-1)
    p[end] = 0.0
    #s1 = TMazeState(1, :north, false)
    #s2 = TMazeState(1, :south, false)
    #d = TMazeInit([s1, s2])
    return TMazeInit(s, p)
end
function rand(rng::AbstractRNG, d::TMazeInit)
    cat = Weights(d.probs)
    idx = sample(rng, cat)
    return d.states[idx]
end
function pdf(d::TMazeInit, s::TMazeState)
    for i = 1:length(d.states)
        if d.states[i] == s
            return d.probs[i]
        end
    end
    return 0.0
end

# observation distribution (deterministic)
mutable struct TMazeObservationDistribution
    current_observation::Int64
end
create_observation_distribution(::TMaze) = TMazeObservationDistribution(1)
iterator(d::TMazeObservationDistribution) = [d.current_observation]

pdf(d::TMazeObservationDistribution, o::Int64) = Float64(o == d.current_observation)
support(d::TMazeObservationDistribution) = d.current_observation
rand(rng::AbstractRNG, d::TMazeObservationDistribution) = d.current_observation

function transition(maze::TMaze, s::TMazeState, a::Int64)
    d=create_transition_distribution(maze)
    d.reset = false
    # check if terminal
    if s.term
        # reset
        d.reset = true
        #copy!(d.current_state, s) # state doesn't change
        return d
    end
    # check if move into terminal move north or south
    if s.x == maze.n + 1
        if a == 1 || a == 3
            d.current_state = TMazeState(1,:none,true) # state now terminal
            return d
        elseif a == 4
            copy!(d.current_state, s)
            d.current_state.x -= 1
            return d
        else
            copy!(d.current_state, s)
            return d
        end
    end
    # check if move along hallway
    if a == 2
        copy!(d.current_state, s)
        d.current_state.x += 1
        return d
    end
    if a == 4
        copy!(d.current_state, s)
        s.x > 1 ? (d.current_state.x -= 1) : (nothing)
        return d
    end
    # if none of the above just stay in place
    copy!(d.current_state, s)
    return d
end

function reward(maze::TMaze, s::TMazeState, a::Int64)
    # check terminal
    s.term ? (return 0.0) : (nothing)
    # check if at junction
    if s.x == maze.n + 1
        # if at junction check action
        if (s.g == :north && a == 1) || (s.g == :south && a == 3)
            return 4.0
        elseif (s.g == :north && a == 3) || (s.g == :south && a == 1)
            return -0.1
        else
            return -0.1
        end
    end
    # if bump against wall
    if s.x < maze.n + 1 && (a == 1 || a == 3)
        return -0.1
    end
    return 0.0
end

# observation mapping
#    1      2        3         4         5
# goal N  goal S  corridor  junction  terminal
function observation(maze::TMaze, sp::TMazeState)
    d::TMazeObservationDistribution = create_observation_distribution(maze)
    sp.term ? (d.current_observation = 5; return d) : (nothing)
    x = sp.x; g = sp.g
    #if x == 1
    if x <= 2
        g == :north ? (d.current_observation = 1) : (d.current_observation = 2)
        return d
    end
    if 1 < x < (maze.n + 1)
        d.current_observation = 3
        return d
    end
    if x == maze.n + 1
        d.current_observation = 4
        return d
    end
    d.current_observation = 5
    return d
end

isterminal(m::TMaze, s::TMazeState) = s.term

discount(m::TMaze) = m.discount

stateindex(maze::TMaze, s::TMazeState) = s.term ? (2*(maze.n+1) + 1) : (2*s.x - (s.g==:north))

function Base.convert(maze::TMaze, s::TMazeState)
    v = Array{Float64}(undef, 2)
    v[1] = s.x
    s.g == :north ? (v[2] = 0.0) : (v[2] = 1.0)
    return v
end

struct MazeBelief
    last_obs::Int64
    mem::Symbol # memory
end
MazeBelief() = MazeBelief(1, :none)

struct MazeUpdater <: Updater end
POMDPs.initialize_belief(bu::MazeUpdater, d::Any) = d

function POMDPs.update(bu::MazeUpdater, b::MazeBelief, a, o)
    mem = b.mem
    if o == 1
        mem = :north
    end
    if o == 2
        mem = :south
    end
    return MazeBelief(o, mem)
end

mutable struct MazeOptimal <: Policy end
POMDPs.updater(p::MazeOptimal) = MazeUpdater()

# 4 actions: go North, East, South, West (1, 2, 3, 4)
# observation mapping
#    1      2        3         4         5
# goal N  goal S  corridor  junction  terminal
function POMDPs.action(p::MazeOptimal, b::MazeBelief)
    # if don't know the goal go back
    if b.mem == :none
        return 4
    end
    if b.mem == :north && b.last_obs == 4
        return 1
    end
    if b.mem == :south && b.last_obs == 4
        return 3
    end
    return 2
end


text:
function demoBurglarDictPot()
    # This demo is the same as demoBurglar.jl, except that here we use the ability to define non-contiguously indexed potentials and variables by demonstrating the use of the Dictoinary of PotArray technique.

    burglar,earthquake,alarm,radio=10,2,30,5 # Variable order and indexing is arbitary (thianks to using DictPot and doesn't need to be contiguous
    yes, no= 1,2 # define states, starting from 1. These do need to be contiguous, starting from 1, since currently the probability tables are defined by arrays of contiguous states.

    # This defines the meaning of the variables and their states.
    # It is not necessary for the computation, but is useful for
    # printing the results in a meaningful form:
    DictVariable=Dict{Int,DiscreteVariable}()
    DictVariable[burglar]=DiscreteVariable("burglar",["yes","no"])
    DictVariable[earthquake]=DiscreteVariable("earthquake",["yes","no"])
    DictVariable[alarm]=DiscreteVariable("alarm",["yes","no"])
    DictVariable[radio]=DiscreteVariable("radio",["yes","no"])

    # Define the discrete probability tables:
    DictPot=Dict{Integer,PotArray}()
    DictPot[burglar]=PotArray(burglar,[0.01 0.99])
    DictPot[earthquake]=PotArray(earthquake,[0.000001 1-0.000001])

    table=zeros(2,2)
    table[yes,yes]=1
    table[no,yes]=0
    table[yes,no]=0
    table[no,no]=1
    DictPot[radio]=PotArray([radio earthquake],table)

    table=zeros(2,2,2)
    table[yes,yes,yes]=0.9999
    table[yes,yes,no]=0.99
    table[yes,no,yes]=0.99
    table[yes,no,no]=0.0001
    table[no,:,:]=1-table[yes,:,:]
    DictPot[alarm]=PotArray([alarm burglar earthquake],table)

    jointpot=prod(DictPot)

    println("p(burglar|alarm=yes):")
    show(condpot(setpot(jointpot,alarm,yes),burglar),DictVariable)

    println("p(burglar|alarm=yes,radio=yes):")
    show(condpot(setpot(jointpot,[alarm radio],[yes yes]),burglar),DictVariable)

end


text:
water_formation = ce"2 H2 + O2 → 2 H2O"
ionic_reaction = ce"Na{+} + Cl{-} > NaCl"
redox = ce"Cr2O7{-2} + H{+1} + {-} = Cr{3+} + H2O"

@testset "ChemEquation" begin
    @test water_formation.tuples ==
        [(cc"H2", 2), (cc"O2", 1), (cc"H2O", -2)]
    @test ionic_reaction.tuples ==
        [(cc"Na{+1}", 1), (cc"Cl{-1}", 1), (cc"ClNa", -1)]
    @test ChemEquation("N2+O2⇌2NO").tuples ==
        [(cc"N2", 1), (cc"O2", 1), (cc"NO", -2)]

    testtuple = [(cc"H2", 0.5), (cc"Cl2", 0.5), (cc"HCl", -1.0)]
    @test ChemEquation(testtuple) == ChemEquation{Float64}(testtuple)
    @test ChemEquation{Rational}("1//2 H2 + 1//2 Cl2 → HCl").tuples == testtuple
    @test ChemEquation{Float64}("0.5 H2 + 0.5 Cl2 → HCl").tuples == testtuple
end

@testset "@ce_str" begin
    @test ce"2 H2 + O2 → 2 H2O" == water_formation
end

@testset "==" begin
    @test ce"H2 = O2" == ce"H2 = O2"
    @test ce"H2+O2=H2O" == ce"H2 +O2 =H2O"
    @test ce"H2+O2=H2O" ≠ ce"2H2+O2=2H2O"
end

@testset "string" begin
    @test string(water_formation) == "2 H2 + O2 = 2 H2O"
    @test string(ionic_reaction) == "Na{+} + Cl{-} = NaCl"
    @test string(redox) == "Cr2O7{-2} + H{+} + e = Cr{+3} + H2O"
end

@testset "show" begin
    @test isa(repr(water_formation), String)
end

@testset "compounds" begin
    @test compounds(water_formation) == [cc"H2", cc"O2", cc"H2O"]
    @test compounds(ionic_reaction) == [cc"Na{+}", cc"Cl{-}", cc"NaCl"]
    @test compounds(redox) == [cc"Cr2O7{-2}", cc"H{+}", cc"{-}", cc"Cr{+3}", cc"H2O"]
end

@testset "elements" begin
    @test elements(water_formation) == ["H", "O"]
    @test elements(ionic_reaction) == ["Na", "Cl"]
    @test elements(redox) == ["Cr", "O", "H", "e"]
end

@testset "hascharge" begin
    @test hascharge(water_formation) == false
    @test hascharge(ionic_reaction) == true
    @test hascharge(redox) == true
end


text:
using Test
using ModelingToolkit, LightGraphs

import ModelingToolkit: value

# use a ReactionSystem to generate systems for testing
@parameters k1 k2 t
@variables S(t) I(t) R(t)

rxs = [Reaction(k1, nothing, [S]),
       Reaction(k1, [S], nothing),
       Reaction(k2, [S,I], [I], [1,1], [2]),
       Reaction(k2, [S,R], [S], [2,1], [2]),
       Reaction(k1*I, nothing, [R]),
       Reaction(k1*k2/(1+t), [S], [R])]
rs = ReactionSystem(rxs, t, [S,I,R], [k1,k2])

#################################
#  testing for Jumps / all dgs
#################################
js = convert(JumpSystem, rs)
S  = value(S); I = value(I); R = value(R)
k1 = value(k1); k2 = value(k2)
# eq to vars they depend on
eq_sdeps  = [Variable[], [S], [S,I], [S,R], [I], [S]]
eq_sidepsf = [Int[], [1], [1,2], [1,3], [2], [1]]
eq_sidepsb = [[2,3,4,6], [3,5],[4]]
deps = equation_dependencies(js)
@test all(i -> isequal(Set(eq_sdeps[i]),Set(deps[i])), 1:length(rxs))
depsbg = asgraph(js)
@test depsbg.fadjlist == eq_sidepsf
@test depsbg.badjlist == eq_sidepsb

# eq to params they depend on
eq_pdeps    = [[k1],[k1],[k2],[k2],[k1],[k1,k2]]
eq_pidepsf = [[1],[1],[2],[2],[1],[1,2]]
eq_pidepsb = [[1,2,5,6],[3,4,6]]
deps = equation_dependencies(js, variables=parameters(js))
@test all(i -> isequal(Set(eq_pdeps[i]),Set(deps[i])), 1:length(rxs))
depsbg2 = asgraph(js, variables=parameters(js))
@test depsbg2.fadjlist == eq_pidepsf
@test depsbg2.badjlist == eq_pidepsb

# var to eqs that modify them
s_eqdepsf = [[1,2,3,6],[3],[4,5,6]]
s_eqdepsb = [[1],[1],[1,2],[3],[3],[1,3]]
ne        = 8
bg        = BipartiteGraph(ne, s_eqdepsf, s_eqdepsb)
deps2     = variable_dependencies(js)
@test isequal(bg,deps2)

# eq to eqs that depend on them
eq_eqdeps = [[2,3,4,6],[2,3,4,6],[2,3,4,5,6],[4],[4],[2,3,4,6]]
dg = SimpleDiGraph(6)
for (eqidx,eqdeps) in enumerate(eq_eqdeps)
    for eqdepidx in eqdeps
       add_edge!(dg, eqidx, eqdepidx)
    end
end
dg3 = eqeq_dependencies(depsbg,deps2)
@test dg == dg3

# var to vars that depend on them
var_vardeps = [[1,2,3],[1,2,3],[3]]
ne = 7
dg = SimpleDiGraph(3)
for (vidx,vdeps) in enumerate(var_vardeps)
    for vdepidx in vdeps
       add_edge!(dg, vidx, vdepidx)
    end
end
dg4 = varvar_dependencies(depsbg,deps2)
@test dg == dg4

#####################################
#       testing for ODE/SDEs
#####################################
os = convert(ODESystem, rs)
deps = equation_dependencies(os)
eq_sdeps  = [[S,I], [S,I], [S,I,R]]
@test all(i -> isequal(Set(eq_sdeps[i]),Set(deps[i])), 1:length(deps))

sdes = convert(SDESystem, rs)
deps = equation_dependencies(sdes)
@test all(i -> isequal(Set(eq_sdeps[i]),Set(deps[i])), 1:length(deps))

deps = variable_dependencies(os)
s_eqdeps = [[1],[2],[3]]
@test deps.fadjlist == s_eqdeps

#####################################
#       testing for nonlin sys
#####################################
@variables x y z
@parameters σ ρ β

eqs = [0 ~ σ*(y-x),
       0 ~ ρ-y,
       0 ~ y - β*z]
ns = NonlinearSystem(eqs, [x,y,z],[σ,ρ,β])
deps = equation_dependencies(ns)
eq_sdeps = [[x,y],[y],[y,z]]
@test all(i -> isequal(Set(deps[i]),Set(value.(eq_sdeps[i]))), 1:length(deps))

using SparseArrays
using ModelingToolkit
using UnPack

# Define some variables
@parameters t L g
@variables x(t) y(t) w(t) z(t) T(t)
D = Differential(t)

# Simple pendulum in cartesian coordinates
eqs = [D(x) ~ w,
       D(y) ~ z,
       D(w) ~ T*x,
       D(z) ~ T*y - g,
       0 ~ x^2 + y^2 - L^2]
pendulum = ODESystem(eqs, t, [x, y, w, z, T], [L, g], name=:pendulum)
sss = SystemStructure(pendulum)
@unpack graph, fullvars, varassoc = sss
@test isequal(fullvars, [x, y, w, z, D(x), D(y), D(w), D(z), T])
@test graph.fadjlist == sort.([[5, 3], [6, 4], [7, 1, 9], [8, 2, 9], [2, 1]])
@test graph.badjlist == sort.([[3, 5], [4, 5], [1], [2], [1], [2], [3], [4], [3, 4]])
@test LightGraphs.ne(graph) == nnz(incidence_matrix(graph)) == 12
@test varassoc == [5, 6, 7, 8, 0, 0, 0, 0, 0]

se = collect(ModelingToolkit.𝑠edges(graph))
@test se == mapreduce(vcat, enumerate(graph.fadjlist)) do (s, d)
    ModelingToolkit.BipartiteEdge.(s, d)
end
de = collect(ModelingToolkit.𝑑edges(graph))
@test de == mapreduce(vcat, enumerate(graph.badjlist)) do (d, s)
    ModelingToolkit.BipartiteEdge.(s, d)
end
ae = collect(ModelingToolkit.edges(graph))
@test ae == vcat(se, de)


text:
# chemical problems
# 22
function haverlys_pooling_problem(x)
    ## Haverly's Pooling Problem
    x1 = x[1]; x2 = x[2]; x3 = x[3];
    x4 = x[4]; x5 = x[5]; x6 = x[6];
    x7 = x[7]; x8 = x[8]; x9 = x[9];
    ## objective function
    f = zeros(2);
    f[1] = -9*x1-15*x2+6*x3+16*x4;
    f[2] = 10 .* (x5+x6);
    ## constraints
    g = zeros(2);
    h = zeros(4);
    g[1] = x9 .* x7+2*x5-2.5*x1;
    g[2] = x9 .* x8+2*x6-1.5*x2;
    h[1] = x7+x8-x4-x3;
    h[2] = x1-x5-x7;
    h[3] = x2-x6-x8;
    h[4] = x9 .* x7+x9 .* x8-3 .* x3-x4;

    return f, g, h
end

# 23
function reactor_network_design(x)
    ## Reactor Network Design
    k1 = 0.09755988; k2 = 0.99*k1;
    k3 = 0.0391908; k4 = 0.9*k3;
    x1 = x[1]; x2 = x[2]; x3 = x[3];
    x4 = x[4]; x5 = x[5]; x6 = x[6];
    ## objective function
    f = zeros(2);
    f[1] = -x4;
    f[2] = x5 .^ (0.5)+x6 .^ (0.5);
    ## constraints
    g = zeros(1);
    h = zeros(4);
    g[1] = f[2]-4;
    h[1] = k1 .* x5 .* x2 + x1 -1;
    h[2] = k3 .* x5 .* x3+x3+x1-1;
    h[3] = k2 .* x6 .* x2 - x1 + x2;
    h[4] = k4 .* x6 .* x4 + x2-x1+x4-x3;

    return f, g, h
end

# 24
function heat_exchanger_network_design(x)
    ## Heat Exchanger Network Design
    x1 = x[1]; x2 = x[2]; x3 = x[3]; x4 = x[4]; x5 = x[5];
    x6 = x[6]; x7 = x[7]; x8 = x[8]; x9 = x[9];
    ## objective function
    f = zeros(3);
    f[1] = 35 .* x1 .^ (0.6)+ 35 .* x2 .^ (0.6);
    f[2] = 200 .* x1 .* x4-x3;
    f[3] = 200 .* x2 .* x6-x5;
    ## constraints
    g = zeros(1);
    h = zeros(6);
    h[1] = x3 - 1e4 .* (x7-100);
    h[2] = x5 - 1e4 .* (300-x7);
    h[3] = x3 - 1e4 .* (600-x8);
    h[4] = x5 - 1e4 .* (900-x9);
    h[5] = x4 .* log(abs(x8-100)+1e-6)-x4 .* log(abs(600-x7)+1e-6)-x8+x7+500;
    h[6] = x6 .* log(abs(x9-x7)+1e-6)-x6 .* log(600)-x9+x7+600;

    return f, g, h
end


text:
# This file was generated, do not modify it. # hide
e1 = rms(y1(), y())
e2 = rms(y2(), y())
emean = rms(0.5*y1() + 0.5*y2(), y())
estack = rms(yhat(), y())
@show e1 e2 emean estack;