text:
"""
    GroupNormV2(channels::Integer, G::Integer, λ=identity;
                initβ=zeros32, initγ=ones32,
                affine=true, track_stats=false,
                ϵ=1f-5, momentum=0.1f0)

Same as [`Flux.GroupNorm`](@ref) but works well with `destructure_parameters` & `destructure_parameters_states`
"""
struct GroupNormV2{F,V,N,W}
    G::Int  # number of groups
    λ::F  # activation function
    β::V  # bias
    γ::V  # scale
    μ::W     # moving mean
    σ²::W    # moving std
    ϵ::N
    momentum::N
    chs::Int # number of channels
    attrs::NormAttributes
end

Flux.hasaffine(gn::GroupNormV2) = Flux.hasaffine(gn.attrs)
Flux._isactive(gn::GroupNormV2) = Flux._isactive(gn.attrs)

Flux.@functor GroupNormV2

Flux.trainable(gn::GroupNormV2) = hasaffine(gn) ? (gn.β, gn.γ) : ()

function GroupNormV2(chs::Int, G::Int, λ=identity; initβ=zeros32, initγ=ones32, affine=true, track_stats=false, ϵ=1f-5,
                     momentum=0.1f0)
    @assert chs % G == 0 "The number of groups ($(G)) must divide the number of channels ($chs)"

    β = affine ? initβ(chs) : nothing
    γ = affine ? initγ(chs) : nothing
    μ = track_stats ? zeros32(G) : nothing
    σ² = track_stats ? ones32(G) : nothing

    return GroupNormV2(G, λ, β, γ, μ, σ², ϵ, momentum, chs, NormAttributes(affine, track_stats, nothing))
end

function (gn::GroupNormV2)(x::AbstractArray{T,N}) where {T,N}
    # Not doing assertion checks
    # @assert N > 2
    # @assert size(x, N - 1) == gn.chs
    sz = size(x)
    x_2 = reshape(x, sz[1:(N - 2)]..., sz[N - 1] ÷ gn.G, gn.G, sz[N])
    N_ = ndims(x_2)
    reduce_dims = 1:(N_ - 2)
    affine_shape = ntuple(i -> i ∈ (N_ - 1, N_ - 2) ? size(x_2, i) : 1, N_)
    return reshape(norm_forward(gn, x_2, reduce_dims, affine_shape), sz)
end

function testmode!(m::GroupNormV2, mode=true)
    return (m.attrs.active = (isnothing(mode) || mode == :auto) ? nothing : !mode; m)
end

function Base.show(io::IO, l::GroupNormV2)
    print(io, "GroupNormV2($(l.chs), $(l.G)")
    l.λ == identity || print(io, ", ", l.λ)
    hasaffine(l) || print(io, ", affine=false")
    return print(io, ")")
end


text:
"""
    BSpline(t, x, w, B, S)

Basis structure for the B-splines generated by the knot set `t`. `x`
and `w` are the associated quadrature roots and weights, respectively,
the columns of `B` correspond to the B-splines resolved on the
quadrature roots, and `S` is the (banded) B-spline overlap matrix.
"""
struct BSpline{T,R<:Real,
               KnotSet<:AbstractKnotSet{<:Any,<:Any,<:Any,R},
               XV<:AbstractVector{T},
               WV<:AbstractVector{R},
               BM<:AbstractMatrix{T},
               SM<:AbstractMatrix{T}} <: Basis{T}
    t::KnotSet
    x::XV
    w::WV
    B::BM
    S::SM
end

const RestrictedBSpline{T} = RestrictedQuasiArray{T,2,<:BSpline{T}}
const AdjointRestrictedBSpline{T} = AdjointRestrictedQuasiArray{T,2,<:BSpline{T}}

const BSplineOrRestricted{T} = BasisOrRestricted{<:BSpline{T}}
const AdjointBSplineOrRestricted{T} = AdjointBasisOrRestricted{<:BSpline{T}}

function basis_functions(t::AbstractKnotSet{k}, x::AbstractVector{T}, m=0) where {k,T}
    nf = numfunctions(t)
    B = spzeros(T, length(x), nf)

    nei = nonempty_intervals(t)
    # This is the amount of quadrature points per interval (assuming
    # the same amount per interval, which is currently the case), this
    # could conceivably be passed as an argument, instead.
    N = findlast(e -> e < t[nei[1]+1], x)

    ix = 1
    for i ∈ nei
        # These are the functions which are non-zero on the interval
        # t[i]..t[i+1].
        for j ∈ max(1,i-k+1):min(i,nf)
            eⱼ = UnitVector{T}(nf, j)
            for l ∈ ix:ix+N-1
                B[l,j] = deBoor(t, eⱼ, x[l], i, m)
            end
        end
        ix += N
    end

    B
end

basis_functions(B::BSpline, args...) = basis_functions(B.t, B.x, args...)
basis_functions(B::RestrictedBSpline, args...) =
    basis_functions(parent(B), args...)[:,indices(B,2)]

function overlap_matrix!(S::Union{BandedMatrix,Tridiagonal}, χ, ξ, w)
    m = min(size(S,1),size(χ,2))
    n = min(size(S,2),size(ξ,2))
    # It is assumed that the bandwidth is correct for the overlap
    # ⟨χ|ξ⟩.
    p = max(bandwidths(S)...)
    W = Diagonal(w)
    for i ∈ 1:m
        for j ∈ max(i-p,1):min(i+p,n)
            S[i,j] = χ[:,i]' * W * ξ[:,j]
        end
    end
    # χ == ξ && isreal(χ) && isreal(W) && isreal(ξ) ?
    #     Symmetric(S) : S
    S
end

function assert_compatible_bases(A::BSplineOrRestricted, B::BSplineOrRestricted)
    A = parent(A)
    B = parent(B)
    A.t.t == B.t.t &&
        A.x == B.x &&
        A.w == B.w ||
        throw(ArgumentError("Can only multiply B-spline bases with identical knot sets, resolved on the same Gauß–Legendre points"))
end

function overlap_matrix!(S::Union{BandedMatrix,Tridiagonal},
                         Ac::AdjointBSplineOrRestricted,
                         B::BSplineOrRestricted,
                         op=I)
    A = parent(Ac)
    assert_compatible_bases(A,B)
    χ = view(parent(A).B, :, indices(A,2))
    ξ = view(parent(B).B, :, indices(B,2))
    overlap_matrix!(S, χ, op*ξ, weights(parent(B)))
end

function BSpline(t::AbstractKnotSet{k}, x::AbstractVector{T}, w::AbstractVector) where {k,T}
    B = basis_functions(t, x)

    nf = numfunctions(t)
    S = BandedMatrix(Zeros{T}(nf, nf), (k-1,k-1))

    BSpline(t, x, w, B, overlap_matrix!(S, B, B, w))
end

"""
    BSpline(t[, N])

Create the B-spline basis corresponding to the knot set `t`. `N` is
the amount of Gauß–Legendre quadrature points per interval.
"""
BSpline(t::AbstractKnotSet, N) = BSpline(t, lgwt(t, N)...)

"""
    BSpline(t[; k′=3])
Create the B-spline basis corresponding to the knot set `t`. `k′` is
the highest polynomial order of operators for which it should be
possible to compute the matrix elements exactly (via Gauß–Legendre
quadrature). The default `k′=3` corresponds to operators O(x²).
"""
BSpline(t::AbstractKnotSet; k′=3) = BSpline(t, num_quadrature_points(order(t), k′))

# * Properties

axes(B::BSpline) = (Inclusion(first(B.t)..last(B.t)), Base.OneTo(numfunctions(B.t)))
size(B::BSpline) = (ℵ₁, numfunctions(B.t))
==(A::BSpline,B::BSpline) = A.t == B.t
Base.hash(B::BSpline, h::UInt) = hash(B.t, h)

distribution(B::BSpline) = distribution(B.t)

order(B::BSplineOrRestricted) = order(unrestricted_basis(B).t)

function show(io::IO, B::BSpline{T}) where T
    write(io, "BSpline{$(T)} basis with $(B.t)")
end

locs(B::BSpline) = B.x
weights(B::BSpline) = B.w

IntervalSets.leftendpoint(B::BSpline) = B.x[1]
IntervalSets.rightendpoint(B::BSpline) = B.x[end]

function centers(B::BSplineOrRestricted)
    x = B'*QuasiDiagonal(axes(B,1))*B
    S = B'B

    diag(x) ./ diag(S)
end

# # * Basis functions

"""
    deBoor(t, c, x[, i[, m=0]])

Evaluate the spline given by the knot set `t` and the set of control
points `c` at `x` using de Boor's algorithm. `i` is the index of the
knot interval containing `x`. If `m≠0`, calculate the `m`th derivative
at `x` instead.
"""
function deBoor(t::AbstractKnotSet, c::AbstractVector{T}, x,
                i=find_interval(t, x), m=0) where T
    isnothing(i) && return zero(eltype(t))
    k = order(t)
    nc = length(c)
    k == 1 && return (nc < i || m > 0) ? zero(T) : c[i]

    α = [r > 0 && r ≤ nc ? c[r] : zero(T)
         for r ∈ i-k+1:i]
    nt = length(t)
    nf = numfunctions(t)
    for j = 1:k-1
        for r = i:-1:max(i-k+j,1)
            jjj = r+k-j
            (jjj > nt || jjj < 1) && continue
            r′ = r - i + k
            r ≠ 1 && r′ == 1 && continue

            a = t[r+k-j]
            b = t[r]

            α[r′] = if j ≤ m
                (k-j)*if r == 1
                    -α[r′]
                elseif r == nf + j
                    α[r′-1]
                else
                    α[r′-1] - α[r′]
                end
            else
                if r == 1
                    (b-x)*α[r′]
                elseif r == nf + j
                    (x-a)*α[r′-1]
                else
                    ((x-a)*α[r′-1] + (b-x)*α[r′])
                end
            end
            α[r′] /= (b-a)
        end
    end

    α[end]
end

getindex(B::BSpline{T}, x::Real, j::Integer) where T =
    deBoor(B.t, UnitVector{T}(size(B,2), j),
           x, find_interval(B.t, x))

function basis_function!(χ, B::BSpline{T}, x::AbstractVector, j) where T
    eⱼ = UnitVector{T}(size(B,2), j)
    for (is,k) ∈ within_support(x, B.t, j)
        for i in is
            χ[i] = deBoor(B.t, eⱼ, x[i], k)
        end
    end
end

function getindex(B::BSpline{T}, x::AbstractVector, sel::AbstractVector) where T
    χ = spzeros(T, length(x), length(sel))
    o = sel[1] - 1
    for j in sel
        basis_function!(view(χ, :, j-o), B, x, j)
    end
    χ
end

function getindex(B::BSplineOrRestricted{T}, x::AbstractVector, j::Integer) where T
    χ = spzeros(T, length(x))
    basis_function!(χ, parent(B), x, j+first(indices(B,2))-1)
    χ
end

getindex(B::BSplineOrRestricted, x, ::Colon) =
    getindex(parent(B), x, indices(B,2))

getindex(B::BSplineOrRestricted, x::AbstractVector, sel::AbstractVector) =
    getindex(parent(B), x, indices(B,2)[sel])

# * Types

const SplineArray{T,N,B<:BSplineOrRestricted} = FuncArray{T,N,B}
const SplineVector{T,B<:BSplineOrRestricted} = FuncVector{T,B}
const SplineMatrix{T,B<:BSplineOrRestricted} = FuncMatrix{T,B}
const SplineVecOrMat{T,B<:BSplineOrRestricted} = FuncVecOrMat{T,B}

const AdjointSplineArray{T,N,B<:BSplineOrRestricted} = AdjointFuncArray{T,N,B}
const AdjointSplineVector{T,B<:BSplineOrRestricted} = AdjointFuncVector{T,B}
const AdjointSplineMatrix{T,B<:BSplineOrRestricted} = AdjointFuncMatrix{T,B}
const AdjointSplineVecOrMat{T,B<:BSplineOrRestricted} = AdjointFuncVecOrMat{T,B}

Base.show(io::IO, spline::SplineVector) =
    write(io, "Spline on $(spline.args[1])")

Base.show(io::IO, spline::SplineMatrix) =
    write(io, "$(size(spline, 2))d spline on $(spline.args[1])")

# * Matrix construction

function Matrix(::UndefInitializer, A::BSplineOrRestricted{T}, B::BSplineOrRestricted{T}, ::Type{U}=T) where {T,U}
    m,n = size(A,2), size(B,2)
    k = max(order(A),order(B))
    if k == 2 && m == n
        dl = Vector{U}(undef, m-1)
        d = Vector{U}(undef, m)
        du = Vector{U}(undef, m-1)
        Tridiagonal(dl, d, du)
    else
        ij = first(indices(B,2)) - first(indices(A,2))
        BandedMatrix{U}(undef, (m,n), (k-1+ij,k-1-ij))
    end
end

Matrix(::UndefInitializer, A::BSplineOrRestricted{T}, ::Type{U}=T) where {T,U} =
    Matrix(undef, A, A, U)

function Base.zeros(A::BSplineOrRestricted{T}, B::BSplineOrRestricted{T}=A, ::Type{U}=T) where {T,U}
    M = Matrix(undef, A, B, U)
    M .= zero(U)
    M
end

# * Mass matrix
@materialize function *(Ac::AdjointBSplineOrRestricted,
                        B::BSplineOrRestricted)
    T -> begin
        Matrix(undef, parent(Ac), B, T)
    end
    dest::AbstractMatrix{T} -> begin
        overlap_matrix!(dest, Ac, B)
    end
end

metric(B::BSpline) = B.S

# * Diagonal operators

@materialize function *(Ac::AdjointBSplineOrRestricted,
                        D::QuasiDiagonal,
                        B::BSplineOrRestricted)
    T -> begin
        Matrix(undef, parent(Ac), B, T)
    end
    dest::AbstractMatrix{T} -> begin
        # Evaluate the quasi-diagonal operator D on the quadrature roots
        # of B.
        op = Diagonal(getindex.(Ref(D.diag), locs(parent(B))))
        overlap_matrix!(dest, Ac, B, op)
    end
end

@materialize function *(Ac::AdjointBSplineOrRestricted,
                        D::QuasiDiagonal,
                        E::QuasiDiagonal,
                        B::BSplineOrRestricted)
    T -> begin
        Matrix(undef, parent(Ac), B, T)
    end
    dest::AbstractMatrix{T} -> begin
        # Evaluate the quasi-diagonal operators D & E on the quadrature roots
        # of B.
        op = Diagonal(getindex.(Ref(D.diag), locs(parent(B))) .*
                      getindex.(Ref(E.diag), locs(parent(B))))
        overlap_matrix!(dest, Ac, B, op)
    end
end

# * Function interpolation

function Base.:(\ )(B::BSpline, f::BroadcastQuasiArray)
    axes(f,1) == axes(B,1) ||
        throw(DimensionMismatch("Function on $(axes(f,1).domain) cannot be interpolated over basis on $(axes(B,1).domain)"))
    B.B \ getindex.(Ref(f), B.x)
end

function Base.:(\ )(B::RestrictedBSpline, f::BroadcastQuasiArray)
    axes(f,1) == axes(B,1) ||
        throw(DimensionMismatch("Function on $(axes(f,1).domain) cannot be interpolated over basis on $(axes(B,1).domain)"))
    # We need to evaluate the basis functions of the restricted
    # B-spline basis on /all/ quadrature points.
    x = locs(parent(B))
    V = B[x,:]
    V \ getindex.(Ref(f), x)
end

export BSpline

ext: jl


text:
@testset "Basic" begin
  Qx, x = PolynomialRing(FlintQQ, "x")
  K, a = NumberField(x^2 - 5, "a")
  OK = maximal_order(K)
  G = matrix(K, 5, 5, [ 1, 0, 0, 0, 0,
                        0, 1, 0, 0, 0,
                        0, 0, 7, 0, 0,
                        0, 0, 0, 1//2*(a + 1), 0,
                        0, 0, 0, 0, 1//2*(7*a + 35) ])

  V = quadratic_space(K, G)
  @test hasse_invariant(V, 2 * OK) == -1
  @test hasse_invariant(V, 7 * OK) == -1
  @test hasse_invariant(V, real_places(K)[1]) == 1
  @test hasse_invariant(V, real_places(K)[2]) == 1

  @test witt_invariant(V, 2 * OK) == -1
  @test witt_invariant(V, 7 * OK) == -1
  @test witt_invariant(V, real_places(K)[1]) == -1
  @test witt_invariant(V, real_places(K)[2]) == -1

  plc = real_places(K)
  _, _, d, H, I = @inferred invariants(V)
  @test d == 147//2*a+245//2
  @test H == Dict(2 * OK => -1, 7 * OK => -1)
  @test length(I) == 2 && (I[1][2], I[2][2]) in [(1, 0), (0, 1)]
  @test is_isometric(V, V)
  @test !is_isometric(V, quadratic_space(K, 5))

  W = quadratic_space(K, G[1:2,1:2])
  f = hom(W, V, K[1 0 0 0 0; 0 1 0 0 0], check=true)
  v = [K(1), K(0)]
  @test f(v) == image(f, v)
  g = hom(W, W, K[0 1; 1 0], check=true)
  h = compose(g, f)
  @test h(v) == f(g(v))
  @test Hecke.is_isotropic(V, infinite_place(K,1))

end


text:
# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

@reexport module SMPS

# Standard library
using LinearAlgebra
using Random
using SparseArrays
using Distributed

# External libraries
using StochasticPrograms
using StochasticPrograms: canonical
using Compat
using Distributions
using Distributions: AliasTable
using OrderedCollections

import Base: read

# Common constants
const Row = Symbol
const Col = Symbol
const Idx = Int
const IdxGroup = Tuple{Idx, Idx, Symbol}
const InclusionType = Symbol
const DistributionType = Symbol
const Period = Symbol
const Block = Symbol
const RowCol = Tuple{Symbol, Symbol}
const RowColVal{T} = Tuple{Symbol, Symbol, T}

const OBJ = :OBJ
const LEQ = :LEQ
const EQ = :EQ
const GEQ = :GEQ
const RANGE = :RANGE
const LOWER = :LOWER
const UPPER = :UPPER
const BINARY = :BINARY
const INTEGER = :INTEGER
const INTEGER_LOWER = :INTEGER_LOWER
const INTEGER_UPPER = :INTEGER_UPPER
const FREE = :FREE
const FIXED = :FIXED

const NAME = :NAME
const END = :ENDATA

# TIM file constants
const PERIODS = :PERIODS
const TIM_MODES = [PERIODS, END]

# COR file constants
const ROWS = :ROWS
const COLUMNS = :COLUMNS
const RHS = :RHS
const RANGES = :RANGES
const BOUNDS = :BOUNDS
const COR_MODES = [NAME, ROWS, COLUMNS, RHS, RANGES, BOUNDS, END]

# STO file constants
const INDEP = :INDEP
const BLOCKS = :BLOCKS
const SCENARIOS = :SCENARIOS
const STO_MODES = [INDEP, BLOCKS, SCENARIOS, END]

const REPLACE = :REPLACE
const ADD = :ADD
const MULTIPLY = :MULTIPLY
const INCLUSIONS = [REPLACE, ADD, MULTIPLY]

const DISCRETE = :DISCRETE
const UNIFORM = :UNIFORM
const NORMAL = :NORMAL
const GAMMA = :GAMMA
const BETA = :BETA
const LOGNORM = :LOGNORM
const DISTRIBUTIONS = [DISCRETE, UNIFORM, NORMAL, GAMMA, BETA, LOGNORM]

export
    SMPSScenario,
    SMPSSampler

include("tim.jl")
include("cor.jl")
include("lp.jl")
include("sto.jl")
include("raw.jl")
include("stage.jl")
include("scenario.jl")
include("sampler.jl")
include("model.jl")
include("reader.jl")

end


text:
using HorizonSideRobots

function x_cross_marking(r::Robot)
    for i in (Nord, Sud)
        for j in (Ost, West)
            diagonal_marking(r, i, j)
        end
    end
end

function diagonal_marking(r::Robot, s1::HorizonSide, s2::HorizonSide)
    k = 0
    while (isborder(r, s1) == false) && (isborder(r, s2) == false)
        putmarker!(r)
        move!(r, s1)
        move!(r, s2)
        k = k + 1
    end
    putmarker!(r)
    i = HorizonSide(mod(Int(s1) + 2, 4))
    j = HorizonSide(mod(Int(s2) + 2, 4))
    while k > 0
        k = k - 1
        move!(r, i)
        move!(r, j)
    end
end


text:
"""
    is_diagonal_noise(prob::DEProblem)

TODO
"""
is_diagonal_noise(prob::DEProblem) = false
is_diagonal_noise(prob::AbstractRODEProblem{uType,tType,isinplace,ND}) where {uType,tType,isinplace,ND} = ND <: Nothing

"""
    isinplace(prob::DEProblem)

Determine whether the function of the given problem operates in place or not.
"""
function isinplace(prob::DEProblem) end
isinplace(prob::AbstractODEProblem{uType,tType,iip}) where {uType,tType,iip} = iip
isinplace(prob::AbstractSteadyStateProblem{uType,iip}) where {uType,iip} = iip
isinplace(prob::AbstractRODEProblem{uType,tType,iip,ND}) where {uType,tType,iip,ND} = iip
isinplace(prob::AbstractDDEProblem{uType,tType,lType,iip}) where {uType,tType,lType,iip} = iip
isinplace(prob::AbstractDAEProblem{uType,duType,tType,iip}) where {uType,duType,tType,iip} = iip
isinplace(prob::AbstractNoiseProblem) = isinplace(prob.noise)
isinplace(::SplitFunction{iip}) where iip = iip


text:
module TestUnits

using Test
using Mimi

import Mimi: verify_units, connect_param!, ComponentReference, @defmodel

# Try directly using verify_units
@test verify_units("kg", "kg")
@test !verify_units("kg", "MT")
@test !verify_units("kg", "")

# Create a model with some matching and some mis-matching units
@defcomp Foo begin
    output = Variable(unit="kg")
end

@defcomp Bar begin
    input = Parameter(unit="MT")
end

@defcomp Baz begin
    input = Parameter(unit="kg")
end

@defmodel m begin
    index[time] = [1]
    component(Foo)
    component(Bar)
    component(Baz)
end

foo = ComponentReference(m, :Foo)
bar = ComponentReference(m, :Bar)
baz = ComponentReference(m, :Baz)

# Check that we cannot connect foo and bar...
@test_throws ErrorException bar[:input] = foo[:output]

# ...unless we pass ignoreunits=true
connect_param!(m, :Bar, :input,  :Foo, :output, ignoreunits=true)

# But we can connect foo and baz
baz[:input] = foo[:output]

end # module


text:
export line, sigmoid, hardsigmoid, relu, elu, selu, tanh, hardtanh, softsign, softplus

function line(x)
    return x
end

line(x, back::Bool) = 1

function sigmoid(x)
    t = exp(-abs(x))
    return (x>=0) ? inv(1 + t) : t / (t + 1)
end

sigmoid(x, back::Bool) = sigmoid(x)*(1-sigmoid(x))

hardsigmoid(x) = max(0, min(1, (x + 2.5) / 6))

function hardsigmoid(x, back::Bool)
    if -2.5 <= x <= 2.5
        return 0.2
    else
        return 0
    end
end

relu(x) = max(0, x)

function relu(x, back::Bool)
    if x <= 0
        return 0
    else
        return 1
    end
end

function elu(x; alpha = 1)
    return (x > 0) ? x : (exp(x) - 1)alpha
end

function elu(x, back::Bool; alpha = 1)
    if x < 0
        return alpha*exp(x)
    else
        return 1
    end
end

function selu(x)
    λ = oftype(float(x), selu_λ)
    α = oftype(float(x), selu_α)
    return ((x > 0) ? x : (exp(x)-1)α)*λ
end

function selu(x, back::Bool)
    λ = oftype(float(x), selu_λ)
    α = oftype(float(x), selu_α)
    if x < 0
        return λ*alpha*exp(x)
    else
        return λ
    end
end

const selu_λ = 1.0507009873554804934193349852946
const selu_α = 1.6732632423543772848170429916717

function tanh(x)
    return exp(x) - exp(-x) / exp(x) + exp(-x)
end

tanh(x, back::Bool) = 4 / (exp(x) + exp(-x)) ^ 2

function hardtanh(x)
    if x > 1
        return 1
    elseif x < -1
        return -1
    else
        return x
    end
end

function hardtanh(x, back::Bool)
    if x < -1 || 1 <= x
        return 0
    else
        return 1
    end
end

function softsign(x)
    return x / (1 + abs(x))
end

softsign(x, back::Bool) = 1 / ((1 + abs(x))^2)

function softplus(x)
    return log(1 + exp(x))
end

softplus(x, back::Bool) = 1 / (1 + exp(-x))


text:
# XPath
# =====

struct _XPathContext
    # type tag
end

struct _NodeSet
    nodeNr::Cint
    nodeMax::Cint
    nodeTab::Ptr{Ptr{_Node}}
end

const XPATH_UNDEFINED   = Cint(0)
const XPATH_NODESET     = Cint(1)
const XPATH_BOOLEAN     = Cint(2)
const XPATH_NUMBER      = Cint(3)
const XPATH_STRING      = Cint(4)
const XPATH_POINT       = Cint(5)
const XPATH_RANGE       = Cint(6)
const XPATH_LOCATIONSET = Cint(7)
const XPATH_USERS       = Cint(8)
const XPATH_XSLT_TREE   = Cint(9)

struct _XPathObject
    typ::Cint
    nodesetval::Ptr{_NodeSet}
    boolval::Cint
    floatval::Cdouble
    stringval::Ptr{UInt8}
    user::Ptr{Cvoid}
    index::Cint
    user2::Ptr{Cvoid}
    index2::Cint
end

"""
    findall(xpath::AbstractString, doc::Document)

Find nodes matching `xpath` XPath query from `doc`.
"""
function Base.findall(xpath::AbstractString, doc::Document)
    return findall(xpath, doc.node)
end

"""
    findfirst(xpath::AbstractString, doc::Document)

Find the first node matching `xpath` XPath query from `doc`.
"""
function Base.findfirst(xpath::AbstractString, doc::Document)
    # string("(", xpath, ")[position()=1]") may be faster
    nodes = findall(xpath, doc)
    return isempty(nodes) ? nothing : first(nodes)
end

"""
    findlast(doc::Document, xpath::AbstractString)

Find the last node matching `xpath` XPath query from `doc`.
"""
function Base.findlast(xpath::AbstractString, doc::Document)
    # string("(", xpath, ")[position()=last()]") may be faster
    nodes = findall(xpath, doc)
    return isempty(nodes) ? nothing : last(nodes)
end

"""
    findall(xpath::AbstractString, node::Node, [ns=namespaces(node)])

Find nodes matching `xpath` XPath query starting from `node`.

The `ns` argument is an iterator of namespace prefix and URI pairs.
"""
function Base.findall(xpath::AbstractString, node::Node, ns=namespaces(node))
    if !ismanaged(node)
        throw(ArgumentError("XPath query on the unmanaged node"))
    end
    context_ptr = new_xpath_context(document(node))
    if context_ptr == C_NULL
        throw_xml_error()
    end
    for (prefix, uri) in ns
        if isempty(prefix)
            @warn "ignored the empty prefix for '$(uri)'; expected to be non-empty"
        else
            ret = register_namespace!(context_ptr, prefix, uri)
            @assert ret == 0
        end
    end
    result_ptr = eval_xpath(node, context_ptr, xpath)
    if result_ptr == C_NULL
        free(context_ptr)
        throw_xml_error()
    end
    try
        result = unsafe_load(result_ptr)
        @assert result.typ == XPATH_NODESET
        if result.nodesetval == C_NULL
            return Vector{Node}()
        end
        nodeset = unsafe_load(result.nodesetval)
        # I don't know why, but this fails to infer the type of elements.
        return Node[Node(unsafe_load(nodeset.nodeTab, i)) for i in 1:nodeset.nodeNr]
    catch
        rethrow()
    finally
        free(context_ptr)
        free(result_ptr)
    end
end

"""
    findfirst(xpath::AbstractString, node::Node, [ns=namespaces(node)])

Find the first node matching `xpath` XPath query starting from `node`.
"""
function Base.findfirst(xpath::AbstractString, node::Node, ns=namespaces(node))
    # string("(", xpath, ")[position()=1]") may be faster
    nodes = findall(xpath, node, ns)
    return isempty(nodes) ? nothing : first(nodes)
end

"""
    findlast(node::Node, xpath::AbstractString, [ns=namespaces(node)])

Find the last node matching `xpath` XPath query starting from `node`.
"""
function Base.findlast(xpath::AbstractString, node::Node, ns=namespaces(node))
    # string("(", xpath, ")[position()=last()]") may be faster
    nodes = findall(xpath, node, ns)
    return isempty(nodes) ? nothing : last(nodes)
end

function new_xpath_context(doc)
    context_ptr = ccall(
        (:xmlXPathNewContext, libxml2),
        Ptr{_XPathContext},
        (Ptr{_XPathObject},),
        doc.node.ptr)
    return context_ptr
end

function register_namespace!(context_ptr, prefix, uri)
    ret = ccall(
        (:xmlXPathRegisterNs, libxml2),
        Cint,
        (Ptr{Cvoid}, Cstring, Cstring),
        context_ptr, prefix, uri)
    return ret
end

function eval_xpath(node, context_ptr, xpath)
    # set the pointer to `node` in the cotnext
    unsafe_store!(Ptr{UInt}(context_ptr), convert(UInt, node.ptr), 2)
    result_ptr = ccall(
        (:xmlXPathEval, libxml2),
        Ptr{_XPathObject},
        (Cstring, Ptr{Cvoid}),
        xpath, context_ptr)
    return result_ptr
end

function free(ptr::Ptr{_XPathContext})
    ccall(
        (:xmlXPathFreeContext, libxml2),
        Cvoid,
        (Ptr{Cvoid},),
        ptr)
end

function free(ptr::Ptr{_XPathObject})
    ccall(
        (:xmlXPathFreeObject, libxml2),
        Cvoid,
        (Ptr{Cvoid},),
        ptr)
end


text:
function handle_error(r::Response)
    r.status != 200 && throw(r)
end

# -------

sunlight_get(auth::String, api::URI, path, args; options...) = sunlight_get(auth, URI(api; path = path), args; options...)

function sunlight_get(auth::String, route::URI, args; query = Dict(), debug = false, options...)
    if auth == nothing || isempty(auth)
        error("Sunlight Foundation API key required! Get one here: http://sunlightfoundation.com/api")
    end

    query["apikey"] = auth
    args != {} && merge!(query, args)

    if debug
        println("URI: ", route)
        println("Params: ", query)
    end

    r = get(route; query = query, options...)

    if debug
        println(r)
    end

    handle_error(r)

    JSON.parse(r.data)
end


text:
module Repeat
using DelimitedFiles

function main()
    values = round.(Int, readdlm("../input/input.txt"))
    prevs = Set()
    total = 0
    done = false
    while !done
        for value in values
            if total in prevs
                println(total)
                done = true
                break
            end
            union!(prevs, [total])
            total += value
        end
        # println("-")
    end
end

main()

end

max_stars_repo_name: tjpalmer/advent

max_issues_repo_name: tjpalmer/advent

max_forks_repo_name: tjpalmer/advent


text:
abstract type
    RenewableGen <: Generator
end

struct RenewableFix <: RenewableGen
    name::String
    available::Bool
    bus::Bus
    tech::TechRenewable
    scalingfactor::TimeSeries.TimeArray
    function RenewableFix(name, status, bus, installedcapacity::Float64, scalingfactor)
        tech = TechRenewable(installedcapacity, nothing, 1.0)
        new(name, status, bus, tech, scalingfactor)
    end
end

RenewableFix(; name="init",
        status = false,
        bus = Bus(),
        installedcapacity = 0.0,
        scalingfactor = TimeSeries.TimeArray(Dates.today(),ones(1))) = RenewableFix(name, status, bus, installedcapacity, scalingfactor)

struct RenewableCurtailment <: RenewableGen
    name::String
    available::Bool
    bus::Bus
    tech::TechRenewable
    econ::Union{EconRenewable,Nothing}
    scalingfactor::TimeSeries.TimeArray # [0-1]
end

function RenewableCurtailment(name::String, status::Bool, bus::Bus, installedcapacity::Float64, econ::Union{EconRenewable,Nothing}, scalingfactor::TimeSeries.TimeArray)
    tech = TechRenewable(installedcapacity, nothing, 1.0)
    return RenewableCurtailment(name, status, bus, tech, econ, scalingfactor)
end

RenewableCurtailment(; name = "init",
                status = false,
                bus= Bus(),
                installedcapacity = 0.0,
                econ = EconRenewable(),
                scalingfactor =  TimeSeries.TimeArray(Dates.today(),ones(1))) = RenewableCurtailment(name, status, bus, installedcapacity, econ, scalingfactor)

struct RenewableFullDispatch <: RenewableGen
    name::String
    available::Bool
    bus::Bus
    tech::TechRenewable
    econ::Union{EconRenewable,Nothing}
    scalingfactor::TimeSeries.TimeArray # [0-1]
end

RenewableFullDispatch(; name = "init",
                status = false,
                bus= Bus(),
                installedcapacity = 0.0,
                econ = EconRenewable(),
                scalingfactor =  TimeSeries.TimeArray(Dates.today(),ones(1))) = RenewableCurtailment(name, status, bus, installedcapacity, econ, scalingfactor)


text:
function _elements(d::T; kwargs...) where T
    all_elements = Vector{Any}(undef, 0)
    _elements!(all_elements, d; kwargs...)
    unique!(all_elements)
    result = _DataElements(all_elements)
    return result
end

# function _elements!(all_elements::Vector{Any}, d::V; kwargs...) where V <: AbstractVector{D} where D <: Data{T} where T
#     # push!(all_elements, d)
#     _elements_fields!(all_elements, d; kwargs...)
#     _elements_iterable!(all_elements, d; kwargs...)
#     _elements_indexable!(all_elements, d; kwargs...)
#     return all_elements
# end

function _elements!(all_elements::Vector{Any}, d::D; kwargs...) where D <: Data{T} where T
    # push!(all_elements, d)
    _elements_fields!(all_elements, d; kwargs...)
    _elements_iterable!(all_elements, d; kwargs...)
    _elements_indexable!(all_elements, d; kwargs...)
    return all_elements
end

function _elements!(all_elements::Vector{Any}, d::T; kwargs...) where T
    push!(all_elements, d)
    _elements_fields!(all_elements, d; kwargs...)
    _elements_iterable!(all_elements, d; kwargs...)
    _elements_indexable!(all_elements, d; kwargs...)
    return all_elements
end

function _elements_fields!(all_elements::Vector{Any}, d::T; kwargs...) where T
    for field in fieldnames(T)
        _elements!(all_elements, _get_property(d, field); kwargs...)
    end
    return all_elements
end

function _elements_iterable!(all_elements::Vector{Any}, d::T; kwargs...) where T
    if _is_iterable(T)
        try
            for object in d
                try
                    _elements!(all_elements, object; kwargs...)
                catch ex_inner
                    # showerror(stderr, ex)
                    # Base.show_backtrace(stderr, catch_backtrace())
                    # @debug("Ignoring exception [inner]", exception=(ex, catch_backtrace()))
                end
            end
        catch ex_outer
            # showerror(stderr, ex)
            # Base.show_backtrace(stderr, catch_backtrace())
            # @debug("Ignoring exception [outer]", exception=(ex, catch_backtrace()))
        end
    end
    return all_elements
end

function _elements_indexable!(all_elements::Vector{Any}, d::T; kwargs...) where T
    if _has_isassigned(T)
        _elements_indexable_with_check_assigned!(all_elements, d; kwargs...)
    else
        _elements_indexable_without_check_assigned!(all_elements, d; kwargs...)
    end
    return all_elements
end

function _elements_indexable_with_check_assigned!(all_elements::Vector{Any}, d::T; kwargs...) where T
    if _is_indexable(T)
        for i = 1:length(d)
            if isassigned(d, i)
                _elements!(all_elements, d[i]; kwargs...)
            end
        end
    end
    return all_elements
end

function _elements_indexable_without_check_assigned!(all_elements::Vector{Any}, d::T; kwargs...) where T
    if _is_indexable(T)
        for i = 1:length(d)
            try
                _elements!(all_elements, d[i])
            catch
            end
        end
    end
    return all_elements
end


text:
# small scrip to compute the solution of Lippmann-Schwinger equation
# We test the types introduced in FastConvolution.jl
# we test that the application is fast and that the construction
# is performed fast.

using PyPlot
using IterativeSolvers

include("../src/FastConvolution.jl")
include("../src/Preconditioner.jl")

# setting the number of threads for the FFT and BLAS
# libraries (please set them to match the number of
# physical cores in your system)
FFTW.set_num_threads(4);
BLAS.set_num_threads(4);

#Defining Omega
h = 0.002
k = 20*2*pi

# size of box
a = 1
x = collect(-a/2:h:a/2)
y = collect(-a/2:h:a/2)
(n,m) = length(x), length(y)
N = n*m
X = repmat(x, 1, m)[:]
Y = repmat(y', n,1)[:]
# we solve \triangle u + k^2(1 + nu(x))u = 0

# We use the modified quadrature in Ruan and Rohklin
(ppw,D) = referenceValsTrapRule();
D0 = D[1];

xs = -0.4; ys = -0.4;                     # point source location
sigma = 0.15;

window(y,alpha, beta) = 1*(abs(y).<=beta) + (abs(y).>beta).*(abs(y).<alpha).*exp(2*exp(-(alpha- beta)./(abs(y)-beta))./ ((abs(y)-beta)./(alpha- beta)-1) )

# version for Jun
# xHet = 0.1;
# yHet = 0.1;
# # Defining the smooth perturbation of the slowness
# nu(x,y) = -0.5*exp( -1/(2*sigma^2)*((x-xHet).^2 + (y-yHet).^2) ).*window(sqrt((x-xHet).^2 + (y-yHet).^2), 0.3,0.38  );

#version for Matthias

xHet = [ -0.131 -0.055  0.052 -0.399 -0.082 -0.170  0.253  -0.192   0.186  0.139 -0.080  -0.213  0.108 -0.235  0.178  0.052  0.179 -0.178  0.244  0.020 ];
yHet = [ -0.417 -0.103  0.371  0.384  0.300  0.121 -0.026   0.352  -0.054 -0.455 -0.262  0.080 -0.053  0.364 -0.210  0.349  0.210 -0.247  0.142 -0.309 ];
sigma = [ 0.162  0.1295 0.175 0.143  0.089  0.169  0.195  0.188  0.101  0.119  0.0365  0.092 0.059 0.182  0.090 0.174  0.082 0.094 0.060 0.184];
amplitude = 18*[-0.0286 -0.0422  0.0261 -0.0067 -0.0538  0.0347  0.040  -0.057 -0.090  0.0262  0.014  0.039 -0.060 -0.083  0.046 -0.041 -0.059  0.059 -0.096 -0.082];

# Defining the smooth perturbation of the slowness
nu(x,y) =  (amplitude[1]*exp( -1/(2*sigma[1]^2)*(2*(x-xHet[1]).^2 + (y-yHet[1]).^2) )+
            amplitude[2]*exp( -1/(2*sigma[2]^2)*((x-xHet[2]).^2 + (y-yHet[2]).^2) )+
            amplitude[3]*exp( -1/(2*sigma[3]^2)*((x-xHet[3]).^2 + (y-yHet[3]).^2) )+
            amplitude[4]*exp( -1/(2*sigma[4]^2)*((x-xHet[4]).^2 + 3*(y-yHet[4]).^2) )+
            amplitude[5]*exp( -1/(2*sigma[5]^2)*((x-xHet[5]).^2 + (y-yHet[5]).^2) )+
            amplitude[6]*exp( -1/(2*sigma[6]^2)*((x-xHet[6]).^2 + (y-yHet[6]).^2) )+
            amplitude[7]*exp( -1/(2*sigma[7]^2)*((x-xHet[7]).^2 + 4*(y-yHet[7]).^2) )+
            amplitude[8]*exp( -1/(2*sigma[8]^2)*((x-xHet[8]).^2 + (y-yHet[8]).^2) )+
            amplitude[9]*exp( -1/(2*sigma[9]^2)*((x-xHet[9]).^2 + (y-yHet[9]).^2) )+
            amplitude[10]*exp( -1/(2*sigma[10]^2)*((x-xHet[10]).^2 + (y-yHet[10]).^2) )+
            amplitude[11]*exp( -1/(2*sigma[11]^2)*((x-xHet[11]).^2 + (y-yHet[11]).^2) )+
            amplitude[12]*exp( -1/(2*sigma[12]^2)*(2*(x-xHet[12]).^2 + 4*(y-yHet[12]).^2) )+
            amplitude[13]*exp( -1/(2*sigma[13]^2)*((x-xHet[13]).^2 + 2*(y-yHet[13]).^2) )+
            amplitude[14]*exp( -1/(2*sigma[14]^2)*((x-xHet[14]).^2 + 5*(y-yHet[14]).^2) )+
            amplitude[15]*exp( -1/(2*sigma[15]^2)*((x-xHet[15]).^2 + (y-yHet[15]).^2) )+
            amplitude[16]*exp( -1/(2*sigma[16]^2)*((x-xHet[16]).^2 + 4*(y-yHet[16]).^2) )+
            amplitude[17]*exp( -1/(2*sigma[17]^2)*((x-xHet[17]).^2 + (y-yHet[17]).^2) )+
            amplitude[18]*exp( -1/(2*sigma[18]^2)*(3*(x-xHet[18]).^2 + 5*(y-yHet[18]).^2) )+
            amplitude[19]*exp( -1/(2*sigma[19]^2)*(3*(x-xHet[19]).^2 + (y-yHet[19]).^2) )+
            amplitude[20]*exp( -1/(2*sigma[20]^2)*((x-xHet[20]).^2 + (y-yHet[20]).^2) ) ).*(window(abs(x),0.45,0.4 ) ).*(window(abs(y),0.48,0.4 ) ).*
            (1 - window(sqrt( (x-xs).^2 + (y-ys).^2),0.15,0.05 ) );

figure(1);clf();
imshow(reshape(sqrt(1 - nu(X,Y)), n,m),extent=[y[1], y[end], x[end], x[1]]);cb =  colorbar();

## You can choose between Duan Rohklin trapezoidal quadrature
#fastconv = buildFastConvolution(x,y,h,k,nu)

# or Greengard Vico Quadrature (this is not optimized and it is 2-3 times slower)
fastconv = buildFastConvolution(x,y,h,k,nu, quadRule = "Greengard_Vico");

# assembling the sparsifiying preconditioner
@time As = buildSparseA(k,X,Y,D0, n ,m);

# assembling As*( I + k^2G*nu)
@time Mapproxsp = As + k^2*(buildSparseAG(k,X,Y,D0, n ,m)*spdiagm(nu(X,Y)));

# defining the preconditioner
precond = SparsifyingPreconditioner(Mapproxsp, As)

# building the RHS from the incident field
u_inc = 1im/4*hankelh1(0, k*sqrt((X-xs).^2 + (Y-ys).^2)+eps(1.0));
rhs = -k^2*FFTconvolution(fastconv, nu(X,Y).*u_inc) ;
figure(2)
clf()
imshow(imag(reshape(rhs,n,m)))
#rhs = u_inc;

# allocating the solution
u = zeros(Complex128,N);

# solving the system using GMRES
@time info =  gmres!(u, fastconv, rhs, precond)
println(info[2].residuals[:])

# plotting the solution
figure(3)
clf()
imshow(imag(reshape(u+u_inc,n,m)))

max_stars_repo_name: Forgotten/Lippmann-Schingwer

max_issues_repo_name: Forgotten/Lippmann-Schingwer

max_forks_repo_name: Forgotten/Lippmann-Schingwer


text:
include("event.jl")

#Initiate an empty list to store the strategies

strategylist = []

mutable struct Strategy

   """

      Class Portfolio:

      Create a portfolio object that is being used to pass order and to follow up on the equity when a strategy is being

      traded.

          Parameters

          ----------

          name: {int64}, optional, default Unknown with number depending on the number of strategies

              The strategy name.

          leverage: {float64}, optional, default 1

              The leverage used for passing a trade.

          Attributes

          ----------

          name: int64

              The strategy name.

          leverage: float64

              The leverage used for passing a trade.

          entrylongsignals: list

              The list of signals used to entry into a long position

          exitlongsignals: list

              The list of signals used to exit a long position

          entryshortsignals: list

              The list of signals used to entry into a short position

          exitshortsignals: list

              The list of signals used to exit a short position

      Example:
         Strategy()

         """
   asset::String
   data::Array
   pipsvalue::Float64
   Strategy(name = "Unknown", leverage = 1) = begin
      if name != "Unknown"
         name = name
      else
         if strategylist != []
            count = []
            for i in strategylist
               if "Unknown" in i.name
                  push!(count, "True")
            if length(count) == 0
               name = "Unknown"
            else
               name = "Unknown" * string(length(count))
         else
            name = name
      leverage = leverage
      entrylongsignals = []
      exitlongsignals = []
      entryshortsignals = []
      exitshortsignals = []
   newstrategy = new(name, leverage, entrylongsignals, exitlongsignals, entryshortsignals, exitshortsignals)
      push!(strategylist, newstrategy)
   end
end

function addentrysignal(strategy, signal, type)
      """ Add a signal to enter into a trade; the signal will be added to the strategy either for "Long" or "Short" trades
      depending on the type given in input. Market has to respect all entry signals before an order is being passed """
   if type == "Long"
      push!(strategy.entrylongsignals, signal)
   elseif type == "Short"
      push!(strategy.entryshortsignals, signal)
   end
end

function addexitsignal(strategy, signal, type)
      """ Add a signal to exit a trade; the signal will be added to the strategy either for "Long" or "Short" trades

      depending on the type given in input. Market has to respect all entry signals before an order is being passed """
   if type == "Long"
      push!(strategy.exitlongsignals, signal)
   elseif type == "Short"
      push!(strategy.exitshortsignals, signal)
   end
end

function entryconditions(strategy, asset)
      """ Determine if the market respects all the entry conditions and if a new order event can be passed. If True a new
      order event is being added to the event queue """
   if count(i -> e.createsignalevent(asset, i)!= false, strategy.entrylongsignals) == length(strategy.entrylongsignals)
      newevent = e.OrderEvent(asset, strategy.leverage, "BUY")
      add_event(newevent)
   elseif count(i -> e.createsignalevent(asset, i)!= false, strategy.entryshortsignals) == length(strategy.entryshortsignals)
      newevent = e.OrderEvent(asset, strategy.leverage, "SELL")
      add_event(newevent)
   end
end

function exitconditions(strategy, asset, portfolio)
      """ Determine if the market respects all the exit conditions and if a new fill event can be passed. If True a new
      fill event is being added to the event queue """

   if portfolio.orderslist != []
      if count(i -> e.createsignalevent(asset, i)!= false, strategy.exitlongsignals) == length(strategy.exitlongsignals)
         newevent = e.FillEvent(asset, strategy.leverage, "BUY")
         add_event(newevent)
      elseif count(i -> e.createsignalevent(asset, i)!= false, strategy.exitshortsignals) == length(strategy.exitshortsignals)
         newevent = e.FillEvent(asset, self.leverage, "SELL")
         add_event(newevent)
      end
   end
end

max_stars_repo_name: SRelieved/JuliaBacktesting

max_issues_repo_name: MartinGenesse/JuliaBacktesting

max_forks_repo_name: MartinGenesse/JuliaBacktesting


text:
export gmsh, init_gmsh, addPoint, addLine, addCurveLoop,
addPlaneSurface, addPhysicalGroup, setPhysicalName, addCircleArc, addSurfaceFilling,
embedPointInLine, embedPointInSurface, embedLine, sync, finalize_gmsh, meshsize, addCircle, compound,
get_entities, addPhysicalGroup

include(joinpath(@__DIR__, "..", "..", "deps", "Gmsh", "gmsh.jl"))
import .gmsh

gmsh_number = 1
is_sync = false

function init_gmsh(name::String = "default")
    global is_sync = false
    gmsh.initialize()
    gmsh.option.setNumber("Mesh.SaveAll", 1)
    gmsh.option.setNumber("General.Terminal", gmsh_number)
    global gmsh_number += 1
    gmsh.model.add(name)
end

function addPoint(x, y, meshSize = 0., tag = -1)
    gmsh.model.geo.addPoint(x, y, 0.0, meshSize, tag)
end

function addLine(startTag, endTag, tag = -1)
    gmsh.model.geo.addLine(startTag, endTag, tag)
end

"""
Add a curve loop (a closed wire) formed by the curves curveTags.
"""
function addCurveLoop(curveTags, tag = -1)
    gmsh.model.geo.addCurveLoop(curveTags, -1)
end

function addPlaneSurface(wireTags, tag = -1)
    gmsh.model.geo.addPlaneSurface(wireTags,  -1)
end

function addPhysicalGroup(dim, tags, tag = -1)
    gmsh.model.addPhysicalGroup(dim, tags,  -1)
end

function setPhysicalName(dim, tag, name)
    gmsh.model.setPhysicalName(dim, tag, name)
end

function addCircleArc(startTag, centerTag, endTag, tag = -1)
    gmsh.model.geo.addCircleArc(startTag, centerTag, endTag, tag)
end

function addSurfaceFilling(wireTags, tag = -1)
    gmsh.model.geo.addSurfaceFilling(wireTags, tag, -1)
end

function embedPointInLine(tags, inTag)
    sync()
    gmsh.model.mesh.embed(0, tags, 1, inTag)
end

function embedPointInSurface(tags, inTag)
    sync()
    gmsh.model.mesh.embed(0, tags, 2, inTag)
end

function embedLine(tags, inTag)
    sync()
    gmsh.model.mesh.embed(1, tags, 2, inTag)
end

function meshsize(fn::String)
    sync()
    field = gmsh.model.mesh.field
    field.add("MathEval", 1)
    field.setString(1, "F", fn)
    field.setAsBackgroundMesh(1)
end

function sync()
    global is_sync = true
    gmsh.model.geo.synchronize()
end

function addCircle(x, y, r, tag=-1)
    pts = [
        addPoint(x, y-r)
        addPoint(x+r, y)
        addPoint(x,y+r)
        addPoint(x-r,y)
        addPoint(x,y)
    ]
    arcs = [
        addCircleArc(pts[1],pts[5],pts[2])
        addCircleArc(pts[2],pts[5],pts[3])
        addCircleArc(pts[3],pts[5],pts[4])
        addCircleArc(pts[4],pts[5],pts[1])
    ]
    addCurveLoop(arcs)
end

function get_entities(tag = -1)
    gmsh.model.getEntities(tag)
end

"""

Forms a physical group of dimension `dim` and tag `tag` with name `name`
"""
function addPhysicalGroup(dim, tags, name, tag=-1)
    if !isa(tags, Array)
        tags = [tags]
    end
    ps = gmsh.model.addPhysicalGroup(dim, tags, tag)
    gmsh.model.setPhysicalName(dim, ps, name)
end

function finalize_gmsh(visualize=false)
    sync()
    ss = get_entities(2)
    for (_, s) in ss
        gmsh.model.mesh.setRecombine(2, s)
    end
    gmsh.model.mesh.generate(2)
    name = tempname()*".msh"
    gmsh.write(name)
    if visualize
        gmsh.fltk.run()
    end
    gmsh.finalize()
    return name
end


text:
using PyPlot
using PyCall
using MAT
using LinearAlgebra
using Statistics
mpl = pyimport("tikzplotlib")

n = 10
m = 2n
h = 0.01
out = 0.5 * ones(n)
out[1:div(n,3)] .= 1.0

out *= 50
rg = 0:20:200
err = zeros(4, length(rg))
for (j,idx) in enumerate([1])
    for (i,k) in enumerate(rg)
        eta = matread("$idx/eta$k.mat")["eta"]
        eta = eta[1:4m:end]
        err[j, i] = mean( ((out-eta) ./ out).^2)
        @info j, i, err[j,i]
    end
end

close("all")
# tp = ["ro-", "g+-", "b*-", "k>-"]
# labels = []
# for k in [1,2,3,5]
# s = length(1:k:m+1)
# push!(labels, "$s Sensors")
# end
# for i = 1:4
#     plot(rg .+ 1, err[i,:], tp[i], label=labels[i])
# end
semilogy(rg .+ 1, err[1,:], tp[1])
legend()
xlabel("Iterations")
ylabel("Error")
mpl.save("space_loss.tex")
savefig("loss.png")

function visualize_inv_eta(X)
    ηmin = 0.5*50.
    ηmax = 50.
    x = LinRange(0.5h,m*h, m)
    y = LinRange(0.5h,n*h, n)
    V = zeros(m, n)
    for i = 1:m
        for j = 1:n
            elem = (j-1)*m + i
            V[i, j] = mean(X[4(elem-1)+1:4elem])
        end
    end
    close("all")
    pcolormesh(x, y, V', vmin=ηmin-(ηmax-ηmin)/4, vmax=ηmax+(ηmax-ηmin)/4, rasterized=true)
    colorbar(shrink=0.5)
    xlabel("x")
    ylabel("y")
    # title("Iteration = $k")
    axis("scaled")
    gca().invert_yaxis()
end

out = repeat(out, 1, 4m)'[:]
visualize_inv_eta(out)
savefig("space_true.pdf")

eta  = matread("1/eta0.mat")["eta"]
visualize_inv_eta(eta)
savefig("space0.pdf")

eta  = matread("1/eta60.mat")["eta"]
visualize_inv_eta(eta)
savefig("space60.pdf")

eta  = matread("1/eta120.mat")["eta"]
visualize_inv_eta(eta)
savefig("space120.pdf")

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages


text:
struct KruskalHeapEntry{T<:Real}
    edge::Edge
    dist::T
end

isless(e1::KruskalHeapEntry, e2::KruskalHeapEntry) = e1.dist < e2.dist

"""
    quick_find!(vs, p, q)

Perform [Quick-Find algorithm](https://en.wikipedia.org/wiki/Disjoint-set_data_structure)
on a given pair of vertices `p`and `q`, and make a connection between them in the vector `vs`.
"""
function quick_find!(vs, p, q)
    pid = vs[p]
    qid = vs[q]
    for i in 1:length(vs)
        if vs[i] == pid
            vs[i] = qid
        end
    end
end

"""
    kruskal_mst(g, distmx=DefaultDistance())

Return a vector of edges representing the minimum spanning tree of a connected, undirected graph `g` with optional
distance matrix `distmx` using [Kruskal's algorithm](https://en.wikipedia.org/wiki/Kruskal%27s_algorithm).
"""
function kruskal_mst end
@traitfn function kruskal_mst{T}(
    g::::(!IsDirected),
    distmx::AbstractMatrix{T} = DefaultDistance()
)

    U = eltype(g)
    edge_list = Vector{KruskalHeapEntry{T}}()
    mst = Vector{Edge}()
    connected_vs = collect(one(U):nv(g))

    sizehint!(edge_list, ne(g))
    sizehint!(mst, ne(g))

    for e in edges(g)
        heappush!(edge_list, KruskalHeapEntry{T}(e, distmx[src(e), dst(e)]))
    end

    while !isempty(edge_list) && length(mst) < nv(g) - 1
        heap_entry = heappop!(edge_list)
        v = src(heap_entry.edge)
        w = dst(heap_entry.edge)

        if connected_vs[v] != connected_vs[w]
            quick_find!(connected_vs, v, w)
            push!(mst, heap_entry.edge)
        end
    end

    return mst
end

max_stars_repo_name: AsileBcd/JohnsonCycles

max_issues_repo_name: AsileBcd/JohnsonCycles

max_forks_repo_name: AsileBcd/JohnsonCycles


text:
"""
$(TYPEDSIGNATURES)

Edit region numbers of grid cells via rectangular mask.
"""
function cellmask!(grid::ExtendableGrid,
                   maskmin::AbstractArray,
                   maskmax::AbstractArray,
                   ireg::Int;
                   tol=1.0e-10)
    xmaskmin=maskmin.-tol
    xmaskmax=maskmax.+tol
    ncells=num_cells(grid)
    cellnodes=grid[CellNodes]
    dim=dim_space(grid)
    cellregions=grid[CellRegions]
    coord=grid[Coordinates]
    for icell=1:ncells
        in_region=true
        for inode=1:num_targets(cellnodes,icell)
            ignode=cellnodes[inode,icell]
            for idim=1:dim
                if coord[idim,ignode]<xmaskmin[idim]
                    in_region=false
                elseif coord[idim,ignode]>xmaskmax[idim]
                    in_region=false
                end
            end
        end
        if in_region
            cellregions[icell]=ireg
        end
    end
    grid[NumCellRegions]=max(num_cellregions(grid),ireg)
    grid
end

"""
$(TYPEDSIGNATURES)

Edit region numbers of grid  boundary facets  via rectangular mask.
For 1D grids, inner boundaries can be added by this method.
"""
function bfacemask!(grid::ExtendableGrid,
                    maskmin::AbstractArray,
                    maskmax::AbstractArray,
                    ireg::Int;
                    tol=1.0e-10)

    xmaskmin=maskmin.-tol
    xmaskmax=maskmax.+tol

    nbfaces=num_bfaces(grid)
    bfacenodes=grid[BFaceNodes]
    Ti=eltype(bfacenodes)
    dim=dim_space(grid)
    bfaceregions=grid[BFaceRegions]
    new_bfacenodes=ElasticArray{Ti,2}(bfacenodes)
    coord=grid[Coordinates]

    function isbface(ix)
        for ibface=1:num_bfaces(grid)
            if bfacenodes[1,ibface]==ix
                return ibface
            end
        end
        return 0
    end

    function isbface(ix,iy)
        for ibface=1:num_bfaces(grid)
            if (bfacenodes[1,ibface] == ix && bfacenodes[2,ibface] == iy) ||
                (bfacenodes[1,ibface] == iy && bfacenodes[2,ibface] == ix)
                return ibface
            end
        end
        return 0
    end

    if dim_space(grid)==1
        for inode=1:num_nodes(grid)
            x=coord[1,inode]
            if x>xmaskmin[1] && x<xmaskmax[1]
                ibface=isbface(inode)
                if ibface>0
                    bfaceregions[ibface]=ireg
                else
                    push!(bfaceregions,ireg)
                    append!(new_bfacenodes,[inode])
                end
            end
        end
    else
        facenodes=grid[FaceNodes]
        for iface=1:size(facenodes,2)
            in_region=true
            for inode=1:num_targets(facenodes,iface)
                ignode=facenodes[inode,iface]
                for idim=1:dim_space(grid)
                    if coord[idim,ignode]<xmaskmin[idim]
                        in_region=false
                    elseif coord[idim,ignode]>xmaskmax[idim]
                        in_region=false
                    end
                end
            end
            if in_region
                ibface=isbface(facenodes[1,iface],facenodes[2,iface])
                if ibface>0
                    bfaceregions[ibface]=ireg
                else
                    push!(bfaceregions,ireg)
                    append!(new_bfacenodes,[facenodes[1,iface],facenodes[2,iface]])
                end
            end
        end
    end
    btype=grid[BFaceGeometries][1]
    grid[BFaceNodes]=Array{Ti,2}(new_bfacenodes)
    grid[BFaceGeometries]=VectorOfConstants(btype,length(bfaceregions))
    grid[NumBFaceRegions]=max(num_bfaceregions(grid),ireg)
    return grid
end

"""
    $(TYPEDSIGNATURES)

Edit region numbers of grid  boundary edges via line mask.
This only works for 3D grids.
"""
function bedgemask!(grid::ExtendableGrid,
                    xa::AbstractArray,
                    xb::AbstractArray,
                    ireg::Int;
                    tol=1.0e-10)
    # Masking of boundary edges makes only sense in 3D
    @assert (dim_space(grid) > 2)

    masked = false

    nbedges        = num_bedges(grid)
    bedgenodes     = grid[BEdgeNodes]
    Ti             = eltype(bedgenodes)
    dim            = dim_space(grid)
    bedgeregions   = grid[BEdgeRegions]
    new_bedgenodes = ElasticArray{Ti,2}(bedgenodes)
    coord          = grid[Coordinates]
    Tv             = eltype(coord)

    # length of boundary edge region
    distsq         = sqrt((xa[1]-xb[1])^2 + (xa[2]-xb[2])^2 + (xa[3]-xb[3])^2)

    bedgenodes = grid[BEdgeNodes]
    # loop over boundary edges
    for ibedge = 1:size(bedgenodes, 2)
        in_region = true

        #loop over nodes of boundary edge
        for inode=1:num_targets(bedgenodes, ibedge)
            ignode = bedgenodes[inode, ibedge]

            # we compute the distance of the boundary edge node to the endpoints
            # if the sum of the distances is larger (with tolerance) than the length
            # of the boundary region, the point does not lie on the edge
            distxa = sqrt((xa[1]-coord[1,ignode])^2
                     + (xa[2]-coord[2,ignode])^2
                     + (xa[3]-coord[3,ignode])^2)
            distxb = sqrt((coord[1,ignode]-xb[1])^2
                     + (coord[2,ignode]-xb[2])^2
                     + (coord[3,ignode]-xb[3])^2)
            diff   = distxa + distxb - distsq
            if (diff > tol)
                in_region = false
                continue
            end
        end

        if in_region
            masked = true
            bedgeregions[ibedge] = ireg
        end
    end
    if !masked
        @warn "Couldn't mask any boundary edges for region $(ireg)"
    end

    grid[NumBEdgeRegions]=max(num_bedgeregions(grid),ireg)
    return grid
end


text:
function make_prelearn(; n_frames, height, width, margin)
	marg_width = width ÷ 2 - margin
	marg_height = height ÷ 2 - margin
	xpos = range(-marg_width, marg_width, length = 3)
	ypos = range(-marg_height, marg_height, length = 5)
	dist = Normal(0, 50)

	circ_start_points = Point.(xpos[1], ypos)
	circ_middle_points = [Point(rand(dist), rand(dist)) for i in circ_start_points]
	middle_shifts = circ_middle_points .- circ_start_points
	circ_final_points = [
		Point(rand(xpos[3] - 2margin:xpos[3]), ypos[i]) for i in 1:length(circ_start_points)
		]
	final_shifts = circ_final_points .- circ_middle_points

	time_steps = floor.(Int, range(1, n_frames, length=6))

	start_lines = [Object(time_steps[1]:time_steps[6], JLine(point + Point(10, 0), point + Point(marg_width/3, 0))) for point in circ_start_points]
	middle_lines = [
		Object(
			time_steps[3]:time_steps[6],
			JLine(p1, p2)
			)
		for (p1, p2) in Base.product(circ_middle_points, circ_middle_points)
		if p1 != p2
		]
	final_lines = [Object(time_steps[5]:time_steps[6], @JShape begin
				arrow(circ_final_points[i], circ_final_points[i+1])
			end) for i in 1:length(circ_final_points)-1]

	points = [Object(time_steps[1]:time_steps[2], JCircle(point, 5, action=:fill)) for point in circ_start_points]
	final_points = [
		Object(time_steps[2]:time_steps[6], JCircle(point, 5, action=:fill, color=randomhue())) for point in circ_start_points
		]

	for (idx, point) in enumerate(final_points)
		act!(
			point,
			Action(time_steps[1]:time_steps[2], anim_translate(middle_shifts[idx]))
		)
		act!(
			point,
			Action(time_steps[3]:time_steps[4], anim_translate(final_shifts[idx]))
		)
	end
	final_points
end


text:
module SaveFemIpState

include("state.jl")

export AbstractIpState, IpStateSingle, IpStateArray

export getIpState!, createIpStateDict, updateIpStateDict!, updateIpStateDict4rmBuffer!, updateStateDict!

end # module


text:
using Printf
using Random
using Statistics
using LinearAlgebra

using Flux
using DifferentialEquations
using DiffEqFlux

using JLD2
using BSON
using Plots

using Flux: @epochs

#####
##### Load data from JLD2 file
#####

file = jldopen("../data/ocean_convection_profiles.jld2")

Is = keys(file["timeseries/t"])

Nz = file["grid/Nz"]
Lz = file["grid/Lz"]
Nt = length(Is)

t = zeros(Nt)
T = T_data = zeros(Nt, Nz)
wT = zeros(Nt, Nz)

for (i, I) in enumerate(Is)
    t[i] = file["timeseries/t/$I"]
    T[i, :] = file["timeseries/T/$I"][1, 1, 2:Nz+1]
    wT[i, :] = file["timeseries/wT/$I"][1, 1, 2:Nz+1]
end

#####
##### Plot animation of T(z,t) from data
#####

z = file["grid/zC"]
#=
anim = @animate for n=1:10:Nt
    t_str = @sprintf("%.2f", t[n] / 86400)
    plot(T[n, :], z, linewidth=2,
         xlim=(19, 20), ylim=(-100, 0), label="",
         xlabel="Temperature (C)", ylabel="Depth (z)",
         title="Deepening mixed layer: $t_str days", show=false)
end

gif(anim, "deepening_mixed_layer.gif", fps=15)
=#
#####
##### Coarse grain data to 32 vertical levels (plus halo regions)
#####

function coarse_grain(data, resolution)
    @assert length(data) % resolution == 0
    s = length(data) / resolution

    data_cs = zeros(resolution)
    for i in 1:resolution
        t = data[Int((i-1)*s+1):Int(i*s)]
        data_cs[i] = mean(t)
    end

    return data_cs
end

coarse_resolution = cr = 32

T_cs = zeros(coarse_resolution+2, Nt)
wT_cs = zeros(coarse_resolution+2, Nt)

z_cs = coarse_grain(collect(z), cr)

for n=1:Nt
    T_cs[2:end-1, n] .= coarse_grain(T[n, :], coarse_resolution)
    wT_cs[2:end-1, n] .= coarse_grain(wT[n, :], coarse_resolution)
end

# Fill halo regions to enforce boundary conditions.
T_cs[1,   :] .= T_cs[2,     :]
T_cs[end, :] .= T_cs[end-1, :]

wT_cs[1,   :] .= wT_cs[2,     :]
wT_cs[end, :] .= wT_cs[end-1, :]

#####
##### Plot coarse temperature and wT profiles
#####
#=
@info "Plotting coarse temperature profile..."

anim = @animate for n=1:10:Nt
    t_str = @sprintf("%.2f", t[n] / 86400)
    plot(T_cs[2:cr+1, n], z_cs, linewidth=2,
         xlim=(19, 20), ylim=(-100, 0), label="",
         xlabel="Temperature (C)", ylabel="Depth (z)",
         title="Deepening mixed layer: $t_str days", show=false)
end

gif(anim, "deepening_mixed_layer_T_coarse.gif", fps=15)

@info "Plotting coarse wT profile..."

anim = @animate for n=1:10:Nt
    t_str = @sprintf("%.2f", t[n] / 86400)
    plot(wT_cs[2:cr+1, n], z_cs, linewidth=2,
         xlim=(-1e-4, 1e-4), ylim=(-100, 0), label="",
         xlabel="Temperature (C)", ylabel="Depth (z)",
         title="Deepening mixed layer: $t_str days", show=false)
end

gif(anim, "deepening_mixed_layer_wT_coarse.gif", fps=15)
=#
#####
##### Generate differentiation matrices
#####

cr_Δz = Lz / cr  # Coarse resolution Δz

# Dzᶠ computes the derivative from cell center to cell (F)aces
Dzᶠ = 1/cr_Δz * Tridiagonal(-ones(cr+1), ones(cr+2), zeros(cr+1))

# Dzᶜ computes the derivative from cell faces to cell (C)enters
Dzᶜ = 1/cr_Δz * Tridiagonal(zeros(cr+1), -ones(cr+2), ones(cr+1))

# Impose boundary condition that derivative goes to zero at top and bottom.
# Dzᶠ[1, 1] = 0
# Dzᶜ[cr, cr] = 0

#####
##### Create training data
#####

Tₙ    = zeros(cr+2, Nt-1)
Tₙ₊₁  = zeros(cr+2, Nt-1)
wTₙ   = zeros(cr+2, Nt-1)
∂zTₙ  = zeros(cr+2, Nt-1)
∂zwTₙ = zeros(cr+2, Nt-1)

for i in 1:Nt-1
       Tₙ[:, i] .=  T_cs[:,   i]
     Tₙ₊₁[:, i] .=  T_cs[:, i+1]
      wTₙ[:, i] .= wT_cs[:,   i]
     ∂zTₙ[:, i] .= Dzᶠ * T_cs[:, i]
    ∂zwTₙ[:, i] .= Dzᶜ * wT_cs[:, i]
end

N_skip = 0  # Skip first N_skip iterations to avoid learning transients?
N = 3  # Number of training data pairs.

# rinds = randperm(Nt-N_skip)[1:N]

# pre_training_data = [(Tₙ[:, i], ∂zwTₙ[:, i]) for i in 1:N]
pre_training_data = [(∂zTₙ[:, i], -wTₙ[:, i]) for i in 1:N]
training_data = [(Tₙ[:, i], Tₙ₊₁[:, i]) for i in 1:N]

#####
##### Create heat flux vector
#####

heat_flux = zeros(cr+2)

Q  = -100 # file["parameters/surface_cooling"]
ρ₀ = file["parameters/density"]
cₚ = file["parameters/specific_heat_capacity"]

heat_flux[2] = Q / (ρ₀ * cₚ * cr_Δz)

#####
##### Create neural network
#####

# Complete black box right-hand-side.
#  dTdt_NN = Chain(Dense(cr+2,  2cr, tanh),
#                  Dense(2cr,  cr+2))

# Use NN to parameterize a diffusivity or κ profile.
#  dTdt_NN = Chain(T -> Dzᶠ*T,
#                Dense(cr+2,  2cr, tanh),
#                Dense(2cr,  cr+2),
#                NNDzT -> Dzᶜ * NNDzT + heat_flux)

# Use NN to parameterize flux.
dTdt_NN = Chain(Dense(cr+2,  2cr, tanh),
                Dense(2cr,  cr+2),
                NN -> Dzᶜ * NN + heat_flux)

NN_params = Flux.params(dTdt_NN)

#####
##### Pre-train the neural network on (T, wT) data pairs
#####

pre_loss_function(∂zTₙ, wTₙ) = sum(abs2, dTdt_NN(∂zTₙ) .- wTₙ)

popt = ADAM(0.01)

function precb()
    loss = sum(abs2, [pre_loss_function(pre_training_data[i]...) for i in 1:N-1])
    println("loss = $loss")
end

pre_train_epochs = 5
for _ in 1:pre_train_epochs
    Flux.train!(pre_loss_function, NN_params, pre_training_data, popt, cb = Flux.throttle(precb, 5))
end

#####
##### Define loss function
#####

tspan = (0.0, 600.0)  # 10 minutes
neural_pde_prediction(T₀) = neural_ode(dTdt_NN, T₀, tspan, Tsit5(), reltol=1e-4, save_start=false, saveat=tspan[2])

loss_function(Tₙ, Tₙ₊₁) = sum(abs2, Tₙ₊₁ .- neural_pde_prediction(Tₙ))

#####
##### Choose optimization algorithm
#####

opt = ADAM(0.1)

#####
##### Callback function to observe training.
#####

function cb()
    train_loss = sum([loss_function(Tₙ[:, i], Tₙ₊₁[:, i]) for i in 1:N])

    nn_pred = neural_ode(dTdt_NN, Tₙ[:, 1], (t[1], t[N]), Tsit5(), saveat=t[1:N], reltol=1e-4) |> Flux.data
    test_loss = sum(abs2, T_cs[:, 1:N] .- nn_pred)

    println("train_loss = $train_loss, test_loss = $test_loss")
    return train_loss
end

cb()

#####
##### Train!
#####

epochs = 10
best_loss = Inf
last_improvement = 0

for epoch_idx in 1:epochs
    global best_loss, last_improvement

    @info "Epoch $epoch_idx"
    Flux.train!(loss_function, NN_params, training_data, opt, cb=cb) # cb=Flux.throttle(cb, 10))

    loss = cb()

    if loss <= best_loss
        @info("Record low loss! Saving neural network out to dTdt_NN.bson")
        BSON.@save "dTdt_NN.bson" dTdt_NN
        best_loss = loss
        last_improvement = epoch_idx
    end

    # If we haven't seen improvement in 2 epochs, drop our learning rate:
    if epoch_idx - last_improvement >= 2 && opt.eta > 1e-6
        opt.eta /= 2.0
        @warn("Haven't improved in a while, dropping learning rate to $(opt.eta)")

        # After dropping learning rate, give it a few epochs to improve
        last_improvement = epoch_idx
    end
end

#####
##### Run the neural PDE forward to see how well it performs just by itself.
#####

nn_pred = neural_ode(dTdt_NN, Tₙ[:, 1], (t[1], t[end]), Tsit5(), saveat=t, reltol=1e-4) |> Flux.data

z_cs = coarse_grain(z, cr)

anim = @animate for n=1:10:Nt
    t_str = @sprintf("%.2f", t[n] / 86400)
    plot(T_cs[2:end-1, n], z_cs, linewidth=2,
         xlim=(19, 20), ylim=(-100, 0), label="Data",
         xlabel="Temperature (C)", ylabel="Depth (z)",
         title="Deepening mixed layer: $t_str days",
         legend=:bottomright, show=false)
    plot!(nn_pred[2:end-1, n], z_cs, linewidth=2, label="Neural PDE", show=false)
end

gif(anim, "deepening_mixed_layer_neural_PDE.gif", fps=15)


text:
julia_version = v"1.6.0"
gcc_version = v"9"
include("../common.jl")

max_stars_repo_name: waralex/Yggdrasil

max_issues_repo_name: waralex/Yggdrasil

max_forks_repo_name: waralex/Yggdrasil


text:
function wvgd(z)
    ns = 0.065-4im
    n1 = 1.5835
    nc = 1.0
    d1 = 1.81e-6
    λ₀ = 0.6328e-6
    k₀ = 2π/λ₀
    k₀d1 = k₀*d1
    κ1 = sqrt(n1^2 - z^2)
    γs = sqrt(z^2 - ns^2)
    γc = sqrt(z^2 - nc^2)
    m11 = cos(κ1*k₀d1)
    m12 = im/κ1*sin(κ1*k₀d1)
    m21 = im*κ1*sin(κ1*k₀d1)
    m22 = cos(κ1*k₀d1)
    w = det([1.0    -m11+im*γc*m12
             1im*γs  -m21+im*γc*m22])
end

# Analysis parameters
xb = 1.0  # real part begin
xe = 2.5  # real part end
yb = -1.0  # imag part begin
ye = 1.0  # imag part end
r = 0.5  # initial mesh step

origcoords = rectangulardomain(complex(xb, yb), complex(xe, ye), r)

# matlab results from https://github.com/PioKow/GRPF for comparison
matlab_zroots = [1.096752543421462 - 0.000197146739811im,
                 1.240454471623525 - 0.000133821833879im,
                 1.353140429314226 - 0.000086139142513im,
                 1.439795544324443 - 0.000052001606673im,
                 1.504169866163284 - 0.000028029270470im,
                 1.548692244058475 - 0.000012100953609im,
                 1.574863045662642 - 0.000002974364907im]

matlab_zpoles = ComplexF64[]

zroots, zpoles = grpf(wvgd, origcoords)

@test length(zroots) == 7
@test length(zpoles) == 0

@test approxmatch(zroots, matlab_zroots)
@test approxmatch(zpoles, matlab_zpoles)

pzroots, pzpoles, quadrants, phasediffs, tess, g2f = grpf(wvgd, origcoords, PlotData())

@test approxmatch(pzroots, matlab_zroots)
@test approxmatch(pzpoles, matlab_zpoles)

# Test with big origcoords
xb = big"1.0"  # real part begin
xe = big"2.5"  # real part end
yb = big"-1.0"  # imag part begin
ye = big"1.0"  # imag part end
r = big"0.5"  # initial mesh step

origcoords = rectangulardomain(complex(xb, yb), complex(xe, ye), r)

bzroots, bzpoles = grpf(wvgd, origcoords)

@test all(isa.(bzroots, Complex{BigFloat}))
@test all(isa.(bzpoles, Complex{BigFloat}))

@test approxmatch(bzroots, zroots)
@test approxmatch(bzpoles, zpoles)


text:
using RDF

using Base.Test
using URIParser

@test Graph(URI("http://graphuri.test")).name == URI("http://graphuri.test")
@test_throws Graph("http://not.an.URI.but.a.string")

g = Graph(URI("http://example.org"))
push!(g, URI("http://test.org/1"), URI("http://test.org/2"), URI("http://test.org/1"))
push!(g, URI("http://test.org/1"), URI("http://test.org/3"), 3.141)
push!(g, URI("http://test.org/1"), URI("http://test.org/4"), URI("http://test.org/1"))
push!(g, URI("http://test.org/1"), URI("http://test.org/4"), 123)
push!(g, URI("http://test.org/2"), URI("http://test.org/2"), URI("http://test.org/1"))
push!(g, URI("http://test.org/2"), URI("http://test.org/2"), "Hullo!")

@test g.size == 6

s = IOBuffer()
ntriples(g, s)
@test length(split(takebuf_string(s), "\n")) == 7

s = IOBuffer()
nquads(g, s)
@test length(split(takebuf_string(s), "\n")) == 7

s = IOBuffer()
turtle(g, s)
@test length(split(takebuf_string(s), "\n")) == 7

@test pop!(g, URI("http://test.org/2")) == 2
@test g.size == 4
@test pop!(g, URI("http://test.org/1"), URI("http://test.org/4")) == 2
@test g.size == 2


text:
# Lebesgue measure
export area, arclength, measure, perimeter, volume

arclength(p::Polytope{1}) = measure(p)
area(     p::Polytope{2}) = measure(p)
volume(   p::Polytope{3}) = measure(p)

perimeter(p::Polytope{2}) = mapreduce(measure, +, facets(p))
area(     p::Polytope{3}) = mapreduce(measure, +, facets(p))

measure(l::LineSegment{<:Point}) = norm(l[2]-l[1])

function measure(q::QuadraticSegment{<:Point})
    # The arc length integral may be reduced to an integral over the square root of a
    # quadratic polynomial using ‖𝘅‖ = √(𝘅 ⋅ 𝘅), which has an analytic solution.
    #     1             1
    # L = ∫ ‖𝗾′(r)‖dr = ∫ √(ar² + br + c) dr
    #     0             0
    P₁ = q[1]
    𝘃₁₃ = q[3] - q[1]
    𝘃₁₂ = q[2] - q[1]
    𝘃₂₃ = q[3] - q[2]
    v₁₂ = norm²(𝘃₁₂)
    𝘃₁₄ = (𝘃₁₃ ⋅ 𝘃₁₂)*inv(v₁₂)*𝘃₁₂
    d = norm(𝘃₁₄ - 𝘃₁₃)
    # If segment is straight
    if d < ϵ_Point
        return √v₁₂ # Distance from P₁ to P₂
    else
        𝘂 = 3𝘃₁₃ + 𝘃₂₃
        𝘃 = -2(𝘃₁₃ + 𝘃₂₃)
        a = 4(𝘃 ⋅ 𝘃)
        b = 4(𝘂 ⋅ 𝘃)
        c = 𝘂 ⋅ 𝘂

        d = √(a + b + c)
        e = 2a + b
        f = 2√a

        l = (d*e - b*√c) / 4a - (b*b - 4a*c) / (4a*f) * log( (d*f + e) / (√c*f + b) )
        return l
    end
end

measure(tri::Triangle{<:Point}) = norm((tri[2] - tri[1]) × (tri[3] - tri[1]))/2

function measure(poly::Polygon{N, Point{2,T}}) where {N,T}
    # Uses the shoelace formula (https://en.wikipedia.org/wiki/Shoelace_formula)
    area = zero(T) # Scalar
    @inbounds @simd for i in Base.OneTo(N-1)
        area += coordinates(poly[i]) × coordinates(poly[i + 1])
    end
    @inbounds area += coordinates(poly[N]) × coordinates(poly[1])
    return norm(area)/2
end

function measure(quad::Quadrilateral{Point{3,T}}) where {T}
    # Hexahedron faces are not necessarily planar, hence we use numerical
    # integration. Gauss-Legendre quadrature over a quadrilateral is used.
    # Let F(r,s) be the interpolation function for the shape. Then,
    #     1 1                          N   N
    # A = ∫ ∫ ‖∂F/∂r × ∂F/∂s‖ ds dr =  ∑   ∑  wᵢwⱼ‖∂F/∂r(rᵢ,sⱼ) × ∂F/∂s(rᵢ,sⱼ)‖
    #     0 0                         i=1 j=1
    # Dispatch on a polynomial order such that the error in computed value when compared to
    # the reference (BigFloat, P=50), yields error approximately equal to eps(T)
    if T === Float32
        N = 8
    elseif T === Float64
        N = 18
    elseif T === BigFloat
        N = 50 # Would go higher if possible
    else
        error("Unsupported type.")
    end
    weights, points = gauss_quadrature(Val(:legendre), RefLine(), Val(N), T)
    area = zero(T)
    for j in Base.OneTo(N)
        @inbounds @simd for i in Base.OneTo(N)
            J = jacobian(quad, points[i][1], points[j][1])
            area += weights[i]*weights[j]*norm(J[:,1] × J[:,2])
        end
    end
    return area
end

function measure(poly::QuadraticPolygon{N, Point{2,T}}) where {N,T}
    # Let F(r,s) be a parameterization of surface S
    # A = ∬ dS = ∬ ‖∂F/∂r × ∂F/∂s‖dr ds
    #     S      T
    # It can be shown that the area of the quadratic polygon is the area of the base
    # linear shape + the area filled/taken away by the quadratic curves
    # outside/inside the linear shape. The area under the quadratic edge is 4/3 the
    # area of the triangle formed by the 3 vertices.
    h = zero(T)
    l = zero(T)
    M = N ÷ 2
    @inbounds @simd for i in Base.OneTo(M-1)
        h += coordinates(poly[i    ]) × coordinates(poly[i + M])
        h -= coordinates(poly[i + 1]) × coordinates(poly[i + M])
        l += coordinates(poly[i    ]) × coordinates(poly[i + 1])
    end
    @inbounds h += coordinates(poly[M]) × coordinates(poly[N])
    @inbounds h -= coordinates(poly[1]) × coordinates(poly[N])
    @inbounds l += coordinates(poly[M]) × coordinates(poly[1])
    return (4h - l)/6
end

# The area integral for 3D quadratic triangles and quadrilaterals appears to have an
# analytic solution, but it involves finding the roots of a quartic polynomial, then
# integrating over the square root of the factored quartic polynomial.
# This has a solution in the form of elliptic integrals (See Byrd and Friedman's
# Handbook of Elliptic Integrals for Engineers and Scientists, 2nd edition,
# equation 251.38), but it's absolutely massive. There may be simplifications after
# the fact that reduce the size of the expression, but for now numerical integration
# is used.
function measure(tri6::QuadraticTriangle{Point{3,T}}) where {T}
    # Gauss-Legendre quadrature over a triangle is used.
    # Let F(r,s) be the interpolation function for tri6,
    #            1 1-r                       N
    # A = ∬ dA = ∫  ∫ ‖∂F/∂r × ∂F/∂s‖ds dr = ∑ wᵢ‖∂F/∂r(rᵢ,sᵢ) × ∂F/∂s(rᵢ,sᵢ)‖
    #     S      0  0                       i=1
    # Dispatch on a polynomial order such that the error in computed value when compared to
    # the reference (BigFloat, P=20), yields error approximately equal to eps(T)
    if T === Float32
        N = 18
    elseif T === Float64
        N = 20 # Would go higher if possible
    elseif T === BigFloat
        N = 20 # Would go higher if possible
    else
        error("Unsupported type.")
    end
    weights, points = gauss_quadrature(Val(:legendre), RefTriangle(), Val(N), T)
    area = zero(T)
    @inbounds @simd for i in eachindex(weights)
        J = jacobian(tri6, points[i][1], points[i][2])
        area += weights[i] * norm(J[:,1] × J[:,2])
    end
    return area
end

function measure(quad8::QuadraticQuadrilateral{Point{3,T}}) where {T}
    # Let F(r,s) be the interpolation function for the shape. Then,
    #     1 1                          N   N
    # A = ∫ ∫ ‖∂F/∂r × ∂F/∂s‖ ds dr =  ∑   ∑  wᵢwⱼ‖∂F/∂r(rᵢ,sⱼ) × ∂F/∂s(rᵢ,sⱼ)‖
    #     0 0                         i=1 j=1
    # Dispatch on a polynomial order such that the error in computed value when compared to
    # the reference (BigFloat, P=50), yields error approximately equal to eps(T)
    if T === Float32
        N = 15
    elseif T === Float64
        N = 38
    elseif T === BigFloat
        N = 50 # Would go higher if possible
    else
        error("Unsupported type.")
    end
    weights, points = gauss_quadrature(Val(:legendre), RefLine(), Val(N),T)
    area = zero(T)
    for j in Base.OneTo(N)
        @inbounds @simd for i in Base.OneTo(N)
            J = jacobian(quad8, points[i][1], points[j][1])
            area += weights[i]*weights[j]*norm(J[:,1] × J[:,2])
        end
    end
    return area
end

max_stars_repo_name: khurrumsaleem/MOCNeutronTransport

max_issues_repo_name: khurrumsaleem/MOCNeutronTransport

max_forks_repo_name: khurrumsaleem/MOCNeutronTransport


text:
using TensorNetworkCodes
using Test

@testset "SimpleCode" begin
    stabilizers = [[1, 3, 3, 1, 0], [0, 1, 3, 3, 1], [1, 0, 1, 3, 3], [3, 1, 0, 1, 3]]
    logicals = [[1, 1, 1, 1, 1], [3, 3, 3, 3, 3]]  # XXXXX, ZZZZZ
    pure_errors = [[0, 1, 0, 0, 0], [0, 0, 0, 0, 3], [0, 0, 3, 0, 0], [1, 0, 0, 0, 0]]
    code = SimpleCode("5-qubit code", stabilizers, logicals, pure_errors)
    @test verify(code)

    code = SimpleCode("5-qubit code", stabilizers, logicals)
    @test verify(code)

    code = SimpleCode()
    @test verify(code)
end

@testset "TensorNetworkCode" begin
    # construct TensorNetworkCode from SimpleCode
    simple_code = five_qubit_code()
    tn_code = TensorNetworkCode(simple_code)
    @test tn_code.stabilizers == simple_code.stabilizers
    @test tn_code.logicals == simple_code.logicals
    @test verify(tn_code)
    @test tn_code.seed_codes[simple_code.name] == simple_code
    # construct SimpleCode from TensorNetworkCode
    new_simple_code = SimpleCode(tn_code)
    @test new_simple_code.stabilizers == simple_code.stabilizers
    @test new_simple_code.logicals == simple_code.logicals
    @test verify(new_simple_code)
end


text:
"""
Function returning the name of a where parameter
"""
function get_param_name(expr) :: Symbol
  @capture(expr, arg_<:arg_type_) && return arg
  @capture(expr, arg_) && return arg
end

"""
Function returning the arguments of a function definition
"""
function get_args(func_def::Dict)
  arg_dict = Dict{Symbol, Any}()
  arg_list = Vector{Symbol}()
  kwarg_list = Vector{Symbol}()
  for arg in (func_def[:args]...,)
    arg_def = splitarg(arg)
    push!(arg_list, arg_def[1])
    arg_dict[arg_def[1]] = arg_def[3] ? Any : arg_dict[arg_def[1]] = arg_def[2]
  end
  for arg in (func_def[:kwargs]...,)
    arg_def = splitarg(arg)
    push!(kwarg_list, arg_def[1])
    arg_dict[arg_def[1]] = arg_def[3] ? Any : arg_dict[arg_def[1]] = arg_def[2]
  end
  arg_list, kwarg_list, arg_dict
end

"""
Function returning the slots of a function definition
"""
function get_slots(func_def::Dict, args::Dict{Symbol, Any}, mod::Module) :: Dict{Symbol, Any}
  slots = Dict{Symbol, Any}()
  func_def[:name] = gensym()
  func_def[:args] = (func_def[:args]..., func_def[:kwargs]...)
  func_def[:kwargs] = []
  body = func_def[:body]
  func_def[:body] = postwalk(transform_yield, func_def[:body])
  nosaves = Set{Symbol}()
  func_def[:body] = postwalk(x->transform_nosave(x, nosaves), func_def[:body])
  func_expr = combinedef(func_def) |> flatten
  @eval(mod, @noinline $func_expr)
  codeinfos = @eval(mod, code_typed($(func_def[:name])))
  for codeinfo in codeinfos
    for (name, type) in collect(zip(codeinfo.first.slotnames, codeinfo.first.slottypes))
      name ∉ nosaves && (slots[name] = type)
    end
  end
  for (argname, argtype) in args
    slots[argname] = argtype
  end
  postwalk(x->remove_catch_exc(x, slots), func_def[:body])
  postwalk(x->make_arg_any(x, slots), body)
  for (key, val) in slots
    if val === Union{}
      slots[key] = Any
    end
  end
  delete!(slots, Symbol("#temp#"))
  delete!(slots, Symbol("_"))
  delete!(slots, Symbol(""))
  delete!(slots, Symbol("#unused#"))
  delete!(slots, Symbol("#self#"))
  slots
end

"""
Function removing the `exc` symbol of a `catch exc` statement of a list of slots.
"""
function remove_catch_exc(expr, slots::Dict{Symbol, Any})
  @capture(expr, (try body__ catch exc_; handling__ end) | (try body__ catch exc_; handling__ finally always__ end)) && delete!(slots, exc)
  expr
end

"""
Function changing the type of a slot `arg` of a `arg = @yield ret` or `arg = @yield` statement to `Any`.
"""
function make_arg_any(expr, slots::Dict{Symbol, Any})
  @capture(expr, arg_ = ex_) || return expr
  _is_yield(ex) || return expr
  slots[arg] = Any
  expr
end

"""
Function checking the use of a return statement with value
"""
function hasreturnvalue(expr)
  @capture(expr, return val_) || return expr
  (val === :nothing || val === nothing) && return expr
  @warn "@resumable function contains return statement with value!"
  expr
end


text:
using StochasticDiffEq, DiffEqCallbacks, DiffEqNoiseProcess, Test, Random

Random.seed!(100)
# Definitions according to vector/matrix representations of operators from QuantumOptics
u0 = ComplexF64[0.0, 1.0]
A = zeros(ComplexF64, 2, 2)
A[2, 1] = 1.0
A[1, 2] = 1.0
As = zeros(ComplexF64, 2, 2)
As[1, 1] = 0.5
As[2, 2] = -0.5

# SDE functions
function f_qm(du, u, p, t)
    du .= -1.0im*A*u
end
function g_qm(du, u, p, t)
    du .= -1.0im*As*u
end

# Set up two identical instances of saving
T = [0:0.1:2;]
fout(u, t, integrator) = abs2(u[1]) - abs2(u[2])

# Define the problems -- note the difference in noise
prob1 = SDEProblem{true}(f_qm, g_qm, u0, (T[1], T[end]),
        noise=StochasticDiffEq.RealWienerProcess(0.0, 0.0,
        rswm = DiffEqNoiseProcess.RSWM(adaptivealg=:RSwM1)))

prob2 = SDEProblem{true}(f_qm, g_qm, u0, (T[1], T[end]),
        noise=StochasticDiffEq.RealWienerProcess(0.0, 0.0,
        rswm = DiffEqNoiseProcess.RSWM(adaptivealg=:RSwM3)))

# Simple averaging
Ntraj = 1000
avg1 = zeros(size(T)...)
avg2 = zeros(size(T)...)
avg3 = zeros(size(T)...)

let
  global avg1,avg2,avg3
  for i=1:Ntraj

    out1 = SavedValues(Float64,ComplexF64)
    scb1 = SavingCallback(fout, out1, saveat=T, save_everystep=false, save_start=false)

    solve(prob1, RKMil{:Stratonovich}(), dt=1e-4, callback=scb1, seed = i, adaptive = false)
    avg1 .+= out1.saveval ./ Ntraj

    out1 = SavedValues(Float64,ComplexF64)
    scb1 = SavingCallback(fout, out1, saveat=T, save_everystep=false, save_start=false)

    solve(prob1, RKMil{:Stratonovich}(), tstops = T, callback=scb1,
          save_everystep=false, save_start=false)
    avg2 .+= out1.saveval ./ Ntraj

    out1 = SavedValues(Float64,ComplexF64)
    scb1 = SavingCallback(fout, out1, saveat=T, save_everystep=false, save_start=false)

    solve(prob2, RKMil{:Stratonovich}(), tstops = T, callback=scb1,
          save_everystep=false, save_start=false)
    avg3 .+= out1.saveval ./ Ntraj
  end
end

#=
using Plots; gr()
plot(T, avg1)
plot!(T, avg2)
plot!(T, avg3)
=#

@test maximum(avg1-avg2) < 0.02
@test maximum(avg1-avg3) < 0.03

#Adding Similar Tests for RKMil_General
# Simple averaging
Ntraj = 1000
avg1 = zeros(size(T)...)
avg2 = zeros(size(T)...)
avg3 = zeros(size(T)...)

let
  global avg1,avg2,avg3
  for i=1:Ntraj

    out1 = SavedValues(Float64,ComplexF64)
    scb1 = SavingCallback(fout, out1, saveat=T, save_everystep=false, save_start=false)

    solve(prob1, RKMil_General(interpretation = :Stratonovich), dt=1e-4, callback=scb1, seed = i, adaptive = false)
    avg1 .+= out1.saveval ./ Ntraj

    out1 = SavedValues(Float64,ComplexF64)
    scb1 = SavingCallback(fout, out1, saveat=T, save_everystep=false, save_start=false)

    solve(prob1, RKMil_General(interpretation = :Stratonovich), tstops = T, callback=scb1,
          save_everystep=false, save_start=false)
    avg2 .+= out1.saveval ./ Ntraj

    out1 = SavedValues(Float64,ComplexF64)
    scb1 = SavingCallback(fout, out1, saveat=T, save_everystep=false, save_start=false)

    solve(prob2, RKMil_General(interpretation = :Stratonovich), tstops = T, callback=scb1,
          save_everystep=false, save_start=false)
    avg3 .+= out1.saveval ./ Ntraj
  end
end

#=
using Plots; gr()
plot(T, avg1)
plot!(T, avg2)
plot!(T, avg3)
=#

@test maximum(avg1-avg2) < 0.02
@test maximum(avg1-avg3) < 0.03


text:
#
# Date created: 2022-03-15
# Author: aradclif
#
#
############################################################################################

# Dᵢ ∀i    : size of each dimension
function sizeblock(N::Int)
    block = Expr(:block)
    for d = 1:N
        ex = Expr(:(=), Symbol(:D_, d), Expr(:call, :size, :A, d))
        push!(block.args, ex)
    end
    block
end

# ∏ᵢ₌₁ᵏ⁻¹Dᵢ    : D₁D₂⋯Dₖ₋₁
offsetk(k::Int) = Expr(:(=), Symbol(:D_, ntuple(identity, k - 1)...),
                       Expr(:call, :*, ntuple(d -> Symbol(:D_, d), k - 1)...))
#                      Expr(:call, :*, ntuple(d -> d == 1 ? 1 : Symbol(:D_, d), k)...))

# ∑ₖ₌₁ᴺ(∏ᵢ₌₁ᵏ⁻¹Dᵢ)    : 1 + D₁ + D₁D₂ + ⋯ + D₁D₂⋯Dₖ₋₁
totaloffsetraw(N::Int) =
    Expr(:(=), :Dstar, Expr(:call, :+, ntuple(d -> d == 1 ? Expr(:call, :*, 1) :
    Expr(:call, :*, ntuple(i -> Symbol(:D_, i), d - 1)...), N)...))
totaloffset(N::Int) =
    Expr(:(=), :Dstar, Expr(:call, :+, 1, ntuple(d -> Symbol(:D_, ntuple(identity, d)...), N - 1)...))

# ∑ₖ₌₁ᴺ(∏ᵢ₌₁ᵏ⁻¹Dᵢ)Iₖ    : I₁ + D₁I₂ + D₁D₂I₃ + ⋯ + D₁D₂⋯Dₖ₋₁Iₖ
dynamictermraw(N::Int) = Expr(:call, :+, ntuple(d -> d == 1 ? :i_1 :
    Expr(:call, :*, ntuple(i -> Symbol(:D_, i), d - 1)..., Symbol(:i_, d)), N)...)
dynamicterm(N::Int) = Expr(:call, :+, ntuple(d -> d == 1 ? :i_1 :
    Expr(:call, :*, Symbol(:D_, ntuple(identity, d - 1)...), Symbol(:i_, d)), N)...)

# Jₒ ≡ index set of outer loop
# Jᵢ ≡ index set of inner loop
# ∑ₖ(∏ᵢ₌₁ᵏ⁻¹Dᵢ)Iₖ    , k ∈ J
partialtermraw(J::NTuple{N, Int}) where {N} =
    Expr(:call, :+, ntuple(d -> J[d] == 1 ? :i_1 :
    Expr(:call, :*, ntuple(i -> Symbol(:D_, i), J[d] - 1)..., Symbol(:i_, J[d])), N)...)

partialterm(J::NTuple{N, Int}) where {N} =
    Expr(:call, :+, ntuple(d -> J[d] == 1 ? :i_1 :
    Expr(:call, :*, Symbol(:D_, ntuple(identity, J[d] - 1)...), Symbol(:i_, J[d])), N)...)

partialtermrawj(J::NTuple{N, Int}) where {N} =
    Expr(:call, :+, ntuple(d -> J[d] == 1 ? :j_1 :
    Expr(:call, :*, ntuple(i -> Symbol(:D_, i), J[d] - 1)..., Symbol(:j_, J[d])), N)...)

singletermraw(k::Int) = k == 1 ? :i_1 :
    Expr(:call, :*, ntuple(i -> Symbol(:D_, i), k - 1)..., Symbol(:i_, k))
singletermrawj(k::Int) = k == 1 ? :j_1 :
    Expr(:call, :*, ntuple(i -> Symbol(:D_, i), k - 1)..., Symbol(:j_, k))
singleterm(k::Int) = k == 1 ? :i_1 :
    Expr(:call, :*, Symbol(:D_, ntuple(identity, k - 1)...), Symbol(:i_, k))

function preexprmax(J::NTuple{N, Int}) where {N}
    block = Expr(:block, Expr(:(=), :m, Expr(:call, :typemin, :T)))
    ex = partialterm(J)
    push!(ex.args, 1, :Dstar)
    push!(block.args, Expr(:(=), :j, ex))
    block
end
# function preexprmax2(J::NTuple{N, Int}) where {N}
#     Expr(:block, Expr(:(=), :m, Expr(:call, :typemin, :T)), Expr(:(=), :j, 0))
# end
function preexpr2(init::Symbol)
    Expr(:block, Expr(:(=), :m, Expr(:call, init, :T)), Expr(:(=), :j, 0))
end

function findmax_quote(N::Int, D)
    b1 = sizeblock(N)
    b2 = Expr(:block, ntuple(d -> offsetk(d + 2), N - 2)...)
    b3 = Expr(:block, totaloffsetraw(N), Expr(:(=), :Dstar, Expr(:call, :-, :Dstar)))
    b4 = preexprmax(tuple((d for d = 1:N if D.parameters[d] !== Val{1})...))
    inner = innerloop(N, D)
    push!(inner.args, maxblock(N, D))
    push!(b4.args, inner)
    push!(b4.args, postexpr(N, D).args...)
    outer = outerloop(N, D)
    push!(outer.args, b4)
    return quote
        $b1
        $b2
        $b3
        $outer
    end
end
function findmax_quote2(N::Int, D)
    b1 = sizeblock(N)
    b2 = Expr(:block, ntuple(d -> offsetk(d + 2), N - 2)...)
    b3 = Expr(:block, totaloffsetraw(N), Expr(:(=), :Dstar, Expr(:call, :-, :Dstar)))
    b4 = preexpr2(:typemin)
    inner = innerloop(N, D)
    push!(inner.args, maxblock2(N, D))
    push!(b4.args, inner)
    push!(b4.args, postexpr(N, D).args...)
    outer = outerloop(N, D)
    push!(outer.args, b4)
    return quote
        $b1
        $b2
        $b3
        @turbo $outer
    end
end
function findmax_quote3(N::Int, D)
    b1 = sizeblock(N)
    b2 = Expr(:block, ntuple(d -> offsetk(d + 2), N - 2)...)
    b3 = Expr(:block, totaloffsetraw(N), Expr(:(=), :Dstar, Expr(:call, :-, :Dstar)))
    b4 = preexpr2(:typemin)
    inner = innerloop(N, D)
    push!(inner.args, maxblock3(N, D))
    push!(b4.args, inner)
    push!(b4.args, postexpr3(N, D).args...)
    outer = outerloop(N, D)
    push!(outer.args, b4)
    return quote
        $b1
        $b2
        $b3
        @turbo $outer
    end
end

D_1 = 4
D_2 = 3
D_3 = 2
N = 3
dd = typeof((D_1, Val(1), D_3))
findmax_quote(N, dd)
findmax_quote2(N, dd)
findmax_quote3(N, dd)
maxblock(N, dd)
maxblock2(N, dd)

A = reshape([1:24;], D_1, D_2, D_3);
dims = (1,)
Dᴮ = ntuple(d -> d ∈ dims ? 1 : size(A, d), N)
Dᴮ′ = ntuple(d -> d ∈ dims ? Val(1) : size(A, d), N)
T = eltype(A)
B = similar(A, Dᴮ);
C = similar(A, Int, Dᴮ);
ex = findmax_quote(N, typeof(Dᴮ′))
ex2 = findmax_quote2(N, typeof(Dᴮ′))
ex3 = findmax_quote3(N, typeof(Dᴮ′))
eval(ex2)
eval(ex3)
CartesianIndices(A)[C] == argmax(A, dims=dims)

@generated _vfindmax2!(B::AbstractArray{T, N}, C::AbstractArray{Tₒ, N},
                       A::AbstractArray{T, N}, dims::D) where {T, Tₒ, N, D} = findmax_quote2(N, D)
function vfindmax2(A::AbstractArray{T, N}, dims::NTuple{M, Int}) where {T, N, M}
    Dᴮ = ntuple(d -> d ∈ dims ? 1 : size(A, d), N)
    Dᴮ′ = ntuple(d -> d ∈ dims ? Val(1) : size(A, d), N)
    B = similar(A, Dᴮ)
    C = similar(A, Int, Dᴮ)
    _vfindmax2!(B, C, A, Dᴮ′)
    B, CartesianIndices(A)[C]
end

@generated _vfindmax3!(B::AbstractArray{T, N}, C::AbstractArray{Tₒ, N},
                       A::AbstractArray{T, N}, dims::D) where {T, Tₒ, N, D} = findmax_quote3(N, D)
function vfindmax3(A::AbstractArray{T, N}, dims::NTuple{M, Int}) where {T, N, M}
    Dᴮ = ntuple(d -> d ∈ dims ? 1 : size(A, d), N)
    Dᴮ′ = ntuple(d -> d ∈ dims ? Val(1) : size(A, d), N)
    B = similar(A, Dᴮ)
    C = similar(A, Int, Dᴮ)
    _vfindmax3!(B, C, A, Dᴮ′)
    B, CartesianIndices(A)[C]
end

dims = (4,)
@benchmark vfindmax2(A, dims)
@benchmark vfindmax3(A, dims)
@benchmark findmax(A, dims=dims)
n = 4
A = rand(50, 50, 50, 50);
A = rand(50, 50, 10000);
A = rand(10,10,10,10,10,10,10,10);

vfindmax2(A, dims) == vfindmax3(A, dims) == findmax(A, dims=dims)
vargmax3(A, dims) == argmax(A, dims=dims)
vfindmax3(A, dims)[2] == vargmax3(A, dims)

@benchmark vfindmax2(A, dims)
@benchmark vfindmax3(A, dims)

ns = [1, 5, 10, 50]
bs = Matrix{NTuple{3, BenchmarkTools.Trial}}(undef, length(ns), 4);
for d = 2:4
    for (i, n) ∈ enumerate(ns)
        A = rand(n, 2n, 3n, 4n)
        bs[i, d] = ((@benchmark findmax(A, dims=($d,))),
                    (@benchmark vfindmax2(A, ($d,))),
                    (@benchmark vfindmax3(A, ($d,))))
    end
end

ts = map(bs) do b
    mean.(b)
end
tso = map(ts) do t
    getproperty.(t, :time)
end

outer = :(for i_3 = axes(A, 3), i_1 = axes(A, 1)
              m = typemin(T)
              j = 0
              for i_2 = axes(A, 2)
                  y = A[i_1, i_2, i_3] > m
                  m = ifelse(y, A[i_1, i_2, i_3], m)
                  j = ifelse(y, +(D_1 * i_2), j)
              end
              B[i_1, 1, i_3] = m
              C[i_1, 1, i_3] = i_1 + D_12 * i_3 + 1 + Dstar + j
          end);
ls = LoopVectorization.LoopSet(outer);
ops = LoopVectorization.operations(ls)
LoopVectorization.loopdependencies.(ops)
@macroexpand @turbo for i_3 = axes(A, 3), i_1 = axes(A, 1)
    m = typemin(T)
    j = 0
    for i_2 = axes(A, 2)
        y = A[i_1, i_2, i_3] > m
        m = ifelse(y, A[i_1, i_2, i_3], m)
        j = ifelse(y, +(D_1 * i_2), j)
    end
    B[i_1, 1, i_3] = m
    C[i_1, 1, i_3] = i_1 + D_12 * i_3 + 1 + Dstar + j
end

outer = :(for i_3 = axes(A, 3), i_2 = axes(A, 2)
              m = typemin(T)
              j = 0
              for i_1 = axes(A, 1)
                  y = A[i_1, i_2, i_3] > m
                  m = ifelse(y, A[i_1, i_2, i_3], m)
                  j = ifelse(y, i_1, j)
              end
              B[1, i_2, i_3] = m
              C[1, i_2, i_3] = j#D_1 * i_2 + D_12 * i_3 + 1 + Dstar + j
          end)
ls = LoopVectorization.LoopSet(outer);
ops = LoopVectorization.operations(ls)
LoopVectorization.loopdependencies.(ops)
ls = LoopVectorization.@turbo_debug for i_3 = axes(A, 3), i_2 = axes(A, 2)
    m = typemin(T)
    j = 0
    for i_1 = axes(A, 1)
        # j = C[1, i_2, i_3]
        # m = B[1, i_2, i_3]
        y = A[i_1, i_2, i_3] > m
        m = ifelse(y, A[i_1, i_2, i_3], m)
        j = ifelse(y, i_1, j)
        # B[1, i_2, i_3] = ifelse(y, A[i_1, i_2, i_3], m)
        # C[1, i_2, i_3] = ifelse(y, i_1, j)
    end
    B[1, i_2, i_3] = m
    C[1, i_2, i_3] = j#D_1 * i_2 + D_12 * i_3 + 1 + Dstar + j
end

############################################################################################
#### 2022-03-18: Experiments applying methodology to reduce, mapreduce
# Alas, performance is worse as the loop ordering is kind of absurd...
function reduce_quote2(OP, I, N::Int, D)
    init = I.instance
    pre = Expr(:block, Expr(:(=), :s, Expr(:call, Symbol(init), :T)))
    inner = innerloop(N, D)
    rₑ = reduceexpr(OP, N)
    push!(inner.args, Expr(:block, rₑ))
    push!(pre.args, inner)
    push!(pre.args, postexpr(:B, :s, N, D))
    outer = outerloop(N, D)
    push!(outer.args, pre)
    # outer
    return quote
        @turbo $outer
    end
end

function mapreduce_quote2(F, OP, I, N::Int, D)
    init = I.instance
    pre = Expr(:block, Expr(:(=), :s, Expr(:call, Symbol(init), :T)))
    inner = innerloop(N, D)
    rₑ = mapreduceexpr(F, OP, N)
    push!(inner.args, rₑ)
    push!(pre.args, inner)
    push!(pre.args, postexpr(:B, :s, N, D))
    outer = outerloop(N, D)
    push!(outer.args, pre)
    outer
end

function reduce_quote2(OP, N::Int, D)
    pre = Expr(:block, Expr(:(=), :s, reduceref(:B, N, D)))
    inner = innerloop(N, D)
    rₑ = reduceexpr(OP, N)
    push!(inner.args, rₑ)
    push!(pre.args, inner)
    push!(pre.args, postexpr(:B, :s, N, D))
    outer = outerloop(N, D)
    push!(outer.args, pre)
    outer
    return quote
        @turbo $outer
    end
end

function _lvreduce2(f, A::AbstractArray{T, N}, dims::NTuple{M, Int}) where {T, N, M}
    if ntuple(identity, Val(N)) ⊆ dims
        B = hvncat(ntuple(_ -> 1, Val(N)), true, lvreduce21(f, A))
    else
        Dᴮ = ntuple(d -> d ∈ dims ? 1 : size(A, d), Val(N))
        Dᴮ′ = ntuple(d -> d ∈ dims ? Val(1) : size(A, d), Val(N))
        B = zeros(Base.promote_op(f, T), Dᴮ)
        _lvreduce2!(f, zero, B, A, Dᴮ′)
        # _lvreduce2!(f, B, A, Dᴮ′)
    end
    return B
end

@generated function _lvreduce2!(f::F, init::I, B::AbstractArray{Tₒ, N}, A::AbstractArray{T, N}, dims::D) where {F, I, Tₒ, T, N, D}
    reduce_quote2(F, I, N, D)
end
@generated function _lvreduce2!(f::F, B::AbstractArray{Tₒ, N}, A::AbstractArray{T, N}, dims::D) where {F, Tₒ, T, N, D}
    reduce_quote2(F, I, N, D)
end

@benchmark _lvreduce2(+, A, (1,2))
@benchmark lvsum(A, dims=(1,2,))
@benchmark vsum(A, dims=(1,2))
_lvreduce2(+, A, (1,2)) == lvsum(A, dims=(1,2)) == vsum(A, dims=(1,2))

############################################################################################
#### Interesting experiment using indices
ex = VectorizedStatistics.staticdim_mean_quote([1,3], 4)
Meta.show_sexpr(ex)
D = typeof((StaticInt(1), 10, StaticInt(1), 10))
N = 4
function reduce_quote3(OP, I, N::Int, D)
    _params = D.parameters
    a = Expr(:ref, :A, ntuple(d -> Symbol(:i_, d), N)...)
    bᵥ = Expr(:call, :view, :B)
    bᵥ′ = Expr(:ref, :Bᵥ)
    rinds = Int[]
    nrinds = Int[]
    for d = 1:N
        if _params[d] === Static.One
            push!(bᵥ.args, Expr(:call, :firstindex, :B, d))
            push!(rinds, d)
        else
            # push!(bᵥ.args, Expr(:call, :Colon))
            push!(bᵥ.args, :)
            push!(nrinds, d)
            push!(bᵥ′.args, Symbol(:i_, d))
        end
    end
    bᵥ = Expr(:(=), :Bᵥ, bᵥ)
    sort!(rinds, rev=true)
    sort!(nrinds, rev=true)
    block = Expr(:block)
    loops = Expr(:for, Expr(:(=), Symbol(:i_, nrinds[1]),
                            Expr(:call, :indices, Expr(:tuple, :A, :B), nrinds[1])), block)
    for i = 2:length(nrinds)
        newblock = Expr(:block)
        push!(block.args,
              Expr(:for, Expr(:(=), Symbol(:i_, nrinds[i]),
                              Expr(:call, :indices, Expr(:tuple, :A, :B), nrinds[i])), newblock))
        block = newblock
    end
    rblock = block
    # Push to before reduction loop
    pre = Expr(:(=), :ξ, Expr(:call, Symbol(I.instance), Expr(:call, :eltype, :Bᵥ)))
    push!(rblock.args, pre)
    # Reduction loop
    for i = 1:length(rinds)
        newblock = Expr(:block)
        push!(block.args,
              Expr(:for, Expr(:(=), Symbol(:i_, rinds[i]),
                              Expr(:call, :axes, :A, rinds[i])), newblock))
        block = newblock
    end
    # Push to inside innermost loop
    reduction = Expr(:(=), :ξ, Expr(:call, Symbol(OP.instance), :ξ, a))
    push!(block.args, reduction)
    # Push to after reduction loop
    post = Expr(:(=), bᵥ′, :ξ)
    push!(rblock.args, post)
    return quote
        $bᵥ
        @tturbo $loops
        return B
    end
end

reduce_quote3(typeof(max), typeof(typemin), N, D)
Dᴮ′ = ntuple(d -> d ∈ dims ? StaticInt(1) : size(A, d), N)
D = typeof(Dᴮ′)
reduce_quote3(typeof(+), typeof(zero), N, D)

function _lvreduce3(op::OP, init::I, A::AbstractArray{T, N}, dims::NTuple{M, Int}) where {OP, I, T, N, M}
    if ntuple(identity, Val(N)) ⊆ dims
        B = hvncat(ntuple(_ -> 1, Val(N)), true, lvreduce1(op, A))
    else
        Dᴮ′ = ntuple(d -> d ∈ dims ? StaticInt(1) : size(A, d), Val(N))
        B = similar(A, Base.promote_op(op, T), Dᴮ′)
        _lvreduce3!(op, init, B, A, Dᴮ′)
    end
    return B
end

@generated function _lvreduce3!(op::OP, init::I, B::AbstractArray{Tₒ, N}, A::AbstractArray{T, N}, dims::D) where {OP, I, Tₒ, T, N, D}
    reduce_quote3(OP, I, N, D)
end

lvsum3(A, dims) = _lvreduce3(+, zero, A, dims)
@benchmark _lvreduce3(+, zero, A, (1,3))
@benchmark lvsum3(A, (1,3))
@benchmark lvsum(A, dims=(1,3), multithreaded=false)
@benchmark vsum(A, dims=(1,3), multithreaded=false)
_lvreduce3(+, zero, A, (1,3)) == lvsum(A, dims=(1,3))

function _lvsum4(A::AbstractArray{T, N}, dims::NTuple{M, Int}) where {T, N, M}
    Dᴮ′ = ntuple(d -> d ∈ dims ? StaticInt(1) : size(A, d), Val(N))
    B = similar(A, Base.promote_op(+, T), Dᴮ′)
    _lvsum4!(B, A, Dᴮ′)
end
function sum_quote(N::Int, D)
    params = D.parameters
    rinds = Int[]
    for d = 1:N
        params[d] === Static.One && push!(rinds, d)
    end
    return VectorizedStatistics.staticdim_sum_quote(rinds, N)
end
@generated _lvsum4!(B, A::AbstractArray{T, N}, dims::D) where {T, N, D} = sum_quote(N, D)

@benchmark _lvsum4(A, (1,3))

function mapreduce_quote3(F, OP, I, N::Int, D)
    _params = D.parameters
    a = Expr(:ref, :A, ntuple(d -> Symbol(:i_, d), N)...)
    bᵥ = Expr(:call, :view, :B)
    bᵥ′ = Expr(:ref, :Bᵥ)
    rinds = Int[]
    nrinds = Int[]
    for d = 1:N
        if _params[d] === Static.One
            push!(bᵥ.args, Expr(:call, :firstindex, :B, d))
            push!(rinds, d)
        else
            # push!(bᵥ.args, Expr(:call, :Colon))
            push!(bᵥ.args, :)
            push!(nrinds, d)
            push!(bᵥ′.args, Symbol(:i_, d))
        end
    end
    bᵥ = Expr(:(=), :Bᵥ, bᵥ)
    sort!(rinds, rev=true)
    sort!(nrinds, rev=true)
    block = Expr(:block)
    loops = Expr(:for, Expr(:(=), Symbol(:i_, nrinds[1]),
                            Expr(:call, :indices, Expr(:tuple, :A, :B), nrinds[1])), block)
    for i = 2:length(nrinds)
        newblock = Expr(:block)
        push!(block.args,
              Expr(:for, Expr(:(=), Symbol(:i_, nrinds[i]),
                              Expr(:call, :indices, Expr(:tuple, :A, :B), nrinds[i])), newblock))
        block = newblock
    end
    rblock = block
    # Push to before reduction loop
    pre = Expr(:(=), :ξ, Expr(:call, Symbol(I.instance), Expr(:call, :eltype, :Bᵥ)))
    push!(rblock.args, pre)
    # Reduction loop
    for i = 1:length(rinds)
        newblock = Expr(:block)
        push!(block.args,
              Expr(:for, Expr(:(=), Symbol(:i_, rinds[i]),
                              Expr(:call, :axes, :A, rinds[i])), newblock))
        block = newblock
    end
    # Push to inside innermost loop
    reduction = Expr(:(=), :ξ, Expr(:call, Symbol(OP.instance), :ξ,
                                    Expr(:call, Symbol(F.instance), a)))
    push!(block.args, reduction)
    # Push to after reduction loop
    post = Expr(:(=), bᵥ′, :ξ)
    push!(rblock.args, post)
    return quote
        $bᵥ
        @tturbo $loops
        return B
    end
end

function _lvmapreduce3(f::F, op::OP, init::I, A::AbstractArray{T, N}, dims::NTuple{M, Int}) where {F, OP, I, T, N, M}
    if ntuple(identity, Val(N)) ⊆ dims
        B = hvncat(ntuple(_ -> 1, Val(N)), true, lvmapreduce1(f, op, A))
    else
        Dᴮ′ = ntuple(d -> d ∈ dims ? StaticInt(1) : size(A, d), Val(N))
        B = similar(A, Base.promote_op(op, Base.promote_op(f, T)), Dᴮ′)
        _lvmapreduce3!(f, op, init, B, A, Dᴮ′)
    end
    return B
end

@generated function _lvmapreduce3!(f::F, op::OP, init::I, B::AbstractArray{Tₒ, N}, A::AbstractArray{T, N}, dims::D) where {F, OP, I, Tₒ, T, N, D}
    mapreduce_quote3(F, OP, I, N, D)
end