module LazyLieManifolds

using LinearAlgebra
using Manifolds
using StaticArrays
using DocStringExtensions

const LLM = LazyLieManifolds

export LLM
##==============================================================================
## Abstract Lie Groups and Manifolds
##==============================================================================

abstract type AbstractLieGroup end
abstract type AbstractLieAlgebra end

Base.getindex(p::AbstractLieGroup) = p.value
Base.getindex(X::AbstractLieAlgebra) = X.value

# extend Manifolds.jl
Manifolds.affine_matrix(p::AbstractLieGroup) = affine_matrix(manifold(p), p[])
Manifolds.screw_matrix(X::AbstractLieAlgebra) = affine_matrix(manifold(X), X[])

"""
    $(SIGNATURES)
"""
function manifold end

"""
    $(SIGNATURES)
"""
function _identity end

##------------------------------------------------------------------------------
## Extending Manifolds.jl
##------------------------------------------------------------------------------

function Manifolds.vee(X::AbstractLieAlgebra)
    M = manifold(X)
    return vee(M, make_identity(M, X[]), X[]) #FIXME
end

# we must use
# seVec = Manifolds.hat(M, p, coords)

function Base.:∘(p::AbstractLieGroup, q::AbstractLieGroup)
    M = manifold(p)
    typeof(p)(compose(M, p[], q[]))
end

function Base.inv(p::AbstractLieGroup)
    M = manifold(p)
    typeof(p)(inv(M, p[]))
end

function Base.:\(p::AbstractLieGroup, q::AbstractLieGroup)
    M = manifold(p)
    typeof(p)(compose(M, inv(M, p[]), q[]))
end

function Base.:/(p::AbstractLieGroup, q::AbstractLieGroup)
    M = manifold(p)
    typeof(p)(compose(M,  p[], inv(M, q[])))
end

function Base.exp(X::AbstractLieAlgebra) #look into retract with RightAction
    M = manifold(X)
    e = identity_element(M, X[]) #TODO this might not be completely correct
    G = convert(AbstractLieGroup, typeof(X))
    return G(exp(M, e, X[]))
end

function Base.log(p::AbstractLieGroup) #look into inverse_retract with RightAction
    M = manifold(p)
    e = identity_element(M, p[])
    g = convert(AbstractLieAlgebra, typeof(p))
    return g(log(M, e, p[]))
end

function Base.log(p::AbstractLieGroup, q::AbstractLieGroup)
    M = manifold(p)
    g = convert(AbstractLieAlgebra, typeof(p))
    return g(log(M, p[], q[]))
end

function Manifolds.group_exp(X::AbstractLieAlgebra) #look into retract with RightAction
    M = manifold(X)
    G = convert(AbstractLieGroup, typeof(X))
    return G(group_exp(M, X[]))
end

function Manifolds.distance(p::AbstractLieGroup, q::AbstractLieGroup)
    M = manifold(p)
    return distance(M, p[], q[])
end

Manifolds.manifold_dimension(p::AbstractLieGroup) = manifold_dimension(manifold(p))

##==============================================================================
## SpecialEuclidean SE{N} and se{N}
##==============================================================================
export se, SE

"""
    $TYPEDEF

Lie Group representation for `Manifolds.SpecialEuclidean(N)`

Related

[`se`](@ref)
"""
struct SE{N, T} <: AbstractLieGroup#TODO {N, T}
  value::T # Group
end

#TODO assert N and/or SE(p::ProductRepr) = ...
SE{N}(p::ProductRepr) where N = SE{N, typeof(p)}(p)

SE{N}() where N = _identity(SE{N})

function SE(R::Matrix{T}, t::Vector{T}) where T<:Real
    N = length(t)
    @assert size(R,1) == size(R,2) == N

    _t = SVector{N}(t)
    _R = SMatrix{N,N}(R)

    p = ProductRepr(_t, _R)

    return SE{N, typeof(p)}(p)
end

manifold(::SE{N,T}) where {N,T} = SpecialEuclidean(N)

getNparameter(::SE{N}) where N = N

#Default SE{N} representation
function _identity(::Type{SE{N}}, ::Type{T}=Float64) where {N, T}
    t = zeros(SVector{N, T})
    R = SMatrix{N,N,T}(one(T)I)
    SE{N}(ProductRepr(t, R))
end

"""
    $TYPEDEF

Lie algebra representation for `Manifolds.SpecialEuclidean(N)`

Related

[`SE`](@ref)
"""
struct se{N, T} <: AbstractLieAlgebra
  value::T # Tangent Vector (Lie algebra, skew simetric matrix)
end

#TODO assert N and/or se(X::ProductRepr) = ...
se{N}(X::ProductRepr) where N = se{N, typeof(X)}(X)

function se{N}(X::AbstractVector) where N
    l = length(X)
    @assert manifold_dimension(SpecialEuclidean(N)) == l "X dimension, $l, does not match manifold dimension, $N"

    M = SpecialEuclidean(N)

    e = _identity(SE{N}, eltype(X))[]
    _X = hat(M, e, X)

    return se{N}(_X)
end

manifold(::se{N,T}) where {N,T} = SpecialEuclidean(N)

Base.convert(::Type{AbstractLieGroup}, ::Type{se{N, T}}) where {N,T} = SE{N}
Base.convert(::Type{AbstractLieAlgebra}, ::Type{SE{N, T}}) where {N,T} = se{N}

# function default_point_type(M::SpecialOrthogonal, ::Type{ET}=Float64, ::Type{AT}=SMatrix) where {ET, AT}
#     s = representation_size(M)
#     return AT{s..., ET, prod(s)}
# end

# function default_point_type(M::SpecialEuclidean, ::Type{ET}=Float64, ::Type{AT}=SArray) where {ET, AT}

#     s1 = representation_size(M.manifold[1])
#     s2 = representation_size(M.manifold[2])
#     return ProductRepr{Tuple{AT{Tuple{s1[1]}, ET}, AT{Tuple{s2[1],s2[2]}, ET}}}
# end

##------------------------------------------------------------------------------
## other operators
##------------------------------------------------------------------------------
#TODO abstract
function ⊕(p::SE{N}, X::se{N}) where N
    M = manifold(p)
    e = identity_element(M, p[])
    typeof(p)(compose(M, p[], exp(M, e, X[])))
end

function ⊕(X::se{N}, p::SE{N}) where N
    M = manifold(p)
    e = identity_element(M, p[])
    typeof(p)(compose(M, exp(M, e, X[]), p[]))
end

# ⊖ - right minus
# X = log(M, e, inv(x) ∘ y)
function ⊖(y::SE{N}, x::SE{N}) where N
    M = manifold(y)
    e = identity_element(M, y[])
    return log(M, e, inv(x) ∘ y )
end

##==============================================================================
## SpecialOrthogonal SO{N} and so{N}
##==============================================================================
export so, SO

"""
    $TYPEDEF

Lie Group representation for `Manifolds.SpecialOrthogonal(N)`

Related

[`so`](@ref)
"""
struct SO{N, T} <: AbstractLieGroup
  value::T # Group
end

#TODO assert N and/or SO(p::ProductRepr) = ...
SO{N}(p::StaticMatrix{N,N,T}) where {N, T} = SO{N, typeof(p)}(p)

SO{N}() where N = _identity(SO{N})

function SO(R::Matrix{T}) where T<:Real
    N = size(R,1)
    @assert size(R,1) == size(R,2)

    return SO{N}(SMatrix{N,N}(R))
end

manifold(::SO{N,T}) where {N,T} = SpecialOrthogonal(N)

#Default SO{N} representation
function _identity(::Type{SO{N}}, ::Type{T}=Float64) where {N, T}
    R = SMatrix{N,N,T}(one(T)I)
    SO{N}(R)
end

"""
    $TYPEDEF

Lie algebra representation for `Manifolds.SpecialOrthogonal(N)`

Related

[`SO`](@ref)
"""
struct so{N, T} <: AbstractLieAlgebra
  value::T # Tangent Vector (Lie algebra, skew simetric matrix)
end

#TODO assert N and/or so(X::ProductRepr) = ...
so{N}(X::StaticMatrix{N,N,T}) where {N,T} = so{N, typeof(X)}(X)

function so{N}(X::AbstractVector) where N
    l = length(X)
    @assert manifold_dimension(SpecialOrthogonal(N)) == l "X dimension, $l, does not match manifold dimension, $N"

    M = SpecialOrthogonal(N)

    e = _identity(SO{N}, eltype(X))[]
    _X = hat(M, e, X)

    return so{N}(_X)
end

manifold(::so{N,T}) where {N,T} = SpecialOrthogonal(N)

Base.convert(::Type{AbstractLieGroup}, ::Type{so{N, T}}) where {N,T} = SO{N}
Base.convert(::Type{AbstractLieAlgebra}, ::Type{SO{N, T}}) where {N,T} = so{N}

##==============================================================================
## TranslationGroup 𝕋{N} and 𝕥{N} feels a bit redundant, but adding for completeness
##==============================================================================
export 𝕥, 𝕋

"""
    $TYPEDEF

Lie Group representation for `Manifolds.TranslationGroup(N)`

Related

[`𝕥`](@ref)
"""
struct 𝕋{N, T} <: AbstractLieGroup
  value::T # Group
end

#TODO assert N and/or 𝕋(p::ProductRepr) = ...
𝕋{N}(p::SVector{N,T}) where {N, T} = 𝕋{N, typeof(p)}(p)

𝕋{N}() where N = _identity(𝕋{N})

function 𝕋(t::Vector{T}) where T<:Real
    N = size(t,1)

    return 𝕋{N}(SVector{N, T}(t))
end

manifold(::𝕋{N,T}) where {N,T} = TranslationGroup(N)

#Default 𝕋{N} representation
function _identity(::Type{𝕋{N}}, ::Type{T}=Float64) where {N, T}
    t = zeros(SVector{N, T})
    𝕋{N}(t)
end

"""
    $TYPEDEF

Lie algebra representation for `Manifolds.TranslationGroup(N)`

Related

[`𝕋`](@ref)
"""
struct 𝕥{N, T} <: AbstractLieAlgebra
  value::T # Tangent Vector (Lie algebra)
end

#TODO assert N and/or 𝕥(X::ProductRepr) = ...
𝕥{N}(X::SVector{N,T}) where {N,T} = 𝕥{N, typeof(X)}(X)

function 𝕥{N}(X::AbstractVector) where N
    l = length(X)
    @assert manifold_dimension(TranslationGroup(N)) == l "X dimension, $l, does not match manifold dimension, $N"

    M = TranslationGroup(N)

    e = _identity(𝕋{N}, eltype(X))[]
    _X = hat(M, e, X)

    return 𝕥{N}(_X)
end

manifold(::𝕥{N,T}) where {N,T} = TranslationGroup(N)

Base.convert(::Type{AbstractLieGroup}, ::Type{𝕥{N, T}}) where {N,T} = 𝕋{N}
Base.convert(::Type{AbstractLieAlgebra}, ::Type{𝕋{N, T}}) where {N,T} = 𝕥{N}

##==============================================================================
## TODO Quaternion Group https://github.com/JuliaManifolds/Manifolds.jl/issues/382
##==============================================================================

##==============================================================================
## RoME like functions sandbox
##==============================================================================
export PosePose, PosePose_x, PosePose_d, Prior_d

# this is in the "local frame" sequential, a similar operation can be done with the group_log (but together)
function PosePose(ⁱmⱼ::SE{N}, ʷxᵢ::SE{N}, ʷxⱼ::SE{N}) where N

    ʷT̂ⱼ = ʷxᵢ ∘ ⁱmⱼ

    ʲT̂ⱼ = ʷxⱼ \ ʷT̂ⱼ

    X = log(ʲT̂ⱼ)

    return vee(X)
end

# this is in the "global frame" e component wise on submanifolds
function PosePose_x(ⁱmⱼ::SE{N}, ʷxᵢ::SE{N}, ʷxⱼ::SE{N}) where N
    M = manifold(ʷxᵢ)

    ʷT̂ⱼ = ʷxᵢ ∘ ⁱmⱼ
    X = se{N}(log(M, ʷxⱼ[], ʷT̂ⱼ[])) # TODO check sign, this (currently) is the vector at ʷT̂ⱼ that points to ʷxⱼ

    return vee(X)
end

# this is in the "global frame" e component wise on submanifolds, but uses distance
function PosePose_d(m, p, q)

    q̂ = p ∘ m
    return distance(q, q̂)

end

#manifolds point to manifold point factor
function PointPoint_d(m, p, q)

    q̂ = p ∘ m
    return distance(q, q̂)

end

function PointPoint_d(X::AbstractLieAlgebra, p, q)
    q̂ = exp(p, X)
    # q̂ = p ⊕ X
    return distance(q, q̂)
end

# manifold prior factor
function Prior_d(meas, p)
#
    return distance(meas, p)
end

function PointPoint_velocity_d(X, dt, p, q)
    q̂ = p ∘ group_exp(X*dt)
    return distance(q, q̂)
end

function Pose2Pose2(meas, wxi, wxj)
    wTjhat = SE2(wxi)*SE2(meas)
    jTjhat = SE2(wxj) \ wTjhat
    return se2vee(jTjhat)
end

end


PowerLawPlateau =  RheoModelClass(
        # Model name
        name="burgerliquid",
        # Model parameters,
        p = (:Gᵩ, :G₀, :τ, :α),
        # Relaxation modulus
        G = quote
                Gᵩ + (G₀ - Gᵩ)./(1 + t/τ).^(α)
            end,
        # Network
        info= "                "
        )


import GreyDecision.Utility: makeminmax
import GreyDecision.GreyNumbers: GreyNumber
import GreyDecision.Vikor: vikor

@testset "Grey Vikor with real numbers" begin
    tol = 0.0001
    w =  [0.110, 0.035, 0.379, 0.384, 0.002, 0.002, 0.010, 0.077]
    #Amat = [
    #  100 92 10 2 80 70 95 80 ;
    #  80  70 8  4 100 80 80 90 ;
    #  90 85 5 0 75 95 70 70 ;
    #  70 88 20 18 60 90 95 85
    #]
    Amat = [
    GreyNumber{Float64}(100.0, 100.0) GreyNumber{Float64}(92.0, 92.0) GreyNumber{Float64}(10.0, 10.0) GreyNumber{Float64}(2.0, 2.0) GreyNumber{Float64}(80.0, 80.0) GreyNumber{Float64}(70.0, 70.0) GreyNumber{Float64}(95.0, 95.0) GreyNumber{Float64}(80.0, 80.0);
    GreyNumber{Float64}(80.0, 80.0) GreyNumber{Float64}(70.0, 70.0) GreyNumber{Float64}(8.0, 8.0) GreyNumber{Float64}(4.0, 4.0) GreyNumber{Float64}(100.0, 100.0) GreyNumber{Float64}(80.0, 80.0) GreyNumber{Float64}(80.0, 80.0) GreyNumber{Float64}(90.0, 90.0);
    GreyNumber{Float64}(90.0, 90.0) GreyNumber{Float64}(85.0, 85.0) GreyNumber{Float64}(5.0, 5.0) GreyNumber{Float64}(0.0, 0.0) GreyNumber{Float64}(75.0, 75.0) GreyNumber{Float64}(95.0, 95.0) GreyNumber{Float64}(70.0, 70.0) GreyNumber{Float64}(70.0, 70.0);
    GreyNumber{Float64}(70.0, 70.0) GreyNumber{Float64}(88.0, 88.0) GreyNumber{Float64}(20.0, 20.0) GreyNumber{Float64}(18.0, 18.0) GreyNumber{Float64}(60.0, 60.0) GreyNumber{Float64}(90.0, 90.0) GreyNumber{Float64}(95.0, 95.0) GreyNumber{Float64}(85.0, 85.0)
    ]

    fns = makeminmax([maximum, maximum, maximum, maximum, maximum, maximum, maximum, maximum])

    vikorresult = vikor(Amat, w, fns)

    @test vikorresult[:bestindex] == 4

    @test isapprox(vikorresult[:scores][1], GreyNumber(0.1975012087551764, 0.1975012087551764), atol=tol)
    @test isapprox(vikorresult[:scores][2], GreyNumber(0.2194064473270817, 0.2194064473270817), atol=tol)
    @test isapprox(vikorresult[:scores][3], GreyNumber(0.3507643203516215, 0.3507643203516215), atol=tol)
    @test isapprox(vikorresult[:scores][4], GreyNumber(-0.16727341435277993, -0.16727341435277993), atol=tol)

end


"""
calculates main overlap metrics
based on https://github.com/Visceral-Project/EvaluateSegmentation/blob/master/source/DiceCoefficientMetric.h
https://github.com/Visceral-Project/EvaluateSegmentation/blob/master/source/JaccardCoefficientMetric.h
https://github.com/Visceral-Project/EvaluateSegmentation/blob/master/source/GlobalConsistencyError.h
"""
module MainOverlap
using ..BasicStructs, Parameters

export calculateBAsicOverlap, dice, jaccard, gce

"""
calculate dice coefficient based on precalulated constants
TnTpFpFns - list of basic metrics - in the order [tn,tp,fp,fn ]
res - ResultMetrics struct holding result of all metrics calculated for given run
return modified ResultMetrics with added Dice, Jaccard and global consistency error coefficients
"""
function calculateBAsicOverlap(TnTpFpFns::Vector{Float64} , res::ResultMetrics) ::ResultMetrics
    tn = TnTpFpFns[1]
    tp = TnTpFpFns[2]
    fp = TnTpFpFns[3]
    fn = TnTpFpFns[4]

    return setproperties(res, (dice=  dice(tn,tp,fp) ,jaccard=  jaccard(tn,tp,fp) ,gce =  gce(tn,tp,fp,fn)  ))

end #calculateDice

```@doc
Calculates Dice Coefficient  on the basis of true negative, true positive, false positive and false negatives
 ```
function dice(tp,fp, fn)::Float64
 return  2*tp/(2*tp + fp + fn)
end#dice

```@doc
Calculates Jaccard Coefficient  on the basis of true negative, true positive, false positive and false negatives
```
function jaccard(tp,fp, fn)::Float64
 return  tp/(tp + fp + fn)
end#jaccard

```@doc
Calculates Global Consistency error on the basis of true negative, true positive, false positive and false negatives
```
function gce(tn,tp,fp, fn)::Float64
    n = tn+fp+fn+tp;
    e1 = ( fn*(fn+ 2*tp)/(tp+fn) + fp*(fp + 2*tn)/(tn+fp) )/n;
    e2 = ( fp*(fp+2*tp)/(tp+fp) + fn*(fn + 2*tn)/(tn+fn) )/n;
    return min(e1, e2);
end#dice

end #MainOverlap


using Persa
using Statistics
using Random

mutable struct Baseline{T<:AbstractFloat} <: Persa.Model
    γ::T
    λ::T
    μ::T
    bu::Vector{T}
    bv::Vector{T}
end

function Baseline(dataset::Dataset; γ::Float64 = 0.001, λ::Float64 = 0.02)
    u = users(dataset)
    v = items(dataset)

    μ = mean(dataset)
    bu = zeros(u)
    bv = zeros(v)

    return Baseline(γ, λ, μ, bu, bv)
end

function Persa.predict(model::Baseline, pref::Preference, user::Int, item::Int)
    model.μ + model.bu[user] + model.bv[item]
end

function Persa.objective(model::Baseline, dataset::Dataset)
    pref = preference(dataset)
    total = 0

    for (u, v, r) in dataset
        total += (r - predict(model, pref, u, v))^2
        total += 0.1 * (model.bu[u]^2 + model.bv[v]^2)
    end

    return total
end

function Persa.update!(model::Baseline, dataset::Dataset)
    γ = model.γ
    λ = model.λ
    pref = preference(dataset)
    idx = shuffle(1:length(dataset))

    for i in 1:length(dataset)
        u, v, r = dataset[idx[i]]

        e = r - predict(model, pref, u, v)

        model.bu[u] += γ * (e .- λ * model.bu[u])
        model.bv[v] += γ * (e .- λ * model.bv[v])
    end
end

data = [
    (user = 1, item = 1, rating = 4),
    (user = 1, item = 3, rating = 2),
    (user = 2, item = 1, rating = 1),
    (user = 2, item = 2, rating = 5),
    (user = 2, item = 3, rating = 3),
    (user = 3, item = 1, rating = 3),
    (user = 3, item = 2, rating = 4),
    (user = 4, item = 2, rating = 5),
    (user = 3, item = 3, rating = 4)
]

dataset = Dataset(data)

model = Baseline(dataset)

mach = machine(model, dataset)

train!(mach, dataset)

mach[1, 3]


function roundTT(mpt::MPT, ϵ::Float64)

    N           = order(mpt);
    cores       = Vector{Array{Float64,3}}(undef,N);
    middlesizes = size(mpt);
    frobnorm    = norm(mpt);
    err2        = 0;

    # Right-to-left orthogonalization
        G = mpt[N];
        for k = N:-1:2
            G        = unfold(G,[1],"left");
            F        = qr(G');
            Q        = Matrix(F.Q);
            R        = Matrix(F.R);
            G        = reshape(Q', ( size(Q,2), middlesizes[k][1], Int(length(Q)/(size(Q,2)*middlesizes[k][1])) ) );
            cores[k] = G;
            G        = nmodeproduct(R,mpt[k-1],3);
        end
        cores[1] = G;
        # Compression of the orthogonalized representation
        δ = ϵ / sqrt(N-1) * frobnorm;
        for k = 1:N-1
            F     = svd(unfold(cores[k],[3],"right"),full=false);
            rcurr = length(F.S);
            sv2   = cumsum(reverse(F.S).^2);
            tr    = Int(findfirst(sv2 .> δ^2))-1;
            if tr > 0
                rcurr = length(F.S) - tr;
                err2 += sv2[tr];
            end
            Utr = F.U[:,1:rcurr];
            Str = F.S[1:rcurr];
            Vtr = F.Vt[1:rcurr,:];

            cores[k]   = reshape( Utr,( Int(length(Utr)/(size(cores[k],2)*size(Vtr,1))), size(cores[k],2), size(Vtr,1) ) );
            cores[k+1] = nmodeproduct(Matrix((Vtr'*Diagonal(Str))'),cores[k+1],1);
        end
        return MPT(cores);
    end

    function roundTT(ttm::MPT{4}, ϵ::Float64)
        tt = mpo2mps(ttm);
        tt = roundTT(tt, ϵ)
        return mps2mpo(tt,size(ttm))
    end


######################################################
#
# Markov-Switching Multifractal (MSM)
#
# Date:   2014-11-19
# Author: Hessel Tuinhof
#
# Function to simulate a MSM process.
#
# Reference:
# - Stein, William: https://raw.githubusercontent.com/sagemath/sagelib/master/sage/finance/markov_multifractal_cython.pyx
#
# Required packages: 	Distributions
#
# Input:
#   num_sim ... number of simulatons, positive integer
#   sim_length ... length of simulation, positive integer
#   m0 ... multiplier, float in [0,2]
#   kbar ... number of frequencies, positive integer
#   sigma ... standard deviation of returns, positve float
#   gamma_kbar ... transition probability, float in [0,1)
#   b ... geometric rate, float in (1,inf)
#
# Output:
#   msm ... simulated log returns
#
# Examples:
#   # Simulation of log returns (1 simulation)
#   plot(MSM(1, 4000, 1.4, 8, 0.55, 0.92, 3))
#
#   # Simulation of log price (10 simulations)
#   plot(cumsum(MSM(10, 4000, 1.4, 8, 0.55, 0.92, 3)))
#
######################################################

#####################   Warning   #####################
# For large simulation purposes the use of the array type
# is not recommendable. Instead one should save the sim-
# ulated returns in an DataFrame or TimeSeries type for
# improved performance.
#####################   Warning   #####################

using Distributions

function MSM(num_sim, sim_length, m0, kbar, sigma, gamma_kbar, b)
	m1 = 2-m0
	gammaone = 1 - exp(log(1-gamma_kbar)/(b^(kbar-1)))
	gamma  = zeros(kbar)
	for i = 1:kbar
		gamma[i] = 1-(1-gammaone)^(b^(i-1))
	end
	markovstatevector = zeros(kbar)
	msm = zeros(sim_length, num_sim)
	for i = 1:num_sim
		for j = 1:kbar
			convert(Bool, rand(Int64) & 1) ? markovstatevector[j] = m0 : markovstatevector[j] = m1
		end
		err = randn(sim_length)
		for j = 1:sim_length
			msm[j,i] = sigma * err[j] * sqrt(prod(markovstatevector))
			for k = 1:kbar
				if rand() <= gamma[k]
					convert(Bool, rand(Int64) & 1) ? markovstatevector[k] = m0 : markovstatevector[k] = m1
				end
			end
		end
	end
	return msm
end


using YaoExtensions, Yao
using Test, Random
using Optim: LBFGS, optimize
using Optim

"""
    learn_u4(u::AbstractMatrix; niter=100)

Learn a general U4 gate. The optimizer is LBFGS.
"""
function learn_u4(u::AbstractBlock; niter=100)
    ansatz = general_U4()
    params = parameters(ansatz)
    println("initial loss = $(operator_fidelity(u,ansatz))")
    optimize(x->-operator_fidelity(u, dispatch!(ansatz, x)),
            (G, x) -> (G .= -operator_fidelity'(u, dispatch!(ansatz, x))[2]),
            parameters(ansatz),
            LBFGS(),
            Optim.Options(iterations=niter))
    println("final fidelity = $(operator_fidelity(u,ansatz))")
    return ansatz
end

using Random
Random.seed!(2)
u = matblock(rand_unitary(4))
c = learn_u4(u; niter=150)


abstract type Surface end

"""
Specify a surface as list of geometric objects.
"""
mutable struct SurfaceAssembly <: Surface
    objects::Array{Surface,1}
end

"""
Geometric sphere surface.
"""
struct Sphere <: Surface
    center::Vector
    radius::Real
    material::Material
end

"""
    hit(sa, ray, tmin, tmax)

Obtain the closest hit record for a ray intersecting the objects stored in a `SurfaceAssembly`.
"""
function hit(sa::SurfaceAssembly, ray::Ray, tmin::Real, tmax::Real)
    # hit record
    rec = nothing
    closest_so_far = tmax
    for obj in sa.objects
        currec, t = hit(obj, ray, tmin, closest_so_far)
        if currec != nothing
            closest_so_far = t;
            rec = currec
        end
    end
    return (rec, closest_so_far)
end

"""
    hit(sphere, ray, tmin, tmax)

Obtain the hit record for a ray intersecting a sphere.
"""
function hit(sphere::Sphere, ray::Ray, tmin::Real, tmax::Real)
    oc = ray.origin - sphere.center
    a = dot(ray.direction, ray.direction)
    b = dot(oc, ray.direction)
    c = dot(oc, oc) - sphere.radius^2;
    discriminant = b^2 - a*c
    if discriminant > 0
        # solutions of the quadratic equation
        t1 = -(b + sign(b)*sqrt(discriminant)) / a
        t2 = c / (a * t1)
        # smaller solution first
        for t in sort([t1, t2])
            if tmin <= t && t < tmax
                point = point_at_parameter(ray, t)
                normal = (point - sphere.center) / sphere.radius
                return (HitRecord(point, normal, sphere.material), t)
            end
        end
    end
    return (nothing, tmax)
end


export SHalley_dir_AD

function SHalley_dir_AD(nlp            :: AbstractNLPModel,
                       nlp_at_x       :: NLPAtX,
                       x              :: AbstractVector,
                       Nwtdirection   :: Function)

    n = length(x)

    dₙ = Nwtdirection(nlp_at_x.Hx, nlp_at_x.gx, verbose = false)

    e₁ = zeros(n)
    e₁[1] = 1.0
    ∇³fdₙ = ∇f³xuv(nlp, x, dₙ, e₁)
    for i = 2:n
        eᵢ = zeros(n)
        eᵢ[i] = 1.0
        ∇³fdₙ = hcat(∇³fdₙ, ∇f³xuv(nlp, x, dₙ, eᵢ))
    end

    hess_plus_tenseur = (nlp_at_x.Hx + ∇³fdₙ)

    dV = Nwtdirection(hess_plus_tenseur, nlp_at_x.gx)

    grad_plus_tuv = -(nlp_at_x.gx + 0.5 .* ∇f³xuv(nlp, x, dₙ, dV))

    dSH = Nwtdirection(nlp_at_x.Hx, -grad_plus_tuv)

    return dSH
end


using UnicodePlots: heatmap
function part1(input)
    coords, folds = parse_input(input)
    coords = fold(coords, first(folds))
    return length(coords)
end

function part2(input)
    coords, folds = parse_input(input)
    for Fold in folds
        coords = fold(coords, Fold)
    end
    m = zeros(Int, 6, 39)
    for coord in coords
        x, y = coord
        m[y + 1, x + 1] = 1
    end
    return heatmap(m[end:-1:1, :])
end

function parse_input(input)
    parts = split(readchomp(input), "\n\n")
    coords = parse_coord.(split(parts[1], "\n"))
    folds = parse_fold.(split(parts[2], "\n"))
    return coords, folds
end

parse_coord(line) = parse.(Int, split(line, ","))
function parse_fold(line)
    parts = split(last(split(line, " ")), "=")
    return parts[1], parse(Int, parts[2])
end

function fold(coords, fold)
    dir, pos = fold
    if dir == "x"
        return unique((mirror(x, pos), y) for (x, y) in coords)
    else
        return unique((x, mirror(y, pos)) for (x, y) in coords)
    end
end

mirror(x, pos) = x < pos ? x : 2 * pos - x


export GMM, get_normalizer

mutable struct GMM{T<:Real}
    # dimension
    d::Int

    # number of components
    K::Int

    # weight, mean, covariance
    # NOTES:
    #   - π is scalar while μ and σ are d-dimensional
    #   - right now, we only support diagonal σ

    π::AbstractArray{T}
    μ::AbstractArray{T}
    σ::AbstractArray{T}

    function GMM{T}(d::Int, K::Int) where T<:Real
        π = CUDA.fill(1.0f0 / K, K)
        μ = CuArray{T}(undef, (d, K))
        σ = CUDA.ones(d, K) # start with unit variances

        new(d, K, π, μ, σ)
    end
end

function Base.show(io::IO, gmm::GMM{T}) where {T<:Real}
    println(io, "$(gmm.d)-dimensional GMM{$T} with $(gmm.K) components:\n")
    println(io, "\tπ: $(gmm.π)")
    println(io, "\tμ: $(gmm.μ)")
    println(io, "\tσ: $(gmm.σ)")
end

function get_normalizer(σ::AbstractArray{T}, d::Int) where {T<:Real}
    # the normalizer of the d-dimensional normal distribution is sqrt((2π)^d |Σ|)
    # since we will always use log-transformed quantities, we thus need
    # log(1/(((2π)^d |Σ|)^(1/2))) = log(((2π)^d |Σ|)^(-1/2)) = -1/2 * (log(2π)*d + log(|Σ|))

    # right now, we only allow diagonal covariances => |Σ| = Π_{i=1}^d σ_i => log(|Σ|) = Σ_{i=1}^d log(σ_i)
    return -0.5f0 .* ((log(2.0f0 * pi) * d) .+ sum(log.(σ); dims=1))
end


using QuasinormalModes
using SymEngine

# ------------------------------------------------------------------
# 1. Analytic Schwarzschild Black Hole
# ------------------------------------------------------------------

struct SchwarzschildData{N,T} <: QuadraticEigenvalueProblem{N,T}
    nIter::N
    x0::T

    vars::Tuple{Basic, Basic}
    exprs::Tuple{Basic, Basic}
end

function SchwarzschildData(nIter::N, x0::T, l::N, s::N) where {N,T}
    vars = @vars x ω

    λ0 = (-1 + (2*im)*ω + x*(4 - 3*x + (4*im)*(-2 + x)*ω))/((-1 + x)^2*x)
    S0 = (l + l^2 + (-1 + s^2)*(-1 + x) + (4*im)*(-1 + x)*ω + 4*(-2 + x)*ω^2)/((-1 + x)^2*x)

    return SchwarzschildData{N,T}(nIter, x0, vars, (λ0, S0))
end

QuasinormalModes.λ0(d::SchwarzschildData{N,T}) where {N,T} = d.exprs[1]
QuasinormalModes.S0(d::SchwarzschildData{N,T}) where {N,T}  = d.exprs[2]

QuasinormalModes.get_niter(d::SchwarzschildData{N,T}) where {N,T} = d.nIter
QuasinormalModes.get_x0(d::SchwarzschildData{N,T}) where {N,T} = d.x0

QuasinormalModes.get_ODEvar(d::SchwarzschildData{N,T}) where {N,T} = d.vars[1]
QuasinormalModes.get_ODEeigen(d::SchwarzschildData{N,T}) where {N,T} = d.vars[2]

# ------------------------------------------------------------------
# 2. Numeric Schwarzschild Black Hole
# ------------------------------------------------------------------

struct NSchwarzschildData{N,T} <: NumericAIMProblem{N,T}
    nIter::N
    x0::T
    l::N
    s::N
end

function NSchwarzschildData(nIter::N, x0::T, l::N, s::N) where {N,T}
    return NSchwarzschildData{N,T}(nIter, x0, l, s)
end

QuasinormalModes.λ0(::NSchwarzschildData{N,T}) where {N,T} = (x,ω) -> (-1 + (2*im)*ω + x*(4 - 3*x + (4*im)*(-2 + x)*ω))/((-1 + x)^2*x)
QuasinormalModes.S0(d::NSchwarzschildData{N,T}) where {N,T} = (x,ω) -> (d.l + d.l^2 + (-1 + d.s^2)*(-1 + x) + (4*im)*(-1 + x)*ω + 4*(-2 + x)*ω^2)/((-1 + x)^2*x)

QuasinormalModes.get_niter(d::NSchwarzschildData{N,T}) where {N,T} = d.nIter
QuasinormalModes.get_x0(d::NSchwarzschildData{N,T}) where {N,T} = d.x0

# ------------------------------------------------------------------
# 3. Constructing problems and caches
# ------------------------------------------------------------------

p_ana = SchwarzschildData(0x00030, Complex(BigFloat("0.43"), BigFloat("0.0")), 0x00000, 0x00000);
p_num = NSchwarzschildData(0x00030, Complex(0.43, 0.0), 0x00000, 0x00000);

c_ana = AIMCache(p_ana)
c_num = AIMCache(p_num)

# ------------------------------------------------------------------
# 4. Computing quasinormal modes
# ------------------------------------------------------------------

m_ana = computeEigenvalues(Serial(), p_ana, c_ana)

function printQNMs(qnms, cutoff, instab)
    println("-"^165)
    println("|", " "^36, "Re(omega)", " "^36, " ", " "^36, "Im(omega)", " "^36, "|")
    println("-"^165)

    for qnm in qnms
        if real(qnm) > cutoff && ( instab ? true : imag(qnm) < big"0.0" )
        println(real(qnm), "    ", imag(qnm))
        end
    end

    println("-"^165)

    return nothing
end

sort!(m_ana, by = x -> imag(x))
printQNMs(m_ana, 1.0e-10, false)

ev = computeEigenvalues(Serial(), p_num, c_num, Complex(0.22, -0.20), nls_xtol = 1.0e-10, nls_ftol = 1.0e-10)


# ---
# title: 891. Sum of Subsequence Widths
# id: problem891
# author: Indigo
# date: 2021-02-01

# difficulty: Hard
# categories: Array, Math
# link: <https://leetcode.com/problems/sum-of-subsequence-widths/description/>
# hidden: true
# ---
#
# Given an array of integers `A`, consider all non-empty subsequences of `A`.
#
# For any sequence S, let the  _width_  of S be the difference between the
# maximum and minimum element of S.
#
# Return the sum of the widths of all subsequences of A.
#
# As the answer may be very large, **return the answer modulo 10^9 + 7**.
#
#
#
# **Example 1:**
#
#
#
#     Input: [2,1,3]
#     Output: 6
#     Explanation: Subsequences are [1], [2], [3], [2,1], [2,3], [1,3], [2,1,3].
#     The corresponding widths are 0, 0, 0, 1, 1, 2, 2.
#     The sum of these widths is 6.
#
#
#
#
# **Note:**
#
#   * `1 <= A.length <= 20000`
#   * `1 <= A[i] <= 20000`
#
#
## @lc code=start
using LeetCode

function sum_subseq_widths(A::Vector{Int})
    MOD = Int(1e9 + 7)
    len = length(A)
    sort!(A)
    pow = fill(2, len)
    for i in 2:len
        pow[i] = (pow[i - 1] << 1) % MOD
    end
    res = 0
    for i in 2:(len - 1)
        res += (A[i]) * ((pow[i - 1] - pow[len - i])) % MOD
    end
    return mod1((res + (A[end] - A[1]) * (pow[len - 1] - 1) % MOD), MOD)
end
## @lc code=end


"""
    function LPHGLET_Synthesis(dvec::Vector{Float64}, GP::GraphPart, BS::BasisSpec, G::GraphSig; gltype::Symbol = :L, ϵ::Float64 = 0.3)

Perform Lapped-HGLET Synthesis transform

### Input Arguments
* `dvec`: the expansion coefficients corresponding to the chosen basis
* `GP`: a GraphPart object
* `BS`: a BasisSpec object
* `G`: a GraphSig object
* `gltype`: :L or :Lsym, indicating which eigenvectors are used
* `ϵ`: relative action bandwidth (default: 0.3)

### Output Argument
* `f`: the reconstructed signal
* `GS`: the reconstructed GraphSig object
"""
function LPHGLET_Synthesis(dvec::Matrix{Float64}, GP::GraphPart, BS::BasisSpec, G::GraphSig; gltype::Symbol = :L, ϵ::Float64 = 0.3)
    # Preliminaries
    W = G.W
    inds = GP.inds
    rs = GP.rs
    N = size(W, 1)
    jmax = size(rs, 2)
    Uf = Matrix{Float64}(I, N, N)
    used_node = Set()

    # fill in the appropriate entries of dmatrix
    dmatrix = dvec2dmatrix(dvec, GP, BS)

    f = zeros(size(dmatrix[:, jmax, :]))

    # Perform the synthesis transform
    for j = 1:jmax
        regioncount = count(!iszero, rs[:,j]) - 1
        # assemble orthogonal folding operator at level j - 1
        keep_folding!(Uf, used_node, W, GP; ϵ = ϵ, j = j - 1)
        for r = 1:regioncount
            # indices of current region
            indr = rs[r, j]:(rs[r + 1, j] - 1)
            # indices of current region's nodes
            indrs = inds[indr, j]
            # number of nodes in current region
            n = length(indrs)

            # only proceed forward if coefficients do not exist
            if (j == jmax || count(!iszero, dmatrix[indr, j + 1, :]) == 0) && count(!iszero, dmatrix[indr, j, :]) > 0
                # compute the eigenvectors
                W_temp = W[indrs,indrs]
                D_temp = Diagonal(vec(sum(W_temp, dims = 1)))
                if gltype == :L
                    # compute the eigenvectors of L ==> svd(L)
                    v = svd(Matrix(D_temp - W_temp)).U
                elseif gltype == :Lsym
                    # check if one can assemble the Lsym
                    if minimum(sum(W[indrs, indrs], dims = 1)) > 10^3 * eps()
                        ### eigenvectors of L_sym ==> svd(L_sym)
                        D_temp_p = Diagonal(vec(sum(W_temp, dims = 1)).^(-0.5))
                        v = svd(Matrix(D_temp_p * (D_temp - W_temp) * D_temp_p)).U
                    else
                        ### eigenvectors of L ==> svd(L)
                        v = svd(Matrix(D_temp - W_temp)).U
                    end
                end
                v = v[:,end:-1:1]

                # standardize the eigenvector signs
                standardize_eigenvector_signs!(v)

                # construct unfolder operator custom to current region
                P = Uf[indrs, :]'

                # reconstruct the signal
                f += (P * v) * dmatrix[indr, j, :]

            end
        end
    end

    # creat a GraphSig object with the reconstructed data
    GS = deepcopy(G)
    replace_data!(GS, f)

    return f, GS
end

"""
    function LPHGLET_Analysis_All(G::GraphSig, GP::GraphPart; ϵ::Float64 = 0.3)

For a GraphSig object 'G', generate the 2 matrices of Lapped-HGLET expansion coefficients
corresponding to the eigenvectors of L and Lsym

### Input Arguments
* `G`:  a GraphSig object
* `GP`: a GraphPart object
* `ϵ`: relative action bandwidth (default: 0.3)

### Output Argument
* `dmatrixlH`:        the matrix of expansion coefficients for L
* `dmatrixlHsym`:     the matrix of expansion coefficients for Lsym
* `GP`:              a GraphPart object
"""
function LPHGLET_Analysis_All(G::GraphSig, GP::GraphPart; ϵ::Float64 = 0.3)
    # Preliminaries
    W = G.W
    inds = GP.inds
    rs = GP.rs
    N = size(W, 1)
    jmax = size(rs, 2)
    fcols = size(G.f, 2)
    Uf = Matrix{Float64}(I, N, N)
    used_node = Set()
    dmatrixlH = zeros(N, jmax, fcols)
    dmatrixlHsym = deepcopy(dmatrixlH)

    for j = 1:jmax
        regioncount = count(!iszero, rs[:,j]) - 1
        # assemble orthogonal folding operator at level j - 1
        keep_folding!(Uf, used_node, W, GP; ϵ = ϵ, j = j - 1)
        for r = 1:regioncount
            # indices of current region
            indr = rs[r, j]:(rs[r + 1, j] - 1)
            # indices of current region's nodes
            indrs = inds[indr, j]
            # number of nodes in current region
            n = length(indrs)

            # compute the eigenvectors
            W_temp = W[indrs,indrs]
            D_temp = Diagonal(vec(sum(W_temp, dims = 1)))
            ## eigenvectors of L ==> svd(L)
            v = svd(Matrix(D_temp - W_temp)).U
            ## eigenvectors of L_sym ==> svd(L_sym)
            if minimum(sum(W[indrs, indrs], dims = 1)) > 10^3 * eps()
                ### eigenvectors of L_sym ==> svd(L_sym)
                D_temp_p = Diagonal(vec(sum(W_temp, dims = 1)).^(-0.5))
                v_sym = svd(Matrix(D_temp_p * (D_temp - W_temp) * D_temp_p)).U
            else
                ### eigenvectors of L ==> svd(L)
                v_sym = deepcopy(v)
            end

            # standardize the eigenvector signs
            v = v[:,end:-1:1]
            standardize_eigenvector_signs!(v)
            v_sym = v_sym[:,end:-1:1]
            standardize_eigenvector_signs!(v_sym)

            # construct unfolding operator custom to current region
            P = Uf[indrs, :]'
            # obtain the expansion coefficients
            dmatrixlH[indr, j, :] = (P * v)' * G.f
            dmatrixlHsym[indr, j, :] = (P * v_sym)' * G.f
        end
    end

    return dmatrixlH, dmatrixlHsym

end

function standardize_eigenvector_signs!(v)
    # standardize the eigenvector signs for HGLET (different with NGWPs)
    for col = 1:size(v, 2)
        row = 1
        standardized = false
        while !standardized
            if v[row, col] > 10^3 * eps()
                standardized = true
            elseif v[row,col] < -10^3 * eps()
                v[:, col] = -v[:, col]
            else
                row += 1
            end
        end
    end
end

"""
    HGLET_dictionary(GP::GraphPart, G::GraphSig; gltype::Symbol = :L)

assemble the whole HGLET dictionary

### Input Arguments
* `GP`: a GraphPart object
* `G`:  a GraphSig object
* `gltype`: `:L` or `:Lsym`

### Output Argument
* `dictionary`: the HGLET dictionary

"""
function HGLET_dictionary(GP::GraphPart, G::GraphSig; gltype::Symbol = :L)
    N = size(G.W, 1)
    jmax = size(GP.rs, 2)
    dictionary = zeros(N, jmax, N)
    for j = 1:jmax
        BS = BasisSpec(collect(enumerate(j * ones(Int, N))))
        dictionary[:, j, :] = HGLET_Synthesis(Matrix{Float64}(I, N, N), GP, BS, G; gltype = gltype)[1]'
    end
    return dictionary
end

"""
    LPHGLET_dictionary(GP::GraphPart, G::GraphSig; gltype::Symbol = :L, ϵ::Float64 = 0.3)

assemble the whole LP-HGLET dictionary

### Input Arguments
* `GP`: a GraphPart object
* `G`:  a GraphSig object
* `gltype`: `:L` or `:Lsym`
* `ϵ`: relative action bandwidth (default: 0.3)

### Output Argument
* `dictionary`: the LP-HGLET dictionary

"""
function LPHGLET_dictionary(GP::GraphPart, G::GraphSig; gltype::Symbol = :L, ϵ::Float64 = 0.3)
    N = size(G.W, 1)
    jmax = size(GP.rs, 2)
    dictionary = zeros(N, jmax, N)
    for j = 1:jmax
        BS = BasisSpec(collect(enumerate(j * ones(Int, N))))
        dictionary[:, j, :] = LPHGLET_Synthesis(Matrix{Float64}(I, N, N), GP, BS, G; gltype = gltype, ϵ = ϵ)[1]'
    end
    return dictionary
end


#   Unit "csp.jl" of the Diagonalization.jl Package for Julia language
#
#   MIT License
#   Copyright (c) 2019, 2020
#   Marco Congedo, CNRS, Grenoble, France:
#   https://sites.google.com/site/marcocongedo/home

# ? CONTENTS :
#   This unit implements Common Spatial Pattern (CSP) filters,
#   that is, filters based on the generalized eigendecomposition

"""
```
(1)
function csp(C₁ :: SorH, C₂ :: SorH;
             eVar     :: TeVaro = ○,
             eVarC    :: TeVaro = ○,
             eVarMeth :: Function = searchsortedfirst,
             selMeth  :: Symbol = :extremal,
             simple   :: Bool = false)

(2)
function csp(X₁ :: Mat, X₂ :: Mat;
             covEst   :: StatsBase.CovarianceEstimator = SCM,
             dims     :: Into = ○,
             meanX₁   :: Tmean = 0,
             meanX₂   :: Tmean = 0,
             wX₁      :: Tw = ○,
             wX₂      :: Tw = ○,
          eVar     :: TeVaro = ○,
          eVarC    :: TeVaro = ○,
          eVarMeth :: Function = searchsortedfirst,
          selMeth  :: Symbol = :extremal,
          simple   :: Bool = false)

(3)
function csp(𝐗₁::VecMat, 𝐗₂::VecMat;
             covEst   :: StatsBase.CovarianceEstimator = SCM,
             dims     :: Into = ○,
             meanX₁   :: Into = 0,
             meanX₂   :: Into = 0,
          eVar     :: TeVaro = ○,
          eVarC    :: TeVaro = ○,
          eVarMeth :: Function = searchsortedfirst,
          selMeth  :: Symbol = :extremal,
          simple   :: Bool = false,
       metric   :: Metric = Euclidean,
       w₁       :: Vector = [],
       w₂       :: Vector = [],
       ✓w       :: Bool = true,
       init₁    :: SorHo = nothing,
       init₂    :: SorHo = nothing,
       tol      :: Real = 0.,
       verbose  :: Bool = false)
```

Return a [LinearFilter](@ref) object:

**(1) Common spatial pattern**
with covariance matrices `C_1` and `C_2` of dimension
``n⋅n`` as input. The subscript of the covariance matrices refers to the `dims`
used to compute it (see above).

`eVar`, `eVarC` and `eVarMeth` are keyword optional arguments
for defining the [subspace dimension](@ref) ``p``. Particularly:
-  By default, the two-step procedure described above is used to find the
   solution. In this case `eVarC` is used for defining the subspace dimension of
   the whitening step. If `eVarC=0.0` is passed (not to be confused with
   `eVarC=0` ), the solution will be find by the generalized
   eigenvalue-eigenvector procedure.
- `eVar` is the keyword optional argument for defining the
   [subspace dimension](@ref) ``p`` using the `.arev` vector
   given by [csp.5].
- `eVarMeth` applies to both `eVarC` and `eVar`. The default value is
   `evarMeth=searchsortedfirst`.

if `selMeth=:extremal` (default) use case **a) Separating two classes**
described above is considered. Any other symbol for `selMeth` will instruct
to consider instead the use case **b) Enhance the signal-to-noise ratio**.

If `simple` is set to `true`, ``p`` is set equal to ``n``
and only the fields `.F` and `.iF` are written in the constructed object.
This option is provided for low-level work when you don't need to define
a subspace dimension or you want to define it by your own methods.

**(2) Common spatial pattern**
with data matrices `X₁` and `X₂` as input.

`X₁` and `X₂` are real or complex data matrices.

`covEst`, `dims`, `meanX₁`, `meanX₂`,  `wX₁` and `wX₂` are optional
keyword arguments to regulate the estimation of the
covariance matrices ``(C_1, C_2)`` of (`X₁`, `X₂`).
Particularly (See [covariance matrix estimations](@ref)),
- `meanX₁` is the `mean` argument for data matrix `X₁`.
- `meanX₂` is the `mean` argument for data matrix `X₂`.
- `wX₁` is the `w` argument for estimating a weighted covariance matrix for `X₁`.
- `wX₂` is the `w` argument for estimating a weighted covariance matrix for `X₂`.
- `covEst` applies to the estimations of both covariance matrices.

Once the two covariance matrices ``C_1`` and ``C_2`` estimated,
method (1) is invoked with optional keyword arguments
`eVar`, `eVarC`, `eVarMeth`, `selMeth` and `simple`.
See method (1) for details.

**(3) Common spatial pattern**
with two vectors of data matrices
`𝐗₁` and `𝐗₂` as input.

`𝐗₁` and `𝐗₂` do not need to hold the same number
of matrices and the number of samples in the matrices they contain
is arbitrary.

`covEst`, `dims`, `meanX₁` and `meanX₂` are optional
keyword arguments to regulate the estimation of the
covariance matrices for all matrices in `𝐗₁` and `𝐗₂`.
See method (2) and [covariance matrix estimations](@ref).

A mean covariance matrix is computed separatedly from the covariance matrices
computed from the data matrices in `𝐗₁` and `𝐗₂`,
using optional keywords arguments
`metric`, `w₁`, `w₂`, `✓w`,
`init₁`, `init₂`, `tol` and `verbose`. Particularly
(see [mean covariance matrix estimations](@ref)),
- `w₁` are the weights for the covariance matrices computed from `𝐗₁`,
- `w₂` are the weights for the covariance matrices computed from `𝐗₂`,
- `init₁` is the initialization for the mean of the covariance matrices computed from `𝐗₁`,
- `init₂` is the initialization for the mean of the covariance matrices computed from `𝐗₂`.
By default, the arithmetic mean is computed.

**See also:** [CSTP](@ref), [PCA](@ref), [AJD](@ref), [mAJD](@ref).

**Examples:**

```
using Diagonalizations, LinearAlgebra, PosDefManifold, Test

# Method (1) real
t, n=50, 10
X1=genDataMatrix(n, t)
X2=genDataMatrix(n, t)
Cx1=Symmetric((X1*X1')/t)
Cx2=Symmetric((X2*X2')/t)
C=Cx1+Cx2
cC=csp(Cx1, Cx2; simple=true)
Dx1=cC.F'*Cx1*cC.F
@test norm(Dx1-Diagonal(Dx1))+1≈1.
Dx2=cC.F'*Cx2*cC.F
@test norm(Dx2-Diagonal(Dx2))+1≈1.
@test cC.F'*C*cC.F≈I
@test norm(Dx1-(I-Dx2))+1≈1.

# Method (1) complex
t, n=50, 10
X1c=genDataMatrix(ComplexF64, n, t)
X2c=genDataMatrix(ComplexF64, n, t)
Cx1c=Hermitian((X1c*X1c')/t)
Cx2c=Hermitian((X2c*X2c')/t)
Cc=Cx1c+Cx2c
cCc=csp(Cx1c, Cx2c; simple=true)
Dx1c=cCc.F'*Cx1c*cCc.F
@test norm(Dx1c-Diagonal(Dx1c))+1. ≈ 1.
Dx2c=cCc.F'*Cx2c*cCc.F
@test norm(Dx2c-Diagonal(Dx2c))+1. ≈ 1.
@test cCc.F'*Cc*cCc.F≈I
@test norm(Dx1c-(I-Dx2c))+1. ≈ 1.

# Method (2) real
c12=csp(X1, X2, simple=true)
Dx1=c12.F'*Cx1*c12.F
@test norm(Dx1-Diagonal(Dx1))+1≈1.
Dx2=c12.F'*Cx2*c12.F
@test norm(Dx2-Diagonal(Dx2))+1≈1.
@test c12.F'*C*c12.F≈I
@test norm(Dx1-(I-Dx2))+1≈1.
@test cC==c12

# Method (2) complex
c12c=csp(X1c, X2c, simple=true)
Dx1c=c12c.F'*Cx1c*c12c.F
@test norm(Dx1c-Diagonal(Dx1c))+1. ≈ 1.
Dx2c=c12c.F'*Cx2c*c12c.F
@test norm(Dx2c-Diagonal(Dx2c))+1. ≈ 1.
@test c12c.F'*Cc*c12c.F≈I
@test norm(Dx1c-(I-Dx2c))+1. ≈ 1.
@test cCc==c12c

# Method (3) real
# CSP of the average covariance matrices
k=10
Xset=[genDataMatrix(n, t) for i=1:k]
Yset=[genDataMatrix(n, t) for i=1:k]

c=csp(Xset, Yset)

# ... selecting subspace dimension allowing an explained variance = 0.9
c=csp(Xset, Yset; eVar=0.9)

# ... subtracting the mean from the matrices in Xset and Yset
c=csp(Xset, Yset; meanX₁=nothing, meanX₂=nothing, eVar=0.9)

# csp on the average of the covariance and cross-covariance matrices
# computed along dims 1
c=csp(Xset, Yset; dims=1, eVar=0.9)

# name of the filter
c.name

using Plots
# plot regularized accumulated eigenvalues
plot(c.arev)

# plot the original covariance matrices and the transformed counterpart
# example when argument `selMeth` is `extremal` (default): 2-class separation
 cC=csp(Cx1, Cx2)
 Cx1Max=maximum(abs.(Cx1));
 h1 = heatmap(Cx1, clim=(-Cx1Max, Cx1Max), title="Cx1", yflip=true, c=:bluesreds);
 h2 = heatmap(cC.F'*Cx1*cC.F, clim=(0, 1), title="F'*Cx1*F", yflip=true, c=:amp);
 Cx2Max=maximum(abs.(Cx2));
 h3 = heatmap(Cx2, clim=(-Cx2Max, Cx2Max), title="Cx2", yflip=true, c=:bluesreds);
 h4 = heatmap(cC.F'*Cx2*cC.F, clim=(0, 1), title="F'*Cx2*F", yflip=true, c=:amp);
 CMax=maximum(abs.(C));
 h5 = heatmap(C, clim=(-CMax, CMax), title="Cx1+Cx2", yflip=true, c=:bluesreds);
 h6 = heatmap(cC.F'*C*cC.F, clim=(0, 1), title="F'*(Cx1+Cx2)*F", yflip=true, c=:amp);
 📈=plot(h1, h3, h5, h2, h4, h6, size=(800,400))
# savefig(📈, homedir()*"\\Documents\\Code\\julia\\Diagonalizations\\docs\\src\\assets\\FigCSP1.png")
```

 ![Figure CSP1](assets/FigCSP1.png)

```
# example when argument `selMeth` is different from `extremal`: enhance snr
 cC=csp(Cx1, Cx2; selMeth=:enhaceSNR)
 Cx1Max=maximum(abs.(Cx1));
 h1 = heatmap(Cx1, clim=(-Cx1Max, Cx1Max), title="Cx1", yflip=true, c=:bluesreds);
 h2 = heatmap(cC.F'*Cx1*cC.F, clim=(0, 1), title="F'*Cx1*F", yflip=true, c=:amp);
 Cx2Max=maximum(abs.(Cx2));
 h3 = heatmap(Cx2, clim=(-Cx2Max, Cx2Max), title="Cx2", yflip=true, c=:bluesreds);
 h4 = heatmap(cC.F'*Cx2*cC.F, clim=(0, 1), title="F'*Cx2*F", yflip=true, c=:amp);
 CMax=maximum(abs.(C));
 h5 = heatmap(C, clim=(-CMax, CMax), title="Cx1+Cx2", yflip=true, c=:bluesreds);
 h6 = heatmap(cC.F'*C*cC.F, clim=(0, 1), title="F'*(Cx1+Cx2)*F", yflip=true, c=:amp);
 📉=plot(h1, h3, h5, h2, h4, h6, size=(800,400))
# savefig(📉, homedir()*"\\Documents\\Code\\julia\\Diagonalizations\\docs\\src\\assets\\FigCSP2.png")

```

 ![Figure CSP2](assets/FigCSP2.png)

```
# Method (3) complex
# CSP of the average covariance matrices
k=10
Xsetc=[genDataMatrix(ComplexF64, n, t) for i=1:k]
Ysetc=[genDataMatrix(ComplexF64, n, t) for i=1:k]

cc=csp(Xsetc, Ysetc)

# ... selecting subspace dimension allowing an explained variance = 0.9
cc=csp(Xsetc, Ysetc; eVar=0.9)

# ... subtracting the mean from the matrices in Xset and Yset
cc=csp(Xsetc, Ysetc; meanX₁=nothing, meanX₂=nothing, eVar=0.9)

# csp on the average of the covariance and cross-covariance matrices
# computed along dims 1
cc=csp(Xsetc, Ysetc; dims=1, eVar=0.9)

# name of the filter
cc.name
```
"""
function csp(C₁ :: SorH, C₂ :: SorH;
             eVar     :: TeVaro = ○,
             eVarC    :: TeVaro = ○,
             eVarMeth :: Function = searchsortedfirst,
             selMeth  :: Symbol = :extremal,
             simple   :: Bool = false)

  # checks
  size(C₁, 1)==size(C₁, 2) || throw(ArgumentError(📌*", csp function: Matrix `C₁` must be square"))
  size(C₂, 1)==size(C₂, 2) || throw(ArgumentError(📌*", csp function: Matrix `C₂` must be square"))
  size(C₁)==size(C₂) || throw(ArgumentError(📌*", csp function: Matrices `C₁` and `C₂` must have the same size"))
  eVar isa Int && eVarC isa Int && eVar>eVarC && throw(ArgumentError(📌*", csp function: `eVar` cannot be larger than `eVarC`"))

  args=("Common Spatial Pattern", false)

  if eVarC≠○ && eVarC≈0. # use gevd, which also actually whitens C1+C2
     λ, U = eig(C₁, C₁+C₂)
     λ = _checkλ(λ) # make sure no imaginary noise is present (for complex data)

     simple ? LF(U, inv(U), Diagonal(λ), ○, ○, ○, args...) :
     begin
        eVar, D, U, p, arev=_ssdcsp!(eVar, λ, U, _minDim(C₁, C₂), eVarMeth, selMeth) # subspace dimension
        LF(U, pinv(U), D, eVar, λ, arev, args...)
     end
  else
     w=whitening(C₁+C₂; eVar=eVarC, eVarMeth=eVarMeth)

     # alert user if eVar passed as an integer exceeds the whitening dim
     whiteDim=size(w.F, 2)
     if eVar isa Int && eVar>whiteDim
        @warn(📌*", csp function: the whitening step reduced the rank to $(whiteDim); `eVar` has been lowered to this value.")
        eVar=whiteDim
     end

     λ, U = eig(Hermitian(w.F'*C₁*w.F)) # get evd of whitened C1
     # Hermitian is necessary for complex data
     λ = _checkλ(λ) # make sure no imaginary noise is present (for complex data)

     simple ? LF(w.F*U, U'*w.iF, Diagonal(λ), ○, ○, ○, args...) :
     begin
        eVar, D, U, p, arev=_ssdcsp!(eVar, λ, U, _minDim(C₁, C₂), eVarMeth, selMeth) # subspace dimension
        LF(w.F*U, U'*w.iF, D, eVar, λ, arev, args...)
     end
  end
end

function csp(X₁ :: Mat, X₂ :: Mat;
             covEst   :: StatsBase.CovarianceEstimator = SCM,
             dims     :: Into = ○,
             meanX₁   :: Tmean = 0,
             meanX₂   :: Tmean = 0,
             wX₁      :: Tw = ○,
             wX₂      :: Tw = ○,
          eVar     :: TeVaro = ○,
          eVarC    :: TeVaro = ○,
          eVarMeth :: Function = searchsortedfirst,
          selMeth  :: Symbol = :extremal,
          simple   :: Bool = false)

   dims===○ && (dims=_set_dims(X₁, X₂))
   _check_data(X₁, dims, covEst, meanX₁, wX₁)===○ && return
   _check_data(X₂, dims, covEst, meanX₂, wX₂)===○ && return

   C₁=_cov(X₁, covEst, dims, meanX₁, wX₁)
   C₂=_cov(X₂, covEst, dims, meanX₂, wX₂)

   csp(C₁, C₂; eVar=eVar, eVarC=eVarC, eVarMeth=eVarMeth,
       selMeth=selMeth, simple=simple)
end

function csp(𝐗₁::VecMat, 𝐗₂::VecMat;
             covEst   :: StatsBase.CovarianceEstimator = SCM,
             dims     :: Into = ○,
             meanX₁   :: Into = 0,
             meanX₂   :: Into = 0,
          eVar     :: TeVaro = ○,
          eVarC    :: TeVaro = ○,
          eVarMeth :: Function = searchsortedfirst,
          selMeth  :: Symbol = :extremal,
          simple   :: Bool = false,
       metric   :: Metric = Euclidean,
       w₁       :: Vector = [],
       w₂       :: Vector = [],
       ✓w       :: Bool = true,
       init₁    :: SorHo = nothing,
       init₂    :: SorHo = nothing,
       tol      :: Real = 0.,
       verbose  :: Bool = false)

   Metric==VonNeumann && throw(ArgumentError(📌*", csp function: A solution for the mean is not available for the Von Neumann metric. Use another metric as `metric` argument"))
   dims===○ && (dims=_set_dims(𝐗₁, 𝐗₂))
   _check_data(𝐗₁, dims, covEst, meanX₁, ○)===○ && return
   _check_data(𝐗₂, dims, covEst, meanX₂, ○)===○ && return

   𝐂₁= _cov(𝐗₁; covEst=covEst, dims=dims, meanX=meanX₁)
   𝐂₂= _cov(𝐗₂; covEst=covEst, dims=dims, meanX=meanX₂)

   csp(mean(metric, 𝐂₁;
            w = w₁, ✓w = ✓w,
            init = init₁===○ ? ○ : Hermitian(init₁), #just init₁ here when you upfate PodDefManifold
            tol = tol, verbose = verbose),
       mean(metric, 𝐂₂;
            w = w₂, ✓w = ✓w,
            init = init₂===○ ? ○ : Hermitian(init₂), #just init₂ here when you upfate PodDefManifold
            tol = tol, verbose = verbose);
       eVar=eVar, eVarC=eVarC, eVarMeth=eVarMeth,
       selMeth=selMeth, simple=simple)
end

"""
```
(1)
function cstp( X :: Mat, C₍₁₎ :: SorH, C₍₂₎ :: SorH;
               eVar     :: TeVaro = ○,
               eVarC    :: TeVaro = ○,
               eVarMeth :: Function = searchsortedfirst,
               simple   :: Bool = false)

(2)
function cstp( 𝐗::VecMat;
               covEst   :: StatsBase.CovarianceEstimator = SCM,
               meanXd₁  :: Into = 0,
               meanXd₂  :: Into = 0,
            eVar     :: TeVaro = ○,
            eVarC    :: TeVaro = ○,
            eVarMeth :: Function = searchsortedfirst,
            simple   :: Bool = false,
         metric   :: Metric = Euclidean,
         w        :: Vector = [],
         ✓w       :: Bool = true,
         init1    :: SorHo = nothing,
         init2    :: SorHo = nothing,
         tol      :: Real = 0.,
         verbose  :: Bool = false)
```

Return a [LinearFilter](@ref) object:

(1)
**Common spatio-temporal pattern**
with ``n⋅m`` mean data matrix `X`,
``m⋅m`` covariance matrices `C₍₁₎` and ``n⋅n`` covariance matrix `C₍₂₎` as input.

`eVar`, `eVarC` and `eVarMeth` are keyword optional arguments
for defining the [subspace dimension](@ref) ``p``. Particularly:
-  `eVarC` is used for defining the subspace dimension of
   the whitening step. The default is 0.999.
- `eVar` is the keyword optional argument for defining the
   [subspace dimension](@ref) ``p`` using the `.arev` vector
   given by [cstp.5]. The default is given in [cstp.6] here above.
- `eVarMeth` applies to both `eVarC` and `eVar`. The default value is
   `evarMeth=searchsortedfirst`.

If `simple` is set to `true`, ``p`` is set equal to ``n``
and only the fields `.F` and `.iF` are written in the constructed object.
This option is provided for low-level work when you don't need to define
a subspace dimension or you want to define it by your own methods.

(2)
**Common spatio-temporal pattern**
with a set of ``k`` data matrices `𝐗` as input.

The ``k`` matrices in `𝐗` are real or complex data matrices.
They must all have the same size.

`covEst`, `meanXd₁` and `meanXd₂` are optional
keyword arguments to regulate the estimation of the
covariance matrices of the data matrices in `𝐗`,
to be used to compute the mean covariance matrices in [cstp.2] here above.
See [covariance matrix estimations](@ref).
`meanXd₁` and `meanXd₂` are the means along dimension 1 and 2, respectively,
of the data matrices in `𝐗`.

The mean covariance matrices ``C_{(1)}`` and ``C_{(1)}`` in [cstp.2]
are computed using optional keywords arguments
`metric`, `w`, `✓w`, `init1`, `init2`, `tol` and `verbose`,
which allow to compute non-Euclidean means.
Particularly (see [mean covariance matrix estimations](@ref)),
- `init1` is the initialization for ``C_{(1)}``,
- `init2` is the initialization for ``C_{(2)}``.
By default, the arithmetic means [cstp.2] are computed.

Once the two covariance matrices ``C_{(1)}`` and ``C_{(2)}`` estimated,
method (1) is invoked with optional keyword arguments
`eVar`, `eVarC`, `eVarMeth` and `simple`.
See method (1) for details.

**Examples:**

```
using Diagonalizations, LinearAlgebra, PosDefManifold, Test

# Method (1) real
t, n, k=10, 20, 10
Xset = [genDataMatrix(t, n) for i = 1:k]
Xfixed=randn(t, n)./1
for i=1:length(Xset) Xset[i]+=Xfixed end
C1=Hermitian( mean((X'*X)/t for X∈Xset) )
C2=Hermitian( mean((X*X')/n for X∈Xset) )
Xbar=mean(Xset)
c=cstp(Xbar, C1, C2; simple=true)
@test c.F[1]'*C2*c.F[1]≈I
@test c.F[2]'*C1*c.F[2]≈I
Z=c.F[1]'*Xbar*c.F[2]
n=minimum(size(Z))
@test norm(Z[1:n, 1:n]-Diagonal(Z[1:n, 1:n]))+1. ≈ 1.
cX=cstp(Xset; simple=true)
@test c==cX

# Method (1) complex
t, n, k=10, 20, 10
Xcset = [genDataMatrix(ComplexF64, t, n) for i = 1:k]
Xcfixed=randn(ComplexF64, t, n)./1
for i=1:length(Xcset) Xcset[i]+=Xcfixed end
C1c=Hermitian( mean((Xc'*Xc)/t for Xc∈Xcset) )
C2c=Hermitian( mean((Xc*Xc')/n for Xc∈Xcset) )
Xcbar=mean(Xcset)
cc=cstp(Xcbar, C1c, C2c; simple=true)
@test cc.F[1]'*C2c*cc.F[1]≈I
@test cc.F[2]'*C1c*cc.F[2]≈I
Zc=cc.F[1]'*Xcbar*cc.F[2]
n=minimum(size(Zc))
@test norm(Zc[1:n, 1:n]-Diagonal(Zc[1:n, 1:n]))+1. ≈ 1.
cXc=cstp(Xcset; simple=true)
@test cc==cXc

# Method (2) real
c=cstp(Xset)

# ... selecting subspace dimension allowing an explained variance = 0.9
c=cstp(Xset; eVar=0.9)

# ... giving weights `w` to the covariance matrices
c=cstp(Xset; w=abs2.(randn(k)), eVar=0.9)

# ... subtracting the means
c=cstp(Xset; meanXd₁=nothing, meanXd₂=nothing, w=abs2.(randn(k)), eVar=0.9)

# explained variance
c.eVar

# name of the filter
c.name

using Plots
# plot the original covariance matrices and the transformed counterpart
c=cstp(Xset)

C1Max=maximum(abs.(C1));
 h1 = heatmap(C1, clim=(-C1Max, C1Max), title="C1", yflip=true, c=:bluesreds);
 D1=c.F[1]'*C2*c.F[1];
 D1Max=maximum(abs.(D1));
 h2 = heatmap(D1, clim=(0, D1Max), title="F[1]'*C2*F[1]", yflip=true, c=:amp);
 C2Max=maximum(abs.(C2));
 h3 = heatmap(C2, clim=(-C2Max, C2Max), title="C2", yflip=true, c=:bluesreds);
 D2=c.F[2]'*C1*c.F[2];
 D2Max=maximum(abs.(D2));
 h4 = heatmap(D2, clim=(0, D2Max), title="F[2]'*C1*F[2]", yflip=true, c=:amp);

XbarMax=maximum(abs.(Xbar));
 h5 = heatmap(Xbar, clim=(-XbarMax, XbarMax), title="Xbar", yflip=true, c=:bluesreds);
 DX=c.F[1]'*Xbar*c.F[2];
 DXMax=maximum(abs.(DX));
 h6 = heatmap(DX, clim=(0, DXMax), title="F[1]'*Xbar*F[2]", yflip=true, c=:amp);
 📈=plot(h1, h3, h5, h2, h4, h6, size=(800,400))
# savefig(📈, homedir()*"\\Documents\\Code\\julia\\Diagonalizations\\docs\\src\\assets\\FigCSTP.png")

```

 ![Figure CSTP](assets/FigCSTP.png)

```

# Method (2) complex
cc=cstp(Xcset)

# ... selecting subspace dimension allowing an explained variance = 0.9
cc=cstp(Xcset; eVar=0.9)

# ... giving weights `w` to the covariance matrices
cc=cstp(Xcset; w=abs2.(randn(k)), eVar=0.9)

# ... subtracting the mean
cc=cstp(Xcset; meanXd₁=nothing, meanXd₂=nothing,
        w=abs2.(randn(k)), eVar=0.9)

# explained variance
c.eVar

# name of the filter
c.name

```

"""
function cstp( X :: Mat, C₍₁₎ :: SorH, C₍₂₎ :: SorH;
               eVar     :: TeVaro = ○,
               eVarC    :: TeVaro = ○,
               eVarMeth :: Function = searchsortedfirst,
               simple   :: Bool = false)

   d₍₂₎, d₍₁₎, dₓ=size(C₍₂₎, 1), size(C₍₁₎, 1), size(X)
   (d₍₁₎==dₓ[2] && d₍₂₎==dₓ[1]) || throw(ArgumentError(📌*", cstp function: For n⋅m matrix X, matrix C₍₁₎ must be m⋅m and matrix C₍₂₎ n⋅n"))
   eVar isa Int && eVarC isa Int && eVar>eVarC && throw(ArgumentError(📌*", cstp function: `eVar` cannot be larger than `eVarC`"))
   args=("Common Spatio-Temporal Pattern", false)

   kwargs=(eVar=eVarC, eVarMeth=eVarMeth, simple=false)
   t=whitening(C₍₁₎; kwargs...)
   s=whitening(C₍₂₎; kwargs...)

   # alert user if eVar passed as an integer exceeds the minimum whitening dim
   whiteDim=min(size(t.F, 2), size(t.F, 2))
   if eVar isa Int && eVar>whiteDim
      @warn(📌*", cstp function: the whitening step reduced the rank to $(whiteDim); `eVar` has been lowered to this value.")
      eVar=whiteDim
   end

   U, λ, V = svd(s.F'*X*t.F; full=true)
   λ = _checkλ(λ) # make sure no imaginary noise is present (for complex data)

   simple ? LF([s.F*U, t.F*V], [U'*s.iF, V'*t.iF], Diagonal(λ), ○, ○, ○, args...) :
   begin
     eVar, D, U, V, p, arev=_ssdcstp!(eVar, λ, U, Matrix(V), _minDim(X), eVarMeth) # subspace dimension
     LF([s.F*U, t.F*V], [U'*s.iF, V'*t.iF], D, eVar, λ, arev, args...)
   end
end

# no dims argument since it is the same for cstp
function cstp( 𝐗::VecMat;
               covEst   :: StatsBase.CovarianceEstimator = SCM,
               meanXd₁  :: Into = 0,
               meanXd₂  :: Into = 0,
            eVar     :: TeVaro = ○,
            eVarC    :: TeVaro = ○,
            eVarMeth :: Function = searchsortedfirst,
            simple   :: Bool = false,
         metric   :: Metric = Euclidean,
         w        :: Vector = [],
         ✓w       :: Bool = true,
         init1     :: SorHo = nothing,
         init2     :: SorHo = nothing,
         tol      :: Real = 0.,
         verbose  :: Bool = false)

   Metric==VonNeumann && throw(ArgumentError(📌*", cstp function: A solution for the mean is not available for the Von Neumann metric. Use another metric as `metric` argument"))
   covEst==SCM && metric ∉ (Euclidean, Wasserstein) && throw(ArgumentError(📌*", cstp function: Only the Euclidean and Wasserstein `metric` can be used if the covariance estimator is `SCM`"))
   _check_data(𝐗, 1, covEst, meanXd₁, ○)===○ && return
   _check_data(𝐗, 2, covEst, meanXd₂, ○)===○ && return

   𝐂₍₁₎=_cov(𝐗; covEst=covEst, dims = 1, meanX = meanXd₁)
   𝐂₍₂₎=_cov(𝐗; covEst=covEst, dims = 2, meanX = meanXd₂)

   cstp(PosDefManifold.fVec(mean, Vector{Matrix}(𝐗)), # multi-threaded Euclidean mean
        mean(metric, 𝐂₍₁₎;
             w = w, ✓w = ✓w,
             init = init1===○ ? ○ : Hermitian(init1), #just init here when you update PodDefManifold
             tol = tol, verbose = verbose),
        mean(metric, 𝐂₍₂₎;
             w = w, ✓w = ✓w,
             init = init2===○ ? ○ : Hermitian(init2), #just init here when you update PodDefManifold.jl
             tol = tol, verbose = verbose),
        eVar     = eVar,
        eVarC    = eVarC,
        eVarMeth = eVarMeth,
        simple   = simple)
end


include("../mpiops.jl")

mpi_init()

function kappa(x, y)
    return 1+x^2+y^2
end

function rhs(x, y)
    -2*x*(1 - x)*(x^2 + y^2 + 1) + 2*x*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) -
             2*y*(1 - y)*(x^2 + y^2 + 1) + 2*y*(-x*y*(1 - x) + x*(1 - x)*(1 - y))
end

function ufunc(x, y)
    x * (1-x) * y * (1-y)
end

mc = MPIConfig(50)
global_to_local, local_to_global = dofmap(mc)
X, Y = get_xy(mc)
f_local = rhs.(X, Y)
using Random; Random.seed!(233)
θ = Variable(fc_init([2,20,20,20,1]))
θ = mpi_bcast(θ)
κ_local = (fc([X'[:] Y'[:]], [20,20,20,1], θ) + 5.0)|>squeeze
κ_local = reshape(κ_local, (mc.n, mc.n))
u_local = poisson_solver(κ_local, f_local, mc)
u = mpi_gather(u_local)[global_to_local]

Uexact = reshape(constant(ufunc.(X, Y)), (-1,))
Uexact = mpi_gather(Uexact)[global_to_local]

sess = Session(); init(sess)

U = run(sess, u)
Uexact = run(sess, Uexact)

if mpi_rank()==0
    close("all")
    figure(figsize = (15,4))
    subplot(131)
    pcolormesh(reshape(U, mc.n * mc.N, mc.n * mc.N))
    colorbar()
    subplot(132)
    pcolormesh(reshape(Uexact, mc.n * mc.N, mc.n * mc.N))
    colorbar()
    subplot(133)
    pcolormesh(reshape(abs.(U - Uexact), mc.n * mc.N, mc.n * mc.N))
    colorbar()
    savefig("poisson_test_$(mpi_size())_$(mc.n).png")
end

if mpi_size()>1
    mpi_finalize()
end


using ModelingToolkit: operation, istree, arguments

# Method of lines discretization scheme

@enum GridAlign center_align edge_align
struct MOLFiniteDifference{T,T2} <: DiffEqBase.AbstractDiscretization
    dxs::T
    time::T2
    upwind_order::Int
    centered_order::Int
    grid_align::GridAlign
end

# Constructors. If no order is specified, both upwind and centered differences will be 2nd order
MOLFiniteDifference(dxs, time; upwind_order = 1, centered_order = 2, grid_align=center_align) =
    MOLFiniteDifference(dxs, time, upwind_order, centered_order, grid_align)

function calculate_weights_cartesian(order::Int, x0::T, xs::AbstractVector, idxs::AbstractVector) where T<:Real
        # Cartesian domain: use Fornberg
        DiffEqOperators.calculate_weights(order, x0, vec(xs[idxs]))
 end
 function calculate_weights_spherical(order::Int, x0::T, x::AbstractVector, idxs::AbstractVector) where T<:Real
        # Spherical domain: see #367
        # https://web.mit.edu/braatzgroup/analysis_of_finite_difference_discretization_schemes_for_diffusion_in_spheres_with_variable_diffusivity.pdf
        # Only order 2 is implemented
        @assert order == 2
        # Only 2nd order discretization is implemented
        # We can't activate this assertion for now because the rules try to create the spherical Laplacian
        # before checking whether there is a spherical Laplacian
        # this could be fixed by dispatching on domain type when we have different domain types
        # but for now everything is an IntervalDomain
        # @assert length(x) == 3
        # TODO: nonlinear diffusion in a spherical domain
        i = idxs[2]
        dx1 = x[i] - x[i-1]
        dx2 = x[i+1] - x[i]
        i0 = i - 1 # indexing starts at 0 in the paper and starts at 1 in julia
        1 / (i0 * dx1 * dx2) * [i0-1, -2i0, i0+1]
 end

function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discretization::DiffEqOperators.MOLFiniteDifference)
    grid_align = discretization.grid_align
    pdeeqs = pdesys.eqs isa Vector ? pdesys.eqs : [pdesys.eqs]
    t = discretization.time
    nottime = filter(x->~isequal(x.val, t.val),pdesys.indvars)
    nspace = length(nottime)
    depvars = pdesys.depvars

    # Discretize space
    space = map(nottime) do x
        xdomain = pdesys.domain[findfirst(d->isequal(x.val, d.variables),pdesys.domain)]
        dx = discretization.dxs[findfirst(dxs->isequal(x.val, dxs[1].val),discretization.dxs)][2]
        dx isa Number ? (xdomain.domain.lower:dx:xdomain.domain.upper) : dx
    end
    dxs = map(nottime) do x
        dx = discretization.dxs[findfirst(dxs->isequal(x.val, dxs[1].val),discretization.dxs)][2]
    end
    # Get tspan
    tdomain = pdesys.domain[findfirst(d->isequal(t.val, d.variables),pdesys.domain)]
    @assert tdomain.domain isa IntervalDomain
    tspan = (tdomain.domain.lower,tdomain.domain.upper)

    # Define the grid on which the dependent variables will be evaluated (see #378)
    # center_align is recommended for Dirichlet BCs
    # edge_align is recommended for Neumann BCs (spatial discretization is conservative)
    if grid_align == center_align
        grid = space
    elseif grid_align == edge_align
        # boundary conditions implementation assumes centered_order=2
        @assert discretization.centered_order==2
        # construct grid including ghost nodes beyond outer edges
        # e.g. space 0:dx:1 goes to grid -dx/2:dx:1+dx/2
        space_ext = map(s -> vcat(2s[1]-s[2],s,2s[end]-s[end-1]), space)
        grid = map(s -> (s[1:end-1]+s[2:end])/2, space_ext)
        # TODO: allow depvar-specific center/edge choice?
    end

    # Build symbolic variables
    space_indices = CartesianIndices(((axes(s)[1] for s in space)...,))
    grid_indices = CartesianIndices(((axes(g)[1] for g in grid)...,))
    depvarsdisc = map(depvars) do u
        [Num(Variable{Symbolics.FnType{Tuple{Any}, Real}}(Base.nameof(ModelingToolkit.operation(u.val)),II.I...))(t) for II in grid_indices]
    end
    spacevals = map(y->[Pair(nottime[i],space[i][y.I[i]]) for i in 1:nspace],space_indices)
    gridvals = map(y->[Pair(nottime[i],grid[i][y.I[i]]) for i in 1:nspace],grid_indices)

    ### INITIAL AND BOUNDARY CONDITIONS ###
    # Build symbolic maps for boundaries
    edges = reduce(vcat,[[vcat([Colon() for j in 1:i-1],1,[Colon() for j in i+1:nspace]),
                          vcat([Colon() for j in 1:i-1],length(space[i]),[Colon() for j in i+1:nspace])] for i in 1:nspace])

    #edgeindices = [indices[e...] for e in edges]
    get_edgevals(i) = [nottime[i]=>first(space[i]),nottime[i]=>last(space[i])]
    edgevals = reduce(vcat,[get_edgevals(i) for i in 1:length(space)])
    edgevars = [[d[e...] for e in edges] for d in depvarsdisc]

    bclocs = map(e->substitute.(pdesys.indvars,e),edgevals) # location of the boundary conditions e.g. (t,0.0,y)
    edgemaps = Dict(bclocs .=> [spacevals[e...] for e in edges])
    initmaps = substitute.(depvars,[t=>tspan[1]])

    # Generate map from variable (e.g. u(t,0)) to discretized variable (e.g. u₁(t))
    subvar(depvar) = substitute.((depvar,),edgevals)
    if grid_align == center_align
        # depvarbcmaps will dictate what to replace the variable terms with in the bcs
        # replace u(t,0) with u₁, etc
        depvarbcmaps = reduce(vcat,[subvar(depvar) .=> edgevar for (depvar, edgevar) in zip(depvars, edgevars)])
    end
    # depvarderivbcmaps will dictate what to replace the Differential terms with in the bcs
    if nspace == 1
        # 1D system
        left_weights(j) = DiffEqOperators.calculate_weights(discretization.upwind_order, grid[j][1], grid[j][1:2])
        right_weights(j) = DiffEqOperators.calculate_weights(discretization.upwind_order, grid[j][end], grid[j][end-1:end])
        central_neighbor_idxs(II,j) = [II-CartesianIndex((1:nspace.==j)...),II,II+CartesianIndex((1:nspace.==j)...)]
        left_idxs = central_neighbor_idxs(CartesianIndex(2),1)[1:2]
        right_idxs(j) = central_neighbor_idxs(CartesianIndex(length(grid[j])-1),1)[end-1:end]
        # Constructs symbolic spatially discretized terms of the form e.g. au₂ - bu₁
        derivars = [[dot(left_weights(1),depvar[left_idxs]), dot(right_weights(1),depvar[right_idxs(1)])]
        for depvar in depvarsdisc]
        # Create list of all the symbolic Differential terms evaluated at boundary e.g. Differential(x)(u(t,0))
        subderivar(depvar,s) = substitute.((Differential(s)(depvar),),edgevals)
        # Create map of symbolic Differential terms with symbolic spatially discretized terms
        depvarderivbcmaps = reduce(vcat,[subderivar(depvar, s) .=> derivars[i]
                                         for (i, depvar) in enumerate(depvars) for s in nottime])

        if grid_align == edge_align
            # Constructs symbolic spatially discretized terms of the form e.g. (u₁ + u₂) / 2
            bcvars = [[dot(ones(2)/2,depvar[left_idxs]), dot(ones(2)/2,depvar[right_idxs(1)])]
                      for depvar in depvarsdisc]
            # replace u(t,0) with (u₁ + u₂) / 2, etc
            depvarbcmaps = reduce(vcat,[subvar(depvar) .=> bcvars[i]
                                  for (i, depvar) in enumerate(depvars) for s in nottime])

        end
    else
        # Higher dimension
        # TODO: Fix Neumann and Robin on higher dimension
        depvarderivbcmaps = []
    end

    # Generate initial conditions and bc equations
    u0 = []
    bceqs = []
    for bc in pdesys.bcs
        if ModelingToolkit.operation(bc.lhs) isa Sym && ~any(x -> isequal(x, t.val), ModelingToolkit.arguments(bc.lhs))
            # initial condition
            # Assume in the form `u(...) ~ ...` for now
            push!(u0,vec(depvarsdisc[findfirst(isequal(bc.lhs),initmaps)] .=> substitute.((bc.rhs,),gridvals)))
        else
            # Algebraic equations for BCs
            i = findfirst(x->occursin(x.val,bc.lhs),first.(depvarbcmaps))
            bcargs = first(depvarbcmaps[i]).val.arguments
            # Replace Differential terms in the bc lhs with the symbolic spatially discretized terms
            # TODO: Fix Neumann and Robin on higher dimension
            lhs = nspace == 1 ? substitute(bc.lhs,depvarderivbcmaps[i]) : bc.lhs

            # Replace symbol in the bc lhs with the spatial discretized term
            lhs = substitute(lhs,depvarbcmaps[i])
            rhs = substitute.((bc.rhs,),edgemaps[bcargs])
            lhs = lhs isa Vector ? lhs : [lhs] # handle 1D
            push!(bceqs,lhs .~ rhs)
        end
    end

    u0 = reduce(vcat,u0)
    bceqs = reduce(vcat,bceqs)

    #---- Count Boundary Equations --------------------
    # Count the number of boundary equations that lie at the spatial boundary on
    # both the left and right side. This will be used to determine number of
    # interior equations s.t. we have a balanced system of equations.

    # get the depvar boundary terms for given depvar and indvar index.
    get_depvarbcs(depvar, i) = substitute.((depvar,),get_edgevals(i))

    # return the counts of the boundary-conditions that reference the "left" and
    # "right" edges of the given independent variable. Note that we return the
    # max of the count for each depvar in the system of equations.
    get_bc_counts(i) =
        begin
            left = 0
            right = 0
            for depvar in depvars
                depvaredges = get_depvarbcs(depvar, i)
                counts = [map(x->occursin(x.val, bc.lhs), depvaredges) for bc in pdesys.bcs]
                left = max(left, sum([c[1] for c in counts]))
                right = max(right, sum([c[2] for c in counts]))
            end
            return [left, right]
        end
    #--------------------------------------------------

    ### PDE EQUATIONS ###
    # Create a stencil in the required dimension centered around 0
    # e.g. (-1,0,1) for 2nd order, (-2,-1,0,1,2) for 4th order, etc
    if discretization.centered_order % 2 != 0
        throw(ArgumentError("Discretization centered_order must be even, given $(discretization.centered_order)"))
    end
    approx_order = discretization.centered_order
    stencil(j, order) = CartesianIndices(Tuple(map(x -> -x:x, (1:nspace.==j) * (order÷2))))

    # TODO: Generalize central difference handling to allow higher even order derivatives
    # The central neighbour indices should add the stencil to II, unless II is too close
    # to an edge in which case we need to shift away from the edge

    # Calculate buffers
    I1 = oneunit(first(grid_indices))
    Imin(order) = first(grid_indices) + I1 * (order÷2)
    Imax(order) = last(grid_indices) - I1 * (order÷2)

    interior = grid_indices[[let bcs = get_bc_counts(i)
                             (1 + first(bcs)):length(g)-last(bcs)
                             end
                             for (i,g) in enumerate(grid)]...]
    eqs = vec(map(Base.product(interior,pdeeqs)) do p
        II,eq = p

        # Use max and min to apply buffers
        central_neighbor_idxs(II,j,order) = stencil(j,order) .+ max(Imin(order),min(II,Imax(order)))
        central_weights_cartesian(d_order,II,j) = calculate_weights_cartesian(d_order, grid[j][II[j]], grid[j], vec(map(i->i[j],
                                                                              central_neighbor_idxs(II,j,approx_order))))
        central_deriv(d_order, II,j,k) = dot(central_weights(d_order, II,j),depvarsdisc[k][central_neighbor_idxs(II,j,approx_order)])

        central_deriv_cartesian(d_order,II,j,k) = dot(central_weights_cartesian(d_order,II,j),depvarsdisc[k][central_neighbor_idxs(II,j,approx_order)])

        # spherical Laplacian has a hardcoded order of 2 (only 2nd order is implemented)
        # both for derivative order and discretization order
        central_weights_spherical(II,j) = calculate_weights_spherical(2, grid[j][II[j]], grid[j], vec(map(i->i[j], central_neighbor_idxs(II,j,2))))
        central_deriv_spherical(II,j,k) = dot(central_weights_spherical(II,j),depvarsdisc[k][central_neighbor_idxs(II,j,2)])

        # get a sorted list derivative order such that highest order is first. This is useful when substituting rules
        # starting from highest to lowest order.
        d_orders(s) = reverse(sort(collect(union(differential_order(eq.rhs, s.val), differential_order(eq.lhs, s.val)))))

        # central_deriv_rules = [(Differential(s)^2)(u) => central_deriv(2,II,j,k) for (j,s) in enumerate(nottime), (k,u) in enumerate(depvars)]
        central_deriv_rules_cartesian = Array{Pair{Num,Num},1}()
        for (j,s) in enumerate(nottime)
            rs = [(Differential(s)^d)(u) => central_deriv_cartesian(d,II,j,k) for d in d_orders(s), (k,u) in enumerate(depvars)]
            for r in rs
                push!(central_deriv_rules_cartesian, r)
            end
        end
        central_deriv_rules_spherical = [(Differential(s))(s^2*Differential(s)(u))/s^2 => central_deriv_spherical(II,j,k)
                                         for (j,s) in enumerate(nottime), (k,u) in enumerate(pdesys.depvars)]

        valrules = vcat([depvars[k] => depvarsdisc[k][II] for k in 1:length(depvars)],
                        [nottime[j] => grid[j][II[j]] for j in 1:nspace])

        # TODO: upwind rules needs interpolation into `@rule`
        forward_weights(II,j) = DiffEqOperators.calculate_weights(discretization.upwind_order, 0.0, grid[j][[II[j],II[j]+1]])
        reverse_weights(II,j) = DiffEqOperators.calculate_weights(discretization.upwind_order, 0.0, grid[j][[II[j]-1,II[j]]])
        # upwinding_rules = [@rule(*(~~a,$(Differential(nottime[j]))(u),~~b) => IfElse.ifelse(*(~~a..., ~~b...,)>0,
        #                         *(~~a..., ~~b..., dot(reverse_weights(II,j),depvars[k][central_neighbor_idxs(II,j)[1:2]])),
        #                         *(~~a..., ~~b..., dot(forward_weights(II,j),depvars[k][central_neighbor_idxs(II,j)[2:3]]))))
        #                         for j in 1:nspace, k in 1:length(pdesys.depvars)]

        ## Discretization of non-linear laplacian.
        # d/dx( a du/dx ) ~ (a(x+1/2) * (u[i+1] - u[i]) - a(x-1/2) * (u[i] - u[i-1]) / dx^2
        b1(II, j, k) = dot(reverse_weights(II, j), depvarsdisc[k][central_neighbor_idxs(II, j, approx_order)[1:2]]) / dxs[j]
        b2(II, j, k) = dot(forward_weights(II, j), depvarsdisc[k][central_neighbor_idxs(II, j, approx_order)[2:3]]) / dxs[j]
        # TODO: improve interpolation of g(x) = u(x) for calculating u(x+-dx/2)
        g(II, j, k, l) = sum([depvarsdisc[k][central_neighbor_idxs(II, j, approx_order)][s] for s in (l == 1 ? [2,3] : [1,2])]) / 2.
        # iv_mid returns middle space values. E.g. x(i-1/2) or y(i+1/2).
        iv_mid(II, j, l) = (grid[j][II[j]] + grid[j][II[j]+l]) / 2.0
        # Dependent variable rules
        r_mid_dep(II, j, k, l) = [depvars[k] => g(II, j, k, l) for k in 1:length(depvars)]
        # Independent variable rules
        r_mid_indep(II, j, l) = [nottime[j] => iv_mid(II, j, l) for j in 1:length(nottime)]
        # Replacement rules: new approach
        rules = [@rule ($(Differential(iv))(*(~~a, $(Differential(iv))(dv), ~~b))) =>
                 dot([Num(substitute(substitute(*(~~a..., ~~b...), r_mid_dep(II, j, k, -1)), r_mid_indep(II, j, -1))),
                      Num(substitute(substitute(*(~~a..., ~~b...), r_mid_dep(II, j, k, 1)), r_mid_indep(II, j, 1)))],
                     [-b1(II, j, k), b2(II, j, k)])
                 for (j, iv) in enumerate(nottime) for (k, dv) in enumerate(depvars)]
        rhs_arg = (SymbolicUtils.operation(eq.rhs) == +) ? SymbolicUtils.arguments(eq.rhs) : [eq.rhs]
        lhs_arg = (SymbolicUtils.operation(eq.lhs) == +) ? SymbolicUtils.arguments(eq.lhs) : [eq.lhs]
        nonlinlap_rules = []
        for t in vcat(lhs_arg,rhs_arg)
            for r in rules
                if r(t) !== nothing
                    push!(nonlinlap_rules, t => r(t))
                end
            end
        end

        rules = vcat(vec(nonlinlap_rules),
                     vec(central_deriv_rules_cartesian),
                     vec(central_deriv_rules_spherical),
                     valrules)

        substitute(eq.lhs,rules) ~ substitute(eq.rhs,rules)

    end)

    # Finalize
    defaults = pdesys.ps === nothing || pdesys.ps === SciMLBase.NullParameters() ? u0 : vcat(u0,pdesys.ps)
    ps = pdesys.ps === nothing || pdesys.ps === SciMLBase.NullParameters() ? Num[] : first.(pdesys.ps)
    # Combine PDE equations and BC equations
    sys = ODESystem(vcat(eqs,unique(bceqs)),t,vec(reduce(vcat,vec(depvarsdisc))),ps,defaults=Dict(defaults))
    sys, tspan
end

function SciMLBase.discretize(pdesys::ModelingToolkit.PDESystem,discretization::DiffEqOperators.MOLFiniteDifference)
    sys, tspan = SciMLBase.symbolic_discretize(pdesys,discretization)
    simpsys = structural_simplify(sys)
    prob = ODEProblem(simpsys,Pair[],tspan)
end


## Metsa-Oy Forest Products Supply Chain
## Edward J. Xu <edxu96@outlook.com>
## Feb 1st, 2020

function optim_1(alpha_t, beta_t, a_i, b_i, e_i, c_i, r_saw, r_plywood, a_j, c_j,
  r_j, a_paper_t, b_paper_j, c_paper, r_paper, gamma_ik, delta_ik, gamma_jk,
  delta_jk, gamma_paper_k, delta_paper_k, p_fuel)

  mod = Model(with_optimizer(Gurobi.Optimizer, Presolve=0, OutputFlag=0))

  @variable(mod, h_t[1:6]>=0, integer=true)
  @variable(mod, y_i[1:5]>=0)
  @variable(mod, y_j[1:2]>=0)
  @variable(mod, y_paper>=0)
  @variable(mod, z_ik[1:5, 1:4]>=0)
  @variable(mod, z_jk[1:2, 1:4]>=0)
  @variable(mod, z_paper_k[1:4]>=0)

  i_t1 = [[1], [2, 4], [3, 5]]
  i_t2 = [[1], [2, 4], [3, 5]]
  j_t2 = [[1], [], [2]]

  @objective(mod, Max,
    -sum(alpha_t[t] * h_t[t] + beta_t[t] * h_t[t]^2 for t = 1:6) +
    -sum(c_i[i] * y_i[i] for i=1:5) +
    -sum(c_j[j] * y_j[j] for j=1:2) - c_paper * y_paper +
    sum(alpha_t[t] * (
      h_t[t] - sum(a_i[i] * y_i[i] for i in i_t1[t])
      ) for t=1:3) +
    sum(alpha_t[t] * (
      h_t[t] + sum(b_i[i] * y_i[i] for i in i_t2[t-3]) -
      sum(a_j[j] * y_j[j] for j in j_t2[t-3]) -
        a_paper_t[t-3] * y_paper) for t=4:6) +
    sum(e_i[i] * y_i[i] for i=1:5) * p_fuel +
    sum(gamma_ik[i, k] * z_ik[i, k] - delta_ik[i, k] * z_ik[i, k]^2
      for i=1:5, k=1:4) +
    sum(gamma_jk[j, k] * z_jk[j, k] - delta_jk[j, k] * z_jk[j, k]^2
      for j=1:2, k=1:4) +
    sum(gamma_paper_k[k] * z_paper_k[k] -
      delta_paper_k[k] * z_paper_k[k]^2
      for k=1:4)
    )

  @constraint(mod, vec_cons_1[t=1:3],
    h_t[t] >= sum(a_i[i] * y_i[i] for i in i_t1[t]))
  @constraint(mod, vec_cons_2[t=4:6],
    h_t[t] + sum(b_i[i] * y_i[i] for i in i_t2[t-3]) >=
    sum(a_j[j] * y_j[j] for j in j_t2[t-3]) + a_paper_t[t-3] * y_paper)
  @constraint(mod, vec_cons_3[j=1:2],
    y_j[j] >= b_paper_j[j] * y_paper)
  @constraint(mod, vec_cons_4[i=1:5],
    y_i[i] >= sum(z_ik[i, k] for k=1:4))
  @constraint(mod, vec_cons_5[j=1:2],
    y_j[j] - b_paper_j[j] * y_paper >= sum(z_jk[j, k] for k=1:4))
  @constraint(mod, y_paper >= sum(z_paper_k[k] for k=1:4))
  @constraint(mod, sum(y_i[i] for i=1:3) <= r_saw)
  @constraint(mod, sum(y_i[i] for i=4:5) <= r_plywood)
  @constraint(mod, vec_cons_6[j=1:2], y_j[j] <= r_j[j])
  @constraint(mod, y_paper <= r_paper)

  optimize!(mod)

  result_h_t = value_vec(h_t)
  result_y_i = value_vec(y_i)
  result_y_j = value_vec(y_j)
  result_y_paper = value(y_paper)
  result_z_ik = value_mat(z_ik)
  result_z_jk = value_mat(z_jk)
  result_z_paper_k = value_vec(z_paper_k)
  result_obj = objective_value(mod)

  println(result_h_t)
  println(result_y_i)
  println(result_y_j)
  println(result_y_paper)
  println(result_z_ik)
  println(result_z_jk)
  println(result_z_paper_k)
  println(result_obj)

  # return result_h_t, result_y_i, result_y_j, result_y_paper, result_z_ik,
  #   result_z_jk, result_z_paper_k, result_obj
end


for N in [Float64, Rational{Int}, Float32]
    # singleton
    S = Singleton(N[1, -2, 3, -4])
    H = Hyperrectangle(zeros(N, 4), N[1, 2, 3, 4])
    @test symmetric_interval_hull(S) == H
    @test box_approximation_symmetric(S) == H

    # linear map of singleton
    M = Diagonal(ones(N, 4))
    @test symmetric_interval_hull(M * S) == H

    # Minkowski sum of singletons
    H2 = Hyperrectangle(zeros(N, 4), 2 * N[1, 2, 3, 4])
    @test symmetric_interval_hull(S + S) == H2

    # linear map of hyperrectangle
    M = Diagonal(N[1, -1, -2, 2])
    H1 = Hyperrectangle(5 * ones(N, 4), N[1, 2, 3, 4])
    H2 = Hyperrectangle(zeros(N, 4), N[6, 7, 16, 18])
    @test symmetric_interval_hull(M * H1) == H2
end

# tests that only work with Float64 and Float32
for N in [Float64, Float32]
    # exponential map of singleton
    M = sparse(diagm(N[1, 1, 1]))
    E = SparseMatrixExp(M) * Singleton(N[1, 2, 3])
    H = Hyperrectangle(zeros(N, 3), N[ℯ, 2ℯ, 3ℯ])
    @test symmetric_interval_hull(E) ≈ H

    # exponential map of hyperrectangle
    M = sparse(diagm(N[1, 1, 1]))
    E = SparseMatrixExp(M) * Hyperrectangle(N[1, 2, 3], N[1, 2, 3])
    H = symmetric_interval_hull(E)
    @test H ≈ Hyperrectangle(zeros(N, 3), N[2ℯ, 4ℯ, 6ℯ])
    # matrix with negative entries
    M = sparse(diagm(N[1, -1]))
    E = SparseMatrixExp(M) * Hyperrectangle(N[1, 1], N[1, 1])
    H = symmetric_interval_hull(E)
    @test H ≈ Hyperrectangle(zeros(N, 2), N[2ℯ, 2ℯ^(-1)])
end


using BSplineKit
using LinearAlgebra
using QuadGK
using SpecialFunctions: airyai

# Solution of Airy differential equation using Galerkin method.
# Problem taken from Olver & Townsend SIAM Rev. 2013.
# The original equation is for a function `u` that satisfies inhomogeneous
# Dirichlet BCs at x = ±1.
# Here we perform the change of variables v(x) = u(x) - u₀(x), such that `v`
# satisfies *homogeneous* Dirichlet BCs.
function test_airy_equation()
    ε = 1e-4
    N = 100
    breaks = collect(range(-1, 1; length = N + 1))
    B = @inferred BSplineBasis(BSplineOrder(6), breaks)
    R = @inferred RecombinedBSplineBasis(Derivative(0), B)  # for Dirichlet BCs

    # Exact solution
    u_exact(x) = airyai(x / cbrt(ε))

    # Affine part of the solution, for change of variables.
    # This is the (unique) affine curve that satisfies the two inhomogeneous BCs.
    u_0(x) = ((1 - x) * u_exact(-1) + (1 + x) * u_exact(1)) / 2

    # This is to obtain coefficients from data at collocation points
    xs = @inferred collocation_points(B)
    Cfact = @inferred lu!(collocation_matrix(B, xs))

    # Variable coefficient operator: c(x) * v(x), with c(x) = x.
    # Note that we could expand c(x) = x using linear splines (k = 2), but we
    # use the same basis as for `v` for simplicity.
    # Here Mc is the linear Galerkin operator associated to multiplication by
    # c(x).
    Mc = let
        cs = @inferred Cfact \ xs  # coefficients for c(x) = x
        Tc = @inferred galerkin_tensor((R, R, B), Derivative.((0, 0, 0)))
        Mc = @inferred Tc * cs
        @test Mc isa BandedMatrix
        @test Mc ≈ Mc'  # the matrix is practically symmetric
        Hermitian(Mc)
    end

    # Second-order term: ε v''(x)
    # Taking advantage of integration by parts, we construct the matrix
    # M′[i, j] = ⟨ ϕ′[i] ϕ′[j] ⟩.
    # The boundary term is zero due to Dirichlet BCs.
    L = let
        L = @inferred galerkin_matrix(R, Derivative.((1, 1)))
        @test L isa Hermitian
        @test L.uplo == 'U'
        parent(L) .*= ε
        L
    end

    # We add the two LHS operators
    # We only need to add the upper triangular part, because of symmetry.
    parent(L) .+= UpperTriangular(Mc)

    # Factorisation
    # Apparently we can't do Cholesky because `Mc` is not SPD.
    # For LU, we can't pass a Hermitian matrix wrapper...
    Lfact = lu(BandedMatrix(L))

    # RHS resulting from change of variables
    f(x) = x * u_0(x)
    rhs = let
        fs = ldiv!(Cfact, f.(xs))     # coefficients of RHS term
        Mf = galerkin_matrix((R, B))  # RHS operator
        -Mf * fs
    end

    # Numerical solution
    vR = Lfact \ rhs  # coefficients of the solution in recombined basis R
    vB = @inferred recombination_matrix(R) * vR  # same in full basis B
    vsol = @inferred Spline(B, vB)  # solution as a spline
    usol(x) = vsol(x) + u_0(x)  # undo change of variables

    # Check error (this will depend on the choices of ε, N, k, …)
    err, _ = quadgk(x -> (usol(x) - u_exact(x))^2, -1, 1; rtol = 1e-8)
    @test err < 1e-6

    nothing
end

@testset "Airy equation" begin
    test_airy_equation()
end


############################################################################################
"""
$(TYPEDEF)
Markov Kernel for particle propagation.

# Fields
$(TYPEDFIELDS)
"""
struct Markov{A,B,C} <: ParticleKernel
    "Initial distribution, function of iter only."
    initial::A
    "Transition distribution, function of full particle trajectory and current iteration count."
    transition::B
    "Data distribution to weight particles. Function of full data, particle trajectory and current iteration count."
    evidence::C
    function Markov(initial::A, transition::B, evidence::C) where {A,B,C}
        return new{A,B,C}(initial, transition, evidence)
    end
end

############################################################################################
function initial(_rng::Random.AbstractRNG, kernel::Markov)
    return rand(_rng, kernel.initial)
end

############################################################################################
function transition(
    _rng::Random.AbstractRNG, kernel::Markov, val::AbstractArray{P}, iter::Integer
) where {P}
    return rand(_rng, kernel.transition(val, iter))
end

############################################################################################
function ℓtransition(
    valₜ::Union{P,AbstractArray{P}}, kernel::Markov, val::AbstractArray{P}, iter::Integer
) where {P}
    return logpdf(kernel.transition(val, iter), valₜ)
end

############################################################################################
export Markov


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

qm = QuArray(m, LabelBasis(4,1), LabelBasis(2,2))
qm_qm = tensor(qm,qm)

@assert qm[(2,0),(1,0)] == qm[3,2]
@assert qm'[(1,0),(2,0)] == qm[3,2]'
@assert QuBase.bases(qm_qm) == (LabelBasis(4,1,4,1), LabelBasis(2,2,2,2))

@assert qm_qm[(2,0,1,0),(1,1,0,1)] == qm_qm[7,12]
qm_qm[(2,0,1,0),(1,1,0,1)] = 123
@assert qm_qm[(2,0,1,0),(1,1,0,1)] == 123

x2subspace = [filter((x...)->sum(x...)==2, QuBase.bases(qm_qm,2))...]
@assert x2subspace == [(1,1,0,0),(1,0,1,0),(0,1,1,0),(1,0,0,1),(0,1,0,1),(0,0,1,1)]

@assert qm_qm[(3,0,3,0), x2subspace] == [0+416im 0+392im 0+420im 0+420im 0+450im 0+416im]
qm_qm[(3,0,3,0), x2subspace] = [1 2 3 4 5 6]
@assert qm_qm[(3,0,3,0), x2subspace] == [1 2 3 4 5 6]


#
# # OnlinePLS: Thomas Breloff (Cointegrated Technologies)
#
#
# # implements an online partial least squares (OPLS) algorithm... modeled on Zeng et al (2014) [1]:
# #		"Incremental partial least squares analysis of big streaming data"
#
#
#
# # Goal: We're looking for a linear combination of the X-components which produces high covariance with y.
# # 			We then fit the reduced-dimension regression y = Zβ + e, where Z = XW has k columns of uncorrelated series.
#
# # Reason:  If X has high multicolinearity between components, there will likely be overfitting
# #					 problems for the model y = Xβ + e
#
#
# # Method:
# # There are 3 phases... model update, weight matrix extraction, and finally least squares fitting.
#
#
# # Model update:
#
# # We are trying to maximize Cov(y, Xw) for some dx1 vector of weights w...
# #	this will approximate the loadings for the first latent vector.
# # Taking the derivative of a Lagrange function gives us the solution w = Xᵀy / ‖Xᵀy‖
# # Our online model updates the estimate of v := Xᵀy, and also tracks the decomposition of X using OnlinePCA
# # in order to estimate the covariance matrix C := XᵀX/n, to reduce cpu/memory usage
#
# # Note: we replace the update of v as in [1] with a smoothed version, which is correct for the exponential case,
# # and approximately correct for the equal weighted case.
# # proof summary for exponential weighting correctness:
# #		vₜ = ∑ᵗ [λ(1-λ)ᵗ⁻ⁱyᵢXᵢ]
# #			 = λyₜXₜ + ∑ᵗ⁻¹ [λ(1-λ)ᵗ⁻¹⁻ⁱyᵢXᵢ]
# #			 = λyₜXₜ + (1-λ)vₜ₋₁
#
# # Weight matrix building:
#
# # When the state is requested: "PLS projection directions are computed from a Gram-Schmidt orthonormalization
# # of the Krylov sequence" [1]
# # This gives us the weighting matrix W:
# #		V := [v, Cv\{v}, C²v\{v,Cv}, ... Cᴷ⁻¹v\{v,Cv,...,Cᴷ⁻²v}]
# #		W := [V₁/‖V₁‖, ..., Vₖ/‖Vₖ‖]
#
# # where "f/{g,h...} refers to the components of f that are orthogonal to the space spanned by {g,h,...}" [1]
#
#
# # Regression:
#
# # Now that we have our weight matrix W, we can project X into a reduced dimensional space,
# #	and should get more robust regression results.
# #		Z = XW
# #		y = Zβ + e
#
#
# # Summary:
# # 	update v estimate
# #		update OnlinePCA estimate (V) of first l pca vectors
# #   build W matrix, and project Z = XW
# #		regress y = Zβ + e
# #		state = (v, V, W, β)
#
#
# #-------------------------------------------------------# Type and Constructors
#
# type OnlinePLS{W<:Weighting} <: OnlineStat
# 	d::Int  			# num dependent vars
# 	l::Int 				# num latent vars in OnlinePCA
# 	k::Int 				# num latent vars in PLS
# 	n::Int
# 	weighting::W
#
# 	ymean::Mean{W}
# 	xmeans::Means{W}
#
# 	v1::VecF				# (d x 1) vector -- estimate of first column of V
# 	W::MatF 				# (d x k) matrix (see fls comment below)
# 	pca::OnlinePCA{W}  # tracks pca decomposition of X
# 	fls::OnlineFLS  # flexible least squares to solve for β: y = (XW)β + e
# end
#
# name(o::OnlinePLS) = "OPLS"
#
#
# function OnlinePLS(d::Integer, l::Integer, k::Integer, δ::Real, wgt::Weighting = default(Weighting))
# 	OnlinePLS(d, l, k, 0, wgt, Mean(wgt), Means(d, wgt), zeros(d), zeros(d, k), OnlinePCA(d, l, wgt), OnlineFLS(k, δ, wgt))
# end
#
# function OnlinePLS(x::AVecF, y::Real, l::Integer, k::Integer, δ::Real, wgt::Weighting = default(Weighting))
# 	d = length(x)
# 	o = OnlinePLS(d, l, k, δ, wgt)
# 	update!(o, x, y)
# 	o
# end
#
# function OnlinePLS(X::AMatF, y::AVecF, l::Integer, k::Integer, δ::Real, wgt::Weighting = default(Weighting))
# 	d = size(X,2)
# 	o = OnlinePLS(d, l, k, δ, wgt)
# 	update!(o, X, y)
# 	o
# end
#
# #-----------------------------------------------------------------------# state
#
# # note: pca and fls refer to the states of the respective objects
# statenames(o::OnlinePLS) = [:v1, :pca, :W, :fls, :nobs]
# state(o::OnlinePLS) = Any[copy(o.v1), state(o.pca), copy(o.W), state(o.fls), nobs(o)]
#
#
# #---------------------------------------------------------------------# update!
#
#
# # # NOTE: assumes X mat is (T x p), where T is the number of observations
# # # TODO: optimize
# # function update!(o::OnlinePLS, X::AMatF, y::AVecF)
# # 	@assert length(y) == size(X,1)
# # 	@inbounds for i in 1:length(y)
# # 		update!(o, row(X,i), y[i])
# # 	end
# # end
#
# function update!(o::OnlinePLS, x::AVecF, y::Real)
#
# 	# center x and y
# 	y = center!(o.ymean, y)
# 	x = center!(o.xmeans, x)
#
# 	# update v1 and pca
# 	# @LOG y x y*x
# 	smooth!(o.v1, y * x, weight(o))
# 	# @LOG o.v1
# 	update!(o.pca, x)
#
# 	# recompute W
# 	V = VecF[]
# 	vi = copy(o.v1)
# 	w = copy(o.v1)
#
# 	# do the gram-schmidt process
# 	@inbounds for i in 1:o.k
# 		if i > 1
# 			# multiply Cw
# 			nextw = zeros(o.d)
# 			for j in 1:o.l
# 				pcaUj = row(o.pca.U, j)  # dx1 vector
# 				pcaVj = row(o.pca.V, j)	 # dx1 vector -- normalized -- eigenvector
# 				nextw += pcaUj * dot(pcaVj, w) / o.n  # TODO:  I don't like this n here... doesn't make sense with exponential weighting... how to change??
# 			end
# 		end
#
# 		# compute vi and add to V
# 		vi = copy(w)
# 		for j in 1:i-1
# 			vi -= vi .* V[j] .* V[j]
# 		end
# 		push!(V, vi)
#
# 		# add normalized vi to W
# 		col!(o.W, i, vi / if0then1(norm(vi)))
# 	end
#
# 	# update fls regression
# 	update!(o.fls, y, o.W' * x)
# 	# @LOG o.fls
#
# 	o.n += 1
# 	return
#
# end
#
#
# function Base.empty!(o::OnlinePLS)
# 	o.n = 0
# 	empty!(o.ymean)
# 	empty!(o.xmeans)
# 	o.v1 = zeros(o.d)
# 	o.W = zeros(o.d, o.k)
# 	empty!(o.pca)
# 	empty!(o.fls)
# 	nothing
# end
#
# function Base.merge!(o1::OnlinePLS, o2::OnlinePLS)
# 	error("Merging undefined for PLS")
# end
#
#
# function StatsBase.coef(o::OnlinePLS)
# 	error("not implemented")
# end
#
# function StatsBase.coeftable(o::OnlinePLS)
# 	error("not implemented")
# end
#
# function StatsBase.confint(o::OnlinePLS, level::Real = 0.95)
# 	error("not implemented")
# end
#
# # predicts yₜ for a given xₜ
# function StatsBase.predict(o::OnlinePLS, x::AVecF)
# 	# center x
# 	# multiply Wx to get a (kx1) vector of latent vars
# 	# y = predict(fls,Wx)
# 	# uncenter y
# 	x = center(o.xmeans, x)
# 	y = predict(o.fls, o.W' * x)
# 	uncenter(o.ymean, y)
# end
#
# # NOTE: uses most recent estimate of βₜ to predict the whole matrix
# function StatsBase.predict(o::OnlinePLS, X::AMatF)
# 	n = size(X,1)
# 	pred = zeros(n)
# 	for i in 1:n
# 		pred[i] = predict(o, row(X,i))
# 	end
# 	pred
# end
#
#


############################################################################
#
#  Quad_Forces.jl
#
#  Purpose:    Computes the forces and moments due to rotors.
#
#  Author:     Edward Daughtery
#
#  Date:       05 August 2016
#
#
############################################################################

function Quad_Forces(motor_array::Motors)

    ## Motor Rate (rad/s)
    wmotor1 = motor_array.motor1.position;
    wmotor2 = motor_array.motor2.position;
    wmotor3 = motor_array.motor3.position;
    wmotor4 = motor_array.motor4.position;

    ## Propeller Forces
    keng = 6.11e-4;
    f1_body = [0;0;-keng*wmotor1^2];
    f2_body = [0;0;-keng*wmotor2^2];
    f3_body = [0;0;-keng*wmotor3^2];
    f4_body = [0;0;-keng*wmotor4^2];

    ## Propeller Moments
    kmeng = 1.5e-5;
    m1_body = [0;0;-kmeng*wmotor1^2];
    m2_body = [0;0;kmeng*wmotor2^2];
    m3_body = [0;0;-kmeng*wmotor3^2];
    m4_body = [0;0;kmeng*wmotor4^2];

    ## Moment arm from CG to Propeller Hub;
    Larm = 0.15;
    r1_arm_body = [ Larm*cosd(45); -Larm*sind(45); 0.0];
    r2_arm_body = [ Larm*cosd(45);  Larm*sind(45); 0.0];
    r3_arm_body = [-Larm*cosd(45);  Larm*sind(45); 0.0];
    r4_arm_body = [-Larm*cosd(45); -Larm*sind(45); 0.0];

    ## Total Moment due to Motors
    m1_body = m1_body + cross(r1_arm_body, f1_body);
    m2_body = m2_body + cross(r2_arm_body, f2_body);
    m3_body = m3_body + cross(r3_arm_body, f3_body);
    m4_body = m4_body + cross(r4_arm_body, f4_body);

    ## Creating Quad Forces Structure
    quadforces = QuadForces(f1_body,
                            f2_body,
                            f3_body,
                            f4_body,
                            m1_body,
                            m2_body,
                            m3_body,
                            m4_body);

    return quadforces;
end


@with_kw struct CarterMethodBL{T} <: AbstractFirstOrderMetricParams{T}
    @deftype T
    M = 1.0
    a = 0.0
    E = 1.0
end

function carter_velocity(u, E, M, a, p)
    let r = u[2], θ = u[3], L = p.L, Q = p.Q
        Σ₀ = CarterMethodBLCoords.Σ(r, a, θ)
        (
            CarterMethodBLCoords.Σδt_δλ(E, L, M, r, a, θ) / Σ₀,
            p.r * CarterMethodBLCoords.Σδr_δλ(E, L, M, Q, r, a) / Σ₀,
            p.θ * CarterMethodBLCoords.Σδθ_δλ(E, L, Q, a, θ) / Σ₀,
            CarterMethodBLCoords.Σδϕ_δλ(E, L, M, r, a, θ) / Σ₀,
        )
    end
end

calc_lq(m::CarterMethodBL{T}, pos, vel) where {T} =
    CarterMethodBLCoords.LQ(m.M, pos[2], m.a, pos[3], vel[3], vel[4])

inner_radius(m::CarterMethodBL{T}) where {T} = m.M + √(m.M^2 - m.a^2)
constrain(::CarterMethodBL{T}, u, v; μ = 0.0) where {T} = v[1]

function integrator_problem(
    m::CarterMethodBL{T},
    pos::StaticVector{S,T},
    vel::StaticVector{S,T},
    time_domain,
) where {S,T}
    L, Q = calc_lq(m, pos, vel)
    ODEProblem{false}(pos, time_domain, make_parameters(L, Q, vel[2], T)) do u, p, λ
        SVector(carter_velocity(u, m.E, m.M, m.a, p)...)
    end
end

function integrator_problem(
    m::CarterMethodBL{T},
    pos::AbstractVector{T},
    vel::AbstractVector{T},
    time_domain,
) where {T}
    L, Q = calc_lq(m, pos, vel)
    ODEProblem{true}(pos, time_domain, make_parameters(L, Q, vel[2], T)) do du, u, p, λ
        du .= carter_velocity(u, m.E, m.M, m.a, p)
    end
end

function alpha_beta_to_vel(m::CarterMethodBL{T}, u, α, β) where {T}
    sinΦ, sinΨ = CarterMethodBLCoords.sinΦsinΨ(m.M, u[2], m.a, u[3], α, β)
    (β < 0.0 ? -1.0 : 1.0, sinΦ, sinΨ)
end

function Vr(m::CarterMethodBL{T}, u, p) where {T}
    L, Q, _, _ = p
    CarterMethodBLCoords.Vr(m.E, L, m.M, Q, u[2], m.a)
end
function Vθ(m::CarterMethodBL{T}, u, p) where {T}
    L, Q, _, _ = p
    CarterMethodBLCoords.Vθ(m.E, L, Q, m.a, u[3])
end

export CarterMethodBL