include("Graph.jl")
include("SimMetric.jl")

using SparseArrays
using LinearAlgebra
import Laplacians

function getSimSelectorByS(S)
    n = size(S, 1)
    f(u, k) = begin
        ret = Vector{Int32}()
        A = Vector{Tuple{Int32, Float64}}()
        for i = 1 : n
            if i != u
                push!(A, (i, S[u, i]))
            end
        end
        sort!(A, by=x->-x[2])
        for i = 1 : k
            push!(ret, A[i][1])
        end
        return ret
    end
    return f
end

function RoleSimSelector(G)
    S = RoleSim(G)
    return getSimSelectorByS(S)
end

function ForestSimSelector(G)
    S = ForestSim(G)
    return getSimSelectorByS(S)
end

function ForestSimApproxSelector(G)
    S = approxForest(G)
    A = Vector{Tuple{Int32, Float64}}()
    foreach(i -> push!(A, (i, S[i])), 1 : G.n)
    sort!(A, by=x->x[2])
    pos = zeros(Int32, G.n)
    foreach(i -> pos[A[i][1]] = i, 1 : G.n)

    sim(x, y) = min(S[x], S[y]) / max(S[x], S[y])

    f(x, k) = begin
        l = pos[x]-1
        r = pos[x]+1
        ret = Vector{Int32}()
        for rep = 1 : k
            if l == 0
                push!(ret, A[r][1])
                r += 1
            elseif r > G.n
                push!(ret, A[l][1])
                l -= 1
            elseif sim(x, A[l][1]) > sim(x, A[r][1])
                push!(ret, A[l][1])
                l -= 1
            else
                push!(ret, A[r][1])
                r += 1
            end
        end
        return ret
    end

    return f
end

function StructSimSelector(G; maxk = 3, wi = 0.5)
    n = G.n
    d = degrees(G)
    B, dk = calcBins(G, maxk)
    maxb = size(B, 3)

    f(x, k) = begin
        S = zeros(n)
        foreach(i -> S[i] = min(d[i], d[x]) / max(d[i], d[x]), 1 : n)
        for k = 1 : maxk
            for i = 1 : n
                delta = 0.0
                foreach(z -> delta += min(B[k, i, z], B[k, x, z]), 1 : maxb)
                delta /= max(dk[k, i], dk[k, x])
                S[i] = (1.0 - wi) * S[i] + wi * delta
            end
        end
        A = Vector{Tuple{Int32, Float64}}()
        for i = 1 : n
            if i != x
                push!(A, (i, S[i]))
            end
        end
        sort!(A, by=x->-x[2])
        ret = Vector{Int32}()
        for i = 1 : k
            push!(ret, A[i][1])
        end
        return ret
    end

    return f
end


include("parallelNNConv.jl")
include("parallelConvolve.jl")
include("parallelImg2colConvolve.jl")

###convolution layers forprop

@doc raw"""
    function layerForProp(
        cLayer::ConvLayer,
        Ai::AbstractArray = Array{Any,1}(undef,0);
        FCache::Dict{Layer,Dict{Symbol, AbstractArray}},
        kwargs...
    )

Perform the layer forward propagation for a `ConvLayer`

# Arguments

- `cLayer::ConvLayer`

- `Ai` := optional activation of the previous layer

- `FCache` := a `Dict` holds the outputs of `layerForProp` of the previous `Layer`(s)

# Returns

- `Dict(:Z => Z, :A => Ao)`
"""
function layerForProp(
    cLayer::CL,
    Ai::AbstractArray = Array{Any,1}(undef,0);
    FCache::Dict{Layer,Dict{Symbol, AbstractArray}},
    kwargs...
) where {CL <: ConvLayer}

    kwargs = Dict{Symbol, Any}(kwargs...)
    NNlib = getindex(kwargs, :NNlib; default=true)
    img2col = getindex(kwargs, :img2col; default=false)
    if length(Ai) == 0
        Ai = FCache[cLayer.prevLayer][:A]
    end
    cLayer.inputS = size(Ai)
    paddedS = paddedSize(cLayer, Ai)[1:end-2]
    s = cLayer.s
    f = cLayer.f
    outputS = outDims(cLayer, Ai)[1:end-2]
    ci, m = cLayer.inputS[end-1:end]
    co = cLayer.channels

    if NNlib
        D = NNConv(cLayer, Ai)
    elseif img2col
        ## in case input different size than previous time
        if (prod((outputS..., co)), prod((paddedS..., ci))) != size(cLayer.K)
            cLayer.K = unroll(cLayer, (paddedS..., ci, m))
        end
        D = img2colConvolve(cLayer, Ai)
    else
        D = convolve(cLayer, Ai)
    end

    cLayer.outputS = size(D[:A])

    # Ai = nothing
    cLayer.forwCount += 1
    # Base.GC.gc()
    return D

end #function layerForProp(cLayer::Conv1D)

### Pooling Layers

include("parallelFastPooling.jl")
include("parallelPooling.jl")

#import only the needed parts not to have conflict
import NNlib.maxpool, NNlib.meanpool, NNlib.maxpool!, NNlib.meanpool!, NNlib.PoolDims

@doc raw"""
    layerForProp(
        cLayer::PoolLayer},
        Ai::AbstractArray = Array{Any,1}(undef,0);
        FCache::Dict{Layer,Dict{Symbol, AbstractArray}},
        kwargs...
    )

Perform the layer forward propagation for a `PoolLayer`

# Arguments

- `cLayer::PoolLayer`

- `Ai` := optional activation of the previous layer

- `FCache` := a `Dict` holds the outputs of `layerForProp` of the previous `Layer`(s)

# Returns

- `Dict(:A => Ao)`
"""
function layerForProp(
    cLayer::PL,
    Ai::AbstractArray = Array{Any,1}(undef,0);
    FCache::Dict{Layer,Dict{Symbol, AbstractArray}},
    kwargs...
) where {PL <: PoolLayer}

    kwargs = Dict{Symbol, Any}(kwargs...)
    fastPool = getindex(kwargs, :fastPool; default=true)
    NNlib = getindex(kwargs, :NNlib; default=true)
    if length(Ai) == 0
        Ai = FCache[cLayer.prevLayer][:A]
    end

    cLayer.inputS = size(Ai)
    paddedS = paddedSize(cLayer, Ai)[1:end-2]
    s = cLayer.s
    f = cLayer.f
    outputS = outDims(cLayer, Ai)[1:end-2]
    ci, m = cLayer.inputS[end-1:end]
    co = cLayer.channels
    cLayer.inputS = size(Ai)
    padS = paddingSize(cLayer, Ai)

    Ao = zeros(eltype(Ai), outputS..., co, m)

    if NNlib
        pooldims = PoolDims(Ai, f, stride = s, padding = padS)
        if cLayer isa MaxPoolLayer
            maxpool!(Ao, Ai, pooldims)
        elseif cLayer isa AveragePoolLayer
            meanpool!(Ao, Ai, pooldims)
        end #if cLayer isa MaxPoolLayer
    elseif s == f && fastPool #to use the built-in reshape and maximum and mean
        fastPooling!(cLayer, Ai, Ao)
    else
        pooling!(cLayer, Ai, Ao)
    end #if NNlibConv

    cLayer.outputS = size(Ao)
    # Ai = nothing
    cLayer.forwCount += 1
    # Base.GC.gc()
    return Dict(:A => Ao)
end #unction layerForProp(cLayer::OneD) where {OneD <: Union{MaxPool1D, AveragePool1D}}


"""
    LaplaceOutput
Return type for Laplace computations.
Fields are `nothing` on return unless requested.
See individual FMM/direct computation function
documentation for specifics.
"""
mutable struct LaplaceOutput <: FMMVals
    pot
    grad
    hess

    pottarg
    gradtarg
    hesstarg

    pgt
    pg
end
function LaplaceOutput(pgt,pg=nothing)
    LaplaceOutput(nothing,nothing,nothing,
                  nothing,nothing,nothing,
                  pgt,pg)
end

"""
```julia
    vals = lfmm3d(eps,sources;charges=nothing,dipvecs=nothing,
                  targets=nothing,pg=0,pgt=0,nd=1)
```
This function computes the N-body Laplace interactions
in three dimensions where the interaction kernel is given
by ``\$1/r\$`` and its gradients.  This is the
``O(N)`` fast multipole code which computes the interactions
to the requested precision.
```math
u(x) = \\sum_{j=1}^{N} c_{j} / \\|x-x_{j}\\| + v_{j}
\\cdot \\nabla( 1/\\|x-x_{j}\\|)  \\, ,
u(x) = \\sum_{j=1}^{N} c_{j} \\frac{1}{\\|x-x_{j}\\|}
- v_{j} \\cdot \\nabla \\left( \\frac{1}{\\|x-x_{j}\\|}
\\right)  \\, ,
```

where ``c_{j}`` are the charge densities,
      ``v_{j}`` are the dipole orientation vectors, and
      ``x_{j}`` are the source locations.
      When ``x=x_{m}``, the term corresponding to
      ``x_{m}`` is dropped from the sum
# Input
# Input
* `eps::Float64` precision requested
* `sources::Array{Float64}` size (3,n) source locations (``x_{j}``)
* `charges::Array{Float64}` size (nd,n) or (n) charge densities (c_{j})
* `dipvecs::Array{Float64}` size (nd,3,n) or (3,n)) dipole orientation vectors (``v_{j}``)
* `targets::Array{Float64}` size (3,nt) target locations (``x``)
* `pg::Integer` source eval flag.
    + Potential (``u``) at sources evaluated if `pg == 1`.
    + Potential and gradient (``\\nabla u``) at sources evaluated if `pg == 2`
* `pgt::Integer` target eval flag.
    + Potential at targets evaluated if `pgt == 1`.
    + Potential and gradient at targets evaluated if `pgt == 2`
* `nd::Integer` number of densities
# Output

`vals::LaplaceOutput` with the fields
Note that the Hessian is returned as a length
6 vector at each point with the second derivatives
in the order: ``\\partial_{xx}``, ``\\partial_{yy}``,
``\\partial_{zz}``, ``\\partial_{xy}``,
``\\partial_{xz}``, ``\\partial_{yz}``.
* `vals.pot::Array{Float64}` size (nd,n) potential at source locations if requested
* `vals.grad::Array{Float64}` size (nd,3,n) gradient at source locations if requested
* `vals.hess::Array{Float64}` size (nd,6,n) Hessian at source locations if requested
* `vals.pottarg::Array{Float64}` size (nd,nt) potential at target locations if requested
* `vals.gradtarg::Array{Float64}` size (nd,3,nt) gradient at target locations if requested
* `vals.hesstarg::Array{Float64}` size (nd,6,nt) Hessian at target locations if requested
"""
function lfmm3d(eps::Float64,sources::Array{Float64};charges::TFN=nothing,
                dipvecs::TFN=nothing,targets::TFN=nothing,
                pg::Integer=0,pgt::Integer=0,nd::Integer=1)

    # check inputs

    @assert size(sources,1) == 3
    @assert nd >= 0
    @assert (0 <= pg  && pg  <= 3) "flag pg not in expected range"
    @assert (0 <= pgt && pgt <= 3) "flag pgt not in expected range"

    # default values

    vals = LaplaceOutput(pgt,pg)

    ifcharge = 0
    ifdipole = 0

    zero = Float64(0)

    pot = zero
    grad = zero
    hess = zero
    pottarg = zero
    gradtarg = zero
    hesstarg = zero

    n = div(length(sources),3)

    nt = 0

    if targets !== nothing
        @assert size(targets,1) == 3
        nt = div(length(targets),3)
    else
        targets = 0.0
    end

    if charges !== nothing
        @assert div(length(charges),nd) == n
        ifcharge = 1
    else
        charges = zero
        ifcharge = 0
    end

    if dipvecs !== nothing
        @assert div(length(dipvecs),nd) == n*3
        ifdipole = 1
    else
        dipvecs = zero
        ifdipole = 0
    end

    if ifcharge == 0 && ifdipole == 0
        @warn "no charges or dipoles provided, doing nothing"
        return vals
    end

    if ((pg != 1 && pg != 2 && pg !=3) &&
        (pgt != 1 && pgt != 2 && pgt != 3))
        @warn "no output requested, doing nothing"
        return vals
    end

    if (pgt == 1 || pgt == 2 || pgt == 3) && targets === nothing
        @warn "target values requested but no targets provided"
    end

    # allocate memory for return values

    if pg == 1 || pg == 2 || pg == 3
        if nd > 1
            pot = zeros(Float64,nd,n)
        else
            pot = zeros(Float64,n)
        end
    end

    if pg == 2 || pg == 3
        if nd > 1
            grad = zeros(Float64,nd,3,n)
        else
            grad = zeros(Float64,3,n)
        end
    end

    if pg == 3
        if nd > 1
            hess = zeros(Float64,nd,6,n)
        else
            hess = zeros(Float64,6,n)
        end
    end

    if pgt == 1 || pgt == 2 || pgt == 3
        if nd > 1
            pottarg = zeros(Float64,nd,nt)
        else
            pottarg = zeros(Float64,nt)
        end
    end

    if pgt == 2 || pgt == 3
        if nd > 1
            gradtarg = zeros(Float64,nd,3,nt)
        else
            gradtarg = zeros(Float64,3,nt)
        end
    end

    if pgt == 3
        if nd > 1
            hesstarg = zeros(Float64,nd,6,nt)
        else
            hesstarg = zeros(Float64,6,nt)
        end
    end

    # actually call the function
    #
    # fortran calling sequence:
    # subroutine lfmm3d(nd,eps,nsource,source,ifcharge,
    #     $    charge,ifdipole,dipvec,ifpgh,pot,grad,hess,ntarg,
    #     $    targ,ifpghtarg,pottarg,gradtarg,hesstarg)
    #

    ccall((:lfmm3d_,libfmm3d),Cvoid,(Fi,Fd,Fi,Fd,Fi,Fd,
                                    Fi,Fd,Fi,Fd,Fd,Fd,Fi,
                                    Fd,Fi,Fd,Fd,Fd),
          nd,eps,n,sources,ifcharge,charges,ifdipole,
          dipvecs,pg,pot,grad,hess,nt,targets,pgt,
          pottarg,gradtarg,hesstarg)

    # load requested values

    if pg == 1 || pg == 2 || pg == 3; vals.pot = pot end
    if pg == 2 || pg == 3; vals.grad = grad end
    if pg == 3; vals.hess = hess end
    if pgt == 1 || pgt == 2 || pgt == 3; vals.pottarg = pottarg end
    if pgt == 2 || pgt == 3; vals.gradtarg = gradtarg end
    if pgt == 3; vals.hesstarg = hesstarg end

    return vals

end

"""
```julia
    vals = l3ddir(sources,targets;charges=nothing,
                dipvecs=nothing,pgt=0,nd=1,
                thresh=1e-16)
```
This function computes the N-body Laplace interactions
in three dimensions where the interaction kernel is given
by ``1/r`` and its gradients.  This is the
``O(N^2)`` direct evaluation code.
```math
 u(x) = \\sum_{j=1}^{N} c_{j} \\frac{1}{\\|x-x_{j}\\|}
- v_{j} \\cdot \\nabla \\left( \\frac{1}{\\|x-x_{j}\\|}
\\right)  \\, ,
```

where ``c_{j}`` are the charge densities,
      ``v_{j}`` are the dipole orientation vectors, and
      ``x_{j}`` are the source locations.
      When ``x=x_{m}``, the term corresponding to
      ``x_{m}`` is dropped from the sum
# Input
* `sources::Array{Float64}` size (3,n) source locations (``x_{j}``)
* `targets::Array{Float64}` size (3,nt) target locations (``x``)
* `charges::Array{ComplexF64}` size (nd,n) or (n) charge densities (c_{j})
* `dipvecs::Array{ComplexF64}` size (nd,3,n) or (3,n)) dipole orientation vectors (``v_{j}``)
* `pgt::Integer` target eval flag.
    + Potential at targets evaluated if `pgt == 1`.
    + Potential and gradient at targets evaluated if `pgt == 2`
* `nd::Integer` number of densities
* `thresh::Float64` threshold for ignoring interactions when ``\\|x-x_{j}\\| \\leq thresh``
# Output

`vals::LaplaceOutput` with the fields
* `vals.pottarg::Array{ComplexF64}` size (nd,nt) or (nt) potential at target locations if requested
* `vals.gradtarg::Array{ComplexF64}` size (nd,3,nt) or (3,nt) gradient at target locations if requested
"""
function l3ddir(sources::Array{Float64},targets::Array{Float64};
                charges::TFN=nothing,dipvecs::TFN=nothing,
                pgt::Integer=0,nd::Integer=1,
                thresh::Float64=1e-16)

    # check inputs

    @assert size(sources,1) == 3
    @assert size(targets,1) == 3
    @assert nd >= 0
    @assert (0 <= pgt && pgt <= 2) "flag pgt not in expected range"

    # default values

    vals = LaplaceOutput(pgt)

    ifcharge = 0
    ifdipole = 0

    zero = Float64(0)

    pottarg = zero
    gradtarg = zero
    hesstarg = zero

    n = div(length(sources),3)
    nt = div(length(targets),3)

    if charges !== nothing
        @assert div(length(charges),nd) == n
        ifcharge = 1
    else
        charges = zero
        ifcharge = 0
    end

    if dipvecs !== nothing
        @assert div(length(dipvecs),nd) == n*3
        ifdipole = 1
    else
        dipvecs = zero
        ifdipole = 0
    end

    if ifcharge == 0 && ifdipole == 0
        @warn "no charges or dipoles provided, doing nothing"
        return vals
    end

    if (pgt != 1 && pgt != 2)
        @warn "no output requested, doing nothing"
        return vals
    end

    # allocate memory for return values

    if pgt == 1 || pgt == 2
        if nd > 1
            pottarg = zeros(Float64,nd,nt)
        else
            pottarg = zeros(Float64,nt)
        end
    end

    if pgt == 2
        if nd > 1
            gradtarg = zeros(Float64,nd,3,nt)
        else
            gradtarg = zeros(Float64,3,nt)
        end
    end

    # dispatch to appropriate wrapper

    if pgt == 1
        if ifcharge == 1
            if ifdipole == 1
                l3ddirectcdp!(nd,sources,charges,
                              dipvecs,n,targets,nt,
                              pottarg,thresh)
                vals.pottarg = pottarg
            else
                l3ddirectcp!(nd,sources,charges,
                             n,targets,nt,pottarg,
                             thresh)
                vals.pottarg = pottarg
            end
        else
            if ifdipole == 1
                l3ddirectdp!(nd,sources,
                             dipvecs,n,targets,nt,
                             pottarg,thresh)
                vals.pottarg = pottarg
            end
        end
    elseif pgt == 2
        if ifcharge == 1
            if ifdipole == 1
                l3ddirectcdg!(nd,sources,charges,
                              dipvecs,n,targets,nt,
                              pottarg,gradtarg,thresh)
                vals.pottarg = pottarg
                vals.gradtarg = gradtarg
            else
                l3ddirectcg!(nd,sources,charges,
                             n,targets,nt,pottarg,
                             gradtarg,thresh)
                vals.pottarg = pottarg
                vals.gradtarg = gradtarg
            end
        else
            if ifdipole == 1
                l3ddirectdg!(nd,sources,
                              dipvecs,n,targets,nt,
                             pottarg,gradtarg,thresh)
                vals.pottarg = pottarg
                vals.gradtarg = gradtarg
            end
        end
    end

    return vals

end

function l3ddirectcp!(nd,sources,charges,n,targets,nt,
                      pot,thresh)
    ccall((:l3ddirectcp_,libfmm3d),Cvoid,(Fi,Fd,Fd,Fi,
                                          Fd,Fi,Fd,Fd),
          nd,sources,charges,n,targets,nt,pot,thresh)
    return
end

function l3ddirectdp!(nd,sources,dipvecs,n,targets,nt,
                      pot,thresh)
    ccall((:l3ddirectdp_,libfmm3d),Cvoid,(Fi,Fd,Fd,Fi,
                                          Fd,Fi,Fd,Fd),
          nd,sources,dipvecs,n,targets,nt,pot,thresh)
    return
end

function l3ddirectcdp!(nd,sources,charges,dipvecs,n,targets,
                       nt,pot,thresh)
    ccall((:l3ddirectcdp_,libfmm3d),Cvoid,(Fi,Fd,Fd,Fd,Fi,
                                          Fd,Fi,Fd,Fd),
          nd,sources,charges,dipvecs,n,targets,nt,pot,thresh)
    return
end

function l3ddirectcg!(nd,sources,charges,n,targets,nt,
                      pot,grad,thresh)
    ccall((:l3ddirectcg_,libfmm3d),Cvoid,(Fi,Fd,Fd,Fi,
                                          Fd,Fi,Fd,Fd,Fd),
          nd,sources,charges,n,targets,nt,pot,grad,thresh)
    return
end

function l3ddirectdg!(nd,sources,dipvecs,n,targets,nt,
                      pot,grad,thresh)
    ccall((:l3ddirectdg_,libfmm3d),Cvoid,(Fi,Fd,Fd,Fi,
                                          Fd,Fi,Fd,Fd,Fd),
          nd,sources,dipvecs,n,targets,nt,pot,grad,thresh)
    return
end

function l3ddirectcdg!(nd,sources,charges,dipvecs,n,targets,
                       nt,pot,grad,thresh)
    ccall((:l3ddirectcdg_,libfmm3d),Cvoid,(Fi,Fd,Fd,Fd,Fi,
                                          Fd,Fi,Fd,Fd,Fd),
          nd,sources,charges,dipvecs,n,targets,nt,pot,grad,thresh)
    return
end


@testset "Spaces" begin
  k = FlintQQ

  q = quadratic_space(k, 2)
  @test sprint(show, q) isa String
  @test sprint(show, Hecke.isometry_class(q)) isa String
  @test sprint(show, Hecke.isometry_class(q, 2)) isa String
  @test is_definite(q)
  v = matrix(k,1,2,[2,1])
  @test inner_product(q,v,v) == matrix(k,1,1,[5])
  @test Hecke._inner_product(lattice(q),v,v) == matrix(k,1,1,[5])

  Qx, x = PolynomialRing(FlintQQ, "x")
  K1, a1 = NumberField(x^2 - 2, "a1")
  K2, a2 = NumberField(x^3 - 2, "a2")

  K1t, t = PolynomialRing(K1, "t")
  F = GF(3)

  Hecke.change_base_ring(::FlintRationalField, ::Hecke.gfp_mat) = error("asd")
  @test_throws ErrorException quadratic_space(FlintQQ, F[1 2; 2 1])

  Hecke.change_base_ring(::FlintRationalField, x::Hecke.gfp_mat) = x
  @test_throws ErrorException quadratic_space(FlintQQ, F[1 2; 2 1])

  L, b = NumberField(t^2 + a1)

  for K in [k, K1, K2, L]
    V = @inferred quadratic_space(K, 2)
    @test_throws ArgumentError quadratic_space(K, -1)

    V = @inferred quadratic_space(K, identity_matrix(FlintZZ, 2))
    @test V == V
    @test V === V
    W = quadratic_space(K, identity_matrix(FlintZZ, 2))
    @test V === W
    W = quadratic_space(K, identity_matrix(FlintZZ, 2), cached = false)
    @test V != W

    @test (@inferred gram_matrix(V)) == identity_matrix(K, 2)
    @test (@inferred dim(V)) == 2
    @test (@inferred rank(V)) == 2
    @test @inferred is_regular(V)
    @test (@inferred involution(V)) === identity

    V = @inferred quadratic_space(K, zero_matrix(K, 2, 2))
    @test (@inferred gram_matrix(V)) == zero_matrix(K, 2, 2)
    @test (@inferred gram_matrix(V)) isa dense_matrix_type(K)
    @test (@inferred dim(V)) == 2
    @test (@inferred rank(V)) == 0
    @test @inferred is_quadratic(V)
    @test @inferred !ishermitian(V)
    @test (@inferred fixed_field(V)) === K

    @test_throws ArgumentError quadratic_space(K, zero_matrix(K, 1, 2))

    M = identity_matrix(K, 2)
    M[1, 2] = 2
    M[2, 1] = 1
    @test_throws ArgumentError quadratic_space(K, M)

    # Determinant & discrimimant

    V = quadratic_space(K, 2)
    @test (@inferred discriminant(V)) == -1
    @test (@inferred det(V)) == 1

    V = quadratic_space(K, 4)
    @test (@inferred discriminant(V)) == 1
    @test (@inferred det(V)) == 1

    M = identity_matrix(K, 2)
    M[1, 2] = 2
    M[2, 1] = 2
    V = quadratic_space(K, M)
    @test (@inferred discriminant(V)) == 3
    @test (@inferred det(V)) == -3

    # Gram matrix

    M = identity_matrix(K, 2)
    M[1, 2] = 2
    M[2, 1] = 2
    V = @inferred quadratic_space(K, M)
    N = zero_matrix(K, 4, 2)
    @test (@inferred gram_matrix(V, N)) == zero_matrix(K, 4, 4)

    N = identity_matrix(QQ, 2)
    @test (@inferred gram_matrix(V, N)) == M

    N = zero_matrix(K, 4, 4)
    @test_throws ArgumentError gram_matrix(V, N)

    v = [[1, 0], [0, 1]]
    @test (@inferred gram_matrix(V, v) == M)

    v = [[1, 0, 0], [0, 1]]
    @test_throws ErrorException gram_matrix(V, v)

    B = @inferred orthogonal_basis(V)
    @test is_diagonal(gram_matrix(V, B))

    D = @inferred diagonal(V)
    @test length(D) == 2
    @test issetequal(D, map(K, [1, -3]))

    M = rand(MatrixSpace(K, 4, 4), -10:10)
    M = M + transpose(M)
    while iszero(det(M))
      M = rand(MatrixSpace(K, 4, 4), -10:10)
      M = M + transpose(M)
    end

    V = quadratic_space(K, M)

    F, S = @inferred Hecke._gram_schmidt(M, identity)
    @test gram_matrix(V, S) == F

    M[1, 1] = 0
    M[1, 2] = 0
    M[1, 3] = 0
    M[1, 4] = 0
    M[2, 1] = 0
    M[3, 1] = 0
    M[4, 1] = 0

    V = quadratic_space(K, M)

    @test_throws ErrorException Hecke._gram_schmidt(M, identity)
    F, S = @inferred Hecke._gram_schmidt(M, identity, false)
    @test gram_matrix(V, S) == F
  end

  fl, T =  Hecke.is_isometric_with_isometry(quadratic_space(QQ, matrix(QQ, 2, 2, [4, 0, 0, 1])),
                                            quadratic_space(QQ, matrix(QQ, 2, 2, [1, 0, 0, 1])))
  @test fl

  Qx, x = PolynomialRing(FlintQQ, "x", cached = false)
  f = x - 1;
  K, a = number_field(f)
  D = matrix(K, 2, 2, [1, 0, 0, 3]);
  V = quadratic_space(K, D)
  fl, T = Hecke.is_isometric_with_isometry(V, V)
  @test fl

  F,a = number_field(x^2 - 2, :a)
  tt = [[-4, 13, -5, 16],
      [-4, 19, 5, -24],
      [1, -1, 0, 1],
      [3, 8, -2, -5],
      [-3, -13, 4, 17],
      [7, 19, -3, -8],
      [3, -17, -1, 6],
      [11, -9, 5, -4],
      [-5, 7, -3, 4],
      [10, 27, -3, -8]]
  dd = [[7//5*a - 5//8, -2//9],
      [-a + 8//3, 4//5*a - 4//3],
      [2*a - 9//4, 10//9*a - 3],
      [-a - 1, -9//7*a - 4],
      [-9//2, 1//2*a - 5//4],
      [a - 1//3, 1//2*a - 7//4],
      [-3//7*a - 5//6, 9//2*a + 5//9],
      [-1//2*a + 3//4, 1//10*a - 10//9],
      [2//3*a + 1//5, 3//10*a + 1//4],
      [3//5*a + 7//9, -1//2*a + 1//9],
      [-1//3, a - 10//7],
      [6//5, -9//2*a - 7//3],
      [5//3*a - 1, -5//7*a + 1//9],
      [4//5*a + 1, -10//9*a + 7//5],
      [2//5*a - 6//5, -3//4*a - 5],
      [-2*a + 2, -1//8*a + 4//7],
      [a - 1//8, 2*a - 10//7],
      [5//4, 9*a + 1],
      [-5//7*a - 1//3, -3//5*a + 4//7],
      [-5//3*a + 1//2, -1//9*a - 8//9]]
  tt = [matrix(F, 2, 2, [F(s) for s in t]) for t in tt]
  dd = [diagonal_matrix([F(s) for s in d]) for d in dd]

  for i in 1:10
    t = rand(tt)
    d = rand(dd)
    q1 = quadratic_space(F, d)
    q2 = quadratic_space(F, t*d*transpose(t))
    fl, T = Hecke.is_isometric_with_isometry(q1, q2)
    @test fl
    @test d == T*gram_matrix(q2)*transpose(T)

    # the above calls _isisometric_with_isometry_rank_2 on small input
    # test _isisometric_with_isometry with small input here
    fl, T = Hecke._isisometric_with_isometry(gram_matrix(q1), gram_matrix(q2))
    @test fl
    @test d == T*gram_matrix(q2)*transpose(T)
  end
end

@testset "quadratic spaces from invariants" begin
  q = Hecke._quadratic_form_with_invariants(8,QQ(1),[ZZ(2),ZZ(3)],0)
  q = quadratic_space(QQ, q)
  @test hasse_invariant(q, 2) == -1
  @test hasse_invariant(q, ideal(ZZ,3)) == -1
  @test hasse_invariant(q, ZZ(3)) == -1
  # small ranks should be covered by the tests of GenusRep
  K, a = rationals_as_number_field()
  OK = maximal_order(K)
  rk = 8
  det = K(1)
  pinf = infinite_place(K, 1)
  for finite in [[ideal(OK, 2),ideal(OK,5)],[ideal(OK, 3),ideal(OK,7)]]
    for neg in [Dict(pinf=>0),Dict(pinf=>4),Dict(pinf=>8)]
      q = Hecke._quadratic_form_with_invariants(rk, det, finite, neg)
      rkq, ker, detq, finiteq, negq = Hecke._quadratic_form_invariants(q)
      @test rkq == rk
      @test ker == 0
      @test is_square(detq*det)[1]
      @test all(finiteq[p] == -1 for p in finite)
      @test Dict(negq) == neg
    end
  end

  R,x = PolynomialRing(QQ,:x)
  F, a = number_field(x^2 - 3)
  OF = maximal_order(F)
  inf1 = infinite_place(F, 1)
  inf2 = infinite_place(F, 2)
  p2 = prime_ideals_over(OF, 2)[1]
  p3 = prime_ideals_over(OF, 3)[1]
  p5 = prime_ideals_over(OF, 5)[1]
  p13a, p13b = prime_ideals_over(OF, 13)

  d = [F(t) for t in [-3//4*a, -2*a - 5//4, -3//10*a + 1//5, -4//3*a - 2//9, -5//7*a - 3//4]]
  q = quadratic_space(F, diagonal_matrix(d))
  s = Hecke.isometry_class(q)
  @test s == Hecke.isometry_class(representative(s))
  rk = 5
  det = F(30)
  neg =
  for (finite,neg) in [([p2,p13a],(0,0)),([p2,p3,p13a],(2,0)), ([p3,p13b],(4,4))]
    neg = Dict(inf1=>neg[1],inf2=>neg[2])
    q = Hecke._quadratic_form_with_invariants(rk, det, finite, neg)
    rkq, ker, detq, finiteq, negq = Hecke._quadratic_form_invariants(q)
    @test rkq == rk
    @test ker == 0
    @test is_square(detq*det)[1]
    @test all(finiteq[p] == -1 for p in finite)
    @test Dict(negq) == neg
  end

  rk = 3
  det = F(-30)
  Hecke._quadratic_form_with_invariants(1, det, [], Dict(inf1=>1,inf2=>1))
  for (finite,neg) in [([p2,p13a,p5],(1,3)),([p3,p13a],(1,1)), ([p3,p13b],(3,3))]
    neg = Dict(inf1=>neg[1],inf2=>neg[2])
    q = Hecke._quadratic_form_with_invariants(rk, det, finite, neg)
    rkq, ker, detq, finiteq, negq = Hecke._quadratic_form_invariants(q)
    @test rkq == rk
    @test ker == 0
    @test is_square(detq*det)[1]
    @test all(finiteq[p] == -1 for p in finite)
    @test Dict(negq) == neg
    Hecke._isisotropic_with_vector(q)
  end

  _Q, = Hecke.rationals_as_number_field()
  K, = quadratic_field(3)
  u = QQ(1)
  v = QQ(2)
  for _A in -3:3
    for _B in -3:3
      for KK in [QQ, _Q, K]
        (iszero(_A) || iszero(_B)) && continue
        A = KK(_A)
        B = KK(_B)
        a = A * u^2 + B * v^2
        fl, _u, _v = Hecke._solve_conic_affine(A, B, a)
        @test fl
        @test a == A * _u^2 + B * _v^2
        A = KK(1//_A)
        B = KK(1//_B)
        a = A * u^2 + B * v^2
        fl, _u, _v = Hecke._solve_conic_affine(A, B, a)
        @test fl
        @test a == A * _u^2 + B * _v^2
      end
    end
  end
end

@testset begin "finding isotropic vectors"
  d  = fmpq[25//21, -1, 37//26, 31//45, -24//25, -9//25]
  q = diagonal_matrix(d)
  b, v = Hecke._isisotropic_with_vector(q)
  q = quadratic_space(QQ, q)
  @test b
  @test inner_product(q, v, v)==0

#  too long even for a long test
#   if long_test
#     K,b = cyclotomic_field(16)
#     F, _a = number_field(minpoly(b+b^-1))
#     d = [36//25*_a - 35//29, -7//2*_a + 26//3, -28//15*_a - 33//28, -12//37*_a + 12, 7//32*_a + 11//3]
#     q = diagonal_matrix(d)
#     Hecke._isisotropic_with_vector(q)
#  end
end

@testset "isometry classes of spaces" begin
  # isometry classes over the rationals
  q = quadratic_space(QQ,QQ[-1 0; 0 1])
  qq = quadratic_space(QQ,QQ[-1 0; 0 3])
  q_deg = quadratic_space(QQ,QQ[-1 0 0; 0 3 0; 0 0 0])
  g = Hecke.isometry_class(q)
  gg = Hecke.isometry_class(qq)
  gg_deg = Hecke.isometry_class(q_deg)
  ggg = Hecke.isometry_class(orthogonal_sum(q,qq)[1])
  @test g + gg == ggg
  @test g + gg - g == gg
  @test g + g + gg - g == gg+ g
  @test represents(gg,-1)
  @test represents(gg,-3)
  @test represents(gg_deg, -3)
  @test Hecke.is_isometric_with_isometry(q, representative(g))[1]
  g2 = Hecke.isometry_class(q,2)
  for p in [2,3,5,7,11]
    @test Hecke.isometry_class(q, p) == local_symbol(g, p)
  end
  @test g2 == local_symbol(g, 2)
  @test Hecke.signature_tuple(q) == Hecke.signature_tuple(g)
  @test hasse_invariant(q,2) == hasse_invariant(g2)
  @test dim(q) == dim(g)
  @test is_square(det(q)*det(g))
  @test witt_invariant(q, 2) == witt_invariant(g2)
  q0 = quadratic_space(QQ,matrix(QQ,0,0,fmpq[]))
  g0 = Hecke.isometry_class(q0)
  g0p = Hecke.isometry_class(q0, 2)
  @test g == g+g0
  @test Hecke.represents(g, g0)
  @test Hecke.isometry_class(representative(gg + gg + g)) == gg + gg + g
  @test Hecke.isometry_class(representative(g+g+gg+gg)) == g + g + gg+gg

  # isometry classes over number fields
  R, x = PolynomialRing(QQ, "x")
  F, a = number_field(x^2 -3)
  infF = infinite_place(F,1)
  infF2 = infinite_place(F,2)
  q = quadratic_space(F, F[1 0; 0 a])
  @test Hecke.is_isotropic(q, infF)
  qq = quadratic_space(F, F[-49 0; 0 a])
  h = quadratic_space(F, F[0 1; 1 a])
  @test Hecke.is_isotropic(qq, infF2)
  @test Hecke._isisotropic_with_vector(gram_matrix(h))[1]
  @test !Hecke._isisotropic_with_vector(gram_matrix(q))[1]
  hh,_,_ = orthogonal_sum(qq,quadratic_space(F,-gram_matrix(qq)))
  i = Hecke._maximal_isotropic_subspace(gram_matrix(hh))
  @test nrows(i)==dim(qq)
  @test i*gram_matrix(hh)*transpose(i) == 0

  g = Hecke.isometry_class(q)
  gg = Hecke.isometry_class(qq)
  @test represents(g, a)
  @test represents(g, 1+a)
  @test represents(gg, -49)
  @test represents(gg, a-49)
  @test represents(gg+g, g)
  @test represents(gg+g, gg)
  p = prime_ideals_over(maximal_order(F),2)[1]
  gp = Hecke.isometry_class(q, p)
  @test g + gg + g - g  == g + gg
  @test Hecke.signature_tuples(q) == Hecke.signature_tuples(g)
  @test Hecke.signature_tuple(q, infF) == Hecke.signature_tuple(g, infF)
  @test hasse_invariant(q,p) == hasse_invariant(gp)
  @test dim(q) == dim(g)
  @test is_square(det(q)*det(g))[1]
  r = quadratic_space(g)
  @test Hecke.is_isometric_with_isometry(q, r)[1]
  @test is_isometric(q,r, p)
  @test is_isometric(q,r, infF)
  @test is_isometric(q,r)
  L = Zlattice(gram=ZZ[1 1; 1 2])
  g = genus(L)
  c1 = Hecke.isometry_class(ambient_space(L))
  c2 = Hecke.rational_isometry_class(g)
  @test c1 == c2

  # More complicated isisotropic_with_isometry
  F = QQ[2 0 0 0 0 0; 0 1 0 0 0 0; 0 0 70//9 0 0 0; 0 0 0 -311//105 64//21 -286//105; 0 0 0 64//21 -67//21 65//21; 0 0 0 -286//105 65//21 -446//105]
  fl, v = Hecke._isisotropic_with_vector(F)
  @test fl
  vm = matrix(QQ, 1, 6, v)
  @test iszero(vm * F * transpose(vm))
end


VcovFormula(::Type{Val{:robust}}) = VcovRobustFormula()

struct VcovRobustFormula <: AbstractVcovFormula  end

struct VcovRobustMethod <: AbstractVcovMethod end
VcovMethod(::AbstractDataFrame, ::VcovRobustFormula) = VcovRobustMethod()

function vcov!(v::VcovRobustMethod, x::VcovData)
    S = shat!(v, x)
    return sandwich(x.crossmatrix, S)
end

# S_{(l-1) * K + k, (l'-1)*K + k'} = \sum_i X[i, k] res[i, l] X[i, k'] res[i, l']
function shat!(::VcovRobustMethod, x::VcovData{T, N}) where {T, N}
    dim = size(x.regressors, 2) * size(x.residuals, 2)
    X2 = fill(zero(Float64), size(x.regressors, 1), dim)
    index = 0
    for k in 1:size(x.residuals, 2)
        for j in 1:size(x.regressors, 2)
            index += 1
            @inbounds @simd for i in 1:size(x.regressors, 1)
                X2[i, index] = x.regressors[i, j] * x.residuals[i, k]
            end
        end
    end
    S2 = X2' * X2
    rmul!(S2, size(x.regressors, 1) / x.dof_residual)
    return S2
end


const DEFAULT_NUM_POLYS = 4

struct StochGalerkinODE{
    odeFuncType,
    dimType,
    varsType,
    numPolysType,
    indsDegType,
    polysType,
    gridType,
    normMapType,
}
    ode_func!::odeFuncType
    dim_::dimType
    vars::varsType
    num_polys::numPolysType
    indices_polys::indsDegType
    poly_basis::polysType
    grid::gridType
    normalization_map::normMapType

    function StochGalerkinODE(ode_func!, dim_, vars; wType = Float64, num_polys = nothing)

        if num_polys === nothing
            num_polys = Tuple(DEFAULT_NUM_POLYS for i in Base.OneTo(length(vars)))
        end

        @assert length(vars) == length(num_polys)

        indices_polys = CartesianIndices(num_polys)
        num_vars = length(vars)

        poly_basis_temp = Vector{OrthogonalPolyBasis}(undef, num_vars)

        for i in Base.OneTo(length(vars))
            var = vars[i]
            if var isa Uniform
                @assert var.a ≈ -1.0 && var.b ≈ 1.0
                poly_basis_temp[i] = LegendrePolyBasis(num_polys[i]; wType)
            elseif var isa Normal
                # Using physicist's hermite polynomial
                @assert var.μ ≈ 0.0 && var.σ ≈ sqrt(2.0) / 2
                poly_basis_temp[i] = HermitePolyBasis(num_polys[i]; wType)
            else
                throw("Not implemented")
            end
        end

        poly_basis = Tuple(poly_basis_temp)

        grid = TensorIntegrationGrid(dim_, num_polys, vars; wType)

        normalization_map = zeros(wType, num_polys)

        for index in CartesianIndices(normalization_map)
            normalization_map[index] = mapreduce(
                i -> dot(
                    grid.quads[i].weights,
                    (x -> poly_basis[i](x, index[i] - 1)^2).(grid.quads[i].nodes),
                ),
                *,
                Base.OneTo(length(index)),
            )
        end

        return new{
            typeof(ode_func!),
            typeof(dim_),
            typeof(vars),
            typeof(num_polys),
            typeof(indices_polys),
            typeof(poly_basis),
            typeof(grid),
            typeof(normalization_map),
        }(
            ode_func!,
            dim_,
            vars,
            num_polys,
            indices_polys,
            poly_basis,
            grid,
            normalization_map,
        )
    end
end

struct StochGalerkinODECache{yHatType,yHatGridType,intEstType,argCacheType,odeCacheType}
    y_hat::yHatType
    y_hat_for_grid::yHatGridType
    integral_est::intEstType
    arg_cache::argCacheType
    ode_cache::odeCacheType

    function StochGalerkinODECache(s::StochGalerkinODE, u)
        @unpack dim_, vars, num_polys = s

        uType = eltype(u)
        num_vars = length(vars)

        y_hat = zeros(uType, dim_)
        y_hat_for_grid = zeros(uType, dim_, num_polys...)
        integral_est = zeros(uType, dim_)
        arg_cache = zeros(uType, num_vars)
        ode_cache = zeros(uType, dim_)

        return new{
            typeof(y_hat),
            typeof(y_hat_for_grid),
            typeof(integral_est),
            typeof(arg_cache),
            typeof(ode_cache),
        }(
            y_hat,
            y_hat_for_grid,
            integral_est,
            arg_cache,
            ode_cache,
        )
    end
end

@inline function compute_poly_prod(
    s::StochGalerkinODE,
    c::StochGalerkinODECache,
    index::CartesianIndex,
)
    @unpack poly_basis = s
    @unpack arg_cache = c
    prod = one(first(arg_cache))
    @inbounds for i in Base.OneTo(length(index))
        prod *= poly_basis[i](arg_cache[i], index[i] - 1)
    end
    return prod
end

@inline function compute_weights_prod(s::StochGalerkinODE, index::CartesianIndex)
    @unpack grid = s
    @unpack quads = grid
    prod = one(eltype(first(quads).weights))
    @inbounds for i in Base.OneTo(length(index))
        prod *= quads[i].weights[index[i]]
    end
    return prod
end

@inline function set_arg!(
    s::StochGalerkinODE,
    c::StochGalerkinODECache,
    index::CartesianIndex,
)
    @unpack grid = s
    @unpack arg_cache = c
    @unpack quads = grid
    @inbounds for i in Base.OneTo(length(index))
        arg_cache[i] = quads[i].nodes[index[i]]
    end
    nothing
end

function y_hat_expand!(s::StochGalerkinODE, c::StochGalerkinODECache, u)
    @unpack indices_polys, poly_basis = s
    @unpack y_hat, arg_cache = c
    yType = eltype(y_hat)
    fill!(y_hat, zero(yType))
    @inbounds for index in indices_polys
        # Compiler can't infer without this
        prod::yType = compute_poly_prod(s, c, index)
        y_hat .+= (@view u[:, index]) .* prod
    end
    nothing
end

function set_inner_prod!(s::StochGalerkinODE, c::StochGalerkinODECache, u)
    @unpack grid = s
    @unpack y_hat, y_hat_for_grid = c
    @unpack quad_indices = grid
    @inbounds for index in quad_indices
        set_arg!(s, c, index)
        y_hat_expand!(s, c, u)
        y_hat_for_grid[:, index] .= y_hat
    end
    nothing
end

function integrate!(
    s::StochGalerkinODE,
    c::StochGalerkinODECache,
    u,
    p,
    t,
    index::CartesianIndex,
)
    @unpack ode_func!, poly_basis, grid = s
    @unpack y_hat, y_hat_for_grid, integral_est, arg_cache, ode_cache = c
    @unpack quad_indices = grid
    fill!(integral_est, zero(eltype(integral_est)))
    @inbounds for index_ in quad_indices
        set_arg!(s, c, index_)
        prod = compute_poly_prod(s, c, index) * compute_weights_prod(s, index_)
        y_hat .= (@view y_hat_for_grid[:, index_])
        ode_func!(ode_cache, y_hat, p, arg_cache, t)
        integral_est .+= ode_cache .* prod
    end
    nothing
end

function stoch_galerkin_ode_dynamics!(s::StochGalerkinODE, du, u, p, t)
    c = StochGalerkinODECache(s, u)
    stoch_galerkin_ode_dynamics!(s, c, du, p, t)
    nothing
end

function stoch_galerkin_ode_dynamics!(
    s::StochGalerkinODE,
    c::StochGalerkinODECache,
    du,
    u,
    p,
    t,
)
    @unpack indices_polys, normalization_map = s
    @unpack integral_est = c
    set_inner_prod!(s, c, u)
    @inbounds for index in indices_polys
        norm_const = normalization_map[index]
        integrate!(s, c, u, p, t, index)
        du[:, index] .= integral_est ./ norm_const
    end
    nothing
end

function (s::StochGalerkinODE)(u0, tspan, p; re_use_cache = true, kwargs...)
    @unpack dim_, num_polys = s
    u_orth_0 = zeros(eltype(u0), dim_, num_polys...)
    u_orth_0[:, CartesianIndex(Tuple(1 for i in num_polys))] .= u0
    if re_use_cache
        local c = StochGalerkinODECache(s, u_orth_0)
        dynamics_iip(du, u, p, t) = stoch_galerkin_ode_dynamics!(s, c, du, u, p, t)
        prob = ODEProblem(dynamics_iip, u_orth_0, tspan, p)
        sol = solve(prob; kwargs...)
    else
        dynamics_oop(du, u, p, t) = stoch_galerkin_ode_dynamics!(s, du, u, p, t)
        prob = ODEProblem(dynamics_oop, u_orth_0, tspan, p)
        sol = solve(prob; kwargs...)
    end
    return sol
end


# Fix EM for Gaussian SPNs... (numerically unstable)

@testset "GSPN parameter learning" begin

    println()
    @info "## Learning Gaussian SPNs"

    # First create model to generate data
    A = Vector([
        SumNode(),     # 1
        ProductNode(), # 2
        ProductNode(), # 3
        ProductNode(), # 4
        GaussianDistribution(1, 2, 18),  # 5
        GaussianDistribution(1, 11, 8), # 6
        GaussianDistribution(2, 3, 10),  # 7
        GaussianDistribution(2, -4, 7)   # 8
    ])
    B = [
        [2, 3, 4], # 1 Sum P(A, B) 0.2,0.5,0.3
        [5, 7],    # 2 Prod P(A, B)=P(A)*P(B)
        [5, 8],    # 3 Prod P(A, B)=P(A)*P(B)
        [6, 8],    # 4 Prod P(A, B)=P(A)*P(B)
    ]
    C = sparse([2,   3,     4],
               [1,   1,     1],
               [4/20, 9/20, 7/20])

    SPN = SumProductNetwork{Int64}(A, B, C)
    #println(GSPN)

    @info " Generating data set..."
    N = 10000
    data = rand(SPN,N)

    # Now learn parameters from data
    SPN2 = SumProductNetwork{Int64}(deepcopy(A), deepcopy(B), deepcopy(C))
    initialize(SPN2) # random parameter initialization

    @info "Running Expectation Maximization until convergence..."

    learner = EMParamLearner()
    learner.minimumvariance = 0.1

    while !converged(learner) && learner.steps < 1000
        SumProductNetworks.step(learner,SPN2,data)
        # mean absolute error
        error = MAE(SPN,SPN2,data)
        println("It: $(learner.steps) \t NLL: $(-learner.score) \t MAE: $error")
    end

    @test learner.score > -60000
    @test MAE(SPN,SPN2,data) < 0.3

    for node in leaves(SPN2)
        println(node)
    end

    println()

 end # End of test set


include("data.jl")

mutable struct supportsresistances
   close::Float64
   date::
   timeframe::String
   asset::String
   regime::String
   pos::String
   supportsresistances(close, date, asset, timeframe) = begin
      close = close
      date = date
      timeframe = timeframe.upper()
      asset = asset
      regime = asset.regime
      pos = length(asset.data[timeframe])
   new(close, date, asset, timeframe)
   end
end

function RSI(Market::market)
   position = length(Market.data)
   if position < 15
      return 0.00
   else
      data = Market.data[position-14:position]
      wins = []
      losses = []
      for i in data
         if i.Close < i.Open
            push!(losses,(i.Open-i.Close))
         elseif i.Close > i.Open
            push!(wins,(i.Close-i.Open))
         end
      end
      winssum = sum(wins)/length(wins)
      lossessum = sum(losses)/length(losses)
      firstrs = winssum/lossessum
      RSI = 100-(100/firstrs)
      return RSI
   end
end

function MACD(ma26::Float64,ma12::Float64,ma9::Float64)
   macdline = ma12 - ma26
   signalline = ma9
   macdhistogram = macdline - signalline
   return macdhistogram
end

function MA9(Market::market)
   position = length(Market.data)
   if position < 10
      return 0.00
   else
      ma9 = sum(Market.data[position-9].Close:Market.data[position].Close)/9
      return ma9
   end
end

function MA12(Market::market)
   position = length(Market.data)
   if position < 13
      return 0.00
   else
      ma12 = sum(Market.data[position-12].Close:Market.data[position].Close)/12
      return ma12
   end
end

function MA26(Market::market)
   position = length(Market.data)
   if position < 27
      return 0.00
   else
      ma26 = sum(Market.data[position-26].Close:Market.data[position].Close)/26
      return ma26
   end
end

function MA52(Market::market)
   position = length(Market.data)
   if position < 53
      return 0.00
   else
      ma52 = sum(Market.data[position-52].Close:Market.data[position].Close)/52
      return ma52
   end
end


using Distributions

# config
Ks = [3, 10, 30, 100, 1000, 10000]
n = 10^5

function run_sample(sampler, n::Int)
	for i = 1:n
		rand(sampler)
	end
end

macro bench_sampler(construct_expr)
	quote
		s = $(esc(construct_expr)) # construct sampler
		tname = typeof(s)

		# warm up
		run_sample(s, 100)

		et = @elapsed run_sample(s, n)

		mps = n * 1.0e-6 / et
		@printf("  %-28s :  elapsed = %12.6f sec  |  %8.4f MPS\n",
			tname, et, mps)
	end
end

for K in Ks
	println("K = $K")
	println("-----------------------")

	p = fill(1/K, K)
	pu = fill(uint64(1), K)

	@bench_sampler DiscreteUniform(1, K)   # samples generated by rand(1:K)
	@bench_sampler Categorical(p)
	# @bench_sampler Distributions.DiscreteDistributionTable(p)
	# @bench_sampler Distributions.huffman(1:K, pu)
	@bench_sampler Distributions.AliasTable(p)

	println()
end


@testset "MPS Tensors and Transfer Matrices" begin
    d1 = 3; d2 = 2; AD1 = 4; AD2 = 5; BD1 = 2; BD2 = 3

    A1 = rand_MPSTensor(Float64,d1,AD1,AD2)
    @test bond_dim(A1) == AD1
    @test bond_dim_R(A1) == AD2
    @test phys_dim(A1) == d1

    A1 = rand_MPSTensor_unitary(Float64,d1,AD1,AD2)
    @test bond_dim(A1) == AD1
    @test bond_dim_R(A1) == AD2
    @test phys_dim(A1) == d1

    for T in (Float32, Float64, ComplexF32, ComplexF64)
        U = puMPS.MPS.randunitary(T,3)
        @test U  * U' ≈ Matrix{T}(I,3,3)
        @test U' * U  ≈ Matrix{T}(I,3,3)

        A1 = rand_MPSTensor_unitary(T,d1,AD1,AD2)
        @test eltype(A1) == T

        TM1 = TM_convert(TM_dense(A1,A1))
        x = Matrix{T}(I,AD2,AD2)
        y = applyTM_r(A1,A1,x)
        @test eltype(y) == T
        @test all(x - I .≈ 0)
        @tensor y2[t,b] := TM1[t,b,ti,bi] * x[ti,bi]
        @test y2 ≈ y

        B1 = rand_MPSTensor(T,d1,BD1,BD2)
        @test eltype(B1) == T
        x = rand(T,AD1,BD1)
        y = applyTM_l(A1,B1,x)
        TM1 = TM_convert(TM_dense(A1,B1))
        @test eltype(TM1) == T
        @tensor y2[t,b] := conj(TM1[ti,bi,t,b]) * x[ti,bi]
        @test y2 ≈ y

        A2 = rand_MPSTensor(T,d2,AD2,AD2)
        B2 = rand_MPSTensor_unitary(T,d2,BD2,BD2)
        TM2 = TM_convert(TM_dense(A2,B2))
        @tensor TM12[t1,b1,t2,b2] := TM1[t1,b1,ti,bi] * TM2[ti,bi,t2,b2]

        TM1 = TM_convert(TM1)
        TM2 = TM_convert(TM2)
        TM12 = TM_convert(TM12)

        TM12_2 = applyTM_r(A1,B1,TM2)
        @test eltype(TM12_2) == T
        @test TM12_2 ≈ TM12

        TM12_2 = applyTM_l(A2,B2,TM1)
        @test eltype(TM12_2) == T
        @test TM12_2 ≈ TM12

        wrk = workvec_applyTM_r(A1,B1,TM2)
        fill!(TM12_2, 0.0)
        res = applyTM_r!(TM12_2,A1,B1,TM2,wrk)
        @test res === TM12_2
        @test TM12_2 ≈ TM12

        wrk = workvec_applyTM_l(A2,B2,TM1)
        fill!(TM12_2, 0.0)
        res = applyTM_l!(TM12_2,A2,B2,TM1,wrk)
        @test res === TM12_2
        @test TM12_2 ≈ TM12
    end
end

@testset "MPS TM eigensolvers" begin
    d = 3; AD = 4; BD = 2
    for T in (Float32, Float64, ComplexF32, ComplexF64)
        A = rand_MPSTensor_unitary(T,d,AD,AD)
        B = rand_MPSTensor(T,d,BD,BD)
        ev, eVmL, eVmR = tm_eigs_dense(A,B)
        srt = sortperm(ev, by=abs, rev=true)
        ev = ev[srt]
        eVmL = eVmL[srt]
        eVmR = eVmR[srt]

        for j = 1:length(ev)
            y = applyTM_r(A,B,eVmR[j])
            @test y ≈ eVmR[j] * ev[j]

            y = applyTM_l(A,B,eVmL[j])
            @test y ≈ eVmL[j] * conj(ev[j])
        end

        evLS, evRS, eVmLS, eVmRS = tm_eigs(A,B,4,D_dense_max=0)
        @test argmax(abs.(evLS)) == argmax(abs.(evRS)) == 1
        @test abs(evLS[1]) ≈ abs(ev[1])
        @test abs(evRS[1]) ≈ abs(ev[1])

        tol = sqrt(eps(real(T)))
        for j = 1:length(evLS)
            ind = findfirst(x->abs(x - evLS[j]) < tol, ev)
            @test ind !== nothing

            y = applyTM_l(A,B,eVmLS[j])
            @test y ≈ eVmLS[j] * conj(evLS[j])

            ind = findfirst(x->abs(x - evRS[j]) < tol, ev)
            @test ind !== nothing

            y = applyTM_r(A,B,eVmRS[j])
            @test y ≈ eVmRS[j] * evRS[j]
        end

        ev1, l, r = tm_dominant_eigs(A,A,D_dense_max=1000)
        ev1S, lS, rS = tm_dominant_eigs(A,A,D_dense_max=0)
        @test ev1 ≈ ev1S
        @test l ≈ lS
        @test r ≈ rS
    end
end

@testset "MPS/MPO Tensors and Transfer Matrices" begin
    d1 = 3; d2 = 2; AD1 = 4; AD2 = 5; BD1 = 2; BD2 = 3
    M1 = 2; M2 = 4
    for T in (Float32, Float64, ComplexF32, ComplexF64)
        O1 = rand_MPOTensor(T,d1,M1,M2)
        A1 = rand_MPSTensor_unitary(T,d1,AD1,AD2)
        B1 = rand_MPSTensor_unitary(T,d1,BD1,BD2)
        TM1 = TM_dense_MPO(A1,B1,O1)

        O2 = rand_MPOTensor(T,d2,M2,M1)
        A2 = rand_MPSTensor(T,d2,AD2,AD1)
        B2 = rand_MPSTensor(T,d2,BD2,BD1)
        TM2 = TM_dense_MPO(A2,B2,O2)

        @tensor TM12[t1,m1,b1,t2,m2,b2] := TM1[t1,m1,b1,ti,mi,bi] * TM2[ti,mi,bi,t2,m2,b2]

        TM12_2 = applyTM_MPO_r(A1,B1,O1,TM2)
        @test TM12_2 ≈ TM12
        TM12_2 = applyTM_MPO_l(A2,B2,O2,TM1)
        @test TM12_2 ≈ TM12

        wrk = workvec_applyTM_MPO_r(A1,B1,O1,TM2)
        fill!(TM12_2, 0.0)
        res = applyTM_MPO_r!(TM12_2,A1,B1,O1,TM2,wrk)
        @test res === TM12_2
        @test TM12_2 ≈ TM12

        wrk = workvec_applyTM_MPO_l(A2,B2,O2,TM1)
        fill!(TM12_2, 0.0)
        res = applyTM_MPO_l!(TM12_2,A2,B2,O2,TM1,wrk)
        @test res === TM12_2
        @test TM12_2 ≈ TM12
    end
end


# Goal: Evaluate impact of alpha and reaction kinetics on modeled stream-reach
# solute dynamics, assuming steady state
# Note: these are the plots in your 2021 SFS talk

τ₋ = 0.1; # minutes
τ₊ = 7200; # minutes
α = 1.5; # unitless
k = 0.000130; # per minute
sHZ = 1.0; # arbitrary units
sCH = 2.0; # same arbitrary units as sHZ
v = 0.03333 #2 meters/min
qPerUnitHZ = qCalc_powerLaw(1, τ₋, τ₊, α)
xspan = (0.0, 1000) # meters
u0 = [100] # ppb nitrate-N

# Base Problem: nitrate removal only
@stream_model("firstOrder", "powerLaw")
p = [τ₋, τ₊, α, k, qPerUnitHZ, sHZ, sCH, v]
prob = ODEProblem(stream_steady_state!, u0, xspan, p)
sol = solve(prob)

# I think this code finds the value of alpha that makes
# total hyporheic exchange equal to total hyporheic exchange
# with alpha = 1.5 and maximum residence time = 7200.
# (Rewrite this if you need it again)
function f!(F, α)
    F[1] = (43200^(-α[1]+1) - 0.1^(-α[1] + 1) - (43200^(-α[1]))*(43200-0.1)*(-α[1]+1))/
        ((1/2)*(43200^(-α[1]))*(1-α[1])*(43200^2-τ₋^2)-(1/(2-α[1]))*(43200^(2-α[1])-0.1^(2-α[1]))-
        (43200^(1-α[1]))*(43200-0.1)*(-α[1])) - 0.049907218626099197
end
# nlsolve(f!, [1.5])

τ₊ = 43200 # minutes
α = 1.607
qPerUnitHZ = qCalc_powerLaw(1, τ₋, τ₊, α)
p = [τ₋, τ₊, α, k, qPerUnitHZ, sHZ, sCH, v]
prob4 = ODEProblem(stream_steady_state!, u0, xspan, p)
sol4 = solve(prob4)

# Solve model assuming first-order nitrification and denitrification
# and two different ammonium concentrations (10 and 20 ppb)
@stream_model("firstOrderProductionAndConsumption", "powerLaw")
NH40 = 10 # ppb ammonium-N
knit = 0.0011 # per minute
p = [τ₋, τ₊, α, qPerUnitHZ, sHZ, sCH, v, knit, k, NH40]
prob2 = ODEProblem(stream_steady_state!, u0, xspan, p)
sol2 = solve(prob2)

NH40 = 20 # ppb ammonium-N
p = [τ₋, τ₊, α, qPerUnitHZ, sHZ, sCH, v, knit, k, NH40]
prob3 = ODEProblem(stream_steady_state!, u0, xspan, p)
sol3 = solve(prob3)

# Plot nitrate concentration along stream channel
# for various assumptions about alpha and reaction kinetics
plot(
    sol,
    label = "First-order Uptake",
    ylabel = "Nitrate Concentration (μg/L N)",
    xlabel = "Distance Along Stream (meters)",
    lw = 4,
    linecolor = colorant"#3D8D59",
    background_color = :black,
    ylim = (15, 105),
    legend = :false,
    yguidefontsize=15,
    xguidefontsize=15
)
savefig("steadyState1.png")

plot!(
    sol4,
    linecolor = colorant"#3D8D59",
    linestyle = :dot,
    label = "Larger tauMax",
    xlabel = "Distance Along Stream (meters)",
    lw = 4,
    ylim = (15, 105),
    legend = :false,
    yguidefontsize=15,
    xguidefontsize=15
)
savefig("steadyState2.png")

plot!(
    sol2,
    linecolor = colorant"#77BDC6",
    label = "First-order Production and Uptake",
    xlabel = "Distance Along Stream (meters)",
    lw = 4,
    ylim = (15, 105),
    legend = :false,
    yguidefontsize=15,
    xguidefontsize=15
)
savefig("steadyState3.png")

plot!(
    sol3,
    linecolor = colorant"#18537E",
    label = "First-order Production and Uptake\n with Elevated NH4+",
    xlabel = "Distance Along Stream (meters)",
    lw = 4,
    ylim = (15, 105),
    legend = :false,
    yguidefontsize=15,
    xguidefontsize=15
)
savefig("steadyStateRFP4.png")

# Plot relationships between residence time and nitrate
# concentration for different assumptions
ts = collect(0.1:0.1:7200)
tsMin = ts./60
NO3A = u0[1].*exp.(-k.*ts)
NH40 = 10
NO30 = u0[1]
NO3B = (NH40*knit/(k - knit)).*(exp.(-knit.*ts) .- exp.(-k.*ts)) .+ NO30.*exp.(-k.*ts)
NH40 = 20
NO3C = (NH40*knit/(k - knit)).*(exp.(-knit.*ts) .- exp.(-k.*ts)) .+ NO30.*exp.(-k.*ts)
ENV["GKS_ENCODING"]="utf-8"

plot(
    tsMin, NO3A,
    linecolor = colorant"#3D8D59",
    lw = 4,
    ylabel = "Nitrate Concentration (μg/L N)",
    xlabel = "Residence Time (τ, minutes)",
    legend = :false,
    yguidefontsize=15,
    xguidefontsize=15,
    ylim = (40,105),
    background_color = :black
)
savefig("RFPs1.png")

plot!(
    tsMin, NO3B,
    lw = 4,
    linecolor = colorant"#77BDC6",
    xlabel = "Residence Time (τ, minutes)",
    legend = :false,
    yguidefontsize=15,
    xguidefontsize=15,
    ylim = (40,105)
)
savefig("RFPs2.png")

plot!(tsMin, NO3C,
lw = 4,
linecolor = colorant"#18537E",
xlabel = "Residence Time (τ, minutes)",
legend = :false,
yguidefontsize=15,
xguidefontsize=15,
ylim = (40,105))
savefig("RFP3s.png")

# Plot residence time distribution for two different
# combinations of alpha and maximum residence time
τ₊ = 7200
ProbA = zeros(length(ts)-1)
for i in 2:length(ts)
    ProbA[i-1] = E(ts[i], 1.5, τ₋, τ₊)
end

τ₊ = 43200
ProbB = zeros(length(ts)-1)
for i in 2:length(ts)
    ProbB[i-1] = E(ts[i], 1.607, τ₋, τ₊)
end

plot(tsMin[2:30], ProbA[1:29],
linecolor = colorant"#3D8D59",
lw = 3,
ylim = (0.0, 2.0),
ylabel = "Probability Density",
background_color = :black,
xlabel = "Residence Time (τ, minutes)",
legend = :false,
yguidefontsize=15,
xguidefontsize=15)
savefig("PDF1.png")

plot!(tsMin[2:30], ProbB[1:29],
lw = 5,
linecolor = colorant"#3D8D59",
linestyle = :dot,
xlabel = "Residence Time (τ, minutes)",
legend = :false,
yguidefontsize=15,
xguidefontsize=15)
savefig("PDF2.png")


## Time performance on VGG16

using BenchmarkTools
using Flux
using CUDA
using Zygote: pullback

function vgg16()
    Chain(
        Conv((3, 3), 3 => 64, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(64),
        Conv((3, 3), 64 => 64, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(64),
        MaxPool((2,2)),
        Conv((3, 3), 64 => 128, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(128),
        Conv((3, 3), 128 => 128, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(128),
        MaxPool((2,2)),
        Conv((3, 3), 128 => 256, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(256),
        Conv((3, 3), 256 => 256, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(256),
        Conv((3, 3), 256 => 256, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(256),
        MaxPool((2,2)),
        Conv((3, 3), 256 => 512, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(512),
        Conv((3, 3), 512 => 512, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(512),
        Conv((3, 3), 512 => 512, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(512),
        MaxPool((2,2)),
        Conv((3, 3), 512 => 512, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(512),
        Conv((3, 3), 512 => 512, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(512),
        Conv((3, 3), 512 => 512, relu, pad=(1, 1), stride=(1, 1)),
        BatchNorm(512),
        MaxPool((2,2)),
        flatten,
        Dense(512, 4096, relu),
        Dropout(0.5),
        Dense(4096, 4096, relu),
        Dropout(0.5),
        Dense(4096, 10)
    )
end

let model=vgg16(), x=rand(Float32, 32, 32, 3, 64)
    println("CPU benchmark")
    run_benchmark(model, x, cuda=false)
    println("CUDA benchmark")
    run_benchmark(model, x, cuda=true)
end


module SmallLassoDemo

using VarianceReducedSG
using ProximalOperators
using LinearAlgebra
using Random
using Test

# Square-norm gradient
sqgrad(x) = x

function setup_data(N, n, mt)
	# Generate data
	A = randn(mt, n,N)
	y = A*randn(mt,N) # No noise

	# Normalize and reformat
	norm_consts = [norm(A[i,:]) for i = 1:size(A,1)]
	data = [A[i,:]./norm_consts[i] for i = 1:size(A,1)]
	b = [y[i]/norm_consts[i] for i = 1:length(y)]

	return data, b
end

################################################################################
# SAGA
################################################################################
@testset "Lasso" begin
	println("Lasso")
	N = 100
	n = 1000
	mt = MersenneTwister(123)

	data, b = setup_data(N,n,mt)

	f(y,x,i) = (y .= data[i].*sqgrad(dot(data[i],x) - b[i]))
	fs(x,i) = sqgrad(x - b[i])
	vrg = VRGradient(f, zeros(N), n)
	vrg_lin = LinearVRG(fs, data, zeros(N))
	vrg_uni = UniformVRG(f, zeros(N), n)
	reg = NormL1(.01)

	# Solve
	niter = n*100
	alg = SAGA(vrg, 10*randn(mt,N), 1/3, mt, prox_f=reg)
	iterate!(alg, niter)

	x_star = primiterates(alg)
	fval(x,y) = norm(x - x_star)
	progress = ShowFuncVal(fval, "|| x - x^*||")

	# Solve again for different params and solvers, all should give same result
	println("Prox-SAGA - VRGradient")
	niter = n*100
	nlogs = 10
	alg = SAGA(vrg, randn(mt,N), 1/10, mt, prox_f=reg)
	iterate!(alg, niter, (progress, ShowNewLine()), nlogs)
	@test isapprox(primiterates(alg), x_star)

	println("Prox-SAGA - LinearVRG - Importance")
	niter = n*100
	nlogs = 10
	alg = SAGA(
		vrg_lin, randn(mt,N), 1/4, mt, prox_f=reg, weights=0.5.+rand(mt,n))
	iterate!(alg, niter, (progress, ShowNewLine()), nlogs)
	@test isapprox(primiterates(alg), x_star)

	println("Prox-SAG - LinearVRG - Importance")
	niter = n*100
	nlogs = 10
	alg = SAG(
		vrg_lin, randn(mt,N), 1/4, mt, prox_f=reg, weights=0.5.+rand(mt,n))
	iterate!(alg, niter, (progress, ShowNewLine()), nlogs)
	@test isapprox(primiterates(alg), x_star)

	println("Prox-SVAG - LinearVRG")
	niter = n*100
	nlogs = 10
	alg = SVAG(vrg_lin, randn(mt,N), 1/4, n/20, mt, prox_f=reg)
	iterate!(alg, niter, (progress, ShowNewLine()), nlogs)
	@test isapprox(primiterates(alg), x_star)

	println("Prox-ASVAG - LinearVRG")
	niter = n*100
	nlogs = 10
	alg = ASVAG(vrg_lin, randn(mt,N), 1/4, mt, prox_f=reg)
	iterate!(alg, niter,
		(progress,
			ShowAlgState(a->a.innoweight[], "IW"),
			ShowNewLine()
		),
		nlogs)
	@test isapprox(primiterates(alg), x_star)

	println("Prox-SVRG - UniformVRG")
	stagelen = n/10
	niter = 50*n
	nlogs = 10
	alg = SVRG(vrg_uni, randn(mt,N), 1/8, stagelen, mt, prox_f=reg)
	iterate!(alg, niter, (progress, ShowNewLine()), nlogs)
	@test isapprox(primiterates(alg), x_star)

	println("Prox-LSVRG - UniformVRG - Importance")
	q = 10/n
	niter = n*50
	nlogs = 10
	alg = LSVRG(
		vrg_uni, randn(mt,N), 1/8, q, mt, prox_f=reg,weights=0.5.+rand(mt,n))
	iterate!(alg, niter, (progress, ShowNewLine()), nlogs)
	@test isapprox(primiterates(alg), x_star)

	println("Prox-ILSVRG - VRGradient - Importance")
	q = 10/n
	niter = n*50
	nlogs = 10
	alg = ILSVRG(
		vrg, randn(mt,N), 1/8, q, mt, prox_f=reg,weights=0.5.+rand(mt,n))
	iterate!(alg, niter, (progress, ShowNewLine()), nlogs)
	@test isapprox(primiterates(alg), x_star)

	println("Prox-QSAGA - LinearVRG - q = 10")
	q = 10
	niter = n*50
	nlogs = 10
	alg = QSAGA(
		vrg_lin, randn(mt,N), 1/8, q, mt, prox_f=reg)
	iterate!(alg, niter, (progress, ShowNewLine()), nlogs)
	@test isapprox(primiterates(alg), x_star)

	println("Prox-QSAGA - VRGadient - q = 10 - Without replacement")
	q = 10
	niter = n*50
	nlogs = 10
	alg = QSAGA(
		vrg, randn(mt,N), 1/8, q, mt, prox_f=reg, replace=false)
	iterate!(alg, niter, (progress, ShowNewLine()), nlogs)
	@test isapprox(primiterates(alg), x_star)
end

end


using Oceananigans
using Statistics
using Random
using UnicodePlots

function two_dimensional_turbulence(ν=1e-4; Nx=128, architecture=CPU(), advection=WENO5())
    Random.seed!(123)
    Lx = 2π

    model = NonhydrostaticModel(timestepper = :RungeKutta3,
                                architecture = architecture,
                                grid = RegularRectilinearGrid(size=(Nx, Nx), halo=(3, 3), extent=(Lx, Lx), topology=(Periodic, Periodic, Flat)),
                                advection = advection,
                                buoyancy = nothing,
                                tracers = nothing,
                                closure = IsotropicDiffusivity(ν=ν))

    u, v, w = model.velocities
    uᵢ = rand(size(u)...)
    vᵢ = rand(size(v)...)

    uᵢ .-= mean(uᵢ)
    vᵢ .-= mean(vᵢ)

    set!(model, u=uᵢ, v=vᵢ)

    progress(sim) = @info "(ν = $ν) iteration: $(sim.model.clock.iteration), time: $(round(Int, sim.model.clock.time))"

    simulation = Simulation(model, Δt=0.1, stop_time=30, iteration_interval=50, progress=progress)

    run!(simulation)

    return Array(interior(model.velocities.u))[:, :, 1]
end

u_reference = two_dimensional_turbulence()

fig = heatmap(u_reference)

display(fig)

viscosities = [1e-5, 2e-5, 5e-5, 8e-5, 9e-5, 1.1e-5, 1.2e-4, 2e-4, 5e-4, 1e-3]
errors = zeros(0)

for ν in viscosities
    u = two_dimensional_turbulence(ν)

    error = norm(u - u_reference) / norm(u_reference)

    push!(errors, error)
end

fig = scatterplot(viscosities, errors)

display(fig)


############################################################
# BinaryCrossEntropyLossLayer
############################################################
@defstruct BinaryCrossEntropyLossLayer Layer (
  name :: AbstractString = "binary-cross-entropy-loss",
  (weight :: AbstractFloat = 1.0, weight >= 0),
  (bottoms :: Vector{Symbol} = Symbol[], length(bottoms) == 2),
)
@characterize_layer(BinaryCrossEntropyLossLayer,
  has_loss => true,
  is_sink  => true,
  can_do_bp => true,
)

type BinaryCrossEntropyLossLayerState{T} <: LayerState
  layer :: BinaryCrossEntropyLossLayer
  loss  :: T
end

function setup(backend::Backend, layer::BinaryCrossEntropyLossLayer, inputs::Vector{Blob}, diffs::Vector{Blob})
  data_type = eltype(inputs[1])
  state = BinaryCrossEntropyLossLayerState(layer, convert(data_type, 0))
  return state
end

function shutdown(backend::Backend, state::BinaryCrossEntropyLossLayerState)
end

function forward(backend::CPUBackend, state::BinaryCrossEntropyLossLayerState, inputs::Vector{Blob})
  pred = vec(inputs[1].data)
  label = vec(inputs[2].data)
  loss = BLAS.dot(log(pred), label) + BLAS.dot(log1p(-pred), (1-label))

  num = get_num(inputs[1])
  state.loss = state.layer.weight * -loss/num
end

function backward(backend::CPUBackend, state::BinaryCrossEntropyLossLayerState, inputs::Vector{Blob}, diffs::Vector{Blob})
  diff = diffs[1]
  if isa(diff, CPUBlob)

    # Diffs is df/dloss
    # we want to multiply this by dloss/dpred and dl

    label = inputs[2].data
    pred  = reshape(inputs[1].data, size(label))

    # l = -w sum_i log(p_i) y_i + (log(1-p_i)(1-y_i)
    # dl/dp_i = -w ( y_i/p_i - ((1-y_i)(1-p_i) )
    # dl/dy_i = -w ( log(p_i) - (log(1-p_i) ) = -w log(p_i/(1-p_i))

    n = length(pred)
    a = convert(eltype(pred), -state.layer.weight/get_num(inputs[1]))

    erase!(diff) # is this correct? square-loss does it - should we allow for any incoming diff?
    dl_dpred = (label ./ pred) - ((1-label) ./ (1-pred)) # dloss/d?
    BLAS.axpy!(n, a, dl_dpred, 1, diff.data, 1)
  end
  diff = diffs[2]
  if isa(diff, CPUBlob)
    dl_dlabel = log(pred ./ (1-pred))
    erase!(diff) # is this correct? square-loss does it
    BLAS.axpy!(n, a, dl_dlabel, 1, diff.data, 1)
  end
end


"""
$(TYPEDEF)
$(TYPEDFIELDS)

"""
struct QuantNormal <: AbstractDistDelayed
    mu::NestedFloat
    sigma::NestedFloat
    q::NestedFloat
end

function quantnormal(mu::Float64, sigma::Float64, q::Float64)::Float64
    return round(normal(mu, sigma) / q) * q
end


using Luxor, Colors

# n-gons
@png begin
    for i in 23:-1:3
        sethue(i/20, 0.7, .99)
        ngonside(O, 75, i, 0, :fill)
        sethue("navy")
        ngonside(O, 75, i, 0, :stroke)
    end
end


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
using FwiFlow
# include("../ops.jl")
# Random.seed!(233)

function saturation(s0,dporodt,pt,perm,poro,qw,qo,muw,muo,sref,dt,h)
    saturation_ = load_op_and_grad("./build/libSaturation","saturation")
    s0,dporodt,pt,perm,poro,qw,qo,muw,muo,sref,dt,h = convert_to_tensor([s0,dporodt,pt,perm,poro,qw,qo,muw,muo,sref,dt,h], [Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64])
    saturation_(s0,dporodt,pt,perm,poro,qw,qo,muw,muo,sref,dt,h)
end

# TODO: specify your input parameters
m = 100
n = 10
h = 0.01
x = zeros(n, m)
y = zeros(n, m)
for i = 1:m
    for j = 1:n
        x[j, i] = h*i
        y[j, i] = h*j
    end
end
t = 3.0
s0 = @. (x^2 + y^2)/(1+x^2+y^2) * exp(-t)
dporodt = exp(t) * zeros(n, m)
pt = @. (x^2+y^2)
perm = rand(n, m)
poro = exp(t) * ones(n, m)
qw = ones(n, m)
qo = ones(n, m)
muw = 2.0
muo = 3.0
sref = s0
dt = 0.01
h = 0.1
u = sat_op2(s0,dporodt,pt,perm,poro,qw,qo,muw,muo,sref,dt,h)
u3 = sat_op(s0,pt,perm,poro,qw,qo,muw, muo,sref,dt,h)

# u3 = sat_op(s0,dporodt,pt,perm,poro,qw,qo,muw,muo,sref,dt,h)
sess = Session(); init(sess)
# @show run(sess, u)

@show run(sess, u3-u)

# uncomment it for testing gradients
error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
# s0 qw qo
function scalar_function(m)
    # return sum(saturation(s0,dporodt,pt,perm,poro,qw,qo,muw,muo,sref,dt,h)^2)
    return sum(sat_op2(m,dporodt,pt,perm,poro,qw,qo,muw,muo,sref,dt,h)^2)
    # return sum(sat_op(m,pt,perm,poro,qw,qo,muw,muo,sref,dt,h)^2)
end

# TODO: change `m_` and `v_` to appropriate values
m_ = constant(0.7*rand(n, m))
v_ = rand(n, m)

# m_ = constant(s0)
# v_ = rand(n, m)

y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


module PulseShapes
# provides unscaled pulse shapes for waveforms

import Base.tanh

"""
	constant(;pulse_length=0.0, sampling_rate=1.2e9)

Constant pulse shape for delays or square pulses.
"""
function constant(;pulse_length=0.0, sampling_rate=1.2e9)
	num_pts = round(UInt, pulse_length*sampling_rate)
	ones(Complex128, num_pts)
end

"""
	gaussian(;pulse_length=0.0, sampling_rate=1.2e9, cutoff=2.0)::Vector{Complex128}

Gaussian shaped pulse of a given pulse_length and sampling_rate going between -cutoff and cutoff σ wide.
Shape is shifted down to avoid finite cutoff step and start and end of pulse.
"""
function gaussian(;pulse_length=0.0, sampling_rate=1.2e9, cutoff=2.0)::Vector{Complex128}
	num_pts = round(UInt, pulse_length*sampling_rate)
	x_pts = linspace(-cutoff, cutoff, num_pts)
	# pull the pulse down so there is no bit step and the start/end of the pulse
	# i.e. find the shift such that the next point in the pulse would be zero
	x_step = x_pts[2] - x_pts[1]
	next_val = exp(-0.5 * (x_pts[1] - x_step)^2)
	shape = 1/(1-next_val) .* (exp.(-0.5 * (x_pts.^2)) - next_val)
	complex(shape)
end

@enum HALF_GAUSSIAN_DIRECTION HALF_GAUSSIAN_RISE HALF_GAUSSIAN_FALL

"""
	half_gaussian(;pulse_length=0.0, sampling_rate=1.2e9, cutoff=2.0, direction=HALF_GAUSSIAN_RISE)

A half gaussian pulse. Can be used as an excitation pulse but also used to round a square pulse.
"""

function half_gaussian(;pulse_length=0.0, sampling_rate=1.2e9, cutoff=2.0, direction::HALF_GAUSSIAN_DIRECTION=HALF_GAUSSIAN_RISE)
	num_pts = round(UInt, pulse_length*sampling_rate)
	x_pts = linspace(-cutoff, 0, num_pts)
	# pull the pulse down so there is no bit step and the start/end of the pulse
	# i.e. find the shift such that the next point in the pulse would be zero
	x_step = x_pts[2] - x_pts[1]
	next_val = exp(-0.5 * (x_pts[1] - x_step)^2)
	shape = 1/(1-next_val) .* (exp.(-0.5 * (x_pts.^2)) - next_val)
	if direction == HALF_GAUSSIAN_FALL
		reverse!(shape)
	end
	complex(shape)
end

"""
	drag(;pulse_length=0.0, sampling_rate=1.2e9, drag_scaling=1.0)::Vector{Complex128}

DRAG gaussian shape with -0.5 * (time derivative of gaussian) on the quadrature. The
derivative is scaled by `drag_scaling`.
"""
function drag(;pulse_length=0.0, sampling_rate=1.2e9, cutoff=2.0, drag_scaling=1.0)::Vector{Complex128}
	inphase = gaussian(;pulse_length=pulse_length, sampling_rate=sampling_rate, cutoff=cutoff)

	# d/dx of exp(-0.5 * x_pts.^2) = - x_pts * exp(-0.5 * x_pts.^2)
	# the time derivative needs a substitution rule from x -> t
	# where t is in units of sampling_rate
	# the pulse pulse_length in σ or x units is 2*cutoff so t = (num_points/(2*cutoff)) x
	num_pts = round(UInt, pulse_length*sampling_rate)
	x_pts = linspace(-cutoff, cutoff, num_pts)
	deriv_scale = 2*cutoff / num_pts
	quadrature = deriv_scale * drag_scaling * x_pts .* inphase
	complex(inphase + 1im*quadrature)
end

"""
	tanh(;pulse_length=0.0, sampling_rate=1.2e9, sigma=0.0, cutoff=2)::Vector{Complex128}

Square pulse with tanh rounded edges.
"""
function tanh(;pulse_length=0.0, sampling_rate=1.2e9, sigma=1e-9, cutoff=2)::Vector{Complex128}
	num_pts = round(UInt, pulse_length*sampling_rate)
	t = linspace(-pulse_length/2, pulse_length/2, num_pts)
	t₁ = -pulse_length/2 + cutoff*sigma
	t₂ = pulse_length/2 - cutoff*sigma
	shape = 0.5 * ( tanh.((t-t₁)/sigma) + tanh.((t₂-t)/sigma) )
	complex(shape)
end

"""
	exp_decay(;pulse_length=0.0, sampling_rate=1.2e9, sigma=1e-9, steady_state=0.5)::Vector{Complex128}

An exponentially decaying pulse to try and populate the cavity as quickly as possible.
But then don't overdrive it.
"""
function exp_decay(;pulse_length=0.0, sampling_rate=1.2e9, sigma=1e-9, steady_state=0.5)::Vector{Complex128}
	num_pts = round(UInt, pulse_length*sampling_rate)
	time_pts = (1.0/sampling_rate) * (0:num_pts-1)
	shape = steady_state + (1-steady_state)*exp(-time_pts / sigma)
	complex(shape)
end

"""
	CLEAR(;pulse_length=0.0, sampling_rate=1.2e9, sigma=0.0, steady_state=0.5, step_length=0.0, amp1=0.0, amp2=0.0)::Vector{Complex128}

Shape to quickly deplete the cavity at the end of a measurement.
Exponentially decaying measurement pulse followed by 2 steps of pulse_length `step_length` and amplitudes `amp1`, `amp2`.

Reference:

Rapid Driven Reset of a Qubit Readout Resonator

D. T. McClure, Hanhee Paik, L. S. Bishop, M. Steffen, Jerry M. Chow, and Jay M. Gambetta
Phys. Rev. Applied 5, 011001 – Published 27 January 2016 https://doi.org/10.1103/PhysRevApplied.5.011001
"""
function CLEAR(;pulse_length=0.0, sampling_rate=1.2e9, sigma=0.0, steady_state=0.5, step_length=0.0, amp1=0.0, amp2=0.0)::Vector{Complex128}
	meas_pulse = exp_decay(;pulse_length=pulse_length-2*step_length, sampling_rate=sampling_rate, sigma=sigma, steady_state=steady_state)
	num_pts_step = round(UInt, step_length*sampling_rate)
	clear_pulse1 = amp1 * ones(Complex128, num_pts_step)
	clear_pulse2 = amp2 * ones(Complex128, num_pts_step)
	[meas_pulse; clear_pulse1; clear_pulse2]
end

"""
	arb_axis_drag(;pulse_length=0.0, sampling_rate=1.2e9, nut_freq=10e6, rot_angle=0.0, polar_angle=0.0, drag_scaling=0.0)

Single-qubit arbitrary axis pulse implemented with phase ramping and frame change.
Parameters
    nutFreq: effective nutation frequency per unit of drive amplitude (Hz)
    rotAngle : effective rotation rotAngle (radians)
    polarAngle : polar angle of rotation axis (radians)
    aziAngle : azimuthal (radians)
"""
function arb_axis_drag(;pulse_length=0.0, sampling_rate=1.2e9, nut_freq=10e6, rot_angle=0.0, Θ=0.0, drag_scaling=0.0)::Vector{Complex128}
	if pulse_length > 0
		# start from a gaussian shaped pulse
		gauss_pulse = gaussian(pulse_length=pulse_length, sampling_rate=sampling_rate)
		# scale to achieve to the desired rotation
		cal_scale = (rot_angle/2/pi)*sampling_rate/sum(gauss_pulse)
		# calculate the phase ramp steps to achieve the desired Z component to the rotation axis
		phase_steps = -2π*cos(Θ)*cal_scale*gauss_pulse/sampling_rate
		# Calculate Z DRAG correction to phase steps
		# β is a conversion between XY drag scaling and Z drag scaling
		β = drag_scaling/sampling_rate
		instantaneous_detuning = β * (2π*cal_scale*sin(Θ)*gauss_pulse).^2
		phase_steps += instantaneous_detuning*(1/sampling_rate)
		#center phase ramp around the middle of the pulse time steps
		phase_ramp = cumsum(phase_steps) - phase_steps/2
		shape = (1/nut_freq)*sin(Θ)*cal_scale*gauss_pulse.*exp(1im*phase_ramp)
	elseif abs(Θ) <1e-10
		#Otherwise assume we have a zero-length Z rotation
		shape = Vector{Complex128}()
	else
		error("Non-zero transverse rotation with zero-length pulse.")
	end
	complex(shape)
end
end


import StringUtils

function lines_to_bigints(a::Array{String,1})
    result::Array{BigInt,1} = []
    for s in a
        big = parse(BigInt, chomp(s))
        push!(result, big)
    end
    return result
end

function main()
    if length(ARGS) > 0
        filename = ARGS[1]
        a = readlines(filename)
        m = lines_to_bigints(a)
        s = string(sum(m))
        @printf("The first 10 digits are %s\n", s[1:10])
    else
        error("this program needs a single numeric argument")
    end
end

main()


# * Term symbols

struct Term
    L::HalfInteger
    S::HalfInteger
    parity::Parity
end
Term(L::Real, S::Real, parity::Integer) =
    Term(convert(HalfInteger, L), convert(HalfInteger, S), convert(Parity, parity))

function Base.parse(::Type{Term}, s::AbstractString)
    m = match(r"([0-9]+)([A-Z]|\[[0-9/]+\])([oe ]{0,1})", s)
    isnothing(m) && throw(ArgumentError("Invalid term string $s"))
    L = lowercase(m[2])
    L = if L[1] == '['
        L = strip(L, ['[',']'])
        if occursin("/", L)
            Ls = split(L, "/")
            length(Ls) == 2 && length(Ls[1]) > 0 && length(Ls[2]) > 0 ||
                throw(ArgumentError("Invalid term string $(s)"))
            Rational(parse(Int, Ls[1]),parse(Int, Ls[2]))
        else
            parse(Int, L)
        end
    else
        findfirst(L, spectroscopic)[1]-1
    end
    denominator(L) ∈ [1,2] || throw(ArgumentError("L must be integer or half-integer"))
    S = (parse(Int, m[1]) - 1)//2
    Term(L, S, m[3] == "o" ? p"odd" : p"even")
end

macro T_str(s::AbstractString)
    parse(Term, s)
end

Base.zero(::Type{Term}) = T"1S"

multiplicity(t::Term) = convert(Int, 2t.S + 1)
weight(t::Term) = (2t.L + 1) * multiplicity(t)

import Base.==
==(t1::Term, t2::Term) = ((t1.L == t2.L) && (t1.S == t2.S) && (t1.parity == t2.parity))

import Base.<
<(t1::Term, t2::Term) = ((t1.S < t2.S) || (t1.S == t2.S) && (t1.L < t2.L)
                         || (t1.S == t2.S) && (t1.L == t2.L) && (t1.parity < t2.parity))
import Base.isless
isless(t1::Term, t2::Term) = (t1 < t2)

Base.hash(t::Term) = hash((t.L,t.S,t.parity))

include("xu2006.jl")

# This function calculates the term symbol for a given orbital ℓʷ
function xu_terms(ℓ::Int, w::Int, p::Parity)
    ts = map(((w//2 - floor(Int, w//2)):w//2)) do S
        S′ = 2S |> Int
        map(L -> repeat([Term(L,S,p)], Xu.X(w, ℓ, S′, L)), 0:w*ℓ)
    end
    vcat(vcat(ts...)...)
end

function terms(orb::Orbital, occ::Int)
    ℓ = orb.ℓ
    g = degeneracy(orb)
    occ > g && throw(ArgumentError("Invalid occupancy $occ for $orb with degeneracy $g"))
    (occ > g/2 && occ != g) && (occ = g - occ)

    p = parity(orb)^occ
    if occ == 1
        [Term(ℓ,1//2,p)] # Single electron
    elseif ℓ == 0 && occ == 2 || occ == degeneracy(orb) || occ == 0
        [Term(0,0,p"even")] # Filled ℓ shell
    else
        xu_terms(ℓ, occ, p) # All other cases
    end
end

function terms(config::Configuration{O}) where {O<:AbstractOrbital}
    ts = map(config) do (orb,occ,state)
        terms(orb,occ)
    end

    final_terms(ts)
end

"""
    count_terms(orb, occ, term)

Count how many times `term` occurs among the valid terms of
`orb`^`occ`. Example:

```
count_terms(o"1s", 2, T"1S") # 1
```
"""
function count_terms(orb::Orbital, occ::Int, term::Term)
    ℓ = orb.ℓ
    g = degeneracy(orb)
    occ > g && throw(ArgumentError("Invalid occupancy $occ for $orb with degeneracy $g"))
    (occ > g/2 && occ != g) && (occ = g - occ)

    p = parity(orb)^occ
    if occ == 1
        term == Term(ℓ,1//2,p) ? 1 : 0
    elseif ℓ == 0 && occ == 2 || occ == degeneracy(orb) || occ == 0
        term == Term(0,0,p"even") ? 1 : 0
    else
        S′ = convert(Int, 2term.S)
        Xu.X(occ, orb.ℓ, S′, convert(Int, term.L))
    end
end

function write_L(io::IO, term::Term)
    if isinteger(term.L)
        write(io, uppercase(spectroscopic_label(convert(Int, term.L))))
    else
        write(io, "[$(numerator(term.L))/$(denominator(term.L))]")
    end
end

function Base.show(io::IO, term::Term)
    write(io, to_superscript(multiplicity(term)))
    write_L(io, term)
    write(io, to_superscript(term.parity))
end

# * Intermediate terms, seniority

struct IntermediateTerm
    term::Term
    seniority::Int
    function IntermediateTerm(term::Term, seniority::Int)
        iseven(multiplicity(term)) ⊻ iseven(seniority) ||
            throw(ArgumentError("Invalid seniority $(seniority) for term $(term)"))
        new(term, seniority)
    end
end

function Base.show(io::IO, iterm::IntermediateTerm)
    # This is the notation by Giulio Racah, p.377:
    # - Racah, G. (1943). Theory of complex spectra. iii. Physical Review,
    #   63(9-10), 367–382. http://dx.doi.org/10.1103/physrev.63.367
    write(io, to_subscript(iterm.seniority))
    show(io, iterm.term)
end

Base.isless(a::IntermediateTerm, b::IntermediateTerm) =
    a.seniority < b.seniority ||
    a.seniority == b.seniority && a.term < b.term

function intermediate_terms(orb::Orbital, w::Int=one(Int))
    ts = terms(orb, w)
    its = map(unique(ts)) do t
        its = IntermediateTerm[]
        previously_seen = 0
        # The seniority number is defined as the minimum occupancy
        # number ν ∈ n:-2:0 for which the term first appears, e.g. the
        # ²D term first occurs in the d¹ configuration, then twice in
        # the d³ configuration (which will then have the terms ₁²D and
        # ₃²D).
        #
        # We have to loop in reverse, since odd occupation numbers
        # should go from 1 and even from 0.
        for ν ∈ reverse(w:-2:0)
            nn = count_terms(orb, ν, t) - previously_seen
            previously_seen += nn
            append!(its, repeat([IntermediateTerm(t, ν)], nn))
        end
        its
    end
    sort(vcat(its...))
end

function intermediate_terms(config::Configuration)
    map(config) do (orb,occ,state)
        intermediate_terms(orb,occ)
    end
end

export Term, @T_str, multiplicity, weight, terms, count_terms, IntermediateTerm, intermediate_terms


# Author: Lukasz Bratos

include("diff.jl")

f(x) = exp(x)
g(x) = x^2.0 * sin(x)

for i in [5, 10, 15]
    Diff.rysujNnfx(f, 0.0, 1.0, i, "plot5a$i")
    Diff.rysujNnfx(g, -1.0, 1.0, i, "plot5b$i")
end


# ---
# title: 292. Nim Game
# id: problem292
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Brainteaser, Minimax
# link: <https://leetcode.com/problems/nim-game/description/>
# hidden: true
# ---
#
# You are playing the following Nim Game with your friend:
#
#   * Initially, there is a heap of stones on the table.
#   * You and your friend will alternate taking turns, and **you go first**.
#   * On each turn, the person whose turn it is will remove 1 to 3 stones from the heap.
#   * The one who removes the last stone is the winner.
#
# Given `n`, the number of stones in the heap, return `true` _if you can win the
# game assuming both you and your friend play optimally, otherwise
# return_`false`.
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 4
#     Output: false
#     Explanation: These are the possible outcomes:
#     1. You remove 1 stone. Your friend removes 3 stones, including the last stone. Your friend wins.
#     2. You remove 2 stones. Your friend removes 2 stones, including the last stone. Your friend wins.
#     3. You remove 3 stones. Your friend removes the last stone. Your friend wins.
#     In all outcomes, your friend wins.
#
#
# **Example 2:**
#
#
#
#     Input: n = 1
#     Output: true
#
#
# **Example 3:**
#
#
#
#     Input: n = 2
#     Output: true
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 231 - 1`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


#! /usr/bin/julia

# Rosetta Code, Order disjoint list items

function order_disjoint{T<:AbstractArray}(m::T, n::T)
    rlen = length(n)
    rdis = zeros(Int, rlen)
    for (i, e) in enumerate(n)
        j = findfirst(m, e)
        while j in rdis && j != 0
            j = findnext(m, e, j+1)
        end
        rdis[i] = j
    end
    if 0 in rdis
        throw(DomainError())
    end
    sort!(rdis)
    p = copy(m)
    p[rdis] = n
    return p
end

testm = {["the", "cat", "sat", "on", "the", "mat"],
         ["the", "cat", "sat", "on", "the", "mat"],
         ["A", "B", "C", "A", "B", "C", "A", "B", "C"],
         ["A", "B", "C", "A", "B", "D", "A", "B", "E"],
         ["A", "B"],
         ["A", "B"],
         ["A", "B", "B", "A"],
         }

testn = {["mat", "cat"],
         ["cat", "mat"],
         ["C", "A", "C", "A"],
         ["E", "A", "D", "A"],
         ["B"],
         ["B", "A"],
         ["B", "A"],
         }

for i in 1:length(testm)
    m = join(testm[i], " ")
    n = join(testn[i], " ")
    p = join(order_disjoint(testm[i], testn[i]), " ")
    println("    (", m, ", ", n, ") => ", p)
end


# ---
# title: 985. Sum of Even Numbers After Queries
# id: problem985
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/sum-of-even-numbers-after-queries/description/>
# hidden: true
# ---
#
# We have an array `A` of integers, and an array `queries` of queries.
#
# For the `i`-th query `val = queries[i][0], index = queries[i][1]`, we add val
# to `A[index]`.  Then, the answer to the `i`-th query is the sum of the even
# values of `A`.
#
# _(Here, the given`index = queries[i][1]` is a 0-based index, and each query
# permanently modifies the array `A`.)_
#
# Return the answer to all queries.  Your `answer` array should have `answer[i]`
# as the answer to the `i`-th query.
#
#
#
# **Example 1:**
#
#
#
#     Input: A = [1,2,3,4], queries = [[1,0],[-3,1],[-4,0],[2,3]]
#     Output: [8,6,2,4]
#     Explanation:
#     At the beginning, the array is [1,2,3,4].
#     After adding 1 to A[0], the array is [2,2,3,4], and the sum of even values is 2 + 2 + 4 = 8.
#     After adding -3 to A[1], the array is [2,-1,3,4], and the sum of even values is 2 + 4 = 6.
#     After adding -4 to A[0], the array is [-2,-1,3,4], and the sum of even values is -2 + 4 = 2.
#     After adding 2 to A[3], the array is [-2,-1,3,6], and the sum of even values is -2 + 6 = 4.
#
#
#
#
# **Note:**
#
#   1. `1 <= A.length <= 10000`
#   2. `-10000 <= A[i] <= 10000`
#   3. `1 <= queries.length <= 10000`
#   4. `-10000 <= queries[i][0] <= 10000`
#   5. `0 <= queries[i][1] < A.length`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


#=

    particles.jl

    Implements the type `Particle` used in ABC SMC.

    Author:     Alexander P. Browning
                ======================
                School of Mathematical Sciences
                Queensland University of Technology
                ======================
                ap.browning@icloud.com
                alexbrowning.me

=#

"""
    Particle(θ::Vector,d::Number,w::Number)

SMC particle with location `θ`, distance `d`, and weight `w`.
"""
mutable struct Particle
    θ::Vector   # Location
    d::Number   # Discrepancy
    w::Number   # Weight
end
Particle(θ,d) = Particle(θ,d,1.0)

##############################################################
## Extend methods to particles
##############################################################

import Base: sort!
import Distributions: insupport
import StatsBase: sample, weights
import Statistics: cov, mean

weights(P::Vector{Particle})  = weights([p.w for p in P])
distances(P::Vector{Particle}) = [p.d for p in P]
locations(P::Vector{Particle}) = hcat([p.θ for p in P]...)
insupport(d::Distribution,p::Particle) = insupport(d,p.θ)

"""
    sort!(p::Vector{Particle})

Sort a vector of `Particle`s, `P` by distance `d`.
"""
sort!(P::Vector{Particle}) = permute!(P,sortperm(distances(P)))

"""
    cov(p::Vector{Particle})

Weighted covariance matrix of particles `P`.
"""
cov(P::Vector{Particle}) = cov(permutedims(locations(P)),weights(P))

mean(P::Vector{Particle}) = mean(locations(P),weights(P),dims=2)[:]

"""
    sample(prior,dist,N=1)

Sample particles from the `prior` distribution and attach distance using function `dist`.
"""
function StatsBase.sample(prior::Distribution,dist::Function,w::Number=1)
    θ = rand(prior)
    d = dist(θ)
    Particle(θ,d,w)
end
StatsBase.sample(prior::Distribution,dist::Function,n::Int) = [sample(prior,dist,1/n) for i in 1:n]

"""
    copy(p::Particle)

Copies a particle.
"""
Base.copy(p::Particle) = Particle(p.θ,p.d,p.w)

"""
    findmin(p::Particle)

Returns (pᵢ,i) where pᵢ is the particle with the smallest discrepancy.
"""
function Base.findmin(p::Vector{Particle})
    D = [p[i].d for i = 1:length(p)]
    idx = findmin(D)[2]
    p[idx],idx
end

"""
    minimum(p::Particle)

Finds the particle with the smallest discrepancy.
"""
Base.minimum(p::Vector{Particle}) = findmin(p)[1]

"""
    findmax(p::Particle)

Returns (pᵢ,i) where pᵢ is the particle with the largest discrepancy.
"""
function findmax(p::Vector{Particle})
    D = [p[i].d for i = 1:length(p)]
    idx = findmin(-D)[2]
    p[idx],idx
end

"""
    maximum(p::Particle)

Finds the particle with the largest discrepancy.
"""
Base.maximum(p::Vector{Particle}) = findmax(p)[1]

Base.ndims(P::Vector{Particle}) = ndims(P[1])
Base.ndims(p::Particle) = length(p.θ)

ess(P::Vector{Particle}) = 1 / sum(weights(P).^2)


using DiagonalSBP
using PGFPlotsX: @pgf, Axis, Plot, pgfsave, Table, LogLogAxis
using ForwardDiff: derivative
using OrdinaryDiffEq: ODEProblem, solve, Tsit5
using SparseArrays: sparsevec
using Printf: @sprintf

# This uses the fourth-order, low-storage, Runge--Kutta scheme of
# Carpenter and Kennedy (1994) (in their notation (5,4) 2N-Storage RK
#                               scheme).
#
# @TECHREPORT{CarpenterKennedy1994,
#             author = {M.~H. Carpenter and C.~A. Kennedy},
#             title = {Fourth-order {2N-storage} {Runge-Kutta}
#                      schemes},
#             institution = {National Aeronautics and Space
#                            Administration},
#             year = {1994},
#             number = {NASA TM-109112},
#             address = {Langley Research
#                        Center, Hampton, VA},
#            }
function timestep!(q, f!, p, dt, (t0, t1), Δq = similar(q), Δq2 = similar(q))
    T = eltype(q)

    RKA = (
        T(0),
        T(-567301805773 // 1357537059087),
        T(-2404267990393 // 2016746695238),
        T(-3550918686646 // 2091501179385),
        T(-1275806237668 // 842570457699),
    )

    RKB = (
        T(1432997174477 // 9575080441755),
        T(5161836677717 // 13612068292357),
        T(1720146321549 // 2090206949498),
        T(3134564353537 // 4481467310338),
        T(2277821191437 // 14882151754819),
    )

    RKC = (
        T(0),
        T(1432997174477 // 9575080441755),
        T(2526269341429 // 6820363962896),
        T(2006345519317 // 3224310063776),
        T(2802321613138 // 2924317926251),
    )

    nstep = ceil(Int, (t1 - t0) / dt)
    dt = (t1 - t0) / nstep

    fill!(Δq, 0)
    fill!(Δq2, 0)
    for step in 1:nstep
        t = t0 + (step - 1) * dt
        for s in 1:length(RKA)
            f!(Δq2, q, p, t + RKC[s] * dt)
            Δq .+= Δq2
            q .+= RKB[s] * dt * Δq
            Δq .*= RKA[s % length(RKA) + 1]
        end
    end

    nothing
end

function bisection(F, a::T, b::T, tol = 100 * eps(T)) where T
    (a, b) = a < b ? (a, b) : (b, a)

    Fa = F(a)
    Fb = F(b)
    if abs(Fa) < abs(Fb)
        abs(Fa) < eps(T) && return a
    else
        abs(Fb) < eps(T) && return b
    end
    @assert Fa * Fb < 0

    V = (a + b) / 2
    while b - a > tol
        FV = F(V)
        if FV == 0
            break
        elseif FV * Fa < 0
            b = V
        else
            a = V
        end
        V = (a + b) / 2
    end

    return V
end

function semianalytic_solution(u0, v0, σ0, tend::T, F) where T

    # Characteristics propagating into interface
    w_m(x, t) = v0(x - t) - σ0(x - t)
    w_p(x, t) = v0(x + t) + σ0(x + t) # zero

    # Find V by root finding
    η(t) = w_p(0, t) - w_m(0, t)
    g(V, t) = 2F(V) + V - η(t)
    V(t) = t > 0 ? bisection(V -> g(V, t), zero(T), η(t)) : zero(t)

    # Set Characteristics propagating out of interface
    Q_m(t) = w_m(0, t) + 2F(V(t))
    Q_p(t) = w_p(0, t) - 2F(V(t))

    # Integrate variables in order to compute displacement
    ∫Q_m_ = solve(ODEProblem((_, _, t) -> Q_m(t), zero(T), (zero(T), tend)),
                  Tsit5(),
                  dtmax = 0.1,
                  abstol=1e-12,
                  reltol=1e-12)
    ∫Q_m(t) = t > 0 ? ∫Q_m_(t) : zero(T)
    ∫Q_p_ = solve(ODEProblem((_, _, t) -> Q_p(t), zero(T), (zero(T), tend)),
                  Tsit5(),
                  dtmax = 0.1,
                  abstol=1e-12,
                  reltol=1e-12)
    ∫Q_p(t) = t > 0 ? ∫Q_p_(t) : zero(T)

    ∫q_m(x, t) = ∫Q_m(t + x)
    ∫q_p(x, t) = ∫Q_p(t - x)

    q_m(x, t) = Q_m(t + x)
    q_p(x, t) = Q_p(t - x)

    u_m(x, t) = u0(x - t) + ∫q_m(x, t) / 2
    v_m(x, t) = (q_m(x, t) + w_m(x, t)) / 2
    σ_m(x, t) = (q_m(x, t) - w_m(x, t)) / 2

    u_p(x, t) = u0(x + t) + ∫q_p(x, t) / 2
    v_p(x, t) =  (q_p(x, t) + w_p(x, t)) / 2
    σ_p(x, t) = -(q_p(x, t) - w_p(x, t)) / 2

    return u_m, u_p, v_m, v_p, σ_m, σ_p
end

function noncharacteristic_operator!(dq, q, D, b0, bN, HI, F)
    Nq = size(D, 1)
    rngs = ntuple(j-> ((j-1) * Nq + 1) : (j * Nq), 4)
    v_m, v_p, u_m, u_p = view.(Ref(q), rngs)

    V_m = v_p[1] - v_m[Nq]
    us_m, us_p = u_m[Nq], u_p[1]
    τs_m = F(V_m)
    τs_p = -τs_m

    operator!(dq, q, D, b0, bN, HI, τs_m, τs_p, us_m, us_p)
    nothing
end

function characteristic_operator!(dq, q, D, b0, bN, HI, F, γ)
    Nq = size(D, 1)
    rngs = ntuple(j-> ((j-1) * Nq + 1) : (j * Nq), 4)
    v_m, v_p, u_m, u_p = view.(Ref(q), rngs)

    us_m, us_p = q[4Nq + 1], q[4Nq + 2]

    # Normal
    n0, nN = -1, 1

    τ_m = nN * bN' * u_m + γ * (us_m - u_m[Nq])
    τ_p = n0 * b0' * u_p + γ * (us_p - u_p[ 1])

    w_m = v_m[Nq] - τ_m
    w_p = v_p[ 1] - τ_p

    η_m = w_p - w_m
    g(V_m) = 2F(V_m) + V_m - η_m
    V_m = bisection(g, zero(η_m), η_m)
    τs_m = F(V_m)
    τs_p = -τs_m

    operator!(dq, q, D, b0, bN, HI, τs_m, τs_p, us_m, us_p)

    dq[4Nq + 1] = w_m + τs_m
    dq[4Nq + 2] = w_p + τs_p

    nothing
end

function operator!(dq, q, D, b0, bN, HI, τs_m, τs_p, us_m, us_p)
    Nq = size(D, 1)
    rngs = ntuple(j-> ((j-1) * Nq + 1) : (j * Nq), 4)
    dv_m, dv_p, du_m, du_p = view.(Ref(dq), rngs)
    v_m, v_p, u_m, u_p = view.(Ref(q), rngs)

    # Set the displacement rates
    du_m .= v_m
    du_p .= v_p

    # Set the velocity rates
    dv_m .= D * u_m
    dv_p .= D * u_p

    # Normal
    n0, nN = -1, 1

    # Set the Neumann bcs
    dv_m[ 1] -= n0 * HI[ 1, 1] * b0' * u_m
    dv_p[Nq] -= nN * HI[Nq,Nq] * bN' * u_p

    # Set the interface values
    dv_m[Nq] += HI[Nq, Nq] * (τs_m - nN * bN' * u_m)
    dv_p[ 1] += HI[ 1,  1] * (τs_p - n0 * b0' * u_p)
    dv_m .-= nN * HI * bN * (us_m - u_m[Nq])
    dv_p .-= n0 * HI * b0 * (us_p - u_p[ 1])
    nothing
end

function convergence_test(p, Ns, β, cfl_c, cfl_nc, output = false)
    T = Float64

    x_m = range(T(-1), stop = 0, length=100)
    x_p = range(T(0), stop = 1, length=100)

    # Pulse center
    μ = -1 / 2

    # Pulse width
    w = 1 / 15

    u0(x) = exp(-((x - μ) / w)^2)
    v0(x) = derivative(t->u0(x-t), 0)
    σ0(x) = derivative(x->u0(x), x)

    tend::T = 1

    F(V) = β * asinh(V)

    ue_m, ue_p, ve_m, ve_p, σe_m, σe_p = semianalytic_solution(u0, v0, σ0, tend, F)

    # storage for error
    err_nc  = zeros(length(Ns))
    err_c   = zeros(length(Ns))

    for (iter, N) = enumerate(Ns)
        if output
            @info @sprintf " level %d with N = %4d" iter N
        end
        Nq = N + 1

        # Setup the SBP operators
        (D, S0, SN, HI, H, x_p) = DiagonalSBP.D2(p, N; xc = (0, 1))
        b0, bN = S0[1, :], SN[end, :]
        θ_R = DiagonalSBP.D2_remainder_parameters(p).θ_R
        θ_H = DiagonalSBP.D2_remainder_parameters(p).θ_H
        γ = N * (1 / θ_R + 1 / θ_H)

        # Create the grid
        x_m = x_p .- 1
        x = [x_m; x_p]

        # Create the solution
        q0 = zeros(T, 4Nq)
        q0[1:2Nq] = v0.(x)
        q0[2Nq .+ (1:2Nq)] = u0.(x)
        qnc = copy(q0)

        fnc!(Δq, q, _, _) = noncharacteristic_operator!(Δq, q, D, b0, bN, HI, F)

        qc = [q0; u0(0); u0(0)]
        fc!(Δq, q, _, _) = characteristic_operator!(Δq, q, D, b0, bN, HI, F, γ)

        h = 1 / N
        dtnc = cfl_nc * h
        dtc = cfl_c * h

        timestep!(qnc, fnc!, nothing, dtnc, (T(0), tend))
        timestep!(qc, fc!, nothing, dtc, (T(0), tend))

        rngs = ntuple(j-> ((j-1) * Nq + 1) : (j * Nq), 4)
        _, _, unc_m, unc_p = view.(Ref(qnc), rngs)
        _, _, uc_m, uc_p = view.(Ref(qc), rngs)
        err_c[iter] = sqrt(sum(H * ((uc_m - ue_m.(x_m, tend)).^2 +
                                    (uc_p - ue_p.(x_p, tend)).^2)))
        err_nc[iter] = sqrt(sum(H * ((unc_m - ue_m.(x_m, tend)).^2 +
                                     (unc_p - ue_p.(x_p, tend)).^2)))
        if output
            if iter > 1
                rate = ((log(err_nc[iter-1]) - log(err_nc[iter])) /
                        (log(Ns[iter]) - log(Ns[iter-1])))
                @info @sprintf """ Non-Characteristic
                error = %.2e
                rate  = %.2e""" err_nc[iter] rate
            else
                @info @sprintf """ Non-Characteristic
                error = %.2e""" err_nc[iter]
            end
            if iter > 1
                rate = ((log(err_c[iter-1]) - log(err_c[iter])) /
                        (log(Ns[iter]) - log(Ns[iter-1])))
                @info @sprintf """ Characteristic
                error = %.2e
                rate  = %.2e""" err_c[iter] rate
            else
                @info @sprintf """ Characteristic
                error = %.2e""" err_c[iter]
            end
        end
    end
    return err_c, err_nc
end

let
    Ns = 17 * 2 .^ (0:5)

    hs = 1 ./ Ns
    tri_x = [hs[end-1], hs[end], hs[end-1], hs[end-1]]
    tri_y = [hs[end-1], hs[end], hs[end  ], hs[end-1]] ./ hs[end-1]

    cfl_nc = Dict()
    cfl_nc[2, 1] = 1 // 2
    cfl_nc[4, 1] = 1 // 2
    cfl_nc[6, 1] = 1 // 2

    cfl_nc[2, 2] = 1 // 2
    cfl_nc[4, 2] = 1 // 4
    cfl_nc[6, 2] = 1 // 4

    cfl_nc[2, 32] = 1 // 32
    cfl_nc[4, 32] = 1 // 64
    cfl_nc[6, 32] = 1 // 64

    cfl_nc[2, 64] = 1 // 64
    cfl_nc[4, 64] = 1 // 128
    cfl_nc[6, 64] = 1 // 128

    cfl_nc[2, 128] = 1 // 128
    cfl_nc[4, 128] = 1 // 256
    cfl_nc[6, 128] = 1 // 256
    for β in (32, 64, 128)
        @show β

        C2, NC2 = convergence_test(2, Ns, β, 1 // 2, cfl_nc[2, β])
        C4, NC4 = convergence_test(4, Ns, β, 1 // 2, cfl_nc[4, β])
        C6, NC6 = convergence_test(6, Ns, β, 1 // 4, cfl_nc[6, β])

        @pgf a = LogLogAxis(
                            {
                             ymax = 1,
                             ymin = 1e-11,
                             xmin = hs[end]/2,
                             xmax = 2hs[1],
                             xlabel = "\$h\$",
                             ylabel = "\$\\|\\Delta u\\|_{H}\$",
                            }
                           )
        tc2 = 0.6*NC2[end-1] * tri_y.^2
        @pgf push!(a, Plot(
                           {
                            color = "black",
                           },
                           Table(tri_x, tc2)
                          )
                  )
        tc4 = 0.6*NC4[end-1] * tri_y.^4
        @pgf push!(a, Plot(
                           {
                            color = "black",
                           },
                           Table(tri_x, 0.6*NC4[end-1] * tri_y.^4)
                          )
                  )
        tc6 = 0.6*NC6[end-1] * tri_y.^6
        @pgf push!(a, Plot(
                           {
                            color = "black",
                           },
                           Table(tri_x, 0.6*NC6[end-1] * tri_y.^5)
                          )
                  )
        #=
        @pgf push!(a,
        [@sprintf("""
        \\node[anchor=north] () at (axis cs:%e, %e){\$s = 2\$};
        \\node[anchor=north] () at (axis cs:%e, %e){\$s = 4\$};
        \\node[anchor=north] () at (axis cs:%e, %e){\$s = 5\$};
        """,
        (hs[end-1] + hs[end]) / 2, tc2[2],
        (hs[end-1] + hs[end]) / 2, tc4[2],
        (hs[end-1] + hs[end]) / 2, tc6[2],
        )])
        =#
        @pgf push!(a,
                   [@sprintf("""
                             \\node[anchor=west] () at (axis cs:%e, %e){\$2\$};
                             \\node[anchor=west] () at (axis cs:%e, %e){\$4\$};
                             \\node[anchor=west] () at (axis cs:%e, %e){\$5\$};
                             """,
                             hs[end-1], exp((log(tc2[1]) + log(tc2[2]))/2),
                             hs[end-1], exp((log(tc4[1]) + log(tc4[2]))/2),
                             hs[end-1], exp((log(tc6[1]) + log(tc6[2]))/2),
                            )])

        @pgf push!(a, Plot(
                           {
                            color = "red",
                            mark = "o",
                           },
                           Table(hs, C2)
                          )
                  )
        @pgf push!(a, Plot(
                           {
                            color = "red",
                            mark = "x",
                           },
                           Table(hs, NC2)
                          )
                  )
        @pgf push!(a, Plot(
                           {
                            color = "blue",
                            mark = "o",
                           },
                           Table(hs, C4)
                          )
                  )
        @pgf push!(a, Plot(
                           {
                            color = "blue",
                            mark = "x",
                           },
                           Table(hs, NC4)
                          )
                  )
        @pgf push!(a, Plot(
                           {
                            color = "green",
                            mark = "o",
                           },
                           Table(hs, C6)
                          )
                  )
        @pgf push!(a, Plot(
                           {
                            color = "green",
                            mark = "x",
                           },
                           Table(hs, NC6)
                          )
                  )
        pgfsave("convergence_interface_$β.tex", a)
        pgfsave("convergence_interface_$β.pdf", a)
    end
end