push!(LOAD_PATH, joinpath(@__DIR__, "..", ".."))

import ClimaCore.Geometry, LinearAlgebra, UnPack
import ClimaCore:
    Fields,
    Domains,
    Topologies,
    Meshes,
    DataLayouts,
    Operators,
    Geometry,
    Spaces

using OrdinaryDiffEq:
    OrdinaryDiffEq,
    ODEProblem,
    ODEFunction,
    solve,
    SSPRK33,
    Rosenbrock23,
    ImplicitEuler

using Logging: global_logger
using TerminalLoggers: TerminalLogger
using LinearAlgebra
global_logger(TerminalLogger())

const FT = Float64

# https://github.com/CliMA/CLIMAParameters.jl/blob/master/src/Planet/planet_parameters.jl#L5
const MSLP = 1e5 # mean sea level pressure
const grav = 9.8 # gravitational constant
const R_d = 287.058 # R dry (gas constant / mol mass dry air)
const γ = 1.4 # heat capacity ratio
const C_p = R_d * γ / (γ - 1) # heat capacity at constant pressure
const C_v = R_d / (γ - 1) # heat capacit at constant volume
const R_m = R_d # moist R, assumed to be dry

domain = Domains.IntervalDomain(
    Geometry.ZPoint{FT}(0.0),
    Geometry.ZPoint{FT}(30e3),
    boundary_tags = (:bottom, :top),
)
#mesh = Meshes.IntervalMesh(domain, Meshes.ExponentialStretching(7.5e3); nelems = 30)
mesh = Meshes.IntervalMesh(domain; nelems = 30)

cspace = Spaces.CenterFiniteDifferenceSpace(mesh)
fspace = Spaces.FaceFiniteDifferenceSpace(cspace)

# https://github.com/CliMA/Thermodynamics.jl/blob/main/src/TemperatureProfiles.jl#L115-L155
# https://clima.github.io/Thermodynamics.jl/dev/TemperatureProfiles/#DecayingTemperatureProfile
function decaying_temperature_profile(z; T_virt_surf = 280.0, T_min_ref = 230.0)
    # Scale height for surface temperature
    H_sfc = R_d * T_virt_surf / grav
    H_t = H_sfc

    z′ = z / H_t
    tanh_z′ = tanh(z′)

    ΔTv = T_virt_surf - T_min_ref
    Tv = T_virt_surf - ΔTv * tanh_z′

    ΔTv′ = ΔTv / T_virt_surf
    p =
        MSLP * exp(
            (
                -H_t *
                (z′ + ΔTv′ * (log(1 - ΔTv′ * tanh_z′) - log(1 + tanh_z′) + z′))
            ) / (H_sfc * (1 - ΔTv′^2)),
        )
    ρ = p / (R_d * Tv)
    ρθ = ρ * Tv * (MSLP / p)^(R_m / C_p)
    return (ρ = ρ, ρθ = ρθ)
end

Π(ρθ) = C_p * (R_d * ρθ / MSLP)^(R_m / C_v)
Φ(z) = grav * z

function discrete_hydrostatic_balance!(
    ρ,
    w,
    ρθ,
    Δz::Float64,
    _grav::Float64,
    Π::Function,
)
    # compute θ such that
    #   I(θ)[i+1/2] = -g / ∂f(Π(ρθ))
    # discretely, then set
    #   ρ = ρθ/θ
    for i in 1:(length(ρ) - 1)
        #  ρ[i+1] = ρθ[i+1]/(-2Δz*_grav/(Π(ρθ[i+1]) - Π(ρθ[i])) - ρθ[i]/ρ[i])
        ρ[i + 1] =
            ρθ[i + 1] /
            (-2 * _grav / ((Π(ρθ[i + 1]) - Π(ρθ[i])) / Δz) - ρθ[i] / ρ[i])

        ρ[i + 1] =
            ρθ[i + 1] /
            (1 / ((-2 * _grav) * (Π(ρθ[i + 1]) - Π(ρθ[i]))Δz) - ρθ[i] / ρ[i])

        ∂Π∂z = (Π(ρθ[i + 1]) - Π(ρθ[i])) / Δz
    end
end

zc = Fields.coordinate_field(cspace)
Yc = decaying_temperature_profile.(zc.z)
w = Geometry.Cartesian3Vector.(zeros(FT, fspace))
zf = parent(Fields.coordinate_field(fspace).z)
Δz = zf[2:end] - zf[1:(end - 1)]
Y_init = copy(Yc)
w_init = copy(w)
Y = Fields.FieldVector(Yc = Yc, w = w)

function tendency!(dY, Y, _, t)
    Yc = Y.Yc
    w = Y.w
    dYc = dY.Yc
    dw = dY.w

    If = Operators.InterpolateC2F()
    ∂ = Operators.DivergenceF2C(
        bottom = Operators.SetValue(Geometry.Cartesian3Vector(zero(FT))),
        top = Operators.SetValue(Geometry.Cartesian3Vector(zero(FT))),
    )
    ∂f = Operators.GradientC2F()
    B = Operators.SetBoundaryOperator(
        bottom = Operators.SetValue(Geometry.Cartesian3Vector(zero(FT))),
        top = Operators.SetValue(Geometry.Cartesian3Vector(zero(FT))),
    )

    @. dYc.ρ = -(∂(w * If(Yc.ρ)))
    @. dYc.ρθ = -(∂(w * If(Yc.ρθ)))
    @. dw = B(
        Geometry.CartesianVector(
            -(If(Yc.ρθ / Yc.ρ) * ∂f(Π(Yc.ρθ))) - ∂f(Φ(zc.z)),
        ),
    )
    return dY
end

dY = tendency!(similar(Y), Y, nothing, 0.0)

function jacobian!(J, Y, p, t)
    # copyto!(J, LinearAlgebra.I)
    # @info length(Y)
    # @info Y[1]

    J .= 0.0

    @info "Jacobian computation!!!!"
    # N cells
    N = div(length(Y) - 1, 3)

    ρ, ρθ, w = Y[1:N], Y[(N + 1):(2N)], Y[(2N + 1):(3N + 1)]

    # construct cell center
    ρh = [ρ[1]; (ρ[1:(N - 1)] + ρ[2:N]) / 2.0; ρ[N]]
    ρθh = [ρθ[1]; (ρθ[1:(N - 1)] + ρθ[2:N]) / 2.0; ρθ[N]]

    Πc = Π.(ρθ)
    Πh = [NaN64; (Πc[1:(N - 1)] + Πc[2:N]) / 2.0; NaN64]
    Δzh = [NaN64; (Δz[1:(N - 1)] + Δz[2:N]) / 2.0; NaN64]

    for i in 1:N
        J[i, i + 2N] = ρh[i] / Δz[i]
        J[i, i + 2N + 1] = -ρh[i + 1] / Δz[i]
    end

    for i in 1:N
        J[i + N, i + 2N] = ρθh[i] / Δz[i]
        J[i + N, i + 2N + 1] = -ρθh[i + 1] / Δz[i]
    end

    # 0 for i = 1, N+1
    for i in 2:N
        J[i + 2N, (i - 1)] = -grav / (2 * ρh[i])
        J[i + 2N, (i - 1) + 1] = -grav / (2 * ρh[i])

        J[i + 2N, (i - 1) + N] = (γ - 1) * Πh[i] ./ (ρh[i] * Δzh[i])
        J[i + 2N, (i - 1) + 1 + N] = -(γ - 1) * Πh[i] ./ (ρh[i] * Δzh[i])
    end

    return J

    # D_ρ = diagm(0=>-ρh/Δz, -1=>ρh/Δz)[1:N, 1:N-1]
    # D_Θ = diagm(0=>-ρθh/Δz, -1=>ρθh/Δz)[1:N, 1:N-1]
    # G_W = (γ - 1) * diagm(0=>Πh./ρh/Δz, 1=>-Πh./ρh/Δz)[1:N-1, 1:N]
    # A_W = diagm(0=>-ones(N-1)./ρh/2, 1=>-ones(N-1)./ρh/2)[1:N-1, 1:N]

    # P = ([zeros(N,N)     zeros(N,N)      D_ρ;
    #       zeros(N,N)     zeros(N,N)      D_Θ
    #       A_W*_grav        G_W          zeros(N+1,N+1)])

end

function linsolve!(::Type{Val{:init}}, f, u0; kwargs...)
    function _linsolve!(x, A, b, update_matrix = false; kwargs...)
        #     _A = RecursiveFactorization.lu(A)
        #   ldiv!(x,_A,b)
        #   @info "norm: ", norm(x), norm(A), norm(b)
        x = copy(b)
        N = div(length(x) - 1, 3)
        J = A[(2N + 1):(3N + 1), (2N + 1):(3N + 1)]
        J +=
            -A[(2N + 1):(3N + 1), 1:N] *
            (Diagonal(A[1:N, 1:N]) \ A[1:N, (2N + 1):(3N + 1)])
        J +=
            -A[(2N + 1):(3N + 1), (N + 1):(2N)] * (
                Diagonal(A[(N + 1):(2N), (N + 1):(2N)]) \
                A[(N + 1):(2N), (2N + 1):(3N + 1)]
            )

        x[(2N + 1):(3N + 1)] +=
            -A[(2N + 1):(3N + 1), 1:N] * (Diagonal(A[1:N, 1:N]) \ b[1:N])
        x[(2N + 1):(3N + 1)] +=
            -A[(2N + 1):(3N + 1), (N + 1):(2N)] *
            (Diagonal(A[(N + 1):(2N), (N + 1):(2N)]) \ b[(N + 1):(2N)])

        x[(2N + 1):(3N + 1)] .= Tridiagonal(J) \ x[(2N + 1):(3N + 1)]
        x[1:N] .=
            Diagonal(A[1:N, 1:N]) \
            (b[1:N] - A[1:N, (2N + 1):(3N + 1)] * x[(2N + 1):(3N + 1)])
        x[(N + 1):(2N)] .=
            Diagonal(A[(N + 1):(2N), (N + 1):(2N)]) \ (
                b[(N + 1):(2N)] -
                A[(N + 1):(2N), (2N + 1):(3N + 1)] * x[(2N + 1):(3N + 1)]
            )

    end
end

Δt = 600.0
ndays = 1.0

# Solve the ODE operator
prob = ODEProblem(
    ODEFunction(
        tendency!,
        jac = jacobian!,
        jac_prototype = zeros(length(Y), length(Y)),
        tgrad = (dT, Y, p, t) -> fill!(dT, 0),
    ),
    Y,
    (0.0, 60 * 60 * 24 * ndays),
)
# 60 * 60 * 24 * ndays
sol = solve(
    prob,
    # ImplicitEuler(),
    Rosenbrock23(linsolve = linsolve!),
    dt = Δt,
    saveat = 60 * 60, # save every hour
    progress = true,
    progress_message = (dt, u, p, t) -> t,
);

ENV["GKSwstype"] = "nul"
import Plots
Plots.GRBackend()

dirname = "hydrostatic_implicit"
path = joinpath(@__DIR__, "output", dirname)
mkpath(path)

z_centers = parent(Fields.coordinate_field(cspace))
z_faces = parent(Fields.coordinate_field(fspace))

function hydrostatic_plot(u, Yc_init, w_init; title = "", size = (1024, 600))
    sub_plt1 = Plots.plot(
        parent(Yc_init.ρ),
        z_centers,
        marker = :circle,
        xlabel = "ρ",
        label = "T=0",
    )
    sub_plt1 = Plots.plot!(sub_plt1, parent(u.Yc.ρ), z_centers, label = "T")

    sub_plt2 = Plots.plot(
        parent(w_init),
        z_faces,
        marker = :circle,
        xlim = (-0.2, 0.2),
        xlabel = "ω",
        label = "T=0",
    )
    sub_plt2 = Plots.plot!(sub_plt2, parent(u.w), z_faces, label = "T")

    sub_plt3 = Plots.plot(
        parent(Yc_init.ρθ),
        z_centers,
        marker = :circle,
        xlabel = "ρθ",
        label = "T=0",
    )
    sub_plt3 = Plots.plot!(sub_plt3, parent(u.Yc.ρθ), z_centers, label = "T")

    return Plots.plot(
        sub_plt1,
        sub_plt2,
        sub_plt3,
        title = title,
        layout = (1, 3),
        size = size,
    )
end

anim = Plots.@animate for (i, u) in enumerate(sol.u)
    hydrostatic_plot(u, Y_init, w_init, title = "Hour $(i)")
end
Plots.mp4(anim, joinpath(path, "hydrostatic.mp4"), fps = 10)

Plots.png(
    hydrostatic_plot(sol[end], Y_init, w_init),
    joinpath(path, "hydrostatic_end.png"),
)

function linkfig(figpath, alt = "")
    # buildkite-agent upload figpath
    # link figure in logs if we are running on CI
    if get(ENV, "BUILDKITE", "") == "true"
        artifact_url = "artifact://$figpath"
        print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
    end
end

linkfig("output/$(dirname)/hydrostatic_end.png", "Hydrostatic End")


#######################################################################
# Spectrum
#######################################################################
struct Spectrum{E,T,A<:AbstractMatrix{T}}
    energies::E
    states::A
end

"""
    spectrum(h; method = defaultmethod(h))

Compute the spectrum of a 0D Hamiltonian `h` (or alternatively of the bounded unit cell of a
finite dimensional `h`) using one of the following `method`s

    method                    diagonalization function
    --------------------------------------------------------------
    LinearAlgebraPackage()     LinearAlgebra.eigen!
    ArpackPackage()            Arpack.eigs (must be `using Arpack`)

The energies and eigenstates in the resulting `s::Spectrum` object can be accessed with
`energies(s)` and `states(s)`

# See also
    `energies`, `states`, `bandstructure`

"""
function spectrum(h; method = defaultmethod(h))
    matrix = similarmatrix(h, method)
    bloch!(matrix, h)
    (ϵk, ψk) = diagonalize(matrix, method)
    return Spectrum(ϵk, ψk)
end

"""
    energies(s::Spectrum)

Return the energies of `s` as a `Vector`

# See also
    `spectrum`, `states`
"""
energies(s::Spectrum) = s.energies

"""
    states(s::Spectrum)

Return the states of `s` as the columns of a `Matrix`

# See also
    `spectrum`, `energies`
"""
states(s::Spectrum) = s.states

#######################################################################
# Bandstructure
#######################################################################
struct Band{M,A<:AbstractVector{M},MD<:Mesh,S<:AbstractArray}
    mesh::MD        # Mesh with missing vertices removed
    simplices::S    # Tuples of indices of mesh vertices that define mesh simplices
    states::A       # Must be resizeable container to build & refine band
    dimstates::Int  # Needed to extract the state at a given vertex from vector `states`
end

function Band(mesh::Mesh{D}, states::AbstractVector{M}, dimstates::Int) where {M,D}
    simps = simplices(mesh, Val(D))
    return Band(mesh, simps, states, dimstates)
end

struct Bandstructure{D,M,B<:Band{M},MD<:Mesh{D}}   # D is dimension of parameter space
    bands::Vector{B}
    kmesh::MD
end

function Base.show(io::IO, b::Bandstructure{D,M}) where {D,M}
    ioindent = IOContext(io, :indent => string("  "))
    print(io,
"Bandstructure: bands for a $(D)D hamiltonian
  Bands        : $(length(b.bands))
  Element type : $(displayelements(M))")
    print(ioindent, "\n", b.kmesh)
end

# API #
"""
    bands(bs::Bandstructure)

Return a vector of all the `Band`s in `bs`.
"""
bands(bs::Bandstructure) = bs.bands

"""
    vertices(bs::Bandstructure, i)

Return the vertices `(k..., ϵ)` of the i-th band in `bs`, in the form of a
`Vector{SVector{L+1}}`, where `L` is the lattice dimension.
"""
vertices(bs::Bandstructure, i) = vertices(bands(bs)[i])

vertices(b::Band) = vertices(b.mesh)

"""
    states(bs::Bandstructure, i)

Return the states of each vertex of the i-th band in `bs`, in the form of a `Matrix` of size
`(nψ, nk)`, where `nψ` is the length of each state vector, and `nk` the number of vertices.
"""
states(bs::Bandstructure, i) = states(bands(bs)[i])

states(b::Band) = reshape(b.states, b.dimstates)

#######################################################################
# bandstructure
#######################################################################
"""
    bandstructure(h::Hamiltonian, mesh::Mesh; minprojection = 0.5, method = defaultmethod(h))

Compute the bandstructure of Bloch Hamiltonian `bloch(h, ϕs...)` with `ϕs` evaluated on the
vertices of `mesh`. It is assumed that `h` is hermitian.

The option `minprojection` determines the minimum projection between eigenstates to connect
them into a common subband. The option `method` is chosen automatically if unspecified, and
can be one of the following

    method                    diagonalization function
    --------------------------------------------------------------
    LinearAlgebraPackage()     LinearAlgebra.eigen!
    ArpackPackage()            Arpack.eigs (must be `using Arpack`)

Options passed to the `method` will be forwarded to the diagonalization function. For example,
`method = ArpackPackage(nev = 8, sigma = 1im)` will use `Arpack.eigs(matrix; nev = 8,
sigma = 1im)` to compute the bandstructure.

    bandstructure(h::Hamiltonian; resolution = 13, shift = missing, kw...)

Same as above with a  uniform `mesh = marchingmesh(h; npoints = resolution, shift = shift)`
of marching tetrahedra (generalized to the lattice dimensions of the Hamiltonian). Note that
`resolution` denotes the number of points along each Bloch axis, including endpoints (can be
a tuple for axis-dependent points).

# Example
```
julia> h = LatticePresets.honeycomb() |> unitcell(3) |> hamiltonian(hopping(-1, range = 1/√3));

julia> bandstructure(h; resolution = 25, method = LinearAlgebraPackage())
Bandstructure: bands for a 2D hamiltonian
  Bands        : 8
  Element type : scalar (Complex{Float64})
  Mesh{2}: mesh of a 2-dimensional manifold
    Vertices   : 625
    Edges      : 3552
```

# See also
    marchingmesh
"""
function bandstructure(h::Hamiltonian{<:Any,L,M}; resolution = 13, shift = missing, kw...) where {L,M}
    checkfinitedim(h)
    mesh = marchingmesh(h; npoints = resolution, shift = shift)
    return bandstructure(h,  mesh; kw...)
end

function bandstructure(h::Hamiltonian, mesh::Mesh; method = defaultmethod(h), minprojection = 0.5)
    checkfinitedim(h)
    ishermitian(h) || throw(ArgumentError("Hamiltonian must be hermitian"))
    d = diagonalizer(h, mesh, method, minprojection)
    matrix = similarmatrix(h, method)
    return bandstructure!(matrix, h, mesh, d)
end

function bandstructure!(matrix::AbstractMatrix, h::Hamiltonian{<:Lattice,<:Any,M}, mesh::MD, d::Diagonalizer) where {M,D,T,MD<:Mesh{D,T}}
    nϵ = 0                           # Temporary, to be reassigned
    ϵks = Matrix{T}(undef, 0, 0)     # Temporary, to be reassigned
    ψks = Array{M,3}(undef, 0, 0, 0) # Temporary, to be reassigned

    dimh = size(h, 1)
    nk = nvertices(mesh)
    # function to apply to eigenvalues when building bands (depends on momenta type)
    by = _maybereal(T)

    p = Progress(nk, "Step 1/2 - Diagonalising: ")
    for (n, ϕs) in enumerate(vertices(mesh))
        bloch!(matrix, h, ϕs)
        # (ϵk, ψk) = diagonalize(Hermitian(matrix), d)  ## This is faster (!)
        (ϵk, ψk) = diagonalize(matrix, d.method)
        resolve_degeneracies!(ϵk, ψk, ϕs, matrix, d.codiag)
        if n == 1  # With first vertex can now know the number of eigenvalues... Reassign
            nϵ = size(ϵk, 1)
            ϵks = Matrix{T}(undef, nϵ, nk)
            ψks = Array{M,3}(undef, dimh, nϵ, nk)
        end
        copyslice!(ϵks, CartesianIndices((1:nϵ, n:n)),
                   ϵk,  CartesianIndices((1:nϵ,)), by)
        copyslice!(ψks, CartesianIndices((1:dimh, 1:nϵ, n:n)),
                   ψk,  CartesianIndices((1:dimh, 1:nϵ)), identity)
        ProgressMeter.next!(p; showvalues = ())
    end

    p = Progress(nϵ * nk, "Step 2/2 - Connecting bands: ")
    pcounter = 0
    bands = Band{M,Vector{M},Mesh{D+1,T,Vector{SVector{D+1,T}}},Vector{NTuple{D+1,Int}}}[]
    vertindices = zeros(Int, nϵ, nk) # 0 == unclassified, -1 == different band, > 0 vertex index
    pending = CartesianIndex{2}[]
    sizehint!(pending, nk)
    while true
        src = findfirst(iszero, vertindices)
        src === nothing && break
        resize!(pending, 1)
        pending[1] = src # source CartesianIndex for band search
        band = extractband(mesh, pending, ϵks, ψks, vertindices, d.minprojection)
        nverts = nvertices(band.mesh)
        nverts > D && push!(bands, band) # avoid bands with no simplices
        pcounter += nverts
        ProgressMeter.update!(p, pcounter; showvalues = ())
    end
    return Bandstructure(bands, mesh)
end

_maybereal(::Type{<:Complex}) = identity
_maybereal(::Type{<:Real}) = real

function extractband(kmesh::Mesh{D,T}, pending, ϵks::AbstractArray{T}, ψks::AbstractArray{M}, vertindices, minprojection) where {D,T,M}
    dimh, nϵ, nk = size(ψks)
    kverts = vertices(kmesh)
    states = eltype(ψks)[]
    sizehint!(states, nk * dimh)
    verts = SVector{D+1,T}[]
    sizehint!(verts, nk)
    adjmat = SparseMatrixBuilder{Bool}()
    vertindices[first(pending)] = 1 # pending starts with a single vertex
    for c in pending
        ϵ, k = Tuple(c) # c == CartesianIndex(ϵ::Int, k::Int)
        vertex = vcat(kverts[k], SVector(ϵks[c]))
        push!(verts, vertex)
        appendslice!(states, ψks, CartesianIndices((1:dimh, ϵ:ϵ, k:k)))
        for edgek in edges(kmesh, k)
            k´ = edgedest(kmesh, edgek)
            proj, ϵ´ = findmostparallel(ψks, k´, ϵ, k)
            if proj >= minprojection
                if iszero(vertindices[ϵ´, k´]) # unclassified
                    push!(pending, CartesianIndex(ϵ´, k´))
                    vertindices[ϵ´, k´] = length(pending) # this is clever!
                end
                indexk´ = vertindices[ϵ´, k´]
                indexk´ > 0 && pushtocolumn!(adjmat, indexk´, true)
            end
        end
        finalizecolumn!(adjmat)
    end
    for (i, vi) in enumerate(vertindices)
        @inbounds vi > 0 && (vertindices[i] = -1) # mark as classified in a different band
    end
    mesh = Mesh(verts, sparse(adjmat))
    return Band(mesh, states, dimh)
end

function findmostparallel(ψks::Array{M,3}, destk, srcb, srck) where {M}
    T = real(eltype(M))
    dimh, nϵ, nk = size(ψks)
    maxproj = zero(T)
    destb = 0
    @inbounds for nb in 1:nϵ
        proj = zero(M)
        for i in 1:dimh
            proj += ψks[i, nb, destk]' * ψks[i, srcb, srck]
        end
        absproj = T(abs(tr(proj)))
        if maxproj <= absproj  # must happen at least once
            destb = nb
            maxproj = absproj
        end
    end
    return maxproj, destb
end


# ---
# title: 661. Image Smoother
# id: problem661
# author: zhwang
# date: 2022-03-24
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/image-smoother/description/>
# hidden: true
# ---
#
# Given a 2D integer matrix M representing the gray scale of an image, you need
# to design a smoother to make the gray scale of each cell becomes the average
# gray scale (rounding down) of all the 8 surrounding cells and itself. If a
# cell has less than 8 surrounding cells, then use as many as you can.
#
# **Example 1:**
#
#
#
#     Input:
#     [[1,1,1],
#      [1,0,1],
#      [1,1,1]]
#     Output:
#     [[0, 0, 0],
#      [0, 0, 0],
#      [0, 0, 0]]
#     Explanation:
#     For the point (0,0), (0,2), (2,0), (2,2): floor(3/4) = floor(0.75) = 0
#     For the point (0,1), (1,0), (1,2), (2,1): floor(5/6) = floor(0.83333333) = 0
#     For the point (1,1): floor(8/9) = floor(0.88888889) = 0
#
#
# **Note:**
#
#   1. The value in the given matrix is in the range of [0, 255].
#   2. The length and width of the given matrix are in the range of [1, 150].
#
#
## @lc code=start
using LeetCode

function image_smoother(img::Matrix{Int})::Matrix{Int}
    presum, (m, n) = similar(img), size(img)
    presum[1, 1] = img[1, 1]
    for i in 2:m
        presum[i, 1] = presum[i - 1, 1] + img[i, 1]
    end
    for j in 2:n
        presum[1, j] = presum[1, j - 1] + img[1, j]
    end
    for i in 2:m, j in 2:n
        presum[i, j] =
            img[i, j] + presum[i - 1, j] + presum[i, j - 1] - presum[i - 1, j - 1]
    end
    ## obtain results from 2-dim presum
    res, indices = similar(img), CartesianIndices(img)
    for i in 1:m, j in 1:n
        x1 = i > 1 ? i - 2 : i - 1
        y1 = j > 1 ? j - 2 : j - 1
        x2 = i < m ? i + 1 : i
        y2 = j < n ? j + 1 : j
        idx1, idx2, idx3, idx4 = CartesianIndex.([(x1, y1), (x1, y2), (x2, y1), (x2, y2)])
        area =
            presum[idx4] - (idx2 ∈ indices && presum[idx2]) -
            (idx3 ∈ indices && presum[idx3]) + (idx1 ∈ indices && presum[idx1])
        res[i, j] = area ÷ prod((idx4 - idx1).I)
    end
    return res
end
## @lc code=end


"""
    MatrixTDist(ν, M, Σ, Ω)
```julia
ν::Real            positive degrees of freedom
M::AbstractMatrix  n x p location
Σ::AbstractPDMat   n x n scale
Ω::AbstractPDMat   p x p scale
```
The [matrix *t*-distribution](https://en.wikipedia.org/wiki/Matrix_t-distribution)
generalizes the multivariate *t*-distribution to ``n\\times p`` real
matrices ``\\mathbf{X}``. If ``\\mathbf{X}\\sim \\textrm{MT}_{n,p}(\\nu,\\mathbf{M},\\boldsymbol{\\Sigma},
\\boldsymbol{\\Omega})``, then its probability density function is

```math
f(\\mathbf{X} ; \\nu,\\mathbf{M},\\boldsymbol{\\Sigma}, \\boldsymbol{\\Omega}) =
c_0 \\left|\\mathbf{I}_n + \\boldsymbol{\\Sigma}^{-1}(\\mathbf{X} - \\mathbf{M})\\boldsymbol{\\Omega}^{-1}(\\mathbf{X}-\\mathbf{M})^{\\rm{T}}\\right|^{-\\frac{\\nu+n+p-1}{2}},
```

where

```math
c_0=\\frac{\\Gamma_p\\left(\\frac{\\nu+n+p-1}{2}\\right)}{(\\pi)^\\frac{np}{2} \\Gamma_p\\left(\\frac{\\nu+p-1}{2}\\right)} |\\boldsymbol{\\Omega}|^{-\\frac{n}{2}} |\\boldsymbol{\\Sigma}|^{-\\frac{p}{2}}.
```

If the joint distribution ``p(\\mathbf{S},\\mathbf{X})=p(\\mathbf{S})p(\\mathbf{X}|\\mathbf{S})``
is given by

```math
\\begin{aligned}
\\mathbf{S}&\\sim \\textrm{IW}_n(\\nu + n - 1, \\boldsymbol{\\Sigma})\\\\
\\mathbf{X}|\\mathbf{S}&\\sim \\textrm{MN}_{n,p}(\\mathbf{M}, \\mathbf{S}, \\boldsymbol{\\Omega}),
\\end{aligned}
```

then the marginal distribution of ``\\mathbf{X}`` is
``\\textrm{MT}_{n,p}(\\nu,\\mathbf{M},\\boldsymbol{\\Sigma},\\boldsymbol{\\Omega})``.
"""
struct MatrixTDist{T <: Real, TM <: AbstractMatrix, TΣ <: AbstractPDMat, TΩ <: AbstractPDMat} <: ContinuousMatrixDistribution
    ν::T
    M::TM
    Σ::TΣ
    Ω::TΩ
    logc0::T
end

#  -----------------------------------------------------------------------------
#  Constructors
#  -----------------------------------------------------------------------------

function MatrixTDist(ν::T, M::AbstractMatrix{T}, Σ::AbstractPDMat{T}, Ω::AbstractPDMat{T}) where T <: Real
    n, p = size(M)
    0 < ν < Inf || throw(ArgumentError("degrees of freedom must be positive and finite."))
    n == dim(Σ) || throw(ArgumentError("Number of rows of M must equal dim of Σ."))
    p == dim(Ω) || throw(ArgumentError("Number of columns of M must equal dim of Ω."))
    logc0 = matrixtdist_logc0(Σ, Ω, ν)
    R = Base.promote_eltype(T, logc0)
    prom_M = convert(AbstractArray{R}, M)
    prom_Σ = convert(AbstractArray{R}, Σ)
    prom_Ω = convert(AbstractArray{R}, Ω)
    MatrixTDist{R, typeof(prom_M), typeof(prom_Σ), typeof(prom_Ω)}(R(ν), prom_M, prom_Σ, prom_Ω, R(logc0))
end

function MatrixTDist(ν::Real, M::AbstractMatrix, Σ::AbstractPDMat, Ω::AbstractPDMat)
    T = Base.promote_eltype(ν, M, Σ, Ω)
    MatrixTDist(convert(T, ν), convert(AbstractArray{T}, M), convert(AbstractArray{T}, Σ), convert(AbstractArray{T}, Ω))
end

MatrixTDist(ν::Real, M::AbstractMatrix, Σ::Union{AbstractMatrix, LinearAlgebra.Cholesky}, Ω::Union{AbstractMatrix, LinearAlgebra.Cholesky}) = MatrixTDist(ν, M, PDMat(Σ), PDMat(Ω))
MatrixTDist(ν::Real, M::AbstractMatrix, Σ::AbstractPDMat, Ω::Union{AbstractMatrix, LinearAlgebra.Cholesky}) = MatrixTDist(ν, M, Σ, PDMat(Ω))
MatrixTDist(ν::Real, M::AbstractMatrix, Σ::Union{AbstractMatrix, LinearAlgebra.Cholesky}, Ω::AbstractPDMat) = MatrixTDist(ν, M, PDMat(Σ), Ω)

#  -----------------------------------------------------------------------------
#  REPL display
#  -----------------------------------------------------------------------------

 show(io::IO, d::MatrixTDist) = show_multline(io, d, [(:ν, d.ν), (:M, d.M), (:Σ, Matrix(d.Σ)), (:Ω, Matrix(d.Ω))])

#  -----------------------------------------------------------------------------
#  Conversion
#  -----------------------------------------------------------------------------

function convert(::Type{MatrixTDist{T}}, d::MatrixTDist) where T <: Real
    MM = convert(AbstractArray{T}, d.M)
    ΣΣ = convert(AbstractArray{T}, d.Σ)
    ΩΩ = convert(AbstractArray{T}, d.Ω)
    MatrixTDist{T, typeof(MM), typeof(ΣΣ), typeof(ΩΩ)}(T(d.ν), MM, ΣΣ, ΩΩ, T(d.logc0))
end

function convert(::Type{MatrixTDist{T}}, ν, M::AbstractMatrix, Σ::AbstractPDMat, Ω::AbstractPDMat, logc0) where T <: Real
    MM = convert(AbstractArray{T}, M)
    ΣΣ = convert(AbstractArray{T}, Σ)
    ΩΩ = convert(AbstractArray{T}, Ω)
    MatrixTDist{T, typeof(MM), typeof(ΣΣ), typeof(ΩΩ)}(T(ν), MM, ΣΣ, ΩΩ, T(logc0))
end

#  -----------------------------------------------------------------------------
#  Properties
#  -----------------------------------------------------------------------------

size(d::MatrixTDist) = size(d.M)

rank(d::MatrixTDist) = minimum( size(d) )

insupport(d::MatrixTDist, X::Matrix) = isreal(X) && size(X) == size(d)

function mean(d::MatrixTDist)
    n, p = size(d)
    d.ν + p - n > 1 || throw(ArgumentError("mean only defined for df + p - n > 1"))
    return d.M
end

mode(d::MatrixTDist) = d.M

cov(d::MatrixTDist, ::Val{true}=Val(true)) = d.ν <= 2 ? throw(ArgumentError("cov only defined for df > 2")) : Matrix(kron(d.Ω, d.Σ)) ./ (d.ν - 2)

cov(d::MatrixTDist, ::Val{false}) = ((n, p) = size(d); reshape(cov(d), n, p, n, p))

var(d::MatrixTDist) = d.ν <= 2 ? throw(ArgumentError("var only defined for df > 2")) : reshape(diag(cov(d)), size(d))

params(d::MatrixTDist) = (d.ν, d.M, d.Σ, d.Ω)

@inline partype(d::MatrixTDist{T}) where {T <: Real} = T

#  -----------------------------------------------------------------------------
#  Evaluation
#  -----------------------------------------------------------------------------

function matrixtdist_logc0(Σ::AbstractPDMat, Ω::AbstractPDMat, ν::Real)
    #  returns the natural log of the normalizing constant for the pdf
    n = dim(Σ)
    p = dim(Ω)
    term1 = logmvgamma(p, (ν + n + p - 1) / 2)
    term2 = - (n * p / 2) * logπ
    term3 = - logmvgamma(p, (ν + p - 1) / 2)
    term4 = (-n / 2) * logdet(Ω)
    term5 = (-p / 2) * logdet(Σ)
    term1 + term2 + term3 + term4 + term5
end

function logkernel(d::MatrixTDist, X::AbstractMatrix)
    n, p = size(d)
    A = X - d.M
    (-(d.ν + n + p - 1) / 2) * logdet( I + (d.Σ \ A) * (d.Ω \ A') )
end

#  -----------------------------------------------------------------------------
#  Sampling
#  -----------------------------------------------------------------------------

#  Theorem 4.2.1 in Gupta and Nagar (1999)

function _rand!(rng::AbstractRNG, d::MatrixTDist, A::AbstractMatrix)
    n, p = size(d)
    S = rand(rng, InverseWishart(d.ν + n - 1, d.Σ) )
    A .= rand(rng, MatrixNormal(d.M, S, d.Ω) )
end

#  -----------------------------------------------------------------------------
#  Relationship with Multivariate t
#  -----------------------------------------------------------------------------
#  if a t-distributed random matrix is in fact just a row or column,
#  it is equivalent to a t-distributed random vector.
#  -----------------------------------------------------------------------------

function MvTDist(MT::MatrixTDist)
    n, p = size(MT)
    all([n, p] .> 1) && throw(ArgumentError("Row or col dim of `MatrixTDist` must be 1 to coerce to `MvTDist`"))
    ν, M, Σ, Ω = params(MT)
    MvTDist(ν, vec(M), (1 / ν) * kron(Σ, Ω))
end

#  -----------------------------------------------------------------------------
#  Test utils
#  -----------------------------------------------------------------------------

function _univariate(d::MatrixTDist)
    check_univariate(d)
    ν, M, Σ, Ω = params(d)
    μ = M[1]
    σ = sqrt( Matrix(Σ)[1] * Matrix(Ω)[1] / ν )
    return LocationScale(μ, σ, TDist(ν))
end

_multivariate(d::MatrixTDist) = MvTDist(d)

function _rand_params(::Type{MatrixTDist}, elty, n::Int, p::Int)
    ν = elty( n + p + 1 + abs(10randn()) )
    M = randn(elty, n, p)
    Σ = (X = 2rand(elty, n, n) .- 1; X * X')
    Ω = (Y = 2rand(elty, p, p) .- 1; Y * Y')
    return ν, M, Σ, Ω
end


using SuperTB, PyPlot, ProgressMeter, LinearAlgebra
close("all")

# set up a model (amorphous system with AII-class hopping)
tb = TBModel(hopAgarwala)
tb.periodicity = trues(2) # periodic boundary conditions
tb.basisdim    = 2

# run over this parameter range
M = range(-4.5, 4.5, length=50)
sites = 12:12:24^2
Nsamples = 1 # number of samples in ensemble

β = zeros(length(M), length(sites), Nsamples)

@showprogress 2 "Building phase diagram (#sites-progress)..." for s = eachindex(sites)
    for g = 1:Nsamples
        tb.pos = rand(sites[s],2)

        for m = eachindex(M)
            tb.hopfun(ri,rj,per) = hopAgarwala(ri, rj, per,
                                               1.0/6,  # R
                                               M[m],   # M
                                               0.25,   # t₂
                                               0.5,    # λ
                                               1.0/24) # a
            # solve the system
            H = Hermitian(Matrix(setHopping!(tb)))
            F = eigen(H)

            # compute Bott indices
            β[m,s,g] = bott(F.vectors,tb,sites[s])[1]
        end
    end
end
βᵉ = dropdims(sum(β,3),dims=3)./Nsamples
ρ = sites./24^2

figure()
surf(ρ,M,βᵉ)


function tss(y::AbstractVector, hasintercept::Bool, weights::AbstractWeights)
    if hasintercept
        m = mean(y, weights)
        sum(@inbounds (y[i] - m)^2 * weights[i] for i in eachindex(y))
    else
        sum(@inbounds y[i]^2 * weights[i] for i in eachindex(y))
    end
end

function Fstat(coef::Vector{Float64}, matrix_vcov::AbstractMatrix{Float64}, has_intercept::Bool)
    coefF = copy(coef)
    # TODO: check I can't do better
    length(coef) == has_intercept && return NaN
    if has_intercept
        coefF = coefF[2:end]
        matrix_vcov = matrix_vcov[2:end, 2:end]
    end
    return (coefF' * (matrix_vcov \ coefF)) / length(coefF)
end


using WeaklyHard
"""
Example for comparing dominance and equivalence of WeaklyHard constraints
"""

# Constraints to test
lambda1 = AnyHitConstraint(48, 54)
lambda2 = AnyMissConstraint(12, 67)
lambda3 = AnyHitConstraint(12, 15)
lambda4 = RowMissConstraint(9)
lambda5 = RowHitConstraint(20, 47)
lambda6 = RowHitConstraint(5, 12)
lambda7 = RowMissConstraint(3)

lambdas = [lambda1, lambda2, lambda3, lambda4, lambda5, lambda6, lambda7]

# Compare individual constraints
for l1 in lambdas
    for l2 in lambdas
        if is_equivalent(l1, l2)
            println("$l1 \t===\t $l2")
        elseif is_dominant(l1, l2)
            println("$l1 \t<=\t $l2")
        else
            println("$l1 \t</=\t $l2")
        end # if
    end # for
    println(" ")
end # for


struct NullInfoState <: CFR.AbstractInfoState end

# remove game from AbstractCFRSolver type parameters?
struct TabularStrategy{K,G} <: CFR.AbstractCFRSolver{K,G,NullInfoState}
    game::G
    d::Dict{K,Vector{Float64}}
end

# function TabularStrategy(game::Game{H,K}, dict::Dict{K,<:AbstractVector}) where {H,K}
#     return TabularStrategy{K, typeof(game)}(dict)
# end

CFR.strategy(sol::TabularStrategy{K}, I::K) where K = sol.d[I]

const α_opt = 0.10

const OptimalKuhnStrategy = TabularStrategy(CFR.Kuhn(), Dict(
    (1,1,SA[-1,-1,-1]) => [1. - α_opt, α_opt],
    (1,2,SA[-1,-1,-1]) => [1., 0.],
    (1,3,SA[-1,-1,-1]) => [1-3α_opt,3α_opt],

    (2,1,SA[0,-1,-1]) => [2/3,1/3],
    (2,2,SA[0,-1,-1]) => [1., 0.],
    (2,3,SA[0,-1,-1]) => [0., 1.],

    (2,1,SA[1,-1,-1]) => [1., 0.],
    (2,2,SA[1,-1,-1]) => [2/3,1/3],
    (2,3,SA[1,-1,-1]) => [0., 1.],

    (1,1,SA[0,1,-1]) => [1., 0.],
    (1,2,SA[0,1,-1]) => [2/3-α_opt,1/3+α_opt],
    (1,3,SA[0,1,-1]) => [0., 1.]
))

const SuboptimalKuhnStrategy = TabularStrategy(CFR.Kuhn(), Dict(
    (1,1,SA[-1,-1,-1]) => normalize!(rand(2),1),
    (1,2,SA[-1,-1,-1]) => normalize!(rand(2),1),
    (1,3,SA[-1,-1,-1]) => normalize!(rand(2),1),

    (2,1,SA[0,-1,-1]) => normalize!(rand(2),1),
    (2,2,SA[0,-1,-1]) => normalize!(rand(2),1),
    (2,3,SA[0,-1,-1]) => normalize!(rand(2),1),

    (2,1,SA[1,-1,-1]) => normalize!(rand(2),1),
    (2,2,SA[1,-1,-1]) => normalize!(rand(2),1),
    (2,3,SA[1,-1,-1]) => normalize!(rand(2),1),

    (1,1,SA[0,1,-1]) => normalize!(rand(2),1),
    (1,2,SA[0,1,-1]) => normalize!(rand(2),1),
    (1,3,SA[0,1,-1]) => normalize!(rand(2),1),
))

const OptimalRPSStrategy = TabularStrategy(CFR.IIEMatrixGame(), Dict(
    0 => fill(1/3,3),
    1 => fill(1/3,3)
))

const SuboptimalRPSStrategy = TabularStrategy(CFR.IIEMatrixGame(), Dict(
    0 => Float64[1,0,0],
    1 => Float64[1,0,0]
))

const SlightlySuboptimalRPSStrategy = TabularStrategy(CFR.IIEMatrixGame(), Dict(
    0 => Float64[0.35,0.30,0.35],
    1 => Float64[0.30,0.35,0.35]
))

@testset "Exploitability" begin
    @testset "RPS" begin
        atol = 1e-3

        RPS = CFR.IIEMatrixGame()
        e_sol1 = ExploitabilitySolver(RPS, 1)
        e_opt1 = exploitability(e_sol1, OptimalRPSStrategy)

        e_sol2 = ExploitabilitySolver(RPS, 2)
        e_opt2 = exploitability(e_sol2, OptimalRPSStrategy)

        @test ≈(e_opt1, 0.0, atol=atol)
        @test ≈(e_opt2, 0.0, atol=atol)

        e_sol1 = ExploitabilitySolver(RPS, 1)
        e_ssopt1 = exploitability(e_sol1, SlightlySuboptimalRPSStrategy)

        e_sol2 = ExploitabilitySolver(RPS, 2)
        e_ssopt2 = exploitability(e_sol2, SlightlySuboptimalRPSStrategy)

        @test e_ssopt1 > 0.0
        @test e_ssopt2 > 0.0

        e_sol1 = ExploitabilitySolver(RPS, 1)
        e_sopt1 = exploitability(e_sol1, SuboptimalRPSStrategy)

        e_sol2 = ExploitabilitySolver(RPS, 2)
        e_sopt2 = exploitability(e_sol2, SuboptimalRPSStrategy)

        @test e_sopt1 > 0.0
        @test e_sopt2 > 0.0

        @test e_sopt1 > e_ssopt1 > e_opt1
        @test e_sopt2 > e_ssopt2 > e_opt2
    end

    @testset "Kuhn" begin
        atol = 1e-3

        kuhn = CFR.Kuhn()
        e_sol1 = ExploitabilitySolver(kuhn, 1)
        e_opt1 = exploitability(e_sol1, OptimalKuhnStrategy)

        e_sol2 = ExploitabilitySolver(kuhn, 2)
        e_opt2 = exploitability(e_sol2, OptimalKuhnStrategy)

        @test ≈(e_opt1, 0.0, atol=atol)
        @test ≈(e_opt2, 0.0, atol=atol)

        e_sol1 = ExploitabilitySolver(kuhn, 1)
        e_sopt1 = exploitability(e_sol1, SuboptimalKuhnStrategy)

        e_sol2 = ExploitabilitySolver(kuhn, 2)
        e_sopt2 = exploitability(e_sol2, SuboptimalKuhnStrategy)

        @test e_sopt1 > 0.0
        @test e_sopt2 > 0.0
    end
end


import ChainRulesCore.Tangent

mutable struct BatchNorm2d
    gamma::AbstractArray
    beta::AbstractArray
    eps::Real
    momentum::Real
    training::Bool
    running_mean::AbstractArray
    running_var::AbstractArray
    count::Int
end

BatchNorm2d(num_features::Int; eps=1e-5, momentum=0.1, training=true) =
    BatchNorm2d(
        ones(1, 1, num_features, 1),
        zeros(1, 1, num_features, 1),
        eps,
        momentum,
        training,
        zeros(1, 1, num_features, 1),
        zeros(1, 1, num_features, 1),
        0
    )

Base.show(io::IO, m::BatchNorm2d) = print(io, "BatchNorm2d($(length(m.beta)))")

trainmode!(m::BatchNorm2d, train::Bool=true) = (m.training = train; nothing)

function batchnorm_impl(gamma::AbstractArray, beta::AbstractArray, x::AbstractArray,
                        mu::AbstractArray, sigma2::AbstractArray, momentum::Real; eps, training)
    x_hat = (x .- mu) ./ sqrt.(sigma2 .+ eps)
    return gamma .* x_hat .+ beta
end

function ∇batchnorm_impl(gamma::AbstractArray, beta::AbstractArray, x::AbstractArray, dy::AbstractArray,
                         mu::AbstractArray, sigma2::AbstractArray, momentum; eps, training)
    # based on:
    # https://kratzert.github.io/2016/02/12/understanding-the-gradient-flow-through-the-batch-normalization-layer.html
    sum_dims = tuple([i for i=1:ndims(x) if size(x, i) != size(mu, i)]...)   # all dims but channel
    P = prod(size(x, i) for i in sum_dims)   # product of all dims but channel
    typ = eltype(gamma)
    _1 = one(typ)        # 1 of eltype(x)
    _05 = _1 / 2         # 0.5 of eltype(x)
    # part of forward pass - we can optimize it later and take from cache
    xmu = x .- mu
    sqrtvar = sqrt.(sigma2 .+ eps)
    ivar = one(typ) ./ sqrtvar
    xhat = xmu .* ivar
    # reverse pass
    dbeta = Yota.unbroadcast(beta, dy)
    dgamma = Yota.unbroadcast_prod_x(gamma, xhat, dy)
    dxhat = dy .* gamma
    divar = sum(dxhat .* xmu; dims=sum_dims)
    dxmu1 = dxhat .* ivar
    dsqrtvar = -(sqrtvar .^ -2) .* divar
    dvar = _05 ./ sqrtvar .* dsqrtvar
    dsq = one(typ) ./ P .* dvar
    dxmu2 = 2 .* xmu .* dsq
    dx1 = (dxmu1 .+ dxmu2)
    dmu = -sum(dxmu1 .+ dxmu2; dims=sum_dims)
    # dx2 = Yota.∇mean(x, dmu, sum_dims)
    dx2 = Yota.∇mean(dmu, mean, x, sum_dims)[2]
    dx = dx1 .+ dx2
    return dgamma, dbeta, dx
end

batch_mean(x, dims) = mean(x, dims=dims)
function batch_var(x, dims)
    # note that this is equivalent to Statistics.var(x; dims=dims, corrected=false)
    # but also works for CuArrays
    mu = mean(x, dims=dims)
    return mean((x .- mu) .^ 2; dims=dims)
end

# BatchNorm layer is somewhat tricky: on one hand, we want to track running mean and var during
# training, and since these are mutating operations, we need to hide them inside primitives -
# something like `forward(m:BatchNorm, x)`. On other hand, Yota doesn't support diffrules
# w.r.t. struct fields, so `forward()` above cannot be a primitive. To overcome it, we introduce
# a helper function batchnorm2d() which pretends to be a normal primitive but:
#  * takes a @nodiff struct
#  * separately, takes learnable gamma and beta parameters
function batchnorm2d(m::BatchNorm2d, gamma, beta, x)
    if m.training
        dims = (1, 2, 4)
        mu = batch_mean(x, dims)
        sigma2 = batch_var(x, dims)
        # track running mean and variance for each channel separately
        # based on: https://github.com/ptrblck/pytorch_misc/blob/master/batch_norm_manual.py#L50-L69
        # note: running stats aren't used during training, but are required for evaluation
        m.count += 1
        exp_avg_factor = eltype(x)(1.0) / m.momentum
        m.running_mean .= @. exp_avg_factor * mu + (1 - exp_avg_factor) * m.running_mean
        m.running_var .= @. exp_avg_factor * sigma2 * m.count / (m.count - 1) +
            (1 - exp_avg_factor) * m.running_var
        return batchnorm_impl(gamma, beta, x, mu, sigma2, m.momentum;
                         eps=m.eps, training=true)
    else
        return batchnorm_impl(gamma, beta, x, m.running_mean, m.running_var, m.momentum;
                         eps=m.eps, training=false)
    end
end

(m::BatchNorm2d)(x) = batchnorm2d(m, m.gamma, m.beta, x)

function ∇batchnorm2d(dy, m::BatchNorm2d, gamma, beta, x)
    if m.training
        dims = (1, 2, 4)
        return ∇batchnorm_impl(gamma, beta, x, dy, batch_mean(x, dims), batch_var(x, dims), m.momentum;
                          eps=m.eps, training=true)
    else
        return ∇batchnorm_impl(gamma, beta, x, dy, m.running_mean, m.running_var, m.momentum;
                          eps=m.eps, training=false)
    end
end

function ChainRulesCore.rrule(::typeof(batchnorm2d), m::BatchNorm2d, gamma, beta, x)
    y = batchnorm2d(m, gamma, beta, x)
    function batchnorm2d_pullback(dy)
        dgamma, dbeta, dx = ∇batchnorm2d(dy, m, gamma, beta, x)
        return NoTangent(), Tangent{BatchNorm2d}(gamma=dgamma, beta=dbeta), dgamma, dbeta, dx
    end
    return y, batchnorm2d_pullback
end

# function register_batchnorm_derivs()
#     # @nodiff batchnorm2d(_m, _gamma, _beta, x) _m
#     # @diffrule batchnorm2d(_m, _gamma, _beta, x) _gamma getindex(∇batchnorm2d(dy, _m, _gamma, _beta, x), 1)
#     # @diffrule batchnorm2d(_m, _gamma, _beta, x) _beta getindex(∇batchnorm2d(dy, _m, _gamma, _beta, x), 2)
#     # @diffrule batchnorm2d(_m, _gamma, _beta, x) x getindex(∇batchnorm2d(dy, _m, _gamma, _beta, x), 3)
# end

function forward(m::BatchNorm2d, x)
    size(x, ndims(x)-1) == length(BN.β) ||
        error("BatchNorm expected $(length(BN.β)) channels, got $(size(x, ndims(x)-1))")
    dims = length(size(x))
    channels = size(x, dims-1)
    affine_shape = ntuple(i->i == ndims(x) - 1 ? size(x, i) : 1, ndims(x))
    m = div(prod(size(x)), channels)
    γ = reshape(BN.γ, affine_shape...)
    β = reshape(BN.β, affine_shape...)
    if !istraining()
        μ = reshape(BN.μ, affine_shape...)
        σ² = reshape(BN.σ², affine_shape...)
        ϵ = BN.ϵ
    else
        T = eltype(x)
        axes = [1:dims-2; dims] # axes to reduce along (all but channels axis)
        μ = mean(x, dims = axes)
        σ² = sum((x .- μ) .^ 2, dims = axes) ./ m
        ϵ = convert(T, BN.ϵ)
        # update moving mean/std
        mtm = BN.momentum
        S = eltype(BN.μ)
        BN.μ  = (1 - mtm) .* BN.μ .+ mtm .* S.(reshape(μ, :))
        BN.σ² = (1 - mtm) .* BN.σ² .+ (mtm * m / (m - 1)) .* S.(reshape(σ², :))
    end

    let λ = BN.λ
        x̂ = (x .- μ) ./ sqrt.(σ² .+ ϵ)
        λ.(γ .* x̂ .+ β)
    end
end


using SearchDiscoveryReplication
using Test

# Required packages
using Distributions,Random

# Test sets
@testset "Reservation values" begin

	# Search value
	F = Distributions.DiscreteNonParametric([-1/2,1/2],[1/2,1/2])
	@test round(calcZs(0,F,0.1)[1] ,digits=5)== 0.3

	F = Distributions.Normal() ;
	@test round(calcZs(0,F,0.1)[1] ,digits=5)== 0.90235
	@test round.(calcZs([0,1],F,0.1) ,digits=5) == [0.90235,1.90235]

	# Expansion value
	nd = 1 ; cs = 0.1 ; cd = 0.1
	G = Distributions.Normal()
	@test round(calcZd(G,cs,cd,nd),digits=5) == 1.20125

	G = Distributions.Normal(0.2,0.4)
	@test round(calcZd(G,cs,cd,nd),digits=5) == 0.80748

	nd = 2 ; G = Distributions.Normal()
	@test round(calcZd(G,cs,cd,nd),digits=5) == 1.58767

end

@testset "Weitzman model" begin

	m = Weitzman(
		β = [2.0,-0.4], cs = 0.2 , cd = 0.4 ,
		nd = 1, nA0 = 2 , dChars =  MvNormal([1,1],1/2), outDummy=false,
		seed = 234,dV= Normal(0,0), dE = Normal(),cfun = "linear", smo = 25,
		options=estimOptions(nDraws =3)
		)
	nCons = 5
	nProdTot = 10
	nProdCons = 3
	d = genData(m,nCons,nProdTot,nProdCons )
	@test d.path[1] == [3]
	@test d.purch[1] == 3
	@test d.search[1:3] == [1,0,1]

	cfun = SearchDiscoveryReplication.getCfun(m.cfun)
	Random.seed!(m.seed)
	maxPos = maximum(d.pos)

	@test SearchDiscoveryReplication.negLogLik(m,vcat(m.β,m.cs,m.cd),length(m.β),d,cfun,maxPos) ≈ 14.5350 atol = 1e-4

	wd = calcWelfareDemand(m,d)
	@test wd[1]		 	≈ 11.8747 atol = 1e-4
	@test wd[2]	 		≈ 12.9414 atol = 1e-4
	@test wd[4][1]	  	≈ 0.1333 atol = 1e-4
	@test wd[4][end]	≈ 0.5333 atol = 1e-4
end

@testset "FullInfo model" begin
	m = FullInfo(β = [2.0,-0.4],dChars = MvNormal([1,1],1/2), outDummy = false, seed = 234,
					dV= Normal(0,0), dE = Normal(),smo=25,options=estimOptions(nDraws =3))
	nCons = 5
	nProdTot = 10
	nProdCons = 3
	d = genData(m,nCons,nProdTot,nProdCons )
	@test d.purch[1] == 3
	Random.seed!(m.seed)
	@test SearchDiscoveryReplication.negLogLik(m,m.β,d) ≈ 42.8430 atol = 1e-4

	wd = calcWelfareDemand(m,d)
	@test wd[1]		 	≈ 12.9784 atol = 1e-4
	@test wd[4][1]		≈ 0.0666 atol = 1e-4
	@test wd[4][end]	≈ 0.4 atol = 1e-4
end

@testset "RandomSearch model" begin
	m = RandomSearch(β = [2.0,-0.4],cs = 0.3,dChars = MvNormal([1,1],1/2), outDummy = false,
			seed = 234, dV= Normal(0,0), dE = Normal(),cfun = "linear",
			smo=5,options=estimOptions(nDraws =3))
	nCons = 5
	nProdTot = 10
	nProdCons = 3
	d = genData(m,nCons,nProdTot,nProdCons )
	@test d.purch[1] == 3
	@test isequal(d.path[1],[2,3])
	Random.seed!(m.seed)
	cfun = SearchDiscoveryReplication.getCfun(m.cfun)
	@test SearchDiscoveryReplication.negLogLik(m,vcat(m. β,m.cs),d,cfun) ≈ 11.1138 atol = 1e-4

	wd = calcWelfareDemand(m,d)
	@test wd[1]			≈ 10.2291 atol = 1e-4
	@test wd[2]		 	≈ 12.2291  atol = 1e-4
	@test wd[4][1]		≈ 0.1333 atol = 1e-4
	@test wd[4][end]	≈ 0.2666 atol = 1e-4
end

@testset "SearchDiscovery model" begin
	m = SearchDiscovery(β = [1.2,0.4],cs = 0.3,cd = 0.3,dze = -0.1, dChars = MvNormal([0.4,0.5],[0.8,1.2]),
		outDummy = false, seed = 23, nd = 1 , nA0 = 1, dV= Normal(0,1), dE = Normal(),cfun = "linear", zdfun = "linear",
		options=estimOptions(nDraws =200))

	nCons = 30
	nProdTot = 10
	nProdCons = 5
	d = genData(m,nCons,nProdTot,nProdCons)
	@test isequal(d.path[1],[2,3])
	@test d.purch[1] == 3
	@test d.stop[1] == 1

	wd = calcWelfareDemand(m,d)
	@test wd[1]			≈ 45.5269 atol = 1e-4
	@test wd[2]			≈ 65.0179 atol = 1e-4
	@test wd[4][1]		≈ 0.1866 atol = 1e-4
	@test wd[4][end]	≈ 0.1195 atol = 1e-4
end


"""Reset the constants to allow for arbitrary number of fraction bits."""
function set_nfrac(n::Int)
    @assert (n > 6) && (n<max_nfrac_supported+1) "Only 6...11 fraction bits supported, $n given."

    # change pointer references based on n
    nfrac[] = n
    nint[] = 15-nfrac[]
    scale[] = 2^nfrac[]
    scale_over_logof2[] = scale[]/log(2f0)
    max_diff_resolvable[] = find_max_diff_res(scale[])

    # change Float32 lookup table for conversion LogFixPoint16 -> Float32
    f32lookup[:] = createF32LookupTable(nint[],nfrac[])

    # change addition and subtraction table
    addTable[:] = createAddLookup(scale[])
    subTable[:] = createSubLookup(scale[])

    if rounding_mode[] == :lin
        c_b[] = rounding_correction(scale[])
    else
        c_b[] = 0
    end

    @warn "LogFixPoint16 was changed to $(nint[]) integer and $(nfrac[]) fraction bits."

    return nothing
end

"""Change the rounding mode from round-to-nearest in either linear or logarithmic space."""
function set_rounding_mode(mode::Symbol=:lin)
    @assert mode in [:lin,:log] "only mode :lin or :log allowed."

    rounding_mode[] = mode

    if rounding_mode[] == :lin
        c_b[] = rounding_correction(scale[])
    else
        c_b[] = 0
    end

    # recalculate lookup tables
    max_diff_resolvable[] = find_max_diff_res(scale[])
    f32lookup[:] = createF32LookupTable(nint[],nfrac[])
    addTable[:] = createAddLookup(scale[])
    subTable[:] = createSubLookup(scale[])

    @warn "LogFixPoint16 rounding mode changed to round to nearest in $(string(mode))-space."
end


# Compute a complete set of sums using a fast recursive algorithm based on a simple combinatorial relationship between the required sums.
# This code has been partially optimized for quick updates, which means that it is more complicated that it would need to be if we were just updating from scratch every time.

# Throughout these functions:
# - N::Dict{Vector{Int64},<:Integer} is a Dict mapping sorted partition vectors to their volume sums.
# - M::Dict{Vector{Int64},<:Integer} is a Dict mapping sorted partition vectors to their *uncorrected* sums. By *uncorrected*, we mean that these sums must be multipled by a set of combinatorial constants prior to including in the modularity function. Each entry of M differs from the corresponding entry of N by a combinatorial constant.
# - p::Vector{Int64} is a partition vector
# - D::Array{Int64, 1} is the degree sequence of the hypergraph, which does not change.
# - Z::Array{Int64, 1} is the group label vector, which does change.
# - ℓ::Int64 is the number of groups.

# ------------------------------------------------------------------------------
# COMPUTATION OF SUMS FROM SCRATCH
# ------------------------------------------------------------------------------

function correctOvercounting(M::Dict{Vector{Int64},<:Integer}, p::Vector{Int64})
    """
    Internal function: no reason to call this outside of evalSums().
    Computes the second term in the recurrence relation given in the paper.
    M::Dict{Vector{Int64},<:Integer}, a Dict mapping sorted partition vectors to their *uncorrected* volume-sums,
    p::Vector{Int64}, the partition for which we are calculating.
    returns S::bigInt, the second term in the recurrence relation.
    """
    pk = p[end]
    S = 0
    for i = 1:length(p)-1
        p_ = copy(p)[1:(end-1)]
        p_[i] += pk
        S += M[-sort(-p_)]
    end
    return(S)
end

function computeMoments(Z::Vector{<:Integer}, D::Vector{<:Integer}, r::Integer, ℓ::Integer=0)
    """
    Compute the vectors V of group volumes and μ of volume-moments.
    No reason to call this outside of evalSums()
    Z::Vector{Int64}, the vector of cluster labels.
    D::Vector{Int64}, the degree sequence
    r::Int64, the size of the largest hyperedge
    ℓ::Int64, the number of groups (defaults to maximum(Z))
    returns V::Array{Int64, 1} the vector of group volumes and μ::Vector{Int64} of volume-moments.
    NOTE: might want to implement bigInts in μ
    """
    if ℓ==0
        ℓ=maximum(Z)
    end

    V = [sum(D[Z .== j]) for j in 1:ℓ]
    μ = [sum(V.^i) for i = 1:r]
    return(V, μ)
end

function evalConstants(r::Int64; ℓ = r, bigInt::Bool=true)
    """
    Evaluate the combinatorial constants required to convert the "convenient" sums M to the required sums N required for modularity calculations.
    r::The size of the largest hyperedge.
    bigInt:Bool, whether to use bigInts for this computation, recommended unless the instance is VERY small.
    return: C::Dict{Array{Int64, 1}, bigInt} a Dict mapping a partition vector to its associated combinatorial constant.
    """

    C =
        if bigInt C = Dict{Vector{Int64},BigInt}()
        else      C = Dict{Vector{Int64},Int64}()
        end

    for i = 1:r, j = 1:minimum([i,ℓ]), p in Combinatorics.partitions(i, j)
        orderCorrection =
            if bigInt prod([factorial(big(c)) for c in values(StatsBase.countmap(p))])
            else      prod([factorial(c) for c in values(StatsBase.countmap(p))])
            end
        # orderCorrection = 1
        C[p] = Combinatorics.multinomial(big.(p)...) ÷ orderCorrection
    end
    return(C)
end;

function evalSums(Z::Vector{<:Integer}, D::Vector{Int64}, r::Integer; constants::Bool=true, ℓ::Integer=maximum(Z), bigInt::Bool=true)
    """
    Z::Array{Int64, 1} the vector of integer group labels
    D::Array{Int64, 1} the vector of degrees
    r::Int64, the largest hyperedge size to compute
    constants::Bool, whether to postmultiply by the combinatorial constants computed by evalConstants prior to returning the values. Choose true if computing modularity directly using the results, choose false if performing incremental updates a la Louvain
    ℓ::Int64, the number of groups (defaults to maximum(Z))
    bigInt::Bool, whether to use bigInt conversions for D and Z. Recommended unless the instance is VERY small.
    return (V::Array{bigInt,1}, μ::Array{bigInt,1}, M::Dict{Vector{Int64},bigInt}, with V and μ as described in computeMoments(), and M::Dict{Vector{Int64},bigInt} is the array of (optionally uncorrected) volume sums.
    """

    if bigInt
        D = convert(Array{BigInt,1}, D)
    end

    V, μ = computeMoments(Z, D, r, ℓ)

    M = bigInt ? Dict{Vector{Int64}, BigInt}() : Dict{Vector{Int64}, Int64}()

    ix = ℓ == 1 ? length(Z) : ℓ

    for i = 1:r, j = 1:minimum([i,ix]), p in Combinatorics.partitions(i, j)
        M[p] = μ[p[end]]*get(M, p[1:(end-1)], 1) - correctOvercounting(M,p)
    end

    if constants
        C = evalConstants(r; bigInt=bigInt, ℓ = ix)
        for p in keys(M)
            M[p] *= C[p]
        end
    end
    return V, μ, M
end

# Extra methods for evalSums for working with degree dicts and the custom hypergraph class.

function evalSums(Z::Vector{<:Integer}, H::hypergraph, ℓ::Integer=maximum(Z), bigInt::Bool=true, constants::Bool = true)
    r = maximum(keys(H.E))
    return evalSums(Z, H.D, r; constants = constants, ℓ = ℓ, bigInt = bigInt)
end

# ------------------------------------------------------------------------------
# INCREMENTS
# ------------------------------------------------------------------------------

function momentIncrements(V::Vector{<:Integer}, μ::Vector{<:Integer}, i::Int64, t::Int64, D::Vector{<:Integer}, Z::Vector{<:Integer})
    """
    Compute the update for the vectors V and μ associated with moving node i from its current group to group t.
    V::Array{Integer, 1}, the vector of group volumes.
    μ::Array{Integer, 1}, the vector of volume-moments.
    i::Int64, the node to move
    t::Int64, the proposed new group for node i
    D::Array{Integer, 1}, the degree vector
    Z::Array{Integer, 1}, the vector of group labels
    returns ΔV, Δμ, the required updates in V and μ
    """

    ΔV = zero(V)
    # update V
    s      = Z[i];
    ΔV[t] += D[i];
    ΔV[s] -= D[i];

    r = length(μ)
    Δμ = [(V[s] + ΔV[s])^j - V[s]^j + (V[t] + ΔV[t])^j - V[t]^j for j in 1:r]

    return ΔV, Δμ
end

function momentIncrements(V::Vector{<:Integer}, μ::Vector{<:Integer}, I_::Vector{<:Integer}, t::Integer, D::Vector{<:Integer}, Z::Vector{<:Integer})
    """
    Compute the update for the vectors V and μ associated with moving all the nodes listed in I from their group to group t.
    THIS CODE ASSUMES THAT THE NODES IN I ARE ALL IN THE SAME GROUP, AND WILL GIVE INCORRECT RESULTS OTHERWISE.
    V::Array{Integer, 1}, the vector of group volumes.
    μ::Array{Integer, 1}, the vector of volume-moments.
    I::Vector{Int64}, the list of multiple nodes to move
    t::Int64, the proposed new group for node i
    D::Array{Integer, 1}, the degree vector
    Z::Array{Integer, 1}, the vector of group labels
    returns ΔV, Δμ, the required updates in V and μ
    """

    ΔV = zero(V)
    # update V
    s  = Z[I_[1]];

    for i ∈ I_
        ΔV[t] += D[i];
        ΔV[s] -= D[i];
    end

    r = length(μ)
    Δμ = [(V[s] + ΔV[s])^j - V[s]^j + (V[t] + ΔV[t])^j - V[t]^j for j in 1:r]

    return ΔV, Δμ
end

function increments(V::Vector{<:Integer}, μ::Vector{<:Integer}, M::Dict{Vector{Int64},<:Integer}, I_::T, t::Integer, D::Vector{<:Integer}, Z::Vector{<:Integer}; bigInt::Bool=true) where T <: Union{<:Integer, Vector{<:Integer}}
    """
    Compute the update for the vectors V and μ, as well as the uncorrected sums M, associated with moving node(s) I from its current group to group t.
    THIS CODE ASSUMES THAT THE NODES IN I ARE ALL IN THE SAME GROUP, AND WILL GIVE INCORRECT RESULTS OTHERWISE.
    V::Array{Integer, 1}, the vector of group volumes.
    μ::Array{Integer, 1}, the vector of volume-moments.
    M::Dict{Array{Int64, 1}, bigInt}, the Dict of uncorrected volume sums.
    i::Int64, the node to move
    t::Int64, the proposed new group for node i
    D::Array{Integer, 1}, the degree vector
    Z::Array{Integer, 1}, the vector of group labels
    returns (ΔV, Δμ, ΔM), the required updates in V, μ, and M
    """

    ΔV, Δμ = momentIncrements(V, μ, I_, t, D, Z)

    # compute increments in M using recursion formula from notes
    if bigInt ΔM = Dict{Vector{Int64}, BigInt}()
    else      ΔM = Dict{Vector{Int64}, Int64}()
    end

    r = maximum([sum(p) for p in keys(M)])

    for i = 1:r, j = 1:i, p in Combinatorics.partitions(i, j)
        ΔM[p] = Δμ[p[end]]*get(M, p[1:(end-1)], 1) + μ[p[end]]*get(ΔM, p[1:(end-1)], 0) + Δμ[p[end]]*get(ΔM, p[1:(end-1)], 0) - correctOvercounting(ΔM,p)
    end
    return(ΔV, Δμ, ΔM)
end

function addIncrements(V::Vector{<:Integer}, μ::Vector{<:Integer}, M::Dict{Vector{Int64},<:Integer}, ΔV::Vector{<:Integer}, Δμ::Vector{<:Integer}, ΔM::Dict{Vector{Int64},<:Integer})
    """
    Add the increments (ΔV, Δμ, ΔM) to (V, μ, M), entrywise in the case of M.
    V::Array{Integer, 1}, the vector of group volumes.
    μ::Array{Integer, 1}, the vector of volume-moments.
    M::Dict{Array{Int64, 1}, bigInt}, the Dict of uncorrected volume sums.
    ΔV, Δμ, ΔM: increments in each of the above quantities returned by increments()
    """
    M̃ = Dict(p => M[p] + ΔM[p] for p in keys(M))
    return(V + ΔV, μ + Δμ, M̃)
end

# ------------------------------------------------------------------------------
# COMPLETE COMPUTATION OF SECOND (VOLUME) TERM IN MODULARITY
# ------------------------------------------------------------------------------

function aggregateSums(M, Ω::IntensityFunction)
    M̂ = Dict()
    for p in keys(M)
        p̂ = Ω.aggregator(p)
        M̂[p̂] = get(M̂, p̂, 0) + M[p]
    end
    return M̂
end

function second_term_eval(H::hypergraph, Z::Vector{<:Integer}, Ω::IntensityFunction; α, ℓ::Int64 = 0, bigInt::Bool=true)
    """
    Naive implementation, computes sums from scratch.
    H::hypergraph
    Z::Array{Int64, 1}, the group label vector.
    ℓ::Int64, number of groups
    Ω: group interation function (e.g., planted partition). Needs to have a mode argument which, when set to value "partition", will cause evaluation on partition vectors rather than label vectors.
    bigInt::Bool, whether to use bigInt conversions. Strongly recommended unless the instance is VERY small.
    """

    obj = 0.0

    if ℓ == 0
        ℓ = maximum(Z)
    end

    V, μ, M = evalSums(Z, H, ℓ, bigInt)
    M̂ = aggregateSums(M, Ω)
    for p̂ in keys(M̂)
        if M̂[p̂] >= 1
            obj += Ω.ω(p̂, α)*M̂[p̂]
        end
    end
    return obj
end


# using Revise
# using Pkg; Pkg.activate(".")
using HyperModularity
using Test

## Reading in a dataset
dataset = "contact-high-school-classes"
maxsize = 5
minsize = 2
return_labels = true
H, L = read_hypergraph_data(dataset,maxsize,minsize,return_labels)

## Run graph Louvain on reduced graph

# weighted = true means that when clique expandin hyperedge e
# weight all edges by 1/(|e|-1)
# weighted = false means they are weighted by 1, but the
# resulting graph is still integer weighted
weighted = true

# binary = true means that after clique expansion, all edges
# are unweighted. If true, "weighted" flag doesn't matter
binary = false
A = CliqueExpansion(H,weighted,binary)

# Run modularity on clique expansion
gamma = 1.0
Z = CliqueExpansionModularity(H,gamma;weighted=true,randflag=false,binary=false,clusterpenalty=0.0,maxits=1000)

# Compute the MLE for the resolution parameter based on Z
(ωᵢ, ωₒ) = computeDyadicResolutionParameter(H, Z; mode = 0)
γ = (ωᵢ - ωₒ)/(log(ωᵢ) - log(ωₒ))

# log-likelihood
loglike = dyadicLogLikelihood(H, Z, ωᵢ, ωₒ)

# Above steps can be iteratively repeated to reach global max for log-likelihood
# Good place to test log likelihood computations


using LinearAlgebra
using Makie
using Combinatorics
using AbstractAlgebra
import Quaternions
H = Quaternions.Quaternion

const TOLERANCE = 1e-7
const GOLDEN_RATIO = (1 + sqrt(5)) / 2
const RADIUS = 0.025
const SEGMENTS = 60

"""
σ(q)

Sends a point on a 3-sphere to a point in the plane x₄=0 with the given point.
This is the stereographic projection of a 3-sphere. S³ ↦ R³
"""
function σ(q)
    Point3f0(q.s, q.v1, q.v2) ./ (1 - q.v3)
end

"""
π(q)

Sends a point on a unit 3-sphere to a point on a unit 2-sphere with the given
point. S³ ↦ S²
"""
function π(q)
    x₁ = q.s
    x₂ = q.v1
    x₃ = q.v2
    x₄ = q.v3
    Point3f0(2(x₁*x₃ + x₂*x₄), 2(x₂*x₃ - x₁*x₄), x₁^2+x₂^2-x₃^2-x₄^2)
end

"""
S¹action(α, q)

Performs a group action corresponding to moving along the circumference of a
circle with the given angle and the point.
"""
function S¹action(α, q)
    z₁ = Complex(q.s, q.v1)
    z₂ = Complex(q.v2, q.v3)
    αz₁, αz₂ = exp(im * α) * z₁, exp(im * α) * z₂
    H(real(αz₁), imag(αz₁), real(αz₂), imag(αz₂))
end

"""
rotate3D(point, q)

Rotates a point in the 3D space with the given point and the unit quaternion.
"""
function rotate3D(point, q)
    p = H(0, point[1], point[2], point[3])
    R = conj(q) * p * q
    Point3f0(R.v1, R.v2, R.v3)
end

"""
get_center(A, B, C)

Finds the center point of the fiber circle under stereographic projection
with the given 3 points on the circle circumference.
"""
function get_center(A, B, C)
    a = norm(B - C)
    b = norm(A - C)
    c = norm(A - B)
    numerator = a^2 * (b^2 + c^2 - a^2) * A +
                b^2 * (a^2 + c^2 - b^2) * B +
                c^2 * (a^2 + b^2 - c^2) * C
    denominator = a^2 * (b^2 + c^2 - a^2) +
                  b^2 * (a^2 + c^2 - b^2) +
                  c^2 * (a^2 + b^2 - c^2)
    numerator / denominator
end

"""
compressed_σ(q)

Sends a point on a 3-sphere to a point in the compressed plane x₄=0 with the
given point. This is the stereographic projection of a 3-sphere. S³ ↦ R³
"""
function compressed_σ(q)
    sigmoid(x) = exp(x) / (1 + exp(x))
    2 .* (sigmoid.(σ(q)) .- 0.5)
end

"""
get_rotation_axis(i, n)

Calculates a unit quaternion corresponding to the axis of a 3D rotation with
the given initial and end unit vectors.
"""
function get_rotation_axis(i, n)
    u = normalize(cross(n, i))
    # The angle of rotation
    dotproduct = dot(n, i)
    if isapprox(dotproduct, 1, atol=TOLERANCE)
        θ = acos(1)
    elseif isapprox(dotproduct, -1, atol=TOLERANCE)
        θ = acos(-1)
    else
        θ = acos(dotproduct)
    end
    H(cos(θ/2), (sin(θ/2).*u)...)
end

"""
get_fiber(point; segments, r=0.025)

Calculates a torus of revolution for building a surface in a specific way with
the given point in the base space, the number of segments, the number of
samples and the radius of the smaller circle in the torus of revolution.
"""
function get_fiber(point; segments = 30, r = 0.025)
    # Find 3 points on the circle
    A = σ(S¹action(pi / 5, point))
    B = σ(S¹action(pi / 4, point))
    C = σ(S¹action(pi / 3, point))
    # The fiber circle center
    Q = get_center(A, B, C)
    # The bigger radius
    R = norm(Q - A)
    A = A ./ R
    B = B ./ R
    C = C ./ R
    Q = Q ./ R
    # The bigger radius
    R = norm(Q - A)
    # The normal to the initial circle
    i = [0.0, 0.0, 1.0]
    # Get the normal to the rotated circle
    n = normalize(normalize(cross(A - Q, B - Q)))
    # The first rotation versor
    q = get_rotation_axis(i, n)
    # The vector pointing to the positive direction of the x axis on the
    # plane of the rotated circle
    i = rotate3D([1.0, 0.0, 0.0], q)
    # The second rotation versor
    q′ = get_rotation_axis(i, normalize(A - Q))
    # Construct a torus of revolution grid
    fiber = Array{Point3f0}(undef, segments, segments)
    for i in 1:segments
        for j in 1:segments
            longitude = j * 2pi / (segments - 1)
            latitude = i * 2pi / (segments - 1)
            x₁ = (R + r * cos(longitude)) * cos(latitude)
            x₂ = (R + r * cos(longitude)) * sin(latitude)
            x₃ = r * sin(longitude)
            fiber[i, j] = rotate3D(Point3f0(x₁, x₂, x₃), q * q′) + Q
        end
    end
    fiber
end

"""
iseven(v, x)

Checks the parity of a permutation to see whether it is even or not with the
given original vector and the permutated vector.
"""
function iseven(v, x)
    # Get a trivial sequence of the given permutation
    g = Perm([findfirst(isequal(x[i]), v) for i in 1:length(x)])
    if AbstractAlgebra.parity(g) == 0
        return true
    else
        return false
    end
end

"""
get_vertices()

Returns the versors corresponding to the 120 vertices of a 600-cell inscribed
inside of a 3-sphere.
"""
function get_vertices()
    set = []
    set1 = [getindex.(Ref([1/2, -1/2]), 1 .+ digits(i-1; base=2, pad=4))
            for i=1:2^4]
    set2 = [i for i in permutations([0, 0, 0, 1]) |> collect]
    set3 = [i for i in permutations([0, 0, 0, -1]) |> collect]
    v = [[GOLDEN_RATIO, 1, 1/GOLDEN_RATIO, 0],
         [GOLDEN_RATIO, 1, -1/GOLDEN_RATIO, 0],
         [GOLDEN_RATIO, -1, 1/GOLDEN_RATIO, 0],
         [GOLDEN_RATIO, -1, -1/GOLDEN_RATIO, 0],
         [-GOLDEN_RATIO, 1, 1/GOLDEN_RATIO, 0],
         [-GOLDEN_RATIO, 1, -1/GOLDEN_RATIO, 0],
         [-GOLDEN_RATIO, -1, 1/GOLDEN_RATIO, 0],
         [-GOLDEN_RATIO, -1, -1/GOLDEN_RATIO, 0]] .* 1/2
    set4 = []
    for i in 1:length(v)
        f(x) = iseven(v[i], x)
        push!(set4, filter(f, [j for j in permutations(v[i]) |> collect])...)
    end
    push!(set, set1..., set2..., set3..., set4...)
    map((x) -> H(x...), unique(set))
end

"""
build_surface(scene, points, color; transparency, shading)

Builds a surface with the given scene, points, color, transparency and shading.
"""
function build_surface(scene,
                       points,
                       color;
                       transparency = false,
                       shading = true)
    x = @lift([$points[i, j][1]
               for i in 1:size($points, 1), j in 1:size($points, 2)])
    y = @lift([$points[i, j][2]
               for i in 1:size($points, 1), j in 1:size($points, 2)])
    z = @lift([$points[i, j][3]
               for i in 1:size($points, 1), j in 1:size($points, 2)])
    surface!(scene,
             x,
             y,
             z,
             color = color,
             transparency = transparency,
             shading = shading)
end

"""
get_faces(vertices, distance)

Calculates the faces in a manifold with the given vertices and the distance
between each vertex in a face.
"""
function get_faces(vertices, distance)
    array = []
    for i in 1:length(vertices)
        for j in i:length(vertices)
            for k in j:length(vertices)
                a = vertices[i]
                b = vertices[j]
                c = vertices[k]
                ab = norm(a - b)
                ac = norm(a - c)
                bc = norm(b - c)
                if isapprox(ab, distance, atol = TOLERANCE) &&
                   isapprox(ac, distance, atol = TOLERANCE) &&
                   isapprox(bc, distance, atol = TOLERANCE)
                    push!(array, [a, b, c])
                end
            end
        end
    end
    unique(array)
end

"""
get_Hopf_vertices(vertices)

Filters the extra vertices that are on the same Hopf fiber with the given
vertices on a 3-sphere.
"""
function get_Hopf_vertices(vertices)
    array = []
    for i in 1:length(vertices)
        a = π(vertices[i])
        exists = false
        for j in 1:length(array)
            b = π(array[j])
            if isapprox(a, b, atol = TOLERANCE)
                exists = true
            end
        end
        if !exists
            push!(array, vertices[i])
        end
    end
    unique(array)
end

# The scene object that contains other visual objects
scene = Scene(backgroundcolor = :white, show_axis=false, resolution = (360, 360))
# Use a slider for rotating the base space in an interactive way
sθ, θ = textslider(0:0.01:2pi, "θ", start = 0)
sϕ, ϕ = textslider(0:0.01:2pi, "ϕ", start = 0)
sψ, ψ = textslider(0:0.01:2pi, "ψ", start = 0)
# Calculate a unit quaternion as the rotation axis
u = @lift([cos($ϕ) * cos($ψ), cos($ϕ) * sin($ψ), sin($ϕ)])
q = @lift(H(cos($θ), sin($θ) * $u[1], sin($θ) * u[2], sin($θ) * $u[3]))
vertices = get_vertices()
# Cut a 3-sphere into 2 pieces and get the first one
#f(x) = abs(Complex(x.s, x.v1)) < abs(Complex(x.v2, x.v3))
#vertices = filter(f, total_vertices)
fiber_vertices = get_Hopf_vertices(vertices)
faces = get_faces(vertices, 1/GOLDEN_RATIO)
α₁ = Node(0.0)
α₂ = Node(1.0)

for vertex in fiber_vertices
    fiber = @lift(get_fiber($q * vertex, segments = SEGMENTS, r = RADIUS) ./ GOLDEN_RATIO)
    color = @lift(fill(RGBAf0(compressed_σ($q * vertex)..., $α₁), SEGMENTS, SEGMENTS))
    build_surface(scene, fiber, color, shading = true, transparency = true)
end

for i in 1:length(faces)
    a, b, c = faces[i]
    rotateda = @lift(compressed_σ($q * a))
    rotatedb = @lift(compressed_σ($q * b))
    rotatedc = @lift(compressed_σ($q * c))
    l = @lift([$rotateda, $rotatedb, $rotatedc])
    indices = [1, 2, 3,   2, 3, 1,   3, 1, 2]
    color = @lift([RGBAf0($rotateda..., rand() * $α₂),
                   RGBAf0($rotatedb..., rand() * $α₂),
                   RGBAf0($rotatedc..., rand() * $α₂)])
    mesh!(scene, l, color = color)
end

# update eye position
eye_position, lookat = Vec3f0(2, 2, 2), Vec3f0(0)
update_cam!(scene, eye_position, lookat)
scene.center = false # prevent scene from recentering on display

#fullscene = hbox(scene, vbox(sθ, sϕ, sψ), parent = Scene(resolution = (500, 500)))

record(scene, "600-cell.gif") do io
    frames = 90
    for i in 1:frames
        recordframe!(io) # record a new frame
        # animate the scene
        eye_position = Vec3f0(2 - 2sin(i / frames * pi) + 0.00001,
                              2 - 2sin(i / frames * pi) + 0.00001,
                              2 + 2sin(i / frames * pi))
        update_cam!(scene, eye_position, lookat)
        α₁[] = cos(i / frames * 2pi)
        α₂[] = 1 - sin(i / frames * 2pi)
        θ[] = i / frames * 2pi
        ψ[] = cos(i / frames * pi/2) * 2pi
            ϕ[] = sin(i / frames * pi/2) * 2pi
    end
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    DirectionPartition(direction; tol=1e-6)

A method for partitioning spatial objects along a given `direction`
with bandwidth tolerance `tol`.
"""
struct DirectionPartition{Dim,T} <: SPredicatePartitionMethod
  direction::Vec{Dim,T}
  tol::Float64

  function DirectionPartition{Dim,T}(direction, tol) where {Dim,T}
    new(normalize(direction), tol)
  end
end

DirectionPartition(direction::Vec{Dim,T}; tol=1e-6) where {Dim,T} =
  DirectionPartition{Dim,T}(direction, tol)

DirectionPartition(direction::NTuple{Dim,T}; tol=1e-6) where {Dim,T} =
  DirectionPartition(Vec(direction), tol=tol)

function (p::DirectionPartition)(x, y)
  δ = x - y
  d = p.direction
  norm(δ - (δ⋅d)*d) < p.tol
end


"""
## fmplat!

This subroutine forms the 2nd derivatives for rectangular
plate bending fin_els.

### Method
```julia
fmplat!(d2x, d2y, d2xy, points, aa, bb, i)
```

### Arguments
```julia
* d2x::Vector{Float64}       : x derivative term (Updated)
* d2y::Vector{Float64}       : y derivative term (Updated)
* d2xy::Vector{Float64}      : x,y derivative term (Updated)
* points::Matrix{Float64}    : Location of Gauss points
* aa::Float64                : Dimension of plate
* bb::Float64                : Dimension of plate
* i::Int                   : Gauss point to use
```
"""
function fmplat!(d2x::Vector{Float64}, d2y::Vector{Float64}, d2xy::Vector{Float64},
  points::Matrix{Float64}, aa::Float64, bb::Float64, i::Int)
  #
  # This subroutine forms the 2nd derivatives for rectangular
  # plate bending fin_els.
  #
  # Updates d2x, d2y and d2xy vectors
  #
 x = points[i,1]
 e = points[i,2]
 xp1 = x+1.0
 xp12 = xp1*xp1
 xp13 = xp12*xp1
 ep1 = e+1.0
 ep12 = ep1*ep1
 ep13 = ep12*ep1
 p1 = 1.0-0.75*xp12+0.25*xp13
 q1 = 1.0-0.75*ep12+0.25*ep13
 p2 = 0.5*aa*xp1*(1.0-xp1+0.25*xp12)
 q2 = 0.5*bb*ep1*(1.0-ep1+0.25*ep12)
 p3 = 0.25*xp12*(3.0-xp1)
 q3 = 0.25*ep12*(3.0-ep1)
 p4 = 0.25*aa*xp12*(0.5*xp1-1.0)
 q4 = 0.25*bb*ep12*(0.5*ep1-1.0)
 dp1 = 1.5*xp1*(0.5*xp1-1.0)
 dq1 = 1.5*ep1*(0.5*ep1-1.0)
 dp2 = aa*(0.5-xp1+0.375*xp12)
 dq2 = bb*(0.5-ep1+0.375*ep12)
 dp3 = 1.5*xp1*(1.0-0.5*xp1)
 dq3 = 1.5*ep1*(1.0-0.5*ep1)
 dp4 = 0.5*aa*xp1*(0.75*xp1-1.0)
 dq4 = 0.5*bb*ep1*(0.75*ep1-1.0)
 d2p1 = 1.5*x
 d2p2 = 0.25*aa*(3.0*x-1.0)
 d2p3 = -d2p1
 d2p4 = 0.25*aa*(3.0*x+1.0)
 d2q1 = 1.5*e
 d2q2 = 0.25*bb*(3.0*e-1.0)
 d2q3 = -d2q1
 d2q4 = 0.25*bb*(3.0*e+1.0)
 d2x[1] = d2p1*q1
 d2x[2] = d2p2*q1
 d2x[3] = d2p1*q2
 d2x[4] = d2p2*q2
 d2x[5] = d2p1*q3
 d2x[6] = d2p2*q3
 d2x[7] = d2p1*q4
 d2x[8] = d2p2*q4
 d2x[9] = d2p3*q3
 d2x[10] = d2p4*q3
 d2x[11] = d2p3*q4
 d2x[12] = d2p4*q4
 d2x[13] = d2p3*q1
 d2x[14] = d2p4*q1
 d2x[15] = d2p3*q2
 d2x[16] = d2p4*q2
 d2y[1] = p1*d2q1
 d2y[2] = p2*d2q1
 d2y[3] = p1*d2q2
 d2y[4] = p2*d2q2
 d2y[5] = p1*d2q3
 d2y[6] = p2*d2q3
 d2y[7] = p1*d2q4
 d2y[8] = p2*d2q4
 d2y[9] = p3*d2q3
 d2y[10] = p4*d2q3
 d2y[11] = p3*d2q4
 d2y[12] = p4*d2q4
 d2y[13] = p3*d2q1
 d2y[14] = p4*d2q1
 d2y[15] = p3*d2q2
 d2y[16] = p4*d2q2
 d2xy[1] = dp1*dq1
 d2xy[2] = dp2*dq1
 d2xy[3] = dp1*dq2
 d2xy[4] = dp2*dq2
 d2xy[5] = dp1*dq3
 d2xy[6] = dp2*dq3
 d2xy[7] = dp1*dq4
 d2xy[8] = dp2*dq4
 d2xy[9] = dp3*dq3
 d2xy[10] = dp4*dq3
 d2xy[11] = dp3*dq4
 d2xy[12] = dp4*dq4
 d2xy[13] = dp3*dq1
 d2xy[14] = dp4*dq1
 d2xy[15] = dp3*dq2
 d2xy[16] = dp4*dq2
 #(d2x, d2y, d2xy)
end


module test_suite_applications

using BasisFunctions, DomainSets, FrameFun, LinearAlgebra
using Test

@testset "FourierExtensionPlatform: DiscreteGramStyle" begin
    F = FourierExtensionPlatform(0.0..0.5)
    opts = (samplingstyle=FrameFun.DiscreteGramStyle(),)
    P = dictionary(F, 4)
    D = azdual_dict(F,4;opts...)
    M = measure(F,4;opts...)
    A = SynthesisOperator(P, M)
    Z = SynthesisOperator(D, M)
    Zt = Z'
    G1 = Zt*A
    G2 = BasisFunctions.default_mixedgram(D, P, M;warnslow=false)
    G3 = BasisFunctions.mixedgram(D, P, M)
    @test Matrix(G1) ≈ Matrix(G2)
    @test Matrix(G2) ≈ Matrix(G3)
end

@testset "FourierExtensionPlatform: GramStyle" begin
    F = FourierExtensionPlatform(0.0..0.5)
    opts = (samplingstyle=FrameFun.GramStyle(),)
    P = dictionary(F, 4)
    D = azdual_dict(F,4;opts...)
    M = measure(F,4;opts...)
    A = SynthesisOperator(P, M)
    Z = SynthesisOperator(D, M)
    @info "Two `Slow computation of Gram matrix entrywise.` warnings follow."
    Zt = Z'
    G1 = Zt*A
    G2 = BasisFunctions.default_mixedgram(D, P, M;warnslow=false)
    G3 = BasisFunctions.mixedgram(D, P, M)
    @test Matrix(G1) ≈ Matrix(G2)
    @test Matrix(G2) ≈ Matrix(G3)

    n = 101
    F = FourierExtensionPlatform(0.0..0.5)
    opts = (samplingstyle=FrameFun.GramStyle(),warnslow=false)
    A = AZ_A(F, n; opts...)
    Z = AZ_Z(F, n; opts...)
    MG = Z'*A
    @test sum(.0001 .< svdvals(MG) .< .9999) == 9
end

@testset "FourierExtensionPlatform: GenericOperatorStyle" begin
    F = FourierExtensionPlatform(0.0..0.5)
    opts = (samplingstyle=FrameFun.GramStyle(),
            problemstyle=FrameFun.GenericOperatorStyle(), atol=1e-4,rtol=1e-4,warnslow=false)
    n = 11
    f = Fun(exp, F, n; opts...)
    x = .1923
    @test abs(f(x) - exp(x)) < 1e-3

    opts = (samplingstyle=FrameFun.DiscreteGramStyle(),
            problemstyle=FrameFun.GenericOperatorStyle(), atol=1e-4,rtol=1e-4,warnslow=false)
    f = Fun(exp, F, n; opts...)
    @test abs(f(x) - exp(x)) < 1e-3
end

@testset "WeightedSumPlatform" begin
    # Simple platform construction test
    n = 12
    domain = 0.0..0.5
    P = FourierExtensionPlatform(domain)
    WP = WeightedSumPlatform(P,x->sqrt(x),x->1)
    f = x->sqrt(x)*(1-x)-exp(x)
    F = Fun(f, WP, (n,n), solverstyle=AZStyle())
    rgrid = randomgrid(domain, 200)
    abserror = sum(abs.(F.(rgrid)-f.(rgrid)))/length(rgrid)
    @test abserror<2e-7
end

end # module


#
#   propagator.jl
#
#   Includes functions and structures for propagating orbits
#   within the two-body problem.
#

"""
Currently not exported. Used for ideal two-body numerical integration.
"""
function RestrictedTwoBodyTic!(∂u, u, p, t)
    ∂u[1:3] .= u[4:6]
    ∂u[4:6] .= -p.μ .* (u[1:3] ./ norm(u[1:3])^3)
    return nothing
end

"""
Currently not exported. Used for ideal two-body numerical integration.
"""
function RestrictedBiasedTwoBodyTic!(∂u, u, p, t)
    ∂u.rᵢ =  u.vᵢ
    ∂u.vᵢ = (-p.μ .* (u.rᵢ ./ norm(u.rᵢ,2)^3)) .+ (normalize(u.vᵢ) .* p.T)
    return nothing
end

function SciMLBase.ODEProblem(orbit::CartesianOrbit, Δt::Real = ustrip(timeunit(orbit.state), period(orbit)))
    u = orbit.state.rv
    t = (orbit.state.t, Δt)
    p = (μ = ustrip(lengthunit(orbit.state)^3 / timeunit(orbit.state)^2, mass_parameter(orbit.system)),)
    return ODEProblem(RestrictedTwoBodyTic!, u, t, p)
end
SciMLBase.ODEProblem(orbit::CartesianOrbit, Δt::Unitful.Time) = ODEProblem(orbit, ustrip(timeunit(orbit.state), Δt))

"""
Uses OrdinaryDiffEq solvers to propagate `orbit` Δt into the future.
All keyword arguments are passed directly to OrdinaryDiffEq solvers.

References:
* [1] https://diffeq.sciml.ai/v4.0/tutorials/ode_example.html
* [2] https://github.com/SciML/DifferentialEquations.jl/issues/393#issuecomment-658210231
* [3] https://discourse.julialang.org/t/smart-kwargs-dispatch/14571/15
"""
function propagate(orbit::CartesianOrbit,
                   Δt = period(orbit);
                   thrust::Unitful.Acceleration = 0.0u"N/kg",
                   kwargs...)

    # Set default kwargs
    defaults = (;  reltol=1e-14, abstol=1e-14)
    options = merge(defaults, kwargs)

    # Initial conditions
    problem = ODEProblem(orbit, Δt)

    # Integrate!
    sols = solve(problem; options...)

    # Return PropagationResult structure
    if sols.retcode != :Success
        @warn "DifferentialEquations solvers returned code $(string(sols.retcode))."
    end

    return [
        CartesianOrbit(sols.u[i][1:3] * lengthunit(orbit.state), sols.u[i][4:6] * velocityunit(orbit.state), orbit.system, sols.t[i] * timeunit(orbit.state))
        for i ∈ 1:length(sols.t)
    ]
end

propagate(::KeplerianOrbit, args...; kwargs...) = throw(ArgumentError("Keplerian orbit propagation isn't supported at this time! Please convert your orbit to a `CartesianOrbit` for propagation."))


using ITensors
using Printf

function ITensors.op(::OpName"expτSS", ::SiteType"S=1/2", s1::Index, s2::Index; τ)
  h =
    1 / 2 * op("S+", s1) * op("S-", s2) +
    1 / 2 * op("S-", s1) * op("S+", s2) +
    op("Sz", s1) * op("Sz", s2)
  return exp(τ * h)
end

function main(; N=10, cutoff=1E-8, δτ=0.1, beta_max=2.0)

  # Make an array of 'site' indices
  s = siteinds("S=1/2", N; conserve_qns=true)

  # Make gates (1,2),(2,3),(3,4),...
  gates = ops([("expτSS", (n, n + 1), (τ=-δτ / 2,)) for n in 1:(N - 1)], s)
  # Include gates in reverse order to complete Trotter formula
  append!(gates, reverse(gates))

  # Initial state is infinite-temperature mixed state
  rho = MPO(s, "Id") ./ √2

  # Make H for measuring the energy
  terms = OpSum()
  for j in 1:(N - 1)
    terms += 1 / 2, "S+", j, "S-", j + 1
    terms += 1 / 2, "S-", j, "S+", j + 1
    terms += "Sz", j, "Sz", j + 1
  end
  H = MPO(terms, s)

  # Do the time evolution by applying the gates
  # for Nsteps steps
  for β in 0:δτ:beta_max
    energy = inner(rho, H)
    @printf("β = %.2f energy = %.8f\n", β, energy)
    rho = apply(gates, rho; cutoff)
    rho = rho / tr(rho)
  end

  return nothing
end


#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modelling langauge for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# Hock-Schittkowski Nonlinear Test Suite  -  HS114
# This file is JuMP implementation of the model described in
#  W. Hock, K. Schittkowski, Test Examples for Nonlinear Programming
#  Codes, Lecture Notes in Economics and Mathematical Systems,
#  Springer, No, 187, 1981
# More information, including original model description, at
# http://www.ai7.uni-bayreuth.de/downloads.htm
#
# This problem has a quadratic objective with quadratic constraints as well
# as a couple constraints that have fractions of variables.
#############################################################################

using JuMP
using Base.Test

let

n = 10
a = 0.99
b = 0.9

lower = [0.0001, 0.0001, 0.0001, 0.0001, 0.0001, 85, 90, 3, 1.2, 145]
upper = [2000, 16000, 120, 5000, 2000, 93, 95, 12, 4, 162]
start = [1745, 12000, 110, 3048, 1974, 89.2, 92.8, 8, 3.6, 145]

m = Model()
@variable(m, lower[i] <= x[i=1:n] <= upper[i], start = start[i])

@NLobjective(m, Min, 5.04*x[1] + .035*x[2] + 10*x[3] + 3.36*x[5] - .063*x[4]*x[7])

@NLconstraint(m, 35.82 - .222*x[10] - b*x[9] >= 0)
@NLconstraint(m, -133 + 3*x[7] - a*x[10] >= 0)
@NLconstraint(m, -(35.82 - .222*x[10] - b*x[9]) + x[9]*(1/b - b) >= 0)
@NLconstraint(m, -(-133 + 3*x[7] - a*x[10]) + (1/a - a)*x[10] >= 0)
@NLconstraint(m, 1.12*x[1] + .13167*x[1]*x[8] - .00667*x[1]*x[8]^2 - a*x[4] >= 0)
@NLconstraint(m, 57.425 + 1.098*x[8] - .038*x[8]^2 + .325*x[6] - a*x[7] >= 0)
@NLconstraint(m, -(1.12*x[1] + .13167*x[1]*x[8] - .00667*x[1]*x[8]^2 - a*x[4]) + (1/a - a)*x[4] >= 0)
@NLconstraint(m, -(57.425 + 1.098*x[8] - .038*x[8]^2 + .325*x[6] - a*x[7]) + (1/a - a)*x[7] >= 0)
@NLconstraint(m, 1.22*x[4] - x[1] - x[5] == 0)
@NLconstraint(m, 98000*x[3]/(x[4]*x[9] + 1000*x[3]) - x[6] == 0)
@NLconstraint(m, (x[2] + x[5])/x[1] - x[8] == 0)

solve(m)

@test_approx_eq_eps getobjectivevalue(m) -1768.80696 1e-3

end


function vertCongruenceAA(W)
	err, i, todelete, vclasses = 10^-6, 1, [], []
	verts = convert(Lar.Points, W')
	kdtree = NearestNeighbors.KDTree(verts)
	newverts = zeros(Int, size(verts,2))
	for vi in 1:size(verts,2)
		if !(vi in todelete)
			nearvs = NearestNeighbors.inrange(kdtree, verts[:,vi], err)
			push!(vclasses,nearvs)
			newverts[nearvs] .= i
			nearvs = setdiff(nearvs, vi)
			todelete = union(todelete, nearvs)
			i += 1
		end
	end
	V = zeros(3,length(vclasses))
	for (k,class) in enumerate(vclasses)
			V[:,k] = sum(W[class,:],dims=1)/length(class)
	end
	return V, vclasses
end

function cellCongruenceAA(Delta,inclasses)
	cellarray = Lar.cop2lar(Delta)
	new_e = Array{Int64,1}(undef,size(Delta,2))
	for (k,class) in enumerate(inclasses)
		for e in class
			new_e[e] = k
		end
	end
  cells = [map(e->new_e[e], face) for face in cellarray]
  outclasses = DefaultOrderedDict{Array{Int64,1},Array{Int64,1}}([])
  for (k,face) in enumerate(cells)
    if outclasses[face] == []
    	outclasses[face] = [k]
    else
    	append!(outclasses[face],[k])
    end
  end
	FEnew = sort(collect(keys(outclasses)))
  outclasses = sort(collect(values(outclasses)))
  return FEnew,outclasses
end

function chainCongruenceAA(W, T)
	V, vclasses = vertCongruenceAA(W)
	EV, eclasses = cellCongruenceAA(T[1],vclasses)
	FE, fclasses = cellCongruenceAA(T[2],eclasses)
	#@show size(V,2) - size(EV,1) + size(FE,1)
	return V,EV,FE
end


using  NonLinearReactionAdvectionDiffusionWithFrontData;
using  NonLinearReactionAdvectionDiffusionWithFrontData: heterogeneity;
using  NonLinearReactionAdvectionDiffusionWithFrontData: f1, f2;
using  NonLinearReactionAdvectionDiffusionWithFrontData: apply_on_dynamic_mesh;

@testset "Градиент на точных данных и статической сетке    " begin

    # Создадим априорную информацию и точное решение для точного q(x)
    a, b, t₀, T, N, M, ε, Xₙ, Tₘ, qₙ, ulₘ, urₘ, u₀ = NonLinearReactionAdvectionDiffusionWithFrontData.dparams();
    #
    u, XX, TP = solve(u₀, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qₙ);
    #' ## Генерация априорной информации
    #########################################################################################
    ϕl      = phidetermination(qₙ, ulₘ, Xₙ, N, Tₘ, M);                      # Левый вырожденный корень
    ϕr      = phidetermination(qₙ, urₘ, Xₙ, N, Tₘ, M, reverseX = true);     # Правый вырожденный корень
    ϕ       = Φ(ϕl, ϕr, N, M);                                              # Серединный корень
    f1_data = f1(ϕ, u, Xₙ, N, M);                                           # Положение переходного слоя
    f2_data = f2(f1_data, u, Xₙ, N, M);                                     # Значение функции на переходном слое
    #########################################################################################
    nothing #hide

    # Функция `heterogeneity` — для внутреннего использования,
    # она принимает аргументы без граничных точек по x.
    Uₙₘ = u[2:N, :];
    X = XX[2:N, :];
    # Создадим замыкания, для удобства
    hmap(z) = [  - heterogeneity(n, m, X[:, m], N, Tₘ, M, Uₙₘ, f1_data, f2_data, z) for n in 1:N-1, m in 1:M+1]
    nnz(arr) = length( findall( x -> ! isapprox(x, 0), arr ))

    # Количество ненулевых элементов — 0
    @test       nnz(hmap(0.0001)) < M/4        # 0, Не рекомендовано
    @test 0 < nnz(hmap(0.0004)) < M/10    # 5, норм!

    # Для визуального контроля, можно использовать следующий код дающий анимацию
    #=
    # {{{
    using Plots, LaTeXStrings;
    pyplot();

    heatmap(h(0.0004)', label=L"- 2 \delta(x - f_1(t)) ( u^s(x,t) - f_2(t))")

    d = [ deltaw( n, f1[m], Xₙ[2:N], N) for n in 1:N-1, m in 1:M+1];    # Значения первого множителя
    uf = [ ( Uₙₘ[n+1, m] - f2[m] ) for n in 1:N-1, m in 1:M+1];         # Значения второго множителя

    yl = extrema( [H; d; uf]  );
    a = Animation()
    for m in 1:M
        nnd = findfirst( x-> x!= 0, d[:,m]);                # Первый ненулевой элемент на m-ом шаге по T
        nnh = findfirst( x-> x!= 0, H[:,m]);                # Первый ненулевой элемент на m-ом шаге по T
        ym = d[nnd, m] > H[nnh, m] ? d[nnd,m] : H[nnh, m]   # Выбор наибольшего ненулевого элемента
                                                            # чтобы нарисовать красивую вертикальную лиинию
        scatter(X,  H[:, m],    label=L"- 2 \delta(x - f_1(t)) ( u^s(x,t) - f_2(t))", ylims = yl)
        scatter!(X, d[:, m],    label=L"\delta(x - f_1(t))")
        scatter!(X, uf[:, m],   label=L"( u^s(x,t) - f_2(t))")
        plot!([Xₙ[nnd+1], Xₙ[nnd+1]], [ym, 0], line=:dash, label="") # Вертикальная линия
        frame(a)
    end
    g = gif(a, "heterogenyry.gif")

    # Карта поверхности неоднородности
    heatmap(X, Tₘ, H',  xlabel = "x", ylabel="t", title=L"- 2 \delta(x - f_1(t)) ( u^s(x,t) - f_2(t))")
    savefig("heterogeneity.png")
    # }}}
    =#

end

@testset "Градиент на точных данных и динамической сетке   " begin

    using NonLinearReactionAdvectionDiffusionWithFrontData;
    using NonLinearReactionAdvectionDiffusionWithFrontData: Φ, apply_on_dynamic_mesh;
    using NonLinearReactionAdvectionDiffusionWithFrontData: f1, f2;

    a, b, t₀, T, N, M, ε, Xₙ, Tₘ, qₙ, ulₘ, urₘ, u₀, mshfrm = NonLinearReactionAdvectionDiffusionWithFrontData.dparams_nonuniform();
    #
    u, XX, TP = solve(u₀, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qₙ, create_mesh = mshfrm);
    #' ## Генерация априорной информации
    #########################################################################################
    ϕl      = phidetermination(qₙ, ulₘ, Xₙ, N, Tₘ, M);                          # Левый вырожденный корень
    ϕr      = phidetermination(qₙ, urₘ, Xₙ, N, Tₘ, M, reverseX = true);         # Правый вырожденный корень
    ϕ       = Φ(ϕl, ϕr, N, M);                                                  # Полуразность вырожденных корней
    ϕ       = apply_on_dynamic_mesh(ϕ, XX, N, M);                               # Аппроксимация на переменную сетку
    ϕl      = apply_on_dynamic_mesh(ϕl, XX, N, M);                              # Аппроксимация на переменную сетку
    ϕr      = apply_on_dynamic_mesh(ϕr, XX, N, M);                              # Аппроксимация на переменную сетку
    f1_data = f1(ϕ, u, XX, N, M);                                               # Положение переходного слоя
    f2_data = f2(f1_data, u, XX, N, M);                                         # Значение функции на переходном слое
    #########################################################################################

    using  NonLinearReactionAdvectionDiffusionWithFrontData: heterogeneity
    # Функция `heterogeneity` — для внутреннего использования,
    # она принимает аргументы без граничных точек по x.
    Uₙₘ = u[2:N, :];
    X = XX[2:N, :];
    hmap(z) = [  - heterogeneity(n, m, X[:, m], N, Tₘ, M, Uₙₘ, f1_data, f2_data, z) for n in 1:N-1, m in 1:M+1]
    nnz(arr) = length( findall( x -> ! isapprox(x, 0), arr ));

    @test nnz(hmap(0.00002)) > M * 0.05         # Много ненулевых элементов,
                                                # значит итерационный процесс продолжится и уйдет от решения
    @test nnz(hmap(0.0000001)) < M * 0.05        # 0, Не рекомендовано
    @test 0 < nnz(hmap(0.000003)) < M * 0.01    # 4, Должно быть норм!

end


# function BB_BB_intersect(tt::TT_Cache, a::AABB, b::AABB)
#     t = tt.R_a_b * b.c + tt.t_a_b - a.c
#     return BB_BB_intersect(a.e, b.e, t, tt.R_a_b, tt.abs_R_a_b)
# end

function BB_BB_intersect(tt::TT_Cache, a::OBB, b::OBB)
    i_dh_a = basic_dh(a.R', -a.R' * a.c)   # transform from frame A to   OBB A
    dh_a_b = basic_dh(tt.R_a_b, tt.t_a_b)  # transform from frame B to frame A
    dh_b   = basic_dh(b.R, b.c)            # transform from   OBB B to frame B

    dh_final = i_dh_a * dh_a_b * dh_b
    R_tot, t = dh_R_t(dh_final)

    abs_R_tot = abs.(R_tot) .+ 1.0e-14
    return BB_BB_intersect(a.e, b.e, t, R_tot, abs_R_tot)
end

s221(r::SVector{3, Float64}) = SVector{3, Float64}(r[3], r[3], r[2])
s100(r::SVector{3, Float64}) = SVector{3, Float64}(r[2], r[1], r[1])

function BB_BB_intersect(e_a::SVector{3, Float64}, e_b::SVector{3, Float64}, t::SVector{3, Float64}, R::SMatrix{3,3,Float64,9}, abs_R::SMatrix{3,3,Float64,9})
    # This implementation of the OBB_OBB_SAT test is based off of the test described
    # in table 4.1 on page 103 of Ericson's Real-Time Collision Detection.

    R0_012 = SVector{3, Float64}(R[1], R[4], R[7])
    R1_012 = SVector{3, Float64}(R[2], R[5], R[8])
    R2_012 = SVector{3, Float64}(R[3], R[6], R[9])
    abs_R0_012 = SVector{3, Float64}(abs_R[1], abs_R[4], abs_R[7])
    abs_R1_012 = SVector{3, Float64}(abs_R[2], abs_R[5], abs_R[8])
    abs_R2_012 = SVector{3, Float64}(abs_R[3], abs_R[6], abs_R[9])

    # face test 1/2
    T_dot_L = abs.(t)
    r_a = e_a
    r_b = abs_R * e_b
    any((r_a + r_b) .< T_dot_L) && (return false)

    # face test 2/2
    T_dot_L = abs.(transpose(R) * t)  # these are transposes
    r_a = transpose(abs_R) * e_a
    r_b = e_b
    any((r_a + r_b) .< T_dot_L) && (return false)

    eb_100     = s100(e_b)
    eb_221     = s221(e_b)
    abs_R0_221 = s221(abs_R0_012)
    abs_R0_100 = s100(abs_R0_012)
    abs_R1_221 = s221(abs_R1_012)
    abs_R1_100 = s100(abs_R1_012)
    abs_R2_221 = s221(abs_R2_012)
    abs_R2_100 = s100(abs_R2_012)

    t0 = t[1]
    t1 = t[2]
    t2 = t[3]
    ea0 = e_a[1]
    ea1 = e_a[2]
    ea2 = e_a[3]

    # cross test 1/3
    T_dot_L = abs.(t2      * R1_012     - t1      * R2_012)
    r_a     =      ea1     * abs_R2_012 + ea2     * abs_R1_012
    r_b     =      eb_100 .* abs_R0_221 + eb_221 .* abs_R0_100
    any((r_a + r_b) .< T_dot_L) && (return false)

    # cross test 2/3
    T_dot_L = abs.(t0  * R2_012     -     t2  * R0_012)
    r_a     =      ea0 * abs_R2_012 +    ea2  * abs_R0_012
    r_b     =  eb_100 .* abs_R1_221 + eb_221 .* abs_R1_100
    any((r_a + r_b) .< T_dot_L) && (return false)

    # cross test 3/3
    T_dot_L =  abs.(t1 * R0_012     -     t0  * R1_012)
    r_a     =      ea0 * abs_R1_012 +    ea1  * abs_R0_012
    r_b     =  eb_100 .* abs_R2_221 + eb_221 .* abs_R2_100
    any((r_a + r_b) .< T_dot_L) && (return false)
    return true
end


Ad = [
   [-1.0/6.0  3.0/6.0 -3.0/6.0 1.0/6.0];
   [ 3.0/6.0 -6.0/6.0  0.0/6.0 4.0/6.0];
   [-3.0/6.0  3.0/6.0  3.0/6.0 1.0/6.0];
   [ 1.0/6.0  0.0/6.0  0.0/6.0 0.0/6.0]
]

dAd = [
   [ 0.0 -0.5  1.0 -0.5];
   [ 0.0  1.5 -2.0  0.0];
   [ 0.0 -1.5  1.0  0.5];
   [ 0.0  0.5  0.0  0.0]
]

d2Ad = [
   [ 0.0 0.0 -1.0  1.0];
   [ 0.0 0.0  3.0 -2.0];
   [ 0.0 0.0 -3.0  1.0];
   [ 0.0 0.0  1.0  0.0]
]

U(s,i) = Symbol(string(s,i))
U(s,i,j) = Symbol(string(s,i,"_",j))

function create_Phi(d, extrap, diff)
    lines = []
    for i=1:d
        block = []
        rhs_1 = U("tp", i,1)
        rhs_2 = U("tp", i,2)
        rhs_4 = U("tp", i,4)
        rhs_3 = U("tp", i,3)
        if extrap == "none"
            for j=1:4
                eq = :($(U("Phi_",i,j)) = ($(Ad[j,1])*$rhs_1 + $(Ad[j,2])*$rhs_2 + $(Ad[j,3])*$rhs_3 + $(Ad[j,4])*$rhs_4) )
                push!(lines,eq)
            end
        elseif extrap == "natural"
            block = quote
                if $(U("t",i))<0
                    $( [ :($(U("Phi_",i,j)) = ($(dAd[j,4])*$rhs_3 + $(Ad[j,4])) ) for j=1:4 ]...)
                elseif $(U("t",i))>1
                    $( [ :($(U("Phi_",i,j)) = (3*$(Ad[j,1]) + 2*$(Ad[j,2]) + $(Ad[j,3]))*($rhs_3-1) + ($(Ad[j,1])+$(Ad[j,2])+$(Ad[j,3])+$(Ad[j,4])) ) for j=1:4 ]...)
                else
                    $( [ :($(U("Phi_",i,j)) = ($(Ad[j,1])*$rhs_1 + $(Ad[j,2])*$rhs_2 + $(Ad[j,3])*$rhs_3 + $(Ad[j,4])*$rhs_4) ) for j=1:4 ]...)
                end
            end
            # for l in block.args
            push!(lines,block.args...)
        end
        if diff
            for j=1:4
                eq = :($(U("dPhi_",i,j)) = ($(dAd[j,1])*$rhs_1 + $(dAd[j,2])*$rhs_2 + $(dAd[j,3])*$rhs_3 + $(dAd[j,3])*$rhs_4 )*$(U("dinv",i)) )
                push!(lines,eq)
            end
        end
    end
    return lines
end

function tensor_prod(symbs, inds, add_index=false)
    if length(symbs)==0
        subscripts = [:($(U("i",i))+$(inds[i])) for i=1:length(inds)]
        if add_index
            return :(C[k,$(subscripts...)])
        else
            return :(C[$(subscripts...)])
        end
    else
        h = symbs[1]
        if length(symbs)>1
            q = symbs[2:end]
        else
            q = []
        end
        exprs = []
        for i = 1:4
            e = Meta.parse( string(h,"_",i,".*(",tensor_prod(q,cat(inds,[i]; dims=1),add_index),")") )
            push!(exprs,e)
        end
        return :(+($(exprs...)))
    end
end

tensor_prod([], Int64[1,2,3,4],)            # C[i1+1,i2+2,i3+3,i4+4]
tensor_prod([], Int64[1,2,3,4], true)      # C[k,i1+1,i2+2,i3+3,i4+4]

tensor_prod(["Phi_1"], Int64[])            # Phi_1_1 * C[i1 + 1] + Phi_1_2 * C[i1 + 2] + Phi_1_3 * C[i1 + 3] + Phi_1_4 * C[i1 + 4]
tensor_prod(["Phi_1", "Phi_2"], Int64[])   # Phi_1_1 * (Phi_2_1 * C[i1 + 1,i2 + 1] + Phi_2_2 * C[i1 + 1,i2 + 2] + Phi_2_3 * C[i1 + 1,i2 + 3] + Phi_2_4 * C[i1 + 1,i2 + 4]) + Phi_1_2 * (Phi_2_1 * C[i1 + 2,i2 + 1] + Phi_2_2 * C[i1 + 2,i2 + 2] + Phi_2_3 * C[i1 + 2,i2 + 3] + Phi_2_4 * C[i1 + 2,i2 + 4]) + Phi_1_3 * (Phi_2_1 * C[i1 + 3,i2 + 1] + Phi_2_2 * C[i1 + 3,i2 + 2] + Phi_2_3 * C[i1 + 3,i2 + 3] + Phi_2_4 * C[i1 + 3,i2 + 4]) + Phi_1_4 * (Phi_2_1 * C[i1 + 4,i2 + 1] + Phi_2_2 * C[i1 + 4,i2 + 2] + Phi_2_3 * C[i1 + 4,i2 + 3] + Phi_2_4 * C[i1 + 4,i2 + 4])
tensor_prod(["Phi_1", "dPhi_2"], Int64[])  # Phi_1_1 * (dPhi_2_1 * C[i1 + 1,i2 + 1] + dPhi_2_2 * C[i1 + 1,i2 + 2] + dPhi_2_3 * C[i1 + 1,i2 + 3] + dPhi_2_4 * C[i1 + 1,i2 + 4]) + Phi_1_2 * (dPhi_2_1 * C[i1 + 2,i2 + 1] + dPhi_2_2 * C[i1 + 2,i2 + 2] + dPhi_2_3 * C[i1 + 2,i2 + 3] + dPhi_2_4 * C[i1 + 2,i2 + 4]) + Phi_1_3 * (dPhi_2_1 * C[i1 + 3,i2 + 1] + dPhi_2_2 * C[i1 + 3,i2 + 2] + dPhi_2_3 * C[i1 + 3,i2 + 3] + dPhi_2_4 * C[i1 + 3,i2 + 4]) + Phi_1_4 * (dPhi_2_1 * C[i1 + 4,i2 + 1] + dPhi_2_2 * C[i1 + 4,i2 + 2] + dPhi_2_3 * C[i1 + 4,i2 + 3] + dPhi_2_4 * C[i1 + 4,i2 + 4])

function create_parameters(d)
    lines = []
    for i=1:d
        block = quote
            $(U("M",i)) = orders[$i]
            $(U("start",i)) =  a[$i]
            $(U("dinv",i)) = (orders[$i]-1.0)/(b[$i]-a[$i])
        end
        for ll in block.args
            push!(lines, ll)
        end
    end
    return lines
end

function create_local_parameters(d)
    lines = []
    for i=1:d
        bl = quote
            $(U("x",i)) = S[n][$i]
            $(U("u",i)) = ($(U("x",i)) - $(U("start",i)))*$(U("dinv",i))
            $(U("i",i)) = (floor(Int,$(U("u",i)) ))
            $(U("i",i)) = max( min($(U("i",i)),$(U("M",i))-2), 0 )
            $(U("t",i)) = $(U("u",i))-$(U("i",i))
            $(U("tp",i,1)) = $(U("t",i))*$(U("t",i))*$(U("t",i))
            $(U("tp",i,2)) = $(U("t",i))*$(U("t",i))
            $(U("tp",i,3)) = $(U("t",i))
            $(U("tp",i,4)) = 1.0;
        end
        for ll in bl.args
            push!(lines, ll)
        end
    end
    return lines
end

function create_function(d,extrap="natural")
    expr = quote
        function $(Symbol(string("eval_UC_spline!")))( a, b, orders, C::Array{T,d}, S::Vector{SVector{d,Float64}}, V::Vector{T}) where d where T
            $(create_parameters(d)...)
            N = size(S,1)
            # @fastmath @inbounds @simd( # doesn't seem to make any difference
            for n=1:N
                $(create_local_parameters(d)...)
                $(create_Phi(d,extrap,false)...)
                V[n] = $( tensor_prod([string("Phi_",i) for i=1:d], Int64[]) )
            end
            # )
        end
    end
    return expr
end


using SparseArrays, LinearAlgebra, Arpack
include("io/io_helper.jl")

function diagonalize(ham::Hamiltonian, param::Dict{Any,Any})
    """ Performs the diagonalization as specified in the paramter dictionary.
        Done either by iterated Arnoldi method with ARPACK (when full_diag is
        set to `false`) or by full diagonalization.

        Notes:
         -  Julia has a matrix-independent iteration number of 300 per default,
            this is potentially problematic. Needs to be adjusted in some cases.
    """
    sparse_ham = sparse(ham.row, ham.col, ham.data, ham.n_fock, ham.n_fock)

    if ham.n_fock > 1
        @info "Starting to diagonalize the Hamiltonian." n_fock=Int(ham.n_fock)
        if get(param, "full_diag", false)
            # Full diagonalization with built-in method.
            time = @elapsed mem = @allocated F = eigen(Matrix(sparse_ham))
            @info "Done computing the full spectrum." time=time allocated=MemoryTag(mem) spectrum=F.values
            return F.values, F.vectors
        else
            # Arnoldi method with ARPACK.
            which_map = Dict([
                ("SA", :SR)
            ])
            # Here, we set the default values that are otherwise overwritten
            # if specified in the parameter dictionary.
            time = @elapsed mem = @allocated ev, est = eigs(
                sparse_ham,
                nev=get(param, "n_eigenvalues", 10),
                which=which_map[get(param, "ev_type", "SA")],
                tol=get(param, "tol", 1e-15),
                maxiter=get(param, "max_iter", 10*ham.n_fock)
            )
            @info "Done computing the lower spectrum." time=time allocated=MemoryTag(mem) spectrum=ev
            return ev, est
        end
    end
    @info "Trivial 1x1 matrix to diagonalize."
    return Matrix(sparse_ham)[1,1], [1]
end