module Functionals

export J_T_ss, J_T_sm, J_T_re
export gate_functional, make_gate_chi
export make_gradient, make_chi

import ..WeightedObjective

using LinearAlgebra
using Zygote: Zygote
using FiniteDifferences: FiniteDifferences

@doc raw"""
Average complex overlap of the target states with forward-propagated states.

```julia
f_tau(ϕ, objectives; τ=nothing)
```

calculates

```math
f_τ = \frac{1}{N} \sum_{k=1}^{N} w_k τ_k
```

with

```math
τ_k = ⟨ϕ_k^\tgt|ϕ_k(T)⟩
```

in Hilbert space, or

```math
τ_k = \tr[ρ̂_k^{\tgt\,\dagger} ρ̂_k(T)]
```

in Liouville space, where ``|ϕ_k⟩`` or ``ρ̂_k`` are the elements
of `ϕ`, and ``|ϕ_k^\tgt⟩`` or ``ρ̂_k^\tgt`` are the
target states from the `target_state` field of the `objectives`. If `τ` is
given as a keyword argument, it must contain the values `τ_k` according to the
above definition. Otherwise, the ``τ_k`` values will be calculated internally.

``N`` is the number of objectives, and ``w_k`` is an optional weight for each
objective. For any objective that has a `weight` attribute (cf.
[`WeightedObjective`](@ref)), the ``w_k`` is taken from that attribute;
otherwise, ``w_k = 1``. The weights, if present, are not automatically
normalized, they are assumed to have values such that the resulting ``f_τ``
lies in the unit circle of the complex plane. Usually, this means that the
weights should sum to ``N``.
"""
function f_tau(ϕ, objectives; τ=nothing)
    N = length(objectives)
    if τ === nothing
        τ = [dot(objectives[k].target_state, ϕ[k]) for k = 1:N]
    end
    f::ComplexF64 = 0
    for k = 1:N
        obj = objectives[k]
        w = isa(obj, WeightedObjective) ? obj.weight : 1.0
        f += w * τ[k]
    end
    return f / N
end

@doc raw"""State-to-state phase-insensitive fidelity.

```julia
F_ss(ϕ, objectives; τ=nothing)
```

calculates

```math
F_{\text{ss}} = \frac{1}{N} \sum_{k=1}^{N} w_k |τ_k|^2 \quad\in [0, 1]
```

with ``N``, ``w_k`` and ``τ_k`` as in [`f_tau`](@ref).
"""
function F_ss(ϕ, objectives; τ=nothing)
    N = length(objectives)
    if τ === nothing
        τ = [dot(objectives[k].target_state, ϕ[k]) for k = 1:N]
    end
    F::ComplexF64 = f_tau(ϕ, objectives; τ=abs.(τ) .^ 2)
    @assert imag(F) < 1e-10
    return real(F)
end

@doc raw"""State-to-state phase-insensitive functional.

```julia
J_T_ss(ϕ, objectives; τ=nothing)
```

calculates

```math
J_{T,\text{ss}} = 1 - F_{\text{ss}} \in [0, 1].
```

All arguments are passed to [`F_ss`](@ref).
"""
function J_T_ss(ϕ, objectives; τ=nothing)
    return 1.0 - F_ss(ϕ, objectives; τ=τ)
end

@doc raw"""Krotov-states χ for functional [`J_T_ss`](@ref).

```julia
chi_ss!(χ, ϕ, objectives; τ=nothing)
```

sets the elements of `χ` according to

```math
|χ_k⟩
= -\frac{∂ J_{T,\text{ss}}}{∂ ⟨ϕ_k(T)|}
= \frac{1}{N} w_k τ_k |ϕ^{\tgt}_k⟩
```

with ``|ϕ^{\tgt}_k⟩``, ``τ_k`` and ``w_k`` as defined in [`f_tau`](@ref).

Note: this function can be obtained with `make_chi(J_T_ss, objectives)`.
"""
function chi_ss!(χ, ϕ, objectives; τ=nothing)
    N = length(objectives)
    if τ === nothing
        τ = [dot(objectives[k].target_state, ϕ[k]) for k = 1:N]
    end
    for k = 1:N
        obj = objectives[k]
        ϕₖ_tgt = obj.target_state
        copyto!(χ[k], ϕₖ_tgt)
        w = isa(obj, WeightedObjective) ? obj.weight : 1.0
        lmul!((τ[k] * w) / N, χ[k])
    end
end

@doc raw"""Square-modulus fidelity.

```julia
F_sm(ϕ, objectives; τ=nothing)
```

calculates

```math
F_{\text{sm}} = |f_τ|^2  \quad\in [0, 1].
```

All arguments are passed to [`f_tau`](@ref) to evaluate ``f_τ``.
"""
function F_sm(ϕ, objectives; τ=nothing)
    return abs(f_tau(ϕ, objectives; τ=τ))^2
end

@doc raw"""Square-modulus functional.

```julia
J_T_sm(ϕ, objectives; τ=nothing)
```

calculates

```math
J_{T,\text{sm}} = 1 - F_{\text{sm}} \quad\in [0, 1].
```

All arguments are passed to [`f_tau`](@ref) while evaluating ``F_{\text{sm}}``
in [`F_sm`](@ref).
"""
function J_T_sm(ϕ, objectives; τ=nothing)
    return 1.0 - F_sm(ϕ, objectives; τ=τ)
end

@doc raw"""Gradient for [`J_T_sm`](@ref).

```julia
grad_J_T_sm!(∇J_T, τ, ∇τ)
```

analytically sets the elements of the gradient `∇J_T` according to

```math
\frac{∂ J_{T,\text{sm}}(\{τ_k\})}{\partial ϵ_{ln}}
= \frac{1}{N^2} \sum_{k=1}^N \sum_{k'=1}^N\left[
        \frac{\partial τ_{k'}^*}{∂ϵ_{ln}} τ_k +
        τ_{k'}^* \frac{\partial τ_k^*}{∂ϵ_{ln}}
   \right]
= -\frac{2}{N} \Re \sum_{k=1}^N \sum_{k'=1}^N
  τ_{k'}^* \frac{∂τ_k}{\partial ϵ_{ln}}
```

with all quantities as defined in [`make_gradient`](@ref).

Note: this function can be obtained with
`make_gradient(J_T_sm, objectives, via=:tau)`.
"""
function grad_J_T_sm!(∇J_T, τ, ∇τ)
    N = length(τ) # number of objectives
    L, N_T = size(∇τ[1])  # number of controls/time intervals
    ∇J_T′ = reshape(∇J_T, L, N_T)  # writing to ∇J_T′ modifies ∇J_T
    for l = 1:L
        for n = 1:N_T
            ∇J_T′[l, n] = real(sum([conj(τ[k′]) * ∇τ[k][l, n] for k′ = 1:N for k = 1:N]))
        end
    end
    lmul!(-2 / N^2, ∇J_T)
    return ∇J_T
end

@doc raw"""Krotov-states χ for functional [`J_T_sm`](@ref).

```julia
chi_sm!(χ, ϕ, objectives; τ=nothing)
```

sets the elements of `χ` according to

```math
|χ_k⟩
= -\frac{\partial J_{T,\text{sm}}}{\partial ⟨ϕ_k(T)|}
= \frac{1}{N^2} w_k \sum_{j}^{N} w_j τ_j |ϕ_k^{\tgt}⟩
```

with ``|ϕ^{\tgt}_k⟩``, ``τ_j`` and ``w_k`` as defined in [`f_tau`](@ref).

Note: this function can be obtained with `make_chi(J_T_sm, objectives)`.
"""
function chi_sm!(χ, ϕ, objectives; τ=nothing)

    N = length(objectives)
    if τ === nothing
        τ = [dot(objectives[k].target_state, ϕ[k]) for k = 1:N]
    end

    w = ones(N)
    for k = 1:N
        obj = objectives[k]
        w[k] = isa(obj, WeightedObjective) ? obj.weight : 1.0
    end

    for k = 1:N
        obj = objectives[k]
        ϕₖ_tgt = obj.target_state
        copyto!(χ[k], ϕₖ_tgt)
        lmul!(w[k] * sum(w .* τ) / N^2, χ[k])
    end

end

@doc raw"""Real-part fidelity.

```julia
F_re(ϕ, objectives; τ=nothing)
```

calculates

```math
F_{\text{re}} = \Re[f_{τ}] \quad\in \begin{cases}
    [-1, 1] & \text{in Hilbert space} \\
    [0, 1] & \text{in Liouville space.}
\end{cases}
```

All arguments are passed to [`f_tau`](@ref) to evaluate ``f_τ``.
"""
function F_re(ϕ, objectives; τ=nothing)
    return real(f_tau(ϕ, objectives; τ=τ))
end

@doc raw"""Real-part functional.

```julia
J_T_re(ϕ, objectives; τ=nothing)
```

calculates

```math
J_{T,\text{re}} = 1 - F_{\text{re}} \quad\in \begin{cases}
    [0, 2] & \text{in Hilbert space} \\
    [0, 1] & \text{in Liouville space.}
\end{cases}
```

All arguments are passed to [`f_tau`](@ref) while evaluating ``F_{\text{re}}``
in [`F_re`](@ref).
"""
function J_T_re(ϕ, objectives; τ=nothing)
    return 1.0 - F_re(ϕ, objectives; τ=τ)
end

@doc raw"""Krotov-states χ for functional [`J_T_re`](@ref).

```julia
chi_re!(χ, ϕ, objectives; τ=nothing)
```

sets the elements of `χ` according to

```math
|χ_k⟩
= -\frac{∂ J_{T,\text{re}}}{∂ ⟨ϕ_k(T)|}
= \frac{1}{2N} w_k |ϕ^{\tgt}_k⟩
```

with ``|ϕ^{\tgt}_k⟩`` and ``w_k`` as defined in [`f_tau`](@ref).

Note: this function can be obtained with `make_chi(J_T_re, objectives)`.
"""
function chi_re!(χ, ϕ, objectives; τ=nothing)
    N = length(objectives)
    if τ === nothing
        τ = [dot(objectives[k].target_state, ϕ[k]) for k = 1:N]
    end
    for k = 1:N
        obj = objectives[k]
        ϕₖ_tgt = obj.target_state
        copyto!(χ[k], ϕₖ_tgt)
        w = isa(obj, WeightedObjective) ? obj.weight : 1.0
        lmul!(w / (2N), χ[k])
    end
end

"""Convert a functional from acting on a gate to acting on propagated states.

```
J_T = gate_functional(J_T_U; kwargs...)
```

constructs a functional `J_T` that meets the requirements for
[`make_gradient`](@ref) and [`make_chi`](@ref). That is, the output `J_T` takes
positional positional arguments `ϕ` and `objectives`. The input functional
`J_T_U` is assumed to have the signature `J_T_U(U; kwargs...)` where `U` is a
matrix with elements ``U_{ij} = ⟨Ψ_i|ϕ_j⟩``, where ``|Ψ_i⟩`` is the
`initial_state` of the i'th `objectives` (assumed to be the i'th canonical
basis state) and ``|ϕ_j⟩`` is the result of forward-propagating ``|Ψ_j⟩``. That
is, `U` is the projection of the time evolution operator into the subspace
defined by the basis in the `initial_states` of the  `objectives`.
"""
function gate_functional(J_T_U; kwargs...)

    function J_T(ϕ, objectives; τ=nothing)
        N = length(objectives)
        U = [(objectives[i].initial_state ⋅ ϕ[j]) for i = 1:N, j = 1:N]
        return J_T_U(U; kwargs...)
    end

    return J_T

end

@doc raw"""
Return a function to evaluate ``|χ_k⟩ = -∂J_T(Û)/∂⟨ϕ_k|`` via the chain rule.

```julia
chi! = make_gate_chi(J_T_U, objectives; use_finite_differences=false, kwargs...)
```

returns a function equivalent to

```julia
chi! = make_chi(gate_functional(J_T_U; kwargs...), objectives)
```

```math
\begin{split}
    |χ_k⟩
    &= -\frac{∂}{∂⟨ϕ_k|} J_T \\
    &= - \frac{1}{2} \sum_i (∇_U J_T)_{ik} \frac{∂ U_{ik}}{∂⟨ϕ_k|} \\
    &= - \frac{1}{2} \sum_i (∇_U J_T)_{ik} |Ψ_i⟩
\end{split}
```

where ``|Ψ_i⟩`` is the basis state stored as the `initial_state` of the i'th
`objective`, see [`gate_functional`](@ref).

The gradient ``∇_U J_T`` is obtained via automatic differentiation, or via
finite differences if `use_finite_differences=true`.

Compared to the more general [`make_chi`](@ref), `make_gate_chi` will generally
have a slightly smaller numerical overhead, as it pushes the use of automatic
differentiation down by one level.

With `use_finite_differences=true`, this routine serves to test and debug
gradients for gate functionals obtained by automatic differentiation.
"""
function make_gate_chi(J_T_U, objectives; use_finite_differences=false, kwargs...)

    N = length(objectives)
    basis = [obj.initial_state for obj in objectives]

    function zygote_gate_chi!(χ, ϕ, objectives; τ=nothing)
        function _J_T(U)
            -J_T_U(U; kwargs...)
        end
        U = [basis[i] ⋅ ϕ[j] for i = 1:N, j = 1:N]
        if use_finite_differences
            fdm = FiniteDifferences.central_fdm(5, 1)
            ∇J = FiniteDifferences.grad(fdm, gate -> _J_T(gate), U)[1]
        else
            ∇J = Zygote.gradient(gate -> _J_T(gate), U)[1]
        end
        for k = 1:N
            χ[k] .= 0.5 * sum([∇J[i, k] * basis[i] for i = 1:N])
        end
    end

    return zygote_gate_chi!

end

@doc raw"""Gradient for an arbitrary functional evaluated via χ-states.

```julia
grad_J_T_via_chi!(∇J_T, τ, ∇τ)
```

sets the (vectorized) elements of the gradient `∇J_T` to the gradient
``∂J_T/∂ϵ_{ln}`` for an arbitrary functional ``J_T=J_T(\{|ϕ_k(T)⟩\})``, under
the assumption that

```math
\begin{aligned}
    τ_k &= ⟨χ_k|ϕ_k(T)⟩ \quad \text{with} \quad |χ_k⟩ &= -∂J_T/∂⟨ϕ_k(T)|
    \quad \text{and} \\
    ∇τ_{kln} &= ∂τ_k/∂ϵ_{ln}\,,
\end{aligned}
```

where ``|ϕ_k(T)⟩`` is a state resulting from the forward propagation of some
initial state ``|ϕ_k⟩`` under the pulse values ``ϵ_{ln}`` where ``l`` numbers
the controls and ``n`` numbers the time slices. The ``τ_k`` are the elements of
`τ` and ``∇τ_{kln}`` corresponds to `∇τ[k][l, n]`.

In this case,

```math
(∇J_T)_{ln} = ∂J_T/∂ϵ_{ln} = -2 \Re \sum_k ∇τ_{kln}\,,
```

see [`make_gradient`](@ref).

Note that the definition of the ``|χ_k⟩`` matches exactly the definition
of the boundary condition for the backward propagation in Krotov's method, see
[`make_chi`](@ref). Specifically, there is a minus sign in front of the
derivative, compensated by the minus sign in the factor ``(-2)`` of the final
``(∇J_T)_{ln}``.
"""
function grad_J_T_via_chi!(∇J_T, τ, ∇τ)
    N = length(τ) # number of objectives
    L, N_T = size(∇τ[1])  # number of controls/time intervals
    ∇J_T′ = reshape(∇J_T, L, N_T)  # writing to ∇J_T′ modifies ∇J_T
    for l = 1:L
        for n = 1:N_T
            ∇J_T′[l, n] = real(sum([∇τ[k][l, n] for k = 1:N]))
        end
    end
    lmul!(-2, ∇J_T)
    return ∇J_T
end

@doc raw"""Return a function that evaluates the gradient ``∇J_T``.

```julia
grad_func! = make_gradient(J_T, objectives; via=:tau, force_zygote=false)
```

creates a function `gradfunc!(∇J_T, τ, ∇τ)` that takes a vector `τ` of values
``τ_k`` and a vector `∇τ` of gradients ``∇τ_k`` where the ``(ln)``'th element
of ``∇τ_k`` is ``∂τ_k/∂ϵ_{ln}``, and writes the element ``(ln)`` of the
gradient `∇J_T` as ``∂J_T/∂ϵ_{ln}``. The definition of ``τ_k`` depends on
`via`, see below. The ``ϵ_{ln}`` are the values of the control control field
discretized to the midpoints of a time grid. The index ``l`` numbers the
control and ``n`` numbers the time slice. The gradient (like the controls) are
assumed to be vectorized. That is, ``∇J_T`` is a vector of values with a
double-index ``(ln)``.

The passed `J_T` parameter corresponding to the functional ``J_T`` must be a
function that takes a vector of forward-propagates states `ϕ` and a vector of
objectives as positional parameters.  It must also accept a vector `τ` as a
keyword argument, which contains the overlaps of the states in `ϕ` and the
`target_state` fields of the `objectives`. If the `objectices` do not define
`target_states`, or if the ``τ``-values are not available, `J_T` must accept
`τ=nothing`. See [`J_T_sm`](@ref) for an example.

## Gradient via τ

For `via=:tau` (default), we define

```math
τ_k ≡ ⟨ϕ_k^\tgt|ϕ_k(T)⟩
```

as the overlap of ``|ϕ_k(T)⟩`` resulting from the forward propagation of the
`initial_state` ``|ϕ_k⟩``  of the k'th objective under the pulse values
``ϵ_{ln}``, and ``|ϕ_k^\tgt⟩`` as the `target_state` of the k'th objective.

We then understand ``J_T`` as a function of the ``τ_k``, and evaluate the
elements of ``∇J_T`` via the chain rule:

```math
(∇J_T)_{ln} ≡ \frac{∂J_T(\{τ_k\})}{∂ϵ_{ln}}
= 2\Re\sum_k
    \frac{∂J_T}{∂τ_k}
    \frac{∂τ_k}{∂ϵ_{ln}}\,.
```

Since the ``τ_k`` are complex numbers,

```math
\frac{∂J_T}{∂τ_k} = \frac{1}{2}\left(
    \frac{∂ J_T}{∂ \Re[τ_k]}
    - i \frac{∂ J_T}{∂ \Im[τ_k]}
\right)
```

is defined as the [Wirtinger
derivative](https://www.ekinakyurek.me/complex-derivatives-wirtinger/), and

```math
\frac{∂τ_k}{∂ϵ_{ln}} = \frac{∂\Re[τ_k]}{∂ϵ_{ln}} + i \frac{∂\Im[τ_k]}{∂ϵ_{ln}}
```

is simply the derivative of a complex number with respect to the real-valued
``ϵ_{ln}``.

Thus, the returned `grad_func!` effectively encodes the outer derivative
``∂J_T/∂τ_k``. For functionals where that derivative is known analytically,
the analytic expression is used, e.g., [`J_T_sm`](@ref) →
[`grad_J_T_sm!`](@ref).

Otherwise, or if `force_zygote=true`, the outer derivative is
determined directly from `J_T`, via automatic differentiation (using
[Zygote](https://fluxml.ai/Zygote.jl)).

## Gradient via χ

For `via=:chi`, the functional ``J_T`` is understood directly as a function of
the forward-propagated states ``|ϕ_k⟩`` instead of a function of overlaps with
the target states. This is useful in particular if the `objectives` do not
define objectives and/or the functional `J_T` cannot be expressed in terms of
overlaps.

Again we apply a chain rule to calculate the elements of the gradient ``∇J_T``:

```math
\begin{split}
(∇J_T)_{ln} &≡ \frac{∂J_T(\{|ϕ_k(T)⟩\})}{∂ϵ_{ln}}\\
&= 2\Re\sum_k
    \frac{∂J_T}{∂|ϕ_k(T)⟩}
    \frac{∂|ϕ_k(T)⟩}{∂ϵ_{ln}} \\
&= -2 \Re \sum_k \frac{∂}{∂ϵ_{ln}} ⟨χ_k(T)|ϕ_k(T)⟩\,,
\end{split}
```

with

```math
|χ_k⟩
= -\frac{J_T}{⟨ϕ_k(T)|}
= -\frac{1}{2}\left(
    \left\vert \frac{∂J_T}{∂\Re[ϕ_k]} \right\rangle
    + i \left\vert \frac{∂J_T}{∂\Im[ϕ_k]} \right\rangle
    \right)
```

as a matrix-calculus extension of the Wirtinger derivative. This definition of
``|χ_k⟩`` (note the minus sign!) matches the definition of the boundary
condition in Krotov's method, and for a given functional `J_T`, the
states ``|χ_k⟩`` can be obtained with [`make_chi`](@ref).

We define

```math
τ_k ≡ ⟨χ_k(T)|ϕ_k(T)⟩
```

and associate the `gradfunc!` argument `∇τ[k][l, n]` with

```math
(∇τ_k)_{ln} = \frac{∂τ_k}{∂ ϵ_{ln}}
```

so that structurally, ``(∇J_T)_{ln}`` is the same as for `via=:tau`, just that
``τ_k`` is now defined with respect to the boundary condition state ``|χ_k⟩``
instead of the target state ``|ϕ_k^\tgt⟩``.

The returned `grad_func!` that encodes the above equations is
[`grad_J_T_via_chi!`](@ref). This is independent of `J_T`, since the dependency
on the functional `J_T` is entirely encoded in the states ``|χ_k(T)⟩``, and
thus the gradient `∇τ`. Also, `force_zygote=true` has no effect for `via=:chi`.
Instead, `force_zygote` should be passed to the underlying [`make_chi`](@ref).

!!! tip

    In order to extend `make_gradient` with an analytic implementation for a
    new `J_T` function, define a new method like so:

    ```julia
    make_gradient(::typeof(J_T_sm), objectives, via::Val{:tau}) = grad_J_T_sm!
    ```

    which links `make_gradient` for [`J_T_sm`](@ref) to [`grad_J_T_sm!`](@ref).
"""
function make_gradient(J_T, objectives; via::Symbol=:tau, force_zygote=false)
    if force_zygote && (via == :tau)
        return make_zygote_gradient(J_T, objectives)
    else
        return make_gradient(J_T, objectives, Val(via))
    end
end

function make_zygote_gradient(J_T, objectives)

    function zygote_gradfunc!(∇J_T, τ, ∇τ)
        ∇J = Zygote.gradient(τ -> J_T(nothing, objectives; τ=τ), τ)[1]
        ∇J_T .= vec(sum(real(∇J) .* real.(∇τ) + imag(∇J) .* imag.(∇τ)))
    end

    return zygote_gradfunc!
end

make_gradient(J_T, objectives, via::Val{:chi}) = grad_J_T_via_chi!
make_gradient(J_T, objectives, via::Val{:tau}) = make_zygote_gradient(J_T, objectives)
make_gradient(::typeof(J_T_sm), objectives, via::Val{:tau}) = grad_J_T_sm!

@doc raw"""Return a function that evaluates ``|χ_k⟩ = -∂J_T/∂⟨ϕ_k|``.

```julia
chi! = make_chi(J_T, objectives; force_zygote=false)
```

creates a function `chi!(χ, ϕ, objectives; τ=nothing)` that sets
``|χ_k⟩ = -∂J_T/∂⟨ϕ_k|``. This is the state used as the boundary condition for
the backward propagation propagation in Krotov's method, as well as GRAPE if
[`grad_J_T_via_chi!`](@ref) is used. It is defined as a [Wirtinger
derivative](https://www.ekinakyurek.me/complex-derivatives-wirtinger/),
see [`make_gradient`](@ref).

The function `J_T` must take a vector of states `ϕ`
and a vector of `objectives` as positional parameters, and a vector `τ` as a
keyword argument, see e.g. [`J_T_sm`](@ref). If all objectives define a
`target_state`, then `τ` will be the overlap of the states `ϕ` with those
target states. The functional `J_T` may or may not use those overlaps.
Likewise, the resulting `chi!` may or may not use the keyword parameter `τ`.

For functionals where ``-∂J_T/∂⟨ϕ_k|`` is known analytically, that analytic
derivative will be returned, e.g.,

* [`J_T_sm`](@ref) → [`chi_sm!`](@ref),
* [`J_T_re`](@ref) → [`chi_re!`](@ref),
* [`J_T_ss`](@ref) → [`chi_ss!`](@ref).

Otherwise, or if `force_zygote=true`, automatic differentiation via Zygote is
used to calculate the derivative directly from `J_T`.

!!! tip

    In order to extend `make_chi` with an analytic implementation for a new
    `J_T` function, define a new method `make_analytic_chi` like so:

    ```julia
    make_analytic_chi(::typeof(J_T_sm), objectives) = chi_sm!
    ```

    which links `make_chi` for [`J_T_sm`](@ref) to [`chi_sm!`](@ref).
"""
function make_chi(J_T, objectives; force_zygote=false)
    if force_zygote
        return make_zygote_chi(J_T, objectives)
    else
        return make_analytic_chi(J_T, objectives)
    end
end

function make_zygote_chi(J_T, objectives)

    N = length(objectives)

    function zygote_chi!(χ, ϕ, objectives; τ=nothing)
        function _J_T(Ψ...)
            -J_T(Ψ, objectives)
        end
        for (k, ∇J) ∈ enumerate(Zygote.gradient(_J_T, ϕ...))
            copyto!(χ[k], ∇J)
            lmul!(0.5, χ[k])
        end
    end

    return zygote_chi!

end

make_analytic_chi(J_T, objectives) = make_zygote_chi(J_T, objectives)
# Well, the above fallback to `make_zygote_chi` is clearly not "analytic", but
# the intent of this name is to allow users to define new analytic
# implementation, cf. the explanation in the doc of `make_chi`.
make_analytic_chi(::typeof(J_T_sm), objectives) = chi_sm!
make_analytic_chi(::typeof(J_T_re), objectives) = chi_re!
make_analytic_chi(::typeof(J_T_ss), objectives) = chi_ss!

end


function _draw_action_from_MvNormal!(a::Vector{Float64}, P::MvNormal)
    # this assumes it is 2x2

    μ::Vector{Float64} = P.μ
    Σ::Matrix{Float64} = P.Σ.mat

    a[1] = randn()*sqrt(Σ[1,1]) + μ[1]

    # compute conditional values
    μ₂ = μ[2] + Σ[1,2]*(a[1] - μ[1])/Σ[1,1]
    Σ₂ = Σ[2,2] - Σ[1,2]*Σ[1,2]/Σ[1,1]
    a[2] = randn()*sqrt(Σ₂) + μ₂

    a
end
function _get_e_and_denom(P::MvNormal, a::Vector{Float64})

    μ::Vector{Float64} = P.μ
    Σ::Matrix{Float64} = P.Σ.mat

    Δx = a[1] - μ[1]
    Δy = a[2] - μ[2]
    det = Σ[1,1]*Σ[2,2] - Σ[1,2]*Σ[1,2]

    e = exp(-0.5*((Σ[2,2]*Δx - Σ[1,2]*Δy)*Δx + (Σ[1,1]*Δy - Σ[1,2]*Δx)*Δy) / det)
    denom = sqrt(4*π*π*det)

    (e,denom)
end
function _pdf_of_MvNormal(P::MvNormal{PDMats.PDMat{Float64,Matrix{Float64}},Vector{Float64}}, a::Vector{Float64})
    # this assumes it is 2x2

    e, denom = _get_e_and_denom(P, a)
    e / denom
end
function _logpdf_of_MvNormal(P::MvNormal{PDMats.PDMat{Float64,Matrix{Float64}},Vector{Float64}}, a::Vector{Float64})
    # this assumes it is 2x2

    e, denom = _get_e_and_denom(P, a)
    log(e) - log(denom)
end


@testset "check parameter order change" begin
    R, (x, y, z) = PolynomialRing(QQ, ["x","y","z"], ordering=:deglex)

    fs = [y^2, x]
    x, y, z = gens(parent(first(Groebner.groebner(fs, ordering=:lex))))
    @test Groebner.groebner(fs, ordering=:lex) == [y^2, x]
    @test ordering(parent(first(Groebner.groebner(fs, ordering=:lex)))) === :lex

    x, y, z = gens(parent(first(Groebner.groebner(fs, ordering=:degrevlex))))
    @test Groebner.groebner(fs, ordering=:degrevlex) == [x, y^2]
    @test ordering(parent(first(Groebner.groebner(fs, ordering=:degrevlex)))) === :degrevlex

    x, y, z = gens(parent(first(fs)))
    @test Groebner.groebner(fs, ordering=:deglex) == [x, y^2]
    @test ordering(parent(first(Groebner.groebner(fs, ordering=:deglex)))) === :deglex

    x, y, z = gens(parent(first(Groebner.groebner(fs, forsolve=true))))
    @test Groebner.groebner(fs, forsolve=true) == [y^2, x]
    @test ordering(parent(first(Groebner.groebner(fs, forsolve=true)))) === :lex

    noon = Groebner.change_ordering(Groebner.noonn(2), :lex)
    gb = Groebner.groebner(noon, ordering=:lex)
    x1, x2 = gens(parent(first(gb)))
    @test gb == [
            x2^5 - 10//11*x2^4 - 11//5*x2^3 + 2*x2^2 + 331//1100*x2 - 11//10,
            x1 + x2^4 - 10//11*x2^3 - 11//10*x2^2 + x2 - 10//11]

    gb = Groebner.groebner(noon, ordering=:degrevlex)
    x1, x2 = gens(parent(first(gb)))
    @test gb == [
            x1^2 - x2^2 - 10//11*x1 + 10//11*x2,
            x2^3 + 10//11*x1*x2 - 10//11*x2^2 - 11//10*x2 + 1,
            x1*x2^2 - 11//10*x1 + 1]

    gb = Groebner.groebner(noon, ordering=:deglex)
    x1, x2 = gens(parent(first(gb)))
    @test gb == [
            x1^2 - x2^2 - 10//11*x1 + 10//11*x2,
            x2^3 + 10//11*x1*x2 - 10//11*x2^2 - 11//10*x2 + 1,
            x1*x2^2 - 11//10*x1 + 1]

end


@testsection "Special tensors" begin
    # Rotations
    θ, ϕ, ψ = symbols("θ ϕ ψ", real = true) ;
    cθ, cϕ, cψ, sθ, sϕ, sψ = symbols("cθ cϕ cψ sθ sϕ sψ", real = true) ;
    d = Dict(cos(θ) => cθ, cos(ϕ) => cϕ, cos(ψ) => cψ, sin(θ) => sθ, sin(ϕ) => sϕ, sin(ψ) => sψ) ;
    R = Tens(subs(rot3(θ, ϕ, ψ),d...)) ;
    R6 = invKM(subs(KM(rot6(θ, ϕ, ψ)),d...)) ;
    @test R6 == R ⊠ˢ R

end


struct Ray
    origin::AbstractVec
    direction::AbstractVec
end

function point_at(ray::Ray, t::AbstractFloat)::AbstractVec
    ray.origin + t * ray.direction
end


using DelimitedFiles
using Test

#=
Refactor after seeing Tony Hallam's solution and wanting to tidy up with
functions and tests. Might try using doctests tomorrow.

I'm sure this can be solved in a really elegant recursive way but I need
to think about it some more. (Recursion is a state of mind, it seems.)
=#

"""
    part_1(data, target)

Search for the number making up to 2020.
"""
function part_1(data::Array, target::Integer=2020)::Integer
    for d in data
        if (target - d) in data
            return convert(Integer, d * (target - d))
        end
    end
end

test = readdlm("./data/day01_test.txt")
@test part_1(test) == 514579

data = readdlm("./data/day01.txt")
println("Part 1: $(part_1(data))")

"""
    part_2(data, target)

I'm sure this could be recursive but I can't quite get it...
"""
function part_2(data::Array, target::Integer=2020)::Integer
    for d in data
        for e in [i for i in data if ((d + i) < target)]
            if (target - d - e) in data
                return convert(Integer, d * e * (target - d - e))
            end
        end
    end
end

@test part_2(test) == 241861950

println("Part 2: $(part_2(data))")


struct SRS{R<:AbstractRNG} <: Imputor
    rng::R
end

# Docstring below uses julia-repl cause the rng may give different result on different
# versions of julia.
"""
    SRS(; rng=Random.GLOBAL_RNG)

Simple Random Sampling (SRS) imputation is a method for imputing both continuous and categorical
variables. Furthermore, it completes imputation while preserving the distributional
properties of the variables (e.g., mean, standard deviation).

The basic idea is that for a given variable, `x`, with missing data, we randomly draw
from the observed values of `x` to impute the missing elements. Since the random draws
from `x` for imputation are done in proportion to the frequency distribution of the values
in `x`, the univariate distributional properties are generally not impacted; this is true
for both categorical and continuous data.

# Keyword Arguments
* `rng::AbstractRNG`: A random number generator to use for observation selection

# Example
```julia-repl
julia> using Random; using Impute: SRS, impute

julia> M = [1.0 2.0 missing missing 5.0; 1.1 2.2 3.3 missing 5.5]
2×5 Matrix{Union{Missing, Float64}}:
 1.0  2.0   missing  missing  5.0
 1.1  2.2  3.3       missing  5.5

julia> impute(M, SRS(; rng=MersenneTwister(1234)); dims=:rows)
2×5 Matrix{Union{Missing, Float64}}:
 1.0  2.0  1.0  2.0  5.0
 1.1  2.2  3.3  3.3  5.5
```
"""
SRS(; rng=Random.GLOBAL_RNG) = SRS(rng)

function _impute!(data::AbstractArray{Union{T, Missing}}, imp::SRS) where T
    obs_values = collect(skipmissing(data))
    if !isempty(obs_values)
        for i in eachindex(data)
            if ismissing(data[i])
                data[i] = rand(imp.rng, obs_values)
            end
        end
    end

    return data
end


# Load Julia packages (libraries) needed  for the snippets in chapter 0

using StatisticalRethinking, CmdStan
#gr(size=(600,600));

# CmdStan uses a tmp directory to store the output of cmdstan

ProjDir = rel_path("..", "scripts", "04")
cd(ProjDir)

# CmdStan uses a tmp directory to store the output of cmdstan

ProjDir = rel_path("..", "scripts", "04")
cd(ProjDir)

#-

howell1 = CSV.read(rel_path("..", "data", "Howell1.csv"), delim=';')
df = convert(DataFrame, howell1);
df2 = filter(row -> row[:age] >= 18, df);
first(df2, 5)

# ### Snippet 4.21

# Sample 20 random heights

n = size(df2, 1)
selected_ind = sample(1:n, 20, replace=false);
df3 = df2[selected_ind, :];

# ### Snippet 4.22

# Generate approximate probabilities

struct Post
  mu::Float64
  sigma::Float64
  ll:: Float64
  prod::Float64
  prob::Float64
end

mu_list = repeat(range(140, 170, length=200), 200);
sigma_list = repeat(range(4, 20, length=200), inner=200);

ll = zeros(40000);
for i in 1:40000
    d1 = Normal(mu_list[i], sigma_list[i])
    ll[i] = sum(log.(pdf.(d1, df3[:height])))
end

d2 = Normal(178.0, 20.0)
d3 = Uniform(0, 50)
prod = ll + log.(pdf.(d2, mu_list)) + log.(pdf.(d3, sigma_list))
prob = exp.(prod .- maximum(prod))
post = DataFrame(mu=mu_list, sigma=sigma_list, ll=ll, prod=prod, prob=prob)
first(post, 10)

# Sample post

samples = post[sample(1:size(post, 1), Weights(post[:prob]), 10000, replace=true), :]

# ### Snippet 4.23

# Density of sigma

density(samples[:sigma])

# End of `clip-21-23.jl`


D=DBayesianNode(:D, ["d1","d2","d_water"])

a1=CBayesianNode(:a1)
a2=CBayesianNode(:a2)
a3=CBayesianNode(:a3)

bn=BayesianNetwork([D,a1,a2])

add_edge!(bn,D,a1)
add_edge!(bn,D,a2)

analysis=[a1,a2]

# Distributions

using Distributions

pdfNodeDict = Dict{Symbol, Array{Function,1}}()

for a in analysis
    a_d1_dist = Normal(rand(1:200), rand(1:15))
    a_d1_func = x -> pdf(a_d1_dist,x)
    a_d2_dist = Normal(rand(200:400), rand(1:15))
    a_d2_func = x -> pdf(a_d2_dist,x)

    samples = [ rand(dist) for i=1:100, dist in [a_d1_dist, a_d2_dist] ]

    a_d_water_pdf = x -> 1/(maximum(samples) - minimum(samples))

    pdfNodeDict[a.label] = [a_d1_func, a_d2_func, a_d_water_pdf]

    add_probability!(bn, P(a.label|:D), ProbabilityDensityDistribution(states(D), [a_d1_func, a_d2_func, a_d_water_pdf]))
end

l_states_d = length(states(D))
add_probability!(bn, P(:D), ProbabilityDistribution(states(D), fill(1/l_states_d, l_states_d)))
inputDict = Dict{Symbol,Any}({:a1 => 30, :a2 => 30})
testCPD = P(:D|inputDict)
diagProbs = query(bn,testCPD)

# Normalize results of query so we can compare
normalized_diagProbs = normalize_factor(convert(Array{Float64,1},collect(values(diagProbs))))

# For all a in analysis, run every function of a and give it the input defined in inputDict. The reducing steps performs an element-wise multiplication, and finally it's normalized.
calculatedDiagProbs = normalize_factor(reduce(.*, [ map(x -> x(inputDict[a.label]), pdfNodeDict[a.label]) for a in analysis ]))

# Test if the two normalized arrays are approximately equal
@test sum(map(x -> isapprox(normalized_diagProbs[x[1]], calculatedDiagProbs[x[1]]), enumerate(calculatedDiagProbs))) == l_states_d

#bn.query(P(:D|{:a1 => 30, :a2 => 250, a_d_water_pdf => 0}))


u1 = UncertainValue(Normal, 0, 0.1)
u2 = UncertainValue(Uniform, rand(100))
u3 = UncertainValue(rand(100))
u4 = UncertainValue(Normal, 5, 0.1)
u5 = UncertainValue(Normal, 10, 1.5)

vals = [u1, u2, u3, u4, u5, 6]
wts = [1, 3, 3, 1, 0.5, 0.1]

pop = UncertainScalarPopulation(vals, wts)
@test resample(pop) isa Real
@test resample(pop, 10) isa Vector{T} where {T <: Real}

x = UncertainScalarPopulation([u1, u2, u3, u4, u5], rand(5))
vals = [pop, u5, x, u4, u5, 6]
wts = rand(length(vals))

pop2 = UncertainScalarPopulation(vals, wts)
@test resample(pop2) isa Real
@test resample(pop2, 100000) isa Vector{T} where {T <: Real}

# Create a population consisting of a mixture of different types of uncertain values,
# both theoretical, fitted from empirical distributions, and discrete populations.
u1 = UncertainValue(Normal, 0, 0.1)
u2 = UncertainValue(Uniform, rand(100))
u3 = UncertainValue(rand(100))
u4 = UncertainValue(Normal, 5, 0.1)
u5 = UncertainValue(Normal, 10, 1.5)

wts = [1, 3, 3, 1, 0.5]
vals = [u1, u2, u3, u4, u5]
pop3 = UncertainScalarPopulation(vals, wts)

@test resample(pop3) isa Real
@test resample(pop3, 100) isa Vector{T} where {T <: Real}


Random.seed!(0)

################################################################################
### Regularized least squares, with two variable blocks to make things weird
################################################################################

println("Testing: regularized least squares, with two variable blocks to make things weird")

m, n1, n2 = 30, 50, 100

A1 = randn(m, n1)
A2 = randn(m, n2)
b = randn(m)

lam1 = 0.2
lam2 = 1.0

# Solve with FPG

x1_fpg = Variable(n1)
x2_fpg = Variable(n2)
expr = ls(A1*x1_fpg + A2*x2_fpg - b) + lam1*norm(x1_fpg, 1) + lam2*norm(x2_fpg, 2)
prob = problem(expr)
@time sol = solve(prob, StructuredOptimization.FPG(tol=1e-10,verbose=0,maxit=20000))
println(sol)

# Solve with ZeroFPR

x1_zerofpr = Variable(n1)
x2_zerofpr = Variable(n2)
expr = ls(A1*x1_zerofpr + A2*x2_zerofpr - b) + lam1*norm(x1_zerofpr, 1) + lam2*norm(x2_zerofpr, 2)
prob = problem(expr)
@time sol = solve(prob, StructuredOptimization.ZeroFPR(tol=1e-10,verbose=0))
println(sol)

# Solve with PANOC

x1_panoc = Variable(n1)
x2_panoc = Variable(n2)
expr = ls(A1*x1_panoc + A2*x2_panoc - b) + lam1*norm(x1_panoc, 1) + lam2*norm(x2_panoc, 2)
prob = problem(expr)
@time sol = solve(prob, StructuredOptimization.PANOC(tol=1e-10,verbose=0))
println(sol)

# Solve with minimize, use default solver/options

x1 = Variable(n1)
x2 = Variable(n2)
@time sol = @minimize ls(A1*x1 + A2*x2 - b) + lam1*norm(x1, 1) + lam2*norm(x2, 2)
println(sol)

@test norm(~x1_fpg - ~x1_zerofpr, Inf)/(1+norm(~x1_zerofpr, Inf)) <= 1e-6
@test norm(~x2_fpg - ~x2_zerofpr, Inf)/(1+norm(~x2_zerofpr, Inf)) <= 1e-6
@test norm(~x1_fpg - ~x1_panoc, Inf)/(1+norm(~x1_panoc, Inf)) <= 1e-6
@test norm(~x2_fpg - ~x2_panoc, Inf)/(1+norm(~x2_panoc, Inf)) <= 1e-6
@test norm(~x1 - ~x1_zerofpr, Inf)/(1+norm(~x1_zerofpr, Inf)) <= 1e-3
@test norm(~x2 - ~x2_zerofpr, Inf)/(1+norm(~x2_zerofpr, Inf)) <= 1e-3

res = A1*~x1_fpg + A2*~x2_fpg - b
grad1 = A1'*res
grad2 = A2'*res
ind1_zero = (~x1_fpg .== 0)
subgr1 = lam1*sign.(~x1_fpg)
subdiff1_low, subdiff1_upp = copy(subgr1), copy(subgr1)
subdiff1_low[ind1_zero] .= -lam1
subdiff1_upp[ind1_zero] .= +lam1
subgr2 = lam2*(~x2_fpg/norm(~x2_fpg, 2))

@test maximum(subdiff1_low + grad1) <= 1e-6
@test maximum(-subdiff1_upp - grad1) <= 1e-6
@test norm(grad2 + subgr2) <= 1e-6

###############################################################################
## Lasso problem with known solution
###############################################################################

println("Testing: lasso problem with known solution")

m, n, nnz_x_star = 200, 100, 10
A = randn(m, n)
lam = 1.0
x_star = randn(n)
x_star[nnz_x_star+1:end] .= 0.0
y_star = lam*sign.(x_star)
b = A*x_star + A'\y_star
@test norm(A'*(A*x_star - b) + lam*sign.(x_star)) <= 1e-12

# Solve with PG

x_pg = Variable(n)
expr = ls(A*x_pg - b) + lam*norm(x_pg, 1)
prob = problem(expr)
@time sol = solve(prob, StructuredOptimization.PG(tol=1e-10,verbose=0))
println(sol)

@test norm(~x_pg - x_star, Inf) <= 1e-8
@test norm(A'*(A*~x_pg - b) + lam*sign.(~x_pg)) <= 1e-6

# Solve with FPG

x_fpg = Variable(n)
expr = ls(A*x_fpg - b) + lam*norm(x_fpg, 1)
prob = problem(expr)
@time sol = solve(prob, StructuredOptimization.FPG(tol=1e-10,verbose=0))
println(sol)

@test norm(~x_fpg - x_star, Inf) <= 1e-8
@test norm(A'*(A*~x_fpg - b) + lam*sign.(~x_fpg)) <= 1e-6

# Solve with ZeroFPR

x_zerofpr = Variable(n)
expr = ls(A*x_zerofpr - b) + lam*norm(x_zerofpr, 1)
prob = problem(expr)
@time sol = solve(prob, StructuredOptimization.ZeroFPR(tol=1e-10,verbose=0))
println(sol)

@test norm(~x_zerofpr - x_star, Inf) <= 1e-8
@test norm(A'*(A*~x_zerofpr - b) + lam*sign.(~x_zerofpr)) <= 1e-5

# Solve with PANOC

x_panoc = Variable(n)
expr = ls(A*x_panoc - b) + lam*norm(x_panoc, 1)
prob = problem(expr)
@time sol = solve(prob, StructuredOptimization.PANOC(tol=1e-10,verbose=0))
println(sol)

@test norm(~x_panoc - x_star, Inf) <= 1e-8
@test norm(A'*(A*~x_panoc - b) + lam*sign.(~x_panoc)) <= 1e-5

################################################################################
### Problem with smooth, non-quadratic term
################################################################################

println("Testing: problem with smooth, non-quadratic term")

m, n, nnz_x_orig = 200, 500, 10
A = randn(m, n)
lam = 1.0
x_orig = randn(n)
x_orig[nnz_x_orig+1:end] .= 0.0
b = A*x_orig + randn(m)

# Solve with PG

x_pg = Variable(n)
expr = smooth(norm(A*x_pg - b, 2)) + lam*norm(x_pg, 1)
prob = problem(expr)
@time sol = solve(prob, PG(tol=1e-8,verbose=0))
println(sol)

# Solve with FPG

x_fpg = Variable(n)
expr = smooth(norm(A*x_fpg - b, 2)) + lam*norm(x_fpg, 1)
prob = problem(expr)
@time sol = solve(prob, FPG(tol=1e-8,verbose=0))
println(sol)

# Solve with ZeroFPR

x_zerofpr = Variable(n)
expr = smooth(norm(A*x_zerofpr - b, 2)) + lam*norm(x_zerofpr, 1)
prob = problem(expr)
@time sol = solve(prob, StructuredOptimization.ZeroFPR(tol=1e-8,verbose=0))
println(sol)

# Solve with PANOC

x_panoc = Variable(n)
expr = smooth(norm(A*x_panoc - b, 2)) + lam*norm(x_panoc, 1)
prob = problem(expr)
@time sol = solve(prob, StructuredOptimization.PANOC(tol=1e-8,verbose=0))
println(sol)

# Solve with minimize, default solver/options

x = Variable(n)
@time sol = @minimize smooth(norm(A*x - b, 2)) + lam*norm(x, 1)
println(sol)

@test norm(~x_fpg - ~x_zerofpr, Inf)/(1+norm(~x_zerofpr, Inf)) <= 1e-6
@test norm(~x_fpg - ~x_panoc, Inf)/(1+norm(~x_panoc, Inf)) <= 1e-6
@test norm(~x - ~x_zerofpr, Inf)/(1+norm(~x_zerofpr, Inf)) <= 1e-3

################################################################################
### Box-constrained least-squares
################################################################################

println("Testing: box-constrained least-squares")

m, n = 500, 200
A = randn(m, n)
lb, ub = -1.0, 1.0
x_orig = 2.0*randn(n)
x_orig = max.(lb, min.(ub, x_orig))
b = A*x_orig + randn(m)

# Solve with PG

x_pg = Variable(n)
expr = ls(A*x_pg - b)
prob = problem(expr, x_pg in [lb, ub])
@time sol = solve(prob, StructuredOptimization.PG(tol=1e-8,verbose=0))
println(sol)

@test norm(~x_pg - max.(lb, min.(ub, ~x_pg)), Inf) <= 1e-12
@test norm(~x_pg - max.(lb, min.(ub, ~x_pg - A'*(A*~x_pg - b))), Inf)/(1+norm(~x_pg, Inf)) <= 1e-8

# Solve with FPG

x_fpg = Variable(n)
expr = ls(A*x_fpg - b)
prob = problem(expr, x_fpg in [lb, ub])
@time sol = solve(prob, StructuredOptimization.FPG(tol=1e-8,verbose=0))
println(sol)

@test norm(~x_fpg - max.(lb, min.(ub, ~x_fpg)), Inf) <= 1e-12
@test norm(~x_fpg - max.(lb, min.(ub, ~x_fpg - A'*(A*~x_fpg - b))), Inf)/(1+norm(~x_fpg, Inf)) <= 1e-8

# Solve with ZeroFPR

x_zerofpr = Variable(n)
expr = ls(A*x_zerofpr - b)
prob = problem(expr, x_zerofpr in [lb, ub])
@time sol = solve(prob, StructuredOptimization.ZeroFPR(tol=1e-8,verbose=0))
println(sol)

@test norm(~x_zerofpr - max.(lb, min.(ub, ~x_zerofpr)), Inf) <= 1e-12
@test norm(~x_zerofpr - max.(lb, min.(ub, ~x_zerofpr - A'*(A*~x_zerofpr - b))), Inf)/(1+norm(~x_zerofpr, Inf)) <= 1e-8

# Solve with PANOC

x_panoc = Variable(n)
expr = ls(A*x_panoc - b)
prob = problem(expr, x_panoc in [lb, ub])
@time sol = solve(prob, StructuredOptimization.PANOC(tol=1e-8,verbose=0))
println(sol)

@test norm(~x_panoc - max.(lb, min.(ub, ~x_panoc)), Inf) <= 1e-12
@test norm(~x_panoc - max.(lb, min.(ub, ~x_panoc - A'*(A*~x_panoc - b))), Inf)/(1+norm(~x_panoc, Inf)) <= 1e-8

# Solve with minimize, default solver/options

x = Variable(n)
@time sol = @minimize ls(A*x - b) st x in [lb, ub]
println(sol)

@test norm(~x - max.(lb, min.(ub, ~x)), Inf) <= 1e-12
@test norm(~x - max.(lb, min.(ub, ~x - A'*(A*~x - b))), Inf)/(1+norm(~x, Inf)) <= 1e-4

################################################################################
### Non-negative least-squares from a known solution
################################################################################

println("Testing: non-negative least-squares from a known solution")

# Lagrangian:
#
#   0.5||Ax-b||^2 + y'(x-z) + [z >= 0]
#
# Optimality conditions:
#
#   A'(Ax-b) + y = 0, or A'b = A'Ax + y
#   x = z
#   z >= 0
#   y <= 0
#   y'z = 0

m, n, nnz_x_star = 500, 200, 100
A = randn(m, n)
x_star = rand(n)
x_star[nnz_x_star+1:end] .= 0.0
y_star = -rand(n)
y_star[1:nnz_x_star] .= 0.0
b = A*x_star + A'\y_star

# Solve with PG

x_pg = Variable(n)
expr = ls(A*x_pg - b)
prob = problem(expr, x_pg >= 0.0)
@time sol = solve(prob, StructuredOptimization.PG(tol=1e-8,verbose=0))
println(sol)

@test all(~x_pg .>= 0.0)
@test norm(~x_pg - x_star, Inf)/(1+norm(x_star, Inf)) <= 1e-8

# Solve with FPG

x_fpg = Variable(n)
expr = ls(A*x_fpg - b)
prob = problem(expr, x_fpg >= 0.0)
@time sol = solve(prob, StructuredOptimization.FPG(tol=1e-8,verbose=0))
println(sol)

@test all(~x_fpg .>= 0.0)
@test norm(~x_fpg - x_star, Inf)/(1+norm(x_star, Inf)) <= 1e-8

# Solve with ZeroFPR

x_zerofpr = Variable(n)
expr = ls(A*x_zerofpr - b)
prob = problem(expr, x_zerofpr >= 0.0)
@time sol = solve(prob, StructuredOptimization.ZeroFPR(tol=1e-8,verbose=0))
println(sol)

@test all(~x_zerofpr .>= 0.0)
@test norm(~x_zerofpr - x_star, Inf)/(1+norm(x_star, Inf)) <= 1e-8

# Solve with PANOC

x_panoc = Variable(n)
expr = ls(A*x_panoc - b)
prob = problem(expr, x_panoc >= 0.0)
@time sol = solve(prob, StructuredOptimization.PANOC(tol=1e-8,verbose=0))
println(sol)

@test all(~x_zerofpr .>= 0.0)
@test norm(~x_zerofpr - x_star, Inf)/(1+norm(x_star, Inf)) <= 1e-8

# Solve with minimize, default solver/options

x = Variable(n)
@time sol = @minimize ls(A*x - b) st x >= 0.0
println(sol)

@test all(~x .>= 0.0)
@test norm(~x - x_star, Inf)/(1+norm(x_star, Inf)) <= 1e-6


using Test

using MultivariateBases

using DynamicPolynomials

function api_test(B::Type{<:AbstractPolynomialBasis}, degree)
    @polyvar x[1:2]
    for basis in [
        maxdegree_basis(B, x, degree),
        basis_covering_monomials(B, monomials(x, 0:degree))
    ]
        n = binomial(2 + degree, 2)
        @test length(basis) == n
        @test typeof(copy(basis)) == typeof(basis)
        @test nvariables(basis) == 2
        @test variables(basis) == x
        @test monomialtype(typeof(basis)) == monomialtype(x[1])
        @test typeof(empty_basis(typeof(basis))) == typeof(basis)
        @test length(empty_basis(typeof(basis))) == 0
        @test polynomialtype(basis, Float64) == polynomialtype(x[1], Float64)
        @test polynomial(i -> 0.0, basis) isa polynomialtype(basis, Float64)
        @test polynomial(zeros(n, n), basis, Float64) isa polynomialtype(basis, Float64)
        @test polynomial(ones(n, n), basis, Float64) isa polynomialtype(basis, Float64)
    end
end

function orthogonal_test(B::Type{<:AbstractMultipleOrthogonalBasis}, univ::Function, even_odd_separated::Bool)
    @test MultivariateBases.even_odd_separated(B) == even_odd_separated
    @polyvar x y
    univariate_x = univ(x)
    univariate_y = univ(y)

    @testset "Univariate $var" for (var, univ) in [(x, univariate_x), (y, univariate_y)]
        basis = maxdegree_basis(B, (var,), 4)
        for i in 1:5
            @test basis.polynomials[length(basis) + 1 - i] == univ[i]
        end
    end

    @testset "basis_covering_monomials" begin
        monos = basis_covering_monomials(B, monovec([x^2 * y, y^2]))
        if even_odd_separated
            exps = [(2, 1), (0, 2), (0, 1), (0, 0)]
        else
            exps = [(2, 1), (2, 0), (1, 1), (0, 2), (1, 0), (0, 1), (0, 0)]
        end
        for i in 1:length(monos)
            @test monos.polynomials[i] == univariate_x[exps[i][1] + 1] * univariate_y[exps[i][2] + 1]
        end
        monos = basis_covering_monomials(B, monovec([x^4, x^2, x]))
        if even_odd_separated
            exps = [4, 2, 1, 0]
        else
            exps = [4, 3, 2, 1, 0]
        end
        for i in 1:length(monos)
            @test monos.polynomials[i] == univariate_x[exps[i] + 1]
        end
    end
end

@testset "Monomial" begin
    include("monomial.jl")
end
@testset "Scaled" begin
    include("scaled.jl")
end
@testset "Fixed" begin
    include("fixed.jl")
end
@testset "Hermite" begin
    include("hermite.jl")
end
@testset "Laguerre" begin
    include("laguerre.jl")
end
@testset "Legendre" begin
    include("legendre.jl")
end
@testset "Chebyshev" begin
    include("chebyshev.jl")
end


function finalpositions(
    pointstofvecs::Dict{Point2D,Array{Float64,1}},
    scaler::Scaler,
    deltat::Float64,
    fd::Function,
    geps::Float64,
    deps::Float64,
    h0::Float64
)::Tuple{Any, Float64}
    newp, dpoints = buildnewpoints(pointstofvecs, scaler, deltat, geps, fd)
    finalp = buildfinalpoints(newp, scaler, fd, deps)
    return finalp, moveindex(dpoints, deltat, h0)
end

function buildnewpoints(
    pointstofvecs::Dict{Point2D,Array{Float64,1}},
    scaler::Scaler,
    deltat::Float64,
    geps::Float64,
    fd::Function
)::Tuple{Array{Array{Float64, 1},1},  Array{Array{Float64, 1},1}}
    newp = Array{Array{Float64, 1},1}()
    dpoints = Array{Array{Float64, 1},1}()
    for (point, force) in pointstofvecs
        p = unscaledpoint2d(point, scaler)
        np = [getx(p), gety(p)] + deltat * force
        push!(newp, [np[1], np[2]])
        if fd(np) < -geps
            push!(dpoints, force)
        end
    end
    return newp, dpoints
end

function buildfinalpoints(
    newp::Array{Array{Float64, 1},1},
    scaler::Scaler,
    fd::Function,
    deps::Float64
)
    finalp = Array{Array{Float64, 1}, 1}()
    for p in newp
        x = p[1]
        y = p[2]
        d = fd([x, y])
        if d > 0
            dgradx = (fd([x + deps, y]) - d) / deps
            dgrady = (fd([x, y + deps]) - d) / deps
            res = [x, y] - [d * dgradx, d * dgrady]
            push!(finalp, res)
        else
            push!(finalp, [x, y])
        end
    end
    finalp = map(p -> scaledpoint([p[1], p[2]], scaler), unique(finalp))
    finalp = transpose(reshape(vcat(finalp...), 2, length(finalp)))
end

function moveindex(
    dpoints::Array{Array{Float64,1},1},
    deltat::Float64,
    h0::Float64
)::Float64
    d = map(row -> sum(deltat * row .^ 2), dpoints)
    push!(d, 0)
    return maximum(sqrt.(d) / h0)
end


#------------------------------
# functions to define and manipulate ODE-VAE model
#------------------------------

#------------------------------
# ODE systems
#------------------------------

# Linear
function linear_2d_system(du,u,p,t)
    a11, a12, a21, a22 = p
    z1,z2 = u
    du[1] = dz1 = a11 * z1 + a12 * z2
    du[2] = dz2 = a21 * z1 + a22 * z2
end

# Lotka-Volterra
function lotka_volterra(du,u,p,t)
    α, β, γ, δ = p
    z1,z2 = u
    du[1] = dz1 = α*z1 - β*z1*z2 #prey; α: growth, β: meeting rate of prey&predator=predation rate
    du[2] = dz2 = γ*z1*z2 - δ*z2 #predator; γ: growth of predator, δ: loss rate of predator
end

#------------------------------
# define and initialize model
#------------------------------

mutable struct odevae
    p
    q
    zdim
    paramNN
    encoder
    encodedμ
    encodedlogσ
    decoder
    decodedμ
    decodedlogσ
    ODEprob
end

downscaled_glorot_uniform(dims...) = (rand(Float32, dims...) .- 0.5f0) .* sqrt(0.01f0/sum(dims)) # smaller weights initialisation

# initialise model
function init_vae(p, q, zdim, nODEparams, ODEprob; seed::Int=1234)
    myinit = ODEprob.f.f == lotka_volterra ? downscaled_glorot_uniform : Flux.glorot_uniform
    shift(arg) = ODEprob.f.f == lotka_volterra ? tanh.(arg) .+ 1 : sigmoid(arg).-0.5
    # seed
    Random.seed!(seed)
    # parameter network
    paramNN = Chain(Dense(q,q,tanh, init=myinit),
                Dense(q, nODEparams, arg ->(shift(arg)), init=myinit),
                #Dense(q, nODEparams, arg -> (2 ./(1 .+ exp.(-arg))) .- 1.0),
                Flux.Diagonal(nODEparams))
    #   VAE encoder
    Dz, Dh = zdim, p
    encoder, encodedμ, encodedlogσ = Dense(p, Dh, arg ->(tanh.(arg) .+ 1)), Dense(Dh, Dz), Chain(Dense(Dh, Dz, arg -> -relu(arg)), Flux.Diagonal(Dz))
    # VAE decoder
    decoder, decodedμ, decodedlogσ = Dense(Dz, Dh, tanh), Dense(Dh, p), Dense(Dh, p)

    model = odevae(p, q, zdim, paramNN, encoder, encodedμ, encodedlogσ, decoder, decodedμ, decodedlogσ, ODEprob)
    return model
end

#------------------------------
# define model loss
#------------------------------

latentz(μ, logσ) = μ .+ sqrt.(exp.(logσ)) .* randn(Float32,size(μ)...) # sample latent z,

kl_q_p(μ, logσ) = 0.5 .* sum(exp.(logσ) + μ.^2 .- 1 .- (logσ),dims=1)

logp_x_z(m::odevae, x, z) = sum(logpdf.(Normal.(m.decodedμ(m.decoder(z)), sqrt.(exp.(m.decodedlogσ(m.decoder(z))))), x),dims=1) # get reconstruction error

sqnorm(x) = sum(abs2, x)
reg(m::odevae) = sum(sqnorm, Flux.params(m.decoder,m.decodedμ,m.decodedlogσ)) # regularisation term in loss

getparams(m::odevae) = Flux.params(m.encoder, m.encodedμ, m.encodedlogσ, m.decoder, m.decodedμ, m.decodedlogσ, m.paramNN) # get parameters of VAE model

function loss(X, Y, t, m::odevae)
    latentμ, latentlogσ = m.encodedμ(m.encoder(X)), m.encodedlogσ(m.encoder(X))
    learnedparams = m.paramNN(Y)
    curparams = [learnedparams[1], m.ODEprob.p[2], m.ODEprob.p[3], learnedparams[2]]
    curts = Int.(floor.(t .*(1 ./dt) .+1))
    smoothμ = Array(solve(m.ODEprob, Tsit5(), u0 = [latentμ[1,1], latentμ[2,1]], p=curparams, saveat=dt))[:,curts]
    combinedμ = hcat(latentμ[:,1],smoothμ)
    combinedz = latentz.(combinedμ, latentlogσ)
    ELBO = 1.0 .* logp_x_z(m, X, combinedz) .- 0.5 .* kl_q_p(combinedμ, latentlogσ)
    lossval = sum(-ELBO) + 0.01*reg(m) + 1.0*sum((smoothμ .- latentμ[:,2:end]).^2)
    return lossval
end

function loss_wrapper(m::odevae)
    return function (X, Y, t) loss(X, Y, t, m) end
end

# to experiment:
# xs, tvals, group1, group2 = generate_xs(n, p, true_u0, sol_group1, sol_group2);
# trainingdata = zip(xs, x_baseline, tvals);
# X, Y, t = first(trainingdata)


# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

# Initialize data structures in element container
function init_elements!(elements, mesh::StructuredMesh{1}, basis::LobattoLegendreBasis)
  @unpack node_coordinates, left_neighbors,
          jacobian_matrix, contravariant_vectors, inverse_jacobian = elements

  # Calculate node coordinates, Jacobian matrix, and inverse Jacobian determinant
  for cell_x in 1:size(mesh, 1)
    calc_node_coordinates!(node_coordinates, cell_x, mesh.mapping, mesh, basis)

    calc_jacobian_matrix!(jacobian_matrix, cell_x, node_coordinates, basis)

    calc_inverse_jacobian!(inverse_jacobian, cell_x, jacobian_matrix)
  end

  # Contravariant vectors don't make sense in 1D, they would be identical to inverse_jacobian
  fill!(contravariant_vectors, NaN)

  initialize_neighbor_connectivity!(left_neighbors, mesh)

  return nothing
end

# Calculate physical coordinates to which every node of the reference element is mapped
# `mesh.mapping` is passed as an additional argument for type stability (function barrier)
function calc_node_coordinates!(node_coordinates, cell_x, mapping, mesh::StructuredMesh{1},
                                basis::LobattoLegendreBasis)
  @unpack nodes = basis

  # Get cell length in reference mesh
  dx = 2 / size(mesh, 1)

  # Calculate node coordinates of reference mesh
  cell_x_offset = -1 + (cell_x-1) * dx + dx/2

  for i in eachnode(basis)
    # node_coordinates are the mapped reference node_coordinates
    node_coordinates[1, i, cell_x] = mapping(cell_x_offset + dx/2 * nodes[i])[1]
  end
end

# Calculate Jacobian matrix of the mapping from the reference element to the element in the physical domain
function calc_jacobian_matrix!(jacobian_matrix, element, node_coordinates::AbstractArray{<:Any,3},
                               basis::LobattoLegendreBasis)
  @views mul!(jacobian_matrix[1, 1, :, element], basis.derivative_matrix, node_coordinates[1, :, element]) # x_ξ

  return jacobian_matrix
end

# Calculate inverse Jacobian (determinant of Jacobian matrix of the mapping) in each node
function calc_inverse_jacobian!(inverse_jacobian::AbstractArray{<:Any, 2}, element, jacobian_matrix)
  @views inverse_jacobian[:, element] .= inv.(jacobian_matrix[1, 1, :, element])

  return inverse_jacobian
end

# Save id of left neighbor of every element
function initialize_neighbor_connectivity!(left_neighbors, mesh::StructuredMesh{1})
  # Neighbors in x-direction
  # Inner elements
  for cell_x in 2:size(mesh, 1)
    left_neighbors[1, cell_x] = cell_x - 1
  end

  if isperiodic(mesh)
    # Periodic boundary
    left_neighbors[1, 1] = size(mesh, 1)
  else
    # Use boundary conditions
    left_neighbors[1, 1] = 0
  end

  return left_neighbors
end

end # @muladd


# ## Data

import CSV

using DataFrames
using Turing
using TuringModels

data_path = joinpath(TuringModels.project_root, "data", "Howell1.csv")
df = CSV.read(data_path, DataFrame; delim=';')
df = filter(row -> row.age >= 18, df);

# For df, see Section [df](#df).

# ## Model

@model function line(height)
    σ ~ Uniform(0, 50)
    μ ~ Normal(178, 20)

    height .~ Normal.(μ, σ)
end

model = line(df.height);

# ## Output

chains = sample(model, NUTS(0.65), 1000)

# \defaultoutput{}

# ## Original output

"""
Iterations = 1:1000
Thinning interval = 1
Chains = 1,2,3,4
Samples per chain = 1000

Empirical Posterior Estimates:
         Mean        SD       Naive SE       MCSE      ESS
alpha 154.597086 0.27326431 0.0043206882 0.0036304132 1000
 beta   0.906380 0.04143488 0.0006551430 0.0006994720 1000
sigma   5.106643 0.19345409 0.0030587777 0.0032035103 1000

Quantiles:
          2.5%       25.0%       50.0%       75.0%       97.5%
alpha 154.0610000 154.4150000 154.5980000 154.7812500 155.1260000
 beta   0.8255494   0.8790695   0.9057435   0.9336445   0.9882981
sigma   4.7524368   4.9683400   5.0994450   5.2353100   5.5090128
""";

# ## df

df


cd(@__DIR__); include("setups/grid7x3.jl");

## ROT2
D = natural_eigdist(𝚽, 𝛌, Q; α = 0.1, distance = :ROT)
E = transform(fit(MDS, D, maxoutdim=2, distances=true))
E[1, :] .*= sign(E[1, 2]); E[2, :] .*= -sign(E[1, 2])
plt = grid7x3_mds_heatmaps(E, 𝚽; backend = :pyplot)
savefig(plt, "../figs/Grid7x3_MDS_ROT2_alpha01.png")
display(plt)


##############################################################
"""
$(SIGNATURES)

Calculation of Bernoulli function via Horner scheme based on Taylor
coefficients around 0.
"""

function bernoulli_horner(x)
    y=x/47_900_160
    y=x*y
    y=x*(-1/1_209_600+y)
    y=x*y
    y=x*(1/30_240+y)
    y=x*y
    y=x*(-1/720+y)
    y=x*y
    y=x*(1/12+y)
    y=x*(-1/2+y)
    y=1+y
end

##############################################################
"""
$(SIGNATURES)

Bernoulli function ``B(x)=\\frac{x}{e^x-1}`` for exponentially
fitted upwinding.

The name `fbernoulli` has been chosen to avoid confusion
with Bernoulli from JuliaStats/Distributions.jl

Returns a real number containing the result.
"""
function fbernoulli(x)
    expx=exp(x)
    if expx<0.99999
        return x/(expx-1)
    elseif expx < 1.00001
        return bernoulli_horner(x)
    else
        return x/(expx-1)
    end
end

##############################################################

const etruncmax = 100.0
const etruncmin = -100.0

function trexp(x)
  if(x<etruncmin)
    return 1.0/(exp(etruncmin)*(-x+etruncmin+1.0));
  elseif (x<etruncmax)
    return exp(x);
  else
      return exp(etruncmax)*(x-etruncmax+1.0);
  end
end

"""
$(SIGNATURES)

Bernoulli function ``B(x)=\\frac{x}{e^x-1}`` for exponentially
fitted upwind, joint evaluation for positive and negative
argument

Usually, we need ``B(x), B(-x)`` togehter,
and it is cheaper to calculate them together.

Returns two real numbers containing the result for argument
`x` and argument `-x`.

"""
function fbernoulli_pm(x)
    expx=exp(x)
    if abs(expx-1)>0.00001
        bp=x/(expx-1)
        bm=x/(1-1/expx)
        return bp,bm
    else
        y=bernoulli_horner(x)
        return y,x+y
    end
end


__precompile__
import FastTransforms
using FastGaussQuadrature

struct Quadrature
    Nq::Int;
    w::Array{Float64,1};
    wTens::Array{Float64,1};
    xi::Array{Float64,1};

   function Quadrature(NqVal,quadType)
        if quadType == "Gauss"
            #xi,w = QuadNodesWeights(NqVal,-1.0,1.0)
            xi,w = gausslegendre(NqVal)
        elseif quadType == "ClenshawCurtis"
            xi,w = QuadNodesWeightsClCu(NqVal,-1.0,1.0)
        end
        xi = xi[end:-1:1];
        w = w[end:-1:1];

        wTens = zeros(NqVal^2);

        for k = 1:NqVal
            for q = 1:NqVal
                wTens[(q-1)*NqVal+k] = w[k]*w[q];
            end
        end

        new(NqVal,w,wTens,xi)
   end

end

function QuadNodesWeightsClCu(Nq::Int64,a::Float64,b::Float64) # transformation to a,b not tested!
    xiQuad,wQuad = FastTransforms.clenshawcurtis(Nq,0.,0.)
    xiQuad=(a*(1-xiQuad)+b*(1+xiQuad))/2.0;
    wQuad = wQuad*(2.0/(b-a));
    return xiQuad,wQuad;
end

function Integral(q::Quadrature, f::Function)
    dot( q.w, f.(q.xi) );
end

function Integral(q::Quadrature, f::Function, a::Float64, b::Float64)
    w = q.w*(b-a)*0.5;
    xi = (a*(ones(Float64,length(w)).-q.xi)+b*(ones(Float64,length(w))+q.xi))/2.0;
    #return q.w'f.(xi);
    return sum(w'f.(xi))
end

function IntegralVec(q::Quadrature, fVec::Array{Float64,1}, a::Float64, b::Float64)
    w = q.w*(b-a)*0.5;
    return sum(w'fVec)
end

function Integral(q::Quadrature, fVec::Array{Float64,1})
    return sum(q.wTens'fVec)
end

function IntegralMat(q::Quadrature, fVec::Array{Float64,2}, a::Float64, b::Float64)
    w = q.w*(b-a)*0.5;
    return fVec*w;
end

function IntegralMat(q::Quadrature, fVec::Array{Float64,2})
    return fVec*q.w;
end


module ReactiveMPModelsAutoregressiveTest

using Test, InteractiveUtils
using Rocket, ReactiveMP, GraphPPL, Distributions
using BenchmarkTools, Random, Plots, Dates, LinearAlgebra, StableRNGs

## Model definition
## -------------------------------------------- ##
@model [ default_factorisation = MeanField() ] function ar_model(n, order)

    x = datavar(Vector{Float64}, n)
    y = datavar(Float64, n)

    γ ~ GammaShapeRate(1.0, 1.0)
    θ ~ MvNormalMeanPrecision(zeros(order), diageye(order))

    for i in 1:n
        y[i] ~ NormalMeanPrecision(dot(x[i], θ), γ)
    end

    return x, y, θ, γ
end

@model [ default_factorisation = MeanField() ] function lar_model(T::Type{ Multivariate }, n, order, c, stype, τ)

    # Parameter priors
    γ  ~ GammaShapeRate(1.0, 1.0)
    θ  ~ MvNormalMeanPrecision(zeros(order), diageye(order))

    # We create a sequence of random variables for hidden states
    x = randomvar(n)
    # As well a sequence of observartions
    y = datavar(Float64, n)

    ct = constvar(c)
    # We assume observation noise to be known
    cτ = constvar(τ)

    # Prior for first state
    x0 ~ MvNormalMeanPrecision(zeros(order), diageye(order))

    x_prev = x0

    # AR process requires extra meta information
    meta = ARMeta(Multivariate, order, stype)

    for i in 1:n
        # Autoregressive node uses structured factorisation assumption between states
        x[i] ~ AR(x_prev, θ, γ) where { q = q(y, x)q(γ)q(θ), meta = meta }
        y[i] ~ NormalMeanPrecision(dot(ct, x[i]), cτ)
        x_prev = x[i]
    end

    return x, y, θ, γ
end
@model [ default_factorisation = MeanField() ] function lar_model(T::Type{Univariate}, n, order, c, stype, τ)

    # Parameter priors
    γ  ~ GammaShapeRate(1.0, 1.0)
    θ  ~ NormalMeanPrecision(0.0, 1.0)

    # We create a sequence of random variables for hidden states
    x = randomvar(n)
    # As well a sequence of observartions
    y = datavar(Float64, n)

    ct = constvar(c)
    # We assume observation noise to be known
    cτ = constvar(τ)

    # Prior for first state
    x0 ~ NormalMeanPrecision(0.0, 1.0)

    x_prev = x0

    # AR process requires extra meta information
    meta = ARMeta(Univariate, order, stype)

    for i in 1:n
        x[i] ~ AR(x_prev, θ, γ) where { q = q(y, x)q(γ)q(θ), meta = meta }
        y[i] ~ NormalMeanPrecision(ct * x[i], cτ)
        x_prev = x[i]
    end

    return x, y, θ, γ
end
## -------------------------------------------- ##
## Inference definition
## -------------------------------------------- ##
function ar_inference(inputs, outputs, order, niter)
    n = length(outputs)

    model, (x, y, θ, γ) = ar_model(n, order, options = (limit_stack_depth = 500, ))

    γ_buffer = keep(Marginal)
    θ_buffer = keep(Marginal)
    fe       = keep(Float64)

    γ_sub = subscribe!(getmarginal(γ), γ_buffer)
    θ_sub = subscribe!(getmarginal(θ), θ_buffer)
    f_sub = subscribe!(score(Float64, BetheFreeEnergy(), model), fe)

    setmarginal!(γ, GammaShapeRate(1.0, 1.0))

    for _ in 1:niter
        update!(x, inputs)
        update!(y, outputs)
    end

    unsubscribe!((γ_sub, θ_sub, f_sub))

    return γ_buffer, θ_buffer, fe
end

function lar_init_marginals!(::Type{ Multivariate }, order, γ, θ)
    setmarginal!(γ, GammaShapeRate(1.0, 1.0))
    setmarginal!(θ, MvNormalMeanPrecision(zeros(order), diageye(order)))
end

function lar_init_marginals!(::Type{ Univariate }, order, γ, θ)
    setmarginal!(γ, GammaShapeRate(1.0, 1.0))
    setmarginal!(θ, NormalMeanPrecision(0.0, 1.0))
end

function lar_inference(data, order, artype, stype, niter, τ)

    # We build a full graph based on nber of observatios
    n = length(data)

    # Depending on the order of AR process `c` is
    # either a nber or a vector
    c = ReactiveMP.ar_unit(artype, order)

    # Note that to run inference for huge model it might be necessary to pass extra
    # options = (limit_stack_depth = 100,) to limit stack depth during recursive inference procedure
    model, (x, y, θ, γ) = lar_model(artype, n, order, c, stype, τ)

    # We are going to keep `γ` and `θ` estimates for all VMP iterations
    # But `buffer` only last posterior estimates for a sequence of hidden states `x`
    # We also will keep Bethe Free Energy in `fe`
    γ_buffer = keep(Marginal)
    θ_buffer = keep(Marginal)
    x_buffer = buffer(Marginal, n)
    fe       = keep(Float64)

    γsub  = subscribe!(getmarginal(γ), γ_buffer)
    θsub  = subscribe!(getmarginal(θ), θ_buffer)
    xsub  = subscribe!(getmarginals(x), x_buffer)
    fesub = subscribe!(score(Float64, BetheFreeEnergy(), model), fe)

    lar_init_marginals!(artype, order, γ, θ)

    # We update data several times to perform several VMP iterations
    for i in 1:niter
        update!(y, data)
    end

    # It is important to unsubscribe from running observables
    unsubscribe!((γsub, θsub, xsub, fesub))

    return γ_buffer, θ_buffer, x_buffer, fe
end

@testset "Autoregressive model" begin

    @testset "Autoregressive model" begin
        ## -------------------------------------------- ##
        ## Data creation
        ## -------------------------------------------- ##
        function ar_ssm(series, order)
            inputs = [reverse!(series[1:order])]
            outputs = [series[order + 1]]
            for x in series[order+2:end]
                push!(inputs, vcat(outputs[end], inputs[end])[1:end-1])
                push!(outputs, x)
            end
            return inputs, outputs
        end
        rng  = StableRNG(1234)
        series          = randn(rng, 1_000)
        ## -------------------------------------------- ##
        ## Inference execution and test inference results
        for order in 1:5
            inputs, outputs = ar_ssm(series, order)
            γ_buffer, θ_buffer, fe = ar_inference(inputs, outputs, order, 15)
            @test length(γ_buffer) === 15
            @test length(θ_buffer) === 15
            @test length(fe) === 15
            @test last(fe) < first(fe)
            @test all(filter(e -> abs(e) > 1e-3, diff(getvalues(fe))) .< 0)
        end
        ## -------------------------------------------- ##
        ## Form debug output
        base_output = joinpath(pwd(), "_output", "models")
        mkpath(base_output)
        timestamp        = Dates.format(now(), "dd-mm-yyyy-HH-MM")
        benchmark_output = joinpath(base_output, "ar_model_benchmark_$(timestamp)_v$(VERSION).txt")
        ## -------------------------------------------- ##
        ## Create output benchmarks
        inputs5, outputs5 = ar_ssm(series, 5)
        benchmark = @benchmark ar_inference($inputs5, $outputs5, 5, 15)#
        open(benchmark_output, "w") do io
            show(io, MIME("text/plain"), benchmark)
            versioninfo(io)
        end
        ## -------------------------------------------- ##
    end

    @testset "Latent autoregressive model" begin
        ## -------------------------------------------- ##
        ## Data creation
        ## -------------------------------------------- ##
        # The following coefficients correspond to stable poles
        coefs_ar_5 = [ 0.10699399235785655, -0.5237303489793305, 0.3068897071844715, -0.17232255282458891, 0.13323964347539288 ]

        function generate_lar_data(rng, n, θ, γ, τ)
            order        = length(θ)
            states       = Vector{Vector{Float64}}(undef, n + 3order)
            observations = Vector{Float64}(undef, n + 3order)

            γ_std = sqrt(inv(γ))
            τ_std = sqrt(inv(γ))

            states[1] = randn(rng, order)

            for i in 2:(n + 3order)
                states[i]       = vcat(rand(rng, Normal(dot(θ, states[i - 1]), γ_std)), states[i-1][1:end-1])
                observations[i] = rand(rng, Normal(states[i][1], τ_std))
            end

            return states[1+3order:end], observations[1+3order:end]
        end
        # Seed for reproducibility
        rng  = StableRNG(123)
        # Number of observations in synthetic dataset
        n = 500
        # AR process parameters
        real_γ = 5.0
        real_τ = 5.0
        real_θ = coefs_ar_5
        states, observations = generate_lar_data(rng, n, real_θ, real_γ, real_τ)
        ## -------------------------------------------- ##
        ## Inference execution

        # AR order 1
        γ, θ, xs, fe = lar_inference(observations, 1, Univariate, ARsafe(), 15, real_τ)
        @test length(xs) === n
        @test length(γ)  === 15
        @test length(θ)  === 15
        @test length(fe) === 15
        @test abs(last(fe) - 518.9182342) < 0.01
        @test all(filter(e -> abs(e) > 1e-3, diff(getvalues(fe))) .< 0)

        for i in 1:4
            γ, θ, xs, fe = lar_inference(observations, i, Multivariate, ARsafe(), 15, real_τ)
            @test length(xs) === n
            @test length(γ)  === 15
            @test length(θ)  === 15
            @test length(fe) === 15
        end

        # AR order 5
        γ, θ, xs, fe = lar_inference(observations, length(real_θ), Multivariate, ARsafe(), 15, real_τ)

        ## -------------------------------------------- ##
        ## Test inference results
        @test length(xs) === n
        @test length(γ)  === 15
        @test length(θ)  === 15
        @test length(fe) === 15
        @test abs(last(fe) - 514.66086) < 0.01
        @test all(filter(e -> abs(e) > 1e-1, diff(getvalues(fe))) .< 0)
        @test (mean(last(γ)) - 3.0std(last(γ)) < real_γ < mean(last(γ)) + 3.0std(last(γ)))
        ## -------------------------------------------- ##
        ## Form debug output
        base_output = joinpath(pwd(), "_output", "models")
        mkpath(base_output)
        timestamp        = Dates.format(now(), "dd-mm-yyyy-HH-MM")
        plot_output      = joinpath(base_output, "lar_model_plot_$(timestamp)_v$(VERSION).png")
        benchmark_output = joinpath(base_output, "lar_model_benchmark_$(timestamp)_v$(VERSION).txt")
        ## -------------------------------------------- ##
        ## Create output plots
        p1 = plot(first.(states), label="Hidden state")
        p1 = scatter!(p1, observations, label="Observations")
        p1 = plot!(p1, first.(mean.(xs)), ribbon = sqrt.(first.(var.(xs))), label="Inferred states", legend = :bottomright)

        p2 = plot(mean.(γ), ribbon = std.(γ), label = "Inferred transition precision", legend = :bottomright)
        p2 = plot!([ real_γ ], seriestype = :hline, label = "Real transition precision")

        p3 = plot(getvalues(fe), label = "Bethe Free Energy")

        p = plot(p1, p2, p3, layout = @layout([ a; b c ]))
        savefig(p, plot_output)
        ## -------------------------------------------- ##
        ## Create output benchmarks
        benchmark = @benchmark lar_inference($observations, length($real_θ), Multivariate, ARsafe(), 15, $real_τ) seconds=15
        open(benchmark_output, "w") do io
            show(io, MIME("text/plain"), benchmark)
            versioninfo(io)
        end
        ## -------------------------------------------- ##
    end

end

end


type BrentSolver <: Solver1D
  solverInfo::SolverInfo
end
BrentSolver(maxEvals::Int = 100, lowerBoundEnforced::Bool = false, upperBoundEnforced::Bool = false, lowerBound::Float64 = 0.0, upperBound::Float64 = 0.0) =
  BrentSolver(SolverInfo(maxEvals, lowerBoundEnforced, upperBoundEnforced, lowerBound, upperBound))

# brent solver function
function _solve(solver::BrentSolver, f::Function, accuracy::Float64, xMin::Float64, xMax::Float64, fxMin::Float64, fxMax::Float64, root::Float64, eval_num::Int)
  froot = f(root)
  eval_num += 1
  max_evals = solver.solverInfo.maxEvals
  if froot * fxMin < 0
    xMax = xMin
    fxMax = fxMin
  else
    xMin = xMax
    fxMin = fxMax
  end

  d = root - xMax
  e = d

  while eval_num < max_evals
    if (froot > 0.0 && fxMax > 0.0) || (froot < 0.0 && fxMax < 0.0)
      # rename xMin, root, xMax, and adjust the bounds
      xMax = xMin
      fxMax = fxMin
      e = d = root - xMin
    end

    if abs(fxMax) < abs(froot)
      xMin = root
      root = xMax
      xMax = xMin
      fxMin = froot
      froot = fxMax
      fxMax = fxMin
    end

    # convergence check
    xAcc1 = 2.0 * eps(Float64) * abs(root) + 0.5 * accuracy
    xMid = (xMax - root) / 2.0
    if abs(xMid) <= xAcc1 || is_close(froot, 0.0)
      _ = f(root)
      eval_num += 1
      return root
    end

    if abs(e) >= xAcc1 && abs(fxMin) > abs(froot)
      # attempting inverse quadratic interpolation
      s = froot / fxMin
      if is_close(xMin, xMax)
        p = 2.0 * xMid * s
        q = 1.0 - s
      else
        q = fxMin / fxMax
        r = froot / fxMax
        p = s * (2.0 * xMid * q * (q - r) - (root - xMin) * (r - 1.0))
        q = (q - 1.0) * (r - 1.0) * (s - 1.0)
      end

      if p > 0.0
        q = -q # check if in bounds
      end

      p = abs(p)
      min1 = 3.0 * xMid * q - abs(xAcc1 * q)
      min2 = abs(e * q)

      if (2.0 * p < (min1 < min2 ? min1 : min2))
        e = d
        d = p / q
      else
        # interpolation failed, used bisection
        d = xMid
        e = d
      end
    else
      # bounds are increasing too slowly
      d = xMid
      e = d
    end
    xMin = root
    fxMin = froot
    if abs(d) > xAcc1
      root += d
    else
      root += sign(xMid) * xAcc1
    end

    froot = f(root)
    eval_num += 1
  end

  error("Maximum number of function evals exceeded!")
end


using Optim

#cumulpdf(x,alpha,xmax) =  x^(1-alpha) - xmax^(1-alpha)

cumulpdf(x,a,xmax) = (x^(1-a)-xmax^(1-a)) / (1-xmax^(1-a))

function loghistogram(ts, nb=20)
# histogram of log-transformed values
   his = fit(Histogram, ts, nbins=nb)
   x1 = his.edges[1][2:end]
   ind = findall(x->x != 0, his.weights)  # remove bins with entry 0
   w = his.weights[ind]
   w = w ./ float(sum(w))
   x1 = x1[ind]
   y1 = log.(w) .- x1
   ex1 = exp.(x1)
   ey1 = exp.(y1)
   x1,y1,ex1,ey1
end

function log2histogram1(ts)
# histogram with logarithmic binning
# bin positions at powers of two, 2,4,8,16,...
    maxt = maximum(ts)
    maxind = ceil(Int,log2(maxt))
    bins = 2.0 .^ (1:maxind)   # bins
    nbins = length(bins)
    w = zeros(nbins)        # weights
    for i=1:length(ts)
       t = ts[i]
       if t == 1
           w[1] += 1
       else
           ind = ceil(Int,log2(t))
           w[ind] +=1
       end
    end
    ind = findall(x->x != 0, w)  # remove bins with entry 0
    w = w[ind]
    bins = bins[ind]
    distr = w ./ bins  # get distribution by normalization
    distr = distr / sum(distr)

    bins, distr
end

function log2histogram(ts)
# histogram with logarithmic binning
# bin positions at powers of two, 1,2,4,8,16,...
    maxt = maximum(ts)
    maxind = ceil(Int,log2(maxt))
    bins = 2.0 .^ (0:maxind)   # bins
    nbins = length(bins)
    weight = zeros(nbins)        # weights
    widths = zeros(nbins)        # widths of bins
    for t in ts
       ind = ceil(Int,log2(t)) + 1
       weight[ind] +=1
    end
    widths[1] = 0.5
    widths[2:end] = bins[1:end-1]
    ind = findall(x->x != 0, weight)  # remove bins with entry 0
    weight = weight[ind]
    widths = widths[ind]
    bins   = bins[ind]
    distr = weight ./ widths  # get distribution by normalization
    distr = distr / sum(distr)

    bins, distr
end

# log likelihood
like(mu,g,nmax) = log((1-mu)/(nmax^(1-mu)-1)) - mu*g

like_inf(mu,g) = log(mu-1) - mu*g

function standard_err(M,mu,nmax)
# standard error of exponent
   mu1 = mu-1.0
   denom = 1-nmax^mu1
   1.0 / sqrt(1.0/(mu1*mu1) - nmax^mu1 * log(log(nmax))/(denom*denom)) /
   sqrt(M)
end

function KS_test(C,n,mu,nmax)
    # Kolmogorov-Smirnov test
   diff = abs.(cumulpdf.(C,mu,nmax) .- (1:n)/n)
   maximum(diff)
end

function estimate(C)
   nmax = maximum(C)
   n = length(C)
   estimate(C,n,nmax)
end

function estimate(C,n,nmax)
   # estimate fitting values for truncated power law
   # critical exponent, standard deviation Kolmogorov-Smirnov distance
   g = sum(log.(C)) / n   # log(geometric mean)
   #g = sum(log.(2*C)) / n   # try correction from Clauset
   # find maximum using Brent's method
   mu = optimize(x->-like(x,g,nmax), 1.01, 3.0).minimizer
  # mu = optimize(x->-like(x,g,nmax), 1.01, 3.0, GoldenSection()).minimizer
   st = standard_err(n,mu,nmax)
   #d = KS_test(C,n,mu,nmax)
   #mu, st, d
   mu, st
end

function randPL(n,mu,nmax)
# generate sample of random numbers taken from a truncated power law
   mu1 = 1.0 - mu
   r = (1.0 .- rand(n) .* (1.0 - nmax^mu1)) .^ (1.0/mu1)
   sort(r, rev=true)
end


using NearestNeighbors
using Distances

function knn_clustering(data::AbstractMatrix;
    k = 10,
    mode = :hard,
    treetype::Type{T} = BallTree,
    metric::M = Euclidean()) where {T <: NearestNeighbors.NNTree, M <: Distances.PreMetric}

    N = size(data)[2]
    tree = treetype(data, metric)
    # centroids = hcat(map(r-> data[:,r], rand(1:N, N))...)
    # centroids = hcat(map(n-> data[:,n], 1:N)...)
    knnq = knn(tree, data, k)
    # @show length(unique(knnq[1]...))

    C = Vector{Array{Int64,1}}()
    used = Set{Int64}()
    if mode == :soft
        for (i,knn) in enumerate(knnq[1])
            points = Vector{Int64}()
            for nn in knn
                push!(used, nn)
                # push!(points, data[:,nn])
                push!(points, nn)
            end
            if length(points) > 0
                push!(C, points)
            end
        end
    elseif mode == :hard
        for (i, knn) in enumerate(knnq[1])
            points = Vector{Int64}()
            for nn in knn
                if !(nn in used)
                    push!(used, nn)
                    # push!(points, data[:,nn])
                    push!(points, nn)
                end
            end
            if length(points) > 0
                push!(C,points)
            end
        end
    else
        error("Unknown mode: $mode")
    end

    C, sort(collect(used))
end


__precompile__()

module CayleyDickson

import Base: +, -, *, /, \, ==
import Base: real, isreal, conj, inv, zero, one, iszero
import Base: show, promote_rule, widen, convert, float
# import Unreal: unreal

"Error when finding the inverse of a zero divisor."
const ZeroDivisorInverse = "inverse of zero divisor"

"Error when the denominator in a quotient is zero."
const ZeroDenominator = "denominator is zero"

include("Construct.jl")
include("Elliptic.jl")
include("Binion.jl")
include("Quaternion.jl")
include("Octonion.jl")
include("Hyperbolic.jl")
include("SplitBinion.jl")
include("SplitQuaternion.jl")
include("SplitOctonion.jl")
include("Parabolic.jl")
include("Exo1Real.jl")
include("Exo2Real.jl")
include("Exo1Binion.jl")
include("Exo1SplitBinion.jl")
include("Exo3Real.jl")
include("Exo2Binion.jl")
include("Exo2SplitBinion.jl")
include("Exo1Quaternion.jl")
include("Exo1SplitQuaternion.jl")
include("Sedenion.jl")
include("Exo4Real.jl")
include("random.jl")
include("real.jl")

# Types
export Construct, Elliptic, Hyperbolic, Parabolic,
       Binion, SplitBinion, Exo1Real,
       Quaternion, SplitQuaternion, Exo2Real, Exo1Binion, Exo1SplitBinion,
       Octonion, SplitOctonion, Exo1Quaternion, Exo1SplitQuaternion, Exo3Real, Exo2Binion, Exo2SplitBinion,
       Sedenion, Exo4Real

# Methods
export unreal, array, iszerodivisor, random, quadrance,
       commutator, associator, jacobiator, alternatorL, alternatorR, flexator,
       anti_commutator, anti_associator, anti_alternatorL, anti_alternatorR, anti_flexator,
       cloak, dagger, star, selfstar, anti_selfstar

end


mutable struct MCEuropeanEngine{S <: AbstractBlackScholesProcess, RSG <: AbstractRandomSequenceGenerator} <: MCVanillaEngine{S, RSG}
  process::S
  timeSteps::Int
  timeStepsPerYear::Int
  requiredSamples::Int
  maxSamples::Int
  requiredTolerance::Float64
  brownianBridge::Bool
  seed::Int
  # mcSimulation::MCSimulation{RSG, T}
  antitheticVariate::Bool
  rsg::RSG
end

function MCEuropeanEngine(process::S; timeSteps::Int = -1, timeStepsPerYear::Int = -1, brownianBridge::Bool = false, antitheticVariate::Bool = false,
                          requiredSamples::Int = -1, requiredTolerance::Float64 = -1.0, maxSamples::Int = typemax(Int), seed::Int = 1,
                          rsg::RSG = InverseCumulativeRSG(seed)) where {S <: AbstractBlackScholesProcess, RSG <: AbstractRandomSequenceGenerator}
  # build mc sim
  # mcSim = MCSimulation{RSG, SingleVariate}(antitheticVariate, false, rsg, SingleVariate())

  return MCEuropeanEngine{S, RSG}(process, timeSteps, timeStepsPerYear, requiredSamples, maxSamples, requiredTolerance, brownianBridge, seed, antitheticVariate, rsg)
end

struct EuropeanPathPricer{P <: PlainVanillaPayoff} <: AbstractPathPricer
  payoff::P
  discount::Float64
end

EuropeanPathPricer(optionType::OT, strike::Float64, disc::Float64) where {OT <: OptionType} = EuropeanPathPricer(PlainVanillaPayoff{OT}(optionType, strike), disc)

(pricer::EuropeanPathPricer)(path::Path) = pricer.payoff(path[end]) * pricer.discount

function time_grid(pe::MCVanillaEngine, opt::VanillaOption)
  lastExerciseDate = opt.exercise.dates[end]
  t = get_time(pe.process, lastExerciseDate)
  if pe.timeSteps != -1
    return TimeGrid(t, pe.timeSteps)
  elseif pe.timeStepsPerYear != -1
    steps = round(Int, floor(pe.timeStepsPerYear * t))
    return TimeGrid(t, max(steps, 1))
  else
    error("time steps not specified")
  end
end

function path_generator(pe::MCVanillaEngine, opt::VanillaOption)
  dimensions = get_factors(pe.process)
  grid = time_grid(pe, opt)
  init_sequence_generator!(pe.rsg, dimensions * (length(grid.times) - 1))

  return PathGenerator(pe.process, grid, pe.rsg, pe.brownianBridge)
end

function path_pricer(pe::MCEuropeanEngine, opt::VanillaOption)
  payoff = opt.payoff
  process = pe.process

  return EuropeanPathPricer(payoff.optionType, payoff.strike, discount(process.riskFreeRate, time_grid(pe, opt)[end]))
end

function _calculate!(pe::MCVanillaEngine, opt::VanillaOption)
  mcSim = MCSimulation(pe, false, opt, SingleVariate())
  _calculate!(mcSim, pe.requiredTolerance, pe.requiredSamples, pe.maxSamples)
  opt.results.value = stats_mean(mcSim.mcModel.sampleAccumulator)
  return pe, opt
end


"""
    SeisAddNoise(d, snr; <keyword arguments>)

Add noise at a given signal-to-noise ratio level `snr` to an N-dimensional input
data `d`. Noise can be band limited using kewyord `L`.

# Arguments
* `d::Array{Real, N}`: N-dimensional data.
* `snr::Real`: signal-to-noise ratio.

# Keyword arguments
* `db::Bool=false`: `db=false` if `snr` is given by amplitude, `db=true` if
snr is given in dB.
* `pdf::AbstractString="gaussian"`: random noise probability distribution:
`"gaussian"` or `"uniform"`.
* `L::Int=1`: averaging operator length to band-limit the random noise.

# Examples
```
julia> w = Ricker(); wn = SeisAddNoise(w, 2); plot(w); plot(wn);
MeasureSNR(w, wn)

julia> d, extent = SeisHypEvents(); dn = SeisAddNoise(d, 1.0, db=true, L=9);
SeisPlot([d dn], extent); MeasureSNR(d, dn, db=true)
```
Credits: Juan I. Sabbione, 2016
"""

function SeisAddNoise{T<:Real, N}(d::Array{T, N}, snr::Real; db::Bool=false,
                                  pdf::AbstractString="gaussian", L::Int=1)

    noise = GenNoise(size(d), pdf, L=L)
    if db==false
        noise = noise/vecnorm(noise) * vecnorm(d)/snr
    elseif db==true
        noise = noise/vecnorm(noise) * vecnorm(d)/10.0^(0.05*snr)
    end
    noisy = d + noise
    assert(abs(MeasureSNR(d, noisy, db=db) - snr) < 1e10*eps(AbstractFloat(snr)))
    return noisy

end

# Generate the noise
function GenNoise(dims::Tuple, pdf::AbstractString; L::Int=1)
    N = length(dims)
    n1 = dims[1]
    N == 1 ? nx = 1 : nx = *(dims[2:end]...)
    hamm = Hamming(L)
    L2 = floor(Int,L/2)
    if pdf=="gaussian"
        noise = randn(n1, nx)
    elseif pdf=="uniform"
        noise = -1 + 2*rand(n1, nx)
    else
        error("pdf must be gussian or uniform")
    end
    for j = 1:nx
     noise[:,j] = conv(noise[:,j], hamm)[L2+1:n1+L2]
    end
    reshape(noise, dims...)
end


using CSFML.LibCSFML

mutable struct CellularAutomata
    matrix
end

function gol_number_of_adjacent_cells(i, j, matrix)
    # faster implementation :

    #   memory estimate:  0 bytes
    #   allocs estimate:  0
    #   --------------
    #   minimum time:     34.048 ns (0.00% GC)
    #   median time:      35.196 ns (0.00% GC)
    #   mean time:        35.230 ns (0.00% GC)
    #   maximum time:     101.239 ns (0.00% GC)
    #   --------------
    #   samples:          10000
    #   evals/sample:     993

    h, w = size(matrix)

    sum = 0
    if i == 1
        if j == 1
            #  __..
            # |0x
            # |xx

            sum = matrix[ 2, 1 ] + matrix[ 1, 2 ] + matrix[ 2, 2 ]
        elseif j == w
            #         ..__
            #           xo|
            #           xx|

            sum = matrix[ i+1, j ] + matrix[ i,j-1 ] + matrix[ i+1, j-1 ]
        else
            #   ..___..
            #     xox
            #     xxx

            sum = matrix[ i+1, j-1 ] + matrix[ i+1, j ] + matrix[ i+1, j+1 ] +
                        matrix[ i, j-1 ] + matrix[ i, j+1 ]
        end
    elseif i == h
        if j == 1
            # :
            # |xx
            # |ox
            #  --

            sum = matrix[ i-1, j ] + matrix[ i-1, j+1 ] + matrix[ i, j+1 ]
        elseif j == w
            #               :
            #             xx|
            #             xo|
            #             --

            sum = matrix[ i, j-1 ] + matrix[ i-1, j ] + matrix[ i-1, j-1 ]
        else
            #
            #      xxx
            #      xox
            #    ..---..

            sum = matrix[ i-1, j-1 ] + matrix[ i-1, j ] + matrix[ i-1, j+1 ] +
                matrix[ i, j-1 ] + matrix[ i, j+1 ]
        end
    else # 1 < i < h
        if j == 1
            #  :
            #  |xx
            #  |ox
            #  |xx
            #  :

            sum = matrix[ i-1, j+1 ] + matrix[ i, j+1 ] + matrix[ i+1, j+1] +
                 matrix[ i-1, j ] + matrix[ i+1, j ]
        elseif j == w
            #          :
            #        xx|
            #        xo|
            #        xx|
            #          :

            sum = matrix[ i-1, j-1 ] + matrix[ i, j-1 ] + matrix[ i+1, j-1 ] +
                matrix[ i-1, j ] + matrix[ i+1, j ]
        else
            #     xxx
            #     xox
            #     xxx

            sum = matrix[ i-1, j-1 ] + matrix[ i-1, j ] + matrix[ i-1, j+1 ] +
                  matrix[ i,   j-1 ]                    + matrix[ i,   j+1 ] +
                  matrix[ i+1, j-1 ] + matrix[ i+1, j ] + matrix[ i+1, j+1 ]
        end
    end
    sum
end

# GOL rules
function gol_compute_next_state(i, j, previous_state_matrix)
	adjacent_live_cell_count = gol_number_of_adjacent_cells(i, j, previous_state_matrix)
    previous_state = previous_state_matrix[i,j]
    next_state = 0
    if previous_state == 1
        if adjacent_live_cell_count < 2
            next_state = 0
        elseif adjacent_live_cell_count <= 3
            next_state = 1
        elseif adjacent_live_cell_count > 3
            next_state = 0
        end
    elseif previous_state == 0
        if adjacent_live_cell_count == 3
            next_state = 1
        else
            next_state = 0
        end
    end
    next_state
end

function update_ca!(ca::CellularAutomata)
    h, w = size(ca.matrix)
    next_matrix = similar(ca.matrix)
    for i=1:w, j=1:h
        next_matrix[i,j] = gol_compute_next_state(i, j, ca.matrix)
    end
    ca.matrix = next_matrix
end

function draw_ca(window,ca,cell_size)
    matrix_h, matrix_w = size(ca.matrix)
    square = sfRectangleShape_create()
    sfRectangleShape_setSize(square, sfVector2f(cell_size,cell_size) )
    sfRectangleShape_setFillColor(square,sfColor_fromRGBA(255,155,105,255) )
    for i in 0:matrix_h-1, j in 0:matrix_w-1
        value = ca.matrix[i + 1, j + 1]
        if value == 1
            sfRectangleShape_setPosition(square, sfVector2f(cell_size*i, cell_size*j) )
            sfRenderWindow_drawRectangleShape(window, square, C_NULL)
        end
    end
end

println("""=========================================================================================

                                ███▄ ▄███▓ ▄▄▄       ██▓ ███▄     █
                                ▓██▒▀█▀ ██▒▒████▄    ▓██▒ ██ ▀█   █
                                ▓██    ▓██░▒██  ▀█▄  ▒██▒▓██  ▀█ ██▒
                                ▒██    ▒██ ░██▄▄▄▄██ ░██░▓██▒  ▐▌██▒
                                ▒██▒   ░██▒ ▓█   ▓██▒░██░▒██░   ▓██░
                                ░ ▒░   ░  ░ ▒▒   ▓▒█░░▓  ░ ▒░   ▒ ▒
                                ░  ░      ░  ▒   ▒▒ ░ ▒ ░░ ░░   ░ ▒░
                                ░      ░     ░   ▒    ▒ ░   ░   ░ ░
                                    ░         ░  ░ ░           ░

===============================================================================================""")

Config = (
    window_size         = 500    ,   # pixels
    ca_matrix_size      = 50      ,   # matrix size
    target_frequency    = 10           # CA updates per seconds.
)

# init window
window_height = Config.window_size
window_width = Config.window_size
mode = sfVideoMode(window_height, window_width, 32)
window = sfRenderWindow_create(mode, "Game Of Life", sfResize | sfClose, C_NULL)
@assert window != C_NULL
sfWindow_setVerticalSyncEnabled(window,true)
sfWindow_setActive(window, sfTrue)

# init font
font = sfFont_createFromFile(joinpath("/julia/playgrounds/resources/NotoSans-Black.ttf"))
@assert font != C_NULL

# init ca
s = Config.ca_matrix_size
w, h = s, s
matrix = [ rand([0,1]) for i = 1:h, j = 1:w ]
ca = CellularAutomata(matrix)

# init matrix rendering engine
cell_size = Config.window_size / s

# init clock
clock = sfClock_create()
# Let's express time in seconds.
function get_time(clock ::Ptr{Nothing})
    sfTime_asSeconds(sfClock_getElapsedTime(clock))
end
function restart(clock ::Ptr{Nothing})
     sfTime_asSeconds(sfClock_restart(clock))
end

event_ref = Ref{sfEvent}()
target_frequency = Config.target_frequency
frame_start_timestamp = -Inf
time_frame_computing = 0

running = true
while (running)

    while Bool(sfRenderWindow_isOpen(window)) && running

        previous_frame_timestamp = frame_start_timestamp

#       # process events
        frame_start_timestamp = get_time(clock)
        frame_time = frame_start_timestamp - previous_frame_timestamp
        actual_frequency = 1 / frame_time

        while Bool(sfRenderWindow_pollEvent(window, event_ref))
            if event_ref[].type == sfEvtClosed
                sfRenderWindow_close(window)
                running = false
                println("Render window closed.")
            elseif event_ref[].type == sfEvtResized
                glViewport(9, 9, event_ref[].size.width, event_ref[].size.height)
                println("Viewport updated.")
            elseif event_ref[].type == sfEvtKeyPressed
                if event_ref[].key.code == sfKeyQ
                    running = false
                    println("Quiting application.")
                end
            end
        end

#       # update CA
        ca_start_timestamp = get_time(clock)
        time_event_processing = ca_start_timestamp - frame_start_timestamp
        #println("Event processing took $time_event_processing seconds.")

        update_ca!(ca)

#       # draw scene
        drawca_start_timestamp = get_time(clock)
        time_update_ca = drawca_start_timestamp - ca_start_timestamp

        sfRenderWindow_clear(window, sfColor_fromRGBA(24,20,18,255))

        # ----------------------------------------------------------------------------
        draw_ca(window,ca,cell_size)
        #draw_scene(window,ca,cell_size)
        # ----------------------------------------------------------------------------

#       # draw hud
        drawhud_start_timestamp = get_time(clock)
        time_drawca = drawhud_start_timestamp - drawca_start_timestamp

        sf_text = sfText_create()
        sfText_setPosition(sf_text, sfVector2f(4,window_height-20))
        sfText_setString(sf_text, "frame computing time : $time_frame_computing seconds.")
        sfText_setFont(sf_text, font)
        sfText_setCharacterSize(sf_text, 14)
        sfRenderWindow_drawText(window, sf_text, C_NULL)

        sf_text2 = sfText_create()
        sfText_setPosition(sf_text2, sfVector2f(4,window_height-40))
        sfText_setString(sf_text2, "fps : $actual_frequency")
        sfText_setFont(sf_text2, font)
        sfText_setCharacterSize(sf_text2, 14)
        sfRenderWindow_drawText(window, sf_text2, C_NULL)

#       # render
        render_start_timestamp = get_time(clock)
        sfRenderWindow_display(window)
#       # computing end
        frame_end_timestamp = get_time(clock)

        time_render = frame_end_timestamp - render_start_timestamp

        time_drawhud = render_start_timestamp - drawhud_start_timestamp
        time_frame_computing = frame_end_timestamp - frame_start_timestamp

        time_to_sleep = 1 / target_frequency - time_frame_computing

        time_sum = time_event_processing +
             time_drawca +
             time_drawhud +
             time_render

        time_report = """
        **
        event processing         : $time_event_processing
        draw Cellular Automata   : $time_drawca
        draw hud                 : $time_drawhud
        render                   : $time_render
        time sum                 : $time_sum
        time computing frame     : $time_frame_computing

        """
        println(time_report)
        if time_to_sleep > 0.005
            sleep(time_to_sleep-0.005)
        end
    end

end

#sfText_destroy(text)
sfFont_destroy(font)
#sfSprite_destroy(sprite)
#sfTexture_destroy(texture)
sfRenderWindow_destroy(window)


@polly function kernel_covariance(float_n, data, cov, mean)
    n,m = size(data)

    for j = 1:m
        mean[j] = zero(eltype(mean))
        for i = 1:n
            mean[j] += data[i,j]
        end
        mean[j] /= float_n
    end

    for i = 1:n, j = 1:m
        data[i,j] -= mean[j]
    end

    for i = 1:m, j = i:m
        cov[i,j] = zero(eltype(cov))
        for k = 1:n
            cov[i,j] += data[k,i] * data[k,j]
        end
        cov[i,j] /= (float_n - one(eltype(cov)))
        cov[j,i] = cov[i,j]
    end
end

let
    m = 1200
    n = 1400

    data = zeros(Float32,n,m)
    cov = zeros(Float32,m,m)
    mean = zeros(Float32,m)

    float_n = Float32(n)

    for i = 1:n, j = 1:m
        data[i,j] = (i*j) / m
    end

    SUITE["covariance"] = @benchmarkable kernel_covariance($float_n, data, cov, mean) setup = (data = copy($data); cov = copy($cov); mean = copy($mean))
end


"""
    SineAudioGenerator

An audio generator for a simple sine wave.  The waveform will be continuous even
if the rate (frequency) is modified during playback.
"""
mutable struct SineAudioGenerator <: AbstractAudioGenerator
    phase::Float64
    amp::Float64
    rate::Float64
end

function reset!(gen::SineAudioGenerator)
    gen.phase = 0
    nothing
end
function has_samples(gen::SineAudioGenerator, sample_rate::Float64)
    true
end
function next_samples!(gen::SineAudioGenerator, max_n::Int,
                       sample_rate::Float64)
    dphase = gen.rate / sample_rate
    pi2 = 2*pi
    samples = zeros(Float64, max_n, 2)
    p = gen.phase
    for i in 1:max_n
        samples[i, 1] = samples[i, 2] = gen.amp * sin(pi2*p)
        p += dphase
    end
    gen.phase = p
    samples
end


#############################################################################
# sumlargesteigs.jl
# Handles top k eigenvalues of a symmetric or Hermitian matrix
# (and imposes the constraint that its argument be symmetric or Hermitian)
# All expressions and atoms are subtypes of AbstractExpr.
# Please read expressions.jl first.
#############################################################################
import LinearAlgebra.eigvals

### sumlargesteigs

struct SumLargestEigs <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr,AbstractExpr}
    size::Tuple{Int,Int}

    function SumLargestEigs(x::AbstractExpr, k::AbstractExpr)
        children = (x, k)
        m, n = size(x)
        if m == n
            return new(:sumlargesteigs, hash(children), children, (1, 1))
        else
            error("sumlargesteigs can only be applied to a square matrix.")
        end
    end
end

function sign(x::SumLargestEigs)
    return NoSign()
end

function monotonicity(x::SumLargestEigs)
    return (Nondecreasing(), NoMonotonicity())
end

function curvature(x::SumLargestEigs)
    return ConvexVexity()
end

function evaluate(x::SumLargestEigs)
    return eigvals(evaluate(x.children[1]))[end-x.children[2]:end]
end

function sumlargesteigs(x::AbstractExpr, k::Int)
    return k == 0 ? Constant(0) : SumLargestEigs(x, Constant(k))
end

# Create the equivalent conic problem:
#   minimize sk + Tr(Z)
#   subject to
#            Z ⪰ 0
#            A ⪰ 0
#            Z + sI ⪰ A
# See Ben-Tal and Nemirovski, "Lectures on Modern Convex Optimization"
# Example 18.c

function conic_form!(x::SumLargestEigs, unique_conic_forms)
    if !has_conic_form(unique_conic_forms, x)
        A = x.children[1]
        k = x.children[2]
        m, n = size(A)
        if sign(A) == ComplexSign()
            Z = ComplexVariable(n, n)
        else
            Z = Variable(n, n)
        end
        s = Variable()
        # The two inequality constraints have the side effect of constraining A to be symmetric,
        # since only symmetric matrices can be positive semidefinite.
        p = minimize(
            s * k + real(tr(Z)),
            Z + s * Matrix(1.0I, n, n) - A ⪰ 0,
            Z ⪰ 0,
        )
        cache_conic_form!(unique_conic_forms, x, p)
    end
    return get_conic_form(unique_conic_forms, x)
end


module Saveresult
    include("saveresult_module.jl")
    using Printf
    using .Saveresult_module

    const rho_csv_filename = "rho.csv"
    const rhoTilde_csv_filename = "rhoTilde.csv"
    const y_csv_filename = "y.csv"

    asymptotic(x) = let
        LAMBDA = 3.886
        K = 0.190785707092222

        return (1.0 + (K * x) ^ (3.0 / LAMBDA)) ^ -LAMBDA
    end

    expansion(x) = let
        B = -1.588071022611375

        return (((1.0 / 3.0 * x + 0.4 * B * sqrt(x)) * x) + 4.0 / 3.0 * sqrt(x) + B) * x + 1.0
    end

    function saveresult(data, xarray, yarray)
        alpha = (128.0 / (9.0 * pi ^ 2) * data.Z) ^ (1.0 / 3.0)

        sr_param = Saveresult_module.Saveresult_param(
            alpha,
            32.0 / (9.0 * pi ^ 3) * data.Z * data.Z,
            xarray[2] - xarray[1],
            get_s(xarray, yarray),
            collect(0.0:xarray[2] - xarray[1]:floor(xarray[end] / alpha)),
            xarray,
            yarray,
            data.Z)

        @printf(" Energy = %.15f (Hartree)\n", makeenergy(sr_param))

        saverho(sr_param)
        saverhoTilde(sr_param)
        savey(sr_param)

        @printf("電子密度を%sと%sに，y(x)を%sに書き込みました．\n", rho_csv_filename, rhoTilde_csv_filename, y_csv_filename)
    end

    exactrho(r, sr_param) = let
        return 4.0 * r * r * sr_param.Z ^ 3 * exp(-2.0 * sr_param.Z * r)
    end

    exactrhoTilde(r, sr_param) = let
        return 4.0 * sr_param.Z ^ 3 * exp(-2.0 * sr_param.Z * r)
    end

    makeenergy(sr_param) = let
        y0_prime = (sr_param.yarray[2] - sr_param.yarray[1]) / (sr_param.xarray[2] - sr_param.xarray[1])

        @printf(" y'(0) = %.14f\n", y0_prime)

        return 3.0 / 7.0 * sr_param.alpha * sr_param.Z ^ (7.0 / 3.0) * y0_prime
    end

    rho(r, sr_param) = let
        return sr_param.s * sr_param.b * (1.0 / sr_param.alpha) ^ 2 * sqrt(r) * y(sr_param, r) * sqrt(y(sr_param, r))
    end

    rhoTilde(r, sr_param) = let
        return sr_param.s * sr_param.b * (y(sr_param, r) / r) * sqrt(y(sr_param, r) / r)
    end

    saverho(sr_param) = let
        open(rho_csv_filename, "w" ) do fp
            for r in sr_param.rarray
                write(fp, @sprintf("%.15f, %.15f, %.15f\n", r, rho(sr_param.alpha * r, sr_param), exactrho(r, sr_param)))
            end
        end
    end

    saverhoTilde(sr_param) = let
        open(rhoTilde_csv_filename, "w" ) do fp
            for r in sr_param.rarray
                write(fp, @sprintf("%.15f, %.15f, %.15f\n", r, rhoTilde(sr_param.alpha * r, sr_param), exactrhoTilde(r, sr_param)))
            end
        end
    end

    savey(sr_param) = let
        open(y_csv_filename, "w" ) do fp

            for x in sr_param.xarray
                write(fp, @sprintf("%.15f, %.15f, %.15f, %.15f\n", x, y(sr_param, x), expansion(x), asymptotic(x)))
            end
        end
    end

    y(sr_param, x) = let
        klo = 1
        max = length(sr_param.xarray)
        khi = max

        # 表の中の正しい位置を二分探索で求める
        @inbounds while khi - klo > 1
            k = (khi + klo) >> 1

            if sr_param.xarray[k] > x
                khi = k
            else
                klo = k
            end
        end

        # yvec_[i] = f(xvec_[i]), yvec_[i + 1] = f(xvec_[i + 1])の二点を通る直線を代入
        return (sr_param.yarray[khi] - sr_param.yarray[klo]) / (sr_param.xarray[khi] - sr_param.xarray[klo]) * (x - sr_param.xarray[klo]) + sr_param.yarray[klo]
    end

    get_s(xarray, yarray) = let
        sum = 0.0
        len = length(xarray)
        max = len - 2

        # Simpsonの公式によって数値積分する
        @inbounds @simd for i = 1:2:max
            f0 = yarray[i] * sqrt(yarray[i]) * sqrt(xarray[i])
            f1 = yarray[i + 1] * sqrt(yarray[i + 1]) * sqrt(xarray[i + 1])
            f2 = yarray[i + 2] * sqrt(yarray[i + 2]) * sqrt(xarray[i + 2])

            sum += (f0 + 4.0 * f1 + f2)
        end

        sum *= (xarray[2] - xarray[1]) / 3.0

        return 4.0 * pi / sum
    end
end