module Bodies
    import LinearAlgebra
    using NBodyLeapFrog: V3

    const mass_solar = 1.988544e30 # in kg
    const mass_mercury = 3.302e23
    const mass_venus = 48.685e23
    const mass_earth = 5.97219e24
    const mass_mars = 6.4185e23
    const mass_jupyter = 1898.13e24
    const mass_saturn = 5.68319e26
    const mass_neptune = 102.41e24
    const mass_uranus = 86.8103e24
    const mass_pluto= 1.307e22

    const mass_planetX = 3.2e24 # Planet refers to an additional planet I can add to the solar
    # system to test if it is maximally packed. (The solar system is not
    # maximally packed. There is room to add another planet and preserve the
    # stable orbit of all planets.)

    #  unit of time -> year, unit of space -> AU
    const year = 3.154e7 #year in seconds
    const AU = 1.496e11 #in m

    const G_standard = 6.67259e-11 # in m^3/(kg-s^2)
    const G_year_AU = G_standard*(1/AU)^3/(1/mass_solar*(1/year)^2)

    struct Body{T}
        r::V3{T}
        v::V3{T}
        m::T
    end

    const mercury = Body(
        V3(8.887985138765460E-02, -4.426150338141062E-01, -4.475716356484761E-02),
        # velocity in AU/day
        V3(2.190877912081542E-02, 7.161568136528000E-03, -1.425929443086507E-03),
        mass_mercury / mass_solar
    )

    const venus = Body(
        V3(4.043738093622098E-02, -7.239789211502183E-01, -1.241560658530024E-02),
        V3(2.005742309538389E-02, 1.141448268256643E-03, -1.142174441569258E-03),
        mass_venus/mass_solar
    )

    const earth =  Body(
        V3(-2.020844529756663E-02, -1.014332737790859E+00, -1.358267619371298E-05),
        V3(1.692836723212859E-02, -3.484006532982474E-04, 6.028542314557626E-07),
        mass_earth/mass_solar
    )

    # target 1998
    #  -3.466334931755365E-02 -1.013773181327570E+00  2.111689861662178E-04
    #   1.692129864984556E-02 -5.252811129268817E-04  3.987686870581435E-07

    const mars = Body(
        V3(7.462481663749645E-01, -1.181663652521456E+00, -4.321921404013512E-02),
        V3(1.235610918162121E-02, 8.680869489377649E-03 , -1.220500608452554E-04),
        mass_mars/mass_solar
    )

    const planetX = Body(
        V3(0.0, 2.06, 0.0),
        V3(1.235610918162121E-02, 0.0, 0.0),
        mass_planetX / mass_solar
    )

    const jupyter = Body(
        V3(3.384805319103406E+00, 3.658805636759595E+00, -9.100441946210819E-02),
        V3(-5.634671617093230E-03, 5.479180979634376E-03, 1.034981407898108E-04),
        mass_jupyter/mass_solar
    )

    const saturn = Body(
        V3(-1.083899692644216E-01, -1.003995196286016E+01, 1.793391553155583E-01),
        V3(5.278410787728323E-03, -7.712342079566598E-05, -2.084447335785041E-04),
        mass_saturn / mass_solar
    )

    const neptune = Body(
        V3(4.675566709791660E+00, -2.985428200863175E+01, 5.070034142531887E-01),
        V3(3.080716380724798E-03, 5.030733458293977E-04, -8.101711269674541E-05),
        mass_neptune / mass_solar
    )

    const uranus = Body(
        V3(-2.693448460292631E-01, -1.927606446869220E+01, -6.808868692550485E-02),
        V3(3.903100242621723E-03, -2.380111092360100E-04, -5.164025224695875E-05),
        mass_uranus / mass_solar
    )

    const pluto = Body(
        V3(-2.129074273328636E+01, -1.896633337434039E+01, 8.187955378677129E+00),
        V3(2.276295756013608E-03, -2.670481848836963E-03, -3.669545371032554E-04),
        mass_pluto / mass_solar
    )

    # target  -1.156541154581570E+01 -2.704864218000164E+01  6.239749761161465E+00
    #   2.964408290188142E-03 -1.722224413824548E-03 -6.839434010481107E-04

    const sun = Body(
        V3(-3.430031536367300E-03, 1.761881027012596E-03, 1.246691303879918E-05),
        V3(3.433119412673547E-06, -5.231300927361546E-06, -2.972974735550750E-08),
        1.0
    )

    # Arrays for modeling the solar system
    const set = [sun, mercury, venus, earth, mars, jupyter, saturn, uranus, neptune, pluto]
    const set_packed = [sun, mercury, venus, earth, mars, planetX, jupyter, saturn, uranus, neptune, pluto]

    const dayToYear = 365.25

    function chunit_day2year(body::Body)
        Body(body.r, body.v * dayToYear, body.m)
    end
end


using Combinatorics

function part1()
	lines = readlines("inputD8.txt")
	easies = Int[2,4,3,7]
	counter = 0

	for line in lines
		afterpipe = split(line, "|")[2]
		words = split(afterpipe)
		counter += sum([length(word) in easies for word in words])
	end

	return counter
end

function part2()
	lines = readlines("inputD8.txt")
	letters_to_digit = Dict{String, Int}("abcefg" => 0, "abcdefg" => 8, "abdefg" => 6, "abcdfg" => 9, "cf" => 1, "abdfg" => 5, "bcdf" => 4, "acdeg" => 2, "acf" => 7, "acdfg" => 3)

	ordered = ["a", "b", "c", "d", "e", "f", "g"]
	perms = collect(permutations(ordered))
	total = 0
	for line in lines
		splitbypipe = split(line, "|")
		beforepipe = split.(split(splitbypipe[1]), "")
		for perm in perms
			mapping = Dict(ordered .=> perm)

			replaced = [[get(mapping, j, j) for j in b] for b in beforepipe]
			sorted = join.(sort.(replaced))

			if all(r in keys(letters_to_digit) for r in sorted)
				# wow, we did it
				afterpipe = split.(split(splitbypipe[2]), "")

				afterpipereplaced = [[get(mapping, j, j) for j in a] for a in afterpipe]
				afterpipesorted = join.(sort.(afterpipereplaced))

				final_number = parse(Int, join([letters_to_digit[i] for i in afterpipesorted]))
				total += final_number
				break
			end
		end
	end

	return total
end

println(part1())
println(part2())


@doc raw"""
    GeneralizedStiefel{n,k,𝔽,B} <: AbstractDecoratorManifold{𝔽}

The Generalized Stiefel manifold consists of all $n\times k$, $n\geq k$ orthonormal
matrices w.r.t. an arbitrary scalar product with symmetric positive definite matrix
$B\in R^{n × n}$, i.e.

````math
\operatorname{St}(n,k,B) = \bigl\{ p \in \mathbb F^{n × k}\ \big|\ p^{\mathrm{H}} B p = I_k \bigr\},
````

where $𝔽 ∈ \{ℝ, ℂ\}$,
$\cdot^{\mathrm{H}}$ denotes the complex conjugate transpose or Hermitian, and
$I_k \in \mathbb R^{k × k}$ denotes the $k × k$ identity matrix.

In the case $B=I_k$ one gets the usual [`Stiefel`](@ref) manifold.

The tangent space at a point $p\in\mathcal M=\operatorname{St}(n,k,B)$ is given by

````math
T_p\mathcal M = \{ X \in 𝔽^{n × k} : p^{\mathrm{H}}BX + X^{\mathrm{H}}Bp=0_n\},
````
where $0_k$ is the $k × k$ zero matrix.

This manifold is modeled as an embedded manifold to the [`Euclidean`](@ref), i.e.
several functions like the [`zero_vector`](@ref) are inherited from the embedding.

The manifold is named after
[Eduard L. Stiefel](https://en.wikipedia.org/wiki/Eduard_Stiefel) (1909–1978).

# Constructor
    GeneralizedStiefel(n, k, B=I_n, F=ℝ)

Generate the (real-valued) Generalized Stiefel manifold of $n\times k$ dimensional
orthonormal matrices with scalar product `B`.
"""
struct GeneralizedStiefel{n,k,𝔽,TB<:AbstractMatrix} <: AbstractDecoratorManifold{𝔽}
    B::TB
end

function GeneralizedStiefel(
    n::Int,
    k::Int,
    B::AbstractMatrix=Matrix{Float64}(I, n, n),
    𝔽::AbstractNumbers=ℝ,
)
    return GeneralizedStiefel{n,k,𝔽,typeof(B)}(B)
end

active_traits(f, ::GeneralizedStiefel, args...) = merge_traits(IsEmbeddedManifold())

@doc raw"""
    check_point(M::GeneralizedStiefel, p; kwargs...)

Check whether `p` is a valid point on the [`GeneralizedStiefel`](@ref) `M`=$\operatorname{St}(n,k,B)$,
i.e. that it has the right [`AbstractNumbers`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/types.html#number-system) type and $x^{\mathrm{H}}Bx$
is (approximately) the identity, where $\cdot^{\mathrm{H}}$ is the complex conjugate
transpose. The settings for approximately can be set with `kwargs...`.
"""
function check_point(M::GeneralizedStiefel{n,k,𝔽}, p; kwargs...) where {n,k,𝔽}
    c = p' * M.B * p
    if !isapprox(c, one(c); kwargs...)
        return DomainError(
            norm(c - one(c)),
            "The point $(p) does not lie on $(M), because x'Bx is not the unit matrix.",
        )
    end
    return nothing
end

# overwrite passing to embedding
function check_size(M::GeneralizedStiefel{n,k,𝔽}, p) where {n,k,𝔽}
    return check_size(get_embedding(M), p) #avoid embed, since it uses copyto!
end
function check_size(M::GeneralizedStiefel{n,k,𝔽}, p, X) where {n,k,𝔽}
    return check_size(get_embedding(M), p, X) #avoid embed, since it uses copyto!
end

@doc raw"""
    check_vector(M::GeneralizedStiefel, p, X; kwargs...)

Check whether `X` is a valid tangent vector at `p` on the [`GeneralizedStiefel`](@ref)
`M`=$\operatorname{St}(n,k,B)$, i.e. the [`AbstractNumbers`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/types.html#number-system) fits,
`p` is a valid point on `M` and
it (approximately) holds that $p^{\mathrm{H}}BX + \overline{X^{\mathrm{H}}Bp} = 0$, where
`kwargs...` is passed to the `isapprox`.
"""
function check_vector(M::GeneralizedStiefel{n,k,𝔽}, p, X; kwargs...) where {n,k,B,𝔽}
    if !isapprox(p' * M.B * X, -conj(X' * M.B * p); kwargs...)
        return DomainError(
            norm(p' * M.B * X + conj(X' * M.B * p)),
            "The matrix $(X) does not lie in the tangent space of $(p) on $(M), since x'Bv + v'Bx is not the zero matrix.",
        )
    end
    return nothing
end

get_embedding(::GeneralizedStiefel{N,K,𝔽}) where {N,K,𝔽} = Euclidean(N, K; field=𝔽)

@doc raw"""
    inner(M::GeneralizedStiefel, p, X, Y)

Compute the inner product for two tangent vectors `X`, `Y` from the tangent space of `p`
on the [`GeneralizedStiefel`](@ref) manifold `M`. The formula reads

````math
(X, Y)_p = \operatorname{trace}(v^{\mathrm{H}}Bw),
````
i.e. the metric induced by the scalar product `B` from the embedding, restricted to the
tangent space.
"""
inner(M::GeneralizedStiefel, p, X, Y) = dot(X, M.B * Y)

@doc raw"""
    manifold_dimension(M::GeneralizedStiefel)

Return the dimension of the [`GeneralizedStiefel`](@ref) manifold `M`=$\operatorname{St}(n,k,B,𝔽)$.
The dimension is given by

````math
\begin{aligned}
\dim \mathrm{St}(n, k, B, ℝ) &= nk - \frac{1}{2}k(k+1) \\
\dim \mathrm{St}(n, k, B, ℂ) &= 2nk - k^2\\
\dim \mathrm{St}(n, k, B, ℍ) &= 4nk - k(2k-1)
\end{aligned}
````
"""
function manifold_dimension(::GeneralizedStiefel{n,k,ℝ}) where {n,k}
    return n * k - div(k * (k + 1), 2)
end
manifold_dimension(::GeneralizedStiefel{n,k,ℂ}) where {n,k} = 2 * n * k - k * k
manifold_dimension(::GeneralizedStiefel{n,k,ℍ}) where {n,k} = 4 * n * k - k * (2k - 1)

@doc raw"""
    project(M::GeneralizedStiefel,p)

Project `p` from the embedding onto the [`GeneralizedStiefel`](@ref) `M`, i.e. compute `q`
as the polar decomposition of $p$ such that $q^{\mathrm{H}}Bq$ is the identity,
where $\cdot^{\mathrm{H}}$ denotes the hermitian, i.e. complex conjugate transposed.
"""
project(::GeneralizedStiefel, ::Any)

function project!(M::GeneralizedStiefel, q, p)
    s = svd(p)
    e = eigen(s.U' * M.B * s.U)
    qsinv = e.vectors * Diagonal(1 ./ sqrt.(e.values))
    q .= s.U * qsinv * e.vectors' * s.V'
    return q
end

@doc raw"""
    project(M:GeneralizedStiefel, p, X)

Project `X` onto the tangent space of `p` to the [`GeneralizedStiefel`](@ref) manifold `M`.
The formula reads

````math
\operatorname{proj}_{\operatorname{St}(n,k)}(p,X) = X - p\operatorname{Sym}(p^{\mathrm{H}}BX),
````

where $\operatorname{Sym}(y)$ is the symmetrization of $y$, e.g. by
$\operatorname{Sym}(y) = \frac{y^{\mathrm{H}}+y}{2}$.
"""
project(::GeneralizedStiefel, ::Any, ::Any)

function project!(M::GeneralizedStiefel, Y, p, X)
    A = p' * M.B' * X
    copyto!(Y, X - p * Hermitian((A + A') / 2))
    return Y
end
@doc raw"""
    retract(M::GeneralizedStiefel, p, X)
    retract(M::GeneralizedStiefel, p, X, ::PolarRetraction)
    retract(M::GeneralizedStiefel, p, X, ::ProjectionRetraction)

Compute the SVD-based retraction [`PolarRetraction`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/retractions.html#ManifoldsBase.PolarRetraction) on the
[`GeneralizedStiefel`](@ref) manifold `M`, which in this case is the same as
the projection based retraction employing the exponential map in the embedding
and projecting the result back to the manifold.

The default retraction for this manifold is the [`ProjectionRetraction`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/retractions.html#ManifoldsBase.ProjectionRetraction).
"""
retract(::GeneralizedStiefel, ::Any...)

default_retraction_method(::GeneralizedStiefel) = ProjectionRetraction()

function retract_polar!(M::GeneralizedStiefel, q, p, X)
    project!(M, q, p + X)
    return q
end
function retract_project!(M::GeneralizedStiefel, q, p, X)
    project!(M, q, p + X)
    return q
end

function Base.show(io::IO, M::GeneralizedStiefel{n,k,𝔽}) where {n,k,𝔽}
    return print(io, "GeneralizedStiefel($(n), $(k), $(M.B), $(𝔽))")
end


# Computes Stein discrepancy between a weighted sample and a target with gradlogdensity
#
# Args:
#   points - n x d array of sample points
#   weights - n x 1 array of real-valued weights associated with sample points
#     (default: equal weights)
#   gradlogdensity - a function taking point x in Rd and outputting grad log p(x)
#   operator - string in {"langevin", "riemannian-langevin"} indicating the Stein operator
#     to use (default: "langevin")
#   method - One of {"graph", "classical", "kernel"}, this uses different
#     methods in order to construct the different discrepancies.

function stein_discrepancy(; points=[],
                           gradlogdensity=nothing,
                           method="graph",
                           operator="langevin",
                           kwargs...)
    # Check arguments
    isempty(points) && error("Must provide non-empty array of sample points")
    isa(gradlogdensity, Function) ||
        error("Must specify gradlogdensity of type Function")

    if method in ["graph", "classical"]
        # call appropriate graph discrepancy
        if method == "graph"
            if operator == "langevin"
                langevin_graph_discrepancy(;points=points, gradlogdensity=gradlogdensity, kwargs...)
            elseif operator == "riemannian-langevin"
                riemannian_langevin_graph_discrepancy(;points=points, gradlogdensity=gradlogdensity, kwargs...)
            else
                error("unrecognized operator: $(operator)")
            end
        elseif method == "classical"
            if operator == "langevin"
                langevin_classical_discrepancy(;points=points, gradlogdensity=gradlogdensity, kwargs...)
            else
                error("unrecognized operator: $(operator)")
            end
        end
    elseif method == "kernel"
        if operator == "langevin"
            langevin_kernel_discrepancy(;points=points, gradlogdensity=gradlogdensity, kwargs...)
        else
            error("unrecognized operator: $(operator)")
        end
    else
        error("unrecognized method: $(method)")
    end
end

# stein_discrepancy aliases
function ksd(;kwargs...)
    stein_discrepancy(; method="kernel", kwargs...)
end

function gsd(;kwargs...)
    stein_discrepancy(; method="graph", kwargs...)
end


@testset "no-alloc static" begin
    function theta(x)
        if x[1] > 0
            return atan(x[2] / x[1]) / (2.0 * pi)
        else
            return (pi + atan(x[2] / x[1])) / (2.0 * pi)
        end
    end

    function fletcher_powell_fg_static(∇f, x)
        T = eltype(x)
        theta_x = theta(x)

        if !(∇f==nothing)
            if x[1]^2 + x[2]^2 == 0
                dtdx1 = T(0)
                dtdx2 = T(0)
            else
                dtdx1 = - x[2] / ( T(2) * pi * ( x[1]^2 + x[2]^2 ) )
                dtdx2 =   x[1] / ( T(2) * pi * ( x[1]^2 + x[2]^2 ) )
            end
            ∇f1 = -2000.0*(x[3]-10.0*theta_x)*dtdx1 +
                200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 )
            ∇f2 = -2000.0*(x[3]-10.0*theta_x)*dtdx2 +
                200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 )
            ∇f3 =  200.0*(x[3]-10.0*theta_x) + 2.0*x[3];
            ∇f = @SVector[∇f1, ∇f2, ∇f3]
        end

        fx = 100.0 * ((x[3] - 10.0 * theta_x)^2 + (sqrt(x[1]^2 + x[2]^2) - 1.0)^2) + x[3]^2

        if ∇f == nothing
            return fx
        else
            return fx, ∇f
        end
    end
    sv3 = @SVector[0.0,0.0,0.0]

    fg_static = OnceDiffed(fletcher_powell_fg_static; infer=true)

    state0 = (@SVector[-0.5, 0.0, 0.0], I+sv3*sv3')
    @allocated minimize(fg_static, state0, BFGS(Inverse()), OptOptions())
    _alloc = @allocated minimize(fg_static, state0, BFGS(Inverse()), OptOptions())
    @test _alloc == 0
    @allocated minimize(fg_static, state0, BFGS(Inverse()))
    _alloc = @allocated minimize(fg_static, state0, BFGS(Inverse()))
    @test _alloc == 0
    @allocated minimize(fg_static, state0, BFGS(Inverse()))
    _alloc = @allocated minimize(fg_static, state0, BFGS(Inverse()))
    @test _alloc == 0

    minimize(fg_static, state0, BFGS(Direct()), OptOptions())
    _alloc = @allocated minimize(fg_static, state0, BFGS(Direct()), OptOptions())
    @test _alloc == 0
    minimize(fg_static, state0, BFGS(Direct()))
    _alloc = @allocated minimize(fg_static, state0, BFGS(Direct()))
    @test _alloc == 0
    minimize(fg_static, state0, BFGS(Direct()))
    _alloc = @allocated minimize(fg_static, state0, BFGS(Direct()))
    @test _alloc == 0

    minimize(fg_static, state0, DFP(Inverse()), OptOptions())
    _alloc = @allocated minimize(fg_static, state0, DFP(Inverse()), OptOptions())
    @test _alloc == 0
    minimize(fg_static, state0, DFP(Inverse()))
    _alloc = @allocated minimize(fg_static, state0, DFP(Inverse()))
    @test _alloc == 0
    minimize(fg_static, state0, DFP(Inverse()))
    _alloc = @allocated minimize(fg_static, state0, DFP(Inverse()))
    @test _alloc == 0

    minimize(fg_static, state0, DFP(Direct()), OptOptions())
    _alloc = @allocated minimize(fg_static, state0, DFP(Direct()), OptOptions())
    @test _alloc == 0
    minimize(fg_static, state0, DFP(Direct()))
    _alloc = @allocated minimize(fg_static, state0, DFP(Direct()))
    @test _alloc == 0
    minimize(fg_static, state0, DFP(Direct()))
    _alloc = @allocated minimize(fg_static, state0, DFP(Direct()))
    @test _alloc == 0

    minimize(fg_static, state0, SR1(Inverse()), OptOptions())
    _alloc = @allocated minimize(fg_static, state0, SR1(Inverse()), OptOptions())
    @test _alloc == 0
    minimize(fg_static, state0, SR1(Inverse()))
    _alloc = @allocated minimize(fg_static, state0, SR1(Inverse()))
    @test _alloc == 0
    minimize(fg_static, state0, SR1(Inverse()))
    _alloc = @allocated minimize(fg_static, state0, SR1(Inverse()))
    @test _alloc == 0

    minimize(fg_static, state0, SR1(Direct()), OptOptions())
    _alloc = @allocated minimize(fg_static, state0, SR1(Direct()), OptOptions())
    @test _alloc == 0
    minimize(fg_static, state0, SR1(Direct()))
    _alloc = @allocated minimize(fg_static, state0, SR1(Direct()))
    @test _alloc == 0
    minimize(fg_static, state0, SR1(Direct()))
    _alloc = @allocated minimize(fg_static, state0, SR1(Direct()))
    @test _alloc == 0

end


# ------------------------------------------------------------------------------------------
# # Intro to Flux.jl
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We have learned how machine learning allows us to classify data as apples or bananas with
# a single neuron. However, some of those details are pretty fiddly! Fortunately, Julia has
# a powerful package that does much of the heavy lifting for us, called
# [`Flux.jl`](https://fluxml.github.io/).
#
# *Using `Flux` will make classifying data and images much easier!*
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Using `Flux.jl`
#
# In the next notebook, we are going to see how Flux allows us to redo the calculations from
# the previous notebook in a simpler way. We can get started with `Flux.jl` via:
# ------------------------------------------------------------------------------------------

using Flux

# ------------------------------------------------------------------------------------------
# #### Helpful built-in functions
#
# When working we'll `Flux`, we'll make use of built-in functionality that we've had to
# create for ourselves in previous notebooks.
#
# For example, the sigmoid function, σ, that we have been using already lives within `Flux`:
# ------------------------------------------------------------------------------------------

?σ

# ------------------------------------------------------------------------------------------
# Importantly, `Flux` allows us to *automatically create neurons* with the **`Dense`**
# function. For example, in the last notebook, we were looking at a neuron with 2 inputs and
# 1 output:
#
#  <img src="data/single-neuron.png" alt="Drawing" style="width: 500px;"/>
#
#  We could create a neuron with two inputs and one output via
# ------------------------------------------------------------------------------------------

model = Dense(2, 1, σ)

# ------------------------------------------------------------------------------------------
# This `model` object comes with places to store weights and biases:
# ------------------------------------------------------------------------------------------

model.W

model.b

typeof(model.W)

# ------------------------------------------------------------------------------------------
# Unlike in previous notebooks, note that `W` is no longer a `Vector` (1D `Array`) and `b`
# is no longer a number! Both are now stored in so-called `TrackedArray`s and `W` is
# effectively being treated as a matrix with a single row. We'll see why in the next
# notebook.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Other helpful built-in functionality includes ways to automatically calculate gradients
# and also the cost function that we've used in previous notebooks -
#
# $$L(w, b) = \sum_i \left[y_i - f(x_i, w, b) \right]^2$$
#
# This is the "mean square error" function, which in `Flux` is named **`Flux.mse`**.
# ------------------------------------------------------------------------------------------

.nbexports/introductory-tutorials/broader-topics-and-ecosystem/intro-to-ml/13. ML - Intro to Flux.jl.jl


using SparseDiffTools
using LinearAlgebra
using BandedMatrices
using BlockBandedMatrices

n=10
dense=fill(1.,(n,n))
uptri=UpperTriangular(dense)
lotri=LowerTriangular(dense)

diagonal=Diagonal(dense)
bidiagonalU=Bidiagonal(dense,:U)
bidiagonalL=Bidiagonal(dense,:L)
tridiagonal=Tridiagonal(dense)
symtridiagonal=SymTridiagonal(dense)

banded=BandedMatrix(dense,(1,2))
blockbanded1=BlockBandedMatrix(dense,([1,2,3,4],[4,3,2,1]),(1,0))
blockbanded2=BlockBandedMatrix(dense,([4,3,2,1],[1,2,3,4]),(1,1))
bandedblockbanded1=BandedBlockBandedMatrix(dense,([1,2,3,4],[4,3,2,1]),(1,0),(1,1))
bandedblockbanded2=BandedBlockBandedMatrix(dense,([4,3,2,1],[1,2,3,4]),(1,1),(1,0))
bandedblockbanded3 = BandedBlockBandedMatrix(Zeros(2 * 4 , 2 * 4), ([4, 4] ,[4, 4]), (1, 1), (1, 1))

@test matrix_colors(dense)==1:n
@test matrix_colors(uptri)==1:n
@test matrix_colors(lotri)==1:n

@test matrix_colors(diagonal)==fill(1,n)
@test matrix_colors(bidiagonalU)==[1,2,1,2,1,2,1,2,1,2]
@test matrix_colors(bidiagonalL)==[1,2,1,2,1,2,1,2,1,2]
@test matrix_colors(tridiagonal)==[1,2,3,1,2,3,1,2,3,1]
@test matrix_colors(symtridiagonal)==[1,2,3,1,2,3,1,2,3,1]

@test matrix_colors(banded)==[1,2,3,4,1,2,3,4,1,2]
@test matrix_colors(blockbanded1)==[1,2,3,4,5,6,7,1,2,5]
@test matrix_colors(blockbanded2)==[1,5,6,7,8,9,1,2,3,4]
@test matrix_colors(bandedblockbanded1)==[1,2,3,1,4,5,6,1,2,4]
@test matrix_colors(bandedblockbanded2)==[1,3,4,5,6,5,1,2,1,2]
@test matrix_colors(bandedblockbanded3)==[1,2,3,1,4,5,6,4]

function _testvalidity(A)
    colorvec=matrix_colors(A)
    ncolor=maximum(colorvec)
    for colorvec in 1:ncolor
        subA=A[:,findall(x->x==colorvec,colorvec)]
        @test maximum(sum(subA,dims=2))<=1.0
    end
end

_testvalidity(dense)
_testvalidity(uptri)
_testvalidity(lotri)
_testvalidity(diagonal)
_testvalidity(bidiagonalU)
_testvalidity(bidiagonalL)
_testvalidity(tridiagonal)
_testvalidity(symtridiagonal)
_testvalidity(banded)
_testvalidity(blockbanded1)
_testvalidity(blockbanded2)
_testvalidity(bandedblockbanded1)
_testvalidity(bandedblockbanded2)


include("../src/hypergraph-affinity-functions.jl")
using MAT
using Plots
M = matread("Total_Ns.mat")
Ng = M["Ng"]
Nf = M["Nf"]
Nw = M["Nw"]

N = Nf
r = 10

## Bootstrap
for k = 2:4
B = 100
for j = 1:3
if  j == 1
    M = Nf
elseif j == 2
    M = Nw
else
    M = Ng
end

H1, H2 = affinities_from_N(M)
alpha = 1/2
B1 = baselines(r,alpha)
B2 = baselines(r,1-alpha)
tR1 = H1./B1
tR2 = H2./B2

# Set of types of hyperedges
Nk = weights(M[k,1:k+1])

TR1 = tR1[k,1:k]
TR0 = tR2[k,1:k]

R0, R1, MR1, MR0, SR1, SR0, A0, A1, MA1, MA0, SA1, SA0 = Bootstrap_From_N(Nk,B,alpha)

# Plot
sc = 1
lw = 2
ms = 9
s1 = 375
s2 = 300
gfs = 18
tfs = 18
titlesize = 18
color1 = :green
color2 = :blue
title = ""
leg = false
if j == 3
    xlab = "Affinity type t"
else
    xlab = ""
end

if k == 2
    yl = [.2,1.7]
elseif k ==3
    yl = [.2,2]
elseif k == 4
    yl = [.1,2.5]
end

if k == 2
    ylab = "Affinity / Baseline"
else
    ylab = ""
end
p = Plots.plot([1; k],[1; 1],yscale = :identity,fontfamily = "helvetica",linewidth = 2,legend = false,grid = false, size = (s1,s2),color = :gray)
title!(p,title,titlefontsize = titlesize)
xaxis!(p,xlab, xticks = 1:k,tickfontsize=tfs,guidefontsize = gfs)
yaxis!(p,ylab,tickfontsize=tfs,guidefontsize = gfs,ylim = yl)
plot!(p,[TR0 TR0], fillrange=[MR0-sc*SR0 MR0+sc*SR0], fillalpha=0.3,color = color2, linewidth = lw, markersize = ms,
    markershape = :circle,markerstrokecolor = color2, label = class1)
plot!(p,1:k,[TR1 TR1], fillrange=[MR1-sc*SR1 MR1+sc*SR1],fillalpha=0.3,color = color1, linewidth = lw, markersize = ms,
    markershape = :circle, markerstrokecolor = color1, label = class2)

savefig("Figures/Ratios_$(k)_$(j)_sc_$sc.pdf")
end
end


#
# Visualizations related to orbital energy
#

# These are not fully implemented!

"""
Plot the zero velocity curves for the Synodic, normalized CR3BP system.
"""
function zerovelocityplot(orbit::CR3BPOrbit;
                          nondimensional_range = range(-2; stop=2, length=1000),
                          kwargs...)

    curves = zerovelocity_curves(orbit; nondimensional_range = nondimensional_range)

    default = (; title  = "Zero Velocity Curves" * (name(system(orbit)) == (:Primary, :Secondary) ? "" : " for the $(name(system(orbit))[1])-$(name(system(orbit))[2]) System"),
                 xlabel = "X ($(lengthunit(orbit)))",
                 ylabel = "Y ($(lengthunit(orbit)))",
                 labels = :none,
                 formatter = :plain,
                 grid      = :on,
                 linewidth = 2,
                 dpi       = 150)

    options = merge(default, kwargs)

    fig = plot(; options...)
    for curve ∈ curves
        plot!(fig, curve[:,1], curve[:,2]; kwargs...)
    end

    return fig
end

"""
Plot the zero velocity curves for the Synodic, normalized CR3BP
system to the last figure.
"""
function zerovelocityplot!(fig, orbit::CR3BPOrbit;
                          nondimensional_range = range(-2; stop=2, length=1000),
                          kwargs...)

    curves = zerovelocity_curves(orbit; nondimensional_range = nondimensional_range)

    default = (; )
    options = merge(default, kwargs)

    fig = plot!(; options...)
    for (i,curve) ∈ zip(1:length(curves), curves)
        plot!(fig, curve[:,1], curve[:,2]; kwargs...)
    end

    return fig
end


@testset "woodbury_pd_mat" begin
    A = randn(4, 2)
    D = Diagonal(randn(2).^2 .+ 1)
    S = Diagonal(randn(4).^2 .+ 1)
    x = randn(size(A, 1))

    W = WoodburyPDMat(A, D, S)
    W_dense = PDMat(Symmetric(Matrix(W)))

    @testset "Valid constructors" begin
        @test W == typeof(W)(W.A, W.D, W.S)
    end

    @testset "invalid constructors error" begin
        @test_throws ArgumentError WoodburyPDMat(randn(5, 2), D, S)
        @test_throws ArgumentError WoodburyPDMat(randn(4, 3), D, S)
        @test_throws ArgumentError WoodburyPDMat(A, Diagonal(.-randn(2).^2), S)
        @test_throws ArgumentError WoodburyPDMat(A, D, Diagonal(.-randn(4).^2))
    end

    @testset "Basic functionality" begin
        # Checks getindex works.
        @test all(isapprox.(W, W_dense))
    end

    @testset "unwhiten!" begin
        @test PDMats.unwhiten!(similar(x), W, x) ≈ PDMats.unwhiten!(similar(x), W_dense, x)
    end

    @testset "logdet" begin
        @test logdet(W) ≈ logdet(W_dense)
        test_ad(randn(), A, D, S) do A, D, S
            logdet(WoodburyPDMat(A, D, S))
        end
    end

    @testset "invquad" begin
        @test invquad(W, x) ≈ invquad(W_dense, x)

        test_ad(randn(), A, D, S, x) do A, D, S, x
            W = WoodburyPDMat(A, D, S)
            return invquad(W, x)
        end
    end

    @testset "*" begin
        c = 2.0
        @test c * W == W * c
        @test c * W_dense ≈ c * W atol=1e-6
        @test (c * W) isa WoodburyPDMat
    end

    @testset "MvNormal logpdf" begin
        m = randn(size(A, 1))
        @test logpdf(MvNormal(m, W), x) ≈ logpdf(MvNormal(m, Symmetric(Matrix(W))), x)

        test_ad(randn(), m, A, D, S, x) do m, A, D, S, x
            W = WoodburyPDMat(A, D, S)
            return logpdf(MvNormal(m, W), x)
        end
    end

    @testset "GenericMvTDist logpdf" begin
        α = 2.1
        m = randn(size(A, 1))
        @test isapprox(
            logpdf(Distributions.GenericMvTDist(α, m, W), x),
            logpdf(Distributions.GenericMvTDist(α, m, W_dense), x),
        )

        test_ad(randn(), α, m, A, D, S, x) do α, m, A, D, S, x
            W = WoodburyPDMat(A, D, S)
            return logpdf(Distributions.GenericMvTDist(α, m, W), x)
        end
    end
end


##########################################################################
#   	Basic functionality for doing MPS calculations with OBC     	#
##########################################################################
using LinearAlgebra
using Arpack

# Define sites as 3 legged tensors filled with elements of some type
Site{T} = Array{T,3}
# An MPS is an array of Sites of a certain type
MPS{T} = Vector{Site{T}}
# Define operators as 4 legged tensors filled with elements of some type
Operator{T} = Array{T,4}
# A MPO is an array of Operators
MPO{T} = Vector{Operator{T}}

"""
    random_mps_obc(N::Int, D::Int, d, tensortype::Type{T}=ComplexF64)::MPS{T} where T

Function to generate a random MPS with OBC
N: Number of sites
D: Bond dimension
d: Vector of physical dimensions, if a single integer dimension is given, it is assumed that all physical dimensions are the same, if varying dimensions are wanted a vector of integers of length N has to be supplied
The convention for the 3 index tensors is that the first two are the virtual ones, the last one is the physical one
   3
   |
1--A--2
"""
function random_mps_obc(N::Int, D::Int, d, tensortype::Type{T} = ComplexF64)::MPS{T} where {T}
    # Ensure the the system size is at least two
    @assert(N > 1)

    # Initialize
    mps = Array{Site{T}}(undef, N)

    if isa(d, Number)
        # Single number, I assume all dimensions are the same
        dim = d * ones(Int64, N)
    else
        # Make sure that the input is really a vector as julia distinguishes between nx1 matrices and column vectors
        dim = vec(d)
        # Make sure input is valid
        @assert(length(dim) == N)
    end

    # Left boundary tensor (row vector)
    mps[1] = rand(tensortype, 1, D, dim[1])

    # Right boundary tensor (column vector)
    mps[N] = rand(tensortype, D, 1, dim[N])

    # Tensors in between
    for i = 2:N-1
        mps[i] = rand(tensortype, D, D, dim[i])
    end

    return mps
end

"""
    basis_state_obc(configuration::Vector{<:Int}, d::Int=2)::MPS

Prepare the a product state corresponding |configuration> on N sites where configuration is an Array containing N elements from 1 to d. Each site is then initialized in the dth canonical basis state
"""
function basis_state_obc(configuration::Vector{<:Int}, d::Int = 2)::MPS{Float64}
    # Some error checking
    if any(x -> (x < 1 || x > d), configuration)
        throw(ArgumentError("configuration must contain integer elements in the range from 1 to d, got d=$(repr(d)), configuration=$(repr(configuration))"))
    end
    # Generate the MPS
    N = length(configuration)
    psi = MPS{Float64}(undef, N)
    tensors = Vector{Array{Float64,3}}(undef, d)
    for i = 1:d
        tmp = zeros(Float64, 1, 1, d)
        tmp[1, 1, i] = 1.0
        tensors[i] = tmp
    end
    for i = 1:N
        psi[i] = tensors[configuration[i]]
    end
    return psi
end

"""
    calculate_overlap(mps1::MPS,mps2::MPS)::Number

Given two mps, compute the overlap <mps1|mps2>.
"""
function calculate_overlap(mps1::MPS, mps2::MPS)::Number
    # Make sure, that the two MPS have the same length
    N1 = length(mps1)
    N2 = length(mps2)
    @assert(N1 == N2)

    # Now compute ther overlap
    overlap = ones(Float64, 1, 1)
    for i = 1:N1
        overlap = contract_tensors(overlap, [2], mps2[i], [1])
        overlap = contract_tensors(conj(mps1[i]), [1; 3], overlap, [1; 3])
    end

    return overlap[1]
end

"""
   expectation_value(mps::MPS, mpo::MPO)::Number

Given a MPS and a MPO compute the expectation value <mps|mpo|mps>.
"""
function expectation_value(mps::MPS, mpo::MPO)::Number
    # Make sure, that the two MPS have the same length
    N1 = length(mps)
    N2 = length(mpo)
    @assert(N1 == N2)

    # Contract zipper like
    val = ones(Float64, 1, 1, 1)
    for i = 1:N1
        val = contract_tensors(val, [1], conj(mps[i]), [1])
        val = contract_tensors(val, [1; 4], mpo[i], [1; 3])
        val = contract_tensors(val, [1; 4], mps[i], [1; 3])
    end

    return val[1]
end

"""
    gaugeMPS(mps::MPS{T}, direction::Symbol=:right, normalize::Bool=false)::MPS{T}

Function to bring an MPS with OBC in canonical form
direction: tells, whether it will be left or right canonical gauge
If normalize is set to true, the resulting state will be normalized.
"""
function gaugeMPS(mps::MPS{T}, direction::Symbol = :right, normalize::Bool = false)::MPS{T} where {T}
    mps_gauged = deepcopy(mps)
    gaugeMPS!(mps_gauged, direction, normalize)
    return mps_gauged
end

"""
    gaugeMPS!(mps::MPS,direction::Dir=right,)

Bring an MPS with OBC in canonical form. Direction tells, whether it will be left or right canonical gauge If normalize is set to true, the resulting state will be
normalized. This function overwrites the input MPS with its gauged version
"""
function gaugeMPS!(mps::MPS, direction::Symbol = :right, normalize::Bool = false)
    # Check that we got a meaningful direction
    if (direction != :left && direction != :right)
        throw(ArgumentError("direction must be :left or :right, got $(repr(direction))"))
    end
    # Start gauging
    N = length(mps)
    if (direction == :left)
        # Case that I want left canonical gauge
        M = mps[1]
        for i = 1:N-1
            mps[i], res = gauge_site(M, direction)
            M = contract_tensors(res, [2], mps[i+1], [1])
            if (i == (N - 1))
                if (normalize)
                    mps[N], _ = gauge_site(M, direction)
                else
                    mps[N] = M
                end
            end
        end
    else
        # Case that I want right canonical gauge
        M = mps[N]
        for i = N:-1:2
            mps[i], res = gauge_site(M, direction)
            M = contract_tensors(mps[i-1], [2], res, [1], [1; 3; 2])
            if (i == 2)
                if (normalize)
                    mps[1], _ = gauge_site(M, direction)
                else
                    mps[1] = M
                end
            end
        end
    end
end

"""
    gauge_site(A::Site{T}, direction::Dir)::Tuple{Site{T},Matrix{T}} where T

Bring a single tensor A of an MPS with OBC in canonical form
direction tells, whether it will be left or right normalized
The new tensor is M, the residual matrix which has to be multiplied
in the next tensor is stored in res
"""
function gauge_site(A::Site{T}, direction::Symbol)::Tuple{Site{T},Matrix{T}} where {T}
    if (direction == :left)
        Dl, Dr, d = size(A)
        M = permutedims(A, [3 1 2])
        M = reshape(M, (d * Dl, Dr))
        U, S, V = svd(M)
        dsv = length(S)
        U = reshape(U, (d, Dl, dsv))
        M = permutedims(U, [2 3 1])
        res = Matrix(Diagonal(S)) * V'
        return M, res
    else
        Dl, Dr, d = size(A)
        M = permutedims(A, [1 3 2])
        M = reshape(M, (Dl, Dr * d))
        U, S, V = svd(M)
        dsv = length(S)
        V = V'
        V = reshape(V, (dsv, d, Dr))
        M = permutedims(V, [1 3 2])
        res = U * Matrix(Diagonal(S))
        return M, res
    end
end

"""
    contract_virtual_indices(mps::MPS)::Vector{<:Number}

Given an MPS contract the virtual indices such that one obtains a dense vector. The indices are ordered such that they are compatible with the standard Julia kronecker product.

Warning: the object constructed will have exponential memory requirements in terms of the number of sites, use with care!
"""
function contract_virtual_indices(mps::MPS)::Vector{<:Number}
    N = length(mps)

    # Since we deal with open boundary conditions, we drop the dummy indices one on the left (right) boundary for the first (last) tensor manually. We start from the right to have the physical indices in the order compatible with Julia's kronecker product
    res = mps[N][:, 1, :]
    for i = N-1:-1:2
        res = contract_tensors(res, [ndims(res) - 1], mps[i], [2])
    end
    res = contract_tensors(res, [ndims(res) - 1], mps[1][1, :, :], [1])

    # Now reshape the result accordingly
    res = reshape(res, prod(size(res)))

    return res
end

"""
    contract_virtual_indices(mps::MPO)::Matrix{<:Number}

Given an MPO contract the virtual indices such that one obtains a dense matrix. The indices are ordered such that they are compatible with the standard Julia kronecker product.

Warning: the object constructed will have exponential memory requirements in terms of the number of sites, use with care!
"""
function contract_virtual_indices(mpo::MPO)::Matrix{<:Number}
    N = length(mpo)

    # Since we deal with open boundary conditions, we drop the dummy indices one on the left (right) boundary for the first (last) tensor manually. We start from the right to have the physical indices in the order compatible with Julia's kronecker product
    res = mpo[N][:, 1, :, :]
    for i = N-1:-1:2
        res = contract_tensors(res, [ndims(res) - 2], mpo[i], [2])
    end
    res = contract_tensors(res, [ndims(res) - 2], mpo[1][1, :, :, :], [1])

    # Reshuffle the indices to be compatible with Julia's standard kronecker product
    res = permutedims(res, [collect(1:2:ndims(res)); collect(2:2:ndims(res))])

    # Now reshape the result accordingly
    dims = size(res)
    dr = prod(dims[1:N])
    dc = prod(dims[N+1:end])
    res = reshape(res, (dr, dc))

    return res
end

"""
    find_groundstate(H::Array{Operator},D::Int64,acc::Float64,max_sweeps::Int64=-1)

Given a Hamiltonian MPO H, find an MPS approximation for its ground state with bond
dimension D converged to a relative accuracy acc. If a positive number for max_sweeps
is given, the maximum amount of iterations is limited to max_sweeps
"""
function find_groundstate(H::Vector{Operator{T}}, D::Int64, d::Int64, acc::Float64, max_sweeps::Int64) where {T}
    N = length(H)

    # Get the physical dimensions from the Hamiltonian
    d = Vector{Int64}(undef, length(H))
    for i = 1:length(H)
        d[i] = size(H[i], 3)
    end

    # Random MPS to start the calculation with
    mps = random_mps_obc(N, D, d)

    # Put it in right canonical gauge
    gaugeMPS!(mps, :right)
    # Precalculate the partial contractions
    LR = setup_R(H, mps)

    # Now start the sweeping
    num_of_sweeps = 0
    E0 = 0
    Eold = 1E5
    while (true)
        num_of_sweeps = num_of_sweeps + 1
        println("Sweep number ", num_of_sweeps)

        # From left to right starting by 1 up to N-1
        E_local = 0
        res = 0
        for i = 1:N-1
            Left = LR[i]
            Right = LR[i+1]
            E_local, M = solve_eigenvalue_problem(Left, Right, H[i])
            mps[i], _ = gauge_site(M, :left)
            LR[i+1] = update_left(Left, mps[i], mps[i], H[i])
        end

        # From left to right starting by 1 up to N-1
        for i = N:-1:2
            Left = LR[i]
            Right = LR[i+1]
            E_local, M = solve_eigenvalue_problem(Left, Right, H[i])
            mps[i], res = gauge_site(M, :right)
            LR[i] = update_right(Right, mps[i], mps[i], H[i])
        end

        # Check convergence criterion
        if (abs((Eold - E_local) / Eold) < acc)
            E0 = E_local
            # Restore normalisation (permutation is needed to bring the indices back into right order after contracting)
            mps[1] = contract_tensors(mps[1], [2], res, [1], [1; 3; 2])
            println("Convergence to desired accuracy achieved")
            break
        elseif (max_sweeps > 0 && num_of_sweeps >= max_sweeps)
            E0 = E_local
            # Restore normalisation (permutation is needed to bring the indices back into right order after contracting)
            mps[1] = contract_tensors(mps[1], [2], res, [1], [1; 3; 2])
            warn("Reached maximum number of iterations before convergence to desired accuracy")
            break
        end
        Eold = E_local
    end

    return E0, mps, num_of_sweeps
end

"""
    getHeff(Left,Right,W::Operator)

Construct the effective Hamiltonian
"""
function getHeff(Left, Right, W::Operator)
    Htemp = contract_tensors(W, [2], Right, [2])
    Htemp = contract_tensors(Left, [2], Htemp, [1])
    Htemp = permutedims(Htemp, [1; 5; 3; 2; 6; 4])
    dim1, dim2, dim3, dim4, dim5, dim6 = size(Htemp)
    Heff = reshape(Htemp, (dim1 * dim2 * dim3, dim4 * dim5 * dim6))

    return Heff
end

"""
    solve_eigenvalue_problem(Left,Right,W)

Function to construct and solve the eigenvalue problem which arises on each site for the effective Hamiltonian
"""
function solve_eigenvalue_problem(Left, Right, W)
    Htemp = contract_tensors(W, [2], Right, [2])
    Htemp = contract_tensors(Left, [2], Htemp, [1])
    Htemp = permutedims(Htemp, [1; 5; 3; 2; 6; 4])
    dim1, dim2, dim3, dim4, dim5, dim6 = size(Htemp)
    Heff = reshape(Htemp, (dim1 * dim2 * dim3, dim4 * dim5 * dim6))

    E_local, M = eigs(Heff, nev = 1, which = :SR)
    M = reshape(M, (dim1, dim5, dim3))

    return real(E_local[1]), M
end

"""
    setup_R(H::MPO{T1},mps::MPS{T2}) where {T1,T2}

Function to calculate the partial contractions needed to form the effective Hamiltonian
"""
function setup_R(H::MPO{T1}, mps::MPS{T2}) where {T1,T2}
    N = length(H)
    Tres = Base.return_types(*, (T1, T2))[1]
    LR = Vector{Array{Tres,3}}(undef, N + 1)

    # We need only N-1 partial contractions, however, we set the edges to dummy values 1 that we can recursively compute every contraction resuing the previous ones
    LR[1] = ones(Tres, 1, 1, 1)
    LR[N+1] = ones(Tres, 1, 1, 1)

    # Now compute the partial contractions starting from the right (as I start sweeping on the left in the ground state search)
    for i = N:-1:2
        tmp = contract_tensors(LR[i+1], [3], mps[i], [2])
        tmp = contract_tensors(H[i], [2; 4], tmp, [2; 4])
        LR[i] = contract_tensors(conj(mps[i]), [2; 3], tmp, [3; 2])
    end

    return LR
end

"""
    update_left(LR, M_left::Site, M_right::Site, W::Operator)

Function to perform an update the partial contractions required for iterative ground state search starting from the left stored in LR
"""
function update_left(LR, M_left::Site, M_right::Site, W::Operator)
    res = contract_tensors(LR, [3], M_right, [1])
    res = contract_tensors(W, [1; 4], res, [2; 4])
    res = contract_tensors(conj(M_left), [1; 3], res, [3; 2])

    return res
end

"""
    update_right(LR, M_left::Site, M_right::Site, W::Operator)

Function to perform an update the partial contractions required for iterative ground state search starting from the right stored in LR
"""
function update_right(LR, M_left::Site, M_right::Site, W::Operator)
    res = contract_tensors(LR, [3], M_right, [2])
    res = contract_tensors(W, [2; 4], res, [2; 4])
    res = contract_tensors(conj(M_left), [2; 3], res, [3; 2])

    return res
end

"""
    apply_operator(operator::MPO{T1}, mps::MPS{T2})::MPS where {T1,T2}

Apply an operator given as MPO to an MPS. The resulting MPS will have a bond dimension that is the product of the bond dimensions of the MPS and the MPO.
"""
function apply_operator(operator::MPO{T1}, mps::MPS{T2})::MPS where {T1,T2}
    N1 = length(mps)
    N2 = length(operator)
    @assert(N1 == N2)

    # Generate a new MPO of the correct type
    Tres = Base.return_types(*, (T1, T2))[1]
    res = MPS{Tres}(undef, N1)

    # Apply the MPO to the MPS and generate new MPS
    for i = 1:N1
        temp = contract_tensors(operator[i], [4], mps[i], [3])
        temp = permutedims(temp, (1, 4, 2, 5, 3))
        dim1, dim2, dim3, dim4, dim5 = size(temp)
        res[i] = reshape(temp, (dim1 * dim2, dim3 * dim4, dim5))
    end
    return res
end

"""
    apply_operator!(operator::MPO{T1}, mps::MPS{T2})::MPS where {T1,T2}

Apply an operator given as MPO to an MPS. The resulting MPS will have a bond dimension that is the product of the bond dimensions of the MPS and the MPO. The input will be overwritten by the result. This requires that the type of the elements of the input MPS is able to accomodate the result of multiplying the MPO tensors into the MPS tensors (e.g. applying a complex MPO to a real MPS cannot be done inpalce as the result will be complex)
"""
function apply_operator!(operator::MPO, mps::MPS)
    N1 = length(mps)
    N2 = length(operator)
    @assert(N1 == N2)

    # Contract the tensors for each site
    for i = 1:N1
        temp = contract_tensors(operator[i], [4], mps[i], [3])
        temp = permutedims(temp, (1, 4, 2, 5, 3))
        dim1, dim2, dim3, dim4, dim5 = size(temp)
        mps[i] = reshape(temp, (dim1 * dim2, dim3 * dim4, dim5))
    end
end

"""
    apply_operator(op1::MPO{T1}, op2::MPO{T2})::MPO where {T1,T2}

Multiply two MPOs together to get an expression for op2 * op1 in MPO form. The resulting MPS will have a bond dimension that is the product of the bond dimensions of both MPOs.
"""
function apply_operator(op1::MPO{T1}, op2::MPO{T2})::MPO where {T1,T2}
    N1 = length(op1)
    N2 = length(op2)
    @assert(N1 == N2)

    # Generate a new MPS of the correct type
    Tres = Base.return_types(*, (T1, T2))[1]
    res = MPO{Tres}(undef, N1)

    # Contract the tensors for each site
    for i = 1:N1
        temp = contract_tensors(op2[i], [4], op1[i], [3])
        temp = permutedims(temp, (1, 4, 2, 5, 3, 6))
        dim1, dim2, dim3, dim4, dim5, dim6 = size(temp)
        res[i] = reshape(temp, (dim1 * dim2, dim3 * dim4, dim5, dim6))
    end
    return res
end

"""
    sum_states(mps1::MPS{T1}, mps2::MPS{T2})::MPS where {T1,T2}

Add two MPSs together to get an expression for mps1 + mps2 in MPS form. The resulting MPS will have a bond dimension that is the sum of the bond dimensions of both MPOs.
"""
function sum_states(mps1::MPS{T1}, mps2::MPS{T2})::MPS where {T1,T2}
    N1 = length(mps1)
    N2 = length(mps2)
    @assert(N1 == N2)

    # Generate a new MPO of the correct type
    Tres = Base.return_types(+, (T1, T2))[1]
    res = MPS{Tres}(undef, N1)

    # The first site needs special treatment
    tensor1 = mps1[1]
    tensor2 = mps2[1]
    _, Dr1, d1 = size(tensor1)
    _, Dr2, d2 = size(tensor2)
    new_tensor = zeros(Tres, 1, Dr1 + Dr2, d1)
    for r = 1:d1
        new_tensor[1, :, r] = [tensor1[1:1, :, r] tensor2[1:1, :, r]]
    end
    res[1] = new_tensor
    # The tensors in between
    for i = 2:N1-1
        tensor1 = mps1[i]
        tensor2 = mps2[i]
        Dl1, Dr1, d1 = size(tensor1)
        Dl2, Dr2, d2 = size(tensor2)
        new_tensor = zeros(Tres, Dl1 + Dl2, Dr1 + Dr2, d1)
        for r = 1:d1
            new_tensor[1:Dl1, 1:Dr1, r] = tensor1[:, :, r]
            new_tensor[Dl1+1:end, Dr1+1:end, r] = tensor2[:, :, r]
        end
        res[i] = new_tensor
    end
    # The last site needs special treatment
    tensor1 = mps1[N1]
    tensor2 = mps2[N1]
    Dl1, _, d1 = size(tensor1)
    Dl2, _, d2 = size(tensor2)
    new_tensor = zeros(Tres, Dl1 + Dl2, 1, d1)
    for r = 1:d1
        new_tensor[:, 1, r] = [tensor1[:, 1, r]; tensor2[:, 1, r]]
    end
    res[N1] = new_tensor

    return res
end

"""
    sum_operators(op1::MPO{T1}, op2::MPO{T2})::MPO where {T1,T2}

Add two MPOs together to get an expression for op2 + op1 in MPO form. The resulting MPO will have a bond dimension that is the sum of the bond dimensions of both MPOs.
"""
function sum_operators(op1::MPO{T1}, op2::MPO{T2})::MPO where {T1,T2}
    N1 = length(op1)
    N2 = length(op2)
    @assert(N1 == N2)

    # Generate a new MPO of the correct type
    Tres = Base.return_types(+, (T1, T2))[1]
    res = MPO{Tres}(undef, N1)

    # The first site needs special treatment
    tensor1 = op1[1]
    tensor2 = op2[1]
    _, Dr1, dr1, dc1 = size(tensor1)
    _, Dr2, dr2, dc2 = size(tensor2)
    new_tensor = zeros(Tres, 1, Dr1 + Dr2, dr1, dc1)
    for r = 1:dr1
        for c = 1:dc1
            new_tensor[1, :, r, c] = [tensor1[1:1, :, r, c] tensor2[1:1, :, r, c]]
        end
    end
    res[1] = new_tensor
    # The tensors in between
    for i = 2:N1-1
        tensor1 = op1[i]
        tensor2 = op2[i]
        Dl1, Dr1, dr1, dc1 = size(tensor1)
        Dl2, Dr2, dr2, dc2 = size(tensor2)
        new_tensor = zeros(Tres, Dl1 + Dl2, Dr1 + Dr2, dr1, dc1)
        for r = 1:dr1
            for c = 1:dc1
                new_tensor[1:Dl1, 1:Dr1, r, c] = tensor1[:, :, r, c]
                new_tensor[Dl1+1:end, Dr1+1:end, r, c] = tensor2[:, :, r, c]
            end
        end
        res[i] = new_tensor
    end
    # The last site needs special treatment
    tensor1 = op1[N1]
    tensor2 = op2[N1]
    Dl1, _, dr1, dc1 = size(tensor1)
    Dl2, _, dr2, dc2 = size(tensor2)
    new_tensor = zeros(Tres, Dl1 + Dl2, 1, dr1, dc1)
    for r = 1:dr1
        for c = 1:dc1
            new_tensor[:, 1, r, c] = [tensor1[:, 1, r, c]; tensor2[:, 1, r, c]]
        end
    end
    res[N1] = new_tensor

    return res
end

"""
    compute_entropy(mps::MPS, n:Int)::Float64

Compute the von Neumann along for the bipartion of the sites into two subsets A={1,...,n} and B={n+1,...,N} where N is the length of the MPS. If n<=0 is supplied, a bipartion of into A. For the result to make sense, mps has to be a normalized quantum state.
"""
function compute_entropy(mps::MPS{T}, n::Int = 0)::Float64 where {T}
    # Extact the length and check if the given position is reasonable
    N = length(mps)
    @assert(n <= N)

    # The entropy of the entire state is simply zero since it is a pure state, so nothing to compute
    if n == N
        return 0.0
    end

    # In case a value n<=0 is given, we assume the bipartion is taken in the center
    if n <= 0
        n = Int(round(N / 2))
    end

    # Get a copy of the input
    mps_loc = deepcopy(mps)

    # Put the sites left to n into left canonical gauge
    M = mps_loc[1]
    for i = 1:n
        mps_loc[i], res = gauge_site(M, :left)
        M = contract_tensors(res, [2], mps_loc[i+1], [1])
    end
    mps_loc[n+1] = M

    # Now start contracting from the right boundary to obtain the reduced density operator
    rdm = ones(T, 1, 1)
    for i = N:-1:n+1
        rdm = contract_tensors(rdm, [2], mps_loc[i], [2])
        rdm = contract_tensors(conj(mps_loc[i]), [2; 3], rdm, [1; 3])
    end

    # Now diagonalize the reduced density matrix and compute the entropy
    ev = real(eigvals(rdm))
    # Eigenvalues that are numerically zero sometimes become -1E-16, to prevent problems with the logarithm we filter them
    ev = filter(x -> x > 0.0, ev)
    # The von Neumann entropy for the reduced density operator
    entropy = -sum(ev .* log2.(ev))

    return entropy
end

"""
    sample_from_mps!(mps::MPS, gauge_input::Bool=true)::Vector{Int64}

Generate a sample from the probability distribution of basis states defined by the MPS following A. Ferris, G. Vidal, PRB 85 165146 (2021). If the flag gague_input is set to true, the input MPS will be put in right canoncial gauge and normalized, which is a requirement for the algorithm to work. In case one is sure that the MPS is already properly gauged and normalized the gauging step can be spared by setting the flag to false. In this case the input will stay untouched.
"""
function sample_from_mps!(mps::MPS, gauge_input::Bool = true)::Vector{Int64}
    # Extact the length and check if the given position is reasonable
    N = length(mps)
    res = zeros(Int64, N)

    # Put the state into right canonical gauge and make sure it is normalized
    if gauge_input
        gaugeMPS!(mps, :right, true)
    end

    # Now sample from the MPS
    A = 0
    p = 0.0
    M = mps[1][1, :, :]
    for i = 1:N
        d = size(M, 2)
        pacc = 0
        r = rand()
        for l = 1:d
            # Prepare the basis state
            basis_state = zeros(d)
            basis_state[l] = 1.0
            # Contract it into the physical index of the tensor
            A = contract_tensors(M, [2], basis_state, [1])
            # Determine the probability for the basis vector e_l
            p = contract_tensors(conj(A), [1], A, [1])
            pacc += real(p[1])
            if r < pacc
                res[i] = l
                break
            end
        end
        if i < N
            # Contract the result into the next tensor
            M = contract_tensors(1 / sqrt(real(p[1])) * A, [1], mps[i+1], [1])
        end
    end
    return res
end

"""
    sample_from_mps(mps::MPS)::Vector{Int64}

Generate a sample from the probability distribution of basis states defined by the MPS following A. Ferris, G. Vidal, PRB 85 165146 (2021).
"""
function sample_from_mps(mps::MPS)::Vector{Int64}
    return sample_from_mps!(deepcopy(mps))
end

"""
    svd_compress_mps(mps::MPS, Dmax::Int,, tol::Real = 0.0)::MPS

Compress a given MPS applying an a singular value decomposition at each bond. If Dmax > 0 is supplied (and simultaneously tol = 0.0), a maximum of Dmax singular values is kept, thus truncating the MPS to one with maximum bond dimension Dmax. If tol > 0 is given (and simultaneously Dmax = 0) then all singular values > tol are kept. If both are specified then at most Dmax singular values > tol are kept.
"""
function svd_compress_mps(mps::MPS, Dmax::Int, tol::Real = 0.0)::MPS
    # One of the two parameters has to be larger than zero
    @assert((Dmax > 0) || (tol > 0))
    # Extract the length and prepare a result
    N = length(mps)
    res = deepcopy(mps)
    # Gauge in both directions that redundant dimensions at the boundaries are removed
    gaugeMPS!(res, :left)
    gaugeMPS!(res, :right)
    Dnew = 0
    for i = 1:N-1
        # Check if the given bond dimension is larger than Dmax
        if size(res[i], 2) > Dmax
            Dl1, _, d1 = size(res[i])
            _, Dr2, d2 = size(res[i+1])
            tmp = contract_tensors(res[i], [2], res[i+1], [1])
            tmp = reshape(tmp, (Dl1 * d1, Dr2 * d2))
            U, S, V = svd(tmp)
            M = diagm(S) * V'
            # Now truncate
            if Dmax > 0 && tol == 0.0
                U = U[:, 1:Dmax]
                M = M[1:Dmax, :]
                Dnew = Dmax
            else
                ind = findall(x -> x > tol, S)
                if Dmax > 0 && length(ind) > Dmax
                    ind = ind[1:Dmax]
                end
                U = U[:, ind]
                M = M[ind, :]
                Dnew = length(ind)
            end
            # Reshape and set new tensors
            res[i] = permutedims(reshape(U, (Dl1, d1, Dnew)), (1, 3, 2))
            res[i+1] = reshape(M, (Dnew, Dr2, d2))
        end
    end
    return res
end

"""
    decompose_into_mpo(M::Matrix{T}, d::Vector{Int}) where T <:Number

Given a many-body operator in form of a dense matrix, decompose it into an MPO. It is assumed that the many-body operator follows the index convention of Julia's built-in kronecker product. The vector d specifies the local dimensions of the Hilbert space.

# Examples
Decomposing a simple two-site operator made up from sum of Pauli terms into MPO form.
```julia-repl
julia>  Id = [1.0 0; 0.0 1.0]
julia>  X = [0.0 1.0; 1.0 0.0]
julia>  Z = [1.0 0.0; 0.0 -1.0]
julia>  H = kron(X,X) + kron(Id,Z) + kron(Z,Id)
julia>  mpo = decompose_into_mpo(H, 2)
```
"""
function decompose_into_mpo(M::Matrix{T}, d::Vector{Int})::MPO{T} where {T<:Number}
    N = length(d)
    dim = prod(d)
    if size(M) != (prod(d), prod(d))
        throw(ArgumentError("matrix not compatible with specified local dimensions, got size(M)=$(repr(size(M))), d=$(repr(d))"))
    end
    # The MPO holding the result
    res = MPO{T}(undef, N)
    # Reshape the matrix and rearrange the indices, such that the row and column index for each site are adjacent and add dummy indices 1 at the boundaries
    A = reshape(M, (1, reverse(d)..., reverse(d)..., 1))
    ind = collect(Iterators.flatten(zip(collect(N:-1:1), collect(2N:-1:N+1))))
    A = permutedims(A, (1, ind .+ 1..., 2 * length(d) + 2))
    # Now split it into tensors using an SVD
    Dl = 1
    for i = 1:N-1
        # Take the first three indices together
        dims = size(A)
        A = reshape(A, (prod(dims[1:3]), prod(dims[4:end])))
        # SVD the matrix representation
        U, S, V = svd(A)
        A = diagm(S) * V'
        Dr = size(U, 2)
        # Extract the tensor and the remaining part
        U = reshape(U, (Dl, d[i], d[i], Dr))
        res[i] = permutedims(U, (1, 4, 2, 3))
        A = reshape(A, (Dr, dims[4:end]...))
        Dl = Dr
    end
    res[N] = permutedims(A, (1, 4, 2, 3))

    return res
end

"""
    decompose_into_mpo(M::Matrix{T}, d::Int) where T <:Number

Simplified interface to the more general method assuming that all local dimensions are equal to d.
"""
function decompose_into_mpo(M::Matrix{T}, d::Int)::MPO{T} where {T<:Number}
    dl, dr = size(M)
    if dl != dr
        throw(ArgumentError("local dimensions must all be the same, obtained a matrix with dimensions $(repr((dl,dr)))"))
    end
    N = Int(round(log(d, dl)))
    return decompose_into_mpo(M, d * ones(Int64, N))
end


"""
Plot example file, and create a matlab file for good measure.  If you are not interested in the
matlab file comment out those sections.
"""

# Pkg.add("Plots")  # uncomment for first run
# Pkg.add("Revise")  # uncomment for first run
# using Revise
using MAT
using HDF5
using Plots
using Gantner

# load file data
name = joinpath(@__DIR__, "data.dat")
(t, time, fs, chanTextLegend) = gantnerread(name)
# t = range(0, length=size(Time, 1), step=1/fs)

plotly()       # use for general purpose plotting and looking for peaks
plot(t, time, label = reshape(chanTextLegend, 1, :))
xlabel!("Time [s]")
ylabel!("Amp")
title!(name)
ylims!(0, 1.2)
xlims!(0,0.3)

# base, ext = split(filename, ".")
# fnamemat = "$base.mat"
# command from .julia/packages/MAT/.../README.md
matwrite(joinpath(@__DIR__, "data.mat"), Dict(
	"t" => collect(t),
	"Time" => time);
	compress = true)

# delete data.mat when done
sleep(10)
rm(joinpath(@__DIR__, "data.mat"))


## Dyad-based network measures ##

"""
    density(g::Graph)

Graph density. It is a proportion of number of existing edges to the maximum
possible number of edges. Density of an empty graph is 0, and density of
a full graph is 1. Graphs having self-loops and/or multiple edges can have
density larger than 1.
"""
density(g::Graph) = edgecount(g) / maxedgecount(g)

"""
    mutuality(g::Graph)

Number of mutual dyads in the graph. A dyad is mutual if for two nodes
`A` and `B` there are both edge `A → B` and `B → A`.
Edges in undirected graphs always constitute mutual dyads.

# References

Moreno, J.L., and Jennings, H.H. (1938). "Statistics of Social Configurations." Sociometry, 1, 342-374.
"""
function mutuality(g::DirectedGraph)
    count = 0
    for n = g.nodes
        i, j = 1, 1
        while i < length(n.forward) && j < length(n.reverse)
            if n.forward[i] < n.reverse[j]
                i += 1
            elseif n.forward[i] > n.reverse[j]
                j += 1
            else
                count += 1
                i += 1
                j += 1
            end
        end
    end
    return count
end

mutuality(g::UndirectedGraph) = edgecount(g)

"""
    reciprocity(g::Graph)

Graph reciprocity. It is a proportion of mutual dyads to a total
number of non-null dyads.
[Wiki](https://en.wikipedia.org/wiki/Reciprocity_(network_science)#Traditional_definition)
"""
reciprocity(g::Graph) = mutuality(g) / edgecount(g)

"""
    dyadcensus(g::Graph)

Dyad census of the graph. It is a 3-tuple of form (M, A, N), standing for
the numbers of mutual, asymmetric, and null dyads.
"""
function dyadcensus(g::Graph)
    m = mutuality(g)
    return (m, edgecount(g) - m, maxedgecount(g) - edgecount(g))
end

"""
    selfloopcount(g::Graph)

Number of self-loops (edges that start and end at the same node) in the graph.
"""
function selfloopcount(g::Graph)
    cnt = 0
    @inbounds for i = nodes(g)
        @inbounds for j = neighbors(g, i)
            if i == j; cnt += 1; end
        end
    end
    return cnt
end

function selfloopnodes(g::Graph)
    ns = Int[]
    @inbounds for i = nodes(g)
        @inbounds for j = neighbors(g, i)
            if i == j
                push!(ns, i)
                break
            end
        end
    end
    return ns
end


module TestOutliernicer

using StatsBase: sample, mean
using Test
using TSML

function test_artificialdata()
  Random.seed!(123)
  mdates = DateTime(2017,1,1):Dates.Hour(1):DateTime(2017,6,1)
  mvals = rand(1:1000,length(mdates))
  # create some outliers
  soutliers = rand([500:10000;-10000:500],div(length(mdates),10))
  soutndx = sample(1:length(mdates),length(soutliers))
  mvals[soutndx] = soutliers
  df = DataFrame(Date=mdates,Value=mvals)
  outnicer = Outliernicer(Dict(:dateinterval => Dates.Hour(1)))
  fit!(outnicer,df)
  resdf = transform!(outnicer,df)
  @test round(mean(resdf.Value),digits=2) == 510.32
end
@testset "Outliernicer: using artificial data" begin
  test_artificialdata()
end

function test_basicoutlier()
  Random.seed!(123)
  fname = joinpath(dirname(pathof(TSML)),"../data/testdata.csv")
  csvfilter = CSVDateValReader(Dict(:filename=>fname,:dateformat=>"dd/mm/yyyy HH:MM"))
  valgator = DateValgator(Dict(:dateinterval=>Dates.Hour(1)))
  valnner = DateValNNer(Dict(:dateinterval=>Dates.Hour(1)))
  stfier = Statifier(Dict(:processmissing=>true))
  mono = Monotonicer(Dict())
  outliernicer = Outliernicer(Dict(:dateinterval=>Dates.Hour(1)))

  mpipeline1 = @pipeline csvfilter |> valgator |> mono |> valnner |> outliernicer |> stfier
  respipe1 = fit_transform!(mpipeline1)
  @test round(sum(respipe1[1,3:20])) == -213862.0

  mpipeline2 = @pipeline csvfilter |> valgator |> mono |> outliernicer |> stfier
  respipe2 = fit_transform!(mpipeline2)
  @test round(sum(respipe2[1,3:20])) == -219595.0

  mpipeline3 = @pipeline csvfilter |> valgator |> valnner |> mono |> outliernicer |> stfier
  respipe3 = fit_transform!(mpipeline3)
  @test round(sum(respipe3[1,3:20])) == -213862.0
end
@testset "Outliernicer: readcsv |> valgator |> valnner |> mono |> outliernicer |> stfier" begin
  test_basicoutlier()
end

function test_typesoutliernicer()
  regularfile = joinpath(dirname(pathof(TSML)),"../data/typedetection/regular.csv")
  monofile = joinpath(dirname(pathof(TSML)),"../data/typedetection/monotonic.csv")
  dailymonofile = joinpath(dirname(pathof(TSML)),"../data/typedetection/dailymonotonic.csv")
  regularfilecsv = CSVDateValReader(Dict(:filename=>regularfile,:dateformat=>"dd/mm/yyyy HH:MM"))
  monofilecsv = CSVDateValReader(Dict(:filename=>monofile,:dateformat=>"dd/mm/yyyy HH:MM"))
  dailymonofilecsv = CSVDateValReader(Dict(:filename=>dailymonofile,:dateformat=>"dd/mm/yyyy HH:MM"))

  valgator = DateValgator(Dict(:dateinterval=>Dates.Hour(1)))
  valnner = DateValNNer(Dict(:dateinterval=>Dates.Hour(1)))
  stfier = Statifier(Dict(:processmissing=>true))
  mono = Monotonicer(Dict())
  outliernicer = Outliernicer(Dict(:dateinterval=>Dates.Hour(1)))

  regpipeline = @pipeline regularfilecsv |> valgator |> valnner |> mono |> outliernicer |> stfier
  fit!(regpipeline)
  regulardf=transform!(regpipeline)
  @test round(sum(regulardf[1,3:20])) == -61184.0

  monopipeline = @pipeline monofilecsv |> valgator |> valnner |> mono |> outliernicer |> stfier
  fit!(monopipeline)
  monodf=transform!(monopipeline)
  @test round(sum(monodf[1,3:20])) == -890049.0

  dailymonopipeline = @pipeline dailymonofilecsv |> valgator |> valnner |> mono |> outliernicer |> stfier
  fit!(dailymonopipeline)
  dailymonodf=transform!(dailymonopipeline)

  @test round(sum(dailymonodf[1,3:20])) == -294446.0

end
@testset "Outliernicer: monotonic type and outlier detections" begin
  test_typesoutliernicer()
end

end


#####
##### Building blocks for traversing a Hamiltonian deterministically, using the leapfrog
##### integrator.
#####

####
#### kinetic energy
####

"""
$(TYPEDEF)

Kinetic energy specifications. Implements the methods

- `Base.size`

- [`kinetic_energy`](@ref)

- [`calculate_p♯`](@ref)

- [`∇kinetic_energy`](@ref)

- [`rand_p`](@ref)

For all subtypes, it is implicitly assumed that kinetic energy is symmetric in
the momentum `p`,

```julia
kinetic_energy(κ, p, q) == kinetic_energy(κ, .-p, q)
```

When the above is violated, the consequences are undefined.
"""
abstract type KineticEnergy end

"""
$(TYPEDEF)

Euclidean kinetic energies (position independent).
"""
abstract type EuclideanKineticEnergy <: KineticEnergy end

"""
$(TYPEDEF)

Gaussian kinetic energy, with ``K(q,p) = p ∣ q ∼ 1/2 pᵀ⋅M⁻¹⋅p + log|M|`` (without constant),
which is independent of ``q``.

The inverse covariance ``M⁻¹`` is stored.

!!! note
    Making ``M⁻¹`` approximate the posterior variance is a reasonable starting point.
"""
struct GaussianKineticEnergy{T <: AbstractMatrix,
                             S <: AbstractMatrix} <: EuclideanKineticEnergy
    "M⁻¹"
    M⁻¹::T
    "W such that W*W'=M. Used for generating random draws."
    W::S
    function GaussianKineticEnergy(M⁻¹::T, W::S) where {T, S}
        @argcheck checksquare(M⁻¹) == checksquare(W)
        new{T,S}(M⁻¹, W)
    end
end

"""
$(SIGNATURES)

Gaussian kinetic energy with the given inverse covariance matrix `M⁻¹`.
"""
GaussianKineticEnergy(M⁻¹::AbstractMatrix) = GaussianKineticEnergy(M⁻¹, cholesky(inv(M⁻¹)).L)

"""
$(SIGNATURES)

Gaussian kinetic energy with the given inverse covariance matrix `M⁻¹`.
"""
GaussianKineticEnergy(M⁻¹::Diagonal) = GaussianKineticEnergy(M⁻¹, Diagonal(.√inv.(diag(M⁻¹))))

"""
$(SIGNATURES)

Gaussian kinetic energy with a diagonal inverse covariance matrix `M⁻¹=m⁻¹*I`.
"""
GaussianKineticEnergy(N::Integer, m⁻¹ = 1.0) = GaussianKineticEnergy(Diagonal(fill(m⁻¹, N)))

function Base.show(io::IO, κ::GaussianKineticEnergy{T}) where {T}
    print(io::IO, "Gaussian kinetic energy ($(nameof(T))), √diag(M⁻¹): $(.√(diag(κ.M⁻¹)))")
end

## NOTE about implementation: the 3 methods are callable without a third argument (`q`)
## because they are defined for Gaussian (Euclidean) kinetic energies.

Base.size(κ::GaussianKineticEnergy, args...) = size(κ.M⁻¹, args...)

"""
$(SIGNATURES)

Return kinetic energy `κ`, at momentum `p`.
"""
kinetic_energy(κ::GaussianKineticEnergy, p, q = nothing) = dot(p, κ.M⁻¹ * p) / 2

"""
$(SIGNATURES)

Return ``p♯ = M⁻¹⋅p``, used for turn diagnostics.
"""
calculate_p♯(κ::GaussianKineticEnergy, p, q = nothing) = κ.M⁻¹ * p

"""
$(SIGNATURES)

Calculate the gradient of the logarithm of kinetic energy in momentum `p`.
"""
∇kinetic_energy(κ::GaussianKineticEnergy, p, q = nothing) = calculate_p♯(κ, p)

"""
$(SIGNATURES)

Generate a random momentum from a kinetic energy at position `q`.
"""
rand_p(rng::AbstractRNG, κ::GaussianKineticEnergy, q = nothing) = κ.W * randn(rng, size(κ.W, 1))

####
#### Hamiltonian
####

struct Hamiltonian{K,L}
    "The kinetic energy specification."
    κ::K
    """
    The (log) density we are sampling from. Supports the `LogDensityProblem` API.
    Technically, it is the negative of the potential energy.
    """
    ℓ::L
    """
    $(SIGNATURES)

    Construct a Hamiltonian from the log density `ℓ`, and the kinetic energy specification
    `κ`. `ℓ` with a vector are expected to support the `LogDensityProblems` API, with
    gradients.
    """
    function Hamiltonian(κ::K, ℓ::L) where {K <: KineticEnergy,L}
        @argcheck capabilities(ℓ) ≥ LogDensityOrder(1)
        @argcheck dimension(ℓ) == size(κ, 1)
        new{K,L}(κ, ℓ)
    end
end

Base.show(io::IO, H::Hamiltonian) = print(io, "Hamiltonian with $(H.κ)")

"""
$(TYPEDEF)

A log density evaluated at position `q`. The log densities and gradient are saved, so that
they are not calculated twice for every leapfrog step (both as start- and endpoints).

Because of caching, a `EvaluatedLogDensity` should only be used with a specific Hamiltonian,
preferably constructed with the `evaluate_ℓ` constructor.

In composite types and arguments, `Q` is usually used for this type.
"""
struct EvaluatedLogDensity{T,S}
    "Position."
    q::T
    "ℓ(q). Saved for reuse in sampling."
    ℓq::S
    "∇ℓ(q). Cached for reuse in sampling."
    ∇ℓq::T
    function EvaluatedLogDensity(q::T, ℓq::S, ∇ℓq::T) where {T <: AbstractVector,S <: Real}
        @argcheck length(q) == length(∇ℓq)
        new{T,S}(q, ℓq, ∇ℓq)
    end
end

# general constructors below are necessary to sanitize input from eg Diagnostics, or an
# initial position given as integers, etc

function EvaluatedLogDensity(q::AbstractVector, ℓq::Real, ∇ℓq::AbstractVector)
    q, ∇ℓq = promote(q, ∇ℓq)
    EvaluatedLogDensity(q, ℓq, ∇ℓq)
end

EvaluatedLogDensity(q, ℓq::Real, ∇ℓq) = EvaluatedLogDensity(collect(q), ℓq, collect(∇ℓq))

"""
$(SIGNATURES)

Evaluate log density and gradient and save with the position. Preferred interface for
creating `EvaluatedLogDensity` instances.
"""
function evaluate_ℓ(ℓ, q)
    ℓq, ∇ℓq = logdensity_and_gradient(ℓ, q)
    if isfinite(ℓq)
        EvaluatedLogDensity(q, ℓq, ∇ℓq)
    else
        EvaluatedLogDensity(q, oftype(ℓq, -Inf), q) # second q used just as a placeholder
    end
end

"""
$(TYPEDEF)

A point in phase space, consists of a position (in the form of an evaluated log density `ℓ`
at `q`) and a momentum.
"""
struct PhasePoint{T <: EvaluatedLogDensity,S}
    "Evaluated log density."
    Q::T
    "Momentum."
    p::S
    function PhasePoint(Q::T, p::S) where {T,S}
        @argcheck length(p) == length(Q.q)
        new{T,S}(Q, p)
    end
end

"""
$(SIGNATURES)

Log density for Hamiltonian `H` at point `z`.

If `ℓ(q) == -Inf` (rejected), skips the kinetic energy calculation.

Non-finite values (incl `NaN`, `Inf`) are automatically converted to `-Inf`. This can happen
if

1. the log density is not a finite value,

2. the kinetic energy is not a finite value (which usually happens when `NaN` or `Inf` got
mixed in the leapfrog step, leading to an invalid position).
"""
function logdensity(H::Hamiltonian{<:EuclideanKineticEnergy}, z::PhasePoint)
    @unpack ℓq = z.Q
    isfinite(ℓq) || return oftype(ℓq, -Inf)
    K = kinetic_energy(H.κ, z.p)
    ℓq - (isfinite(K) ? K : oftype(K, Inf))
end

function calculate_p♯(H::Hamiltonian{<:EuclideanKineticEnergy}, z::PhasePoint)
    calculate_p♯(H.κ, z.p)
end

"""
    leapfrog(H, z, ϵ)

Take a leapfrog step of length `ϵ` from `z` along the Hamiltonian `H`.

Return the new phase point.

The leapfrog algorithm uses the gradient of the next position to evolve the momentum. If
this is not finite, the momentum won't be either, `logdensity` above will catch this and
return an `-Inf`, making the point divergent.
"""
function leapfrog(H::Hamiltonian{<: EuclideanKineticEnergy}, z::PhasePoint, ϵ)
    @unpack ℓ, κ = H
    @unpack p, Q = z
    @argcheck isfinite(Q.ℓq) "Internal error: leapfrog called from non-finite log density"
    pₘ = p + ϵ/2 * Q.∇ℓq
    q′ = Q.q + ϵ * ∇kinetic_energy(κ, pₘ)
    Q′ = evaluate_ℓ(H.ℓ, q′)
    p′ = pₘ + ϵ/2 * Q′.∇ℓq
    PhasePoint(Q′, p′)
end


Random.seed!(10)
for i = 1:20
    local n = rand(15:40)
    local m = rand(5:15)
    tol = randn()
    local A, B = gencontrollable(n,m,tol=tol)
    @test iscontrollable(A,B)
    @test isstabled(A) == (tol >= 0)
end

for i = 1:20
    local n = rand(15:30)
    local m = rand(5:15)
    tol = randn()
    local A, B = gencontrollable(n,m,:continuous,tol=tol,maxiter=500)
    @test iscontrollable(A,B)
    @test isstablec(A) == (tol >= 0)
end


struct DBFGS{T1, T2} <: QuasiNewton{T1}
   approx::T1
   theta::T2
end

DBFGS(approx) = DBFGS(;approx=approx)
DBFGS(;approx=Inverse(), theta=0.2) = DBFGS(approx, theta)

function update!(scheme::DBFGS{<:Direct, <:Any}, B, s, y)
   # We could write this as
   #     B .+= (y*y')/dot(s, y) - (B*s)*(s'*B)/(s'*B*s)
   #     B .+= (y*y')/dot(s, y) - b*b'/dot(s, b)
   # where b = B*s
   # But instead, we split up the calculations. First calculate the denominator
   # in the first term
   σ = dot(s, y)
   ρ = inv(σ) # scalar
   # Then calculate the vector b
   b = B*s # vector temporary
   sb = dot(s, b)
   if σ ≥ scheme.theta*sb
      θ = 1.0
      # Calculate one vector divided by dot(s, b)
      ρbb = inv(sb)*b
      # And calculate
      B .+= (inv(σ)*y)*y' .- ρbb*b'
   else
      θ = 0.8*sb/(sb-σ)
      r = y*θ + (1-θ)*b
      # Calculate one vector divided by dot(s, b)
      ρbb = inv(dot(s, b))*b
      # And calculate
      B .+= (inv(dot(s, r))*r)*r' .- ρbb*b'
   end
end
function update(scheme::DBFGS{<:Direct, <:Any}, B, s, y)
   # As above, but out of place

   σ = dot(s, y)
   b = B*s
   sb = dot(s, b)
   if σ ≥ scheme.theta*sb
      θ = 1.0
      # Calculate one vector divided by dot(s, b)
      ρbb = inv(sb)*b
      # And calculate
      return B .+= (inv(σ)*y)*y' .- ρbb*b'
   else
      θ = 0.8*sb/(sb-σ)
      r = y*θ + (1-θ)*b
      # Calculate one vector divided by dot(s, b)
      ρbb = inv(dot(s, b))*b
      # And calculate
      return B .+= (inv(dot(s, r))*r)*r' .- ρbb*b'
   end
end
function update(scheme::DBFGS{<:Inverse, <:Any}, H, s, y)
   σ = dot(s, y)
   ρ = inv(σ)
   #   if isfinite(ρ)
   C = (I - ρ*s*y')
   H = C*H*C' + ρ*s*s'
   #   end
   H
end
function update!(scheme::DBFGS{<:Inverse, <:Any}, H, s, y)
   σ = dot(s, y)
   ρ = inv(σ)

   if isfinite(ρ)
      Hy = H*y
      H .= H .+ ((σ+y'*Hy).*ρ^2)*(s*s')
      Hys = Hy*s'
      Hys .= Hys .+ Hys'
      H .= H .- Hys.*ρ
   end
   H
end

function update!(scheme::DBFGS{<:Inverse, <:Any}, A::UniformScaling, s, y)
   update(scheme, A, s, y)
end
function update!(scheme::DBFGS{<:Direct, <:Any}, A::UniformScaling, s, y)
   update(scheme, A, s, y)
end


# ------------------------------------------------------------------------------------------
# # Plotting / Gráficas
#
# ## Introducción
#
# Hay muchas manera de graficar en Julia (además de usar PyPlot).

#
# Aquí veremos como usar `Plots.jl`.
# ------------------------------------------------------------------------------------------

#Pkg.add("Plots")
using Plots

# ------------------------------------------------------------------------------------------
# Una de las ventajas de `Plots.jl` Es que permite cambiar los backends sin costo alguno. En
# este notebook, vamos a intentar usar `gr()` y `plotlyjs()` como backends.

#
# Primero vamos a generar datos artificiales para graficar
# ------------------------------------------------------------------------------------------

x = -3:0.1:3
f(x) = x^2

y = f.(x)

# ------------------------------------------------------------------------------------------
# **Carguemos el backend de GR**
# ------------------------------------------------------------------------------------------

gr()

plot(x, y, label="linea")
scatter!(x, y, label="puntos")

# ------------------------------------------------------------------------------------------
# El `!` al final de `scatter!` indica que sea una función mutante,  indicando que los
# puntos se van a agregar a la gráfica preexistente.
#
# Por contraste, en vez de usar `scatter!`, usa `scatter` para ver como funciona.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# **Sin cambiar de sintaxis, cambiamos al backend de `plotlyjs()`**
# ------------------------------------------------------------------------------------------

plotlyjs()

plot(x, y, label="line")
scatter!(x, y, label="points")

# ------------------------------------------------------------------------------------------
# Y nos fijamos como cambia la primera gráfica de la segunda
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Subiendo de nivel
#
# La sintaxis para agregar títulos y líneas es bastante directa.
#
# Ahora, en el nombre de la ciencia, vamos a examinar la relación entre la temperatura
# global y el número de piratas entre 1860 y 2000.
# ------------------------------------------------------------------------------------------

globaltemperatures = [14.4, 14.5, 14.8, 15.2, 15.5, 15.8]
numpirates = [45000, 20000, 15000, 5000, 400, 17]

# Primero plotteamos los datos
plot(numpirates, globaltemperatures, legend=false)
scatter!(numpirates, globaltemperatures, legend=false)

# Agregamos Títulos y etiquetas/labels.
xlabel!("Número de piratas [Apróx.]")
ylabel!("Temperatura Global (C)")
title!("Influencia de población de piratas en calentamiento global")

# Primero graficamos datos
plot(numpirates, globaltemperatures, legend=false)
scatter!(numpirates, globaltemperatures, legend=false)

# Este comando invierte el eje x para que podamos ver los cambios hacia adelante en el tiempo, de 1860 a 2000
xflip!()

# Add titles and labels
xlabel!("Number of Pirates [Approximate]")
ylabel!("Global Temperature (C)")
title!("Influence of pirate population on global warming")

# ------------------------------------------------------------------------------------------
# Para crear una gráfica con subráficas, sólo nombramos a cada una de las subgráficas y las
# ponemos junto con la especificación de diseño en una sóla llamada a `plot`.
# ------------------------------------------------------------------------------------------

p1 = plot(x, x)
p2 = plot(x, x.^2)
p3 = plot(x, x.^3)
p4 = plot(x, x.^4)
plot(p1,p2,p3,p4,layout=(2,2),legend=false)

# ------------------------------------------------------------------------------------------
# ### Ejercicios
#
# 8.1 Grafica y vs x para `y = x^2` usando el backend de PyPlot.
# ------------------------------------------------------------------------------------------


# recipes
@recipe function f(pt::Point2D)
    seriestype := :scatter
    label --> :none
    [(pt.x, pt.y)]
end

@recipe function f(pt::Point3D)
    seriestype := :scatter
    label --> :none
    [pt.x], [pt.y], [pt.z]
end

@recipe function f(c::Circle2D)
    seriestype := :shape
    linecolor --> :black
    aspect_ratio --> :equal
    fillalpha --> 0.25
    label --> :none
    θ = 0:0.1:2π
    c.pt.x .+ c.r .* cos.(θ), c.pt.y .+ c.r .* sin.(θ)
end

@recipe function f(particle::Particle2D)
    particle.c
end

@recipe function f(c::Circle3D)
    seriestype := :scatter
    label --> :none
    markersize --> 10c.r
    [c.pt.x], [c.pt.y], [c.pt.z]
end

@recipe function f(geomobjs::AbstractArray{<:ConcreteGeometricTypes})
    for obj in geomobjs
        @series begin
            obj
        end
    end
end

@recipe f(geomobjs::T...) where {T<:ConcreteGeometricTypes} = collect(geomobjs)


using SDE
using Base.Test
require(Pkg.dir("SDE","src", "Randm.jl"))
#include(Pkg.dir("SDE","src", "quad.jl"))
import SDE.eps2

function mc(M::Vector)
  m = mean(M)
  ste = std(M)/sqrt(length(M))
  m, ste
end
qu = sqrt(2)*erfinv(2*0.995-1)

macro repeat(ex, n)
    quote
       r1 = $ex
       y = zeros(size(r1)..., $n)
       s = prod(size(r1))
       y[1:s] = r1
       for i in 1:$n-1
        y[i*s+1:(i+1)*s] = $ex
       end
       y
    end
end

for d in [1:3]
    N = 5000
    xi = 0.3
	global T, t, v, h, P, B, A, beta, s
	println("Test dimension $d xi $xi")

	if (d== 1) srand(7)
	elseif (d== 2) srand(8)
	else srand(6)
	end
	B = xi.*Randm.randstable(d)
#	A = Randm.randposdef(d)
	Si = rand(d,d)
	A = Si*Si'
	beta = 0.5*randn(d)

	P = MvLinPro(B, beta, Si)
	P0 = MvAffPro(beta, Si)
	P00 = MvLinPro(eye(d)*1E-7, beta, Si)

	function aa(s)
	 A
	end

	T = 1.7
	t = 0.5
	v = zeros(d)

	println("Test SDE.H, SDE.K")

    H1 = SDE.H(t,T, P)
    H2 = SDE.varH(t, T, B, aa)
    K = SDE.K(t, T, P)
    H3 = expm((T-t)*B')*inv(K)*expm((T-t)*B)

    H0 = SDE.H(t, T, P0)
    K0 = SDE.K(t, T, P0)
    H00 = SDE.H(t, T, P00)
    @test norm(H0 - H00) < 1E-5*d^2

	h = T-t
	phi = expm(h*B)
	phim = expm(-h*B)
	K2 = P.lambda - phi*P.lambda*phi'
	H4 = phi'*inv(P.lambda - phi*P.lambda*phi')*phi

	@test norm(H1 - H2) < 1E-10
	@test norm(H1 - H3) < 1E-10
	@test norm(H1 - H4) < 1E-10

	@test norm(K-K2) < d*24*eps()

	function varr3(t, x, T, v, P)
	    h = T - t
		binv = inv(P.B)
		phi = expm(h*P.B)
		phim = expm(-h*P.B)
		vt =  phim*( v) + phim * binv * beta -binv*beta
		SDE.H(t, T, P) * (vt-x )

	end

	println("Test SDE.r")

	srand(5)
	x0 = x = randn(d)/d
	v = randn(d)
	r1 = SDE.r(t, x, T, v, P)
	r2 = SDE.varr(t, x, T, v, P)
	@test norm(r1 - r2) < 1E-10

	r3 = varr3(t, x, T, v, P)
	@test norm(r1 - r3) < 1E-10

	r10 = SDE.r(t, x, T, v, P0)
	r100 = SDE.r(t, x, T, v, P00)
	@test norm(r10 - r100) < 1E-5*d^2

	println("Test SDE.mu")

	mu = SDE.mu(t, x, T, P)
	mu0 = SDE.mu(t, x, T, P0)
	mu2 = SDE.varmu(t, x, T, P)
	@test norm(mu - mu2) < 1E-13

    println("Testing mean and variance of samplep")

    y = @repeat(SDE.samplep(t, x, T, P), N)
    mcy = mc(sum(y,1)[:])
	println(repr(mcy), sum(mu))
    @test abs(mcy[1]- sum(mu)) < qu*mcy[2]
    muhat = mean(y,2)
	Khat = cov((y.-x)')
	println(norm(Khat - K))
	@test norm(Khat-K) < 0.04*d

    y = @repeat(SDE.samplep(t, x, T, P0), N)
    mcy = mc(sum(y,1)[:])
	println(repr(mcy), sum(mu0))
    @test abs(mcy[1]- sum(mu0)) < qu*mcy[2]
    muhat = mean(y,2)
	Khat = cov((y.-x)')
	println(norm(Khat - K0))
	@test norm(Khat-K0) < 0.04*d

	#checks if transition density p integrates to 1 by monte carlo integration
	println("Test SDE.lp")

	function testp(h, x0, N, B, beta, Si)
	    local P, P0
		mu =  B*x0+beta
		P = MvLinPro(B, beta, Si)
		P0 = MvAffPro(mu, Si)
    	# monte carlo integration of the transition density p with proposals according to p0
		w1 = zeros(N)
		for n in 1:N
			y = SDE.samplep(0., x0, h, P0 )
			w1[n] =  exp(SDE.lp(0., x0, h,  y, P)-SDE.lp(0., x0, h, y, P0))
		end
		# monte carlo integration of the transition density p0 with proposals according to p
	    w2 = zeros(N)
		for n in 1:N
			y = SDE.samplep(0., x0, h, P )
			w2[n] =  exp(SDE.lp(0., x0, h,  y, P0)-SDE.lp(0., x0, h, y, P))
		end

		return mc(w1), mc(w2), w1, w2
	end

	srand(5)
	n1, n2, w1, w2 = testp(0.05, x, N, B, beta, Si)
	println(repr(n1), repr(n2))
	@test abs(n1[1]-1.) < qu*n1[2]
	@test abs(n2[1]-1.) < qu*n2[2]

	@test abs(SDE.lp(t, x,T, v, P) - SDE.varlp(t, x,T, v, (t,s) -> t-s, B, beta, s->A)) < 1E-10

    tmin = t/2
    tmax = T
    T = tmax-tmin

	s = SDE.soft(t, tmin, T)
	@test norm(t - SDE.tofs(s, tmin, T)) < 1E-13
	@test norm(t - (tmax - (T - s)^2/T)) < 1E-13

	v1 = SDE.Vs(s, T, v, P0)
	v2 = SDE.V(t, tmax, v, P0)
#    println("v10, v20", repr(v1), repr(v2))
	@test norm(v1-v2) < d^2*1E-13

	v1 = SDE.Vs(s, T, v, P)
	v2 = SDE.V(t, tmax, v, P)
#    println("v1, v2", repr(v1), repr(v2))
 	@test norm(v1-v2) < d^2*1E-13

    v1 = SDE.dotVs(s, T, v, P0)
    v2 = SDE.dotVs(s, T, v, P00)
    println ("dotv ", repr(v1) ,repr(v2))
    @test norm(v1 - v2) < 1E-5

	us =  SDE.uofx(s, x, T, v,  P)
	xt =  SDE.xofu(s, us, T, v,  P)
	@test norm(x - xt) < 1E-10

	r1 =  SDE.r(t, x, tmax, v, P)
 	r2 =  SDE.J(s, T, P) *us*T/(T-s)
	@test norm(r1 .- r2) < 1E-10

end


export batchqromp

# Write a batch version of the QR-OMP code

# Should not be considered production-ready

function batchqromp(ψ::AbstractMatrix{T}, u::AbstractVector{T}; invert::Bool=true, verbose::Bool = true, ϵrel::Float64 = 1e-1, batchterms::Int64 = 1, maxterms::Int64=typemax(Int64)) where {T}
    m, n = size(ψ)

    residue = copy(u)
    idxset = Int64[]
    dict = collect(1:n)
    ϵu = norm(u)
    ϵhist = [ϵu]
    ϵrel *= ϵu
    ϕ = zeros(T, n)
    ratio = zeros(T,n)
    cache = zeros(T, m, batchterms)

    @inbounds for j=1:n
        for i=1:m
        ϕ[j] += ψ[i,j]^2
        end
    end

    cache = zeros(T, m, batchterms)
    F = elasticqrfact(cache)
    ek = zeros(T, n, batchterms)
    q = zeros(T, m, batchterms)
    new_idx = zeros(Int64, batchterms)

    @inbounds for k=1:div(n, batchterms)
        # Step  7 choose best candidate

        @fastmath @inbounds for (j, idx) in enumerate(dict)
            ψj = view(ψ,:,idx)
            ratio[idx] = dot(residue, ψj)^2/ϕ[idx]

            # if ratio >= entry
            #     entry = copy(ratio)
            #     new_idx = copy(idx)
            # end
        end
        new_idx .= partialsortperm(view(ratio, dict), 1:batchterms, rev = true)

        # Step 8 Update set of selected basis
        append!(idxset, new_idx)

        # Step 9 Update candidate dictionary
        filter!(x-> x != new_idx, dict)

        # Step 10 Update Q(k-1) and R(k-1) with ψi(k)
        if k>1
            # F = updateqrfactUnblocked!(F, view(ψ,:,new_idx))
            @time updateelasticqrfact!(F, view(ψ,:,new_idx), cache)
            @inbounds for i=1:batchterms
                ek[(k-2)*batchterms+i,i] = 0.0
            end
        else
            # F = qrfactUnblocked(ψ[:,new_idx:new_idx])
            F = elasticqrfact(view(ψ,:,new_idx))
        end

        # Step 11 Update residual
        @inbounds for i=1:batchterms
            ek[(k-1)*batchterms+i,i] = 1.0
        end

        mul!(q, F.Q, view(ek,1:k,1:batchterms))

        @inbounds for (j, idx) in enumerate(dict)
            ψj = view(ψ,:,idx)
            ϕj = dot(q, ψj)
            ϕ[idx] -= ϕj^2
        end
        factor = dot(q, u)
        residue -= factor*q

        # Step 12 Calculate stopping critera
        ϵ = norm(residue)

        if verbose == true
            push!(ϵhist, copy(ϵ))
        end
        if ϵ < ϵrel || k == maxterms #|| size(idxset,1) + batchterms > n
            break
        end
    end

    # Solve the system with the set of indices
    if verbose == true
        if invert == true
            c = zeros(length(idxset))
            ldiv!(c, F, u)
            return idxset, c, ϵhist
        else
            return idxset, ϵhist
        end
    else
        if invert == true
            c = zeros(length(idxset))
            ldiv!(c, F, u)
            return idxset, c
        else
            return idxset
        end
    end
end


"""
```julia
    Conjugate( ϵ::T, scaled::Bool = true ) where {T <: Number}
```

Terminate the Fast Gradient Method when the conjugate residual value is below `ϵ`.
The value of `ϵ` is divided by the number of variables in the problem when `scaled`
is true, producing a scaled termination criteria.

This criteria is based on the results in §6.4 of [1].

[1] S. Richter, ‘Computational complexity certification of gradient methods for real-time model predictive control’, Doctoral Thesis, ETH Zurich, 2012.

See also: [`Gradient`](@ref) and [`Best`](@ref)
"""
mutable struct Conjugate{T} <: AbstractStopCondition where {T <: Number}
    ϵ::T

    scale::Bool   # Should a scaled termination criteria be used
    n::Int        # Number of variables in the problem
    ϵₛ::T         # Scaled termination criteria

    function Conjugate( ϵ::T; scaled::Bool = true ) where {T <: Number}
        return new{T}( ϵ, scaled, 1, ϵ )
    end
end

function init!( cond::Conjugate, n, L, μ )
    cond.n = n

    if( cond.scale )
        cond.ϵₛ = cond.ϵ * 1/n
    else
        cond.ϵₛ = cond.ϵ
    end
end

function (cond::Conjugate)( state::FGMState )
    return compute_conjugate( state ) < cond.ϵₛ
end

function compute_conjugate( state::FGMState )
    return abs( state.xₙ'*state.∇xₙ + norm( state.∇xₙ, 1 ) )
end

"""
```julia
    Gradient( ϵ::T, scaled::Bool = true ) where {T <: Number}
```

Terminate the Fast Gradient Method when the gradient value is below `ϵ`.
The value of `ϵ` is divided by the number of variables in the problem when `scaled`
is true, producing a scaled termination criteria.

This criteria is based on the results in §6.4 of [1].

[1] S. Richter, ‘Computational complexity certification of gradient methods for real-time model predictive control’, Doctoral Thesis, ETH Zurich, 2012.

See also: [`Conjugate`](@ref) and [`Best`](@ref)
"""
mutable struct Gradient{T} <: AbstractStopCondition where {T <: Number}
    ϵ::T
    L::T
    μ::T
    c::T

    scale::Bool   # Should a scaled termination criteria be used
    n::Int        # Number of variables in the problem
    ϵₛ::T         # Scaled termination criteria

    function Gradient( ϵ::T; scaled::Bool = true ) where {T <: Number}
        return new{T}( ϵ, scaled, 0.0, 0.0, scaled, 1, ϵ )
    end
end

function init!( cond::Gradient, n, L, μ )
    cond.μ = μ     #
    cond.L = L     #
    cond.c = 0.5 * ( 1/cond.μ - 1/cond.L )  # Coefficient of the computation

    cond.n = n

    if( cond.scale )
        cond.ϵₛ = cond.ϵ * 1/n
    else
        cond.ϵₛ = cond.ϵ
    end
end

function (cond::Gradient)( state::FGMState )
    return abs( cond.c * norm( cond.L* ( state.y - state.xₙ ) )^2 ) < cond.ϵₛ
end

"""
```julia
    Best( ϵ₁::T, ϵ₂:T = ϵ₁; scaled::Bool = true ) where {T <: Number}
```

Terminate the Fast Gradient Method when either the gradient value is below `ϵ₁` or the
conjugate residual value is below `ϵ₂`. `ϵ₁` and `ϵ₂` are divided by the number of variables
in the problem to produce a scaled termination criteria when `scaled` is true.

This criteria is based on the results in §6.4 of [1].

[1] S. Richter, ‘Computational complexity certification of gradient methods for real-time model predictive control’, Doctoral Thesis, ETH Zurich, 2012.

See also: [`Gradient`](@ref) and [`Conjugate`](@ref)
"""
struct Best <: AbstractStopCondition
    grad::Gradient
    conj::Conjugate

    function Best( ϵ₁::T, ϵ₂::T = ϵ₁; scaled::Bool = true ) where {T <: Number}
        return new( Gradient( ϵ₁; scaled = scaled ), Conjugate( ϵ₂; scaled = scaled ) )
    end
end

function init!( cond::Best, n, L, μ )
    init!( cond.grad, n, L, μ )
    init!( cond.conj, n, L, μ )
end

function (cond::Best)( state::FGMState )
    return cond.grad( state ) || cond.conj( state )
end


using CSDP
const CSDPOPT = with_optimizer(CSDP.Optimizer, printlevel = 0)

@testset "CSDP SDP triangle Problems" begin
    list_of_sdp_triang_problems = [
        # sdpt1_test, # CSDP is returning SLOW_PROGRESS
        # sdpt2_test, # CSDP is returning SLOW_PROGRESS
        sdpt3_test,
        sdpt4_test
    ]
    test_strong_duality(list_of_sdp_triang_problems, CSDPOPT)
end


# ==========================================================================
# Unprocessed Solution for the dual problem

immutable DualSolution
  minimizer::Vector{Float64}
  bias::Float64
  minimum::Float64
  iterations::Int
  isconverged::Bool
end

function Base.show(io::IO, r::DualSolution)
  println(io, typeof(r))
  _printvariable(io, 14, "minimizer()", minimizer(r))
  _printvariable(io, 14, "bias()", bias(r))
  _printvariable(io, 14, "minimum()", minimum(r))
  _printvariable(io, 14, "iterations()", iterations(r))
  _printvariable(io, 14, "isconverged()", isconverged(r), newline = false)
end

coef(s::DualSolution) = s.minimizer
minimizer(s::DualSolution) = s.minimizer
Base.minimum(s::DualSolution) = s.minimum
bias(s::DualSolution) = s.bias
iterations(s::DualSolution) = s.iterations
isconverged(s::DualSolution) = s.isconverged


using
  PyPlot,
  JLD2,
  Printf,
  FourierFlows

using FFTW: ifft

import GeophysicalFlows.BarotropicQG
import GeophysicalFlows.BarotropicQG: energy, meanenergy, enstrophy, meanenstrophy

# Numerical parameters and time-stepping parameters
nx  = 512      # 2D resolution = nx^2
stepper = "FilteredETDRK4"   # timestepper
dt  = 2e-2     # timestep
nsteps = 20000 # total number of time-steps
nsubs  = 500   # number of time-steps for plotting
               # (nsteps must be multiple of nsubs)

# Physical parameters
Lx  = 2π       # domain size
nu  = 8.0e-10  # viscosity
nnu = 2        # viscosity order
f0  = -1.0     # Coriolis parameter
beta = 1.4015  # the y-gradient of planetary PV
mu   = 1.0e-2  # linear drag
   F = 0.0012  # normalized wind stress forcing on domain-averaged
               # zonal flow U(t) flow

# Topographic PV
topoPV(x, y) = @. 2*cos(10x)*cos(10y)

# Forcing on the domain-averaged U equation
calcFU(t) = F

# Initialize problem
prob = BarotropicQG.ForcedProblem(nx=nx, Lx=Lx, f0=f0, beta=beta, eta=topoPV,
                  calcFU=calcFU, nu=nu, nnu=nnu, mu=mu, dt=dt, stepper=stepper)
sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

x, y = gridpoints(g)

# Files
filepath = "."
plotpath = "./plots_acctopo"
plotname = "snapshots"
filename = joinpath(filepath, "acctopo.jl.jld2")

# File management
if isfile(filename); rm(filename); end
if !isdir(plotpath); mkdir(plotpath); end

# Initialize with zeros
BarotropicQG.set_zeta!(prob, 0*x)

# Create Diagnostics
E = Diagnostic(energy, prob; nsteps=nsteps)
Q = Diagnostic(enstrophy, prob; nsteps=nsteps)
Emean = Diagnostic(meanenergy, prob; nsteps=nsteps)
Qmean = Diagnostic(meanenergy, prob; nsteps=nsteps)
diags = [E, Emean, Q, Qmean]

# Create Output
get_sol(prob) = sol # extracts the Fourier-transformed solution
get_u(prob) = irfft(im*g.lr.*g.invKrsq.*sol, g.nx)
out = Output(prob, filename, (:sol, get_sol), (:u, get_u))

function plot_output(prob, fig, axs; drawcolorbar=false)

  # Plot the PV field and the evolution of energy and enstrophy.

  sol, v, p, g = prob.sol, prob.vars, prob.params, prob.grid
  BarotropicQG.updatevars!(prob)

  sca(axs[1])
  pcolormesh(x, y, v.q)
  axis("square")
  xlim(0, 2)
  xticks(0:0.5:2)
  ylim(0, 2)
  yticks(0:0.5:2)
  title(L"$\nabla^2\psi + \eta$ (part of the domain)")
  if drawcolorbar==true
    colorbar()
  end

  sca(axs[2])
  cla()
  plot(mu*E.t[1:E.i], E.data[1:E.i], label=L"$E_{\psi}$")
  plot(mu*E.t[1:Emean.i], Emean.data[1:Emean.i], label=L"$E_U$")

  xlabel(L"\mu t")
  ylabel(L"E")
  legend()

  sca(axs[3])
  cla()
  plot(mu*Q.t[1:Q.i], Q.data[1:Q.i], label=L"$Q_{\psi}$")
  plot(mu*Qmean.t[1:Qmean.i], Qmean.data[1:Qmean.i], label=L"$Q_U$")
  xlabel(L"\mu t")
  ylabel(L"Q")
  legend()
  pause(0.001)
end

fig, axs = subplots(ncols=3, nrows=1, figsize=(15, 4))
plot_output(prob, fig, axs; drawcolorbar=true)

# Step forward
startwalltime = time()

while cl.step < nsteps
  stepforward!(prob, diags, nsubs)

  # Message
  log = @sprintf("step: %04d, t: %d, E: %.4f, Q: %.4f, τ: %.2f min",
    cl.step, cl.t, E.data[E.i], Q.data[Q.i],
    (time()-startwalltime)/60)

  println(log)

  plot_output(prob, fig, axs; drawcolorbar=false)

end
println((time()-startwalltime))

plot_output(prob, fig, axs; drawcolorbar=false)

savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
savefig(savename, dpi=240)


using BackwardsLinalg
using Random
using Test

@testset "symeigen real" begin
    A = randn(4,4)
    A = A+A'
    op = randn(4, 4)
    op += op'
    function f(A)
        E, U = symeigen(A)
        E |> sum
    end
    function g(A)
        E, U = symeigen(A)
        v = U[:,1]
        (v'*op*v)[]|>real
    end
    @test gradient_check(f, A)
    @test gradient_check(g, A)
end

@testset "symeigen complex" begin
    Random.seed!(6)
    A = randn(ComplexF64, 4,4)
    A = A+A'
    op = randn(ComplexF64, 4, 4)
    op += op'
    function f(A)
        E, U = symeigen(A)
        E |> sum
    end
    function g(A)
        E, U = symeigen(A)
        v = U[:,1]
        (v'*op*v)[]|>real
    end
    @test gradient_check(f, A)
    @test gradient_check(g, A)
end


using Yao
using Yao.Intrinsics
################## Benchmarks #################
using BenchmarkTools
v = collect(0:1<<16-1)
orders = randperm(16)
#bres = @benchmark  reorder($v, $orders)
#@code_warntype  collect((reordered_basis(16, orders)))
#bres = @benchmark  collect((reordered_basis(16, orders)))
#bres = @benchmark  collect((reordered_basis(16, orders)))
Pm = pmrand(1<<16)
#@benchmark $Pm |> invorder
Dv = Diagonal(v)
@benchmark $Dv |> invorder