"""
Automatic differentiation with per-variable storage of derivative information (as a tree structure) in a new `Type` [`ADval`](@ref)` <: Real`.
This provides a way of detecting Jacobian (dy/dx) sparsity patterns for problems of form y = f(x)
where x and y are vectors.

 Notes:
  - Will fail if y = f(x) produces a tree with loops, eg by inverting a matrix to solve a linear system.
  - Very slow as each arithmetic operation generates a memory allocation (in order to form the derivative tree structure).
    Speed is approx 10 Mop/s (on a ~4Ghz laptop core) so 100 - 1000x slower than `y = f(x)` with primitive types. Intent
    is that the sparsity information is calculated once and then used to enable a separate fast AD calculator.
  - Also calculates derivative information (slowly).
  - The tree of derivative information grows without limit, so currently limited to cases `y = f(x)` where f does
    not generate too many operations. It would be straightforward to fix this by limiting tree to size `~4*length(x)`
    per element and simplifying it when this limit exceeded (as tree size = 2*length(x) per element is enough for a
    dense Jacobian)
  - Uses DiffRules.jl to autogenerate operations.
  - See ForwardDiff.Dual https://github.com/JuliaDiff/ForwardDiff.jl/blob/master/src/dual.jl
    for examples of definitions required.

# Example usage:
    x_ad = SparsityTracing.create_advec(x)           # x_ad takes values from x,  derivatives set to identity
    y_ad = f(x_ad)                              # tree of derivative information is accumulated into each element of y_ad
    jac = SparsityTracing.jacobian(y_ad, length(x))  # walk trees and generate sparse Jacobian
 """
module SparsityTracing

using SparseArrays
using DiffRules

import SpecialFunctions

"base type for tree of derivative information"
abstract type AbstractDerivNode
end

"""
    ADval

Scalar variable type T with derivatives
"""
struct ADval{T<:Real} <: Real
    val::T
    derivnode::Union{Nothing, AbstractDerivNode}
end

"""
    create_advec(x::Vector{T}) -> Vector{ADval{T, length(x)}}

Create a vector of `ADval` scalars, each with value initialised from `x`
and `length(x)` derivatives initialised to identity (each element derivative 1 wrt itself and 0 elsewhere).
Usual case for an independent vector variable `x` (eg to use as input with `y = f(x)`).

# Examples:
 ```jldoctest; setup = :(import SparsityTracing)
julia> v_ad = SparsityTracing.create_advec([1.0, 5.0])
2-element Vector{SparsityTracing.ADval{Float64}}:
 SparsityTracing.ADval{Float64}(1.0, <derivnode>)
 SparsityTracing.ADval{Float64}(5.0, <derivnode>)

julia> x_ad, y_ad = v_ad;

julia> x_ad        # deriv is sparse vector
SparsityTracing.ADval{Float64}
  val=1.0
  deriv:
  [1]  =  1.0

julia> y_ad       # deriv is sparse vector
SparsityTracing.ADval{Float64}
  val=5.0
  deriv:
  [2]  =  1.0

julia> SparsityTracing.jacobian(v_ad, 2)  # sparse 2x2 identity matrix.
2×2 SparseArrays.SparseMatrixCSC{Float64, Int64} with 2 stored entries:
 1.0   ⋅
  ⋅   1.0

julia> z_ad = x_ad*y_ad^2    # d(x*y^2)/dx = y^2 = 5^2 = 25,  d(x*y^2)/dy = 2*x*y = 2*1*5 = 10
SparsityTracing.ADval{Float64}
  val=25.0
  deriv:
  [1]  =  25.0
  [2]  =  10.0
```
"""
function create_advec(x::Vector{T}) where T

    n = length(x)
    ADvec = Vector{ADval{T}}(undef, n)

    for i = 1:n
        ADvec[i] = ADval{T}(x[i], DerivLeaf(i))
    end

    return ADvec
end

###############################
# Construction and conversion #
###############################

"create from a non-AD value, set derivatives to zero"
ADval{T}(v) where {T} = ADval{T}(v, nothing)
ADval{T}(v::ADval{T}) where {T} = v

Base.convert(::Type{ADval{T}}, v) where {T} = ADval{T}(v)
Base.convert(::Type{ADval{T}}, v::Number) where {T} = ADval{T}(v)
Base.convert(::Type{ADval{T}}, d::ADval{T}) where {T} = d

Base.promote_rule(::Type{T}, ::Type{ADval{T}} ) where {T} = ADval{T}

"get value. NB: no automatic conversion to Float64 etc"
value(d::ADval) = d.val

#####################
# Generic functions #
#####################

Base.zero(d::ADval{T}) where {T}        = zero(typeof(d))
Base.zero(::Type{ADval{T}}) where {T}   = ADval{T}(zero(T), nothing)

Base.one(d::ADval{T}) where {T}         = one(typeof(d))
Base.one(::Type{ADval{T}}) where {T}    = ADval{T}(one(T), nothing)

Base.sign(d::ADval{T}) where {T} = sign(d.val)

###################
# Pretty printing #
###################

"compact form"
function Base.show(io::IO, d::ADval)
    print(io, typeof(d), "(", d.val, ", <derivnode>)")
end

"multiline form"
function Base.show(io::IO, ::MIME"text/plain", d::ADval)
    println(io, typeof(d))
    println(io, "  val=", d.val)
    println(io, "  deriv:\n", deriv(d))
end

###################
# Get derivatives #
###################

"evaluate derivnode by traversing tree, and return as a sparse vector"
function deriv(d::ADval{T}, N) where {T}
    I = Vector{Int64}()
    V = Vector{T}()
    reduce_deriv!(I, V, 1.0, d.derivnode)
    # reduce_deriv_tree!(I, V, 1.0, d.derivnode)
    return sparsevec(I, V, N)
end

function deriv(d::ADval{T}) where {T}
    I = Vector{Int64}()
    V = Vector{T}()
    reduce_deriv!(I, V, 1.0, d.derivnode)
    # reduce_deriv_tree!(I, V, 1.0, d.derivnode)
    return sparsevec(I, V)
end

"""
    jacobian(advec::Vector{ADval}, N) -> SparseMatrixCSC

Aggregate derivatives to form sparse Jacobian.
"""
function jacobian(advec::Vector{ADval{T}}, N) where {T}

    Nvar = length(advec)

    # @info "SparsityTracing.jacobian ($Nvar, $N)"

    # define arrays for sparse matrix creation
    I = Vector{Int64}()
    J = Vector{Int64}()
    V = Vector{T}()

    for i = 1:Nvar
        # @info "  i=$i"
        derivJ, derivV = findnz(deriv(advec[i], N))
        # @info "  length(derivJ) = $(length(derivJ))"
        for j = 1:length(derivJ)
            push!(I, i)
            push!(J, derivJ[j])
            push!(V, derivV[j])
        end
    end

    return sparse(I, J, V, Nvar, N)
end

"doesn't need to be mutable, but slightly faster (with Julia 1.5)"
mutable struct DerivLeaf <: AbstractDerivNode
    index::Int64
end

"doesn't need to be mutable, but slightly faster (with Julia 1.5)"
mutable struct DerivNode{T} <: AbstractDerivNode
    lmult::T
    lnode::Union{Nothing, AbstractDerivNode}
    rmult::T
    rnode::Union{Nothing, AbstractDerivNode}
end

function reduce_deriv!(I, v, mult, node::Nothing)
    return nothing
end

function reduce_deriv!(I, V, mult, node::DerivLeaf)
    push!(I, node.index)
    push!(V, mult)
    return nothing
end

function reduce_deriv!(I, V, mult, node::DerivNode)
    reduce_deriv!(I, V, mult*node.lmult, node.lnode)
    reduce_deriv!(I, V, mult*node.rmult, node.rnode)
    return nothing
end

function reduce_deriv_tree!(I, V, mult, root::Union{Nothing, DerivNode, DerivLeaf})

    stack = Vector{Tuple{Float64, Union{Nothing, DerivNode, DerivLeaf}}}()

    curr = (mult, root)
    npush = 0

    while curr[2] isa DerivNode || !isempty(stack)
        # reach the left-most node of the curr node
        while curr[2] isa DerivNode
            # put curr on stack and traverse left subtree
            push!(stack, curr)
            mult, node = curr
            curr = (mult*node.lmult, node.lnode)
            npush += 1
            if npush % 1000000 == 0
                @info "    npush $npush length(stack, I, V) $(length(stack)) $(length(I)) $(length(V))"
            end
        end
        # curr[2] must be either Nothing or a DerivLeaf
        if curr[2] isa DerivLeaf
            mult, node = curr
            push!(I, node.index)
            push!(V, mult)
            # @info "    $(node.index)   $mult"
        end
        mult, node = pop!(stack)
        # we have visited the node and its left subtree
        # visit the right subtree
        curr = (mult*node.rmult, node.rnode)
    end

    # curr[2] must be either Nothing or a DerivLeaf
    if curr[2] isa DerivLeaf
        mult, node = curr
        push!(I, node.index)
        push!(V, mult)
        # @info "    $(node.index)   $mult"
    end

    # @info " reduce_deriv_tree! length(I) = $(length(I))"
    return nothing

end

##########################################
# Autogenerate functions using DiffRules #
##########################################

"display DiffRules rules (diagnostic for development)"
function list_diffrules(arity)
    for key in DiffRules.diffrules()
        if key[3] in arity
            println(key)
        end
    end
end

"generate a unary AD rule
   ns       eg :Base
   op       eg :sin"
function gen_unary_rule(ns, op)

    deriv = DiffRules.diffrule(ns, op, :(u.val)) # eg :(cos(u.val))

    @eval $ns.$op(u::ADval{T}) where {T} = ADval{T}($ns.$op(u.val), DerivNode(convert(T, $deriv), u.derivnode, zero(T), nothing))
end

"generate a binary AD rule
   ns       eg :Base
   op       eg :+"
function gen_binary_rule(ns, op)

    u_deriv, v_deriv = DiffRules.diffrule(ns, op, :(u.val), :(v.val))
    @eval $ns.$op(u::ADval{T}, v::ADval{T}) where {T} = ADval{T}($ns.$op(u.val, v.val),
        DerivNode(convert(T, $u_deriv), u.derivnode, convert(T, $v_deriv), v.derivnode))

    u_deriv, dummy = DiffRules.diffrule(ns, op, :(u.val), :v)
    @eval $ns.$op(u::ADval{T}, v::Real) where {T} = ADval{T}($ns.$op(u.val, v),
        DerivNode(convert(T, $u_deriv), u.derivnode, zero(T), nothing))

    dummy, v_deriv = DiffRules.diffrule(ns, op, :u, :(v.val))
    @eval $ns.$op(u::Real, v::ADval{T}) where {T} = ADval{T}($ns.$op(u, v.val),
        DerivNode(convert(T, $v_deriv), v.derivnode, zero(T), nothing))
end

"generate a logical rule (ignore derivatives, just test value)"
function gen_logical_rule(op)
    @eval Base.$op(u::ADval{T}, v::ADval{T}) where {T} = $op(u.val, v.val)
    @eval Base.$op(u::ADval{T}, v::Real) where {T} = $op(u.val, v)
    @eval Base.$op(u::Real, v::ADval{T}) where {T} = $op(u, v.val)
end

function gen_unary_rules(namespaces, excludeops)
    for (ns, op, arity) in DiffRules.diffrules()
        if ns in namespaces && arity == 1 && !(op in excludeops)
            gen_unary_rule(ns, op)
        end
    end
end

function gen_binary_rules(namespaces, excludeops)
    for (ns, op, arity) in DiffRules.diffrules()
        if ns in namespaces && arity == 2 && !(op in excludeops)
            gen_binary_rule(ns, op)
        end
    end
end

function gen_logical_rules()
    for op in (:<, :>, :>=, :<=, :(==))
        gen_logical_rule(op)
    end
end

# add functions to Base
gen_unary_rules((:Base, :SpecialFunctions), ())
gen_binary_rules((:Base, :SpecialFunctions), ())
gen_logical_rules()

###############
# TODO fixups #
###############
# Errors Base.:^(d::ADval{T}, n::Integer) is ambiguous ?
# adding literal_pow apparently fixes this but might not be the right fix
Base.literal_pow(::typeof(^), d::ADval{T}, ::Val{0}) where {T} = one(d)
Base.literal_pow(::typeof(^), d::ADval{T}, ::Val{1}) where {T} = d
Base.literal_pow(::typeof(^), d::ADval{T}, ::Val{2}) where {T} = d*d
Base.literal_pow(::typeof(^), d::ADval{T}, ::Val{3}) where {T} = d*d*d
Base.literal_pow(::typeof(^), d::ADval{T}, ::Val{4}) where {T} = d*d*d*d

end


# # [Introduction to periodic problems and plane-wave discretisations](@id periodic-problems)
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/guide/@__NAME__.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/guide/@__NAME__.ipynb)

# In this example we want to show how DFTK can be used to solve simple one-dimensional
# periodic problems. Along the lines this notebook serves as a concise introduction into
# the underlying theory and jargon for solving periodic problems using plane-wave
# discretizations.

# ## Periodicity and lattices
# A periodic problem is characterized by being invariant to certain translations.
# For example the ``\sin`` function is periodic with periodicity ``2π``, i.e.
# ```math
#    \sin(x) = \sin(x + 2πm) \quad ∀ m ∈ \mathbb{Z},
# ```
# This is nothing else than saying that any translation by an integer multiple of ``2π``
# keeps the ``\sin`` function invariant. In a more formal one can use the
# translation operator ``T_{-2πm}`` to write this as:
# ```math
#    T_{-2πm} \, \sin(x) = \sin(x + 2πm) = \sin(x).
# ```
#
# Whenever such periodicity exists one can exploit it to save computational work.
# Consider a problem in which we want to find a function ``f : \mathbb{R} → \mathbb{R}``,
# but *a priori* the solution is known to be periodic with periodicity ``a``. As a consequence
# of said periodicity it is sufficient to determine the values of ``f`` for all ``x`` from the
# interval ``[-a/2, a/2)`` to uniquely define ``f`` on the full real axis. Naturally exploiting
# periodicity in a computational procedure thus greatly reduces the required amount of work.
#
# Let us introduce some jargon: The periodicity of our problem implies that we may define
# a **lattice**
# ```
#         -3a/2      -a/2      +a/2     +3a/2
#        ... |---------|---------|---------| ...
#                 a         a         a
# ```
# with lattice constant ``a``. Each cell of the lattice is an identical periodic image of
# any of its neighbors. For finding ``f`` it is thus sufficient to consider only the
# problem inside a **unit cell** ``[-a/2, a/2)``. In passing we note that the definition
# of the unit cell is itself only unique up to translations. A choice ``[0, a)``,
# for example, would have done just as well.
#
# ## Periodic operators and the Bloch transform
# Not only functions, but also operators can feature periodicity.
# Consider for example the **free-electron Hamiltonian**
# ```math
#     H = -\frac12 Δ.
# ```
# In free-electron model (which gives rise to this Hamiltonian) electron motion is only
# by their own kinetic energy. As this model features no potential which could make one point
# in space more preferred than another, we would expect this model to be periodic.
# If an operator is periodic with respect to a lattice such as the one defined above,
# than it commutes with all lattice translations. For the free-electron case ``H``
# one can easily show exactly that, i.e.
# ```math
#    T_{ma} H = H T_{ma} \quad  ∀ m ∈ \mathbb{Z}.
# ```
# We note in passing that the free-electron model is actually very special in the sense that
# the choice of ``a`` is completely arbitrary here. In other words ``H`` is periodic
# with respect to any translation. In general, however, periodicity is only
# attained with respect to a finite number of translations ``a`` and we will take this
# viewpoint here.
#
# **Bloch's theorem** now tells us that for periodic operators,
# the solutions to the eigenproblem
# ```math
#     H ψ_{kn} = ε_{kn} ψ_{kn}
# ```
# satisfy a factorization
# ```math
#     ψ_{kn}(x) = e^{i k⋅x} u_{kn}(x)
# ```
# into a plane wave ``e^{i k⋅x}`` and a lattice-periodic function
# ```math
#    T_{ma} u_{kn}(x) = u_{kn}(x - ma) = u_{kn}(x) \quad ∀ m ∈ \mathbb{Z}.
# ```
# In this ``n`` is a labeling integer index and ``k`` is a real number,
# whose details will be clarified in the next section.
# The index ``n`` is sometimes also called the **band index** and
# functions ``ψ_{kn}`` satisfying this factorization are also known as
# **Bloch functions** or **Bloch states**.
#
# Consider the application of ``2H = -Δ = - \frac{d^2}{d x^2}``
# to such a Bloch wave. First we notice for any function ``f``
# ```math
#    -i∇ \left( e^{i k⋅x} f \right)
#    = -i\frac{d}{dx} \left( e^{i k⋅x} f \right)
#    = k e^{i k⋅x} f -i∇ e^{i k⋅x} f = (-i∇ + k) e^{i k⋅x} f.
# ```
# Using this result twice one shows that applying ``-Δ`` yields
# ```math
# \begin{aligned}
#    -\Delta \left(e^{i k⋅x} u_{kn}(x)\right)
#    &= -i∇ ⋅ \left[-i∇ \left(u_{kn}(x) e^{i k⋅x} \right) \right] \\
#    &= -i∇ ⋅ \left[(-i∇ + k) u_{kn}(x) e^{i k⋅x} \right] \\
#    &= (-i∇ + k)^2 u_{kn}(x) e^{i k⋅x} \\
#    &= e^{i k⋅x} 2H_k u_{kn}(x),
# \end{aligned}
# ```
# where we defined
# ```math
#     H_k = \frac12 (-i∇ + k)^2.
# ```
# The action of this operator on a function ``u_{kn}`` is given by
# ```math
#     H_k u_{kn} = e^{-i k⋅x} H e^{i k⋅x} u_{kn},
# ```
# which in particular implies that
# ```math
#    H_k u_{kn} = ε_{kn} u_{kn} \quad ⇔ \quad H (e^{i k⋅x} u_{kn}) = ε_{kn} (e^{i k⋅x} u_{kn}).
# ```
# To seek the eigenpairs of ``H`` we may thus equivalently
# find the eigenpairs of *all* ``H_k``.
# The point of this is that the eigenfunctions ``u_{kn}`` of ``H_k``
# are periodic (unlike the eigenfunctions ``ψ_{kn}`` of ``H``).
# In contrast to ``ψ_{kn}`` the functions ``u_{kn}`` can thus be fully
# represented considering the eigenproblem only on the unit cell.
#
# A detailed mathematical analysis shows that the transformation from ``H``
# to the set of all ``H_k`` for a suitable set of values for ``k`` (details below)
# is actually a unitary transformation, the so-called **Bloch transform**.
# This transform brings the Hamiltonian into the symmetry-adapted basis for
# translational symmetry, which are exactly the Bloch functions.
# Similar to the case of choosing a symmetry-adapted basis for other kinds of symmetries
# (like the point group symmetry in molecules), the Bloch transform also makes
# the Hamiltonian ``H`` block-diagonal[^1]:
# ```math
#     T_B H T_B^{-1} ⟶ \left( \begin{array}{cccc} H_1&&&0 \\ &H_2\\&&H_3\\0&&&\ddots \end{array} \right)
# ```
# with each block ``H_k`` taking care of one value ``k``.
# This block-diagonal structure under the basis of Bloch functions lets us
# completely describe the spectrum of ``H`` by looking only at the spectrum
# of all ``H_k`` blocks.
#
# [^1]: Notice that block-diagonal is a bit an abuse of terms here, since the Hamiltonian
#       is not a matrix but an operator and the number of blocks is essentially infinite.
#       The mathematically precise term is that the Bloch transform reveals the fibers
#       of the Hamiltonian.
#
# ## The Brillouin zone
#
# We now consider the parameter ``k`` of the Hamiltonian blocks in detail.
#
# - As discussed ``k`` is a real number. It turns out, however, that some of
#   these ``k∈\mathbb{R}`` give rise to operators related by unitary transformations
#   (again due to translational symmetry).
# - Since such operators have the same eigenspectrum, only one version needs to be considered.
# - The smallest subset from which ``k`` is chosen is the **Brillouin zone** (BZ).
#
# - The BZ is the unit cell of the **reciprocal lattice**, which may be constructed from
#   the **real-space lattice** by a Fourier transform.
# - In our simple 1D case the reciprocal lattice is just
#   ```
#     ... |--------|--------|--------| ...
#            2π/a     2π/a     2π/a
#   ```
#   i.e. like the real-space lattice, but just with a different lattice constant
#   ``b = 2π / a``.
# - The BZ in our example is thus ``B = [-π/a, π/a)``. The members of ``B``
#   are typically called ``k``-points.
#
# ## Discretization and plane-wave basis sets
#
# With what we discussed so far the strategy to find all eigenpairs of a periodic
# Hamiltonian ``H`` thus reduces to finding the eigenpairs of all ``H_k`` with ``k ∈ B``.
# This requires *two* discretisations:
#
#   - ``B`` is infinite (and not countable). To discretize we first only pick a finite number
#     of ``k``-points. Usually this **``k``-point sampling** is done by picking ``k``-points
#     along a regular grid inside the BZ, the **``k``-grid**.
#   - Each ``H_k`` is still an infinite-dimensional operator.
#     Following a standard Ritz-Galerkin ansatz we project the operator into a finite basis
#     and diagonalize the resulting matrix.
#
# For the second step multiple types of bases are used in practice (finite differences,
# finite elements, Gaussians, ...). In DFTK we currently support only plane-wave
# discretizations.
#
# For our 1D example normalized plane waves are defined as the functions
# ```math
# e_{G}(x) = \frac{e^{i G x}}{\sqrt{a}}  \qquad G \in b\mathbb{Z}
# ```
# and typically one forms basis sets from these by specifying a
# **kinetic energy cutoff** ``E_\text{cut}``:
# ```math
# \left\{ e_{G} \, \big| \, (G + k)^2 \leq 2E_\text{cut} \right\}
# ```
#
# ## Correspondence of theory to DFTK code
#
# Before solving a few example problems numerically in DFTK, a short overview
# of the correspondence of the introduced quantities to data structures inside DFTK.
#
# - ``H`` is represented by a `Hamiltonian` object and variables for hamiltonians are usually called `ham`.
# - ``H_k`` by a `HamiltonianBlock` and variables are `hamk`.
# - ``ψ_{kn}`` is usually just called `ψ`.
#   ``u_{kn}`` is not stored (in favor of ``ψ_{kn}``).
# - ``ε_{kn}`` is called `eigenvalues`.
# - ``k``-points are represented by `Kpoint` and respective variables called `kpt`.
# - The basis of plane waves is managed by `PlaneWaveBasis` and variables usually just called `basis`.
#
# ## Solving the free-electron Hamiltonian
#
# One typical approach to get physical insight into a Hamiltonian ``H`` is to plot
# a so-called **band structure**, that is the eigenvalues of ``H_k`` versus ``k``.
# In DFTK we achieve this using the following steps:
#
# Step 1: Build the 1D lattice. DFTK is mostly tailored for 3D problems.
# Therefore quantities related to the problem space are have a fixed
# dimension 3. The convention is that for 1D / 2D problems the
# trailing entries are always zero and ignored in the computation.
# For the lattice we therefore construct a 3x3 matrix with only one entry.
using DFTK

lattice = zeros(3, 3)
lattice[1, 1] = 20.;

# Step 2: Select a model. In this case we choose a free-electron model,
# which is the same as saying that there is only a Kinetic term
# (and no potential) in the model. The `n_electrons` is dummy here.
model = Model(lattice; n_electrons=0, terms=[Kinetic()])

# Step 3: Define a plane-wave basis using this model and a cutoff ``E_\text{cut}``
# of 300 Hartree. The ``k``-point grid is given as a regular grid in the BZ
# (a so-called **Monkhorst-Pack** grid). Here we select only one ``k``-point (1x1x1),
# see the note below for some details on this choice.
basis = PlaneWaveBasis(model; Ecut=300, kgrid=(1, 1, 1))

# Step 4: Plot the bands! Select a density of ``k``-points for the ``k``-grid to use
# for the bandstructure calculation, discretize the problem and diagonalize it.
# Afterwards plot the bands.

using Unitful
using UnitfulAtomic
using Plots

plot_bandstructure(basis; n_bands=6, kline_density=30)

# !!! note "Selection of k-point grids in `PlaneWaveBasis` construction"
#     You might wonder why we only selected a single ``k``-point (clearly a very crude
#     and inaccurate approximation). In this example the `kgrid` parameter specified
#     in the construction of the `PlaneWaveBasis`
#     is not actually used for plotting the bands. It is only used when solving more
#     involved models like density-functional theory (DFT) where the Hamiltonian is
#     non-linear. In these cases before plotting the bands the self-consistent field
#     equations (SCF) need to be solved first. This is typically done on
#     a different ``k``-point grid than the grid used for the bands later on.
#     In our case we don't need this extra step and therefore the `kgrid` value passed
#     to `PlaneWaveBasis` is actually arbitrary.

# ## Adding potentials
# So far so good. But free electrons are actually a little boring,
# so let's add a potential interacting with the electrons.
#
# - The modified problem we will look at consists of diagonalizing
#   ```math
#   H_k = \frac12 (-i \nabla + k)^2 + V
#   ```
#   for all ``k \in B`` with a periodic potential ``V`` interacting with the electrons.
#
# - A number of "standard" potentials are readily implemented in DFTK and
#   can be assembled using the `terms` kwarg of the model.
#   This allows to seamlessly construct
#
#   * density-functial theory (DFT) models for treating electronic structures
#     (see the [Tutorial](@ref)).
#   * Gross-Pitaevskii models for bosonic systems
#     (see [Gross-Pitaevskii equation in one dimension](@ref))
#   * even some more unusual cases like anyonic models.
#
# In this tutorial we will go a little more low-level and directly provide
# an analytic potential describing the interaction with the electrons to DFTK.
#
# First we define a custom Gaussian potential as a new "element" inside DFTK:

struct ElementGaussian <: DFTK.Element
    α  # Prefactor
    L  # Extend
end

## Some default values
ElementGaussian() = ElementGaussian(0.3, 10.0)

## Real-space representation of a Gaussian
function DFTK.local_potential_real(el::ElementGaussian, r::Real)
    -el.α / (√(2π) * el.L) * exp(- (r / el.L)^2 / 2)
end

## Fourier-space representation of the Gaussian
function DFTK.local_potential_fourier(el::ElementGaussian, q::Real)
    ## = ∫ -α exp(-(r/L)^2 exp(-ir⋅q) dr
    -el.α * exp(- (q * el.L)^2 / 2)
end

# A single potential looks like:

using Plots
using LinearAlgebra
nucleus = ElementGaussian()
plot(r -> DFTK.local_potential_real(nucleus, norm(r)), xlims=(-50, 50))

# With this element at hand we can easily construct a setting
# where two potentials of this form are located at positions
# ``20`` and ``80`` inside the lattice ``[0, 100]``:

using LinearAlgebra

## Define the 1D lattice [0, 100]
lattice = diagm([100., 0, 0])

## Place them at 20 and 80 in *fractional coordinates*,
## that is 0.2 and 0.8, since the lattice is 100 wide.
nucleus = ElementGaussian()
atoms = [nucleus => [[0.2, 0, 0], [0.8, 0, 0]]]

## Assemble the model, discretize and build the Hamiltonian
model = Model(lattice; atoms=atoms, terms=[Kinetic(), AtomicLocal()])
basis = PlaneWaveBasis(model; Ecut=300, kgrid=(1, 1, 1));
ham   = Hamiltonian(basis)

## Extract the total potential term of the Hamiltonian and plot it
potential = DFTK.total_local_potential(ham)[:, 1, 1]
rvecs = collect(r_vectors_cart(basis))[:, 1, 1]  # slice along the x axis
x = [r[1] for r in rvecs]                        # only keep the x coordinate
plot(x, potential, label="", xlabel="x", ylabel="V(x)")

# This potential is the sum of two "atomic" potentials (the two "Gaussian" elements).
# Due to the periodic setting we are considering interactions naturally also occur
# across the unit cell boundary (i.e. wrapping from `100` over to `0`).
# The required periodization of the atomic potential is automatically taken care,
# such that the potential is smooth across the cell boundary at `100`/`0`.
#
# With this setup, let's look at the bands:

using Unitful
using UnitfulAtomic

plot_bandstructure(basis; n_bands=6, kline_density=15)

# The bands are noticeably different.
#  - The bands no longer overlap, meaning that the spectrum of $H$ is no longer continuous
#    but has gaps.
#
#  - The two lowest bands are almost flat. This is because they represent
#    two tightly bound and localized electrons inside the two Gaussians.
#
#  - The higher the bands are in energy, the more free-electron-like they are.
#    In other words the higher the kinetic energy of the electrons, the less they feel
#    the effect of the two Gaussian potentials. As it turns out the curvature of the bands,
#    (the degree to which they are free-electron-like) is highly related to the delocalization
#    of electrons in these bands: The more curved the more delocalized. In some sense
#    "free electrons" correspond to perfect delocalization.


# Henon gridding for dimension explanation
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

fig = figure()
ax = subplot(1,2,1)
ax2 = subplot(1,2,2)
he = Systems.henon()
N = 10000
tr = trajectory(he, N, Ttr = 100)
es = [0.2, 0.05, 0.01]
ns = zero(es)
ax.grid(false)
mini = minima(tr)

rectcols = ["C0", "C1", "C3"]

for (i, e) ∈ enumerate(es)
    if i ≤ 2
        for x in mini[1]:e:1.4
            ax.axvline(x; color = "k", lw = 0.5, alpha = 1/(3^i))
        end
        for y in mini[2]:e:0.4
            ax.axhline(y; color = "k", lw = 0.5, alpha = 1/(3^i))
        end
    end
    p, bins = binhist(tr, e)
    for b in bins
        r = matplotlib.patches.Rectangle(b, e, e; color = rectcols[i], ec = "k", lw = 1/i)
        ax.add_artist(r)
    end
    ns[i] = length(probabilities(tr, e))
end

ax.plot(tr[:, 1], tr[:, 2], ls = "None", marker = ".", color = COLORS[3], ms = 0.75, zorder = 99)
ax.set_yticks([])
ax.set_ylabel("\$y\$")
ax.set_xticks([])
ax.set_xlabel("\$x\$")

ax2.scatter(log.(1 ./ es), log.(ns); c = rectcols, s = 200, zorder = 99)
s = linreg(log.(1 ./ es), log.(ns))[2]
ax2.plot(log.(1 ./ es), log.(1 ./ es) .* s .+ 2, color = "C2")
ax2.text(3.5, 6, "\$\\Delta = $(rdspl(s))\$"; color = "C2")
ax2.set_xlabel("\$\\log ( 1/\\varepsilon)\$"; labelpad = -15)
ax2.set_ylabel("\$\\log ( M)\$")
ax2.set_xticks([2, 3, 4])
ax2.set_xticklabels([2, "", 4])
fig.tight_layout(pad = 0.3)
fig.subplots_adjust(wspace = 0.2)
wsave(plotsdir("5", "henon_gridding"), fig)


using Base.Test
using FastGaussQuadrature

include("test_gausslegendre.jl")
include("test_gaussjacobi.jl")
include("test_gaussradau.jl")
include("test_gausslobatto.jl")


function correction_field(N, V)
    pl, wl = tri_quadrature(N)
    pf, wf = triface_quadrature(N)

    Np = (N + 1) * (N + 2) ÷ 2
    ψf = zeros(3, N + 1, Np)
    for i = 1:3
        ψf[i, :, :] .= vandermonde_matrix(Tri, N, pf[i, :, 1], pf[i, :, 2])
    end

    σ = zeros(3, N + 1, Np)
    for k = 1:Np
        for j = 1:N+1
            for i = 1:3
                σ[i, j, k] = wf[i, j] * ψf[i, j, k]
            end
        end
    end

    V = vandermonde_matrix(Tri, N, pl[:, 1], pl[:, 2])

    ϕ = zeros(3, N + 1, Np)
    for f = 1:3, j = 1:N+1, i = 1:Np
        ϕ[f, j, i] = sum(σ[f, j, :] .* V[i, :])
    end

    return ϕ
end


export
    q_pdf,
    q_pdf!

real_tr_mul(𝐚, 𝐛) = sum(real(𝐚[i, :]' * 𝐛[:, i]) for i in 1:size(𝐚, 1))

"""
    q_pdf(state::AbstractState, θ::Real, x::Real; T=Float64)

Quadrature prabability at point (θ, x)
"""
function q_pdf(T::Type{<:Real}, state::AbstractState, θ::Real, x::Real)
    𝛑̂_res = Matrix{Complex{T}}(undef, state.dim, state.dim)

    return q_pdf!(𝛑̂_res, state, θ, x)
end

q_pdf(state::AbstractState, θ::Real, x::Real) = q_pdf(Float64, state, θ, x)

function q_pdf!(𝛑̂_res::AbstractMatrix, state::StateMatrix, θ::Real, x::Real)
    return real_tr_mul(𝛑̂!(𝛑̂_res, θ, x, dim=state.dim), state.𝛒)
end

function q_pdf!(𝛑̂_res::AbstractMatrix, state::StateVector, args...; kwargs...)
    return q_pdf!(𝛑̂_res, StateMatrix(state), args...; kwargs...)
end

"""
    q_pdf(state::AbstractState, θ::Real, x::Real; T=Float64)

Quadrature prabability at points (θs, xs)
"""
function q_pdf(T::Type{<:Real}, state::AbstractState, θs, xs)
    𝛑̂_res_vec = [Matrix{Complex{T}}(undef, state.dim, state.dim) for _ in 1:Threads.nthreads()]
    𝐩 = Matrix{T}(undef, length(θs), length(xs))

    return q_pdf!(𝛑̂_res_vec, 𝐩, state, θs, xs)
end

q_pdf(state::AbstractState, θs, xs) = q_pdf(Float64, state, θs, xs)

function q_pdf!(𝛑̂_res_vec::Vector{Matrix{Complex{T}}}, 𝐩::Matrix{T}, state::StateMatrix, θs, xs) where {T}
    @sync for (j, x) in enumerate(xs)
        for (i, θ) in enumerate(θs)
            Threads.@spawn 𝐩[i, j] = q_pdf!(𝛑̂_res_vec[Threads.threadid()], state, θ, x)
        end
    end

    return 𝐩
end

function q_pdf!(𝛑̂_res_vec::Vector{Matrix{Complex{T}}}, 𝐩::Matrix{T}, state::StateVector, args...; kwargs...) where {T}
    q_pdf!(𝛑̂_res_vec, 𝐩, StateMatrix(state), args...; kwargs...)
end


using OrdinaryDiffEq
using Debugger
f(u,p,t) = 1.01*u
u0=1/2
tspan = (0.0,1.0)
prob = ODEProblem(f,u0,tspan)
@bp
sol = solve(prob,Midpoint(),reltol=1e-8,abstol=1e-8)


using LinearAlgebra

begin const
    n = 22
    m = 16
    mr = 3
    md = m - mr

    orig = [ 1 2 3 4 5 6 7 8 8 9 10 11 13 1 2 4 5 7 8 14 2 10];
    dest = [ 4 16 15 5 6 10 16 9 12 10 11 14 15 16 6 8 9 11 13 15 4 13];

    absn = [11 18 38 4 8 15 26 4 10 19 26 7 21 33 33 16];
    ordn = [28 21 8 21 17 17 26 9 13 13 18 4 9 18 12 24];

    r = Float64[100, 10,1000, 100, 100, 10, 1000, 100, 1000, 100,
                1000, 1000, 1000, 10, 10, 100, 100, 1000, 100, 1000, 100, 10];
    pr = Float64[105, 104, 110];
    fd = Float64[0.08, -1.30, 0.13, 0.09, 0.16, 0.14, 0.12, 0.07, 0.17, 0.11,
                0.25, 0.01, 0.13];

    ################################################################################
    A = zeros(m, n);

    for l = 1:n
        A[orig[l], l] = +1
        A[dest[l], l] = -1
    end

    Ar = A[1:mr, :]
    Ad = A[mr+1:m, :];

    AdT = Ad[:, 1:md]
    AdI = inv(AdT);

    AdC = Ad[:, md+1:n];

    # Here I stands for identity matrix (former eye function)
    B = [-AdI * AdC ; I];

    # Vecteur des débits admissibles
    q0 = [AdI * fd; zeros(n-md)];
end

function hydrauliqueP(qc)
    # débit des arcs
    q = q0 + B*qc
    # pertes de charge des arcs
    z = r .* abs.(q) .* q
    # flux aux noeuds
    f = [Ar*q ; fd]
    # pressions aux noeuds
    temp = (Ar'*pr) + z
    p = [pr; -AdI'*temp[1:md]]
    return q, z, f, p
end

function hydrauliqueD(pd)
    # pressions aux noeuds
    p = [pr ; pd]
    # pertes de charge aux arcs
    z = -A' * p;
    # débit des arcs
    q = z ./ sqrt(r .* abs.(z))
    # flux aux noeuds
    f = [Ar*q; fd]
    return q, z, f, p
end

# vérification des équations d'équilibre d'un réseau de
# distribution d'eau
function verification(q, z, f, p)
    # écarts maximaux sur les lois de Kirchhoff
    tol1 = max(abs.(A*q-f))
    tol2 = max(abs.(A'*p+z))
    println("Vérification des équations d'équilibre du réseau")
    println("sur les débits: ", tol1, "\t sur les pressions: ", tol2)
end


function arclength(φ₀)
# A terms
    Acoef =  [  -3/2*[ 1.0  0.0  -1/8      0.0      -1/64]; # A₁
               15/16*[ 0.0  1.0   0.0     -1/4        0.0]; # A₂
              -35/48*[ 0.0  0.0   1.0      0.0      -5/16]; # A₃
                     [ 0.0  0.0   0.0  315/512        0.0]; # A₄
                     [ 0.0  0.0   0.0      0.0  -693/1280]; # A₅
             ]
    Aj = Acoef * [n^k for k=1:5]

    ΣA = sum([Aj[j]*sin(2j*φ₀) for j=1:5])
    S̅φ₀ = (m₀*a)/(1+n) * (A₀*φ₀ + ΣA)

    return S̅φ₀
end


using Test

@testset "(32)(32) no-signaling regression" begin

using BellScenarip: LocalPolytope, ConvexPolytope

# Description:
#   Parses inequalities published at http://www.itp.tuwien.ac.at/~svozil/publ/3-3.html
#   from .txt file. Inequalities are represented in the no-signaling subspace for
#   3 inputs and 2 outputs.
#
# Input: None.
#
# Output: Array, contains 16-dim inequality vectors representing a facet inequality
#
function svozil_pitowsky_inequalities()
    filepath = "./test/regression/files/32-32_svozil_pitowsky_facets.txt"

    (inequalities) = open(filepath) do file
        lines = readlines(file)

        ids = Dict(
            "a1" => 2,
            "a2" => 3,
            "a3" => 4,
            "b1" => 5,
            "b2" => 6,
            "b3" => 7,
            "a1b1" => 8,
            "a1b2" => 9,
            "a1b3" => 10,
            "a2b1" => 11,
            "a2b2" => 12,
            "a2b3" => 13,
            "a3b1" => 14,
            "a3b2" => 15,
            "a3b3" => 16,
        )

        inequalities = []

        for id in 7:2:size(lines)[1]
            line = lines[id]

            lhs_match = match(r"^#\d+:\s(\d)\s\?", line)
            rhs_matches = collect( eachmatch(r"([+-])\s?(\d)?(a\db\d|a\d|b\d|)", line) )

            if size(rhs_matches)[1] > 0
                ineq = zeros(Int64, (1,16))
                ineq[1] = parse(Int64, "-" * lhs_match.captures[1])

                for match in rhs_matches
                    index = ids[match[3]]
                    sign = match[1]
                    multiplier = match.captures[2] === nothing ? "1" : match.captures[2]

                    ineq[index] = parse(Int64, sign * multiplier )
                end

                push!(inequalities, ineq)
            end
        end

        (inequalities)
    end

    inequalities
end

@testset "verifying polytope" begin
    dir = "./test/regression/files/"

    vertices = LocalPolytope.vertices((3,2),(3,2))

    @test size(vertices[1]) == (15,1)

    constraints = ConvexPolytope.facet_constraints(vertices, dir=dir)
    @test constraints[2] == []
    @test constraints[3] == []

    inequalities = constraints[1]
    @test size(inequalities) == (684,)
    @test size(inequalities[1]) == (1,16)
    @test issetequal(inequalities, svozil_pitowsky_inequalities())
end

end


using Plots
using Wavelets
WT.paul10
supertype(typeof(WT.paul10))
wavelet(WT.morl)
wavelet(WT.morl,4)
plotlyjs()
J=11; n = 2^J
x = testfunction(n,"Bumps")
supertype(typeof(WT.dog0))
supertype(typeof(WT.paul5))
supertype(typeof(WT.morl))
EmphasizeFrequencyInfo = WT.Morlet(20)
y = cwt(x, wavelet(WT.morl))
heatmap(abs.(y)); plot!(x+90,label="")
y = cwt(x, wavelet(WT.dog1))
heatmap(abs.(y)); plot!(x+90,label="")
y = cwt(x, wavelet(WT.paul4))
heatmap(abs.(y)); plot!(x+90,label="")


@testset "inverse.jl" begin
    InverseFunctions.test_inverse(log1pexp, randn())
    InverseFunctions.test_inverse(logexpm1, randexp())

    InverseFunctions.test_inverse(log1mexp, -randexp())

    InverseFunctions.test_inverse(log2mexp, log(2) - randexp())

    InverseFunctions.test_inverse(logistic, randn())
    InverseFunctions.test_inverse(logit, rand())
end


using LinearAlgebra
import Base: *
import LinearAlgebra: mul!

mutable struct LBFGSOperator{M,R,I,T}
    currmem::I
    curridx::I
    s::T
    y::T
    s_M::Vector{T}
    y_M::Vector{T}
    ys_M::Vector{R}
    alphas::Vector{R}
    H::R
end

function LBFGSOperator{M}(x::T) where {M,T}
    R = real(eltype(x))
    s_M = [zero(x) for i = 1:M]
    y_M = [zero(x) for i = 1:M]
    s = zero(x)
    y = zero(x)
    ys_M = zeros(R, M)
    alphas = zeros(R, M)
    LBFGSOperator{M,R,typeof(0),T}(0, 0, s, y, s_M, y_M, ys_M, alphas, one(R))
end

LBFGSOperator(M, x) = LBFGSOperator{M}(x)

function update!(L::LBFGSOperator{M}, s, y) where M
    L.s .= s
    L.y .= y
    ys = real(dot(L.s, L.y))
    if ys > 0
        L.curridx += 1
        if L.curridx > M
            L.curridx = 1
        end
        L.currmem += 1
        if L.currmem > M
            L.currmem = M
        end
        L.ys_M[L.curridx] = ys
        copyto!(L.s_M[L.curridx], L.s)
        copyto!(L.y_M[L.curridx], L.y)
        yty = real(dot(L.y, L.y))
        L.H = ys / yty
    end
    return L
end

function reset!(L::LBFGSOperator{M,R,I}) where {M,R,I}
    L.currmem, L.curridx = zero(I), zero(I)
    L.H = one(R)
end

function (*)(L::LBFGSOperator, v)
    w = similar(v)
    mul!(w, L, v)
end

# Two-loop recursion

function mul!(d, L::LBFGSOperator, v)
    d .= v
    idx = loop1!(d, L)
    d .*= L.H
    loop2!(d, idx, L)
    return d
end

function loop1!(d, L::LBFGSOperator{M}) where M
    idx = L.curridx
    for i = 1:L.currmem
        L.alphas[idx] = real(dot(L.s_M[idx], d)) / L.ys_M[idx]
        d .-= L.alphas[idx] .* L.y_M[idx]
        idx -= 1
        if idx == 0
            idx = M
        end
    end
    return idx
end

function loop2!(d, idx, L::LBFGSOperator{M}) where M
    for i = 1:L.currmem
        idx += 1
        if idx > M
            idx = 1
        end
        beta = real(dot(L.y_M[idx], d)) / L.ys_M[idx]
        d .+= (L.alphas[idx] - beta) .* L.s_M[idx]
    end
    return d
end

struct LBFGS{M} end

LBFGS(M) = LBFGS{M}()

acceleration_style(::Type{<:LBFGS}) = QuasiNewtonStyle()

function initialize(::LBFGS{M}, x) where M
    return LBFGSOperator{M}(x)
end


module HessianTest

import Calculus

using Base.Test
using ForwardDiff

include(joinpath(dirname(@__FILE__), "utils.jl"))

#############################
# rosenbrock hardcoded test #
#############################

x = [0.1, 0.2, 0.3]
v = rosenbrock(x)
g = [-9.4, 15.6, 52.0]
h = [-66.0  -40.0    0.0;
     -40.0  130.0  -80.0;
       0.0  -80.0  200.0]

for c in (Chunk{1}(), Chunk{2}(), Chunk{3}())

    # single-threaded #
    #-----------------#
    @test_approx_eq h ForwardDiff.hessian(rosenbrock, x, c)

    out = similar(x, 3, 3)
    ForwardDiff.hessian!(out, rosenbrock, x, c)
    @test_approx_eq out h

    out = HessianResult(x)
    ForwardDiff.hessian!(out, rosenbrock, x, c)
    @test_approx_eq ForwardDiff.value(out) v
    @test_approx_eq ForwardDiff.gradient(out) g
    @test_approx_eq ForwardDiff.hessian(out) h

    # multithreaded #
    #---------------#
    if ForwardDiff.IS_MULTITHREADED_JULIA
        @test_approx_eq h ForwardDiff.hessian(rosenbrock, x, c; multithread = true)

        out = similar(x, 3, 3)
        ForwardDiff.hessian!(out, rosenbrock, x, c; multithread = true)
        @test_approx_eq out h

        out = HessianResult(x)
        ForwardDiff.hessian!(out, rosenbrock, x, c; multithread = true)
        @test_approx_eq ForwardDiff.value(out) v
        @test_approx_eq ForwardDiff.gradient(out) g
        @test_approx_eq ForwardDiff.hessian(out) h
    end
end

########################
# test vs. Calculus.jl #
########################

for f in VECTOR_TO_NUMBER_FUNCS
    v = f(X)
    g = ForwardDiff.gradient(f, X)
    h = ForwardDiff.hessian(f, X)
    # finite difference approximation error is really bad for Hessians...
    @test_approx_eq_eps h Calculus.hessian(f, X) 0.01
    for c in CHUNK_SIZES
        for usecache in (true, false)
            println("  ...testing $f with (chunk size = $c) and (usecache = $usecache)")
            chunk = Chunk{c}()

            # single-threaded #
            #-----------------#
            out = ForwardDiff.hessian(f, X, chunk; usecache = usecache)
            @test_approx_eq out h

            out = similar(X, length(X), length(X))
            ForwardDiff.hessian!(out, f, X, chunk; usecache = usecache)
            @test_approx_eq out h

            out = HessianResult(X)
            ForwardDiff.hessian!(out, f, X, chunk; usecache = usecache)
            @test_approx_eq ForwardDiff.value(out) v
            @test_approx_eq ForwardDiff.gradient(out) g
            @test_approx_eq ForwardDiff.hessian(out) h

            # multithreaded #
            #---------------#
            if ForwardDiff.IS_MULTITHREADED_JULIA
                out = ForwardDiff.hessian(f, X, chunk; multithread = true, usecache = usecache)
                @test_approx_eq out h

                out = similar(X, length(X), length(X))
                ForwardDiff.hessian!(out, f, X, chunk; multithread = true, usecache = usecache)
                @test_approx_eq out h

                out = HessianResult(X)
                ForwardDiff.hessian!(out, f, X, chunk; multithread = true, usecache = usecache)
                @test_approx_eq ForwardDiff.value(out) v
                @test_approx_eq ForwardDiff.gradient(out) g
                @test_approx_eq ForwardDiff.hessian(out) h
            end
        end
    end
end

end # module


using Flux, Plots
using Base.Iterators: repeated
using Flux: onehotbatch, onecold, crossentropy, throttle, @epochs
using LinearAlgebra

gridsize = 150;
tstart=0.0
tend=20.0
myfunc(x) = sin(x)+0.5*sin(2*x)+0.24*cos(x);
x = collect(range(tstart,stop=tend,length=gridsize));
y= myfunc.(x);

xtest=collect(tend:0.1:2*tend)
xpred=collect(tstart:0.1:2*tend)

data = []
for i in 1:length(x)
    push!(data, ([x[i]], y[i])) #have to input a list
end

Q = 20;
ann = Chain(Dense(1,Q,sin),Dense(Q,1));

function loss(x, y)
    pred=ann(x)
    loss=Flux.mse(ann(x), y)
    return loss
end

opt = ADAM()

ps=params(ann)

for j=1:20
    ann[1].W[j]=j
end
delete!(ps, ann[1].W) #Make the first layers weights (frequencies not trainable)
@epochs 3000 Flux.train!(loss,ps, data, opt)

plot(xpred, ann(xpred')', color=:red, lw=2.0, label="")
scatter!(x,myfunc.(x), color=:blue, legend=false)
scatter!(xtest,myfunc.(xtest), color=:green, lw=2.0, label="", markershape=:cross)
ylims!((-1.5,1.5))
xlims!((tstart,2*tend))
title!("Curve Fitting with Neural Networks")
xlabel!("Input")
ylabel!("Output")


type D2Function
    value::Float64
    gradient::Array{Float64,1}
    hessian::Array{Float64,2}
end

###{{{ Newton-Raphson

function NR(f, x; iter=200,tol=1e-12,verbose=false,method=0,gamma=1,trace=0)
    val = f(x)::D2Function
    # S = val.grad["grad"]';  H = val["hess"]; loglik=val["val"]
    S = val.gradient::Vector{Float64};  H = val.hessian::Matrix{Float64}; loglik=val.value::Float64
    p = size(x,1)
    counter = 1
    for j=1:iter
        if verbose print("S=",S') end
        s = s0 = mean(S.*S)
        if isnan(s) error("Lack of convergence (bad starting values?)") end
        if s<tol break end
        if method==1
            W = pinv(-H + gamma*(S'S)[1]*eye(p))
        elseif method==2
            W = pinv(-H + gamma*S*S')
        else
            W = pinv(-H)
        end
        if trace>0
            print(S')
        end
        delta = W*S
        x0 = x; loglik0 = loglik
        lambda=1;
        while s>=s0 || loglik<loglik0
            x = x0+lambda*delta
            val = f(x)
            counter += 1
            ##S = val["grad"]';  H = val["hess"]; loglik=val["val"]
            S = val.gradient::Vector{Float64};  H = val.hessian::Matrix{Float64}; loglik=val.value::Float64
            s = mean(S.*S)
            if isnan(s) s=s0 end
            lambda = lambda/2.0
            if (lambda<1e-12)
                error("Lack of convergence (bad starting values?)")
            end
        end
    end
    return x,val,counter
end

###}}} Newton-Raphson


function center_of_mass(
    pos::Vector{T},
    box::MolecularBox{T,N,P};
    weights = ones(Float64,length(pos))
) where {T,N,P}
    if sum(P) != N
        error("Calculating the center of mass is currently only supported for"*
            "fully-periodic systems")
    end
    if length(pos)==1 ; return pos[1] ; end

    L = box.lengths

    length(weights) == length(pos) || error("Weights and positions mismatch")

    invtotweight = 1.0/sum(weights)
    i2pi = 0.5/pi

    a = zero(eltype(pos))
    b = zero(eltype(pos))

    for i in eachindex(pos)
        t = pos[i]./L*2pi
        a -= cos.(t) * weights[i]
        b -= sin.(t) * weights[i]
    end
    a = a*invtotweight
    b = b*invtotweight

    tcom = atan.(b,a)+pi
    com = tcom.*i2pi.*L
end


# This is code for generating random trees of various sorts.
# The hope is that one might have low stretch,
# although none is guaranteed to.

using DataStructures

# sample edges with probability proportional to their weight

"""
    tree = randishKruskal(A)

A heuristic for computing low-stretch spanning trees.  Where Kruskal's MST algorithm adds edges in order of weight, this algorithm adds them at random with probability proportional to their weight.
"""
function randishKruskal(mat::SparseMatrixCSC)
    n = size(mat)[1]
    (ai,aj,av) = findnz(triu(mat))
    m = length(av)

    s = Sampler(av)

    comps = IntDisjointSets(n)

    treeinds = zeros(Int64,n-1)
    numintree = 0
    for j in 1:m
        i = pop!(s)
        if !DataStructures.in_same_set(comps,ai[i],aj[i])
            numintree = numintree+1
            treeinds[numintree] = i
            DataStructures.union!(comps,ai[i],aj[i])
        end
    end

    tree = sparse(ai[treeinds],aj[treeinds],av[treeinds],n,n)
    tree = tree + tree'

    return tree

end

"""
    tree = randishPrim(A)

A heuristic for computing low-stretch spanning trees.  Where Prim's MST algorithm grows a cluster by always adding the edge on the boundary of maximum weight, this algorithm adds a boundary edge with probability proportional to its weight.
"""
function randishPrim(mat::SparseMatrixCSC{Tval,Tind}) where {Tval,Tind}

  n = mat.n
  m = nnz(mat)

  (ai, aj, av) = findnz(mat)
  flipInd = flipIndex(mat)

  visited = zeros(Bool,n)

  s = Sampler(m)

  treeEdges = zeros(Tind,n-1)
  numEdges = 0

  # should randomize the initial choice
  v = one(Tind)

  # always add edge by highest vertex number first
  for ind in mat.colptr[v]:(mat.colptr[v+1]-1)
      wt = mat.nzval[ind]
      push!(s, ind, wt)
  end

  visited[v] = true

  while s.nitems > 0

    edge = pop!(s)

    v = aj[edge]
    if visited[v]
        v = ai[edge]
    end

    if !visited[v]

        numEdges += 1
        treeEdges[numEdges] = edge

        for ind in mat.colptr[v]:(mat.colptr[v+1]-1)
            u = mat.rowval[ind]

            wt = mat.nzval[ind]
            useind = v < u ? ind : flipInd[ind]
            if !visited[u]
                push!(s, useind, wt)

            end
        end

        visited[v] = true
    end

  end

  treeEdges = treeEdges[1:numEdges]

  tr = sparse(ai[treeEdges],aj[treeEdges],av[treeEdges],n,n)
  tr = tr + tr';

  return tr

end # randishPrim


"""
    WoodburyPDMat(
        A::AbstractMatrix{T}, D::Diagonal{T}, S::Diagonal{T},
    ) where {T<:Real}

Lazily represents matrices of the form
```julia
W = A * D * A' + S
```
`D` and `S` must have only non-negative entries.

Using this matrix type is a good idea if `size(A, 1) > size(A, 2)` as the structure in the
matrix can be exploited to accelerate operations involving `W`'s inverse [1], such as
`invquad`, and it's determinant [2], such as `logdet`.

You probably don't want to use this matrix type if `size(A, 1) < size(A, 2)`.

[1] - https://en.wikipedia.org/wiki/Woodbury_matrix_identity
[2] - https://en.wikipedia.org/wiki/Matrix_determinant_lemma
"""
struct WoodburyPDMat{
    T<:Real, TA<:AbstractMatrix{T}, TD<:Diagonal{T}, TS<:Diagonal{T},
} <: AbstractPDMat{T}
    A::TA
    D::TD
    S::TS
    function WoodburyPDMat(
        A::AbstractMatrix{T}, D::Diagonal{T}, S::Diagonal{T},
    ) where {T<:Real}
        validate_woodbury_arguments(A, D, S)
        return new{T, typeof(A), typeof(D), typeof(S)}(A, D, S)
    end
end

function WoodburyPDMat(
    A::AbstractMatrix{T}, D::AbstractVector{T}, S::AbstractVector{T},
) where {T<:Real}
    return WoodburyPDMat(A, Diagonal(D), Diagonal(S))
end

WoodburyPDMat{T, TA, TD, TS}(vals...) where {T, TA, TD, TS} = WoodburyPDMat(vals...)

PDMats.dim(W::WoodburyPDMat) = size(W.A, 1)

# Convesion method. Primarily useful for testing purposes.
Base.Matrix(W::WoodburyPDMat) = W.A * W.D * W.A' + W.S

Base.getindex(W::WoodburyPDMat, inds...) = getindex(Matrix(W), inds...)

function validate_woodbury_arguments(A, D, S)
    if size(A, 1) != size(S, 1)
        throw(ArgumentError("size(A, 1) != size(S, 1)"))
    end
    if size(A, 2) != size(D, 1)
        throw(ArgumentError("size(A, 2) != size(D, 1)"))
    end
    if any(x -> x < 0, diag(D))
        throw(ArgumentError("Detected negative element on diagonal of D: $(D)"))
    end
    if any(x -> x < 0, diag(S))
        throw(ArgumentError("Detected negative element on diagonal of S: $(S)"))
    end
end

@non_differentiable validate_woodbury_arguments(A, D, S)

function LinearAlgebra.logdet(W::WoodburyPDMat)
    C_S = cholesky(W.S)
    B = C_S.U' \ (W.A * cholesky(W.D).U')
    return logdet(C_S) + logdet(cholesky(Symmetric(I + B'B)))
end

# Utilises the matrix inversion lemma to produce an efficient implementation.
function PDMats.invquad(W::WoodburyPDMat{<:Real}, x::AbstractVector{<:Real})
    C_S = cholesky(W.S)
    B = C_S.U' \ (W.A * cholesky(W.D).U')
    α = C_S.U' \ x
    β = B' * α
    return α'α - sum(abs2, cholesky(Symmetric(I + B'B)).U' \ β)
end

# This doesn't get us the computational wins, but it's unclear how to construct a
# root for a Woodbury matrix. Consequently, if performance is very important when sampling,
# it's necessary to implement a method of `rand` or `_rand` that explicitly uses ancestral
# sampling to exploit the approximately low-rank structre in a WoodburyPDMat.
function PDMats.unwhiten!(r::DenseVecOrMat, W::WoodburyPDMat{<:Real}, x::DenseVecOrMat)
    return unwhiten!(r, PDMat(Symmetric(Matrix(W))), x)
end

# NOTE: the parameterisation to scale up the Woodbury matrix is not unique. Here we
# implement one way to scale it.
*(a::WoodburyPDMat, c::Real) = WoodburyPDMat(a.A, a.D * c, a.S * c)
*(c::Real, a::WoodburyPDMat) = a * c


@doc "Calculate the distance between point A (latA,lonA) and Point B(latB,lonB)"
function distanceInKm(latA::AbstractFloat,lonA::AbstractFloat,latB::AbstractFloat,lonB::AbstractFloat)
  return 2 * 6372.8 * asin(sqrt(sind((latB-latA)/2)^2 + cosd(latA) * cosd(latB) * sind((lonB - lonA)/2)^2))
end

@doc "Calculate the distance between point A(latA,lonA) and a vector of Points(latB,lonB)"
function distanceInKm(latA::AbstractFloat,lonA::AbstractFloat,latB::Vector{AbstractFloat},lonB::Vector{AbstractFloat})
    n = size(latB,1)
    resp::Vector{AbstractFloat}=Vector{AbstractFloat}(n)
    if (size(lonB,1)!=n)
     error("latB and lonB have different sizes")
    end
    @simd for i in 1:n  # @simd vectoraze for loop
      @inbounds resp[i] = distanceInKm(latA,lonA,latB[i],lonB[i]) # @inbounds to not check bounds
    end
end

@doc "Calculate the distance between a vector of points"
function distanceInKm(latA::Vector{AbstractFloat},lonA::Vector{AbstractFloat},latB::Vector{AbstractFloat},lonB::Vector{AbstractFloat})
    n = size(latA,1)
    isEqualSize::Bool = n == size(lonA,1) == size(latB,1)==size(lonB,1)

    if (!isEqualSize)
     error("the vectors have different sizes")
    end
    resp::Vector{AbstractFloat}=Vector{AbstractFloat}(n)
    @simd for i in 1:n  # @simd vectoraze for loop
      @inbounds resp[i] = distanceInKm(latA[i],lonA[i],latA[i],lonA[i])
    end
end

@doc "Calculate the Azimute between A and B"
function azimuthLineAtoB(latA::AbstractFloat,lonA::AbstractFloat,latB::AbstractFloat,lonB::AbstractFloat)
     numer = cosd(latB) * sind(lonB - lonA)
     denom = sind(latB) * cosd(latA) - cosd(latB) * sind(latA) * cosd(lonB - lonA)
     return rad2deg((atan2(numer, denom) + 2π) % 2π)
end


using IterTools
using Combinatorics

## Tools: functions for multiple iterators
function comb2col(comb, n)
  #given a possible combination of 1:n, returns a matrix with 1 at the choosen values
  temp = zeros(Int, n, 1)
  temp[comb] = 1
  return temp
end

function matcheck(mat, m)
  #returns false if any row of the matrix is zero or any row or column has a sum greater than m
  a = sum(mat, 2) #sum of each row
  if minimum(a) == 0 #zero row
    return false
  end
  if maximum(a) > m #row with sum over m
    return false
  end
  if maximum(sum(mat,1)) > m #col with a sum over m
    return false
  end
  return true
end

## Angle matrix 1st row
function acol1(n)
  #returns an iterator for the first (straight edge) column of the angle matrix
  iter = [zeros(Int, n, 1)] #first element
  for i = 1:n-3 #loop through number of angles containing edge
    comb = combinations(1:n-3, i) #combinations 1:n, choose i
    temp = imap(x->comb2col(x .+ 3, n), comb) #convert to 1s and 0s. the .+3 shifts the combination to ensure that the first 3 elements are zero
    iter = chain(iter, temp) #build the chain
  end
  return iter
end

## Angle matrix top half
function rowpadzeros(comb, k, numzeros)
  #returns the combination with 1 at the choosen values, then padded with zeros
  rowbase = comb2col(comb, k)' #non-zero part of the rowbase. transposed to make a row
  return [zeros(Int, 1, numzeros) rowbase]
end

function aupperrow(n, k)
  #returns the iterator for a angle matrix row with total lenth n.
  #the first n-k elements are always zero
  iter1 = [zeros(Int, 1, n)] #first element
  comb = combinations(1:k, 1) #choose one of the last n-k rows
  iter = imap(x->rowpadzeros(x, k, n-k), comb)
  return chain(iter1, iter)
end

function aupper(n)
  #returns an iterator for the upper triangular half of the angle matrix
  allrows = (aupperrow(n, n-i+1) for i in 1:n) #each row
  a = product(allrows...)
  b = imap(x -> vcat(x...), a)
  return b
end

## Angle matrix bottom half
immutable AngRowGen
  #generates the possible row with sum less than m
  n::Int #total row length
  m::Int #row maximum sum
  l::Int #non-zero row length
end

function Base.start(a::AngRowGen)
  row = zeros(Int, 1, a.n)
  row[1] = -1 #so the first iteration gives all zeros
  return row
end

function Base.next(a::AngRowGen, row)
  for i in 1:a.l #iterate through element, starting with first non-zero element
    if sum(row) < a.m #still room in the raw
      row[i] += 1 #increment and return
      return row, row
    else #no room
      row[i] = 0
    end
  end
  return row, row
end

function Base.done(a::AngRowGen, row)
  return row[a.l] == a.m
end

function alower(n, m)
  #returns an iterator for the upper triangular half of the angle matrix
  allrows = (AngRowGen(n, m, i) for i in 1:n-1)
  a = product(allrows...)
  b = imap(x->vcat(zeros(Int, 1, n), x...), a) #build matrix with zero row on top
  return chain([zeros(Int, n, n)], b) #because product is removing the all zero option
end

## Build anlge matrix
function anglemat(n, m)
  #returns an iterator of the possible angle matrices
  a = product(aupper(n), alower(n, m)) #all upper and lower half combinations
  b = imap(x->x[1] + x[2], a) #add upper and lower half
  c = Iterators.filter(x->matcheck(x, m), b) #filter out zero rows and rows and columns summing over m
  d = product(acol1(n), c) #combine first column and rest of matrix
  f = imap(x->hcat(x...), d)
  return f
end

## Side matrix row
immutable SideRowGen
  #generates the possible row with sum less than m
  n::Int #total row length
  m::Int #row maximum sum
  l::Int #sum of elements from l to end is less than floor(Int, sqrt(m))
end

function Base.start(d::SideRowGen)
  row = zeros(Int, 1, d.n)
  return row
end

function Base.next(d::SideRowGen, row)
  flsqm = floor(Int, sqrt(d.m))
  for i in 1:d.n #iterate through each element
    #(still room in the row) and (element is in beginning of or there is room at the end)
    if (sum(row) < d.m) && ((i < d.l) || sum(row[d.l:end]) < flsqm)
      row[i] += 1
      return row, row
    else #no room
      row[i] = 0
    end
  end
  return row, row
end

function Base.done(d::SideRowGen, row)
  flsqm = floor(Int, sqrt(d.m))
  if row[end] == flsqm #last element at max
    if d.l == 1 #avoid looking for zero index
      return row[end] == flsqm #last unconstrained value + last value
    else
      return row[d.l-1] + row[end] == d.m #last unconstrained value + last value
    end
  end
  return false
end

## Build side matrix
function sidemat(n, m)
  #returns an iterator of the possible side matricies

  ###### correct version. unexpected behavior in product is skipping the first value
  #allrows = (SideRowGen(n, m, i) for i in 1:n) #build all row iterators
  #a = product(allrows...) #combine all row iterators
  ###### end correct version

  ###### begin workaround
  allrows = (SideRowGen(n, m, i) for i in 2:n) #build all row iterators
  firstrow = ([zeros(Int, 1, i) 1 zeros(Int, 1, n-i-1)] for i = 0:n-1)
  a = product(firstrow, allrows...) #combine all row iterators
  ###### end workaround

  b = imap(x->vcat(x...),a)
  c = Iterators.filter(x->matcheck(x,m), b) ##filter out zero rows and rows and columns summing over m
  #(zero rows are impossible, this could be sped up)
  return c
end


function calc_submatrices(
  nbset  # tuple with the following elements.
    # nodes: array of Strings containing independent node names.
    # num_nodes: number of nodes.
    # branches: array of struct Branch.
    # num_branches: sum of branch phases.
    # num_dyns: number of dynamic branches.
  )  #=

calculates submatrices to build the sparse-tableau equations of the given circuit. The following variables are returned as the elements of the tuple "smset".
- Ac: branch-node connectivity matrix.
- Z: coefficient matrix of i.
- Y: coefficient matrix of v.
- B1: incidence matrix of alpha.
- B2: incidence matrix of [ i v ].

Author: Taku Noda
Started on Jan. 6, 2021
Update on June 29, 2021

Copyright 2021 Central Research Institute of Electric Power Industry

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

=#

( nodes, num_nodes, branches, num_branches, num_dyns ) = nbset
Ac = build_Ac( num_nodes, branches, num_branches )
Z, Y, B1, B2 = build_ZYB!( branches, num_branches, num_dyns )
smset = ( Ac, Z, Y, B1, B2 )

return smset

end


using Turing, Distributions, MCMCChains, KernelDensity,AdvancedHMC
using StatsBase, CSV, DataFrames, Plots, StatsPlots
using ReverseDiff
using Combinatorics
Turing.setadbackend(:reversediff)

# functions
function top(x,n = 3, test = false)
    """
    Generates list of every possible order of n picks from the neural net prbabilities
    e.g. the 3 most likely species (in order) for a given data point.
    Adds up the instances of each combination, in a given data set x.
    If x is training data subset for instances of a given species, then
    return proportions.
    If x is test data, then return the actual counts.
    """
top3 = []
conf = []
tops = collect(permutations([1,2,3,4,5],n))
for i in 1:size(x)[1]
    row = Array(x[i,1:end])

    push!(topx,sortperm(row, rev = true)[1:n])
    push!(conf,maximum(row))
end

    topxs = [sum([j == k for j in topx]) for k in tops]
    if test = true

        return topxs
    else
        return topxs ./ length(top3)
end

@model top3_model(a,o,p,s,v, y, n) = begin
        """
        a - v are the top n values from the 'top' function.
        y is the top n values for the test data.
        n is the number of observations in the test data
        """
    props ~ Dirichlet(5,1)
    x = (a .* props[1] .* n) .+ (o .* props[2] .* n) .+ (p .* props[3] .* n) .+ (s .* props[4] .* n) .+ (v .* props[5] .* n)
    y .~ Poisson.(x)
end;

y5 = CSV.read("5y.csv")[1:end,3:end]# actualy values
names!(y5,[:aa,:od,:pc,:sr,:vg])
pred5 = CSV.read("5pred.csv")[1:end,3:end] # y value predictions (probabilities) from neural net
names!(pred5,[:aa,:od,:pc,:sr,:vg])

n = size(y5)[1] #number of observations

# split data into training and test sets
train_pred5 = pred5[1:n÷2,1:end]
train_y5 = y5[1:n÷2,1:end]
test_pred5 = pred5[n÷2:end,1:end]
test_y5= y5[n÷2:end,1:end]

choices = train_pred5[train_y5[:aa] .== 1,1:end]
aatop3 = top(choices,3)
choices = train_pred5[train_y5[:od] .== 1,1:end]
odtop3 = top(choices,3)
choices = train_pred5[train_y5[:pc] .== 1,1:end]
pctop3 = top(choices,3)
choices = train_pred5[train_y5[:sr] .== 1,1:end]
srtop3 = top(choices,3)
choices = train_pred5[train_y5[:vg] .== 1,1:end]
vgtop3 = top(choices,3)

#r = sample(1:size(test_y5)[1],Weights(1.0 .- Array(test_y5[1:end,1]) ), n÷4)
r = sample(1:size(test_y5)[1], n÷4)
pred_sub = test_pred5[r,1:end]
y_sub5 = test_y5[r,1:end]
y_vals1 = top(pred_sub,1,test = true)
y_vals2 = top(pred_sub,2,test = true)
y_vals3 = top(pred_sub,3, test = true)

# draw samples
top_chain1 = sample(top3_model(aatop,odtop,pctop,srtop,vgtop, y_vals1, n÷4), MH(), 1000000)
top_chain2 = sample(top3_model(aatop2,odtop2,pctop2,srtop2,vgtop2, y_vals2, n÷4), MH(), 1000000)
top_chain3 = sample(top3_model(aatop3,odtop3,pctop3,srtop3,vgtop3, y_vals3, n÷4), MH(), 1000000)

# plotting
density(top_chain1["props[1]"], label = "aa1")
density!(top_chain2["props[1]"], label = "aa2")
density!(top_chain3["props[1]"], label = "aa3")
vline!([sum(y_sub5[:aa])/(n÷4)], label = "aa")

density(top_chain1["props[2]"], label = "od1")
density!(top_chain2["props[2]"], label = "od2")
density!(top_chain3["props[2]"], label = "od3")
vline!([sum(y_sub5[:od])/(n÷4)], label = "od")

density(top_chain1["props[3]"], label = "pc1")
density!(top_chain2["props[3]"], label = "pc2")
density!(top_chain3["props[3]"], label = "pc3")
vline!([sum(y_sub5[:pc])/(n÷4)], label = "pc")

density(top_chain1["props[4]"], label = "sr1")
density!(top_chain2["props[4]"], label = "sr2")
density!(top_chain3["props[4]"], label = "sr3")
vline!([sum(y_sub5[:sr])/(n÷4)], label = "sr")

density(top_chain1["props[5]"], label = "vg1")
density(top_chain2["props[5]"], label = "vg2")
density(top_chain3["props[5]"], label = "vg3")
vline!([sum(y_sub5[:vg])/(n÷4)], label = "vg")


#___  ____       ___ _____   _   _       _ _
#|  \/  | |     |_  /  __ \ | | | |     (_| |
#| .  . | |       | | /  \/ | | | |_ __  _| |_ ___
#| |\/| | |       | | |     | | | | '_ \| | __/ _ \
#| |  | | |___/\__/ | \__/\ | |_| | | | | | || (_) |
#_______\_____\____/ \____/  _____|_|___|_____\_____ _____ _____
#| ___ \        (_)         | | \ \ / / / __  |  _  / __  |  _  |
#| |_/ _ __ ___  _  ___  ___| |_ \ V /  `' / /| |/' `' / /| |/' |
#|  __| '__/ _ \| |/ _ \/ __| __|/   \    / / |  /| | / / |  /| |
#| |  | | | (_) | |  __| (__| |_/ /^\ \ ./ /__\ |_/ ./ /__\ |_/ /
#\_|  |_|  \___/| |\___|\___|\__\/   \/ \_____/\___/\_____/\___/
#              _/ |
#             |__/
#
# This code is part of the proposal of the team "MLJC UniTo" - University of Turin
# for "ProjectX 2020" Climate Change for AI.
# The code is licensed under MIT 3.0
# Please read readme or comments for credits and further information.

# Compiler: Julia 1.5.3

# Short description of this file: Test Neural PDE low level

using Pkg
Pkg.add("ModelingToolkit")
Pkg.add("DiffEqFlux")
Pkg.add("Flux")
Pkg.add("Plots")
Pkg.add("Test")
Pkg.add("Optim")
Pkg.add("CUDA")

using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using Plots #GalacticOptim
using PyPlot

@parameters t, x, y, θ
@variables u(..)
@derivatives Dt'~t
@derivatives Dx'~x
@derivatives Dy'~y
@derivatives Dxx''~x
@derivatives Dyy''~y

#2D PDE
eq  = Dt(u(t,x,y,θ)) + u(t,x,y,θ)*Dx(u(t,x,y,θ)) + u(t,x,y,θ)*Dy(u(t,x,y,θ)) ~ Dxx(u(t,x,y,θ)) + Dyy(u(t,x,y,θ))

# Initial and boundary conditions
bcs = [u(0,x,y,θ) ~ exp(x+y)*cos(x+y) ,
       u(2,x,y,θ) ~ exp(x+y)*cos(x+y+4*2) ,
       u(t,0,y,θ) ~ exp(y)*cos(y+4t),
       u(t,2,y,θ) ~ exp(2+y)*cos(2+y+4t) ,
       u(t,x,0,θ) ~ exp(x)*cos(x+4t),
       u(t,x,2,θ) ~ exp(x+2)*cos(x+2+4t)]

# Space and time domains
domains = [x ∈ IntervalDomain(0.0,2.0),
           y ∈ IntervalDomain(0.0,2.0),
           t ∈ IntervalDomain(0.0,2.0)]# Discretization

dx = 0.125; dy= 0.125; dt = 0.125

# Neural network
chain = FastChain(FastDense(3,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))

#strategy = GridTraining()
strategy = StochasticTraining()
discretization = PhysicsInformedNN([dt,dx,dy],chain,strategy=strategy)

indvars = [t,x,y]
depvars = [u]
dim = length(domains)

expr_pde_loss_function = build_loss_function(eq,indvars,depvars)
expr_bc_loss_functions = [build_loss_function(bc,indvars,depvars) for bc in bcs]
train_sets = generate_training_sets(domains,dx,bcs,indvars,depvars)

train_domain_set, train_bound_set, train_set = train_sets

phi = discretization.phi
autodiff = discretization.autodiff
derivative = discretization.derivative
initθ = discretization.initθ

pde_loss_function = get_loss_function(eval(expr_pde_loss_function),
                                      train_domain_set,
                                      phi,
                                      derivative,
                                      strategy)
bc_loss_function = get_loss_function(eval.(expr_bc_loss_functions),
                                     train_bound_set,
                                     phi,
                                     derivative,
                                     strategy)

function loss_function(θ,p)
    return pde_loss_function(θ) + bc_loss_function(θ)
end

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

f = NeuralPDE.OptimizationFunction(loss_function, GalacticOptim.AutoZygote())

prob = GalacticOptim.OptimizationProblem(f, initθ)

# optimizer
opt = GalacticOptim.ADAM(0.01)
res = GalacticOptim.solve(prob, opt; cb = cb, maxiters=30)
prob
phi = discretization.phi

ts,xs,ys = [domain.domain.lower:dx:domain.domain.upper for domain in domains]
#u_predict_contourf = reshape([first(phi([t,x,y],res.minimizer)) for t in ts for x in xs for y in ys] ,length(ys),length(xs),length(ts))
#Plots.plot(ts, xs, ys, u_predict_contourf, linetype=:contourf,title = "predict")

#u_predict = [[first(phi([t,x,y],res.minimizer)) for x in xs for y in ys] for t in ts ]
u_predict = [reshape([first(phi([t,x,y],res.minimizer)) for x in xs for y in ys], (length(xs),length(ys))) for t in ts ]
p1= Plots.plot(xs, u_predict[2],title = "t = 0.1");
p2= Plots.plot(xs, u_predict[6],title = "t = 0.5");
p3= Plots.plot(xs, u_predict[end],title = "t = 1");
Plots.plot(p1,p2,p3)


# Tests for Profit Profit Efficiency Enhanced Russell Graph Slack Based Measure DEA Models
@testset "ProfitERGDEAModel" begin

    # Test using Book data
    X = [2; 4; 8; 12; 6; 14; 14; 9.412];
    Y = [1; 5; 8; 9; 3; 7; 9; 2.353];
    W = [1; 1; 1; 1; 1; 1; 1; 1];
    P = [2; 2; 2; 2; 2; 2; 2; 2];

    # Profit Enhanced Russell Graph Slack Based Measure
    profiterg = deaprofiterg(X, Y, W, P)

    @test typeof(profiterg) == ProfitERGDEAModel

    @test nobs(profiterg) == 8
    @test ninputs(profiterg) == 1
    @test noutputs(profiterg) == 1
    @test ismonetary(profiterg) == false

    @test efficiency(profiterg) == efficiency(profiterg, :Economic)
    @test efficiency(profiterg, :Economic)   ≈ [4.0; 0.5; 0.0; 1/6; 0.8; 4/7; 2/7; 1.2706] atol = 1e-5
    @test efficiency(profiterg, :Technical)  ≈ [0.0; 0.0; 0.0; 0.0; 0.6; 11/21; 1/7; 0.8] atol = 1e-5
    @test efficiency(profiterg, :Allocative) ≈ [4.0; 0.5; 0.0; 1/6; 0.2; 1/21; 1/7; 0.4706] atol = 1e-5
    @test normfactor(profiterg) ≈ [2.0; 4.0; 8.0; 12.0; 10.0; 14.0; 14.0; 10.0]

    # With monetary = true
    profitergmon = deaprofiterg(X, Y, W, P, monetary = true)

    @test ismonetary(profitergmon) == true

    @test efficiency(profitergmon) ≈ efficiency(profiterg) .* normfactor(profiterg)
    @test efficiency(profitergmon, :Technical) ≈ efficiency(profiterg, :Technical) .* normfactor(profiterg)
    @test efficiency(profitergmon, :Allocative) ≈ efficiency(profiterg, :Allocative) .* normfactor(profiterg)

    # Print
    show(IOBuffer(), profiterg)

    # Test errors
    @test_throws DimensionMismatch deaprofiterg([1; 2 ; 3], [4; 5], [1; 1; 1], [1; 1]) #  Different number of observations
    @test_throws DimensionMismatch deaprofiterg([1; 2; 3], [4; 5; 6], [1; 2; 3; 4], [1; 1; 1]) # Different number of observation in input prices
    @test_throws DimensionMismatch deaprofiterg([1; 2; 3], [4; 5; 6], [1; 1; 1], [1; 2; 3; 4]) # Different number of observation in output prices
    @test_throws DimensionMismatch deaprofiterg([1 1; 2 2; 3 3 ], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3], [1; 1; 1]) # Different number of input prices and inputs
    @test_throws DimensionMismatch deaprofiterg([1; 2; 3 ], [4; 5; 6], [1; 1; 1], [1 1 1; 2 2 2; 3 3 3]) # Different number of output prices and outputs

end


using ImageMagick
using Images, TestImages, StaticArrays, ImageTracking, ImageView, LinearAlgebra, CoordinateTransformations, Gtk.ShortNames
#=Image Credit:  C. Liu. Beyond Pixels: Exploring New Representations and
#Applications for Motion Analysis. Doctoral Thesis. Massachusetts Institute of
#Technology. May 2009. =#
img1 = load("./error_evaluation_demo_data/frame1.png")
img2 = load("./error_evaluation_demo_data/frame2.png")

algorithm = Farneback(50, estimation_window = 11,
                         σ_estimation_window = 9.0,
                         expansion_window = 6,
                         σ_expansion_window = 5.0)

estimated_flow = optical_flow(Gray{Float32}.(img1), Gray{Float32}.(img2), algorithm)

# read the ground truth flow
file_path =  "./error_evaluation_demo_data/example.flo"
ground_truth_flow = read_flow_file(file_path)

# Endpoint error statistics
endpoint_error = evaluate_flow_error(ground_truth_flow, estimated_flow, EndpointError())
ep_mean, ep_sd, ep_RX, ep_AX = calculate_statistics(endpoint_error)
println("Endpoint error statistics :")
println("Mean = ", ep_mean, ", standard deviation = ", ep_sd, ", RX Dict = ", ep_RX, ", AX Dict = ", ep_AX)

# Angular error statistics
angular_error = evaluate_flow_error(ground_truth_flow, estimated_flow, AngularError())
ang_mean, ang_sd, ang_RX, ang_AX = calculate_statistics(angular_error)
println("Angular error statistics :")
println("Mean = ", ang_mean, ", standard deviation = ", ang_sd, ", RX Dict = ", ang_RX, ", AX Dict = ", ang_AX)


using Random, Statistics, StatsBase
struct FunnelDistribution{T<:Real,U<:Integer} <: Distribution{Multivariate,Continuous}
    "Variance of the dominant normal distribution."
    a::T

    "Variance of the supporting normal distributions."
    b::T

    "Number of dimensions."
    n::U
end

function FunnelDistribution(; a::Real = 1.0, b::Real = 0.5, n::Integer = 3)
    a_pr, b_pr = promote(a, b)
    FunnelDistribution(a_pr, b_pr, n)
end

Base.length(d::FunnelDistribution) = d.n
Base.eltype(d::FunnelDistribution) = Base.eltype(d.a)

Statistics.mean(d::FunnelDistribution) = zeros(d.n)

function Statistics.cov(dist::FunnelDistribution)
    cov(nestedview(rand(rng, sampler(dist), 10^5)))
end

StatsBase.params(d::FunnelDistribution) = (d.a, d.b, d.n)

function Distributions._logpdf(d::FunnelDistribution, x::AbstractArray)
    idxs = eachindex(x)
    s = logpdf(Normal(0, d.a^2), x[idxs[1]])
    @inbounds for i in idxs[2:end]
        s += logpdf(Normal(0, exp(2 * d.b * x[1])), x[i])
    end
    s
end

function Distributions.gradlogpdf(d::FunnelDistribution{T, U}, x::AbstractVector{T}) where {T<:Real, U<:Integer}
    FillArrays.Zeros(length(d))
end

function Distributions._rand!(rng::AbstractRNG, d::FunnelDistribution, x::AbstractVector)
    idxs = eachindex(x)
    x[idxs[1]] = rand(rng, Normal(0, d.a^2))
    @inbounds for i in idxs[2:end]
        x[i] = rand(rng, Normal(0, exp(2 * d.b * x[1])))
    end
    return x
end
#
#
# dist = FunnelDistribution(1.0, 0.5, 2)
#
# s=rand(dist,200)
#
# scatter(s[1,:],s[2,:])


abstract type AbstractBC{T} <: AbstractDiffEqLinearOperator{T} end

abstract type AtomicBC{T} <: AbstractBC{T} end

"""
Robin, General, and in general Neumann, Dirichlet and Bridge BCs
are not necessarily linear operators.  Instead, they are affine
operators, with a constant term Q*x = Qa*x + Qb.
"""
abstract type AffineBC{T} <: AtomicBC{T} end

struct NeumannBC{N} end
struct Neumann0BC{N} end
struct DirichletBC{N} end
struct Dirichlet0BC{N} end
struct PeriodicBC{T} <: AtomicBC{T}
    PeriodicBC(T::Type) = new{T}()
end
struct RobinBC{T, V<:AbstractVector{T}} <: AffineBC{T}
    a_l::V
    b_l::T
    a_r::V
    b_r::T
    function RobinBC(l::NTuple{3,T}, r::NTuple{3,T}, dx::T, order = 1) where {T}
        αl, βl, γl = l
        αr, βr, γr = r

        s = calculate_weights(1, one(T), Array(one(T):convert(T,order+1))) #generate derivative coefficients about the boundary of required approximation order

        a_l = -s[2:end]./(αl*dx/βl + s[1])
        a_r = s[end:-1:2]./(αr*dx/βr - s[1]) # for other boundary stencil is flippedlr with *opposite sign*

        b_l = γl/(αl+βl*s[1]/dx)
        b_r = γr/(αr-βr*s[1]/dx)

        return new{T, typeof(a_l)}(a_l, b_l, a_r, b_r)
    end
    function RobinBC(l::AbstractVector{T}, r::AbstractVector{T}, dx::AbstractVector{T}, order = 1) where {T}
        αl, βl, γl = l
        αr, βr, γr = r

        s_index = Array(one(T):convert(T,order+1))
        dx_l, dx_r = dx[1:length(s_index)], dx[(end-length(s_index)+1):end]

        s = calculate_weights(1, one(T), s_index) #generate derivative coefficients about the boundary of required approximation order
        denom_l = αl+βl*s[1]/dx_l[1]
        denom_r = αr-βr*s[1]/dx_r[end]

        a_l = -βl.*s[2:end]./(denom_l*dx_l[2:end])
        a_r = βr.*s[end:-1:2]./(denom_r*dx_r[1:(end-1)]) # for other boundary stencil is flippedlr with *opposite sign*

        b_l = γl/denom_l
        b_r = γr/denom_r

        return new{T, typeof(a_l)}(a_l, b_l, a_r, b_r)
    end
end

"""
q = GeneralBC(α_leftboundary, α_rightboundary, dx::T, approximation_order)

-------------------------------------------------------------------------------------

Implements a generalization of the Robin boundary condition, where α is a vector of coefficients.
Represents a condition of the form α[1] + α[2]u[0] + α[3]u'[0] + α[4]u''[0]+... = 0
Implemented in a similar way to the RobinBC (see above).
This time there are multiple stencils for multiple derivative orders - these can be written as a matrix S.
All components that multiply u(0) are factored out, turns out to only involve the first column of S, s̄0. The rest of S is denoted S`. the coeff of u(0) is s̄0⋅ᾱ[3:end] + α[2].
the remaining components turn out to be ᾱ[3:end]⋅(S`ū`) or equivalently (transpose(ᾱ[3:end])*S`)⋅ū`. Rearranging, a stencil q_a to be dotted with ū` upon extension can readily be found, along with a constant component q_b
"""
struct GeneralBC{T, L<:AbstractVector{T}, R<:AbstractVector{T}} <:AffineBC{T}
    a_l::L
    b_l::T
    a_r::R
    b_r::T
    function GeneralBC(αl::AbstractVector{T}, αr::AbstractVector{T}, dx::T, order = 1) where {T}
        nl = length(αl)
        nr = length(αr)
        S_l = zeros(T, (nl-2, order+nl-2))
        S_r = zeros(T, (nr-2, order+nr-2))

        for i in 1:(nl-2)
            S_l[i,:] = [transpose(calculate_weights(i, one(T), Array(one(T):convert(T, order+i)))) transpose(zeros(T, Int(nl-2-i)))]./(dx^i) #am unsure if the length of the dummy_x is correct here
        end

        for i in 1:(nr-2)
            S_r[i,:] = [transpose(calculate_weights(i, convert(T, order+i), Array(one(T):convert(T, order+i)))) transpose(zeros(T, Int(nr-2-i)))]./(dx^i)
        end
        s0_l = S_l[:,1] ; Sl = S_l[:,2:end]
        s0_r = S_r[:,end] ; Sr = S_r[:,(end-1):-1:1]

        denoml = αl[2] .+ αl[3:end] ⋅ s0_l
        denomr = αr[2] .+ αr[3:end] ⋅ s0_r

        a_l = -transpose(transpose(αl[3:end]) * Sl) ./denoml
        a_r = reverse(-transpose(transpose(αr[3:end]) * Sr) ./denomr)

        b_l = -αl[1]/denoml
        b_r = -αr[1]/denomr
        new{T, typeof(a_l), typeof(a_r)}(a_l,b_l,a_r,b_r)
    end

    function GeneralBC(αl::AbstractVector{T}, αr::AbstractVector{T}, dx::AbstractVector{T}, order = 1) where {T}

        nl = length(αl)
        nr = length(αr)
        dx_l, dx_r = (dx[1:(order+nl-2)], reverse(dx[(end-order-nr+3):end]))
        S_l = zeros(T, (nl-2, order+nl-2))
        S_r = zeros(T, (nr-2, order+nr-2))

        for i in 1:(nl-2)
            S_l[i,:] = [transpose(calculate_weights(i, one(T), Array(one(T):convert(T, order+i)))) transpose(zeros(T, Int(nl-2-i)))]./(dx_l.^i)
        end

        for i in 1:(nr-2)
            S_r[i,:] = [transpose(calculate_weights(i, convert(T, order+i), Array(one(T):convert(T, order+i)))) transpose(zeros(T, Int(nr-2-i)))]./(dx_r.^i)
        end
        s0_l = S_l[:,1] ; Sl = S_l[:,2:end]
        s0_r = S_r[:,end] ; Sr = S_r[:,(end-1):-1:1]

        denoml = αl[2] .+ αl[3:end] ⋅ s0_l
        denomr = αr[2] .+ αr[3:end] ⋅ s0_r

        a_l = -transpose(transpose(αl[3:end]) * Sl) ./denoml
        a_r = reverse(-transpose(transpose(αr[3:end]) * Sr) ./denomr)

        b_l = -αl[1]/denoml
        b_r = -αr[1]/denomr
        new{T,  typeof(a_l), typeof(a_r)}(a_l,b_l,a_r,b_r)
    end
end

#implement Neumann and Dirichlet as special cases of RobinBC
NeumannBC(α::NTuple{2,T}, dx::Union{AbstractVector{T}, T}, order = 1) where T = RobinBC((zero(T), one(T), α[1]), (zero(T), one(T), α[2]), dx, order)
DirichletBC(αl::T, αr::T) where T = RobinBC((one(T), zero(T), αl), (one(T), zero(T), αr), 1.0, 2.0 )
#specialized constructors for Neumann0 and Dirichlet0
Dirichlet0BC(T::Type) = DirichletBC(zero(T), zero(T))
Neumann0BC(dx::Union{AbstractVector{T}, T}, order = 1) where T = NeumannBC((zero(T), zero(T)), dx, order)

# other acceptable argument signatures
#RobinBC(al::T, bl::T, cl::T, dx_l::T, ar::T, br::T, cr::T, dx_r::T, order = 1) where T = RobinBC([al,bl, cl], [ar, br, cr], dx_l, order)

Base.:*(Q::AffineBC, u::AbstractVector) = BoundaryPaddedVector(Q.a_l ⋅ u[1:length(Q.a_l)] + Q.b_l, Q.a_r ⋅ u[(end-length(Q.a_r)+1):end] + Q.b_r, u)
Base.:*(Q::PeriodicBC, u::AbstractVector) = BoundaryPaddedVector(u[end], u[1], u)

Base.size(Q::AtomicBC) = (Inf, Inf) #Is this nessecary?

gettype(Q::AbstractBC{T}) where T = T


# advent of code 2020 day 3

function load_data(path:: String):: AbstractArray
    file = open(path)
    data_array = readlines(file)
    close(file)
    return data_array
end

function walker_acc(data_array:: AbstractArray, x_step=1, y_step=3):: Int
    size_x, size_y = size(data_array)[1], length(data_array[1])
    numb_tree, pos_x, pos_y = 0, 1 + x_step, 1 + y_step
    while (pos_x <= size_x)
        if data_array[pos_x][(pos_y - 1)%end + 1] == '#'
            numb_tree += 1
        end
        pos_x += x_step
        pos_y += y_step
    end

    return numb_tree
end

data_array = load_data("./data/day3.txt")
numb_tree = walker_acc(data_array)
println("Numb of tree is: $numb_tree")

slope_to_test = [[1, 1], [3, 1], [5, 1], [7, 1], [1, 2]]
let prod_path = 1
    for slope in slope_to_test
        prod_path *= walker_acc(data_array, slope[2], slope[1])
    end
    println("product of multipel slopes is: $prod_path")
end


"""
    QuarticNonconvexPeriodicSemidiscretization(D, Di, split_form)

A semidiscretization of the quartic nonconvex conservation law
    \$\\partial_t u(t,x) + \\partial_x ( u(t,x)^4 - 10 u(t,x)^2 + 3 u(t,x) ) = 0\$
with periodic boundary conditions.

`D` is a first-derivative SBP operator, `Di` an associated dissipation operator
or `nothing`, and `split_form::Union{Val(true), Val(false)}` determines whether
the canonical split form or the conservative form is used.
"""
struct QuarticNonconvexPeriodicSemidiscretization{T,Derivative<:AbstractDerivativeOperator{T},
                                                  Dissipation,
                                                  SplitForm<:Union{Val{false}, Val{true}}} <: AbstractSemidiscretization
    derivative::Derivative
    dissipation::Dissipation
    tmp1::Vector{T}
    tmp2::Vector{T}
    split_form::SplitForm

    function QuarticNonconvexPeriodicSemidiscretization(derivative::Derivative, dissipation::Dissipation, split_form::SplitForm=Val{false}()) where {T, Derivative<:AbstractDerivativeOperator{T}, Dissipation, SplitForm<:Union{Val{false}, Val{true}}}
        if dissipation !== nothing
            @argcheck size(derivative) == size(dissipation) DimensionMismatch
            @argcheck grid(derivative) == grid(dissipation) ArgumentError
        end
        N = size(derivative, 2)
        tmp1 = Array{T}(undef, N)
        tmp2 = Array{T}(undef, N)
        new{T,Derivative,Dissipation,SplitForm}(derivative, dissipation, tmp1, tmp2, split_form)
    end
end

function Base.show(io::IO, semi::QuarticNonconvexPeriodicSemidiscretization)
    if get(io, :compact, false)
        print(io, "Semidiscretization of the quartic nonconvex conservation law (periodic)")
    else
        print(io, "Semidiscretization of the quartic nonconvex conservation law\n")
        print(io, "  \$ \\partial_t u(t,x) + \\partial_x ( u(t,x)^4 - 10 u(t,x)^2 + 3 u(t,x) ) = 0 \$ \n")
        print(io, "with periodic boundaries using")
        if semi.split_form == Val{true}()
            print(io, " a split form and: \n")
        else
            print(io, " no split form and: \n")
        end
        print(io, semi.derivative)
        print(io, semi.dissipation)
    end
end

function (disc::QuarticNonconvexPeriodicSemidiscretization)(du, u, p, t)
    @unpack tmp1, tmp2, derivative, dissipation, split_form = disc
    @boundscheck begin
        @argcheck length(u) == length(tmp1)
        @argcheck length(u) == length(tmp2)
        @argcheck length(u) == length(du)
    end

    # volume terms
    if typeof(split_form) <: Val{true}
        m_2_5 = -2*one(eltype(u)) / 5
        m_20_3 = -20*one(eltype(u)) / 3
        m_3 = -3*one(eltype(u))

        # # D u^4
        @. tmp2 = u^4
        mul!(tmp1, derivative, tmp2)
        @. du = m_2_5 * tmp1

        @. tmp2 = u^3
        mul!(tmp1, derivative, tmp2)
        @. du += m_2_5 * u * tmp1

        @. tmp2 = u^2
        mul!(tmp1, derivative, tmp2)
        @. du += m_2_5 * u^2 * tmp1

        mul!(tmp1, derivative, u)
        @. du += m_2_5 * u^3 * tmp1

        # D u^2
        @. tmp2 = u^2
        mul!(tmp1, derivative, tmp2)
        @. du -= m_20_3 * tmp1

        mul!(tmp1, derivative, u)
        @. du -= m_20_3 * u * tmp1

        # D u
        mul!(tmp1, derivative, u)
        @. du += m_3 * tmp1
    else
        @. tmp2 = u^4 - 10 * u^2 + 3 * u
        mul!(tmp1, derivative, tmp2)
        @. du = -tmp1
    end

    # dissipation
    if dissipation !== nothing
        mul!(tmp1, dissipation, u)
        @. du += tmp1
    end

    nothing
end


@assert begin
    bfun=_get_gradfunc(:(1.0-xi-eta-zeta), [:xi, :eta, :zeta], 2)

    grad=bfun(zeros(Float64, 3)...)

    all(
        @. isapprox(grad, -1.0)
    )
end

@assert begin
    intp=InterpolationFunction(
        :(a0+xi*a1+eta*a2),
        [
            :xi, :eta
        ],
        [
            :a0, :a1, :a2
        ],
        [
            [0, 0]
        ]
    )

    dmn=Domain(
        intp,
        [
            (
                0.0, 1.0
            ),
            (
                0.0, :(1.0-xi)
            )
        ],
        [
            [0.0, 0.0],
            [1.0, 0.0],
            [0.0, 1.0]
        ]
    )

    vals=[1.0, 0.0, 0.0]
    coefs=regress(dmn, vals)

    isapprox(coefs, [1.0, -1.0, -1.0])
end


#=
formulates and solves the (dual of the) polynomial envelope problem described in:
D. Papp and S. Yildiz. Sum-of-squares optimization without semidefinite programming
available at https://arxiv.org/abs/1712.01792
=#

using SparseArrays

struct PolyEnvelopeNative{T <: Real} <: ExampleInstanceNative{T}
    n::Int
    rand_halfdeg::Int
    num_polys::Int
    env_halfdeg::Int
    primal_wsos::Bool # use primal formulation, else use dual
end

function build(inst::PolyEnvelopeNative{T}) where {T <: Real}
    (n, num_polys) = (inst.n, inst.num_polys)
    @assert inst.rand_halfdeg <= inst.env_halfdeg
    domain = PolyUtils.BoxDomain{T}(-ones(T, n), ones(T, n))

    # generate interpolation
    (U, pts, Ps, _, w) = PolyUtils.interpolate(domain,
        inst.env_halfdeg, get_quadr = true)

    # generate random data
    L = binomial(n + inst.rand_halfdeg, n)
    c_or_h = vec(Ps[1][:, 1:L] * rand(T(-9):T(9), L, num_polys))

    if inst.primal_wsos
        # WSOS cone in primal
        c = -w
        A = zeros(T, 0, U)
        b = T[]
        G = repeat(sparse(one(T) * I, U, U), outer = (num_polys, 1))
        h = c_or_h
    else
        # WSOS cone in dual
        c = c_or_h
        A = repeat(sparse(one(T) * I, U, U), outer = (1, num_polys))
        b = w
        G = Diagonal(-one(T) * I, num_polys * U) # TODO uniformscaling
        h = zeros(T, num_polys * U)
    end

    cones = Cones.Cone{T}[Cones.WSOSInterpNonnegative{T, T}(U, Ps,
        use_dual = !inst.primal_wsos) for k in 1:num_polys]

    model = Models.Model{T}(c, A, b, G, h, cones)
    return model
end


export getMassMatrix

"""
	M = getMassMatrix(m::Vector,gamma::Vector,Mesh::RegularMesh)

	Input:
		m     - model
		gamma - attenuation
		Mesh  - regular mesh

	Output
		M::SparseMatrixCSC
"""
function getMassMatrix(m::Vector,gamma::Vector,Mesh::RegularMesh)

    An2cc = getNodalAverageMatrix(Mesh)
    M     = sdiag(An2cc'*(m.*(1 .- 1im*gamma)))

    return M
end


module HartreeFockModule
export computeEnergyHartreeFock, evaluateSCF, evaluateHartreeFock
using ..BaseModule
using ..BasisModule
using ..BasisSetModule
using ..GeometryModule
using ..SpecialMatricesModule
using ..IntegralsModule
using ..ShellModule
using ..LibInt2Module
using ..InitialGuessModule

function computeEnergyHartreeFock(
  density::Matrix,
  kinetic::Matrix,
  nuclearAttraction::Matrix,
  coulomb::Matrix,
  exchange::Matrix,
  interatomicRepulsion::Float64;
  info::Bool=true)

  P = density
  info && println("=================================")
  eT = 2*trace(kinetic*P)
  info && println("Kinetic Energy: $eT")
  eV = 2*trace(nuclearAttraction*P)
  info && println("Nuclear Attraction Energy: $eV")
  info && println("Sum 1-Electron Energy: $(eT+eV)")
  info && println("---------------------------------")
  # G = 2J-K
  eJ = 2*trace(coulomb*P)
  info && println("Coulomb Energy: $eJ")
  eK = -trace(exchange*P)
  info && println("Exchange Energy: $eK")
  info && println("Total Electronic Energy: $(eT+eV+eJ+eK)")
  info && println("Total Energy: $(eT+eV+eJ+eK+interatomicRepulsion)")
  info && println("=================================")

  return eT+eV+eJ+eK
end

function computeEnergyHartreeFock(
  basis::GaussianBasis,
  geometry::Geometry,
  density::Matrix;
  info::Bool=true)

  P = density
  T = computeMatrixKinetic(basis)
  V = computeMatrixNuclearAttraction(basis,geometry)
  J = computeMatrixCoulomb(basis,density)
  K = computeMatrixExchange(basis,density)
  Vnn = computeEnergyInteratomicRepulsion(geometry)
  computeEnergyHartreeFock(P,T,V,J,K,Vnn;info=info)
end

function computeEnergyHartreeFock(
  density::Matrix,
  matrixKinetic::Matrix,
  matrixNuclearAttraction::Matrix,
  electronRepulsionIntegrals::Array{Float64,4},
  energyInteratomicRepulsion::Float64;
  info::Bool=true,
  # inceptions
  electronRepulsionIntegralsExchange::Array{Float64,4}=electronRepulsionIntegrals)

  J = computeMatrixCoulomb(electronRepulsionIntegrals,density)
  K = computeMatrixExchange(electronRepulsionIntegralsExchange,density)
  computeEnergyHartreeFock(density, matrixKinetic, matrixNuclearAttraction, J, K, energyInteratomicRepulsion;info=info)
end

function evaluateSCFStep(
  fock::Matrix,
  overlap::Matrix,
  electronNumber::Integer)

  eig = eigfact(Symmetric(fock),Symmetric(overlap)) # because Symmetric is faster and allows for sorted eigenvalues
  moEnergies = eig.values
  Cocc = eig.vectors[:,1:electronNumber]
  P = Cocc * Cocc.'

  return (moEnergies,P)
end

computeMatrixCoefficients(fock::Matrix, overlap::Matrix) = eigvecs(Symmetric(fock),Symmetric(overlap))

"""
    evaluateSCF(initialGuessDensity::Matrix, overlap::Matrix, kinetic::Matrix, nuclearAttraction::Matrix,
                coulomb::Function, exchange::Function,
                interatomicRepulsion::Float64, electronNumber::Integer)

All evaluateSCF calls take the following optional named arguments:
 - `info::Bool=true`,
 - `detailedinfo::Bool=info`,
 - `energyConvergenceCriterion::Float64 = 1e-8`

This is the fundamental routine to evaluate the Self-consistent field cycle of
Hartree-Fock theory and returns the converged (orbital energies, total energy,
density). This function call requires each element in its most primitive form
and is intended for use in unusual contexts or for performance reasons. For
more usable versions see below.
"""
function evaluateSCF(
  initialGuessDensity::Matrix,
  overlap::Matrix,
  kinetic::Matrix,
  nuclearAttraction::Matrix,
  coulomb::Function,
  exchange::Function,
  interatomicRepulsion::Float64,
  electronNumber::Integer;
  info::Bool=true,
  detailedinfo::Bool=info,
  energyConvergenceCriterion::Float64 = 1e-8)

  P::Matrix = initialGuessDensity
  totalEnergy = computeEnergyHartreeFock(P,kinetic,nuclearAttraction,coulomb(P),exchange(P),interatomicRepulsion;info=false)
  energies = Vector{Float64}()

  energyConverged = false
  i = 0
  J = coulomb(P)
  K = exchange(P)
  while (!energyConverged)
    i+=1
    F = kinetic+nuclearAttraction+2J-K
    energies,P = evaluateSCFStep(F,overlap,electronNumber)
    J = coulomb(P) # By requiring J,K to contain the new density, the variational principle guarantees no energy below the exact energy is obtained
    K = exchange(P) # think: E2 = P II P = tr(P*G[P]) != tr(P*G[P'])
    oldEnergy = totalEnergy
    totalEnergy = computeEnergyHartreeFock(P,kinetic,nuclearAttraction,J,K,interatomicRepulsion,info=detailedinfo)
    info && println("E[$i]: $totalEnergy")
    if (abs(totalEnergy-oldEnergy) < energyConvergenceCriterion)
      info && println("Converged!")
      energyConverged = true
    end
  end
  return energies, totalEnergy, P
end

"""
    evaluateSCF(basis, geometry, initialGuess, electronNumber)

allows to evaluate the Self-consistent field cycle of Hartree-Fock theory in the most
accessible way. Essentially, this function call combines its arguments in such a way
that it can then call upon the more fundamental `evaluateSCF` that takes only primitive
inputs. To this end, the arguments are very flexible:
- `basis`: One of `Basis`, `BasisSet`, Vector{Shell}, Vector{LibInt2Shell}
- `geo`: One of `Geometry`, `String` (the name of a file containing the geometry)
- `initialGuess`: Either a `Matrix` that is used as initial density matrix, or an `InitialGuess` (ZeroGuess, SADGuess)
- `electronNumber`: number of occupied closed shell orbitals (as the other forms are not yet implemented) or `Ionization` of the system

This evaluateSCF method additionally introduces the following optional named arguments:
- `computeTensorElectronRepulsionIntegralsCoulomb`: f: basis/shells/... -> ERIs (used for Coulomb part (and Exchange part if not overridden))
- `computeTensorElectronRepulsionIntegralsExchange`: f: basis/shells/... -> ERIs (used for Exchange part)
"""
function evaluateSCF(
  basis::Union{GaussianBasis,Vector{Shell},Vector{LibInt2Shell},BasisSet},
  geometry::Union{Geometry,String},
  initialGuessDensity::Union{Matrix,InitialGuess},
  electronNumber::Union{Integer,Ionization}=Ionization(0);
  detailedinfo::Bool=true,
  info::Bool=true,
  energyConvergenceCriterion::Float64 = 1e-8,
  computeTensorElectronRepulsionIntegralsCoulomb = computeTensorElectronRepulsionIntegrals,
  computeTensorElectronRepulsionIntegralsExchange = computeTensorElectronRepulsionIntegrals)

  if isa(geometry,String)
    geometry = Geometry(geometry)
  end

  input = basis
  if isa(input,Vector{LibInt2Shell}) && (LibInt2Shell != Shell)
    # keep working with LibInt2Shell
  elseif isa(input,BasisSet) && (LibInt2Shell != Shell)
    basis = computeBasisShellsLibInt2(input,geometry)
  elseif isa(input,BasisSet) && (LibInt2Shell == Shell)
    basis = computeBasis(input,geometry)
  elseif isa(input,GaussianBasis)
    # At the moment we simply use implicit Shell->Basis conversion which avoids ERI recomputation in every step and is therefore faster for the moment
    # later we will want to switch over to working with Shells directly as shell-structure allows for sparsity
    # evaluateSCF(Pinit, S, T, V, P->computeMatrixCoulomb(shells,P), P->computeMatrixExchange(shells,P), Vnn, electronNumber; energyConvergenceCriterion=energyConvergenceCriterion, info=info, detailedinfo=detailedinfo)
    basis = convert(GaussianBasis,input)
  end

  if isa(electronNumber,Ionization)
    electronNumber::Int64 = computeNumberElectrons(geometry,electronNumber.charge)/2
  end

  S = computeMatrixOverlap(basis)
  T = computeMatrixKinetic(basis)
  V = computeMatrixNuclearAttraction(basis,geometry)
  Pinit = InitialGuessModule.computeDensityMatrixGuess(initialGuessDensity,size(S)[1])
  Vnn = computeEnergyInteratomicRepulsion(geometry)

  if (LibInt2Shell != Shell) && isa(basis,Vector{LibInt2Shell})
    Xarg = basis
    Carg = basis
  else
    ERIsCoulomb = computeTensorElectronRepulsionIntegralsCoulomb(basis)
    ERIsExchange = computeTensorElectronRepulsionIntegralsCoulomb==computeTensorElectronRepulsionIntegralsExchange ? ERIsCoulomb : computeTensorElectronRepulsionIntegralsExchange(basis)
    Xarg = ERIsExchange
    Carg = ERIsCoulomb
  end

  evaluateSCF(Pinit, S, T, V, P->computeMatrixCoulomb(Carg,P), P->computeMatrixExchange(Xarg,P), Vnn, electronNumber; energyConvergenceCriterion=energyConvergenceCriterion, info=info, detailedinfo=detailedinfo)
end

evaluateHartreeFock = evaluateSCF

end # module


function lag!{T}(Yl::Array{T, 2}, Y::Array{T, 1}, p::Int64)
    ## Given an Array, return the matrix
  ## of lagged values of Y
  ## [y_t-p, y_{t-p-1}, ..., y_{t-1}]
  N  = length(Y)
  Ts, pl = size(Yl)
  for i = 1:p
    for j = 1:N-p
      @inbounds Yl[j, i] = Y[p-i+j]
    end
  end
  if pl==p+1
    for j=1:N-p
      @inbounds Yl[j,p+1] = 1.0
    end
  end
end

function olsvar{T}(y::Array{T, 2})
    # Input : data y, lag order 1
    # Output: coefficient estimate β, residual U
    N, K = size(y)
    Y = y[2:N,:]  # LHS variable
    X = y[1:N-1,:]
    A = X\Y    # OLS estimator
    U = Y-X*A  # estimated residuals
    return U, A'
end

function ar{T}(Y::Array{T, 2}, lag::Int64)
    N, p = size(Y)
    Yl = Array(T, N-lag, lag+1)
	   ρ  = Array(T, p, lag)
	   σ² = Array(T, p)
	   for j = 1:p
		      lag!(Yl, Y[:,j], lag)
		      Yt          = Y[lag+1:end,j]
		      θ           = (Yl\Yt)
		      ρ[j, 1:lag] = θ[1:lag]
		      ε           = Yt - Yl*θ
		      σ²[j]       = (ε'ε/N)[1]
	   end
	   return ρ, σ²
end

function arma{T}(Y::Array{T,2})
    ## Estimate an ARMA(1,1) for each column of Y
end

ar{T}(Y::Array{T, 2}) = ar(Y, 1)

d_bw_andrews = @compat Dict(:TruncatedKernel         => :(0.6611*(a2*N)^(0.2)),
                            :BartlettKernel          => :(1.1447*(a1*N)^(1/3)),
                            :ParzenKernel            => :(2.6614*(a2*N)^(0.2)),
                            :QuadraticSpectralKernel => :(1.3221*(a2*N)^(0.2)))

for tty in [:TruncatedKernel, :BartlettKernel, :ParzenKernel, :QuadraticSpectralKernel]
    @eval  $:(bw_andrews)(k::($tty), a1, a2, N) = $(d_bw_andrews[tty])
end

function pre_white(X::AbstractMatrix)
    X, D = olsvar(X)
    (X, inv(eye(size(X, 2))-D))
end

function getalpha(X::AbstractMatrix, approx::Symbol)
    ## @assert approx == :ar ## || approx == :arma
    ## if approx == :ar
    ρ, σ² = ar(X)
    σ⁴    = (σ²).^2
    nm    = 4.*ρ.^2.*σ⁴./((1-ρ).^6.*(1+ρ).^2)
    dn    = σ⁴./(1-ρ).^4
    α₁    = sum(nm)/sum(dn)
    nm    = 4.*ρ.^2.*σ⁴./(1-ρ).^8
    α₂    = sum(nm)/sum(dn)
## elseif approx == :arma  [TODO]
## end
    return α₁, α₂
end

function optimalbw_ar_one(X::AbstractMatrix, k::TruncatedKernel)
    a1, a2 = getalpha(X, :ar)
    T, p   = size(X)
    return .6611*(a2*T)^(1/5)
end

function optimalbw_ar_one(X::AbstractMatrix, k::BartlettKernel)
    a1, a2 = getalpha(X, :ar)
    T, p   = size(X)
    return 1.1447*(a1*T)^(1/3)
end

function optimalbw_ar_one(X::AbstractMatrix, k::ParzenKernel)
    a1, a2 = getalpha(X, :ar)
    T, p   = size(X)
    return 2.6614*(a2*T)^(1/5)
end

function optimalbw_ar_one(X::AbstractMatrix, k::QuadraticSpectralKernel)
    a1, a2 = getalpha(X, :ar)
    T, p   = size(X)
    return 1.3221*(a2*T)^(1/5)
end

function bwAndrews{T}(X::Array{T, 2}, k::HAC; prewhite::Bool = false, approx::Symbol = :ar)
    N, p  = size(X)
    !prewhite || ((X, D) = pre_white(X))
    a1, a2 = getalpha(X, approx)
    return bw_andrews(k, a1, a2, N)
end


@testset "Regressive product" begin
    @test dual(dual(A) ∧ dual(B)) == A ∨ B
end