text:
using SparseDiffTools
using ForwardDiff: Dual, jacobian
using SparseArrays, Test
using LinearAlgebra
using BlockBandedMatrices
using BandedMatrices
using StaticArrays

fcalls = 0
function f(dx,x)
    global fcalls += 1
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx[end] = x[end-1] - 2x[end]
    nothing
end

function oopf(x)
    global fcalls += 1
    dx = zero(x)
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx[end] = x[end-1] - 2x[end]
    dx
end

function nsqf(x)#length(dx)<length(x)
    global fcalls +=1
    dx = zero(x)[1:div(length(x),2)]
    for i in 2:length(dx)
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx
end

function nsqf2(x)#length(dx)>length(x)
    global fcalls +=1
    dx = zeros(eltype(x),length(x)*2)
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx
end

function nsqf!(dx,x)
    global fcalls +=1
    for i in 2:length(dx)
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    nothing
end

function nsqf2!(dx,x)
    global fcalls +=1
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    nothing
end

function staticf(x,N=length(x))
    global fcalls += 1
    SVector{N}([i == 1 ? -2x[1]+x[2] : (i == N ? x[N-1]-2x[N] : x[i-1]-2x[i]+x[i+1]) for i in 1:N])
end

function staticnsqf(x,N=div(length(x),2))
    global fcalls += 1
    SVector{N}(vcat([-2x[1]+x[2]],[x[i-1]-2x[i]+x[i+1] for i in 2:N]))
end

function second_derivative_stencil(N)
    A = zeros(N,N)
    for i in 1:N, j in 1:N
        (j-i==-1 || j-i==1) && (A[i,j]=1)
        j-i==0 && (A[i,j]=-2)
    end
    A
end

@info "ended definitions"

x = rand(30)
dx = rand(30)

J = jacobian(f, dx, x)
@test J ≈ second_derivative_stencil(30)
_J = sparse(J)
@test fcalls == 3

fcalls = 0
_J1 = similar(_J)
forwarddiff_color_jacobian!(_J1, f, x, colorvec = repeat(1:3,10))
@test _J1 ≈ J
@test fcalls == 1

@info "second passed"

fcalls = 0
_J1 = forwarddiff_color_jacobian(oopf, x, colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = _J)
@test _J1 ≈ J
@test typeof(_J1) == typeof(_J)
@test fcalls == 1

@info "third passed"

fcalls = 0
_J1 = forwarddiff_color_jacobian(oopf, x, colorvec = repeat(1:3,10), sparsity = _J)
@test _J1 ≈ J
@test fcalls == 1

#oop with in-place Jacobian
fcalls = 0
_oop_jacout = sparse(1.01 .* J) # want to be nonzero to check that the pre-allocated matrix is overwritten properly
forwarddiff_color_jacobian(_oop_jacout, oopf, x; colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = _J)
@test _oop_jacout ≈ J
@test typeof(_oop_jacout) == typeof(_J)
@test fcalls == 1

# BandedMatrix
_oop_jacout = BandedMatrix(-1 => diag(J, -1) .* 1.01, 0 => diag(J, 0) .* 1.01,
                           1 => diag(J, 1) .* 1.01) # check w/BandedMatrix instead of sparse
fcalls = 0
forwarddiff_color_jacobian(_oop_jacout, oopf, x; colorvec = repeat(1:3,10), sparsity = _J)
@test _oop_jacout ≈ J
@test isa(_oop_jacout, BandedMatrix)
@test fcalls == 1

@info "4th passed"

fcalls = 0
_J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = SMatrix{30,30}(_J))
@test _J1 ≈ J
@test fcalls == 1

@info "5"

_J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), jac_prototype = SMatrix{30,30}(_J))
@test _J1 ≈ J
_J1 = forwarddiff_color_jacobian(oopf, x, jac_prototype = similar(_J))
@test _J1 ≈ J
_J1 = forwarddiff_color_jacobian(oopf, x)
@test _J1 ≈ J

#Non-square Jacobian
#length(dx)<length(x)
nsqJ = jacobian(nsqf,x)
spnsqJ = sparse(nsqJ)
_nsqJ = forwarddiff_color_jacobian(nsqf, x, dx = nothing)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, colorvec = repeat(1:3,10), sparsity = spnsqJ)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, jac_prototype = similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, colorvec = repeat(1:3,10), sparsity = spnsqJ, jac_prototype = similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, jac_prototype = SMatrix{15,30}(nsqJ))
@test _nsqJ ≈ nsqJ
@test typeof(_nsqJ) == typeof(SMatrix{15,30}(nsqJ))
_nsqJ = forwarddiff_color_jacobian(staticnsqf, SVector{30}(x), jac_prototype = SMatrix{15,30}(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(staticnsqf, SVector{30}(x), jac_prototype = SMatrix{15,30}(nsqJ), colorvec = repeat(1:3,10), sparsity = spnsqJ)
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf!, x)
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf!, x, colorvec = repeat(1:3,10), sparsity = spnsqJ )
@test _nsqJ ≈ nsqJ

#length(dx)>length(x)
nsqJ = jacobian(nsqf2,x)
spnsqJ = sparse(nsqJ)
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, dx = nothing)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, colorvec = repeat(1:3,10), sparsity = spnsqJ)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, jac_prototype = similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, colorvec = repeat(1:3,10), sparsity = spnsqJ, jac_prototype = similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, jac_prototype = SMatrix{60,30}(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf2!, x)
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf2!, x, colorvec = repeat(1:3,10), sparsity = spnsqJ )
@test _nsqJ ≈ nsqJ

fcalls = 0
_J1 = similar(_J)
jac_cache = ForwardColorJacCache(f,x,colorvec = repeat(1:3,10), sparsity = _J1)
forwarddiff_color_jacobian!(_J1, f, x, jac_cache)
@test _J1 ≈ J
@test fcalls == 1

fcalls = 0
_J1 = similar(_J)
_denseJ1 = collect(_J1)
forwarddiff_color_jacobian!(_denseJ1, f, x, colorvec = repeat(1:3,10), sparsity = _J1)
@test _denseJ1 ≈ J
@test fcalls == 1

fcalls = 0
_J1 = similar(_J)
_denseJ1 = collect(_J1)
jac_cache = ForwardColorJacCache(f,x,colorvec = repeat(1:3,10), sparsity = _J1)
forwarddiff_color_jacobian!(_denseJ1, f, x, jac_cache)
@test _denseJ1 ≈ J
@test fcalls == 1

_Jt = similar(Tridiagonal(J))
forwarddiff_color_jacobian!(_Jt, f, x, colorvec = repeat(1:3,10), sparsity = _Jt)
@test _Jt ≈ J

#https://github.com/JuliaDiff/FiniteDiff.jl/issues/67#issuecomment-516871956
function f(out, x)
	x = reshape(x, 100, 100)
	out = reshape(out, 100, 100)
	for i in 1:100
		for j in 1:100
			out[i, j] = x[i, j] + x[max(i -1, 1), j] + x[min(i+1, size(x, 1)), j] +  x[i, max(j-1, 1)]  + x[i, min(j+1, size(x, 2))]
		end
	end
	return vec(out)
end
x = rand(10000)
J = BandedBlockBandedMatrix(Ones(10000, 10000), fill(100, 100), fill(100, 100), (1, 1), (1, 1))
Jsparse = sparse(J)
colors = matrix_colors(J)
forwarddiff_color_jacobian!(J, f, x, colorvec=colors)
forwarddiff_color_jacobian!(Jsparse, f, x, colorvec=colors)
@test J ≈ Jsparse

# Non vector input
x = rand(2,2)
oopf(x) = x
iipf(fx,x) = (fx.=x)
J = forwarddiff_color_jacobian(oopf,x)
@test J ≈ Matrix(I,4,4)
J = zero(J)
forwarddiff_color_jacobian!(J,iipf,x,dx=similar(x))
@test J ≈ Matrix(I,4,4)

#1x1 SVector test
x = SVector{1}([1.])
f(x) = x
J = forwarddiff_color_jacobian(f,x)
@test J isa SArray
@test J ≈ SMatrix{1,1}([1.])


text:
module MyFunctions

export σ, γi_a, γi_l, γi_u
export FL, fL, FΓ, fΓ, IΓ

using SpecialFunctions

"""
    σ(x)

Evaluates the sigmoid function defined as
``σ(x) = 1/\\big(1+e^{-x}\\big)``.

# Examples
```

julia> σ(0)
0.5

julia> σ(1)
0.7310585786300049

julia> σ(Inf)
1.0

julia> σ([-Inf,-1,-1.,-1/2,-0.5,0,0.,sqrt(2),π,Inf])
10-element Array{Float64,1}:
 0.0
 0.2689414213699951
 0.2689414213699951
 0.3775406687981454
 0.3775406687981454
 0.5
 0.5
 0.8044296825069569
 0.9585761678336371
 1.0
```
"""
@. function σ(x)
    return 1.0/ ( 1.0 +exp( -x ) )
end

σ(1)
σ(Inf)
σ([-Inf,-1,-1.,-1/2,-0.5,0,0.,sqrt(2),π,Inf])

"""
    γi_a(a,x[;IND=0])

Compute the lower incomple normalized Gamma function.

Invokes SpecialFunctions gamma_inc(a,x,IND) and returns the array [l,u]

# Sources:
1. https://dlmf.nist.gov/8.2

# Examples
```julia-repl
julia> γi_a( 2, 3 )
2-element Array{Float64,1}:
 0.8008517265285442
 0.19914827347145578

julia> γi_a( 2, 0.301, IND=1 ) -γi_a( 2, 0.301, IND=0 )
2-element Array{Float64,1}:
 -4.2643048814294815e-11
  4.264311126433995e-11

julia> γi_a.( 2, [3,4,5,6,7] )
5-element Array{Array{Float64,1},1}:
 [0.8008517265285442, 0.19914827347145578]
 [0.9084218055563291, 0.0915781944436709]
 [0.9595723180054871, 0.040427681994512805]
 [0.9826487347633355, 0.01735126523666451]
 [0.9927049442755639, 0.00729505572443613]
```
"""
function γi_a(a,x;IND=0)
    l::Float64,u::Float64 = gamma_inc(a,x,IND)
    return [l,u]
end

γi_a( 2, 3 )
γi_a( 2, 0.301, IND=1 ) -γi_a( 2, 0.301, IND=0 )
γi_a.( 2, [3,4,5,6,7] )

"""
    γi_l(a,x[;IND=0])

Compute the lower incomple normalized Gamma function.

Invokes SpecialFunctions gamma_inc(a,x,IND) and returns only l

# Sources:
1. https://dlmf.nist.gov/8.2

# Examples
```julia-repl
julia> γi_l( 2, 3 )
0.8008517265285442

julia> γi_l( 2, 0.301, IND=1 ) -γi_l( 2, 0.301, IND=0 )
-4.2643048814294815e-11

julia> γi_l.( 2, [3,4,5,6,7] )
5-element Array{Array{Float64,1},1}:
 [0.8008517265285442]
 [0.9084218055563291]
 [0.9595723180054871]
 [0.9826487347633355]
 [0.9927049442755639]
```
"""
function γi_l(a,x;IND=0)
    l::Float64,u::Float64 = gamma_inc(a,x,IND)
    return l
end

γi_l( 2, 3 )
γi_l( 2, 0.301, IND=1 ) -γi_l( 2, 0.301, IND=0 )
γi_l.( 2, [3,4,5,6,7] )

"""
    γi_u(a,x[;IND=0])

Compute the lower incomple normalized Gamma function.

Invokes SpecialFunctions gamma_inc(a,x,IND) and returns only l

# Sources:
1. https://dlmf.nist.gov/8.2

# Examples
```julia-repl
julia> γi_u( 2, 3 )
0.19914827347145578

julia>  γi_u( 2, 0.301, IND=1 ) -γi_u( 2, 0.301, IND=0 )
4.264311126433995e-11

julia> γi_u.( 2, [3,4,5,6,7] )
5-element Array{Float64,1}:
 0.19914827347145578
 0.0915781944436709
 0.040427681994512805
 0.01735126523666451
 0.00729505572443613
```
"""
function γi_u(a,x;IND=0)
    l::Float64,u::Float64 = gamma_inc(a,x,IND)
    return l
end

γi_u( 2, 3 )
γi_u( 2, 0.301, IND=1 ) -γi_u( 2, 0.301, IND=0 )
γi_u.( 2, [3,4,5,6,7] )

# notation for distributions
# fX => probability density function (PDF) for distribution X
# FX => cummulative distribution function (CDF) for distribution X
# IX => integral of CDF for distribution X

#   -----------------------------------------------------------------------   #
#
#   Logistic distribution
#       PDF, CDF (Logistic function)
#
#   -----------------------------------------------------------------------   #

# CDF
"""
    FL(x[; p])

Compute the CDF for a variable with logistic distribution,
    ``F_L(x;x_0,k) = 1/\\big[1+e^{-k(x-x_0)}\\big]``
    parametrized via keyword argument ``p = [x_0,k]``, with default
    values ``[ x_0, k] = [0,1]``.

Alternatively, the CDF can be cast in a closed-form expression
    ``f_L(x;x_0,k) = σ( k(x-x_0) )``.

# Sources:
1. https://en.wikipedia.org/wiki/Logistic_distribution
2. https://mathworld.wolfram.com/LogisticDistribution.html

# Examples
```

julia> FL( 1 )
0.7310585786300049

julia> FL( 1.0 )
0.7310585786300049

julia> FL( [1,3/2,2.5] )
3-element Array{Float64,1}:
 0.7310585786300049
 0.8175744761936437
 0.9241418199787566

julia> FL( 1; p=[2,1] )
0.2689414213699951

julia> FL( 1.0; p=[2.0,1.0] )
0.2689414213699951

julia> FL( [1,3/2,2.5]; p=[2.0,1.0] )
3-element Array{Float64,1}:
 0.2689414213699951
 0.3775406687981454
 0.6224593312018546
```
"""
@. function FL(x; p=[ 0.0, 1.0])
    return σ( p[2]*(x-p[1]) )
end

FL( 1 )
FL( 1.0 )
FL( [1,3/2,2.5] )

FL( 1; p=[2,1] )
FL( 1.0; p=[2.0,1.0] )
FL( [1,3/2,2.5]; p=[2.0,1.0] )

# PDF: full form and condensed form using the CDF, ie σ
"""
    fL(x[; p])

Compute the PDF for a variable with logistic distribution,
    ``f_L(x;x_0,k) = k e^{-k(x-x_0)}/\\Big[ \\Big( 1 +e^{-k(x-x_0)} \\big)^2 \\Big]``
    parametrized via keyword argument ``p = [x_0,k]``, with default values
    ``[ x_0, k] = [0,1]``.
    Alternatively, the PDF can be cast in a closed-form expression
    ``f_L(x;x_0,k) = k σ( k(x-x_0) )σ( -k(x-x_0) )``.

# Sources:
1. https://en.wikipedia.org/wiki/Logistic_distribution
2. https://mathworld.wolfram.com/LogisticDistribution.html

# Examples
```julia-repl
julia> fL(0.)
0.25

julia> fL([3/2,2,2.5]; p=[2,7/2])
3-element Array{Float64,1}:
 0.4414522881532932
 0.875
 0.44145228815329324
```
"""
@. function fL(x; p=[0.0, 1.0])
    return  p[2] *σ( p[2]*(x-p[1]) ) *σ( -p[2]*(x-p[1]) )
end

fL( 1 )
fL( 1.0 )
fL( [1,3/2,2.5] )

fL( 1; p=[2,1] )
fL( 1.0; p=[2.0,1.0] )
fL( [1,3/2,2.5]; p=[2.0,1.0] )

#   -----------------------------------------------------------------------   #
#
#   Gamma distribution
#       CDF, PDF and integral of CDF
#
#   -----------------------------------------------------------------------   #

# I could not implement directly a vectorize function with multiple arguments,
# while also using an if statement, so I define a non-nonvectorized version
# which is invoked by the vectorize version
function FΓ_def(x; p=[2,1], IND=0)::Float64
    if x >= 0
        return γi_l( p[1], p[2] *x, IND=IND )
    else
        return 0.0
    end
end

"""
    FΓ(x[; p=[2,1], IND=0])

Compute the CDF for a variable with gamma distribution,
    ``F_Γ(x;α,β) = γ(α,β x) \\big/ Γ(α)``
    parametrized via keyword argument ``p = [α,β]``, with default
    values ``[α,β] = [2,1]``. Keyword IND=0 argument required for
    gamma_inc function, defaults to maximum accuracy.

``γ(α,β x)`` and ``Γ(α)`` are the lower incomplete and the complete
    Gamma functions, respectively

To facilitate implementation of code using this function we extend
    this the domain to ℝ by setting ``F_Γ(x;α,β) = 0 ∀ x < 0``.

Requires SpecialFunctions package and makes use of gamma_inc
    function.

FΓ invokes the source function FΓ_def to allow for trivial broadcast
    over the variable x. FΓ_def is not exported by default.

# Sources:
1. https://en.wikipedia.org/wiki/Gamma_distribution
2. https://mathworld.wolfram.com/GammaDistribution.html
3. https://dlmf.nist.gov/5.2
4. https://dlmf.nist.gov/8.2

# Examples
```julia-repl
julia> FΓ(0.)
0.0

julia> FΓ([-1,3/2,2.5]; p=[2,1])
3-element Array{Float64,1}:
 0.0
 0.4421745996289254
 0.7127025048163542

julia> FΓ([1,3/2,2.5]; p=[2,1], IND=1) - FΓ([1,3/2,2.5]; p=[2,1])
3-element Array{Float64,1}:
 -8.316106692163316e-10
 -7.803298285313787e-9
  0.0
```
"""
@. function FΓ(x; p=[2,1], IND=0)
    FΓ_def(x; p=p, IND=IND)
end

FΓ( 1 )
FΓ( 1.0 )
FΓ( [1,3/2,2.5] )

FΓ( 1; p=[2,1] )
FΓ( 1.0; p=[2.0,1.0] )
FΓ( [1,3/2,2.5]; p=[2.0,1.0] )

FΓ( 1; p= [2,1], IND=0 )
FΓ( 1.0; p= [2.0,1.0], IND=0 )
FΓ( [1,3/2,2.5]; p= [2.0,1.0], IND=0 )

# PDF: full form and condensed form using the CDF, ie σ
function fΓ_def(x; p=[2,1])
    if x >=0
        return p[2]^p[1] *x^(p[1]-1.0) *exp( -p[2]*x ) /gamma(p[1])
    else
        return 0.0
    end
end

"""
    fΓ(x[; p])

Compute the PDF at for a variable with gamma distribution,
    ``f_L(x;α,β) = β^α x^{α-1} e^{-β x} \\big/  Γ(α)``
    parametrized by p = [α,β], with default values
    ``[α,β] = [2,1]`` and ``Γ(α)`` is the complete Gamma
    function.

To facilitate implementation of code using this function
    we extend this the domain to ℝ by setting
    ``f_Γ(x;α,β) = 0 ∀ x < 0``.

fΓ invokes the source function fΓ_def to allow for trivial broadcast
    over the variable x. fΓ_def is not exported by default.

# Sources:
1. https://en.wikipedia.org/wiki/Gamma_distribution
2. https://mathworld.wolfram.com/GammaDistribution.html
3. https://dlmf.nist.gov/5.2

# Examples
```julia-repl
julia> fΓ(0)
0.0

julia> fΓ([-1,3/2,2.5]; p=[2,1])
3-element Array{Float64,1}:
 0.0
 0.33469524022264474
 0.205212496559747
```
"""
@. function fΓ(x; p=[2,1])
    fΓ_def(x; p=p)
end

fΓ( 1 )
fΓ( 1.0 )
fΓ([1.0, 1.3, 2.8])

fΓ( 1; p=[2,1] )
fΓ( 1.0; p=[2.0,1.0] )
fΓ([1.0, 1.3, 2.8]; p=[2.0, 1.0])

function IΓ_def(t; ρ=[2.0,1.0], IND=0)::Float64
    if t >= 0.0
        α, β = ρ
        a = ( t -α/β )*γi_l(α,β*t;IND=IND)
        b = β^(α-1.0) *t^α *exp( -β*t ) /gamma(α)
        return a+b
    else
        return 0.0
    end
end;
"""
    IΓ(x[; ρ, IND])

Compute the integral of the CDF at for a variable with gamma
    distribution
    ``I_Γ(x;α,β) = \\big[ (x-α/β)γ(α,βt) +β^{α-1} t^α e^{-βt} \\big] \\big/ Γ(α)``
    parametrized via keyword argument ρ = [α,β], with default values
    ``[α,β] = [2,1]`` and ``Γ(α)`` is the gamma function. Keyword
    IND=0 argument required for gamma_inc function, defaults to
    maximum accuracy.

``γ(α,β x)`` and ``Γ(α)`` are the lower incomplete and the complete
    Gamma functions, respectively.

To facilitate implementation of code using this function we extend
    this the domain to ℝ by setting ``I_Γ(x<0;α,β) = 0``.

Requires SpecialFunctions package and makes use of gamma_inc
    function.

IΓ invokes the source function IΓ_def to allow for trivial broadcast
    over the variable x. IΓ_def is not exported by default.

To facilitate implementation of code using this function
    we extend its the domain to ℝ by setting
    ``I_Γ(x;α,β) = 0 ∀ x < 0``.

# Sources:
1. https://en.wikipedia.org/wiki/Gamma_distribution
2. https://mathworld.wolfram.com/GammaDistribution.html
3. https://dlmf.nist.gov/5.2
4. https://dlmf.nist.gov/8.2

# Examples
```julia-repl
julia> IΓ(0)
0.0

julia> IΓ([-1,3/2,2.5]; ρ=[2,1])
3-element Array{Float64,1}:
 0.0
 0.2809555605195043
 0.8693824938075446

julia> IΓ([1,3/2,2.5]; ρ=[2,1], IND=1) -IΓ([1,3/2,2.5]; ρ=[2,1])
3-element Array{Float64,1}:
 8.316106692163316e-10
 3.901649170412469e-9
 0.0
```
"""
@. function IΓ(x;ρ=[2,1],IND=0)
    return  IΓ_def(x; ρ=ρ,IND=IND)
end

IΓ( 1 )
IΓ( 1.0 )
IΓ( [1,3/2,2.5] )

IΓ( 1; ρ=[2,1] )
IΓ( 1.0; ρ=[2.0,1.0] )
IΓ( [1,3/2,2.5]; ρ=[2.0,1.0] )

IΓ( 1; ρ=[2,1], IND=0 )
IΓ( 1.0; ρ=[2.0,1.0], IND=0 )
IΓ( [1,3/2,2.5]; ρ=[2.0,1.0], IND=0 )

IΓ([1,3/2,2.5]; ρ=[2,1], IND=1) -IΓ([1,3/2,2.5]; ρ=[2,1])

end

ext: jl

max_stars_repo_name: f-hipolito/PopulationDynamics

max_issues_repo_name: f-hipolito/PopulationDynamics

max_forks_repo_name: f-hipolito/PopulationDynamics


text:
using YAML
using LinearAlgebra
using Dates
using Random
using Parameters

#push!(LOAD_PATH, ".")  # includeの代わり．includeダメ絶対

include("utils.jl")
include("rmp.jl")
include("rmp_tree.jl")
include("static_environment.jl")
include("plot_using_Plots.jl")
include("kinematics.jl")
include("lagrange.jl")

# using .RMP
# using .RMPTree
#using .StaticEnvironment
using .Kinematics: q_neutral, q_min, q_max
#using .Utilis
using .Dynamics

### yaml関係 ###

"""動かない障害物をセット

obs_param : yamlを読んで作ったObsParam_の入ったリスト
"""
function set_obs(obs_param)
    if isnothing(obs_param)
        return nothing
    else
        obs = Vector{State{Float64}}()
        for param in obs_param
            p = param["data"] |> keytosymbol
            if param["name"] == "point"
                arg = ObsParam_point(;p...)
            elseif param["name"] == "cylinder"
                arg = ObsParam_cylinder(;p...)
            elseif param["name"] == "plane"
                arg = ObsParam_plane(;p...)
            end
            append!(obs, _set_obs(arg))
        end
    end
    return obs
end

"""静止目標をセット"""
function set_goal(goal_param)
    p = goal_param["data"] |> keytosymbol
    if goal_param["name"] == "static"
        arg = GoalParam_point(;p...)
    end
    _set_goal(arg)
end

"""rmp制御器をセット"""
function set_rmp(rmp_param)
    goal = []
    if rmp_param[1]["collision_avoidance"]["name"] == "Original"
        obs = OriginalRMPCollisionAvoidance{Float64}[]
    elseif rmp_param[1]["collision_avoidance"]["name"] == "fromGDS"
        obs = RMPfromGDSCollisionAvoidance{Float64}[]
    end
    jl = []

    for p in rmp_param
        #println(p)

        # 目標到達
        if haskey(p, "goal_attractor") && !isnothing(p["goal_attractor"])
            pa = p["goal_attractor"]
            _p = pa["data"] |> keytosymbol
            if pa["name"] == "Original"
                push!(goal, OriginalRMPAttractor(;_p...))
            elseif pa["name"] == "fromGDS"
                push!(goal, RMPfromGDSAttractor(;_p...))
            end
        end

        # 障害物回避
        if haskey(p, "collision_avoidance") && !isnothing(p["collision_avoidance"])
            po = p["collision_avoidance"]
            _p = po["data"] |> keytosymbol
            if po["name"] == "Original"
                push!(obs, OriginalRMPCollisionAvoidance(;_p...))
            elseif po["name"] == "fromGDS"
                push!(obs, RMPfromGDSCollisionAvoidance(;_p...))
            end
        end

        # ジョイント制限回避
        if haskey(p, "joint_limit_avoidance")
            pjl = p["joint_limit_avoidance"]
            _p = pjl["data"] |> keytosymbol
            if pjl["name"] == "Original"
                push!(jl, OriginalJointLimitAvoidance(;_p...))
            elseif pjl["name"] == "fromGDS"
                push!(jl, RMPfromGDSJointLimitAvoidance(;_p...))
            end
        end
    end

    return (
        joint_limit_avoidance=jl[1],
        attractor=goal[1],
        obs_avoidance=obs,
    )
end

"""ジョイント制限を守れてるかチェック"""
function check_JointLimitation(q::Vector{T}) where T
    return (q .< q_min) .| (q_max .< q)
end

"""制御点と障害物点の最小距離を計算"""
function calc_min_dis_to_obs(
    nodes::Vector{Vector{Node{T}}}, obs::Vector{State{T}}
    ) where T
    dis_to_obs = 1.0  # 障害物への最小近接距離

    _temp_dis_to_obs = Vector{T}(undef, length(obs))

    for i in 2:9
        for j in 1:length(nodes[i])
            for k in 1:length(obs)
                _temp_dis_to_obs[k] = norm(obs[k].x .- nodes[i][j].x)  # 障害物との距離
            end
            _d = minimum(_temp_dis_to_obs)
            #println(dis_to_obs, _d)
            if i == 1
                dis_to_obs = _d
            elseif dis_to_obs > _d
                #println("k")
                dis_to_obs = _d
            end
        end

    end
    return dis_to_obs
end

"""シミュレーションのデータ

t : 時刻
q : ジョイントベクトルのリスト
dq :
ddq :
desired_ddq :
u :
"""
@with_kw struct Data{T}
    t::StepRangeLen{T}  # 時刻
    q::Vector{Vector{T}}  # ジョイントベクトル
    dq::Vector{Vector{T}}  # ジョイント角速度ベクトル
    ddq::Vector{Vector{T}}  # ジョイント角加速度ベクトル
    desired_ddq::Vector{Vector{T}}  # 制御入力ベクトル
    u::Vector{Vector{T}}  # トルクベクトル
    error::Vector{T}  # eeと目標位置との誤差
    dis_to_obs::Vector{T}
    nodes::Vector{Vector{Vector{Node{T}}}}  # ノード
    goal::Vector{State{T}}
    obs::Vector{Vector{State{T}}}
    jl::Vector{BitVector}
    F::Vector{Vector{T}}  # 外力トルクベクトル
    Fc::Vector{Vector{T}}  # 実際に対象に加えた力
end

"""初期に行うやつ"""
function init_Data(
    q₀::Vector{T}, dq₀::Vector{T}, TIME_SPAN::T, Δt::T, obs, goal,
    ;isWithMass::Bool,
    ) where T

    t = range(0.0, TIME_SPAN, step = Δt)  # 時間軸

    if isWithMass
        _u = Vector{Vector{T}}(undef, length(t))
        _F = Vector{Vector{T}}(undef, length(t))
        _Fc = Vector{Vector{T}}(undef, length(t))
    else
        _u = [zeros(T, 7) for i in 1:length(t)]
        _F = [zeros(T, 7) for i in 1:length(t)]
        _Fc = [zeros(T, 3) for i in 1:length(t)]
    end

    data = Data(
        t = t,
        q = Vector{Vector{T}}(undef, length(t)),
        dq = Vector{Vector{T}}(undef, length(t)),
        ddq = Vector{Vector{T}}(undef, length(t)),
        desired_ddq = Vector{Vector{T}}(undef, length(t)),
        u = _u,
        error = Vector{T}(undef, length(t)),
        dis_to_obs = Vector{T}(undef, length(t)),
        nodes = Vector{Vector{Vector{Node{T}}}}(undef, length(t)),
        goal = Vector{State{T}}(undef, length(t)),
        obs = Vector{Vector{State{T}}}(undef, length(t)),
        jl = Vector{BitVector}(undef, length(t)),
        F = _F,
        Fc = _Fc,
    )

    # 初期値代入
    nodes₀ = update_nodes(nothing, q₀, dq₀)  # 初期ノード

    data.q[1] = q₀
    data.dq[1] = dq₀
    data.ddq[1] = zeros(T, 7)
    data.desired_ddq[1] = zeros(T, 7)
    data.error[1] = norm(goal.x - nodes₀[end][end].x)
    data.dis_to_obs[1] = 0.0
    data.nodes[1] = nodes₀
    data.goal[1] = goal
    data.obs[1] = obs
    data.jl[1] = check_JointLimitation(q₀)

    if isWithMass
        data.u[1] = zeros(T, 7)
        data.F[1] = zeros(T, 7)
        data.Fc[1] = zeros(T, 3)
    end

    return data
end

"""
オイラー法で（質量考えずに）シミュレーション
"""
function whithout_mass(
    q₀::Vector{T}, dq₀::Vector{T}, TIME_SPAN::T, Δt::T, obs, goal, rmp_param
    ) where T

    data = init_Data(
        q₀, dq₀, TIME_SPAN, Δt, obs, goal, isWithMass=false
        )

    # ぐるぐる回す
    for i in 1:length(data.t)-1
        #println("i = ", i)
        data.nodes[i+1] = update_nodes!(
            nodes = deepcopy(data.nodes[i]),
            q = data.q[i],
            dq = data.dq[i],
            rmp_param = rmp_param,
            goal = data.goal[i],
            obs = data.obs[i],
        )
        data.error[i+1] = norm(data.goal[i].x .- data.nodes[i][end][end].x)
        data.dis_to_obs[i+1] = calc_min_dis_to_obs(data.nodes[i+1], obs)

        data.desired_ddq[i+1] = calc_desired_ddq(
            data.nodes[i]
            )
        data.ddq[i+1] = data.desired_ddq[i+1]
        data.q[i+1] = data.q[i] .+ data.dq[i]*Δt
        data.dq[i+1] = data.dq[i] .+ data.ddq[i]*Δt
        data.goal[i+1] = goal
        data.obs[i+1] = obs
        data.jl[i+1] = check_JointLimitation(data.q[i+1])

        #println(data.nodes[i][1][1].x - data.nodes[i+1][1][1].x)
        #println(data.nodes[i][1] == data.nodes[i+1][1])

    end
    data
end

"""外力を計算

xdのx座標が左に来たらそれに比例した反力
"""
function externalF_at_ee(
    x::Vector{T}, xd::Vector{T}, dx::Vector{T}, t
    ) where T

    if x[1] >= xd[1]
        return zero(x)
    else
        println("t = ", t)
        println("接触!!!")
        xe = [x[1]-xd[1], x[2], x[3]]  # コンプライアンス中心
        Pe = diagm([5.0, 0.0, 0.0])
        De = diagm([1.0, 0.0, 0.0])
        println("反発力 = ", -Pe*(x .- xe) .- De*dx)
        return -Pe*(x .- xe) .- De*dx |> vec
    end
end

"""外力を計算

そっくりそのまま返す（難しい）
"""
function externalF_at_ee(;
    x::Vector{T}, xd::Vector{T},
    u::Vector{T}, Jend::Matrix{T},
    ) where T

    if norm(x - xd) > 0.005
        #println("範囲外")
        return zero(x)
    else
        println("接触!!!")
        println("物体からの反力の大きさ = ", norm(pinv(transpose(Jend)) * u))
        return -pinv(transpose(Jend)) * u |> vec
    end
end

"""
ルンゲクッタとオイラー用

q : 関節角度ベクトル
dq : 関節角速度ベクトル
u : 入力トルクベクトル
F : 外力ベクトル
"""
function dx(;
    q::Vector{T}, dq::Vector{T}, u::Vector{T},
    F::Vector{T}, Fc::Vector{T},
    Jend::Matrix{T}
    ) where T
    ddq = calc_real_ddq(
        u=u, q=q, dq=dq, F=F, Fc=Fc, Jend=Jend
    )
    return (dq = dq , ddq = ddq)
end

"""オイラー法"""
function euler_onestep(q, dq, nodes, u, goal, Δt, F)
    # _exF = externalF_at_ee(
    #     x = nodes[end][end].x,
    #     xd = goal.x,
    #     u = u,
    #     Jend = nodes[end][end].Jo,
    # )

    _exF = externalF_at_ee(
        nodes[end][end].x,
        goal.x,
        nodes[end][end].dx,
        0.0
    )

    k1 = dx(
        q = q,
        dq = dq,
        u = u,
        F = F,
        Fc = _exF,
        Jend = nodes[end][end].Jo,
    )

    ddq = k1.ddq
    dq = dq .+ k1.ddq .* Δt
    q = q .+ k1.dq .*Δt

    return ddq, dq, q
end

"""ルンゲクッタのワンステップ"""
function runge_kutta_onestep(q, dq, nodes, u, goal, Δt, F)
    # 1
    _exF = externalF_at_ee(
        x = nodes[end][end].x,
        xd = goal.x,
        u = u,
        Jend = nodes[end][end].Jo,
    )
    # println(_exF)
    # println(size(nodes[end][end].Jo))
    k1 = dx(
        q = q,
        dq = dq,
        u = u,
        F = F,
        Fc = _exF,
        Jend = nodes[end][end].Jo,
    )

    # 2
    _q = q .+ k1.dq .* Δt/2
    _dq = dq .+ k1.ddq .* Δt/2
    _, _,
    _, _, _, _,
    _, Jos_cpoint_all,
    cpoints_x_global, cpoints_dx_global,
    _, _ = calc_all(_q, _dq)

    _exF = externalF_at_ee(
        x = cpoints_x_global[end][end],
        xd = goal.x,
        u = u,
        Jend = Jos_cpoint_all[end][end],
    )
    k2 = dx(
        q=_q,
        dq=_dq,
        u=u,
        F=F,
        Fc=_exF,
        Jend=Jos_cpoint_all[end][end]
    )

    # 3
    _q = _q .+ k2.dq .* Δt/2
    _dq = _dq .+ k2.ddq .* Δt/2
    _, _,
    _, _, _, _,
    _, Jos_cpoint_all,
    cpoints_x_global, cpoints_dx_global,
    _, _ = calc_all(_q, _dq)
    _exF = externalF_at_ee(
        x = cpoints_x_global[end][end],
        xd = goal.x,
        u = u,
        Jend = Jos_cpoint_all[end][end],
    )
    k3 = dx(
        q = _q,
        dq = _dq,
        u = u,
        F = F,
        Fc = _exF,
        Jend = Jend=Jos_cpoint_all[end][end]
    )

    # 4
    _q = _q .+ k3.dq .* Δt
    _dq = _dq .+ k3.ddq .* Δt
    _, _,
    _, _, _, _,
    _, Jos_cpoint_all,
    cpoints_x_global, cpoints_dx_global,
    _, _ = calc_all(_q, _dq)
    _exF = externalF_at_ee(
        x = cpoints_x_global[end][end],
        xd = goal.x,
        u = u,
        Jend = Jos_cpoint_all[end][end],
    )
    k4 = dx(
        q = _q,
        dq = _dq,
        u = u,
        F = F,
        Fc = _exF,
        Jend = Jos_cpoint_all[end][end]
    )

    ddq = k1.ddq .+ 2 .* k2.ddq .+ 2 .* k3.ddq .+ k4.ddq
    dq = dq .+ ddq .* Δt/6
    q = q .+ (k1.dq .+ 2 .* k2.dq .+ 2 .* k3.dq .+ k4.dq) .* Δt/6

    return ddq, dq, q
end

"""
質量を考えてシミュレーション実行

ルンゲクッタを使用
"""
function with_mass(
    q₀::Vector{T}, dq₀::Vector{T}, TIME_SPAN::T, Δt::T,
    obs::Vector{State{T}}, goal::State{T}, rmp_param
    ) where T

    data = init_Data(
        q₀, dq₀, TIME_SPAN, Δt, obs, goal, isWithMass=true
    )

    ## 一定外乱
    F = zeros(T, 7)  # なし
    #F = rand(T, 7) * 0.001

    # ぐるぐる回す
    for i in 1:length(data.t)-1
        #println("t = ", data.t[i])
        data.nodes[i+1] = update_nodes!(
            nodes = deepcopy(data.nodes[i]),
            q = data.q[i],
            dq = data.dq[i],
            rmp_param = rmp_param,
            goal = data.goal[i],
            obs = data.obs[i],
        )
        data.error[i+1] = norm(data.goal[i].x .- data.nodes[i][end][end].x)
        data.dis_to_obs[i+1] = calc_min_dis_to_obs(data.nodes[i], data.obs[i])

        data.desired_ddq[i+1] = calc_desired_ddq(data.nodes[i])
        data.u[i+1] = calc_torque(data.q[i], data.dq[i], data.desired_ddq[i+1])

        # # ルンゲクッタで次の値を決定
        # data.ddq[i+1], data.dq[i+1], data.q[i+1] = runge_kutta_onestep(
        #     data.q[i], data.dq[i], data.nodes[i], data.u[i], data.goal[i], Δt, F
        # )

        # オイラー法で次の値を決定
        data.ddq[i+1], data.dq[i+1], data.q[i+1] = euler_onestep(
            data.q[i], data.dq[i], data.nodes[i], data.u[i], data.goal[i], Δt, F
        )

        data.F[i+1] = F
        # data.Fc[i+1] = externalF_at_ee(
        #     x = data.nodes[i][end][end].x,
        #     xd = data.goal[i].x,
        #     u = data.u[i],
        #     Jend = data.nodes[i][end][end].Jo,
        # )

        data.Fc[i+1] = externalF_at_ee(
            data.nodes[i][end][end].x,
            goal.x,
            data.nodes[i][end][end].dx,
            data.t[i]
        )

        data.goal[i+1] = goal
        data.obs[i+1] = obs
        data.jl[i+1] = check_JointLimitation(data.q[i+1])
    end

    data
end

"""ひとまずシミュレーションやってみｓる"""
function run_simulation(;
    saveRmpData::Bool, isWithMas::Bool, TIME_SPAN::T, Δt::T,
    rmps, obs, goal, save_path
) where T

    q₀ = q_neutral
    dq₀ = zeros(T, 7)

    if isWithMas
        @time data = with_mass(q₀, dq₀, TIME_SPAN, Δt, obs, goal, rmps)
    else
        @time data = whithout_mass(q₀, dq₀, TIME_SPAN, Δt, obs, goal, rmps)
    end

    @time plot_simulation_data(data, save_path)
    @time plot_rmp(data, save_path)
    @time make_animation(data, save_path)

    #println(typeof(length(data.q)))

    return data
end

"""ランナー

設定を読み込みシミュレーションを実行
config : yamlファイルのパス
path : 結果保存先のパス
"""
function runner(config, save_path)
    params = YAML.load_file(config)
    sim_param = params["sim_param"]
    rmp_param = params["rmp_param"]
    env_param = params["env_param"]
    obs = set_obs(env_param["obstacle"])
    goal = set_goal(env_param["goal"])
    rmps = set_rmp(rmp_param)

    data = run_simulation(
        saveRmpData = sim_param["saveRmpData"],
        isWithMas = sim_param["isWithMass"],
        TIME_SPAN = sim_param["TIME_SPAN"],
        Δt = sim_param["TIME_INTERVAL"],
        rmps = rmps,
        obs = obs,
        goal = goal,
        save_path = save_path
    )
    return data
end

"""データ保存先のパス

シミュレーション結果を保存するディレクトリのパスを返す
ないなら新規作成
"""
function get_time_string()
    _nd = now()  # 時刻取得
    _Y = _nd |> Dates.year |> string
    _M = _nd |> Dates.month |> string
    _D = _nd |> Dates.day |> string
    _h = _nd |> Dates.hour |> string
    _m = _nd |> Dates.minute |> string
    _s = _nd |> Dates.second |> string
    t = _h * "_" * _m * "_" * _s

    linpath = "../result_of_manipulator_dynamics/" *
    _Y * "_" * _M * "_" * _D * "/" *
    t * "/"

    mkpath(linpath)
    path = linpath
    return path
end

#config = "./config/sice.yaml"  # シミュレーション設定のパス
config = "./config/sice_2.yaml"

path = get_time_string()  # 実行時のデータ保存パス

println("hoge...")
data = runner(config, path)
println("hoge!")


text:
#
# fallback choice model functions for parameter handling
#

# paramranges(cm::ChoiceModel) = Tuple{Float64,Float64}[]

numparams(cm::ChoiceModel) = length(paramranges(cm::ChoiceModel))

# reset any internal state of the choice model
resetstate!(cm::ChoiceModel) = nothing

# setparams!(cm::ChoiceModel, params) = nothing

# getparams(cm::ChoiceModel) = (Float64)[]

# estimateparams!(cm::ChoiceModel, cptraces) = nothing

include("simple_choice_model.jl")
include("sampler_choice_model.jl")
include("nmcs_choice_model.jl")
include("mcts_choice_model.jl")
include("ge_choice_model.jl")
include("switch_choice_model.jl")


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule Binutils_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("Binutils")
JLLWrappers.@generate_main_file("Binutils", UUID("489e263e-5428-50b0-a723-147a141b401e"))
end  # module Binutils_jll


text:
function binom(n::Integer, k::Integer)
    n ≥ k || return 0 # short circuit base cases
    (n == 1 || k == 0) && return 1

    n * binom(n - 1, k - 1) ÷ k
end

@show binom(5, 3)

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
using LinearAlgebra
using SparseArrays

function laplacian()
    bc = (dirichlet, periodic)
    st = stencil(bc)
    n = (3, 4)

    n̅, x̅, x̅ₘ, dx̅ = mesh(bc, st, n)

    δ̅₋, δ̅₊ = backwarddiff(bc, n̅), forwarddiff(bc, n̅)
    σ̅₋, σ̅₊ = backwardinterp(bc, n̅), forwardinterp(bc, n̅)

    Δ̅ = -mapreduce(*, +, δ̅₋, δ̅₊)
    Δ = restrict(Δ̅, n̅, n) |> sparse

    x = restrict.(x̅, Ref(n̅), Ref(n))
    xₘ = restrict.(x̅ₘ, Ref(n̅), Ref(n))

    f(u, v) = π * u ^ 2 * v - √2 * v ^ 2 - √3u + 1

    y = f.(x...) |> collect
    b = Δ * y
    yₕ = Δ \ b

    norm(y - yₕ)
end

@testset "rectangular" begin
    @test isapprox(laplacian(), 0, atol=1e-15)
end


text:
module RandomLinearModels

using LinearAlgebra

function isstabled(A)
    all(x->abs(x)<=1, eigvals(A))
end

function isstablec(A)
    all(x->real(x)<=0, eigvals(A))
end

function controllability(A,B)
    n,m = size(B)
    R = zeros(n,n*m)
    Ak = zero(A) .+ I(n)
    for k = 1:n
        Rk = view(R, :, (k-1)*m .+ (1:m))
        Rk .= Ak*B
        Ak = Ak * A
    end
    return R
end

iscontrollable(A,B) = rank(controllability(A,B)) == size(A,1)

function genA(v, n, m)
    # Generate a random orthogonal matrix
    X = randn(n,n)
    Q = qr(X).Q
    A = Q*Diagonal(v)*Q'
end

genB(n,m) = randn(n,m)

function gendiscrete(n, m, tol::Real=1e-4)
    v = randn(n)
    v = v ./ (norm(v,Inf) + tol)     # make it (marginally) discrete stable
    A = genA(v, n, m)
    B = genB(n,m)
    A,B
end

function gencontinuous(n, m, tol::Real=1e-4)
    v = randn(n)
    v = v .- (maximum(v) + tol)     # make it (marginally) continuous stable
    A = genA(v, n, m)
    B = genB(n,m)
    A,B
end

function gencontrollable(n, m, type::Symbol=:discrete; tol::Real=1e-4, maxiter=20)
    cnt = 0
    while true
        if type == :discrete
            A,B = gendiscrete(n, m, tol)
        elseif type == :continuous
            A,B = gencontinuous(n, m, tol)
        end
        iscontrollable(A,B) && (return A,B)
        if cnt > maxiter
            throw(ErrorException("exceeded max number of attempt to find a controllable system."))
        end
        cnt += 1
    end
end

end


text:
# test parts
# include("pressure.jl")

resize!(ARGS, 1)
ARGS[1] = "input_vals_airfoil.jl"
# ARGS[1] = "input_vals_vortex_parallel.jl"

#----  Initialize EulerEquationMod for all the global variables necessary  ----#
include("startup.jl")

#----  Write the surface pressure coefficient for the concerned edge  ----#
#=
g_edges = [3]
# Create data structure for storing coefficient of pressure
nface_arr = zeros(Int,length(g_edges))
for i = 1:length(g_edges)
  g_edge_number = g_edges[i] # Extract geometric edge number
  nface_arr[i] = EulerEquationMod.getnFaces(mesh, g_edge_number)
end
pressCoeffArrWrite = Array{Array{Float64,2}}(length(g_edges))
for i = 1:length(g_edges)
  pressCoeffArrWrite[i] = zeros(Float64, mesh.sbpface.numnodes, nface_arr[i])
end
EulerEquationMod.writeSurfacePressureCoeff(mesh, sbp, eqn, opts, g_edges, pressCoeffArrWrite)
=#

# Create the objective function data object
functional_faces = opts["geom_faces_objective"]
objective = EulerEquationMod.BoundaryForceData{Complex128, :drag}(mesh, sbp, eqn, opts, functional_faces)
objective = EulerEquationMod.createObjectiveFunctionalData(mesh, sbp, eqn, opts)

#=
#----  Read in the target surface pressure coefficients  ----#
EulerEquationMod.readSurfacePressureCoeff(mesh, sbp, eqn, opts, g_edges,
objective.pressCoeff_obj.targetCp_arr)

# Check if the pressure coefficients are being read in correctly

ctr = 0
for i = 1:length(g_edges)
  for j = 1:size(pressCoeffArrWrite[i],2)
    for k = 1:size(pressCoeffArrWrite[i],1)
      err = objective.pressCoeff_obj.targetCp_arr[i][k,j] - pressCoeffArrWrite[i][k,j]
      if abs(err) > 1e-12
        ctr += 1
        println("something is wrong")
      end
    end
  end
end

if ctr == 0
  println("all good")
end
=#

#---- Get the objective function value  ----#
EulerEquationMod.evalFunctional(mesh, sbp, eqn, opts, objective)
println("Objective function name = ", opts["objective_function"])
println("objective.drag_val = $(objective.drag_val)")
println("objective.lift_val = $(objective.lift_val)")

EulerEquationMod.evalFunctional(mesh, sbp, eqn, opts, objective)
println("Objective function name = ", opts["objective_function"])
println("objective.drag_val = $(objective.drag_val)")
println("objective.lift_val = $(objective.lift_val)")

# Calculate the adjoint vector
adjoint_vec = zeros(Complex128, mesh.numDof)
EulerEquationMod.calcAdjoint(mesh, sbp, eqn, opts, objective, adjoint_vec)

PdePumiInterface.saveSolutionToMesh(mesh, real(adjoint_vec))
PdePumiInterface.writeVisFiles(mesh, "adjoint_field")

if MPI.Initialized()
  MPI.Finalize()
end


text:
module TestGLM

# using Revise
using Test

using MLJBase
import MLJModels
import GLM
using MLJModels.GLM_

###
### OLSREGRESSOR
###

task = load_boston()
X, y = task();

train, test = partition(eachindex(y), 0.7)

atom_ols = OLSRegressor()

Xtrain = selectrows(X, train)
ytrain = selectrows(y, train)
Xtest  = selectrows(X, test)

fitresult, _, report = fit(atom_ols, 1, Xtrain, ytrain)
MLJBase.fitted_params(atom_ols, fitresult)

p = predict_mean(atom_ols, fitresult, Xtest)

# hand made regression to compare

Xa = MLJBase.matrix(X) # convert(Matrix{Float64}, X)
Xa1 = hcat(Xa, ones(size(Xa, 1)))
coefs = Xa1[train, :] \ y[train]
p2 = Xa1[test, :] * coefs

@test p ≈ p2

info(atom_ols)

p_distr = predict(atom_ols, fitresult, selectrows(X, test))

###
### GLMCOUNT
###

using Random: seed!
using Distributions

seed!(0)

X = randn(100, 3) .* randn(3)'
Xtable = table(X)

α = 0.1
β = [-0.3, 0.2, -0.1]
λ = exp.(α .+ X * β)

y = [rand(Poisson(λᵢ)) for λᵢ ∈ λ]

atom_glmcount = GLMCount()

fitresult, _, _ = fit(atom_glmcount, 1, Xtable, y)

p = predict_mean(atom_glmcount, fitresult, Xtable)

p_distr = predict(atom_glmcount, fitresult, Xtable)

end # module
true


text:
"""
Calculate value, gradient, and hessian of the variational ELBO.
"""
module DeterministicVI

using Base.Threads: threadid, nthreads

import ..Config
using ..BivariateNormals: BivariateNormalDerivatives, BvnComponent,
                          GalaxySigmaDerivs, get_bvn_cov, eval_bvn_pdf!,
                          get_bvn_derivs!, transform_bvn_derivs!
using ..Model
using ..Model: ImagePatch, BvnBundle
import ..Celeste: Const, @aliasscope, @unroll_loop
using ..SensitiveFloats
import ..Log
using ..Transform
import DataFrames
import Optim
import ForwardDiff.Dual
using StaticArrays
import Base: convert

export ElboArgs, generic_init_source, catalog_init_source, init_sources,
       VariationalParams, elbo, ElboIntermediateVariables

function init_thread_pool!(pool::Vector, create)
    if length(pool) != Base.Threads.nthreads()
        empty!(pool)
        for i = 1:Base.Threads.nthreads()
            push!(pool, create())
        end
    end
end

"""
Return a default-initialized VariationalParams instance.
"""
function generic_init_source(init_pos::Vector{Float64})
    ret = Vector{Float64}(length(CanonicalParams))
    ret[ids.is_star] = 0.5
    ret[ids.pos] = init_pos
    ret[ids.flux_loc] = log(2.0)
    ret[ids.flux_scale] = 1e-3
    ret[ids.gal_frac_dev] = 0.5
    ret[ids.gal_axis_ratio] = 0.5
    ret[ids.gal_angle] = 0.0
    ret[ids.gal_radius_px] = 1.0
    ret[ids.k] = 1.0 / size(ids.k, 1)
    ret[ids.color_mean] = 0.0
    ret[ids.color_var] =  1e-2
    ret
end

"""
Return VariationalParams instance initialized form a catalog entry
"""
function catalog_init_source(ce::CatalogEntry; max_gal_radius_px=Inf)
    # TODO: sync this up with the transform bounds
    ret = generic_init_source(ce.pos)

    # TODO: don't do this thresholding for background sources,
    # just for sources that are being optimized
    ret[ids.is_star[1]] = ce.is_star ? 0.8: 0.2
    ret[ids.is_star[2]] = ce.is_star ? 0.2: 0.8

    ret[ids.flux_loc[1]] = log(max(0.1, ce.star_fluxes[3]))
    ret[ids.flux_loc[2]] = log(max(0.1, ce.gal_fluxes[3]))

    function get_color(color_var, color_mean)
        color_var > 0 && color_mean > 0 ? min(max(log(color_var / color_mean), -9.), 9.) :
            color_var > 0 && color_mean <= 0 ? 3.0 :
                color_var <= 0 && color_mean > 0 ? -3.0 : 0.0
    end

    function get_colors(raw_fluxes)
        [get_color(raw_fluxes[c+1], raw_fluxes[c]) for c in 1:4]
    end

    ret[ids.color_mean[:, 1]] = get_colors(ce.star_fluxes)
    ret[ids.color_mean[:, 2]] = get_colors(ce.gal_fluxes)

    ret[ids.gal_frac_dev] = min(max(ce.gal_frac_dev, 0.015), 0.985)

    ret[ids.gal_axis_ratio] = ce.is_star ? .8 : min(max(ce.gal_axis_ratio, 0.015), 0.985)
    ret[ids.gal_angle] = ce.gal_angle
    ret[ids.gal_radius_px] = ce.is_star ? 0.2 : min(max_gal_radius_px, max(ce.gal_radius_px, 0.2))

    ret
end

function init_sources(target_sources::Vector{Int}, catalog::Vector{CatalogEntry})
    ret = Vector{Vector{Float64}}(length(catalog))
    for s in 1:length(catalog)
        ret[s] = catalog_init_source(catalog[s])
    end
    for s in target_sources
        ret[s][:] = generic_init_source(catalog[s].pos)
    end
    ret
end

include("deterministic_vi/elbo_args.jl")
include("deterministic_vi/elbo_kl.jl")
include("deterministic_vi/source_brightness.jl")
include("deterministic_vi/elbo_objective.jl")
include("deterministic_vi/ConstraintTransforms.jl")
include("deterministic_vi/ElboMaximize.jl")

end


text:
module DCellC

# --------------------------------------------------------------------------- #
# Import external functionality

using Knet

import Colors
import FileIO
import Images
import ImageFiltering
import JLD2

# --------------------------------------------------------------------------- #
# Feature library

include("image.jl")
include("label.jl")
include("imageop.jl")

include("model.jl")
include("training.jl")
include("lesson.jl")

include("synthetic.jl")
include("io.jl")
include("util.jl")

# --------------------------------------------------------------------------- #
# Exports

export Image, GreyscaleImage, RGBImage,
       ordered_patches, random_patches,
       imgdata, imgsize, imgchannels, imgconvert,
       imgtype, crop, greyscale

export Label, adjacency, proxymap

export ImageOp, RandomImageOp, Pipeline,
       Id, FlipV, FlipH, Jitter, Soften,
       PixelNoise, StretchV, StretchH,
       apply

export Model, FCModel, UNetLike, Multiscale3, FCRNA,
       weights, state, save, load, density, label,
       density_patched, hasbatchnorm

export train!, train, loss

export Lesson, Selection, lessonsave, lessonload

# legacy stuff
export SharpCircleCell, SharpEllipticCell,
       synthesize

export fileext, imgsave, imgload, lblsave, lblload,
       lmgsave, lmgload, modelsave, modelload

export declutter, merge

end # DCellC

max_stars_repo_name: tscode/DCellC

max_issues_repo_name: tscode/DCellC


text:
# ------------------------------------------------------------------------------------------
# # 条件判断
#
# ## `if` 关键字
#
# Julia 中 `if` 语法为：
#
# ```julia
# if <condition 1>
#     <option 1>
# elseif <condition 2>
#     <option 2>
# else
#     <option 3>
# end
# ```
#
# 让我们按照条件来执行不同的语句。
#
# 举个例子，我们也许想要实现FizzBuzz测试：
#
# 给定一个数字N，
# + 如果N可以被3整除则打印"Fizz"
# + 如果N可以被5整除则打印"Buzz"
# + 如果N可以同时被3和5整除则打印"FizzBuzz"
# + 其他情况则打印这个数字本身！
#
# 在下面给N取一个值：
# ------------------------------------------------------------------------------------------

N = 10 # 改变这里的值试试

if (N % 3 == 0) && (N % 5 == 0) # `&&`指的是"AND"； % 是取余运算
    println("FizzBuzz")
elseif N % 3 == 0
    println("Fizz")
elseif N % 5 == 0
    println("Buzz")
else
    println(N)
end

# ------------------------------------------------------------------------------------------
# ## 三元运算符
#
# 我们可以用三元运算符语法重写上一个代码块
#
# ```julia
# a ? b : c
# ```
#
# 它等价于
#
# ```julia
# if a
#     b
# else
#     c
# end
# ```
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# 现在来试试返回两个数中较大的数。在下面给 `x` 和 `y` 赋值：
# ------------------------------------------------------------------------------------------

x = 10
y = 0

# ------------------------------------------------------------------------------------------
# 用 `if` 和 `else` 关键字的话可以这么写：
# ------------------------------------------------------------------------------------------

if x > y
    x
else
    y
end

# ------------------------------------------------------------------------------------------
# 用三元运算符的话这个条件判断语句应该像这样：
# ------------------------------------------------------------------------------------------

(x > y) ? x : y

# ------------------------------------------------------------------------------------------
# ## 短路运算
#
# 我们已经见过这样的语法
# ```julia
# a && b
# ```
# 当 `a` 和 `b` 都为真的时候返回真。
#
# 当然了，如果 `a` 是假的话，Julia 不需要判断 `b` 的真假就可以判断整体为假；
# 这就“短路”了，直接返回 `false`。
#
# 第二个参数 `b` 可能是个复杂的表达式，比如说一个有副作用的函数调用，这样的话这个函数就不会被调用：
# ------------------------------------------------------------------------------------------

false && (println("hi"); true)

true && (println("hi"); true)

# ------------------------------------------------------------------------------------------
# 另一方面，如果 `a` 是真，Julia 只要返回 `b` 的值作为整个表达式的值。
# 这就意味着不需要判断 `b` 的真假！`b` 甚至可以是个报错：
# ------------------------------------------------------------------------------------------

(x > 0) && error("x cannot be greater than 0")

# ------------------------------------------------------------------------------------------
# 类似的，检验一下逻辑“或”的 `||` 运算符，它也有短路运算。
# ------------------------------------------------------------------------------------------

true || println("hi")

# ------------------------------------------------------------------------------------------
# 以及
# ------------------------------------------------------------------------------------------

false || println("hi")

# ------------------------------------------------------------------------------------------
# ## 练习
#
# ### 5.1
# 写一个函数 `f`：当一个数为偶数时返回自身的字符串形式，为奇数时返回"odd"。
# ------------------------------------------------------------------------------------------

function f(n)
    "$n"
end

@assert f(4) == "4"
@assert f(5) == "odd"

# ------------------------------------------------------------------------------------------
# ### 5.2
# 将 5.1 中的代码用三元运算符重写。
# ------------------------------------------------------------------------------------------

f(n) = "odd"

@assert f(8) == "4"
@assert f(9) == "odd"

max_stars_repo_name: NahsiN/JuliaTutorials

max_issues_repo_name: NahsiN/JuliaTutorials

max_forks_repo_name: NahsiN/JuliaTutorials


text:
using Pkg
Pkg.activate(@__DIR__)
using RiceBPH
debug = let d = get(ENV, "DEBUG", "false")
    try
        d = parse(Bool, d)
    catch
        d = false
    end
    d
end
RiceBPH.Dashboard.start(; debug=debug)

max_stars_repo_name: ndgnuh/RiceBPH

max_issues_repo_name: ndgnuh/RiceBPH

max_forks_repo_name: ndgnuh/RiceBPH


text:
# Is included in latest CUBLAS version, can be removed in the future!

## getriBatched - performs batched matrix inversion

# Loop over all these CUBLAS functions and create custom overloads of getri_batched!
for (fname, elty) in
    ((:cublasDgetriBatched, :Float64),
     (:cublasSgetriBatched, :Float32),
     (:cublasZgetriBatched, :ComplexF64),
     (:cublasCgetriBatched, :ComplexF32))
    @eval begin
    # cublasStatus_t cublasDgetriBatched(
    #   cublasHandle_t handle, int n, double **A,
    #   int lda, int *PivotArray, double **C,
    #   int ldc, int *info, int batchSize)
        function getri_batched!(n, Aptrs::CuVector{CuPtr{$elty}},
                          lda, Cptrs::CuVector{CuPtr{$elty}},ldc,
                          pivotArray::CuArray{Cint})
            batchSize = length(Aptrs)
            info = CuArray{Cint}(undef, batchSize)
            CUBLAS.$fname(CUBLAS.handle(), n, Aptrs, lda, pivotArray, Cptrs, ldc, info, batchSize)
            CUBLAS.unsafe_free!(Cptrs)
            CUBLAS.unsafe_free!(Aptrs)
            return info
        end
    end
end

# Loop over all these CUBLAS functions and create custom overloads of getri_batched!
for (fname, elty) in
        ((:cublasDgetriBatched, :Float64),
         (:cublasSgetriBatched, :Float32),
         (:cublasZgetriBatched, :ComplexF64),
         (:cublasCgetriBatched, :ComplexF32))
    @eval begin
        # cublasStatus_t cublasDgetriBatched(
        #   cublasHandle_t handle, int n, double **A,
        #   int lda, int *PivotArray, double **C,
        #   int ldc, int *info, int batchSize)
        function getri_batched!(A::Vector{<:CuMatrix{$elty}},
                              C::Vector{<:CuMatrix{$elty}},
                              pivotArray::CuMatrix{Cint})
            n = size(A[1])[1]
            lda = max(1, stride(A[1], 2))
            ldc = max(1, stride(C[1], 2))
            Aptrs = CUBLAS.unsafe_batch(A)
            Cptrs = CUBLAS.unsafe_batch(C)
            info = CuArrays.zeros(Cint, length(A))
            CUBLAS.$fname(CUBLAS.handle(), n, Aptrs, lda, pivotArray, Cptrs, ldc, info, length(A))
            CUBLAS.unsafe_free!(Cptrs)
            CUBLAS.unsafe_free!(Aptrs)

            return info
        end
    end
end

# CUDA has no strided batched getri, but we can at least avoid constructing costly views (copied this over from gertf)
function getri_strided_batched!(A::CuArray{<:Any,3}, C::CuArray{<:Any,3}, pivot::CuArray{Cint})
    m, n = size(A, 1), size(A, 2)
    if m != n
        throw(DimensionMismatch("All matrices must be square!"))
    end
    ldc = max(1, stride(C, 2))
    lda = max(1, stride(A, 2))
    Cptrs = CUBLAS.unsafe_strided_batch(C)
    Aptrs = CUBLAS.unsafe_strided_batch(A)
    return getri_batched!(n, Aptrs, lda, Cptrs, ldc, pivot)
end


text:
"""
# Subsetting based on Index

```julia
subset(ts::TS, from::T, to::T) where {T<:Union{Int, TimeType}}
```

Create a subset of `ts` based on the `Index` starting `from`
(inclusive) till `to` (inclusive).

# Examples
```jldoctest; setup = :(using TSx, DataFrames, Dates, Random, Statistics)
julia> using Random;
julia> random(x) = rand(MersenneTwister(123), x);
julia> dates = Date("2022-02-01"):Week(1):Date("2022-02-01")+Month(9);
julia> ts = TS(random(length(dates)), dates)
julia> show(ts)
(40 x 1) TS with Date Index

 Index       x1
 Date        Float64
───────────────────────
 2022-02-01  0.768448
 2022-02-08  0.940515
 2022-02-15  0.673959
 2022-02-22  0.395453
 2022-03-01  0.313244
 2022-03-08  0.662555
 2022-03-15  0.586022
 2022-03-22  0.0521332
 2022-03-29  0.26864
 2022-04-05  0.108871
 2022-04-12  0.163666
 2022-04-19  0.473017
 2022-04-26  0.865412
 2022-05-03  0.617492
 2022-05-10  0.285698
 2022-05-17  0.463847
 2022-05-24  0.275819
 2022-05-31  0.446568
 2022-06-07  0.582318
 2022-06-14  0.255981
 2022-06-21  0.70586
 2022-06-28  0.291978
 2022-07-05  0.281066
 2022-07-12  0.792931
 2022-07-19  0.20923
 2022-07-26  0.918165
 2022-08-02  0.614255
 2022-08-09  0.802665
 2022-08-16  0.555668
 2022-08-23  0.940782
 2022-08-30  0.48
 2022-09-06  0.790201
 2022-09-13  0.356221
 2022-09-20  0.900925
 2022-09-27  0.529253
 2022-10-04  0.031831
 2022-10-11  0.900681
 2022-10-18  0.940299
 2022-10-25  0.621379
 2022-11-01  0.348173

julia> subset(ts, Date(2022, 03), Date(2022, 07))
(18 x 1) TS with Date Index

 Index       x1
 Date        Float64
───────────────────────
 2022-03-01  0.313244
 2022-03-08  0.662555
 2022-03-15  0.586022
 2022-03-22  0.0521332
 2022-03-29  0.26864
 2022-04-05  0.108871
 2022-04-12  0.163666
 2022-04-19  0.473017
 2022-04-26  0.865412
 2022-05-03  0.617492
 2022-05-10  0.285698
 2022-05-17  0.463847
 2022-05-24  0.275819
 2022-05-31  0.446568
 2022-06-07  0.582318
 2022-06-14  0.255981
 2022-06-21  0.70586
 2022-06-28  0.291978

julia> subset(TS(1:20, -9:10), -4, 5)
(10 x 1) TS with Int64 Index

 Index  x1
 Int64  Int64
──────────────
    -4      6
    -3      7
    -2      8
    -1      9
     0     10
     1     11
     2     12
     3     13
     4     14
     5     15

julia> subset(ts,:,Date("2022-04-12"))
(11 x 1) TS with Date Index

 Index       x1
 Date        Float64
───────────────────────
 2022-02-01  0.768448
 2022-02-08  0.940515
 2022-02-15  0.673959
 2022-02-22  0.395453
 2022-03-01  0.313244
 2022-03-08  0.662555
 2022-03-15  0.586022
 2022-03-22  0.0521332
 2022-03-29  0.26864
 2022-04-05  0.108871
 2022-04-12  0.163666

julia> subset(ts,Date("2022-9-27"),:)
(6 x 1) TS with Date Index

 Index       x1
 Date        Float64
──────────────────────
 2022-09-27  0.529253
 2022-10-04  0.031831
 2022-10-11  0.900681
 2022-10-18  0.940299
 2022-10-25  0.621379
 2022-11-01  0.348173

```

"""
function subset(ts::TS, from::T, to::T) where {T<:Union{Int, TimeType}}
    TS(DataFrames.subset(ts.coredata, :Index => x -> x .>= from .&& x .<= to))
end

function subset(ts::TS, ::Colon, to::T) where {T<:Union{Int, TimeType}}
    TS(DataFrames.subset(ts.coredata, :Index => x -> x .<= to))
end

function subset(ts::TS, from::T, ::Colon) where {T<:Union{Int, TimeType}}
    TS(DataFrames.subset(ts.coredata, :Index => x -> x .>= from))
end


text:
"""
NNPT fileformat
See EUROCONTROL NEST Manual for NNPT fileformat description

Example NNPT file format
17189
*01BP;DB;47.6513888888889;19.1944444444444;_
*01DC;DB;47.57;21.7191666666667;_
*01RD;DB;52.1555555555556;4.78694444444444;_
*02BP;DB;47.3927777777778;19.5241666666667;_
*02DC;DB;47.3327777777778;21.5022222222222;_
*03BP;DB;47.1725;19.6636111111111;_
*03DC;DB;47.3191666666667;21.6580555555556;_
*03RD;DB;52.0022222222222;4.57416666666667;_
*04BP;DB;47.2441666666667;19.7613888888889;_
*04DC;DB;47.3402777777778;21.7294444444444;_
*04RD;DB;51.8711111111111;4.55694444444444;_
*05BP;DB;47.7327777777778;18.7886111111111;_
"""
module Nnpt

export read

using Format
using CSV
using Dates
using DataFrames

const fileformat = Dict(1=>String, 2=>String, 3=>Float64, 4=>Float64, 5=>String)

const header_format = ["NAV_ID", "NAV_TYPE", "LAT_DEG", "LON_DEG", "NAV_NAME"]

function read(file)
    df = CSV.read(file, delim=";", types=fileformat, header=header_format,
    copycols=true, datarow=2)
    reformat!(df)
    return df
end

function reformat!(df)
    df[:,:TEMP] = stringreformat.(df[:,:NAV_NAME])
    select!(df, Not(:NAV_NAME))
    df[:,:NAV_NAME] =  df[:,:TEMP]
    select!(df, Not(:TEMP))
end

function stringreformat(str::String)
    if str == "_"
        return missing
    else
        return str
    end
end

end # module


text:
module MockFieldsBenchs

using InplaceArrays.Fields
using InplaceArrays.Fields: MockField, MockBasis

@inline function repeat(n,f,args...)
  for i in 1:n
    f(args...)
  end
  nothing
end

function bench1(n)
  np = 4
  p = Point(1,2)
  x = fill(p,np)
  v = 3.0
  d = 2
  f = MockField{d}(v)
  cache = field_cache(f,x)
  @time repeat(n,evaluate!,cache,f,x)
  ∇f = gradient(f)
  ∇cache = field_cache(∇f,x)
  @time repeat(n,evaluate!,∇cache,∇f,x)
end

function bench2(n)
  np = 4
  p = Point(1,2)
  x = fill(p,np)
  v = 3.0
  d = 2
  ndof = 8
  f = MockBasis{d}(v,ndof)
  cache = field_cache(f,x)
  @time repeat(n,evaluate!,cache,f,x)
  ∇f = gradient(f)
  ∇cache = field_cache(∇f,x)
  @time repeat(n,evaluate!,∇cache,∇f,x)
end

for n in (1,1,10,1000,100000)
  @eval begin
    println("+++ runing suite for n = $($n) +++")
    bench1($n)
    bench2($n)
  end
end

end # module


text:
filename = ARGS[1]

const Sentence = Vector{String}
const Vocab = Dict{String, Int}

function getSentences(filename::String)::Vector{Sentence}
    fid = open(filename)
    sentences = map(x -> split(chomp(x)," "), readlines(fid))
    close(fid)
    sentences
end

function getVocab(sentences::Vector{Sentence})::Vocab
    vocab = Dict{String, Int}()
    foldl((vocab, s) -> foldl((vocab, w) -> begin vocab[w] = get(vocab, w, 0) + 1; vocab end, vocab, s), vocab, sentences)
end

function getDocCounts(sentences::Vector{Sentence})::Dict{String, Int}
    getDocCounts(getVocab(sentences), sentences)
end

function getDocCounts(vocab::Vocab, sentences::Vector{Sentence})::Dict{String, Int}
    counts = Dict(map(k -> k => 0, keys(vocab)))
    #foldl((counts,s) -> foldl((counts, w) -> begin counts[w] += 1; counts end, counts, s), counts, sentences)
    Dict(map(w -> (w, length(filter(s -> w in s, sentences))), keys(vocab)))
end

function getWords(sentences::Vector{Sentence})::Vector{String}
    getWords(getVocab(sentences))
end

function getWords(vocab::Vocab)::Vector{String}
    vocab |> keys |> collect |> sort
end

function getFreqMatrix(sentences::Vector{Sentence})::Matrix{Float64}
    getFreqMatrix(getVocab(sentences), sentences)
end

#word x sent
function getFreqMatrix(vocab::Vocab, sentences::Vector{Sentence})::Matrix{Float64}
    nSentences = length(sentences)
    wordMap = Dict(map(p -> (p[2], p[1]), getWords(vocab) |> enumerate))
    nWords = length(wordMap)
    freqMat = zeros(nWords, nSentences)
    for (j, s) in enumerate(sentences)
        for w in s
            i = wordMap[w]
            freqMat[i,j] += 1
        end
    end
    freqMat
end

function getTF(sentences::Vector{Sentence})::Matrix{Float64}
    getTF(getVocab(sentences), sentences)
end

function getTF(vocab::Vocab, sentences::Vector{Sentence})::Matrix{Float64}
    freqMat = getFreqMatrix(vocab, sentences)
    getTF(freqMat, sentences)
end

function getTF(freqMat::Matrix{Float64}, sentences::Vector{Sentence})::Matrix{Float64}
    tfMat = zeros(Float64, size(freqMat))
    for (i, s) in enumerate(sentences)
        tfMat[:, i] = freqMat[:, i] ./ length(s)
    end
    tfMat
end

function getIDF(sentences::Vector{Sentence})::Dict{String, Float64}
    function getIDF(getVocab(sentences), sentences)
end

function getIDF(vocab::Vocab, sentences::Vector{Sentence})::Dict{String, Float64}
    counts = getDocCounts(vocab, sentences)
    Dict([(w, log(length(sentences) / c)) for (w, c) in counts])
end

function getTFIDF(sentences::Vector{Sentence})::Matrix{Float64}
    getTFIDF(getVocab(sentences), sentences)
end

function getTFIDF(vocab::Vocab, sentences::Vector{Sentence})::Matrix{Float64}
    tfMat = getTF(vocab, sentences)
    idf = getIDF(vocab, sentences)
    idfs = map(w -> idf[w], getWords(vocab))
    for i in 1:size(tfMat, 1)
        tfMat[i, :] *= idfs[i]
    end
    tfMat
end


text:
using Dice
using Dice: num_flips, num_nodes, to_dice_ir

function code_arith(f, b1::Int, b2::Int, res::Int)
    code = @dice begin
        function uniform(b::Int, w::Int) # b is the bits for uniform, w is the bitwidth
            x = Vector(undef, b)
            for i = b:-1:1
                x[i] = flip(0.5)
            end
            return add_bits(DistInt(x), w - b)
        end
        a = (uniform(b1, b1))
        b = (uniform(b2, b2))
        y = f(a, b)
        prob_equals(y[1], res) & !y[2]
    end
    code
end

# BDD analysis
code = code_arith(+, 2, 2, 0)
bdd = compile(code)
@assert infer(code, :bdd) ≈ 0.0625

code = code_arith(+, 2, 3, 0)
bdd = compile(code)
@assert infer(code, :bdd) ≈ 0.03125

code = code_arith(+, 3, 2, 0)
bdd = compile(code)
@assert infer(code, :bdd) ≈ 0.03125

code = code_arith(-, 2, 2, 0)
bdd = compile(code)
@assert infer(code, :bdd) ≈ 0.25

code = code_arith(-, 1, 2, 0)
bdd = compile(code)
@assert infer(code, :bdd) ≈ 0.25

code = code_arith(-, 3, 2, 0)
bdd = compile(code)
@assert infer(code, :bdd) ≈ 0.125

code = code_arith(*, 2, 2, 0)
bdd = compile(code)
@assert infer(code, :bdd) ≈ 0.4375

code = code_arith(*, 1, 2, 0)
bdd = compile(code)
@assert infer(code, :bdd) ≈ 0.625

code = code_arith(*, 3, 2, 0)
bdd = compile(code)
@assert infer(code, :bdd) ≈ 0.34375

# #         bdd = compile(code)
# # num_flips(bdd)
# # num_nodes(bdd)
# # @assert infer(code, :bdd) == 0.0625

# IR analysis
# to_dice_ir(code)
# has_dice_binary() && rundice(code)
# has_dice_binary() && infer(code, :ocaml)


text:
"""
Default Data Loaders for ULMFiT training for Sentiment Analysis
 - WikiText-103 corpus is to pre-train the Language model
 - IMDB movie review dataset - unsup data is used for fine-tuning Language Mode for Sentiment Analysis
 - IMDB movie review dataset - labelled data is used for training classifier for Sentiment Analysis
"""
# using CorpusLoaders
using Random

# WikiText-103 corpus loader
cd(@__DIR__)
include("WikiText-103")

# this custom preprocessing is based on the observation on the IMDB dataset
# these preprocessings will significantly reduce the "_unk_" (UNKNOWN) tokens in the corpus
function imdb_preprocess(text::AbstractString)
    ## Edit here if any preprocessing step is needed ##
    function put(en, symbol)
        l = length(en)
        (l == 1) && return en
        for i=1:l-1
            insert!(en, i*2, string(symbol))
        end
        return en
    end
    function split_word(word, symbol)
        length(word) == 1 && return [word]
        return split(word, symbol)
    end
    text = collect(text)
    deleteat!(text, findall(x -> !isvalid(x), text))
    text = join(text)
    text = replace(text, "

" => '\n')
    text = replace(text, "
" => '\n')
    tokens = intern.(lowercase.(tokenize(text)))
    for symbol in [',', '.', '-', '/', "'s"]
        tokens = split_word.(tokens, symbol)
        temp = []
        for token in tokens
            try
                append!(temp, put(token, symbol))
            catch
                append!(temp, token)
            end
        end
        tokens = temp
    end
    deleteat!(tokens, findall(x -> isequal(x, "")||isequal(x, " "), tokens))
    return tokens
end

# Loads WikiText-103 corpus and output a Channel to give a mini-batch at each call
function load_wikitext_103(batchsize::Integer, bptt::Integer)
    corpuspath = joinpath(datadep"WikiText-103", "wiki.valid.tokens")
    corpus = read(open(corpuspath, "r"), String)
    corpus = intern.(tokenize(corpus))
    return Channel(x -> generator(x, corpus; batchsize = batchsize, bptt = bptt));
end

# IMDB Data loaders for Sentiment Analysis specifically
# IMDB data loader for fine-tuning Language Model
function imdb_fine_tune_data(batchsize::Integer, bptt::Integer, num_examples::Integer=50000)
    imdb_dataset = IMDB("train_unsup")
    dataset = []
    for path in imdb_dataset.filepaths   #extract data from the files in directory and put into channel
        open(path) do fileio
            cur_text = read(fileio, String)
            append!(dataset, imdb_preprocess(cur_text))
        end #open
    end #for
    return Channel(x -> generator(x, dataset; batchsize=batchsize, bptt=bptt))
end

# IMDB data loader for training classifier
function imdb_classifier_data(batchsize::Integer)
    filepaths = IMDB("train_neg").filepaths
    append!(filepaths, IMDB("train_pos").filepaths)
    [shuffle!(filepaths) for _=1:10]
    Channel(csize=1) do docs
        n_batches = Int(floor(length(filepaths)/batchsize))
        put!(docs, n_batches)
        for i=1:n_batches
            X, Y = [], []
            for j=1:batchsize
                path = filepaths[(i-1)*batchsize+j]
                open(path) do fileio
                    cur_text = read(fileio, String)
                    tokens = imdb_preprocess(cur_text)
                    push!(X, tokens)
                    push!(Y, Float32((parse(Int32, split(path, '_')[2][1:end-4]))/10))
                end #open
            end#for
            X = pre_pad_sequences(X, "_pad_")
            put!(docs, [Flux.batch(X[k][j] for k=1:batchsize) for j=1:length(X[1])])
            put!(docs, Y[:, :])
        end #for
    end #channel
end


text:
# Copyright (c) 2019 Mathieu Besançon, Oscar Dowson, and contributors
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

using Clang.Generators
import HiGHS_jll

highs = joinpath(HiGHS_jll.artifact_dir, "include")
c_api = joinpath(highs, "interfaces", "highs_c_api.h")

build!(
    create_context(
        [c_api, joinpath(highs, "util", "HighsInt.h")],
        vcat(get_default_args(), "-I$highs"),
        load_options(joinpath(@__DIR__, "generate.toml")),
    ),
)

open(joinpath(@__DIR__, "..", "src", "gen", "libhighs.jl"), "a") do io
    for line in readlines(c_api)
        m = match(r"const HighsInt kHighs([a-zA-Z]+) = (-?[0-9]+);", line)
        if m === nothing
            continue
        end
        println(io, "const kHighs$(m[1]) = HighsInt($(m[2]))")
    end
end


text:
"""
spin29.jl - unscented transform robustness for the δfq problem, this is the old working
spin23.jl
"""

WDIR = joinpath(@__DIR__, "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

using Altro
using Distributions
using HDF5
using Hyperopt
using ForwardDiff
using LinearAlgebra
using Random
using RobotDynamics
using StaticArrays
using Zygote
using TrajectoryOptimization
const RD = RobotDynamics
const TO = TrajectoryOptimization

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin29"
const SAVE_PATH = joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME)

# problem
const CONTROL_COUNT = 1
const STATE_COUNT = 2
const ASTATE_SIZE_BASE = STATE_COUNT * HDIM_ISO + 3 * CONTROL_COUNT
# state indices
const STATE1_IDX = SVector{HDIM_ISO}(1:HDIM_ISO)
const STATE2_IDX = SVector{HDIM_ISO}(STATE1_IDX[end] + 1:STATE1_IDX[end] + HDIM_ISO)
const INTCONTROLS_IDX = STATE2_IDX[end] + 1:STATE2_IDX[end] + CONTROL_COUNT
const CONTROLS_IDX = INTCONTROLS_IDX[end] + 1:INTCONTROLS_IDX[end] + CONTROL_COUNT
const DCONTROLS_IDX = CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + CONTROL_COUNT
const S1_IDX = SVector{HDIM_ISO}(DCONTROLS_IDX[end] + 1:DCONTROLS_IDX[end] + HDIM_ISO)
const S2_IDX = SVector{HDIM_ISO}(S1_IDX[end] + 1:S1_IDX[end] + HDIM_ISO)
const S3_IDX = SVector{HDIM_ISO}(S2_IDX[end] + 1:S2_IDX[end] + HDIM_ISO)
const S4_IDX = SVector{HDIM_ISO}(S3_IDX[end] + 1:S3_IDX[end] + HDIM_ISO)
const S5_IDX = SVector{HDIM_ISO}(S4_IDX[end] + 1:S4_IDX[end] + HDIM_ISO)
const S6_IDX = SVector{HDIM_ISO}(S5_IDX[end] + 1:S5_IDX[end] + HDIM_ISO)
const S7_IDX = SVector{HDIM_ISO}(S6_IDX[end] + 1:S6_IDX[end] + HDIM_ISO)
const S8_IDX = SVector{HDIM_ISO}(S7_IDX[end] + 1:S7_IDX[end] + HDIM_ISO)
const S9_IDX = SVector{HDIM_ISO}(S8_IDX[end] + 1:S8_IDX[end] + HDIM_ISO)
const S10_IDX = SVector{HDIM_ISO}(S9_IDX[end] + 1:S9_IDX[end] + HDIM_ISO)
# control indices
const D2CONTROLS_IDX = 1:CONTROL_COUNT
# misc
const SAMPLE_STATES = [IS3_ISO]
const SAMPLE_STATE_COUNT = 1
const SAMPLES_PER_STATE = 10
const ASTATE_SIZE = ASTATE_SIZE_BASE + SAMPLE_STATE_COUNT * SAMPLES_PER_STATE * HDIM_ISO
const ACONTROL_SIZE = CONTROL_COUNT

# model
module Data
using RobotDynamics
using StaticArrays
const RD = RobotDynamics
const HDIM_ISO = 4
mutable struct Model <: RD.AbstractModel
    fq_cov::Float64
    alpha::Float64
end
end
Model = Data.Model
@inline RD.state_dim(model::Model) = ASTATE_SIZE
@inline RD.control_dim(model::Model) = ACONTROL_SIZE
@inline astate_sample_inds(sample_state_index::Int, sample_index::Int) = (
    SVector{HDIM_ISO}((
        ASTATE_SIZE_BASE + (sample_state_index - 1) * SAMPLES_PER_STATE * HDIM_ISO
        + (sample_index - 1) * HDIM_ISO + 1
    ):(
        ASTATE_SIZE_BASE + (sample_state_index - 1) * SAMPLES_PER_STATE * HDIM_ISO
        + sample_index * HDIM_ISO
    ))
)
const SAMPLE_IDXS = [astate_sample_inds(i, j)
                     for i = 1:SAMPLE_STATE_COUNT for j = 1:SAMPLES_PER_STATE]

# dynamics
function RD.discrete_dynamics(::Type{RK3}, model::Model, astate::SVector{ASTATE_SIZE},
                              acontrol::SVector{ACONTROL_SIZE}, time::Real, dt::Real)
    # base dynamics
    negi_hc = astate[CONTROLS_IDX[1]] * NEGI_H1_ISO
    h_prop = exp(dt * (FQ_NEGI_H0_ISO + negi_hc))
    state1 = h_prop * astate[STATE1_IDX]
    state2 = h_prop * astate[STATE2_IDX]
    intcontrols = astate[INTCONTROLS_IDX[1]] + dt * astate[CONTROLS_IDX[1]]
    controls = astate[CONTROLS_IDX[1]] + dt * astate[DCONTROLS_IDX[1]]
    dcontrols = astate[DCONTROLS_IDX[1]] + dt * acontrol[D2CONTROLS_IDX[1]]

    # unscented transform
    s1 = astate[S1_IDX]
    s2 = astate[S2_IDX]
    s3 = astate[S3_IDX]
    s4 = astate[S4_IDX]
    s5 = astate[S5_IDX]
    s6 = astate[S6_IDX]
    s7 = astate[S7_IDX]
    s8 = astate[S8_IDX]
    s9 = astate[S9_IDX]
    s10 = astate[S10_IDX]
    # compute state mean
    sm = 1//10 .* (
        s1 + s2 + s3 + s4 + s5
        + s6 + s7 + s8 + s9 + s10
    )
    # compute state covariance
    d1 = s1 - sm
    d2 = s2 - sm
    d3 = s3 - sm
    d4 = s4 - sm
    d5 = s5 - sm
    d6 = s6 - sm
    d7 = s7 - sm
    d8 = s8 - sm
    d9 = s9 - sm
    d10 = s10 - sm
    s_cov = 1 / (2 * model.alpha^2) .* (
        d1 * d1' + d2 * d2' + d3 * d3' + d4 * d4'
        + d5 * d5' + d6 * d6' + d7 * d7' + d8 * d8'
        + d9 * d9' + d10 * d10'
    )
    # perform cholesky decomposition on joint covariance
    cov = zeros(eltype(s_cov), HDIM_ISO + 1, HDIM_ISO + 1)
    cov[1:HDIM_ISO, 1:HDIM_ISO] .= s_cov
    cov[HDIM_ISO + 1, HDIM_ISO + 1] = model.fq_cov
    # TOOD: cholesky! requires writing zeros in upper triangle
    cov_chol = model.alpha * cholesky(Symmetric(cov)).L
    s_chol1 = cov_chol[1:HDIM_ISO, 1]
    s_chol2 = cov_chol[1:HDIM_ISO, 2]
    s_chol3 = cov_chol[1:HDIM_ISO, 3]
    s_chol4 = cov_chol[1:HDIM_ISO, 4]
    fq_chol1 = cov_chol[HDIM_ISO + 1, 1]
    fq_chol2 = cov_chol[HDIM_ISO + 1, 2]
    fq_chol3 = cov_chol[HDIM_ISO + 1, 3]
    fq_chol4 = cov_chol[HDIM_ISO + 1, 4]
    fq_chol5 = cov_chol[HDIM_ISO + 1, 5]
    # propagate transformed states
    s1 = exp(dt * ((FQ + fq_chol1) * NEGI_H0_ISO + negi_hc)) * (sm + s_chol1)
    s2 = exp(dt * ((FQ + fq_chol2) * NEGI_H0_ISO + negi_hc)) * (sm + s_chol2)
    s3 = exp(dt * ((FQ + fq_chol3) * NEGI_H0_ISO + negi_hc)) * (sm + s_chol3)
    s4 = exp(dt * ((FQ + fq_chol4) * NEGI_H0_ISO + negi_hc)) * (sm + s_chol4)
    s5 = exp(dt * ((FQ + fq_chol5) * NEGI_H0_ISO + negi_hc)) * sm
    s6 = exp(dt * ((FQ - fq_chol1) * NEGI_H0_ISO + negi_hc)) * (sm - s_chol1)
    s7 = exp(dt * ((FQ - fq_chol2) * NEGI_H0_ISO + negi_hc)) * (sm - s_chol2)
    s8 = exp(dt * ((FQ - fq_chol3) * NEGI_H0_ISO + negi_hc)) * (sm - s_chol3)
    s9 = exp(dt * ((FQ - fq_chol4) * NEGI_H0_ISO + negi_hc)) * (sm - s_chol4)
    s10 = exp(dt * ((FQ - fq_chol5) * NEGI_H0_ISO + negi_hc)) * sm
    # normalize
    s1 = s1 ./sqrt(s1's1)
    s2 = s2 ./sqrt(s2's2)
    s3 = s3 ./sqrt(s3's3)
    s4 = s4 ./sqrt(s4's4)
    s5 = s5 ./sqrt(s5's5)
    s6 = s6 ./sqrt(s6's6)
    s7 = s7 ./sqrt(s7's7)
    s8 = s8 ./sqrt(s8's8)
    s9 = s9 ./sqrt(s9's9)
    s10 = s10 ./sqrt(s10's10)

    astate_ = [
        state1; state2; intcontrols; controls; dcontrols;
        s1; s2; s3; s4; s5; s6; s7; s8; s9; s10
    ]

    return astate_
end

# This cost puts a gate error cost on
# the sample states and a LQR cost on the other terms.
# The hessian w.r.t the state and controls is constant.
struct Cost{N,M,T} <: TO.CostFunction
    Q::Diagonal{T, SVector{N,T}}
    R::Diagonal{T, SVector{M,T}}
    q::SVector{N, T}
    c::T
    target_state1::SVector{HDIM_ISO, T}
    q_ss::T
    active_samples::Array{Int,1}
end

function Cost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}},
              xf::SVector{N,T}, target_state1::SVector{HDIM_ISO,T}, q_ss::T,
              active_samples::Array{Int,1}) where {N,M,T}
    q = -Q * xf
    c = 0.5 * xf' * Q * xf
    return Cost{N,M,T}(Q, R, q, c, target_state1, q_ss, active_samples)
end

@inline TO.state_dim(cost::Cost{N,M,T}) where {N,M,T} = N
@inline TO.control_dim(cost::Cost{N,M,T}) where {N,M,T} = M
@inline Base.copy(cost::Cost{N,M,T}) where {N,M,T} = Cost{N,M,T}(
    cost.Q, cost.R, cost.q, cost.c,
    cost.target_state1, cost.q_ss, cost.active_samples
)

function TO.stage_cost(cost::Cost{N,M,T}, astate::SVector{N}) where {N,M,T}
    cost_ = 0.5 * astate' * cost.Q * astate + cost.q'astate + cost.c
    for i = 1:SAMPLE_STATE_COUNT
        for j in cost.active_samples
            sample_idx = astate_sample_inds(i, j)
            cost_ = cost_ + gate_error_iso2(astate, cost.target_state1; s1o=sample_idx[1] - 1)
        end
    end
    return cost_
end

@inline TO.stage_cost(cost::Cost{N,M,T}, astate::SVector{N}, acontrol::SVector{M}) where {N,M,T} = (
    TO.stage_cost(cost, astate) + 0.5 * acontrol' * cost.R * acontrol
)

function TO.gradient!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T},
                      astate::SVector{N,T}) where {N,M,T}
    E.q = cost.Q * astate + cost.q
    for i = 1:SAMPLE_STATE_COUNT
        for j in cost.active_samples
            sample_idx = astate_sample_inds(i, j)
            E.q[sample_idx] = (cost.q_ss * jacobian_gate_error_iso2(astate, cost.target_state1;
                                                                   s1o=sample_idx[1] - 1))
        end
    end
    return false
end

function TO.gradient!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T},
                      acontrol::SVector{M,T}) where {N,M,T}
    TO.gradient!(E, cost, astate)
    E.r = cost.R * acontrol
    E.c = 0
    return false
end

function TO.hessian!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T},
                     astate::SVector{N,T}) where {N,M,T}
    hess_sample = -1 * cost.q_ss * hessian_gate_error_iso2(cost.target_state1)
    for i = 1:SAMPLE_STATE_COUNT
        for j in cost.active_samples
            sample_idx = astate_sample_inds(i, j)
            E.Q[sample_idx, sample_idx] = hess_sample
        end
    end
    E.Q = E.Q + cost.Q
    return true
end

function TO.hessian!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T},
                     acontrol::SVector{M,T}) where {N,M,T}
    TO.hessian!(E, cost, astate)
    E.R = cost.R
    E.H .= 0
    return true
end

# main
function run_traj(;gate_type=xpiby2, evolution_time=60., solver_type=altro,
                  sqrtbp=false, integrator_type=rk3, qs=[1e0, 1e0, 1e0, 1e-1, 1e0, 1e-1],
                  dt_inv=Int64(1e1), smoke_test=false, constraint_tol=1e-8, al_tol=1e-4,
                  pn_steps=2, max_penalty=1e11, verbose=true, save=true,
                  fq_cov=FQ * 1e-2, max_iterations=Int64(2e5), gradient_tol_int=1,
                  dJ_counter_limit=Int(1e2), state_cov=1e-2, seed=0, alpha=1.,)
    Random.seed!(seed)
    model = Model(fq_cov, alpha)
    n = RD.state_dim(model)
    m = RD.control_dim(model)
    t0 = 0.

    # initial state, target state
    x0 = zeros(n)
    xf = zeros(n)
    target_states = Array{SVector{HDIM_ISO},1}(undef, SAMPLE_STATE_COUNT)
    gate = GT_GATE_ISO[gate_type]
    x0[STATE1_IDX] = IS1_ISO_
    x0[STATE2_IDX] = IS2_ISO_
    xf[STATE1_IDX] = gate * IS1_ISO_
    xf[STATE2_IDX] = gate * IS2_ISO_
    state_dist = Distributions.Normal(0., state_cov)
    for i = 1:SAMPLE_STATE_COUNT
        sample_state = SAMPLE_STATES[i]
        target_state = gate * sample_state
        target_states[i] = target_state
        for j = 1:SAMPLES_PER_STATE
            sample_idx = astate_sample_inds(i, j)
            sample = sample_state .+ rand(state_dist, HDIM_ISO)
            x0[sample_idx] = sample ./ sqrt(sample'sample)
            xf[sample_idx] = target_state
        end
    end
    x0 = SVector{n}(x0)
    xf = SVector{n}(xf)

    # control amplitude constraint
    x_max = fill(Inf, n)
    x_max[CONTROLS_IDX] .= MAX_CONTROL_NORM_0
    x_min = fill(-Inf, n)
    x_min[CONTROLS_IDX] .= -MAX_CONTROL_NORM_0
    x_max = SVector{n}(x_max)
    x_min = SVector{n}(x_min)

    # control amplitude constraint at boundary
    x_max_boundary = fill(Inf, n)
    x_max_boundary[CONTROLS_IDX] .= 0
    x_min_boundary = fill(-Inf, n)
    x_min_boundary[CONTROLS_IDX] .= 0
    x_max_boundary = SVector{n}(x_max_boundary)
    x_min_boundary = SVector{n}(x_min_boundary)

    # initial trajectory
    dt = dt_inv^(-1)
    N = Int(floor(evolution_time * dt_inv)) + 1
    U0 = [SVector{m}([
        fill(1e-4, CONTROL_COUNT);
    ]) for k = 1:N-1]
    X0 = [SVector{n}([
        fill(NaN, n);
    ]) for k = 1:N]
    Z = Traj(X0, U0, dt * ones(N))

    # cost function
    Q = Diagonal(SVector{n}([
        fill(qs[1], STATE_COUNT * HDIM_ISO); # ψ1, ψ2
        fill(qs[2], CONTROL_COUNT); # ∫a
        fill(qs[3], CONTROL_COUNT); # a
        fill(qs[4], CONTROL_COUNT); # ∂a
        fill(0, SAMPLES_PER_STATE * HDIM_ISO);
    ]))
    Qf = Q * N
    R = Diagonal(SVector{m}([
        fill(qs[6], CONTROL_COUNT); # ∂2a
    ]))
    # objective = LQRObjective(Q, R, Qf, xf, N)
    active_samples = Array(1:SAMPLES_PER_STATE)
    cost_k = Cost(Q, R, xf, target_states[1], qs[5], active_samples)
    cost_f = Cost(Qf, R, xf, target_states[1], N * qs[5], active_samples)
    objective = TO.Objective(cost_k, cost_f, N)

    # constraints
    # must satisfy control amplitude bound
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # must statisfy conrols start and end at 0
    control_bnd_boundary = BoundConstraint(n, m, x_max=x_max_boundary, x_min=x_min_boundary)
    # must reach target state, must have zero net flux
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX; INTCONTROLS_IDX])
    # must obey unit norm
    norm_idxs = copy(SAMPLE_IDXS)
    push!(norm_idxs, STATE1_IDX)
    push!(norm_idxs, STATE2_IDX)
    norm_constraints = [NormConstraint(n, m, 1, TO.Equality(), idx) for idx in norm_idxs]
    constraints = ConstraintList(n, m, N)
    add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, control_bnd_boundary, N-1:N-1)
    add_constraint!(constraints, target_astate_constraint, N:N)
    for norm_constraint in norm_constraints
        add_constraint!(constraints, norm_constraint, 2:N-1)
    end

    # solve problem
    prob = Problem{IT_RDI[integrator_type]}(model, objective, constraints,
                                            x0, xf, Z, N, t0, evolution_time)
    solver = ALTROSolver(prob)
    verbose_pn = verbose ? true : false
    verbose_ = verbose ? 2 : 0
    projected_newton = solver_type == altro ? true : false
    constraint_tolerance = solver_type == altro ? constraint_tol : al_tol
    iterations_inner = smoke_test ? 1 : 300
    iterations_outer = smoke_test ? 1 : 30
    n_steps = smoke_test ? 1 : pn_steps
    set_options!(
        solver, square_root=sqrtbp, constraint_tolerance=constraint_tolerance,
        projected_newton_tolerance=al_tol, n_steps=n_steps,
        penalty_max=max_penalty, verbose_pn=verbose_pn, verbose=verbose_,
        projected_newton=projected_newton, iterations_inner=iterations_inner,
        iterations_outer=iterations_outer, iterations=max_iterations,
        gradient_tolerance_intermediate=gradient_tol_int,
        dJ_counter_limit=dJ_counter_limit,
    )
    Altro.solve!(solver)

    # Post-process.
    acontrols_raw = TO.controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    cidx_arr = Array(CONTROLS_IDX)
    d2cidx_arr = Array(D2CONTROLS_IDX)
    cmax = TO.max_violation(solver)
    cmax_info = TO.findmax_violation(TO.get_constraints(solver))
    iterations_ = iterations(solver)

    result = Dict(
        "acontrols" => acontrols_arr,
        "controls_idx" => cidx_arr,
        "d2controls_dt2_idx" => d2cidx_arr,
        "evolution_time" => evolution_time,
        "astates" => astates_arr,
        "qs" => qs,
        "cmax" => cmax,
        "cmax_info" => cmax_info,
        "dt" => dt,
        "solver_type" => Integer(solver_type),
        "sqrtbp" => Integer(sqrtbp),
        "max_penalty" => max_penalty,
        "constraint_tol" => constraint_tol,
        "al_tol" => al_tol,
        "gradient_tolerance_intermediate" => gradient_tol_int,
        "dJ_counter_limit" => dJ_counter_limit,
        "integrator_type" => Integer(integrator_type),
        "gate_type" => Integer(gate_type),
        "save_type" => Integer(jl),
        "iterations" => iterations_,
        "seed" => seed,
        "fq_cov" => fq_cov,
        "max_iterations" => max_iterations,
    )

    # save
    if save
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            for key in keys(result)
                write(save_file, key, result[key])
            end
        end
        result["save_file_path"] = save_file_path
    end

    return result
end

max_stars_repo_name: SchusterLab/rbqoc

max_issues_repo_name: SchusterLab/rbqoc

max_forks_repo_name: SchusterLab/rbqoc


text:
#=
This file is auto-generated. Do not edit.
=#

#! format: off

"""
    mutable struct ESAC6A <: AVR
        Tr::Float64
        Ka::Float64
        Ta::Float64
        Tk::Float64
        Tb::Float64
        Tc::Float64
        Va_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
        Vr_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
        Te::Float64
        VFE_lim::Float64
        Kh::Float64
        VH_max::Float64
        Th::Float64
        Tj::Float64
        Kc::Float64
        Kd::Float64
        Ke::Float64
        E_sat::Tuple{Float64, Float64}
        Se::Tuple{Float64, Float64}
        V_ref::Float64
        saturation_coeffs::Tuple{Float64, Float64}
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int
        states_types::Vector{StateTypes}
        internal::InfrastructureSystemsInternal
    end

Modified AC6A. Used to represent field-controlled alternator-rectifier excitation systems with system-supplied electronic voltage regulators.
Parameters of IEEE Std 421.5 Type AC6A Excitacion System. ESAC6A in PSSE and PSLF.

# Arguments
- `Tr::Float64`: Regulator input filter time constant in s, validation range: `(0, 0.5)`, action if invalid: `warn`
- `Ka::Float64`: Regulator output gain, validation range: `(0, 1000)`, action if invalid: `warn`
- `Ta::Float64`: Regulator output lag time constant in s, validation range: `(0, 10)`, action if invalid: `warn`
- `Tk::Float64`: Voltage Regulator lead time constant, validation range: `(0, 10)`, action if invalid: `warn`
- `Tb::Float64`: Regulator denominator (lag) time constant in s, validation range: `(0, 20)`, action if invalid: `warn`
- `Tc::Float64`: Regulator numerator (lead) time constant in s, validation range: `(0, 20)`, action if invalid: `warn`
- `Va_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}`: Limits for regulator output `(Va_min, Va_max)`
- `Vr_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}`: Limits for exciter field voltage `(Vr_min, Vr_max)`
- `Te::Float64`: Exciter field time constant, validation range: `(eps(), 2)`, action if invalid: `error`
- `VFE_lim::Float64`: Exciter field current limiter reference, validation range: `(-5, 20)`, action if invalid: `warn`
- `Kh::Float64`: Exciter field current regulator feedback gain, validation range: `(0, 100)`, action if invalid: `warn`
- `VH_max::Float64`: Exciter field current limiter maximum output, validation range: `(0, 100)`, action if invalid: `warn`
- `Th::Float64`: Exciter field current limiter denominator (lag) time constant, validation range: `(0, 1)`
- `Tj::Float64`: Exciter field current limiter numerator (lead) time constant, validation range: `(0, 1)`, action if invalid: `warn`
- `Kc::Float64`: Rectifier loading factor proportional to commutating reactance, validation range: `(0, 1)`
- `Kd::Float64`: Demagnetizing factor, function of exciter alternator reactances, validation range: `(0, 2)`, action if invalid: `warn`
- `Ke::Float64`: Exciter field proportional constant, validation range: `(0, 2)`, action if invalid: `warn`
- `E_sat::Tuple{Float64, Float64}`: Exciter output voltage for saturation factor: (E1, E2)
- `Se::Tuple{Float64, Float64}`: Exciter saturation factor at exciter output voltage: (Se(E1), Se(E2))
- `V_ref::Float64`: Reference Voltage Set-point, validation range: `(0, nothing)`
- `saturation_coeffs::Tuple{Float64, Float64}`: Coefficients (A,B) of the function: Se(V) = B(V - A)^2/V
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`: The states are:
	Vm: Sensed terminal voltage,
	Vr1: Lead-lag state,
	Vr2: Regulator output state,
	Ve: Integrator output state,
	Vr3: Feedback output state
- `n_states::Int`: ESAC6A has 5 states
- `states_types::Vector{StateTypes}`: ESAC6A has 5 states
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct ESAC6A <: AVR
    "Regulator input filter time constant in s"
    Tr::Float64
    "Regulator output gain"
    Ka::Float64
    "Regulator output lag time constant in s"
    Ta::Float64
    "Voltage Regulator lead time constant"
    Tk::Float64
    "Regulator denominator (lag) time constant in s"
    Tb::Float64
    "Regulator numerator (lead) time constant in s"
    Tc::Float64
    "Limits for regulator output `(Va_min, Va_max)`"
    Va_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
    "Limits for exciter field voltage `(Vr_min, Vr_max)`"
    Vr_lim::NamedTuple{(:min, :max), Tuple{Float64, Float64}}
    "Exciter field time constant"
    Te::Float64
    "Exciter field current limiter reference"
    VFE_lim::Float64
    "Exciter field current regulator feedback gain"
    Kh::Float64
    "Exciter field current limiter maximum output"
    VH_max::Float64
    "Exciter field current limiter denominator (lag) time constant"
    Th::Float64
    "Exciter field current limiter numerator (lead) time constant"
    Tj::Float64
    "Rectifier loading factor proportional to commutating reactance"
    Kc::Float64
    "Demagnetizing factor, function of exciter alternator reactances"
    Kd::Float64
    "Exciter field proportional constant"
    Ke::Float64
    "Exciter output voltage for saturation factor: (E1, E2)"
    E_sat::Tuple{Float64, Float64}
    "Exciter saturation factor at exciter output voltage: (Se(E1), Se(E2))"
    Se::Tuple{Float64, Float64}
    "Reference Voltage Set-point"
    V_ref::Float64
    "Coefficients (A,B) of the function: Se(V) = B(V - A)^2/V"
    saturation_coeffs::Tuple{Float64, Float64}
    ext::Dict{String, Any}
    "The states are:
	Vm: Sensed terminal voltage,
	Vr1: Lead-lag state,
	Vr2: Regulator output state,
	Ve: Integrator output state,
	Vr3: Feedback output state"
    states::Vector{Symbol}
    "ESAC6A has 5 states"
    n_states::Int
    "ESAC6A has 5 states"
    states_types::Vector{StateTypes}
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function ESAC6A(Tr, Ka, Ta, Tk, Tb, Tc, Va_lim, Vr_lim, Te, VFE_lim, Kh, VH_max, Th, Tj, Kc, Kd, Ke, E_sat, Se, V_ref=1.0, saturation_coeffs=PowerSystems.get_avr_saturation(E_sat, Se), ext=Dict{String, Any}(), )
    ESAC6A(Tr, Ka, Ta, Tk, Tb, Tc, Va_lim, Vr_lim, Te, VFE_lim, Kh, VH_max, Th, Tj, Kc, Kd, Ke, E_sat, Se, V_ref, saturation_coeffs, ext, [:Vm, :Vr1, :Vr2, :Ve, :Vr3], 5, [StateTypes.Hybrid, StateTypes.Hybrid, StateTypes.Hybrid, StateTypes.Differential, StateTypes.Hybrid], InfrastructureSystemsInternal(), )
end

function ESAC6A(; Tr, Ka, Ta, Tk, Tb, Tc, Va_lim, Vr_lim, Te, VFE_lim, Kh, VH_max, Th, Tj, Kc, Kd, Ke, E_sat, Se, V_ref=1.0, saturation_coeffs=PowerSystems.get_avr_saturation(E_sat, Se), ext=Dict{String, Any}(), states=[:Vm, :Vr1, :Vr2, :Ve, :Vr3], n_states=5, states_types=[StateTypes.Hybrid, StateTypes.Hybrid, StateTypes.Hybrid, StateTypes.Differential, StateTypes.Hybrid], internal=InfrastructureSystemsInternal(), )
    ESAC6A(Tr, Ka, Ta, Tk, Tb, Tc, Va_lim, Vr_lim, Te, VFE_lim, Kh, VH_max, Th, Tj, Kc, Kd, Ke, E_sat, Se, V_ref, saturation_coeffs, ext, states, n_states, states_types, internal, )
end

# Constructor for demo purposes; non-functional.
function ESAC6A(::Nothing)
    ESAC6A(;
        Tr=0,
        Ka=0,
        Ta=0,
        Tk=0,
        Tb=0,
        Tc=0,
        Va_lim=(min=0.0, max=0.0),
        Vr_lim=(min=0.0, max=0.0),
        Te=0,
        VFE_lim=0,
        Kh=0,
        VH_max=0,
        Th=0,
        Tj=0,
        Kc=0,
        Kd=0,
        Ke=0,
        E_sat=(0.0, 0.0),
        Se=(0.0, 0.0),
        V_ref=0,
        saturation_coeffs=(0.0, 0.0),
        ext=Dict{String, Any}(),
    )
end

"""Get [`ESAC6A`](@ref) `Tr`."""
get_Tr(value::ESAC6A) = value.Tr
"""Get [`ESAC6A`](@ref) `Ka`."""
get_Ka(value::ESAC6A) = value.Ka
"""Get [`ESAC6A`](@ref) `Ta`."""
get_Ta(value::ESAC6A) = value.Ta
"""Get [`ESAC6A`](@ref) `Tk`."""
get_Tk(value::ESAC6A) = value.Tk
"""Get [`ESAC6A`](@ref) `Tb`."""
get_Tb(value::ESAC6A) = value.Tb
"""Get [`ESAC6A`](@ref) `Tc`."""
get_Tc(value::ESAC6A) = value.Tc
"""Get [`ESAC6A`](@ref) `Va_lim`."""
get_Va_lim(value::ESAC6A) = value.Va_lim
"""Get [`ESAC6A`](@ref) `Vr_lim`."""
get_Vr_lim(value::ESAC6A) = value.Vr_lim
"""Get [`ESAC6A`](@ref) `Te`."""
get_Te(value::ESAC6A) = value.Te
"""Get [`ESAC6A`](@ref) `VFE_lim`."""
get_VFE_lim(value::ESAC6A) = value.VFE_lim
"""Get [`ESAC6A`](@ref) `Kh`."""
get_Kh(value::ESAC6A) = value.Kh
"""Get [`ESAC6A`](@ref) `VH_max`."""
get_VH_max(value::ESAC6A) = value.VH_max
"""Get [`ESAC6A`](@ref) `Th`."""
get_Th(value::ESAC6A) = value.Th
"""Get [`ESAC6A`](@ref) `Tj`."""
get_Tj(value::ESAC6A) = value.Tj
"""Get [`ESAC6A`](@ref) `Kc`."""
get_Kc(value::ESAC6A) = value.Kc
"""Get [`ESAC6A`](@ref) `Kd`."""
get_Kd(value::ESAC6A) = value.Kd
"""Get [`ESAC6A`](@ref) `Ke`."""
get_Ke(value::ESAC6A) = value.Ke
"""Get [`ESAC6A`](@ref) `E_sat`."""
get_E_sat(value::ESAC6A) = value.E_sat
"""Get [`ESAC6A`](@ref) `Se`."""
get_Se(value::ESAC6A) = value.Se
"""Get [`ESAC6A`](@ref) `V_ref`."""
get_V_ref(value::ESAC6A) = value.V_ref
"""Get [`ESAC6A`](@ref) `saturation_coeffs`."""
get_saturation_coeffs(value::ESAC6A) = value.saturation_coeffs
"""Get [`ESAC6A`](@ref) `ext`."""
get_ext(value::ESAC6A) = value.ext
"""Get [`ESAC6A`](@ref) `states`."""
get_states(value::ESAC6A) = value.states
"""Get [`ESAC6A`](@ref) `n_states`."""
get_n_states(value::ESAC6A) = value.n_states
"""Get [`ESAC6A`](@ref) `states_types`."""
get_states_types(value::ESAC6A) = value.states_types
"""Get [`ESAC6A`](@ref) `internal`."""
get_internal(value::ESAC6A) = value.internal

"""Set [`ESAC6A`](@ref) `Tr`."""
set_Tr!(value::ESAC6A, val) = value.Tr = val
"""Set [`ESAC6A`](@ref) `Ka`."""
set_Ka!(value::ESAC6A, val) = value.Ka = val
"""Set [`ESAC6A`](@ref) `Ta`."""
set_Ta!(value::ESAC6A, val) = value.Ta = val
"""Set [`ESAC6A`](@ref) `Tk`."""
set_Tk!(value::ESAC6A, val) = value.Tk = val
"""Set [`ESAC6A`](@ref) `Tb`."""
set_Tb!(value::ESAC6A, val) = value.Tb = val
"""Set [`ESAC6A`](@ref) `Tc`."""
set_Tc!(value::ESAC6A, val) = value.Tc = val
"""Set [`ESAC6A`](@ref) `Va_lim`."""
set_Va_lim!(value::ESAC6A, val) = value.Va_lim = val
"""Set [`ESAC6A`](@ref) `Vr_lim`."""
set_Vr_lim!(value::ESAC6A, val) = value.Vr_lim = val
"""Set [`ESAC6A`](@ref) `Te`."""
set_Te!(value::ESAC6A, val) = value.Te = val
"""Set [`ESAC6A`](@ref) `VFE_lim`."""
set_VFE_lim!(value::ESAC6A, val) = value.VFE_lim = val
"""Set [`ESAC6A`](@ref) `Kh`."""
set_Kh!(value::ESAC6A, val) = value.Kh = val
"""Set [`ESAC6A`](@ref) `VH_max`."""
set_VH_max!(value::ESAC6A, val) = value.VH_max = val
"""Set [`ESAC6A`](@ref) `Th`."""
set_Th!(value::ESAC6A, val) = value.Th = val
"""Set [`ESAC6A`](@ref) `Tj`."""
set_Tj!(value::ESAC6A, val) = value.Tj = val
"""Set [`ESAC6A`](@ref) `Kc`."""
set_Kc!(value::ESAC6A, val) = value.Kc = val
"""Set [`ESAC6A`](@ref) `Kd`."""
set_Kd!(value::ESAC6A, val) = value.Kd = val
"""Set [`ESAC6A`](@ref) `Ke`."""
set_Ke!(value::ESAC6A, val) = value.Ke = val
"""Set [`ESAC6A`](@ref) `E_sat`."""
set_E_sat!(value::ESAC6A, val) = value.E_sat = val
"""Set [`ESAC6A`](@ref) `Se`."""
set_Se!(value::ESAC6A, val) = value.Se = val
"""Set [`ESAC6A`](@ref) `V_ref`."""
set_V_ref!(value::ESAC6A, val) = value.V_ref = val
"""Set [`ESAC6A`](@ref) `saturation_coeffs`."""
set_saturation_coeffs!(value::ESAC6A, val) = value.saturation_coeffs = val
"""Set [`ESAC6A`](@ref) `ext`."""
set_ext!(value::ESAC6A, val) = value.ext = val
"""Set [`ESAC6A`](@ref) `states_types`."""
set_states_types!(value::ESAC6A, val) = value.states_types = val


text:
using Test
using RenyiEntropyOfMagic
using LinearAlgebra
using OMEinsum

@testset "boundary_decomposition" begin
    A = rand(5, 2)
    L, U = boundary_decomposition(A)
    @test kron(U, U) * L ≈ boundary_vector(A)

    L, U = boundary_decomposition(A; rank=8)
    @test size(L) == (64, )
    @test size(U) == (625, 8)
end

@testset "hosvd3" begin
    T = rand(3, 5, 2)
    C, U1, U2 = hosvd3(T)
    @test ein"ijk,pi,qj->pqk"(C, U1, U2) ≈ T

    # TODO: construct a test core tensor
    # C, U1, U2 = hosvd3(T; left=3, right=4)
    # norm(ein"ijk,pi,qj->pqk"(C, U1, U2) - T)
end

@testset "exact decomposition" begin
    A = rand(2, 3, 2)
    T = left_right_kron(A)
    L = ein"abc,def,ghi,jkl->adgjbehkcfil"(A, A, A, A);
    L = reshape(L, (size(A, 1)^4, size(A, 2)^4, 16))
    @test L ≈ T

    @test transfer_matrix(A) ≈ sandwitch(T)

    C, U1, U2 = hosvd3(T)
    T3 = sandwitch(C)
    @test kron(U1, U1) * T3 * kron(transpose(U2), transpose(U2)) ≈ transfer_matrix(A)

    T, U1, U2 = middle_decomposition(A)
    @test kron(U1, U1) * T * kron(transpose(U2), transpose(U2)) ≈ transfer_matrix(A)
end

@testset "multiple sites" begin
    A1, A2, A3 = rand(3, 2), rand(3, 3, 2), rand(3, 2)
    T1, T2, T3 = boundary_vector(A1), transfer_matrix(A2), boundary_vector(A3)

    K1, U1 = boundary_decomposition(A1)
    K2, U2, U3 = middle_decomposition(A2)
    K3, U4 = boundary_decomposition(A3)

    @test kron(U1, U1) * K1 ≈ T1
    @test kron(U2, U2) * K2 * kron(transpose(U3), transpose(U3)) ≈ T2
    @test kron(U4, U4) * K3 ≈ T3

    M1 = transpose(U1) * U2
    M2 = transpose(U3) * U4

    @test transpose(K1) * kron(M1, M1) * K2 * kron(M2, M2) * K3 ≈ dot(T1, T2 * T3)

    @test replicated_expectation([A1, A2, A3]) ≈ dot(T1, T2 * T3)

    A1, A2, A3, A4 = rand(2, 2), rand(2, 3, 2), rand(3, 2, 2), rand(2, 2)
    T1, T2, T3, T4 = boundary_vector(A1), transfer_matrix(A2), transfer_matrix(A3), boundary_vector(A4)

    K1, U1 = boundary_decomposition(A1)
    K2, U2, U3 = middle_decomposition(A2)
    K3, U4, U5 = middle_decomposition(A3)
    K4, U6 = boundary_decomposition(A4)

    M1 = transpose(U1) * U2
    M2 = transpose(U3) * U4
    M3 = transpose(U5) * U6

    @test transpose(K1) * kron(M1, M1) * K2 * kron(M2, M2) * K3 * kron(M3, M3) * K4 ≈
        transpose(T1) * T2 * T3 * T4

    @test replicated_expectation([A1, A2, A3, A4]) ≈ transpose(T1) * T2 * T3 * T4
end

@testset "truncated rank" begin
    A1, A2, A3, A4 = rand(2, 2), rand(2, 3, 2), rand(3, 2, 2), rand(2, 2)
    ranks=Ranks(
        left_boundary=8,
        right_boundary=8,
        middle=[
            (8, 16),
            (16, 8),
        ]
    )
    @test replicated_expectation([A1, A2, A3, A4]; ranks) ≈ replicated_expectation([A1, A2, A3, A4]) rtol=1e-1
end


text:
###############################################################################
#
# Regrid the data to lower resolutions
#
###############################################################################
"""
    regrid_LUT(ds::AbstractDataset{FT},
               zoom::Int;
               nan_weight::Bool=false
    ) where {FT<:AbstractFloat}

Regrid the data from high to low resolution and return the struct, given
- `ds` [`GriddedDataset`](@ref) type struct
- `zoom` Integer geophysical zoom factor (>=2)
- `nan_weight` Optional. If true, assuming nan = 0; if false, neglecting it.
"""
function regrid_LUT(
            ds::GriddedDataset{FT},
            zoom::Int;
            nan_weight::Bool=false
) where {FT<:AbstractFloat}
    # unpack data
    @unpack data = ds;

    # create a matrix to store the sum and a matrix to store the counts
    mat_size  = (Int(size(data,1)/zoom), Int(size(data,2)/zoom), size(data,3));
    mat_data  = zeros(FT , mat_size);
    mat_count = zeros(Int, mat_size);

    # loop through the datasets
    for _lon_i in 1:mat_size[1], _lat_i in 1:mat_size[2]
        # select a region to map into target matrix
        _lon_j = (_lon_i-1)*zoom+1 : _lon_i*zoom;
        _lat_j = (_lat_i-1)*zoom+1 : _lat_i*zoom;

        # re-mapping the data per layer
        for _layer in 1:mat_size[3]
            for _lon_x in _lon_j, _lat_x in _lat_j
                _tmp_data = data[_lon_x, _lat_x, _layer];
                if !isnan(_tmp_data)
                    mat_data[ _lon_i, _lat_i, _layer] += _tmp_data;
                    mat_count[_lon_i, _lat_i, _layer] += 1;
                elseif nan_weight
                    mat_count[_lon_i, _lat_i, _layer] += 1;
                end
            end
        end
    end

    # refactor the matrix by taking the mean
    mat_data ./= mat_count;

    return GriddedDataset{FT}(data     = mat_data   ,
                              res_time = ds.res_time,
                              var_name = ds.var_name,
                              var_attr = ds.var_attr)
end