""" 2-category of categories, functors, and natural transformations.

Categories in mathematics appear in the large, often as categories of sets with
extra structure, and in the small, as algebraic structures that generalize
groups, monoids, preorders, and graphs. This division manifests in Catlab as
well. Large categories (in spirit, if not in the [technical
sense](https://ncatlab.org/nlab/show/large+category)) occur throughout Catlab as
`@instance`s of the theory of categories. For computational reasons, small
categories are usually presented by generators and relations.

This module provides a minimal interface to accomodate both situations. Category
instances are supported through the wrapper type [`TypeCat`](@ref). Finitely
presented categories are provided by another module, [`FinCats`](@ref).
"""
module Categories
export Cat, TypeCat, Functor, Transformation, dom, codom, compose, id,
  ob, hom, is_hom_equal, ob_map, hom_map, dom_ob, codom_ob, component,
  OppositeCat, op, co

using AutoHashEquals

using ...GAT
import ...Theories: Category2, ob, hom, dom, codom, compose, ⋅, ∘, id,
  composeH, *

# Categories
############

""" Size of a category, used for dispatch and subtyping purposes.

A [`Cat`](@ref) type having a particular `CatSize` means that categories of that
type are *at most* that large.
"""
abstract type CatSize end

""" Size of a large category, such as Set.

To the extent that they form a category, we regard Julia types and functions
([`TypeCat`](@ref)) as forming a large category.
"""
struct LargeCatSize <: CatSize end

""" Abstract base type for a category.

The objects and morphisms in the category have Julia types `Ob` and `Hom`,
respectively. Note that these types do *not* necessarily form an `@instance` of
the theory of categories, as they may not meaningfully form a category outside
the context of this object. For example, a finite category regarded as a
reflexive graph with a composition operation might have type `Cat{Int,Int}`,
where the objects and morphisms are numerical identifiers for vertices and edges
in the graph.

The basic operations available in any category are: [`dom`](@ref),
[`codom`](@ref), [`id`](@ref), [`compose`](@ref).
"""
abstract type Cat{Ob,Hom,Size<:CatSize} end

""" Coerce or look up object in category.

Converts the input to an object in the category, which should be of type `Ob` in
a category of type `Cat{Ob,Hom}`. How this works depends on the category, but a
common case is to look up objects, which might be integers or GAT expressions,
by their human-readable name, usually a symbol.

See also: [`hom`](@ref).
"""
function ob end

""" Coerce or look up morphism in category.

See also: [`ob`](@ref).
"""
function hom end

""" Domain of morphism in category.
"""
dom(C::Cat, f) = dom(f)

""" Codomain of morphism in category.
"""
codom(C::Cat, f) = codom(f)

""" Identity morphism on object in category.
"""
id(C::Cat, x) = id(x)

""" Compose morphisms in a category.
"""
compose(C::Cat, fs...) = compose(fs...)

""" Are two morphisms in a category equal?

By default, just checks for equality of Julia objects using ``==``. In some
categories, checking equality of morphisms may involve nontrivial reasoning.
"""
is_hom_equal(C::Cat, f, g) = is_hom_equal(f, g)
is_hom_equal(f, g) = f == g

# Instances
#----------

""" Pair of Julia types regarded as a category.

The Julia types should form an `@instance` of the theory of categories
(`Theories.Category`).
"""
struct TypeCat{Ob,Hom} <: Cat{Ob,Hom,LargeCatSize} end

TypeCat(Ob::Type, Hom::Type) = TypeCat{Ob,Hom}()

# FIXME: Type conversion isn't practical because types are often too tight.
#ob(::TypeCat{Ob,Hom}, x) where {Ob,Hom} = convert(Ob, x)
#hom(::TypeCat{Ob,Hom}, f) where {Ob,Hom} = convert(Hom, f)
ob(::TypeCat, x) = x
hom(::TypeCat, f) = f

Base.show(io::IO, ::TypeCat{Ob,Hom}) where {Ob,Hom} =
  print(io, "TypeCat($Ob, $Hom)")

# Functors
##########

""" Abstract base type for a functor between categories.

A functor has a domain and a codomain ([`dom`](@ref) and [`codom`](@ref)), which
are categories, and object and morphism maps, which can be evaluated using
[`ob_map`](@ref) and [`hom_map`](@ref). The functor object can also be called
directly when the objects and morphisms have distinct Julia types. This is
sometimes but not always the case (see [`Cat`](@ref)), so when writing generic
code one should prefer the `ob_map` and `hom_map` functions.
"""
abstract type Functor{Dom<:Cat,Codom<:Cat} end

""" Evaluate functor on object.
"""
@inline ob_map(F::Functor, x) = do_ob_map(F, x)

""" Evaluate functor on morphism.
"""
@inline hom_map(F::Functor, f) = do_hom_map(F, f)

""" Identity functor on a category.
"""
@auto_hash_equals struct IdentityFunctor{Dom<:Cat} <: Functor{Dom,Dom}
  dom::Dom
end

codom(F::IdentityFunctor) = F.dom

do_ob_map(F::IdentityFunctor, x) = ob(F.dom, x)
do_hom_map(F::IdentityFunctor, f) = hom(F.dom, f)

function Base.show(io::IO, F::IdentityFunctor)
  print(io, "id(")
  show_domains(io, F, codomain=false)
  print(io, ")")
end

""" Composite of functors.
"""
@auto_hash_equals struct CompositeFunctor{Dom,Codom,
    F<:Functor{Dom,<:Cat},G<:Functor{<:Cat,Codom}} <: Functor{Dom,Codom}
  fst::F
  snd::G
end
Base.first(F::CompositeFunctor) = F.fst
Base.last(F::CompositeFunctor) = F.snd

dom(F::CompositeFunctor) = dom(first(F))
codom(F::CompositeFunctor) = codom(last(F))

do_ob_map(F::CompositeFunctor, x) = ob_map(F.snd, ob_map(F.fst, x))
do_hom_map(F::CompositeFunctor, f) = hom_map(F.snd, hom_map(F.fst, f))

function Base.show(io::IO, F::CompositeFunctor)
  print(io, "compose(")
  show(io, first(F))
  print(io, ", ")
  show(io, last(F))
  print(io, ")")
end

show_type_constructor(io::IO, ::Type{<:Functor}) = print(io, "Functor")

function show_domains(io::IO, f; domain::Bool=true, codomain::Bool=true,
                      recurse::Bool=true)
  if get(io, :hide_domains, false)
    print(io, "…")
  else
    if domain
      show(IOContext(io, :compact=>true, :hide_domains=>!recurse), dom(f))
    end
    if codomain
      domain && print(io, ", ")
      show(IOContext(io, :compact=>true, :hide_domains=>!recurse), codom(f))
    end
  end
end

# Callables
#----------

Functor(f::Function, g::Function, C::Cat, D::Cat) = FunctorCallable(f, g, C, D)

""" Functor defined by two Julia callables, an object map and a morphism map.
"""
@auto_hash_equals struct FunctorCallable{Dom,Codom} <: Functor{Dom,Codom}
  ob_map::Any
  hom_map::Any
  dom::Dom
  codom::Codom
end

dom(F::FunctorCallable) = F.dom
codom(F::FunctorCallable) = F.codom
do_ob_map(F::FunctorCallable, x) = F.ob_map(x)
do_hom_map(F::FunctorCallable, f) = F.hom_map(f)

# Instances
#----------

(F::Functor{TypeCat{Ob,Hom}})(x::Ob) where {Ob,Hom} = ob_map(F, x)
(F::Functor{TypeCat{Ob,Hom}})(f::Hom) where {Ob,Hom} = hom_map(F, f)

# Natural transformations
#########################

""" Abstract base type for a natural transformation between functors.

A natural transformation ``α: F ⇒ G`` has a domain ``F`` and codomain ``G``
([`dom`](@ref) and [`codom`](@ref)), which are functors ``F,G: C → D`` having
the same domain ``C`` and codomain ``D``. The transformation consists of a
component ``αₓ: Fx → Gx`` in ``D`` for each object ``x ∈ C``, accessible using
[`component`](@ref) or indexing notation (`Base.getindex`).
"""
abstract type Transformation{C<:Cat,D<:Cat,Dom<:Functor,Codom<:Functor} end

""" Component of natural transformation.
"""
function component end

@inline Base.getindex(α::Transformation, c) = component(α, c)

""" Domain object of natural transformation.

Given ``α: F ⇒ G: C → D``, this function returns ``C``.
"""
dom_ob(α::Transformation) = dom(dom(α)) # == dom(codom(α))

""" Codomain object of natural transformation.

Given ``α: F ⇒ G: C → D``, this function returns ``D``.
"""
codom_ob(α::Transformation) = codom(dom(α)) # == codom(codom(α))

@auto_hash_equals struct IdentityTransformation{C<:Cat,D<:Cat,Dom<:Functor{C,D}} <:
     Transformation{C,D,Dom,Dom}
  dom::Dom
end

codom(α::IdentityTransformation) = α.dom

function component(α::IdentityTransformation, x)
  F = dom(α)
  id(codom(F), ob_map(F, x))
end

const IdIdTransformation{C<:Cat} = IdentityTransformation{C,C,IdentityFunctor{C}}

# 2-category of categories
##########################

@instance Category2{Cat,Functor,Transformation} begin
  dom(F::Functor) = F.dom
  codom(F::Functor) = F.codom
  id(C::Cat) = IdentityFunctor(C)

  function compose(F::Functor, G::Functor; strict::Bool=true)
    !strict || codom(F) == dom(G) ||
      error("Domain mismatch in composition $F ⋅ $G")
    compose_id(F, G)
  end

  dom(α::Transformation) = α.dom
  codom(α::Transformation) = α.codom
  id(F::Functor) = IdentityTransformation(F)

  function compose(α::Transformation, β::Transformation; strict::Bool=true)
    !strict || codom(α) == dom(β) ||
      error("Domain mismatch in vertical composition $α ⋅ $β")
    compose_id(α, β)
  end
  function composeH(α::Transformation, β::Transformation; strict::Bool=true)
    !strict || codom_ob(α) == dom_ob(β) ||
      error("Domain mismatch in horizontal composition $α * $β")
    composeH_id(α, β)
  end
  function composeH(α::Transformation, H::Functor; strict::Bool=true)
    !strict || codom_ob(α) == dom(H) ||
      error("Domain mismatch in whiskering $α * $H")
    composeH_id(α, H)
  end
  function composeH(F::Functor, β::Transformation; strict::Bool=true)
    !strict || codom(F) == dom_ob(β) ||
      error("Domain mismatch in whiskering $F * $β")
    composeH_id(F, β)
  end
end

# XXX: Is this normalization of identities using multiple dispatch a good idea?
# Unlike in `Sets`, it doesn't feel great since it requires so much boilerplate.

@inline compose_id(F::Functor, G::Functor) = do_compose(F, G)
@inline compose_id(F::Functor, ::IdentityFunctor) = F
@inline compose_id(::IdentityFunctor, G::Functor) = G
@inline compose_id(F::IdentityFunctor, ::IdentityFunctor) = F

@inline compose_id(α::Transformation, β::Transformation) = do_compose(α, β)
@inline compose_id(α::Transformation, ::IdentityTransformation) = α
@inline compose_id(::IdentityTransformation, β::Transformation) = β
@inline compose_id(α::IdentityTransformation, ::IdentityTransformation) = α

@inline composeH_id(α::Transformation, β::Transformation) = do_composeH(α, β)
@inline composeH_id(α::Transformation, ::IdIdTransformation) = α
@inline composeH_id(::IdIdTransformation, β::Transformation) = β
@inline composeH_id(α::IdIdTransformation, ::IdIdTransformation) = α

@inline composeH_id(α::Transformation, H::Functor) = do_composeH(α, H)
@inline composeH_id(α::Transformation, ::IdentityFunctor) = α
@inline composeH_id(α::IdentityTransformation, H::Functor) =
  id(compose_id(dom(α), H))
@inline composeH_id(α::IdentityTransformation, ::IdentityFunctor) = α

@inline composeH_id(F::Functor, β::Transformation) = do_composeH(F, β)
@inline composeH_id(::IdentityFunctor, β::Transformation) = β
@inline composeH_id(F::Functor, β::IdentityTransformation) =
  id(compose_id(F, dom(β)))
@inline composeH_id(::IdentityFunctor, β::IdentityTransformation) = β

do_compose(F::Functor, G::Functor) = CompositeFunctor(F, G)

@inline function do_composeH(α::Transformation, β::Transformation)
  do_composeH(α, β, Val{:covariant})
end
function do_composeH(α::Transformation, β::Transformation,
                     ::Type{Val{:covariant}})
  G, H = codom(α), dom(β)
  compose_id(composeH_id(α, H), composeH_id(G, β))
end
function do_composeH(α::Transformation, β::Transformation,
                     ::Type{Val{:contravariant}})
  F, K = dom(α), codom(β)
  compose_id(composeH_id(F, β), composeH_id(α, K))
end

# Oppositization 2-functor
#-------------------------

""" Opposite category, where morphism are reversed.

Call `op(::Cat)` instead of directly instantiating this type.
"""
@auto_hash_equals struct OppositeCat{Ob,Hom,Size<:CatSize,C<:Cat{Ob,Hom,Size}} <:
    Cat{Ob,Hom,Size}
  cat::C
end

ob(C::OppositeCat, x) = ob(C.cat, x)
hom(C::OppositeCat, f) = hom(C.cat, f)

dom(C::OppositeCat, f) = codom(C.cat, f)
codom(C::OppositeCat, f) = dom(C.cat, f)
id(C::OppositeCat, x) = id(C.cat, x)
compose(C::OppositeCat, f, g) = compose(C.cat, g, f)

""" Opposite functor, given by the same mapping between opposite categories.

Call `op(::Functor)` instead of directly instantiating this type.
"""
@auto_hash_equals struct OppositeFunctor{C,D,F<:Functor{C,D}} <: Functor{C,D}
    # XXX: Requires more type parameters: ObC, HomC, ObD, HomD.
    #Functor{OppositeCat{C},OppositeCat{D}}
  func::F
end

dom(F::OppositeFunctor) = op(dom(F.func))
codom(F::OppositeFunctor) = op(codom(F.func))

do_ob_map(F::OppositeFunctor, x) = ob_map(F.func, x)
do_hom_map(F::OppositeFunctor, f) = hom_map(F.func, f)

do_compose(F::OppositeFunctor, G::OppositeFunctor) =
  OppositeFunctor(do_compose(F.func, G.func))

#= Not yet needed because the only natural transformations we currently support
#are `FinTransformationMap`, for which can just implement `op` directly.

""" Opposite natural transformation between opposite functors.

Call `op(::Transformation)` instead of directly instantiating this type.
"""
@auto_hash_equals struct OppositeTransformation{C,D,F,G,T<:Transformation{C,D,F,G}} <: Transformation{C,D,F,G}
    # XXX: Requires more type parameters: ObC, HomC, ObD, HomD.
    #Transformation{OppositeCat{C},OppositeCat{D},OppositeFunctor{C,D,G},OppositeFunctor{C,D,F}}
  trans::T
end

dom(α::OppositeTransformation) = op(codom(α.trans))
codom(α::OppositeTransformation) = op(dom(α.trans))

component(α::OppositeTransformation, x) = component(α.trans, x)

do_compose(α::OppositeTransformation, β::OppositeTransformation) =
  OppositeTransformation(do_compose(β.trans, α.trans))
do_composeH(α::OppositeTransformation, β::OppositeTransformation) =
  OppositeTransformation(do_composeH(α.trans, β.trans))
do_composeH(F::OppositeFunctor, β::OppositeTransformation) =
  OppositeTransformation(do_composeH(F.func, β.trans))
do_composeH(α::OppositeTransformation, H::OppositeFunctor) =
  OppositeTransformation(do_composeH(α.trans, H.func))
=#

""" Oppositization 2-functor.

The oppositization endo-2-functor on Cat, sending a category to its opposite, is
covariant on objects and morphisms and contravariant on 2-morphisms, i.e., is a
2-functor ``op: Catᶜᵒ → Cat``. For more explanation, see the
[nLab](https://ncatlab.org/nlab/show/opposite+category).
"""
op(C::Cat) = OppositeCat(C)
op(F::Functor) = OppositeFunctor(F)
#op(α::Transformation) = OppositeTransformation(α)
op(C::OppositeCat) = C.cat
op(F::OppositeFunctor) = F.func
#op(α::OppositeTransformation) = α.trans

""" 2-cell dual of a 2-category.
"""
function co end

end


using NonlinearEigenproblems: nep_gallery, compute_Mder, contour_block_SS, contour_beyn
using FEASTSolver

using SuiteSparse

SuiteSparse.UMFPACK.umf_ctrl[8] = 0;

function info(Λ, X, residuals, c, r)
    inside(x) = in_contour(x, c, r)
    in_eig = Λ[inside.(Λ)]
    in_res = residuals[inside.(Λ)]
    print("number inside : ")
    println(size(Λ[inside.(Λ)])[1])
    if sum(inside.(Λ)) > 0
        in_res_conv = in_res[in_res .<= 1e-5]
        in_eig_conv = in_eig[in_res .<= 1e-5]
        print("number inside converged : ")
        println(size(in_eig_conv)[1])
        print("max res inside: ")
        println(maximum(residuals[inside.(Λ)]))
        if size(in_res_conv, 1) > 0
            print("max res inside non spurious: ")
            println(maximum(in_res_conv))
        end
    end
    # display(in_eig_conv)
end

nep = nep_gallery("nlevp_native_gun")
T(x) = compute_Mder(nep, x)

n = size(nep, 1)
C = complex(62500.0, 0.0)
R = 50000

e, v, res = nlfeast!(T, rand(ComplexF64,n,32), 2^3, 10, c=C, r=R, debug=true, ϵ=10e-10, store=true)
# e, v, res = nlfeast!(T, rand(ComplexF64,n,50), 2^4, 10, c=C, r=2*R, debug=true, ϵ=10e-16, store=false)
#
# e, v, res = nlfeast_moments!(T, rand(ComplexF64,n,4), 2^5, 0, moments=8, c=C, r=R, debug=true, ϵ=10e-16, store=false)
info(e,v,res,C,R)

# e, v = contour_beyn(nep, radius=R, σ=C, N=2^8, k=80)
# e, v = contour_block_SS(nep, radius=2*R, σ=C, N=2^8, k=80, K=4)
# print(e)


struct UKFTUIntermediate{T,X,TS,AS<:Union{Matrix{T}, Augmented{Matrix{T}, Matrix{T}}}}
    P_chol::AS
    xi_temp::X
    transformed_sigma_points::TS
    unbiased_sigma_points::TS
    x_apri::Vector{T}
    p_apri::Matrix{T}
end

UKFTUIntermediate(T::Type, num_x::Number) =
    UKFTUIntermediate(
        Matrix{T}(undef, num_x, num_x),
        Vector{T}(undef, num_x),
        TransformedSigmaPoints(Vector{T}(undef, num_x), Matrix{T}(undef, num_x, 2 * num_x), MeanSetWeightingParameters(0.0)), # Weighting parameters will be reset
        TransformedSigmaPoints(Vector{T}(undef, num_x), Matrix{T}(undef, num_x, 2 * num_x), MeanSetWeightingParameters(0.0)),
        Vector{T}(undef, num_x),
        Matrix{T}(undef, num_x, num_x)
    )

UKFTUIntermediate(num_x::Number) = UKFTUIntermediate(Float64, num_x)

struct UKFMUIntermediate{T,X,TS,AS<:Union{Matrix{T}, Augmented{Matrix{T}, Matrix{T}}}}
    P_chol::AS
    xi_temp::X
    y_est::Vector{T}
    transformed_sigma_points::TS
    unbiased_sigma_points::TS
    ỹ::Vector{T}
    innovation_covariance::Matrix{T}
    cross_covariance::Matrix{T}
    s_chol::Matrix{T}
    kalman_gain::Matrix{T}
    x_posterior::Vector{T}
    p_posterior::Matrix{T}
end

function UKFMUIntermediate(T::Type, num_x::Number, num_y::Number)
    UKFMUIntermediate(
        Matrix{T}(undef, num_x, num_x),
        Vector{T}(undef, num_x),
        Vector{T}(undef, num_y),
        TransformedSigmaPoints(Vector{T}(undef, num_y), Matrix{T}(undef, num_y, 2 * num_x), MeanSetWeightingParameters(0.0)), # Weighting parameters will be reset
        TransformedSigmaPoints(Vector{T}(undef, num_y), Matrix{T}(undef, num_y, 2 * num_x), MeanSetWeightingParameters(0.0)),
        Vector{T}(undef, num_y),
        Matrix{T}(undef, num_y, num_y),
        Matrix{T}(undef, num_x, num_y),
        Matrix{T}(undef, num_y, num_y),
        Matrix{T}(undef, num_x, num_y),
        Vector{T}(undef, num_x),
        Matrix{T}(undef, num_x, num_x)
    )
end

UKFMUIntermediate(num_x::Number, num_y::Number) = UKFMUIntermediate(Float64, num_x, num_y)

sigmapoints(tu::SPTimeUpdate) = tu.χ
sigmapoints(tu::SPMeasurementUpdate) = tu.𝓨

function time_update(x, P, f, Q, weight_params::AbstractWeightingParameters = WanMerweWeightingParameters())
    χₖ₋₁ = calc_sigma_points(x, P, weight_params)
    χₖ₍ₖ₋₁₎ = transform(f, χₖ₋₁)
    x_apri = mean(χₖ₍ₖ₋₁₎)
    unbiased_χₖ₍ₖ₋₁₎ = substract_mean(χₖ₍ₖ₋₁₎, x_apri)
    P_apri = cov(unbiased_χₖ₍ₖ₋₁₎, Q)
    SPTimeUpdate(x_apri, P_apri, χₖ₍ₖ₋₁₎)
end

function time_update!(tu::UKFTUIntermediate, x, P, f!, Q, weight_params::AbstractWeightingParameters = WanMerweWeightingParameters())
    χₖ₋₁ = calc_sigma_points!(tu.P_chol, x, P, weight_params)
    χₖ₍ₖ₋₁₎ = transform!(tu.transformed_sigma_points, tu.xi_temp, f!, χₖ₋₁)
    x_apri = mean!(tu.x_apri, χₖ₍ₖ₋₁₎)
    unbiased_χₖ₍ₖ₋₁₎ = substract_mean!(tu.unbiased_sigma_points, χₖ₍ₖ₋₁₎, x_apri)
    P_apri = cov!(tu.p_apri, unbiased_χₖ₍ₖ₋₁₎, Q)
    SPTimeUpdate(x_apri, P_apri, χₖ₍ₖ₋₁₎)
end

function measurement_update(
    x,
    P,
    y,
    h,
    R;
    weight_params::AbstractWeightingParameters = WanMerweWeightingParameters(),
    consider = nothing
)
    χₖ₍ₖ₋₁₎ = calc_sigma_points(x, P, weight_params)
    𝓨 = transform(h, χₖ₍ₖ₋₁₎)
    y_est = mean(𝓨)
    unbiased_𝓨 = substract_mean(𝓨, y_est)
    S = cov(unbiased_𝓨, R)
    ỹ = y .- y_est
    Pᵪᵧ = cov(χₖ₍ₖ₋₁₎, unbiased_𝓨)
    K, P_posterior = calc_kalman_gain_and_posterior_covariance(P, Pᵪᵧ, S, consider)
    x_posterior = calc_posterior_state(x, K, ỹ, consider)
    SPMeasurementUpdate(x_posterior, P_posterior, 𝓨, ỹ, S, K)
end

function calc_kalman_gain_and_posterior_covariance(P, Pᵪᵧ, S, consider)
    K = calc_kalman_gain(Pᵪᵧ, S, consider)
    P_posterior = calc_posterior_covariance(P, Pᵪᵧ, K, consider)
    K, P_posterior
end

function measurement_update!(mu::UKFMUIntermediate, x, P, y, h!, R; weight_params::AbstractWeightingParameters = WanMerweWeightingParameters(1e-3, 2, 0))
    χₖ₍ₖ₋₁₎ = calc_sigma_points!(mu.P_chol, x, P, weight_params)
    𝓨 = transform!(mu.transformed_sigma_points, mu.xi_temp, h!, χₖ₍ₖ₋₁₎)
    y_est = mean!(mu.y_est, 𝓨)
    unbiased_𝓨 = substract_mean!(mu.unbiased_sigma_points, 𝓨, y_est)
    S = cov!(mu.innovation_covariance, unbiased_𝓨, R)
    Pᵪᵧ = cov!(mu.cross_covariance, χₖ₍ₖ₋₁₎, unbiased_𝓨)
    mu.ỹ .= y .- y_est
    K, P_posterior = calc_kalman_gain_and_posterior_covariance!(mu.s_chol, mu.kalman_gain, mu.p_posterior, P, Pᵪᵧ, S)
    x_posterior = calc_posterior_state!(mu.x_posterior, x, K, mu.ỹ)
    SPMeasurementUpdate(x_posterior, P_posterior, 𝓨, mu.ỹ, S, K)
end

function calc_kalman_gain_and_posterior_covariance!(s_chol, kalman_gain, p_post, P, Pᵪᵧ, S)
    K = calc_kalman_gain!(s_chol, kalman_gain, Pᵪᵧ, S)
    P_posterior = calc_posterior_covariance!(p_post, P, Pᵪᵧ, K)
    K, P_posterior
end


import Base: +, -, *, /
import Base.Broadcast: broadcasted

function +(x1::CuArray{T,N}, x2::CuArray{T,N}) where {T,N}
    @assert length(x1) == length(x2)
    y = fill!(similar(x1), 0)
    axpy!(T(1), x1, y)
    axpy!(T(1), x2, y)
    y
end

function -(x1::CuArray{T,N}, x2::CuArray{T,N}) where {T,N}
    @assert length(x1) == length(x2)
    y = fill!(similar(x1), 0)
    axpy!(T(1), x1, y)
    axpy!(T(-1), x2, y)
    y
end

function -(x::CuArray{T}) where T
    y = fill!(similar(x1), 0)
    axpy!(T(-1), x, y)
    y
end

*(A::CuMatrix{T}, x::CuVector{T}) where T = gemv('N', T(1), A, x)
*(A::CuMatrix{T}, B::CuMatrix{T}) where T = gemm('N', 'N', T(1), A, B)

function broadcasted(::typeof(+), x1::CuArray{T}, x2::CuArray{T}) where T
    x1, x2 = length(x1) >= length(x2) ? (x1,x2) : (x2,x1)
    y = copy(x1)
    CUDNN.add!(1, x2, 1, y)
    y
end

function broadcasted(::typeof(-), x1::CuArray{T}, x2::CuArray{T}) where T
    x1, x2 = length(x1) >= length(x2) ? (x1,x2) : (x2,x1)
    y = copy(x1)
    CUDNN.add!(-1, x2, 1, y)
    y
end

function broadcast_addto!(dest::CuArray{T}, src::CuArray{T}) where T
    CUDNN.add!(1, src, 1, dest)
end

function broadcasted(::typeof(*), x1::CuArray{T}, x2::CuArray{T}) where T
    elemtimes(x1, x2)
end

@generated function elemtimes(x1::CuArray{T,N}, x2::CuArray{T,N}) where {T,N}
    Ct = cstring(T)
    k = Kernel("""
    __global__ void elemtimes($Ct *y, Array<$Ct,$N> x1, Array<$Ct,$N> x2) {
        int idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx >= x1.length()) return;
        int ndidxs[$N];
        x1.ndindex(ndidxs, idx);
        y[idx] = x1[idx] * x2(ndidxs);
    }
    """)
    quote
        length(x1) < length(x2) && return elemtimes(x2, x1)
        y = similar(x1)
        gdims, bdims = cudims(length(y))
        $k(gdims, bdims, pointer(y), x1, x2)
        y
    end
end


"""
`Transition{T<:Real}`

The transition equation of the state-space model takes the form

    `s_t = TTT*s_{t-1} + RRR*ϵ_t + CCC`

The `Transition` type stores the coefficient `Matrix{T}`s (`TTT`, `RRR`) and constant `Vector{T} CCC`.
"""
mutable struct Transition{T<:Real}
    TTT::Matrix{T}
    RRR::Matrix{T}
    CCC::Vector{T}
end
function Transition(TTT::Matrix{T}, RRR::Matrix{T}) where T<:Real
    CCC = zeros(eltype(TTT), size(TTT, 1))
    Transition{T}(TTT, RRR, CCC)
end
function Transition(TTT::Matrix{T}, RRR::Matrix{T}, CCC::Matrix{T}) where T<:Real
    Transition{T}(TTT, RRR, collect(CCC))
end
function Base.getindex(eq::Transition, d::Symbol)
    if d in (:TTT, :RRR, :CCC)
        return getfield(eq, d)
    else
        throw(KeyError(d))
    end
end

"""
`Measurement{T<:Real}`

The measurement equation of the state-space model takes the form

   `y_t = ZZ*s_t + DD + u_t`

where the error `u_t` is the measurement error, which is uncorrelated with the
shocks in the transition equation `ϵ_t`.

### Fields

If `Ns` is the number of states `s_t`, `Ny` is the number of
observables `y_t`, and `Ne` is the number of shocks `ϵ_t`:

- `ZZ`: the `Ny` x `Ns`  measurement matrix
- `DD`: the `Ny` x 1 constant vector
- `QQ`: the `Ne` x `Ne` covariance matrix for the shocks `ϵ_t`
- `EE`: the `Ny` x `Ny` covariance matrix for the measurement error `η_t`
"""
mutable struct Measurement{T<:Real}
    ZZ::Matrix{T}
    DD::Vector{T}
    QQ::Matrix{T}
    EE::Matrix{T}
end

function Base.getindex(M::Measurement, d::Symbol)
    if d in (:ZZ, :DD, :QQ, :EE)
        return getfield(M, d)
    else
        throw(KeyError(d))
    end
end

function measurement(m::AbstractDSGEModel, trans::Transition; shocks::Bool=true)
    TTT = trans[:TTT]
    RRR = trans[:RRR]
    CCC = trans[:CCC]
    measurement(m, TTT, RRR, CCC; shocks=shocks)
end

"""
```
PseudoMeasurement{T<:Real}
```

The pseudo-measurement equation of the state-space model takes the form

   `x_t = ZZ_pseudo*s_t + DD_pseudo`

### Fields

Let `Ns` be the number of states `s_t` and `Nx` be the number of
pseudo-observables `x_t`:

- `ZZ_pseudo`: the `Nx` x `Ns` pseudo-measurement matrix
- `DD_pseudo`: the `Nx` x 1 constant vector
"""
mutable struct PseudoMeasurement{T<:Real}
    ZZ_pseudo::Matrix{T}
    DD_pseudo::Vector{T}
end

function Base.getindex(M::PseudoMeasurement, d::Symbol)
    if d in (:ZZ_pseudo, :DD_pseudo)
        return getfield(M, d)
    else
        throw(KeyError(d))
    end
end

abstract type AbstractSystem{T} end

"""
`System{T <: Real}`

A mutable struct containing the transition and measurement equations for a
state-space model. The matrices may be directly indexed: `sys[:TTT]`
returns `sys.transition.TTT`, `sys[:ZZ]` returns `sys.measurement.ZZ`, etc.
"""
mutable struct System{T <: Real} <: AbstractSystem{T}
    transition::Transition{T}
    measurement::Measurement{T}
    pseudo_measurement::PseudoMeasurement{T}
end

function System(transition::Transition{T}, measurement::Measurement{T}) where T<:Real
    # Initialize empty pseudo-measurement equation
    _n_states = size(transition.TTT, 1)
    _n_pseudo = 0
    ZZ_pseudo = zeros(_n_pseudo, _n_states)
    DD_pseudo = zeros(_n_pseudo)
    pseudo_measurement = PseudoMeasurement(ZZ_pseudo, DD_pseudo)

    return System(transition, measurement, pseudo_measurement)
end

function Base.getindex(system::System, d::Symbol)
    if d in (:transition, :measurement, :pseudo_measurement)
        return getfield(system, d)
    elseif d in fieldnames(typeof(system.transition))
        return getfield(system.transition, d)
    elseif d in fieldnames(typeof(system.measurement))
        return getfield(system.measurement, d)
    elseif d in fieldnames(typeof(system.pseudo_measurement))
        return getfield(system.pseudo_measurement, d)
    else
        throw(KeyError(d))
    end
end

function Base.copy(system::System)
    trans = Transition(system[:TTT], system[:RRR], system[:CCC])
    meas  = Measurement(system[:ZZ], system[:DD], system[:QQ], system[:EE])
    pseudo_meas = PseudoMeasurement(system[:ZZ_pseudo], system[:DD_pseudo])
    return System(trans, meas, pseudo_meas)
end

"""
`RegimeSwitchingSystem{T <: Real}`

A mutable struct containing the transition and measurement equations for a
state-space model with regime-switching.
The matrices may be directly indexed: `sys[1, :TTT]`
returns `sys.regime[1].transition.TTT`, etc.
"""
mutable struct RegimeSwitchingSystem{T <: Real} <: AbstractSystem{T}
    transitions::Vector{Transition{T}}
    measurements::Vector{Measurement{T}}
    pseudo_measurements::Vector{PseudoMeasurement{T}}

    RegimeSwitchingSystem{T}(transitions, measurements, pseudo_measurements) where {T <: Real} = length(transitions) == length(measurements) ?
        new(transitions, measurements, pseudo_measurements) :
        error("The number of Transition (n = $(length(transitions)))" *
              " and Measurement (n = $(length(measurements))) objects must match.")
end

RegimeSwitchingSystem(transitions::Vector{Transition{T}}, measurements::Vector{Measurement{T}},
                      pseudo_measurements::Vector{PseudoMeasurement{T}}) where {T <: Real} =
                          RegimeSwitchingSystem{T}(transitions, measurements, pseudo_measurements)

function RegimeSwitchingSystem(transitions::Vector{Transition{T}},
                               measurements::Vector{Measurement{T}}) where {T<:Real}

    # Initialize empty pseudo-measurement equation
    _n_states = size(transitions[1].TTT, 1)
    _n_pseudo = 0
    ZZ_pseudo = zeros(_n_pseudo, _n_states)
    DD_pseudo = zeros(_n_pseudo)
    pseudo_measurement = PseudoMeasurement(ZZ_pseudo, DD_pseudo)
    pseudo_measurements = [pseudo_measurement for i in 1:length(transitions)]

    return RegimeSwitchingSystem(transitions, measurements, pseudo_measurements)
end

function RegimeSwitchingSystem(systems::Vector{System{T}}) where {T <: Real}
    n_regimes           = length(systems)
    transitions         = Vector{Transition{T}}(undef, n_regimes)
    measurements        = Vector{Measurement{T}}(undef, n_regimes)
    pseudo_measurements = Vector{PseudoMeasurement{T}}(undef, n_regimes)
    for i in 1:n_regimes
        transitions[i]         = systems[i].transition
        measurements[i]        = systems[i].measurement
        pseudo_measurements[i] = systems[i].pseudo_measurement
    end
    return RegimeSwitchingSystem(transitions, measurements, pseudo_measurements)
end

# Retrieve System correspond to a regime
function System(system::RegimeSwitchingSystem, regime::Int)
    return System(system.transitions[regime],
                  system.measurements[regime],
                  system.pseudo_measurements[regime])
end

# Get fields of the RegimeSwitchingSystem
function Base.getindex(system::RegimeSwitchingSystem{T},
                       d::Symbol) where {T<:Real}
    if d in (:transitions, :measurements, :pseudo_measurements)
        return getfield(system, d)
    elseif d == :regimes
        return 1:length(system.transitions)
    else
        throw(KeyError(d))
    end
end

# Get a specific regime
function Base.getindex(system::RegimeSwitchingSystem{T},
                       d::Int) where {T<:Real}
    if d < 1 || d > length(system.transitions)
        throw(BoundsError(system.transitions,d))
    else
        return System(system, d)
    end
end

# Get specific matrix or type from specific regime
function Base.getindex(system::RegimeSwitchingSystem{T},
                       regime::Int, d::Symbol) where {T <: Real}
    if d == :transition
        system.transitions[regime]
    elseif d == :measurement
        system.measurements[regime]
    elseif d == :pseudo_measurement
        system.pseudo_measurements[regime]
    elseif d in (:TTT, :RRR, :CCC)
        system.transitions[regime][d]
    elseif d in (:ZZ, :DD, :QQ, :EE)
        system.measurements[regime][d]
    elseif d in (:ZZ_pseudo, :DD_pseudo)
        system.pseudo_measurements[regime][d]
    else
        throw(KeyError(d))
    end
end

function n_regimes(system::RegimeSwitchingSystem{T}) where {T<:Real}
    return length(system.transitions)
end

function Base.copy(system::RegimeSwitchingSystem{T}) where {T<:Real}
    n_reg = n_regimes(system)
    transitions         = Vector{Transition{T}}(undef, n_reg)
    measurements        = Vector{Measurement{T}}(undef, n_reg)
    pseudo_measurements = Vector{PseudoMeasurement{T}}(undef, n_reg)

    for i in 1:length(system[:transitions])
        trans = Transition(system[i,:TTT], system[i,:RRR], system[i,:CCC])
        meas  = Measurement(system[i,:ZZ], system[i,:DD],
                            system[i,:QQ], system[i,:EE])
        pseudo_meas = PseudoMeasurement(system[i,:ZZ_pseudo], system[i,:DD_pseudo])

        transitions[i]         = trans
        measurements[i]        = meas
        pseudo_measurements[i] = pseudo_meas
    end
    return RegimeSwitchingSystem(transitions, measurements, pseudo_measurements)
end

function Base.deepcopy(system::RegimeSwitchingSystem{T}) where {T<:Real}
    n_reg = n_regimes(system)
    transitions         = Vector{Transition{T}}(undef, n_reg)
    measurements        = Vector{Measurement{T}}(undef, n_reg)
    pseudo_measurements = Vector{PseudoMeasurement{T}}(undef, n_reg)

    for i in 1:length(system[:transitions])
        trans = Transition(deepcopy(system[i,:TTT]), deepcopy(system[i,:RRR]), deepcopy(system[i,:CCC]))
        meas  = Measurement(deepcopy(system[i,:ZZ]), deepcopy(system[i,:DD]),
                            deepcopy(system[i,:QQ]), deepcopy(system[i,:EE]))
        pseudo_meas = PseudoMeasurement(deepcopy(system[i,:ZZ_pseudo]), deepcopy(system[i,:DD_pseudo]))

        transitions[i]         = trans
        measurements[i]        = meas
        pseudo_measurements[i] = pseudo_meas
    end
    return RegimeSwitchingSystem(transitions, measurements, pseudo_measurements)
end

"""
`TimeVaryingInformationSetSystem`

A mutable struct containing the transition and measurement equations for a
state-space model with regime-switching and time-varying information sets.
"""
mutable struct TimeVaryingInformationSetSystem{T <: Real} <: AbstractSystem{T}
    transitions::Vector{Vector{Transition{T}}}
    measurements::Vector{Measurement{T}}
    pseudo_measurements::Vector{PseudoMeasurement{T}}
    information_set::Vector{UnitRange{Int}}
    select::Vector{Int}
end

function TimeVaryingInformationSetSystem(transitions::Vector{Vector{Transition{T}}},
                                         measurements::Vector{Measurement{T}},
                                         information_set::Vector{UnitRange{Int}},
                                         select::Vector{Int}) where {T <: Real}

    # Initialize empty pseudo-measurement equation
    _n_states = size(transitions[select[1]][1].TTT, 1)
    _n_pseudo = 0
    ZZ_pseudo = zeros(_n_pseudo, _n_states)
    DD_pseudo = zeros(_n_pseudo)
    pseudo_measurement = PseudoMeasurement(ZZ_pseudo, DD_pseudo)
    pseudo_measurements = [pseudo_measurement for i in 1:length(measurements)]

    return TimeVaryingInformationSetSystem(transitions, measurements, pseudo_measurements, information_set, select)
end

function TimeVaryingInformationSetSystem(transitions::Vector{Vector{Transition{T}}},
                                         measurements::Vector{Measurement{T}},
                                         pseudo_measurements::Vector{PseudoMeasurement{T}},
                                         information_set::Vector{UnitRange{Int}},
                                         select::Vector{Int}) where {T <: Real}

    # Initialize empty pseudo-measurement equation
    return TimeVaryingInformationSetSystem{T}(transitions, measurements, pseudo_measurements, information_set, select)
end

function System(system::TimeVaryingInformationSetSystem, select::Int, regime::Int)
    return System(system.transitions[select][regime],
                  system.measurements[regime],
                  system.pseudo_measurements[regime])
end

function Base.getindex(system::TimeVaryingInformationSetSystem, d::Symbol)
    if d in (:transitions, :measurements, :pseudo_measurements, :information_set, :select)
        return getfield(system, d)
    elseif d == :regimes
        return 1:length(system.measurements)
    else
        throw(KeyError(d))
    end
end

function Base.getindex(system::TimeVaryingInformationSetSystem, regime::Int,
                       d::Symbol)
    if d == :measurement
        system.measurements[regime]
    elseif d == :pseudo_measurement
        system.pseudo_measurements[regime]
    elseif d in (:ZZ, :DD, :QQ, :EE)
        system.measurements[regime][d]
    elseif d in (:ZZ_pseudo, :DD_pseudo)
        system.pseudo_measurements[regime][d]
    else
        throw(KeyError(d))
    end
end

# Get a specific regime
function Base.getindex(system::TimeVaryingInformationSetSystem{T},
                       select::Int, d::Int) where {T <: Real}
    if d < 1 || d > length(system.measurements)
        throw(BoundsError(system.measurements, d))
    else
        return System(system, select, d)
    end
end

# Get specific matrix or type from specific regime
function Base.getindex(system::TimeVaryingInformationSetSystem{T},
                       select::Int, regime::Int, d::Symbol) where {T <: Real}
    if d == :transition
        system.transitions[select][regime]
    elseif d in (:TTT, :RRR, :CCC)
        system.transitions[select][regime][d]
    else
        throw(KeyError(d))
    end
end

function n_regimes(system::TimeVaryingInformationSetSystem)
    return length(system.measurements)
end


function convertFromGensysIn(g0,g1,pi)

    gDim  = size(g0, 1)
    piCol = size(pi, 2)

    theHM = hcat(-g1, zeros(gDim, piCol))
    theHM = vcat(theHM, zeros(piCol, gDim + piCol))

    theH0 = hcat(g0, -pi)
    theH0 = vcat(theH0, zeros(piCol, gDim + piCol))
    #theH0=sparse([...
    #              g0,-pi;...
    #              zeros(piCol,gDim+piCol)]);

    matrix = Compat.Matrix(I, piCol, piCol)
    theHP = hcat(zeros(piCol, gDim), matrix)
    theHP = vcat(zeros(gDim, gDim + piCol), theHP)
    #theHP=sparse([...
    #             zeros(gDim,gDim+piCol);...
    #              zeros(piCol,gDim),eye(piCol)]);

    return theHM, theH0, theHP

end # function


import MathOptCartesians
const MOC = MathOptCartesians

using Test

import MathOptInterface
const MOI = MathOptInterface

@testset "Scalar set construction" begin
    set = MOC.ScalarCartesianSet(MOI.LessThan(0.0), MOI.GreaterThan(1.3))
    @test MOI.dimension(set) == 2
end

@testset "Vector set construction and dimensions" begin
    set = MOC.VectorCartesianSet(MOI.Zeros(4), MOI.Nonnegatives(2))
    @test MOI.dimension(set) == 6
    @test set.start_indices == (1, 5)
    @test set.dims == (4, 2)
end

@testset "Cartesian dual" begin
    set = MOC.VectorCartesianSet(MOI.Zeros(4), MOI.Nonnegatives(2))
    dset = MOI.dual_set(set)
    @test MOI.dimension(dset) == MOI.dimension(set)
    @test dset.sets[1] == MOI.Reals(4)
    @test dset.sets[2] == MOI.Nonnegatives(2)
end


# Set up the environment to run this example. Make sure you're within the folder that this
# file lives in.
using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

using AbstractGPs, Plots, Random, Stheno
gr();

# Define model.
σ², ω, T = 1e0, 1.0, 25.0;
f = @gppp let
    f1 = periodic(GP(SEKernel()), ω)
    f2 = GP(0.1 * SEKernel())
    f3 = f1 + f2
end;

# Sample from marginal process to generate toy data.
rng = MersenneTwister(123456);
S = 25;
x = GPPPInput(:f3, range(0.0, T; length=300));
fx = f(x, σ²);
y = rand(rng, fx);

xp_ = range(-2.5, T + 2.5; length=500);

# Plots for stuff.
f₁′_plot, f₂′_plot, f₃′_plot = plot(), plot(), plot();

posterior_plot = plot();
approx_in_marginal_posterior_plot = plot();
approx_in_latents_posterior_plot = plot();

##################################### Exact Inference ######################################

# Compute the posterior processes, sample from them, and compute marginals.
@show logpdf(fx, y)
f′ = posterior(fx, y);
xp_f1 = GPPPInput(:f1, xp_);
xp_f2 = GPPPInput(:f2, xp_);
xp_f3 = GPPPInput(:f3, xp_);
xp = BlockData(xp_f1, xp_f2, xp_f3);
f′xp = rand(f′(xp, 1e-9), S);
f₁′xp, f₂′xp, f₃′xp = split(xp, f′xp);

items = [
    (xp_f1, f₁′xp, :red, "exact", f₁′_plot),
    (xp_f2, f₂′xp, :red, "", f₂′_plot),
    (xp_f3, f₃′xp, :red, "", f₃′_plot),
];

# Plot posterior marginals and samples.
for (x, samples, colour, name, plt) in items
    plot!(plt, xp_, f′(x); fillalpha=0.2, color=colour, label=name);
    plot!(plt, xp_, samples; alpha=0.3, color=colour, label="", linewidth=1);
end

# Plot observations.
scatter!(posterior_plot, x.x, y;
    markercolor=:blue,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.7,
    label="",
);

##################### Approximate inference with pseudo-data in f3 #########################

# Compute approximate posterior process.
M3 = 25;
z = GPPPInput(:f3, collect(range(0, T; length=M3)));
u = f(z, 1e-9);
f′_approx = approx_posterior(VFE(), fx, y, u);

@show elbo(fx, y, u);

items = [
    (xp_f1, :green, "Z in f3", f₁′_plot),
    (xp_f2, :green, "", f₂′_plot),
    (xp_f3, :green, "", f₃′_plot),
];

# Posterior marginal variance.
for (x, colour, name, plt) in items
    plot!(plt, xp_, f′_approx(x); fillalpha=0.3, fillcolor=colour, label=name);
end

# Plot observations and pseudo-input locations.
scatter!(approx_in_marginal_posterior_plot, x.x, y;
    markercolor=:blue,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.7,
    label="",
);
scatter!(approx_in_marginal_posterior_plot, z.x, zeros(M3);
    markercolor=:black,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.8,
    label="Z",
);

############# Perform approximate inference by placing pseudo-data in f1 and f2 ############

# Compute approximate posterior process.
M1, M2 = 15, 10;
z1 = GPPPInput(:f1, collect(range(0.0; step=1 / ω, length=M1)));
z2 = GPPPInput(:f2, collect(range(0.0; step=T, length=M2)));
z12 = BlockData(z1, z2);
u12 = f(z12, 1e-9);
f′_approx_12 = approx_posterior(VFE(), fx, y, u12);

@show elbo(fx, y, u12);

items = [
    (xp_f1, :blue, "Z in f1 and f2", f₁′_plot),
    (xp_f2, :blue, "", f₂′_plot),
    (xp_f3, :blue, "", f₃′_plot),
];

# Plot posterior marginals.
for (x, colour, name, plt) in items
    plot!(plt, xp_, f′_approx_12(x); fillalpha=0.3, fillcolor=colour, label=name)
end

# Plot observations and pseudo-input locations.
scatter!(approx_in_latents_posterior_plot, x.x, y;
    markercolor=:blue,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.7,
    label="",
);
scatter!(approx_in_latents_posterior_plot, z1.x, zeros(M1);
    markercolor=:black,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.8,
    label="Z₁",
);
scatter!(approx_in_latents_posterior_plot, z2.x, zeros(M2);
    markercolor=:magenta,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.8,
    label="Z₂",
);

# display(posterior_plot);
# display(approx_in_marginal_posterior_plot);
# display(approx_in_latents_posterior_plot);

plot(f₁′_plot, f₂′_plot, f₃′_plot; layout=(3, 1))


######################################################
#
# MSM Plot
#
# Date:   2014-12-02
# Author: Hessel Tuinhof
#
# Simulate and plot a Markov-Switching Multifractal (MSM)
#
# Required packages: 	PyPlot
# Required functions:	MSM.jl
#
######################################################

# load packages & functions
using PyPlot
require("../src/msm/MSM.jl")

# define colors
rouge = (0.79, 0.0, 0.1)

# configure parameters
num_sim = 1
sim_length = 5000
kbar = 8
m0 = 1.4
sigma = 0.5
gamma_kbar = 0.95
b = 3

# set seed
srand(314159265358979323846)

# plot MSM (returns & log price)
figure()
msm = MSM(num_sim, sim_length, m0, kbar, sigma, gamma_kbar, b)
plot(msm, linewidth = 0.5, color = rouge)
xlim(xmax = sim_length)
figure()
plot(cumsum(msm), linewidth = 1, color = rouge)
xlim(xmax = sim_length)


using TenPuzzle
using Test

@testset "Basic fetures" begin
  @testset "Typical problems" begin
      @test maketen(10) == "10"
      @test maketen(5, 5) == "5 + 5"

      @test maketen(1, 1, 9, 9) in (
          "(1 + 1 / 9) * 9", "(1 / 9 + 1) * 9",
          "9 * (1 + 1 / 9)", "9 * (1 / 9 + 1)"
      )
  end

end


#=
#define FMPR_RND_DOWN  0    RoundingMode{:Down}()
#define FMPR_RND_UP    1
#define FMPR_RND_FLOOR 2
#define FMPR_RND_CEIL  3
#define FMPR_RND_NEAR  4
=#

function round{P}(x::ArbFloat{P}, sig::Int=P, base::Int=10)
    sig=abs(sig); base=abs(base)
    sigbits = min(P, ceil(Int, (sig * log(base)/log(2.0))))
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_set_round), Void,  (Ptr{ArbFloat}, Ptr{ArbFloat}, Int), &z, &x, sigbits)
    z
end

function ceil{P}(x::ArbFloat{P}, sig::Int=P, base::Int=10)
    sig=abs(sig); base=abs(base)
    sigbits = min(P, ceil(Int, (sig * log(base)/log(2.0))))
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_ceil), Void,  (Ptr{ArbFloat}, Ptr{ArbFloat}, Int), &z, &x, sigbits)
    z
end

function floor{P}(x::ArbFloat{P}, sig::Int=P, base::Int=10)
    sig=abs(sig); base=abs(base)
    sigbits = min(P, ceil(Int, (sig * log(base)/log(2.0))))
    z = initializer(ArbFloat{P})
    ccall(@libarb(arb_floor), Void,  (Ptr{ArbFloat}, Ptr{ArbFloat}, Int), &z, &x, sigbits)
    z
end

function trunc{P}(x::ArbFloat{P}, sig::Int=P, base::Int=10)
    sig=abs(sig); base=abs(base)
    sigbits = min(P, ceil(Int, (sig * log(base)/log(2.0))))
    z = initializer(ArbFloat{P})
    y = abs(x)
    ccall(@libarb(arb_floor), Void,  (Ptr{ArbFloat}, Ptr{ArbFloat}, Int), &z, &y, sigbits)
    signbit(x) ? -z : z
end

function round{T,P}(::Type{T}, x::ArbFloat{P}, sig::Int=P, base::Int=10)
    z = round(x, sig, base)
    convert(T, z)
end
function ceil{T,P}(::Type{T}, x::ArbFloat{P}, sig::Int=P, base::Int=10)
    z = ceil(x, sig, base)
    convert(T, z)
end
function floor{T,P}(::Type{T}, x::ArbFloat{P}, sig::Int=P, base::Int=10)
    z = floor(x, sig, base)
    convert(T, z)
end
function trunc{T,P}(::Type{T}, x::ArbFloat{P}, sig::Int=P, base::Int=10)
    z = trunc(x, sig, base)
    convert(T, z)
end

fld{P}(x::ArbFloat{P}, y::ArbFloat{P}) = convert(Int, floor(x/y))
cld{P}(x::ArbFloat{P}, y::ArbFloat{P}) = convert(Int, ceil(x/y))
div{P}(x::ArbFloat{P}, y::ArbFloat{P}) = convert(Int, trunc(x/y))

rem{P}(x::ArbFloat{P}, y::ArbFloat{P}) = convert(Int, x - div(x,y)*y)
mod{P}(x::ArbFloat{P}, y::ArbFloat{P}) = convert(Int, x - fld(x,y)*y)

function divrem{P}(x::ArbFloat{P}, y::ArbFloat{P})
   d = div(x,y)
   r = convert(Int, x - d*y)
   d,r
end

function fldmod{P}(x::ArbFloat{P}, y::ArbFloat{P})
   d = fld(x,y)
   r = convert(Int, x - d*y)
   d,r
end


export
    OptimalRHS,
    optimize_rhs,
    DPSolution,
    value_iteration,
    value_residual,
    solve_residual,
    nonlinear_solve_value,
    bellman_residual_plot

"""
Object characterizing the optimal solution (a value and a policy) to
the right hand side of a Bellman equation at a given state.
"""
immutable OptimalRHS{TV,TP}
    value::TV
    policy::TP
end

"""
Accessor function for the value.
"""
value(o::OptimalRHS) = o.value

"""
Accessor function for the optimal policy.
"""
policy(o::OptimalRHS) = o.policy

"""
`optimize_rhs(model, value, state)` is the core function for solving
dynamic programming problems.

The user should implement this function for each type of `model` that
will be solved by the library. For a given value function `value` and
state `state`, it should return an OptimalRHS object, containing the
optimal value and the optimal policy.
"""
function optimize_rhs end

"""
Solution to a dynamic programming problem.
"""
@with_kw immutable DPSolution
    model
    value
    policy
    iterations
    converged
end

value(dp::DPSolution) = dp.value
policy(dp::DPSolution) = dp.policy

"""
Solve a dynamic programming problem by value iteration.
"""
function value_iteration(model, V, options = IterationOptions())
    @unpack max_iter = options
    basis = V.basis
    states = collocation_points(basis)
    for iteration in 1:max_iter
        VP = [optimize_rhs(model, V, s) for s in states]
        V′points = value.(VP)
        Δ = V′points - collocation_values(V)
        V′ = basis \ V′points
        converged = _converged(Δ, options)
        if converged || iteration == max_iter
            return DPSolution(model = model,
                              value = V′,
                              policy = basis \ policy.(VP),
                              iterations = iteration,
                              converged = converged)
        else
            V = V′
        end
    end
end

function value_iteration(solution::DPSolution,
                         options = IterationOptions())
    @unpack model, value = solution
    value_iteration(model, value, options)
end

"""
Calculate the residual between the right- and left-hand sides of a
Bellman equation at the given state.
"""
function value_residual(model, value, state)
    optimize_rhs(model, value, state).value - value(state)
end

function value_residual(solution::DPSolution, state)
    @unpack model, value = solution
    value_residual(model, value, state)
end

"""
Solve for `residual(f) = 0` in a linearly interpolated space of
functions.

The collocation points are determined from `initial_f`, which is also
used as a starting point.

`residual_function(model, f, s)` will be called with all `s`
collocation points, and the result is used as the residual. A
nonlinear solver is used to find `f`.

`f` and the result of `nlsolve` are returned as two values.
"""
function solve_residual(model,
                        residual_function,
                        initial_f;
                        nlsolve_options...)
    @unpack basis, α = initial_f
    states = collocation_points(basis)
    function f!(α, residual)
        f = basis * α
        residual .= [residual_function(model, f, state) for state in states]
    end
    result = nlsolve(f!, α; nlsolve_options...)
    basis * result.zero, result
end

function nonlinear_solve_value(solution::DPSolution; nlsolve_options...)
    value, root = solve_residual(solution.model, value_residual,
                                 solution.value; nlsolve_options...)
    policy(state) = optimize_rhs(solution.model, value, state).policy
    DPSolution(model = solution.model,
               value = value,
               policy = value.basis \ policy,
               iterations = root.iterations,
               converged = NLsolve.converged(root))
end

"""
Plot the residual of the Bellman equation of `solution`, using `N`
points, for a univariate solution. `statename` labels the x axis.
"""
function bellman_residual_plot(solution; statename="state", N=100)
    basis = solution.value.basis
    dense_z = linspace(basis, N)
    residual = [value_residual(solution, z) for z in dense_z]
    plot(dense_z, residual, xlab=statename, ylab="residual of Bellman equation",
         legend = false)
    scatter!(collocation_points(basis), zeros(degf(basis)))
end


module TensND

import Base: @pure, eltype
import LinearAlgebra: normalize, dot, tr

using LinearAlgebra, SymPy, Tensors, OMEinsum, Rotations

include("array_utils.jl")
include("bases.jl")
include("tens.jl")
include("tens_isotropic.jl")
include("special_tens.jl")
include("coorsystems.jl")

# function __init__()
# end

end # module


## Bessel
@registerN Bessel AbstractCCOP1 α

"""
    Bessel{α}

Implements the [Bessel](https://dlmf.nist.gov/18.34) polynomials, introduced by [Krall and Frink](https://www.ams.org/journals/tran/1949-065-01/S0002-9947-1949-0028473-1/S0002-9947-1949-0028473-1.pdf) (with `b=2`). The  case `a=2` corresponds to the
[Bessel](https://en.wikipedia.org/wiki/Bessel_polynomials) polynomials of Wikipedia. The Bessel  polynomials are not orthogonal over  a domain of the real  line, rather over an arbitray curve in the complex plane enclosing the  origin.  The weight  function is `ρ(x)=(2πi)^(-1)∑Γ(α)/Γ(α+n-1)(-β/x)^n`,   where `β=2`.

```jldoctest
julia> using Polynomials, SpecialPolynomials

julia> 𝐐 = Rational{Int}
Rational{Int64}

julia> x = variable(Polynomial{𝐐})
Polynomials.Polynomial(x)

julia> [basis(Bessel{3//2, 𝐐}, i)(x) for i in 0:5]
6-element Vector{Polynomial{Rational{Int64}, :x}}:
 Polynomials.Polynomial(1//1)
 Polynomials.Polynomial(1//1 + 3//4*x)
 Polynomials.Polynomial(1//1 + 5//2*x + 35//16*x^2)
 Polynomials.Polynomial(1//1 + 21//4*x + 189//16*x^2 + 693//64*x^3)
 Polynomials.Polynomial(1//1 + 9//1*x + 297//8*x^2 + 1287//16*x^3 + 19305//256*x^4)
 Polynomials.Polynomial(1//1 + 55//4*x + 715//8*x^2 + 10725//32*x^3 + 182325//256*x^4 + 692835//1024*x^5)
```

"""
Bessel
export Bessel
basis_symbol(::Type{<:Bessel{α}}) where {α} = "Cᵅ"
Polynomials.domain(::Type{<:Bessel}) = Polynomials.Interval(0, Inf)

abcde(::Type{<:Bessel{α}}) where {α} = NamedTuple{(:a, :b, :c, :d, :e)}((1, 0, 0, α, 2))

function k1k0(::Type{P}, k::Int) where {α,P<:Bessel{α}}
    k < 0 && return zero(eltype(P)) / 1
    iszero(k) && return (one(eltype(P)) * α) / 2

    val = one(eltype(P))
    val *= (2k + α) * (2k + α - 1)
    val /= (k + α - 1) * 2
    val
end

norm2(::Type{<:Bessel{α}}, n) where {α} =
    -1^(n + α - 1) * Γ(1 + n) * 2^(α - 1) / (Γ(n + α - 2) * (2n + α - 1))

# From https://www.ams.org/journals/tran/1949-065-01/S0002-9947-1949-0028473-1/S0002-9947-1949-0028473-1.pdf we use
# kn = (n + α - 1)_n / 2^n not (n+α+1)_n/2^n
# Koepf suggests kn = (n + α + 1)_n/2^n
function leading_term(P::Type{<:Bessel{α}}, n::Int) where {α}
    one(eltype(P)) / 2^n * Pochhammer(n + α - 1, n)
end

weight_function(::Type{<:Bessel{α}}) where {α} =
    z -> (2π * im)^(-1) * z^(α - 2) * exp(-2 / z)
generating_function(::Type{<:Bessel}) = (t, x) -> error("XXX")
function classical_hypergeometric(::Type{<:Bessel{α}}, n, x) where {α}
    as = (-n, n + α - 1)
    bs = ()
    pFq(as, bs, -x / 2) #/ kn
end

## Overrides XXX fails wih 1 and 2
function B̃n(P::Type{<:Bessel{α}}, n::Int) where {α}
    val = one(eltype(P))
    (iszero(n) && α == 2) && return val
    val *= 2 * (α - 2)
    val /= (2 * n + α) * (2 * n + α - 2)
    val
end

function C̃n(P::Type{<:Bessel{α}}, n::Int) where {α}
    val = one(eltype(P))
    n == 1 && return -(val * 4) / (α^2 * (α + 1))

    val *= -4 * n * (n + α - 2)
    val /= (2 * n + α - 3) * (2 * n + α - 2)^2 * (2 * n + α - 1)
    val
end

#Bn(::Type{<:Bessel{1}}, n::Int, ::Type{S}) where {S} = error("α=1 is not correct")
#B̃n(P::Type{<:Bessel{2}}, ::Val{0}) where {α} =  one(eltype(P))  # 0  otherwise
#iszero(N) #?  one(eltype(P)) : zero(eltype(P)))
#C̃n(P::Type{<:Bessel{α}}, ::Val{1}) where {α} =  -(one(eltype(P))*4)/(α^2*(α + 1))

b̂̃n(::Type{<:Bessel{2}}, n::Int) = (one(eltype(P)) * 2) / (n * (2n + 2))
b̂̃n(::Type{<:Bessel{2}}, ::Val{0}) = one(eltype(P)) * Inf


cd(@__DIR__); include("setups/rgc100.jl")
gr(dpi = 200)

## construct full HGLET and Lapped-HGLET dictionaries
@time HGLET_dic = HGLET_dictionary(GP, G_Sig; method = :L)
@time LPHGLET_dic = LPHGLET_dictionary(GP, G_Sig; method = :L, ϵ = 0.3)

## generate figures
# pre-selected (j, k, l)s
T = [(2, 1, 1), (2, 1, 12), (3, 1, 3), (3, 1, 5),
    (4, 5, 6), (6, 23, 3), (7, 5, 2), (8, 15, 1)]

for (j, k, l) in T
    𝚽H = HGLET_dic[GP.rs[k, j]:(GP.rs[k + 1, j] - 1), j, :]'
    𝚽lH = LPHGLET_dic[GP.rs[k, j]:(GP.rs[k + 1, j] - 1), j, :]'
    p1 = gplot(W, X; width = 1)
        scatter_gplot!(X; marker = 𝚽H[:, l], plotOrder = :s2l)
        plot!(xlims = [-180, 220], xlabel = "x(μm)", ylabel = "y(μm)", frame = :box, cbar = false)
        plot!(left_margin = 3mm, bottom_margin = 3mm, clims = (-0.05, 0.05))
    p2 = gplot(W, X; width = 1)
        scatter_gplot!(X; marker = 𝚽lH[:, l], plotOrder = :s2l)
        plot!(xlims = [-180, 220], xlabel = "x(μm)", ylabel = "y(μm)", frame = :box, cbar = false)
        plot!(left_margin = 3mm, bottom_margin = 3mm, clims = (-0.05, 0.05))
    plt = plot(p1, p2, size = (800, 400), layout = Plots.grid(1, 2))
    savefig(plt, "../figs/RGC100_HGLET_LPHGLET_j$(j-1)k$(k-1)l$(l-1).png")
end


@doc raw"""
    cost_acceleration_bezier(
        M::AbstractManifold,
        B::AbstractVector{P},
        degrees::AbstractVector{<:Integer},
        T::AbstractVector{<:AbstractFloat},
    ) where {P}

compute the value of the discrete Acceleration of the composite Bezier curve

```math
\sum_{i=1}^{N-1}\frac{d^2_2 [ B(t_{i-1}), B(t_{i}), B(t_{i+1})]}{\Delta_t^3}
```

where for this formula the `pts` along the curve are equispaced and denoted by
``t_i``, ``i=1,…,N``, and ``d_2`` refers to the second order absolute difference [`costTV2`](@ref)
(squared). Note that the Beziér-curve is given in reduces form as a point on a `PowerManifold`,
together with the `degrees` of the segments and assuming a differentiable curve, the segments
can internally be reconstructed.

This acceleration discretization was introduced in[^BergmannGousenbourger2018].

# See also

[`grad_acceleration_bezier`](@ref), [`cost_L2_acceleration_bezier`](@ref), [`grad_L2_acceleration_bezier`](@ref)

[^BergmannGousenbourger2018]:
    > Bergmann, R. and Gousenbourger, P.-Y.: A variational model for data fitting on
    > manifolds by minimizing the acceleration of a Bézier curve.
    > Frontiers in Applied Mathematics and Statistics (2018).
    > doi [10.3389/fams.2018.00059](http://dx.doi.org/10.3389/fams.2018.00059),
    > arXiv: [1807.10090](https://arxiv.org/abs/1807.10090)
"""
function cost_acceleration_bezier(
    M::AbstractManifold,
    B::AbstractVector{P},
    degrees::AbstractVector{<:Integer},
    T::AbstractVector{<:AbstractFloat},
) where {P}
    Bt = get_bezier_segments(M, B, degrees, :differentiable)
    p = de_casteljau(M, Bt, T)
    n = length(T)
    f = p[[1, 3:n..., n]]
    b = p[[1, 1:(n - 2)..., n]]
    d = distance.(Ref(M), p, shortest_geodesic.(Ref(M), f, b, Ref(0.5))) .^ 2
    samplingFactor = 1 / (((max(T...) - min(T...)) / (n - 1))^3)
    return samplingFactor * sum(d)
end
@doc raw"""
    cost_L2_acceleration_bezier(M,B,pts,λ,d)

compute the value of the discrete Acceleration of the composite Bezier curve
together with a data term, i.e.

````math
\frac{λ}{2}\sum_{i=0}^{N} d_{\mathcal M}(d_i, c_B(i))^2+
\sum_{i=1}^{N-1}\frac{d^2_2 [ B(t_{i-1}), B(t_{i}), B(t_{i+1})]}{\Delta_t^3}
````

where for this formula the `pts` along the curve are equispaced and denoted by
``t_i`` and ``d_2`` refers to the second order absolute difference [`costTV2`](@ref)
(squared), the junction points are denoted by ``p_i``, and to each ``p_i`` corresponds
one data item in the manifold points given in `d`. For details on the acceleration
approximation, see [`cost_acceleration_bezier`](@ref).
Note that the Beziér-curve is given in reduces form as a point on a `PowerManifold`,
together with the `degrees` of the segments and assuming a differentiable curve, the
segments can internally be reconstructed.

# See also

[`grad_L2_acceleration_bezier`](@ref), [`cost_acceleration_bezier`](@ref), [`grad_acceleration_bezier`](@ref)
"""
function cost_L2_acceleration_bezier(
    M::AbstractManifold,
    B::AbstractVector{P},
    degrees::AbstractVector{<:Integer},
    T::AbstractVector{<:AbstractFloat},
    λ::AbstractFloat,
    d::AbstractVector{P},
) where {P}
    Bt = get_bezier_segments(M, B, degrees, :differentiable)
    p = get_bezier_junctions(M, Bt)
    return cost_acceleration_bezier(M, B, degrees, T) +
           λ / 2 * sum((distance.(Ref(M), p, d)) .^ 2)
end

@doc raw"""
    costIntrICTV12(M, f, u, v, α, β)

Compute the intrinsic infimal convolution model, where the addition is replaced
by a mid point approach and the two functions involved are [`costTV2`](@ref)
and [`costTV`](@ref). The model reads

```math
E(u,v) =
  \frac{1}{2}\sum_{i ∈ \mathcal G}
    d_{\mathcal M}\bigl(g(\frac{1}{2},v_i,w_i),f_i\bigr)
  +\alpha\bigl( β\mathrm{TV}(v) + (1-β)\mathrm{TV}_2(w) \bigr).
```
"""
function costIntrICTV12(M::AbstractManifold, f, u, v, α, β)
    IC = 1 / 2 * distance(M, shortest_geodesic(M, u, v, 0.5), f)^2
    TV12 = β * costTV(M, u) + (1 - β) * costTV2(M, v)
    return IC + α * TV12
end

@doc raw"""
    costL2TV(M, f, α, x)

compute the ``ℓ^2``-TV functional on the `PowerManifold manifold `M` for given
(fixed) data `f` (on `M`), a nonnegative weight `α`, and evaluated at `x` (on `M`),
i.e.

```math
E(x) = d_{\mathcal M}^2(f,x) + \alpha \operatorname{TV}(x)
```

# See also
[`costTV`](@ref)
"""
costL2TV(M, f, α, x) = 1 / 2 * distance(M, f, x)^2 + α * costTV(M, x)

@doc raw"""
    costL2TVTV2(M, f, α, β, x)

compute the ``ℓ^2``-TV-TV2 functional on the `PowerManifold` manifold `M` for
given (fixed) data `f` (on `M`), nonnegative weight `α`, `β`, and evaluated
at `x` (on `M`), i.e.

```math
E(x) = d_{\mathcal M}^2(f,x) + \alpha\operatorname{TV}(x)
  + β\operatorname{TV}_2(x)
```

# See also
[`costTV`](@ref), [`costTV2`](@ref)
"""
function costL2TVTV2(M::PowerManifold, f, α, β, x)
    return 1 / 2 * distance(M, f, x)^2 + α * costTV(M, x) + β * costTV2(M, x)
end

@doc raw"""
    costL2TV2(M, f, β, x)

compute the ``ℓ^2``-TV2 functional on the `PowerManifold` manifold `M`
for given data `f`, nonnegative parameter `β`, and evaluated at `x`, i.e.

```math
E(x) = d_{\mathcal M}^2(f,x) + β\operatorname{TV}_2(x)
```

# See also
[`costTV2`](@ref)
"""
function costL2TV2(M::PowerManifold, f, β, x)
    return 1 / 2 * distance(M, f, x)^2 + β * costTV2(M, x)
end

@doc raw"""
    costTV(M, x, p)

Compute the ``\operatorname{TV}^p`` functional for a tuple `pT` of pointss
on a [Manifold](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.Manifold) `M`, i.e.

```math
E(x_1,x_2) = d_{\mathcal M}^p(x_1,x_2), \quad x_1,x_2 ∈ \mathcal M
```

# See also

[`grad_TV`](@ref), [`prox_TV`](@ref)
"""
function costTV(M::AbstractManifold, x::Tuple{T,T}, p::Int=1) where {T}
    return distance(M, x[1], x[2])^p
end
@doc raw"""
    costTV(M,x [,p=2,q=1])

Compute the ``\operatorname{TV}^p`` functional for data `x`on the `PowerManifold`
manifold `M`, i.e. ``\mathcal M = \mathcal N^n``, where ``n ∈ \mathbb N^k`` denotes
the dimensions of the data `x`.
Let ``\mathcal I_i`` denote the forward neighbors, i.e. with ``\mathcal G`` as all
indices from ``\mathbf{1} ∈ \mathbb N^k`` to ``n`` we have
``\mathcal I_i = \{i+e_j, j=1,…,k\}\cap \mathcal G``.
The formula reads

```math
E^q(x) = \sum_{i ∈ \mathcal G}
  \bigl( \sum_{j ∈  \mathcal I_i} d^p_{\mathcal M}(x_i,x_j) \bigr)^{q/p}.
```

# See also
[`grad_TV`](@ref), [`prox_TV`](@ref)
"""
function costTV(M::PowerManifold, x, p=1, q=1)
    power_size = power_dimensions(M)
    R = CartesianIndices(Tuple(power_size))
    d = length(power_size)
    maxInd = last(R)
    cost = fill(0.0, Tuple(power_size))
    for k in 1:d # for all directions
        ek = CartesianIndex(ntuple(i -> (i == k) ? 1 : 0, d)) #k th unit vector
        for i in R # iterate over all pixel
            j = i + ek # compute neighbor
            if all(map(<=, j.I, maxInd.I)) # is this neighbor in range?
                cost[i] += costTV(M.manifold, (x[M, Tuple(i)...], x[M, Tuple(j)...]), p)
            end
        end
    end
    cost = (cost) .^ (1 / p)
    if q > 0
        return sum(cost .^ q)^(1 / q)
    else
        return cost
    end
end
@doc raw"""
    costTV2(M,(x1,x2,x3) [,p=1])

Compute the ``\operatorname{TV}_2^p`` functional for the 3-tuple of points
`(x1,x2,x3)`on the [Manifold](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.Manifold) `M`. Denote by

```math
  \mathcal C = \bigl\{ c ∈  \mathcal M \ |\ g(\tfrac{1}{2};x_1,x_3) \text{ for some geodesic }g\bigr\}
```

the set of mid points between ``x_1`` and ``x_3``. Then the function reads

```math
d_2^p(x_1,x_2,x_3) = \min_{c ∈ \mathcal C} d_{\mathcal M}(c,x_2).
```

# See also
[`grad_TV2`](@ref), [`prox_TV2`](@ref)
"""
function costTV2(M::MT, x::Tuple{T,T,T}, p=1) where {MT<:AbstractManifold,T}
    # note that here mid_point returns the closest to x2 from the e midpoints between x1 x3
    return 1 / p * distance(M, mid_point(M, x[1], x[3]), x[2])^p
end
@doc raw"""
    costTV2(M,x [,p=1])

compute the ``\operatorname{TV}_2^p`` functional for data `x` on the
`PowerManifold` manifoldmanifold `M`, i.e. ``\mathcal M = \mathcal N^n``,
where ``n ∈ \mathbb N^k`` denotes the dimensions of the data `x`.
Let ``\mathcal I_i^{\pm}`` denote the forward and backward neighbors, respectively,
i.e. with ``\mathcal G`` as all indices from ``\mathbf{1} ∈ \mathbb N^k`` to ``n`` we
have ``\mathcal I^\pm_i = \{i\pm e_j, j=1,…,k\}\cap \mathcal I``.
The formula then reads

```math
E(x) = \sum_{i ∈ \mathcal I,\ j_1 ∈  \mathcal I^+_i,\ j_2 ∈  \mathcal I^-_i}
d^p_{\mathcal M}(c_i(x_{j_1},x_{j_2}), x_i),
```

where ``c_i(⋅,⋅)`` denotes the mid point between its two arguments that is
nearest to ``x_i``.

# See also
[`grad_TV2`](@ref), [`prox_TV2`](@ref)
"""
function costTV2(M::PowerManifold, x, p::Int=1, Sum::Bool=true)
    Tt = Tuple(power_dimensions(M))
    R = CartesianIndices(Tt)
    d = length(Tt)
    minInd, maxInd = first(R), last(R)
    cost = fill(0.0, Tt)
    for k in 1:d # for all directions
        ek = CartesianIndex(ntuple(i -> (i == k) ? 1 : 0, d)) #k th unit vector
        for i in R # iterate over all pixel
            jF = i + ek # compute forward neighbor
            jB = i - ek # compute backward neighbor
            if all(map(<=, jF.I, maxInd.I)) && all(map(>=, jB.I, minInd.I)) # are neighbors in range?
                cost[i] += costTV2(
                    M.manifold,
                    (x[M, Tuple(jB)...], x[M, Tuple(i)...], x[M, Tuple(jF)...]),
                    p,
                )
            end
        end # i in R
    end # directions
    if p != 1
        cost = (cost) .^ (1 / p)
    end
    if Sum
        return sum(cost)
    else
        return cost
    end
end


#=
The following codes are based on old pmm related functions written by sammy-suyama.
https://github.com/sammy-suyama/BayesBook/blob/master/src/PoissonMixtureModel.jl
=#

const PMM = MixtureModel{Univariate, Discrete, Poisson}

struct BayesianPMM
    α
    prior::Vector{Gamma}
    BayesianPMM(α, prior) = new(α, prior)
end

function rand(model::BayesianPMM)
    K = length(model.prior)
    a = rand(Dirichlet(K, model.α))
    λ = rand.(model.prior)
    MixtureModel(Poisson, λ, a)
end

function one_hot(model::PMM, N::Int)
    K = length(model.components)
    a = probs(model) # mixture weight

    # one hot vector of (K, N)
    onehotbatch(rand(Categorical(a), N), collect(1:K))
end

function generate_toy_data(model::PMM, N::Int)
    D = length(model)

    data = Array{eltype(model.components[1])}(undef, N, D)
    # convert to vector if D == 1
    D == 1 ? data = dropdims(data, dims = 2) : data
    S = one_hot(model, N)
    @inbounds for n = 1:N
        s = argmax(S[:, n])
        data[n, :] .= rand(model.components[s])
    end
    data, S
end

function init_S(data, prior_model::BayesianPMM)
    N = size(data, 1)
    K = length(prior_model.prior)
    onehotbatch(rand(Categorical(ones(K)./K), N), collect(1:K))
end

function add_stats(prior_model::BayesianPMM, data, S)
    @argcheck length(prior_model.prior) == size(S, 1)
    D = length(prior_model.prior[1])
    K = size(S, 1)
    sum_S = sum(S, dims=2)
    α = [prior_model.α + sum_S[k] for k = 1:K]
    XS = data'*S'

    a = Vector{Float64}(undef, K)
    b = Vector{Float64}(undef, K)
    for k = 1:K
        a[k] = prior_model.prior[k].α + XS[d, k] for d = 1:D
        b[k] = prior_model.prior[k].θ + sum_S[k]
    end
    BayesianPMM(α, Gamma.(a, b))
end


using VectorizedArgMinMax
using Test

@testset "VectorizedArgMinMax.jl" begin
    @testset "vargmax.jl" begin
        n = 10
        A = rand(n, n, n);
        @test argmax(A) == vargmax(A)
        @test argmax(A, dims=1) == vargmax(A, dims=1)
        @test argmax(A, dims=2) == vargmax(A, dims=2)
        @test argmax(A, dims=3) == vargmax(A, dims=3)
    end
    @testset "vfindmax.jl" begin
        n = 10
        A = rand(n, n, n);
        @test findmax(A) == vfindmax(A)
        @test findmax(A, dims=1) == vfindmax(A, dims=1)
        @test findmax(A, dims=2) == vfindmax(A, dims=2)
        @test findmax(A, dims=3) == vfindmax(A, dims=3)
    end
end


using Revise
using DataFrames, BenchmarkTools, LinearAlgebra, Distributions, StatsBase, Statistics, Random
Random.seed!(11)

#set up a siimilar regression situation
function crosssection(N,K; loud::Bool=true)

  #generate our data
  ϵ = rand(Laplace(0,1), N)
  Ξ = hcat((k->rand(Laplace(0,2),N)).(1:K)...)
  þ = 1 ./ (þᵢ->abs(þᵢ) ≤ 10^-5 ? 10.0^-5 * sign(þᵢ) : þᵢ).(rand(Chisq(1), N)./4)
  #þ .*= rand(Bernoulli(),N)
  X = Ξ .* þ
  #abstotX = sum(abs.(X), dims=1)

  function normalizemat!(X)
    X .-= mean(X, dims=1)
    abstot = sum(abs.(X), dims=1)
    #abstot₊ = sum((x->x≥0 ? x : 0.).(X), dims=1)
    #abstot₋ = sum((x->x<0 ? -x : 0.).(X), dims=1)
    #X .= hcat(((t)-> (x->x≥0 ? x/(abstot₊[t[1]]) : x/(abstot₋[t[1]])).(t[2])).(enumerate(eachcol(X)))...)
    X ./= abstot .* 2
  end
  #normalize the unobservable actual
  normalizemat!(X)

  #normalize the observable
  normalizemat!(Ξ)

  #σx = std(X, dims=1)
  #X ./= σx
  loud && display(abstotX)
  #isplay(X)
  X = hcat(ones(N),X)
  β = [50.0; (Float64).(1:K) .* 300  ...;]

  y = X*β  .+ ϵ

  #display(β)
  function err(Xqr,y,b)
    u = y.-X*Xb
    Λ = u .* u
    RinvQp = Xqr.R\I * Xqr.Q'
    Σ = RinvQp * Diagonal(Λ) * (RinvQp')
    return diag(Σ).^0.5
  end

  #errchk(X,y,b) = diag(((X'*X)\I) * X'* Diagonal((y.-X*Xb) .* (y.-X*Xb)) * X * ((X'*X)\I)).^0.5

  loud && @info "Test non-confounded version (β=$β)"
  Xqr = qr(X)
  Xrinv = Xqr.R\I
  Xb =  Xrinv * Xqr.Q' * y
  eX = (y.-X*Xb)
  errX = err(Xqr,y,Xb)#diag(eX' * eX ./ (N-K-1) .* (Xqr.R'*Xqr.R)\I) .^ 0.5
  #errXchk = errchk(X,y,Xb)

  loud && @info "Test observed version"
  Ξ1 =hcat(ones(N),Ξ)
  Ξqr = qr(Ξ1)
  Ξrinv = Ξqr.R\I
  Ξb =  Ξrinv * Ξqr.Q' * y
  errΞ = err(Ξqr,y,Ξb)
  loud && @info "Non-confounded version: $(Xb). ErrX = $errX"

  ΣΞbabs = sum(abs.(Ξ1 .* Ξb'), dims=1)
  loud && @info "Results: Raw b: $Ξb, raw error: $errΞ \n"
  loud && @info "Abs total: $(ΣΞbabs)"

  #=@info "Test confounded version"
  Xqr = qr(X)
  Xrinv = Xqr.R\I
  Xb =  Xrinv * Xqr.Q' * y
  @info "Xb: $(Xb) (vs non-confounded of $(Ξb))"
  ΣXb = vec(sum(, dims=1))=#
  return (Ξb=Ξb, ΣΞbabs=ΣΞbabs)
end

tups = vcat((crosssection(2_000,2, loud=false) for i in 1:1000)...)
#display(tups)
Ξb, As = vcat((t->t.Ξb').(tups)...), vcat((t->t.ΣΞbabs).(tups)...)
display(Ξb)
@info "mean raw: $(mean(Ξb, dims=1))"
@info "mean: $(mean(As, dims=1))"
@info "std: $(std(As, dims=1))"

#crosssection(2_000,2)


using Test, MLToolkit
using MLToolkit.DistributionsX: u2gumbel
using Distributions: Dirichlet
using Statistics: mean
using Tracker: gradient, data
using FiniteDifferences: central_fdm
using StatsFuns: logit, logistic
using Flux: gpu

@testset "Gumbel" begin
    n_randtests = 5
    n_samples = 5_000
    d, n = 2, 10
    atol = 0.015

    @testset "GumbelSoftmax" begin
        for _ = 1:n_randtests
            # Vector `p`
            p = rand(Dirichlet(ones(d))) |> gpu
            gs = GumbelSoftmax(p)
            @test mean(gs) == p
            x = rand(gs, n_samples)
            @test vec(mean(x; dims=2)) ≈ mean(gs) atol=atol * d

            # Matrix `p`
            p = rand(Dirichlet(ones(d)), n) |> gpu
            gs = GumbelSoftmax(p)
            xs = [rand(gs) for _ = 1:n_samples]
            @test mean(xs) ≈ mean(gs) atol=atol * d * n

            logpdf(gs, xs[1])
            @warn "Correctness of `logpdf(::GumbelSoftmax, x)` is not tested."
        end
    end

    @testset "GumbelSoftmax2D" begin
        for _ = 1:n_randtests
            # Vector `p`
            p1 = rand()
            gs = GumbelSoftmax2D(p1)
            @test mean(gs) == [p1, 1 - p1]
            x = rand(gs, n_samples)
            @test vec(mean(x; dims=2)) ≈ mean(gs) atol=atol * 2

            # Matrix `p`
            p1 = rand(n) |> gpu
            gs = GumbelSoftmax2D(p1)
            @test mean(gs) == [p1'; 1 .- p1']
            xs = [rand(gs) for _ = 1:n_samples]
            @test mean(xs) ≈ mean(gs) atol=atol * 2 * n
        end
    end

    @testset "track(u2gumbel, u)" begin
        for _ = 1:n_randtests
            u = rand() / 2 + 1 / 4
            f(x) = sum(u2gumbel(x))
            g_ad, = data.(gradient(f, u))
            g_fd = central_fdm(5, 1)(f, u)
            @test g_ad ≈ g_fd
        end
    end

    @testset "GumbelBernoulli" begin
        for _ = 1:n_randtests
            p = rand(d, n) |> gpu
            gb = GumbelBernoulli(p)
            @test mean(gb) == p
            xs = [rand(gb) for _ = 1:n_samples]
            @test mean(xs) ≈ mean(gb) atol=atol * d * n

            # Consistency between two versions of log density
            x = rand(gb)
            @test logpdf(gb, x) ≈ logpdfCoV(gb, x)
        end
    end

    @testset "GumbelBernoulliLogit" begin
        for _ = 1:n_randtests
            p = rand(d, n) |> gpu
            gbl = GumbelBernoulliLogit(logit.(p))
            @test mean(gbl) ≈ p
            xs = [rand(gbl) for _ = 1:n_samples]
            @test mean(xs) ≈ mean(gbl) atol=atol * d * n
        end
    end
end


using MultistartOptimization
using Test
using NLopt: NLopt

include("test_functions.jl")

@testset "test function sanity checks" begin
    for F in TEST_FUNCTIONS
        @test F(minimum_location(F, 10)) ≈ 0
    end
end

@testset "global optimization" begin
    for F in setdiff(TEST_FUNCTIONS, (RASTRIGIN, )) # Rastrigin disabled for now
        n = 10
        P = MinimizationProblem(F, lower_bounds(F, n), upper_bounds(F, n))
        local_method = NLoptLocalMethod(NLopt.LN_BOBYQA)
        multistart_method = TikTak(100)
        p = multistart_minimization(multistart_method, local_method, P)
        x₀ = minimum_location(F, n)
        @test p.location ≈ x₀ atol = 1e-5
        @test p.value ≈ F(x₀) atol = 1e-10
    end
end

@testset "custom local minimization and infeasibility" begin
    P = MinimizationProblem(x -> sum(abs2, x), [-1, -1], [1, 1])
    x0 = [0.0, 0.0]
    function _local_method(minimization_problem, x)
        if sum(abs2, x) < 0.25
            (location = x0, value = 0.0)
        else
            nothing
        end
    end
    p = multistart_minimization(TikTak(100), _local_method, P)
    @test p.location == x0
    @test p.value == 0.0
end


# "in this version, set data and smoothness terms using arrays.
#  grid neighborhood structure is assumed.#
#
using GCoptimization

width = 10
height = 5
num_pixels = width * height
num_labels = 7
result = zeros(num_pixels)    # "stores result of optimization"

immutable ForDataFn
    data::Vector{Cint}
    numLab::Cint
end

function dataFn(p::GCOSiteID, l::GCOLabelID, data::Ptr{Void})::GCOEnergyTermType
    myData = unsafe_pointer_to_objref(data)
    numLab = myData.numLab
    return myData.data[p*numLab+l+1]
end

smoothFn(p1::GCOSiteID, p2::GCOSiteID, l1::GCOLabelID, l2::GCOLabelID)::GCOEnergyTermType = (l1-l2)*(l1-l2) <= 4 ? (l1-l2)*(l1-l2) : 4

# "first set up the array for data costs"
data = Vector{Cint}(num_pixels * num_labels)
for i = 0:num_pixels-1, l = 0:num_labels-1
    if i < 25
        if l == 0
            data[i*num_labels+l+1] = 0
        else
            data[i*num_labels+l+1] = 10
        end
    else
        if l == 5
            data[i*num_labels+l+1] = 0
        else
            data[i*num_labels+l+1] = 10
        end
    end
end

gco = GCoptimizationGridGraph(width, height, num_labels)

# "set up the needed data to pass to function for the data costs"
toFn = ForDataFn(data, num_labels)
setDataCost(gco, dataFn, pointer_from_objref(toFn))

# "smoothness comes from function pointer"
setSmoothCost(gco, smoothFn)

println("Before optimization energy is ", compute_energy(gco))
expansion(gco, 2)    # "run expansion for 2 iterations."
println("After optimization energy is ", compute_energy(gco))

for i = 0:num_pixels-1
    result[i+1] = whatLabel(gco, i)
end

showall(result)


module Dating

using Distributions, Roots

export DatingEstimate, coaltime, SimpleEstimate, SpeedDating, SDResult, upper,
       lower

"""
An abstract dating estimate type.

Types inheriting from DatingEstimate should have the following methods defined:

* lower: Returns the lower bound of the dating estimate.
* upper: Returns the upper bound of the dating estimate.

They may define their own show methods.
"""
abstract DatingEstimate

Base.print(io::IO, de::DatingEstimate) = println(io, "$(lower(de)) ... $(upper(de))")
Base.show(io::IO, de::DatingEstimate) = println(io, "Coalesence time estimate between two sequences:\nt lies between $(lower(de)) and $(upper(de))")
Base.in(val::Float64, de::DatingEstimate) = val <= upper(de) && val >= lower(de)

# Different coalescence time estimation algorithms.
# method coaltime is dispatched according to arguments or type of
immutable SimpleEstimate end

"""
    coaltime(len::Int, nmut::Int, mu::Float64, ::Type{SimpleEstimate})

Compute the coalescence time between two sequences by assuming a mutation rate,
and then assuming the divergence between two sequences is 2 * mu * t

# Examples
```julia
coaltime(50, 17, 10e-9, SimpleEstimate)
```
"""
function coaltime(len::Int, nmut::Int, mu::Float64, ::Type{SimpleEstimate})
    return (nmut / len) / (2 * mu)
end

immutable SpeedDating end

"""
SDResult is a simple datatype to store the results of the coaltime
function.

coaltime estimates date range that may be considered a 95% confidence range in
which the true coalsescence time between two sequences lies.

The type store the 5%, 50%, and 95% values for the range.
"""
immutable SDResult <: DatingEstimate
    lower::Float64
    middle::Float64
    upper::Float64
end

lower(x::Float64) = x
lower(x::SDResult) = x.lower
upper(c::Float64) = x
upper(x::SDResult) = x.upper

function Base.show(io::IO, de::SDResult)
    println(io, "Coalescence time estimate:\n5%: $(de.lower), 95%: $(de.upper)")
end

function binomzero(p0::Float64, N::Int, B::Int)
    f(p::Float64) = cdf(Binomial(N, p), B) - p0
    return fzero(f, 0, 1)
end

"""
    coaltime(len::Int, nmut::Int, mu::Float64, ::Type{SpeedDating})

Compute the coalescence time between two sequences by modelling the process of
mutation accumulation between two sequences as a bernoulli process.

This method was first described in the paper:
Ward, B. J., & van Oosterhout, C. (2016). Hybridcheck: Software for the rapid detection, visualization and dating of recombinant regions in genome sequence data. Molecular Ecology Resources, 16(2), 534–539.

len is the length of the two aligned sequences, nmut is the number of mutations
and mu is the assumed mutation rate.

# Examples
```julia
coaltime(50, 17, 10e-9, SpeedDating)
```
"""
function coaltime(len::Int, nmut::Int, mu::Float64, ::Type{SpeedDating})
    ninetyfive = ceil(binomzero(0.05, len, nmut) / (2 * mu))
    fifty = ceil(binomzero(0.5, len, nmut) / (2 * mu))
    five = ceil(binomzero(0.95, len, nmut) / (2 * mu))
    return SDResult(five, fifty, ninetyfive)
end

end


# the loop version of loop
function compact!(x)
    @inbounds for i in 1:3
        if x[i] == 0
            for j in i+1:4
                if x[j] != 0
                    tmp = x[j:4]
                    x[j:4] .= 0
                    x[i:i+4-j] .= tmp
                    break;
                end
            end
        end
    end
    x
end

function move_loop!(x)
    # firstly, compact x so that there are no gaps
    # however cells are not yet merged
    compact!(x)

    # now the data is compact; we combine it
    @inbounds if (x[1] != 0) && (x[1] == x[2])
        x[1] += 1
        if (x[3] != 0)  && (x[3] == x[4])
            x[2], x[3], x[4] = x[3] + 1, 0 , 0
        else
            x[2], x[3], x[4] = x[3], x[4], 0
        end
    elseif (x[2] != 0) && (x[2] == x[3])
        x[2] += 1
        x[3], x[4] = x[4], 0
    elseif (x[3] != 0) && (x[3] == x[4])
        x[3] += 1
        x[4] = 0
    end

    x
end


function evaluate_function_at_nodes(func::Function, unnormalised_nodes::Array{Float64,1}, limits::OrderedDict{Symbol,Tuple{Float64,Float64}}, function_takes_Dict::Bool)
    nodes = length(unnormalised_nodes)
    dimensions = length(limits)
    labels = collect(keys(limits))
    normalised_node_array = Array{Float64,2}(undef,nodes, dimensions)
    for i in 1:length(labels)
        dimen = labels[i]
        normalised_nodes = ((unnormalised_nodes .+ 1) .* ((limits[dimen][2]-limits[dimen][1])/2)) .+ limits[dimen][1]
        normalised_node_array[:,i] = normalised_nodes
    end
    y = Array{Float64,dimensions}(undef,repeat([nodes],dimensions)...)
    y_indicies = CartesianIndices(y)
    for i in y_indicies
        index = Tuple(i)
        coords = Array{Float64,1}()
        for j in 1:length(index)
            append!(coords, normalised_node_array[index[j],j])
        end
        if function_takes_Dict
            coordsDict = Dict{Symbol,Float64}(labels .=> coords)
            y[i] = func(coordsDict)
        else
            y[i] = func(coords...)
        end
    end
    return y
end

# The operations of the next two functions are not optimal. If we have the same
# vector in two columns then we could more efficiently multiply them by noting that
# the outer product will be symmetric.
# At some stage it would be good to make better use of symmetry. Come back to this
# wehn the tensor
function expand_array(arr::Array{T} where T, vecc::Array{T,1} where T)
    sze = size(arr)
    dim = length(sze)
    new_arr = Array{typeof(vecc[1]),dim + 1}(undef, vcat(sze..., sze[1])...)
    for i in 1:length(vecc)
        new_arr[repeat([:],dim)...,i] = arr .* vecc[i]
    end
    return new_arr
end
function tensor_outer_product(arr::Array{T,2} where T)
    dim = size(arr)[2]
    if dim == 1
        return Array{typeof(arr[1,1]),1}(arr[:,1])
    end
    big_array = arr[:,1] #* hcat(arr[:,2]...) #transpose(arr[:,2]) as transpose doesnt work with nonnumeric types for T.
    for i in 2:dim
        big_array = expand_array(big_array, arr[:,i])
    end
    return big_array
end

function get_cholesky_coefficients(evaluated_chebyshevs_on_sum_squared::Array{Float64,2}, y::Array{Float64})
    degree = size(evaluated_chebyshevs_on_sum_squared)[2]
    nodes = size(evaluated_chebyshevs_on_sum_squared)[1]
    dimensions = length(size(y))
    coefficients = zeros(repeat([degree], dimensions)...)
    perms = unique(collect(permutations(1:dimensions)))
    indices = CartesianIndices(coefficients)
    for i in indices
        if (coefficients[i] ≂ 0.0)
            reduced_chebs = evaluated_chebyshevs_on_sum_squared[:,[Tuple(i)...]]
            chebs_mat = tensor_outer_product(reduced_chebs) # This will return a nodes^dim size array like y.
            # We coudl get the coefficient by immediately elementwise multiplying by y and then summing but note
            # that if we permute chebs_mat we get the correct chebs_mat for other parts. So we do that first.
            for j in perms
                ind = Tuple(i)[j]
                if (coefficients[CartesianIndex(ind...)] ≂ 0.0)
                    permuted_chevs = Base.permutedims(chebs_mat, j)
                    coefficients[CartesianIndex(ind...)] = sum(permuted_chevs .* y)
                end
            end
        end
    end
    return coefficients
end

function convert_chebyshevs(chebyshevs::Array{Sum_Of_Functions,1}, limits::OrderedDict{Symbol,Tuple{Float64,Float64}})
    dimensions = length(limits)
    degree = length(chebyshevs)
    ks     = collect(keys(limits))
    rescaled_chebs = Array{Sum_Of_Functions,2}(undef, degree, dimensions)
    for i in 1:dimensions
        k = ks[i]
        left, right = limits[k]
        transformed_chebyshevs = convert_to_linearly_rescale_inputs.(chebyshevs, (right-left)/2, +(right+left)/2)
        conversion_dict        = Dict{Symbol,Symbol}([default_symbol] .=> k)
        rescaled_chebs[:,i]    = rebadge.(transformed_chebyshevs, Ref(conversion_dict))
    end
    rescaled_chebs_outer_product = tensor_outer_product(rescaled_chebs)
    return rescaled_chebs_outer_product
end

"""
    create_chebyshev_approximation(f::Function, nodes::Int, degree::Int, limits::OrderedDict{Symbol,Tuple{Float64,Float64}}, function_takes_Dict::Bool = false)
Creates a Sum_Of_Functions that approximates a function, f, with a set of chebyshevs of a particular degree. The nodes input specifies at how many locations the
function is to be evaluated for approximation purposes in each dimension. The limits OrderedDict specifies the domain of where the function is to be approximated.

If function_takes_Dict is true then the function will be evaluated by inputting a  Dict{Symbol,Float64}. Otherwise the function will be evaluated with f(values(limits)...)
Note that the order of the OrderedDict specifies the order of inputs to the function in this case.
"""
function create_chebyshev_approximation(f::Function, nodes::Int, degree::Int, limits::OrderedDict{Symbol,Tuple{Float64,Float64}}, function_takes_Dict::Bool = false)
    # This is all after Algorithm 6.2 from Judd (1998) Numerical Methods in Economics.
    if nodes <= degree
        error("Need to have more nodes than degree to use a chebyshev approximation")
    end
    k = 1:nodes
    unnormalised_nodes = -cos.( (((2 .* k) .- 1) ./ (2 * nodes)) .* pi    )

    chebyshevs = get_chevyshevs_up_to(degree, true)
    evaluated_chebyshevs_on_sum_squared = Array{Float64,2}(undef,nodes, degree)
    for i in 1:degree
        cheb_evaluated_at_each_node = evaluate.(Ref(chebyshevs[i]),unnormalised_nodes)
        evaluated_chebyshevs_on_sum_squared[:,i] = cheb_evaluated_at_each_node ./ sum( cheb_evaluated_at_each_node .^ 2 )
    end
    y = evaluate_function_at_nodes(f, unnormalised_nodes, limits, function_takes_Dict)
    a = get_cholesky_coefficients(evaluated_chebyshevs_on_sum_squared, y)
    transformed_chebyshevs = convert_chebyshevs(chebyshevs, limits)

    all_terms = a .* transformed_chebyshevs
    final_func = Sum_Of_Functions([all_terms...])
    return final_func
end


abstract type Sequence end
value(d::Sequence) = d.val
step!(d::Sequence) = d.val

struct ConstantSequence{T} <: Sequence
    val::T
end

mutable struct ExponentialSequence{T} <: Sequence
    val::T
    multiplier::T
end
step!(ED::ExponentialSequence) = ED.val *= ED.multiplier

mutable struct HeavisideSequence{T} <: Sequence
    origVal::T  # Value to start with
    lateVal::T  # Value to switch to
    switchIter::Int
    iter::Int
end
HeavisideSequence(origVal, lateVal, switchIter, iter = 1) = HeavisideSequence(origVal, lateVal, switchIter, iter)
value(HD::HeavisideSequence) = HD.iter < HD.switchIter ? HD.origVal : HD.lateVal
function step!(HD::HeavisideSequence)
    HD.iter += 1
    HD.iter < HD.switchIter ? HD.origVal : HD.lateVal
end

mutable struct ArithmeticSequence{T} <: Sequence
    val::T
    addVal::T
end
step!(AD::ArithmeticSequence) = AD.val += AD.addVal

mutable struct LinearSequence{T} <: Sequence
    startVal::T
    endVal::T
    startIter::Int
    endIter::Int
    currIter::Int
    step::T
    val::T
end
function LinearSequence(startVal::T, endVal::T, endIter::Int; startIter=1, iter = 1) where {T}
    step = (endVal - startVal)/(endIter - startIter)
    val = iter < endIter ? startVal + step*max(iter - startIter, 0) : endVal
    LinearSequence(startVal, endVal, startIter, endIter, 1, step, val)
end

function step!(LD::LinearSequence)
    LD.currIter += 1
    LD.val = LD.currIter < LD.endIter ? LD.startVal + LD.step*max(LD.currIter - LD.startIter, 0) : LD.endVal
end


using Plots
using FFTW
using LinearAlgebra

include("fdtd.jl")

function A2b()
    # Spatial discretization
    n = 2000 + 1
    Δx = 1.0e-3

    # Time discretization
    Δt = Δx

    #steps to take
    steps = 2000

    #result at first and last step
    snapshots = Int.(round.(collect(range(1, stop = steps, length = 20))))

    #simulate
    result = FDTD1D(steps, snapshots, n, Δx, Δt)

    #save
    SaveFDTDResult("A2b.csv", result)

    #plot result
    ps = []
    for s = 1:length(snapshots)
        p = plot(result.xe, result.efield[s,:], ylims = [-1.5, 1.5])
        plot!(p, result.xh, result.hfield[s,:])
        push!(ps, p)
    end
    plot(ps..., layout = (length(snapshots), 1), show = true, size = [1200, 600])
end

function A2c()
    # Spatial discretization
    n = 2000 + 1
    Δx = 1.0e-3

    # Time discretization
    Δt = 1.01 * Δx
    steps = 130
    snapshots = [130, 140]
    result = FDTD1D(steps, snapshots, n, Δx, Δt)
    SaveFDTDResult("A2c.csv", result)
    p = plot(result.xe, result.efield[1,:], ylims = [-1.5, 1.5])
    plot!(p, result.xh, result.hfield[1,:])
end

function A2d()
    n = 2000 + 1
    Δx = 1.0e-3
    Δt = Δx
    steps = 2000
    snapshots = collect(1:steps)
    result = FDTD1D(steps, snapshots, n, Δx, Δt, 1.0)
    e = result.efield[:, 1000]
    h = result.hfield[:, 1000]
    ef = fft(e .* ((-1.0).^collect(1:length(e))))
    hf = fft(h .* ((-1.0).^collect(1:length(h))))
    s = 0.5 .* real.(ef .* hf)
    plot(abs.(s[950:1050]))
end

function A2f()
    n = 2000 + 1
    Δx = 1.0e-3
    Δt = Δx
    steps = 1500
    snapshots = Int.(round.(collect(range(1, stop = steps, length = 20))))
    result = FDTD1D(steps, snapshots, n, Δx, Δt, 1.0)
    SaveFDTDResult("A2f.csv", result)
end

function main()
    #A2b()
    #A2c()
    A2d()
    #A2f()
end

main()


"""
    ProjectedPointDistribution(m::Manifold, d, proj!, x)

Generates a random point in ambient space of `m` and projects it to `m`
using function `proj!`. Generated arrays are of type `TResult`, which can be
specified by providing the `x` argument.
"""
struct ProjectedPointDistribution{TResult, TM<:Manifold, TD<:Distribution, TProj} <: MPointDistribution{TM}
    manifold::TM
    d::TD
    proj!::TProj
end

function ProjectedPointDistribution(M::Manifold, d::Distribution, proj!, x::TResult) where TResult
    return ProjectedPointDistribution{TResult, typeof(M), typeof(d), typeof(proj!)}(M, d, proj!)
end

function support(d::ProjectedPointDistribution)
    return MPointSupport(d.manifold)
end

function rand(rng::AbstractRNG, d::ProjectedPointDistribution{TResult}) where TResult
    x = convert(TResult, rand(rng, d.d))
    d.proj!(d.manifold, x)
    return x
end

function _rand!(rng::AbstractRNG, d::ProjectedPointDistribution, x::AbstractArray{<:Number})
    _rand!(rng, d.d, x)
    d.proj!(d.manifold, x)
    return x
end

"""
    ProjectedFVectorDistribution(type::VectorBundleFibers, x, d, project_vector!)

Generates a random vector from ambient space of manifold `type.manifold`
at point `x` and projects it to vector space of type `type` using function
`project_vector!`.
Generated arrays are of type `TResult`.
"""
struct ProjectedFVectorDistribution{TResult, TSpace<:VectorBundleFibers, TX, TD<:Distribution, TProj} <: FVectorDistribution{TSpace, TX}
    type::TSpace
    x::TX
    d::TD
    project_vector!::TProj
end

function ProjectedFVectorDistribution(type::VectorBundleFibers, x, d::Distribution, project_vector!, xt::TResult) where TResult
    return ProjectedFVectorDistribution{TResult, typeof(type), typeof(x), typeof(d), typeof(project_vector!)}(type, x, d, project_vector!)
end

function support(tvd::ProjectedFVectorDistribution)
    return FVectorSupport(tvd.type, tvd.x)
end

function rand(rng::AbstractRNG, d::ProjectedFVectorDistribution{TResult}) where TResult
    v = convert(TResult, reshape(rand(rng, d.d), size(d.x)))
    d.project_vector!(d.type, v, d.x, v)
    return v
end

function _rand!(rng::AbstractRNG, d::ProjectedFVectorDistribution, v::AbstractArray{<:Number})
    # calling _rand!(rng, d.d, v) doesn't work for all arrays types
    copyto!(v, rand(rng, d))
    return v
end


using Test, YaoArrayRegister

@testset "test fidelity" begin
    reg = rand_state(3)
    reg_ = rand_state(3)
    reg2 = repeat(reg, 3)
    @test fidelity(reg, reg) ≈ [1]
    @test fidelity(reg, reg_)[] < 1
    @test fidelity(reg2, reg2) ≈ [1, 1, 1]

    # mix
    reg4 = join(reg, reg)
    reg5 = join(reg_, reg_)
    focus!(reg4, 1:3)
    focus!(reg5, 1:3)
    @test isapprox(fidelity(reg, reg_)[], fidelity(reg4, reg5)[], atol = 1e-5)
    @test isapprox.(fidelity(reg, reg_)[], fidelity(repeat(reg4, 3), repeat(reg5, 3)), atol = 1e-5) |>
          all
end

@testset "test trace distance" begin
    reg = rand_state(3)
    reg_ = rand_state(3)
    reg2 = repeat(reg, 3)
    dm = ρ(reg)
    dm_ = ρ(reg_)
    dm2 = ρ(reg2)
    @test reg |> probs ≈ dm |> probs
    @test isapprox(tracedist(dm, dm), tracedist(reg, reg), atol = 1e-5)
    @test isapprox(tracedist(dm, dm_), tracedist(reg, reg_), atol = 1e-5)
    @test isapprox(tracedist(dm2, dm2), tracedist(reg2, reg2), atol = 1e-5)

    # mix
    reg4 = join(reg, reg)
    reg5 = join(reg_, reg_)
    focus!(reg4, 1:3)
    focus!(reg5, 1:3)
    dm4 = reg4 |> density_matrix
    dm5 = reg5 |> density_matrix
    @test isapprox(tracedist(dm, dm_)[], tracedist(dm4, dm5)[], atol = 1e-5)
    @test isapprox.(
        tracedist(dm, dm_)[],
        tracedist(repeat(reg4, 3) |> density_matrix, repeat(reg5, 3) |> density_matrix),
        atol = 1e-5,
    ) |> all
end

@testset "purify" begin
    reg = rand_state(6)
    reg_p = purify(reg |> ρ)
    @test reg_p |> isnormalized
    @test reg_p |> exchange_sysenv |> probs |> maximum ≈ 1
    reg_p = purify(reg |> ρ; nbit_env = 0)
    @test Yao.fidelity(reg, reg_p) ≈ [1]

    reg = rand_state(6; nbatch = 10)
    reg_p = purify(reg |> ρ)
    @test reg_p |> isnormalized
    @test reg_p |> exchange_sysenv |> probs |> maximum ≈ 1
    reg_p = purify(reg |> ρ; nbit_env = 0)
    @test Yao.fidelity(reg, reg_p) ≈ ones(10)
    reg_p = purify(reg |> ρ; nbit_env = 2)
    @test reg_p |> nqubits == 8
end


#=
Create efficient approximation function

Ignacio Quintero Mächler

t(-_-t)

Created 16 09 2017
Updated 17 01 2019
=#

"""
    linpred(val::Float64, x1::Float64, x2::Float64, y1::Float64, y2::Float64)

Estimate val according to linear interpolation for a range.
"""
linpred(val::Float64, x1::Float64, x2::Float64, y1::Float64, y2::Float64) =
  (y1 + (val - x1)*(y2 - y1)/(x2 - x1))::Float64

"""
    idxrange(x::Array{Float64,1}, val::Float64)

Get indexes in sorted vector `x` corresponding to the range in which
`val` is in using a sort of uniroot algorithm.
"""
function idxrange(x::Array{Float64,1}, val::Float64)

  @inbounds begin

    a::Int64 = 1

    if x[a] > val
      return a, false
    end

    b::Int64 = lastindex(x)

    if x[b] < val
      return b, false
    end

    mid::Int64 = div(b,2)

    while b-a > 1
      val < x[mid] ? b = mid : a = mid
      mid = div(b + a, 2)
    end

    if x[a] == val
      return a, false
    elseif x[b] == val
      return b, false
    else
      return a, true
    end

  end
end

"""
    @generated function approxf_full(t ::Float64,
                                     r ::Array{Float64,1},
                                     x ::Array{Float64,1},
                                     y ::Array{Float64,N},
                                     ::Val{nc}) where {N, nc}
Returns the values of `y` at `t` using an approximation function.
"""
@generated function approxf_full(t ::Float64,
                                 r ::Array{Float64,1},
                                 x ::Array{Float64,1},
                                 y ::Array{Float64,N},
                                 ::Val{nc}) where {N, nc}

  lex1 = quote end
  pop!(lex1.args)

  if N == 1
    push!(lex1.args, :(r[1] = linpred(t, x[a], x[a+1], y[a], y[a+1])::Float64))
  else
    # unroll loop
    for i = Base.OneTo(nc)
      push!(lex1.args, :(r[$i] = linpred(t, xa, xap1, y[a,$i], y[a+1, $i])::Float64))
    end

    # add one assignment
    pushfirst!(lex1.args, :(xap1 = x[a+1]::Float64))
    pushfirst!(lex1.args, :(xa   = x[a]::Float64))
  end

  lex2 = quote end
  pop!(lex2.args)

  if N == 1
    push!(lex2.args, :(r[1] = y[a]::Float64))
  else
    # unroll loop
    for i = Base.OneTo(nc)
      push!(lex2.args, :(r[$i] = y[a,$i]::Float64))
    end
  end

  lex = quote
    a, lp = idxrange(x, t)::Tuple{Int64, Bool}
    if lp
      $lex1
    else
      $lex2
    end
  end

  # aesthetic cleaning
  deleteat!(lex.args,[1,3])

  popfirst!(lex.args[2].args[2].args)
  lex.args[2].args[2] = lex.args[2].args[2].args[1]

  popfirst!(lex.args[2].args[3].args)
  lex.args[2].args[3] = lex.args[2].args[3].args[1]

  return quote
    @inbounds begin
      $lex
    end
    return nothing
  end
end

"""
  make_af(x::Array{Float64,1}, y::Array{Float64,N}, ::Val{ny})

make approximate function closure
"""
function make_af(x::Array{Float64,1}, y::Array{Float64,N}, ::Val{ny}) where {N, ny}

  af! = (t::Float64, r::Array{Float64,1}) ->
    begin
      approxf_full(t::Float64, r::Array{Float64,1}, x::Array{Float64,1}, y::Array{Float64,N}, Val(ny))
      return nothing
    end

  return af!
end

"""
    @generated function approxf_full(t ::Float64,
                                     r ::Array{Float64,1},
                                     x ::Array{Float64,1},
                                     y ::Array{Float64,1},
                                     ::Val{nc}) where {N, nc}

Returns the values of `y` at `t` using an approximation function
when `y` is an array of arrays.
"""
@generated function approxf_full(t ::Float64,
                                 r ::Array{Float64,1},
                                 x ::Array{Float64,1},
                                 y ::Array{Array{Float64,1},1},
                                 ::Val{nc}) where {nc}

  lex1 = quote end
  pop!(lex1.args)

  # unroll loop
  for i = Base.OneTo(nc)
    push!(lex1.args, :(r[$i] = linpred(t, xa, xap1, y[a][$i], y[a+1][$i])::Float64))
  end

  # add one assignment
  pushfirst!(lex1.args, :(xap1 = x[a+1]::Float64))
  pushfirst!(lex1.args, :(xa   = x[a]::Float64))

  lex2 = quote end
  pop!(lex2.args)

  # unroll loop
  for i = Base.OneTo(nc)
    push!(lex2.args, :(r[$i] = y[a][$i]::Float64))
  end

  lex = quote
    a, lp = idxrange(x, t)::Tuple{Int64, Bool}
    if lp
      $lex1
    else
      $lex2
    end
  end

  # aesthetic cleaning
  deleteat!(lex.args,[1,3])

  popfirst!(lex.args[2].args[2].args)
  lex.args[2].args[2] = lex.args[2].args[2].args[1]

  popfirst!(lex.args[2].args[3].args)
  lex.args[2].args[3] = lex.args[2].args[3].args[1]

  return quote
    @inbounds begin
      $lex
    end
    return nothing
  end
end

"""
  make_af(x::Array{Float64,1}, y::Array{Array{Float64,1},1}, ::Val{ny})

make approximate function closure
"""
function make_af(x::Array{Float64,1}, y::Array{Array{Float64,1},1}, ::Val{ny}) where {ny}

  af! = (t::Float64, r::Array{Float64,1}, y::Array{Array{Float64,1},1}) ->
    begin
      approxf_full(t::Float64, r::Array{Float64,1}, x::Array{Float64,1}, y::Array{Array{Float64,1},1}, Val(ny))
      return nothing
    end

  return af!
end


@testset "Limited Length Cycles" begin
    completedg = complete_digraph(4)
    pathdg = path_digraph(5)
    cycledg = cycle_digraph(5)

    @testset "complete digraph" for g in testgraphs(completedg)
        @test length(simplecycles_limited_length(g, 0)) == 0
        @test length(simplecycles_limited_length(g, 1)) == 0
        @test length(simplecycles_limited_length(g, 2)) == 6
        @test length(simplecycles_limited_length(g, 3)) == 14
        @test length(simplecycles_limited_length(g, 4)) == 20
        @test length(simplecycles_limited_length(g, 4, 10)) == 10
        @test length(simplecycles_limited_length(g, 4, typemax(Int))) == 20
    end

    @testset "path digraph" for g in testgraphs(pathdg)
        @test length(simplecycles_limited_length(g, 1)) == 0
        @test length(simplecycles_limited_length(g, 2)) == 0
        @test length(simplecycles_limited_length(g, 3)) == 0
        @test length(simplecycles_limited_length(g, 4)) == 0
        @test length(simplecycles_limited_length(g, 5)) == 0
    end

    @testset "cycle digraph" for g in testgraphs(cycledg)
        @test length(simplecycles_limited_length(g, 1)) == 0
        @test length(simplecycles_limited_length(g, 2)) == 0
        @test length(simplecycles_limited_length(g, 3)) == 0
        @test length(simplecycles_limited_length(g, 4)) == 0
        @test length(simplecycles_limited_length(g, 5)) == 1
    end

    @testset "self loops" begin
        selfloopg = DiGraph([
            0 1 0 0;
            0 0 1 0;
            1 0 1 0;
            0 0 0 1;
        ])
        cycles = simplecycles_limited_length(selfloopg, nv(selfloopg))
        @test [3] in cycles
        @test [4] in cycles
        @test [1, 2, 3] in cycles || [2, 3, 1] in cycles || [3, 1, 2] in cycles
        @test length(cycles) == 3
    end

    @testset "octahedral graph" begin
        octag = smallgraph(:octahedral)
        octadg = DiGraph(octag)
        octalengths, _ = simplecycleslength(octadg)
        for k = 1:6
            @test sum(octalengths[1:k]) == length(simplecycles_limited_length(octag, k))
            @test sum(octalengths[1:k]) == length(simplecycles_limited_length(octadg, k))
        end
    end
end