export ETKF, rdnortho

"""
$(TYPEDSIGNATURES)

A routine to generate mean-preserving random rotations.

References:

Nerger, L., Janjić, T., Schröter, J., & Hiller, W. (2012). A unification of ensemble square root Kalman filters. Monthly Weather Review, 140(7), 2335-2345.
Tödter, J., & Ahrens, B. (2015). A second-order exact ensemble square root filter for nonlinear data assimilation. Monthly Weather Review, 143(4), 1347-1367.
"""
function rdnortho(N::Int64)
    Ω = (svd(randn(N-1,N-1)).V)';
    b1 = 1/√(N)
    b1sgn = b1*ones(N)
    b1sgn[end] += sign(b1)
    B = zeros(N,N)
    B[:,1] = b1*ones(N);
    B[:,2:end] = (-(I - 1/(abs(b1)+1.0)*b1sgn*b1sgn'))[:,1:end-1]
    Λb = zeros(N,N)
    Λb[1,1] = 1.0
    Λb[2:end,2:end] = Ω
    return B*Λb*B'
end

"""
$(TYPEDEF)

A structure for the ensemble transform Kalman filter (ETKF)

References:

Bishop, C. H., Etherton, B. J., & Majumdar, S. J. (2001). Adaptive sampling with the ensemble transform Kalman filter. Part I: Theoretical aspects. Monthly weather review, 129(3), 420-436.

## Fields
$(TYPEDFIELDS)
"""

struct ETKF<:SeqFilter
    "Filter function"
    G::Function

    "Standard deviations of the measurement noise distribution"
    ϵy::AdditiveInflation

    "Time step dynamic"
    Δtdyn::Float64

    "Time step observation"
    Δtobs::Float64

	"Random shuffle"
	Δtshuff::Float64

	"Boolean: is the covariance matrix localized"
	islocal::Bool

    "Boolean: is state vector filtered"
    isfiltered::Bool
end

function ETKF(G::Function, ϵy::AdditiveInflation,
    Δtdyn, Δtobs, Δtshuff; islocal = false, isfiltered = false)
    @assert norm(mod(Δtobs, Δtdyn))<1e-12 "Δtobs should be an integer multiple of Δtdyn"

    return ETKF(G, ϵy, Δtdyn, Δtobs, Δtshuff, islocal, isfiltered)
end

# If no filtering function is provided, use the identity in the constructor.
function ETKF(ϵy::AdditiveInflation,
    Δtdyn, Δtobs, Δtshuff; islocal = false)
    @assert norm(mod(Δtobs, Δtdyn))<1e-12 "Δtobs should be an integer multiple of Δtdyn"

    return ETKF(x-> x, ϵy, Δtdyn, Δtobs, Δtshuff, islocal, false)
end

function Base.show(io::IO, enkf::ETKF)
	println(io,"ETKF  with filtered = $(enkf.isfiltered)")
end

# """
#     Define action of ETKF on EnsembleStateMeas
# """
# Bocquet Data assimilation p.160 Chapter 6 ETKF

function (enkf::ETKF)(X, ystar, t)
	Ny = size(ystar,1)
	Nx = size(X,1)-Ny
	Ne = size(X, 2)

	meas  = viewmeas(X,Ny,Nx)
	state = viewstate(X,Ny,Nx)

	x̄ = copy(mean(state,dims=2)[:,1])

	Xf = copy(state)
	Xf .-= x̄
	rmul!(Xf, 1/sqrt(Ne-1))

	ȳ = copy(mean(meas,dims=2)[:,1])

	S = enkf.ϵy.σ^(-1)*(meas .- ȳ)

	rmul!(S, 1/√(Ne-1))

	δH = enkf.ϵy.σ^(-1)*(ystar - ȳ)

	# Th = Hermitian(inv(S'*S+I))
	λ, ϕ = eigen(Hermitian(S'*S + I))

	# w = Th*(S'*δH)
	w = ϕ * Diagonal(1 ./ λ) * ϕ'*(S'*δH)

	if mod(t, enkf.Δtshuff) ==0
		U = rdnortho(Ne)
	else
		U = I
	end

	# ens.state.S .= x̄ .+ Xf*(w .+ √(Ne-1)*sqrt(Th)*U)
	state .= x̄ .+ Xf*(w .+ √(Ne-1)*ϕ * Diagonal( 1 ./ sqrt.(λ)) * ϕ'*U)
	return X
end

function (enkf::ETKF)(X, ystar, ȳf, t)

	Ny = size(ystar,1)
	Nx = size(X,1)-Ny
	Ne = size(X, 2)

    meas  = viewmeas(X,Ny,Nx)
	state = viewstate(X,Ny,Nx)

	x̄ = copy(mean(state,dims=2)[:,1])

	Xf = copy(state)
	Xf .-= x̄
	rmul!(Xf, 1/sqrt(Ne-1))

	S = enkf.ϵy.σ^(-1)*(meas .- ȳf)

	δH = enkf.ϵy.σ^(-1)*(ystar - ȳf)

	# Th = inv(Hermitian(I + S'*S))
	λ, ϕ = eigen(Hermitian(S'*S + I))

	# w = Th*(S'*δH)
	w = ϕ * Diagonal(1 ./ λ) * ϕ'*(S'*δH)

	if mod(t, enkf.Δtshuff) == 0
		U = rdnortho(Ne)
	else
		U = I
	end

	# ens.state.S .= x̄ + Xf*(w + √(Ne-1)*sqrt(Th)*U')
	state .= x̄ .+ Xf*(w .+ √(Ne-1)*ϕ * Diagonal( 1 ./ sqrt.(λ)) * ϕ'*U)
	return X
end


module TestBazykin85
include("preamble_plots.jl")

using DiffEqBase: ODEProblem
using StaticArrays: SVector

using Bifurcations: resolved_points
using Bifurcations.BifurcationsBase: contkind, HopfCont
using Bifurcations.Codim2: NormalizingAS, BackReferencingAS
using Bifurcations.Examples.Bazykin85

# Bifurcation points calculated in some version of Bifurcations.jl
KNOWN_POINTS = Dict(
    # Bazykin85.Bazykin85Param => points
    Bazykin85.Bazykin85Param(
        ϵ = 0.01,
    ) => [
        (Codim2.PointTypes.cusp,
         [32.5936, 20.4743, 0.956835, 0.290633, 0.901233, 0.00356842]),
        (Codim2.PointTypes.bogdanov_takens,
         [42.688, 21.6306, 0.993354, 0.1151, 0.860704, 0.00605879]),
        (Codim2.PointTypes.bogdanov_takens,
         [6.26577, 3.60155, 0.932668, 0.360736, 0.453624, 0.175128]),
        (Codim2.PointTypes.cusp,
         [25.8231, 2.44209, 0.999458, 0.0329046, 0.0887673, 2.80236]),
    ],
)

@testset "smoke Bazykin85 codim-2 (ϵ=$ϵ, $(nameof(ASType)))" for
        ϵ in [0.01, 0.001],
        ASType in [NormalizingAS, BackReferencingAS]

    prob = Bazykin85.make_prob(
        Bazykin85.Bazykin85Param(
            ϵ = ϵ,
        ),
        phase_space = (SVector(0.0, 0.0),  # u_min
                       SVector(Inf, Inf)), # u_max
    )
    ode = prob.p.de_prob

    codim1_solver = init(prob)
    solve!(codim1_solver)

    resolved = resolved_points(codim1_solver)
    if ϵ > 0.001
        @test length(resolved) == 5
    else
        @test_broken length(resolved) == 5
    end

    sn_point, = sort!(
        special_intervals(codim1_solver, Codim1.PointTypes.saddle_node),
        by=p->p.u0[1])
    hopf_point, = special_intervals(codim1_solver, Codim1.PointTypes.hopf)

    hopf_solver1 = nothing
    sn_solver1 = nothing
    for point in [sn_point, hopf_point]
        codim2_prob = BifurcationProblem(
            point,
            codim1_solver,
            (@lens _.δ),
            (0.0, 10.0),
            augmented_system = ASType(),
        )
        codim2_solver = init(
            codim2_prob;
            # h0 = 0.001,
            nominal_angle_rad = 0.01,
            max_samples = 1000,
        )
        solve!(codim2_solver)

        @testset "show" begin
            smoke_test_solver_show(codim2_solver)
        end

        @testset "plot" begin
            smoke_test_solver_plot(codim2_solver)
        end

        @testset "resolve $p" for p in [Codim2.PointTypes.bogdanov_takens,
                                        Codim2.PointTypes.fold_hopf,
                                        Codim2.PointTypes.hopf_hopf]
            @test_nothrow resolved_points(codim2_solver, p)
        end

        if contkind(codim2_solver) isa HopfCont
            hopf_solver1 = codim2_solver
        else
            sn_solver1 = codim2_solver
        end
    end

    @testset "switch to saddle-node" begin
        @assert hopf_solver1 !== nothing
        # Find the right-most Bogdanov-Takens bifurcation:
        point = first(sort(
            special_intervals(hopf_solver1,
                           Codim2.PointTypes.bogdanov_takens);
            by = p -> p.u0[end - 1],  # α
            rev = true))
        sn_prob = BifurcationProblem(
            point, hopf_solver1;
            augmented_system = ASType(),
        )
        sn_solver2 = init(sn_prob)
        @test_nothrow solve!(sn_solver2)
    end

    @testset "switch to hopf" begin
        @assert sn_solver1 !== nothing
        # Find the right-most Bogdanov-Takens bifurcation:
        point = first(sort(
            special_intervals(sn_solver1,
                           Codim2.PointTypes.bogdanov_takens);
            by = p -> p.u0[end - 1],  # α
            rev = true))
        hopf_prob = BifurcationProblem(
            point, sn_solver1;
            augmented_system = ASType(),
        )
        hopf_solver2 = init(hopf_prob)
        @test_nothrow solve!(hopf_solver2)

        bpoints = special_intervals(hopf_solver2, Codim2.PointTypes.bautin)
        @test length(bpoints) == 1
        @test_nothrow resolved_points(hopf_solver2, Codim2.PointTypes.bautin)
    end

    if haskey(KNOWN_POINTS, ode.p)
        desired_points = KNOWN_POINTS[ode.p]
        @testset "known points" begin
            # TODO: do this for all solvers (hopf_solver1 etc.)
            uniquified = sort(
                uniquify_points(resolved_points(sn_solver1)),
                by = point -> point.u[end],
            )
            @test length(uniquified) == length(desired_points)
            for (actual, desired) in zip(uniquified, desired_points)
                type_desired, u_desired = desired
                if ASType == NormalizingAS
                    @test actual.u ≈ u_desired  rtol=1e-4
                else
                    # Compare DS state and parameters; ignore the eigenvector
                    # TODO: don't hard-code offsets:
                    @test actual.u[1:2] ≈ u_desired[1:2]  rtol=1e-4
                    @test actual.u[end-1:end] ≈ u_desired[end-1:end]  rtol=1e-4
                end
                @test actual.point_type == type_desired
            end
        end
    end
end

end  # module


function test_thermalLoad_1()

    T = 2

    # With these values, the heater uses 0.2 power to keep a ΔT of 1
    l = DR.DER.ThermalLoad(
        index=1, T=2,
        temp_min=[20.0, 20.0],
        temp_max=[25.0, 25.0],
        temp_ext=[0.0, 0.0],
        temp_init=20.0,
        pwr_min=[0.0, 0.0],
        pwr_max=[10.0, 10.0],
        C=1.0, η=1.0, μ=0.2,
        binflag=false
    )

    # GLPK does not support -in-Interval constraints, so we bridge
    model = MOI.Bridges.full_bridge_optimizer(GLPK.Optimizer(), Float64)

    lmin = [0.0, 0.0]
    lmax = [10.0, 10.0]
    price = [2.0, 1.0]

    # Instantiate initial model
    h = DR.DER.House(
        0, T,
        lmin, lmax, price,
        [l]
    )

    model, var2idx, con2idx = DR.DER.build_model!(h, model)

    MOI.optimize!(model)

    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 12.0

    # Temperature should be kept at minimum value, i.e., 20
    T1 = var2idx[(:thermal, l.index, :temp, 1)]
    T2 = var2idx[(:thermal, l.index, :temp, 2)]
    @test MOI.get(model, MOI.VariablePrimal(), T1) ≈ 20.0
    @test MOI.get(model, MOI.VariablePrimal(), T2) ≈ 20.0

    # Check net load
    pnet1 = var2idx[(:thermal, l.index, :pnet, 1)]
    pnet2 = var2idx[(:thermal, l.index, :pnet, 2)]
    @test MOI.get(model, MOI.VariablePrimal(), pnet1) ≈ 4.0
    @test MOI.get(model, MOI.VariablePrimal(), pnet2) ≈ 4.0

    return nothing

end

function test_thermalLoad_2()

    T = 2

    # With these values, the heater uses 0.2 power to keep a ΔT of 1
    l = DR.DER.ThermalLoad(
        index=1, T=2,
        temp_min=[20.0, 20.0],
        temp_max=[25.0, 25.0],
        temp_ext=[0.0, 0.0],
        temp_init=20.0,
        pwr_min=[0.0, 0.0],
        pwr_max=[10.0, 10.0],
        C=1.0, η=1.0, μ=0.2,
        binflag=false
    )

    # GLPK does not support -in-Interval constraints, so we bridge
    model = MOI.Bridges.full_bridge_optimizer(GLPK.Optimizer(), Float64)

    lmin = [0.0, 0.0]
    lmax = [10.0, 10.0]
    price = [1.0, 2.0]  # Now it's worth heating more in the first period

    # Instantiate initial model
    h = DR.DER.House(
        0, T,
        lmin, lmax, price,
        [l]
    )

    model, var2idx, con2idx = DR.DER.build_model!(h, model)

    MOI.optimize!(model)

    # T1 should be higher
    # T2 should be at minimum, i.e., 20
    T1 = var2idx[(:thermal, l.index, :temp, 1)]
    T2 = var2idx[(:thermal, l.index, :temp, 2)]
    @test MOI.get(model, MOI.VariablePrimal(), T1) ≈ 25.0
    @test MOI.get(model, MOI.VariablePrimal(), T2) ≈ 20.0

    # Check net load
    pnet1 = var2idx[(:thermal, l.index, :pnet, 1)]
    pnet2 = var2idx[(:thermal, l.index, :pnet, 2)]
    @test MOI.get(model, MOI.VariablePrimal(), pnet1) ≈ 9.0
    @test MOI.get(model, MOI.VariablePrimal(), pnet2) ≈ 0.0

    # Check objective value
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 9.0
    return nothing

end

@testset "Thermal" begin
    test_thermalLoad_1()
    test_thermalLoad_2()
end


# -*- coding: utf-8 -*-
#
# ANISODIFF - Anisotropic diffusion
#
# Usage:
#   diff = anisodiff(img, niter, kappa, lambda, option)
#
# Arguments:
# | img    - input image (2D grayscale)
# | niter  - number of iterations
# | kappa  - conduction coefficient (gradient modulus threshold)
# |          This parameter controls conduction as a function of gradient.
# |          If kappa is low, small intensity gradients are able to block
# |          conduction and hence diffusion across step edges.  A large value
# |          reduces the influence of intensity gradients on conduction.
# | lambda - integration constant for stability (0 <= lambda <= 0.25)
# |          This parameter controls the diffusion speed, and you
# |          usually want it at the maximum value of 0.25.
# |          default value: 0.25
# | option - conduction coefficient functions proposed by Perona & Malik:
# |          1: c(x, y, t) = exp(-(nablaI/kappa).^2)
# |             privileges high-contrast edges over low-contrast ones
# |          2: c(x, y, t) = 1 ./ (1 + (nablaI/kappa).^2)
# |             privileges wide regions over smaller ones
# |          default value: 1
#
# Returns:
# | diff   - anisotropic diffused image
#
# Reference:
# [1] P. Perona and J. Malik.
#     Scale-space and edge detection using ansotropic diffusion.
#     IEEE Transactions on Pattern Analysis and Machine Intelligence,
#     12(7):629-639, July 1990.
#     https://dx.doi.org/10.1109%2F34.56205
#
# Credits:
# [1] Peter Kovesi
#     pk@peterkovesi.com
#     MATLAB and Octave Functions for Computer Vision and Image Processing
#     http://www.peterkovesi.com/matlabfns/Spatial/anisodiff.m
#     --
#     June 2000  original version
#     March 2002 corrected diffusion eqn No 2.
# [2] Daniel Lopes
#     Anisotropic Diffusion (Perona & Malik)
#     http://www.mathworks.com/matlabcentral/fileexchange/14995-anisotropic-diffusion--perona---malik-
#
#
# Aaron LI <aaronly.me@gmail.com>
# 2015/07/17
#

include("calc_k_percentile.jl");

function anisodiff(img, niter, k=calc_k_percentile, lambda=0.25, option=1)
    diff = float(img)
    rows, cols = size(diff)

    for i = 1:niter
        println("anisodiff - iteration: ", i)

        # Construct diffl which is the same as diff but
        # has an extra padding of zeros around it.
        diffl = zeros(rows+2, cols+2)
        diffl[2:rows+1, 2:cols+1] = diff

        # North, South, East and West differences
        deltaN = diffl[1:rows,   2:cols+1] - diff
        deltaS = diffl[3:rows+2, 2:cols+1] - diff
        deltaE = diffl[2:rows+1, 3:cols+2] - diff
        deltaW = diffl[2:rows+1, 1:cols]   - diff

        # Calculate the kappa
        if isa(k, Function)
            kappa = k(diff)
        else
            kappa = k
        end

        println("  kappa: ", kappa)

        # Conduction
        if option == 1
            cN = exp(-(deltaN/kappa).^2)
            cS = exp(-(deltaS/kappa).^2)
            cE = exp(-(deltaE/kappa).^2)
            cW = exp(-(deltaW/kappa).^2)
        elseif option == 2
            cN = 1 ./ (1 + (deltaN/kappa).^2)
            cS = 1 ./ (1 + (deltaS/kappa).^2)
            cE = 1 ./ (1 + (deltaE/kappa).^2)
            cW = 1 ./ (1 + (deltaW/kappa).^2)
        end

        diff += lambda * (cN.*deltaN + cS.*deltaS + cE.*deltaE + cW.*deltaW)
    end

    return diff
end


module Test_ordred

using DescriptorSystems
using LinearAlgebra
using Polynomials
using Random
using Test

println("Test_ordred")
Random.seed!(2123)

@testset "Order Reduction Tools" begin

@testset "gss2ss" begin

sys = rdss(0,0,0);
sys1, r = gss2ss(sys);
@test iszero(sys-sys1) && r == 0

# Example: Kunkel & Mehrmann 2006 (two simple infinite eigenvalues, one uncontrollable, one non-dynamic)
a = [0 1 0 0; 1 0 0 0; -1 0 0 1; 0 1 1 1]; e = [1 0 0 0;0 0 1 0;0 0 0 0; 0 0 0 0];
b = [0;0;0;-1]; c = [0 0 1 0]; d = [0]; x0 = [1,2,0,1]; u0 = [1,];
sysc = dss(a,e,b,c,d);

# inconsistent initial condition
sysrc, xt0, Mx, Mu = dss2ss(sysc, x0; state_mapping = true, simple_infeigs=false);
@test iszero(sysc-sysrc,atol=1.e-7) && norm(Mx*xt0+Mu*u0-x0) > 0.1
x1 = Mx*xt0+Mu*u0;
sysrc, xt2, Mx, Mu = dss2ss(sysc, x1; state_mapping = true, simple_infeigs=false);
@test iszero(sysc-sysrc,atol=1.e-7) && norm(Mx*xt2+Mu*u0-x1) < 0.00001

# inconsistent initial condition
sysrc, xt1, Mx1, Mu1 = dss2ss(sysc, x0; state_mapping = true, simple_infeigs=true);
@test iszero(sysc-sysrc,atol=1.e-7) && norm(Mx1*xt1+Mu1*u0-x0) > 0.1
x1 = Mx1*xt1+Mu1*u0;
sysrc, xt3, Mx1, Mu1 = dss2ss(sysc, x1; state_mapping = true, simple_infeigs=true);
@test iszero(sysc-sysrc,atol=1.e-7) && norm(Mx1*xt3+Mu1*u0-x1) < 0.00001

sys1, r = gss2ss(sysc);
@test iszero(sysc-sys1) && r == 2

a = [8.872223171059933e-01     1.082089117437485e-01    -1.536042862271242e+00
                         0    -5.047254200920694e-01    -5.841082955835114e-02
                         0                         0     1.966838801162540e-01];
b = [ 1.688311161078364e-02     1.129239736565561e+00
     -9.928988024046033e-01     5.713290486116640e-02
      5.766407371577509e-02     8.543665825437958e-01];
c = [ -1.092400401898314e-01     1.061567864551283e+00    -1.535723490367775e+00
       3.312197206498771e-01     4.039678086153220e-02     3.924365487077013e+00];

d = [0     0
     0     0];
e = [4.992358310049688e-01     6.088864643615568e-02    -8.643241045903737e-01
                         0                         0                         0
                         0                         0                         0];
sysc = dss(a,e,b,c,d);
x0 = [1,2,-1]; u0 = [1,-1];

# inconsistent initial condition
sysrc, xt0, Mx, Mu = dss2ss(sysc, x0; state_mapping = true, simple_infeigs=false);
@test iszero(sysc-sysrc,atol=1.e-7) && norm(Mx*xt0+Mu*u0-x0) > 0.1
x1 = Mx*xt0+Mu*u0;
sysrc, xt2, Mx, Mu = dss2ss(sysc, x1; state_mapping = true, simple_infeigs=false);
@test iszero(sysc-sysrc,atol=1.e-7) && norm(Mx*xt2+Mu*u0-x1) < 0.00001

# inconsistent initial condition
sysrc, xt1, Mx1, Mu1 = dss2ss(sysc, x0; state_mapping = true, simple_infeigs=true);
@test iszero(sysc-sysrc,atol=1.e-7) && norm(Mx1*xt1+Mu1*u0-x0) > 0.1
x1 = Mx1*xt1+Mu1*u0;
sysrc, xt3, Mx1, Mu1 = dss2ss(sysc, x1; state_mapping = true, simple_infeigs=true);
@test iszero(sysc-sysrc,atol=1.e-7) && norm(Mx1*xt3+Mu1*u0-x1) < 0.00001

sys1, r = gss2ss(sysc)
@test iszero(sysc-sys1,atol=1.e-7) && r == 1

Ty = Float64
for Ty in (Float64,Complex{Float64})
sys = rdss(T = Ty,7,2,3);
sys1, r = gss2ss(sys);
@test iszero(sys-sys1,atol=1.e-7) && r == 7 && sys1.E == I

sys2, xt0, Mx, Mu = dss2ss(sys; state_mapping = true, simple_infeigs=false);
@test iszero(sys2-sys,atol=1.e-7)

sys1, r = gss2ss(sys,Eshape="triu");
@test iszero(sys-sys1,atol=1.e-7) && r == 7 && istriu(sys1.E)

sys2, r2 = gss2ss(sys1,Eshape="triu");
@test iszero(sys-sys2,atol=1.e-7) && r2 == 7 && istriu(sys2.E)

sys3, r3 = gss2ss(sys1,Eshape="ident");
@test iszero(sys-sys3,atol=1.e-7) && r3 == 7 && sys3.E == I

sys1, r = gss2ss(sys,Eshape="diag");
@test iszero(sys-sys1,atol=1.e-7) && r == 7 && isdiag(sys1.E)

sys = rdss(T = Ty,3,2,3,id=ones(Int,2));
sys1, r = gss2ss(sys);
@test iszero(sys-sys1,atol=1.e-7) && r == 3 && sys1.E == I

sys2, xt0, Mx, Mu = dss2ss(sys; state_mapping = true, simple_infeigs=false);
@test iszero(sys2-sys,atol=1.e-7)

sys1, r = gss2ss(sys,Eshape="triu");
@test iszero(sys-sys1,atol=1.e-7) && r == 3 && istriu(sys1.E)

sys1, r = gss2ss(sys,Eshape="diag");
@test iszero(sys-sys1,atol=1.e-7) && r == 3 && isdiag(sys1.E)

sys = rdss(T = Ty,3,2,3,id=[ones(Int,2); 2*ones(Int,1)]);
sys1, r = gss2ss(sys);
@test iszero(sys-sys1,atol=1.e-7) && r == 4 && sys1.E[1:r,1:r] == I

sys1, r = gss2ss(sys,Eshape="triu");
@test iszero(sys-sys1,atol=1.e-7) && r == 4 && istriu(sys1.E)

sys1, r = gss2ss(sys,Eshape="diag");
@test iszero(sys-sys1,atol=1.e-7) && r == 4 && isdiag(sys1.E)

sys = rdss(T = Ty,disc=true,3,2,3,id=[ones(Int,2); 2*ones(Int,1)])';
sys1, r = gss2ss(sys,Eshape="diag");
@test iszero(sys1-sys,atol=1.e-7) && r == 7 && isdiag(sys1.E)

sys = rdss(T = Ty,3,2,2,id=[ones(Int,2); 2*ones(Int,1)],disc=true);
sys1  = gminreal(gss2ss(sys/sys)[1],atol=1.e-7)
@test iszero(sys1-I,atol=1.e-7)

sys = rdss(T = Ty,3,2,2,id = [ones(Int,2); 2*ones(Int,1)],disc=true);
sys1  = gminreal(gss2ss(sys/sys)[1],atol=1.e-7);
@test iszero(sys1-I,atol=1.e-7)

sys = rdss(T = Ty,3,2,2,id = ones(Int,2), iduc=[ones(Int,2); 2*ones(Int,1)],disc=true);
sys2, xt0, Mx, Mu = dss2ss(sys; fast = true, atol=1.e-7, state_mapping = true, simple_infeigs=false);
@test iszero(sys2-sys,atol=1.e-7)

sys = rdss(T = Ty,3,2,2,id = ones(Int,2), iduc=[ones(Int,2); 2*ones(Int,1)],disc=true);
sys2, xt0, Mx, Mu = dss2ss(sys; fast = false, atol=1.e-7, state_mapping = true, simple_infeigs=false);
@test iszero(sys2-sys,atol=1.e-7)

end #Ty

end #gss2ss

@testset "gminreal and gir" begin

for fast in (true, false)

A2 = zeros(0,0); E2 = zeros(0,0); C2 = zeros(0,0); B2 = zeros(0,0); D2 = zeros(0,0);
sys = dss(A2,E2,B2,C2,D2);

sys1 = gminreal(sys, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)

sys1 = gir(sys, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)

A2 = rand(3,3); E2 = zeros(3,3); C2 = zeros(0,3); B2 = zeros(3,0); D2 = zeros(0,0);
sys = dss(A2,E2,B2,C2,D2);

sys1 = gminreal(sys, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)

sys1 = gminreal(sys, fast = fast, contr = false)
@test iszero(sys-sys1,atol=1.e-7)

sys1 = gir(sys, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)

sys1 = gir(sys, fast = fast, contr = false)
@test iszero(sys-sys1,atol=1.e-7)

# B and D vectors
A2 = rand(3,3); E2 = zeros(3,3); C2 = zeros(1,3); B2 = zeros(3,1); D2 = zeros(1);
sys = dss(A2,E2,B2,C2,D2);

sys1 = gminreal(sys, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)

sys1 = gminreal(sys, fast = fast, contr = false)
@test iszero(sys-sys1,atol=1.e-7)

sys1 = gir(sys, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)

sys1 = gir(sys, fast = fast, contr = false)
@test iszero(sys-sys1,atol=1.e-7)

# Example 1: DeTeran, Dopico, Mackey, ELA 2009

A2 = [1.0  0.0  0.0  0.0  0.0  0.0
0.0  1.0  0.0  0.0  0.0  0.0
0.0  0.0  1.0  0.0  0.0  0.0
0.0  0.0  0.0  1.0  0.0  0.0
0.0  0.0  0.0  0.0  1.0  0.0
0.0  0.0  0.0  0.0  0.0  1.0];

E2 = [0.0  0.0  1.0  0.0  0.0  0.0
0.0  0.0  0.0  1.0  0.0  0.0
0.0  0.0  0.0  0.0  1.0  0.0
0.0  0.0  0.0  0.0  0.0  1.0
0.0  0.0  0.0  0.0  0.0  0.0
0.0  0.0  0.0  0.0  0.0  0.0];

B2 = [0.0  0.0
0.0  0.0
0.0  1.0
1.0  0.0
1.0  0.0
0.0  0.0];

C2 = [-1.0   0.0  0.0  0.0  0.0  0.0
0.0  -1.0  0.0  0.0  0.0  0.0];

D2 = [0.0  0.0
0.0  1.0];

sys = dss(A2,E2,B2,C2,D2);

# compute minimal realization
@time sys1 = gminreal(sys, fast = fast)
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 3 #&& nuo == 0 && nse == 1
# an order reduction without enforcing controllability and observability may not be possible
@time sys1 = gminreal(sys,contr=false,obs=false, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 0
# compute an irreducible realization which still contains a non-dynamic mode
@time sys1 = gminreal(sys,noseig=false, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 2

sys = dss(E2,A2,B2,C2,D2);
# compute minimal realization for a standard system (i.e., irreducible realization)
@time sys1 = gminreal(sys, fast = fast)
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 2

sys = dss(A2,E2,B2,C2,D2);
# compute irreducible realization which still contains a non-dynamic mode using only orthogonal transformations
sys1 = gir(sys, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 2
@time sys1, Q, Z =  gir_lrtran(sys, fast = fast, ltran = true, rtran = true);
L = Q[:,1:sys1.nx]'; R = Z[:,1:sys1.nx];
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 2 &&
      L*sys.A*R ≈ sys1.A && L*sys.E*R ≈ sys1.E && L*sys.B ≈ sys1.B && sys.C*R ≈ sys1.C

# minimal realization requires elimination of non-dynamic modes
sys1 = gir(sys, noseig=true, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 3
@time sys1, Q, Z =  gir_lrtran(sys, noseig=true, fast = fast, ltran = true, rtran = true);
L = Q[:,1:sys1.nx]'; R = Z[:,1:sys1.nx];
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 3 &&
      L*sys.A*R ≈ sys1.A && L*sys.E*R ≈ sys1.E && L*sys.B ≈ sys1.B && sys.C*R ≈ sys1.C

# order reduction may results even when applying the infinite controllability/observability algorithm
sys1 = gir(sys, finite = false, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 2
# order reduction may results even when applying the finite controllability/observability algorithm
sys1 = gir(sys,infinite = false, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)   && order(sys)-order(sys1) == 1
# an order reduction without enforcing controllability and observability may not be possible
sys1 = gir(sys,contr=false,obs=false, noseig=true, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)   && order(sys)-order(sys1) == 0

sys = dss(E2,A2,B2,C2,D2);
# compute minimal realization for a standard system (i.e., irreducible realization)
sys1 = gir(sys, fast = fast)
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 2
# an order reduction without enforcing controllability is not be possible
sys1 = gir(sys,contr=false,fast = fast)
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 0

# Example Van Dooren & Dewilde, LAA 1983.
# P = zeros(3,3,3)
# P[:,:,1] = [1 2 -2; 0 -1 -2; 0 0 0]
# P[:,:,2] = [1 3 0; 1 4 2; 0 -1 -2]
# P[:,:,3] = [1 4 2; 0 0 0; 1 4 2]

# observable realization with A2 = I
A2 = [ 1.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
0.0  1.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
0.0  0.0  1.0  0.0  0.0  0.0  0.0  0.0  0.0
0.0  0.0  0.0  1.0  0.0  0.0  0.0  0.0  0.0
0.0  0.0  0.0  0.0  1.0  0.0  0.0  0.0  0.0
0.0  0.0  0.0  0.0  0.0  1.0  0.0  0.0  0.0
0.0  0.0  0.0  0.0  0.0  0.0  1.0  0.0  0.0
0.0  0.0  0.0  0.0  0.0  0.0  0.0  1.0  0.0
0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  1.0];

E2 = [ 0.0  0.0  0.0  1.0  0.0  0.0  0.0  0.0  0.0
0.0  0.0  0.0  0.0  1.0  0.0  0.0  0.0  0.0
0.0  0.0  0.0  0.0  0.0  1.0  0.0  0.0  0.0
0.0  0.0  0.0  0.0  0.0  0.0  1.0  0.0  0.0
0.0  0.0  0.0  0.0  0.0  0.0  0.0  1.0  0.0
0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  1.0
0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0];

B2 = [ 0.0   0.0   0.0
0.0   0.0   0.0
0.0   0.0   0.0
1.0   3.0   0.0
1.0   4.0   2.0
0.0  -1.0  -2.0
1.0   4.0   2.0
0.0   0.0   0.0
1.0   4.0   2.0];

C2 = [ -1.0   0.0   0.0  0.0  0.0  0.0  0.0  0.0  0.0
0.0  -1.0   0.0  0.0  0.0  0.0  0.0  0.0  0.0
0.0   0.0  -1.0  0.0  0.0  0.0  0.0  0.0  0.0];

D2 = [ 1.0   2.0  -2.0
0.0  -1.0  -2.0
0.0   0.0   0.0];

sys = dss(A2,E2,B2,C2,D2);
# build a strong (least order) minimal realization
@time sys1  = gminreal(sys, fast = fast);
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 6

# the system is observable
@time sys1 = gminreal(sys,obs=false, fast = fast);
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 6

# irreducible realization is not minimal
@time sys1 = gir(sys, fast = fast);
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 5

# minimal realization requires removing of non-dynamic modes
@time sys1 = gir(sys, atol = 1.e-7, obs=false, finite = false, noseig = true, fast = fast);
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 6
sys1, Q, Z =  gir_lrtran(sys, atol = 1.e-7, obs=false, finite = false, noseig=true, fast = fast, ltran = true, rtran = true);
L = Q[:,1:sys1.nx]'; R = Z[:,1:sys1.nx];
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 6 &&
      L*sys.A*R ≈ sys1.A && L*sys.E*R ≈ sys1.E && L*sys.B ≈ sys1.B && sys.C*R ≈ sys1.C

# minimal realization requires removing of non-dynamic modes
@time sys1 = gir(sys, atol = 1.e-7, obs=false, finite = false, noseig = false, fast = fast);
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 5
sys1, Q, Z =  gir_lrtran(sys, atol = 1.e-7, obs=false, finite = false, noseig=false, fast = fast, ltran = true, rtran = true);
L = Q[:,1:sys1.nx]'; R = Z[:,1:sys1.nx];
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 5 &&
      L*sys.A*R ≈ sys1.A && L*sys.E*R ≈ sys1.E && L*sys.B ≈ sys1.B && sys.C*R ≈ sys1.C && Q*Q' ≈ I && Z*Z' ≈ I

# Example 1 - (Varga, Kybernetika, 1990)
A2 = [
1 0 0 0 -1 0 0 0
0 1 0 0 0 -1 0 0
0 0 1 0 0 0 0 0
0 0 0 1 0 0 0 0
0 0 0 0 -1 0 0 0
0 0 0 0 0 -1 0 0
0 0 0 0 3 0 1 0
0 0 0 0 0 2 0 1
]
E2 = [
0 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0
1 0 0 0 0 0 0 0
0 1 0 0 0 0 0 0
0 0 1 0 0 0 0 0
0 0 0 1 0 0 0 0
]
B2 = [
      -1 1
      0 0
      0 0
      0 0
      1 -2
      -2 3
      0 0
      3 -3
]
C2 = [
      0 0 0 0 0 0 -1 0
      0 0 0 0 0 0 0 -1
]
D2 = zeros(Int,2,2);

sys = dss(A2,E2,B2,C2,D2);
# build a strong (least order) minimal realization
@time sys1  = gminreal(sys, fast = fast);
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 7

# irreducible realization is not minimal
@time sys1 = gir(sys, atol = 1.e-7, fast = fast);
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 6
@time sys1, Q, Z =  gir_lrtran(sys, atol = 1.e-7, noseig=true, fast = fast, ltran = true, rtran = true)
L = Q[:,1:sys1.nx]'; R = Z[:,1:sys1.nx];
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 7 &&
      L*sys.A*R ≈ sys1.A && L*sys.E*R ≈ sys1.E && L*sys.B ≈ sys1.B && sys.C*R ≈ sys1.C

@time sys1, Q, Z =  gir_lrtran(sys, atol = 1.e-7, noseig=false, fast = fast, ltran = true, rtran = true);
L = Q[:,1:sys1.nx]'; R = Z[:,1:sys1.nx];
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 6 &&
      L*sys.A*R ≈ sys1.A && L*sys.E*R ≈ sys1.E && L*sys.B ≈ sys1.B && sys.C*R ≈ sys1.C && Q*Q' ≈ I && Z*Z' ≈ I

# SISO standard system, B and D vectors
A2 = [
      0 0 0 -24 0 0 0 0 0 0 0
      1 0 0 -50 0 0 0 0 0 0 0
      0 1 0 -35 0 0 0 0 0 0 0
      0 0 1 -10 0 0 0 0 0 0 0
      0 0 0 0 0 0 0 -30 0 0 0
      0 0 0 0 1 0 0 -61 0 0 0
      0 0 0 0 0 1 0 -41 0 0 0
      0 0 0 0 0 0 1 -11 0 0 0
      0 0 0 0 0 0 0 0 0 0 -15
      0 0 0 0 0 0 0 0 1 0 -23
      0 0 0 0 0 0 0 0 0 1 -9
]
E2 = I;
B2 = [18; 42; 30; 6; 10;17;8;1;0;-10;-2;]
C2 = [0 0 0 0 0 0 0 1 0 0 0]
D2 = [0]

sys = dss(A2,B2,C2,D2);
# build a strong (least order) minimal realization
@time sys1  = gminreal(sys, fast = fast, atol=1.e-7);
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 10

# irreducible realization is also minimal
@time sys1 = gir(sys, atol = 1.e-7, fast = fast);
@test iszero(sys-sys1,atol=1.e-7) && order(sys)-order(sys1) == 10
sys1, Q, Z =  gir_lrtran(sys, atol = 1.e-7, fast = fast, ltran = true, rtran = true)
L = Q[:,1:sys1.nx]'; R = Z[:,1:sys1.nx];
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == 10 &&
      L*sys.A*R ≈ sys1.A && L*sys.E*R ≈ sys1.E && L*sys.B ≈ sys1.B && sys.C*R ≈ sys1.C && Q*Q' ≈ I && Z*Z' ≈ I

for Ty in (Float64, Complex{Float64})

#fast = true; Ty = Complex{Float64}
n = 10; m = 5; p = 6;
nuc = 3; nuo = 4;
sys = rss(n, p, m; T = Ty, nuc = 3, nuo = 4);

@time sys1  = gminreal(sys, atol = 1.e-7, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuc+nuo

@time sys1  = gir(sys, atol = 1.e-7, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuc+nuo
@time sys1, Q, Z =  gir_lrtran(sys, atol = 1.e-7, fast = fast, ltran = true, rtran = true)
L = Q[:,1:sys1.nx]'; R = Z[:,1:sys1.nx];
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == nuc+nuo &&
      L*sys.A*R ≈ sys1.A && L*sys.E*R ≈ sys1.E && L*sys.B ≈ sys1.B && sys.C*R ≈ sys1.C

@time sys1  = gminreal(sys, atol = 1.e-7, contr = false, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuo

@time sys1  = gir(sys, atol = 1.e-7, contr = false, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuo

@time sys1  = gminreal(sys, atol = 1.e-7, obs = false, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuc

@time sys1  = gir(sys, atol = 1.e-7, obs = false, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuc

n = 10; m = 5; p = 6;
nfuc = 3; nfuo = 4; Ty = Float64; fast=true;
iduc = ones(Int,3); ; iduo = [ones(Int,3); 2*ones(Int,2)];
nuc = nfuc+sum(iduc); nuo = nfuo+sum(iduo);
sys = rdss(n, p, m; T = Ty, nfuc, nfuo, iduc, iduo);

@time sys1  = gminreal(sys, atol = 1.e-7, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuc+nuo

@time sys1  = gir(sys, atol = 1.e-7, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuc+nuo
@time sys1, Q, Z =  gir_lrtran(sys,  atol = 1.e-7, fast = fast, ltran = true, rtran = true)
L = Q[:,1:sys1.nx]'; R = Z[:,1:sys1.nx];
@test iszero(sys-sys1,atol=1.e-7)  && order(sys)-order(sys1) == nuc+nuo  &&
      L*sys.A*R ≈ sys1.A && L*sys.E*R ≈ sys1.E && L*sys.B ≈ sys1.B && sys.C*R ≈ sys1.C

@time sys1  = gminreal(sys, atol = 1.e-7, contr = false, noseig = false, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuo

@time sys1  = gir(sys, atol = 1.e-7, contr = false, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuo

@time sys1  = gminreal(sys, atol = 1.e-7, obs = false, noseig = false, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuc

@time sys1  = gir(sys, atol = 1.e-7, obs = false, fast = fast);
@test iszero(sys-sys1, atol = 1.e-7) && order(sys)-order(sys1) == nuc

end

end
end

@testset "gbalmr" begin

sys = rss(0,0,0);
@time sysr, hsv = gbalmr(sys)
@test iszero(sys-sysr,atol=1.e-7) && hsv == Float64[]

sys = rdss(0,0,0);
@time sysr, hsv = gbalmr(sys)
@test iszero(sys-sysr,atol=1.e-7) && hsv == Float64[]

n = 5; m = 3; p = 2;

Ty = Float64; fast = true;
for Ty in (Float64, Complex{Float64})

for fast in (true, false)

# standard continuous-time
sys = rss(n,p,m,stable = true, T = Ty);

@time sysr, hsv = gbalmr(sys,fast = fast, balance = true)
@test iszero(sys-sysr,atol=1.e-7)

@time sysr, hsv = gbalmr(sys,fast = fast, balance = true, matchdc = true, ord = 3)
@test dcgain(sys) ≈ dcgain(sysr)

@time sysr, hsv = gbalmr(sys,fast = fast, matchdc = true, ord = 3)
@test dcgain(sys) ≈ dcgain(sysr)

try
    @time sysr, hsv = gbalmr(sys',fast = fast)
    @test false
catch
    @test true
end

@time sysr, hsv = gbalmr(sys-sys,fast = fast, atolhsv = 1.e-7)
@test hsv[1] < 1.e-7 && order(sysr) == 0

@time sysr, hsv = gbalmr(sys+sys,fast = fast, atolhsv = 1.e-7, balance = true)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n && iszero(2*sys-sysr, atol=1.e-7)

@time sysr, hsv = gbalmr(sys+sys,fast = fast, atolhsv = 1.e-7)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n && iszero(2*sys-sysr, atol=1.e-7)

@time sysr, hsv = gbalmr([sys sys],fast = fast, atolhsv = 1.e-7, balance = true)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n && iszero([sys sys] - sysr, atol=1.e-7)

@time sysr, hsv = gbalmr([sys sys],fast = fast, atolhsv = 1.e-7)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n && iszero([sys sys] - sysr, atol=1.e-7)

# standard discrete-time
sys = rss(n,p,m,T = Ty, disc = true, stable = true);
@time sysr, hsv = gbalmr(sys,fast = fast, balance = true)
@test iszero(sys-sysr,atol=1.e-7)

@time sysr, hsv = gbalmr(sys,fast = fast, balance = true, matchdc = true, ord = 3)
@test dcgain(sys) ≈ dcgain(sysr)

@time sysr, hsv = gbalmr(sys,fast = fast, matchdc = true, ord = 3)
@test dcgain(sys) ≈ dcgain(sysr)

try
    @time sysr, hsv = gbalmr(sys',fast = fast)
    @test false
catch
    @test true
end

@time sysr, hsv = gbalmr(sys-sys,fast = fast, atolhsv = 1.e-7)
@test hsv[1]< 1.e-7 && order(sysr) == 0

@time sysr, hsv = gbalmr(sys+sys,fast = fast, atolhsv = 1.e-7)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n && iszero(2*sys-sysr,atol=1.e-7)

@time sysr, hsv = gbalmr([sys sys],fast = fast)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n

# descriptor continuous-time non-singular E
sys = rdss(n,p,m,stable = true, T = Ty);

@time sysr, hsv = gbalmr(sys,fast = fast, balance = true)
@test iszero(sys-sysr, atol = 1.e-7)

@time sysr, hsv = gbalmr(sys,fast = fast, balance = true, matchdc = true, ord = 3)
@test dcgain(sys) ≈ dcgain(sysr)

@time sysr, hsv = gbalmr(sys,fast = fast, matchdc = true, ord = 3)
@test dcgain(sys) ≈ dcgain(sysr)

try
    @time sysr, hsv = gbalmr(sys',fast = fast)
    @test false
catch
    @test true
end

@time sysr, hsv = gbalmr(sys-sys,fast = fast, atolhsv = 1.e-7)
@test hsv[1] < 1.e-7 && order(sysr) == 0

@time sysr, hsv = gbalmr(sys+sys,fast = fast, atolhsv = 1.e-7)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n && iszero(2*sys-sysr,atol=1.e-7)

@time sysr, hsv = gbalmr([sys sys],fast = fast, atolhsv = 1.e-7)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n

# descriptor discrete-time non-singular E
sys = rdss(n,p,m,T = Ty, disc = true, stable = true);
@time sysr, hsv = gbalmr(sys,fast = fast, balance = true)
@test iszero(sys-sysr,atol=1.e-7)

@time sysr, hsv = gbalmr(sys,fast = fast, balance = true, matchdc = true, ord = 3)
@test dcgain(sys) ≈ dcgain(sysr)

@time sysr, hsv = gbalmr(sys,fast = fast, matchdc = true, ord = 3)
@test dcgain(sys) ≈ dcgain(sysr)

try
    @time sysr, hsv = gbalmr(sys',fast = fast)
    @test false
catch
    @test true
end

@time sysr, hsv = gbalmr(sys-sys,fast = fast, atolhsv = 1.e-7)
@test hsv[1]< 1.e-7 && order(sysr) == 0

@time sysr, hsv = gbalmr(sys+sys,fast = fast, atolhsv = 1.e-7)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n && iszero(2*sys-sysr,atol=1.e-7)

@time sysr, hsv = gbalmr([sys sys],fast = fast)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n

# proper descriptor continuous-time singular E
sys = rdss(n,p,m,T = Ty, stable = true,id=ones(Int,3));

@time sysr, hsv = gbalmr(sys,atol=1.e-7,balance = true)
@test iszero(sys-sysr,atol=1.e-7)

@time sysr, hsv = gbalmr(sys,fast = fast, balance = true, matchdc = true, ord = 3)
@test dcgain(sys) ≈ dcgain(sysr)

try
    @time sysr, hsv = gbalmr(sys',fast = fast)
    @test false
catch
    @test true
end

@time sysr, hsv = gbalmr(sys+sys, fast = fast, atolhsv = 1.e-7, atol = 1.e-7)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n && iszero(2*sys-sysr,atol=1.e-7)

@time sysr, hsv = gbalmr([sys sys], ord = n, fast = fast, atol = 1.e-7)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n

# proper descriptor discrete-time singular E
sys = rdss(n,p,m,T = Ty, stable = true, disc = true, id=ones(Int,3));

@time sysr, hsv = gbalmr(sys,atol=1.e-7,balance = true)
@test iszero(sys-sysr,atol=1.e-7)

@time sysr, hsv = gbalmr(sys,fast = fast, balance = true, matchdc = true, ord = 3)
@test dcgain(sys) ≈ dcgain(sysr)

try
    @time sysr, hsv = gbalmr(sys',fast = fast)
    @test false
catch
    @test true
end

@time sysr, hsv = gbalmr(sys+sys,fast = fast, atolhsv = 1.e-7,atol = 1.e-7)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n && iszero(2*sys-sysr,atol=1.e-7)

@time sysr, hsv = gbalmr([sys sys],fast = fast,atol = 1.e-7)
@test norm(hsv[n+1:end]) < 1.e-7 && order(sysr) == n

end # fast
end # Ty
end # gbalmr

end
end # module


# # Decaying barotropic QG turbulence over topography
#
#md # This example can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/literated/singlelayerqg_decaying_topography.ipynb).
#
# An example of decaying barotropic quasi-geostrophic turbulence over topography.
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add GeophysicalFlows, Plots, Printf, Random, Statistics"
# ```

# ## Let's begin
# Let's load `GeophysicalFlows.jl` and some other needed packages.
#
using GeophysicalFlows, Plots, Printf, Random

using Statistics: mean

# ## Choosing a device: CPU or GPU

dev = CPU()     # Device (CPU/GPU)
nothing # hide

# ## Numerical parameters and time-stepping parameters

      n = 128            # 2D resolution = n²
stepper = "FilteredRK4"  # timestepper
     dt = 0.05           # timestep
 nsteps = 2000           # total number of time-steps
 nsubs  = 10             # number of time-steps for intermediate logging/plotting (nsteps must be multiple of nsubs)
nothing # hide

# ## Physical parameters

L = 2π        # domain size
nothing # hide

# Define the topographic potential vorticity, ``\eta = f_0 h(x, y)/H``. The topography here is
# an elliptical mount at ``(x, y) = (1, 1)``, and an elliptical depression at ``(x, y) = (-1, -1)``.
σx, σy = 0.4, 0.8
topographicPV(x, y) = 3exp(-(x-1)^2/(2σx^2) -(y-1)^2/(2σy^2)) - 2exp(-(x+1)^2/(2σx^2) -(y+1)^2/(2σy^2))
nothing # hide

# ## Problem setup
# We initialize a `Problem` by providing a set of keyword arguments. Not providing a viscosity
# coefficient `ν` leads to the module's default value: `ν=0`. In this example numerical instability
# due to accumulation of enstrophy in high wavenumbers is taken care with the `FilteredTimestepper`
# we picked. Thus, we choose not to do any dealiasing by providing `aliased_fraction=0`.
#
# The topophic PV is prescribed via keyword argument `eta`.
prob = SingleLayerQG.Problem(dev; nx=n, Lx=L, eta=topographicPV,
                                  dt=dt, stepper=stepper, aliased_fraction=0)
nothing # hide

# and define some shortcuts
sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
x, y = grid.x, grid.y
nothing # hide

# and let's plot the topographic PV. Note that when plotting, we decorate the variable to be
# plotted with `Array()` to make sure it is brought back on the CPU when the variable lives
# on the GPU.
contourf(grid.x, grid.y, Array(params.eta'),
          aspectratio = 1,
            linewidth = 0,
               levels = 10,
                    c = :balance,
                 clim = (-3, 3),
                xlims = (-grid.Lx/2, grid.Lx/2),
                ylims = (-grid.Ly/2, grid.Ly/2),
               xticks = -3:3,
               yticks = -3:3,
               xlabel = "x",
               ylabel = "y",
                title = "topographic PV η=f₀h/H")

# ## Setting initial conditions

# Our initial condition consist of a flow that has power only at wavenumbers with
# ``6 < \frac{L}{2\pi} \sqrt{k_x^2 + k_y^2} < 12`` and initial energy ``E_0``.
# `ArrayType()` function returns the array type appropriate for the device, i.e., `Array` for
# `dev = CPU()` and `CuArray` for `dev = GPU()`.

E₀ = 0.04 # energy of initial condition

K = @. sqrt(grid.Krsq)                             # a 2D array with the total wavenumber

Random.seed!(1234)
qih = ArrayType(dev)(randn(Complex{eltype(grid)}, size(sol)))
@. qih = ifelse(K < 6  * 2π/L, 0, qih)
@. qih = ifelse(K > 12 * 2π/L, 0, qih)
qih *= sqrt(E₀ / SingleLayerQG.energy(qih, vars, params, grid))  # normalize qi to have energy E₀
qi = irfft(qih, grid.nx)

SingleLayerQG.set_q!(prob, qi)
nothing # hide

# Let's plot the initial vorticity field:

p1 = heatmap(x, y, Array(vars.q'),
         aspectratio = 1,
              c = :balance,
           clim = (-8, 8),
          xlims = (-grid.Lx/2, grid.Lx/2),
          ylims = (-grid.Ly/2, grid.Ly/2),
         xticks = -3:3,
         yticks = -3:3,
         xlabel = "x",
         ylabel = "y",
          title = "initial vorticity ∂v/∂x-∂u/∂y",
     framestyle = :box)

p2 = contourf(x, y, Array(vars.ψ'),
        aspectratio = 1,
             c = :viridis,
        levels = range(-0.25, stop=0.25, length=11),
          clim = (-0.25, 0.25),
         xlims = (-grid.Lx/2, grid.Lx/2),
         ylims = (-grid.Ly/2, grid.Ly/2),
        xticks = -3:3,
        yticks = -3:3,
        xlabel = "x",
        ylabel = "y",
         title = "initial streamfunction ψ",
    framestyle = :box)

layout = @layout Plots.grid(1, 2)
p = plot(p1, p2, layout=layout, size = (800, 360))

# ## Diagnostics

# Create Diagnostics -- `energy` and `enstrophy` functions are imported at the top.
E = Diagnostic(SingleLayerQG.energy, prob; nsteps=nsteps)
Z = Diagnostic(SingleLayerQG.enstrophy, prob; nsteps=nsteps)
diags = [E, Z] # A list of Diagnostics types passed to "stepforward!" will  be updated every timestep.
nothing # hide

# ## Output

# We choose folder for outputing `.jld2` files.
filepath = "."
filename = joinpath(filepath, "decayingbetaturb.jld2")
nothing # hide

# Do some basic file management,
if isfile(filename); rm(filename); end
nothing # hide

# and then create Output.
get_sol(prob) = sol # extracts the Fourier-transformed solution
out = Output(prob, filename, (:sol, get_sol))
nothing # hide

# ## Visualizing the simulation

# We define a function that plots the vorticity and streamfunction and
# their corresponding zonal mean structure.

function plot_output(prob)
  q = prob.vars.q
  ψ = prob.vars.ψ
  η = prob.params.eta

  pq = heatmap(x, y, Array(q'),
       aspectratio = 1,
            legend = false,
                 c = :balance,
              clim = (-6, 6),
             xlims = (-grid.Lx/2, grid.Lx/2),
             ylims = (-grid.Ly/2, grid.Ly/2),
            xticks = -3:3,
            yticks = -3:3,
            xlabel = "x",
            ylabel = "y",
             title = "vorticity ∂v/∂x-∂u/∂y",
        framestyle = :box)

  contour!(pq, x, y, Array(η'),
          levels=0.5:0.5:3,
          lw=2, c=:black, ls=:solid, alpha=0.7)

  contour!(pq, x, y, Array(η'),
          levels=-2:0.5:-0.5,
          lw=2, c=:black, ls=:dash, alpha=0.7)

  pψ = contourf(x, y, Array(ψ'),
       aspectratio = 1,
            legend = false,
                 c = :viridis,
            levels = range(-0.75, stop=0.75, length=31),
              clim = (-0.75, 0.75),
             xlims = (-grid.Lx/2, grid.Lx/2),
             ylims = (-grid.Ly/2, grid.Ly/2),
            xticks = -3:3,
            yticks = -3:3,
            xlabel = "x",
            ylabel = "y",
             title = "streamfunction ψ",
        framestyle = :box)

  l = @layout Plots.grid(1, 2)
  p = plot(pq, pψ, layout = l, size = (800, 360))

  return p
end
nothing # hide

# ## Time-stepping the `Problem` forward

# We time-step the `Problem` forward in time.

startwalltime = time()

p = plot_output(prob)

anim = @animate for j = 0:round(Int, nsteps/nsubs)

  if j % (1000 / nsubs) == 0
    cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy])

    log = @sprintf("step: %04d, t: %d, cfl: %.2f, E: %.4f, Q: %.4f, walltime: %.2f min",
      clock.step, clock.t, cfl, E.data[E.i], Z.data[Z.i], (time()-startwalltime)/60)

    println(log)
  end

  p[1][1][:z] = Array(vars.q)
  p[1][:title] = "vorticity, t="*@sprintf("%.2f", clock.t)
  p[2][1][:z] = Array(vars.ψ)

  stepforward!(prob, diags, nsubs)
  SingleLayerQG.updatevars!(prob)
end

mp4(anim, "singlelayerqg_decaying_topography.mp4", fps=12)


using DynamicPolynomials
# include("D:\\Programs\\blockpop\\TSSOS\\src\\TSSOS.jl")
using TSSOS

# unconstrained optimization using the TSSOS hierarchy
@polyvar x[1:3]
f = 1+x[1]^4+x[2]^4+x[3]^4+x[1]*x[2]*x[3]+x[2]
opt,sol,data = tssos_first(f, x, newton=true, reducebasis=true, TS="block", solution=true, QUIET=true)
# optimum = 0.475274778453039

opt,sol,data = tssos_higher!(data, TS="MD", solution=true, QUIET=true)
# optimum = 0.4752747815682939

# constrained optimization using the TSSOS hierarchy
@polyvar x[1:3]
f = 1+x[1]^4+x[2]^4+x[3]^4+x[1]*x[2]*x[3]+x[2]
g = 1-x[1]^2-2*x[2]^2
h = x[2]^2+x[3]^2-1
pop = [f, g, h]
opt,sol,data = tssos_first(pop, x, 2, numeq=1, TS="block", solution=true, QUIET=true)
# optimum = 0.6861925732538557

opt,sol,data = tssos_higher!(data, TS="MD", solution=true, QUIET=true)
# optimum = 0.6861925611033903

# constrained optimization using the CS-TSSOS hierarchy
n = 6
@polyvar x[1:n]
f = 1+sum(x.^4)+x[1]*x[2]*x[3]+x[3]*x[4]*x[5]+x[3]*x[4]*x[6]+x[3]*x[5]*x[6]+x[4]*x[5]*x[6]
pop = [f, 1-sum(x[1:3].^2), 1-sum(x[3:6].^2)]
d = 2
opt,sol,data = cs_tssos_first(pop, x, d, numeq=1, TS="block", solution=true, QUIET=true)
# optimum = 0.7174253329322346

opt,sol,data = cs_tssos_higher!(data, TS="MD", solution=true, QUIET=true)
# optimum = 0.7174253409780793

println("Run successfully!")


using FourierFlows, PyPlot, JLD2, Printf, Random, FFTW

using Random: seed!

import GeophysicalFlows.TwoDTurb
import GeophysicalFlows.TwoDTurb: energy, enstrophy
import GeophysicalFlows: peakedisotropicspectrum

dev = CPU()     # Device (CPU/GPU)

# Parameters
  n = 256
  L = 2π
 nν = 2
  ν = 0.0
 dt = 1e-2
nsteps = 5000
 nsubs = 200

# Files
filepath = "."
plotpath = "./plots_McWilliams1984"
plotname = "snapshots"
filename = joinpath(filepath, "McWilliams1984.jld2")

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

# Initialize problem
prob = TwoDTurb.Problem(; nx=n, Lx=L, ny=n, Ly=L, ν=ν, nν=nν, dt=dt, stepper="FilteredRK4", dev=dev)

sol, cl, vs, gr, filter = prob.sol, prob.clock, prob.vars, prob.grid, prob.timestepper.filter
x, y = gridpoints(gr)

# Initial condition closely following pyqg barotropic example
# that reproduces the results of the paper by McWilliams (1984)
seed!(1234)
k0, E0 = 6, 0.5
zetai  = peakedisotropicspectrum(gr, k0, E0, mask=filter)
TwoDTurb.set_zeta!(prob, zetai)

# Create Diagnostic -- energy and enstrophy are functions imported at the top.
E = Diagnostic(energy, prob; nsteps=nsteps)
Z = Diagnostic(enstrophy, prob; nsteps=nsteps)
diags = [E, Z] # A list of Diagnostics types passed to "stepforward!" will
# be updated every timestep.

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

function plot_output(prob, fig, axs; drawcolorbar=false)
  # Plot the vorticity field and the evolution of energy and enstrophy.
  TwoDTurb.updatevars!(prob)
  sca(axs[1])
  pcolormesh(x, y, vs.zeta)
  clim(-40, 40)
  axis("off")
  axis("square")
  if drawcolorbar==true
    colorbar()
  end

  sca(axs[2])
  cla()
  plot(E.t[1:E.i], E.data[1:E.i]/E.data[1])
  plot(Z.t[1:Z.i], Z.data[1:E.i]/Z.data[1])
  xlabel(L"t")
  ylabel(L"\Delta E, \, \Delta Z")

  pause(0.01)
end

# Step forward

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

startwalltime = time()
while cl.step < nsteps
  stepforward!(prob, diags, nsubs)
  saveoutput(out)

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

  println(log)
  plot_output(prob, fig, axs; drawcolorbar=false)
end
println("finished")
plot_output(prob, fig, axs; drawcolorbar=false)

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


module PSO

using Printf, CUDA
using PV

export initialize, evolve!, pso

"""
    initialize(x_lb::AbstractVector{T}, x_ub::AbstractVector{T},
        d::Integer, np::Integer) where T <: AbstractFloat

Initialize a population of `np` particles with random positions `X` and velocities `V`.
`d` is the dimension (i.e., number of variables). `x_lb` and `x_ub` are lower bound and upper bound
of variables.

Two matrices `X` and `V` of size `d × np` are returned, each column of which denotes a position
or velocity.
"""
function initialize(x_lb::AbstractVector{T}, x_ub::AbstractVector{T},
        d::Integer, np::Integer) where T <: AbstractFloat
    @assert d >= 1 && np >= 1 && length(x_lb) == d && length(x_ub) == d
    dx = x_ub .- x_lb
    X = rand(T, d, np) .* dx .+ x_lb
    V = rand(T, d, np) .* 2 .* dx .- dx
    return X, V
end

@inline function _clamp!(x::AbstractVector, l::AbstractVector, u::AbstractVector)
    @inbounds for i in eachindex(x)
        x[i] = clamp(x[i], l[i], u[i])
    end
    nothing
end

@inline function bounce_back!(x, v, l, u)
    T = eltype(x)
    @inbounds for i in eachindex(x)
        t = x[i] + v[i]
        if t < l[i]
            t = l[i] + rand(T) * (x[i] - l[i])
        elseif t > u[i]
            t = x[i] + rand(T) * (u[i] - x[i])
        end
        x[i] = t
    end
end

"""
    evolve!(f::Function, X::Matrix{T}, V::Matrix{T}, x_lb::AbstractVector{T},
        x_ub::AbstractVector{T}, niter::Integer=100; w, c1, c2, log=false) where T <: AbstractFloat

Evolve particles with initial positions `X` and velocities `V` for at most `niter` iterations
to minimize the function `f`. `x_lb` and `x_ub` are lower bound and upper bound of variables.

`w, c1, c2` are parameters of PSO. If `log` is `true`, then progress information is printed in each
iteration.

The best particle's position, its fitness, and the best fitness history for `niter` iterations are
returned.

This method is designed to work on CPU to validate the correctness of PSO.
"""
function evolve!(f::Function, X::Matrix{T}, V::Matrix{T}, x_lb::AbstractVector{T},
        x_ub::AbstractVector{T}, niter::Integer=100;
        w, c1, c2, log=false) where T <: AbstractFloat
    @assert size(X) == size(V)
    @assert size(X, 1) == length(x_lb) == length(x_ub)
    @assert w_min <= w_max
    np = size(X, 2)
    fitness = zeros(T, np)
    pbest = copy(X)
    pbest_fitness = fill(T(Inf), np)
    gbest = X[:, 1]
    gbest_fitness = Inf
    gbest_fitness_hist = zeros(T, niter)

    if log
        @printf("%-10s%-10s\n", "iter", "sse")
    end

    ws = range(w_max, w_min; length=niter_w_decay)  # inertia weight decay

    w = w_max
    for iter = 1:niter
        # w = iter <= niter_w_decay ? ws[iter] : w_min
        # w *= 0.998
        w = max(w, w_min)
        # evaluate fitness and update pbest
        @inbounds for i = 1:np
            fitness[i] = f(@view X[:, i])
            if fitness[i] <= pbest_fitness[i]
                pbest_fitness[i] = fitness[i]
                pbest[:, i] .= @view X[:, i]
            end
        end
        # update gbest
        pbest_idx = argmin(pbest_fitness)
        @inbounds if pbest_fitness[pbest_idx] <= gbest_fitness
            gbest_fitness = pbest_fitness[pbest_idx]
            gbest .= @view pbest[:, pbest_idx]
        end
        gbest_fitness_hist[iter] = gbest_fitness
        if log
            @printf("%-10d%-10.5e%10f\n", iter, gbest_fitness,w)
        end
        # terminate?
        if iter == niter
            break
        end
        # update velocity/position: parallel
        @inbounds @views for i = 1:np
            V[:, i] .= w .* V[:, i] .+ c1 .* rand() .* (pbest[:, i] .- X[:, i]) .+
                        c2 .* rand() .* (gbest .- X[:, i])
            # X[:, i] .+= V[:, i]
            # _clamp!(X[:, i], x_lb, x_ub)
            bounce_back!(X[:, i], V[:, i], x_lb, x_ub)
        end
    end
    return gbest, gbest_fitness, gbest_fitness_hist
end

@inline function _copy!(dest::AbstractVector, src::AbstractVector)
    @inbounds for i = 1:length(src)
        dest[i] = src[i]
    end
end

"""
    evolve!(evaluator::Function, data::CuMatrix{T}, temperature::T,
        X::CuMatrix{T}, V::CuMatrix{T}, x_lb::CuVecOrMat{T},
        x_ub::CuVecOrMat{T}, niter::Integer=100;
        w, c1, c2, log=false, threads::Integer=160,
        min_delta=0.0f0, patience=niter) where T <: AbstractFloat

This method is designed to work on GPU.
Refer to the CPU-based [`evolve!`](@ref) for details. Note the two parameters for early termination:
- `min_delta`: a threshold to decide whether the best fitness (i.e., fitness of *gbest*) has reduced
- `patience`: if the best fitness has not reduced for continuous `patience` iterations, terminate.

A `NamedTuple` is returned with three fields
- `sse`: SSE of the best solution
- `solution`: best solution (*gbest*)
- `stopped_iter`: the iteration at which PSO is terminated
"""
function evolve!(evaluator::Function, data::CuMatrix{T}, temperature::T,
        X::CuMatrix{T}, V::CuMatrix{T}, x_lb::CuVecOrMat{T},
        x_ub::CuVecOrMat{T}, niter::Integer=100;
        w, c1, c2, log=false, threads::Integer=160,
        min_delta=0.0f0, patience=niter) where T <: AbstractFloat
    np = size(X, 1)
    nv = size(X, 2)
    fitness = CUDA.zeros(T, np)
    pbest = copy(X)
    pbest_fitness = CUDA.fill(T(Inf), np)
    gbest = X[1, :]  # copied
    gbest_fitness = T(Inf)
    gbest_fitness_hist = CUDA.zeros(T, niter)
    R = CUDA.rand(T, 4 * np, niter)   # pre-generate random numbers

    function evaluate!(iter::Integer)
        i = threadIdx().x + (blockIdx().x - 1) * blockDim().x
        if i > np
            return nothing
        end
        @inbounds fitness[i] = fit = evaluator(data, temperature, @view X[i, :])
        @inbounds if fit <= pbest_fitness[i]
            pbest_fitness[i] = fit
            @views _copy!(pbest[i, :], X[i, :])
        end
        nothing
    end

    function update_vel_pos!(iter)
        i = threadIdx().x + (blockIdx().x - 1) * blockDim().x
        if i > np
            return nothing
        end
        # update velocity and position
        @inbounds r1 = R[i, iter]
        @inbounds r2 = R[i + np, iter]
        @inbounds for j = 1:nv  # update each component individually
            V[i, j] = new_v = w * V[i, j] + c1 * r1 * (pbest[i, j] - X[i, j]) + c2 * r2 * (gbest[j] - X[i, j])
            old_x = X[i, j]
            t = old_x  + new_v
            l = x_lb[j]
            u = x_ub[j]
            if t < l
                t = l + R[i + 2 * np, iter] * (old_x - l)
            elseif t > u
                t = old_x + R[i + 3 * np, iter] * (u - old_x)
            end
            X[i, j] = t
        end
        nothing
    end

    # kernel config
    blocks = np ÷ threads + (np % threads > 0)
    if log
        @show threads, blocks
        @printf("%-10s%-10s\n", "iter", "sse")
    end

    wait = 0
    stopped_iter = 0

    for iter = 1:niter
        # fitness
        @cuda threads = threads blocks = blocks evaluate!(iter)
        # update gbest： parallel reduction
        pbest_idx = argmin(pbest_fitness)
        min_pbest_fitness = pbest_fitness[pbest_idx]
        if min_pbest_fitness <= gbest_fitness
            if gbest_fitness - min_pbest_fitness >= min_delta
                wait = 0
            else
                wait += 1
            end
            gbest_fitness = min_pbest_fitness
            @inbounds _copy!(gbest, @view pbest[pbest_idx, :])
        end
        if log
            @printf("%-10d%-10.5e\n", iter, gbest_fitness)
        end
        # terminate?
        if wait >= patience
            stopped_iter = iter
            if log
                @printf("Early stopped at iteration %d\n", iter)
            end
            break
        end
        if iter == niter
            stopped_iter = niter
            break
        end

        # update velocity/position: parallel
        @cuda threads = threads blocks = blocks update_vel_pos!(iter)
    end
    return (sse = gbest_fitness, solution = Array(gbest), stopped_iter = stopped_iter)
end

"""
    pso(data::AbstractMatrix{T}, temperature::T, model::Symbol, x_lb::AbstractVector{T},
        x_ub::AbstractVector{T}; niter::Integer=400, np::Integer=1600,
        w::T, c1::T, c2::T, log=false, threads::Integer=160,
        min_delta::T=0.0, patience=niter + 1) where T <: AbstractFloat

Parallel PSO for PV parameter identification. Each column of `data` is a V-I point. `model` can be either
`:SDM` or `:DDM`. `x_lb` and `x_ub` denote the lower bound and upper bound of each variable to
confine the search range.
- `niter`: number of iterations
- `np`: number of particles
- `w, c1, c2`: PSO velocity update parameters
- `log`: print progress information if `true`
- `threads`: number of threads per block in CUDA kernel
- `min_delta, patience`: early stopping. Similar to *keras* API:
    https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/EarlyStopping
"""
function pso(data::AbstractMatrix{T}, temperature::T, model::Symbol, x_lb::AbstractVector{T},
        x_ub::AbstractVector{T}; niter::Integer=400, np::Integer=1600,
        w::T, c1::T, c2::T, log=false, threads::Integer=160,
        min_delta::T=0.0, patience=niter + 1) where T <: AbstractFloat
    @assert model in (:SDM, :DDM)
    nv = model === :SDM ? 5 : 7
    @assert length(x_lb) == length(x_ub) == nv
    @assert size(data, 1) == 2
    @assert threads % 32 == 0

    # initialize: each row is a particle (for GPU)
    dx = reshape(x_ub .- x_lb, 1, :) |> CuArray
    x_lb = reshape(x_lb, 1, :) |> CuArray
    x_ub = reshape(x_ub, 1, :) |> CuArray
    X = CUDA.rand(T, np, nv) .* dx .+ x_lb
    V = CUDA.rand(T, np, nv) .* 2 .* dx .- dx

    # evolve particles
    if model === :SDM
        return evolve!(PV.sse_sdm, data |> CuArray, temperature,
            X, V, x_lb |> CuArray, x_ub |> CuArray, niter;
            w, c1, c2, log, threads, min_delta, patience)
    else
        return evolve!(PV.sse_ddm, data |> CuArray, temperature,
            X, V, x_lb |> CuArray, x_ub |> CuArray, niter;
            w, c1, c2, log, threads, min_delta, patience)
    end
end

end # module


# implements a rollout policy
module RolloutPolicies

using DiscreteMDPs

export RolloutPolicy
export action

import DiscreteMDPs.Policy
import DiscreteMDPs.action

type RolloutPolicy <: Policy
    d::Int64               # search depth
    n::Int64               # number of iterations
    c::Float64             # explortaion constant
    Q::Vector{Float64}     # Q values at the root state
    N::Vector{Int64}       # Number of visits to each state
    actions::Vector{Int64} # Vector of allowable actions
    mdp::DiscreteMDP
end

function RolloutPolicy(mdp::DiscreteMDP; d::Int64=20, n::Int64=100, c::Float64=0.1)
    nActions = numActions(mdp)
    Q = zeros(nActions)
    N = zeros(Int64, nActions)
    actions = [1:nActions]
    return RolloutPolicy(d, n, c, Q, N, actions, mdp)
end

function action(policy::RolloutPolicy, sIdx::Int64)
    n = policy.n
    fill!(policy.Q, 0.0)
    fill!(policy.N, 0)
    for i = 1:n
        simulate(policy, sIdx)
    end
    a = policy.actions
    return a[indmax(policy.Q)]
end

function simulate(policy::RolloutPolicy, s::Int64)
    Q = policy.Q
    N = policy.N
    c = policy.c
    d = policy.d
    actions = policy.actions
    mdp = policy.mdp
    discount = mdp.discount
    i = indmax(Q + c .* real(sqrt(complex(log(sum(N)) ./ N))))
    a = actions[i]
    sp = nextState(mdp, s, a)
    q = reward(mdp, s, a) + discount*rollout(policy, sp, d)
    N[i] += 1
    Q[i] += (q - Q[i]) / N[i]
    return
end

function rollout(policy::RolloutPolicy, s::Int64, depth::Int64)
    mdp = policy.mdp
    discount = mdp.discount
    if depth == 0
        return 0.0
    else
        a  = getAction(mdp, s)
        sp = nextState(mdp, s, a)
        return reward(mdp, s, a) + discount*rollout(policy, sp, depth-1)
    end
end

end # module


@doc doc"""
    Metric

Abstract type for the pseudo-Riemannian metric tensor $g$, a family of smoothly
varying inner products on the tangent space. See [`inner`](@ref).
"""
abstract type Metric end

@doc doc"""
    RiemannianMetric <: Metric

Abstract type for Riemannian metrics, a family of positive definite inner
products. The positive definite property means that for $v \in T_x M$, the
inner product $g(v, v) > 0$ whenever $v$ is not the zero vector.
"""
abstract type RiemannianMetric <: Metric end

@doc doc"""
    LorentzMetric <: Metric

Abstract type for Lorentz metrics, which have a single time dimension. These
metrics assume the spacelike convention with the time dimension being last,
giving the signature $(++...+-)$.
"""
abstract type LorentzMetric <: Metric end

"""
    MetricManifold{M<:Manifold,G<:Metric} <: Manifold

Equip a [`Manifold`](@ref) explicitly with a [`Metric`](@ref) `G`.

For a Metric Manifold, by default, assumes, that you implement the linear form
from [`local_metric`](@ref) in order to evaluate the exponential map.

If the corresponding [`Metric`](@ref) `G` yields closed form formulae for e.g.
the exponential map and this is implemented directly (without solving the ode),
you can of course still implement that directly.

# Constructor
    MetricManifold(manifold, metric)

    """
struct MetricManifold{M<:Manifold,G<:Metric} <: Manifold
    manifold::M
    metric::G
end

is_decorator_manifold(M::MMT) where {MMT <: MetricManifold}= Val(true)

"""
    is_default_metric(M,G)

indicate, whether the [`Metric`](@ref) `G` is the default metric for
the [`Manifold`](@ref) `M`. This means that any occurence of
[`MetricManifold`](@ref)(M,G) where `typeof(is_default_metric(M,G)) = Val{true}`
falls back to just be called with `M` such that the [`Manifold`](@ref) `M`
implicitly has this metric, for example if this was the first one implemented
or is the one most commonly assumed to be used.
"""
is_default_metric(M::Manifold,G::Metric) = Val(false)
"""
    is_default_metric(MM)

indicate, whether the [`Metric`](@ref) `MM.G` is the default metric for
the [`Manifold`](@ref) `MM.M` within the [`MetricManifold`](@ref) `MM`.
This means that any occurence of
[`MetricManifold`](@ref)`(MM.M,MM.G)` where `typeof(is_default_metric(MM.M,MM.G)) = Val{true}`
falls back to just be called with `MM.MM` such that the [`Manifold`](@ref) `MM.M`
implicitly has the metric `MM.G`, for example if this was the first one
implemented or is the one most commonly assumed to be used.
"""
is_default_metric(M::MMT) where {MMT <: MetricManifold} = is_default_metric(base_manifold(M),metric(M))
# this automatically undecorates
# this should automatically decorate at least for simple cases
convert(T::Type{MetricManifold{MT,GT}},M::MT) where {MT,GT} = _convert_with_default(M,GT,is_default_metric(M,GT()))
_convert_with_default(M::MT,T::Type{<:Metric},::Val{true}) where {MT <: Manifold} = MetricManifold(M,T())
_convert_with_default(M::MT,T::Type{<:Metric},::Val{false}) where {MT <: Manifold} = error("Can not convert $(M) to a MetricManifold{$(MT),$(T)}, since $(T) is not the default metric.")

@doc doc"""
    metric(M::MetricManifold)

Get the metric $g$ of the manifold `M`.
"""
metric(M::MetricManifold) = M.metric

@doc doc"""
    local_metric(M::MetricManifold, x)

Local matrix representation at the point `x` of the metric tensor $g$ on the
manifold `M`, usually written $g_{ij}$. The matrix has the property that
$g(v, w)=v^T [g_{ij}] w = g_{ij} v^i w^j$, where the latter expression uses
Einstein summation convention.
"""
function local_metric(M::MetricManifold, x)
    error("Local metric not implemented on $(typeof(M)) for point $(typeof(x))")
end

@doc doc"""
    inverse_local_metric(M::MetricManifold, x)

Local matrix representation of the inverse metric (cometric) tensor, usually
written $g^{ij}$
"""
inverse_local_metric(M::MetricManifold, x) = inv(local_metric(M, x))

@doc doc"""
    det_local_metric(M::MetricManifold, x)

Determinant of local matrix representation of the metric tensor $g$
"""
det_local_metric(M::MetricManifold, x) = det(local_metric(M, x))

@doc doc"""
    log_local_metric_density(M::MetricManifold, x)

Return the natural logarithm of the metric density $\rho$ of `M` at `x`, which
is given by $\rho=\log \sqrt{|\det [g_{ij}]|}$.
"""
log_local_metric_density(M::MetricManifold, x) = log(abs(det_local_metric(M, x))) / 2

@doc doc"""
    local_metric_jacobian(M::MetricManifold, x)

Get partial derivatives of the local metric of `M` at `x` with respect to the
coordinates of `x`, $\frac{\partial}{\partial x^k} g_{ij} = g_{ij,k}$. The
dimensions of the resulting multi-dimensional array are ordered $(i,j,k)$.
"""
function local_metric_jacobian(M, x)
    error("local_metric_jacobian not implemented on $(typeof(M)) for point $(typeof(x)). For a suitable default, enter `using ForwardDiff`.")
end

@doc doc"""
    christoffel_symbols_first(M::MetricManifold, x)

Compute the Christoffel symbols of the first kind in local coordinates.
The Christoffel symbols are (in Einstein summation convention)

$\Gamma_{ijk} = \frac{1}{2} \left[g_{kj,i} + g_{ik,j} - g_{ij,k}\right],$

where $g_{ij,k}=\frac{\partial}{\partial x^k} g_{ij}$ is the coordinate
derivative of the local representation of the metric tensor. The dimensions of
the resulting multi-dimensional array are ordered $(i,j,k)$.
"""
function christoffel_symbols_first(M::MetricManifold, x)
    ∂g = local_metric_jacobian(M, x)
    n = size(∂g, 1)
    Γ = similar(∂g, Size(n, n, n))
    @einsum Γ[i,j,k] = 1 / 2 * (∂g[k,j,i] + ∂g[i,k,j] - ∂g[i,j,k])
    return Γ
end

@doc doc"""
    christoffel_symbols_second(M::MetricManifold, x)

Compute the Christoffel symbols of the second kind in local coordinates.
The Christoffel symbols are (in Einstein summation convention)

$\Gamma^{l}_{ij} = g^{kl} \Gamma_{ijk},$

where $\Gamma_{ijk}$ are the Christoffel symbols of the first kind, and
$g^{kl}$ is the inverse of the local representation of the metric tensor.
The dimensions of the resulting multi-dimensional array are ordered $(l,i,j)$.
"""
function christoffel_symbols_second(M::MetricManifold, x)
    ginv = inverse_local_metric(M, x)
    Γ₁ = christoffel_symbols_first(M, x)
    Γ₂ = similar(Γ₁)
    @einsum Γ₂[l,i,j] = ginv[k,l] * Γ₁[i,j,k]
    return Γ₂
end

@doc doc"""
    christoffel_symbols_second_jacobian(M::MetricManifold, x)

Get partial derivatives of the Christoffel symbols of the second kind
for manifold `M` at `x` with respect to the coordinates of `x`,
$\frac{\partial}{\partial x^l} \Gamma^{k}_{ij} = \Gamma^{k}_{ij,l}.$
The dimensions of the resulting multi-dimensional array are ordered $(i,j,k,l)$.
"""
function christoffel_symbols_second_jacobian(M, x)
    error("christoffel_symbols_second_jacobian not implemented on $(typeof(M)) for point $(typeof(x)). For a suitable default, enter `using ForwardDiff`.")
end

@doc doc"""
    riemann_tensor(M::MetricManifold, x)

Compute the Riemann tensor $R^l_{ijk}$, also known as the Riemann curvature
tensor, at the point `x`. The dimensions of the resulting multi-dimensional
array are ordered $(l,i,j,k)$.
"""
function riemann_tensor(M::MetricManifold, x)
    n = size(x, 1)
    Γ = christoffel_symbols_second(M, x)
    ∂Γ = christoffel_symbols_second_jacobian(M, x) ./ n
    R = similar(∂Γ, Size(n, n, n, n))
    @einsum R[l,i,j,k] = ∂Γ[l,i,k,j] - ∂Γ[l,i,j,k] + Γ[s,i,k] * Γ[l,s,j] - Γ[s,i,j] * Γ[l,s,k]
    return R
end

"""
    ricci_tensor(M::MetricManifold, x)

Compute the Ricci tensor, also known as the Ricci curvature tensor,
of the manifold `M` at the point `x`.
"""
function ricci_tensor(M::MetricManifold, x)
    R = riemann_tensor(M, x)
    n = size(R, 1)
    Ric = similar(R, Size(n, n))
    @einsum Ric[i,j] = R[l,i,l,j]
    return Ric
end

"""
    ricci_curvature(M::MetricManifold, x)

Compute the Ricci scalar curvature of the manifold `M` at the point `x`.
"""
function ricci_curvature(M::MetricManifold, x)
    ginv = inverse_local_metric(M, x)
    Ric = ricci_tensor(M, x)
    S = sum(ginv .* Ric)
    return S
end

"""
    gaussian_curvature(M::MetricManifold, x)

Compute the Gaussian curvature of the manifold `M` at the point `x`.
"""
gaussian_curvature(M::MetricManifold, x) = ricci_curvature(M, x) / 2

"""
    einstein_tensor(M::MetricManifold, x)

Compute the Einstein tensor of the manifold `M` at the point `x`.
"""
function einstein_tensor(M::MetricManifold, x)
    Ric = ricci_tensor(M, x)
    g = local_metric(M, x)
    ginv = inverse_local_metric(M, x)
    S = sum(ginv .* Ric)
    G = Ric - g .* S / 2
    return G
end

@doc doc"""
    solve_exp_ode(M::MetricManifold, x, v, tspan; solver=AutoVern9(Rodas5()), kwargs...)

Approximate the exponential map on the manifold over the provided timespan
assuming the Levi-Civita connection by solving the ordinary differential
equation

$\frac{d^2}{dt^2} x^k + \Gamma^k_{ij} \frac{d}{dt} x_i \frac{d}{dt} x_j = 0,$

where $\Gamma^k_{ij}$ are the Christoffel symbols of the second kind, and
the Einstein summation convention is assumed. The arguments `tspan` and
`solver` follow the `OrdinaryDiffEq` conventions. `kwargs...` specify keyword
arguments that will be passed to `OrdinaryDiffEq.solve`.

Currently, the numerical integration is only accurate when using a single
coordinate chart that covers the entire manifold. This excludes coordinates
in an embedded space.
"""
function solve_exp_ode(M, x, v, tspan; kwargs...)
    error("solve_exp_ode not implemented on $(typeof(M)) for point $(typeof(x)), vector $(typeof(v)), and timespan $(typeof(tspan)). For a suitable default, enter `using OrdinaryDiffEq`.")
end
exp(M::MMT, x, v, T::AbstractVector{T} where T) where {MMT <: MetricManifold} = exp(M, is_default_metric(M), x, v, T)
exp(M::MMT, ::Val{true}, x, v, T::AbstractVector{T} where T) where {MMT<:MetricManifold} = exp(base_manifold(M), x, v, T)
function exp(M::MMT, ::Val{false}, x, v, T::AbstractVector{T} where T) where {MMT<:MetricManifold}
    sol = solve_exp_ode(M, x, v, extrema(T); dense=false, saveat=T)
    n = length(x)
    return map(i -> sol.u[i][n+1:end], 1:length(T))
end

exp!(M::MMT, y, x, v) where {MMT <: MetricManifold}= exp!(M, is_default_metric(M), y, x, v)
exp!(M::MMT, ::Val{true}, y, x, v) where {MMT<:MetricManifold} = exp!(base_manifold(M),y,x,v)
"""
    exp(MM::MetricManifold, x, v, args...)

Numerically integrate the exponential map assuming the Levi-Civita connection.
See [`solve_exp_ode`](@ref) if the [`Metric`](@ref) `G` of the [`MetricManifold`](@ref){M,G}`
is not marked as the default one.

Currently, the numerical integration is only accurate when using a single
coordinate chart that covers the entire manifold. This excludes coordinates
in an embedded space.
"""
function exp!(M::MMT, ::Val{false}, y, x, v) where {MMT<:MetricManifold}
    tspan = (0.0, 1.0)
    sol = solve_exp_ode(M, x, v, tspan; dense=false, saveat=[1.0])
    n = length(x)
    y .= sol.u[1][n+1:end]
    return y
end

inner(M::MMT, x, v, w) where {MMT <: MetricManifold}= inner(M, is_default_metric(M), x, v, w)
inner(M::MMT, ::Val{false}, x, v, w) where {MMT <: MetricManifold}= dot(v, local_metric(M, x) * w)
inner(M::MMT, ::Val{true}, x, v, w) where {MMT <: MetricManifold} = inner(base_manifold(M), x,v,w)

log!(M::MMT, w, x, y) where {MMT <: MetricManifold} = log!(M, is_default_metric(M), w, x, y)
log!(M::MMT, ::Val{true}, w, x, y) where {MMT <: MetricManifold} = log!(base_manifold(M), w, x, y)
log!(M::MMT, ::Val{false}, w, x, y) where {MMT <: MetricManifold} = error("Logarithmic map not implemented on $(typeof(M)) for points $(typeof(x)) and $(typeof(y)).")

tangent_orthonormal_basis(M::MMT, x, v) where {MMT <: MetricManifold} = tangent_orthonormal_basis(M, is_default_metric(M), x, v)
tangent_orthonormal_basis(M::MMT, ::Val{true}, x, v) where {MMT <: MetricManifold} = tangent_orthonormal_basis(base_manifold(M), x, v)
tangent_orthonormal_basis(M::MMT, ::Val{false}, x, v) where {MMT <: MetricManifold} = error("tangent_orthogonal_basis not implemented on $(typeof(M)) for point $(typeof(x)) and tangent vector $(typeof(v)).")

project_point!(M::MMT, y, x) where {MMT <: MetricManifold} = project_point!(M, is_default_metric(M), y, x)
project_point!(M::MMT, ::Val{true}, y, x) where {MMT <: MetricManifold} = project_point!(base_manifold(M), y, x)
project_point!(M::MMT, ::Val{false}, y, x) where {MMT <: MetricManifold} = error("project_point! not implemented on $(typeof(M)) for point $(typeof(x))")

project_tangent!(M::MMT, w, x, v) where {MMT <: MetricManifold} = project_tangent!(M, is_default_metric(M), w, x, v)
project_tangent!(M::MMT, ::Val{true}, w, x, v) where {MMT <: MetricManifold} = project_tangent!(base_manifold(M), w, x, v)
project_tangent!(M::MMT, ::Val{false}, w, x, v) where {MMT <: MetricManifold} = error("project_tangent! not implemented for a $(typeof(M)) and tangent $(typeof(v)) at point $(typeof(x)).")

vector_transport_to!(M::MMT, vto, x, v, y, m::AbstractVectorTransportMethod) where {MMT <: MetricManifold} = vector_transport_to!(M, is_default_metric(M), vto, x, v, y, m)
vector_transport_to!(M::MMT, ::Val{true}, vto, x, v, y, m)  where {MMT <: MetricManifold}= vector_transport_to!(base_manifold(M), vto, x, v, y, m)
vector_transport_to!(M::MMT, ::Val{false}, vto, x, v, y, m)  where {MMT <: MetricManifold}= error("vector transport from a point of type $(typeof(x)) to a type $(typeof(y)) on a $(typeof(M)) for a vector of type $(v) and the $(typeof(m)) not yet implemented.")

projected_distribution(M::MMT, d, x) where {MMT <: MetricManifold} = projected_distribution(M, is_default_metric(M), d,x)
projected_distribution(M::MMT, ::Val{true}, d, x) where {MMT <: MetricManifold} = projected_distribution(base_manifold(M), d,x)
projected_distribution(M::MMT, ::Val{false}, d, x) where {MMT <: MetricManifold} =error("projected_distribution not implemented for a $(typeof(M)) and with $(typeof(d)) at point $(typeof(x)).")

projected_distribution(M::MMT, d) where {MMT <: MetricManifold} = projected_distribution(M, is_default_metric(M), d)
projected_distribution(M::MMT, ::Val{true}, d) where {MMT <: MetricManifold} = projected_distribution(base_manifold(M), d)
projected_distribution(M::MMT, ::Val{false}, d) where {MMT <: MetricManifold} =error("projected_distribution not implemented for a $(typeof(M)) with $(typeof(d)).")

normal_tvector_distribution(M::MMT, x, σ) where {MMT <: MetricManifold} = normal_tvector_distribution(M, is_default_metric(M), x, σ)
normal_tvector_distribution(M::MMT, ::Val{true}, x, σ) where {MMT <: MetricManifold} = normal_tvector_distribution(base_manifold(M), x, σ)
normal_tvector_distribution(M::MMT, ::Val{false}, x, σ) where {MMT <: MetricManifold} =error("normal_tvector_distribution not implemented for a $(typeof(M)) at point $(typeof(x)) with standard deviation $(typeof(σ)).")

function inner(B::VectorBundleFibers{<:CotangentSpaceType, MMT}, x, v, w) where {MMT<:MetricManifold}
    ginv = inverse_local_metric(B.M, x)
    return dot(v, ginv * w)
end

function flat!(M::MMT, v::FVector{CotangentSpaceType}, x, w::FVector{TangentSpaceType}) where {MMT<:MetricManifold}
    g = local_metric(M, x)
    copyto!(v.data, g*w.data)
    return v
end

function sharp!(M::MMT, v::FVector{TangentSpaceType}, x, w::FVector{CotangentSpaceType}) where {MMT<:MetricManifold}
    ginv = inverse_local_metric(M, x)
    copyto!(v.data, ginv*w.data)
    return v
end

# These are independent of the metric and hence can always fall back to M
injectivity_radius(M::MMT, args...) where {MMT <: MetricManifold} = injectivity_radius(base_manifold(M), args...)
zero_tangent_vector!(M::MMT, v, x) where {MMT <: MetricManifold} = zero_tangent_vector!(base_manifold(M), v, x)
check_manifold_point(M::MMT, x; kwargs...) where {MMT <: MetricManifold} = check_manifold_point(base_manifold(M), x; kwargs...)
check_tangent_vector(M::MMT, x, v; kwargs...) where {MMT <: MetricManifold} = check_tangent_vector(base_manifold(M), x, v; kwargs...)


using Printf
using Statistics
using Random
using Oceananigans
using Oceananigans.Units
using GLMakie

function baroclinic_adjustment(kskew, ksymmetric, tapering, scale; dt = 20minutes, gradient = "y")

# Architecture
architecture = CPU()

# Domain
Lx = 1000kilometers
Ly = 1000kilometers  # north-south extent [m]
Lz = 1kilometers     # depth [m]
Nx = 20
Ny = 20
Nz = 20
save_fields_interval = 0.5day
stop_time = 30days
Δt = dt

zfaces = -Lz .+ (-0.5 * (cos.(range(0.0, π, length=Nz + 1))) .+ 0.5) .* Lz # range(-Lz, 0.0, length = Nz+1)
zfaces = range(-Lz, 0.0, length=Nz + 1)

grid = RectilinearGrid(architecture;
                       topology = (Bounded, Bounded, Bounded),
                       size = (Nx, Ny, Nz),
                       x = (-Lx/2, Lx/2),
                       y = (-Ly/2, Ly/2),
                       z = zfaces,
                       halo = (3, 3, 3))

coriolis = FPlane(latitude = -45)

println("The diffusive timescale is ", (zfaces[1] - zfaces[2])^2 / 1e3)

Δy = Ly/Ny
@show κh = νh = Δy^4 / 10days
vertical_closure = VerticalScalarDiffusivity(ν=1e-2, κ=1e-4)
horizontal_closure = HorizontalScalarBiharmonicDiffusivity(ν=νh, κ=κh)

gerdes_koberle_willebrand_tapering = FluxTapering(tapering)
gent_mcwilliams_diffusivity = IsopycnalSkewSymmetricDiffusivity(κ_skew=kskew,
    κ_symmetric=ksymmetric,
    slope_limiter=gerdes_koberle_willebrand_tapering)

closures = (vertical_closure, horizontal_closure, gent_mcwilliams_diffusivity)

@info "Building a model..."

model = HydrostaticFreeSurfaceModel(grid = grid,
                                    coriolis = coriolis,
                                    buoyancy = BuoyancyTracer(),
                                    closure = closures,
                                    tracers = (:b, :c),
                                    momentum_advection = WENO5(),
                                    tracer_advection = WENO5(),
                                    free_surface = ImplicitFreeSurface())

@info "Built $model."

"""
Linear ramp from 0 to 1 between -Δy/2 and +Δy/2.

For example:

y < y₀           => ramp = 0
y₀ < y < y₀ + Δy => ramp = y / Δy
y > y₀ + Δy      => ramp = 1
"""
function ramp(x, y, Δ)
    gradient == "x" && return min(max(0, x / Δ + 1/2), 1)
    gradient == "y" && return min(max(0, y / Δ + 1/2), 1)
    gradient == "xy" && return 0.5 *( min(max(0, x / Δ + 1/2), 1) + min(max(0, y / Δ + 1/2), 1) )
end

# Parameters
N² = 4e-6 # [s⁻²] buoyancy frequency / stratification
M² = 8e-8 # [s⁻²] horizontal buoyancy gradient

Δy = 100kilometers * 1.0
Δz = 100

Δc = 2Δy
Δb = Δy * M²
ϵb = 0e-2 * Δb # noise amplitude

bᵢ(x, y, z) = N² * z + scale * Δb * ramp(x, y, Δy)
cᵢ(x, y, z) = 00.25 # exp(-x^2 / 2Δc^2) * exp(-y^2 / 2Δc^2) * exp(-(z + Lz/4)^2 / 2Δz^2)

set!(model, b=bᵢ, c=cᵢ)

#####
##### Simulation building
#####

simulation = Simulation(model; Δt, stop_time)
wall_clock = [time_ns()]

function print_progress(sim)
    @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): (%6.3e, %6.3e, %6.3e) m/s, next Δt: %s\n",
            100 * (sim.model.clock.time / sim.stop_time),
            sim.model.clock.iteration,
            prettytime(sim.model.clock.time),
            prettytime(1e-9 * (time_ns() - wall_clock[1])),
            maximum(abs, sim.model.velocities.u),
            maximum(abs, sim.model.velocities.v),
            maximum(abs, sim.model.velocities.w),
            prettytime(sim.Δt))

    wall_clock[1] = time_ns()

    return nothing
end

simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(50))

@info "Running the simulation..."

run!(simulation)

@info "Simulation completed in " * prettytime(simulation.run_wall_time)
println("done with gradient in ", gradient)

loss = extrema(interior(simulation.model.tracers.c)[:, :, :, end])
return loss, simulation
end

klist = [0.0, 1e2, 1e3, 2e3]
taperinglist = [0.0, 1e-2, 2e-2, 0.5e-2, 10.0]
scalelist = [0.5, 1.0, 2.0]
namelist = (; κskew = 1e3, κsymmetric = 1e3, tapering = 1e-2, scale = 2.0)

clist = []
plist = []
for k in klist, tapering in taperinglist, scale in scalelist
    namelist = (; κskew=k, κsymmetric=k, tapering=tapering, scale=scale)
    println("current on ", namelist)
    cextrema, simulation = baroclinic_adjustment(namelist...)
    push!(clist, cextrema)
    push!(plist, namelist)
end

for i in eachindex(clist)
    if (abs(clist[i]) > 0.3) | (abs(clist[i]) < 0.2)
        println("-----------------------")
        println("For parameter values ", plist[i])
        println("the simulation was not good")
        println("----------------------")
    end
end

# (; κskew = 2000.0, κsymmetric = 2000.0, tapering = 0.01, scale = 2.0) failed


struct ROCSAlignmentResult{D,S,T,F<:AbstractAffineMap,X<:AbstractGMM{D,S},Y<:AbstractGMM{D,T}} <: AlignmentResults
    x::X
    y::Y
    minimum::T
    tform::F
end

"""
    com = center_of_mass(gmm)

Returns the center of mass of `gmm`, where its first order moments are equal to 0.
"""
function center_of_mass(positions::AbstractMatrix{<:Real}, weights=ones(eltype(positions),size(positions,2)))
    return SVector{size(positions,1)}([sum(hcat([weights[i]*positions[:,i] for i=1:size(positions,2)]...), dims=2)...]/sum(weights))
end

center_of_mass(gaussians::AbstractVector{<:AbstractIsotropicGaussian}) =
    return center_of_mass(hcat([g.μ for g in gaussians]...), [g.ϕ for g in gaussians])

center_of_mass(gmm::AbstractIsotropicGMM) = center_of_mass(gmm.gaussians)

center_of_mass(mgmm::AbstractMultiGMM) = center_of_mass(collect(Iterators.flatten([gmm.gaussians for (k,gmm) in mgmm])))

"""
    m = second_moment(gmm, center, dim1, dim2)

Returns the second order moment of `gmm`
"""
function mass_matrix(positions::AbstractMatrix{<:Real},
                     widths=zeros(eltype(positions),size(positions,2)),
                     weights=ones(eltype(positions),size(positions,2)),
                     center=center_of_mass(positions,weights))
    t = eltype(positions)
    npts = size(positions,2)
    M = fill(zero(t), 3, 3)
    dists = [positions[:,i].-center for i=1:npts]
    # diagonal terms
    for i=1:3
        for j=1:npts
            M[i,i] += weights[j] * (dists[j][i]^2 + widths[j]^2)
        end
    end
    # off-diagonal terms
    for i=1:3
        for j=i+1:3
            for k=1:npts
                inc = weights[k] * dists[k][i] * dists[k][j]
                M[i,j] += inc
                M[j,i] += inc
            end
        end
    end
    return M
end

"""
    tforms = inertial_transforms(gmm)

Returns 10 transformations to put `gmm` in an inertial frame. That is, the mass matrix of the GMM
is made diagonal, and the GMM center of mass is made the origin.
"""
function inertial_transforms(positions::AbstractMatrix{<:Real},
                             widths=zeros(eltype(positions),size(positions,2)),
                             weights=ones(eltype(positions),size(positions,2));
                             invert = false)
    com = center_of_mass(positions, weights)
    massmat = mass_matrix(positions, widths, weights, com)
    evecs = eigvecs(massmat)

    T = eltype(com)
    N = length(com)

    # obtain all rotations that align the mass matrix to the coordinate system
    tforms = AffineMap{SMatrix{N,N,T,N^2},SVector{N,T}}[]

    # first, align the the eigenvectors to the coordinate system axes
    # make sure that a reflection is not performed
    if det(evecs) < 0
        evecs[:,end] = -evecs[:,end]
    end

    push!(tforms, inv(LinearMap(SMatrix{length(com),length(com)}(evecs))) ∘ Translation(-com))

    # then rotate the first transformation about each coordinate axis by -π/2, π/2, and π
    for i=1:3
        axis = zeros(N)
        axis[i] = one(T)
        for angle in (π)
            push!(tforms, LinearMap(AngleAxis(angle, axis...)) ∘  tforms[1])
        end
    end
    if invert
        return [inv(tform) for tform in tforms]
    else
        return tforms
    end
end

inertial_transforms(gaussians; kwargs...
    ) = inertial_transforms(hcat([g.μ for g in gaussians]...), [g.σ for g in gaussians], [g.ϕ for g in gaussians]; kwargs...)

inertial_transforms(gmm::AbstractIsotropicGMM; kwargs...
    ) = inertial_transforms(gmm.gaussians; kwargs...)

inertial_transforms(mgmm::AbstractMultiGMM; kwargs...
    ) = inertial_transforms(collect(Iterators.flatten([gmm.gaussians for (k,gmm) in mgmm])); kwargs...)

"""
    score, tform, nevals = rocs_align_gmms(gmmfixed, gmmmoving; maxevals=1000)

Finds the optimal alignment between the two supplied GMMs using steric multipoles,
based on the [ROCS alignment algorithm.](https://docs.eyesopen.com/applications/rocs/theory/shape_shape.html)
"""
function rocs_align(gmmmoving::AbstractGMM, gmmfixed::AbstractGMM; kwargs...)
    # align both GMMs to their inertial frames
    tformfixed = inertial_transforms(gmmfixed)[1]    # Only need one alignment for the fixed GMM
    tformsmoving = inertial_transforms(gmmmoving)    # Take all 4 rotations for the GMM to be aligned

    # perform local alignment starting at each inertial rotation for the "moving" GMM
    results = Tuple{Float64, NTuple{6,Float64}}[]
    for tformm in tformsmoving
        push!(results, local_align(tformm(gmmmoving), tformfixed(gmmfixed); kwargs...))
    end

    # combine the inertial transform with the subsequent alignment transform for the best result
    minoverlap, mindex = findmin([r[1] for r in results])
    tformmoving = AffineMap(results[mindex][2]...) ∘ tformsmoving[mindex]

    # Apply the inverse of `tformfixed` to the optimized transformation
    alignment_tform = inv(tformfixed) ∘ tformmoving

    # return the result
    return ROCSAlignmentResult(gmmmoving, gmmfixed, minoverlap, alignment_tform)
end


"""
    tensorsvd(A::AbstractTensor{T,2}; svdtrunc)
returns the `svd` of an AbstractTensor. `svdtrunc` is a function that has the
singular values as input and returns a number specifying how many of them to keep.
The default `svdtrunc` is `svdtrunc_default` and keeps all of them.

Other options are:
svdtrunc_discardzero
svdtrunc_maxχ
svdtrunc_maxcumerror
svdtrunc_maxerror
"""
function tensorsvd end

function _tensorsvd(A::Array{<:Any,2}; svdtrunc = svdtrunc_default,
        helper::Bool = false)
    F = LA.svd(A)
    svals = F.S
    cutoff = svdtrunc(svals)
    U = F.U[:, 1:cutoff]
    S = LA.diagm(0=>svals[1:cutoff])
    Vt = F.Vt[1:cutoff, :]
    helper && return (U, S, Vt, cutoff)
    return (U, S, Vt)
end

"""
    tensorsvd(A::AbstractTensor, indexes; svdtrunc)
works like `tensorsvd` except that `A` can have arbitrary rank.
`indexes` specifies which indices the fuse for `A` to be a rank-2 tensor
as in `fuselegs`. The tensor is then fused, `tensorsvd` applied and split again.
"""
function tensorsvd(A::AbstractTensor, indexes; svdtrunc = svdtrunc_default)
    fA, rs = fuselegs(A, indexes)
    U, S, Vt = tensorsvd(fA, svdtrunc=svdtrunc)
    li1 = length(indexes[1])
    li2 = length(indexes[2])
    if indexes[1] isa Tuple
        indxs = (ntuple(x -> (1,1,x), li1)..., 2)
        U = splitlegs(U, indxs, rs...)
    end
    if indexes[2] isa Tuple
        indxs = (1, ntuple(x -> (2,2,x), li2)...)
        Vt = splitlegs(Vt, indxs, rs...)
    end
    return (U,S,Vt)
end

svdtrunc_default(x) = length(x)
"""
    svdtrunc_discardzero(s)
return the number of non-zero values in s
"""
svdtrunc_discardzero(x) = length(filter(!iszero, x))

"""
    svdtrunc_maxχ(χ)
return a function that given singular values `s` returns the min of the
length of `s` and χ.
"""
svdtrunc_maxχ(χ) = x -> min(length(x), χ)

"""
    svdtrunc_maxcumerror(ϵ::Real; χ::Int = typemax(Int))
return a function that returns then min of χ and l where l is the number
of singular values that need to be kept to have the truncated sum up to not
more than ϵ.
"""
svdtrunc_maxcumerror(ϵ::Real; χ::Int = typemax(Int)) = x -> _maxcumerror(x, ϵ, χ)

"""
    svdtrunc_maxerror(ϵ::Real; χ::Int = typemax(Int))
return a function that returns then min of χ and l where l is the number
of singular values that need to be kept such that the largest discarded singular
value is below ϵ.
"""
svdtrunc_maxerror(ϵ::Real; χ::Int = typemax(Int)) = x -> _maxerror(x, ϵ, χ)

function _maxcumerror(xs, ϵ, χ)
    acc, lxs = 0, length(xs)
    for i in lxs:-1:1
        acc += xs[i]/xs[1]
        acc >= ϵ && return min(i+1,χ)
    end
    return min(lxs, χ)
end

function _maxerror(xs, ϵ, χ)
    index = findfirst(x -> x/xs[1] < ϵ, xs)
    index == nothing && return min(length(xs), χ)
    return min(index-1, χ)
end

#DTensor
function tensorsvd(A::DTensor{T,2}; svdtrunc = svdtrunc_default) where T
    DTensor.(_tensorsvd(A.array, svdtrunc = svdtrunc))
end

#DASTensor
function tensorsvd(A::DASTensor{T,N,SYM,CHARGES,SIZES,CHARGE};
                   svdtrunc = svdtrunc_default) where
                   {T,N,SYM,CHARGES,SIZES,CHARGE}
    N == 2 || throw(ArgumentError("SVD only works on rank 2 tensors"))
    lch = connectingcharge(A)
    ld =  sizes(A,2)[[chargeindex(c,charges(A,2)) for c in lch]]

    U  = DASTensor{T,N}(SYM, (charges(A,1), lch),
            deepcopy.((sizes(A,1), ld)),
            in_out(A), charge(A))
    S  = DASTensor{real(T),N}(SYM, (lch, lch),
            deepcopy.((ld, ld)),
            vcat(inv(in_out(A,2)), in_out(A,2)))
    Vd = DASTensor{T,N}(SYM, (lch, charges(A,2)),
            deepcopy((ld,sizes(A,2))),
            vcat(inv(in_out(A,2)), in_out(A,2)))

    for (k, degen) in tensor(A)
        in, out = k[1], k[2]
        U[DASSector(in, out)], S[DASSector(out, out)], Vd[DASSector(out, out)], cutoff =
                _tensorsvd(degen, svdtrunc = svdtrunc, helper = true)
        U.dims[2][chargeindex(out,charges(U,2))] = cutoff
        S.dims[1][chargeindex(out,charges(S,1))] = cutoff
        S.dims[2][chargeindex(out,charges(S,2))] = cutoff
        Vd.dims[1][chargeindex(out,charges(Vd,1))] = cutoff
    end
    return (U, S, Vd)
end

"""
    connectingcharge(A)
where A is a rank-2 DASTensor,
returns the charges on the second index that can be realized
given the charges on the first index and the charge of A.
"""
function connectingcharge(A::DASTensor{<:Any,2})
    chs1, chs2 = charges(A)
    io = in_out(A)
    io[1] == InOut(-1) && (chs1 = inv(chs1))
    io[2] == InOut( 1) && (chs2 = inv(chs2))
    #                         (1)ch
    #                          ↓
    #canonical form: chs1 →(1)→A→(-1)→ lch = chs1 + ch ∩ chs2
    chs1 += charge(A)
    lch = chs1 ∩ chs2
    io[2] == InOut(1) && return inv(lch)
    return lch
end

"""
    tensorqr(A::AbstractTensor)

returns tensor Q,R such that A = QR and Q is an orthogonal/unitary matrix
and R is an upper triangular matrix.
"""
function tensorqr end

"""
    tensorqr(A::AbstractTensor, inds)

returns the `tensorqr` of `A` fused according to `inds`.
"""
function tensorqr(A::AbstractTensor, inds)
    fA, rs = fuselegs(A, inds)
    Q,R = tensorqr(fA)
    li1, li2 = length.(inds)
    if inds[1] isa Tuple
        indxs = (ntuple(x -> (1,1,x), li1)..., 2)
        Q = splitlegs(Q, indxs, rs...)
    end
    if inds[2] isa Tuple
        indxs = (1, ntuple(x -> (2,2,x), li2)...)
        R = splitlegs(R, indxs, rs...)
    end
    return (Q,R)
end

function tensorqr(a::DTensor{<:Any,2})
    f = LA.qr(a.array)
    return DTensor(Matrix(f.Q)), DTensor(Matrix(f.R))
end

function tensorqr(a::DASTensor{T,2,SYM}) where {T,SYM}
    lch = connectingcharge(a)
    ld =  sizes(a,2)[[chargeindex(c,charges(a,2)) for c in lch]]

    Q = DASTensor{T,2}(SYM, (charges(a,1), lch),
            deepcopy.((sizes(a,1), ld)),
            in_out(a), charge(a))
    R = DASTensor{T,2}(SYM, (lch, charges(a,2)),
            deepcopy((ld,sizes(a,2))),
            vcat(inv(in_out(a,2)), in_out(a,2)))

    for (k, degen) in tensor(a)
        in, out = k[1], k[2]
        Q[DASSector(in, out)], R[DASSector(out, out)] = LA.qr(degen)
    end
    return (Q,R)
end

"""
    tensorrq(A::AbstractTensor)

returns tensor R,Q such that A = RQ and Q is
obeys Q*Q' = 1 and R is triangular.
This is simply a wrapper for `tensorqr`, useful for e.g. MPS canonicalization.
"""
function tensorrq end

"""
    tensorrq(A::AbstractTensor, inds)

returns the `tensorrq` of `A` fused according to `inds`.
"""
function tensorrq(c, inds = ((1,),(2,)))
    q, r = tensorqr(c, reverse(inds))
    nq, nr = ndims.((q,r))
    q = TO.tensorcopy(q, circshift(1:nq,-1),  1:nq)
    r = TO.tensorcopy(r, circshift(1:nr,1), 1:nr)
    return r, q
end

function tensoreig end

function tensoreig(A::AbstractTensor, indexes; truncfun = svdtrunc_default)
    fA, rs = fuselegs(A, indexes)
    E, D, Ep = tensoreig(fA, truncfun=svdtrunc)
    li1, li2 = length.(indexes)

    if indexes[1] isa Tuple
        indxs = (ntuple(x -> (1,1,x), li1)..., 2)
        E = splitlegs(E, indxs, rs...)
    end
    if indexes[2] isa Tuple
        indxs = (1, ntuple(x -> (2,2,x), li2)...)
        Ep = splitlegs(Ep, indxs, rs...)
    end
    return (E, D, Ep)
end

function _tensoreig(a::Matrix; truncfun = length, helper=false)
    if LA.ishermitian(a)
        evals, evecs = LA.eigen(LA.Hermitian(a))
    else
        evals, evecs = LA.eigen(a)
    end

    p = sortperm(evals, by=abs, rev=true)
    evals = evals[p]
    evecs = evecs[:,p]
    cutoff = truncfun(evals)
    E = evecs[:,1:cutoff]
    d = LA.diagm(0 => evals[1:cutoff])
    helper && return E, d, E', cutoff
    return E, d, E'
end

function tensoreig(a::DTensor{T,2}; truncfun = length) where T
    E, d, Ep = _tensoreig(a.array, truncfun = truncfun)
    DTensor(E), DTensor(d), DTensor(collect(Ep))
end

function tensoreig(A::DASTensor{T,N,SYM,CHARGES,SIZES,CHARGE};
                   truncfun = svdtrunc_default) where
                   {T,N,SYM,CHARGES,SIZES,CHARGE}
    N == 2 || throw(ArgumentError("EIG only works on rank 2 tensors"))
    lch = connectingcharge(A)
    ld =  sizes(A,2)[[chargeindex(c,charges(A,2)) for c in lch]]

    E  = DASTensor{T,N}(SYM, (charges(A,1), lch),
            deepcopy.((sizes(A,1), ld)),
            in_out(A), charge(A))
    D  = DASTensor{T,N}(SYM, (lch, lch),
            deepcopy.((ld, ld)),
            vcat(inv(in_out(A,2)), in_out(A,2)))
    Ep = DASTensor{T,N}(SYM, (lch, charges(A,2)),
            deepcopy((ld,sizes(A,2))),
            vcat(inv(in_out(A,2)), in_out(A,2)))

    for (k, degen) in tensor(A)
        in, out = k[1], k[2]
        E[DASSector(in, out)], D[DASSector(out, out)], Ep[DASSector(out, out)], cutoff =
                _tensoreig(degen, truncfun = truncfun, helper = true)
        E.dims[2][chargeindex(out,charges(E,2))] = cutoff
        D.dims[1][chargeindex(out,charges(D,1))] = cutoff
        D.dims[2][chargeindex(out,charges(D,2))] = cutoff
        Ep.dims[1][chargeindex(out,charges(Ep,1))] = cutoff
    end
    return (E, D, Ep)
end