text:
using ArrayInterface, StructArrays, StaticArrays, HybridArrays, LinearAlgebra, LoopVectorization, BenchmarkTools, ChainRules, Zygote, ForwardDiff, Tullio #, CUDA

## instantiate arrays of various types to test sub-components of Maxwell Operator function
Nx = 100
Ny = 100
Nz = 1
# vanilla Arrays
H = rand(Float64,2,Nx,Ny,Nz)
d = rand(Float64,3,Nx,Ny,Nz)
ei = rand(Float64,3,3,Nx,Ny,Nz)
e = rand(Float64,3,Nx,Ny,Nz) #similar(d)
# Array{SVec} & Array{SMatrix}
Hs = [SVector{2,Float64}(H[:,i,j,k]) for i=1:Nx,j=1:Ny,k=1:Nz]
ds = [SVector{3,Float64}(d[:,i,j,k]) for i=1:Nx,j=1:Ny,k=1:Nz]
eis = [SMatrix{3,3,Float64,9}(ei[:,:,i,j,k]) for i=1:Nx,j=1:Ny,k=1:Nz]
es = similar(ds)

Hs2 = reinterpret(reshape, SVector{2,Float64}, H)
ds2 = reinterpret(reshape, SVector{3,Float64}, d)
eis2 = reinterpret(reshape, SMatrix{3,3,Float64,9}, reshape(ei,(9,Nx,Ny,Nz)))
es2 = reinterpret(reshape, SVector{3,Float64}, e)

# CuArrays
# eic = cu(ei)
# dc = cu(d)
# ec = similar(dc)
# HybridArrays
Hh = HybridArray{Tuple{StaticArrays.Dynamic(),StaticArrays.Dynamic(),StaticArrays.Dynamic(),2}}(copy(permutedims(d,(2,3,4,1))))
dh = HybridArray{Tuple{3,StaticArrays.Dynamic(),StaticArrays.Dynamic(),StaticArrays.Dynamic()}}(d)
eih = HybridArray{Tuple{3,3,StaticArrays.Dynamic(),StaticArrays.Dynamic(),StaticArrays.Dynamic()}}(ei)
eh = similar(dh)

dh2 = HybridArray{Tuple{StaticArrays.Dynamic(),StaticArrays.Dynamic(),StaticArrays.Dynamic(),3}}(copy(permutedims(d,(2,3,4,1))))
eih2 = HybridArray{Tuple{StaticArrays.Dynamic(),StaticArrays.Dynamic(),StaticArrays.Dynamic(),3,3}}(copy(permutedims(eih,(3,4,5,1,2))))
eh2 = similar(dh2)

@btime @tullio $e[a,i,j,k] =  $ei[a,b,i,j,k] * $d[b,i,j,k]; # 34.657 μs (0 allocations: 0 bytes)
@btime @tullio $eh[a,i,j,k] =  $eih[a,b,i,j,k] * $dh[b,i,j,k] ; # 192.462 μs (6 allocations: 144 bytes)
@btime @tullio $eh.data[a,i,j,k] =  $eih.data[a,b,i,j,k] * $dh.data[b,i,j,k] ; # 192.462 μs (6 allocations: 144 bytes)
@btime @tullio $eh2[i,j,k,a] = $eih2[i,j,k,a,b] * $dh2[i,j,k,b] ; # 187.520 μs (5 allocations: 128 bytes)
@btime @tullio $eh2.data[i,j,k,a] = $eih2.data[i,j,k,a,b] * $dh2.data[i,j,k,b] ; # 187.520 μs (5 allocations: 128 bytes)
@btime @tullio $es[i,j,k] = $eis[i,j,k] * $ds[i,j,k] ; # 42.807 μs (0 allocations: 0 bytes)
@btime @. $es = $eis * $ds; # 22.849 μs (0 allocations: 0 bytes)
@btime @avx @. $es = $eis * $ds; # 21.716 μs (0 allocations: 0 bytes)
@btime map!(*,$es,$eis,$ds); # 21.329 μs (0 allocations: 0 bytes)
@btime vmap!(*,$es,$eis,$ds); # 21.329 μs (0 allocations: 0 bytes)
@btime vmapt!(*,$es,$eis,$ds); # 21.327 μs (0 allocations: 0 bytes)
@btime vmapnt!(*,$es,$eis,$ds); # 21.327 μs (0 allocations: 0 bytes)
@btime vmapntt!(*,$es,$eis,$ds); # 21.318 μs (0 allocations: 0 bytes)

@btime @. $es2 = $eis2 * $ds2; # 1.420 ms (20008 allocations: 625.48 KiB)
@btime @avx @. $es2 = $eis2 * $ds2; # 1.376 ms (20008 allocations: 625.48 KiB)
@btime map!(*,$es2,$eis2,$ds2); # 21.374 μs (0 allocations: 0 bytes)
@btime vmap!(*,$es2,$eis2,$ds2); # 21.329 μs (0 allocations: 0 bytes)
@btime vmapt!(*,$es2,$eis2,$ds2); # 21.327 μs (0 allocations: 0 bytes)v
@btime vmapnt!(*,$es2,$eis2,$ds2); # 21.327 μs (0 allocations: 0 bytes)
@btime vmapntt!(*,$es2,$eis2,$ds2); # 21.318 μs (0 allocations: 0 bytes)

function many3x3inverts_avx2!(Y, X)
    @assert size(Y) === size(X)
    @assert size(Y,4) == size(Y,5) === 3
    @avx for k ∈ axes(Y,3), j ∈ axes(Y,2), i ∈ axes(Y,1), h in 0:0, l in 0:0

        X₁₁ = X[i,j,k,1+l,1+h]
        X₂₁ = X[i,j,k,2+l,1+h]
        X₃₁ = X[i,j,k,3+l,1+h]
        X₁₂ = X[i,j,k,1+l,2+h]
        X₂₂ = X[i,j,k,2+l,2+h]
        X₃₂ = X[i,j,k,3+l,2+h]
        X₁₃ = X[i,j,k,1+l,3+h]
        X₂₃ = X[i,j,k,2+l,3+h]
        X₃₃ = X[i,j,k,3+l,3+h]

        Y₁₁ = X₂₂*X₃₃ - X₂₃*X₃₂
        Y₂₁ = X₂₃*X₃₁ - X₂₁*X₃₃
        Y₃₁ = X₂₁*X₃₂ - X₂₂*X₃₁

        Y₁₂ = X₁₃*X₃₂ - X₁₂*X₃₃
        Y₂₂ = X₁₁*X₃₃ - X₁₃*X₃₁
        Y₃₂ = X₁₂*X₃₁ - X₁₁*X₃₂

        Y₁₃ = X₁₂*X₂₃ - X₁₃*X₂₂
        Y₂₃ = X₁₃*X₂₁ - X₁₁*X₂₃
        Y₃₃ = X₁₁*X₂₂ - X₁₂*X₂₁

        d = 1 / ( X₁₁*Y₁₁ + X₁₂*Y₂₁ + X₁₃*Y₃₁ )

        Y[i,j,k,1+l,1+h] = Y₁₁ * d
        Y[i,j,k,2+l,1+h] = Y₂₁ * d
        Y[i,j,k,3+l,1+h] = Y₃₁ * d
        Y[i,j,k,1+l,2+h] = Y₁₂ * d
        Y[i,j,k,2+l,2+h] = Y₂₂ * d
        Y[i,j,k,3+l,2+h] = Y₃₂ * d
        Y[i,j,k,1+l,3+h] = Y₁₃ * d
        Y[i,j,k,2+l,3+h] = Y₂₃ * d
        Y[i,j,k,3+l,3+h] = Y₃₃ * d

    end
end

eh2 = similar(eih2)
es = similar(eis)
vmap!(inv,es,eis)
ei2 = permutedims(ei,(3,4,5,1,2))
eii2 = similar(ei2)

many3x3inverts_avx2!(eh2, eih2)
many3x3inverts_avx2!(eii2, ei2)
@btime vmap!(inv,$es, $eis) # 70.073 μs (0 allocations: 0 bytes)
@btime many3x3inverts_avx2!($eh2, $eih2) # 24.790 μs (0 allocations: 0 bytes)
@btime many3x3inverts_avx2!($eii2, $ei2) # 24.663 μs (0 allocations: 0 bytes)

##
using BenchmarkTools, StaticArrays, LoopVectorization
using OptiMode: calc_kpg, MaxwellGrid
Δx = 6.0
Δy = 4.0
Δz = 1.0
Nx = 100
Ny = 100
Nz = 1
k = 1.6
Neigs = 1

T = Float64
ε⁻¹ = rand(Float64,3,3,Nx,Ny,Nz)
# H⃗ = randn(Complex{T},(2*Nx*Ny*Nz,Neigs))
# H = randn(Complex{T},(2,Nx,Ny,Nz))
# e = randn(Complex{T},(3,Nx,Ny,Nz))
# d = randn(Complex{T},(3,Nx,Ny,Nz))
H⃗ = randn(T,(2*Nx*Ny*Nz,Neigs))
H = randn(T,(2,Nx,Ny,Nz))
e = randn(T,(3,Nx,Ny,Nz))
d = randn(T,(3,Nx,Ny,Nz))
mag,mn = calc_kpg(k,Δx,Δy,Δz,Nx,Ny,Nz)
m = copy(mn[:,1,:,:,:])
n = copy(mn[:,2,:,:,:])
@btime calc_kpg($k,$Δx,$Δy,$Δz,$Nx,$Ny,$Nz) # 1.452 ms (30047 allocations: 4.89 MiB)

Hs = reinterpret(reshape, SVector{2,T}, H)
ds = reinterpret(reshape, SVector{3,T}, d)
ε⁻¹s = reinterpret(reshape, SMatrix{3,3,T,9}, reshape(ε⁻¹,(9,Nx,Ny,Nz)))
es = reinterpret(reshape, SVector{3,T}, e)
ms = reinterpret(reshape, SVector{3,T}, m)
ns = reinterpret(reshape, SVector{3,T}, n)

# H⃗2 = permutedims(H⃗,(2,3,4,1))
ε⁻¹2 = permutedims(ε⁻¹,(3,4,5,1,2))
H2 = permutedims(H,(2,3,4,1))
e2 = permutedims(e,(2,3,4,1))
d2 = permutedims(d,(2,3,4,1))
mn2 = permutedims(mn,(3,4,5,2,1))
m2 = copy(mn2[:,:,:,1,:])
n2 = copy(mn2[:,:,:,2,:])

function kx_tc1!(d,H,mn,mag)
    @fastmath @inbounds for i ∈ axes(d,2), j ∈ axes(d,3), k ∈ axes(d,4)
        @fastmath @inbounds d[1,i,j,k] = ( H[1,i,j,k] * mn[1,2,i,j,k] - H[2,i,j,k] * mn[1,1,i,j,k] ) * -mag[i,j,k]
        @fastmath @inbounds d[2,i,j,k] = ( H[1,i,j,k] * mn[2,2,i,j,k] - H[2,i,j,k] * mn[2,1,i,j,k] ) * -mag[i,j,k]
        @fastmath @inbounds d[3,i,j,k] = ( H[1,i,j,k] * mn[3,2,i,j,k] - H[2,i,j,k] * mn[3,1,i,j,k] ) * -mag[i,j,k]
    end
    return d
end

function kx_tc2!(d,H,mn,mag)
    # @assert size(Y) === size(X)
    @assert size(d,4) == 3
    @assert size(H,4) === 2
    @avx for k ∈ axes(d,3), j ∈ axes(d,2), i ∈ axes(d,1), l in 0:0
        d[i,j,k,1+l] = ( H[i,j,k,1] * mn[i,j,k,2,1+l] - H[i,j,k,2] * mn[i,j,k,1,1+l] ) * -mag[i,j,k]
        d[i,j,k,2+l] = ( H[i,j,k,1] * mn[i,j,k,2,2+l] - H[i,j,k,2] * mn[i,j,k,1,2+l] ) * -mag[i,j,k]
        d[i,j,k,3+l] = ( H[i,j,k,1] * mn[i,j,k,2,3+l] - H[i,j,k,2] * mn[i,j,k,1,3+l] ) * -mag[i,j,k]
    end
    return d
end

function kx_tc3!(d,H,mn,mag)
    # @assert size(Y) === size(X)
    @assert size(d,1) == 3
    @assert size(H,1) === 2
    @avx for k ∈ axes(d,4), j ∈ axes(d,3), i ∈ axes(d,2), l in 0:0
        d[1+l,i,j,k] = ( H[1,i,j,k] * mn[1+l,2,i,j,k] - H[2,i,j,k] * mn[1+l,1,i,j,k] ) * -mag[i,j,k]
        d[2+l,i,j,k] = ( H[1,i,j,k] * mn[2+l,2,i,j,k] - H[2,i,j,k] * mn[2+l,1,i,j,k] ) * -mag[i,j,k]
        d[3+l,i,j,k] = ( H[1,i,j,k] * mn[3+l,2,i,j,k] - H[2,i,j,k] * mn[3+l,1,i,j,k] ) * -mag[i,j,k]
    end
    return d
end

function kx_tc4!(d,H,m,n,mag)
    # @assert size(Y) === size(X)
    @assert size(d,4) == 3
    @assert size(H,4) === 2
    @avx for k ∈ axes(d,3), j ∈ axes(d,2), i ∈ axes(d,1), l in 0:0
        d[i,j,k,1+l] = ( H[i,j,k,1] * n[i,j,k,1+l] - H[i,j,k,2] * m[i,j,k,1+l] ) * -mag[i,j,k]
        d[i,j,k,2+l] = ( H[i,j,k,1] * n[i,j,k,2+l] - H[i,j,k,2] * m[i,j,k,2+l] ) * -mag[i,j,k]
        d[i,j,k,3+l] = ( H[i,j,k,1] * n[i,j,k,3+l] - H[i,j,k,2] * m[i,j,k,3+l] ) * -mag[i,j,k]
    end
    return d
end

function kx_ct1!(H,e,m,n,mag)
    # @assert size(Y) === size(X)
    @assert size(e,4) == 3
    @assert size(H,4) === 2
    @avx for k ∈ axes(H,3), j ∈ axes(H,2), i ∈ axes(H,1), l in 0:0
        H[i,j,k,1+l] =  (	e[i,j,k,1+l] * n[i,j,k,1+l] + e[i,j,k,2+l] * n[i,j,k,2+l] + e[i,j,k,3+l] * n[i,j,k,3+l]	) * -mag[i,j,k]
		H[i,j,k,2+l] =  (	e[i,j,k,1+l] * m[i,j,k,1+l] + e[i,j,k,2+l] * m[i,j,k,2+l] + e[i,j,k,3+l] * m[i,j,k,3+l]	) * mag[i,j,k]
    end
    return H
end

function eid1!(e,ei,d)
    @assert size(e,4) === 3
    @assert size(d,4) === 3
    @assert size(ei,4) === 3
    @assert size(ei,5) === 3
    @avx for k ∈ axes(e,3), j ∈ axes(e,2), i ∈ axes(e,1), l in 0:0, h in 0:0
        e[i,j,k,1+h] =  ei[i,j,k,1+h,1+l]*d[i,j,k,1+l] + ei[i,j,k,2+h,1+l]*d[i,j,k,2+l] + ei[i,j,k,3+h,1+l]*d[i,j,k,3+l]
        e[i,j,k,2+h] =  ei[i,j,k,1+h,2+l]*d[i,j,k,1+l] + ei[i,j,k,2+h,2+l]*d[i,j,k,2+l] + ei[i,j,k,3+h,2+l]*d[i,j,k,3+l]
        e[i,j,k,3+h] =  ei[i,j,k,1+h,3+l]*d[i,j,k,1+l] + ei[i,j,k,2+h,3+l]*d[i,j,k,2+l] + ei[i,j,k,3+h,3+l]*d[i,j,k,3+l]
    end
    return e
end

kx_tc1!(d,H,mn,mag)
kx_tc2!(d2,H2,mn,mag)
kx_tc3!(d,H,mn,mag)
kx_tc4!(d2,H2,m2,n2,mag)

kx_ct1!(H2,e2,m2,n2,mag)

eid1!(e2,ε⁻¹2,d2)

# kx_t2c
@btime kx_tc1!($d,$H,$mn,$mag) # 74.970 μs (0 allocations: 0 bytes) real only: 52.685 μs
@btime kx_tc2!($d2,$H2,$mn2,$mag) # 26.336 μs (0 allocations: 0 bytes) real only: 14.938 μs
@btime kx_tc3!($d,$H,$mn,$mag) # real only: 30.481 μs
@btime kx_tc4!($d2,$H2,$m2,$n2,$mag) # real only: 14.872 μs (0 allocations: 0 bytes)

# kx_c2t
@btime kx_ct1!($H2,$e2,$m2,$n2,$mag) # real only: 13.784 μs (0 allocations: 0 bytes)

# eps_inv_dot
vmap!(*,es,ε⁻¹s,ds)
@btime vmap!(*,$es,$ε⁻¹s,$ds) # real only: 22.090 μs (0 allocations: 0 bytes)
@btime eid1!($e2,$ε⁻¹2,$d2) # real only: 15.297 μs (0 allocations: 0 bytes)

# FFTs
using LinearAlgebra, FFTW

Nx,Ny,Nz = 100,100,1
FFTW.set_num_threads(4)
T = Float64
e = randn(Complex{T},(3,Nx,Ny,Nz))
d = randn(Complex{T},(3,Nx,Ny,Nz))
e2 = permutedims(e,(2,3,4,1))
d2 = permutedims(d,(2,3,4,1))
e3 = copy(e2[:,:,1,:])
d3 = copy(d2[:,:,1,:])
𝓕4! = plan_fft!(e2,(1:3),flags=FFTW.PATIENT)
𝓕⁻¹4! = plan_bfft!(d2,(1:3),flags=FFTW.PATIENT)
mul!(d2,𝓕4!,d2)
mul!(e2,𝓕⁻¹4!,e2)
@btime mul!($d2,$𝓕4!,$d2)      # 112.357 μs (0 allocations: 0 bytes)
@btime mul!($e2,$𝓕⁻¹4!,$e2)

function fftw_benchmark(Nx,Ny,Nz;T=Float64)
    e = randn(Complex{T},(Nx,Ny,Nz,3))
    d = randn(Complex{T},(Nx,Ny,Nz,3))
    for nthread in [1 4 8 12 16 24]
        println("Threads: $nthread")
        FFTW.set_num_threads(nthread)
        𝓕! = plan_fft!(e,(1:3),flags=FFTW.PATIENT)
        𝓕⁻¹! = plan_ifft!(d,(1:3),flags=FFTW.PATIENT)
        mul!(d,𝓕!,d)
        mul!(e,𝓕⁻¹!,e)
        @benchmark mul!($d,$𝓕!,$d)
        @benchmark mul!($e,$𝓕⁻¹!,$e)
    end
end

fftw_benchmark(100,100,1)

### fftw_benchmak results for Nx=Ny=100,Nz=1, T=Float64
# Threads: 1
#   110.145 μs (0 allocations: 0 bytes)
#   132.275 μs (0 allocations: 0 bytes)
#   105.332 μs (0 allocations: 0 bytes)
#   120.216 μs (0 allocations: 0 bytes)
#   120.237 μs (0 allocations: 0 bytes)
#   132.937 μs (0 allocations: 0 bytes)
#   111.051 μs (0 allocations: 0 bytes)
#   122.812 μs (0 allocations: 0 bytes)
# Threads: 2
#   122.391 μs (52 allocations: 3.00 KiB)
#   179.158 μs (52 allocations: 3.00 KiB)
#   126.739 μs (26 allocations: 1.50 KiB)
#   173.404 μs (26 allocations: 1.50 KiB)
#   93.651 μs (53 allocations: 3.03 KiB)
#   117.219 μs (52 allocations: 3.00 KiB)
#   105.269 μs (0 allocations: 0 bytes)
#   116.246 μs (26 allocations: 1.50 KiB)
# Threads: 4
#   64.103 μs (80 allocations: 5.00 KiB)
#   96.021 μs (81 allocations: 5.03 KiB)
#   47.492 μs (33 allocations: 2.00 KiB)
#   86.047 μs (33 allocations: 2.00 KiB)
#   60.454 μs (80 allocations: 5.00 KiB)
#   92.712 μs (80 allocations: 5.00 KiB)
#   45.197 μs (33 allocations: 2.00 KiB)
#   85.238 μs (33 allocations: 2.00 KiB)
# Threads: 6
#   100.160 μs (189 allocations: 11.00 KiB)
#   145.808 μs (345 allocations: 20.00 KiB)
#   52.482 μs (189 allocations: 11.00 KiB)
#   98.934 μs (345 allocations: 20.00 KiB)
#   51.032 μs (109 allocations: 7.12 KiB)
#   84.613 μs (110 allocations: 7.16 KiB)
#   47.348 μs (33 allocations: 2.00 KiB)
#   84.243 μs (33 allocations: 2.00 KiB)
# Threads: 8
#   96.731 μs (231 allocations: 14.00 KiB)
#   150.332 μs (429 allocations: 26.00 KiB)
#   60.577 μs (231 allocations: 14.00 KiB)
#   107.479 μs (429 allocations: 26.00 KiB)
#   58.915 μs (138 allocations: 9.16 KiB)
#   78.913 μs (138 allocations: 9.16 KiB)
#   51.880 μs (33 allocations: 2.00 KiB)
#   86.810 μs (132 allocations: 8.00 KiB)
# Threads: 10
#   92.914 μs (273 allocations: 17.00 KiB)
#   156.300 μs (513 allocations: 32.00 KiB)
#   69.380 μs (273 allocations: 17.00 KiB)
#   112.121 μs (513 allocations: 32.00 KiB)
#   52.422 μs (168 allocations: 11.44 KiB)
#   91.838 μs (168 allocations: 11.44 KiB)
#   58.995 μs (153 allocations: 9.50 KiB)
#   96.359 μs (153 allocations: 9.50 KiB)
# Threads: 12
#   88.757 μs (273 allocations: 17.00 KiB)
#   143.094 μs (513 allocations: 32.00 KiB)
#   66.251 μs (273 allocations: 17.00 KiB)
#   116.689 μs (513 allocations: 32.00 KiB)
#   65.331 μs (195 allocations: 13.41 KiB)
#   87.204 μs (196 allocations: 13.44 KiB)
#   56.527 μs (153 allocations: 9.50 KiB)
#   96.995 μs (153 allocations: 9.50 KiB)

## full Helmholtz Operator
using ArrayInterface, LinearAlgebra, LoopVectorization, FFTW, LinearMaps, IterativeSolvers, BenchmarkTools
using OptiMode: calc_kpg

Nx,Ny,Nz = 100,100,1
Δx,Δy,Δz = 6.0,4.0,1.0
Ninv = 1.0 / ( Nx*Ny*Nz )
k = 1.6

FFTW.set_num_threads(1)
T = Float64

Hin = randn(Complex{T},(Nx,Ny,Nz,2))
Hout = similar(Hin)
e = randn(Complex{T},(Nx,Ny,Nz,3))
d = randn(Complex{T},(Nx,Ny,Nz,3))
ε⁻¹ = randn(T,(Nx,Ny,Nz,3,3))
mag,mn = calc_kpg(k,Δx,Δy,Δz,Nx,Ny,Nz)
mn2 = permutedims(mn,(3,4,5,2,1))
m = copy(mn2[:,:,:,1,:])
n = copy(mn2[:,:,:,2,:])
𝓕! = plan_fft!(copy(e),(1:3),flags=FFTW.PATIENT)
𝓕⁻¹! = plan_bfft!(copy(d),(1:3),flags=FFTW.PATIENT)
𝓕 = plan_fft(copy(e),(1:3),flags=FFTW.PATIENT)
𝓕⁻¹ = plan_bfft(copy(d),(1:3),flags=FFTW.PATIENT)

function kx_tc!(d,H,m,n,mag)::Array{ComplexF64,4}
    # @assert size(Y) === size(X)
    # @assert size(d,4) == 3
    # @assert size(H,4) === 2
    @avx for k ∈ axes(d,3), j ∈ axes(d,2), i ∈ axes(d,1), l in 0:0
        d[i,j,k,1+l] = ( H[i,j,k,1] * n[i,j,k,1+l] - H[i,j,k,2] * m[i,j,k,1+l] ) * -mag[i,j,k]
        d[i,j,k,2+l] = ( H[i,j,k,1] * n[i,j,k,2+l] - H[i,j,k,2] * m[i,j,k,2+l] ) * -mag[i,j,k]
        d[i,j,k,3+l] = ( H[i,j,k,1] * n[i,j,k,3+l] - H[i,j,k,2] * m[i,j,k,3+l] ) * -mag[i,j,k]
    end
    return d
end

function kx_ct!(H,e,m,n,mag,Ninv)::Array{ComplexF64,4}
    # @assert size(Y) === size(X)
    # @assert size(e,4) == 3
    # @assert size(H,4) === 2
    @avx for k ∈ axes(H,3), j ∈ axes(H,2), i ∈ axes(H,1), l in 0:0
        scale = mag[i,j,k] * Ninv
        H[i,j,k,1+l] =  (	e[i,j,k,1+l] * n[i,j,k,1+l] + e[i,j,k,2+l] * n[i,j,k,2+l] + e[i,j,k,3+l] * n[i,j,k,3+l]	) * -scale  # -mag[i,j,k] * Ninv
		H[i,j,k,2+l] =  (	e[i,j,k,1+l] * m[i,j,k,1+l] + e[i,j,k,2+l] * m[i,j,k,2+l] + e[i,j,k,3+l] * m[i,j,k,3+l]	) * scale   # mag[i,j,k] * Ninv
    end
    return H
end

function eid!(e,ei,d)::Array{ComplexF64,4}
    # @assert size(e,4) === 3
    # @assert size(d,4) === 3
    # @assert size(ei,4) === 3
    # @assert size(ei,5) === 3
    @avx for k ∈ axes(e,3), j ∈ axes(e,2), i ∈ axes(e,1), l in 0:0, h in 0:0
        e[i,j,k,1+h] =  ei[i,j,k,1+h,1+l]*d[i,j,k,1+l] + ei[i,j,k,2+h,1+l]*d[i,j,k,2+l] + ei[i,j,k,3+h,1+l]*d[i,j,k,3+l]
        e[i,j,k,2+h] =  ei[i,j,k,1+h,2+l]*d[i,j,k,1+l] + ei[i,j,k,2+h,2+l]*d[i,j,k,2+l] + ei[i,j,k,3+h,2+l]*d[i,j,k,3+l]
        e[i,j,k,3+h] =  ei[i,j,k,1+h,3+l]*d[i,j,k,1+l] + ei[i,j,k,2+h,3+l]*d[i,j,k,2+l] + ei[i,j,k,3+h,3+l]*d[i,j,k,3+l]
    end
    return e
end

# function M!(Hout::AbstractArray{T,4},Hin::AbstractArray{T,4},e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv) where {T<:Real}
function M!(Hout,Hin,e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)
    kx_tc!(d,Hin,m,n,mag);
    mul!(d,𝓕!,d);
    eid!(e,ε⁻¹,d);
    mul!(d,𝓕⁻¹!,d);
    kx_ct!(Hout,e,m,n,mag,Ninv);
end

function M2!(Hout,Hin,e,ef,d,df,ε⁻¹,m,n,mag,𝓕,𝓕⁻¹,Ninv)
    kx_tc!(df,Hin,m,n,mag);
    mul!(d,𝓕,df);
    eid!(e,ε⁻¹,d);
    mul!(ef,𝓕⁻¹,e);
    kx_ct!(Hout,ef,m,n,mag,Ninv);
end

function Mnofft!(Hout,Hin,e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)
    kx_tc!(d,Hin,m,n,mag);
    eid!(e,ε⁻¹,d);
    kx_ct!(Hout,e,m,n,mag,Ninv);
end

# function M!(Hout::AbstractVector,Hin::AbstractVector,e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv) where {T<:Real}
#     s = size(mag)
#     vec(M!(reshape(Hout,(s[1],s[2],s[3],2)),reshape(Hin,(s[1],s[2],s[3],2),e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)))
# end

ef = similar(e)
df = similar(d)

Hin = randn(ComplexF64,size(Hin))
Hout = randn(ComplexF64,size(Hin))
# Mop * Hin

kx_tc!(d,Hin,m,n,mag);
mul!(d,𝓕!,d);
eid!(e,ε⁻¹,d);
mul!(e,𝓕⁻¹!,e);
kx_ct!(Hout,e,m,n,mag,Ninv);
M!(Hout,Hin,e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)
M2!(Hout,Hin,e,ef,d,df,ε⁻¹,m,n,mag,𝓕,𝓕⁻¹,Ninv)
Mnofft!(Hout,Hin,e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)

@benchmark kx_tc!($d,$Hin,$m,$n,$mag)
@benchmark mul!($d,$𝓕!,$d)
@benchmark eid!($e,$ε⁻¹,$d)
@benchmark mul!($d,$𝓕⁻¹!,$d)
@benchmark kx_ct!(Hout,e,m,n,mag,Ninv)
@benchmark M!($Hout,$Hin,$e,$d,$ε⁻¹,$m,$n,$mag,$𝓕!,$𝓕⁻¹!,$Ninv)

# FFTW threads = 4
# @btime M!($Hout,$Hin,$e,$d,$ε⁻¹,$m,$n,$mag,$𝓕!,$𝓕⁻¹!,$Ninv) # 216.260 μs (144 allocations: 8.50 KiB)
# @btime M2!($Hout,$Hin,$e,$ef,$d,$df,$ε⁻¹,$m,$n,$mag,$𝓕,$𝓕⁻¹,$Ninv) # 214.032 μs (144 allocations: 8.50 KiB)
# @btime Mnofft!($Hout,$Hin,$e,$d,$ε⁻¹,$m,$n,$mag,$𝓕!,$𝓕⁻¹!,$Ninv) # 84.440 μs (0 allocations: 0 bytes)

# FFTW threads = 1
@btime M!($Hout,$Hin,$e,$d,$ε⁻¹,$m,$n,$mag,$𝓕!,$𝓕⁻¹!,$Ninv) # 295.195 μs (0 allocations: 0 bytes)
@btime M2!($Hout,$Hin,$e,$ef,$d,$df,$ε⁻¹,$m,$n,$mag,$𝓕,$𝓕⁻¹,$Ninv) # 82.470 μs (0 allocations: 0 bytes)
@btime Mnofft!($Hout,$Hin,$e,$d,$ε⁻¹,$m,$n,$mag,$𝓕!,$𝓕⁻¹!,$Ninv) # 273.492 μs (0 allocations: 0 bytes)

##
using LinearMaps

HinV = copy(vec(Hin))
HoutV = copy(vec(Hout))
Mop = LinearMap{ComplexF64}((y,x)->vec(M!(reshape(y,(Nx,Ny,Nz,2)),reshape(x,(Nx,Ny,Nz,2)),e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)),2*length(mag),ishermitian=true,ismutating=true)
Mop_pd = LinearMap{ComplexF64}((y,x)->vec(M!(reshape(y,(Nx,Ny,Nz,2)),reshape(x,(Nx,Ny,Nz,2)),e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)),2*length(mag),ishermitian=true,ismutating=true,isposdef=true)
M̂1!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv) = LinearMap{ComplexF64}((y,x)->vec(M!(reshape(y,(Nx,Ny,Nz,2)),reshape(x,(Nx,Ny,Nz,2)),e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)),2*length(mag),ishermitian=true,ismutating=true)
M̂2!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv,Nx,Ny,Nz) = LinearMap{ComplexF64}((y,x)->vec(M!(reshape(y,(Nx,Ny,Nz,2)),reshape(x,(Nx,Ny,Nz,2)),e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)),2*length(mag),ishermitian=true,ismutating=true)
M̂3!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv) = LinearMap{ComplexF64}((y,x)->vec(M!(reshape(y,(size(mag)...,2)),reshape(x,(size(mag)...,2)),e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)),2*length(mag),ishermitian=true,ismutating=true)
M̂5!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv) = LinearMap{ComplexF64}((y::Vector{ComplexF64},x::Vector{ComplexF64})->vec(Mnofft!(reshape(y,(size(mag)...,2)),reshape(x,(size(mag)...,2)),e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv))::Vector{ComplexF64},2*length(mag),ishermitian=true,ismutating=true)
function M̂4!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)
    let (Nx,Ny,Nz) = size(mag), Mmap(y,x) = vec(M!(reshape(y,(Nx,Ny,Nz,2)),reshape(x,(Nx,Ny,Nz,2)),e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)), N = (2*Nx*Ny*Nz)
        Mop = LinearMap{ComplexF64}(Mmap,N,ishermitian=true,ismutating=true)
    end
end
f_Mop!(y::Vector{ComplexF64},x::Vector{ComplexF64})::Vector{ComplexF64} =  vec(Mnofft!(reshape(y,(Nx,Ny,Nz,2)),reshape(x,(Nx,Ny,Nz,2)),e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv))
N_Mop = 2*Nx*Ny*Nz
Mop5 = LinearMap{ComplexF64}(f_Mop!,f_Mop!,N_Mop,N_Mop,ishermitian=true,ismutating=true)

# function LinearAlgebra.mul!(y::AbstractVecOrMat, A::LinearMaps.FunctionMap, x::AbstractVector)
#     # LinearMaps.check_dim_mul(y, A, x)
#     return fill!(y, iszero(A.λ) ? zero(eltype(y)) : A.λ*sum(x))
# end

Mop8 = let Nx=Nx,Ny=Ny,Nz=Nz,Δx=Δx,Δy=Δy,Δz=Δz,k=k
    Ninv = 1.0 / ( Nx*Ny*Nz )
    N = 2*Nx*Ny*Nz
    shp = Nx,Ny,Nz,2
    # Hin = randn(Complex{T},(Nx,Ny,Nz,2))
    # Hout = similar(Hin)
    e = randn(Complex{T},(Nx,Ny,Nz,3))
    d = randn(Complex{T},(Nx,Ny,Nz,3))
    ε⁻¹ = randn(T,(Nx,Ny,Nz,3,3))
    mag,mn = calc_kpg(k,Δx,Δy,Δz,Nx,Ny,Nz)
    mn2 = permutedims(mn,(3,4,5,2,1))
    m = copy(mn2[:,:,:,1,:])
    n = copy(mn2[:,:,:,2,:])
    F! = plan_fft!(copy(e),(1:3),flags=FFTW.PATIENT)
    F⁻¹! = plan_bfft!(copy(d),(1:3),flags=FFTW.PATIENT)
    function M(y,x)
        kx_tc!(d,x,m,n,mag);
        mul!(d,F!,d);
        eid!(e,ε⁻¹,d);
        mul!(d,F⁻¹!,d);
        kx_ct!(y,e,m,n,mag,Ninv);
    end
    # LinearMap{ComplexF64}((2*Nx*Ny*Nz),ishermitian=true,ismutating=true) do (y,x)
    #     vec(M!(reshape(y,(Nx,Ny,Nz,2)),reshape(x,(Nx,Ny,Nz,2)),e,d,ε⁻¹,m,n,mag,F!,F⁻¹!,Ninv))
    LinearMap{ComplexF64}((y,x)->vec(M!(reshape(y,shp),reshape(x,shp),e,d,ε⁻¹,m,n,mag,F!,F⁻¹!,Ninv)),N,ishermitian=true,ismutating=true)
end

Mop1 = M̂1!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)
Mop2 = M̂2!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv,Nx,Ny,Nz)
Mop3 = M̂3!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)
Mop4 = M̂4!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)
Mop5 = M̂5!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)
mul!(HoutV,Mop,HinV)
mul!(HoutV,Mop1,HinV)
mul!(HoutV,Mop2,HinV)
mul!(HoutV,Mop3,HinV)
mul!(HoutV,Mop4,HinV)
mul!(HoutV,Mop_pd,HinV)
mul!(HoutV,Mop5,HinV)
mul!(HoutV,Mop8,HinV)

struct MaxwellMap{T, F} <: LinearMaps.LinearMap{Complex{T}}
    e::Array{Complex{T}, 4}
    d::Array{Complex{T}, 4}
    ε⁻¹::Array{T, 5}
    m::Array{T, 4}
    n::Array{T, 4}
    mag::Array{T, 3}
    𝓕!::FFTW.cFFTWPlan{Complex{T}}
    𝓕⁻¹!::FFTW.cFFTWPlan{Complex{T}}
    Ninv::T
    f::F
    N::Int
    _ismutating::Bool
    _issymmetric::Bool
    _ishermitian::Bool
    _isposdef::Bool
end

mutable struct mMaxwellMap{T, F} <: LinearMaps.LinearMap{Complex{T}}
    e::Array{Complex{T}, 4}
    d::Array{Complex{T}, 4}
    ε⁻¹::Array{T, 5}
    m::Array{T, 4}
    n::Array{T, 4}
    mag::Array{T, 3}
    𝓕!::FFTW.cFFTWPlan{Complex{T}}
    𝓕⁻¹!::FFTW.cFFTWPlan{Complex{T}}
    Ninv::T
    f::F
    N::Int
    _ismutating::Bool
    _issymmetric::Bool
    _ishermitian::Bool
    _isposdef::Bool
end

function mMaxwellMap{T}(f::F, N::Int;
    ismutating::Bool  = _ismutating(f),
    issymmetric::Bool = false,
    ishermitian::Bool = (T<:Real && issymmetric),
    isposdef::Bool    = false) where {T, F1, F2}
    FunctionMap{T, F1, F2}(f, fc, M, N, ismutating, issymmetric, ishermitian, isposdef)
end

# function (f::mMaxwellMap)(i, nm, TT)
#     pos_i, x_i = readvalue(f.buf, f.pos, f.len, TT; f.kw...)
#     f.pos = pos_i
#     return x_i
# end

# FFTW threads = 4
# @btime mul!($HoutV,$Mop,$HinV)  # 224.553 μs (152 allocations: 8.92 KiB)
# @btime mul!($HoutV,$Mop1,$HinV) # 217.570 μs (153 allocations: 8.94 KiB)
# @btime mul!($HoutV,$Mop2,$HinV) # 222.795 μs (150 allocations: 8.83 KiB)
# @btime mul!($HoutV,$Mop3,$HinV) # 218.957 μs (150 allocations: 8.83 KiB)
# @btime mul!($HoutV,$Mop4,$HinV) # 226.071 μs (150 allocations: 8.83 KiB)
# @btime mul!($HoutV,$Mop_pd,$HinV)  # 224.553 μs (152 allocations: 8.92 KiB)
# @btime mul!($HoutV,$Mop5,$HinV) # 217.802 μs (152 allocations: 8.92 KiB)

# FFTW threads = 1
@btime mul!($HoutV,$Mop,$HinV)  # 302.679 μs (8 allocations: 432 bytes)
@btime mul!($HoutV,$Mop1,$HinV) # 287.758 μs (9 allocations: 448 bytes)
@btime mul!($HoutV,$Mop2,$HinV) # 285.625 μs (6 allocations: 336 bytes)
@btime mul!($HoutV,$Mop3,$HinV) # 287.651 μs (6 allocations: 336 bytes)
@btime mul!($HoutV,$Mop4,$HinV) # 289.960 μs (6 allocations: 336 bytes)
@btime mul!($HoutV,$Mop_pd,$HinV) # 295.120 μs (8 allocations: 432 bytes)
@btime mul!($HoutV,$Mop5,$HinV) # 295.057 μs (8 allocations: 432 bytes)
@btime mul!($HoutV,$Mop8,$HinV)

function mymul!(y::AbstractVecOrMat, A::LinearMaps.FunctionMap, x::AbstractVector)
    # LinearMaps.ismutating(A) ? A.f(y, x) : copyto!(y, A.f(x))
    A.f(y, x)
    # return y
end

mymul!(HoutV,Mop3,HinV)
@btime mymul!($HoutV,$Mop3,$HinV)
Mop3.f(HoutV,HinV)
@btime @inbounds $Mop5.f($HoutV,$HinV)
Mop4.f
Mop6 = LinearMap{ComplexF64}((y,x)->vec(deepcopy(M!)(reshape(y,(100,100,1,2)),reshape(x,(100,100,1,2)),e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)),2*100*100*1,ishermitian=true,ismutating=true)

x = Vector{ComplexF64}(randn(ComplexF64,20000))
y = Vector{ComplexF64}(randn(ComplexF64,20000))
mymul!(y,Mop6,x)
@btime mymul!($y,$Mop6,$x)
# struct MyFnMap{T} <: LinearMaps.FunctionMap{T}
#     λ::T
#     size::Dims{2}
#     function MyFnMap(λ::T, dims::Dims{2}) where {T}
#         # all(≥(0), dims) || throw(ArgumentError("dims of MyFillMap must be non-negative"))
#         promote_type(T, typeof(λ)) == T || throw(InexactError())
#         return new{T}(λ, dims)
#     end
# end

LinearMaps._unsafe_mul!(HoutV,Mop3,HinV)
LinearMaps._unsafe_mul!(HoutV,Mop4,HinV)

@btime LinearMaps._unsafe_mul!($HoutV,$Mop3,$HinV) # 287.651 μs (6 allocations: 336 bytes)
@btime LinearMaps._unsafe_mul!($HoutV,$Mop4,$HinV)

function mul_test()
    Nx,Ny,Nz = 100,100,1
    Δx,Δy,Δz = 6.0,4.0,1.0
    Ninv = 1.0 / ( Nx*Ny*Nz )
    k = 1.6
    Hin = randn(Complex{Float64},(Nx,Ny,Nz,2))
    Hout = similar(Hin)
    e = randn(Complex{Float64},(Nx,Ny,Nz,3))
    d = randn(Complex{Float64},(Nx,Ny,Nz,3))
    ε⁻¹ = randn(T,(Nx,Ny,Nz,3,3))
    mag,mn = calc_kpg(k,Δx,Δy,Δz,Nx,Ny,Nz)
    mn2 = permutedims(mn,(3,4,5,2,1))
    m = copy(mn2[:,:,:,1,:])
    n = copy(mn2[:,:,:,2,:])
    𝓕! = plan_fft!(copy(e),(1:3)) #,flags=FFTW.PATIENT)
    𝓕⁻¹! = plan_bfft!(copy(d),(1:3)) #,flags=FFTW.PATIENT)
    Mop_test = M̂3!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)
    HinV = copy(vec(Hin))
    HoutV = copy(vec(Hout))
    mul!(HoutV,Mop_test,HinV)
    @benchmark mul!(HoutV,$(M̂3!(e,d,ε⁻¹,m,n,mag,𝓕!,𝓕⁻¹!,Ninv)),HinV)
end
mul_test()

##
function fnew(y,x)
    @. y = 3 + 2.0im
end

Mop4c.f = fnew

Mop4c.f(HoutV,HinV) ≈ Mop3.f(HoutV,HinV)

##
"""
    t2c: v (transverse vector) → a (cartesian vector)
"""
function t2c(v::SVector{2,ComplexF64},k::KVec)::SVector{3,ComplexF64}
    return v[1] * k.m + v[2] * k.n
end

"""
    c2t: a (cartesian vector) → v (transverse vector)
"""
function c2t(a::SVector{3,ComplexF64},k::KVec)::SVector{2,ComplexF64}
    v0 = a ⋅ k.m
    v1 = a ⋅ k.n
    return SVector(v0,v1)
end

"""
    kcross_t2c: a (cartesian vector) = k × v (transverse vector)
"""
function kcross_t2c(v::SVector{2,ComplexF64},k::KVec)::SVector{3,ComplexF64}
    return ( v[1] * k.n - v[2] * k.m ) * k.mag
end

"""
    kcross_c2t: v (transverse vector) = k × a (cartesian vector)
"""
function kcross_c2t(a::SVector{3,ComplexF64},k::KVec)::SVector{2,ComplexF64}
    at1 = a ⋅ k.m
    at2 = a ⋅ k.n
    v0 = -at2 * k.mag
    v1 = at1 * k.mag
    return SVector(v0,v1)
end

"""
    kcrossinv_t2c: compute a⃗ (cartestion vector) st. v⃗ (cartesian vector from two trans. vector components) ≈ k⃗ × a⃗
    This neglects the component of a⃗ parallel to k⃗ (not available by inverting this cross product)
"""
function kcrossinv_t2c(v::SVector{2,ComplexF64},k::KVec)::SVector{3,ComplexF64}
    return ( v[1] * k.n - v[2] * k.m ) * ( -1 / k.mag )
end

"""
    kcrossinv_c2t: compute  v⃗ (transverse 2-vector) st. a⃗ (cartestion 3-vector) = k⃗ × v⃗
    This cross product inversion is exact because v⃗ is transverse (perp.) to k⃗
"""
function kcrossinv_c2t(a::SVector{3,ComplexF64},k::KVec)::SVector{2,ComplexF64}
    at1 = a ⋅ k.m
    at2 = a ⋅ k.n
    v0 = -at2 * (-1 / k.mag )
    v1 = at1 * ( -1 / k.mag )
    return SVector(v0,v1)
end

"""
    ucross_t2c: a (cartesian vector) = u × v (transverse vector)
"""
function ucross_t2c(u::SVector{3,ComplexF64},v::SVector{2,ComplexF64},k::KVec)::SVector{3,ComplexF64}
    return cross(u,t2c(v,k))
end

##

function eid1!(e::AbstractArray{T},ei::AbstractArray{T},d::AbstractArray{T},Nx,Ny,Nz)::AbstractArray{T} where T
    @fastmath @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
        @fastmath @inbounds e[1,i,j,k] =  ei[1,1,i,j,k]*d[1,i,j,k] + ei[2,1,i,j,k]*d[2,i,j,k] + ei[3,1,i,j,k]*d[3,i,j,k]
        @fastmath @inbounds e[2,i,j,k] =  ei[1,2,i,j,k]*d[1,i,j,k] + ei[2,2,i,j,k]*d[2,i,j,k] + ei[3,2,i,j,k]*d[3,i,j,k]
        @fastmath @inbounds e[3,i,j,k] =  ei[1,3,i,j,k]*d[1,i,j,k] + ei[2,3,i,j,k]*d[2,i,j,k] + ei[3,3,i,j,k]*d[3,i,j,k]
    end
    return e
end

function eid2!(e::AbstractArray{T},ei::AbstractArray{T},d::AbstractArray{T},Nx,Ny,Nz)::AbstractArray{T} where T
    @fastmath @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
         @simd for a = 1:3
            @fastmath @inbounds @views e[:,i,j,k] .=  ei[a,:,i,j,k]*d[a,i,j,k]
        end
    end
    return e
end

function eid3!(e::AbstractArray{SVector{3,T}},ei::AbstractArray{SMatrix{3,3,T,9}},d::AbstractArray{SVector{3,T}})::AbstractArray{SVector{3,T}} where T
    @fastmath broadcast!(*,es,eis,ds)
end

function eid4(ei::AbstractArray{T},d::AbstractArray{T})::AbstractArray{T} where T
    @tullio e[a,i,j,k] :=  ei[a,b,i,j,k] * d[b,i,j,k]
end

function eid4!(e::AbstractArray{T},ei::AbstractArray{T},d::AbstractArray{T})::AbstractArray{T} where T
    @tullio e[a,i,j,k] =  ei[a,b,i,j,k] * d[b,i,j,k] fastmath=true
end

function eid5!(es,eis, ds)
    @tullio es[i,j,k] = eis[i,j,k] * ds[i,j,k]
end

function eid5(eis, ds)
    @tullio es[i,j,k] := eis[i,j,k] * ds[i,j,k]
end

### Warmups

# Array + hand written for loop
eid1!(e,ei,d,Nx,Ny,Nz)
# Array + @simd + @views + broadcast! + for loop
eid2!(e,ei,d,Nx,Ny,Nz) ≈ eid1!(e,ei,d,Nx,Ny,Nz)
# StaticArrays+broadcast!
eid3!(es,eis,ds)
# normal Array+Tullio
eid4(ei,d)
eid4!(e,ei,d)
# CuArrays + Tullio
# eid4(eic,dc)
# eid4!(ec,eic,dc)
# Array(eid4(eic,dc))
# HybridArrays+Tullio
eid4(eih,dh)
eid4!(eh,eih,dh)
# StaticArrays+Tullio
eid5(eis,ds)
eid5!(es,eis,ds)

### Benchmarks

# Array + hand written for loop
@btime eid1!($e,$ei,$d,$Nx,$Ny,$Nz) #74.750 μs (0 allocations: 0 bytes)
# Array + @simd + @views + broadcast! + for loop
@btime eid2!($e,$ei,$d,$Nx,$Ny,$Nz) #1.308 ms (30000 allocations: 3.20 MiB)
# StaticArrays+broadcast!
@btime eid3!($es,$eis,$ds) # 22.703 μs (1 allocation: 48 bytes)
# normal Array+Tullio
@btime eid4!($e,$ei,$d) # 32.310 μs (0 allocations: 0 bytes)
@btime eid4($ei,$d) # 33.921 μs (3 allocations: 234.50 KiB)
# CuArrays + Tullio
# @btime eid4($eic,$dc) # 19.833 μs (137 allocations: 5.53 KiB), 1.069 s (2228391 allocations: 93.97 MiB) w/o KernelAbstractions
# @btime eid4!($ec,$eic,$dc) # 84.793 μs (127 allocations: 5.31 KiB)
# @btime Array(eid4($eic,$dc)) # 2.039 ms (142 allocations: 122.88 KiB)
# HybridArrays+Tullio
@btime eid4($eih,$dh) # 180.503 μs (11 allocations: 234.67 KiB)
@btime eid4!($eh,$eih,$dh) # 197.167 μs (6 allocations: 144 bytes)
@btime eid4($(eih.data),$(dh.data)) # 33.650 μs (3 allocations: 234.50 KiB)
@btime eid4!($(eh.data),$(eih.data),$(dh.data)) # 33.003 μs (0 allocations: 0 bytes)
# StaticArrays+Tullio
@btime eid5($eis,$ds) # 180.503 μs (11 allocations: 234.67 KiB)
@btime eid5!($es,$eis,$ds) # 197.167 μs (6 allocations: 144 bytes)
## grads
Zygote.@adjoint (T::Type{<:SArray})(xs::Number...) = T(xs...), dv -> (nothing, dv...)
Zygote.@adjoint (T::Type{<:SArray})(x::AbstractArray) = T(x), dv -> (nothing, dv)
Zygote.@adjoint (T::Type{<:SMatrix})(xs::Number...) = T(xs...), dv -> (nothing, dv...)
Zygote.@adjoint (T::Type{<:SMatrix})(x::AbstractMatrix) = T(x), dv -> (nothing, dv)
Zygote.@adjoint (T::Type{<:SVector})(xs::Number...) = T(xs...), dv -> (nothing, dv...)
Zygote.@adjoint (T::Type{<:SVector})(x::AbstractVector) = T(x), dv -> (nothing, dv)
Zygote.@adjoint enumerate(xs) = enumerate(xs), diys -> (map(last, diys),)
_ndims(::Base.HasShape{d}) where {d} = d
_ndims(x) = Base.IteratorSize(x) isa Base.HasShape ? _ndims(Base.IteratorSize(x)) : 1
Zygote.@adjoint function Iterators.product(xs...)
                    d = 1
                    Iterators.product(xs...), dy -> ntuple(length(xs)) do n
                        nd = _ndims(xs[n])
                        dims = ntuple(i -> i<d ? i : i+nd, ndims(dy)-nd)
                        d += nd
                        func = sum(y->y[n], dy; dims=dims)
                        ax = axes(xs[n])
                        reshape(func, ax)
                    end
                end
#

function sum2(op,arr)
    return sum(op,arr)
end

function sum2adj( Δ, op, arr )
    n = length(arr)
    g = x->Δ*Zygote.gradient(op,x)[1]
    return ( nothing, map(g,arr))
end

Zygote.@adjoint function sum2(op,arr)
    return sum2(op,arr),Δ->sum2adj(Δ,op,arr)
end

Zygote.refresh()
# using ReverseDiff: TrackedArray
# @inline Base.:+(x::TrackedArray{V,D}, y::StaticArray) where {V,D} = record_plus(x, Array(y), D)
# @inline Base.:+(x::StaticArray, y::TrackedArray{V,D}) where {V,D} = record_plus(Array(x), y, D)
# @inline Base.:-(x::TrackedArray{V,D}, y::StaticArray) where {V,D} = record_minus(x, Array(y), D)
# @inline Base.:-(x::StaticArray, y::TrackedArray{V,D}) where {V,D} = record_minus(Array(x), y, D)
# deriv!(t::StaticArray, v::AbstractArray) = deriv!(Tuple(t), Tuple(v))
# `StaticArray`s don't support mutation unless the eltype is a bits type (`isbitstype`).
# capture(t::SA) where SA <: StaticArray = istracked(t) ? SA(map(capture, t)) : copy(t)
##
function foo!(a,b,c)
    sum(sum(broadcast!(*,a,b,c)))
end

function foo(b,c)
    sum(sum(broadcast(*,b,c)))
end

function foot(ei,d)
    # @tullio esum :=  ei[a,b,i,j,k] * d[b,i,j,k]
    @tullio e :=  ei[a,b,i,j,k] * d[b,i,j,k]
end

function foots(eis::Array{SArray{Tuple{3,3},Float64,2,9},3},ds::Array{SArray{Tuple{3},Float64,1,3},3})::Float64
    # @tullio esum :=  ei[a,b,i,j,k] * d[b,i,j,k]
    @tullio es_sum :=  eis[i,j,k] * ds[i,j,k]
    sum(es_sum)
    # es = vmap(*,eis,ds)
    # sum(sum(es))
end

foo(eis,ds)
foo!(es,eis,ds)
foot(ei,d)
# foot(cu(ei),cu(dc))
foot(eih,dh)
foots(eis,ds)

@btime foo($eis,$ds) # 33.774 μs (2 allocations: 234.45 KiB)
@btime foo!($es,$eis,$ds) # 31.154 μs (0 allocations: 0 bytes)
@btime foot($ei,$d) # 125.603 μs (1 allocation: 16 bytes)
# @btime foot($eic,$dc) # 1.025 s (2198390 allocations: 91.22 MiB)
@btime foot($eih,$dh) # 80.442 μs (10 allocations: 224 bytes)
@btime foots($eis,$ds) # 32.728 μs (2 allocations: 234.45 KiB)
###### Grads!

using ReverseDiff: GradientTape, GradientConfig, gradient, gradient!, compile, DiffResults

##Zygote.gradient(foo,eis,ds)
Zygote.gradient(foo,eis,ds)
Zygote.gradient(foot,ei,d)
Zygote.gradient(foot,eih,dh)
Zygote.gradient(foot,eih.data,dh.data)
Zygote.gradient(foots,eis,ds)
# Zygote.gradient(foot,eic,dc)
const f_tape = GradientTape(foot, (rand(size(ei)...), rand(size(d)...)))

inputs = (ei,d)
results = (similar(ei),similar(d));
all_results = map(DiffResults.GradientResult, results);
cfg = GradientConfig(inputs)

Tracker.gradient(foot,eih,dh)

# ReverseDiff
const compiled_f_tape = compile(GradientTape(foot, (rand(size(ei)...), rand(size(d)...))))
inputs = (ei, d)
results = (similar(ei), similar(d))
all_results = map(DiffResults.GradientResult, results)
cfg = GradientConfig(inputs)
ReverseDiff.gradient!(results, compiled_f_tape, inputs)
ReverseDiff.gradient!(all_results, compiled_f_tape, inputs) # the same as the above but primal is loaded into the the provided `DiffResult` instances (see DiffResults.jl documentation).

### Benchmarks
@btime ReverseDiff.gradient!($results, $compiled_f_tape, $inputs) # 33.863 ms (11 allocations: 512 bytes)
@btime ReverseDiff.gradient!($all_results, $compiled_f_tape, $inputs) # 34.298 ms (11 allocations: 512 bytes)
@btime Zygote.gradient(foo,$eis,$ds) # 1.337 ms (50048 allocations: 6.11 MiB)
@btime Zygote.gradient(foot,$ei,$d) # 456.707 μs (83 allocations: 941.67 KiB)
# @btime Zygote.gradient(foot,$eic,$dc) # 40.263 ms (394 allocations: 16.59 KiB)
@btime Zygote.gradient(foot,$eih,$dh) # 434.074 μs (96 allocations: 941.98 KiB)
@btime Zygote.gradient(foots,$eis,$ds) # 434.074 μs (96 allocations: 941.98 KiB)

ReverseDiff.gradient(foot,eih,dh)
foot(cu(ei),cu(d))

##
mul(A, B) = @tullio C[i,k] := A[i,j] * B[j,k]

A = rand(3,40); B = rand(40,500);
A * B ≈ mul(A, B) # true

using Tracker # or Zygote
ΔA = Tracker.gradient((A,B) -> sum(mul(A, B)), A, B)[1]
ΔA ≈ ones(3,500) * B' # true

using CUDA, KernelAbstractions # Now defined with a GPU version:
mul(A, B) = @tullio C[i,k] := A[i,j] * B[j,k]

cu(A * B) ≈ mul(cu(A), cu(B)) # true

cu(ΔA) ≈ Tracker.gradient((A,B) -> sum(mul(A, B)), cu(A), cu(B))[1] ; # true

# Reduction over min/max:
Tracker.gradient(x -> (@tullio (max) res := x[i]^3), [1,2,3,-2,-1,3])[1]
f_cuda = (X,Y) -> sum(mul(X, Y))
@btime Tracker.gradient(f_cuda, $(cu(A)), $(cu(B))) # 435.266 μs (524 allocations: 16.63 KiB)
@btime Zygote.gradient(f_cuda, $(cu(A)), $(cu(B))) # 439.961 μs (450 allocations: 16.94 KiB)

eic

dc


text:
"Tell whether there are too foo items in the array."
foo(xs::Array) = ...

ext: jl

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
include("dipole.jl")

const c = 299792458.
const mu0 = 4*pi*1e-7
const eps0 = 1/(mu0*c^2)

#dipoles:
const N=1000;
const distance=1000+1000*rand(N)
const theta=2*acos(rand(N)-1)
const phi=2*pi*rand(N)
xyz = hcat(distance.*sin(theta).*cos(phi), distance.*sin(theta).*sin(phi), distance.*cos(theta))
#moments:
const th_dip=2*acos(rand(N)-1)
const phi_dip=2*pi*rand(N)

#dipole moment
#total time averaged radiated power P= 1 W dipole moment => |p|=sqrt(12πcP/µOω⁴)
const Pow=1
const amp_max=sqrt(12*pi*c*Pow/(mu0*(2*pi*freq)^4))/N
const amplitude=amp_max*rand(N)
const phase=2*pi*rand(N)
I=hcat(amplitude.*sin(th_dip).*cos(phi_dip), amplitude.*sin(th_dip).*sin(phi_dip), amplitude.*cos(th_dip))

const freq=1e9

#observation points
M=500
#observation points
x=rand(M)*(10)-5
y=rand(M)*(10)-5
z=rand(M)*(10)-5

println("Computing the radiation...")
Z=zeros(M)
for i=1:M
  perc=round(i/M*100)
  println("$perc %")
  r=[x[i],y[i],z[i]]
  E=zeros(Complex128,3)
  B=zeros(Complex128,3)
  for m=1:N
    p=vec(I[m,:])
    R=vec(xyz[m,:])
    Ed,Bd=Hertz_dipole (r, p, R, phase[m], freq)
    E+=Ed
    B+=Bd
  end
  Z[i]=sqrt(sum(abs(E).^2))/sqrt(sum(abs(B).^2))*mu0 #real(E).^2#0.5*numpy.cross(E.T,conjugate(B.T))
end

using PyPlot
plt.hist(Z)

max_stars_repo_name: manuamador/Z

max_issues_repo_name: manuamador/Z

max_forks_repo_name: manuamador/Z


text:
module lmp_str

    import Dates
    using Printf
    using Statistics
    import LinearAlgebra.norm

    abstract type Str end
    abstract type Potential end
    abstract type Potential_Unit end
    abstract type Data end
    abstract type Unit end

    # Main Files
    include("Data.jl")
    export Atom, Bond, Angle, Dihedral, Improper
    export Data_Cell, Data_Basic, Data_Unit, Data_Sum

    include("useful_funcs.jl")
    export max, min, get_data, diag, conv, dist, norm_vec, rot_mat
    export add!, set!, central_point_atom, central_point_box, copy_array

    include("Constants.jl")
    export const_list
    #=
    export Const_k_b, Const_n_a, Const_density_wat,
        Const_kcal2j, Const_kcalm2j, Const_kcalm2t,
        Const_g2kg, Const_kg2g, Const_gm2g, Const_gm2kg,
        Const_an2m, Const_an2nm, Const_nm2m, Const_cm2an, Const_dm2m, Const_m2dm,
        Const_fs2s, Const_ps2s, Const_ns2s
    =#

    # Structure Files
    include("./Structure/Structure_Basic.jl")
    export Structure_Basic, transform
    export structure_sc, structure_bcc, structure_fcc, structure_dc

    include("./Structure/Structure_Wat.jl")
    export Structure_Wat
    export structure_tip3p, structure_spc, structure_spce
    export structure_tip4p_cut, structure_tip4p_long
    export structure_tip4p_2005, structure_tip4p_ice
    export structure_tip4p_ew, structure_tip4p_fq
    export structure_tip5p, structure_tip4p_2005, structure_tip5p_2018
    export structure_tip7p

    include("./Structure/Structure_Si.jl")
    export Structure_Si
    export structure_si, structure_sio2
    export structure_si3n4, structure_si3n4_ort, structure_si3n4_pdb, structure_si3n4_ort_pdb

    include("./Structure/Structure_C.jl")
    export Structure_C
    export structure_graphene, structure_graphene_ort

    include("./Structure/Structure_Ti.jl")
    export Structure_Ti
    export structure_tio2_anatase, structure_tio2_rutile

    include("./Structure/Structure_Ca.jl")
    export Structure_Ca
    export structure_caco3

    include("./Structure/Structure_Ion.jl")
    export Structure_Ion

    include("./Structure/Structure_VMD.jl")
    export Structure_VMD

    # generator Files
    include("./generator/genr.jl")
    export genr, genr_cell, genr_atom
    export genr_bond, genr_angle, genr_dihedral, genr_improper

    include("./generator/modify.jl")
    export move, find, select, remove!, remove_else!

    include("./generator/genr_ions.jl")
    export genr_ions

    include("./generator/combine.jl")
    export cat_data, sort_data!

    # Potential Files
    include("./Potential/Potential_Charmm36.jl")
    export Potential_Charmm36, Potential_Charmm36_Unit, potential_charmm36

    # converter Files
    include("./converter/useful_funcs.jl")

    include("converter/converter_vmd.jl")
    export converter_vmd

    # writeer Files
    include("./writer/write_data.jl")
    export write_data, write_info

    include("./writer/write_xyz.jl")
    export write_xyz

    include("./writer/write_pdb.jl")
    export write_pdb

end # module


text:
## ------------------------------------------------------------
# gen proj test
_TO_REMOVE = []
for fun_flag in [true, false]

    plot = ProjAssistant.ImgTools.Plots.plot
    RTAG = string(rand(UInt))
    TopMod = Symbol("Top", RTAG)
    Sub1Mod = Symbol("Sub1")
    Sub2Mod = Symbol("Sub2")
    Sub3Mod = Symbol("Sub3")
    TopDir = joinpath(tempdir(), string(TopMod))
    rm(TopDir; force = true, recursive = true)
    push!(_TO_REMOVE, TopDir)
    try
        # ---------------------------------------------------------------------
        # verbose
        ProjAssistant.global_conf(:VERBOSE, fun_flag)

        # ---------------------------------------------------------------------
        # prepare test "project"
        @info("Test project", TopMod, TopDir)
        @eval module $(TopMod)

            import ProjAssistant
            if $(fun_flag); ProjAssistant.gen_top_proj(@__MODULE__, $(TopDir))
                else; ProjAssistant.@gen_top_proj dir=$(TopDir)
            end

            function __init__()
                if $(fun_flag); ProjAssistant.create_proj_dirs(@__MODULE__)
                    else; ProjAssistant.@create_proj_dirs
                end
            end

            module $(Sub1Mod)

                import ProjAssistant
                if $(fun_flag); ProjAssistant.gen_sub_proj(@__MODULE__)
                    else; ProjAssistant.@gen_sub_proj
                end

                module $(Sub2Mod)

                    module $(Sub3Mod)

                        import Main.$(TopMod).$(Sub1Mod)
                        import ProjAssistant

                        if $(fun_flag); ProjAssistant.gen_sub_proj(@__MODULE__)
                            else; ProjAssistant.@gen_sub_proj(parent=$(Sub1Mod))
                        end

                        function __init__()
                            if $(fun_flag); ProjAssistant.create_proj_dirs(@__MODULE__)
                                else; ProjAssistant.@create_proj_dirs
                            end

                        end
                    end
                end # module $(Sub2Mod)

                function __init__()
                    if $(fun_flag); ProjAssistant.create_proj_dirs(@__MODULE__)
                        else; ProjAssistant.@create_proj_dirs
                    end

                end
            end # module $(Sub1Mod)
        end

        Top = getproperty(Main, TopMod)
        Sub1 = getproperty(Top, Sub1Mod)
        Sub2 = getproperty(Sub1, Sub2Mod)
        Sub3 = getproperty(Sub2, Sub3Mod)

        ## ------------------------------------------------------------
        # Test
        @testset "gen_projects" begin

            ## ------------------------------------------------------------
            # static funs
            for Proj in [Top, Sub1, Sub3]
                for funname in (
                        :projdir,
                        :devdir, :datdir, :srcdir, :plotsdir, :scriptsdir, :papersdir,
                        :procdir, :rawdir, :cachedir,
                    )
                    fun = getproperty(ProjAssistant, funname)

                    df = fun(Proj, "bla", (;A = 1), "jls")
                    @test ProjAssistant.isvalid_dfname(df)
                    @test df == fun(Proj, basename(df))
                end

            end

            @test procdir(Top) != datdir(Top)
            @test procdir(Sub1) == datdir(Sub1)
            @test procdir(Sub3) == datdir(Sub3)

            @test istop_proj(Top)
            @test !istop_proj(Sub1)
            @test !istop_proj(Sub3)

            ## ------------------------------------------------------------
            @info("dirtree")
            for dirfun in [
                    devdir, datdir, srcdir, plotsdir, scriptsdir, papersdir,
                    procdir, rawdir, cachedir
                ]
                @test dirfun(Sub1) == dirname(dirfun(Sub3))
            end

            ## ------------------------------------------------------------
            # save/load data
            for Mod in [Top, Sub1, Sub3]

                @info("save/load data", Mod)
                mkdir = true
                for (sfun, lfun) in [
                        (sprocdat, lprocdat),
                        (srawdat, lrawdat),
                        (sdat, ldat)
                    ]

                    dat0 = rand(10, 10)

                    for fargs in [
                                ("test_file", (;h = hash(dat0)), "jls"),
                                (["subdir"], "test_file", (;h = hash(dat0)), "jls"),
                            ]

                        cfile1 = sfun(Mod, dat0, fargs...; mkdir)
                        @test isfile(cfile1)
                        dat1 = lfun(Mod, fargs...)
                        @test all(dat0 .== dat1)
                        dat1 = lfun(Mod, cfile1)
                        @test all(dat0 .== dat1)
                        cfile2 = sfun(Mod, dat0, basename(cfile1); mkdir)
                        @test basename(cfile1) == basename(cfile2)
                        dat1 = lfun(Mod, cfile2)
                        @test all(dat0 .== dat1)
                    end

                    @test lfun(Mod, "NOT_A_FILE", hash(tempname()), ".jls") do
                        true
                    end
                end # for (sfun, lfun)
            end # for Mod

            ## ------------------------------------------------------------
            # save/load cache
            for Proj in [Top, Sub1, Sub3]

                @info("save/load cache", Proj)

                dat0 = rand(10, 10)

                cid = (:TEST, :CACHE, hash(dat0))
                cfile = scache(Proj, dat0, cid)
                @test isfile(cfile)
                dat1 = lcache(Proj, cid)
                @test all(dat0 .== dat1)

                cfile = scache(Proj, dat0)
                @test isfile(cfile)
                dat1 = lcache(Proj, cfile)
                @test all(dat0 .== dat1)

            end

            ## ------------------------------------------------------------
            # sglob/lglob
            for Proj in [Sub1, Sub3]

                @info("sglob/lglob", Proj)

                gfile = globfile(Proj, :test, :glob)
                rm(gfile; force = true)
                @assert !isfile(gfile)

                # funs
                dat0 = rand(10,10)
                sglob(Proj, dat0, :test, :glob)
                @test isfile(gfile)

                dat1 = lglob(Proj, :test, :glob)
                @test all(dat0 .== dat1)

                # macros
                rm(gfile; force = true)
                @assert !isfile(gfile)

                dat0 = rand(10,10)
                @sglob Proj test.glob=dat0
                @test isfile(gfile)

                @lglob Proj dat1=test.glob
                @test all(dat0 .== dat1)

            end

            ## ------------------------------------------------------------
            # save imgs
            for Proj in [Top, Sub1, Sub3]

                @info("save imgs", Proj)
                figfile0 = plotsdir(Proj, "test", (;A = 1), ".png")

                p = plot(rand(100))
                figfile1 = sfig(Proj, p, "test", (;A = 1), ".png")
                @test figfile0 == figfile1
                @test isfile(figfile1)

                ps = map((_) -> plot(rand(100)), 1:10)
                figfile1 = sfig(Proj, ps, "test", (;A = 1), ".png")
                @test figfile0 == figfile1
                @test isfile(figfile1)

                figfile0 = plotsdir(Proj, "test", (;A = 1), ".gif")
                figfile1 = sgif(Proj, ps, "test", (;A = 1), ".gif")
                @test figfile0 == figfile1
                @test isfile(figfile1)

            end

        end

    finally
        rm.(_TO_REMOVE; force = true, recursive = true)
    end
end


text:
module LearnTask

export producer

function producer()
    produce("start")
    for n = 1:4
        produce(2n)
    end
    produce("stop")
end

function consumer()
    print_with_color(:red, 1)
end

end


text:
@rule typeof(-)(:in1, Marginalisation) (m_out::Any, m_in2::Any, meta::Any) = begin
    @call_rule typeof(+)(:out, Marginalisation) (m_in1 = m_out, m_in2 = m_in2, meta = meta)
end


text:
module BumpStdlibs

import Dates
import Downloads
import GitHub
import HTTP
import JSON3
import TimeZones
import URIs

export StdlibInfo
export bump_stdlibs

include("types.jl")

include("assert.jl")
include("bump-stdlibs.jl")
include("checksums.jl")
include("filter-stdlib-list.jl")
include("git.jl")
include("github.jl")
include("inputs.jl")
include("stdlib-list.jl")
include("util.jl")

end # end module BumpStdlibs


text:
debug = false
let
    for use_autodiff in (false, true)
        for (name, prob) in Optim.UnconstrainedProblems.examples
            if prob.isdifferentiable && !(name in ("Polynomial", "Large Polynomial", "Himmelblau")) # it goes in a direction of ascent -> f_converged == true
                debug && @ show "** Name: $name"
                f_prob = prob.f
                iterations = name == "Powell" ? 2000 : 1000
                res = Optim.optimize(f_prob, prob.initial_x, MomentumGradientDescent(),
                                     OptimizationOptions(autodiff = use_autodiff,
                                                         iterations = iterations,
                                                         show_trace = debug))
                debug && @show res.minimum
                debug && @show prob.solutions
                @assert norm(res.minimum - prob.solutions, Inf) < 1e-2
            end
        end
    end
end


text:
include("../test.jl")
using GerryChain
using Test

@testset "Loading a graph" begin
    @testset "2x2 graph" begin
        graph = load_graph_from_edge_list("data/2x2_grid.dat", "data/2x2_grid_weight.dat")
        @test graph.num_nodes == 4
        @test graph.num_edges == 4
        @test graph.populations == [1, 2, 4, 8]
        @test graph.total_pop == 15
        @test graph.adj_matrix == [0 1 2 0; 1 0 0 3; 2 0 0 4; 0 3 4 0]
        @test graph.neighbors == [[2, 3], [1, 4], [1, 4], [2, 3]]
    end
    @testset "2x2 triangular lattice" begin
        graph = load_graph_from_edge_list("data/2x2_triangular_lattice.dat", "data/2x2_triangular_lattice_weight.dat")
        @test graph.num_edges == 9
        @test graph.num_nodes == 6
        @test graph.populations == [1, 2, 4, 8, 10, 12]
        @test graph.total_pop == 37
        @test graph.neighbors == [[2, 3], [1, 3, 4], [1, 2, 4, 5, 6], [2, 3, 6], [3,6], [3, 4, 5]]
    end
    @testset "6x6 grid" begin
        graph = load_graph_from_edge_list("data/6x6_grid.dat", "data/6x6_grid_weight.dat")
        @test graph.num_edges == 60
        @test graph.num_nodes == 36
        @test graph.populations == collect(1:36)
        @test graph.total_pop == sum(collect(1:36))
    end
    @testset "Disconnected graph" begin
        @test_throws ArgumentError load_graph_from_edge_list("data/6x6_disconnected_grid.dat", "data/6x6_grid_weight.dat")
    end
end

max_stars_repo_name: mggg/benchmarks

max_issues_repo_name: mggg/benchmarks

max_forks_repo_name: mggg/benchmarks


text:
module TestStrayCopse
using Test
using StableRNGs
using StrayCopses

@testset "basic stray copse" begin
    X = collect(reshape([1., 2., 3., 4., 1., 2., 2., 4.], :, 2))
    y = [1, 1, 2, 2]
    rng = StableRNG(2020)

    res = fit(StrayCopse(), X, y, rng, n_features_per_node = 2, n_trees = 2)
    @test length(res.sc) == 2
end

end # module


text:
using ProgressMeter

include("./plot_ic_reward.jl")
include("./plot_ic_goal_visitation.jl")

# Experiment parameters
log_interval = 100
num_steps = 40000

# Getting baseline TB
folder = "M:/globus/InterestTB/OneDTMaze_RR_dpi/"
ic = ItemCollection(folder)

search_ic = search(ic, Dict("demon_learner" => "SR", "demon_opt" => "Auto"))
tb_ic = get_best(search_ic, ["demon_eta"], :oned_tmaze_dpi_error)

configs = [
    # Dict(
    #     :search => Dict("demon_learner" => "Q", "demon_update" => "ETB"),
    #     :label => "Q-ETB"
    # )
    # Dict(
    #     :search => Dict("demon_learner" => "Q", "demon_update" => "EmphESARSA"),
    #     :label => "Q-EmphESARSA"
    # )
    # Dict(
    #     :search => Dict("demon_learner" => "Q", "demon_update" => "InterestTB"),
    #     :label => "Q-InterestTB"
    # )
    Dict(
        :search => Dict("demon_learner" => "SR", "demon_update" => "ETB"),
        :label => "SR-ETB"
    )
    Dict(
        :search => Dict("demon_learner" => "SR", "demon_update" => "EmphESARSA"),
        :label => "SR-EmphESARSA"
    )
    Dict(
        :search => Dict("demon_learner" => "SR", "demon_update" => "InterestTB"),
        :label => "SR-InterestTB"
    )
]

folder = "M:/globus/InterestTB/Experiment3/"
ic = ItemCollection(folder)

# println(diff(ic))

# asdfsadf

# Plotting sum reward for all GVF
p = plot()
@showprogress for config in configs
    local search_ic = search(ic, config[:search])
    local best_ic = get_best(search_ic, ["alpha_init", "eta"], :oned_tmaze_dpi_error)
    plot_ic_reward!(p, best_ic, :oned_tmaze_dpi_error, log_interval; smooth_step=10, label=config[:label])
end

ylims!(p, (0, 1.5))
plot_ic_reward!(p, tb_ic, :oned_tmaze_dpi_error, log_interval; smooth_step=10, label="SR TB")

savefig("plotting/chunlok/generated_plots/Experiment3_rewards.svg")

# Plotting reward per GVF
# best_ics = []
# for config in configs
#     local search_ic = search(ic, config[:search])
#     local best_ic = get_best(search_ic, ["alpha_init", "eta"], :oned_tmaze_dpi_error)
#     push!(best_ics, best_ic)
# end

# for gvf_i in 1:4
#     p = plot()
#     for (config_i, config) in enumerate(configs)
#         plot_ic_reward_per_gvf!(p, best_ics[config_i], :oned_tmaze_dpi_error, log_interval, num_steps, gvf_i; smooth_step=1, label=config[:label])
#     end
#     plot!(title="GVF $(gvf_i)")
#     savefig("plotting/chunlok/generated_plots/Experiment4_rewards_gvf_$(gvf_i)")
# end

# Plotting sum reward for all GVF

# for config in configs
#     p = plot()
#     local search_ic = search(ic, config[:search])
#     local best_ic = get_best(search_ic, ["alpha_init", "eta"], :oned_tmaze_dpi_error)
#     plot_ic_goal_visitation!(p, best_ic)
#     plot!(title=config[:label])
#     savefig("plotting/chunlok/generated_plots/Experiment4_goal_visitation_$(config[:label])")
# end

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
using Documenter, SpectralLibraries

makedocs(modules = [SpectralLibraries],
    format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
    pages = [
        "Home" => "index.md",
        "Libraries" => [
            "PHOENIX ACES" => "phoenix.md",
        ],
    ],
    strict = true,
    repo = "https://github.com/mileslucas/SpectralLibraries.jl/blob/{commit}{path}#L{line}",
    sitename = "SpectralLibraries.jl",
    authors = "Miles Lucas <mdlucas@hawaii.edu>", )

deploydocs(repo = "github.com/mileslucas/SpectralLibraries.jl")


text:
# Raytracer.jl
# Raytracing for the generation of photorealistic images in Julia
# Copyright (c) 2021 Samuele Colombo, Paolo Galli

# Renderer implementations

"""
    Renderer

Abstract type for functors that map [`Ray`](@ref) to `RGB{Float32}`.

Each subtype of this type must be a callable like `(r::Renderer)(ray::Ray)` and must return a `RGB{Float32}`.
Each subtype of this type should have a member of type [`World`](@ref) to check for intersections of the given [`Ray`](@ref).

See also: [`OnOffRenderer`](@ref), [`FlatRenderer`](@ref), [`PathTracer`](@ref), [`PointLightRenderer`](@ref)
"""
abstract type Renderer end

################
# OnOffRenderer

"""
    OnOffRenderer <: Renderer

A basic bichrome [`Renderer`](@ref) that checks whether a [`Ray`](@ref) has collided or not.

This renderer returns its field `off_color` when the given `Ray` is `nothing`, else it returns its field `on_color`.

# Fields

- `world::World`: the [`World`](@ref) to render.
- `on_color::RGB{Float32}`: color if the ray collide.
- `off_color::RGB{Float32}`: color if the ray do not collide.
"""
struct OnOffRenderer <: Renderer
    world::World
    on_color::RGB{Float32}
    off_color::RGB{Float32}
end

@doc """
    OnOffRenderer(world::World, on_color::RGB{Float32}, off_color::RGB{Float32})

Constructor for a [`OnOffRenderer`](@ref) instance.
""" OnOffRenderer(::World, ::RGB{Float32}, ::RGB{Float32})

"""
    OnOffRenderer(world::World
                  ; on_color::RGB{Float32} = WHITE,
                    off_color::RGB{Float32} = BLACK)

Constructor for a [`OnOffRenderer`](@ref) instance.

If no color is specified, it will default on [`WHITE`](@ref) and [`BLACK`](@ref).
"""
OnOffRenderer(world::World; on_color::RGB{Float32} = WHITE, off_color::RGB{Float32} = BLACK) = OnOffRenderer(world, on_color, off_color)

"""
    (oor::OnOffRenderer)(ray::Ray)

Render a [`Ray`](@ref) and return a `RBG{Float32}`.
"""
function (oor::OnOffRenderer)(ray::Ray)
    !any(quick_ray_intersection.(Ref(ray), oor.world)) ? oor.off_color : oor.on_color
end

###############
# FlatRenderer

"""
    FlatRenderer <: Renderer

A basic [`Renderer`](@ref) that returns the color of the [`Shape`](@ref) first hit by a given [`Ray`](@ref).

This renderer returns the color stored in the `material` field of the [`Shape`](@ref) first hit by the given [`Ray`](@ref) at the hit point.
To this renderer there is no difference between radiated light and reflected color. There are no shades, diffusions or reflections.
If there are no hits this renderer returns the value of its field `background_color`.

# Fields

- `world::World`: the [`World`](@ref) to render.
- `background_color::RGB{Float32}`: color if the ray do not collide.
"""
struct FlatRenderer <: Renderer
    world::World
    background_color::RGB{Float32}
end

@doc """
    FlatRenderer(world::World, background_color::RGB{Float32})

Constructor for a [`FlatRenderer`](@ref) instance.
""" FlatRenderer(::World, ::RGB{Float32})

"""
    FlatRenderer(world::World; background_color::RGB{Float32} = BLACK)

Constructor for a [`FlatRenderer`](@ref) instance.

If no color is specified, it will default on [`BLACK`](@ref).
"""
FlatRenderer(world::World; background_color::RGB{Float32} = BLACK) = FlatRenderer(world, background_color)

"""
    (oor::FlatRenderer)(ray::Ray)

Render a [`Ray`](@ref) and return a `RBG{Float32}`.
"""
function (fr::FlatRenderer)(ray::Ray)
    hit = ray_intersection(ray, fr.world)
    isnothing(hit) ? fr.background_color : hit.material.emitted_radiance(hit.surface_point) + hit.material.brdf.pigment(hit.surface_point)
end

#############
# PathTracer

"""
    PathTracer <: Renderer

A path-tracing [`Renderer`](@ref) that considers the optical path of a [`Ray`](@ref) from the observer to a light source.

# Fields

- `world::World`: the [`World`](@ref) to render.
- `background_color::RGB{Float32}`: color if the ray do not collide.
- `rng::PCG`: a [`PCG`](@ref) random number generator to appropriately scatter rays.
- `n::Int`: how many scattered rays should be generated for the mc integration.
- `max_depth::Int`: the maximum number of scatters a ray should be subjected to before stopping.
- `roulette_depth::Int`: the depth at which the russian roulette algorithm should start (if > 'max_depth` then it will never start).
"""
struct PathTracer <: Renderer
    world::World
    background_color::RGB{Float32}
    rng::PCG
    n::Int
    max_depth::Int
    roulette_depth::Int
end

@doc """
    PathTracer(world::World, background_color::RGB{Float32}, rng::PCG, n::Int, max_depth::Int, roulette_depth::Int)

Constructor for a [`PathTracer`](@ref) instance.
""" PathTracer(::World, ::RGB{Float32}, ::PCG, ::Int, ::Int, ::Int)

"""
    PathTracer(world::World
               ; background_color::RGB{Float32} = BLACK,
                 rng::PCG = PCG(),
                 n::Int = 10,
                 max_depth::Int = 2,
                 roulette_depth::Int = 3)

Constructor for a [`PathTracer`](@ref) instance.
"""
function PathTracer(world::World; background_color::RGB{Float32} = BLACK, rng::PCG = PCG(), n::Int = 10, max_depth::Int = 2, roulette_depth::Int = 3)
    PathTracer(world, background_color, rng, n, max_depth, roulette_depth)
end

"""
    (oor::PathTracer)(ray::Ray)

Render a [`Ray`](@ref) and return a `RBG{Float32}`.
"""
function (pt::PathTracer)(ray::Ray)
    ray.depth > pt.max_depth && return BLACK

    hit_record = ray_intersection(ray, pt.world)
    isnothing(hit_record) && return pt.background_color

    hit_material = hit_record.material
    hit_color = hit_material.brdf.pigment(hit_record.surface_point)
    emitted_radiance = hit_material.emitted_radiance(hit_record.surface_point)

    hit_color_lum = max(hit_color...)

    # Russian roulette
    if ray.depth >= pt.roulette_depth
        # q = max(0.05f0, 1f0 - hit_color_lum)
        # if rand(pt.rng, Float32) > q
        #     # Keep the recursion going, but compensate for other potentially discarded rays
        #     hit_color *= 1f0 / (1f0 - q)
        if rand(pt.rng, Float32) > hit_color_lum
            hit_color *= 1f0 / (1f0 - hit_color_lum)
        else
            # Terminate prematurely
            return emitted_radiance
        end
    end

    # Monte Carlo integration
    cum_radiance = BLACK
    # Only do costly recursions if it's worth it
    if hit_color_lum > 0f0
        for _ ∈ 1:pt.n
            new_ray = scatter_ray(
                hit_material.brdf,
                pt.rng,
                hit_record.ray.dir,
                hit_record.world_point,
                hit_record.normal,
                ray.depth + 1
            )
            # Recursive call
            new_radiance = pt(new_ray)
            cum_radiance += hit_color * new_radiance
        end
    end

    emitted_radiance + cum_radiance * (1f0 / pt.n)
end

#####################
# PointLightRenderer

"""
    PointLightRenderer <: Renderer

Point-light tracing [`Renderer`](@ref). This renderer is similar to what POV-Ray provides by default.

# Fields

- `world::World`: the [`World`](@ref) to render.
- `lights::Lights`: a [`Lights`](@ref) instance that contain a list of lights.
- `background_color::RGB{Float32}`: color if the ray do not collide.
- `ambient_color::RGB{Float32}`: the ambient color of the scene.
"""
struct PointLightRenderer <: Renderer
    world::World
    lights::Lights
    background_color::RGB{Float32}
    ambient_color::RGB{Float32}
end

@doc """
    PointLightRenderer(world::World, lights::Lights, background_color::RGB{Float32}, ambient_color::RGB{Float32})

Constructor for a [`PointLightRenderer`](@ref) instance.
""" PointLightRenderer(::World, ::Lights, ::RGB{Float32}, ::RGB{Float32})

"""
    PointLightRenderer(world::World, lights::Lights
                       ; background_color::RGB{Float32} = BLACK,
                         ambient_color::RGB{Float32} = WHITE * 1f-3)

Constructor for a [`PointLightRenderer`](@ref) instance.
"""
function PointLightRenderer(world::World, lights::Lights; background_color::RGB{Float32} = BLACK, ambient_color::RGB{Float32} = WHITE * 1f-3)
    PointLightRenderer(world, lights, background_color, ambient_color)
end

"""
    (oor::PointLightRenderer)(ray::Ray)

Render a [`Ray`](@ref) and return a `RBG{Float32}`.
"""
function (plr::PointLightRenderer)(ray::Ray)
    hit_record = ray_intersection(ray, plr.world)
    isnothing(hit_record) && return plr.background_color

    hit_material = hit_record.material

    result_color = plr.ambient_color
    for cur_light ∈ plr.lights
        is_point_visible(plr.world, cur_light.position, hit_record.world_point) || continue

        distance_vec = hit_record.world_point - cur_light.position
        distance = norm(distance_vec)
        in_dir = distance_vec * (1f0 / distance)
        cos_theta = max(0f0, normalized_dot(-distance_vec, hit_record.normal))

        distance_factor = cur_light.linear_radius > 0 ? (cur_light.linear_radius / distance)^2 : 1f0

        emitted_color = hit_material.emitted_radiance(hit_record.surface_point)
        brdf_color = at(hit_material.brdf, hit_record.normal, in_dir, -ray.dir, hit_record.surface_point)

        result_color += (emitted_color + brdf_color) * cur_light.color * cos_theta * distance_factor
    end

    return result_color
end


text:
# train.jl
# This script takes the saved feature data and will train a Random Forest to
# classify bird syllables in spectrograms
# author: Timothy Evans
# date: 02/08/14

push!(LOAD_PATH, pwd())
using DecisionTree, JLD
importall config, dataprocessing

#load paths from config.jl
specpath,labelledpath,trainpath,modelpath = loadpaths()
print("Loading in training data...\n")
#import data
data = readcsv(string(modelpath,datafile));
(m,n) = size(data);

#use 70% of data in training set
ind = randperm(m);
ntrain = Int(ceil(0.7*m));
trainind = ind[1:ntrain];
testind = ind[ntrain+1:end];

xtrain = data[trainind,1:end-1];
ytrain = toint(data[trainind,end]);
xtest = data[testind,1:end-1];
ytest = toint(data[testind,end]);

print("Training Random Forest...\n")
# build the random forest using config paramaters.
model = build_forest(ytrain, xtrain, Int(ceil(sqrt(n))), ntrees, ratio)

print("Testing model accuracy...\n")
#pixelwise accuracy:
testpred = toint(apply_forest(model,xtest));
# accuracy = 1-sum(abs(ytest-testpred))/length(ytest);

confm = confusmat(ytest,testpred);
acc = trace(confm)./length(ytest);
C = confm./sum(confm,2);

@printf "\nConfusion Matrix:\n\n\t[%.3f , %.3f\n\t%.3f , %.3f]\n\nConfusion Accuracies:\n\n\t[%d , %d\n\t %d , %d]\n\nTest Accuracy: %.3f\n" C[1,1] C[1,2] C[2,1] C[2,2] confm[1,1] confm[1,2] confm[2,1] confm[2,2] acc

print("Writing Random Forest model to disk (if ntrees is big this may take some time)...\n")
# Save the random forest structure for future use
jldopen(string(modelpath,"model.jld"),"w") do file
  write(file,"model",model)
end

max_stars_repo_name: kokakodetect/kokako-detection

max_issues_repo_name: kokakodetect/kokako-detection

max_forks_repo_name: kokakodetect/kokako-detection


text:
f = x -> x^2 + 2x - 1

# ---------------
using Plots
plot(f, 0, 10, xlabel="x", ylabel="y")

function myplot(x, y; style=:solid, width=1, color=:black)
    plot(x, y, style=style, width=width, color=color)
end

x = collect(0:10)
y = f.(x)
myplot(x,y)
myplot(x, y, style=:dot, color=:blue)

# ---------- let Blocks --------

x, y, z = -1, -1, -1

@show x y z;

let x = 1, z
    @show x y z;
end

@show x y z;

# ---------- Tasks (aka Coroutines) --------

#=
    A task is a control structure that can pass control cooperatively
    without returning.  A task is implemented as a function having as its first
    argument a Channel object.  A Channel is used to pass values from the
    function to a calee.
=#

"""
    fib(c::Channel)

    Generates a Fibonacci sequence.  Suspended after each call to put!.
"""
function fib(c::Channel)
    a = 0
    b = 1
    put!(c, a)        # suspends, waits for channel to clear
    while true
        put!(c, b)    #suspends
        (a, b) = (b, a+b)
    end
end

fibgen = Channel(fib);
take!(fibgen)
take!(fibgen)
take!(fibgen)
take!(fibgen)
take!(fibgen)
take!(fibgen)
take!(fibgen)
take!(fibgen)

for val in Channel(fib)
    print(val, " ")
    val > 30 && break
end

# ------- Types --------

#=
    A concrete type consisting of just bits is a primitive type.
    Julia allows you to declare your own primitive types.

    The standard types are defined as:

    primitive type Float64 <: AbstractFloat 64 end
    primitive type  Bool <: Integer 8 end
    primitive type Char <: AbstractChar 32 end
    primitive type Int64 <: Signed 64 end
=#

# Create primitive type Byte and constructor
primitive type Byte 8 end
Byte(val::UInt8) = reinterpret(Byte,val)

b = Byte(0x1d);
typeof(b)
@show b;

# c = Byte(0xfff) # ERROR: MethodError: no method matching Byte(::UInt16)

# Parametric types

struct Point{T<:Real}
    x::T
    y::T
    Point{T}(x,y) where {T<:Real} = new(x,y)   # same as default inner constructor
end

"""
    default outer constructor.
    x, y have same type
"""
Point(x::T, y::T) where {T<:Real} = Point{T}(x,y)

p = Point(0.0, 0.0)
q = Point(3, 4)

#=
struct Point{T<:Int64}
    x::T
    y::T
end     #invalid redefinition of constant Point
=#

mp = IntPoint(3, 4)
# Type Unions

#  Efficient code can be generated when a type union has a small number of types
IntOrString = Union{Int64, String}
150 :: IntOrString
"Hello World" :: IntOrString
# 150 :: String      # TypeError

# -------- Methods ----------

# -------- Parametric Methods -------
isintpoint(p::Point{T}) where {T} = (T === Int64)
q
p
isintpoint(q)
isintpoint(p)

# -------- Functors ---------

struct Polynomial{R}
    coeff::Vector{R}
end

# Define a functor on instances of Polynomial
function (p::Polynomial)(x)
    val = p.coeff[end]
    for coeff in p.coeff[end-1:-1:1]
        val = val * x + coeff
    end
    val
end

# p(x) = 100x² + 10x + 1
p = Polynomial([1,10,100])
p(3)

#= -------- Conversion and Promotion ---------- =#

# Conversion
x = 12
typeof(x)

y = convert(UInt8, x)

typeof(y)
y
typeof(x)
x

Point{Int64} isa Type{Point{Int64}}
#Point{Int64} isa Point{Int64}
#typeof(Point{Int64})

Base.convert(::Type{Point{T}}, x::Array{T, 1}) where {T<:Real} = Point(x...)
c = convert(Point{Int64}, [1,2]);
c

# Promotion

promote(1, 2.5, 3)

# specify promotion behavior using promote_rule

# promote_rule(::Type{Float64}, ::Type{Int32}) = Float64


text:
function foo(x, y)
    x + y
end

foobar(x::Int) = 2*x

@test Matte.argument_names(foo) == [:x, :y]
@test Matte.argument_names(foobar) == [:x]


text:
@testset "CategoricalLikelihood" begin
    lik = CategoricalLikelihood()
    IN_DIM = 3
    OUT_DIM = 4
    N = 10
    X = MOInput([rand(IN_DIM) for _ in 1:N], OUT_DIM)
    test_interface(lik, IndependentMOKernel(SqExponentialKernel()), X)
end


text:
using LazySets: _leq, _geq, isapproxzero, _isapprox, ABSZTOL

# approximate <= and
@test _leq(2e-15, 1e-15) && _leq(1e-15, 2e-15)
@test _leq(1//100, 1//99) && _leq(1//100, 0.099)

# approximate <= with ztol
@test !_leq(2e-15, 1e-15, ztol=1e-15) && _leq(1e-15, 2e-15, ztol=1e-15)

# approximate >=
@test _geq(2e-15, 1e-15) && _geq(1e-15, 2e-15)

# approximate >= with ztol
@test _geq(2e-15, 1e-15, ztol=1e-15) && !_geq(1e-15, 2e-15, ztol=1e-15)

# default absolute zero tolerance for rational
ABSZTOL(eltype(1/100)) == zero(Rational{Int})

# default absolute zero tolerance for FP
ABSZTOL(eltype(0.01)) == sqrt(eps(Float64))

# approximately zero tests
@test isapproxzero(0//1)
@test isapproxzero(1e-8) && !isapproxzero(1e-8, ztol=1e-10)

# approximate equality
@test !_isapprox(1//1, 1//1 + 1//1000000000000)
@test _isapprox(2e-15, 1e-15)
@test !_isapprox(2e-6, 1e-6)


text:
module MCMC

using BangBang
using Distributions
using ProgressBars
using StatsFuns
import Random
using Bijectors
import LinearAlgebra

abstract type Model end
abstract type InferenceAlgorithm end
const OneOrMoreSymbols = Union{Symbol, NTuple{N, Symbol} where N}

abstract type Metropolis end
# Implement _update(::Random.AbstractRNG, ::Metropolis, curr::Real, logprob::Function)
function _update(met::Metropolis, curr::T, logprob::Function) where T
  return _update(Random.GLOBAL_RNG, met, curr, logprob)
end
function update(rng::Random.AbstractRNG, met::Metropolis, curr::T, logprob::Function) where T
  return _update(rng, met, curr, logprob)[1]
end
function update(met::Metropolis, curr::T, logprob::Function) where T
  return update(Random.GLOBAL_RNG, met, curr, logprob)
end

# Utilities.
include("misc.jl")

# Metropolis.
include("static_rwm.jl")
include("rwm.jl")
include("arwm.jl")
include("mvarwm.jl")

# General stuff.
include("inference/gibbs.jl")
include("mcmc-functions.jl")
include("metrics.jl")

export mcmc, wsample_logprob, Gibbs, Conditional
export dic, deviance, log_bayes_factor
export RWM, mvarwm
export quantiles, logmeanexp, hellinger

end # module


text:
"""
        _cols_by_rows(rows_index,cols_index)

Returns a vector of rows where each row contains
a vector of its column indices.
"""
function _cols_by_rows(rows_index,cols_index)
    nrows = maximum(rows_index)
    cols_by_rows = [eltype(rows_index)[] for _ in 1:nrows]
    for (i,j) in zip(rows_index,cols_index)
        push!(cols_by_rows[i],j)
    end
    return cols_by_rows
end

"""
        _rows_by_cols(rows_index,cols_index)

Returns a vector of columns where each column contains
a vector of its row indices.
"""
function _rows_by_cols(rows_index,cols_index)
    return _cols_by_rows(cols_index,rows_index)
end

"""
        matrix2graph(sparse_matrix, [partition_by_rows::Bool=true])

A utility function to generate a graph from input
sparse matrix, columns are represented with vertices
and 2 vertices are connected with an edge only if
the two columns are mutually orthogonal.

Note that the sparsity pattern is defined by structural nonzeroes, ie includes
explicitly stored zeros.
"""
function matrix2graph(sparse_matrix::SparseMatrixCSC{<:Number, Int}, partition_by_rows::Bool=true)

    (rows_index, cols_index, _) = findnz(sparse_matrix)

    ncols = size(sparse_matrix, 2)
    nrows = size(sparse_matrix, 1)

    num_vtx = partition_by_rows ? nrows : ncols

    inner = SimpleGraph(num_vtx)

    if partition_by_rows
        rows_by_cols = _rows_by_cols(rows_index,cols_index)
        @inbounds for (cur_row,cur_col) in zip(rows_index,cols_index)
            if !isempty(rows_by_cols[cur_col])
                for next_row in rows_by_cols[cur_col]
                    if next_row < cur_row
                        add_edge!(inner, cur_row, next_row)
                    end
                end
            end
        end
    else
        cols_by_rows = _cols_by_rows(rows_index,cols_index)
        @inbounds for (cur_row,cur_col) in zip(rows_index,cols_index)
            if !isempty(cols_by_rows[cur_row])
                for next_col in cols_by_rows[cur_row]
                    if next_col < cur_col
                        add_edge!(inner, cur_col, next_col)
                    end
                end
            end
        end
    end
    return VSafeGraph(inner)
end


text:
println("	PROJECT")
project(C="15/15", T="85/40", G="1/110", L="-20/60");	# Fails in GMT5
project(nothing, C="15/15", T="85/40", G="1/110", L="-20/60");	# bit of cheating

println("	SAMPLE1D")
d = [-5 74; 38 68; 42 73; 43 76; 44 73];
sample1d(d, T="2c", A=:r);

println("	SPECTRUM1D")
D = gmt("gmtmath -T0/10239/1 T 10240 DIV 360 MUL 400 MUL COSD");
spectrum1d(D, S=256, W=true, par=(GMT_FFT=:brenner,), N=true, i=1);

println("	SPHTRIANGULATE")
D = sphtriangulate(rand(10,3), V=:q);		# One dataset per triangle????

println("	SPHINTERPOLATE")
sphinterpolate(rand(10,3), I=0.1, R="0/1/0/1");

println("	SPHDISTANCE")
# SPHDISTANCE  (would fail with: Could not obtain node-information from the segment headers)
G = sphdistance(R="0/10/0/10", I=0.1, Q=D, L=:k, Vd=dbg2);	# But works with data from sph_3.sh test
@test sphdistance(nothing, R="0/10/0/10", I=0.1, Q="D", L=:k, Vd=dbg2) == "sphdistance  -I0.1 -R0/10/0/10 -Lk -QD"

println("	SURFACE")
G = surface(rand(100,3) * 150, R="0/150/0/150", I=1, Ll=-100, upper=100, V=:q);
@assert(size(G.z) == (151, 151))

println("	SPLITXYZ")
# SPLITXYZ (fails?)
splitxyz([-14.0708 35.0730 0; -13.7546 35.5223 0; -13.7546 35.5223 0; -13.2886 35.7720 0; -13.2886 35.7720 0; -12.9391 36.3711 0], C=45, A="45/15", f="g", Vd=dbg2)

println("	TRIANGULATE")
G = triangulate(rand(100,3) * 150, R="0/150/0/150", I=1, grid=[]);
triangulate(rand(5,3), R="0/150/0/150", voronoi=:pol, Vd=dbg2);

println("	NEARNEIGHBOR")
G = nearneighbor(rand(100,3) * 150, R="0/150/0/150", I=1, N=4, S=10, r=true);

println("	XYZ2GRD")
D=grd2xyz(G); # Use G of previous test
xyz2grd(D, R="0/150/0/150", I=1, r=true);
xyz2grd(D, xlim=(0,150), ylim=(0,150), I=1, r=true);

println("	TREND1D")
D = gmt("gmtmath -T10/110/1 T 50 DIV 2 POW 2 MUL T 60 DIV ADD 4 ADD 0 0.25 NRAND ADD T 25 DIV 2 MUL PI MUL COS 2 MUL 2 ADD ADD");
trend1d(D, N="p2,F1+o0+l25", F=:xm);

println("	TREND2D")
trend2d(D, F=:xyr, N=3);


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

using Markdown
using InteractiveUtils

# ╔═╡ 0c9c514e-598c-418f-90e5-4a74ffb5bdc0
begin
    import Pkg
    # activate a temporary environment
    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
        Pkg.PackageSpec(name="Distributions"),
        Pkg.PackageSpec(name="DataFrames", version="1"),
        Pkg.PackageSpec(name="StanSample", version="5.2"),
        Pkg.PackageSpec(name="StatisticalRethinking", version="4.4"),
    ])
    using Plots, PlutoUI
	using Distributions, DataFrames
	using StanSample, StatisticalRethinking
end

# ╔═╡ f2224d36-7972-4e21-badc-3a401bf6f0bb
pwd()

# ╔═╡ 7d6d3864-3bff-414d-b077-ab59673f6d25
@__DIR__

# ╔═╡ eb07b03c-4617-4e07-b68a-476b903460e0
Pkg.pkg"status"

# ╔═╡ ed6fe9e3-df1d-4c60-b3f0-9cab2c81eb78
stan0_0 = "
data {
	int N;
	vector[N] D;
}
parameters {
	real a;
	real sigma;
}
model {
	a ~ normal(0, 2);
	sigma ~ exponential(1);
	D ~ normal(a, sigma);
}";

# ╔═╡ a022fd36-61aa-4e85-92e0-0ab55339468b
data = (N=1000, D=rand(Normal(0,1), 1000));

# ╔═╡ 3104104b-b90b-4279-a79e-a94502869172
begin
	m0_0s = SampleModel("m0_0s", stan0_0)
	rc0_0s = stan_sample(m0_0s; data)
end;

# ╔═╡ 3dab9dbd-2955-4a11-816f-d986ce0529bd
if success(rc0_0s)
	df0_0s = read_samples(m0_0s, :dataframe)
	PRECIS(df0_0s)
end

# ╔═╡ Cell order:
# ╠═f2224d36-7972-4e21-badc-3a401bf6f0bb
# ╠═7d6d3864-3bff-414d-b077-ab59673f6d25
# ╠═0c9c514e-598c-418f-90e5-4a74ffb5bdc0
# ╠═eb07b03c-4617-4e07-b68a-476b903460e0
# ╠═ed6fe9e3-df1d-4c60-b3f0-9cab2c81eb78
# ╠═a022fd36-61aa-4e85-92e0-0ab55339468b
# ╠═3104104b-b90b-4279-a79e-a94502869172
# ╠═3dab9dbd-2955-4a11-816f-d986ce0529bd


text:
# Utilities to deal with fetching/serving actual Pkg resources

const REGISTRIES = Dict(
    "23338594-aafe-5451-b93e-139f81909106" =>
        "https://github.com/JuliaRegistries/General",
)
const STORAGE_SERVERS = [
    "http://127.0.0.1:8080",
]
sort!(STORAGE_SERVERS)

const uuid_re = raw"[0-9a-z]{8}-[0-9a-z]{4}-[0-9a-z]{4}-[0-9a-z]{4}-[0-9a-z]{12}(?-i)"
const hash_re = raw"[0-9a-f]{40}"
const meta_re     = Regex("^/meta\$")
const registry_re = Regex("^/registry/($uuid_re)/($hash_re)\$")
const resource_re = Regex("""
    ^/registries\$
  | ^/registry/$uuid_re/$hash_re\$
  | ^/package/$uuid_re/$hash_re\$
  | ^/artifact/$hash_re\$
""", "x")
const hash_part_re = Regex("/($hash_re)\$")

function get_registries(server::String)
    regs = Dict{String,String}()
    response = HTTP.get("$server/registries")
    for line in eachline(IOBuffer(response.body))
        m = match(registry_re, line)
        if m !== nothing
            uuid, hash = m.captures
            uuid in keys(REGISTRIES) || continue
            regs[uuid] = hash
        else
            @error "invalid response" server=server resource="/registries" line=line
        end
    end
    return regs
end

"""
    write_atomic(f::Function, path::String)

Performs an atomic filesystem write by writing out to a file on the same
filesystem as the given `path`, then `move()`'ing the file to its eventual
destination.  Requires write access to the file and the containing folder.
Currently stages changes at "<path>.tmp.<randstring>".  If the return value
of `f()` is `false` or an exception is raised, the write will be aborted.
"""
function write_atomic(f::Function, path::String)
    temp_file = path * ".tmp." * randstring()
    try
        retval = open(temp_file, "w") do io
            f(temp_file, io)
        end
        if retval !== false
            mv(temp_file, path; force=true)
        end
    catch e
        rm(temp_file; force=true)
        rethrow(e)
    end
end

function url_exists(url::String)
    response = HTTP.request("HEAD", url, status_exception = false)
    response.status == 200
end

function verify_registry_hash(uuid::String, hash::String)
    isfile(joinpath("cache", "registry", uuid, hash)) && return true
    url = Pkg.Operations.get_archive_url_for_version(REGISTRIES[uuid], hash)
    return url === nothing || url_exists(url)
end

# current registry hashes and servers that know about them
const REGISTRY_HASHES = Dict{String,String}()
const REGISTRY_SERVERS = Dict{String,Vector{String}}()

function update_registries()
    # collect current registry hashes from servers
    regs = Dict(uuid => Dict{String,Vector{String}}() for uuid in keys(REGISTRIES))
    servers = Dict(uuid => Vector{String}() for uuid in keys(REGISTRIES))
    for server in STORAGE_SERVERS
        for (uuid, hash) in get_registries(server)
            push!(get!(regs[uuid], hash, String[]), server)
            push!(servers[uuid], server)
        end
    end
    # for each hash check what other servers know about it
    changed = false
    for (uuid, hash_info) in regs
        isempty(hash_info) && continue # keep serving what we're serving
        for (hash, hash_servers) in hash_info
            for server in servers[uuid]
                server in hash_servers && continue
                url_exists("$server/registry/$uuid/$hash") || continue
                push!(hash_servers, server)
            end
        end
        hashes = sort!(collect(keys(hash_info)))
        sort!(hashes, by = hash -> length(hash_info[hash]))
        for hash in hashes
            # first check if the origin repo knows about this hash
            verify_registry_hash(uuid, hash) || continue
            # try hashes known to fewest servers first, ergo newest
            servers = sort!(hash_info[hash])
            fetch("/registry/$uuid/$hash", servers=servers) !== nothing || continue
            if get(REGISTRY_HASHES, uuid, nothing) != hash
                @info "new current registry hash" uuid=uuid hash=hash servers=servers
                changed = true
            end
            REGISTRY_HASHES[uuid] = hash
            REGISTRY_SERVERS[uuid] = servers
            break # we've got a new registry hash to server
        end
    end
    # write new registry info to file
    if changed
        write_atomic(joinpath("cache", "registries")) do temp_file, io
            for uuid in sort!(collect(keys(REGISTRIES)))
                hash = REGISTRY_HASHES[uuid]
                println(io, "/registry/$uuid/$hash")
            end
            return true
        end
    end
    return changed
end

const fetch_locks = 1024
const FETCH_SEED = rand(UInt)
const FETCH_LOCKS = [ReentrantLock() for _ = 1:fetch_locks]
const FETCH_FAILS = [Set{String}() for _ = 1:fetch_locks]
const FETCH_DICTS = [Dict{String,Event}() for _ = 1:fetch_locks]

function fetch(resource::String; servers=STORAGE_SERVERS)
    path = "cache" * resource
    isfile(path) && return path
    isempty(servers) && throw(@error "fetch called with no servers" resource=resource)
    # make sure only one thread fetches path
    i = (hash(path, FETCH_SEED) % fetch_locks) + 1
    fetch_lock = FETCH_LOCKS[i]
    lock(fetch_lock)
    # check if this has failed to download recently
    fetch_fails = FETCH_FAILS[i]
    if resource in fetch_fails
        @debug "skipping recently failed download" resource=resource
        unlock(fetch_lock)
        return nothing
    end
    # see if any other thread is already downloading
    fetch_dict = FETCH_DICTS[i]
    if path in keys(fetch_dict)
        # another thread is already downloading path
        @debug "waiting for in-progress download" resource=resource
        fetch_event = fetch_dict[path]
        unlock(fetch_lock)
        wait(fetch_event)
        # TODO: try again if path doesn't exist?
        return ispath(path) ? path : nothing
    end
    fetch_dict[path] = Event()
    unlock(fetch_lock)
    # this is the only thread fetching path
    mkpath(dirname(path))
    if length(servers) == 1
        download(servers[1], resource, path)
    else
        race_lock = ReentrantLock()
        @sync for server in servers
            @spawn begin
                response = HTTP.head(server * resource, status_exception = false)
                if response.status == 200
                    # the first thread to get here downloads
                    if trylock(race_lock)
                        download(server, resource, path)
                        unlock(race_lock)
                    end
                end
                # TODO: cancel any hung HEAD requests
            end
        end
    end
    success = isfile(path)
    success || @warn "download failed" resource=resource
    # notify other threads and remove from fetch dict
    lock(fetch_lock)
    success || push!(fetch_fails, resource)
    notify(pop!(fetch_dict, path))
    unlock(fetch_lock)
    # done at last
    return success ? path : nothing
end

function forget_failures()
    for i = 1:fetch_locks
        fetch_lock = FETCH_LOCKS[i]
        lock(fetch_lock)
        empty!(FETCH_FAILS[i])
        unlock(fetch_lock)
    end
end

function tarball_git_hash(tarball::String)
    local tree_hash
    mktempdir() do tmp_dir
        run(`tar -C $tmp_dir -zxf $tarball`)
        tree_hash = bytes2hex(Pkg.GitTools.tree_hash(tmp_dir))
        chmod(tmp_dir, 0o777, recursive=true)
    end
    return tree_hash
end

function download(server::String, resource::String, path::String)
    @info "downloading resource" server=server resource=resource
    hash = let m = match(hash_part_re, resource)
        m !== nothing ? m.captures[1] : nothing
    end

    write_atomic(path) do temp_file, io
        response = HTTP.get(
            status_exception = false,
            response_stream = io,
            server * resource,
        )
        # Raise warnings about bad HTTP response codes
        if response.status != 200
            @warn "response status $(response.status)"
            return false
        end

        # If we're given a hash, then check tarball git hash
        if hash !== nothing
            tree_hash = tarball_git_hash(temp_file)
            # Raise warnings about resource hash mismatches
            if hash != tree_hash
                @warn "resource hash mismatch" server=server resource=resource hash=tree_hash
                return false
            end
        end

        return true
    end
end

function serve_file(http::HTTP.Stream, path::String)
    HTTP.setheader(http, "Content-Length" => string(filesize(path)))
    # We assume that everything we send is gzip-compressed (since they're all tarballs)
    HTTP.setheader(http, "Content-Encoding" => "gzip")
    startwrite(http)

    # Open the path, write it out directly to the HTTP stream
    open(io -> write(http, read(io, String)), path)
end


text:
# TODO under development - experimenting with type to work with manifolds

## ======================================================================================
## Possible type piracy, but also standardizing to common API across repos
## ======================================================================================

DFG.getDimension(::Distributions.Uniform) = 1
DFG.getDimension(::Normal) = 1
DFG.getDimension(Z::MvNormal) = Z |> cov |> diag |> length

DFG.getDimension(Z::FluxModelsDistribution) = length(Z.outputDim)==1 ? Z.outputDim[1] : error("can only do single index tensor at this time, please open an issue with Caesar.jl")
DFG.getDimension(Z::ManifoldKernelDensity) = getManifold(Z) |> getDimension
# TODO deprecate
DFG.getDimension(Z::BallTreeDensity) = Ndim(Z)

## ======================================================================================
## Generic manifold cost functions
## ======================================================================================
"""
    $SIGNATURES
Generic function that can be used in binary factors to calculate distance between points on Lie Groups with measurements.
"""
function distancePoint2Point(M::AbstractGroupManifold, m, p, q)
  q̂ = Manifolds.compose(M, p, m)
  # return log(M, q, q̂)
  return vee(M, q, log(M, q, q̂))
  # return distance(M, q, q̂)
end

#::MeasurementOnTangent
function distanceTangent2Point(M::AbstractGroupManifold, X, p, q)
  q̂ = Manifolds.compose(M, p, exp(M, identity_element(M, p), X)) #for groups
  # return log(M, q, q̂)
  return vee(M, q, log(M, q, q̂))
  # return distance(M, q, q̂)
end

# ::MeasurementOnTangent
function distanceTangent2Point(M::AbstractManifold, X, p, q)
  q̂ = exp(M, p, X)
  # return log(M, q, q̂)
  return vee(M, q, log(M, q, q̂))
  # return distance(M, q, q̂)
end

"""
    $SIGNATURES
Generic function that can be used in prior factors to calculate distance on Lie Groups.
"""
function distancePrior(M::AbstractManifold, meas, p)
  return log(M, p, meas)
  # return distance(M, p, meas)
end

## ======================================================================================
## ManifoldFactor
## ======================================================================================

export ManifoldFactor
# DEV NOTES
# For now, `Z` is on the tangent space in coordinates at the point used in the factor.
# For groups just the lie algebra
# As transition it will be easier this way, we can reevaluate
struct ManifoldFactor{M <: AbstractManifold, T <: SamplableBelief} <: AbstractManifoldMinimize #AbstractFactor
  M::M
  Z::T
end

DFG.getManifold(f::ManifoldFactor) = f.M

function getSample(cf::CalcFactor{<:ManifoldFactor{M,Z}}) where {M,Z}
  #TODO @assert dim == cf.factor.Z's dimension
  #TODO investigate use of SVector if small dims
  if M isa ManifoldKernelDensity
    ret = sample(cf.factor.Z.belief)[1]
  else
    ret = rand(cf.factor.Z)
  end
  #return coordinates as we do not know the point here #TODO separate Lie group
  return ret
end

# function (cf::CalcFactor{<:ManifoldFactor{<:AbstractGroupManifold}})(Xc, p, q)
function (cf::CalcFactor{<:ManifoldFactor})(Xc, p, q)
  # function (cf::ManifoldFactor)(X, p, q)
  M = cf.factor.M
  # M = cf.M
  X = hat(M, p, Xc)
  return distanceTangent2Point(M, X, p, q)
end

## ======================================================================================
## ManifoldPrior
## ======================================================================================
export ManifoldPrior, PackedManifoldPrior

# `p` is a point on manifold `M`
# `Z` is a measurement at the tangent space of `p` on manifold `M`
struct ManifoldPrior{M <: AbstractManifold, T <: SamplableBelief, P, B <: AbstractBasis} <: AbstractPrior
  M::M
  p::P #NOTE This is a fixed point from where the measurement `Z` is made in coordinates on tangent TpM
  Z::T
  basis::B
  retract_method::AbstractRetractionMethod
end

ManifoldPrior(M::AbstractGroupManifold, p, Z) = ManifoldPrior(M, p, Z, ManifoldsBase.VeeOrthogonalBasis(), ExponentialRetraction())

DFG.getManifold(f::ManifoldPrior) = f.M

#TODO
# function ManifoldPrior(M::AbstractGroupManifold, Z::SamplableBelief)
#     # p = identity_element(M, #TOOD)
#     # similar to getPointIdentity(M)
#     return ManifoldPrior(M, Z, p)
# end

# ManifoldPrior{M}(Z::SamplableBelief, p) where M = ManifoldPrior{M, typeof(Z), typeof(p)}(Z, p)

function getSample(cf::CalcFactor{<:ManifoldPrior})
  Z = cf.factor.Z
  p = cf.factor.p
  M = cf.factor.M
  basis = cf.factor.basis
  retract_method = cf.factor.retract_method
  point = samplePoint(M, Z, p, basis, retract_method)

  return point
end

#TODO investigate SVector if small dims, this is slower
# dim = manifold_dimension(M)
# Xc = [SVector{dim}(rand(Z)) for _ in 1:N]

function (cf::CalcFactor{<:ManifoldPrior})(m, p)
  M = cf.factor.M
  # return log(M, p, m)
  return vee(M,p,log(M, p, m))
  # return distancePrior(M, m, p)
end

# dist²_Σ = ⟨X, Σ⁻¹*X'⟩
function mahalanobus_distance2(M, p, q, inv_Σ)
  Xc = log(M, p, q)
  return mahalanobus_distance2(Xc, inv_Σ)
end

function mahalanobus_distance2(M, X, inv_Σ)
  #TODO look to replace with inner(MM, p, X, inv_Σ*X)
  # Xc = get_coordinates(M, p, X, DefaultOrthogonalBasis())
  Xc = vee(M, p, X)
  return Xc' * inv_Σ * Xc
end

Base.@kwdef mutable struct PackedManifoldPrior <: AbstractPackedFactor
  varType::String
  p::Vector{Float64}  #NOTE This is a fixed point from where the measurement `Z` likely stored as a coordinate
  Z::PackedSamplableBelief
end

function convert(::Union{Type{<:AbstractPackedFactor}, Type{<:PackedManifoldPrior}},
                obj::ManifoldPrior )
  #

  varT = DFG.typeModuleName(getVariableType(obj.M))

  c = AMP.makeCoordsFromPoint(obj.M, obj.p)

  # TODO convert all distributions to JSON
  Zst = convert(PackedSamplableBelief, obj.Z) # String

  PackedManifoldPrior(varT, c, Zst)
end

function convert(::Union{Type{<:AbstractFactor}, Type{<:ManifoldPrior}},
                obj::PackedManifoldPrior )
  #

  # piggy back on serialization of InferenceVariable rather than try serialize anything Manifolds.jl
  M = DFG.getTypeFromSerializationModule(obj.varType) |> getManifold

  # TODO this is too excessive
  e0 = identity_element(M)
  # u0 = getPointIdentity(obj.varType)
  p = AMP.makePointFromCoords(M, obj.p, e0) #, u0)

  Z = convert(SamplableBelief, obj.Z)

  ManifoldPrior(M, p, Z)
end

## ======================================================================================
## Generic Manifold Partial Prior
## ======================================================================================

function samplePointPartial(M::AbstractGroupManifold,
                            z::Distribution,
                            partial::Vector{Int},
                            p=identity_element(M),
                            retraction_method::AbstractRetractionMethod=ExponentialRetraction())
  dim = manifold_dimension(M)
  Xc = zeros(dim)
  Xc[partial] .= rand(z)
  X =  hat(M, p, Xc)
  return retract(M, p, X, retraction_method)
end

struct ManifoldPriorPartial{M <: AbstractManifold, T <: SamplableBelief, P <: Tuple} <: AbstractPrior
    M::M
    Z::T
    partial::P
end

DFG.getManifold(f::ManifoldPriorPartial) = f.M

function getSample(cf::CalcFactor{<:ManifoldPriorPartial})
    Z = cf.factor.Z
    M = getManifold(cf.factor)
    partial = collect(cf.factor.partial)

    return (samplePointPartial(M, Z, partial), )
end


text:
#include("CGP.jl")

#
function integration( XI::Vector{Vector{MyInt}}, numinputs )
  X = vcat( XI... )   # Combines all of the lists in XI into a long list
  sum( entropy_lst(XI[i],numinputs) for i = 1:length(XI) ) - entropy_lst(X,numinputs)
end

function entropy_lst( X::Vector{MyInt}, numinputs )
  to_bin = to_binary( X, 2^numinputs )
  sum( entropy( to_bin[:,i]) for i = 1:(2^numinputs) )
end

# Mutual information of the subset of XI whose indices are subset_inds and the complement of this subset
function MI( XI::Vector{Vector{MyInt}}, subset_inds::Vector{Int64} )
  Xinds = collect(1:length(XI))
  entropy_lst( XI[subset_inds] ) + entropy_lst( XI[setdiff(Xinds,subset_inds)] ) - entropy_lst(XI)
end

# Mutual information of X and Y where X and Y are both lists of lists.
function MI( X::Vector{Vector{MyInt}}, Y::Vector{Vector{MyInt}} )
  entropy_lst( X ) + entropy_lst( Y ) - entropy_lst( vcat(X,Y) )
end

function complexity1( XI::Vector{Vector{MyInt}} )
  n = length(XI)
  IntXI = integration(XI)
  Xinds = collect(1:length(XI))
  ssum = 0.0
  for k = 1:(n-1)
    subsets = [XI[setdiff(Xinds,s)] for s in combinations(Xinds,k)]
    println("subsets: ",subsets)
    IXk = map(integration, subsets )
    println("k: ",k,"  IXk: ",IXk)
    ssum += (k/n)*IntXI - sum( IXk )/length(IXk)
  end
  #println("result: ",ssum)
  ssum
end

function complexity2( XI::Vector{Vector{MyInt}} )
  n = length(XI)
  Xinds = collect(1:length(XI))
  ssum = 0.0
  for k = 1:(n-1)
    subsets = [[setdiff(Xinds,s),s] for s in combinations(Xinds,k)]
    println("subsets: ",subsets)
    Xsubsets = [[XI[setdiff(Xinds,s)],XI[s]] for s in combinations(Xinds,k)]
    println("Xsubsets: ",Xsubsets)
    MIlist = [MI(Xsubsets[i][1],Xsubsets[i][2]) for i = 1:length(Xsubsets)]
    println("MIlist: ",MIlist)
  end
  sum(MIlist)/2
end

function complexity5i( )
end

function to_binary( x::MyInt, numbits::Int64 )
  result = MyInt[]
  shift = numbits-1
  mask = MyInt(1) << (numbits-1)
  for i = 1:numbits
    push!(result, (mask & x)>>shift )
    shift -= 1
    mask >>= 1
  end
  result
end

function to_binary( X::Vector{MyInt}, numbits::Int64 )
  result = zeros(MyInt,length(X),numbits)
  for j in 1:length(X)
    shift = numbits-1
    mask = MyInt(1) << (numbits-1)
    for i = 1:numbits
      #push!(result, (mask & X[j])>>shift )
      result[j,i] = (mask & X[j])>>shift
      shift -= 1
      mask >>= 1
    end
  end
  result
end

function entropy( X::Vector{Vector{MyInt}} )
  sum( entropy(x) for x in X )
end