using FFTW
using StaticArrays
using Statistics

function msd_for_corr_time(traj_msd, corr_time, unwrapped = false)
    a = 0
    for k in 1:size(traj_msd)[3]
        for j in 1:size(traj_msd)[2]
            for i in 1:size(traj_msd)[1]-corr_time
                a += (traj_msd[i+corr_time, j, k] - traj_msd[i, j, k])^2
            end
        end
    end
    return a /= (size(traj_msd)[1] - corr_time) * size(traj_msd)[3]
end

function msd_direct(traj::Trajectory, targetatom, resolution, timerange)
    if !traj.unwrapped
        println("Trajectory is not unwrapped -> MSD might be faulty.")
    end

    return msd_direct(traj.coords, traj.timestep_in_fs, traj.atomlabels, targetatom, resolution, timerange)
end

function msd_direct(coords, timestep_fs, atom, targetatom, resolution, timerange)
    timesteps = size(coords)[3]
    corr_time_list = range(0, length = resolution, step = round(Int, timesteps * timerange / resolution))
    # Permuting the trajectory to shape (timesteps, dimensions, atoms) is faster, because it allows the innermost loop over the timesteps to be along the first dimension
    traj_msd = permutedims(coords[:, atom .== targetatom, :], [3, 1, 2])
    msd_result = Array{Float64, 1}(undef, resolution)
    Threads.@threads for i in 1:resolution
        msd_result[i] = msd_for_corr_time(traj_msd, corr_time_list[i])
    end
    # correlation time in ps, MSD in A^2
    return corr_time_list * timestep_fs / 1000, msd_result
end

function msd_fft(traj, timestep_fs, atom, targetatom, timerange)
    traj_msd = @view traj[:, atom .== targetatom, :]
    traj_msd = permutedims(traj_msd, [3, 1, 2])

    # compute S2
    fft_result = rfft(cat(dims = 1, traj_msd, zeros(Float64, size(traj_msd))), 1)
    fft_result = abs2.(fft_result) * size(traj_msd)[1]
    fft_result = irfft(fft_result, size(traj_msd)[1] * 2, 1)
end

function rdf(traj::Trajectory, atomtype1, atomtype2, d_min, d_max, bins)
    coord, atom, pbc = traj.coords, traj.atomlabels, traj.mdbox.pbc_matrix
    # Each bin i contains distances from distance_limits[i] to distance_limits[i+1]; distnace_mean[i] will be the average of both values
    distance_limits = range(d_min, d_max, length = bins + 1)
    distance_mean = (distance_limits[1:end-1] + distance_limits[2:end]) / 2
    # Shell volume of each bin i and overall ideal density used to normalize RDF
    shell_volumes = 4/3 * pi * (distance_limits[2:end].^3 - distance_limits[1:end-1].^3)
    cell_volume = det(pbc) #pbc[1,:] ⋅ (pbc[2,:] × pbc[3,:])
    ideal_density = count(atom .== atomtype1) * count(atom .== atomtype2) / cell_volume
    # Create distance histogram; delegate calculation of all distances and binning to create_dist_histogram_multi
    dist_histogram = create_dist_histogram_multi(traj, atomtype1, atomtype2, distance_limits)
    rdf = dist_histogram ./ (shell_volumes .* ideal_density .* size(coord)[3])
    return distance_mean, rdf
end

function create_dist_histogram_multi(traj::Trajectory, atomtype1, atomtype2, distance_limits)
    coord, atom, mdbox = traj.coords, traj.atomlabels, traj.mdbox
    dist_histogram = zeros(Int, distance_limits.len - 1)
    # Copy trajectories of specified atom types for later use; using @view here causes a lot more memory usage
    coord1 = coord[:, atom .== atomtype1, :]
    coord2 = coord[:, atom .== atomtype2, :]
    d_min, d_max, d_step = distance_limits[1], distance_limits[end], convert(Float64, distance_limits.step)
    # Split length of trajecgory $size(coord)[3] into roughly equal chunks, one for each available thread
    thread_limits = convert.(Int, floor.(range(1, size(coord)[3] + 1, length = Threads.nthreads() + 1)))
    # Lock to prevent two threads from editing dist_histogram
    Rlock = ReentrantLock()
    # Each thread will process one iteration of this loop; each will produce its own dist_histogram_thread and add it to the overall dist_histogram at the end
    Threads.@threads for i_thread in 1:Threads.nthreads()
        dist_histogram_thread = zeros(Int, distance_limits.len - 1)
        mdbox_thread = deepcopy(mdbox)
        i_limits = thread_limits[Threads.threadid()]:(thread_limits[Threads.threadid() + 1] - 1)
        # Iterate through chunk of timesteps set aside for this thread
        @views for i in i_limits
            # Iterate through atoms of both atom types
            for x in eachcol(coord1[:, :, i])
                for y in eachcol(coord2[:, :, i])
                    distance = pbc_dist(x, y, mdbox_thread)
                    bin!(dist_histogram_thread, distance, d_min, d_max, d_step)
                end
            end
        end
        # Add thread-local dist_histogram_thread to overall dist_histogram; lock it in the meantime in order to prevent threads from interfering with each other
        lock(Rlock)
        dist_histogram .+= dist_histogram_thread
        unlock(Rlock)
    end
    return dist_histogram
end

function bin!(dist_histogram, distance, d_min, d_max, d_step)
    # each bin i ∈ [1, bins] covers distances from (d_min + d_step * (i-1)) to (d_min + d_step * i)
    if d_min < distance < d_max
        index = floor(Int, (distance - d_min) / d_step + 1)
        dist_histogram[index] += 1
    end
end

function autocorrFFT(inputarray)
    arraylength = size(inputarray)[end]
    paddedlength = arraylength * 2 + 1
    paddedarray = zeros(eltype(inputarray), size(inputarray)[1:end-1]..., paddedlength)
    @views paddedarray[:, :, 1:arraylength] = inputarray

    FTresult = rfft(paddedarray, 3)
    FTresult .= FTresult .* conj(FTresult)
    autocorr = irfft(FTresult, paddedlength, 3)[:, :, 1:arraylength]
    for i in 1:arraylength
        @views autocorr[:, :, i] ./= (arraylength - i + 1)
    end
    return autocorr
end

function oacf(traj::Trajectory, atomtype1, atomtype2)
    atomno1 = count(traj.atomlabels .== atomtype1)
    atomno2 = count(traj.atomlabels .== atomtype2)
    frameno = size(traj.coords)[3]

    neighborsfirstframe = zeros(Int64, atomno1)
    frame1type2 = traj.coords[:, traj.atomlabels .== atomtype2, 1]
    @views for (i, pointtype1) in enumerate(eachcol(traj.coords[:, traj.atomlabels .== atomtype1, 1]))
        neighborsfirstframe[i] = TrajTools.next_neighbor(pointtype1, frame1type2, traj.mdbox)[1]
    end
    coordstype1 = @view traj.coords[:, traj.atomlabels .== atomtype1, :]
    coordstype2 = @view traj.coords[:, traj.atomlabels .== atomtype2, :]
    bondvectors = zeros(Float64, 3, atomno1, frameno)
    @views for frame in 1:size(traj.coords)[3]
        for atom1 in 1:count(traj.atomlabels .== atomtype1)
            neighborindex = neighborsfirstframe[atom1]
            coord1 = coordstype1[:, atom1, frame]
            coord2 = coordstype2[:, neighborindex, frame]
            bondvectors[:, atom1, frame] .= TrajTools.minimum_image_vector(coord1, coord2, traj.mdbox)
            normalize!(bondvectors[:, atom1, frame])
        end
    end
    auto1 = autocorrFFT(bondvectors)
    # orientational auto correlation
    oacf = vec(mean(sum(auto1, dims=1), dims = 2))
    # time in ps
    time_oacf = (1:frameno) * traj.timestep_in_fs / 1E3
    return time_oacf, oacf
end


""
function run_hpf_iv(file, model_type::Type, optimizer; kwargs...)
    return _PMs.run_model(file, model_type, optimizer, build_hpf_iv; ref_extensions=[ref_add_xfmr!],  solution_processors=[ _HPM.sol_data_model!], multinetwork=true, kwargs...)
end

""
function build_hpf_iv(pm::AbstractPowerModel)
    bounded = false

    for (n, network) in _PMs.nws(pm)
        _PMs.variable_bus_voltage(pm, nw=n, bounded=bounded)
        _PMs.variable_bus_voltage_magnitude_sqr(pm, nw=n, bounded=bounded)
        variable_transformer_voltage(pm, nw=n, bounded=bounded)

        _PMs.variable_branch_current(pm, nw=n, bounded=bounded)
        _PMs.variable_dcline_current(pm, nw=n)
        variable_transformer_current(pm, nw=n, bounded=bounded)

        variable_load_current(pm, nw=n, bounded=bounded)
        variable_gen_current(pm, nw=n, bounded=bounded)

    end

    for (n, network) in _PMs.nws(pm)
        for i in _PMs.ids(pm, :ref_buses, nw=n)
            constraint_ref_bus(pm, i, nw=n)
        end

        for i in _PMs.ids(pm, :bus, nw=n)
            constraint_current_balance(pm, i, nw=n)
            constraint_vm_auxiliary_variable(pm, i, nw=n)
            # constraint_voltage_harmonics_relative_magnitude(pm, i, nw=n)
        end

        for g in _PMs.ids(pm, :gen, nw=n)
            # _PMs.constraint_gen_active_bounds(pm, g, nw=n)
            # _PMs.constraint_gen_reactive_bounds(pm, g, nw=n)
        end

        for i in _PMs.ids(pm, :load, nw=n)
            constraint_load_power(pm, i, nw=n)
        end

        for b in _PMs.ids(pm, :branch, nw=n)
            _PMs.constraint_current_from(pm, b, nw=n)
            _PMs.constraint_current_to(pm, b, nw=n)

            _PMs.constraint_voltage_drop(pm, b, nw=n)

            # _PMs.constraint_current_limit(pm, b, nw=n)
        end

        for t in _PMs.ids(pm, :xfmr, nw=n)
            constraint_transformer_core_excitation(pm, t, nw=n)
            constraint_transformer_core_voltage_drop(pm, t, nw=n)
            constraint_transformer_core_voltage_balance(pm, t, nw=n)
            constraint_transformer_core_current_balance(pm, t, nw=n)

            constraint_transformer_winding_config(pm, t, nw=n)
            constraint_transformer_winding_current_balance(pm, t, nw=n)
        end

        # for d in _PMs.ids(pm, :dcline, nw=n)
        #     _PMs.constraint_dcline_power_losses(pm, d, nw=n)
        # end
    end

    #constraints across harmonics
    fundamental = 1
    # for i in _PMs.ids(pm, :bus, nw=fundamental)
    #     constraint_voltage_magnitude_rms(pm, i)
    #     constraint_voltage_thd(pm, i, fundamental=fundamental)
    # end

    # for b in _PMs.ids(pm, :branch, nw=fundamental)
    #     constraint_current_limit_rms(pm, b)
    # end

    for g in _PMs.ids(pm, :gen, nw=fundamental)
        constraint_active_filter(pm, g, fundamental=fundamental)
    end

    _PMs.objective_min_fuel_and_flow_cost(pm)
    # objective_distortion_minimization(pm)
end


# External packages
using Revise, DelimitedFiles, SymPy

# Add Exasim to Julia search path
cdir = pwd(); ii = findlast("Exasim", cdir);
include(cdir[1:ii[end]] * "/Installation/setpath.jl");

# Exasim packages
using Preprocessing, Mesh, Gencode, Postprocessing

# create pde structure and mesh structure
pde, mesh = Preprocessing.initializeexasim();

# Define PDE model: governing equations, initial solutions, and boundary conditions
pde.model = "ModelD";            # ModelC, ModelD, ModelW
include("pdemodel.jl");          # include the PDE model file

# Set discretization parameters, physical parameters, and solver parameters
pde.porder = 3;          # polynomial degree
pde.torder = 2;          # time-stepping order of accuracy
pde.nstage = 2;          # time-stepping number of stages
pde.dt = 0.05*ones(50);   # time step sizes
pde.soltime = collect(10:10:length(pde.dt)); # steps at which solution are collected
pde.visdt = 0.05; # visualization timestep size

gam = 1.4;                      # specific heat ratio
Re = 1600;                      # Reynolds number
Pr = 0.71;                      # Prandtl number
Minf = 0.2;                     # Mach number
pde.physicsparam = [gam Re Pr Minf];
pde.tau = [2.0];                # DG stabilization parameter
pde.GMRESrestart=40;
pde.linearsolvertol=0.0001;
pde.linearsolveriter=41;
pde.precMatrixType=2;
pde.NLiter=2;

# Choose computing platform and set number of processors
#pde.platform = "gpu";           # choose this option if NVIDIA GPUs are available
pde.mpiprocs = 4;                # number of MPI processors

# create a linear mesh for a cube domain
mesh.p, mesh.t = Mesh.cubemesh(16,16,16,1);
mesh.p = 2*pi*mesh.p;
# expressions for disjoint boundaries
mesh.boundaryexpr = [p -> (p[2,:] .< 1e-3), p -> (p[1,:] .> 2*pi-1e-3), p -> (p[2,:] .> 2*pi-1e-3), p -> (p[1,:] .< 1e-3), p -> (p[3,:] .< 1e-3), p -> (p[3,:] .> 2*pi-1e-3)];
mesh.boundarycondition = [1 1 1 1 1 1]; # Set boundary condition for each disjoint boundary
mesh.periodicexpr = [2 p->p[2:3,:] 4 p->p[2:3,:]; 1 p->p[[1,3],:] 3 p->p[[1,3],:]; 5 p->p[[1,2],:] 6 p->p[[1,2],:]];

# call exasim to generate and run C++ code to solve the PDE model
sol, pde, mesh,~,~,~,~  = Postprocessing.exasim(pde,mesh);


"""
    Response(Y)

Type for storing response matrix

"""
mutable struct Response

    Y::AbstractArray{Float64,2} # Response matrix Y
end

"""
    Predictors(X, Z, hasXIntercept, hasZIntercept)

Type for storing predictor (covariate) matrices. Also stores boolean
variables hasXIntercept and hasZIntercept (if they are not supplied, they
default to false).

"""
mutable struct Predictors

    # Assumes all categorical variables are coded in appropriate contrasts
    # Matrix of row covariates
    X::AbstractArray{Float64,2}
    # Matrix of column covariates
    Z::AbstractArray{Float64,2}

    # Boolean flag indicating whether X has an intercept
    hasXIntercept::Bool
    # Boolean flag indicating whether Z has an intercept
    hasZIntercept::Bool

    # Usual constructor
    Predictors(X, Z, hasXIntercept, hasZIntercept) =
        new(X, Z, hasXIntercept, hasZIntercept)

    # Modified constructor that sets hasXIntercept and hasZIntercept to false
    # by default
    Predictors(X, Z) = new(X, Z, false, false)
end

"""
    RawData(response, predictors)

Type for storing response and predictor matrices

Also stores dimensions of matrices as n, m, p, and q.
- n = number of rows of X = number of rows of Y
- m = number of rows of Z = number of columns of Y
- p = number of columns of X
- q = number of columns of Z

The constructor will compute n, m, p, and q based on the response and
predictor matrices and assert that they are consistent.

"""
mutable struct RawData

    # Response
    response::Response
    # Predictors
    predictors::Predictors

    # Data dimensions.
    # These will be generated and should not be supplied as input to the
    # constructor.
    n::Int64 # Rows of X = rows of Y
    m::Int64 # Rows of Z = columns of Y
    p::Int64 # Columns of X
    q::Int64 # Columns of Z

    RawData(response, predictors) =
        # Assert that the dimensions in response and predictors match
        (size(predictors.X,1) != size(response.Y,1)) ||
            (size(predictors.Z,1) != size(response.Y,2)) ?
        error("Dimension mismatch.") :
        # Construct RawData object
        new(response, predictors, size(predictors.X,1), size(predictors.Z,1),
            size(predictors.X,2), size(predictors.Z,2))
end

"""
    get_X(data)

Extract X matrix from RawData object

# Arguments

- data = RawData object

# Value

2d array

"""
function get_X(data::RawData)

    return data.predictors.X
end

"""
    get_Z(data)

Extract Z matrix from RawData object

# Arguments

- data = RawData object

# Value

2d array

"""
function get_Z(data::RawData)

    return data.predictors.Z
end

"""
    get_Y(data)

Extract Y matrix from RawData object

# Arguments

- data = RawData object

# Value

2d array

"""
function get_Y(data::RawData)

    return data.response.Y
end


struct ARParams <: AbstractParams
    c::Bool
    p::Vector{<:Integer}
end

ARParams(c::Bool, p::Integer) = ARParams(c, collect(1:p))
ARParams(p::Integer) = ARParams(true, p)

function fit(params::ARParams, z::AbstractVector{T}) where T
    p = params.p
    P = maximum(p)
    if p == collect(1:P)
        ϕ, σ2 = levinson_durbin(z, P)
        μ = params.c ? mean(z) * (1 - sum(ϕ)) : zero(T)
    else
        μ, ϕ, θ, β, σ2 = least_squares(z, params.c, params.p)
    end
    return ARModel{P}(μ, ϕ, σ2)
end

const ARModel{p,T} = ARMAXModel{p,0,0,T}
ARModel{p}(μ::T,ϕ::AbstractVector{T},σ2::T) where {p,T} = ARMAXModel{p,0,0,T}(μ, ϕ, T[], T[], σ2)

isstationary(model::ARModel) = all(norm.(roots(Polynomial([1;.-model.ϕ]))) .> 1)

isinvertible(model::ARModel) = true

function forecast(model::M, z::AbstractVector{T}) where {p,T,M <: ARModel{p,T}}
    N = length(z)
    zhat = model.μ
    for i in 1:min(N, p)
        zhat += model.ϕ[i] * z[N - i + 1]
    end
    return zhat
end

forecast(model::M, z::AbstractVector{T}, a::AbstractVector{T}) where {p,T,M <: ARModel{p,T}} = forecast(model, z)

least_squares(z::AbstractVector{T}, c::Bool, p::Vector{<:Integer}) where T = least_squares(z, T[], zeros(T, 0, 0), c, p, Int[])
least_squares(z::AbstractVector{T}, c::Bool, p::Integer) where T = least_squares(z, c, collect(1:p))
least_squares(z::AbstractVector, p::Integer) = least_squares(z, true, p)

function yule_walker(z::AbstractVector, p::Integer)
    ρ = map(i -> autocorrelation(z, i), 1:p)
    Γ = autocorrelation_matrix(z, p)
    ϕ = cholesky(Γ) \ ρ
    m = mean(z)
    γ0 = autocovariance(z, 0)
    σ2 = (1 - dot(ϕ, ρ)) * γ0
    return ϕ, σ2
end

function levinson_durbin!(ϕ::AbstractVector{T}, σ2::Ref{T}, ρ::AbstractVector{T}, p::Integer) where {T}
    num = ρ[p]
    den = one(T)
    for j in 1:p - 1
        num -= ϕ[j] * ρ[p - j]
        den -= ϕ[j] * ρ[j]
    end
    ϕ[p] = num / den
    for j in 1:div(p - 1, 2)
        tmp = ϕ[j]
        ϕ[j] -= ϕ[p] * ϕ[p - j]
        ϕ[p - j] -= ϕ[p] * tmp
    end
        if isodd(p - 1) ϕ[div(p - 1, 2) + 1] *= (one(T) - ϕ[p]) end
    σ2[] *= (one(T) - ϕ[p]^2)
end

function levinson_durbin(z::AbstractVector, p::Integer)
    ρ = map(i -> autocorrelation(z, i), 1:p)
    T = typeof(zero(eltype(z)) / 1)
    ϕ = zeros(T, p)
    σ2 = Ref{T}(autocovariance(z, 0))
    for i in 1:p
        levinson_durbin!(ϕ, σ2, ρ, i)
    end
    return ϕ, σ2[]
end


#using Revise

using Adaptive
using Distributed

@info "Testing ask! and tell! methods"
function simpledriver!(f, learner,step)
    while !(learner.loss() < step)
        xi = ask!(learner,true)
        yi = f(xi)
        tell!(learner,(xi,yi))
    end
end

learner1d = AdaptiveLearner1D((0,1))
simpledriver!(x->exp(-x^2),learner1d,0.1)

learner2d = AdaptiveLearner2D([(0,1),(0,1)])
simpledriver!(x->exp(-x[1]^2 - x[2]^2),learner2d,0.1)

@info "Testing atribute access"

@show learner1d.loss
@show learner1d.x
@show learner1d.y
@show learner1d.data

@show learner2d.loss
@show learner2d.points
@show learner2d.vertices
@show learner2d.values
@show learner2d.ip


using Meganet
using Base.Test

for TYPE=[Float64,Float32]
    S = getSoftMaxLoss(TYPE)
    display(S)

    @testset "softmax (with bias) $TYPE" begin
        testLossFunction(S,out=false)
    end

    S = getSoftMaxLoss(TYPE)
    S.addBias = false
    @testset "softmax (no bias) $TYPE" begin
        testLossFunction(S,out=false)
    end
end


#####
##### Tests copied from IterativeSolvers.jl
##### https://github.com/JuliaLinearAlgebra/IterativeSolvers.jl/blob/v0.9.2/src/lsmr.jl
#####

# Type used in Dampenedtest
# solve (A'A + diag(v).^2 ) x = A'b
# using LSMR in the augmented space Ã = [A ; diag(v)] b̃ = [b; zeros(size(A, 2)]
struct DampenedMatrix{Tv,TA<:AbstractMatrix{Tv},TD<:AbstractVector{Tv}} <: AbstractMatrix{Tv}
    A::TA
    diagonal::TD
end

function Base.size(A::DampenedMatrix)
    m, n = size(A.A)
    l = length(A.diagonal)
    (m + l, n)
end

function Base.size(A::DampenedMatrix, dim::Integer)
    m, n = size(A.A)
    l = length(A.diagonal)
    dim == 1 ? (m + l) : (dim == 2 ? n : 1)
end

function LinearAlgebra.mul!(y::AbstractVector{Tv}, mw::DampenedMatrix, x::AbstractVector{Tv}) where {Tv}
    m₁ = size(mw.A, 1)
    m₂ = size(mw, 1)
    mul!(view(y, 1:m₁), mw.A, x)
    y[m₁+1:m₂] .= mw.diagonal .* x
    return y
end

function LinearAlgebra.mul!(y::AbstractVector, mw::Adjoint{Tv,<:DampenedMatrix}, x::AbstractVector) where {Tv}
    m₁ = size(mw.parent.A, 1)
    m₂ = size(mw.parent, 1)
    mul!(y, adjoint(mw.parent.A), view(x, 1:m₁))
    y .+= mw.parent.diagonal .* view(x, m₁+1:m₂)
    return y
end

"""
Produces the m × n submatrix from
A = [ 1
      1 2
        2 3
          3 4
            ...
              n ]
suitably padded by zeros.
"""
function sol_matrix(m, n)
    mn = min(m, n)
    I, J, V = SparseArrays.spdiagm_internal(-1 => 1.0 : mn - 1, 0 => 1.0 : mn)
    sparse(I, J, V, m, n)
end

@testset "Small dense matrix" for T = (Float32, Float64, ComplexF32, ComplexF64)
    Random.seed!(501)
    A = rand(T, 10, 5)
    b = rand(T, 10)
    x = QSM.lsmr(A, b)
    @test norm(x - A\b) ≤ √eps(real(T))
end

@testset "SOL test" for (m, n, damp) = ((10, 10, 0), (20, 10, 0), (20, 10, 0.1))
    # Test adapted from the BSD-licensed Matlab implementation at
    #    http://www.stanford.edu/group/SOL/software/lsqr.html
    #              Michael Saunders, Systems Optimization Laboratory,
    #              Dept of MS&E, Stanford University.
    #-----------------------------------------------------------------------
    # 11 Apr 1996: First version for distribution with lsqr.m.
    #              Michael Saunders, Dept of EESOR, Stanford University.

    A = sol_matrix(m, n)
    x = float(n : -1 : 1)
    b = A * x
    x_lsmr = QSM.lsmr(A, b, atol = 1e-7, btol = 1e-7, conlim = 1e10, maxit = 10n)
    @test norm(b - A * x) ≤ 1e-4
end

@testset "Dampened test" for (m, n) = ((10, 10), (20, 10))
    Random.seed!(501)
    # Test used to make sure A, b can be generic matrix / vector
    b = rand(m)
    A = rand(m, n)
    v = rand(n)
    A′ = DampenedMatrix(A, v)
    b′ = [b; zeros(n)]
    x = QSM.lsmr(A′, b′)
    @test norm((A'A + Matrix(Diagonal(v)) .^ 2)x - A'b) ≤ 1e-3
end


function lag0(x::AbstractVector, p::Int = 1)
    T = length(x);
    return [zeros(eltype(x), p); x[1:T - p]]
end

function lag0(x::AbstractArray, p::Int = 1)
    T, n = size(x);
    return [zeros(eltype(x), p, n); x[1:T - p,:]]
end

function eye(m::Int, n::Int = m)
    return Matrix{Float64}(I, m, n)
end

maxabs(x) = maximum(abs, x)

function vech(A::AbstractMatrix, k::Int = 0)
    # return A[tril!(trues(size(A)), 0)]
    return A[tril(trues(size(A)), k)]
end

function vech(A::AbstractArray{T,3}, k::Int = 0) where {T}
    B = vech(A[:,:,1], k)
    a = size(A, 3)

    if a > 1
        for i = 2:a
            append!(B, vech(A[:,:,i], k))
        end
    end

    return B
end

function vechRshp(A::AbstractArray{T,3}, k::Int = 0) where {T}
    a = size(A, 3)
    return reshape(vech(A, k), (:, a))
end

# this one is type unstable
function vec2ltri2(v::AbstractVector)
    n = length(v)
    s = round(Int, (sqrt(8n + 1) - 1) / 2)
    @assert s * (s + 1) / 2 == n "length of vector is not triangular"
    k = 0; [ i >= j ? (k += 1; v[k]) : zero(eltype(v)) for i = 1:s, j = 1:s ]
end

function vec2ltri(v::AbstractVector)
    n = length(v)
    s = round(Int, (sqrt(8n + 1) - 1) / 2)
    @assert s * (s + 1) / 2 == n "length of vector is not triangular"
    A = zeros(eltype(v), s, s)
    k = 0
    for j in 1:s, i in j:s
        k += 1
        A[i,j] = v[k]
    end
    return A
end

function vec2ltriW1(v::AbstractVector{T}, dm::Int = 1)::Array{T,2} where {T}
    n = length(v)
    s = round(Int, (sqrt(8n + 1) + 1) / 2)
    @assert s * (s - 1) / 2 == n "length of vector is not triangular below diagonal "
    A = Matrix{T}(I, s, s)
    k = 0
    if dm == 1
        for j in 1:s, i in j:s
            if i > j
                k += 1
                A[i,j] = v[k]
            end
        end
    elseif dm == 2
        for i in 1:s, j in 1:i
            if i > j
                k += 1
                A[i,j] = v[k]
            end
        end
    end
    return A
end

function vec2sym(v::AbstractVector)
    return Symmetric(vec2ltri(v), :L)
end

function vec2sym(v::AbstractMatrix, dm::Int = 1)
    if dm == 2
        v = v'
    end

    a, b = size(v)
    s = round(Int, (sqrt(8a + 1) - 1) / 2)

    A = similar(v, s, s, b)
    for i = 1:b
        A[:,:,i] = vec2sym(v[:,i])
    end
    return A
end

function companionf(beta::AbstractMatrix, n::Int, p::Int, c::Bool = false)
    return [beta[:, 1:end - c];eye(n * p - n) zeros(n * p - n, n)]
end

function companionf(beta::AbstractVector, n::Int, p::Int, c::Bool = false)
    @assert n == 1 throw(ArgumentError("provided beta is univariate"))
    if c || p > 1
        return [reshape(beta[1:end - c], 1, :); eye(p - 1) zeros(p - 1)]
    else
        return Matrix{Float64}(reshape(beta, 1, 1))
    end
end

function stabcheck(beta::AbstractVecOrMat, n::Int, p::Int, c::Bool = false)::Bool
    # beta is (Nx(N*L+c))
    FF = companionf(beta, n, p, c)
    return maximum(abs, eigvals(FF)) >= 1.
end

function stabcheckC(beta::AbstractVecOrMat, n::Int, p::Int, c::Bool = false)::Bool
    # beta is (Nx(N*L+c))
    # stability check which ignores complex numbers
    FF = companionf(beta, n, p, c)
    D = eigvals(FF)
    if isreal(D)
        S = maximum(abs, D) >= 1.
    else
        S = true
    end
    return S
end

function stabcheck(beta::Number, n::Int = 1, p::Int = 1, c::Bool = false)::Bool
    return abs(beta) >= one(eltype(beta))
end
stabcheckC(beta::Number, n::Int = 1, p::Int = 1, c::Bool = false) = stabcheck(beta, n, p, c)

function preparexy(data::AbstractVector, p::Int, c::Bool = false)
    T = length(data);
    Y = data[p + 1:end];
    X = similar(data, T, p + c);
    for i = 1:p
        X[:,(i - 1) + 1:i] = lag0(data, i);
    end
    if c
        X[:,end] = ones(eltype(data), T);
    end
    X = X[p + 1:end,:]
    return Y, X
end

function preparexy(data::AbstractMatrix, p::Int, c::Bool = false)
    T, n = size(data);
    Y = data[p + 1:end,:];
    X = similar(data, T, n * p + c);
    for i = 1:p
        X[:,n * (i - 1) + 1:n * i] = lag0(data, i);
    end
    if c
        X[:,end] = ones(eltype(data), T);
    end
    X = X[p + 1:end,:]
    return Y, X
end

function sumsqr(a::AbstractArray)
    return mapreduce(x -> x^2, +, a)
end

# function wish(h, n::Int)
#     A = cholesky!(h).L * randn(size(h, 1), n)
#     A = A * A'
#     return Symmetric(A)
# end

function wish(h::AbstractArray{S,2}, n::Int)   where {S <: AbstractFloat}
    A = randn(n, size(h, 1)) * cholesky(h).U
    return A'A
end

function wish(h::AbstractArray{S,1}, n::Int)   where {S <: AbstractFloat}
    A = randn(n) .* sqrt.(h)
    return dot(A', A)
end

function iwish(h::AbstractMatrix, n::Int)
    B = randn(n, size(h, 1))
    A = cholesky(h).L / qr!(B).R
    return A * A'
end

function iwish(h::AbstractVector, n::Int)
    B = randn(n, 1)
    A = sqrt.(h) / qr!(B).R
    return dot(A, A')
end

function genPSDmat(K::Int = rand(2:9))
    A = randn(K, K)
    return A'A
end

function genPSDmatStrict(K::Int = rand(2:9); no::Int = rand(1:K - 1))
    A = randn(K, K)
    A = A'A
    vals, vecs = eigen(A)
    vals[1:no] .= 0.
    return vecs * Diagonal(vals) * vecs'
end

function genPDmat(K::Int = rand(2:9))
    A = zeros(K, K)
    while isposdef(A) == false
        A .= randn(K, K)
        A .= A'A
    end
    return A
end

function cholPSD(A::AbstractMatrix{T})::Array{T,2}  where {T <: AbstractFloat}
    B = copy(Symmetric(A));
    if isposdef(B)

        # B = cholesky(B)
        # return convert(Array{Float64,2}, B.U)

        return cholesky!(B).U

    else
        # A = sqrt(A)
        # # D,U = eig(A);

        # A = eigen(A);

        # tol = length(A.values) * eps(maximum(A.values));
        # # D1 = A.values[A.values .> tol];
        # A = diagm(sqrt.(A.vectors)) * A.vectors[:, A.values .> tol]';

        # M = eigen(B)
        # D = M.values
        # U = M.vectors

        D, U = eigen(B)
        tol = length(D) * eps(maximum(D));
        # D1 = D[D .> tol];
        # # A = diagm(0 => sqrt.(D1)) * U[:, D .> tol]';
        # C = Diagonal(sqrt.(D1)) * U[:, D .> tol]';

        # return Diagonal(sqrt.(filter(z -> z > tol, D))) * U[:, D .> tol]'
        cord = findall(z -> z > tol, D)
        # return Diagonal(sqrt.(filter(z -> z > tol, D))) * U[:, findall(z -> z > tol, D)]'
        return Diagonal(sqrt.(D[cord])) * U[:, cord]'

    end
end

"the function regularises the matrix to be positive definite given that there are no eigenvalues of zeros"
function regMat2PD!(A::Matrix)
    D, U = eigen(A)
    A .= Symmetric(U * Diagonal(abs.(D)) * U')
    return nothing
end

function randnPSD(mu::AbstractMatrix{T}, sigma::AbstractMatrix{T})::Array{T}  where {T <: AbstractFloat}
    # input:
    # mu is row vector
    #
    # output:
    # draw is a row vector
    sigmaN = cholPSD(sigma)
    draw = mu + randn(1, size(sigmaN, 1)) * sigmaN
    return draw
end

function randnPSD(mu::AbstractVector{T}, sigma::AbstractMatrix{T})::Array{T}  where {T <: AbstractFloat}
    # input:
    # mu is row vector
    #
    # output:
    # draw is a row vector
    sigmaN = cholPSD(sigma)
    draw = mu' + randn(1, size(sigmaN, 1)) * sigmaN
    return draw
end

function mcSS(PI::Matrix)
    a = eigen(Matrix(PI'))
    b = a.vectors[:,a.values .≈ 1.]
    return vec(b) ./ sum(b)
end

function mcSS2(PI::Matrix, tol::Float64 = eps())
    l = size(PI, 1)
    q = ones(l) ./ l
    q1 = similar(q)
    a = 2.
    while a > tol
        q1 .= PI' * q
        a = maximum(abs, q1 .- q)
        q .= q1
    end
    return q
end

function quantileArr(v::AbstractArray{T,2}, p, dms::Int = 1) where {T}
    a, b = size(v)
    w = length(p)
    p = vec(p)
    if dms == 1
        A = similar(v, w, b)
        for i = 1:b
            A[:,i] = quantile(v[:,i], p)
        end
    elseif dms == 2
        A = similar(v, a, w)
        for i = 1:a
            A[i,:] = quantile(v[i,:], p)
        end
    end
    return A
end

function quantileArr(v::AbstractArray{T,3}, p, dms::Int = 1) where {T}
    a, b, c = size(v)
    w = length(p)
    p = vec(p)

    if dms == 1
        A = similar(v, w, b, c)
        for i in 1:b, j in 1:c
            A[:,i,j] = quantile(v[:,i,j], p)
        end
    elseif dms == 2
        A = similar(v, a, w, c)
        for i in 1:a, j in 1:c
            A[i,:,j] = quantile(v[i,:,j], p)
        end
    elseif dms == 3
        A = similar(v, a, b, w)
        for i in 1:a, j in 1:b
            A[i,j,:] = quantile(v[i,j,:], p)
        end
    end

    return A
end

ismyapprox(a, b, dgts::Int) = isapprox(a, b, atol = dgts)

function normpdf(x::Float64, mu::Float64 = 0.0, sigma::Float64 = 1.0)
    return exp(-0.5 * ((x - mu) ./ sigma).^2) ./ (sqrt(2. * pi) .* sigma)
end

function acf(y::AbstractVector, q::Int)
# function assumes that there are no missings

    z = y .- mean(y)
    T = length(y)

    ac0 = dot(z, z) ./ T
    ac = dot(z[1:end - q], z[q + 1:end]) ./ T
    ac /= ac0; # Normalize to correllation

    return ac
end

function acf(y::AbstractVector, q::AbstractRange = 1:10)
    # function assumes that there are no missings

    z = y .- mean(y)
    T = length(y)

    ac0 = dot(z, z) ./ T
    ac = Array{Float64}(undef, 0);
    for j = q
        push!(ac, dot(z[1:end - j], z[j + 1:end]) ./ T)
    end
    ac ./= ac0; # Normalize to correllation

    return ac
end

function acf(y::AbstractMatrix, q = 1:10)

    T, n = size(y)

    ac = Array{Float64}(undef, 0);

    for i = 1:n
        append!(ac, acf(y[:,i], q))
    end

    ac = reshape(ac, length(q), n)

    return ac
end

function decVCV(VCV::AbstractVector, getAinv::Bool = false)
    Omega = vec2sym(VCV)
    Omega = cholesky(Omega).L
    sigma = diag(Omega)
    A = Omega / Diagonal(sigma)
    if getAinv
        A = inv(A)
    end
    A = vech(A, -1)

    return A, sigma
end

function decVCV(VCV::Matrix{Float64}, getAinv::Bool = false)

    a, T = size(VCV)

    s = round(Int, (sqrt(8a + 1) - 1) / 2)
    s * (s + 1) / 2 == a || error("fromVCV: length of vector is not triangular")

    sigma = Array{Float64}(undef, s, T)
    A = Array{Float64}(undef, round(Int, s * (s - 1) / 2), T)

    for i = 1:T
        A[:,i], sigma[:,i] = decVCV(VCV[:,i], getAinv)
    end

    return A, sigma
end

function ols1(y::AbstractVector, x::AbstractMatrix)
    nobs, nvar = size(x)
    bhat = x \ y
    yhat = x * bhat
    res = y - yhat
    sig2hat = (res'res) / (nobs - nvar)
    sigbhat = sig2hat .* eye(nvar) / (x' * x)
    R2 = var(yhat) / var(y)

    return (nobs = nobs, nvar = nvar, bhat = bhat, yhat = yhat, res = res, sig2hat = sig2hat, sigbhat = sigbhat, R2 = R2)
end

function VARols1(data::AbstractVecOrMat{Float64}, cnst::Bool, P::Int)
    y, x = preparexy(data, P, cnst)
    nobs, nvar = size(x)
    bhat = x \ y
    yhat = x * bhat
    res = y - yhat
    sig2hat = (res'res) / (nobs - nvar)

    return (nobs = nobs, nvar = nvar, bhat = bhat, yhat = yhat, res = res, sig2hat = sig2hat)
end

function nanmean(x::AbstractArray{T}) where T <: AbstractFloat
    sum = zero(eltype(x))
    count = 0
    for i in x
        if !isnan(i)
            sum += i
            count += 1
        end
    end
    return sum / count
end

nanmean(x::AbstractArray, y::Int) = mapslices(nanmean, x; dims = y)

function inbetween(a::Matrix, b::AbstractArray{T,3})  where {T}
    return (view(b, :, :, 1) .<= a) .& (a .<= view(b, :, :, 2))
end

"The function multiplies matrices A and B and extracts the diagonal to a prealoccated vector v"
function getmultdiag!(v::Vector, A::Matrix, B::Matrix)
    # v .= dot.(eachrow(A), eachcol(B))
    tmp = A .* B'
    sum!(v, tmp)
    return nothing
end

"the function that transforms the value from domain (min,max) to range with minus/plus infinity"
function transf1To(a::Float64, minn::Float64, maxx::Float64)
   return -log((maxx - a) ./ (a - minn))
end

"the function that transforms the value back from minus/plus infinity  to range with domain (min,max)"
function transf1Back(a::Float64, minn::Float64, maxx::Float64)
    return minn + (maxx - minn) / (1. + exp(-a))
end

"Function to split samples."
function split_data(df, at = 0.70)
    r = size(df, 1)
    index = Int(round(r * at))
    train = df[1:index, :]
    test  = df[(index + 1):end, :]
    return train, test
end

    "A handy helper function to standardize the data"
function standardize(x;dims::Int = 1)
    return (x .- mean(x, dims = dims)) ./ std(x, dims = dims), x
end

"A handy helper function to unstandardize the data"
function unstandardize(x, orig;dims::Int = 1)
    return (x .+ mean(orig, dims = dims)) .* std(orig, dims = dims)
end

"a function that prints the loop progress"
function showprog(iter::Int, Nreps::Int; it_print::Int = floor(Nreps / iter))
    if mod(iter, it_print) == 0
        println("Iter: $iter / $Nreps ")
    end
end


using Statistics

"""
    psislw(lw, [wcpp, wtrunc])

Compute the Pareto smoothed importance sampling (PSIS).

# Arguments
* `lw::AbstractArray`: Array of size n x m containing m sets of n log weights. It is also possible to provide one dimensional array of length n.
* `wcpp::Real`: Percentage of samples used for GPD fit estimate (default is 20).
* `wtrunc::Float64`: Positive parameter for truncating very large weights to n^wtrunc. Providing False or 0 disables truncation. Default values is 3/4.

# Returns
* `lw_out::AbstractArray`: Smoothed log weights
* `kss::AbstractArray`: Pareto tail indices
"""
function psislw(lw::AbstractArray, wcpp::Int64=20, wtrunc::Float64=3/4)

    lw_out = copy(lw)

    if ~(1 <= ndims(lw_out) <= 2)
        throw(DimensionMismatch("Argument `lw` must be 1 or 2 dimensional."))
    end
    if size(lw_out,1) <= 1
        error("More than one log-weight needed.")
    end
    return _psislw(lw_out, wcpp, wtrunc)
end

function _psislw(lw_out::Array{Float64}, wcpp::Int64, wtrunc::Float64)

    if ndims(lw_out) == 2
        n, m = size(lw_out)
    elseif ndims(lw_out) == 1
        n = length(lw_out)
        m = 1
    end
    kss = zeros(Float64, m)

    # precalculate constants
    cutoffmin = log(eps(0.5))
    logn = log(n)

    # loop over sets of log weights
    for i in 1:m
        x = view(lw_out, :, i)
        x .-= maximum(x)
        xcutoff = max(quantile(x, 1 - wcpp/100), cutoffmin)
        expxcutoff = exp(xcutoff)
        x2 = x[x .> xcutoff]
        n2 = length(x2)
        if n2 <= 4
            # not enough tail samples for gpdfitnew
            k = Inf
        else
            # order of tail samples
            x2si = sortperm(x2)
            # fit generalized Pareto distribution to the right tail samples
            x2 .= exp.(x2) .- expxcutoff
            k, sigma = gpdfitnew(x2)
            # compute ordered statistic for the fit
            sti = collect(0.5:n2) / n2
            qq = gpinv(sti, k, sigma)
            qq .= log.(qq .+ expxcutoff)
            # place the smoothed tail into the output array
            x[x .> xcutoff] = qq[x2si]
        end
        if wtrunc > 0
            # truncate too large weights
            lwtrunc = wtrunc * logn - logn + logsumexp(x)
            x[x .> lwtrunc] .= lwtrunc
        end
        # renormalize weights
        x .-= logsumexp(x)
        # store tail index k
        kss[i] = k
    end
    # If the provided input array is one dimensional, return kss as scalar.
    if ndims(lw_out) == 1
        kss = kss[1]
    end
    return lw_out, kss
end

export
    psislw


# Run JITL, SITL, and PITL first to generate their logs. Then run this script to
# compare them.

module Temp

using Plots
using HDF5

include(joinpath(Pkg.dir("Eratosthenes"), "src", "EratosthenesRotations.jl"))
using .EratosthenesRotations

# Nominal
t, q_BI = HDF5.h5open("out/basic-cube.h5", "r") do logs
    (read(logs, "/cube1/truth/time"),
     read(logs, "/cube1/truth/data/q_BI"))
end

# SITL
t_sitl, q_BI_sitl = HDF5.h5open("out/basic-cube-sitl.h5", "r") do logs
    (read(logs, "/cube1/truth/time"),
     read(logs, "/cube1/truth/data/q_BI"))
end

# PITL
t_pitl, q_BI_pitl = HDF5.h5open("out/basic-cube-pitl.h5", "r") do logs
    (read(logs, "/cube1/truth/time"),
     read(logs, "/cube1/truth/data/q_BI"))
end

# Scalars are logged as 1-by-n (don't know why). Convert to just n for
# the sake of plotting.
t      = dropdims(t, dims=1)
t_sitl = dropdims(t_sitl, dims=1)
t_pitl = dropdims(t_pitl, dims=1)

# Compare the literal values in the quaternion (not the underlying rotation);
# they should be identical.
Δq_sitl = q_BI - q_BI_sitl
display(maximum(Δq_sitl, 2))
Δq_pitl = q_BI - q_BI_pitl
display(maximum(Δq_pitl, 2))
Δq_pitl_wrt_sitl = q_BI_sitl - q_BI_pitl
display(maximum(Δq_pitl_wrt_sitl, 2))

# Ok, also compare the angular differences.
θ_sitl = map(k -> q2aa(qpos(qdiff(q_BI[:,k], q_BI_sitl[:,k])))[1], 1:size(q_BI, 2))
display(maximum(θ_sitl) * 180/pi)
θ_pitl = map(k -> q2aa(qpos(qdiff(q_BI[:,k], q_BI_pitl[:,k])))[1], 1:size(q_BI, 2))
display(maximum(θ_pitl) * 180/pi)
θ_pitl_wrt_sitl = map(k -> q2aa(qpos(qdiff(q_BI_sitl[:,k], q_BI_pitl[:,k])))[1], 1:size(q_BI_sitl, 2))
display(maximum(θ_pitl_wrt_sitl) * 180/pi)

# Show the divergence.
# pyplot()
plotlyjs()
# display(plot(t, transpose(Δq),
#              label  = ["Δq1" "Δq2" "Δq3" "Δq4"],
#              ylabel = "Δq"))
display(plot(t, θ_pitl * (180./pi * 3600.),
             xlabel = "Time (s)",
             ylabel = "θ (arcseconds)"))

end


using .RCall
import .RCall: rcopy, rcopytype, RClass
function rcopy(::Type{PointPattern}, s::Ptr{VecSxp})
    x = rcopy(Vector, s[:x])
    y = rcopy(Vector, s[:y])
    window = rcopy(Dict{Symbol, Any}, s[:window])

    PointPattern(collect(zip(x,y)), (x=window[:xrange], y=window[:yrange]))
end

rcopytype(::Type{RClass{:ppp}}, ::Ptr{VecSxp}) = PointPattern

import .RCall: sexp, protect, unprotect, setclass!, RClass, sexpclass
function sexp(::Type{RClass{:ppp}}, pp::PointPattern)
    # win = protect(sexp(Dict(:type => "regtangle",
    #     :xrange => [pp.window.x...],
    #     :yrange => [pp.window.y...])))
    # setclass!(win, sexp("owin"))
    # r = protect(sexp(Dict(
    #     :n => length(pp.data),
    #     :x => getx.(pp.data),
    #     :y => gety.(pp.data),
    #     :window => win)))
    # setclass!(r, sexp("ppp"))
    # unprotect(2)
    r = R"spatstat.geom::ppp"(getx.(pp.data), gety.(pp.data),
        window=R"spatstat.geom::owin"([pp.window.x...], [pp.window.y...]))
    r
end

sexpclass(f::PointPattern) = RClass{:ppp}

function rcopy(::Type{NamedTuple{(:x, :y), Tuple{Float64, Float64}}}, s::Ptr{VecSxp})
    if !rcopy(R"spatstat.geom::is.rectangle"(s))
        @warn "Non-rectangle windows not supported. Using the bounding box as the window."
    end

    window = rcopy(Dict{Symbol, Any}, s)

    (x=window[:xrange], y=window[:yrange])
end

rcopytype(::Type{RClass{:owin}}, ::Ptr{VecSxp}) = NamedTuple{(:x, :y), Tuple{Float64, Float64}}


using Test
#using DelimitedFiles, LinearAlgebra

#using StableRNGs
#rng = StableRNG(123)

import MLJBase
const Mlj = MLJBase
import Distributions
using BetaML

TESTRNG = FIXEDRNG # This could change...
#TESTRNG = StableRNG(123)

println("*** Testing Clustering...")

# ==================================
# New test
# ==================================
println("Testing initRepreserntative...")

Z₀ = initRepresentatives([1 10.5;1.5 10.8; 1.8 8; 1.7 15; 3.2 40; 3.6 32; 3.6 38],2,initStrategy="given",Z₀=[1.7 15; 3.6 40])
@test isapprox(Z₀,[1.7  15.0; 3.6  40.0])

# ==================================
# New test
# ==================================
println("Testing kmeans...")

(clIdx,Z) = kmeans([1 10.5;1.5 10.8; 1.8 8; 1.7 15; 3.2 40; 3.6 32; 3.3 38; 5.1 -2.3; 5.2 -2.4],3,initStrategy="grid",rng=copy(TESTRNG))
@test clIdx == [2, 2, 2, 2, 3, 3, 3, 1, 1]
#@test (clIdx,Z) .== ([2, 2, 2, 2, 3, 3, 3, 1, 1], [5.15 -2.3499999999999996; 1.5 11.075; 3.366666666666667 36.666666666666664])

# ==================================
# New test
# ==================================
println("Testing kmedoids...")
(clIdx,Z) = kmedoids([1 10.5;1.5 10.8; 1.8 8; 1.7 15; 3.2 40; 3.6 32; 3.3 38; 5.1 -2.3; 5.2 -2.4],3,initStrategy="shuffle",rng=copy(TESTRNG))
@test clIdx == [1, 1, 1, 1, 2, 2, 2, 3, 3]

# ==================================
# New test
# ==================================
println("Testing mixture initialisation and log-pdf...")

m1 = SphericalGaussian()
m2 = SphericalGaussian([1.1,2,3])
m3 = SphericalGaussian(nothing,10.2)
mixtures = [m1,m2,m3]
X = [1 10 20; 1.2 12 missing; 3.1 21 41; 2.9 18 39; 1.5 15 25]
initMixtures!(mixtures,X,minVariance=0.25,rng=copy(TESTRNG))
@test sum([sum(m.μ) for m in mixtures]) ≈ 102.2
@test sum([sum(m.σ²) for m in mixtures]) ≈ 19.651086419753085
mask = [true, true, false]
@test lpdf(m1,X[2,:][mask],mask) ≈ -3.461552516784797

m1 = DiagonalGaussian()
m2 = DiagonalGaussian([1.1,2,3])
m3 = DiagonalGaussian(nothing,[0.1,11,25.0])
mixtures = [m1,m2,m3]
initMixtures!(mixtures,X,minVariance=0.25,rng=copy(TESTRNG))
@test sum([sum(m.σ²) for m in mixtures]) ≈ 291.27933333333334
@test lpdf(m1,X[2,:][mask],mask) ≈ -3.383055441795939

m1 = FullGaussian()
m2 = FullGaussian([1.1,2,3])
m3 = FullGaussian(nothing,[0.1 0.2 0.5; 0 2 0.8; 1 0 5])
mixtures = [m1,m2,m3]
initMixtures!(mixtures,X,minVariance=0.25,rng=copy(TESTRNG))
@test sum([sum(m.σ²) for m in mixtures]) ≈ 264.77933333333334
@test lpdf(m1,X[2,:][mask],mask) ≈ -3.383055441795939

# ==================================
# New test
# ==================================
println("Testing gmm...")
X = [1 10.5;1.5 missing; 1.8 8; 1.7 15; 3.2 40; missing missing; 3.3 38; missing -2.3; 5.2 -2.4]
clusters = gmm(X,3,verbosity=NONE, initStrategy="grid",rng=copy(TESTRNG))
@test isapprox(clusters.BIC,114.1492467835965)
#clusters.pₙₖ
#clusters.pₖ
#clusters.mixtures
#clusters.BIC

# ==================================
# New test
# ==================================
println("Testing predictMissing...")
X = [1 10.5;1.5 missing; 1.8 8; 1.7 15; 3.2 40; missing missing; 3.3 38; missing -2.3; 5.2 -2.4]
out = predictMissing(X,3,mixtures=[SphericalGaussian() for i in 1:3],verbosity=NONE, initStrategy="grid",rng=copy(TESTRNG))
@test isapprox(out.X̂[2,2],14.155186593170251)

X = [1 10.5;1.5 missing; 1.8 8; 1.7 15; 3.2 40; missing missing; 3.3 38; missing -2.3; 5.2 -2.4]
out2 = predictMissing(X,3,mixtures=[DiagonalGaussian() for i in 1:3],verbosity=NONE, initStrategy="grid",rng=copy(TESTRNG))
@test out2.X̂[2,2] ≈ 14.588514438886131

X = [1 10.5;1.5 missing; 1.8 8; 1.7 15; 3.2 40; missing missing; 3.3 38; missing -2.3; 5.2 -2.4]
out3 = predictMissing(X,3,mixtures=[FullGaussian() for i in 1:3],verbosity=NONE, initStrategy="grid",rng=copy(TESTRNG))
@test out3.X̂[2,2] ≈ 11.166652292936876

# ==================================
# NEW TEST
println("Testing MLJ interface for Clustering models....")
X, y                           = Mlj.@load_iris

model                          = KMeans(rng=copy(TESTRNG))
modelMachine                   = Mlj.machine(model, X)
(fitResults, cache, report)    = Mlj.fit(model, 0, X)
distances                      = Mlj.transform(model,fitResults,X)
yhat                           = Mlj.predict(model, fitResults, X)
acc = BetaML.accuracy(Mlj.levelcode.(yhat),Mlj.levelcode.(y),ignoreLabels=true)
@test acc > 0.8

model                          = KMedoids(rng=copy(TESTRNG))
modelMachine                   = Mlj.machine(model, X)
(fitResults, cache, report)    = Mlj.fit(model, 0, X)
distances                      = Mlj.transform(model,fitResults,X)
yhat                           = Mlj.predict(model, fitResults, X)
acc = BetaML.accuracy(Mlj.levelcode.(yhat),Mlj.levelcode.(y),ignoreLabels=true)
@test acc > 0.8

model                       =  GMMClusterer(mixtures=:diag_gaussian,rng=copy(TESTRNG))
modelMachine                =  Mlj.machine(model, X) # DimensionMismatch
(fitResults, cache, report) =  Mlj.fit(model, 0, X)
yhat_prob                   =  Mlj.predict(model, fitResults, X)  # Mlj.transform(model,fitResults,X)
# how to get this ??? Mlj.predict_mode(yhat_prob)
@test Distributions.pdf(yhat_prob[end],2) ≈ 0.5937443601647852

X = [1 10.5;1.5 missing; 1.8 8; 1.7 15; 3.2 40; missing missing; 3.3 38; missing -2.3; 5.2 -2.4]
X = Mlj.table(X)
model                       =  MissingImputator(rng=copy(TESTRNG))
modelMachine                =  Mlj.machine(model,X)
(fitResults, cache, report) = Mlj.fit(model, 0, X)
XD                          =  Mlj.transform(model,fitResults,X)
XDM                         =  Mlj.matrix(XD)
@test isapprox(XDM[2,2],15.441553354222702)
# Use the previously learned structure to imput missings..
Xnew_withMissing            = Mlj.table([1.5 missing; missing 38; missing -2.3; 5.1 -2.3])
XDNew                       = Mlj.transform(model,fitResults,Xnew_withMissing)
XDMNew                      =  Mlj.matrix(XDNew)
@test isapprox(XDMNew[1,2],13.818691793037452)

#=
# Marginally different
XD  == XDNew
Mlj.matrix(XD)  ≈ Mlj.matrix(XDNew)
XDM  == XDMNew

for r in 1:size(XDM,1)
    for c in 1:size(XDM,2)
        if XDM[r,c] != XDMNew[r,c]
            println("($r,$c): $(XDM[r,c]) - $(XDMNew[r,c])")
        end
    end
end
=#

#=
@test Mlj.mean(Mlj.LogLoss(tol=1e-4)(yhat_prob, y)) < 0.0002
Mlj.predict_mode(yhat_prob)
N = size(Mlj.matrix(X),1)
nCl = size(fitResults[2],1)
yhat_matrix = Array{Float64,2}(undef,N,nCl)
[yhat_matrix[n,c]= yhat_prob[n][c] for n in 1:N for c in 1:nCl]
yhat_prob[2]
 Mlj.matrix(yhat_prob)
acc = accuracy(Mlj.levelcode.(yhat),Mlj.levelcode.(y),ignoreLabels=true)
ynorm = Mlj.levelcode.(y)
accuracy(yhat_prob,ynorm,ignoreLabels=true)
@test acc > 0.8
=#

#=
using MLJBase, BetaML
y, _                        =  make_regression(1000, 3, rng=123);
ym                          =  MLJBase.matrix(y)
model                       =  GMM(rng=copy(BetaML.TESTRNG))
(fitResults, cache, report) =  MLJBase.fit(model, 0, nothing, y)
yhat_prob                   =  MLJBase.transform(model,fitResults,y)
yhat_prob                   =  MLJBase.predict(model, fitResults, y)

modelMachine                =  MLJBase.machine(model, nothing, y)
mach                        =  MLJBase.fit!(modelMachine)
yhat_prob                   =  MLJBase.predict(mach, nothing)
=#


"""
```
    SubpixelNonmaximaSuppression <: AbstractEdgeThinningAlgorithm
    SubpixelNonmaximaSuppression(; threshold::Union{Number, Percentile} = Percentile(20))

    f = SubpixelNonmaximaSuppression()
    f(out₁::AbstractArray, out₂::Matrix{<:AbstractArray}, mag::AbstractArray, g₁::AbstractArray, g₂::AbstractArray, f::SubpixelNonmaximaSuppression)
```

Isolates local maxima of gradient magnitude `mag` along the local gradient
direction to subpixel precision.  The arguments `g₁` and `g₂` represent the
gradient in the first spatial dimension (y), and the second spatial dimension
(x), respectively.

The integer components of the local maxima correspond to non-zero row and column
entries `out₁`. The accompanying subpixel offsets  are stored in a 2-D array
`out₂` as length-2 vectors. One can recover the  subpixel coordinates by adding
the subpixel offsets to the integer components.

# Details

TODO

# Example

```julia

using TestImages, FileIO, ImageView, ImageEdgeDetection, ImageFiltering

img =  testimage("mandril_gray")

# Gradient in the first and second spatial dimension
g₁, g₂ = imgradients(img, KernelFactors.scharr)

# Gradient magnitude
mag = hypot.(g₁, g₂)

nms = zeros(eltype(mag), axes(mag))
subpixel_offsets = zeros(SVector{2,Float64}, axes(mag))

# Instantiate the NonmaximaSuppression functor.
f = SubpixelNonmaximaSuppression()

# Suppress the non-maximal gradient magnitudes and store the result in `nms`.
f(nms, subpixel_offsets, mag, g₁, g₂)

imshow(img)
imshow(mag)
imshow(nms)
```

# References
1. J. Canny, "A Computational Approach to Edge Detection," in IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. PAMI-8, no. 6, pp. 679-698, Nov. 1986, doi: 10.1109/TPAMI.1986.4767851.
2. F. Devernay, "A non-maxima suppression method for edge detection with sub-pixel accuracy", Tech. Report RR-2724, INRIA, 1995.
"""
@with_kw struct SubpixelNonmaximaSuppression{T <: Union{Real,AbstractGray, Percentile}} <: AbstractEdgeThinningAlgorithm
    threshold::T = Percentile(20)
end

function (f::SubpixelNonmaximaSuppression)(out₁::AbstractArray, out₂::AbstractArray{<:StaticVector}, mag::AbstractArray, g₁::AbstractArray, g₂::AbstractArray)
    @unpack threshold = f

    low_threshold =  typeof(threshold) <: Percentile ? StatsBase.percentile(vec(mag), threshold.p) : threshold

    # Isolate local maxima of gradient magnitude by “non-maximum suppression”
    # along the local gradient direction.
    suppress_subpixel_non_maxima!(out₁, out₂, mag, g₁, g₂, low_threshold)

    return out₁, out₂
end

"""
```
    suppress_subpixel_non_maxima!(out₁::AbstractArray, out₂::AbstractArray{<:StaticVector}, mag::AbstractArray, gx::AbstractArray, gy::AbstractArray, threshold::Number)
```

Isolates local maxima of gradient magnitude by “non-maximum suppression” along the local gradient direction.

"""
function suppress_subpixel_non_maxima!(out₁::AbstractArray, out₂::AbstractArray{<:StaticVector},  mag::AbstractArray, g₁::AbstractArray, g₂::AbstractArray, threshold::Number)
    itp = interpolate(mag, BSpline(Linear()))
    etp = extrapolate(itp, Flat())
    m₁ = zero(eltype(mag))
    m₂ = zero(eltype(mag))
    @inbounds for i in CartesianIndices(mag)
        r, c = i.I
        d₁ = gray(g₁[i])
        d₂ = gray(g₂[i])
        mc = mag[i]
        if mc < threshold || mc == 0 || isnan(mc)
            out₁[r,c] = zero(eltype(mag))
        else
            # Ensure the vector 𝐝 = [d₁, d₂] has unit norm.
            d₁ = d₁ / mc
            d₂ = d₂ / mc
            m₁ = etp(r + d₁, c + d₂)
            m₂ = etp(r - d₁, c - d₂)
            if ((m₁ <= mc) && (m₂ < mc)) || ((m₁ < mc) && (m₂ <= mc))
                out₁[r,c] = mag[r,c]
                # To obtain subpixel precision we fit a parabola through the
                # three points:
                # (x₁ = -1, y₁ = m₂) , (x₂ = 0, y₂ = mc), (x₃ = 1, y₃ = m₂)
                # and then solve for the vertex of the parabola.
                λ = (m₂ - m₁) / (2*(m₂ - 2*mc + m₁))
                out₂[r,c] = SVector(λ*d₁, λ*d₂)
            else
                out₁[r,c] = zero(eltype(mag))
                out₂[r,c] = zero(eltype(out₂))
            end
        end
    end
    return nothing
end


const superscripts = Dict(0=>"⁰",1=>"¹",2=>"²",3=>"³",4=>"⁴",5=>"⁵",6=>"⁶",7=>"⁷",8=>"⁸",9=>"⁹")

## JacobiQ and JacobiQWeight

function Base.show(io::IO,S::JacobiQ)
    S.a == S.b == 0 ? print(io,"LegendreQ(") : print(io,"JacobiQ($(S.a),$(S.b),")
    show(io,domain(S))
    print(io,")")
end

function Base.show(io::IO,s::JacobiQWeight)
    d=domain(s)
    #TODO: Get shift and weights right
    if s.α==s.β
        print(io,"(x^2-1)^$(s.α)[")
    elseif s.α==0
        print(io,"(x-1)^$(s.β)[")
    elseif s.β==0
        print(io,"(x+1)^$(s.α)[")
    else
        print(io,"(x+1)^$(s.α)*(x-1)^$(s.β)[")
    end

    show(io,s.space)
    print(io,"]")
end

## CauchyWeight

function Base.show(io::IO,s::CauchyWeight)
    #TODO: Get shift and weights right
    if order(s) == 0
        print(io,"π⁻¹log|y-x|[")
    elseif order(s) ≥ 1
        print(io,"π⁻¹|y-x|⁻"*mapreduce(x->superscripts[x],*,reverse!(digits(order(s))))*"[")
    else
        print(io,"π⁻¹|y-x|"*mapreduce(x->superscripts[x],*,reverse!(digits(-order(s))))*"[")
    end
    show(io,s.space)
    print(io,"]")
end

## GreensFun

function Base.show(io::IO,G::GreensFun)
    print(io,"GreensFun with kernels:\n{")
    Kernels = kernels(G)
    for i in 1:length(G)
        print(io,"\n ")
        show(io,Kernels[i])
    end
    print(io,"\n}")
end


include("Segment.jl")
include("UnionDomain.jl")
include("PiecewiseSegment.jl")
include("Point.jl")

convert(::Type{Space},d::ClosedInterval) = Space(Domain(d))

#issubset between domains

issubset(a::IntervalOrSegment{T}, b::PiecewiseSegment{T}) where {T<:Real} =
    a⊆Segment(first(b.points),last(b.points))

function setdiff(b::Segment,a::Point)
    if !(a ⊆ b)
        b
    elseif leftendpoint(b) == a.x  || rightendpoint(b) == a.x
        b
    else
        Segment(leftendpoint(b),a.x) ∪ Segment(a.x,rightendpoint(b))
    end
end

# sort

function Base.setdiff(d::SegmentDomain,p::Point)
    x = Number(p)
    (x ∉ d || x ≈ leftendpoint(d) || x ≈ rightendpoint(d)) && return d
    DifferenceDomain(d,p)
end

# multivariate domainxs

include("multivariate.jl")

affine_setdiff(d::Domain, ptsin::UnionDomain) =
    _affine_setdiff(d, Number.(elements(ptsin)))

affine_setdiff(d::Domain, ptsin::WrappedDomain{<:AbstractVector}) =
    _affine_setdiff(d, ptsin.domain)

function _affine_setdiff(d::Domain, p::Number)
    isempty(d) && return d
    p ∉ d  && return d
    a,b = endpoints(d)
    if leftendpoint(d) > rightendpoint(d)
        a,b = b,a
    end
    UnionDomain(a..p, p..b) # TODO: Clopen interval
end

function _affine_setdiff(d::Domain, pts)
    isempty(pts) && return d
    tol=sqrt(eps(arclength(d)))
    da=leftendpoint(d)
    isapprox(da,pts[1];atol=tol) && popfirst!(pts)
    isempty(pts) && return d
    db=rightendpoint(d)
    isapprox(db,pts[end];atol=tol) && pop!(pts)

    sort!(pts)
    leftendpoint(d) > rightendpoint(d) && reverse!(pts)
    filter!(p->p ∈ d,pts)

    isempty(pts) && return d
    length(pts) == 1 && return d \ pts[1]

    ret = Array{Domain}(undef, length(pts)+1)
    ret[1] = Domain(leftendpoint(d) .. pts[1])
    for k = 2:length(pts)
        ret[k] = Domain(pts[k-1]..pts[k])
    end
    ret[end] = Domain(pts[end] .. rightendpoint(d))
    UnionDomain(ret)
end


@testset "mmd" begin
    T = Float32
    k = GaussianKernel(1f0)
    x = rand(T, 5, 2) |> gpu
    y = rand(T, 5, 2) |> gpu
    m = mmd(k, x, y)
    @test typeof(m) == T

    gs = Flux.gradient(mmd, k, x, y)

    T = Float64
    k = GaussianKernel(1)
    x = rand(T, 5, 2)
    y = rand(T, 5, 2)
    m = mmd(k, x, y)
    @test typeof(m) == T
end

@testset "mmd2 and variance" begin
    o = ones(Float64, 3, 3)
    z = zeros(Float64, 3, 3)
    @test mmd2_and_variance(o, o, o) == (0, 0)
    @test mmd2_and_variance(z, z, z) == (0, 0)
    @test mmd2_and_variance(o, z, z) == (1, 0)
    @test mmd2_and_variance(z, z, o) == (1, 0)
    @test mmd2_and_variance(o, o, z) == (-1, 0)
    @test mmd2_and_variance(o, z, z) == (1, 0)
    @test mmd2_and_variance(z, o, z) == (-2, 0)

    a = [1 2 3; 4 5 6; 7 8 9]
    b = (a .+ 1)'
    c = (a .+ 2)'
    a = a'
    m, v = mmd2_and_variance(a, b, c)
    @test m ≈ 0
    @test v ≈ -3.1111111

    a = [1 1 1; 1 1 1; 0 1 0]'
    b = [1 1 1; 1 0 1; 0 1 0]'
    c = [0 1 1; 1 0 1; 0 1 0]'
    m, v = mmd2_and_variance(a, b, c)
    @test m ≈ 0.33333333
    @test v ≈ -0.0185185185185

    a = [0.31067363 0.17995522 0.69170826
        0.10956897 0.43168305 0.13072019
        0.98321756 0.66711701 0.4277054]'

    b = [0.6308643 0.72541977 0.26667574
        0.24141601 0.56088218 0.2872887
        0.12801543 0.43956629 0.0703913]'

    c = [0.68779085 0.68050291 0.47736484
        0.49306397 0.07206057 0.10580517
        0.99297355 0.90194579 0.32241892]'

    m, v = mmd2_and_variance(a, b, c)
    @test m ≈ 0.32443064
    @test v ≈ 0.33610499
end


function parity(s::Int, i::Int, j::Int)
    c = 0
    for k=min(i, j)+1:max(i, j)-1
        c += (s>>k)&1
    end
    return isodd(c) ? -1. : 1.
end


## Tuning Feature
x = -3:0.002:3
y = gaussianf.(x,σ=1)
hw = halfwidth(y,circ=false,x=x)

x = -5:0.002:1
y = gaussianf.(x,σ=1)
hw = halfwidth(y,circ=false,x=x)
hw = halfwidth(y,circ=true,x=x)

x = -1:0.002:5
y = gaussianf.(x,σ=1)
hw = halfwidth(y,circ=false,x=x)
hw = halfwidth(y,circ=true,x=x)

## Circular Tuning Regression

# save(joinpath(@__DIR__,"circdata.jld2"),"y",y,"x",x,"y1",y1,"x1",x1)
y,x,y1,x1=load(joinpath(@__DIR__,"circdata.jld2"),"y","x","y1","x1")

# Tuning Curve, von Mises and Generalized von Mises
plotcondresponse(y,DataFrame(Ori=rad2deg.(x)),u=0,projection=:polar)
plotcondresponse(y1,DataFrame(Ori=rad2deg.(x1)),u=1,projection=:polar)

plot([vmf,gvmf],0,2π,label=["vmf" "gvmf"],projection=:polar)

# GvM and vM Fitting
mfit = fitmodel(:gvm,x,y)
plot(x->mfit.fun(x,mfit.param),0,2π,proj=:polar)

mfit = fitmodel(:vmn2,x1,y1)
plot(x->mfit.fun(x,mfit.param),0,2π,proj=:polar)

## Tuning Properties
f=factorresponsefeature(rad2deg.(x),y,factor=:Ori)

plot(x->f.fit.gvm.fun(x,f.fit.gvm.param),0,2π,lw=2)
scatter!(mod.(x.+0.5π,2π),y)

plot(x->f.fit.vmn2.fun(x,f.fit.vmn2.param),0,2π,lw=2)
scatter!(x,y)

f1=factorresponsefeature(rad2deg.(x1),y1,factor=:Ori)

plot(x->f1.fit.gvm.fun(x,f1.fit.gvm.param),0,2π,lw=2)
scatter!(mod.(x1.+0.5π,2π),y1)

plot(x->f1.fit.vmn2.fun(x,f1.fit.vmn2.param),0,2π,lw=2)
scatter!(x1,y1)


#=
This file implements the algorithm described in:
"A Schur-Parlett Algorithm for Computing Matrix Functions"
(Davies and Higham, 2003)
(with some modifications to handle real matrices more efficiently).
=#

#=
blockpattern groups eigenvalues in sets, implementing Algorithm 4.1 of the paper.
The function returns S and p. S is the pattern, where S[i] = s means that
vals[i] has been assigned to set s. p is the number of sets identified.
=#
delta = 0.1
function blockpattern(vals::Vector{Complex128}, schurtype::Type)::Tuple{Vector{Int64}, Int64}
	unassigned = -1
	S = fill(unassigned, length(vals))
	p = 0

	# assign vals[i] to set s.
	function assign(i::Int64, s::Int64)
		lambda = vals[i]
		for k = i:length(vals)
			if vals[k] == lambda
				S[k] = s
			end
		end
	end
	function mergesets(s::Int64, t::Int64)
		if s == t
			return
		end
		smin, smax = minmax(s, t)
		for k = 1:length(vals)
			if S[k] == smax
				S[k] = smin
			elseif S[k] > smax
				S[k] -= 1
			end
		end
		p -= 1
	end

	for i = 1:length(vals)
		if S[i] == unassigned
			p += 1
			assign(i, p)
		end
		for j = i+1:length(vals)
			if abs(vals[i]-vals[j]) <= delta
				if S[j] == unassigned
					assign(j, S[i])
				else
					mergesets(S[i], S[j])
				end
			end
		end
	end

	if schurtype <: Real
		# complex conjugate eigenvalues can't be separated in the
		# real Schur factorization, so we merge the sets involving them:
		i = 1
		while i < length(vals)
			if imag(vals[i]) != 0
				mergesets(S[i], S[i+1])
				i += 2
			else
				i += 1
			end
		end
	end
	return S, p
end

# analogous to ordschur/trsen:
function ordvec!(v::Vector{T}, select::B) where {T, B<:Union{Vector{Bool}, BitVector}}
	ilst = 1
	for ifst = 1:length(v)
		if select[ifst]
			if ifst != ilst
				v[ilst], v[ilst+1:ifst] = v[ifst], v[ilst:ifst-1]
			end
			ilst += 1
		end
	end
end

#=
reorder reorders the Schur decomposition (T, Q, vals) according to pattern S,
using the swapping strategy described in Algorithm 4.2.
The returned vector, blockend, contains the indices at which each block ends
(after the reordering).
=#
function reorder!(T::Matrix{N}, Q::Matrix{N}, vals::Vector{Complex{R}}, S::Vector{Int64}, p::Int64)::Vector{Int64} where {
	R<:Real, N<:Union{R, Complex{R}}}

	# for each set, calculate its mean position in S:
	pos = zeros(Float64, p)
	cou = zeros(Int64, p)
	for i = 1:length(S)
		set = S[i]
		pos[set] += i
		cou[set] += 1
	end
	pos ./= cou

	blockend = zeros(Int64, p)
	for set = 1:p
		numordered = (set == 1) ? 0 : blockend[set-1]
		minset = indmin(pos)
		select = [i <= numordered || S[i] == minset for i = 1:length(S)]
		ordschur!(T, Q, select)
		ordvec!(vals, select)
		ordvec!(S, select)
		blockend[set] = count(select)
		pos[minset] = Inf
	end
	return blockend
end

#=
taylorf computes f(T) using Taylor as described in Algorithm 2.6.
=#
function taylorf(f::Func, T::Mat, shift::N)::Mat where {
	Func, N<:Number, Mat<:Union{Matrix{N}, UpperTriangular{N, Matrix{N}}}}

	maxiter = 300
	lookahead = 10
	tay = f(shift + Taylor1(N, 20+lookahead))
	M = T - shift*Mat(eye(T))
	P = M
	F = tay.coeffs[1]*Mat(eye(T))
	for k = 1:maxiter
		needorder = k + lookahead
		@assert needorder <= tay.order + 1
		if needorder > tay.order
			tay = f(shift + Taylor1(N, min(maxiter+lookahead, tay.order*2)))
		end

		Term = tay.coeffs[k+1] * P
		F += Term
		normP = norm(P, Inf)
		P *= M
		small = eps()*norm(F, Inf)
		if norm(Term, Inf) <= small
			#=
			The termination condition is loosely inspired to the one in the paper.
			We check that the next lookahead terms in the series are small,
			estimating ||M^(k+r)|| with max(||M^k||, ||M^(k+1)||).
			Works well in practice, at least for exp, log and sqrt.
			=#
			delta = maximum(abs.(tay.coeffs[k+2:k+1+lookahead]))
			if delta*max(normP, norm(P, Inf)) <= small
				return F
			end
		end
	end
	error("Taylor did not converge.")
end

#=
blockf computes f(T) where T is an "atomic" block.
=#
function blockf(f::Func, T::Matrix{N}, vals::Vector{Complex{R}})::Matrix{N} where {
	Func, R<:Real, N<:Union{R, Complex{R}}}

	n = size(T, 1)
	if n == 1
		return f.(T)
	end
	if N <: Complex
		# T is complex and triangular with one cluster of eigenvalues.
		return taylorf(f, UpperTriangular(T), mean(vals))
	end
	if all(imag(vals) .== 0)
		# T is real and triangular with one cluster of real eigenvalues.
		return taylorf(f, UpperTriangular(T), mean(real(vals)))
	end
	if any(abs.(imag(vals)) .<= delta/2)
		# T is real and quasi-triangular with one cluster of complex
		# conjugate eigenvalues (and possibly some real eigenvalue).
		return taylorf(f, T, mean(real(vals)))
	end
	if n == 2
		#=
		T is real and 2x2 with two well-separated complex conjugate eigenvalues.
		We use the Hermite interpolating polynomial obtained with the
		Lagrange-Hermite formula, simplified assuming f(conj(x)) == conj(f(x)).
		=#
		v1, v2 = vals[1], vals[2]
		@assert v1 == conj(v2)
		psi1 = f(v1)/(v1 - v2)
		return 2.0*real(psi1*(T - [v2 0; 0 v2]))
	end
	#=
	T is real with two well-separated, complex conjugate clusters of eigenvalues.
	We use the complex Schur factorization to further divide T.
	Doing complex Schur at the block level is more efficient than using complex
	arithmetic for the whole matrix.
	=#
	U, Z, vals = schur(Matrix{Complex{R}}(T))::Tuple{Matrix{Complex{R}}, Matrix{Complex{R}}, Vector{Complex{R}}}
	select = imag(vals) .> 0
	ordschur!(U, Z, select)
	ordvec!(vals, select)
	@assert count(select) == n÷2
	F = Z*recf(f, U, vals, [n÷2, n])*Z'
	realF = real(F)
	if norm(imag(F), Inf) > sqrt(eps())*norm(realF, Inf)
		warn("T is real but f(T) has non-negligible imaginary part.")
		# note: we need f(T) to be real, because trsyl can't handle
		# quasi-triangular complex matrices.
	end
	return realF
end

#=
recf computes f(T) using the Parlett recurrence (T is block-triangular).
The paper suggests to do the recurrence iteratively (Algorithm 5.1),
we do it recursively as this version performs better.
=#
function recf(f::Func, T::Matrix{N}, vals::Vector{Complex{R}}, blockend::Vector{Int64})::Matrix{N} where {
	Func, R<:Real, N<:Union{R, Complex{R}}}

	@assert length(blockend) > 0
	if length(blockend) == 1
		return blockf(f, T, vals)
	end

	# split T in 2x2 superblocks of size ~n/2:
	n = size(T, 1)
	b = indmin(abs.(blockend .- n/2))
	bend = blockend[b]
	T11, T12 = T[1:bend, 1:bend], view(T, 1:bend, bend+1:n)
	T21, T22 = view(T, bend+1:n, 1:bend), T[bend+1:n, bend+1:n]

	F11 = recf(f, T11, vals[1:bend], blockend[1:b])
	F22 = recf(f, T22, vals[bend+1:n], blockend[b+1:end] .- bend)

	# T11*F12 - F12*T22 = F11*T12 - T12*F22
	F12, scale = LAPACK.trsyl!('N', 'N', T11, T22, F11*T12 - T12*F22, -1)

	return [F11 F12/scale; zeros(T21) F22]
end

#=
Computes f(A) using the Schur-Parlett algorithm.
When A is a real matrix, computation will be done mostly in real arithmetic
and the algorithm will assume f(conj(x)) == conj(f(x)).
=#
function schurparlett(f::Func, A::Matrix{N})::Matrix{N} where {Func, N<:Union{Float64, Complex128}}
	T, Q, vals = schur(A)
	return schurparlett(f, T, Q, vals)
end

#=
Analogous to schurparlett(f, A), but A is provided as its Schur decomposition.
=#
function schurparlett(f::Func, T::Matrix{Float64}, Q::Matrix{Float64}, vals::Vector{Float64})::Matrix{Float64} where {Func}
	return schurparlett(f, T, Q, Vector{Complex128}(vals))
end
function schurparlett(f::Func, T::Matrix{N}, Q::Matrix{N}, vals::Vector{Complex128})::Matrix{N} where {
	Func, N<:Union{Float64, Complex128}}

	if size(T, 1) != size(T, 2)
		throw(DimensionMismatch("T must be square"))
	end
	if size(Q) != size(T) || length(vals) != size(T, 1)
		throw(DimensionMismatch("T, Q, vals dimension mismatch"))
	end
	if size(T, 1) == 0
		error("T is empty")
	end

	d = diag(T)
	D = diagm(d)
	if norm(T-D, Inf) <= 1000*eps()*norm(D, Inf) # T is diagonal
		return Q*diagm(f.(d))*Q'
	end

	S, p = blockpattern(vals, N)
	blockend = reorder!(T, Q, vals, S, p)
	return Q*recf(f, T, vals, blockend)*Q'
end


#handle case with 1 link index at edges.
function fix_autoMPO1(W::ITensor)::ITensor
    ils=filterinds(W,"Link")
    iss=filterinds(W,"Site")
    @assert length(ils)==1
    @assert length(iss)==2
    il=ils[1] #link index
    Dw=dim(il)
    p=collect(1:Dw)
    p[2],p[Dw]=p[Dw],p[2]
    W1=ITensor(il,iss...)
    for js in eachindval(iss)
        for jl in eachindval(il)
            Sl=slice(W,jl)
            assign!(W1,Sl,il=>p[jl.second])
        end
    end
    return W1
end

function fix_autoMPO2(W::ITensor)::ITensor
    ils=filterinds(W,"Link")
    if length(ils)==1
        return fix_autoMPO1(W)
    end
    iss=filterinds(W,"Site")
    @assert length(ils)==2
    @assert length(iss)==2
    d,n,r,c=parse_links(W)
    Dw1,Dw2=dim(r),dim(c)
    #
    #  set up perm arrays to swap row and col 2 with N
    #
    pr=collect(1:Dw1)
    pc=collect(1:Dw2)
    pr[2],pr[Dw1]=pr[Dw1],pr[2]
    pc[2],pc[Dw2]=pc[Dw2],pc[2]
    W1=ITensor(r,c,iss...)
    for jr in eachindval(r)
        for jc in eachindval(c)
            sl=slice(W,jr,jc)
            assign!(W1,sl,r=>pr[jr.second],c=>pc[jc.second])
        end
    end
    return W1
end

@doc """
    fix_autoMPO!(H::MPO)

Convert AutoMPO output into lower regular form.  If `Dw` x `Dw'` are the dimensions of
the Link indices for any site then this routine simply does 2 swaps:
1. Swap row 2 with row `Dw`
2. Swap column 2 with column `Dw'`

# Arguments
- `H::MPO` : MPO to be characterized.

MPO is modified in place into lower regular form

# Examples

```julia
julia> using ITensors
julia> using ITensorMPOCompression
julia> N=5
julia> sites = siteinds("S=1/2",5);
julia> ampo = OpSum();
julia> for j=1:N-1
julia>     add!(ampo,1.0, "Sz", j, "Sz", j+1);
julia>     add!(ampo,0.5, "S+", j, "S-", j+1);
julia>     add!(ampo,0.5, "S-", j, "S+", j+1);
julia> end
julia> H=MPO(ampo,sites);
julia> pprint(H[2])
I 0 0 0 0
0 I S S S
S 0 0 0 0
S 0 0 0 0
S 0 0 0 0

julia> fix_autoMPO!(H)
julia> pprint(H[2])
I 0 0 0 0
S 0 0 0 0
S 0 0 0 0
S 0 0 0 0
0 S S S I
```
"""
function fix_autoMPO!(H::MPO)
    N=length(H)
    for n in 1:N
        H[n]=fix_autoMPO2(H[n])
    end
end

make_Heisenberg_AutoMPO(sites,NNN::Int64,hz::Float64,ul::reg_form,J::Float64=1.0)::MPO =
    make_Heisenberg_AutoMPO(sites,NNN,hz,J)

@doc """
    make_Heisenberg_AutoMPO(sites,[NNN=1[,hz=0.0[,J=1.0]]])

Use `ITensor.autoMPO` to build up a Heisenberg model Hamiltonian with up to `NNN` neighbour
interactions.  The interactions are hard coded to decay like J/(i-j) between sites `i` and `j`.
The MPO is returned in lower regular form.

# Arguments
- `sites` : Site set defining the lattice of sites.
- `NNN::Int64` : Number of nearest neighbour interactions to include in `H`
- `hz::Float64=0.0` : External magnetic field in `z` direction.
- `J::Float64=1.0` : Nearest neighbour interaction strength.

"""
function make_Heisenberg_AutoMPO(sites,NNN::Int64=1,hz::Float64=0.0,J::Float64=1.0)::MPO
    N=length(sites)
    @assert(N>=NNN)
    ampo = OpSum()
    for j=1:N
        add!(ampo, hz   ,"Sz", j)
    end
    for dj=1:NNN
        f=J/dj
        for j=1:N-dj
            add!(ampo, f    ,"Sz", j, "Sz", j+dj)
            add!(ampo, f*0.5,"S+", j, "S-", j+dj)
            add!(ampo, f*0.5,"S-", j, "S+", j+dj)
        end
    end
    mpo=MPO(ampo,sites)
    fix_autoMPO!(mpo) #swap row[2]<->row[Dw] and col[2]<->col[Dw]
    return mpo
end

make_transIsing_AutoMPO(sites,NNN::Int64,hx::Float64,ul::reg_form,J::Float64=1.0)::MPO =
     make_transIsing_AutoMPO(sites,NNN,hx,J)

@doc """
    make\\_transIsing\\_AutoMPO(sites,NNN[,hx=0.0[,J=1.0]])

 Use `ITensor.autoMPO` to build up a transverse Ising model Hamiltonian with up to `NNN` neighbour
 interactions.  The interactions are hard coded to decay like J/(i-j) between sites `i` and `j`.
 The MPO is returned in lower regular form.

 # Arguments
 - `sites` : Site set defining the lattice of sites.
 - `NNN::Int64` : Number of nearest neighbour interactions to include in `H`
 - `hx::Float64=0.0` : External magnetic field in `x` direction.
 - `J::Float64=1.0` : Nearest neighbour interaction strength.

 """
function make_transIsing_AutoMPO(sites,NNN::Int64=1,hx::Float64=0.0,J::Float64=1.0)::MPO
    do_field = hx!=0.0
    N=length(sites)
    @assert(N>NNN)
    ampo = OpSum()
    if do_field
        for j=1:N
            add!(ampo, hx   ,"Sx", j)
        end
    end
    for dj=1:NNN
        f=J/dj
        for j=1:N-dj
            add!(ampo, f    ,"Sz", j, "Sz", j+dj)
        end
    end
    mpo=MPO(ampo,sites)
    fix_autoMPO!(mpo) #swap row[2]<->row[Dw] and col[2]<->col[Dw]
    return mpo
end

@doc """
    make_transIsing_MPO(sites[,NNN=1[,hx=0.0[,ul=lower[,J=1.0]]]])

Directly coded build up of a transverse Ising model Hamiltonian with up to `NNN` neighbour
interactions.  The interactions are hard coded to decay like J/(i-j) between sites `i` and `j`.

# Arguments
- `sites` : Site set defining the lattice of sites.
- `NNN::Int64=1` : Number of nearest neighbour interactions to include in `H`
- `hx::Float64=0.0` : External magnetic field in `x` direction.
- `ul::reg_form=lower` : build H with `lower` or `upper` regular form.
- `J::Float64=1.0` : Nearest neighbour interaction strength.

"""
function make_transIsing_MPO(sites,NNN::Int64=1,hx::Float64=0.0,ul::reg_form=lower,J::Float64=1.0;kwargs...)::MPO
    pbc::Bool=get(kwargs,:pbc,false)
    mpo=MPO(sites) #make and MPO only to get the indices
    prev_link=Nothing
    for n in 1:length(sites)
        mpo[n]=make_transIsing_op(mpo[n],prev_link,NNN,J,hx,ul)
        prev_link=filterinds(mpo[n],tags="Link,l=$n")[1]
    end
    if pbc
        return mpo #leave end Ws as full matrices, i.e. periodic boundary conditions.
    else
        return ITensorMPOCompression.to_openbc(mpo) #contract with l* and *r at the edges.
    end
end

#  It turns out that the trans-Ising model
#  Dw = 2+Sum(i,i=1..NNN) =2 + NNN*(NNN+1)/2
#
function transIsing_Dw(NNN::Int64)::Int64
    return 2+NNN*(NNN+1)/2
end

function make_Ising_index(Dw::Int64,tags::String,use_qn::Bool,dir)
    if (use_qn)
        ind=Index(QN("Sz",0)=>Dw;dir=dir,tags=tags)
    else
        ind=Index(Dw,tags)
    end
    return ind
end

# NNN = Number of Nearest Neighbours, for example
#    NNN=1 corresponds to nearest neighbour
#    NNN=2 corresponds to nearest and next nearest neighbour
function make_transIsing_op(Wref::ITensor,prev_link,NNN::Int64,J::Float64,hx::Float64=0.0,ul::reg_form=lower)::ITensor
    @assert NNN>=1
    do_field = hx!=0.0
    Dw::Int64=transIsing_Dw(NNN)
    is=filterinds(Wref,tags="Site")[1] #get any site index for generating operators
    use_qn=hasqns(is)
    d,n,r,c=parse_links(Wref)
    if tags(r)==TagSet("")
        r=make_Ising_index(Dw,"Link,l=$(n-1)",use_qn,ITensors.In)
    else
        r=prev_link
    end
    if tags(c)==TagSet("")
        c=make_Ising_index(Dw,"Link,l=$n",use_qn,ITensors.Out)
    else
        c=redim(c,Dw)
    end
    iblock=1;

    W=ITensor(r,dag(c),is,dag(is'))
    Id=op(is,"Id")
    Sz=op(is,"Sz")
    if do_field
        Sx=op(is,"Sx")
    end
    assign!(W,Id,r=>1 ,c=>1 )
    assign!(W,Id,r=>Dw,c=>Dw)

    if ul==lower
        if do_field
            assign!(W ,hx*Sx,r=>Dw,c=>1); #add field term
        end
        #very hard to explain this without a diagram.
        for iNN in 1:NNN
            assign!(W,Sz,r=>iblock+1,c=>1)
            for jNN in 1:iNN-1
                assign!(W,Id,r=>iblock+1+jNN,c=>iblock+jNN)
            end
            Jn=J/(iNN) #interactions need to decay with distance in order for H to extensive
            assign!(W,Jn*Sz,r=>Dw,c=>iblock+iNN)
            iblock+=iNN
        end
    else
        if do_field
            assign!(W,hx*Sx,r=>1,c=>Dw ); #add field term
        end
        #very hard to explain this without a diagram.
        for iNN in 1:NNN
            assign!(W,Sz,r=>1,c=>iblock+1)
            for jNN in 1:iNN-1
                assign!(W,Id,r=>iblock+jNN,c=>iblock+1+jNN)
            end
            Jn=J/(iNN) #interactions need to decay with distance in order for H to extensive
            assign!(W,Jn*Sz,r=>iblock+iNN,c=>Dw)
            iblock+=iNN
        end
    end
    return W
end

function to_openbc(mpo::MPO)::MPO
    N=length(mpo)
    l,r=get_lr(mpo)
    mpo[1]=l*mpo[1]
    mpo[N]=mpo[N]*r
    @assert length(filterinds(inds(mpo[1]),tags="Link"))==1
    @assert length(filterinds(inds(mpo[N]),tags="Link"))==1
    return mpo
end

function get_lr(mpo::MPO)::Tuple{ITensor,ITensor}
    ul::reg_form = is_lower_regular_form(mpo,1e-14) ? lower : upper

    N=length(mpo)
    W1=mpo[1]
    llink=filterinds(inds(W1),tags="l=0")[1]
    l=ITensor(0.0,dag(llink))

    WN=mpo[N]
    rlink=filterinds(inds(WN),tags="l=$N")[1]
    r=ITensor(0.0,dag(rlink))
    if ul==lower
        l[llink=>dim(llink)]=1.0
        r[rlink=>1]=1.0
    else
        l[llink=>1]=1.0
        r[rlink=>dim(rlink)]=1.0
    end

    return l,r
end

function fast_GS(H::MPO,sites)::Tuple{Float64,MPS}
    psi0  = randomMPS(sites,length(H))
    sweeps = Sweeps(5)
    setmaxdim!(sweeps, 2,4,8,16,32)
    setcutoff!(sweeps, 1E-10)
    E,psi= dmrg(H,psi0, sweeps;outputlevel=0)
    return E,psi
end


## winston.jl  Winston plotting routines for ROC and DET
## (c) 2014--2015 David A. van Leeuwen
##
## Licensed under the MIT software license, see LICENSE.md

@require Winston begin
    import Winston

    function rocplot(r::Roc, nr=1; ch=false, title="ROC", xlabel="Pfa", ylabel="Pmiss")
        col = string("krgmcb"[(nr-1) % 6 + 1])
        lty = ["-", "--", ";"][div(nr-1,6) + 1]
        if ch
            chi = find(r.ch)
            x, y = r.pfa[chi], r.pmiss[chi]
        else
            x, y = r.pfa, r.pmiss
        end
        if nr==1
            p = Winston.plot(x, y, lty*col)
            Winston.xlabel(xlabel)
            Winston.ylabel(ylabel)
            Winston.title(title)
        else
            p = Winston.oplot(x, y, lty * col)
        end
        return p
    end
    Winston.plot(r::Roc, nr=1; ch=false, title="ROC", xlabel="Pfa", ylabel="Pmiss") = rocplot(r::Roc, nr; title=title, xlabel=xlabel, ylabel=ylabel)

    function detplot(r::Roc, nr=1; xlabel="Pfa (%)", ylabel="Pmiss (%)", title="DET plot")
        if nr==1
            grid = [0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 40]
            ticklabels = AbstractString[]
            for x in grid
                if x<1
                    push!(ticklabels, @sprintf("%3.1f", x))
                else
                    push!(ticklabels, @sprintf("%1.0f", x))
                end
            end
            grid ./= 100
            p = Winston.plot()
            Winston.setattr(p, :aspect_ratio, 1) # essential for DET plots!
            Winston.title(title)
            Winston.xlim(qnorm(0.8e-3), qnorm(0.55))
            Winston.ylim(qnorm(0.8e-3), qnorm(0.55))
            for axis in (p.x1, p.y1)
                Winston.setattr(axis, "draw_subticks", false)
                Winston.setattr(axis, "ticks", qnorm(grid))
                Winston.setattr(axis, "ticklabels", ticklabels)
                Winston.setattr(axis, "tickdir", 1)
                Winston.setattr(axis, "draw_grid", true)
            end
            for axis in (p.x2, p.y2)
                Winston.setattr(axis, "draw_ticks", false)
            end
            Winston.xlabel(xlabel)
            Winston.ylabel(ylabel)
        end
        col = string("krgmcb"[(nr-1) % 6 + 1])
        lty = ["-", "--", ";"][div(nr-1,6) + 1]
        p = Winston.oplot(qnorm(r.pfa), qnorm(r.pmiss), lty * col)
        return p
    end

    function llrplot(r::Roc; title="LLR plot", xlabel="score", ylabel="LLR")
        mi = max(minimum(r.llr), minimum(r.θ))
        ma = min(maximum(r.llr), maximum(r.θ))
        ran = [mi, ma]
        p = Winston.plot(r.θ, r.llr)
        Winston.xlim(ran)
        Winston.ylim(ran)
        Winston.title(title)
        Winston.xlabel(xlabel)
        Winston.ylabel(ylabel)
        return p
    end

    function apeplot(r::Roc; xmin=-7, xmax=7, title="Applied Probability of Error", xlabel="prior log odds", ylabel="Bayes error rate")
        lo = collect(xmin:0.01:xmax)
        be = ber(r, lo)
        mbe = minber(r, lo)
        defbe = 1 ./ normfactor(lo)
        p = Winston.plot(lo, be, "-r")
        Winston.oplot(lo, mbe, "-g")
        Winston.oplot(lo, defbe, "--k")
        Winston.ylim(0, 1.1*maximum(be))
        Winston.title(title)
        Winston.xlabel(xlabel)
        Winston.ylabel(ylabel)
        return p
    end

    function nbeplot(r::Roc; xmin = -10, xmax = 5, title="Normalized Bayes' Error", xlabel="prior log odds", ylabel="normalized Bayes error")
        lo = collect(xmin:0.01:xmax)
        norm = normfactor(lo)
        bfa, bmiss = [be .* norm for be in ber_famiss(r, lo)]
        nbe = bfa + bmiss
        mbfa, mbmiss = [mbe .* norm for mbe in minber_famiss(r, lo)]
        mnbe = minber(r, lo) .* norm
        p = Winston.plot(lo, nbe, "-r")
        Winston.oplot(lo, mnbe, "-g")
        Winston.oplot(lo, bfa, "--r")
        Winston.oplot(lo, bmiss, "-.r")
        Winston.oplot(lo, mbfa, "--g")
        Winston.oplot(lo, mbmiss, "-.g")
        Winston.ylim(0, 1.1)
        Winston.title(title)
        Winston.xlabel(xlabel)
        Winston.ylabel(ylabel)
        return p
    end

end


#=
    Tiled Cholesky factorization in pure Julia. The serial performance
    essentially comes from the `TriangularSolve` and `LoopVectorization`
    packages. The parallelizatino is handled by `DataFlowTask`s.
=#

cholesky(A::Matrix,args...) = cholesky!(deepcopy(A),args...)

function cholesky!(A::Matrix,s=TILESIZE[],tturbo::Val{T}=Val(false)) where {T}
    _cholesky!(PseudoTiledMatrix(A,s),tturbo)
end

# tiled cholesky factorization
function _cholesky!(A::PseudoTiledMatrix,tturbo::Val{T}=Val(false)) where {T}
    m,n = size(A) # number of blocks
    for i in 1:m
        # _chol!(A[i,i],UpperTriangular)
        @dspawn _chol!(A[i,i],UpperTriangular,tturbo) (A[i,i],) (RW,)
        Aii = A[i,i]
        U = UpperTriangular(Aii)
        L = adjoint(U)
        for j in i+1:n
            Aij = A[i,j]
            # TriangularSolve.ldiv!(L,Aij,tturbo)
            @dspawn TriangularSolve.ldiv!(L,Aij,tturbo) (Aii,Aij) (R,RW)
        end
        for j in i+1:m
            Aij = A[i,j]
            for k in j:n
                # TODO: for k = j, only the upper part needs to be updated,
                # dividing the cost of that operation by two
                Ajk = A[j,k]
                Aji = adjoint(Aij)
                Aik = A[i,k]
                # schur_complement!(Ajk,Aji,Aik,tturbo)
                @dspawn schur_complement!(Ajk,Aji,Aik,tturbo) (Ajk,Aij,Aik) (RW,R,R)
            end
        end
    end
    # wait for all computations before returning
    DataFlowTasks.sync()
    return Cholesky(A.data,'U',zero(BlasInt))
end

# a fork-join approach for comparison with the data-flow parallelism
function _cholesky_forkjoin!(A::PseudoTiledMatrix,tturbo::Val{T}=Val(false)) where {T}
    m,n = size(A) # number of blocks
    for i in 1:m
        _chol!(A[i,i],UpperTriangular,tturbo)
        Aii = A[i,i]
        U = UpperTriangular(Aii)
        L = adjoint(U)
        Threads.@threads for j in i+1:n
            Aij = A[i,j]
            TriangularSolve.ldiv!(L,Aij,tturbo)
        end
        # spawn m*(m+1)/2 tasks and sync them at the end
        @sync for j in i+1:m
            Aij = A[i,j]
            for k in j:n
                Ajk = A[j,k]
                Aji = adjoint(Aij)
                Aik = A[i,k]
                Threads.@spawn schur_complement!(Ajk,Aji,Aik,tturbo)
            end
        end
    end
    return Cholesky(A.data,'U',zero(Int32))
end

# Modified from the generic version from LinearAlgebra (MIT license).
function _chol!(A::AbstractMatrix{<:Real}, ::Type{UpperTriangular},tturbo::Val{T}=Val(false)) where {T}
    Base.require_one_based_indexing(A)
    n = LinearAlgebra.checksquare(A)
    @inbounds begin
        for k = 1:n
            Akk = A[k,k]
            for i = 1:k - 1
                Akk -= A[i,k]*A[i,k]
            end
            A[k,k] = Akk
            Akk, info = _chol!(Akk, UpperTriangular)
            if info != 0
                return UpperTriangular(A), info
            end
            A[k,k] = Akk
            AkkInv = inv(Akk')
            if T
                @tturbo warn_check_args=false for j = k + 1:n
                    for i = 1:k - 1
                        A[k,j] -= A[i,k]*A[i,j]
                    end
                end
                @tturbo warn_check_args=false for j in k+1:n
                    A[k,j] = AkkInv*A[k,j]
                end
            else
                @turbo warn_check_args=false for j = k + 1:n
                    for i = 1:k - 1
                        A[k,j] -= A[i,k]*A[i,j]
                    end
                end
                @turbo warn_check_args=false for j in k+1:n
                    A[k,j] = AkkInv*A[k,j]
                end
            end
        end
    end
    return UpperTriangular(A), convert(Int32, 0)
end
## Numbers
function _chol!(x::Number, uplo)
    rx = real(x)
    rxr = sqrt(abs(rx))
    rval =  convert(promote_type(typeof(x), typeof(rxr)), rxr)
    rx == abs(x) ? (rval, convert(Int32, 0)) : (rval, convert(Int32, 1))
end


#
# Grid data access
# ================
#
# You may want to use the this package to obtain grid data for your application.
# The [ExtendableGrids](https://github.com/j-fu/ExtendableGrids.jl) package provides a lightweight grid data structure which
# in its core is a `Dict{Type,Any}` and the possibility to [dispatch the return type on the key](https://j-fu.github.io/ExtendableGrids.jl/stable/tdict/).
#
function extract_2d()
    builder=SimplexGridBuilder(Generator=Triangulate)

    p1=point!(builder,0,0)
    p2=point!(builder,1,0)
    p3=point!(builder,1,1)
    p4=point!(builder,0,1)

    facetregion!(builder,1)
    facet!(builder,p1,p2)
    facetregion!(builder,2)
    facet!(builder,p2,p3)
    facetregion!(builder,3)
    facet!(builder,p3,p4)
    facetregion!(builder,4)
    facet!(builder,p4,p1)

    grid=simplexgrid(builder,maxvolume=0.25)

    @show grid[Coordinates]
    @show grid[CellNodes]
    @show grid[CellRegions]
    @show grid[BFaceNodes]
    @show grid[BFaceRegions]
    grid
end
# The output of this call is this:
# ```julia
# grid[Coordinates] = [0.0 1.0 1.0 0.0 0.5; 0.0 0.0 1.0 1.0 0.5]
# grid[CellNodes] = Int32[2 4 5 1; 3 1 3 2; 5 5 4 5]
# grid[CellRegions] = Int32[1, 1, 1, 1]
# grid[BFaceNodes] = Int32[2 3 4 1; 1 2 3 4]
# grid[BFaceRegions] = Int32[1, 2, 3, 4]
# ```
# Thus, the grid is described by five arrays:
# - `2 x npoints` coordinates,
# - `3 x ntriangles` connectivity,
# - `ntriangles` region markers,
# - `2 x nfacets` boundary faces,
# - `nfacets` boundary face markers.


module Tests

include("utilities.jl")

const linear_tests = Dict{String, Function}()
const soc_tests = Dict{String, Function}()
const sd_tests = Dict{String, Function}()

include("term.jl")
include("term_fixed.jl")
include("quartic_constant.jl")
include("quadratic.jl")
include("quartic_ideal.jl")
include("univariate_sum.jl")
include("rearrangement.jl")
include("choi.jl")
include("simple_matrix.jl")
include("horn.jl")
include("concave_then_convex_cubic.jl")
include("lyapunov_switched_system.jl")
include("motzkin.jl")
include("BPT12e399.jl")
include("maxcut.jl")
include("chebyshev.jl")
include("quartic_comparison.jl")
include("options_pricing.jl")
include("sosdemo5.jl")
include("sosdemo9.jl")
include("sosdemo10.jl")

@test_suite linear
@test_suite soc
@test_suite sd

end


# Only the affine constraints

using SparseArrays

include("../other_utils/feasibleCheck.jl")
include("projections.jl")

# ---------------------
# Equality constraints
# ---------------------
"""
    AL_AffineEquality(A, b)

Creates an Affine Equality Constraint of the form `Ax = b`.

To check constraint satisfaction, use:
[`satisified`](@ref) and [`whichSatisfied`](@ref)

To evaluate the constraint, use:
- [`getRaw`](@ref) to evaluate `Ax - b`
- [`getNormToProjVals`](@ref) to evaluate the projection to `Ax = b`
"""
struct AL_AffineEquality <: constraint
    A
    b
end

# Check that a constraint is satisfied (Ax - b = 0)
"""
    satisfied(r::AL_AffineEquality, x) = (r.A * x == r.b)

Check constraint satisfaction. See also [`whichSatisfied`](@ref)
"""
satisfied(r::AL_AffineEquality, x) = (r.A * x == r.b)

"""
    whichSatisfied(r::AL_AffineEquality, x) = (r.A * x .== r.b)

Check constraint satisfaction. See also [`satisfied`](@ref)
"""
whichSatisfied(r::AL_AffineEquality, x) = (r.A * x .== r.b)

# Evaluate the constraint.
"""
    getRaw(r::AL_AffineEquality, x) = r.A * x - r.b

Evaluate the constraint without projection
"""
getRaw(r::AL_AffineEquality, x) = r.A * x - r.b

"""
    getProjVecs(r::AL_AffineEquality, x)

For every row, get the projection onto the line `a'x = b`.

See [`projAffineEq`](@ref) for more information.
"""
function getProjVecs(r::AL_AffineEquality, x)
    #=
    We want to project each constraint
    We write Ax = b
    But this is equivalent to a1'x = b1, a2'x = b2, ..., am'x = bm
    Where ak is the row k in A
    =#
    projVecs = []
    for ind in 1:size(r.b, 1)
        aI = r.A[ind, :]
        bI = r.b[ind]
        push!(projVecs, projAffineEq(aI, bI, x))
    end

    return projVecs
end

"""
    getNormToProjVals(r::AL_AffineEquality, x)

Get the row by row constraint violation based on projections. This is better
than [`getRaw`](@ref) since it doesn't depend on the magnitude of the elements
in `r.A` and `r.b`.

See also [`AL_AffineEquality`](@ref), [`getProjVecs`](@ref),
[`projAffineEq`](@ref)
"""
function getNormToProjVals(r::AL_AffineEquality, x)
    #=
    We want to get the projected vector and calculate the distance between the
    original point and the constraint.
    =#
    projVecs = getProjVecs(r, x)
    pvsDiff = [pv - x for pv in projVecs]

    # NaN removal due to floating point error near 0.0
    return replace(norm.(pvsDiff, 2), NaN=>0.0)
end

# Lastly, we do some calculus
"""
    getGradC(r::AL_AffineEquality, x)

Calculates the gradient of a constraint.

For `r::AL_AffineEquality`, `∇c(x) = A`

The variable "x" is passed to the header to maintain the same function header
but it is not used in the function.
"""
function getGradC(r::AL_AffineEquality, x)
    #=
    We want the gradient of the constraints. This is piecewise in the
    inequality case. But it is a single function in the equality case
    =#
    return r.A
end

"""
    getHessC(r::AL_AffineEquality)

Calculates the hessian of a constraint.

For `r::AL_AffineEquality`, `H(c(x)) = 0`
"""
function getHessC(r::AL_AffineEquality)
    #=
    We want the hessian of the constraints. This is just zero for affine
    constraints, but we need to match the dimensions.
    =#
    return spzeros(size(r.A, 2), size(r.A, 2)) # nxn
end

"""
    getHessC_ALTerm(r::AL_AffineEquality, x, rho = 1)

Calculates the hessian of the augmented lagrangian constraint term.

For `r::AL_AffineEquality`, `H(ρ c(x)'c(x) + λ c(x)) = ρ A'A`

The variable "x" is passed to the header to maintain the same function header
but it is not used in the function.
"""
function getHessC_ALTerm(r::AL_AffineEquality, x, rho = 1)
    #=
    The augmented lagrangian constraint term is of the form:
    ρ c(x)'c(x) + λ c(x)

    where c(x) = Ax - b.

    The Hessian matrix is then
    Htot = ρ (c.H + J.J + H.c) + λ H

    But H = 0 for Affine equalities, so
    Htot = ρ (J.J) = ρ A'A
    =#

    jacob = getGradC(r, x)
    return rho * jacob'jacob
end

# -----------------------
# Inequality constraints
# -----------------------

"""
    AL_AffineInequality(A, b)

Creates an Affine Inequality Constraint of the form `Ax ≤ b`.

To check constraint satisfaction, use:
[`satisified`](@ref) and [`whichSatisfied`](@ref)

To evaluate the constraint, use:
- [`getRaw`](@ref) to evaluate `Ax - b`
- [`getNormToProjVals`](@ref) to evaluate the projection to `Ax ≤ b`
"""
struct AL_AffineInequality <: constraint
    A
    b
end

# Check that a constraint is satisfied (Ax - b ≤ 0)
"""
    satisfied(r::AL_AffineInequality, x)

Check constraint satisfaction. See also [`whichSatisfied`](@ref)
"""
satisfied(r::AL_AffineInequality, x) = isFeasiblePolyHedron(r.A, r.b, x)

"""
    whichSatisfied(r::AL_AffineInequality, x)

Check constraint satisfaction. See also [`satisfied`](@ref)
"""
whichSatisfied(r::AL_AffineInequality, x) = (r.A * x .≤ r.b)

# Evaluate the constraint.
"""
    getRaw(r::AL_AffineInequality, x) = r.A * x - r.b

Evaluate the constraint without projection
"""
getRaw(r::AL_AffineInequality, x) = r.A * x - r.b

"""
    getProjVecs(r::AL_AffineInequality, x)

For every row, get the projection onto the line `a'x = b`.

See [`projAffineEq`](@ref) for more information.
"""
function getProjVecs(r::AL_AffineInequality, x)
    #=
    We want to project each constraint
    We write Ax = b
    But this is equivalent to a1'x = b1, a2'x = b2, ..., am'x = bm
    Where ak is the row k in A
    =#
    projVecs = []
    for ind in 1:size(r.b, 1)
        aI = r.A[ind, :]
        bI = r.b[ind]
        push!(projVecs, projAffineIneq(aI, bI, x))
    end

    return projVecs
end

"""
    getNormToProjVals(r::AL_AffineInequality, x)

Get the row by row constraint violation based on projections. This is better
than [`getRaw`](@ref) since it doesn't depend on the magnitude of the elements
in `r.A` and `r.b`. It also automatically accounts for the inequality.

See also [`AL_AffineInequality`](@ref), [`getProjVecs`](@ref),
[`projAffineEq`](@ref)
"""
function getNormToProjVals(r::AL_AffineInequality, x)
    #=
    We want to get the projected vector and calculate the distance between the
    original point and the constraint.
    =#
    projVecs = getProjVecs(r, x)
    pvsDiff = [pv - x for pv in projVecs]
    return norm.(pvsDiff, 2)
end

# Lastly, we do some calculus
"""
    getGradC(r::AL_AffineInequality, x)

Calculates the gradient of a constraint.

For `r::AL_AffineInequality`, `∇c(x) = A` for rows where the constraint is NOT
satisfied.

Unlike the [`AL_AffineEquality`](@ref) case, the variable "x" is passed to
is needed to determine if the constraint is satisfied.
"""
function getGradC(r::AL_AffineInequality, x)
    #=
    We want the gradient of the constraints. This is piecewise in the
    inequality case. But it is a single function in the equality case
    =#
    APost = zeros(size(r.A))
    rowPassed = whichSatisfied(r, x)
    for row in 1:size(r.A, 1)
        if !rowPassed[row]
            # Constraint is not met
            APost[row, :] = r.A[row, :]
        end
    end
    return APost
end

"""
    getHessC(r::AL_AffineEquality)

Calculates the hessian of a constraint.

For `r::AL_AffineInequality`, `H(c(x)) = 0`
"""
function getHessC(r::AL_AffineInequality)
    #=
    We want the hessian of the constraints. This is just zero for affine
    constraints, but we need to match the dimensions.
    =#
    return spzeros(size(r.A, 2), size(r.A, 2)) # nxn
end

"""
    getHessC_ALTerm(r::AL_AffineInequality, x, rho = 1)

Calculates the hessian of the augmented lagrangian constraint term.

For `r::AL_AffineEquality`, `H(ρ c(x)'c(x) + λ c(x)) = ρ A+'A+`

The variable "x" is passed to the header to maintain the same function header
but it is not used in the function.
"""
function getHessC_ALTerm(r::AL_AffineInequality, x, rho = 1)
    #=
    The augmented lagrangian constraint term is of the form:
    ρ c(x)'c(x) + λ c(x)

    where c(x) = Ax - b.

    The Hessian matrix is then
    Htot = ρ (c.H + J.J + H.c) + λ H

    But H = 0 for Affine equalities, so
    Htot = ρ (J.J) = ρ A'A
    =#

    jacob = getGradC(r, x)
    return sparse(rho * jacob'jacob)
end