text:
module NNlib
using Requires

# Include APIs
include("dim_helpers.jl")

# NNPACK support
if Sys.islinux() || Sys.isapple()
    include("nnpack/NNPACK.jl")
else
    is_nnpack_available() = false
end

include("activation.jl")
include("softmax.jl")
include("gemm.jl")
include("conv.jl")
include("pooling.jl")

## Include implementations
include("impl/padding_edges.jl")

# Direct implementations of convolutional and depthwise-convolutional algorithms
include("impl/conv_direct.jl")
include("impl/depthwiseconv_direct.jl")
# im2col implementations of convolutional and depthwise-convolutional algorithms
include("impl/conv_im2col.jl")
include("impl/depthwiseconv_im2col.jl")

# Direct implementations of pooling
include("impl/pooling_direct.jl")

end # module NNlib


text:
# Copyright 2020 Digital Domain 3.0
#
# Licensed under the Apache License, Version 2.0 (the "Apache License")
# with the following modification; you may not use this file except in
# compliance with the Apache License and the following modification to it:
# Section 6. Trademarks. is deleted and replaced with:
#
# 6. Trademarks. This License does not grant permission to use the trade
#    names, trademarks, service marks, or product names of the Licensor
#    and its affiliates, except as required to comply with Section 4(c) of
#    the License and to reproduce the content of the NOTICE file.
#
# You may obtain a copy of the Apache License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the Apache License with the above modification is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the Apache License for the specific
# language governing permissions and limitations under the Apache License.

"""
Extension of Blades into a vector space.

KVectors are containers for Blades of the same grade.  New Blade elements are added with the + operator
"""
module KVectors

export
KVector,
NullKVector,
sortbasis,
isnull,
coords,
prune,
normalize_safe,
basis_1vector

using Blades
using StaticArrays
using LinearAlgebra

import Blades.grade
import Blades.dual
import Blades.⟂
import Blades.∧
import Blades.pseudoscalar
import Blades.⋆

"Container of Blades with the same grade and element type"
struct KVector{T,K,N} <: AbstractArray{T,1}
  k::SVector{N,Blade{T,K}}
end

Base.getindex(B::BT, i::T) where {BT<:KVector, T<:Integer} = B.k[i]
Base.firstindex(k::K) where {K<:KVector} = firstindex(k.k)
Base.lastindex(k::K) where {K<:KVector} = lastindex(k.k)

const NullKVector{T,K} = KVector{T,K,0}

NullKVector{T,K}() where {T,K} = KVector{T,K,0}(SVector{0,T}())
isnull(b::B) where {B<:NullKVector} = true
isnull(b::B) where {B<:KVector} = false

Base.first(k::K) where {K<:KVector} = isempty(k) ? nothing : k.k[1]
Base.length(k::K) where {K<:KVector} = length(k.k)
Base.size(k::K) where {K<:KVector} = (length(k.k),)
Base.isempty(k::K) where {K<:KVector} = isempty(k.k)
Base.iterate(A::K) where {K<:KVector} = isempty(A) ? nothing : (A[1], 2)
Base.iterate(A::K, i::T) where {K<:KVector, T} = length(A) < i ? nothing : (A[i], i+1)
Base.eltype(::Type{K}) where {T,P,K<:KVector{T,P}} = Blade{T,P}

Base.conj(k::K) where {K<:KVector} = K(map(conj, k))

KVector{T}(v::Blade{T,K}) where {T,K} = KVector{T,K,1}(SArray{Tuple{1},Blade{T,K},1,1}(v))
KVector(v::Blade{T,K}) where {T,K} = KVector{T}(v)
Base.convert(::Type{B}, k::K) where {T,N,B<:KVector, K<:Blade{T,N}} = B(k)
KVector(kv::V) where {T, KN, K<:Blade{T,KN}, V<:AbstractVector{K}} = KVector{T,KN,length(kv)}(kv)
KVector(kv::KV) where KV<:KVector = kv

KVector(uv::C, 𝐼::Blade{KT,N}) where {T, N, KT, C<:SVector{N,T}} = KVector(uv.*(basis_1blades(𝐼)))

function KVector(uv::C) where {T, N, C<:SVector{N,T}}
  @warn "resolving Blade types in top-level module via dual(1)"
  KVector(uv.*(basis_1blades(dual(1))))
end

KVector(uv::A) where {T<:Real, A<:Vector{T}} = KVector(SVector{length(uv)}(uv))
KVector(uv::A, 𝐼::Blade{KT,N}) where
  {T<:Real, A<:Vector{T}, N, KT} =
  KVector(SVector{length(uv)}(uv), 𝐼)

KVector(uv, ::Type{I}) where {I<:Blade} = KVector(uv, 1I)

Base.iszero(b::B) where {T,K,N,B<:KVector{T,K,0}} = true
Base.iszero(b::B) where {T,K,N,B<:KVector{T,K,N}} = iszero(sum((x->x*x).((k->k.x).(b.k))))

Base.:+(b::KA, c::KB) where {T<:Number,K, KA<:Blade{T,K}, KB<:Blade{T,K}} =
  KVector{T,K,2}(SArray{Tuple{2},Blade{T,K},1,2}(b,c))
Base.:-(b::KA, c::KB) where {T<:Number,K, KA<:Blade{T,K}, KB<:Blade{T,K}} =
  KVector{T,K,2}(SArray{Tuple{2},Blade{T,K},1,2}(b,-c))

function Base.:+(b::V, k::KVector{T,K,N}) where {T,K, V<:Blade{T,K},N}
  i = findfirst(v->typeof(v)==V, k.k)
  if i == nothing
    KVector{T,K,N+1}(vcat(Blade{T,K}[b],k.k))
  else
    KVector(setindex(k.k,k.k[i]+b,i))
  end
end

Base.:+(b::B, k::V) where {T,K,B<:KVector{T,K},V<:Blade{T,K}} = k+b
Base.:+(b::KVector{T,K}, c::KVector{T,K}) where {T,K} = reduce( (bc,v)->v+bc, b.k; init=c )
Base.:-(b::B, k::V) where {T,K,B<:KVector{T,K},V<:Blade{T,K}} = -k+b
Base.:-(k::V, b::B) where {T,K,B<:KVector{T,K},V<:Blade{T,K}} = k+(-b)
Base.:-(a::B) where {T, B<:KVector{T}} = -one(T)*a
Base.:-(b::KVector{T,K}, c::KVector{T,K}) where {T,K} = b+(-c)

"apply reverse to all Blades in this KVector"
Base.reverse(b::B) where B<:KVector = B(reverse.(b.k))

"""
    grade(a)

the grade k of a k-vector.  roughly speaking, the dimension of the subspace spanned by the k-vector
"""
grade(b::KVector{T,K}) where {T,K} = K
grade(b::KVector{T,K}, i::Integer) where {T,K} = i==K ? b : zero(T)

"apply dual to all Blades in this KVector"
dual(b::KVector{T,K}) where {T,K} = KVector(dual.(b))
dual(b::KVector{T,K,0}) where {T,K} = b
"apply ⟂ to all Blades in this KVector"
⟂(b::KVector{T,K})  where {T,K} = KVector((⟂).(b))
⟂(b::KVector{T,K,0})  where {T,K} = b
"apply ! to all Blades in this KVector"
Base.:!(b::B) where {B<:KVector} = dual(b)

Base.zero(b::Type{B}) where {T,K,N,B<:KVector{T,K,N}} = NullKVector{T,K}()
Base.zero(b::KVector) = zero(typeof(b))

Base.:*(nb::NullKVector{T,K}, nb2::NullKVector{T,K}) where {T,K} = nb
Base.:*(nb::NullKVector{T,K}, b::KVector{T,K}) where {T,K} = nb
Base.:*(b::KVector{T,K}, nb::NullKVector{T,K}) where {T,K} = nb*b
Base.:+(nb::NullKVector{T,K}, nb2::NullKVector{T,K}) where {T,K} = nb
Base.:+(nb::NullKVector{T,K}, b::KVector{T,K}) where {T,K} = b
Base.:+(b::KVector{T,K}, nb::NullKVector{T,K}) where {T,K} = nb+b

Base.:*(s::T, b::B) where {T<:Real, B<:KVector} = B(s*b.k)
Base.:*(b::B, s::T) where {T<:Real, B<:KVector} = s*b
Base.:/(b::B, s::T) where {T<:Real, B<:KVector} = b*(one(T)/s)

pseudoscalar(k::K) where {K<:KVector} = pseudoscalar(first(k))

#==
"""
    inv(k)

Inverse of a KVector when it exists.  Left inverse.  k⋅inv(k) = 1
"""
Base.inv(k::K) where {K<:KVector} = reverse(k)/(k⋅reverse(k))
==#

"""
    ⋆(k)

Hodge star operator mapping k to it's Hodge dual.
"""
⋆(k::K) where {K<:KVector} = mapreduce(⋆,+,k)

"""
    sortbasis(b)

sort blades by bases indices in ascending order within the k-vector
"""
sortbasis(B::BT) where {BT<:KVector} = BT(sort(Vector(B.k); by=Blades.subspace∘typeof))
sortbasis(B::BT) where {BT<:Blade} = B
Base.:(==)(B::BT, B2::BT) where {BT<:KVector} = sortbasis(B).k == sortbasis(B2).k
Base.:(==)(B::BT, B2::BT2) where {T,K,K2,BT<:KVector{T,K}, BT2<:KVector{T,K2}} = false

Base.promote_rule(::Type{B}, ::Type{K}) where {B<:KVector,K<:Blade} = KVector

prune(k::KVector, epsi = eps()) = sum(Iterators.filter(x->abs(scalar(x)) > epsi, k))

"""
    ∧(a,b)

Wedge product between two k-vectors
"""
function ∧(A::KV,B::KV2) where {T,K,N,KV<:KVector{T,K,N}, K2,N2,KV2<:KVector{T,K2,N2}}
  AwB = Blade{T, K+K2}[]
  sizehint!(AwB, N*N2)
  for aᵢ in A.k
    for bⱼ in B.k
      awb = aᵢ∧bⱼ
      if !iszero(awb)
        push!(AwB, awb)
      end
    end
  end
  sum(AwB)
end

"""
    ∧(a,b)

Wedge product between blade and k-vector
"""
function ∧(a::BT,B::KV2) where {T,K,BT<:Blade{T,K}, K2,N2,KV2<:KVector{T,K2,N2}}
  AwB = Blade{T, K+K2}[]
  sizehint!(AwB, N2)
  for bⱼ in B.k
    awb = a∧bⱼ
    if !iszero(awb)
      push!(AwB, awb)
    end
  end
  sum(AwB)
end

∧(B::M,a::N) where {M<:KVector,N<:Blade} = -a∧B
∧(s::T,B::N) where {T<:Real,N<:KVector} = s*B
∧(B::N,s::T) where {T<:Real,N<:KVector} = s*B

"""
    gram(a, u)

The Gram matrix, a length(a)xlength(a) matrix where each entry is the dot product aᵢ⋅uⱼ

Entries are sorted by acending basis indices from a and u.

Example:
for length(a) = 2, gram( a, u ) = [ a₁⋅u₁ a₁⋅u₂ ;
                                    a₂⋅u₁ a₂⋅u₂ ]
"""
function gram(a::K,u::L) where {K<:KVector, L<:KVector}
  a = sortbasis(prune(a))
  u = sortbasis(prune(u))
  n = length(a)
  if n != length(u)
    zero(eltype(eltype(a)))
  else
    [aᵢ⋅uⱼ for uⱼ in u for aᵢ in a] |> A->(reshape(A,n,n))
  end
end

import LinearAlgebra: det

"Determinant of Gram matrix formed from k-vectors a and u"
det(a::K,u::L) where {K<:KVector, L<:KVector} = det(gram(a,u))

import LinearAlgebra: norm,normalize,norm_sqr

norm_sqr(k::K) where {K<:KVector} = k⋅reverse(k)
norm(k::K) where {K<:KVector} = sqrt(norm_sqr(k))
normalize(k::K) where {K<:KVector} = k/norm(k)
normalize_safe(k::K) where {K<:KVector} = (kk = norm_sqr(k); abs(kk) > 0 ? k/norm(k) : k)

import Base: cos

"""
    cos(a, b)

cosine of angle between blades a and be spanned by the k-vectors a and b
"""
cos( a::A, b::B ) where { A<:Union{KVector, Blade}, B<:Union{KVector, Blade}} = norm(a⋅b)/(norm(a)*norm(b))

Blades.basis_1blades( k::K ) where {K<:KVector} = basis_1blades(first(k))

"""A KVector formed from all basis 1-vectors set to magnitude 1.0. Useful as a test direction"""
basis_1vector( k::K ) where {K<:KVector} = KVector(one.(basis_1blades(k)))

"""
    coords(k)

coordinate scalar values for ordered k-vector basis of k.
"""
function coords( k::KV ) where {T,K,KV<:KVector{T,K}}
  be = basis_kblades(first(k),K)
  kb = untype.(k)
  [ (i = findfirst(kbᵢ->beᵢ==kbᵢ, kb); isnothing(i) ? zero(T) : scalar(k[i])) for beᵢ in be ]
end

coords( b::Blade ) = coords(KVector(b))

end # module

ext: jl


text:
using ContinuousObservationToyPOMDPs
using SparseSampling
using QMDP
using POMDPs
using BeliefUpdaters
using POMDPPolicies
using ProfileView
using Profile

m = COTigerPOMDP()
qp = solve(QMDPSolver(), m)
b = initialize_belief(updater(qp), initialstate_distribution(m))
@show qp.action_map
@show actionvalues(qp, b)

function f(width)
    opt = POWSSOptions(5, width)
    p = solve(SparseSamplingSolver(opt), m)
    collect(valuepairs(p, b))
end

# @code_warntype observation(m, :left, :listen_1)

@time f(10)

Profile.clear()
@profile f(10)
ProfileView.view()


text:
## PART 3

include("autocorrelation.jl");

println("\nPart 3");
println("\n\tCalculating autocorrelations");

# Create x and y arrays for plotting
xcorr = [0:10:400];
ycorr = zeros( length( xcorr ), 5 );
for a in [1:5]
  # Calculate autocorrelations
  ycorr[:,a] = autocorrelation(M, v, a, xcorr);
  # Normalize by dividing through by Ĉ[1,a]
  ycorr[:,a] /= ycorr[1,a];
end

# Make plots for each alpha
println("\n\tMaking autocorrelation plots and placing in 'plots' vector");
using Gadfly
plots = Plot[];
for i in [1:5]
  # Plot
  push!(plots, Gadfly.plot ( x = xcorr, y = ycorr[:,i], Geom.line, Geom.point,
      Guide.xlabel("Separation"), Guide.ylabel("Autocorrelation"),
      Guide.title(join(["Autocorrelation for Dataset",string(i)]," "))))
end


text:
struct UniformPeriodicReconstructedFV1D{BalanceLaw, T, Fnum, Reconstruction,
                                        EdgeU, Fluxes, Parallel} <: FiniteVolumeSemidiscretisation
    balance_law::BalanceLaw
    meshx::UniformPeriodicMesh1D{T}
    fnum::Fnum
    reconstruction::Reconstruction

    edge_u::EdgeU
    fluxes::Fluxes
    parallel::Parallel

    function UniformPeriodicReconstructedFV1D(
            balance_law::BalanceLaw, meshx::UniformPeriodicMesh1D{T},
            fnum::Fnum, reconstruction::Reconstruction,
            edge_u::EdgeU, fluxes::Fluxes, parallel::Parallel) where {BalanceLaw, T, Fnum, Reconstruction, EdgeU, Fluxes, Parallel}
        @assert typeof(fluxes) <: AbstractArray{variables(balance_law), 1}
        @assert typeof(edge_u) <: AbstractArray{variables(balance_law), 2}
        @assert size(fluxes, 1) == size(edge_u, 2) == numcells(meshx)
        @assert size(edge_u, 1) == 2
        new{BalanceLaw, T, Fnum, Reconstruction, EdgeU, Fluxes, Parallel}(
            balance_law, meshx, fnum, reconstruction, edge_u, fluxes, parallel)
    end
end

function UniformPeriodicReconstructedFV1D(balance_law, meshx, fnum, reconstruction,
                                          parallel=Val{:serial}())
    fluxes = Array{variables(balance_law)}(undef, numcells(meshx))
    edge_u = Array{variables(balance_law)}(undef, 2, numcells(meshx))
    UniformPeriodicReconstructedFV1D(balance_law, meshx, fnum, reconstruction,
                                     edge_u, fluxes, parallel)
end

function Base.show(io::IO, fv::UniformPeriodicReconstructedFV1D)
    print(io, "Finite volume method ",
            "\n  Balance law:    ", fv.balance_law,
            "\n  Mesh:           ", fv.meshx,
            "\n  Numerical flux: ", fv.fnum,
            "\n  Reconstruction: ", fv.reconstruction)
end

function evaluate_coefficients(u, fv::UniformPeriodicReconstructedFV1D,
                               npoints=2*order(fv.reconstruction))
    evaluate_coefficients(u, fv.balance_law, fv.meshx, fv.reconstruction, npoints)
end

@noinline function (semidisc::UniformPeriodicReconstructedFV1D)(du, u, p, t)
    @boundscheck begin
        if size(u) != size(du)
            error("size(u) = $(size(u)) != $(size(du)) = size(du)")
        end
        @assert length(u) == numcells(semidisc.meshx)
        @assert length(u) >= stencil_width(semidisc.reconstruction) - 1
        if eltype(u) != variables(semidisc.balance_law)
            error("eltype(u) == $(eltype(u)) != $(variables(semidisc.balance_law)) == variables(semidisc.balance_law)")
        end
    end

    fill!(du, zero(eltype(du)))
    add_numerical_fluxes!(du, u, semidisc, t)

    nothing
end

function add_numerical_fluxes!(du, u, fv::UniformPeriodicReconstructedFV1D, t)
    @unpack balance_law, meshx, fnum, reconstruction, edge_u, fluxes, parallel = fv

    reconstruct!(edge_u, u, balance_law, meshx, reconstruction, stencil_width_val(reconstruction), parallel)
    compute_numerical_fluxes!(fluxes, edge_u, balance_law, meshx, fnum, parallel)
    compute_du!(du, fluxes, meshx, parallel)
end

function compute_numerical_fluxes!(fluxes, edge_u, balance_law, meshx::UniformPeriodicMesh1D,
                                   fnum, parallel)
    fluxes[1] = fnum(edge_u[2,end], edge_u[1,1], balance_law)
    for edge in 2:length(fluxes)
        fluxes[edge] = fnum(edge_u[2,edge-1], edge_u[1,edge], balance_law)
    end
    nothing
end

function compute_numerical_fluxes!(fluxes, edge_u, balance_law, meshx::UniformPeriodicMesh1D,
                                   fnum, parallel::Val{:threads})
    @inbounds fluxes[1] = fnum(edge_u[2,end], edge_u[1,1], balance_law)
    @inbounds Threads.@threads for edge in 2:length(fluxes)
        fluxes[edge] = fnum(edge_u[2,edge-1], edge_u[1,edge], balance_law)
    end
    nothing
end

function compute_du!(du, fluxes, meshx::UniformPeriodicMesh1D, parallel::Val{:serial})
    @inbounds for cell in 1:numcells(meshx)-1
        du[cell] -= ( fluxes[cell+1] - fluxes[cell] ) / volume(cell, meshx)
    end
    @inbounds du[end] -= ( fluxes[1] - fluxes[end] ) / volume(numcells(meshx), meshx)
    nothing
end

function compute_du!(du, fluxes, meshx::UniformPeriodicMesh1D, parallel::Val{:threads})
    @inbounds Threads.@threads for cell in 1:numcells(meshx)-1
        du[cell] -= ( fluxes[cell+1] - fluxes[cell] ) / volume(cell, meshx)
    end
    @inbounds du[end] -= ( fluxes[1] - fluxes[end] ) / volume(numcells(meshx), meshx)
    nothing
end

################################################################################

function semidiscretise(fv::UniformPeriodicReconstructedFV1D, u₀func, tspan)
    u₀ = compute_coefficients(u₀func, fv.meshx, order(fv.reconstruction))

    ODEProblem(fv, u₀, tspan)
end

"""
    max_dt(t, u, fv::UniformPeriodicReconstructedFV1D, cfl=0.5)

Compute the maximal time step `dt` satisfying the CFL condition `dt <= cfl * dx / 2λ`,
where `dx` is the length of a cell and `λ` the greatest absolute value of the
speed in this cell.
"""
function max_dt(t, u, fv::UniformPeriodicReconstructedFV1D, cfl=0.5)
    @unpack balance_law, meshx = fv
    dt = mapreduce(cell->volume(cell,meshx) / max_abs_speed(u[cell],balance_law),
                    min, cell_indices(meshx), init=typemax(t))
    if dt == Inf
        dt = mapreduce(cell->volume(cell,meshx), min, cell_indices(meshx), init=dt)
    end
    cfl * dt
end


text:
@testset "resampling.jl" begin
    D = [0.3, 0.4, 0.3]
    num_samples = Int(1e6)
    rng = Random.GLOBAL_RNG

    resSystematic = AdvancedPS.resample_systematic(rng, D, num_samples)
    resStratified = AdvancedPS.resample_stratified(rng, D, num_samples)
    resMultinomial = AdvancedPS.resample_multinomial(rng, D, num_samples)
    resResidual = AdvancedPS.resample_residual(rng, D, num_samples)
    AdvancedPS.resample_systematic(rng, D)

    @test sum(resSystematic .== 2) ≈ (num_samples * 0.4) atol = 1e-3 * num_samples
    @test sum(resStratified .== 2) ≈ (num_samples * 0.4) atol = 1e-3 * num_samples
    @test sum(resMultinomial .== 2) ≈ (num_samples * 0.4) atol = 1e-2 * num_samples
    @test sum(resResidual .== 2) ≈ (num_samples * 0.4) atol = 1e-2 * num_samples
end


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

using Markdown
using InteractiveUtils

# ╔═╡ 8f377998-a17a-11eb-1277-ad9277acfc62
md"# Day 18: Operation Order"

# ╔═╡ 1b9a5f93-3633-491a-b6f4-3da9692e7480
md"""
As you look out the window and notice a heavily-forested continent slowly appear over the horizon, you are interrupted by the child sitting next to you. They're curious if you could help them with their math homework.

Unfortunately, it seems like this "math" follows different rules than you remember.

The homework (your puzzle input) consists of a series of expressions that consist of addition (+), multiplication (*), and parentheses ((...)). Just like normal math, parentheses indicate that the expression inside must be evaluated before it can be used by the surrounding expression. Addition still finds the sum of the numbers on both sides of the operator, and multiplication still finds the product.

However, the rules of *operator precedence* have changed. Rather than evaluating multiplication before addition, the operators have the *same precedence*, and are evaluated left-to-right regardless of the order in which they appear.

For example, the steps to evaluate the expression `1 + 2 * 3 + 4 * 5 + 6` are as follows:

```
1 + 2 * 3 + 4 * 5 + 6
  3   * 3 + 4 * 5 + 6
      9   + 4 * 5 + 6
         13   * 5 + 6
             65   + 6
                 71
```

Parentheses can override this order; for example, here is what happens if parentheses are added to form `1 + (2 * 3) + (4 * (5 + 6))`:

```
1 + (2 * 3) + (4 * (5 + 6))
1 +    6    + (4 * (5 + 6))
     7      + (4 * (5 + 6))
     7      + (4 *   11   )
     7      +     44
            51
```

Here are a few more examples:

- `2 * 3 + (4 * 5)` becomes **`26`**.
- `5 + (8 * 3 + 9 + 3 * 4 * 3)` becomes **`437`**
- `5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))` becomes **`12240`**
- `((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2` becomes **`13632`**

Before you can help with the homework, you need to understand it yourself. *Evaluate the expression on each line of the homework; what is the sum of the resulting values?*
"""

# ╔═╡ eb826757-a608-4eaa-bbe3-df732916f11a
function find_matching_index(s, pos)
	level = 1
	if s[pos] != '('
		throw(ArgumentError("Expected position to refer to a '(' character"))
	end
	for (i, c) ∈ enumerate(s[pos+1:end])
		if c == '('
			level += 1
		elseif c == ')'
			level -= 1
		end

		if level == 0
			return i + pos
		end
	end
end

# ╔═╡ d4e1b7fe-4679-40c7-bd77-e80ebcfc747d
ops = Set(['+', '*'])

# ╔═╡ 37dd5443-2b4f-45b7-afbe-429bf47e7477
function parse_expression(s)
	i = 1
	n = length(s)
	current_num = Char[]
	parts = []
	while i <= n
		c = s[i]
		if c == '('
			j = find_matching_index(s, i)
			push!(parts, parse_expression(s[i+1:j-1]))
			i = j+1
		elseif c ∈ ops
			push!(parts, Symbol(c))
			i += 1
		elseif c >= '0' && c <= '9'
			push!(current_num, c)
			i += 1
		else
			if length(current_num) > 0
				push!(parts, parse(Int, String(current_num)))
				empty!(current_num)
			end
			i += 1
		end
	end
	if length(current_num) > 0
		push!(parts, parse(Int, String(current_num)))
	end
	parts
end

# ╔═╡ c1df810b-cf5d-475a-94e9-ad6901a17b1e
parse_expression("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))")

# ╔═╡ c05d4712-f583-4ff5-86dc-b606244461c6
function evaluate(expr)
	n = length(expr)
	if n == 1
		if length(expr[1]) > 1
			return evaluate(expr[1])
		else
			return expr[1]
		end
	elseif n == 2
		throw(ArgumentError("Unexpected 2 argument expression"))
	elseif n == 3
		left, op, right = evaluate(expr[1]), eval(expr[2]), evaluate(expr[3])
		return op(left, right)
	else
		left, op, right = evaluate(expr[1]), eval(expr[2]), evaluate(expr[3])
		res = op(left, right)
		return evaluate(vcat(res, expr[4:end]))
	end
end

# ╔═╡ 062ffb66-a6fb-41fa-8745-d6002b7c94ed
function evaluate(expr::String)
	return evaluate(parse_expression(expr))
end

# ╔═╡ c5353f1b-03b0-4438-8d69-ab5e57678f65
function part1(filename="day18input.txt")
	lines = readlines(filename)
	sum(map(evaluate, lines))
end

# ╔═╡ 0c639002-d70c-48ef-82e9-da6b9881ec8f
part1()

# ╔═╡ 004d6fdf-0be3-480f-8489-2a4cbc52b7d9
md"""
# Part Two

You manage to answer the child's questions and they finish part 1 of their homework, but get stuck when they reach the next section: advanced math.

Now, addition and multiplication have different precedence levels, but they're not the ones you're familiar with. Instead, addition is evaluated before multiplication.

For example, the steps to evaluate the expression `1 + 2 * 3 + 4 * 5 + 6` are now as follows:

```
1 + 2 * 3 + 4 * 5 + 6
  3   * 3 + 4 * 5 + 6
  3   *   7   * 5 + 6
  3   *   7   *  11
     21       *  11
         231
```
Here are the other examples from above:

- `1 + (2 * 3) + (4 * (5 + 6))` still becomes **`51`**.
- `2 * 3 + (4 * 5)` becomes **`46`**.
- `5 + (8 * 3 + 9 + 3 * 4 * 3)` becomes **`1445`**.
- `5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))` becomes **`669060`**.
- `((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2` becomes **`23340`**.
What do you get if you add up the results of evaluating the homework problems using these new rules?
"""

# ╔═╡ 86b86adc-41a9-444a-ba9d-8fa690c2c4b0
function evaluate_additions(expr)
	new_expr = []
	n = length(expr)
	i = 1
	j = 0
	while i <= n
		if expr[i] == :+
			op = eval(expr[i])
			left = new_expr[j]
			right = expr[i+1]
			new_expr[j] = op(left, right)
			i += 1 # jump over the right element
		else
			push!(new_expr, expr[i])
			j += 1
		end
		i += 1
	end
	new_expr
end

# ╔═╡ f9a522f6-d56b-41c7-81d2-99cba876dfab
function evaluate_multiplications(expr)
	new_expr = []
	n = length(expr)
	i = 1
	j = 0
	while i <= n
		if expr[i] == :*
			op = eval(expr[i])
			left = new_expr[j]
			right = expr[i+1]
			new_expr[j] = op(left, right)
			i += 1 # jump over the right element
		else
			push!(new_expr, expr[i])
			j += 1
		end
		i += 1
	end
	new_expr
end

# ╔═╡ 77e23e48-1c9f-4cfb-9056-e357674d1f92
function evaluate_with_precedence(expr)
	if typeof(expr) == Symbol
		return expr
	end
	n = length(expr)
	if n == 1
		if typeof(expr[1]) == Symbol || length(expr[1]) > 1
			return evaluate_with_precedence(expr[1])
		else
			return expr[1]
		end
	elseif n == 2
		throw(ArgumentError("Unexpected 2 argument expression"))
	elseif n == 3
		left = evaluate_with_precedence(expr[1])
		op = eval(expr[2])
		right = evaluate_with_precedence(expr[3])
		return op(left, right)
	else
		simplified = map(evaluate_with_precedence, expr)
		added = evaluate_additions(simplified)
		return evaluate_multiplications(added)[1]
	end
end

# ╔═╡ 88c9a70d-0dc1-43e4-a809-0670cf45fdb9
function evaluate_with_precedence(expr::String)
	evaluate_with_precedence(parse_expression(expr))
end

# ╔═╡ 263700b4-fc9c-4191-9807-fc5f7095826c
function part2(filename="day18input.txt")
	lines = readlines(filename)
	sum(map(evaluate_with_precedence, lines))
end

# ╔═╡ 4f266104-128a-4812-a73a-81b1564c55bd
part2()

# ╔═╡ Cell order:
# ╟─8f377998-a17a-11eb-1277-ad9277acfc62
# ╟─1b9a5f93-3633-491a-b6f4-3da9692e7480
# ╠═eb826757-a608-4eaa-bbe3-df732916f11a
# ╠═d4e1b7fe-4679-40c7-bd77-e80ebcfc747d
# ╠═37dd5443-2b4f-45b7-afbe-429bf47e7477
# ╠═c1df810b-cf5d-475a-94e9-ad6901a17b1e
# ╠═c05d4712-f583-4ff5-86dc-b606244461c6
# ╠═062ffb66-a6fb-41fa-8745-d6002b7c94ed
# ╠═c5353f1b-03b0-4438-8d69-ab5e57678f65
# ╠═0c639002-d70c-48ef-82e9-da6b9881ec8f
# ╟─004d6fdf-0be3-480f-8489-2a4cbc52b7d9
# ╠═86b86adc-41a9-444a-ba9d-8fa690c2c4b0
# ╠═f9a522f6-d56b-41c7-81d2-99cba876dfab
# ╠═77e23e48-1c9f-4cfb-9056-e357674d1f92
# ╠═88c9a70d-0dc1-43e4-a809-0670cf45fdb9
# ╠═263700b4-fc9c-4191-9807-fc5f7095826c
# ╠═4f266104-128a-4812-a73a-81b1564c55bd


text:
table_σ_o31D_jx = [[4.842, 4.922, 5.071, 5.228, 6.040, 6.803, 7.190, 7.549, 9.000, 8.989, 8.929, 9.000, 8.901, 4.130, 8.985*0.1, 6.782*0.1, 0.000, 0.000]*0.1 [4.843, 4.922, 5.072, 5.229, 6.040, 6.802, 7.189, 7.549, 9.000, 8.989, 8.929, 9.000, 8.916, 4.656, 1.417, 6.995*0.1, 0.000, 0.000]*0.1 [4.843, 4.922, 5.072, 5.229, 6.040, 6.805, 7.189, 7.550, 9.000, 8.989, 8.929, 9.000, 8.967, 5.852, 2.919, 7.943, 0.000, 0.000]*0.1]
# first column: 200K
# second column: 260K
# third column: 320K

function calc_J_o31D(actinic_flux,t)
	r = zeros(18)
	for i in 1:18
        if t <= 200
		    r[i] = actinic_flux[i] * table_σ_o31D_jx[i,1] * 1.000
        elseif t > 320
            r[i] = actinic_flux[i] * table_σ_o31D_jx[i,3] * 1.000
        elseif 200 < t <= 260
            k = (t-200)/(260-200)
            σ = (table_σ_o31D_jx[i,2]-table_σ_o31D_jx[i,1])*k + table_σ_o31D_jx[i,1]
            r[i] = actinic_flux[i] * σ * 1.000
        elseif 260 < t <= 320
            k = (t-260)/(320-260)
            σ = (table_σ_o31D_jx[i,3]-table_σ_o31D_jx[i,2])*k + table_σ_o31D_jx[i,2]
            r[i] = actinic_flux[i] * σ * 1.000
        end
	end
	return r
end
# for all 18 bins
# for one grid? -> given actinic flux at one bin
# O3 -> O2 + O(1D)
# superfast: O3 -> 2OH  include O(1D) + H2O -> 2OH

table_σ_H2O2_jx =[[2.325, 4.629, 5.394, 5.429, 4.447, 3.755, 3.457, 3.197, 5.346*0.1, 4.855*0.1, 3.423*0.1, 8.407*0.01, 5.029*0.01, 3.308*0.01, 2.221*0.01, 8.598*0.001, 1.807*0.0001,0]*10^-19.0 [2.325, 4.629, 5.394, 5.429, 4.447, 3.755, 3.457, 3.197, 5.465*0.1, 4.966*0.1, 3.524*0.1, 9.354*0.01, 5.763*0.01, 3.911*0.01, 2.718*0.01, 1.138*0.01, 2.419*0.0001, 0]*10^-19.0]
# first column : 200K
# second column: 300K

function calc_J_H2O2(actinic_flux,t)
	r = zeros(18)
	for i in 1:18
        if t <= 200
		    r[i] = actinic_flux[i] * table_σ_H2O2_jx[i,1] * 1.000
        elseif t >= 300
            r[i] = actinic_flux[i] * table_σ_H2O2_jx[i,2] * 1.000
        elseif 200 < t < 300
            k = (t-200)/(300-200)
            r[i] = (table_σ_H2O2_jx[i,2]-table_σ_H2O2_jx[i,1])*k + table_σ_H2O2_jx[i,1]
        end
	end
	return r
end
# for all bins
# for one grid? -> given actinic flux at one bin
# H2O2 -> OH + OH same with superfast

table_σ_CH2Oa_jx = [[0.0, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 3.143*0.1, 1.021, 1.269, 2.323, 2.498, 1.133, 2.183, 4.746*0.1, 0.000, 0.000]*10^-20.0 [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 3.147*0.1, 1.018, 1.266, 2.315, 2.497, 1.131, 2.189, 4.751*0.1, 0.000, 0.000]*10^-20.0]
# first column : 223K
# second column: 298K

function calc_J_CH2Oa(actinic_flux, t)
	r = zeros(18)
	for i in 1:18
        if t <= 223
		    r[i] = actinic_flux[i] * table_σ_CH2Oa_jx[i,1] * 1.000
        elseif t >= 298
            r[i] = actinic_flux[i] * table_σ_CH2Oa_jx[i,2] * 1.000
        elseif 223 < t < 298
            k = (t-223)/(298-223)
            r[i] = (table_σ_CH2Oa_jx[i,2]-table_σ_CH2Oa_jx[i,1])*k + table_σ_CH2Oa_jx[i,1]
        end
	end
	return r
end
# for all bins
# for one grid? -> given actinic flux at one bin
# CH2O -> H + HO2 + CO
# superfast: CH2O -> CO + 2HO2; CH2O -> CO

table_σ_CH3OOH_jx = [0.000, 0.000, 0.000, 0.000, 0.000, 3.120, 2.882, 2.250, 2.716*0.1, 2.740*0.1, 2.143*0.1, 5.624*0.01, 3.520*0.01, 2.403*0.01, 1.697*0.01, 7.230*0.001, 6.973*0.0001, 0.000]*10^-19
#298K

function calc_J_CH3OOH(actinic_flux,t)
	r = zeros(18)
	for i in 1:18
		r[i,1] = actinic_flux[i] * table_σ_CH3OOH_jx[i] * 1.000
	end
	return r
end
# for all bins
# for one grid? -> given actinic flux at one bin
# CH3OOH -> OH + HO2 + CH2O same with superfast

table_σ_NO2_jx = [[0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 1.835*0.1, 4.693*0.1, 7.705*0.1, 1.078, 1.470, 1.832, 2.181, 3.138, 4.321, 1.386*0.001]*10^-19.0 [0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 0.000, 2.313*0.1, 4.694*0.1, 7.553*0.1, 1.063, 1.477, 1.869, 2.295, 3.448, 4.643, 4.345*0.001]*10^-19]
# first column: 200K
# second column:294K

function calc_J_NO2(actinic_flux, t)
    r = zeros(18)
	for i in 1:18
		if t <= 200
		    r[i] = actinic_flux[i] * table_σ_NO2_jx[i,1] * 1.000
        elseif t >= 294
            r[i] = actinic_flux[i] * table_σ_NO2_jx[i,2] * 1.000
        elseif 200 < t < 294
            k = (t-200)/(294-200)
            r[i] = (table_σ_NO2_jx[i,2]-table_σ_NO2_jx[i,1])*k + table_σ_NO2_jx[i,1]
        end
	end
	return r
end
# for all bins
#for one grid? -> given actinic flux at one bin
# NO2 -> NO + O
# superfast: NO2 -> NO + O3  include O + O2 -> O3

act = [0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1]
# sample_actinic_flux, not realistic

#Unit Test 1 -- O3 -> O2 + O(1D)
function u1()
    n = 0
    m = ""
    act = [0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1]
    function calc_J_o31Du1(actinic_flux)
        r = zeros(18,3)
        for i in 1:18, j in 1:3
            r[i,j] = actinic_flux[i] * table_σ_o31D_jx[i,j] * 1.000
        end
        return r
    end
    u_1 = calc_J_o31Du1(act)
    test1 = [calc_J_o31D(act,200) calc_J_o31D(act,260) calc_J_o31D(act,320)]
	if isapprox(test1, u_1, atol=0.001) == true
	    m = "unit test 1 passed"
	else
		m = "unit test 1 failed"
	end
return m
end

#Unit Test 2 -- H2O2 -> OH + OH
function u2()
    n = 0
    m = ""
    act = [0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1]
    function calc_J_H2O2u2(actinic_flux)
        r = zeros(18,2)
        for i in 1:18, j in 1:2
            r[i,j] = actinic_flux[i] * table_σ_H2O2_jx[i,j] * 1.000
        end
        return r
    end
    u_2 = calc_J_H2O2u2(act)
    test2 = [calc_J_H2O2(act,200) calc_J_H2O2(act,300)]
	if isapprox(test2, u_2, atol=0.001) == true
	    m = "unit test 2 passed"
	else
		m = "unit test 2 failed"
	end
return m
end

#Unit Test 3 -- CH2O -> H + HO2 + CO
function u3()
    n = 0
    m = ""
    act = [0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1]
    function calc_J_CH2Oau3(actinic_flux)
        r = zeros(18,2)
        for i in 1:18, j in 1:2
            r[i,j] = actinic_flux[i] * table_σ_CH2Oa_jx[i,j] * 1.000
        end
        return r
    end
    u_3 = calc_J_CH2Oau3(act)
    test3 = [calc_J_CH2Oa(act,223) calc_J_CH2Oa(act,298)]
	if isapprox(test3, u_3, atol=0.001) == true
	    m = "unit test 3 passed"
	else
		m = "unit test 3 failed"
	end
return m
end

##Unit Test 4 -- CH3OOH -> OH + HO2 + CH2O
function u4()
    n = 0
    m = ""
    act = [0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1]
    function calc_J_CH3OOHu4(actinic_flux)
        r = zeros(18)
        for i in 1:18,
            r[i,1] = actinic_flux[i] * table_σ_CH3OOH_jx[i,1] * 1.000
        end
        return r
    end
    u_4 = calc_J_CH3OOHu4(act)
    test4 = calc_J_CH3OOH(act,350)
	if isapprox(test4, u_4, atol=0.001) == true
	    m = "unit test 4 passed"
	else
		m = "unit test 4 failed"
	end
return m
end

#Unit Test 5 -- NO2 -> NO + O3
function u5()
    n = 0
    m = ""
    act = [0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1]
    function calc_J_NO2u5(actinic_flux)
        r = zeros(18,2)
        for i in 1:18, j in 1:2
            r[i,j] = actinic_flux[i] * table_σ_NO2_jx[i,j] * 1.000
        end
        return r
    end
    u_5 = calc_J_NO2u5(act)
    test5 = [calc_J_NO2(act,200) calc_J_NO2(act,294)]
	if isapprox(test5, u_5, atol=0.001) == true
	    m = "unit test 5 passed"
	else
		m = "unit test 5 failed"
	end
return m
end


text:
function place_grid_order!(config, trade_grid, price_str, client, bot_path)
    #
    ## Test Bot?
    test_bot = config["testing"] == "1" ? true : false

    info = trade_grid[price_str]

    if info["side"] == client.SIDE_BUY
        #
        ## Calculate BUY qty depending on thes. volume
        qty = Pinguu.get_buy_value(config, trade_grid) / parse(Float64, price_str)
        qty = floor(qty, digits=config["precision"])
    else
        #
        ## Get SELL qty
        qty = typeof(info["qty"]) != Float64 ? parse(Float64, info["qty"]) : info["qty"]
    end
    tmp = Dict(
        "status" => "check",
        "side" => info["side"],
        "price" => price_str,
        "type" => info["type"],
        "qty" => qty
    )

    current_volume = Pinguu.get_current_volume(config)

    if info["side"] == client.SIDE_BUY
        tmp["current_volume"] = current_volume
    end
    #
    ## Sell of extra coin volume, take buy_order
    if info["type"] == "normal" && info["side"] == client.SIDE_SELL
        tmp["buy_order"] = info["buy_order"]
        if haskey(info, "current_volume")
            tmp["current_volume"] = info["current_volume"]
        else
            tmp["current_volume"] = current_volume
        end
    end

    if test_bot
        #
        ## PLACE BUY TESTING
        order = try client.create_test_order(
                symbol=config["symbol"],
                side=info["side"],
                type=client.ORDER_TYPE_LIMIT,
                timeInForce=client.TIME_IN_FORCE_GTC,
                quantity=tmp["qty"],
                price=tmp["price"]
            )
            catch
            return
        end
        tmp["order"] = Dict(
            "price" => price_str,
            "orderId" => 42,
            "executedQty" => "0.00",
            "origQty" => tmp["qty"],
            "side" => tmp["side"],
            "type" => "limit",
            "symbol" => config["symbol"],
            "status" => "NEW",
            "cummulativeQuoteQty" => "0.00000000",
            "executedQty" => "0.00",
            "time" => string(now())
        )
        ##
        #
    else
        #
        ## PLACE ORDER
        if tmp["side"] == client.SIDE_BUY
            order = try client.order_limit_buy(
                    symbol=config["symbol"],
                    quantity=tmp["qty"],
                    price=tmp["price"]
                );
            catch
                return
            end
        else
            order = try client.order_limit_sell(
                    symbol=config["symbol"],
                    quantity=tmp["qty"],
                    price=tmp["price"]
                );
            catch
                return
            end
        end
        tmp["order"] = order
        tmp["order"]["time"] = string(now())
        ##
        #
    end

    trade_grid[price_str] = tmp
    open(joinpath(bot_path, "trade_grid.json"), "w") do f
        JSON.print(f, trade_grid, 4)
    end
end
export place_grid_order!


text:
"""
Describes a function used in `runsimulation!` with:

Inputs:

* `aircraft::Aircraft` : aircraft system struct
* `parameters <: Parameters` : inherits from the `Parameters` type; object containing data required by `<: Action` functions
* `freestream::Freestream` : freestream object
* `environment::Environment` : environment object
* `step_range::AbstractArray` : range of steps defining simulation
* `stepi::Int` : index of the current step
* `step_symbol::String` : defines the step, e.g. `alpha` or `time`

Modifies:

* `aircraft::Aircraft`
* `parameters <: Parameters`

Outputs:

* `flag::Bool` : true if an action experiences errors

Dispatch is as follows:

```
function myaction!(aircraft, parameters, freestream, environment, step_range, stepi, step_symbol) <: Action
    ...
    return flag
end
```
"""
abstract type Action end

max_stars_repo_name: byuflowlab/AircraftSystems

max_issues_repo_name: byuflowlab/AircraftSystems

max_forks_repo_name: byuflowlab/AircraftSystems


text:
function tron_dcopy(n,dx,incx,dy,incy)
    # Copies a vector, x, to a vector, y.
    # It uses unrolled loops for increments equal to one.
    # Jack Dongarra, LINPACK, 3/11/78.
    # Modified 12/3/93, array(1) declarations changed to array(*)

    if n <= 0
        return
    end
    if incx != 1 || incy != 1

        # Code for unequal increments or equal increments
        # not equal to 1

        ix = 1
        iy = 1
        if incx < 0
            ix = (-n+1)*incx + 1
        end
        if incy < 0
            iy = (-n+1)*incy + 1
        end
        for i=1:n
            dy[iy] = dx[ix]
            ix = ix + incx
            iy = iy + incy
        end

        return
    else

        # Code for both increments equal to 1
        # Clean-up loops

        m = mod(n,7)
        if m != 0
            for i=1:m
                dy[i] = dx[i]
            end
            if n < 7
                return
            end
        end
        mp1 = m + 1
        for i=mp1:7:n
            dy[i] = dx[i]
            dy[i + 1] = dx[i + 1]
            dy[i + 2] = dx[i + 2]
            dy[i + 3] = dx[i + 3]
            dy[i + 4] = dx[i + 4]
            dy[i + 5] = dx[i + 5]
            dy[i + 6] = dx[i + 6]
        end

        return
    end
end

if isequal(BLAS_LIBRARY, :OpenBlas)
    dcopy(n,dx,incx,dy,incy) = copyto!(dy, 1, dx, 1, n)
else
    dcopy(n,dx,incx,dy,incy) = tron_dcopy(n,dx,incx,dy,incy)
end

@inline function dcopy(n::Int,dx::CuDeviceArray{Float64,1},incx::Int,
                       dy::CuDeviceArray{Float64,1},incy::Int)
    tx = threadIdx().x
    ty = threadIdx().y

    # Ignore incx and incy for now.
    if tx <= n && ty == 1
        @inbounds dy[tx] = dx[tx]
    end
    CUDA.sync_threads()

    return
end


text:
"""
    Q-table is arranged as (Actions × States)

              s1  s2 ⋅ ⋅ ⋅
        a1 ┌ q11 q12 ⋅ ⋅ ⋅ ┐
        a2 │ q21 q22 ⋅ ⋅ ⋅ │
         ⋅ │  ⋅   ⋅  ⋅     │
         ⋅ │  ⋅   ⋅    ⋅   │
         ⋅ └  ·   ⋅      ⋅ ┘

    , which is updated by
        nextA = argmax Q(A, nextS) ∀ A
    for
        Q(prevA, prevS) += α * [prevR + γ * Q(nextA, nextS) - Q(prevA, prevS)]
"""
function update_Qtable!(Q::MatI, prevS::Int, prevA::Int, prevR::Real, nextS::Int, α::Real, γ::Real) # @code_warntype ✓
    nextA = pick_action(Q, nextS)
    @inbounds Q[prevA, prevS] *= 1.0 - α
    @inbounds Q[prevA, prevS] += α * (prevR + γ * Q[nextA, nextS])
    return nothing
end
"""
    The Q-learning agent picks an action of maximum Q-value of current state.

    params:
    -------
    * Q   := Q-table ∈ (Actions × States)
    * Sdx := index of current state
    * Adx := index of picked action (return)
"""
function pick_action(Q::MatI, Sdx::Int) # @code_warntype ✓
    Adx = 0
    tmp = -Inf
    @inbounds for idx in axes(Q, 1)
        _tmp = Q[idx, Sdx]
        if _tmp ≠ _tmp || tmp < _tmp
            tmp = _tmp
            Adx = idx
        end
    end
    return Adx
end
"""
    The Q-learning agent picks an action based on ε-greedy strategy.

    params:
    -------
    * Q   := Q-table ∈ (States × Actions)
    * Sdx := index of current state
    * Adx := index of picked action (return)
    * ε   := greedy factor for exploration, ε ∈ (0, 1], ε(t = 0) = 1
"""
function pick_action(Q::MatI, Sdx::Int, ε::Real) # @code_warntype ✓
    Adx = pick_action(Q, Sdx)
    rand() < ε ? sample(axes(Q, 1), Adx) : Adx
end


text:
using PolynomialBasis
using ImplicitDomainQuadrature
using Revise
using CutCell
include("../../../../test/useful_routines.jl")

polyorder = 2
numqp = required_quadrature_order(polyorder)
nelmts = 1
width = 1.0
x0 = [0.0, 0.0]
corner = [0.5, 0.5]

basis = TensorProductBasis(2, polyorder)
mesh = CutCell.Mesh(x0, [width, width], [nelmts, nelmts], basis)
levelset = InterpolatingPolynomial(1, basis)
levelsetcoeffs =
    CutCell.levelset_coefficients(x -> -corner_distance_function(x, corner), mesh)

cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
interfacequads = CutCell.InterfaceQuadratures(levelset,levelsetcoeffs,cutmesh,numqp)


text:
module TestFreeDiagrams
using Test

using Catlab.Theories, Catlab.CategoricalAlgebra
using Catlab.CategoricalAlgebra.FinSets: FinSet, FinFunction

A, B, C, D = Ob(FreeCategory, :A, :B, :C, :D)

# Diagrams of flexible shape
############################

# General free diagrams
#----------------------

f, g, h = Hom(:f, A, C), Hom(:g, B, C), Hom(:h, A, B)
diagram = FreeDiagram([A,B,C], [(f,1,3),(g,2,3),(h,1,2)])
@test ob(diagram) == [A,B,C]
@test hom(diagram) == [f,g,h]
@test (src(diagram), tgt(diagram)) == ([1,2,1], [3,3,2])
@test_throws Exception FreeDiagram([A,B,C], [(f,1,2),(g,2,3),(h,1,2)])

# Bipartite free diagrams
#------------------------

bd = BipartiteFreeDiagram([A,B], [C], [(f,1,1),(g,2,1)])
@test (ob₁(bd), ob₂(bd)) == ([A,B], [C])
@test hom(bd) == [f,g]
@test (src(bd), tgt(bd)) == ([1,2], [1,1])
@test FreeDiagram(bd) == FreeDiagram([A,B,C], [(f,1,3),(g,2,3)])

# Diagrams of fixed shape
#########################

# Empty diagrams.
@test isempty(EmptyDiagram{FreeCategory.Ob}())

# Object pairs.
pair = ObjectPair(A,B)
@test length(pair) == 2
@test (first(pair), last(pair)) == (A, B)

# Discrete diagrams.
discrete = DiscreteDiagram([A,B,C])
@test length(discrete) == 3
@test ob(discrete) == [A,B,C]

diagram = FreeDiagram(discrete)
@test ob(diagram) == [A,B,C]
@test isempty(hom(diagram))

# Spans.
f, g = Hom(:f, C, A), Hom(:g, C, B)
span = Span(f,g)
@test apex(span) == C
@test legs(span) == [f,g]
@test feet(span) == [A,B]
@test (left(span), right(span)) == (f, g)

f = Hom(:f, A, A)
@test legs(Span(id(A), f)) == [id(A),f]

f = Hom(:f, A, B)
@test_throws ErrorException Span(f,g)

# Multispans.
f, g, h = Hom(:f, C, A), Hom(:g, C, B), Hom(:h, C, A)
span = Multispan([f,g,h])
@test apex(span) == C
@test legs(span) == [f,g,h]
@test feet(span) == [A,B,A]

diagram = BipartiteFreeDiagram(span)
@test (ob₁(diagram), ob₂(diagram)) == ([C], [A,B,A])
@test hom(diagram) == [f,g,h]
@test (src(diagram), tgt(diagram)) == ([1,1,1], [1,2,3])

diagram = FreeDiagram(span)
@test ob(diagram) == [C,A,B,A]
@test hom(diagram) == [f,g,h]
@test (src(diagram), tgt(diagram)) == ([1,1,1], [2,3,4])

span = Multispan([ id(FinSet(2)) for i in 1:3 ])
span = bundle_legs(span, [1, (2,3)])
@test apex(span) == FinSet(2)
@test codom.(legs(span)) == [FinSet(2), FinSet(4)]

# Cospans.
f, g = Hom(:f, A, C), Hom(:g, B, C)
cospan = Cospan(f,g)
@test apex(cospan) == C
@test legs(cospan) == [f,g]
@test feet(cospan) == [A,B]
@test (left(cospan), right(cospan)) == (f, g)

f = Hom(:f, A, A)
@test legs(Cospan(f, id(A))) == [f,id(A)]

f = Hom(:f, A ,B)
@test_throws ErrorException Cospan(f,g)

# Multicospans.
f, g, h = Hom(:f, A, C), Hom(:g, B, C), Hom(:h, A, C)
cospan = Multicospan([f,g,h])
@test apex(cospan) == C
@test legs(cospan) == [f,g,h]
@test feet(cospan) == [A,B,A]

diagram = BipartiteFreeDiagram(cospan)
@test (ob₁(diagram), ob₂(diagram)) == ([A,B,A], [C])
@test hom(diagram) == [f,g,h]
@test (src(diagram), tgt(diagram)) == ([1,2,3], [1,1,1])

diagram = FreeDiagram(cospan)
@test ob(diagram) == [A,B,A,C]
@test hom(diagram) == [f,g,h]
@test (src(diagram), tgt(diagram)) == ([1,2,3], [4,4,4])

cospan = Multicospan([FinFunction([i],3) for i in 1:3])
cospan = bundle_legs(cospan, [(1,3), 2])
@test apex(cospan) == FinSet(3)
@test legs(cospan) == [FinFunction([1,3], 3), FinFunction([2], 3)]

# Parallel pairs.
f, g = Hom(:f, A, B), Hom(:g, A, B)
pair = ParallelPair(f,g)
@test (dom(pair), codom(pair)) == (A, B)
@test (first(pair), last(pair)) == (f, g)

# Parallel morphisms.
f, g, h = Hom(:f, A, B), Hom(:g, A, B), Hom(:h, A, B)
para = ParallelMorphisms([f,g,h])
@test (dom(para), codom(para)) == (A, B)
@test hom(para) == [f,g,h]

diagram = BipartiteFreeDiagram(para)
@test (ob₁(diagram), ob₂(diagram)) == ([A], [B])
@test hom(diagram) == [f,g,h]
@test (src(diagram), tgt(diagram)) == ([1,1,1], [1,1,1])

diagram = FreeDiagram(para)
@test ob(diagram) == [A,B]
@test hom(diagram) == [f,g,h]
@test (src(diagram), tgt(diagram)) == ([1,1,1], [2,2,2])

end


text:
"""
    GraphNode{T<:Int64} <: AbstractDiscreteWindow{T}

Structure with unique field `idx` representing the index of the vertex of a graph.
"""
struct GraphNode{T<:Int64} <: AbstractDiscreteWindow{T}
    # Corner
    idx::T
end


text:
    # Local compute version

# add more julia processes
nprocs() < 5 ? addprocs(5-nprocs()) : nothing

using Caesar, RoME, Distributions
using YAML, JLD, HDF5
# for drawing
using RoMEPlotting, Gadfly

# const IIF = IncrementalInference

include(joinpath(Pkg.dir("Caesar"),"examples","wheeled","racecar","racecarUtils.jl"))

cfg = loadConfig()

# datafolder = ENV["HOME"]*"/data/racecar/smallloop/"
datafolder = ENV["HOME"]*"/data/racecar/thursday/"
@load datafolder*"tag_det_per_pose.jld" tag_bag
# tag_bag[0]

# Figure export folder
currdirtime = now()
imgdir = joinpath(ENV["HOME"], "Pictures", "racecarimgs", "$(currdirtime)")
mkdir(imgdir)

# Factor graph construction
N = 75
fg = initfg()

psid = 0
pssym = Symbol("x$psid")
# first pose with zero prior
addVariable!(fg, pssym, Pose2)
addFactor!(fg, [pssym], PriorPose2(MvNormal(zeros(3),diagm([0.01;0.01;0.001].^2))))

addApriltags!(fg, pssym, tag_bag[psid])

writeGraphPdf(fg)

# quick solve as sanity check
tree = wipeBuildNewTree!(fg, drawpdf=true)
inferOverTree!(fg,tree, N=N)

#
# plotKDE(getVertKDE(fg, :x0), dims=[1;2])
# plotKDE(getVertKDE(fg, :x0), dims=[1;2])

delete!(tag_bag[61], 18)
delete!(tag_bag[77], 18)
delete!(tag_bag[86], 16)
delete!(tag_bag[110], 6)
delete!(tag_bag[123], 0)
delete!(tag_bag[124], 16)

@async run(`evince bt.pdf`)

prev_psid = 0
# add other positions
for psid in 1:1:115 #length(tag_bag)
  addnextpose!(fg, prev_psid, psid, tag_bag[psid])
  # writeGraphPdf(fg)
  if psid % 25 == 0 || psid == 115
    tree = wipeBuildNewTree!(fg, drawpdf=true)
    inferOverTree!(fg,tree, N=N)

    pl = drawPosesLandms(fg, spscale=0.1, drawhist=false,   meanmax=:mean,xmin=-3,xmax=6,ymin=-5,ymax=2);
    Gadfly.draw(PNG(joinpath(imgdir,"x$(psid).png"),30cm, 25cm),pl)
    pl = drawPosesLandms(fg, spscale=0.1,   meanmax=:mean,xmin=-3,xmax=3,ymin=-2,ymax=2);
    Gadfly.draw(PNG(joinpath(imgdir,"hist_x$(psid).png"),30cm, 25cm),pl)
  end
  prev_psid = psid
end

IncrementalInference.savejld(fg, file="racecar_fg_$(currdirtime).jld")

# fg, = IncrementalInference.loadjld(file="30jul18_9AM.jld")
# addFactor!(fg, [:l5;:l16], Point2Point2Range(Normal(0.5,1.0)))
# ls(fg,:l5)

const KDE = KernelDensityEstimate

fid = open("results.csv","w")
for sym in [ls(fg)[1]...;ls(fg)[2]...]
  p = getVertKDE(fg, sym)
  val = string(KDE.getKDEMax(p))
  println(fid, "$sym, $(val[2:(end-1)])")
end
close(fid)

0

# debugging
#
# using IncrementalInference
# isInitialized(fg,:x0)
# IncrementalInference.doautoinit!(fg, :x0, N=N)
#
# x0l0 = ls(fg, :l0)[1]
# approxConv(fg, x0l0, :l0)
#
#
#
# fct = getVert(fg, x0l0, nt=:fnc)
# ppbr = getfnctype(fct)
#
# X0 = getVal(fg, :x0)
# L0 = getVal(fg, :l0)
#
#
#
#
#
#
#
# IncrementalInference.doautoinit!(fg, :l0, N=N)
# # IncrementalInference.doautoinit!(fg, :l1, N=N)
#
#
# isInitialized(fg, :l0)
#
#
# X1 = getVal(fg,:x0)
#
# plot(x=X1[1,:],y=X1[2,:], Geom.hexbin)
#

#

#


text:
using CSV
using DataFrames

function export_losses_agents(iteration_info_agents, step_sizes, file_path)

    losses = [
        [iter_info.objective for iter_info in info_agent] for
        info_agent in iteration_info_agents
    ]
    iter_max = iteration_info_agents[1][end].iteration
    # df = DataFrame([[0:iter_max]; losses], :auto)
    df = DataFrame(losses, :auto)
    step_sizes = map(string, step_sizes)
    rename!(df, step_sizes)
    # push!(df, [["iter"]; step_sizes])
    CSV.write(file_path, df)
end

function export_statistics(solver_output, file_path)
    df = DataFrame(
        iter = [iter_stats.iteration for iter_stats in solver_output.iteration_stats],
        objective = [iter_stats.objective for iter_stats in solver_output.iteration_stats],
        step_size = [iter_stats.step_size for iter_stats in solver_output.iteration_stats],
        null = [iter_stats.is_null for iter_stats in solver_output.iteration_stats],
    )
    CSV.write(file_path, df)
end


text:
using Test

open("learn_IO.jl", "r") do io
    line = readline(io)
    @show line

    for line in eachline(io)

    end

    # See what happens to readline at eof
    line = readline(io)
    @test "" == line
end


text:
module TestStemmer
    using Base.Test
    using Languages
    using TextAnalysis
    using Compat

    algs = stemmer_types()
    @test !isempty(algs)

    for alg in algs
        stmr = Stemmer(alg)
        TextAnalysis.release(stmr)
    end

    test_cases = @compat Dict{Compat.ASCIIString, Any}(
        "english" => @compat Dict{AbstractString, AbstractString}(
            "working" => "work",
            "worker" => "worker",
            "aβc" => "aβc",
            "a∀c" => "a∀c"
        )
    )

    for (alg, test_words) in test_cases
        stmr = Stemmer(alg)
        for (n,v) in test_words
            @test v == stem(stmr, n)
        end
    end
end


text:
"""
## Indexing
`pa` is an instance of `SeisInvExpt`.
* `pa.mediumm` : inverted medium on the modeling grid (after `update!`)
* `pa.mediumi` : inverted medium on the inversion grid (after `update!`)
* `pa.ageom` : acquisition
* `pa.srcwav`: `srcwav` used to create `pa`
* `pa[:dobs]` : observed data
* `pa[:dcal]` : modeled synthetic data
"""
function Base.getindex(pa::PFWI, s::Symbol)
	@assert s in [:dobs, :dcal]
	if(s==:dobs)
		return pa.paTD.y
	elseif(s==:dcal)
		return pa.paTD.x
	end

	#=
	if(s==:snaps)
		nzd,nxd=length.(pa.c.model.mgrid)
		snaps_all=zeros(nzd,nxd,length(pa.c.itsnaps))
		ip=findall(in.(iss,pa.sschunks))[1]
		issp=findall(x->x==iss,pa.sschunks[ip])[]
		@sync begin
			p=procs(pa.p)[ip]
			@sync remotecall_wait(p) do
				pap=localpart(pa.p)
				for i in eachindex(snaps_all)
					snaps_all[i]=pap.ss[issp].snaps[i]
				end
			end
		end
		return snaps_all
	elseif(s==:data)
		(iss==0) ?  (return pa.c.data[1]) :  (return pa.c.data[1][iss])
	elseif(s==:ageom)
		return pa.c.ageom
	elseif(s==:srcwav)
		return pa.c.srcwav
	end
	=#
end


text:
function genData(m::Weitzman,nCons::Int,nProdTot::Int,
					nProdCons::Int; fname::String = "" , zsMat = []  )
	Random.seed!(m.seed)
	# Initial checks
	@argcheck round((nProdCons-1-m.nA0)/m.nd)==(nProdCons-1-m.nA0)/m.nd  "Number of products (per consumer) does not match m.nA0 and s.nd combo."

	# Search cost specification
	cfun = 	getCfun(m.cfun)

	# Consumer index into full data
	seli = [ 1+ii*nProdCons:(ii+1)*nProdCons for ii = 0:nCons-1]

	# Products' positions (same for all consumers)
	pos = fill([zeros(Int64,1+m.nA0) ; repeat(collect(Int64,1:(nProdCons-1-m.nA0)/m.nd),inner=m.nd)],nCons)

	# Generate products
	pid,chars = genProducts(m::Model,seli,nProdTot,nProdCons)

	# Gather data in NamedTuple for convenience
	d = (seli = seli,pid = vcat(pid...),chars = chars,pos = vcat(pos...))

	# Get ξ value
	ξ   = isempty(zsMat) ?
			[calcZs(0.0,m.dE,cfun(m.cs + m.cd * pp))[1] for pp in 0:nProdCons] :
			[calcZs(zsMat,cfun(m.cs + m.cd * pp)) for pp in 0:nProdCons]

	# Draw random shocks
	v = rand(m.dV,size(d.chars,1))
	e   = rand(m.dE,size(d.chars,1))

	# Get utilities and search values
	u = zeros(size(d.chars,1))
	zs = similar(u)
	fillUZS!(u,zs,m,d,m.β,ξ,e,v)

	path,purch,search = genSearchPaths(m,d,nProdCons,zs,u)

	path = [path[path[:,i].>0,i] for i =1:nCons]

	# Gather again data
	d = (seli = seli,pid = vcat(pid...),chars = chars,pos = vcat(pos...),path = path, purch = purch, search = search, u=u , zs = zs)

	if !isempty(fname)
		@info "Saved data in bson format to $fname."
		@save fname d
	end

	return d
end

###################################################
function genSearchPaths(m::Weitzman,d,maxProd,zs,u)
	n = length(d.seli)

	# Same for all consumers
	path = zeros(Int64,maxProd,n)
	purch = fill(0,n)
	search = fill(false,maxProd,n)

  @inbounds Threads.@threads for k = 1:Threads.nthreads()
		threadRange = getrange(n)
		locPathi 	= zeros(Int64,maxProd)
		locSearchi  = fill(true,maxProd)
		for ii in threadRange # Loop over consumers allocated to specific thread
			# Select consumer
			r = d.seli[ii]
			nProd = length(d.seli[ii])
			# Reset locals
			locPathi .= 0
			locSearchi .= true ; locSearchi[1] = false  # outside option not searchable
			# generate search paths (updates locPathi/locSearchi and returns purchi)
			purchi = genSearchPathi!(m,locPathi,locSearchi,zs[r],u[r])
			# Fill in values
			path[:,ii] = locPathi
			purch[ii] = purchi
			search[:,ii] = locSearchi
		end

	end
	search .= search .== false
	return path,purch,search
end

###################################################
function genSearchPathi!(m::Weitzman,path,searchable,zs,u)

	# Initial status
	cc_s = 0
	ind_p = 1 # index of maximum utility (initially outside option)
	u_max = u[ind_p] # max utility
	# Generate search paths
	while true
		# Current largest search value
		val_s,ind_s = findmax_subsel(zs,searchable)

		if val_s > u_max # search
			cc_s  += 1
			path[cc_s] = ind_s
			searchable[ind_s] = false
			u_max = if u[ind_s] > u_max
						ind_p = ind_s
						u[ind_s]
					else
						u_max
					end
		else # Purchase and end search
			purch = ind_p
			return purch
			break # end with purchase
		end
	end
end

###################################################
function negLogLik(m::Weitzman,θ::Vector{T},nBeta,d,cfun,maxPos;
					e = [], v = [],
					zsMat = [] ) where T<:Real

	# Extract Parameters
	β = θ[1:nBeta]
	cs = θ[nBeta+1]
	cd = if length(θ) > nBeta + 1
			θ[nBeta+2]
		else
			zero(T)
	end

	# Draw random numbers if not provided
	if isempty(e)
		Random.seed!(m.seed+1)
		e = rand(m.dE,size(d.chars,1),m.options.nDraws)
	end
	if isempty(v)
		Random.seed!(m.seed+298)
		v = rand(m.dV,size(d.chars,1),m.options.nDraws)
	end

	u = zeros(T,size(d.chars,1),m.options.nDraws)
	zs = similar(u)

	ξ   = isempty(zsMat) ?
			[calcZs(0.0,m.dE,cfun(cs + cd * pp))[1] for pp in 0:maxPos] :
			[calcZs(zsMat,cfun(cs + cd * pp)) for pp in 0:maxPos]

	fillUZS!(u,zs,m,d,β,ξ,e,v)

	# Loop through consumers
	logLik = zeros(T,Threads.nthreads())

	@inbounds Threads.@threads for k = 1:Threads.nthreads()
		locSumLL = zeros(T,1)
		locV = zeros(T,3,size(e,2))
		locV2 = zeros(T,2)
		for ii in getrange(length(d.seli))
			# v1: must have continued searching until nth searched
			# v2: must have stopped after nth search
			# v3: must have chosen best out of consideration set
			locV .= 0.0
			fillV!(locV,d.path[ii],d.purch[ii],u,zs,d.seli[ii],m.smo)
			addLL!(locSumLL,locV,locV2)
		end
		logLik[k] = locSumLL[1]
	end
	sum_logLik = abs(sum(logLik)) == Inf ? -1e10 : sum(logLik)
	return -sum_logLik
end

function fillUZS!(u,zs,m::Weitzman,d,β,ξ,e,v)
	@inbounds Threads.@threads for i in 1:size(u,1)
		for j in 1:size(u,2)
			u[i,j] = e[i,j] + v[i,j]
			zs[i,j] = v[i,j]
			for k in 1:length(β)
				u[i,j] += d.chars[i,k] * β[k]
				zs[i,j] += d.chars[i,k] * β[k]
			end
			zs[i,j] += ξ[d.pos[i]+1]
		end
	end
	return nothing
end

function addLL!(locSumLL,locV,locV2)
	nDraws = size(locV,2)
	locV2[2] = 0.0
	@inbounds for e in 1:nDraws
		locV2[1] = 0.0
		for k in 1:size(locV,1)
			locV2[1] += locV[k,e]
		end
		locV2[2] += 1 /(1 + locV2[1])  / nDraws
	end
	locSumLL[1] += log(locV2[2])
	return nothing
end

#=
Fills liklihood contributions for a given consumer into matrix `v`.

`v[1,:]`: Continuation rule: Reservation value of search in t needs to be larger than max utility of consideration set in t.

`v[2,:]`: Stopping rule: Max utility in consideration set at purchase in t needs to be larger than search values of all unsearched products in t.

`v[3,:]`: Purchase rule: Utility of purchased needs to be max in consideration set at time of purchase.
### Arguments:
`v`: 3 × nDraws array to be filled \\
`u`: nProd x nDraws array of product utilities of given consumer (u = xb + ϵ ). \\
`zs`: nProd x 1 array of search values for each product \\
`nSearch`: # searches \\
`u_p`: 1 x nDraws array of utility of product ultimately purchased by consumer \\
`smo`: Smoothing parameter
=#
function fillV!(v,path,purch,u,zs,r,smo)
	nSearch = length(path)
	@inbounds for ee = 1:size(u,2)
		u_p = u[r[purch],ee]
		if nSearch == 0  # no searches
			# v1 not defined, since no searches before
			v2!(v,u,zs,path,nSearch,r,ee,smo)
			# v3 not defined, since consideration set only contains outside option
		elseif nSearch == size(u,1)-1 # searched all products
			v1!(v,u,zs,path,nSearch,r,ee,smo)
			# v2 not defined, since no products left to search
			v3!(v,u_p,u,path,nSearch,r,ee,smo)
		else
			v1!(v,u,zs,path,nSearch,r,ee,smo)
			v2!(v,u,zs,path,nSearch,r,ee,smo)
			v3!(v,u_p,u,path,nSearch,r,ee,smo)
		end
	end
  	return
end

function v1!(v,u,zs,path,nSearch,r,ee,smo)
	um = u[r[1],ee] # outside option
	@inbounds for kk = 1:nSearch
		ip = r[path[kk]]
		v[1,ee] += exp(-smo*(zs[ip,ee] - um))
		um = max(um,u[ip,ee])
	end
end

function v2!(v,u,zs,path,nSearch,r,ee,smo)
	@inbounds for ll = 2:length(r)
		if ll in path
			continue
		end
		if nSearch > 0
			@views um = max(u[r[1],ee],findmax_range(u[r,ee],path)[1])
		else
			um = u[r[1],ee]
		end

		v[2,ee] += exp(-smo*(um-zs[r[ll],ee]))
	end
end

function v3!(v,u_p,u,path,nSearch,r,ee,smo)
	v[3,ee] += exp(-smo*(u_p-u[r[1],ee])) # outside option
	@inbounds for ll = 1:nSearch
		v[3,ee] += exp(-smo*(u_p-u[r[path[ll]],ee]))
	end
	v[3,ee] -= 1.0
end

###################################################
function estimateModel(m::Weitzman,d; startvals = [] , pars_true = [], fzsMat =[] )

	# Preliminaries
	cfun = 	getCfun(m.cfun)

	nBeta = size(d.chars,2)
	nCpar = if m.cd == 0
				1
			else
				2
			end

	if isempty(pars_true)
		pars_true = fill(NaN,nBeta+nCpar)
	end

	if isempty(startvals)
		startvals = rand(Uniform(),nBeta+nCpar)
	end

	zsMat = isempty(fzsMat) ? [] : getIntExtPolZs(fzsMat)

	# Random draws, pre-allocate so not drawn in loglik function
	Random.seed!(m.seed+651168) # ensures draws not same as in dataGen
	e = rand(m.dE,size(d.chars,1),m.options.nDraws)
	v = rand(m.dV,size(d.chars,1),m.options.nDraws)

	# Maximum Position
	maxPos = maximum(d.pos)

	############################################################################
	# Optimization
	# Define objective function
	obj_fun(θ,p) = negLogLik(m,θ,nBeta,d,cfun,maxPos;v,e,zsMat)
	obj = OptimizationFunction(obj_fun,m.options.autodiff)
	prob = OptimizationProblem(obj,startvals)

	# Run optimization
	res = solve(prob,m.options.algo;
				show_trace = true , show_every=m.options.showEvery,
				maxiters = m.options.maxIter,
				x_tol = m.options.x_tol,
				f_tol = m.options.f_tol,
				g_tol = m.options.g_tol)

	pars_opt = res.minimizer
	fval_opt = res.minimum

	# Calculate obj. fun at start and true val
	startVal = obj_fun(startvals,[])
	trueVal = isnan(pars_true[1]) ? NaN : obj_fun(pars_true,[])

	# Extract results
	β = pars_opt[1:nBeta]
	cs = pars_opt[nBeta+1]
	if nCpar > 1
			cd = pars_opt[nBeta+2]
	end
	finalVal = fval_opt

	infotxt = "
	####################################
	# Results of Estimation:
	Estimation model:	$(typeof(m))
	Final Objective Value: 		$finalVal
	Start Objective Value:		$startVal
	True Objective Value:		$trueVal
	Parameters:
		β  = $(round.(β,digits=2))
		β0 = $(round.(pars_true[1:nBeta],digits=2))
		cs = $(round(cs,digits=2)) | cs0 = $(round(pars_true[nBeta+1],digits=2))"
	if nCpar >1
		infotxt *= "
		cd = $(round(cd,digits=2)) | cd0 = $(round(pars_true[nBeta+2],digits=2)) "
	end
	infotxt *= "
	#################################### "

	@info infotxt

	if nCpar > 1
		return β,cs,cd,finalVal
	else
		return β,cs,finalVal
	end

end

######################################################
function vecP(m::Weitzman)
	return vcat(m.β,m.cs,m.cd)
end

######################################################
function calcWelfareDemand(m::Weitzman,d; e = [] , v = [], zsMat = [])

	# Define values + cost function
	cfun = getCfun(m.cfun)
	maxPos = maximum(d.pos)
	maxProd = maximum(length.(d.seli))
	nCons = length(d.seli)

	# Draw random numbers if not provided
	if isempty(e)
		Random.seed!(m.seed+1)
		e = rand(m.dE,size(d.chars,1),m.options.nDraws)
	end
	if isempty(v)
		Random.seed!(m.seed+298)
		v = rand(m.dV,size(d.chars,1),m.options.nDraws)
	end

	u = zeros(eltype(m.β),size(d.chars,1),size(e,2))
	zs = similar(u)

	# position-specific costs and search values
	cs = [cfun(m.cs + m.cd * pp) for pp in 0:maxPos]
	ξ   = isempty(zsMat) ?
			[calcZs(0.0,m.dE,cs[i])[1] for i in eachindex(cs)] :
			[calcZs(zsMat,cs[i]) for i in eachindex(cs)]

	fillUZS!(u,zs,m,d,m.β,ξ,e,v)

	# Pre allocate
	sumU = zeros(Threads.nthreads())
	sumC = zeros(Threads.nthreads())
	demand = zeros(Int64,maxProd,Threads.nthreads())
	@inbounds Threads.@threads for k = 1:Threads.nthreads()
		threadRange = getrange(nCons)
		locPathi 	= zeros(Int64,maxProd)
		locSearchi  = fill(true,maxProd)
		locSearchi[1] = false

		locU = 0.0
		locC = 0.0
		locD = zeros(Int64,maxProd)
		for ii in threadRange
			r = d.seli[ii]
			@views posi = d.pos[r]
			for dd in 1:size(e,2)
				@views zsd 	= zs[r,dd]
				@views ud 	= u[r,dd]
				locPathi 	.= 0
				locSearchi[2:end]  .= true
				purchi 	= genSearchPathi!(m,locPathi,locSearchi,zsd,ud)
				locU 	+= ud[purchi]
				for ss = 1:maxProd
					if locPathi[ss] == 0
						break
					else
						locC += cs[posi[locPathi[ss]]+1]
					end
				end
				locD[purchi] += 1
			end
		end
		sumU[k] = locU / size(e,2)
		sumC[k] = locC / size(e,2)
		demand[:,k] .= locD
	end
	return sum(sumU) - sum(sumC), sum(sumU), sum(sumC), sum(demand,dims=2) ./ size(e,2) ./ nCons
end


text:
using QuadGK
using SpecialFunctions
using Roots
using RCall

function Gammadist(gamma)
	return ((adap.be^adap.al)/SpecialFunctions.gamma(adap.al))*(gamma^(adap.al-1))*exp(-adap.be*gamma)
end

function pip0(gamma)
	U = 4*adap.theta_f*adap.Lf/(2.0*adap.NN)
	R = 2*adap.Lf*adap.rho/(2.0*adap.NN)
	return Gammadist(gamma)*exp(-(Gammadist(gamma)*U/(2.0*adap.NN))/(gamma/(adap.NN+0.0)+R/(2.0*adap.NN)))
end

function intpip0()

	f(gam) = pip0(gam)
	return QuadGK.quadgk(f,0,1000)[1]
end

function calcBGam(L,alpha,beta,theta)
	# not sure whats going on here, seems to overestimate or under sometimes
	u = 2*theta/(2.0*adap.NN)
	r = adap.rho/(2.0*adap.NN)

	a = -1.0*(2^(alpha+1))*exp(L*adap.NN*r*beta)*L*adap.N*((1.0/(L*adap.N*r))^(1-alpha))*u*(beta^alpha)
	b =  convert(Float64,SpecialFunctions.gamma_inc(1-alpha,L*adap.NN*r*beta,1)[2])
	#fudge = 1-gamInt.cdf(0.2,a=adap.al2,scale=1./adap.be2)
	#fudge = 1.
	fudge = 0.25

	c = exp(a*b*fudge)

	return c
end

function set_theta_f_gam()
	i(theta) = calcBGam(adap.Lf,adap.al2,adap.be2,theta)-adap.B
	# theta_f  = fsolve(lambda theta ,0.0000001)
	theta_f  = Roots.find_zero(i,0.0000001)
	adap.theta_f = theta_f
end

function calcB(L,theta)
	t = -1.0*adap.gam_neg/(adap.NN+0.)
	u = 2.0*theta/(2.0*adap.NN)
	r = adap.rho/(2.0*adap.NN)

	#return u*t/((t+r*L)^2)
	#return u*t/(2*(t+(1.-np.exp(-2*r*L))/2.)^2)

	# Nordborg models
	#####return u/(2*r*t*(1+((1.-np.exp(-2*r*L))/2.)*(1-t)/t)^2)
	#return u/(t*(1+((1-np.exp(-2.0*r*L))/2.)*(1-t)/t)^2)
	#return u/(t*(1+r*L*(1-t)/t)^2)
	#####return u/(t*(1+(np.exp(-2*r*L)/2)/t)^2)
end

# Br() needs L and theta, not working right now
function get_B_vals()
	ret = Array{Float64}(undef,30)
	for i in 20:50
		L = convert(Int64,round(1.3^i,digits=0))
		ret[i] = (Br(L),L)
	end
	return ret
end

"""
	Estimating and plotting MAP using locfit and ggplot2 in R. It assume your folder contains the posterior estimated through ABCreg
"""
function plotMap(;analysis::String,output::String)

	try
		@eval using RCall
		@eval R"""library(ggplot2);library(abc)"""

		out = filter(x -> occursin("post",x), readdir(analysis,join=true))
		out = filter(x -> !occursin(".1.",x),out)

		open(x) = Array(CSV.read(GZip.open(x),DataFrame,header=false))
		posteriors = open.(out)

		maxp = DataFrame(Array{Float64,2}(undef,size(posteriors,1),5),[:aw,:as,:a,:gamNeg,:shape])
		R"""getmap <- function(df){
				temp = as.data.frame(df)
			    d <-locfit(~temp[,1],temp);
			    map<-temp[,1][which.max(predict(d,newdata=temp))]
			}"""
		getmap(x) = rcopy(R"""matrix(apply($x,2,getmap),nrow=1)""")
		tmp = getmap.(posteriors)
		maxp = DataFrame(vcat(tmp...),[:aw,:as,:a,:gamNeg,:shape])

		al  = maxp[:,1:3]
		gam  = maxp[:,4:end]
		p = rcopy(R"""al = $al
			dal = melt(al)
			pal = ggplot(dal) + geom_density(aes(x=value,fill=variable),alpha=0.5) + scale_fill_manual(values=c('#30504f', '#e2bd9a', '#ab2710'))
			ggsave(pal,filename=paste0($output))
			""")
		return(maxp)
	catch
		println("Please install R, ggplot2 and abc in your system before execute this function")
	end
end

#=function sampledAlpha(;divergence::Array,afs::Array,observed::Matrix,λdiv::Array,λpol::Array,pol::Array,dac::Array)

	## Outputs
	alphas, expDn, expDs = poisson_fixation(observedValues=divergence,λds=λdiv[1],λdn=λdiv[2],λweak=λdiv[3],λstrong=λdiv[4])
	expPn, expPs         = poisson_polymorphism(observedValues=afs,λps=λpol[1],λpn=λpol[2],pol=pol)

	θ_pn = @. observed[1,1] / (expPn[1,:]);
	θ_ps = @. observed[1,2] / (expPs[1,:]);

	r_pn = ones(size(expPn));
	r_ps = ones(size(expPs));
	for j = 2:size(dac,1)
		r_pn[j,:] .= @. (observed[j,1]) / (expPn[j,:]);
		r_pn[j,:] = @. r_pn[j,:] / θ_pn;

		r_ps[j,:] .= @. (observed[j,2]) / (expPs[j,:]);
		r_ps[j,:] .= @. r_ps[j,:] / θ_ps;
	end

	expPn_rj = zeros(size(expPn))
	expPs_rj = zeros(size(expPs))

	expPn_rj[1,:] = @. expPn[1,:] * θ_pn;
	expPs_rj[1,:] = @. expPs[1,:] * θ_ps;
	for j = 2:size(dac,1)
		expPn_rj[j,:] = @. expPn[j,:] * θ_pn * r_pn[j,:];
		expPs_rj[j,:] = @. expPs[j,:] * θ_ps * r_ps[j,:];
	end

	## Alpha from expected values. Used as summary statistics
	αS = @. round(1 - ((expDs/expDn) * (expPn/expPs)'),digits=5)
	αS_v2 = @. round(1 - ((expDs/expDn) * (expPn_rj/expPs_rj)'),digits=5)

	return αS,alphas,expDn,expDs,expPn,expPs
end
=#


text:
using Test

@testset "JuliaPrograms" begin
  include("JuliaPrograms.jl")
end

@testset "AlgebraicNets" begin
  include("AlgebraicNets.jl")
end

@testset "ExpressionTrees" begin
  include("ExpressionTrees.jl")
end


text:
if true

	dset = Dataset()
	indA = 0
	add_int_var(dset, "A")
	add_int_var(dset, "B")

	mat = [0 0;
	       0 0;
	       1 1;
	       1 1;
	       0 1]
	n,m = size(mat)
	set_number_of_records(dset, n)
	dset[0,0:n-1] = mat[:,1]
	dset[1,0:n-1] = mat[:,2]

	net1 = learn(dset, DSL_BayesianSearch)
	net2 = learn(dset, DSL_GreedyThickThinning)
	net3 = learn(dset, LearnParams_GreedyThickThinning() )
	net4 = learn(dset, LearnParams_PC() )
	net5 = learn(dset, LearnParams_NaiveBayes("A") )
	net6 = learn(dset, LearnParams_TreeAugmentedNaiveBayes("A") )

end


text:
using packageTrial
using Documenter

makedocs(;
    modules=[packageTrial],
    authors="Jin Xu",
    repo="https://github.com/sunnyXu/packageTrial.jl/blob/{commit}{path}#L{line}",
    sitename="packageTrial.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://sunnyXu.github.io/packageTrial.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
    doctest = true
)

deploydocs(;
    deps = Deps.pip("mkdocs","python-markdown-math"),
    repo="github.com/sunnyXu/packageTrial.jl",
)


text:
# LS-RTM of the 2D Marmousi model using linearized bregmann
# Author: mlouboutin3@gatech.edu
# Date: April 2022

using Statistics, Random, LinearAlgebra, JOLI
using JUDI, SegyIO, HDF5, PyPlot, SlimOptim

# Load migration velocity model
if ~isfile("$(JUDI.JUDI_DATA)/marmousi_model.h5")
    ftp_data("ftp://slim.gatech.edu/data/SoftwareRelease/Imaging.jl/2DLSRTM/marmousi_model.h5")
end
n, d, o, m0 = read(h5open("$(JUDI.JUDI_DATA)/marmousi_model.h5", "r"), "n", "d", "o", "m0")

# Set up model structure
model0 = Model((n[1], n[2]), (d[1], d[2]), (o[1], o[2]), m0)

# Load data
if ~isfile("$(JUDI.JUDI_DATA)/marmousi_2D.segy")
    ftp_data("ftp://slim.gatech.edu/data/SoftwareRelease/Imaging.jl/2DLSRTM/marmousi_2D.segy")
end
block = segy_read("$(JUDI.JUDI_DATA)/marmousi_2D.segy")
d_lin = judiVector(block)   # linearized observed data

# Set up wavelet
src_geometry = Geometry(block; key = "source", segy_depth_key = "SourceDepth")
wavelet = ricker_wavelet(src_geometry.t[1], src_geometry.dt[1], 0.03)    # 30 Hz wavelet
q = judiVector(src_geometry, wavelet)

###################################################################################################
# Infer subsampling based on free memory
mem = Sys.free_memory()/(1024^3)
t_sub = max(1, ceil(Int, 40/mem))
# Setup operators
opt = Options(subsampling_factor=t_sub, isic=true)  # ~40 GB of memory per source without subsampling
M = judiModeling(model0, q.geometry, d_lin.geometry; options=opt)
J = judiJacobian(M, q)

# Right-hand preconditioners (model topmute)
Mr = judiTopmute(model0.n, 52, 10)
# Sparsity
C = joEye(prod(model0.n); DDT=Float32, RDT=Float32)
# If available use curvelet instead for better result

# Setup linearized bregman
batchsize = 5 * parse(Int, get(ENV, "NITER", "$(q.nsrc ÷ 5)"))
niter = parse(Int, get(ENV, "NITER", "10"))
g_scale = 0

function obj(x)
    flush(stdout)
    dm = PhysicalParameter(x, model0.n, model0.d, model0.o)
    inds = randperm(q.nsrc)[1:batchsize]
    # Preconditionners
    Ml = judiMarineTopmute2D(30, d_lin[inds].geometry)

    residual = Ml*J[inds]*Mr*dm - Ml*d_lin[inds]
    # grad
    G = reshape(Mr'J[inds]'*Ml'*residual, model0.n)
    g_scale == 0 && (global g_scale = .05f0/maximum(G))
    G .*= g_scale
    return .5f0*norm(residual)^2, G[:]
end

# Bregman
bregopt = bregman_options(maxIter=niter, verbose=2, quantile=.9, alpha=.1, antichatter=false, spg=true)
solb = bregman(obj, zeros(Float32, prod(model0.n)), bregopt, C);

# Save final velocity model, function value and history
h5open("lsrtm_marmousi_breg_result.h5", "w") do file
    write(file, "x", reshape(solb.x, model0.n), "z", reshape(solb.z, model0.n), "fval", Float32.(solb.ϕ_trace))
end

# Plot final image
figure(); imshow(copy(adjoint(reshape(solb.x, model0.n))), extent = (0, 7.99, 3.19, 0), cmap = "gray", vmin = -3e-2, vmax = 3e-2)
title("LS-RTM with SGD"); xlabel("Lateral position [km]"); ylabel("Depth [km]")


text:
#=
Very heavily inspired by Calculus.jl, but with an emphasis on performance and DiffEq API convenience.
=#

#=
Compute the finite difference interval epsilon.
Reference: Numerical Recipes, chapter 5.7.
=#
@inline function compute_epsilon(::Type{Val{:forward}}, x::T, eps_sqrt=sqrt(eps(real(T)))) where T<:Number
    eps_sqrt * max(one(real(T)), abs(x))
end

@inline function compute_epsilon(::Type{Val{:central}}, x::T, eps_cbrt=cbrt(eps(real(T)))) where T<:Number
    eps_cbrt * max(one(real(T)), abs(x))
end

@inline function compute_epsilon(::Type{Val{:complex}}, x::T, ::Union{Nothing,T}=nothing) where T<:Real
    eps(T)
end

@inline function compute_epsilon_factor(fdtype::DataType, ::Type{T}) where T<:Number
    if fdtype==Val{:forward}
        return sqrt(eps(real(T)))
    elseif fdtype==Val{:central}
        return cbrt(eps(real(T)))
    else
        return one(real(T))
    end
end

function fdtype_error(funtype::Type{T}=Float64) where T
    if funtype<:Real
        error("Unrecognized fdtype: valid values are Val{:forward}, Val{:central} and Val{:complex}.")
    elseif funtype<:Complex
        error("Unrecognized fdtype: valid values are Val{:forward} or Val{:central}.")
    else
        error("Unrecognized returntype: should be a subtype of Real or Complex.")
    end
end


text:
import Base.isapprox
import FactCheck.roughly

using QSpice.State

roughly(q1::QuantumState) = (q2::QuantumState) -> isapprox(q1, q2)
roughly(q1::Vector{Complex{Float64}}) = (q2::QuantumState) -> isapprox(q1, q2.vector)

function isapprox(q1::Vector{Complex{Float64}}, q2::Vector{Complex{Float64}}, atol=0.0001)
    for (x, y) in zip(q1, q2)
        if !isapprox(x, y, atol = atol)
            return false
        end
    end
    return true
end

max_stars_repo_name: leios/quantum

max_issues_repo_name: leios/quantum

max_forks_repo_name: leios/quantum


text:
#=
BellmanFordExample:
- Julia version: 1.6.2
- Author: samuel
- Date: 2021-09-11
=#

using SuiteSparseGraphBLAS
using SparseArrays
using LinearAlgebra

include("/Users/samuel/Documents/Jupyter/GraphBlasAlgorithmsFinal/src/BellmanFord.jl")
dim = 100
matrix =  GBMatrix(sprand(dim, dim, 1.0))
@time bellmanford(matrix, 1, dim)


text:
type LmResp{V<:FPVector} <: ModResp  # response in a linear model
    mu::V                                  # mean response
    offset::V                              # offset added to linear predictor (may have length 0)
    wts::V                                 # prior weights (may have length 0)
    y::V                                   # response
    function LmResp(mu::V, off::V, wts::V, y::V)
        n = length(y)
        length(mu) == n || error("mismatched lengths of mu and y")
        ll = length(off)
        ll == 0 || ll == n || error("length of offset is $ll, must be $n or 0")
        ll = length(wts)
        ll == 0 || ll == n || error("length of wts is $ll, must be $n or 0")
        new(mu, off, wts, y)
    end
end
convert{V<:FPVector}(::Type{LmResp{V}}, y::V) =
    LmResp{V}(zeros(y), similar(y, 0), similar(y, 0), y)

function updateμ!{V<:FPVector}(r::LmResp{V}, linPr::V)
    n = length(linPr)
    length(r.y) == n || error("length(linPr) is $n, should be $(length(r.y))")
    length(r.offset) == 0 ? copy!(r.mu, linPr) : broadcast!(+, r.mu, linPr, r.offset)
    deviance(r)
end
updateμ!{V<:FPVector}(r::LmResp{V}, linPr) = updateμ!(r, convert(V, vec(linPr)))

function deviance(r::LmResp)
    y = r.y
    mu = r.mu
    wts = r.wts
    v = zero(eltype(y)) + zero(eltype(y)) * zero(eltype(wts))
    if isempty(wts)
        @inbounds @simd for i = eachindex(y,mu)
            v += abs2(y[i] - mu[i])
        end
    else
        @inbounds @simd for i = eachindex(y,mu,wts)
            v += abs2(y[i] - mu[i])*wts[i]
        end
    end
    v
end

function nulldeviance(r::LmResp)
    y = r.y
    m = mean(y)
    wts = r.wts
    v = zero(eltype(y))*zero(eltype(wts))
    if isempty(wts)
        @inbounds @simd for i = 1:length(y)
            v += abs2(y[i] - m)
        end
    else
        @inbounds @simd for i = 1:length(y)
            v += abs2(y[i] - m)*wts[i]
        end
    end
    v
end

function loglikelihood(r::LmResp)
    n = length(r.y)
    wts = r.wts
    sw = zero(log(one(eltype(wts))))
    for w in wts
        sw += log(w)
    end
    -n/2 * (log(2π * deviance(r)/n) + 1 - sw)
end

function nullloglikelihood(r::LmResp)
    n = length(r.y)
    wts = r.wts
    sw = zero(log(one(eltype(wts))))
    for w in wts
        sw += log(w)
    end
    -n/2 * (log(2π * nulldeviance(r)/n) + 1 - sw)
end

function residuals(r::LmResp)
    y = r.y
    mu = r.mu
    if isempty(r.wts)
        y - mu
    else
        wts = r.wts
        resid = similar(y)
        @simd for i = eachindex(resid,y,mu,wts)
            @inbounds resid[i] = (y[i] - mu[i]) * sqrt(wts[i])
        end
        resid
    end
end

type LinearModel{L<:LmResp,T<:LinPred} <: LinPredModel
    rr::L
    pp::T
end

cholfact(x::LinearModel) = cholfact(x.pp)

function fit{LmRespT<:LmResp,LinPredT<:LinPred, T<:FP}(::Type{LinearModel{LmRespT,LinPredT}},
    X::AbstractMatrix{T}, y::FPVector)
    rr = LmRespT(y)
    pp = LinPredT(X)
    installbeta!(delbeta!(pp, rr.y))
    updateμ!(rr, linpred(pp, 0.0))
    LinearModel(rr, pp)
end
function fit(::Type{LinearModel}, X::AbstractMatrix, y::Vector)
    yy = float(y)
    T = eltype(yy)
    return fit(LinearModel{LmResp{typeof(yy)}, DensePredQR{T}}, float(X), yy)
end

lm(X, y) = fit(LinearModel, X, y)
lmc(X, y) = fit(LinearModel{DensePredChol}, X, y)

dof(x::LinearModel) = length(coef(x)) + 1

"""
    deviance(obj::LinearModel)

For linear models, the deviance is equal to the residual sum of squares (RSS).
"""
deviance(obj::LinearModel) = deviance(obj.rr)

"""
    nulldeviance(obj::LinearModel)

For linear models, the deviance of the null model is equal to the total sum of squares (TSS).
"""
nulldeviance(obj::LinearModel) = nulldeviance(obj.rr)
loglikelihood(obj::LinearModel) = loglikelihood(obj.rr)
nullloglikelihood(obj::LinearModel) = nullloglikelihood(obj.rr)

r2(obj::LinearModel) = 1 - deviance(obj)/nulldeviance(obj)

function adjr2(obj::LinearModel)
    n = nobs(obj)
    # dof() includes the dispersion parameter
    p = dof(obj) - 1
    1 - (1 - r²(obj))*(n-1)/(n-p)
end

function dispersion(x::LinearModel, sqr::Bool=false)
    ssqr = deviance(x.rr)/dof_residual(x)
    return sqr ? ssqr : sqrt(ssqr)
end

function coeftable(mm::LinearModel)
    cc = coef(mm)
    se = stderr(mm)
    tt = cc ./ se
    CoefTable(hcat(cc,se,tt,ccdf(FDist(1, dof_residual(mm)), abs2.(tt))),
              ["Estimate","Std.Error","t value", "Pr(>|t|)"],
              ["x$i" for i = 1:size(mm.pp.X, 2)], 4)
end

predict(mm::LinearModel, newx::Matrix) =  newx * coef(mm)

function confint(obj::LinearModel, level::Real)
    hcat(coef(obj),coef(obj)) + stderr(obj) *
    quantile(TDist(dof_residual(obj)), (1. - level)/2.) * [1. -1.]
end
confint(obj::LinearModel) = confint(obj, 0.95)

max_stars_repo_name: werbac/jlib

max_issues_repo_name: werbac/jlib

max_forks_repo_name: werbac/jlib


text:
"""
    GEBT <: AbstractModel

Geometrically exact beam theory model, as implemented by the GXBeam package.
State variables are as defined by GXBeam.  Inputs correspond to the external
forces ``F_{x,i}, F_{y,i}, F_{z,i}, M_{x,i}, M_{y,i}, M_{z,i}`` or
displacements ``u_{x,i}, u_{y,i}, u_{z,i}, \\theta_{x,i}, \\theta_{y,i},
\\theta_{z,i}`` applied to each node, followed by the distributed loads
``f_{x,i}, f_{y,i}, f_{z,i}, m_{x,i}, m_{y,i}, m_{z,i}`` applied to each beam
element, followed by the linear and angular velocity of the system. Parameters
correspond to the location ``p_{x}, p_{y}, p_{z}`` of each node followed by
each beam element's properties. Each beam element's properties are defined by a
triad which defines the orientation of the beam element ``e_{1,x}, e_{1,y},
e_{1,z}, e_{2,x}, e_{2,y}, e_{2,z}, e_{3,x}, e_{3,y}, e_{3,z}``, followed by
the 21 independent entries of the compliance matrix ``C_{11}, C_{12}, C_{13},
C_{14}, C_{15}, C_{16}, C_{22}, C_{23}, C_{24}, C_{25}, C_{26}, C_{33}, C_{34},
C_{35}, C_{36}, C_{44}, C_{45}, C_{46}, C_{55}, C_{56}, C_{66}``, followed by
the beam element's inertial properties ``\\mu, x_{m,2}, x_{m,3}, i_{22}, i_{33},
i_{23}``.
"""
struct GEBT{TF,TI,TC,TB} <: AbstractModel
    # element connectivity
    start::TC
    stop::TC
    # boolean matrix indicating location and degree of freedom of displacement constraints
    displacement::TB
    # scaling parameters for the system
    force_scaling::TF
    mass_scaling::TF
    # indices for accessing governing equations for each point and element
    irow_point::TI
    irow_elem::TI
    irow_elem1::TI
    irow_elem2::TI
    # indices for accessing state variables for each point and beam
    icol_point::TI
    icol_elem::TI
end

# --- Constructors --- #

"""
    GEBT(start, stop, displacement; kwargs...)

Construct a geometrically exact beam theory structural model with beam elements
which extend from the point indices in `start` to the point indices in
`stop` and points with prescribed displacements as specified in
`displacement`.

# Arguments
 - `start`: Vector containing point index where each beam element starts
 - `stop`: Vector containing point index where each beam element stops
 - `displacement`: Boolean matrix indicating the point index and degree of
    freedom of the system's displacement constraints.  Rows correspond to the
    degrees of freedom ``x, y, z, \\theta_x, \\theta_y, \\theta_z`` and
    columns correspond to the point indices ``1, 2, \\dots, N_p`` where ``N_p``
    is the total number of points.

# Keyword Arguments
 - `force_scaling = 1.0`: Factor used to scale system forces/moments internally
 - `moment_scaling = 1.0`: Factor used to scale system mass/inertia internally
"""
function GEBT(start, stop, displacement; force_scaling=1.0, mass_scaling=1.0)

    # system matrix pointers
    N, irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem =
        GXBeam.system_indices(start, stop, false)

    return GEBT(start, stop, displacement, force_scaling, mass_scaling,
        irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem)
end

"""
    GEBT(assembly, prescribed; kwargs...)

Construct a geometrically exact beam theory structural model with connectivity
as specified in `assembly` and displacement constraints as specified in
`prescribed`.

# Keyword Arguments
 - `force_scaling`: Factor used to scale system forces/moments internally.  If
    not specified, a suitable default will be chosen based on the entries of the
    compliance matrix.
 - `mass_scaling`: Factor used to scale system mass/inertia internally.  If not
    specified, a suitable default will be chosen based on the entries of the
    mass matrix.
"""
function GEBT(assembly, prescribed;
    force_scaling = GXBeam.default_force_scaling(assembly),
    mass_scaling = GXBeam.default_mass_scaling(assembly))

    # system dimensions
    npoint = length(assembly.points)
    nelem = length(assembly.elements)

    # include linear and angular momentum
    static = false

    # initialize system pointers
    N, irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem =
        GXBeam.system_indices(assembly.start, assembly.stop, static)

    # set displacement constraint indices and degree of freedom
    displacement = zeros(Bool, 6, npoint)
    for key in keys(prescribed)
        displacement[:,key] .= prescribed[key].force .== false
    end

    return GEBT(assembly.start, assembly.stop, displacement, force_scaling,
        mass_scaling, irow_point, irow_elem, irow_elem1, irow_elem2, icol_point,
        icol_elem)
end

# --- Traits --- #

number_of_states(model::GEBT) = 18*length(model.icol_elem) + 6*length(model.icol_point)

function number_of_inputs(model::GEBT)
    return 6*length(model.icol_point) + 6*length(model.icol_elem) + 6
end

number_of_parameters(model::GEBT) = 3*length(model.icol_point) + 36*length(model.icol_elem)

inplaceness(::Type{<:GEBT}) = InPlace()

rate_jacobian_type(::Type{<:GEBT}) = Linear()

state_jacobian_type(::Type{<:GEBT}) = Nonlinear()

input_jacobian_type(::Type{<:GEBT}) = Linear()

parameter_jacobian_type(::Type{<:GEBT}) = Nonlinear()

time_gradient_type(::Type{<:GEBT}) = Zeros()

# --- Methods --- #

function get_residual!(resid, model::GEBT, dx, x, y, p, t)

    # number of points and elements
    np = length(model.icol_point)
    ne = length(model.icol_elem)

    # index of beginning and end of each beam element
    start = model.start
    stop = model.stop

    # scaling parameters
    force_scaling = model.force_scaling
    mass_scaling = model.mass_scaling

    # model pointers
    irow_point = model.irow_point
    irow_elem = model.irow_elem
    irow_elem1 = model.irow_elem1
    irow_elem2 = model.irow_elem2
    icol_point = model.icol_point
    icol_elem = model.icol_elem

    # location and degree of freedom of displacement constraints
    displacement = model.displacement

    # construct assembly from parameters
    assembly = gxbeam_assembly(p, np, ne, start, stop)

    # extract newly created beam elements
    elements = assembly.elements

    # construct point and element loads from inputs
    prescribed, distributed = gxbeam_loads(y, np, ne, displacement, elements)

    # origin, linear velocity, and angular velocity
    x0 = @SVector zeros(3)
    v0 = @SVector [y[end-5], y[end-4], y[end-3]]
    ω0 = @SVector [y[end-2], y[end-1], y[end]]

    # return mass matrix multiplied state rates
    return GXBeam.dynamic_system_residual!(resid, x, dx, assembly,
        prescribed, distributed, force_scaling, mass_scaling,
        irow_point, irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem,
        x0, v0, ω0)
end

# --- Performance Overloads --- #

function get_rate_jacobian!(J, model::GEBT, dx, x, y, p, t)

    # zero out jacobian entries
    J .= 0

    # number of points and elements
    np = length(model.icol_point)
    ne = length(model.icol_elem)

    # index of beginning and end of each beam element
    start = model.start
    stop = model.stop

    # scaling parameters
    force_scaling = model.force_scaling
    mass_scaling = model.mass_scaling

    # construct assembly from parameters
    assembly = gxbeam_assembly(p, np, ne, start, stop)

    # extract model pointers
    irow_point = model.irow_point
    irow_elem = model.irow_elem
    irow_elem1 = model.irow_elem1
    irow_elem2 = model.irow_elem2
    icol_point = model.icol_point
    icol_elem = model.icol_elem

    # populate mass matrix entries
    return GXBeam.system_mass_matrix!(J, x, assembly, force_scaling, mass_scaling, irow_point,
        irow_elem, irow_elem1, irow_elem2, icol_point, icol_elem)
end

function get_state_jacobian!(J, model::GEBT, dx, x, y, p, t)

    # zero out jacobian entries
    J .= 0

    # number of points and elements
    np = length(model.icol_point)
    ne = length(model.icol_elem)

    # index of beginning and end of each beam element
    start = model.start
    stop = model.stop

    # scaling parameters
    force_scaling = model.force_scaling
    mass_scaling = model.mass_scaling

    # model pointers
    irow_point = model.irow_point
    irow_elem = model.irow_elem
    irow_elem1 = model.irow_elem1
    irow_elem2 = model.irow_elem2
    icol_point = model.icol_point
    icol_elem = model.icol_elem

    # location and degree of freedom of displacement constraints
    displacement = model.displacement

    # construct assembly from parameters
    assembly = gxbeam_assembly(p, np, ne, start, stop)

    # extract newly created beam elements
    elements = assembly.elements

    # construct point and element loads from inputs
    prescribed, distributed = gxbeam_loads(y, np, ne, displacement, elements)

    # origin, linear velocity, and angular velocity
    x0 = @SVector zeros(3)
    v0 = @SVector [y[end-5], y[end-4], y[end-3]]
    ω0 = @SVector [y[end-2], y[end-1], y[end]]

    # return jacobian of right hand side with respect to state variables
    return GXBeam.dynamic_system_jacobian!(J, x, dx, assembly, prescribed,
        distributed, force_scaling, mass_scaling, irow_point, irow_elem, irow_elem1,
        irow_elem2, icol_point, icol_elem, x0, v0, ω0)
end

# TODO: Input and parameter jacobians

# --- Convenience Functions --- #

function set_states!(x, model::GEBT;
    u_e = nothing, theta_e = nothing, F_e = nothing, M_e = nothing,
    P_e = nothing, H_e = nothing, u_p = nothing, theta_p = nothing,
    F_p = nothing, M_p = nothing)

    icol_elem = model.icol_elem
    icol_point = model.icol_point

    displacement = model.displacement

    force_scaling = model.force_scaling
    mass_scaling = model.mass_scaling

    nelem = length(icol_elem)
    npoint = length(icol_point)

    for ielem = 1:nelem

        icol = icol_elem[ielem]

        if !isnothing(u_e)
            GXBeam.set_element_deflections!(x, icol, u_e[ielem])
        end

        if !isnothing(theta_e)
            GXBeam.set_element_rotations!(x, icol, theta_e[ielem])
        end

        if !isnothing(F_e)
            GXBeam.set_element_forces!(x, icol, F_e[ielem], force_scaling)
        end

        if !isnothing(M_e)
            GXBeam.set_element_moments!(x, icol, M_e[ielem], force_scaling)
        end

        if !isnothing(P_e)
            GXBeam.set_element_linear_momenta!(x, icol, P_e[ielem], mass_scaling)
        end

        if !isnothing(H_e)
            GXBeam.set_element_angular_momenta!(x, icol, H_e[ielem], mass_scaling)
        end
    end

    for ipoint = 1:npoint

        icol = icol_point[ipoint]

        prescribed_forces = SVector{6}(view(displacement, :, ipoint)) .== false

        if !isnothing(u_p)
            GXBeam.set_point_deflections!(x, icol, u_p[ipoint], prescribed_forces)
        end

        if !isnothing(theta_p)
            GXBeam.set_point_rotations!(x, icol, theta_p[ipoint], prescribed_forces)
        end

        if !isnothing(F_p)
            GXBeam.set_point_forces!(x, icol, F_p[ipoint], prescribed_forces, force_scaling)
        end

        if !isnothing(M_p)
            GXBeam.set_point_moments!(x, icol, M_p[ipoint], prescribed_forces, force_scaling)
        end
    end

    return x
end

function set_inputs!(y, model::GEBT; point_conditions, element_loads, V, Omega)

    np = length(model.icol_point)
    ne = length(model.icol_elem)

    for ip = 1:np
        y[6*(ip-1) + 1] = point_conditions[1,ip]
        y[6*(ip-1) + 2] = point_conditions[2,ip]
        y[6*(ip-1) + 3] = point_conditions[3,ip]
        y[6*(ip-1) + 4] = point_conditions[4,ip]
        y[6*(ip-1) + 5] = point_conditions[5,ip]
        y[6*(ip-1) + 6] = point_conditions[6,ip]
    end

    for ie = 1:ne
        y[6*np + 6*(ie-1) + 1] = element_loads[1,ie]
        y[6*np + 6*(ie-1) + 2] = element_loads[2,ie]
        y[6*np + 6*(ie-1) + 3] = element_loads[3,ie]
        y[6*np + 6*(ie-1) + 4] = element_loads[4,ie]
        y[6*np + 6*(ie-1) + 5] = element_loads[5,ie]
        y[6*np + 6*(ie-1) + 6] = element_loads[6,ie]
    end

    y[6*np + 6*ne + 1] = V[1]
    y[6*np + 6*ne + 2] = V[2]
    y[6*np + 6*ne + 3] = V[3]

    y[6*np + 6*ne + 4] = Omega[1]
    y[6*np + 6*ne + 5] = Omega[2]
    y[6*np + 6*ne + 6] = Omega[3]

    return y
end

function set_parameters!(p, model::GEBT; assembly)

    np = length(model.icol_point)
    ne = length(model.icol_elem)

    for ip = 1:np
        p[3*(ip-1) + 1] = assembly.points[ip][1]
        p[3*(ip-1) + 2] = assembly.points[ip][2]
        p[3*(ip-1) + 3] = assembly.points[ip][3]
    end

    for ie = 1:ne
        element = assembly.elements[ie]
        minv11, minv12, minv22 = element.minv11, element.minv12, element.minv22
        Msub = inv((@SMatrix [
            minv11[1,1] minv12[1,2] minv12[1,3];
            minv12[1,2] minv22[2,2] minv22[2,3];
            minv12[1,3] minv22[3,2] minv22[3,3];
            ]))
        p[3*np + 36*(ie-1) + 1] = element.Cab[1,1]
        p[3*np + 36*(ie-1) + 2] = element.Cab[2,1]
        p[3*np + 36*(ie-1) + 3] = element.Cab[3,1]
        p[3*np + 36*(ie-1) + 4] = element.Cab[1,2]
        p[3*np + 36*(ie-1) + 5] = element.Cab[2,2]
        p[3*np + 36*(ie-1) + 6] = element.Cab[3,2]
        p[3*np + 36*(ie-1) + 7] = element.Cab[1,3]
        p[3*np + 36*(ie-1) + 8] = element.Cab[2,3]
        p[3*np + 36*(ie-1) + 9] = element.Cab[3,3]
        p[3*np + 36*(ie-1) + 10] = element.C11[1,1]
        p[3*np + 36*(ie-1) + 11] = element.C11[1,2]
        p[3*np + 36*(ie-1) + 12] = element.C11[1,3]
        p[3*np + 36*(ie-1) + 13] = element.C12[1,1]
        p[3*np + 36*(ie-1) + 14] = element.C12[1,2]
        p[3*np + 36*(ie-1) + 15] = element.C12[1,3]
        p[3*np + 36*(ie-1) + 16] = element.C11[2,2]
        p[3*np + 36*(ie-1) + 17] = element.C11[2,3]
        p[3*np + 36*(ie-1) + 18] = element.C12[2,1]
        p[3*np + 36*(ie-1) + 19] = element.C12[2,2]
        p[3*np + 36*(ie-1) + 20] = element.C12[2,3]
        p[3*np + 36*(ie-1) + 21] = element.C11[3,3]
        p[3*np + 36*(ie-1) + 22] = element.C12[3,1]
        p[3*np + 36*(ie-1) + 23] = element.C12[3,2]
        p[3*np + 36*(ie-1) + 24] = element.C12[3,3]
        p[3*np + 36*(ie-1) + 25] = element.C22[1,1]
        p[3*np + 36*(ie-1) + 26] = element.C22[1,2]
        p[3*np + 36*(ie-1) + 27] = element.C22[1,3]
        p[3*np + 36*(ie-1) + 28] = element.C22[2,2]
        p[3*np + 36*(ie-1) + 29] = element.C22[2,3]
        p[3*np + 36*(ie-1) + 30] = element.C22[3,3]
        p[3*np + 36*(ie-1) + 31] =  Msub[1,1] # μ
        p[3*np + 36*(ie-1) + 32] = -Msub[1,3]/Msub[1,1] # xm2
        p[3*np + 36*(ie-1) + 33] =  Msub[1,2]/Msub[1,1] # xm3
        p[3*np + 36*(ie-1) + 34] =  Msub[2,2] # i22
        p[3*np + 36*(ie-1) + 35] =  Msub[3,3] # i33
        p[3*np + 36*(ie-1) + 36] = -Msub[2,3] # i23
    end

    return p
end

function separate_states(model::GEBT, x;
    point_conditions = fill(NaN, 6, length(model.icol_point)))

    TF = eltype(x)
    np = length(model.icol_point)
    ne = length(model.icol_elem)

    u_p = Vector{SVector{3,TF}}(undef, np)
    theta_p = Vector{SVector{3,TF}}(undef, np)
    F_p = Vector{SVector{3,TF}}(undef, np)
    M_p = Vector{SVector{3,TF}}(undef, np)

    u_e = Vector{SVector{3,TF}}(undef, ne)
    theta_e = Vector{SVector{3,TF}}(undef, ne)
    F_e = Vector{SVector{3,TF}}(undef, ne)
    M_e = Vector{SVector{3,TF}}(undef, ne)
    P_e = Vector{SVector{3,TF}}(undef, ne)
    H_e = Vector{SVector{3,TF}}(undef, ne)

    for ip = 1:np

        icol = model.icol_point[ip]

        prescribed_forces = SVector{6}(view(model.displacement, :, ip)) .== false

        # get the displacement and rotations of the point
        u_p[ip] = SVector(ifelse(prescribed_forces[1], x[icol  ], point_conditions[1,ip]),
                    ifelse(prescribed_forces[2], x[icol+1], point_conditions[2,ip]),
                    ifelse(prescribed_forces[3], x[icol+2], point_conditions[3,ip]))
        theta_p[ip] = SVector(ifelse(prescribed_forces[4], x[icol+3], point_conditions[4,ip]),
                    ifelse(prescribed_forces[5], x[icol+4], point_conditions[5,ip]),
                    ifelse(prescribed_forces[6], x[icol+5], point_conditions[6,ip]))

        # overwrite external forces/moments with solved for forces/moments
        F_p[ip] = SVector(ifelse(prescribed_forces[1], point_conditions[1,ip], x[icol  ] * model.force_scaling),
                    ifelse(prescribed_forces[2], point_conditions[2,ip], x[icol+1] * model.force_scaling),
                    ifelse(prescribed_forces[3], point_conditions[3,ip], x[icol+2] * model.force_scaling))
        M_p[ip] = SVector(ifelse(prescribed_forces[4], point_conditions[4,ip], x[icol+3] * model.force_scaling),
                    ifelse(prescribed_forces[5], point_conditions[5,ip], x[icol+4] * model.force_scaling),
                    ifelse(prescribed_forces[6], point_conditions[6,ip], x[icol+5] * model.force_scaling))

        # convert rotation parameter to Wiener-Milenkovic parameters
        scaling = GXBeam.rotation_parameter_scaling(theta_p[ip])
        theta_p[ip] *= scaling

    end

    for ie = 1:ne

        icol = model.icol_elem[ie]

        u_e[ie] = SVector(x[icol], x[icol+1], x[icol+2])
        theta_e[ie] = SVector(x[icol+3], x[icol+4], x[icol+5])
        F_e[ie] = SVector(x[icol+6], x[icol+7], x[icol+8]) .* model.force_scaling
        M_e[ie] = SVector(x[icol+9], x[icol+10], x[icol+11]) .* model.force_scaling
        P_e[ie] = SVector(x[icol+12], x[icol+13], x[icol+14]) .* model.mass_scaling
        H_e[ie] = SVector(x[icol+15], x[icol+16], x[icol+17]) .* model.mass_scaling

        # convert rotation parameter to Wiener-Milenkovic parameters
        scaling = GXBeam.rotation_parameter_scaling(theta_e[ie])
        theta_e[ie] *= scaling

    end

    return (u_p = u_p, theta_p = theta_p, F_p = F_p, M_p = M_p,
        u_e = u_e, theta_e = theta_e, F_e = F_e, M_e = M_e, P_e = P_e, H_e = H_e)
end

function separate_inputs(model::GEBT, y)

    np = length(model.icol_point)
    ne = length(model.icol_elem)

    point_conditions = reshape(view(y, 1 : 6*np), 6, np)

    element_loads = reshape(view(y, 6*np + 1 : 6*np + 6*ne), 6, ne)

    V = view(y, 6*np + 6*ne + 1 : 6*np + 6*ne + 3)

    Omega = view(y, 6*np + 6*ne + 4 : 6*np + 6*ne + 6)

    return (point_conditions = point_conditions, element_loads = element_loads,
        V = V, Omega = Omega)
end

function separate_parameters(model::GEBT, p)
    np = length(model.icol_point)
    ne = length(model.icol_elem)
    start = model.start
    stop = model.stop
    return (assembly = gxbeam_assembly(p, np, ne, start, stop),)
end

# --- Internal Methods --- #

function gxbeam_assembly(p, np, ne, start, stop)
    points = [gxbeam_point(view(p, 3*(ip-1) + 1 : 3*(ip-1) + 3)) for ip = 1:np]
    elements = [gxbeam_element(view(p, 3*np + 36*(ie-1) + 1 : 3*np + 36*(ie-1) + 36),
        points, start[ie], stop[ie]) for ie = 1:ne]
    return GXBeam.Assembly(points, start, stop, elements)
end

gxbeam_point(p) = SVector(p[1], p[2], p[3])

function gxbeam_element(p, points, start, stop)
    # separate element parameters
    e1x, e1y, e1z, e2x, e2y, e2z, e3x, e3y, e3z, C11, C12, C13, C14, C15, C16,
        C22, C23, C24, C25, C26, C33, C34, C35, C36, C44, C45, C46, C55, C56,
        C66, μ, xm2, xm3, i22, i33, i23 = p
    # element length
    ΔL = norm(points[stop] - points[start])
    # element location
    x = (points[start] + points[stop])/2
    # element compliance matrix
    C = @SMatrix [
        C11 C12 C13 C14 C15 C16;
        C12 C22 C23 C24 C25 C26;
        C13 C23 C33 C34 C35 C36;
        C14 C24 C34 C44 C45 C46;
        C15 C25 C35 C45 C55 C56;
        C16 C26 C36 C46 C56 C66
        ]
    # element inverse mass matrix
    minv = inv(@SMatrix [
            μ       0     0       0 μ*xm3 -μ*xm2;
            0       μ     0  -μ*xm3     0      0;
            0       0     μ   μ*xm2     0      0;
            0  -μ*xm3 μ*xm2 i22+i33     0      0;
         μ*xm3      0     0       0   i22   -i23;
        -μ*xm2      0     0       0  -i23    i33;
    ])
    # element triad
    Cab = @SMatrix [
        e1x e2x e3x;
        e1y e2y e3y;
        e1z e2z e3z;
    ]

    return GXBeam.Element(ΔL, x, C, minv, Cab)
end

function gxbeam_loads(y, np, ne, d, elements)
    prescribed = Dict(ip => gxbeam_point_load(view(y,
        6*(ip-1) + 1 : 6*(ip-1) + 6), view(d, :, ip)) for ip = 1:np)
    distributed = Dict(ie => gxbeam_distributed_load(view(y,
        6*np + 6*(ie-1) + 1 : 6*np + 6*(ie-1) + 6), elements[ie].L) for ie = 1:ne)
    return prescribed, distributed
end

function gxbeam_point_load(y, d)
    displacement = SVector{6}(d)
    force = displacement .== false
    value = SVector{6}(y)
    follower = @SVector zeros(6)
    return PrescribedConditions(force, value, follower)
end

function gxbeam_distributed_load(y, ΔL)
    f = SVector(y[1], y[2], y[3])
    m = SVector(y[4], y[5], y[6])
    f1 = f2 = ΔL*f/2
    m1 = m2 = ΔL*m/2
    f1_follower = f2_follower = @SVector zeros(3)
    m1_follower = m2_follower = @SVector zeros(3)
    return DistributedLoads(f1, f2, m1, m2, f1_follower, f2_follower, m1_follower,
        m2_follower)
end