text:
using DataFrames, Dash, DashHtmlComponents, DashCoreComponents, PlotlyJS, UrlDownload

df2 = DataFrame(urldownload("https://raw.githubusercontent.com/plotly/datasets/master/country_indicators.csv"))

dropmissing!(df2)

rename!(df2, Dict(:"Year" => "year"))

available_indicators = unique(df2[!, "Indicator Name"])
years = unique(df2[!, "year"])

app = dash()

app.layout = html_div() do
    html_div(
        children = [
            dcc_dropdown(
                id = "xaxis-column",
                options = [
                    (label = i, value = i) for i in available_indicators
                ],
                value = "Fertility rate, total (births per woman)",
            ),
            dcc_radioitems(
                id = "xaxis-type",
                options = [(label = i, value = i) for i in ["linear", "log"]],
                value = "linear",
            ),
        ],
        style = (width = "48%", display = "inline-block"),
    ),
    html_div(
        children = [
            dcc_dropdown(
                id = "yaxis-column",
                options = [
                    (label = i, value = i) for i in available_indicators
                ],
                value = "Life expectancy at birth, total (years)",
            ),
            dcc_radioitems(
                id = "yaxis-type",
                options = [(label = i, value = i) for i in ["linear", "log"]],
                value = "linear",
            ),
        ],
        style = (width = "48%", display = "inline-block", float = "right"),
    ),
    dcc_graph(id = "indicator-graphic"),
    dcc_slider(
        id = "year-slider-2",
        min = minimum(years),
        max = maximum(years),
        marks = Dict([Symbol(v) => Symbol(v) for v in years]),
        value = minimum(years),
        step = nothing,
    )
end

callback!(
    app,
    Output("indicator-graphic", "figure"),
    Input("xaxis-column", "value"),
    Input("yaxis-column", "value"),
    Input("xaxis-type", "value"),
    Input("yaxis-type", "value"),
    Input("year-slider-2", "value"),
) do xaxis_column_name, yaxis_column_name, xaxis_type, yaxis_type, year_value
    df2f = df2[df2.year .== year_value, :]
    return Plot(
        df2f[df2f[!, Symbol("Indicator Name")] .== xaxis_column_name, :Value],
        df2f[df2f[!, Symbol("Indicator Name")] .== yaxis_column_name, :Value],
        Layout(
            xaxis_type = xaxis_type == "Linear" ? "linear" : "log",
            xaxis_title = xaxis_column_name,
            yaxis_title = yaxis_column_name,
            yaxis_type = yaxis_type == "Linear" ? "linear" : "log",
            hovermode = "closest",
        ),
        kind = "scatter",
        text = df2f[
            df2f[Symbol("Indicator Name")] .== yaxis_column_name,
            Symbol("Country Name"),
        ],
        mode = "markers",
        marker_size = 15,
        marker_opacity = 0.5,
        marker_line_width = 0.5,
        marker_line_color = "white"
    )
end

run_server(app, "0.0.0.0", debug = true)

max_stars_repo_name: joelostblom/dash-docs

max_issues_repo_name: joelostblom/dash-docs

max_forks_repo_name: joelostblom/dash-docs


text:
@testset "Rate limiting" begin
    @testset "Queues" begin
        results = Nothing[]
        f() = (push!(results, nothing); nothing)
        l = Limiter()

        # We start with no queues.
        @test isempty(l.queues)
        # But when we need one, it gets created for us.
        enqueue!(f, l, :GET, "/foo")
        sleep(Millisecond(50))
        @test collect(keys(l.queues)) == ["/foo"]
        @test length(results) == 1

        n = now(UTC)
        q = l.queues["/foo"]
        q.remaining = 1
        q.reset = n + Millisecond(200)

        # If remaining > 0, we don't wait.
        enqueue!(f, l, :GET, "/foo")
        # But if remaining == 0, we do.
        sleep(Millisecond(50))
        q.remaining = 0
        enqueue!(f, l, :GET, "/foo")
        sleep(Millisecond(50))

        @test length(results) == 2
        sleep(Millisecond(250))
        @test length(results) == 3
    end

    @testset "parse_endpoint" begin
        # The variable parameter doesn't matter.
        @test parse_endpoint("/users/1", :GET) == "/users"

        # Unless it's one of these three.
        @test parse_endpoint("/channels/1", :GET) == "/channels/1"
        @test parse_endpoint("/guilds/1", :GET) == "/guilds/1"
        @test parse_endpoint("/webhooks/1", :GET) == "/webhooks/1"

        # Without a numeric parameter at the end, we get the whole endpoint.
        @test parse_endpoint("/users/@me/channels", :GET) == "/users/@me/channels"

        # Special case 1: Deleting messages.
        @test ==(
            parse_endpoint("/channels/1/messages/1", :DELETE),
            "/channels/1/messages DELETE",
        )
        @test parse_endpoint("/channels/1/messages/1", :GET) == "/channels/1/messages"

        # Special case 2: Invites.
        @test parse_endpoint("/invites/abcdef", :GET) == "/invites"
    end
end

ext: jl


text:
using Pkg: @pkg_str
# For reproducability only dependency this has is Luxor,
# and it was created with Luxor v1.6.0
pkg"add Luxor@v1.6"

using Luxor
using Random

const bridge_len = 50

function chain(jiggle=0)
    shaky_rotate(θ) = rotate(θ + jiggle*(rand()-0.5))

    ### 1
    shaky_rotate(0)
    sethue(Luxor.julia_red)
    link()
    m1 = getmatrix()

    ### 2
    sethue(Luxor.julia_green)
    translate(-50, 130);
    shaky_rotate(π/3);
    link()
    m2 = getmatrix()

    setmatrix(m1)
    sethue(Luxor.julia_red)
    overlap(-1.3π)
    setmatrix(m2)

    ### 3
    shaky_rotate(-π/3);
    translate(-120,80);
    sethue(Luxor.julia_purple)
    link()

    setmatrix(m2)
    setcolor(Luxor.julia_green)
    overlap(-1.5π)
end

function link()
    sector(50, 90, π, 0, :fill)
    sector(Point(0, bridge_len), 50, 90, 0, -π, :fill)

    rect(50,-3,40, bridge_len+6, :fill)
    rect(-50-40,-3,40, bridge_len+6, :fill)

    sethue("black")
    move(Point(-50, bridge_len))
    arc(Point(0,0), 50, π, 0, :stoke)
    arc(Point(0, bridge_len), 50, 0, -π, :stroke)

    move(Point(-90, bridge_len))
    arc(Point(0,0), 90, π, 0, :stoke)
    arc(Point(0, bridge_len), 90, 0, -π, :stroke)
    strokepath()
end

function overlap(ang_end)
    sector(Point(0, bridge_len), 50, 90, -0., ang_end, :fill)
    sethue("black")
    arc(Point(0, bridge_len), 50, 0, ang_end, :stoke)
    move(Point(90, bridge_len))
    arc(Point(0, bridge_len), 90, 0, ang_end, :stoke)

    strokepath()
end

# Actually draw it

Random.seed!(16)
Drawing(450,450, "logo.svg")
origin()
translate(50, -130);
chain(0.5)
finish()
preview()


text:
struct Layer{F<:ActivationFunction, N<:Real}
    weights::Matrix{N}
    bias::Vector{N}
    activation::F
end

abstract type AbstractNetwork end

struct Network <: AbstractNetwork
    layers::Vector{Layer}
end

Base.length(L::Layer) = length(L.bias)


text:
function build_cache!(cache::TimeStatusChange, op_problem::OperationsProblem)
    build_cache!(cache, op_problem.canonical)
end

function build_cache!(cache::TimeStatusChange, canonical::Canonical)
    parameter = get_value(canonical, cache.ref)
    value_array = JuMP.Containers.DenseAxisArray{Dict{Symbol, Float64}}(undef, axes(parameter)...)

    for name in parameter.axes[1]
        status = PJ.value(parameter[name])
        value_array[name] = Dict(:count => 999.0, :status => status)
    end

    cache.value = value_array

    return

end


text:
#=

Utility functions used throughout the library.

=#

function pmute(a::AbstractVector)
    # Return all unique permutations of the vector a, which must be a 1d
    # numerical array.

    sort!(a)  # Sort so we can deal with repeated elements

    # just initializing here so these are available in all `while` scopes
    i = 0
    first = 1
    alen = length(a)

    while true
        i = alen
        while true
            i -= 1

            if a[i] < a[i + 1]
                j = alen

                while a[i] >= a[j]
                    j -= 1  # j--
                end

                a[i], a[j] = a[j], a[i]  # swap(a[j], a[i])
                t = a[(i + 1):end]
                reverse!(t)
                a[(i + 1):end] = t

                # Output current.
                produce(copy(a))

                break  # next.
            end

            if i == first
                reverse!(a)
                produce(copy(a))
                return
            end
        end
    end
end

function cartprod(arrs, out=Array(eltype(arrs[1]),
                                  prod([length(a) for a in arrs]),
                                  length(arrs)))
    sz = Int[length(a) for a in arrs]
    k = 1
    for v in product(arrs...)
        i = 1
        for el in v
            out[k, i] = el
            i += 1
        end
        k += 1
    end

    return out
end

# Note the routines below are taken from Combinatorics.jl
# Because that package does not support Julia 0.4 I have
# included them here.  When 0.5 is released these will be
# removed in favor of the package.
immutable WithReplacementCombinations{T}
    a::T
    t::Int
end

Base.eltype{T}(::Type{WithReplacementCombinations{T}}) = Vector{eltype(T)}

Base.length(c::WithReplacementCombinations) = binomial(length(c.a)+c.t-1, c.t)

"generate all combinations with replacement of size t from an array a."
with_replacement_combinations(a, t::Integer) = WithReplacementCombinations(a, t)

Base.start(c::WithReplacementCombinations) = [1 for i in 1:c.t]
function Base.next(c::WithReplacementCombinations, s)
    n = length(c.a)
    t = c.t
    comb = [c.a[si] for si in s]
    if t > 0
        s = copy(s)
        changed = false
        for i in t:-1:1
            if s[i] < n
                s[i] += 1
                for j in (i+1):t
                    s[j] = s[i]
                end
                changed = true
                break
            end
        end
        !changed && (s[1] = n+1)
    else
        s = [n+1]
    end
    (comb, s)
end
Base.done(c::WithReplacementCombinations, s) =
    !isempty(s) && s[1] > length(c.a) || c.t < 0

max_issues_repo_name: EconForge/Smolyak

max_forks_repo_name: EconForge/Smolyak


text:
## Caesar.jl, need an easy way to accumulate ground truth from factors only (like rigid transform tree)
# e.g., PriorX0 + Odometry + Odometry

# using Revise

using RoME
using DistributedFactorGraphs

using Test

##

@testset "Test parametric mean accumulation..." begin
##

fg = initfg()

addVariable!(fg, :x0, Pose2)
addFactor!(fg, [:x0;], PriorPose2(MvNormal(zeros(3),0.001*diagm([1;1;1]))))

addVariable!(fg, :x1, Pose2)
addFactor!(fg, [:x0;:x1], Pose2Pose2(MvNormal([10;0;0.0],0.001*diagm([1;1;1]))))

# drawGraph(fg)

# add parametric means
val = accumulateFactorMeans(fg, [:x0f1; :x0x1f1])

@test isapprox(val, [10;0;0], atol=1e-3)

##
end

#


text:
module CoolDicts

include("types.jl")
include("interface.jl")

export CoolDict

end # module


text:
using SQP, Ipopt
using JuMP

ipopt_solver = optimizer_with_attributes(
    Ipopt.Optimizer,
    "print_level" => 0,
    "warm_start_init_point" => "yes",
)
optimizer = optimizer_with_attributes(
    SQP.Optimizer,
    "external_optimizer" => ipopt_solver,
    "max_iter" => 100,
    "algorithm" => "SQP-TR",
)
model = Model(optimizer)

@variable(model, X);
@variable(model, Y);
@objective(model, Min, X^2 + X);
@NLconstraint(model, X^2 - X == 2);
@NLconstraint(model, X*Y == 1);
@NLconstraint(model, X*Y >= 0);
@constraint(model, X >= -2);

println("________________________________________");
print(model);
println("________________________________________");
JuMP.optimize!(model);

xsol = JuMP.value.(X)
ysol = JuMP.value.(Y)
status = termination_status(model)

println("Xsol = ", xsol);
println("Ysol = ", ysol);

println("Status: ", status);


text:
# ------------------------------------------------------------------------------
# finite element bases
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# utility functions for generating polynomial bases
# ------------------------------------------------------------------------------

"""
```julia
gaussquadrature(q)
```

Construct a Gauss-Legendre quadrature

# Arguments:
- `q`:  number of Gauss-Legendre points

# Returns:
- Gauss-Legendre quadrature points and weights

# Example:
```jldoctest
# generate Gauss-Legendre points and weights
quadraturepoints, quadratureweights = LFAToolkit.gaussquadrature(5);

# verify
truepoints = [
    -√(5 + 2*√(10/7))/3,
    -√(5 - 2*√(10/7))/3,
    0.0,
    √(5 - 2*√(10/7))/3,
    √(5 + 2*√(10/7))/3
];
@assert truepoints ≈ quadraturepoints

trueweights = [
    (322-13*√70)/900,
    (322+13*√70)/900,
    128/225,
    (322+13*√70)/900,
    (322-13*√70)/900
];
@assert trueweights ≈ quadratureweights

# output

```
"""
function gaussquadrature(q::Int)
    quadraturepoints = zeros(Float64, q)
    quadratureweights = zeros(Float64, q)

    if q < 1
        throw(DomainError(q, "q must be greater than or equal to 1")) # COV_EXCL_LINE
    end

    # build qref1d, qweight1d
    for i = 0:floor(Int, q / 2)
        # guess
        xi = cos(π * (2 * i + 1.0) / (2 * q))

        # Pn(xi)
        p0 = 1.0
        p1 = xi
        p2 = 0.0
        for j = 2:q
            p2 = ((2 * j - 1.0) * xi * p1 - (j - 1.0) * p0) / j
            p0 = p1
            p1 = p2
        end

        # first Newton step
        dp2 = (xi * p2 - p0) * q / (xi * xi - 1.0)
        xi = xi - p2 / dp2

        # Newton to convergence
        iter = 0
        maxiter = q^2 * 100
        while iter < maxiter && abs(p2) > 1e-15
            p0 = 1.0
            p1 = xi
            for j = 2:q
                p2 = ((2 * j - 1.0) * xi * p1 - (j - 1.0) * p0) / j
                p0 = p1
                p1 = p2
            end
            dp2 = (xi * p2 - p0) * q / (xi * xi - 1.0)
            xi = xi - p2 / dp2
            iter += 1
        end

        # save xi, wi
        quadraturepoints[i+1] = -xi
        quadraturepoints[q-i] = xi
        wi = 2.0 / ((1.0 - xi * xi) * dp2 * dp2)
        quadratureweights[i+1] = wi
        quadratureweights[q-i] = wi
    end
    quadraturepoints[abs.(quadraturepoints).<10*eps()] .= 0

    # return
    return quadraturepoints, quadratureweights
end

"""
```julia
gausslobattoquadrature(q, weights)
```

Construct a Gauss-Legendre-Lobatto quadrature

# Arguments:
- `q`:        number of Gauss-Legendre-Lobatto points
- `weights`:  boolean flag indicating if quadrature weights are desired

# Returns:
- Gauss-Legendre-Lobatto quadrature points or points and weights

# Example:
```jldoctest
# generate Gauss-Legendre-Lobatto points
quadraturepoints = LFAToolkit.gausslobattoquadrature(5, false);

# verify
truepoints = [-1.0, -√(3/7), 0.0, √(3/7), 1.0];
@assert truepoints ≈ quadraturepoints

# generate Gauss-Legendre-Lobatto points and weights
quadraturepoints, quadratureweights = LFAToolkit.gausslobattoquadrature(5, true);

# verify
trueweights = [1/10, 49/90, 32/45, 49/90, 1/10];
@assert trueweights ≈ quadratureweights

# output

```
"""
function gausslobattoquadrature(q::Int, weights::Bool)
    quadraturepoints = zeros(Float64, q)
    quadratureweights = zeros(Float64, q)

    if q < 2
        throw(DomainError(q, "q must be greater than or equal to 2")) # COV_EXCL_LINE
    end

    # endpoints
    quadraturepoints[1] = -1.0
    quadraturepoints[q] = 1.0
    if weights
        wi = 2.0 / (q * (q - 1.0))
        quadratureweights[1] = wi
        quadratureweights[q] = wi
    end

    # build qref1d, qweight1d
    for i = 1:floor(Int, (q - 1) / 2)
        # guess
        xi = cos(π * i / (q - 1.0))

        # Pn(xi)
        p0 = 1.0
        p1 = xi
        p2 = 0.0
        for j = 2:q-1
            p2 = ((2 * j - 1.0) * xi * p1 - (j - 1.0) * p0) / j
            p0 = p1
            p1 = p2
        end

        # first Newton step
        dp2 = (xi * p2 - p0) * q / (xi * xi - 1.0)
        d2p2 = (2 * xi * dp2 - q * (q - 1.0) * p2) / (1.0 - xi * xi)
        xi = xi - dp2 / d2p2

        # Newton to convergence
        iter = 0
        maxiter = q^2 * 100
        while iter < maxiter && abs(dp2) > 1e-15
            p0 = 1.0
            p1 = xi
            for j = 2:q-1
                p2 = ((2 * j - 1.0) * xi * p1 - (j - 1.0) * p0) / j
                p0 = p1
                p1 = p2
            end
            dp2 = (xi * p2 - p0) * q / (xi * xi - 1.0)
            d2p2 = (2 * xi * dp2 - q * (q - 1.0) * p2) / (1.0 - xi * xi)
            xi = xi - dp2 / d2p2
            iter += 1
        end

        # save xi, wi
        quadraturepoints[i+1] = -xi
        quadraturepoints[q-i] = xi
        if weights
            wi = 2.0 / (q * (q - 1.0) * p2 * p2)
            quadratureweights[i+1] = wi
            quadratureweights[q-i] = wi
        end
    end
    quadraturepoints[abs.(quadraturepoints).<10*eps()] .= 0

    # return
    if weights
        return quadraturepoints, quadratureweights
    else
        return quadraturepoints
    end
end

"""
```julia
buildinterpolationandgradient(
    nodes,
    quadraturepoints,
)
```

Build one dimensional interpolation and gradient matrices, from Fornberg 1998

# Arguments:
- `nodes`:             1d basis nodes
- `quadraturepoints`:  1d basis quadrature points

# Returns:
- One dimensional interpolation and gradient matrices

# Example:
```jldoctest
# get nodes, quadrature points, and weights
numbernodes = 3;
numberquadraturepoints = 4;
nodes = LFAToolkit.gausslobattoquadrature(numbernodes, false);
quadraturepoints, quadratureweights1d = LFAToolkit.gaussquadrature(numberquadraturepoints);

# build interpolation, gradient matrices
interpolation, gradient = LFAToolkit.buildinterpolationandgradient(nodes, quadraturepoints);

# verify
for i in 1:numberquadraturepoints
    total = sum(interpolation[i, :]);
    @assert total ≈ 1.0

    total = sum(gradient[i, :]);
    @assert abs(total) < 1e-14
end

# output

```
"""
function buildinterpolationandgradient(
    nodes1d::AbstractArray{Float64},
    quadraturepoints1d::AbstractArray{Float64},
)
    # check inputs
    numbernodes1d = length(nodes1d)
    numberquadratuepoints1d = length(quadraturepoints1d)
    if numbernodes1d < 2
        # COV_EXCL_START
        throw(
            DomanError(
                numbernodes1d,
                "length of nodes1d must be greater than or equal to 2",
            ),
        )
        # COV_EXCL_STOP
    end
    if numbernodes1d < 2
        # COV_EXCL_START
        throw(
            DomanError(
                numbernodes1d,
                "length of quadraturepoints1d must be greater than or equal to 2",
            ),
        )
        # COV_EXCL_STOP
    end

    # build interpolation, gradient matrices
    # Fornberg, 1998
    interpolation1d = zeros(Float64, numberquadratuepoints1d, numbernodes1d)
    gradient1d = zeros(Float64, numberquadratuepoints1d, numbernodes1d)
    for i = 1:numberquadratuepoints1d
        c1 = 1.0
        c3 = nodes1d[1] - quadraturepoints1d[i]
        interpolation1d[i, 1] = 1.0
        for j = 2:numbernodes1d
            c2 = 1.0
            c4 = c3
            c3 = nodes1d[j] - quadraturepoints1d[i]
            for k = 1:j-1
                dx = nodes1d[j] - nodes1d[k]
                c2 *= dx
                if k == j - 1
                    gradient1d[i, j] =
                        c1 * (interpolation1d[i, k] - c4 * gradient1d[i, k]) / c2
                    interpolation1d[i, j] = -c1 * c4 * interpolation1d[i, k] / c2
                end
                gradient1d[i, k] = (c3 * gradient1d[i, k] - interpolation1d[i, k]) / dx
                interpolation1d[i, k] = c3 * interpolation1d[i, k] / dx
            end
            c1 = c2
        end
    end

    # return
    return interpolation1d, gradient1d
end

# ------------------------------------------------------------------------------
# user utility constructors
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# -- single element bases
# ------------------------------------------------------------------------------

"""
```julia
TensorH1LagrangeBasis(
    numbernodes1d,
    numberquadraturepoints1d,
    numbercomponents,
    dimension,
    collocatedquadrature = false,
)
```

Tensor product basis on Gauss-Legendre-Lobatto points with Gauss-Legendre (default)
  or Gauss-Legendre-Lobatto quadrature points

# Arguments:
- `numbernodes1d`:             number of Gauss-Legendre-Lobatto nodes in 1 dimension
- `numberquadraturepoints1d`:  number of quadrature points in 1 dimension
- `numbercomponents`:          number of components
- `dimension`:                 dimension of basis

# Keyword Arguments:
- `collocatedquadrature = false`:  Gauss-Legendre or Gauss-Legendre-Lobatto quadrature points,
                                       default: false, Gauss-Legendre-Lobatto

# Returns:
- H1 Lagrange tensor product basis object

# Example:
```jldoctest
# generate H1 Lagrange tensor product basis
basis = TensorH1LagrangeBasis(4, 4, 3, 2);

# generate basis with Gauss-Legendre quadrature points
basis = TensorH1LagrangeBasis(4, 4, 3, 2; collocatedquadrature=true);

# verify
println(basis)

# output
tensor product basis:
    numbernodes1d: 4
    numberquadraturepoints1d: 4
    numbercomponents: 3
    dimension: 2
```
"""
function TensorH1LagrangeBasis(
    numbernodes1d::Int,
    numberquadraturepoints1d::Int,
    numbercomponents::Int,
    dimension::Int;
    collocatedquadrature::Bool = false,
)
    # check inputs
    if numbernodes1d < 2
        throw(DomanError(numbernodes1d, "numbernodes1d must be greater than or equal to 2")) # COV_EXCL_LINE
    end
    if numberquadraturepoints1d < 1
        # COV_EXCL_START
        throw(
            DomanError(
                numberquadraturepoints1d,
                "numberquadraturepoints1d must be greater than or equal to 1",
            ),
        )
        # COV_EXCL_STOP
    end
    if dimension < 1 || dimension > 3
        throw(DomanError(dimension, "only 1D, 2D, or 3D bases are supported")) # COV_EXCL_LINE
    end

    # get nodes, quadrature points, and weights
    nodes1d = gausslobattoquadrature(numbernodes1d, false)
    quadraturepoints1d = []
    quadratureweights1d = []
    if collocatedquadrature
        quadraturepoints1d, quadratureweights1d =
            gausslobattoquadrature(numberquadraturepoints1d, true)
    else
        quadraturepoints1d, quadratureweights1d = gaussquadrature(numberquadraturepoints1d)
    end

    # build interpolation, gradient matrices
    interpolation1d, gradient1d = buildinterpolationandgradient(nodes1d, quadraturepoints1d)

    # use basic constructor
    return TensorBasis(
        numbernodes1d,
        numberquadraturepoints1d,
        numbercomponents,
        dimension,
        nodes1d,
        quadraturepoints1d,
        quadratureweights1d,
        interpolation1d,
        gradient1d,
    )
end

"""
```julia
TensorH1UniformBasis(
    numbernodes1d,
    numberquadraturepoints1d,
    numbercomponents,
    dimension,
)
```

Tensor product basis on uniformly spaced points with Gauss-Legendre quadrature points

# Arguments:
- `numbernodes1d`:             number of uniformly spaced nodes in 1 dimension
- `numberquadraturepoints1d`:  number of Gauss-Legendre quadrature points in 1 dimension
- `numbercomponents`:          number of components
- `dimension`:                 dimension of basis

# Returns:
- H1 Lagrange tensor product basis on uniformly spaced nodes object

# Example:
```jldoctest
# generate H1 Lagrange tensor product basis on uniformly spaced nodes
basis = TensorH1UniformBasis(4, 3, 2, 1);

# verify
println(basis)

# output
tensor product basis:
    numbernodes1d: 4
    numberquadraturepoints1d: 3
    numbercomponents: 2
    dimension: 1
```
"""
function TensorH1UniformBasis(
    numbernodes1d::Int,
    numberquadraturepoints1d::Int,
    numbercomponents::Int,
    dimension::Int,
)
    # check inputs
    if numbernodes1d < 2
        throw(DomanError(numbernodes1d, "numbernodes1d must be greater than or equal to 2")) # COV_EXCL_LINE
    end
    if numberquadraturepoints1d < 1
        # COV_EXCL_START
        throw(
            DomanError(
                numberquadraturepoints1d,
                "numberquadraturepoints1d must be greater than or equal to 1",
            ),
        )
        # COV_EXCL_STOP
    end
    if dimension < 1 || dimension > 3
        throw(DomanError(dimension, "only 1D, 2D, or 3D bases are supported")) # COV_EXCL_LINE
    end

    # get nodes, quadrature points, and weights
    nodes1d = [-1.0:(2.0/(numbernodes1d-1)):1.0...]
    quadraturepoints1d, quadratureweights1d = gaussquadrature(numberquadraturepoints1d)

    # build interpolation, gradient matrices
    interpolation1d, gradient1d = buildinterpolationandgradient(nodes1d, quadraturepoints1d)

    # use basic constructor
    return TensorBasis(
        numbernodes1d,
        numberquadraturepoints1d,
        numbercomponents,
        dimension,
        nodes1d,
        quadraturepoints1d,
        quadratureweights1d,
        interpolation1d,
        gradient1d,
    )
end

# ------------------------------------------------------------------------------
# -- marco element bases
# ------------------------------------------------------------------------------

"""
```julia
TensorMacroElementBasisFrom1D(
    numbernodes1d,
    numberquadraturepoints1d,
    numbercomponents,
    dimension,
    numberelements1d,
    basis1dmicro,
    overlapquadraturepoints = false,
)
```

Tensor product macro-element basis from 1d single element tensor product basis

# Arguments:
- `numbernodes1d`:             number of basis nodes in 1 dimension
- `numberquadraturepoints1d`:  number of quadrature points in 1 dimension
- `numbercomponents`:          number of components
- `dimension`:                 dimension of basis
- `numberelements1d`:          number of elements in macro-element
- `basis1dmicro`:              1d micro element basis to replicate

# Keyword Arguments:
- `overlapquadraturepoints`:  Overlap quadrature points between elements, for prolongation
                                  default: false

# Returns:
- Tensor product macro-element basis object
"""
function TensorMacroElementBasisFrom1D(
    numbernodes1d::Int,
    numberquadraturepoints1d::Int,
    numbercomponents::Int,
    dimension::Int,
    numberelements1d::Int,
    basis1dmicro::TensorBasis;
    overlapquadraturepoints::Bool = false,
)
    if numberelements1d < 2
        # COV_EXCL_START
        throw(
            DomanError(numberelements1d, "macro-elements must contain at least 2 elements"),
        )
        # COV_EXCL_STOP
    end

    # compute dimensions
    numbernodes1dmacro = (numbernodes1d - 1) * numberelements1d + 1
    numberquadraturepoints1dmacro =
        overlapquadraturepoints ? (numberquadraturepoints1d - 1) * numberelements1d + 1 :
        numberquadraturepoints1d * numberelements1d

    # basis nodes
    lower = minimum(basis1dmicro.nodes1d)
    width = basis1dmicro.volume
    nodes1dmacro = zeros(numbernodes1dmacro)
    micronodes = (basis1dmicro.nodes1d .- lower) ./ numberelements1d
    for i = 1:numberelements1d
        nodes1dmacro[(i-1)*numbernodes1d-i+2:i*numbernodes1d-i+1] =
            micronodes .+ lower .+ width / numberelements1d * (i - 1)
    end

    # basis quadrature points and weights
    quadratureweights1dmacro = []
    if overlapquadraturepoints
        quadratureweights1dmacro = zeros(numberquadraturepoints1dmacro)
    else
        quadratureweights1dmacro =
            kron(ones(numberelements1d), basis1dmicro.quadratureweights1d)
    end
    quadraturepoints1dmacro = zeros(numberquadraturepoints1dmacro)
    microquadraturepoints = (basis1dmicro.quadraturepoints1d .- lower) ./ numberelements1d
    for i = 1:numberelements1d
        offset = overlapquadraturepoints ? -i + 1 : 0
        quadraturepoints1dmacro[(i-1)*numberquadraturepoints1d+offset+1:i*numberquadraturepoints1d+offset] =
            microquadraturepoints .+ lower .+ width / numberelements1d * (i - 1)
    end

    # basis operations
    interpolation1dmacro = spzeros(numberquadraturepoints1dmacro, numbernodes1dmacro)
    gradient1dmacro = spzeros(numberquadraturepoints1dmacro, numbernodes1dmacro)
    for i = 1:numberelements1d
        offset = overlapquadraturepoints ? -i + 1 : 0
        interpolation1dmacro[
            (i-1)*numberquadraturepoints1d+offset+1:i*numberquadraturepoints1d+offset,
            (i-1)*numbernodes1d-i+2:i*numbernodes1d-i+1,
        ] = basis1dmicro.interpolation1d
        gradient1dmacro[
            (i-1)*numberquadraturepoints1d+offset+1:i*numberquadraturepoints1d+offset,
            (i-1)*numbernodes1d-i+2:i*numbernodes1d-i+1,
        ] = basis1dmicro.gradient1d
    end

    # use basic constructor
    return TensorBasis(
        numbernodes1dmacro,
        numberquadraturepoints1dmacro,
        numbercomponents,
        dimension,
        nodes1dmacro,
        quadraturepoints1dmacro,
        quadratureweights1dmacro,
        interpolation1dmacro,
        gradient1dmacro,
        numberelements1d = numberelements1d,
    )
end

"""
```julia
TensorH1LagrangeMacroBasis(
    numbernodes1d,
    numberquadraturepoints1d,
    numbercomponents,
    dimension,
    numberelements1d,
    collocatedquadrature = false,
)
```

Tensor product macro-element basis on Gauss-Legendre-Lobatto points with
  Gauss-Legendre (default) or Gauss-Legendre-Lobatto quadrature points

# Arguments:
- `numbernodes1d`:                 number of Gauss-Legendre-Lobatto nodes in 1 dimension
- `numberquadraturepoints1d`:      number of quadrature points in 1 dimension
- `numbercomponents`:              number of components
- `dimension`:                     dimension of basis
- `numberelements1d`:              number of elements in macro-element

# Keyword Arguments:
- `collocatedquadrature = false`:  Gauss-Legendre or Gauss-Legendre-Lobatto quadrature points,
                                     default: false, Gauss-Legendre-Lobatto

# Returns:
- H1 Lagrange tensor product macro-element basis object

# Example:
```jldoctest
# generate H1 Lagrange tensor macro-element product basis
basis = TensorH1LagrangeMacroBasis(4, 4, 1, 2, 2);

# generate basis with Gauss-Legendre quadrature points
basis = TensorH1LagrangeMacroBasis(4, 4, 1, 2, 2; collocatedquadrature=true);

# verify
println(basis)

# output
macro-element tensor product basis:
    numbernodes1d: 7
    numberquadraturepoints1d: 8
    numbercomponents: 1
    numberelements1d: 2
    dimension: 2
```
"""
function TensorH1LagrangeMacroBasis(
    numbernodes1d::Int,
    numberquadraturepoints1d::Int,
    numbercomponents::Int,
    dimension::Int,
    numberelements1d::Int;
    collocatedquadrature::Bool = false,
)
    basis1dmicro = TensorH1LagrangeBasis(
        numbernodes1d,
        numberquadraturepoints1d,
        numbercomponents,
        1,
        collocatedquadrature = collocatedquadrature,
    )
    # use common constructor
    return TensorMacroElementBasisFrom1D(
        numbernodes1d,
        numberquadraturepoints1d,
        numbercomponents,
        dimension,
        numberelements1d,
        basis1dmicro,
    )
end

"""
```julia
TensorH1UniformMacroBasis(
    numbernodes1d,
    numberquadraturepoints1d,
    numbercomponents,
    dimension,
    numberelements1d,
)
```

Tensor product macro-element basis on uniformly points with Gauss-Legendre quadrature

# Arguments:
- `numbernodes1d`:             number of uniformly spaced nodes in 1 dimension
- `numberquadraturepoints1d`:  number of Gauss-Legendre quadrature points in 1 dimension
- `numbercomponents`:          number of components
- `dimension`:                 dimension of basis
- `numberelements1d`:          number of elements in macro-element

# Returns:
- H1 Lagrange tensor product macro-element basis on uniformly space nodes object

# Example:
```jldoctest
# generate H1 Lagrange tensor product macro-element basis on uniformly spaced nodes
basis = TensorH1UniformMacroBasis(4, 3, 1, 2, 2);

# verify
println(basis)

# output
macro-element tensor product basis:
    numbernodes1d: 7
    numberquadraturepoints1d: 6
    numbercomponents: 1
    numberelements1d: 2
    dimension: 2
```
"""
function TensorH1UniformMacroBasis(
    numbernodes1d::Int,
    numberquadraturepoints1d::Int,
    numbercomponents::Int,
    dimension::Int,
    numberelements1d::Int,
)
    basis1dmicro =
        TensorH1UniformBasis(numbernodes1d, numberquadraturepoints1d, numbercomponents, 1)
    # use common constructor
    return TensorMacroElementBasisFrom1D(
        numbernodes1d,
        numberquadraturepoints1d,
        numbercomponents,
        dimension,
        numberelements1d,
        basis1dmicro,
    )
end

# ------------------------------------------------------------------------------
# -- p multigrid interpolation bases
# ------------------------------------------------------------------------------

"""
```julia
TensorH1LagrangePProlongationBasis(
    numbercoarsenodes1d,
    numberfinenodes1d,
    numbercomponents,
    dimension,
)
```

Tensor product p-prolongation basis on Gauss-Legendre-Lobatto points

# Arguments:
- `numbercoarsenodes1d`:  number of coarse grid Gauss-Legendre-Lobatto nodes in 1 dimension
- `numberfinenodes1d`:    number of fine grid Gauss-Legendre-Lobatto nodes in 1 dimension
- `numbercomponents`:     number of components
- `dimension`:            dimension of basis

# Returns:
- H1 Lagrange tensor product p-prolongation basis object

# Example:
```jldoctest
# generate H1 Lagrange tensor product p-prolongation basis
basisctof = TensorH1LagrangePProlongationBasis(2, 3, 1, 2);

# verify
println(basisctof)

# output
tensor product basis:
    numbernodes1d: 2
    numberquadraturepoints1d: 3
    numbercomponents: 1
    dimension: 2
```
"""
function TensorH1LagrangePProlongationBasis(
    numbercoarsenodes1d::Int,
    numberfinenodes1d::Int,
    numbercomponents::Int,
    dimension::Int,
)
    return TensorH1LagrangeBasis(
        numbercoarsenodes1d,
        numberfinenodes1d,
        numbercomponents,
        dimension,
        collocatedquadrature = true,
    )
end

# ------------------------------------------------------------------------------
# -- h multigrid interpolation bases
# ------------------------------------------------------------------------------

"""
```julia
TensorHProlongationBasis(
    coarsenodes1d,
    finenodes1d,
    numbercomponents,
    dimension,
    numberfineelements1d,
)
```

Tensor product h-prolongation basis

# Arguments:
- `coarsenodes1d`:         coarse grid node coordinates in 1 dimension
- `finenodes1d`:           fine grid node coordinates in 1 dimension
- `numbercomponents`:      number of components
- `dimension`:             dimension of basis
- `numberfineelements1d`:  number of fine grid elements

# Returns:
- H1 Lagrange tensor product h-prolongation basis object
"""
function TensorHProlongationBasis(
    coarsenodes1d::AbstractArray{Float64,1},
    finenodes1d::AbstractArray{Float64,1},
    numbercomponents::Int,
    dimension::Int,
    numberfineelements1d::Int,
)
    # compute dimensions
    numbercoarsenodes1d = maximum(size(coarsenodes1d))
    numberfinenodes1d = maximum(size(finenodes1d))

    # form coarse to fine basis
    interpolation1d, gradient1d = buildinterpolationandgradient(coarsenodes1d, finenodes1d)
    quadratureweights1d = zeros(numberfinenodes1d)
    return TensorBasis(
        numbercoarsenodes1d,
        numberfinenodes1d,
        numbercomponents,
        dimension,
        coarsenodes1d,
        finenodes1d,
        quadratureweights1d,
        interpolation1d,
        gradient1d,
    )
end

"""
```julia
TensorH1LagrangeHProlongationBasis(
    numbernodes1d,
    numbercomponents,
    dimension,
    numberfineelements1d,
)
```

Tensor product h-prolongation basis on Gauss-Legendre-Lobatto points

# Arguments:
- `numbernodes1d`:         number of Gauss-Legendre-Lobatto nodes in 1 dimension per element
- `numbercomponents`:      number of components
- `dimension`:             dimension of basis
- `numberfineelements1d`:  number of fine grid elements

# Returns:
- H1 Lagrange tensor product h-prolongation basis object

# Example:
```jldoctest
# generate H1 Lagrange tensor product h-prolongation basis
basis = TensorH1LagrangeHProlongationBasis(4, 3, 2, 2);

# verify
println(basis)

# output
tensor product basis:
    numbernodes1d: 4
    numberquadraturepoints1d: 7
    numbercomponents: 3
    dimension: 2
```
"""
function TensorH1LagrangeHProlongationBasis(
    numbernodes1d::Int,
    numbercomponents::Int,
    dimension::Int,
    numberfineelements1d::Int,
)
    # generate nodes
    nodescoarse1d = gausslobattoquadrature(numbernodes1d, false)
    nodesfine1d = zeros((numbernodes1d - 1) * numberfineelements1d + 1)
    for i = 1:numberfineelements1d
        nodesfine1d[(i-1)*numbernodes1d-i+2:i*numbernodes1d-i+1] =
            nodescoarse1d ./ numberfineelements1d .+
            ((2 * (i - 1) + 1) / numberfineelements1d - 1)
    end

    # single coarse to multiple fine elements
    return TensorHProlongationBasis(
        nodescoarse1d,
        nodesfine1d,
        numbercomponents,
        dimension,
        numberfineelements1d,
    )
end

"""
```julia
TensorH1UniformHProlongationBasis(
    numbernodes1d,
    numbercomponents,
    dimension,
    numberfineelements1d,
)
```

Tensor product h-prolongation basis on uniformly spaced points

# Arguments:
- `numbernodes1d`:         number of uniformly spaced nodes per element
- `numbercomponents`:      number of components
- `dimension`:             dimension of basis
- `numberfineelements1d`:  number of fine grid elements

# Returns:
- H1 Lagrange tensor product h-prolongation basis on uniformly spaced nodes object

# Example:
```jldoctest
# generate H1 Lagrange tensor product h-prolongation basis on uniformly spaced nodes
basis = TensorH1UniformHProlongationBasis(4, 3, 2, 2);

# verify
println(basis)

# output
tensor product basis:
    numbernodes1d: 4
    numberquadraturepoints1d: 7
    numbercomponents: 3
    dimension: 2
```
"""
function TensorH1UniformHProlongationBasis(
    numbernodes1d::Int,
    numbercomponents::Int,
    dimension::Int,
    numberfineelements1d::Int,
)
    # generate nodes
    nodescoarse1d = [-1.0:(2.0/(numbernodes1d-1)):1.0...]
    nodesfine1d = [-1.0:(2.0/((numbernodes1d-1)*numberfineelements1d)):1.0...]
    for i = 1:numberfineelements1d
        nodesfine1d[(i-1)*numbernodes1d-i+2:i*numbernodes1d-i+1] =
            nodescoarse1d ./ numberfineelements1d .+
            ((2 * (i - 1) + 1) / numberfineelements1d - 1)
    end

    # single coarse to multiple fine elements
    return TensorHProlongationBasis(
        nodescoarse1d,
        nodesfine1d,
        numbercomponents,
        dimension,
        numberfineelements1d,
    )
end

"""
```julia
TensorH1LagrangeHProlongationMacroBasis(
    numbernodes1d,
    numbercomponents,
    dimension,
    numbercoarseelements1d,
    numberfineelements1d,
)
```

Tensor product macro-element h-prolongation basis on Gauss-Legendre-Lobatto points

# Arguments:
- `numbernodes1d`:           number of Gauss-Legendre-Lobatto nodes in 1 dimension per element
- `numbercomponents`:        number of components
- `dimension`:               dimension of basis
- `numbercoarseelements1d`:  number of coarse grid elements in macro-element
- `numberfineelements1d`:    number of fine grid elements in macro-element

# Returns:
- H1 Lagrange tensor product h-prolongation macro-element basis object

# Example:
```jldoctest
# generate H1 Lagrange tensor product h-prolongation macro-element basis
basis = TensorH1LagrangeHProlongationMacroBasis(4, 1, 2, 2, 4);

# verify
println(basis)

# output
macro-element tensor product basis:
    numbernodes1d: 7
    numberquadraturepoints1d: 13
    numbercomponents: 1
    numberelements1d: 2
    dimension: 2
```
"""
function TensorH1LagrangeHProlongationMacroBasis(
    numbernodes1d::Int,
    numbercomponents::Int,
    dimension::Int,
    numbercoarseelements1d::Int,
    numberfineelements1d::Int,
)
    # validate inputs
    if numberfineelements1d % numbercoarseelements1d != 0
        # COV_EXCL_START
        throw(
            DomanError(
                numberfineelements1d,
                "numberfineelements1d must be a multiple of numbercoarseelements1d",
            ),
        )
        # COV_EXCL_STOP
    end

    # single coarse to multiple fine elements
    scale = Int(numberfineelements1d / numbercoarseelements1d)
    hprolongationbasis1dmicro =
        TensorH1LagrangeHProlongationBasis(numbernodes1d, numbercomponents, 1, scale)

    # use common constructor
    return TensorMacroElementBasisFrom1D(
        numbernodes1d,
        (numbernodes1d - 1) * scale + 1,
        numbercomponents,
        dimension,
        numbercoarseelements1d,
        hprolongationbasis1dmicro,
        overlapquadraturepoints = true,
    )
end

"""
```
TensorH1UniformHProlongationMacroBasis(
    numbernodes1d,
    numbercomponents,
    dimension,
    numbercoarseelements1d,
    numberfineelements1d,
)
```

Tensor product macro-element h-prolongation basis on uniformly spaced points

# Arguments:
- `numbernodes1d`:           number of uniformly spaced nodes per element
- `numbercomponents`:        number of components
- `dimension`:               dimension of basis
- `numbercoarseelements1d`:  number of coarse grid elements in macro-element
- `numberfineelements1d`:    number of fine grid elements in macro-element

# Returns:
- H1 Lagrange tensor product h-prolongation macro-element basis on uniformly spaced nodes object

# Example:
```jldoctest
# generate H1 Lagrange tensor product h-prolongation macro-element basis on uniformly spaced nodes
basis = TensorH1UniformHProlongationMacroBasis(4, 1, 2, 2, 4);

# verify
println(basis)

# output
macro-element tensor product basis:
    numbernodes1d: 7
    numberquadraturepoints1d: 13
    numbercomponents: 1
    numberelements1d: 2
    dimension: 2
```
"""
function TensorH1UniformHProlongationMacroBasis(
    numbernodes1d::Int,
    numbercomponents::Int,
    dimension::Int,
    numbercoarseelements1d::Int,
    numberfineelements1d::Int,
)
    # validate inputs
    if numberfineelements1d % numbercoarseelements1d != 0
        # COV_EXCL_START
        throw(
            DomanError(
                numberfineelements1d,
                "numberfineelements1d must be a multiple of numbercoarseelements1d",
            ),
        )
        # COV_EXCL_STOP
    end

    # single coarse to multiple fine elements
    scale = Int(numberfineelements1d / numbercoarseelements1d)
    hprolongationbasis1dmicro =
        TensorH1UniformHProlongationBasis(numbernodes1d, numbercomponents, 1, scale)

    # use common constructor
    return TensorMacroElementBasisFrom1D(
        numbernodes1d,
        (numbernodes1d - 1) * scale + 1,
        numbercomponents,
        dimension,
        numbercoarseelements1d,
        hprolongationbasis1dmicro,
        overlapquadraturepoints = true,
    )
end

# ------------------------------------------------------------------------------


text:
using Test, WAVI

@testset "Models" begin
    @info "Testing models...."

    @testset "testing model construction" begin
    @info "Testing generic model construction..."
    grid = Grid()
    bed_elevation = zeros(grid.nx, grid.ny)
    model = Model(grid = grid, bed_elevation = bed_elevation)
    @test model isa Model

    #check that a scalar weertman c passed to model emerges as an array
    model = Model(grid = grid, bed_elevation = bed_elevation, params = Params(weertman_c = 1.0))
    @test model.params.weertman_c == 1.0 *ones(grid.nx, grid.ny)

    #check that an array weertman_c works
    model = Model(grid = grid, bed_elevation = bed_elevation, params = Params(weertman_c =1.0*ones(grid.nx, grid.ny)))
    @test model isa Model

    #check that a scalar weertman c passed to model emerges as an array
    model = Model(grid = grid, bed_elevation = bed_elevation, params = Params(accumulation_rate = 1.0))
    @test model.params.accumulation_rate == 1.0 *ones(grid.nx, grid.ny)

    #check that an array weertman_c works
    model = Model(grid = grid, bed_elevation = bed_elevation, params = Params(accumulation_rate =1.0*ones(grid.nx, grid.ny)))
    @test model isa Model
    end

    @testset "Initial Conditions" begin
    @info "Testing initial conditions"
    #passing no initial conditions should revert to default values
    grid = Grid()
    params = Params()
    bed_elevation = zeros(grid.nx, grid.ny)
    model = Model(grid = grid, bed_elevation = bed_elevation)
    @test model.initial_conditions.initial_thickness == params.default_thickness*ones(grid.nx, grid.ny)
    @test model.initial_conditions.initial_damage == params.default_damage*ones(grid.nx, grid.ny, grid.nσ)
    @test model.initial_conditions.initial_viscosity == params.default_viscosity*ones(grid.nx, grid.ny, grid.nσ)
    @test model.initial_conditions.initial_temperature == params.default_temperature*ones(grid.nx, grid.ny, grid.nσ)

    #test passing arrays of thickness, damage, temperature, and viscosity
    initial_conditions = InitialConditions(initial_thickness = 3.14159 * ones(grid.nx,grid.ny),
                                            initial_viscosity = 2.7182818 * ones(grid.nx, grid.ny, grid.nσ),
                                            initial_temperature = 1.618034 * ones(grid.nx, grid.ny, grid.nσ),
                                            initial_damage = 1.414213 * ones(grid.nx, grid.ny, grid.nσ))
    model = Model(grid = grid, bed_elevation = bed_elevation, initial_conditions = initial_conditions)
    @test model isa Model
    @test model.initial_conditions.initial_thickness == 3.14159 * ones(grid.nx,grid.ny)
    @test model.initial_conditions.initial_viscosity == 2.7182818 * ones(grid.nx, grid.ny, grid.nσ)
    @test model.initial_conditions.initial_temperature == 1.618034 * ones(grid.nx, grid.ny, grid.nσ)
    @test model.initial_conditions.initial_damage == 1.414213 * ones(grid.nx, grid.ny, grid.nσ)

    #check fields successfully passed to fields
    @test model.fields.gh.h == 3.14159 * ones(grid.nx,grid.ny)
    @test model.fields.g3d.η == 2.7182818 * ones(grid.nx, grid.ny, grid.nσ)
    @test model.fields.g3d.θ == 1.618034 * ones(grid.nx, grid.ny, grid.nσ)
    @test model.fields.g3d.Φ ==  1.414213 * ones(grid.nx, grid.ny, grid.nσ)

    end

    @testset "Initial Conditions Errors" begin
    @info "Testing initial conditions errors"
    #size incompatibilities:
    grid = Grid()
    bed_elevation = zeros(grid.nx, grid.ny)
    initial_conditions = InitialConditions(initial_thickness = ones((grid.nx + 1), grid.ny))
    @test_throws DimensionMismatch Model(grid = grid, bed_elevation = bed_elevation, initial_conditions = initial_conditions)

    initial_conditions = InitialConditions(initial_viscosity = ones((grid.nx + 1), grid.ny, grid.nσ))
    @test_throws DimensionMismatch Model(grid = grid, bed_elevation = bed_elevation, initial_conditions = initial_conditions)

    initial_conditions = InitialConditions(initial_temperature = ones((grid.nx + 1), grid.ny, grid.nσ))
    @test_throws DimensionMismatch Model(grid = grid, bed_elevation = bed_elevation, initial_conditions = initial_conditions)

    initial_conditions = InitialConditions(initial_damage = ones((grid.nx + 1), grid.ny, grid.nσ))
    @test_throws DimensionMismatch Model(grid = grid, bed_elevation = bed_elevation, initial_conditions = initial_conditions)

    end

    @testset "Passing sigma levels" begin
    @info "Testing σ level passing"
    σ = [0.0, 0.12, 0.13, 0.5, 0.7]
    nσ = length(σ)
    grid = Grid(nσ = nσ, σ = σ)
    bed_elevation = zeros(grid.nx, grid.ny)
    model = Model(grid = grid, bed_elevation = bed_elevation)
    @test model.grid.σ == σ
    @test model.fields.g3d.σ == σ
    @test_throws DimensionMismatch Model(grid = Grid(nσ = 6, σ = σ), bed_elevation = bed_elevation) #size incompatiblity in the sigma levels
    end

    @testset "testing model construction errors"  begin
    @info "Testing model construction errors"
    #test dimnesion mismatch if size of input weertman_c incompatible
    grid = Grid()
    bed_elevation = zeros(grid.nx, grid.ny)
    @test_throws DimensionMismatch Model(grid = grid, bed_elevation = bed_elevation, params = Params(weertman_c =1.0*ones(grid.nx -1, grid.ny -1)))

    end

    @testset "Testing generic velocity solve on a model" begin
    @info "Testing generic model velocity solve..."
    grid = Grid()
    bed_elevation = zeros(grid.nx, grid.ny)
    model = Model(grid = grid, bed_elevation = bed_elevation)
    update_state!(model)
    @test model isa Model
    end

end


text:
module Realtimeplot
    using Plots

    # prep the plots
    plt1 = histogram([0],label=nothing) #plot1
    plt2 = plot([], [],label=nothing) #plot2
    ylabel!("Plaquette")
    xlabel!("MC time")
    plt3 = histogram([0],label=nothing) #plot3
    plt4 = plot([], [],label=nothing) #plot4
    ylabel!("Polyakov loop")
    xlabel!("MC time")
    plot(plt1, plt2, plt3, plt4, layout = 4)

    function mock() #mock data
        plaq = 0.6+0.1*randn()
        if rand()<0.5
            poly = -1.5+randn()/2
        else
            poly = 1.5+randn()/2
        end
        sleep(0.01)
        return plaq,poly
    end

    hist_plaq=[]
    hist_poly=[]

    function plot_refresh!(plt1,plt2,plt3,plt4,hist_plaq,hist_poly, plaq, poly,itrj)
        bins=round(Int,log(itrj)*5+1)
        append!(hist_plaq, plaq)
        append!(hist_poly, poly)
        #
        plt1 = histogram(hist_plaq,bins=bins,label=nothing) #plot1
        xlabel!("Plaquette")
        plt3 = histogram(hist_poly,bins=bins,label=nothing) #plot3
        xlabel!("Polyakov loop")
        #
        push!(plt2, 1, itrj, plaq)
        push!(plt4, 1, itrj, poly)
        #
        plot(plt1, plt2, plt3, plt4, layout = 4)
        gui()
    end

    Ntrj=1000
    for itrj in 1:Ntrj
        plaq,poly=mock()
        #
        plot_refresh!(plt1,plt2,plt3,plt4,hist_plaq,hist_poly, plaq, poly,itrj)
        #
    end
#=
plt = plot([0,0.1], Any[rand(2),sin])
for x in 0.2:0.1:π
    push!(plt, 1, x, rand())
    push!(plt, 2, x, sin(x))
    gui(); sleep(0.5)
end
=#

#=
import PyPlot: plt

x = [1,2,3]
y = [1,2,3]

plt.plot(x,y)
plt.show()
=#
end


text:
# p28b.jl - eigenmodes of Laplacian on the disk

# r coordinate, ranging from -1 to 1 (N must be odd)
N = 25; N2 = Int((N-1)/2);
(D,r) = cheb(N); D2 = D^2;
D1 = D2[2:N2+1,2:N2+1]; D2 = D2[2:N2+1,N:-1:N2+2];
E1 =  D[2:N2+1,2:N2+1]; E2 =  D[2:N2+1,N:-1:N2+2];

# t = theta coordinate, ranging from 0 to 2*pi (M must be even):
M = 20; dt = 2*pi/M; t = dt*(1:M); M2 = Int(M/2);
D2t = toeplitz([-pi^2/(3*dt^2)-1/6; @. .5*(-1)^(2:M)/sin(dt*(1:M-1)/2)^2]);

# Laplacian in polar coordinates:
R = diagm(1./r[2:N2+1]);
Z = zeros(M2,M2); I = eye(M2);
L = kron(D1+R*E1,eye(M)) + kron(D2+R*E2,[Z I;I Z]) + kron(R^2,D2t);

# Compute 25 eigenmodes:
index = 1:25;
(Lam,V) = eig(-L); ii = sortperm(abs.(Lam))[index];
Lam = Lam[ii]; V = V[:,ii];
Lam = sqrt.(real(Lam/Lam[1]));

# Plot nodal lines:
(rr,tt) = (r[1:N2+1],[0;t]);
(xx,yy) = @. (cos(tt)*rr',sin(tt)*rr');
z = exp.(1im*pi*(-100:100)/100);  clf();
for i = 1:25
    subplot(5,5,i);
    u = reshape(real(V[:,i]),M,N2);
    u = [zeros(M+1) u[[M;1:M],:]];
    u = u/norm(u[:],Inf);
    plot(real(z),imag(z));
    xlim(-1.07,1.07); ylim(-1.07,1.07); axis("off"); axis("equal");
    contour(xx,yy,u,levels=[0]);
    title("$(signif(Lam[i],5))",fontsize=8);
end


text:
# Generation of lar model (V,EV) of a random set of 2D lines

using LAR

V = [[0.,1] [1.,1] [-1,0] [2,0] [0,-1] [1,-1]]
EV = [[1,5] [2,6] [3,4]]
lineArray = V,EV
lineFrags = lineIntersection(lineArray)
W,EW = lines2lar(lineArray)
larview(W,EW)
viewexploded(W,EW)

lineArray = randomLines(200,0.4);
V,EV = lines2lar(lineArray);
hpcLines = lar2hpc(V,EV);
marker = p.CIRCLE(.002)([3,1]);
verts = PyObject(PyObject(V')[:tolist]());
markers = p.STRUCT(p.CONS(p.AA(p.T([1,2]))(verts))(marker));
redMarkers = p.COLOR(p.RED)(markers);
p.VIEW(p.STRUCT([hpcLines,redMarkers]))
viewexploded(V,EV)


text:
module NewtonsMethod

using ForwardDiff, LinearAlgebra

export newtonroot

function iterativesolver(f, x₀, checkconvergence, updatex)
    xold = x₀.*10
    fold = 10*f(x₀)
    xnew = x₀
    fnew = f(xnew)
    iter = 1
    while !checkconvergence(xold, xnew, fold, fnew, iter)
        xold = xnew
        xnew = updatex(f, xold)
        fold = fnew
        fnew = f(xnew)
        iter += 1
    end
    return(value = xnew, iter=iter)
end

function newtonroot(f, f′; x₀, tol=1e-7, maxiter=1000)
    res = iterativesolver(f,
                          x₀,
                          (xold, xnew, fold, fnew, iter) -> norm(xnew - xold) < 1e-4 || iter > maxiter,
                          (f, x) -> x - f(x) / f′(x))
    if res.iter <= maxiter
        return res
    else
        return (value = nothing, iter=res.iter)
    end
end

function newtonroot(f; x₀, tol=1e-7, maxiter=1000)
    return newtonroot(f, x -> ForwardDiff.derivative(f, x); x₀, tol, maxiter)
end

end


text:
using StaticArrays
using Parameters
using OrdinaryDiffEq
using LaserTypes
using LinearAlgebra
using Statistics
using ThreadsX
using Plots
using LaTeXStrings
using Distributions
using RecursiveArrayTools
using Serialization

function folder_setup(path)
    path = replace(path,"/" => "")

    if split(pwd(),"/")[end] == path
        println("already in $path")
        return
    end

    if !isdir(path)
        mkdir(path)
        println("making folder $path")
    end

    cd(path)
    println("moved in $path")
end

cd(@__DIR__)
folder_setup("output5")

const N = 28_750 # number of particles, for each ring, 4N for radius averaging
const c = 137.036 # speed of light
const ω = 0.057; const T₀ = 2π/ω # angular frequency, period
const k = ω/c # wavenumber
const λ = c*T₀ # wavelength
const w₀ = 5 * λ # is the waist radius
const a₀ = 2.; const AA = a₀*c # (qA₀)/(mc) = a₀
const w = 10.; const τ = w/ω; # temporal pulse decay time given in terms of number of oscilations and frequency
const n = 5 # number of periods to integrate before and after pulse collides with particles
const π₀ = a₀*c # in atomic units, a₀mc has units of linear momentum and sets the scale for linear momentum transfered to the particle, not that m = 1 for our particles
const maxR = 3.25w₀ # maximum radius for distributing praticles

const g = @SMatrix [1  0  0  0;
                    0 -1  0  0;
                    0  0 -1  0;
                    0  0  0 -1] # Minkowski space metric

@with_kw struct InterPars{qq,mm,typo,rewts}
    q::qq = -1
    m::mm = 1
    type::typo
    roots::rewts = Float64[]
end

################################################################################
#Right Circularly Polarized

# p = 0
CPLGTYPE00 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 0, p = 0,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
CPLGTYPE0_1 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = -1, p = 0,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
CPLGTYPE01 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 1, p = 0,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))

# p = 1
CPLGTYPE10 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 0, p = 1,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
CPLGTYPE1_1 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = -1, p = 1,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
CPLGTYPE11 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 1, p = 1,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))

# p = 2
CPLGTYPE2_1 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = -1, p = 2,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
CPLGTYPE21 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 1, p = 2,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
CPLGTYPE20 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 0, p = 2,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
CPLGTYPE2_2 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = -2, p = 2,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
CPLGTYPE22 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 2, p = 2,  ξx = (1. + 0im)/√2, ξy = (1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))

################################################################################
#Left Circularly Polarized

# p = 0
MCPLGTYPE00 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 0, p = 0,  ξx = (1. + 0im)/√2, ξy = (-1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
MCPLGTYPE0_1 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = -1, p = 0,  ξx = (1. + 0im)/√2, ξy = (-1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
MCPLGTYPE01 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 1, p = 0,  ξx = (1. + 0im)/√2, ξy = (-1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))

# p = 1
MCPLGTYPE10 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 0, p = 1,  ξx = (1. + 0im)/√2, ξy = (-1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
MCPLGTYPE1_1 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = -1, p = 1,  ξx = (1. + 0im)/√2, ξy = (-1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
MCPLGTYPE11 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 1, p = 1,  ξx = (1. + 0im)/√2, ξy = (-1im)/√2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))

################################################################################
#Linearly Polarized

# p = 0
LPLGTYPE00 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 0, p = 0, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
LPLGTYPE0_1 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = -1, p = 0, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
LPLGTYPE01 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 1, p = 0, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))

# p = 1
LPLGTYPE10 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 0, p = 1, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
LPLGTYPE1_1 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = -1, p = 1, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
LPLGTYPE11 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 1, p = 1, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))

# p = 2
LPLGTYPE20 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 0, p = 2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
LPLGTYPE21 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 1, p = 2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
LPLGTYPE22 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = 2, p = 2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
LPLGTYPE2_1 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = -1, p = 2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))
LPLGTYPE2_2 = LaguerreGaussLaser(c = c, m_q = 1., q = -1., λ = λ, a₀ = a₀, ϕ₀ = -π/2,
    w₀ = w₀, m = -2, p = 2, profile = GaussProfile(c = c, τ = w/ω, z₀ = 0.))

field_param = (
    CPLGTYPE0_1 = InterPars(type = CPLGTYPE0_1, roots = [0., 1.954]),
    CPLGTYPE00 = InterPars(type = CPLGTYPE00, roots = [0., 1.518]),
    CPLGTYPE01 = InterPars(type = CPLGTYPE01, roots = [0., 1.954]),
    CPLGTYPE1_1 = InterPars(type = CPLGTYPE1_1, roots = [0., 1., 2.566]),
    CPLGTYPE10 = InterPars(type = CPLGTYPE10, roots = [ 0., 0.707, 2.323]),
    CPLGTYPE11 = InterPars(type = CPLGTYPE11, roots = [0., 1., 2.566]),
    CPLGTYPE2_2 = InterPars(type = CPLGTYPE2_2, roots = [0., 1.000, 1.731, 3.176]),
    CPLGTYPE2_1 = InterPars(type = CPLGTYPE2_1, roots = [0., 0.797, 1.538, 3.005]),
    CPLGTYPE20 = InterPars(type = CPLGTYPE20, roots = [0., 0.541, 1.306, 2.807]),
    CPLGTYPE21 = InterPars(type = CPLGTYPE21, roots = [0., 0.797, 1.538, 3.005]),
    CPLGTYPE22 = InterPars(type = CPLGTYPE22, roots = [0., 1.000, 1.731, 3.176]),
    MCPLGTYPE0_1 = InterPars(type = MCPLGTYPE0_1, roots = [0., 1.954]),
    MCPLGTYPE00 = InterPars(type = MCPLGTYPE00, roots = [0., 1.518]),
    MCPLGTYPE01 = InterPars(type = MCPLGTYPE01, roots = [0., 1.954]),
    MCPLGTYPE1_1 = InterPars(type = MCPLGTYPE1_1, roots = [0., 1., 2.566]),
    MCPLGTYPE10 = InterPars(type = MCPLGTYPE10, roots = [0., 0.707, 2.323]),
    MCPLGTYPE11 = InterPars(type = MCPLGTYPE11, roots = [0., 1., 2.566]),
    LPLGTYPE0_1 = InterPars(type = LPLGTYPE0_1, roots = [0., 1.954]),
    LPLGTYPE00 = InterPars(type = LPLGTYPE00, roots = [0., 1.518]),
    LPLGTYPE01 = InterPars(type = LPLGTYPE01, roots = [0., 1.954]),
    LPLGTYPE1_1 = InterPars(type = LPLGTYPE1_1, roots = [0., 1., 2.566]),
    LPLGTYPE10 = InterPars(type = LPLGTYPE10, roots = [0., 0.707, 2.323]),
    LPLGTYPE11 = InterPars(type = LPLGTYPE11, roots = [0., 1., 2.566]),
    LPLGTYPE2_2 = InterPars(type = LPLGTYPE2_2, roots = [0., 1.000, 1.731, 3.176]),
    LPLGTYPE2_1 = InterPars(type = LPLGTYPE2_1, roots = [0., 0.797, 1.538, 3.005]),
    LPLGTYPE20 = InterPars(type = LPLGTYPE20, roots = [0., 0.541, 1.306, 2.807]),
    LPLGTYPE21 = InterPars(type = LPLGTYPE21, roots = [0., 0.797, 1.538, 3.005]),
    LPLGTYPE22 = InterPars(type = LPLGTYPE22, roots = [0., 1.000, 1.731, 3.176])
)

function dynEq(v,param,τ)
    @unpack q, m, type = param
    x = v[1:4]
    u = v[5:8]

    dx = u
    du = q/m * Fμν(x, type) * g * u

    return vcat(dx,du)
end

function Mμν(v)
    x⁰, x¹, x², x³, u⁰, u¹, u², u³ = v
    Lx = x²*u³ - x³*u²
    Ly = x³*u¹ - x¹*u³
    Lz = x¹*u² - x²*u¹
    # Nx = x¹*u⁰ - x⁰*u¹
    # Ny = x²*u⁰ - x⁰*u²
    # Nz = x³*u⁰ - x⁰*u³
    L² = Lx^2 + Ly^2 + Lz^2
    # N² = Nx^2 + Ny^2 + Nz^2
    #the physical quantities are m* [1,2,3] , m^2*[4], m/c* [5,6,7] , m^2/c^2*[8],
    return (Lx = Lx, Ly = Ly, Lz = Lz, L² = L², #Nx = Nx, Ny = Ny, Nz = Nz, N² = N²,
            x⁰ = x⁰, x¹ = x¹, x² = x², x³ = x³, u⁰ = u⁰, u¹ = u¹, u² = u², u³ = u³)
end

function σA(A,A²)
    .√abs.((A²-A.^2))
end

function Mμντ(X)
    [Mμν(x) for x in X.u]
end

function UniformAnnulus(r₁,r₂)
    θ = rand(Uniform(0,2*π))
    u = rand(Uniform(r₁,r₂))
    v = rand(Uniform(0,r₁+r₂))
    ρ = v<u ? u : r₁+r₂ - u
    [SVector{3}(ρ*cos(θ), ρ*sin(θ),0), ρ]
end

RHO = [UniformAnnulus(0,maxR) for i in 1:4N]
RHOVec = VectorOfArray(RHO)
idx = sortperm(RHO, lt=(x,y)->x[2]<y[2])
R1 = RHOVec[1,:][idx]

τi = - n * τ
τf =   n * τ

const time_samples = 2000
const Δτ = range(0, stop = 2*n*τ/T₀, length = time_samples+1)

function xu_cat(X,U,T)
    x0 = c*T
    x = vcat([x0],X)
    u0 = √(c*c + U⋅U)
    u = vcat([u0],U)
    SVector{8}(vcat(x,u))
end

function problem_set(R0,p,U0 = [0,0,0])
    V0 = [xu_cat(r0,U0,τi) for r0 in R0]
    prob = ODEProblem(dynEq, V0[1], (τi, τf), p)
    function Ret_ODEProb(prob,i,repeat)
        ODEProblem(dynEq, V0[i], (τi, τf), p) #return
    end
    function output_func(sol,i)
        (Mμντ(sol), false) # put sol as argument, if you want the actual trajectories
    end
    EnsembleProblem(prob,prob_func = Ret_ODEProb, output_func = output_func)
end

function Particle_Avg(M,J)
    [ThreadsX.sum(getproperty(M[i][t], J)
    for i in 1:N)/N for t in 1:time_samples+1]
end

function Particle_Avg2(M,J)
    [ThreadsX.sum(getproperty(M[i][t], J)^2
    for i in 1:N)/N for t in 1:time_samples+1]
end

function Particle_AvgR(M)
    [ThreadsX.sum(hypot(getproperty(M[i][t], :x¹), getproperty(M[i][t], :x²))
    for i in 1:N)/N for t in 1:time_samples+1]
end

function Particle_AvgR²(M)
    [ThreadsX.sum(getproperty(M[i][t], :x¹)^2 + getproperty(M[i][t], :x²)^2
    for i in 1:N)/N for t in 1:time_samples+1]
end

# function replot!(plt, p)
#     plot!(plt, [], [], label="\$a_0=$a₀\$", color=nothing)
#     plot!(plt, [], [], label="\$p=$(p.p)\$", color=nothing)
#     plot!(plt, [], [], label="\$m=$(p.m)\$", color=nothing)
# end

rectangle(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])
circ = Shape(Plots.partialcircle(0, 2π))

function add_info!(plt, p, isscatter = false)
    #there is either a bug or confusing behaviour with ylims and aspect_ratio
    xl, yl = (isscatter ? ((-5.658368025329326, 5.667117711980884), (-3.662004592945377, 3.613120699860083)) : (xlims(plt), ylims(plt)))
    frame_width = xl[2] - xl[1]
    frame_height =yl[2] - yl[1]
    box_ratio = (w = 0.14 + 0.05 * abs(p.type.ξy*√2), h = 0.31) # hardcoded, I don't know how to get display length

    # lower left corner coordinates
    box = (x = (0.99 - box_ratio.w)*frame_width + xl[1],
           y = (0.99 - box_ratio.h)*frame_height + yl[1],
           w = box_ratio.w*frame_width,
           h = box_ratio.h*frame_height)

    # compute text positions
    text_x = box.x + box.w/2 # always centered on the box
    text_y = box.y + box.h # starting from the top, mind the minus
    δy = box.h/8

    # find the exact strings to print, related to polarization
    sgn = (imag(p.type.ξy) < 0 ? "-" : "")
    xi = (p.type.ξy == 0 ? "\$\\xi_y=0\$" : "\$\\xi_y= $sgn\\mathrm{i}/\\sqrt{2}\$")

    annotate!(plt, [(text_x, text_y - 1δy, Plots.text("\$a_0=$a₀\$",       :center, 10, "computer modern")),
                    (text_x, text_y - 3δy, Plots.text(xi,                  :center, 10, "computer modern")),
                    (text_x, text_y - 5δy, Plots.text("\$p=$(p.type.p)\$", :center, 10, "computer modern")),
                    (text_x, text_y - 7δy, Plots.text("\$m=$(p.type.m)\$", :center, 10, "computer modern"))])

    # add a rectangle without changing the frame size
    plot!(plt, rectangle(box.w, box.h, box.x, box.y);
                        xlims = xl, ylims = yl,
                        linewidth = 1.5,
                        fillcolor = :white, label = false)
end

function add_rect!(plt, X, Y, roots; kwargs...)
        h = ylims(plt)[2]-ylims(plt)[1]
        y₀ = ylims(plt)[1]
        underlay = plot([],[], label = false)
        for i in 1:(length(roots) - 1)
            a = roots[i]
            b = roots[i+1]
            plot!(underlay,rectangle(b-a, h, a, y₀),
            color_palette = :Set1_4, seriescolor = i,
            fillalpha = .4, linecolor = nothing, label = false)
        end
        plot!(underlay, X, Y;
            xlims = xlims(plt), ylims = ylims(plt), kwargs...)
end

common_settings = ( fontfamily = "computer modern",
                    dpi = 360,
                    label = false,
                    fillalpha=.3,
                    color_palette = :Set1_4,
                    linewidth = 2)

io = open("output_log.txt","w")

write(io,"""Simulation Paramters

Number of Particles per Ring,           N = $N
Speed of Light [Atomic Units],          c = $c
Angular Frequency [Atomic Units],       ω = $ω
Oscillation Period [Atomic Units],      T₀ = $T₀
Wavenumber [Atomic Units],              k = $k
Wavelength [Atomic Units],              λ = $λ
Beam Waist Radius [Atomic Units],       w₀ = $w₀
Reduced Vector Potential,               a₀ = $a₀
Number of Oscillations in the Pulse,    w = $w
Envelope Decay Time,                    τ = $τ
Maximum Linear Momentum Transfered,     π₀ = $π₀
Disc Radius for Particle Distribution,  Rₘₐₓ = $maxR\n
""") # save out important parameters

for (sim_name, p) in pairs(field_param)
    println("$sim_name")
    folder_setup("$sim_name") # change folder for each simulation, as there are many files
    folder_setup("particles")
    write(io, "$sim_name\n")

    plot_data = [plot([],[], label = false) for i in 1:7]

    for j in 1:(length(p.roots) - 1)

        R0 = [UniformAnnulus(w₀*p.roots[j], w₀*p.roots[j+1])[1] for i in 1:N]

        eprob = problem_set(R0, p)
        M = solve(eprob, Vern9(), EnsembleThreads(), abstol=1e-9, reltol=1e-9,
                saveat = (τf-τi)/time_samples, trajectories = N)

        z,  Lz,  L² = Particle_Avg.((M,),(:x³, :Lz, :L²))
        z², Lz², L⁴ = Particle_Avg2.((M,),(:x³, :Lz, :L²))

        ρ  = Particle_AvgR(M)
        ρ² = Particle_AvgR²(M)

        σz  = σA(z, z²)/w₀
        # σρ  = σA(ρ, ρ²)/w₀ # not worth computing because it's too wide and breaks plots
        σLz = σA(Lz, Lz²)/w₀/π₀
        σL² = σA(L², L⁴)/w₀/π₀/w₀/π₀

        write(io,
            "from $(p.roots[j]/w₀) to $(p.roots[j+1]/w₀)\n"*
            "z ± σz $(z[end]/w₀) ± $(σz[end])\n"*
            "Lz ± σLz $(Lz[end]/w₀/π₀) ± $(σLz[end])\n"*
            "L² ± σL² $(L²[end]/w₀/π₀/w₀/π₀) ± $(σL²[end])\n"
             )

        plot_data[1] = plot!(plot_data[1], Δτ, z/w₀;
                                common_settings...,
                                ribbon = σz,
                                seriescolor = j, # these cannot be moved outside due to variable scope
                                xlabel = L"\tau/T_0",
                                ylabel = L"\overline{z}/w_0")

        mean_ρ = 2/3*(p.roots[j+1]^3 - p.roots[j]^3)/(p.roots[j+1]^2 - p.roots[j]^2)

        plot_data[2] = plot!(plot_data[2], Δτ, ρ/w₀ .- mean_ρ;
                                common_settings...,
                                # ribbon = σρ, too big relative to ρ changes, makes plot unreadable
                                seriescolor = j,
                                xlabel = L"\tau/T_0",
                                ylabel = L"\Delta\overline{\rho}/w_0")

        plot_data[3] = plot!(plot_data[3], Δτ, Lz/w₀/π₀;
                                common_settings...,
                                ribbon = σLz,
                                seriescolor = j,
                                xlabel = L"\tau/T_0",
                                ylabel = L"\overline{L_z}/(w_0 \pi_0)")

        plot_data[4] = plot!(plot_data[4], Δτ, L²/w₀/π₀/w₀/π₀;
                                common_settings...,
                                ribbon = (L²/w₀/π₀/w₀/π₀, σL²),
                                seriescolor = j,
                                xlabel = L"\tau/T_0",
                                ylabel = L"\overline{L^2}/(w_0^2 \pi_0^2)")

    end
    write(io,"\n")

    add_info!(plot_data[1], p)
    savefig(plot_data[1], "z_t_$(sim_name).png")
    add_info!(plot_data[2], p)
    savefig(plot_data[2], "rho_t_$(sim_name).png")
    add_info!(plot_data[3], p)
    savefig(plot_data[3], "Lz_t_$(sim_name).png")
    add_info!(plot_data[4], p)
    savefig(plot_data[4], "L2_t_$(sim_name).png")

    eprob = problem_set(R1, p)
    M = solve(eprob, Vern9(), EnsembleThreads(), abstol = 1e-9, reltol = 1e-9,
            saveat = (τf-τi)/time_samples, trajectories = 4N)

    X = [M[i][begin].x¹ for i in 1:4N]/w₀ # pulling out the initial x coordinates
    Y = [M[i][begin].x² for i in 1:4N]/w₀ # pulling out the initial y coordinates
    Lz = [M[i][end].Lz for i in 1:4N] # pulling out the final Lz for colouring
    Lzmax = max(abs.(extrema(Lz))...) # finding maximum
    colours = Lz/Lzmax # doing a normalization

    plot_data[7] = scatter(X,Y;
        marker_z = colours,
        markershape = circ,
        markerstrokealpha = 0,
        markersize = 1.2,
        seriescolor = :PuOr_9,
        clims = (-1., 1.),
        fontfamily = "computer modern",
        dpi = 360,
        label = false,
        xlabel = L"x/w_0", ylabel = L"y/w_0",
        aspect_ratio = 1,
        xticks = -4:1:4)
    add_info!(plot_data[7], p, true)
    savefig(plot_data[7], "XYLz.png")

    xDens = 100
    LzStats = zeros(xDens, 3)
    ρStats  = zeros(xDens, 3)
    Midx = 1
    for i in 1:xDens
        ρ_avg = 0.
        Lz_   = 0.
        Lz²_  = 0.
        ρ_    = 0.
        ρ²_   = 0.
        tally = 0
        m = M[Midx]
        ρ = hypot(m[begin].x¹, m[begin].x²)
        while ρ < i * maxR/xDens
            ρ_avg += ρ
            Lz_   += m[end].Lz
            Lz²_  += m[end].Lz^2
            ρ_    += hypot(m[end].x¹, m[end].x²)
            ρ²_   += m[end].x¹^2 + m[end].x²^2
            tally += 1
            if Midx < 4N
                Midx += 1
            else
                break
            end
            m = M[Midx]
            ρ = hypot(m[begin].x¹, m[begin].x²)
        end
        LzStats[i,:] .= ρ_avg/tally, Lz_/tally, Lz²_/tally
        ρStats[i,:] .= ρ_avg/tally, ρ_/tally,  ρ²_/tally
    end

    plot_data[5] = plot(LzStats[:,1]/w₀, LzStats[:,2]/π₀/w₀,
                        label = false,
                        ribbon = σA(LzStats[:,2],LzStats[:,3])/π₀/w₀)
    plot_data[5] = add_rect!(plot_data[5], LzStats[:,1]/w₀, LzStats[:,2]/π₀/w₀, p.roots;
                        common_settings...,
                        ribbon = σA(LzStats[:,2],LzStats[:,3])/π₀/w₀,
                        seriescolor = 4,
                        xlabel = L"\rho_0/w_0",
                        ylabel = L"\overline{L_z}(\rho_0)/(w_0 \pi_0)",
                        )
    add_info!(plot_data[5], p)
    savefig(plot_data[5], "Lz_rho_$(sim_name).png")

    plot_data[6] = plot(ρStats[:,1]/w₀, ρStats[:,2]/w₀ .- ρStats[:,1]/w₀;
                        label = false,
                        ribbon = σA(ρStats[:,2],ρStats[:,3])/w₀)
    plot_data[6] = add_rect!(plot_data[6], ρStats[:,1]/w₀, ρStats[:,2]/w₀ .- ρStats[:,1]/w₀, p.roots;
                        common_settings...,
                        ribbon = σA(ρStats[:,2],ρStats[:,3])/w₀,
                        seriescolor = 4,
                        xlabel = L"\rho_0/w_0",
                        ylabel = L"\Delta\overline{\rho}_{\mathrm{final}}(\rho_0)/w_0")
    add_info!(plot_data[6], p)
    savefig(plot_data[6], "rho_rho_$(sim_name).png")

    serialize("$(sim_name)", plot_data)
    cd("..")
    cd("..")
    println("done")
end

close(io)
cd("..")

using Telegram
using Telegram.API: sendAnimation, sendMessage, sendSticker
using ConfigEnv
dotenv("../.env")

sendAnimation(animation = "CgACAgQAAxkBAAMWYI1V67io3ILX4lFH-PzAZmoGfugAAi0CAAIvno1SOEmh8Ay-BCYfBA", caption = "Hei, tu!", disable_notification = true)
sendMessage(text = "Fișierul `$(basename(@__FILE__))`, de pe $(ENV["HOSTNAME"]) a fost executat cu succes\\! 💯", parse_mode = "MarkdownV2")
sendSticker(sticker ="CAACAgIAAxkBAAMXYI1V9sq0HUcFb7jdWxkWC5bboqoAAg4AA-nYEygTpj1DX_hIHx8E", caption = "Totul a mers ok!", disable_notification = true)

max_stars_repo_name: toma-vlad/electrodynamic-simulations

max_issues_repo_name: toma-vlad/electrodynamic-simulations

max_forks_repo_name: toma-vlad/electrodynamic-simulations


text:
const prolog = "https://pubchem.ncbi.nlm.nih.gov/rest/pug/"

"""
    cid = get_cid(name="glucose")
    cid = get_cid(smiles="C([C@@H]1[C@H]([C@@H]([C@H](C(O1)O)O)O)O)O")

Return the PubChem **c**ompound **id**entification number for the specified compound.
"""
function get_cid(; name=nothing, smiles=nothing,                   # inputs
                   kwargs...)
    input = "compound/"
    name !== nothing && (input *= "name/$(HTTP.escapeuri(name))/")
    smiles !== nothing && (input *= "smiles/$((smiles))/")
    url = prolog * input * "cids/TXT"
    r = HTTP.request("GET", url; kwargs...)
    return parse(Int, chomp(String(r.body)))
end

"""
    msg = query_substructure(;cid=nothing, smiles=nothing, smarts=nothing,           # specifier for the substructure to search for
                              properties="MolecularFormula,MolecularWeight,XLogP,",  # properties to retrieve
                              output="CSV")                                          # output format

Perform a substructure search of the entire database. You can specify the target via its `cid`, the SMILES string, or a SMARTS string.
Specify the `properties` you want to retrieve as a comma-separated list from among the choices in
http://pubchemdocs.ncbi.nlm.nih.gov/pug-rest, "Compound Property Tables".
Requesting more properties takes more time.

The output is a `Vector{UInt8}`. For `output="CSV"`, a good choice to generate a manipulable result is
`DataFrame(CSV.File(msg))` from the DataFrames and CSV packages, respectively.
Alternatively `String(msg)` will convert it to a string, which you can write to a file.

# Example

```
julia> using PubChemCrawler, CSV, DataFrames

julia> cid = get_cid(name="estriol")
5756

julia> df = CSV.File(query_substructure(;cid)) |> DataFrame      # on Julia 1.0, use `(;cid=cid)`
11607×4 DataFrame
│ Row  │ CID       │ MolecularFormula │ MolecularWeight │ XLogP    │
│      │ $Int     │ String           │ Float64         │ Float64? │
├──────┼───────────┼──────────────────┼─────────────────┼──────────┤
│ 1    │ 5756      │ C18H24O3         │ 288.4           │ 2.5      │
│ 2    │ 5281904   │ C24H32O9         │ 464.5           │ 1.1      │
│ 3    │ 27125     │ C18H24O4         │ 304.4           │ 1.5      │
...
```

will query for derivatives of [estriol](https://en.wikipedia.org/wiki/Estriol).

!!! info
    For complex queries that risk timing out, consider [`query_substructure_pug`](@ref) in combination with [`get_for_cids`](@ref).
"""
function query_substructure(;cid=nothing, smiles=nothing, smarts=nothing,          # inputs
                             properties="MolecularFormula,MolecularWeight,XLogP,", # http://pubchemdocs.ncbi.nlm.nih.gov/pug-rest, "Compound Property Tables"
                             output="CSV",
                             kwargs...)
    input = "compound/fastsubstructure/"
    if cid !== nothing
        (smiles === nothing && smarts === nothing) || throw(ArgumentError("only one of cid, smiles, or smarts can be specified"))
        input *= "cid/$cid/"
    elseif smiles !== nothing
        smarts === nothing || throw(ArgumentError("only one of cid, smiles, or smarts can be specified"))
        input *= "smiles/$smiles/"
    elseif smarts !== nothing
        input *= "smarts/$smarts/"
    else
        throw(ArgumentError("one of cid, smiles, or smarts must be specified"))
    end
    props = canonicalize_properties("property/" * properties)
    url = prolog * input * props * output * "?StripHydrogen=true"
    r = HTTP.request("GET", url; kwargs...)
    return r.body
end

"""
    msg = get_for_cids(cids; properties|xrefs|cids_type|record_type, output="CSV")

Retrieve the given `properties`, `xrefs`, CIDs, or records, respectively, for a list of compounds specified by their `cids`.
The documentation for these traits can be found at http://pubchemdocs.ncbi.nlm.nih.gov/pug-rest; this URL will be referred to as
PUGREST below.

- `properties` include structural features like the molecular formula, number of undefined stereocenters, and so on.
  Specify these as a comma-separated list from among the choices in PUGREST under "Compound Property Tables".
- `xrefs` ("cross-references") include identifiers used by other databases, e.g., the CAS (Registry) number, PubMedID, and so on.
  The supported values for `xrefs` are available at PUGREST under "XRefs".
- `cids_type` is used to retrieve CIDs for compounds related to those specified in `cids`; see PUGREST under "SIDS / CIDS / AIDS".
- `record_type` is used to retrieve data files and to specify options for these files, e.g., 2d or 3d SDF files.
  See PUGREST under "Full-record Retrieval".

`output` specifies the output format. Not all options are applicable to all queries; for example, "CSV" is appropriate for
`properties` queries but "SDF" might be used for a `record_type` query. See PUGREST, "Output".

# Examples

```
julia> using PubChemCrawler, CSV, DataFrames, JSON3

julia> cids = [get_cid(name="cyclic guanosine monophosphate"), get_cid(name="aspirin")]
2-element Array{$Int,1}:
 135398570
      2244

julia> CSV.File(get_for_cids(cids; properties="MolecularFormula,XLogP", output="CSV")) |> DataFrame
2×3 DataFrame
 Row │ CID        MolecularFormula  XLogP
     │ $Int      String            Float64
─────┼──────────────────────────────────────
   1 │ 135398570  C10H12N5O7P          -3.4
   2 │      2244  C9H8O4                1.2

julia> open("/tmp/aspirin_3d.sdf", "w") do io    # save the 3d SDF file for aspirin (CID 2244)
           write(io, get_for_cids(2244; record_type="3d", output="SDF"))
       end
4055

julia> dct = JSON3.read(get_for_cids(cids; xrefs="RN,", output="JSON"));   # get the Registry Number(s) (CAS)

julia> dct[:InformationList][:Information]
2-element JSON3.Array{JSON3.Object,Array{UInt8,1},SubArray{$UInt,1,Array{$UInt,1},Tuple{UnitRange{$Int}},true}}:
 {
   "CID": 135398570,
    "RN": [
            "40732-48-7",
            "7665-99-8"
          ]
}
 {
   "CID": 2244,
    "RN": [
            "11126-35-5",
            "156865-15-5",
            "50-78-2",
            "52080-78-1",
            "921943-73-9",
            "98201-60-6",
            "99512-66-0"
          ]
}
```
"""
function get_for_cids(cids;
                      properties=nothing,
                      xrefs=nothing,
                      cids_type=nothing,
                      record_type=nothing,
                      output="CSV",
                      kwargs...)
    url = prolog * "compound/cid/"
    if xrefs === nothing
        if properties !== nothing
            url *= canonicalize_properties("property/" * properties)
        elseif cids_type !== nothing
            url *= "cids/"
        end
    else
        properties === nothing || error("cannot specify both xref and properties in a single query")
        url *= canonicalize_properties("xrefs/" * xrefs)
    end
    url = joinpath(url, output)
    if cids_type !== nothing
        url *= "?cids_type=" * cids_type
    end
    if record_type !== nothing
        url *= "?record_type=" * record_type
    end
    r = HTTP.request("POST", url, ["Content-Type"=>"application/x-www-form-urlencoded"], "cid="*join(string.(cids), ","); kwargs...)
    return r.body
end

get_for_cids(cid::Int; kwargs...) = get_for_cids([cid]; kwargs...)


text:
using BSON: @save

function get_replay_batch(pos_buffer, π_buffer, res_buffer; batch_size = 32)
  while true
    shfl = shuffle(1:length(pos_buffer))
    pos_buffer, π_buffer, res_buffer = pos_buffer[shfl], π_buffer[shfl], res_buffer[shfl]
    idxs = sample(1:length(pos_buffer), batch_size, replace=false)
    pos_replay = pos_buffer[idxs]
    π_replay = hcat(π_buffer[idxs]...)
    res_replay = res_buffer[idxs]
    if -5 < sum(res_replay) < 5
      return pos_replay, π_replay, res_replay
    end
  end
end

#TODO: Model save path

function save_model(nn::NeuralNet)
  bn,value, policy = cpu.((nn.base_net, nn.value, nn.policy))
  @save "../models/agz_base.bson" bn
  @save "../models/agz_value.bson" value
  @save "../models/agz_policy.bson" policy

  # saving weights
  bn_weights = cpu.(Tracker.data.(params(bn)))
  val_weights = cpu.(Tracker.data.(params(value)))
  pol_weights = cpu.(Tracker.data.(params(policy)))

  @save "../models/weights/agz_base.bson" bn_weights
  @save "../models/weights/agz_value.bson" val_weights
  @save "../models/weights/agz_policy.bson" pol_weights
end

function train(env::GameEnv; num_games::Int = 25000, memory_size::Int = 500000,
  batch_size::Int = 32, epochs = 1, ckp_freq::Int = 1000, readouts::Int = 800,
  tower_height::Int = 19, model = nothing, start_training_after = 50000)

  # @assert 0 ≤ tower_height ≤ 19
  cur_nn = nothing
  if model == nothing
    cur_nn = NeuralNet(env; tower_height = tower_height)
  else
    cur_nn = model
  end

  # prev_nn = deepcopy(cur_nn)

  pos_buffer = Vector{Position}()
  π_buffer = Vector{Vector{Float32}}()
  res_buffer = Vector{Int}()

  for i = 1:num_games
    player = selfplay(env, cur_nn, readouts)
    p, π, v = extract_data(player)

    pos_buffer = vcat(pos_buffer, p)
    π_buffer   = vcat(π_buffer, π)
    res_buffer = vcat(res_buffer, v)

    if length(pos_buffer) > memory_size
      pos_buffer = pos_buffer[end-memory_size+1:end]
      π_buffer   = π_buffer[end-memory_size+1:end]
      res_buffer = res_buffer[end-memory_size+1:end]
    end

    if length(pos_buffer) >= start_training_after
      replay_pos, replay_π, replay_res = get_replay_batch(pos_buffer, π_buffer, res_buffer; batch_size = batch_size)
      loss = train!(cur_nn, (replay_pos, replay_π, replay_res); epochs = epochs)
      result = player.result_string
      num_moves = player.root.position.n
      println("Episode $i over. Loss: $loss. Winner: $result. Moves: $num_moves.")
    end

  # cur_is_winner = evaluate(cur_nn, prev_nn; num_games = eval_games, ro = readouts)
  # print("Evaluated. ")
  # if cur_is_winner
  #   prev_nn = deepcopy(cur_nn)
  #   print("Model updated. ")
  # else
  #   cur_nn = deepcopy(prev_nn)
  #   print("Model retained. ")
  # end

    if i % ckp_freq == 0
      save_model(cur_nn)
      print("Model saved. ")
    end
  end
  return cur_nn
end


text:
module GCNN

export RotGroupConv, RotGroupConvTranspose

using Flux
using Flux: conv, ∇conv_data, convfilter, glorot_uniform, @functor, DenseConvDims

"""
Add an extra dimension to the end of an array.
"""
expand_dims(x) = reshape(x, size(x)..., 1)

"""
A conveniently bundled group of rotation matrices.
"""
struct RotGroup{RT<:AbstractMatrix}
    Rs::Vector{RT}
end

@functor RotGroup

"""
Compute a sparse matrix that simultaneously does `nrots` interpolated rotations
of a m×n matrix. To compute the rotations, the image must be flattened, then
sparse matrix times dense vector multiply, then reshape the output vector into
a nrots×m×n array.
"""
function RotGroup(
        m::Integer, n::Integer, nrots::Integer, periodicity=2π, mask_disk::Bool=true)

    # TODO: until we test m != n case, I wouldn't trust it
    @assert m == n

    ci, cj = div(m, 2)+1, div(n, 2)+1
    c = min(ci, cj)

    Rs = Vector{Matrix{Float32}}(undef, nrots)

    for (r, θ) in enumerate(range(0, periodicity, length=nrots+1)[1:end-1])
        cosθ, sinθ = cos(θ), sin(θ)

        # COO sparse matrix values
        R = zeros(Float32, (m*n, m*n))

        for i in 1:m, j in 1:n
            if mask_disk && (i-ci)^2 + (j-cj)^2 > (c+0.5)^2
                continue
            end

            # flattened index
            l_input = (j-1)*m + i

            # compute fractional rotated indexes
            ir = cosθ * (i - ci) + sinθ * (j - cj) + ci
            jr = -sinθ * (i - ci) + cosθ * (j - cj) + cj

            # interpolate across 2x2 square
            ir1 = floor(Int, ir)
            ir2 = ir1 + 1
            jr1 = floor(Int, jr)
            jr2 = jr1 + 1

            wi = ir - ir1
            wj = jr - jr1

            for (irk, wik) in [(ir1, 1-wi), (ir2, wi)], (jrk, wjk) in [(jr1, 1-wj), (jr2, wj)]
                if 1 <= irk <= m && 1 <= jrk <= n
                    l_output = (jrk-1)*m + irk
                    R[l_output, l_input] = wik * wjk
                end
            end

            Rs[r] = R
        end
    end

    # TODO: It's probably better to use dense matrices since these will be
    # small. (E.g. to rotate a 5x5 rotation kernel, we only need a 25x25 matrix).

    return RotGroup(Rs)
end

"""
Apply each rotation matrix to an array w of shape [width, height, in_channels, out_channels]
produces a vector of rotated arrays each of the same shape.
"""
function (rg::RotGroup)(w::AbstractArray{T,4}) where {T}
    width, height, in_channels, out_channels = size(w)
    w_flat = reshape(w, (width*height, in_channels*out_channels))
    return [reshape(R*w_flat, (width, height, in_channels, out_channels)) for R in rg.Rs]
end

Base.length(rg::RotGroup) = length(rg.Rs)

"""
Group convolutional layer.
"""
struct RotGroupConv{F, WT<:AbstractArray, BT, RT<:AbstractMatrix, M, N}
    rg::RotGroup{RT}
    weight::WT
    bias::BT
    σ::F
    stride::NTuple{N,Int}
    pad::NTuple{M,Int}
    dilation::NTuple{N,Int}
end

@functor RotGroupConv

"""
    RotGroupConv(nrots, filter, in => out, σ = identity)

Construct a Group CNN (G-CNN) layer with `nrots` equally spaced rotations, and
a `(filter, filter)` convolution kernel, with `in` input channels and `out`
output channels, with an optional elementwise nonlinearity `σ`.
"""
function RotGroupConv(
        nrots::Integer, filter::Integer, ch::Pair{<:Integer,<:Integer}, σ=identity;
        init=glorot_uniform, stride::Int=1, pad=0, dilation=1, use_bias::Bool=true)

    weight = convfilter((filter, filter), ch; init)
    # bias = create_bias(weght, use_bias, size(weight, N))
    rg = RotGroup(filter, filter, nrots)

    return RotGroupConv(rg, weight, σ, stride, pad, dilation, use_bias)
end

function RotGroupConv(
        rg::RotGroup, weight::AbstractArray{T, N}, σ, stride, pad, dilation, use_bias) where {T, N}

    stride = Flux.expand(Val(N-2), stride)
    dilation = Flux.expand(Val(N-2), dilation)
    pad = Flux.calc_padding(Conv, pad, size(weight)[1:N-2], dilation, stride)
    bias = Flux.create_bias(weight, use_bias, size(weight, N))
    return RotGroupConv(rg, weight, bias, σ, stride, pad, dilation)
end

"""
Rotation group convolution applied to an array with shape [width, height, in_channels, batches] and
producing an array of shape [width, height, out_channels, batches, nrotations].
"""
function (lyr::RotGroupConv)(x::AbstractArray{T,4}) where {T}
    # Rotate the weights according to each rotation matrix, apply each matrix, then concatenate.
    bias = reshape(lyr.bias, ntuple(_ -> 1, length(lyr.stride))..., :, 1)
    cdims = DenseConvDims(x, lyr.weight; stride=lyr.stride, padding=lyr.pad, dilation=lyr.dilation)
    return cat([expand_dims(lyr.σ.(conv(x, Rw, cdims) .+ bias)) for Rw in lyr.rg(lyr.weight)]..., dims=5)
end

"""
Rotation group convolution applied to an array with shape `[width, height,
in_channels, batches, nrotations]` and producing an array of shape `[width,
height, out_channels, batches, nrotations]`.
"""
function (lyr::RotGroupConv)(x::AbstractArray{T,5}) where {T}
    # TODO:
    nrots = length(lyr.rg)
    @assert size(x, 5) == nrots

    bias = reshape(lyr.bias, ntuple(_ -> 1, length(lyr.stride))..., :, 1)
    cdims = DenseConvDims(size(x)[1:end-1], size(lyr.weight); stride=lyr.stride, padding=lyr.pad, dilation=lyr.dilation)

    Rws = lyr.rg(lyr.weight)
    Rxs = [view(x, :, :, :, :, i) for i in 1:nrots]

    # gradients don't work with zip
    #return cat([expand_dims(lyr.σ.(conv(Rx, Rw))) for (Rx, Rw) in zip(Rxs, Rws)]..., dims=5)

    return cat([expand_dims(lyr.σ.(conv(Rxs[k], Rws[k], cdims) .+ bias)) for k in 1:nrots]..., dims=5)
end

"""
Group transposed convolutional layer.
"""
struct RotGroupConvTranspose{F, WT<:AbstractArray, BT, RT<:AbstractMatrix, M, N}
    rg::RotGroup{RT}
    weight::WT
    bias::BT
    σ::F
    # TODO: Make this a tuple so everythig works.
    stride::NTuple{N,Int}
    pad::NTuple{M,Int}
    dilation::NTuple{N,Int}
end

@functor RotGroupConvTranspose

"""
    RotGroupConvTranspose(nrots, filter, in => out, σ = identity)

Construct a Group CNN (G-CNN) transposed convolution layer with `nrots` equally
spaced rotations, and a `(filter, filter)` convolution kernel, with `in` input
channels and `out` output channels, with an optional elementwise nonlinearity
`σ`.
"""
function RotGroupConvTranspose(
        nrots::Integer, filter::Integer, ch::Pair{<:Integer,<:Integer}, σ=identity;
        init=glorot_uniform, stride::Int=1, pad=0, dilation=1, use_bias::Bool=true)

    weight = convfilter((filter, filter), reverse(ch); init)

    # bias = create_bias(weght, use_bias, size(weight, N))
    rg = RotGroup(filter, filter, nrots)

    return RotGroupConvTranspose(rg, weight, σ, stride, pad, dilation, use_bias)
end

function RotGroupConvTranspose(
        rg::RotGroup, weight::AbstractArray{T, N}, σ, stride, pad, dilation, use_bias) where {T, N}
    stride = Flux.expand(Val(N-2), stride)
    dilation = Flux.expand(Val(N-2), dilation)
    pad = Flux.calc_padding(ConvTranspose, pad, size(weight)[1:N-2], dilation, stride)
    bias = Flux.create_bias(weight, use_bias, size(weight, N-1))
    return RotGroupConvTranspose(rg, weight, bias, σ, stride, pad, dilation)
end

function conv_transpose_dims(c::RotGroupConvTranspose, xsize::NTuple{N, Int}) where {N}
    # Calculate size of "input", from ∇conv_data()'s perspective...
    combined_pad = (c.pad[1:2:end] .+ c.pad[2:2:end])
    I = (xsize[1:end-2] .- 1).*c.stride .+ 1 .+ (size(c.weight)[1:end-2] .- 1).*c.dilation .- combined_pad
    C_in = size(c.weight)[end-1]
    batch_size = xsize[end]
    # Create DenseConvDims() that looks like the corresponding conv()
    return DenseConvDims((I..., C_in, batch_size), size(c.weight);
                        stride=c.stride,
                        padding=c.pad,
                        dilation=c.dilation,
    )
end

"""
Rotation group transposed convolution applied to an array with shape [width,
height, in_channels, batches] and producing an array of shape [width, height,
out_channels, batches, nrotations].
"""
function (lyr::RotGroupConvTranspose)(x::AbstractArray{T,4}) where {T}
    # Rotate the weights according to each rotation matrix, apply each matrix, then concatenate.
    cdims = conv_transpose_dims(lyr, size(x))
    bias = reshape(lyr.bias, map(_->1, lyr.stride)..., :, 1)

    # return cat([expand_dims(lyr.σ.(∇conv_data(x, Rw, cdims) .+ lyr.bias)) for Rw in lyr.rg(lyr.weight)]..., dims=5)
    return cat([expand_dims(lyr.σ.(∇conv_data(x, Rw, cdims) .+ bias)) for Rw in lyr.rg(lyr.weight)]..., dims=5)
end

"""
Rotation group transposed convolution applied to an array with shape `[width,
height, in_channels, batches, nrotations]` and producing an array of shape
`[width, height, out_channels, batches, nrotations]`.
"""
function (lyr::RotGroupConvTranspose)(x::AbstractArray{T,5}) where {T}
    nrots = length(lyr.rg)
    @assert size(x, 5) == nrots

    Rws = lyr.rg(lyr.weight)
    Rxs = [view(x, :, :, :, :, i) for i in 1:nrots]
    cdims = conv_transpose_dims(lyr, size(x)[1:end-1])
    bias = reshape(lyr.bias, map(_->1, lyr.stride)..., :, 1)

    # gradients don't work with zip
    #return cat([expand_dims(lyr.σ.(conv(Rx, Rw))) for (Rx, Rw) in zip(Rxs, Rws)]..., dims=5)

    return cat([expand_dims(lyr.σ.(∇conv_data(Rxs[k], Rws[k], cdims) .+ bias)) for k in 1:nrots]..., dims=5)
end

end # module


text:
using Knet

# if mode is 1, then fanin will be used, otherwise fanout
# Mostly taken from Knet source code
function kaiming_uniform(a...; mode=1, gain=1)
    w = rand(a...)
    if ndims(w) == 1
        fanout = 1
        fanin = length(w)
    elseif ndims(w) == 2
        fanout = size(w,1)
        fanin = size(w,2)
    else
        fanin = div(length(w),  a[end])
        fanout = div(length(w), a[end-1])
    end

    fan = fanin
    if mode != 1 fan = fanout end

    s = convert(eltype(w), gain*sqrt(3 / (fan)))
    return 2s .* w .- s
end

max_stars_repo_name: barisbatuhan/retinaface

max_issues_repo_name: barisbatuhan/retinaface

max_forks_repo_name: barisbatuhan/retinaface


text:
using TrajectoryOptimization: get_nG, gen_usrfun_ipopt, update_derivatives!, get_traj_points!, get_traj_points_derivatives!, update_jacobians!
model, obj0 = Dynamics.cartpole_analytical
obj = copy(obj0)
n,m = model.n, model.m
dt = 0.1

function test_ipopt_funcs(method)
    # Set up problem
    solver = Solver(model,ConstrainedObjective(obj),dt=dt,integration=:rk3_foh)
    N,N_ = TrajectoryOptimization.get_N(solver,method)
    NN = N*(n+m)
    nG, = TrajectoryOptimization.get_nG(solver,method)
    U0 = ones(1,N)*1
    X0 = line_trajectory(obj.x0, obj.xf, N)
    Z = TrajectoryOptimization.packZ(X0,U0)

    # Init vars
    g = zeros((N-1)n)
    grad_f = zeros(NN)
    rows = zeros(nG)
    cols = zeros(nG)
    vals = zeros(nG)

    # Get functions and evaluate
    eval_f, eval_g, eval_grad_f, eval_jac_g = gen_usrfun_ipopt(solver,method)
    J = eval_f(Z)
    eval_g(Z,g)
    eval_grad_f(Z,grad_f)
    eval_jac_g(Z,:Structure,rows,cols,vals)
    eval_jac_g(Z,:vals,rows,cols,vals)
    jac_g = sparse(rows,cols,vals)

    function eval_all(Z)
        eval_f(Z)
        eval_g(Z,g)
        eval_grad_f(Z,grad_f)
        eval_jac_g(Z,:Structure,rows,cols,vals)
        eval_jac_g(Z,:vals,rows,cols,vals)
    end
    eval_all(Z)

    # SNOPT method
    results = DircolResults(solver,method)
    results.Z .= Z
    update_derivatives!(solver,results,method)
    get_traj_points!(solver,results,method)
    get_traj_points_derivatives!(solver,results,method)
    update_jacobians!(solver,results,method)
    J_snopt = cost(solver,results)
    g_snopt = collocation_constraints(solver,results,method)
    results.X_
    grad_f_snopt = cost_gradient(solver,results,method)
    jac_g_snopt = constraint_jacobian(solver,results,method)

    function eval_all_snopt(Z)
        results.Z .= Z
        update_derivatives!(solver,results,method)
        get_traj_points!(solver,results,method)
        get_traj_points_derivatives!(solver,results,method)
        update_jacobians!(solver,results,method)
        J_snopt = cost(solver,results)
        g_snopt = collocation_constraints(solver,results,method)
        grad_f_snopt = cost_gradient(solver,results,method)
        jac_g_snopt = constraint_jacobian(solver,results,method)
        return nothing
    end
    eval_all_snopt(Z)

    # @time eval_all(Z)
    # @time eval_all_snopt(Z)

    @test J_snopt ≈ J
    @test g_snopt == g
    @test grad_f_snopt == grad_f

    dropzeros!(jac_g_snopt)
    dropzeros!(jac_g)
    @test nonzeros(jac_g_snopt) == nonzeros(jac_g)
end
test_ipopt_funcs(:midpoint)
test_ipopt_funcs(:trapezoid)
test_ipopt_funcs(:hermite_simpson)
test_ipopt_funcs(:hermite_simpson_separated)

# Try IPOPT
method = :hermite_simpson_separated
u_bnd = 10
dt = 0.01
solver = Solver(model,ConstrainedObjective(obj,u_max=u_bnd,u_min=-u_bnd),dt=dt,integration=:rk3_foh)
N, = get_N(solver.N,method)
U0 = zeros(m,N)
X0 = line_trajectory(solver,method)
Z = packZ(X0,U0)
@time sol, stats, prob = solve_dircol(solver,X0,U0,method=method,nlp=:ipopt)
@test norm(sol.X[:,end]-obj.xf) < 1e-6


text:
using .Symbolics

include("../Continuous/symbolics.jl")