# Regular functions for evaluating Chebyshev polynominals

function chebyshev_evaluate(weights::AbstractArray{T,N},x::AbstractArray{R,1},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,R<:Number,N,S<:Integer}

  poly = Array{Array{R,2},1}(undef,N)
  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order[i],normalize_node(x[i],domain[:,i]))
  end

  yhat = zero(T)
  @inbounds for i in CartesianIndices(weights)
    poly_product = poly[1][i[1]]
    @inbounds for j = 2:N
      poly_product *= poly[j][i[j]]
    end
    yhat += weights[i]*poly_product
  end

  return yhat

end

function chebyshev_evaluate(weights::AbstractArray{T,N},x::AbstractArray{R,1},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,R<:Number,N,S<:Integer}

  poly = Array{Array{R,2},1}(undef,N)
  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order,normalize_node(x[i],domain[:,i]))
  end

  yhat = zero(T)
  @inbounds for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N
      poly_product = poly[1][i[1]]
      @inbounds for j = 2:N
        poly_product *= poly[j][i[j]]
      end
      yhat += weights[i]*poly_product
    end
  end

  return yhat

end

function chebyshev_evaluate(cheb_poly::ChebPoly,x::AbstractArray{R,1}) where {R<:Number}

  yhat = chebyshev_evaluate(cheb_poly.weights,x,cheb_poly.order,cheb_poly.domain)

  return yhat

end

function cheb_interp(cheb::ChebInterpRoots,x::AbstractArray{R,1}) where {R<:Number}

  weights = chebyshev_weights_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  yhat    = chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  return yhat

end

function cheb_interp(cheb::ChebInterpExtrema,x::AbstractArray{R,1}) where {R<:Number}

  weights = chebyshev_weights_extrema_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  yhat    = chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  return yhat

end

function cheb_interp(cheb::ChebInterpExtended,x::AbstractArray{R,1}) where {R<:Number}

  weights = chebyshev_weights_extended_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  yhat    = chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  return yhat

end

function cheb_interp(cheb::ChebInterpVertesi,x::AbstractArray{R,1}) where {R<:Number}

  weights = chebyshev_weights_vertesi_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)
  yhat    = chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  return yhat

end

function chebyshev_evaluate(weights::AbstractArray{T,N},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T <: AbstractFloat,N,S <: Integer}

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,order,domain)

  end

  return chebeval

end

function chebyshev_evaluate(weights::AbstractArray{T,N},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T <: AbstractFloat,N,S <: Integer}

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,order,domain)

  end

  return chebeval

end

function chebyshev_evaluate(cheb_poly::ChebPoly)

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(cheb_poly,x)

  end

  return chebeval

end

function cheb_interp(cheb::ChebInterpRoots)

  weights = chebyshev_weights_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  end

  return chebeval

end

function cheb_interp(cheb::ChebInterpExtrema)

  weights = chebyshev_weights_extrema_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  end

  return chebeval

end

function cheb_interp(cheb::ChebInterpExtended)

  weights = chebyshev_weights_extended_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  end

  return chebeval

end

function cheb_interp(cheb::ChebInterpVertesi)

  weights = chebyshev_weights_vertesi_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  function chebeval(x::AbstractArray{R,1}) where {R <: Number}

    return chebyshev_evaluate(weights,x,cheb.order,cheb.domain)

  end

  return chebeval

end


using NonLinearReactionAdvectionDiffusionWithFrontData;
using NonLinearReactionAdvectionDiffusionWithFrontData: f1, f2;
using NonLinearReactionAdvectionDiffusionWithFrontData: ARP_y, ∂ARP_∂y;
using LinearAlgebra;

#' Сначала, сгенирируем априорные данные, на увеличенном количестве узлов.
u_l(t) = -8 #+ cos(2*π * t);
u_r(t) =  4 #+ (1 + sin(2*π * t));
qf(x) = 4*sin(3 * π * x);        # Коэффициент линейного усиления, который в обратной
                                # задаче необходимо определить, но при генерации априорной
                                # информации мы задаем некоторый коэффициент, который,
                                # собственно, после имея априорную информацию и будем определять.
ε = 0.2;                        # Малый параметр при старшей производной
a, b = 0, 1;                    # Область по X
t₀, T = 0.0, 0.28;              # Область по T
N, M = 50, 80;                  # Увеличенное Кол-во разбиений по X, T
h = (b-a)/N;                    # шаг по X
τ = (T-t₀)/M;                   # шаг по T
Xₙ = [a  + n*h for n in 0:N];   # Сетка по Х
Tₘ = [t₀ + m*τ for m in 0:M];   # Сетка по Т
qₙ =      qf.(Xₙ);               # Сеточные значения коэффициента линейного усиления
ulₘ=    u_l.(Tₘ);               # Сеточные значения левого  ГУ
urₘ=    u_r.(Tₘ);               # Сеточные значения правого ГУ
u₀ = u_init.(Xₙ);               # Начальные условия
#
u, XX, TP = solve(u₀, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qₙ);
nothing #hide

#' ## Генерация априорной информации
ϕl = phidetermination(qₙ, ulₘ, Xₙ, N, Tₘ, M);                       # Левый вырожденный корень
ϕr = phidetermination(qₙ, urₘ, Xₙ, N, Tₘ, M, reverseX = true);      # Правый вырожденный корень
ϕ = Φ(ϕl, ϕr, N, M);                                                # Серединный корень
f1_data = f1(ϕ, u, Xₙ, N, M);                                       # Положение переходного слоя
f2_data = f2(f1_data, u, Xₙ, N, M);                                 # Значение функции на переходном слое

# Подготовим массивы, выбросив граничные точки
# ведь тестируемая функция — для внутреннего использования
# И развернем массивы по времени
qq = qₙ[2:N];
y = u₀[2:N];
U = u[2:N, :];
U = reverse(U, dims=2);
f1_data = reverse(f1_data);
f2_data = reverse(f2_data);
Tₘ = reverse(Tₘ);

dl, d, du = NonLinearReactionAdvectionDiffusionWithFrontData.ARP_y(y, 1, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qq, U, f1_data, f2_data)

@test length(dl) == N - 2
@test length(d)  == N - 1
@test length(du) == N - 2

@test count(x -> x == 0, dl) == 0
@test count(x -> x == 0, d)  == 0
@test count(x -> x == 0, du) == 0

tdjac = NonLinearReactionAdvectionDiffusionWithFrontData.∂ARP_∂y(y, 1, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qq, U, f1_data, f2_data)
@test typeof(tdjac) <: Tridiagonal
@test tdjac == Tridiagonal( dl, d, du )

# Сравним с якобианом автодифференцирования
# Передадим **любой** априорный параметр для аппроксимации дельта-функции
# он никак не участвует в вычислении якобиана,
# но при использовании якобиана через автодифференцирование, требуется, чтобы он вызывался ТОЧНО так же
# как и функция, якобиан от которой вычисляется.
jac = NonLinearReactionAdvectionDiffusionWithFrontData.∂adjointRP_∂y(y, 1, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qq, U, f1_data, f2_data, 0.0001)
@test isapprox(tdjac, jac)


function approxextrema(I)
    arr = sample(I)
    return (minimum(arr), maximum(arr))
end

function estimatequantile(array, p)
    try
        return quantile(sample(array; n=1e5), p)
    catch
        @warn "quantile could not be estimated! (maybe only NaNs)"
        return 0
    end
end

function sample(I; n=10000)
    n = min(n, length(I))
    len = ceil(Int, √n) # take len blocks of len elements
    startindices = round.(Int, range(firstindex(I) - 1, lastindex(I) - len; length=len))
    indices = vcat((i .+ (1:len) for i in startindices)...)
    ret = filter(isfinite, I[indices])
    if isempty(ret)
        ret = filter(isfinite, I)
    end
    return ret
end

function get_corner_indices(I; max_length=10)
    d = size(I)
    n = min.(max_length, ceil.(Int, d ./ 3)) # n voxels for each dim
    getrange(num, len) = [1:num, (len-num+1):len] # first and last voxels
    return collect(Iterators.product(getrange.(n, d)...))
end

# estimate noise parameters from corner without signal
function estimatenoise(image)
    corners = get_corner_indices(image)
    (lowestmean, ind) = findmin(mean.(filter(isfinite, image[I...]) for I in corners))
    sigma = std(filter(isfinite, image[corners[ind]...]))
    if isnan(sigma) # no corner available
        # estimation that is correct if half the image is signal and half noise
        sigma = 2estimatesigma_from_quantile(image, 1/4)
        lowestmean = sigma / 2
    end
    return lowestmean, sigma
end

# sigma is only calculated for quantile (biased)
function estimatesigma_from_quantile(image, quantile)
    q = estimatequantile(image, quantile)
    samples = filter(x -> x < q, sample(image))
    return std(samples)
end

function robustmask!(image; maskedvalue=if eltype(image) <: AbstractFloat NaN else 0 end)
    image[.!robustmask(image)] .= maskedvalue
    image
end
function robustmask(weight)
    μ, σ = estimatenoise(weight)
    m = mean(filter(isfinite, weight[weight .> 5σ]))
    maximum((5σ, m/5))
    mask = weight .> maximum((5σ, m/5))
    # remove small holes and minimally grow
    boxsizes=[[3,3] for i in 1:ndims(weight)]
    return gaussiansmooth3d(mask; nbox=2, boxsizes) .> 0.55
end

getcomplex(fnmag::AbstractString, fnphase::AbstractString) = getcomplex(niread(fnmag), niread(fnphase))

function getcomplex(mag, phase)
    minp, maxp = approxextrema(phase)
    mag .* exp.((2im * pi / (maxp - minp)) .* phase)
end

function readfromtextheader(filename, searchstring)
    for line in readlines(open(filename, "r"))
        if occursin(searchstring, line)
            # regex to search for "= " or ": " and return the following non-whitespace characters
            return match(r"(?<=(= |: ))(\S+)", line).match
        end
    end
end

# root sum of squares combination
RSS(mag; dims=ndims(mag)) = dropdims(.√sum(mag.^Float32(2); dims); dims)

function getscaledimage(array, div::Number, offset = 0, type::Symbol = :trans)
    array = reshape(array, size(array)[1:2]) # drops trailing singleton dimensions
    scaled = if offset != 0
        (array .- offset) .* (1 / div) .+ 0.5
    else
        array .* (1 / div)
    end
    scaled[isnan.(scaled) .| (scaled .< 0)] .= 0
    scaled[scaled .> 1] .= 1
    if type == :trans
        scaled = reverse(permutedims(scaled, [2 1]); dims = 1)
    else
    end
    scaled
end

function getscaledimage(array, type::Symbol = :trans)
    scaled = robustrescale(array, 0, 1, threshold=true)
    getscaledimage(scaled, 1, 0, type)
end

robustrescale(array, newmin, newmax; threshold=false, mask=trues(size(array)), datatype=Float64) =
    robustrescale!(datatype.(array), newmin, newmax; threshold, mask)

function robustrescale!(array, newmin, newmax; threshold=false, mask=trues(size(array)))
    array[isnan.(array)] .= minimum(filter(isfinite, array))
    q = [0.01, 0.99] # quantiles
    oldq = estimatequantile(array[mask], q)
    oldrange = (oldq[2] - oldq[1]) / (q[2] - q[1])
    oldmin = oldq[1] - q[1] * oldrange
    newrange = newmax - newmin

    array .= (array .- oldmin) .* (newrange / oldrange) .+ newmin

    if threshold
        array[array .< newmin] .= newmin
        array[array .> newmax] .= newmax
    end
    array
end

function rescale(array, newmin, newmax; datatype=eltype(array))
    rescale!(datatype.(array), newmin, newmax)
end

function rescale!(array, newmin, newmax)
    oldmin, oldmax = approxextrema(array)
    factor = (newmax - newmin) / (oldmax - oldmin)
    array .= (array .- oldmin) .* factor .+ newmin
end

to_dim(a::Real, dim::Int) = to_dim([a], dim)
to_dim(V::AbstractVector, dim::Int) = reshape(V, ones(Int, dim-1)..., :)

"""
    getHIP(mag, phase; echoes=[1,2])
return the hermitian inner product between the specified echoes.
"""
function getHIP(mag, phase; echoes=[1,2])
    e1, e2 = echoes
    compl = zeros(ComplexF64, size(mag)[1:3])
    for iCha in 1:size(mag, 5)
        compl .+= exp.(1.0im .* (phase[:,:,:,e2,iCha] .- phase[:,:,:,e1,iCha])) .* mag[:,:,:,e1,iCha] .* mag[:,:,:,e2,iCha]
    end
    compl
end

function getHIP(compl; echoes=[1,2])
    e1, e2 = echoes
    c = zeros(eltype(compl), size(compl)[1:3])
    for iCha in 1:size(compl, 5)
        c .+=  compl[:,:,:,e2,iCha] .* conj.(compl[:,:,:,e1,iCha])
    end
    return c
end


module SpinPackage

using Printf: @printf, @sprintf
using ..Spatials: PID
using ..DegreesOfFreedom: IID, Internal, Index, FilteredAttributes, OID, Operator, LaTeX, latexformat
using ..Terms: wildcard, constant, Subscript, Subscripts, subscriptexpr, Coupling, Couplings, couplingcenters, Term, TermCouplings, TermAmplitude, TermModulate
using ...Interfaces: rank, kind
using ...Prerequisites: Float, decimaltostr, delta
using ...Mathematics.VectorSpaces: VectorSpace, IsMultiIndexable, MultiIndexOrderStyle
using ...Mathematics.AlgebraOverFields: SimpleID, ID

import ..DegreesOfFreedom: script, otype, isHermitian
import ..Terms: couplingcenter, statistics, abbr
import ...Interfaces: dims, inds, expand, matrix, permute

export SID, Spin, SIndex
export usualspinindextotuple
export SOperator, soptdefaultlatex
export SCID, SpinCoupling
export Heisenberg, Ising, Gamma, DM, Sˣ, Sʸ, Sᶻ
export @sc_str, @heisenberg_str, @ising_str, @gamma_str, @dm_str, @sˣ_str, @sʸ_str, @sᶻ_str
export SpinTerm

const sidtagmap = Dict(1=>'x', 2=>'y', 3=>'z', 4=>'+', 5=>'-')
const sidseqmap = Dict(v=>k for (k, v) in sidtagmap)
const sidajointmap = Dict('x'=>'x', 'y'=>'y', 'z'=>'z', '+'=>'-', '-'=>'+')

"""
    SID <: IID

The spin id.
"""
struct SID <: IID
    orbital::Int
    spin::Float
    tag::Char
    function SID(orbital::Int, spin::Real, tag::Char)
        @assert tag in ('x', 'y', 'z', '+', '-') "SID error: not supported tag($tag)."
        new(orbital, spin, tag)
    end
end
Base.fieldnames(::Type{<:SID}) = (:orbital, :spin, :tag)

"""
    SID(; orbital::Int=1, spin::Real=0.5, tag::Char='z')

Create a spin id.
"""
SID(; orbital::Int=1, spin::Real=0.5, tag::Char='z') = SID(orbital, spin, tag)

"""
    adjoint(sid::SID) -> SID

Get the adjoint of a spin id.
"""
Base.adjoint(sid::SID) = SID(sid.orbital, sid.spin, sidajointmap[sid.tag])

"""
    matrix(sid::SID, dtype::Type{<:Number}=Complex{Float}) -> Matrix{dtype}

Get the matrix representation of a sid.
"""
function matrix(sid::SID, dtype::Type{<:Number} = Complex{Float})
    N = Int(2*sid.spin+1)
    result = zeros(dtype, (N, N))
    for i = 1:N, j = 1:N
        row, col = N+1-i, N+1-j
        m, n = sid.spin+1-i, sid.spin+1-j
        result[row, col] = (sid.tag == 'x') ? (delta(i+1, j)+delta(i, j+1))*sqrt(sid.spin*(sid.spin+1)-m*n)/2 :
                           (sid.tag == 'y') ? (delta(i+1, j)-delta(i, j+1))*sqrt(sid.spin*(sid.spin+1)-m*n)/2im :
                           (sid.tag == 'z') ? delta(i, j)*m :
                           (sid.tag == '+') ? delta(i+1, j)*sqrt(sid.spin*(sid.spin+1)-m*n) :
                                              delta(i, j+1)*sqrt(sid.spin*(sid.spin+1)-m*n)
    end
    return result
end

"""
    Spin <: Internal{SID}

The spin interanl degrees of freedom.
"""
struct Spin <: Internal{SID}
    atom::Int
    norbital::Int
    spin::Float
end
IsMultiIndexable(::Type{Spin}) = IsMultiIndexable(true)
MultiIndexOrderStyle(::Type{Spin}) = MultiIndexOrderStyle('C')
dims(sp::Spin) = (sp.norbital, length(sidtagmap))
inds(sid::SID, ::Spin) = (sid.orbital, sidseqmap[sid.tag])
SID(inds::NTuple{2, Int}, sp::Spin) = SID(inds[1], sp.spin, sidtagmap[inds[2]])

"""
    Spin(; atom::Int=1, norbital::Int=1, spin::Real=0.5)

Construct a spin degrees of freedom.
"""
Spin(; atom::Int=1, norbital::Int=1, spin::Real=0.5) = Spin(atom, norbital, spin)

"""
    SIndex{S} <: Index{PID{S}, SID}

The spin index.
"""
struct SIndex{S} <: Index{PID{S}, SID}
    scope::S
    site::Int
    orbital::Int
    spin::Float
    tag::Char
end
Base.fieldnames(::Type{<:SIndex}) = (:scope, :site, :orbital, :spin, :tag)

"""
    union(::Type{P}, ::Type{SID}) where {P<:PID}

Get the union type of `PID` and `SID`.
"""
Base.union(::Type{P}, ::Type{SID}) where {P<:PID} = SIndex{fieldtype(P, :scope)}

"""
    script(oid::OID{<:SIndex}, ::Val{:site}) -> Int
    script(oid::OID{<:SIndex}, ::Val{:orbital}) -> Int
    script(oid::OID{<:SIndex}, ::Val{:spin}) -> Float
    script(oid::OID{<:SIndex}, ::Val{:tag}) -> Char

Get the required script of a spin oid.
"""
script(oid::OID{<:SIndex}, ::Val{:site}) = oid.index.site
script(oid::OID{<:SIndex}, ::Val{:orbital}) = oid.index.orbital
script(oid::OID{<:SIndex}, ::Val{:spin}) = oid.index.spin
script(oid::OID{<:SIndex}, ::Val{:tag}) = oid.index.tag

"""
    usualspinindextotuple

Indicate that the filtered attributes are `(:scope, :site, :orbital)` when converting a spin index to tuple.
"""
const usualspinindextotuple = FilteredAttributes(:scope, :site, :orbital)

"""
    SOperator(value::Number, id::ID{OID}=ID())

Spin operator.
"""
struct SOperator{V<:Number, I<:ID{OID}} <: Operator{V, I}
    value::V
    id::I
    SOperator(value::Number) = new{typeof(value), Tuple{}}(value, ())
    SOperator(value::Number, id::ID{OID}) = new{typeof(value), typeof(id)}(value, id)
end

"""
    soptdefaultlatex

The default LaTeX format of the oids of a spin operator.
"""
const soptdefaultlatex = LaTeX{(:tag,), (:site, :orbital)}('S')
latexformat(SOperator, soptdefaultlatex)

"""
    statistics(opt::SOperator) -> Char
    statistics(::Type{<:SOperator}) -> Char

Get the statistics of SOperator.
"""
statistics(opt::SOperator) = opt |> typeof |> statistics
statistics(::Type{<:SOperator}) = 'B'

"""
    permute(::Type{<:SOperator}, id1::OID{<:SIndex}, id2::OID{<:SIndex}, table) -> Tuple{Vararg{SOperator}}

Permute two fermionic oid and get the result.
"""
function permute(::Type{<:SOperator}, id1::OID{<:SIndex}, id2::OID{<:SIndex}, table)
    @assert (id1.index ≠ id2.index) || (id1.rcoord ≠ id2.rcoord) || (id1.icoord ≠ id2.icoord) "permute error: permuted ids should not be equal to each other."
    if (usualspinindextotuple(id1.index) == usualspinindextotuple(id2.index)) && (id1.rcoord == id2.rcoord) && (id1.icoord == id2.icoord)
        @assert id1.index.spin == id2.index.spin "permute error: noncommutable ids should have the same spin attribute."
        if id1.index.tag == 'x'
            (id2.index.tag == 'y') && return (SOperator(+1im, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == 'z') && return (SOperator(-1im, ID(permutesoid(id1, 'y', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == '+') && return (SOperator(-1, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == '-') && return (SOperator(+1, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
        elseif id1.index.tag == 'y'
            (id2.index.tag == 'x') && return (SOperator(-1im, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == 'z') && return (SOperator(+1im, ID(permutesoid(id1, 'x', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == '+') && return (SOperator(-1im, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == '-') && return (SOperator(-1im, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
        elseif id1.index.tag == 'z'
            (id2.index.tag == 'x') && return (SOperator(+1im, ID(permutesoid(id1, 'y', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == 'y') && return (SOperator(-1im, ID(permutesoid(id1, 'x', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == '+') && return (SOperator(+1, ID(id2)), SOperator(1, ID(id2, id1)))
            (id2.index.tag == '-') && return (SOperator(-1, ID(id2)), SOperator(1, ID(id2, id1)))
        elseif id1.index.tag == '+'
            (id2.index.tag == 'x') && return (SOperator(+1, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == 'y') && return (SOperator(+1im, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == 'z') && return (SOperator(-1, ID(id1)), SOperator(1, ID(id2, id1)))
            (id2.index.tag == '-') && return (SOperator(+2, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
        elseif id1.index.tag == '-'
            (id2.index.tag == 'x') && return (SOperator(-1, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == 'y') && return (SOperator(1im, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
            (id2.index.tag == 'z') && return (SOperator(+1, ID(id1)), SOperator(1, ID(id2, id1)))
            (id2.index.tag == '+') && return (SOperator(-2, ID(permutesoid(id1, 'z', table))), SOperator(1, ID(id2, id1)))
        end
    else
        return (SOperator(1, ID(id2, id1)),)
    end
end
function permutesoid(id::OID{<:SIndex}, tag::Char, table)
    index = replace(id.index, tag=tag)
    (fieldtype(id|>typeof, :seq) <: Nothing) && return replace(id, index=index)
    seq = get(table, index, nothing)
    (seq === nothing) && error("permutesoid error: incomplete index-sequence table with $index not found.")
    return replace(id, index=index, seq=seq)
end

"""
    SCID(; center=wildcard, atom=wildcard, orbital=wildcard, tag='z', subscript=wildcard)

The id of a spin coupling.
"""
struct SCID{C, A, O, S} <: SimpleID
    center::C
    atom::A
    orbital::O
    tag::Char
    subscript::S
    function SCID(center, atom, orbital, tag::Char, subscript)
        @assert tag in ('x', 'y', 'z', '+', '-') "SCID error: not supported tag($tag)."
        new{typeof(center), typeof(atom), typeof(orbital), typeof(subscript)}(center, atom, orbital, tag, subscript)
    end
end
SCID(; center=wildcard, atom=wildcard, orbital=wildcard, tag='z', subscript=wildcard) = SCID(center, atom, orbital, tag, subscript)
Base.fieldnames(::Type{<:SCID}) = (:center, :atom, :orbital, :tag, :subscript)

"""
    SpinCoupling(value::Number, id::ID{SCID}, subscripts::Subscripts)
    SpinCoupling{N}(value::Number=1;
                    tags::NTuple{N, Char},
                    centers::Union{NTuple{N, Int}, Nothing}=nothing,
                    atoms::Union{NTuple{N, Int}, Nothing}=nothing,
                    orbitals::Union{NTuple{N, Int}, Subscript, Nothing}=nothing
                    ) where N

Spin coupling.
"""
struct SpinCoupling{V, I<:ID{SCID}, S<:Subscripts} <: Coupling{V, I}
    value::V
    id::I
    subscripts::S
    function SpinCoupling(value::Number, id::ID{SCID}, subscripts::Subscripts)
        new{typeof(value), typeof(id), typeof(subscripts)}(value, id, subscripts)
    end
end
function SpinCoupling{N}(value::Number=1;
                        tags::NTuple{N, Char},
                        centers::Union{NTuple{N, Int}, Nothing}=nothing,
                        atoms::Union{NTuple{N, Int}, Nothing}=nothing,
                        orbitals::Union{NTuple{N, Int}, Subscript, Nothing}=nothing
                        ) where N
    (centers === nothing) && (centers = ntuple(i->wildcard, Val(N)))
    (atoms === nothing) && (atoms = ntuple(i->wildcard, Val(N)))
    subscript = isa(orbitals, NTuple{N, Int}) ? Subscript(orbitals) : isa(orbitals, Nothing) ? Subscript{N}() : orbitals
    orbitals = subscript.opattern
    subscripts = ntuple(i->subscript.identifier, Val(N))
    return SpinCoupling(value, ID(SCID, centers, atoms, orbitals, tags, subscripts), Subscripts(subscript))
end

"""
    show(io::IO, sc::SpinCoupling)

Show a spin coupling.
"""
function Base.show(io::IO, sc::SpinCoupling)
    @printf io "SpinCoupling(value=%s" decimaltostr(sc.value)
    cache = []
    for attrname in (:centers, :atoms, :orbitals, :tags)
        attrvalue = getproperty(sc.id, attrname)
        all(attrvalue .≠ wildcard) && @printf io ", %s=(%s)" attrname join(attrvalue, ", ")
    end
    subscripts = sc.id.subscripts
    all(subscripts .== wildcard) || all((subscripts .== constant)) || @printf io ", subscripts=(%s)" join(subscripts, ", ")
    @printf io ")"
end

"""
    repr(sc::SpinCoupling) -> String

Get the repr representation of a spin coupling.
"""
function Base.repr(sc::SpinCoupling)
    cache = []
    for (attrname, abbr) in zip((:atoms, :orbitals), ("sl", "ob"))
        attrvalue = getproperty(sc.id, attrname)
        all(attrvalue .≠ wildcard) && push!(cache, @sprintf "%s(%s)" abbr join(attrvalue, ", "))
    end
    result = @sprintf "%s S%s" decimaltostr(sc.value) join(sc.id.tags, 'S')
    (length(cache) > 0) && (result = @sprintf "%s %s" result join(cache, "⊗"))
    centers = sc.id.centers
    any(centers .≠ wildcard) && (result = @sprintf "%s%s@(%s)" result (length(cache) > 0 ? "" : " ") join(centers, ", "))
    subscripts = sc.id.subscripts
    all(subscripts .== wildcard) || all(subscripts .== constant) || (result = @sprintf "%s with %s" result join(subscripts, " && "))
    return result
end

"""
    *(sc1::SpinCoupling, sc2::SpinCoupling) -> SpinCoupling

Get the multiplication between two spin couplings.
"""
Base.:*(sc1::SpinCoupling, sc2::SpinCoupling) = SpinCoupling(sc1.value*sc2.value, sc1.id*sc2.id, sc1.subscripts*sc2.subscripts)

"""
    expand(sc::SpinCoupling, pid::PID, spin::Spin, kind::Union{Val{K}, Nothing}=nothing) where K -> Union{SCExpand, Tuple{}}
    expand(sc::SpinCoupling, pids::NTuple{N, PID}, spins::NTuple{N, Spin}, kind::Union{Val{K}, Nothing}=nothing) where {N, K} -> Union{SCExpand, Tuple{}}

Expand a spin coupling with the given set of point ids and spin degrees of freedom.
"""
expand(sc::SpinCoupling, pid::PID, spin::Spin, kind::Union{Val{K}, Nothing}=nothing) where {K} = expand(sc, (pid,), (spin,), kind)
function expand(sc::SpinCoupling, pids::NTuple{N, PID}, spins::NTuple{N, Spin}, kind::Union{Val{K}, Nothing}=nothing) where {N, K}
    centers = couplingcenters(typeof(sc), sc.id.centers, Val(N))
    rpids = NTuple{rank(sc), eltype(pids)}(pids[centers[i]] for i = 1:rank(sc))
    rspins = NTuple{rank(sc), eltype(spins)}(spins[centers[i]] for i = 1:rank(sc))
    sps = NTuple{rank(sc), Float}(rspins[i].spin for i = 1:rank(sc))
    for (i, atom) in enumerate(sc.id.atoms)
        isa(atom, Int) && (atom ≠ rspins[i].atom) && return ()
    end
    sbexpands = expand(sc.subscripts, NTuple{rank(sc), Int}(rspins[i].norbital for i = 1:rank(sc))) |> collect
    return SCExpand(sc.value, rpids, sbexpands, sps, sc.id.tags)
end
couplingcenter(::Type{<:SpinCoupling}, i::Int, n::Int, ::Val{2}) = (i % 2 == 1) ? 1 : 2
struct SCExpand{V, N, S} <: VectorSpace{Tuple{V, NTuple{N, SIndex{S}}}}
    value::V
    pids::NTuple{N, PID{S}}
    sbexpands::Vector{NTuple{N, Int}}
    spins::NTuple{N, Float}
    tags::NTuple{N, Char}
end
IsMultiIndexable(::Type{<:SCExpand}) = IsMultiIndexable(true)
MultiIndexOrderStyle(::Type{<:SCExpand}) = MultiIndexOrderStyle('C')
dims(sce::SCExpand) = (length(sce.sbexpands),)
@generated function Tuple(index::Tuple{Int}, sce::SCExpand{V, N}) where {V, N}
    exprs = [:(SIndex(sce.pids[$i], SID(sce.sbexpands[index[1]][$i], sce.spins[$i], sce.tags[$i]))) for i = 1:N]
    return Expr(:tuple, :(sce.value), Expr(:tuple, exprs...))
end

"""
    Heisenberg( mode::String="+-z";
                centers::Union{NTuple{2, Int}, Nothing}=nothing,
                atoms::Union{NTuple{2, Int}, Nothing}=nothing,
                orbitals::Union{NTuple{2, Int}, Subscript, Nothing}=nothing
                ) -> Couplings{ID{SCID, 2}, SpinCoupling{Rational{Int}/Int, <:ID{SCID, 2}}}

The Heisenberg couplings.
"""
function Heisenberg(mode::String = "+-z";
                    centers::Union{NTuple{2, Int}, Nothing}=nothing,
                    atoms::Union{NTuple{2, Int}, Nothing}=nothing,
                    orbitals::Union{NTuple{2, Int}, Subscript, Nothing}=nothing)
    @assert (mode == "+-z") || (mode == "xyz") "Heisenberg error: not supported mode($mode)."
    if mode == "+-z"
        sc1 = SpinCoupling{2}(1//2, centers=centers, atoms=atoms, orbitals=orbitals, tags=('+', '-'))
        sc2 = SpinCoupling{2}(1//2, centers=centers, atoms=atoms, orbitals=orbitals, tags=('-', '+'))
        sc3 = SpinCoupling{2}(1//1, centers=centers, atoms=atoms, orbitals=orbitals, tags=('z', 'z'))
    else
        sc1 = SpinCoupling{2}(1, centers=centers, atoms=atoms, orbitals=orbitals, tags=('x', 'x'))
        sc2 = SpinCoupling{2}(1, centers=centers, atoms=atoms, orbitals=orbitals, tags=('y', 'y'))
        sc3 = SpinCoupling{2}(1, centers=centers, atoms=atoms, orbitals=orbitals, tags=('z', 'z'))
    end
    return Couplings(sc1, sc2, sc3)
end

"""
    Ising(  tag::Char;
            centers::Union{NTuple{2, Int}, Nothing}=nothing,
            atoms::Union{NTuple{2, Int}, Nothing}=nothing,
            orbitals::Union{NTuple{2, Int}, Subscript, Nothing}=nothing
            ) -> Couplings{ID{SCID, 2}, SpinCoupling{Rational{Int}/Int, <:ID{SCID, 2}}}

The Ising couplings.
"""
function Ising(tag::Char;centers::Union{NTuple{2, Int}, Nothing}=nothing, atoms::Union{NTuple{2, Int}, Nothing}=nothing, orbitals::Union{NTuple{2, Int}, Subscript, Nothing}=nothing)
    @assert tag in ('x', 'y', 'z') "Ising error: not supported input tag($tag)."
    return Couplings(SpinCoupling{2}(1, centers=centers, atoms=atoms, orbitals=orbitals, tags=(tag, tag)))
end

"""
    Gamma(  tag::Char;
            centers::Union{NTuple{2, Int}, Nothing}=nothing,
            atoms::Union{NTuple{2, Int}, Nothing}=nothing,
            orbitals::Union{NTuple{2, Int}, Subscript, Nothing}=nothing
            ) -> Couplings{ID{SCID, 2}, SpinCoupling{Rational{Int}/Int, <:ID{SCID, 2}}}

The Gamma couplings.
"""
function Gamma(tag::Char; centers::Union{NTuple{2, Int}, Nothing}=nothing, atoms::Union{NTuple{2, Int}, Nothing}=nothing, orbitals::Union{NTuple{2, Int}, Subscript, Nothing}=nothing)
    @assert tag in ('x', 'y', 'z') "Gamma error: not supported input tag($tag)."
    t1, t2 = (tag == 'x') ? ('y', 'z') : (tag == 'y') ? ('z', 'x') : ('x', 'y')
    sc1 = SpinCoupling{2}(1, centers=centers, atoms=atoms, orbitals=orbitals, tags=(t1, t2))
    sc2 = SpinCoupling{2}(1, centers=centers, atoms=atoms, orbitals=orbitals, tags=(t2, t1))
    return Couplings(sc1, sc2)
end

"""
    DM( tag::Char;
        centers::Union{NTuple{2, Int}, Nothing}=nothing,
        atoms::Union{NTuple{2, Int}, Nothing}=nothing,
        orbitals::Union{NTuple{2, Int}, Subscript, Nothing}=nothing
        ) -> Couplings{ID{SCID, 2}, SpinCoupling{Rational{Int}/Int, <:ID{SCID, 2}}}
"""
function DM(tag::Char; centers::Union{NTuple{2, Int}, Nothing}=nothing, atoms::Union{NTuple{2, Int}, Nothing}=nothing, orbitals::Union{NTuple{2, Int}, Subscript, Nothing}=nothing)
    @assert tag in ('x', 'y', 'z') "DM error: not supported input tag($tag)."
    t1, t2 = (tag == 'x') ? ('y', 'z') : (tag == 'y') ? ('z', 'x') : ('x', 'y')
    sc1 = SpinCoupling{2}(+1, centers=centers, atoms=atoms, orbitals=orbitals, tags=(t1, t2))
    sc2 = SpinCoupling{2}(-1, centers=centers, atoms=atoms, orbitals=orbitals, tags=(t2, t1))
    return Couplings(sc1, sc2)
end

scsinglewrapper(::Nothing) = nothing
scsinglewrapper(value::Int) = (value,)

"""
    Sˣ(; atom::Union{Int, Nothing}=nothing, orbital::Union{Int, Nothing}=nothing) -> Couplings{ID{SCID, 1}, SpinCoupling{Rational{Int}/Int, <:ID{SCID, 1}}}

The single Sˣ coupling.
"""
function Sˣ(; atom::Union{Int, Nothing}=nothing, orbital::Union{Int, Nothing}=nothing)
    Couplings(SpinCoupling{1}(1, tags=('x',), atoms=scsinglewrapper(atom), orbitals=scsinglewrapper(orbital)))
end

"""
    Sʸ(; atom::Union{Int, Nothing}=nothing, orbital::Union{Int, Nothing}=nothing) -> Couplings{ID{SCID, 1}, SpinCoupling{Rational{Int}/Int, <:ID{SCID, 1}}}

The single Sʸ coupling.
"""
function Sʸ(; atom::Union{Int, Nothing}=nothing, orbital::Union{Int, Nothing}=nothing)
    Couplings(SpinCoupling{1}(1, tags=('y',), atoms=scsinglewrapper(atom), orbitals=scsinglewrapper(orbital)))
end

"""
    Sᶻ(; atom::Union{Int, Nothing}=nothing, orbital::Union{Int, Nothing}=nothing) -> Couplings{ID{SCID, 1}, SpinCoupling{Rational{Int}/Int, <:ID{SCID, 1}}}

The single Sᶻ coupling.
"""
function Sᶻ(; atom::Union{Int, Nothing}=nothing, orbital::Union{Int, Nothing}=nothing)
    Couplings(SpinCoupling{1}(1, tags=('z',), atoms=scsinglewrapper(atom), orbitals=scsinglewrapper(orbital)))
end

"""
    sc"..." -> SpinCoupling

Construct a SpinCoupling from a literal string.
"""
macro sc_str(str::String)
    ps = split(str, " with ")
    condition = (length(ps) == 2) ? Meta.parse(ps[2]) : :nothing
    fpos = findfirst(r"S[+-xyz0]", ps[1]).start
    lpos = 1 + lastindex(ps[1]) - findfirst(r"[+-xyz0]S", ps[1]|>reverse).start
    coeff = eval(Meta.parse(ps[1][1:fpos-1]))
    tags = Tuple(replace(ps[1][fpos:lpos], "S"=>""))
    attrpairs = Any[:tags=>tags]
    temp = replace(ps[1][lpos+1:end], " "=>"")
    if length(temp) > 0
        components = split(temp, '@')
        centers = (length(components) == 2) ? couplingcenters(components[2]) : nothing
        (centers === nothing) || @assert length(centers) == length(tags) "@sc_str error: dismatched ranks."
        push!(attrpairs, :centers=>centers)
        if length(components[1]) > 0
            for component in split(components[1], '⊗')
                attrname, attrvalue = sccomponent(component)
                if isa(attrvalue, Expr)
                    @assert attrname == :orbitals "@sc_str error: wrong input pattern."
                    @assert length(tags) == length(attrvalue.args) "@sc_str error: dismatched ranks."
                    push!(attrpairs, attrname=>eval(subscriptexpr(attrvalue, condition)))
                else
                    @assert length(tags) == length(attrvalue) "@sc_str error: dismatched ranks."
                    push!(attrpairs, attrname=>attrvalue)
                end
            end
        end
    end
    return SpinCoupling{length(tags)}(coeff; attrpairs...)
end
function sccomponent(str::AbstractString)
    @assert (str[3] == '(') && (str[end] == ')') "sccomponent error: wrong input pattern."
    attrname = (str[1:2] == "sl") ? :atoms : (str[1:2] == "ob") ? :orbitals : error("sccomponent error: wrong input pattern.")
    expr = Meta.parse(str[3:end])
    attrvalue = isa(expr, Expr) ? (all(isa(arg, Int) for arg in expr.args) ? Tuple(expr.args) : expr) : (expr,)
    return attrname => attrvalue
end

function scpairs(str::AbstractString, ::Val{R}) where R
    str = replace(str, " "=>"")
    attrpairs = Pair{Symbol, NTuple{R, Int}}[]
    if length(str) > 0
        components = split(str, '@')
        (length(components) == 2) && push!(attrpairs, :centers=>couplingcenters(components[2]))
        if length(components[1]) > 0
            for component in split(components[1], '⊗')
                attrname, attrvalue = sccomponent(component)
                @assert isa(attrvalue, NTuple{R, Int}) "scpairs error: wrong input pattern."
                push!(attrpairs, attrname=>attrvalue)
            end
        end
    end
    return attrpairs
end

"""
    heisenberg"sl(a₁, a₂)⊗ob(o₁, o₂)@(c₁, c₂)" -> Couplings

The Heisenberg couplings.
"""
macro heisenberg_str(str::String)
    slice = findfirst(r"xyz|\+\-z", str)
    mode = isnothing(slice) ? "+-z" : str[slice]
    isnothing(slice) || (str = str[slice.stop+1:end])
    @assert (mode == "+-z") || (mode == "xyz") "heisenberg_str error: not supported mode($mode)."
    attrpairs = scpairs(str, Val(2))
    if mode == "+-z"
        sc1 = SpinCoupling{2}(1//2; :tags=>('+', '-'), attrpairs...)
        sc2 = SpinCoupling{2}(1//2; :tags=>('-', '+'), attrpairs...)
        sc3 = SpinCoupling{2}(1//1; :tags=>('z', 'z'), attrpairs...)
    else
        sc1 = SpinCoupling{2}(1; :tags=>('x', 'x'), attrpairs...)
        sc2 = SpinCoupling{2}(1; :tags=>('y', 'y'), attrpairs...)
        sc3 = SpinCoupling{2}(1; :tags=>('z', 'z'), attrpairs...)
    end
    return Couplings(sc1, sc2, sc3)
end

"""
    ising"x sl(a₁, a₂)⊗ob(o₁, o₂)@(c₁, c₂)" -> Couplings
    ising"y sl(a₁, a₂)⊗ob(o₁, o₂)@(c₁, c₂)" -> Couplings
    ising"z sl(a₁, a₂)⊗ob(o₁, o₂)@(c₁, c₂)" -> Couplings

The Ising couplings.
"""
macro ising_str(str::String)
    @assert str[1] ∈ ('x', 'y', 'z') "@ising_str error: wrong input pattern."
    attrpairs = (length(str) > 1) ? (@assert str[2] == ' ' "@ising_str error: wrong input pattern."; scpairs(str[3:end], Val(2))) : Pair{Symbol, NTuple{2, Int}}[]
    return Couplings(SpinCoupling{2}(1; :tags=>(str[1], str[1]), attrpairs...))
end

"""
    gamma"x sl(a₁, a₂)⊗ob(o₁, o₂)@(c₁, c₂)" -> Couplings
    gamma"y sl(a₁, a₂)⊗ob(o₁, o₂)@(c₁, c₂)" -> Couplings
    gamma"z sl(a₁, a₂)⊗ob(o₁, o₂)@(c₁, c₂)" -> Couplings

The Gamma couplings.
"""
macro gamma_str(str::String)
    @assert str[1] in ('x', 'y', 'z') "@gamma_str error: wrong input pattern."
    t1, t2 = (str[1] == 'x') ? ('y', 'z') : (str[1] == 'y') ? ('z', 'x') : ('x', 'y')
    attrpairs = (length(str) > 1) ? (@assert str[2] == ' ' "@gamma_str error: wrong input pattern."; scpairs(str[3:end], Val(2))) : Pair{Symbol, NTuple{2, Int}}[]
    sc1 = SpinCoupling{2}(1; :tags=>(t1, t2), attrpairs...)
    sc2 = SpinCoupling{2}(1; :tags=>(t2, t1), attrpairs...)
    return Couplings(sc1, sc2)
end

"""
    dm"x sl(a₁, a₂)⊗ob(o₁, o₂)@(c₁, c₂)" -> Couplings
    dm"y sl(a₁, a₂)⊗ob(o₁, o₂)@(c₁, c₂)" -> Couplings
    dm"z sl(a₁, a₂)⊗ob(o₁, o₂)@(c₁, c₂)" -> Couplings

The DM couplings.
"""
macro dm_str(str::String)
    @assert str[1] in ('x', 'y', 'z') "@dm_str error: wrong input pattern."
    t1, t2 = (str[1] == 'x') ? ('y', 'z') : (str[1] == 'y') ? ('z', 'x') : ('x', 'y')
    attrpairs = (length(str) > 1) ? (@assert str[2] == ' ' "@dm_str error: wrong input pattern."; scpairs(str[3:end], Val(2))) : Pair{Symbol, NTuple{2, Int}}[]
    sc1 = SpinCoupling{2}(+1; :tags=>(t1, t2), attrpairs...)
    sc2 = SpinCoupling{2}(-1; :tags=>(t2, t1), attrpairs...)
    return Couplings(sc1, sc2)
end

"""
    sˣ"sl(a)⊗ob(o)" -> Couplings
    sʸ"sl(a)⊗ob(o)" -> Couplings
    sᶻ"sl(a)⊗ob(o)" -> Couplings

The single Sˣ/Sʸ/Sᶻ coupling.
"""
macro sˣ_str(str::String) Couplings(SpinCoupling{1}(1; :tags=>('x',), scpairs(str, Val(1))...)) end
macro sʸ_str(str::String) Couplings(SpinCoupling{1}(1; :tags=>('y',), scpairs(str, Val(1))...)) end
macro sᶻ_str(str::String) Couplings(SpinCoupling{1}(1; :tags=>('z',), scpairs(str, Val(1))...)) end

"""
    SpinTerm{R}(id::Symbol, value::Any, bondkind::Any;
                couplings::Union{Function, Coupling, Couplings},
                amplitude::Union{Function, Nothing}=nothing,
                modulate::Union{Function, Bool}=false,
                ) where R

Spin term.

Type alias for `Term{'B', :SpinTerm, R, id, V, <:Any, <:TermCouplings, <:TermAmplitude, <:Union{TermModulate, Nothing}}`.
"""
const SpinTerm{R, id, V, B<:Any, C<:TermCouplings, A<:TermAmplitude, M<:Union{TermModulate, Nothing}} = Term{'B', :SpinTerm, R, id, V, B, C, A, M}
function SpinTerm{R}(id::Symbol, value::Any, bondkind::Any;
                    couplings::Union{Function, Coupling, Couplings},
                    amplitude::Union{Function, Nothing}=nothing,
                    modulate::Union{Function, Bool}=false
                    ) where R
    isa(couplings, TermCouplings) || (couplings = TermCouplings(couplings))
    isa(amplitude, TermAmplitude) || (amplitude = TermAmplitude(amplitude))
    isa(modulate, TermModulate) || (modulate = (modulate === false) ? nothing : TermModulate(id, modulate))
    Term{'B', :SpinTerm, R, id}(value, bondkind, couplings, amplitude, modulate, 1)
end
abbr(::Type{<:SpinTerm}) = :sp
isHermitian(::Type{<:SpinTerm}) = true

"""
    otype(T::Type{<:SpinTerm}, I::Type{<:OID})

Get the operator type of a spin term.
"""
otype(T::Type{<:SpinTerm}, I::Type{<:OID}) = SOperator{T|>valtype, ID{I, T|>rank}}

end # module


function multisolve(m::Model, mdata::MultiData, ::WeightedSum, ::LinearProblem)
    objectives = mdata.objectives
    sensemap = Dict(MOI.MIN_SENSE => 1.0, MOI.MAX_SENSE => -1.0)

    vararr = all_variables(m)

    numobj = length(objectives)
    Phi = zeros(numobj, numobj)

    # Individual minimisations
    for (i, objective) in enumerate(objectives)
        @objective(m, objective.sense, objective.f)
        for (key, value) in objective.initialvalue
            set_start_value(variable_by_name(m, key), value)
        end
        optimize!(m, ignore_optimize_hook=true);
        if !(termination_status(m) in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED])
            return termination_status(m)
        end

        push!(mdata.utopiavarvalues, [value(var) for var in vararr])
        push!(mdata.paretovarvalues, [value(var) for var in vararr])

        Phi[:,i] = sensevalue.(objectives)

        push!(mdata.paretofront, value.(objectives))
    end
    Fmax = maximum(Phi, dims=2)
    Fmin = minimum(Phi, dims=2) # == diag(Phi)?
    if Fmax == Fmin
        error("The Nadir and Utopia points are equal") # I think that's what this means?
    end

    mdata.Phi = Phi

    beta = zeros(numobj); beta[end] = 1.0

    betatree = betas(numobj, mdata.pointsperdim - 1)

    for betaval in betatree
        if count(t -> t != 0, betaval) == 1
            # Skip individual optimisations as
            # they are already performed
            continue
        end
        @objective(m, MOI.MIN_SENSE,
            sum(betaval[i] * (sensemap[objectives[i].sense] * objectives[i].f -
            Fmin[i]) / (Fmax[i] - Fmin[i]) for i in Base.OneTo(numobj))
        )

        optimize!(m, ignore_optimize_hook=true);
        if !(termination_status(m) in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED])
            return termination_status(m)
        end

        push!(mdata.paretofront, value.(objectives))
        push!(mdata.paretovarvalues, [value(var) for var in vararr])
    end
    return MOI.OPTIMAL
end

function multisolve(m::Model, mdata::MultiData, ::EpsilonCons, ::LinearProblem)
    objectives = mdata.objectives
    sensemap = Dict(MOI.MIN_SENSE => 1.0, MOI.MAX_SENSE => -1.0)

    vararr = all_variables(m)

    numobj = length(objectives)
    Phi = zeros(numobj, numobj)

    if numobj > 2
        # TODO:
        # The logic  here becomes difficult, as the feasible region will require
        # a dependency between the constraints
        Base.error(":EPS is thought through for > 2 objectives yet")
    end

    # Individual minimisations
    for (i, objective) in enumerate(objectives)
        @objective(m, objective.sense, objective.f)
        for (key, value) in objective.initialvalue
            set_start_value(variable_by_name(m, key), value)
        end
        optimize!(m, ignore_optimize_hook=true);
        if !(termination_status(m) in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED])
            return termination_status(m)
        end

        push!(mdata.utopiavarvalues, [value(var) for var in vararr])
        push!(mdata.paretovarvalues, [value(var) for var in vararr])

        Phi[:,i] = sensevalue.(objectives)

        push!(mdata.paretofront, value.(objectives))
    end
    Fmax = maximum(Phi, dims=2)
    Fmin = minimum(Phi, dims=2) # == diag(Phi)?

    mdata.Phi = Phi
    # maintain last objective, other as constraints
    @objective(m, objectives[end].sense, objectives[end].f)

    @constraint(m, objconstr,
                  sensemap[objectives[1].sense] * objectives[1].f
                  <= Fmax[1])

    for betaval in 0.0:0.001:1.0

        set_normalized_rhs(objconstr, betaval * Fmin[1] + (1 - betaval) * Fmax[1])

        optimize!(m, ignore_optimize_hook=true);
        if !(termination_status(m) in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED])
            return termination_status(m)
        end

        push!(mdata.paretofront, value.(objectives))
        push!(mdata.paretovarvalues, [value(var) for var in vararr])
    end

    return MOI.OPTIMAL
end


module roessler
using ReachabilityAnalysis
@taylorize function roessler!(du, u, p, t)
    local a, b, c = 0.2, 0.2, 5.7

    du[1] = -u[2] - u[3]
    du[2] = u[1] + (a * u[2])
    du[3] = b + (u[3]*(u[1] - c))

    return du
end

function model(X0)
    S = @system(x' = roessler!(x), dim:3)
    return IVP(S, X0)
end
end # module


include("../src/lstm_g/network.jl")
include("../src/lstm_g/viz.jl")

nin = 3
nhidden = 2
nout = 1

# basic ANN
inputlayer  = gatedlayer(3, tag=:input)
hiddenlayer = gatedlayer(2, tag=:hidden)
outputlayer = gatedlayer(1, tag=:output)
@show inputlayer hiddenlayer outputlayer

ci = connect(inputlayer, hiddenlayer)
co = connect(hiddenlayer, outputlayer)
@show ci co

# now add a gating gatedlayer
gatinglayer = gatedlayer(2, tag=:gater)
cg = connect(inputlayer, gatinglayer)

# gate ci
gate!(ci, gatinglayer)

# peephole uses SAME type... node i maps to node i
peep = connect(hiddenlayer, gatinglayer, SAME)

@show gatinglayer

net = GatedNetwork(inputlayer, hiddenlayer, GatedLayer{Float64}[inputlayer, gatinglayer, hiddenlayer, outputlayer])


using LinearAlgebra
using BlockFactorizations
using BenchmarkTools

d = 512
n = 16
# testing Diagonal BlockFactorization
A = Diagonal([(randn(d, d)) for i in 1:n])
B = BlockFactorization(A)

# x = [randn(d) for _ in 1:n]
# y = [zeros(d) for _ in 1:n]
# @btime A*x; # this doesn't work in 1.6
# @btime mul!(y, A, x);

x = randn(d*n)
y = zeros(d*n)
@btime B*x;
@btime mul!(y, B, x);


import EventHistory
##reload("EventHistory")

##################################################
### EventClass
##################################################

## Right-censored event-time
stop   = [2,3,3];
status = [false,true,true];
e1     = EventHistory.Event(stop,status)

EventHistory.Time(e1)
EventHistory.Status(e1)

## Right-censored+left truncation
start = [0,1,2];
e2    = EventHistory.Event(start,stop,status)

EventHistory.Entry(e2)

## Competing risks model
cause = [0,2,1];
e3    = EventHistory.Event(start,stop,cause)

EventHistory.Status(e3)
EventHistory.Cause(e3)

## Interval censoring
right = [2,3,Inf];
left  = [1,-Inf,1];
e4    = EventHistory.Event(left,right,"interval")

## Formula syntax (see also below)
d = DataFrames.DataFrame(start=start,stop=stop,status=status);
EventHistory.Event([:stop,:status],d)
EventHistory.Event([:start,:stop,:status],d)

##################################################
### Cox regression
##################################################
using RDatasets
ovarian = dataset("survival", "ovarian");
ovarian[:Group] = ovarian[:Rx].-1;
S = EventHistory.Event([:FUTime,:FUStat], ovarian);
ovarian[:S] = EventHistory.Event([:FUTime,:FUStat], ovarian);

mm = EventHistory.phreg(@formula(S~Age+Group), ovarian)

## Prediction
EventHistory.predict(mm, surv=false, X=[0 0]); ## Baseline
s56 = EventHistory.predict(mm, X=[56 1], order=true); ## Survival probabilities age 40, group 1
EventHistory.predict(mm, X=[56 0], time=[100,400,600]); ## ... at time 100,400,600
EventHistory.predict(mm, X=[56 1; 56 0], time=[600,100,400]) ## ... both groups

using Winston
plot(s56[:,1], s56[:,2])

s = EventHistory.predict(mm, X=[56 1; 56 0], order=true)
pr = DataFrame(Time=[s[:,1];s[:,1]], S=[s[:,2];s[:,3]], Group=repeat(["Group1","Group2"], inner=size(s,1)))

using Gadfly
p = Gadfly.plot(pr, x="Time", y="S", color="Group",
                Geom.step, Geom.point,
                Guide.ylabel("Survival probability"), Guide.title("Age 56"))

draw(PNG("surv.png",7inch,7inch),p)

##################################################
### Cox regression with left truncation
##################################################
d = DataFrame(start=[1,2,5,2,1,7,3,4,8,8],
              stop=[2,3,6,7,8,9,9.5,10,14,17],
              event=[1,1,1,1,1,1,1,0,0,0],
              x=[1,0,0,1,0,1,1,1,0,0],
              z=[1.0,0,2.0,0,3.0,0,4.0,0,5.0,0]);
d[:S] = EventHistory.Event([:start,:stop,:event], d);
mm = EventHistory.phreg(@formula(S~x), d)

mm = EventHistory.phreg(@formula(S~x+z), d)


balancedbrackets(str::AbstractString) = foldl((x, y) -> x < 0 ? -1 : x + y, 0, collect((x == '[') - (x == ']') for x in str)) == 0


export InterpolationFunction

"""
Struct defining an interpolation basis

* `ex`: expression defining the basis function format
* `coords`: vector of symbols for coordinates in 3D space
* `linpars`: vector of symbols for linear parameters, defining the interpolation
    as a linear transformation
* `DOFs`: `Vector{Vector{Int64}}`. Each inner vector defines a set of derivative orders
    for a degree of freedom of the basis function, in relationship to the variables at `coords`.

    Example:
        a degree of freedom defined by `[0, 1, 0]` characterizes the first derivative, in the second
        cartesian axis at a node in the three-dimensional space.
"""
mutable struct InterpolationFunction
    ex::Expr
    coords::Vector{Symbol}
    linpars::Vector{Symbol}
    DOFs::Vector{
        Vector{Int64}
    }
    _DOF_analysis_functions::Matrix{Function}
    _funcs::Array{Function}
    _npts::Int64
end

"""
Constructor for `InterpolationFunction` struct.

* `ex`: expression defining the basis function format
* `coords`: vector of symbols for coordinates in 3D space
* `linpars`: vector of symbols for linear parameters, defining the interpolation
    as a linear transformation
* `DOFs`: `Vector{Vector{Int64}}`. Each inner vector defines a set of derivative orders
    for a degree of freedom of the basis function, in relationship to the variables at `coords`.

    Example:
        a degree of freedom defined by `[0, 1, 0]` characterizes the first derivative, in the second
        cartesian axis at a node in the three-dimensional space.

* `maximum_derivative`: maximum derivative order to be kept in function form
"""
function InterpolationFunction(
    ex::Expr,
    coords::Vector{Symbol},
    linpars::Vector{Symbol},
    DOFs::Vector{
        Vector{Int64}
    };
    maximum_derivative::Int64=1
)
    _DOF_analysis_functions=Matrix{Function}(undef, length(DOFs), length(linpars))

    for (i, DOF) in enumerate(DOFs)
        for (j, lp) in enumerate(linpars)
            _DOF_analysis_functions[i, j]=get_linpar_contribution(ex, coords, linpars, lp, DOF)
        end
    end

    nd=length(DOFs[1])

    _funcs=Array{Function, nd}(undef, [maximum_derivative+1 for i=1:nd]...)

    for fcinds in CartesianIndices(_funcs)
        _funcs[fcinds]=get_basefun_derivative(
            ex,
            coords,
            linpars,
            collect(Tuple(fcinds)).-1
        )
    end

    InterpolationFunction(
        ex,
        coords,
        linpars,
        DOFs,
        _DOF_analysis_functions,
        _funcs,
        length(linpars)/length(DOFs)
    )
end

"""
Gets a matrix defining DOFs at given sample points as a linear transformation of `linpars`
"""
function get_interp_linmat(interp::InterpolationFunction, coord_values::Vector{Vector{Float64}})
    z=zeros(Float64, length(interp.linpars))

    vcat(
        [
            [Base.invokelatest(interp._DOF_analysis_functions[i, j], pt..., z...)
                for i=1:size(interp._DOF_analysis_functions, 1), j=1:size(interp._DOF_analysis_functions, 2)]
                    for pt in coord_values
        ]...
    )
end

"""
Get interpolation function of given derivative order from interpolation object
"""
function get_interpfun(interp::InterpolationFunction; deriv_order::Union{Nothing, Vector{Int64}}=nothing)
    return interp._funcs[
        (isnothing(deriv_order) ? 1 : (deriv_order.+1))...
    ]
end


"""
The output is in different form from the other functions
Here voxels are tranposed and faces are in 2-D array.
"""
function get_surface_grid_per_voxel(data3d::Array; voxelsize_mm::Array=[1,1,1])
    # Get vertices and count the voxesl (*4 face number)
    verts, nvoxels = LarSurf.vertices_and_count_voxels_as_the_square_face_preprocessing(data3d, voxelsize_mm)

    # Create square faces

    threshold = 0
    faces = LarSurf.create_square_faces(data3d, nvoxels, threshold)

    new_verts, new_faces = LarSurf.keep_surface_faces(verts, faces)

    verts = new_verts
    faces = new_faces

    return verts, faces

end

# function continue(verts, faces)
    # # Triangulation
    #
    # trifaces = LarSurf.triangulation(faces)
    #
    # # Visualization
    # LarSurf.check_vf(verts, trifaces)
    #
    # V, EV, FE = LarSurf.to_lar(verts, trifaces)
    #
    # EVc = Lar.cop2lar(EV)
    # Vnew = convert(Array, transpose(V))
    # VV = [i for i=1:length(Vnew)]
    # return Vnew, [VV, EV]
# end


using LinearAlgebra
using SparseArrays

function degree_selector(t, M, U, p)
    C = ceil.(abs.(t)*M)'*U
    C = zero_to_inf.(C)

    # idx is a CarthesianIndex if C' is a Matrix, or a scalar if C' is a row
    # vector. idx[1] extract the first index, i.e. row, of the CarthesianIndex
    cost, idx = findmin(C')
    m = idx[1]

    if cost == Inf
        cost = 0
    end
    s = max(cost/m,1);
    return (m, s)
end

function zero_to_inf(x::Number)
    if x == 0.
        Inf
    else
        x
    end
end


module PartitionModel

using DataFrames
using JuMP
using Cbc

export Solve

type PartitionData
  n::Int                                    # Number of entities
  df::DataFrame                             # Dataframe containing entity data
  Ids::DataArray                            # Array of entities
  vNum::Vector{Symbol}                      # Vector containing numeric variables
  vCat::Vector{Symbol}                      # Vector containing categoric variables
  nMet::Dict{Symbol, Dict{String, Float64}} # Numerical metrics
  cMet::Dict{Symbol, Dict{String, Float64}} # Categorical metrics
end

function Solve(fname::String, n_groups::Int, timelimit::Float64=999.9, vars=[])
  return Solve(CreateData(fname, vars), n_groups, timelimit)
end

function GroupSize(m::Int, n::Int)
  k = floor(n / m)
  y = Array(Int, m)
  for g = 1:(n % m)
    y[g] = k + 1
  end
  for g = (n % m + 1) : m
    y[g] = k
  end
  return y
end

function Solve(data::PartitionData, n_groups::Int, timelimit::Float64=999.9)
  groups = GroupSize(n_groups, data.n)

  m = Model(solver=CbcSolver(seconds=timelimit))

  @defVar(m, x[1:length(data.Ids), 1:n_groups], Bin, string("x", i, j))
  @defVar(m, mmin[data.vNum])           # Minimum mean of group
  @defVar(m, mmax[data.vNum])           # Maximum mean of group
  @defVar(m, vmin[data.vNum])           # Minimum variance of group
  @defVar(m, vmax[data.vNum])           # Maximum variance of group
  @defVar(m, cvio[data.vCat] >= 0)           # Violation of categorical constraint

  k1 = 1.0
  k2 = 1.0
  k3 = 1000.0
  @setObjective(m, :Min,
                k1 * sum([(mmax[v] - mmin[v]) / data.nMet[v]["mean"] for v in data.vNum])
                + k2 * sum([(vmax[v] - vmin[v]) / data.nMet[v]["var"] for v in data.vNum])
                + k3 * sum(cvio)
                )

  # Each entity can only be allocated to single group
  @addConstraint(m, c1[i=1:data.n], sum([x[i,j] for j=1:n_groups]) == 1)

  # Set number of people per group
  @addConstraint(m, c2[j=1:n_groups], sum([x[i,j] for i=1:data.n]) == groups[j])

  # Bound numerical variables
  @addConstraint(m, cMeanMax[v=data.vNum, j=1:n_groups], sum([data.df[v][i] * x[i, j] for i=1:data.n]) >= groups[j] * mmin[v])
  @addConstraint(m, cMeanMin[v=data.vNum, j=1:n_groups], sum([data.df[v][i] * x[i, j] for i=1:data.n]) <= groups[j] * mmax[v])
  @addConstraint(m, cVarMax[v=data.vNum, j=1:n_groups], sum([(data.df[v][i] - data.nMet[v]["mean"])^2 * x[i, j] for i=1:data.n]) >= groups[j] * vmin[v])
  @addConstraint(m, cVarMax[v=data.vNum, j=1:n_groups], sum([(data.df[v][i] - data.nMet[v]["mean"])^2 * x[i, j] for i=1:data.n]) <= groups[j] * vmax[v])

  status = solve(m)

  if status == :Optimal || status == :UserLimit
    solution = get_solution(x, data)
    data.df[:Group] = solution
    return data.df
  elseif status == :Unbounded
    error("Unable to solve. Model is unbounded, which means something is wrong with your data")
  elseif status == :Error
    error("Cbc couldn't solve this model.")
  end
end

function get_solution(x, data::PartitionData)
  solution = Array(Int, data.n)
  for i = 1:data.n
    for j = 1:size(x, 2)
      if getValue(x[i, j]) > 0.5
        solution[i] = j
      end
    end
  end
  return solution
end

function CreateData(fname::String, vars::Vector=[])
  df = readtable(fname)
  Ids = df[names(df)[1]]
  vNum = Array(Symbol, 0)
  vCat = Array(Symbol, 0)
  vMet = Dict{Symbol, Dict{String, Float64}} ()
  cMet = Dict{Symbol, Dict{String, Float64}} ()
  for v in names(df)[2:end]
  	if length(vars) > 0
	   	if !(string(v) in vars)
  			# Excluded variable
  			continue
  		end
  	end
    if typeof(df[v]) == DataArray{UTF8String, 1}
      push!(vCat, v)
      cMet[v] = categorical_metrics(df[v])
    else
      push!(vNum, v)
      vMet[v] = numerical_metrics(df[v])
    end
  end
  return PartitionData(length(Ids), df, Ids, vNum, vCat, vMet, cMet)
end

function numerical_metrics(c::DataArray{Float64, 1})
#   return Dict("mean" => mean(c), "var" => var(c))
  return {string("mean") => mean(c), string("var") => var(c)}

end

function categorical_metrics(c::DataArray{UTF8String, 1})
  metrics = Dict{String, Float64}()
  for i in c
    if haskey(metrics, string(i))
      metrics[string(i)] += 1.0
    else
      metrics[string(i)] = 1.0
    end
  end

  for (k, v) in metrics
    metrics[k] /= length(c)
  end
  return metrics
end

end


using ReservoirComputing
using MLJLinearModels
#model parameters
const approx_res_size = 30
const radius = 1.2
const activation = tanh
const degree = 6
const sigma = 0.1
const beta = 0.0
const alpha = 1.0
const nla_type = NLADefault()
const in_size = 3
const out_size = 3
const extended_states = false
const delta = 0.5

const train_len = 50
const predict_len = 12
data = ones(Float64, in_size, 100)
train = data[:, 1:1+train_len-1]
test = data[:, train_len:train_len+predict_len-1]

#constructor 1
esn = ESN(approx_res_size,
    train,
    degree,
    radius,
    activation = activation,
    sigma = sigma,
    alpha = alpha,
    nla_type = nla_type,
    extended_states = extended_states)

#test train Ridge
t = Ridge(beta, Analytical())
@test isequal(t.lambda, beta)
@test isequal(t.solver, Analytical())
W_out = ESNtrain(t, esn)
@test size(W_out) == (out_size, esn.res_size)

#test train Lasso
t = Lasso(beta, ProxGrad())
@test isequal(t.lambda, beta)
@test isequal(t.solver, ProxGrad())
W_out = ESNtrain(t, esn)
@test size(W_out) == (out_size, esn.res_size)

#test train ElastNet
t = ElastNet(beta, beta, ProxGrad())
@test isequal(t.lambda, beta)
@test isequal(t.gamma, beta)
@test isequal(t.solver, ProxGrad())
W_out = ESNtrain(t, esn)
@test size(W_out) == (out_size, esn.res_size)

#test train RobustHuber
t = RobustHuber(delta, beta, 0.0, Newton())
@test isequal(t.delta, delta)
@test isequal(t.lambda, beta)
@test isequal(t.gamma, 0.0)
@test isequal(t.solver, Newton())
W_out = ESNtrain(t, esn)
@test size(W_out) == (out_size, esn.res_size)


struct Cylinder <: Shape
    id::Int64
    radius::Float64
    zmin::Float64
    zmax::Float64
    obj_to_world::Transformation
    world_to_obj::Transformation
end

Cylinder(T::Transformation) = Cylinder(1, 1.0, -1.0, 1.0, T, inv(T))
Cylinder(id::Int64, radius::Real, zmin::Real, zmax::Real) = Cylinder(id, radius, zmin, zmax, Transformation(), Transformation())

can_intersect(C::Cylinder) = true
area(C::Cylinder) = 2π * C.radius * (C.zmax - C.zmin)

function obj_bounds(C::Cylinder)
    pMin = Point3(-C.radius, -C.radius, C.zmin)
    pMax = Point3( C.radius,  C.radius, C.zmax)
    return BoundingBox(pMin, pMax)
end

function world_bounds(C::Cylinder)
    C.obj_to_world(obj_bounds(C))
end

function (T::Transformation)(C::Cylinder)
    T2 = T * C.obj_to_world
    Cylinder(C.id, C.radius, C.zmin, C.zmax, T2, inv(T2))
end

function cylinder_coefs(ray::Ray, cyl::Cylinder)
    A = ray.direction.x^2 + ray.direction.y^2
    B = 2.0 * (ray.direction.x * ray.origin.x + ray.direction.y * ray.origin.y)
    C = ray.origin.x^2 + ray.origin.y^2 - cyl.radius^2
    return A, B, C
end

function shape_intersect(R::Ray, cyl::Cylinder)
    ray = cyl.world_to_obj(R)

    A, B, C = cylinder_coefs(ray, cyl)
    if A ≈ 0.0
        return nothing, NaN, NaN
    end
    hit, tnear, tfar = quadratic(A,B,C)

    if !hit || tnear > ray.tmax || tfar < ray.tmin || (tnear < ray.tmin && tfar > ray.tmax)
        # The ray didn't even hit the infinitely long extension of the cylinder.
        return nothing, NaN, NaN
    end

    P1 = ray(tnear)
    P2 = ray(tfar)

    if cyl.zmin <= P1.z <= cyl.zmax && tnear >= ray.tmin
        P = P1
        t = tnear
    elseif cyl.zmin <= P2.z <= cyl.zmax && tnear <= ray.tmax
        P = P2
        t = tfar
    else
        return nothing, NaN, NaN
    end

    n = normalize(Normal3(P.x, P.y, 0.0))
    s = Vector3(0.0, 0.0, 1.0)

    n = flip_normal(n, ray.direction)

    DG = DifferentialGeometry(cyl.obj_to_world(P), cyl.obj_to_world(n), cyl.obj_to_world(s))

    return DG, t, 5e-4 * t
end

function intersectP(r::Ray, cyl::Cylinder)
    ray = cyl.world_to_obj(r)
    A, B, C = cylinder_coefs(ray, cyl)
    if A ≈ 0.0
        return false
    end
    hit, tnear, tfar = quadratic(A,B,C)

    if !hit || tnear > ray.tmax || tfar < ray.tmin || (tnear < ray.tmin && tfar > ray.tmax)
        return false
    end

    P1 = ray(tnear)
    P2 = ray(tfar)

    return (cyl.zmin <= P1.z <= cyl.zmax && tnear >= ray.tmin) || (cyl.zmin <= P2.z <= cyl.zmax && tnear <= ray.tmax)

end


##############################################################################
##
## Non Linear solver using Pseudo-Transient Continuation Method
##
##############################################################################

# Implicit time step
function implicit_timestep(F!, ypost, Δ; is_algebraic = fill(false, size(ypost)...), iterations = 100, verbose = true, method = :newton, autodiff = :forward, maxdist = 1e-9, J0c = (nothing, nothing))
    F_helper!(ydot, y) = (F!(ydot, y) ; ydot .+= .!is_algebraic .* (ypost .- y) ./ Δ)
    J0, colorvec = J0c
    if J0 == nothing
        result = nlsolve(F_helper!, ypost; iterations = iterations, show_trace = verbose, ftol = maxdist, method = method, autodiff = autodiff)
    else
        if autodiff == :forward
            jac_cache = ForwardColorJacCache(F_helper!, deepcopy(ypost); colorvec = colorvec, sparsity = J0)
            j_helper! = (J, y) -> forwarddiff_color_jacobian!(J, F_helper!, y, jac_cache)
        else
            j_helper! = (J, y) -> DiffEqDiffTools.finite_difference_jacobian!(J, F_helper!, y; colorvec = colorvec)
        end
        result = nlsolve(OnceDifferentiable(F_helper!, j_helper!, deepcopy(ypost), deepcopy(ypost), J0), ypost; iterations = iterations, show_trace = verbose, ftol = maxdist, method = method)
    end
    return result.zero, result.residual_norm
end

# Solve for steady state
function finiteschemesolve(F!, y0; Δ = 1.0, is_algebraic = fill(false, size(y0)...), iterations = 100, inner_iterations = 10, verbose = true, inner_verbose = false, method = :newton, autodiff = :forward, maxdist = 1e-9, scale = 10.0, J0c = (nothing, nothing))
    ypost = y0
    ydot = zero(y0)
    F!(ydot, ypost)
    distance = norm(ydot) / length(ydot)
    isnan(distance) && throw("F! returns NaN with the initial value")
    if Δ == Inf
        ypost, distance = implicit_timestep(F!, y0, Δ; is_algebraic = is_algebraic, verbose = verbose, iterations = iterations,  method = method, autodiff = autodiff, maxdist = maxdist, J0c = J0c)
    else
        coef = 1.0
        olddistance = distance
        iter = 0
        while (iter < iterations) & (Δ >= 1e-12) & (distance > maxdist)
            iter += 1
            y, nldistance = implicit_timestep(F!, ypost, Δ; is_algebraic = is_algebraic, verbose = inner_verbose, iterations = inner_iterations, method = method, autodiff = autodiff, maxdist = maxdist, J0c = J0c)
            F!(ydot, y)
            distance, olddistance = norm(ydot) / length(ydot), distance
            distance = isnan(distance) ? Inf : distance
            if nldistance <= maxdist
                # if the implicit time step is correctly solved
                verbose && @show iter, Δ, distance
                coef = (distance <= olddistance) ? scale * coef : 1.0
                Δ = Δ * coef * olddistance / distance
                ypost, y = y, ypost
            else
                verbose && @show iter, Δ, NaN
                # if the implict time step is not solved
                # revert and diminish the time step
                coef = 1.0
                Δ = Δ / 10
                distance = olddistance
            end
        end
    end
    verbose && (distance > maxdist) && @warn "Iteration did not converge"
    return ypost, distance
end


##
using ACE, Printf, Test, LinearAlgebra
using ACE: evaluate, evaluate_d, read_dict, write_dict
using ACE.Testing

verbose = false
maxdeg = 10

##

@info("Testing Transforms and TransformedPolys")
for p in 2:4
   @info("p = $p, random transform")
   trans = PolyTransform(1+rand(), 1+rand())
   @info("      test (de-)dictionisation")
   @test all(ACE.Testing.test_fio(trans))
   B1 = transformed_jacobi(maxdeg, trans, 3.0; pcut = p)
   B2 = transformed_jacobi(maxdeg, trans, 3.0, 0.5, pin = p, pcut = p)
   for B in [B1, B2]
      B == B1 && @info("basis = 1s")
      B == B2 && @info("basis = 2s")
      for r in [3 * rand(10); [3.0]]
         P = evaluate(B, r)
         dP = evaluate_d(B, r)
         errs = []
         verbose && @printf("     h    |     error  \n")
         for p = 2:10
            local h = 0.1^p
            dPh = (evaluate(B, r+h) - P) / h
            push!(errs, norm(dPh - dP, Inf))
            verbose && @printf(" %.2e | %2e \n", h, errs[end])
         end
         print_tf(@test (/(extrema(errs)...) < 1e-3) || (minimum(errs) < 1e-10) )
      end
      println()
   end
end

##

@info("Testing PolyTransforms")
for p = 2:4
   r0 = 1+rand()
   trans = PolyTransform(p, r0)
   ACE.Testing.test_transform(trans, [r0/2, 3*r0])
end
println()

##

@info("Testing Morse Transform")
for lam = 1.0:3.0
   r0 = 1+rand()
   trans = ACE.Transforms.MorseTransform(lam, r0)
   ACE.Testing.test_transform(trans, [r0/2, 3*r0])
end
println()

##

@info("Testing Agnesi Transform")
for p = 2:4
   r0 = 1+rand()
   trans = ACE.Transforms.AgnesiTransform(r0, p)
   ACE.Testing.test_transform(trans, [r0/2, 3*r0])
end
println()

##

# TODO: This could be moved to some Tools package
#       to visualize the good transforms

# using Plots
# r0 = 1.0
# rr = range(0.0, 3*r0, length=200)
# plot(; size = (500, 300))
# for p = 2:4
#    tpoly = PolyTransform(p, r0)
#    tagnesi = ACE.Transforms.AgnesiTransform(r0, p)
#    plot!(rr, tagnesi.(rr), lw=2, c=p-1, label = "p = $p")
#    plot!(rr, tpoly.(rr), lw=2, c=p-1, ls = :dash, label = "")
# end
# xlabel!("r")
# ylabel!("x")
# title!("solid = Agnesi, dashed = Poly")
# vline!([1.0], lw=2, c=:black, label = "r0")
# ylims!(0.0, 2.0)
# #---
#
# function visualize_transform(T, rrange; nrays = 50, c = 1,
#                              inverse = false, straightr0 = false)
#    rnn = T.r0
#    xnn = T(rnn)
#    r0, r1 = extrema(rrange)
#    x0, x1 = T(r0), T(r1)
#    # plot r -> x
#    if !inverse
#       rr = range(r0, r1, length=nrays)
#       xx = T.(rr)
#    else
#       xx = range(x0, x1, length=nrays)
#       rr = ACE.Transforms.inv_transform.(Ref(T), xx)
#    end
#    rnn = (rnn - rr[1]) / (rr[end] - rr[1])
#    xnn = (xnn - xx[1]) / (xx[end] - xx[1])
#    rr = (rr .- rr[1]) / (rr[end] - rr[1])
#    xx = (xx .- xx[1]) / (xx[end] - xx[1])
#    if straightr0
#       rr = rr .- rnn; rnn = 0.0
#       xx = xx .- xnn; xnn = 0.0
#    end
#    plot(; size = (500, 200), ylims = (-1, 2), xticks = [])
#    for (r, x) in zip(rr, xx)
#       plot!([r, x], [0, 1], lw=1, label = "", c = c, )
#    end
#    plot!([rnn, xnn], [0, 1], lw=3, c = :red, label = "")
#    plot!([rr[1], rr[end]], [0, 0], lw=3, label = "", c = :black)
#    plot!([xx[1], xx[end]], [1, 1], lw=3, label = "", c = :black)
#    yticks!([0, 1], ["r", "x"])
# end
#
#
#
# tagnesi = AgnesiTransform(1.0, 3)
# P1 = visualize_transform(tagnesi, (0.3*r0, 3*r0), straightr0=true)
# title!(P1, "Agnesi-3")
#
# tagnesi4 = AgnesiTransform(1.0, 4)
# P3 = visualize_transform(tagnesi4, (0.3*r0, 3*r0), straightr0=true)
# title!(P3, "Agnesi-4")
#
# tpoly = PolyTransform(2, 1.0)
# P2 = visualize_transform(tpoly, (0.3*r0, 3*r0), straightr0=true)
# title!(P2, "Poly-2")
# plot(P2, P1, P3, layout = (3,1), size = (500, 600))
#

#---


using LinearAlgebra, ForwardDiff, Distributions
using Convex, SCS, ECOS
using IterativeSolvers

# problem setup
p = 60
N = 30
x_sol = max.(0.0, randn(p))
G = rand(N, p)
h = G * x_sol
ν_sol = randn(N)
λ_sol = rand(p)
λ_sol[x_sol .> 0.0] .= 0.0
c = -G' * ν_sol + λ_sol

"Convex.jl"
x = Variable(p)
prob = minimize(c' * x)
prob.constraints += h - G * x == 0.0
prob.constraints += x >= 0.0

@time solve!(prob, ECOS.Optimizer)

@show prob.status
@show prob.constraints
@show x.value
# @show prob.constraints[1].dual
# @show prob.constraints[2].dual
A = [-1.0 * Diagonal(ones(p)); G; -G]
b = [zeros(p); h; -h]
m = 2 * N + p
k = p + m + 1

Q = Array([zeros(p, p) A' c;
           -A zeros(m, m) b;
           -c' -b'      0.0])

function κ(z)
    # z_proj = zero(z)
    # z_proj[1:N] = z[1:N]
    # z_proj[N .+ (1:p)] = max.(0.0, z[N .+ (1:p)])

    z_proj = max.(0.0, z)
    return z_proj
    # return z
end

κ(randn(10))

function Pc(z)
    z_proj = zero(z)
    z_proj[1:p] = z[1:p]
    z_proj[p .+ (1:m)] = κ(z[p .+ (1:m)])
    z_proj[p + m + 1] = max(0.0, z[p + m + 1])

    return z_proj
end

function F(z)
    ũ = z[1:k]
    u = z[k .+ (1:k)]
    v = z[2 * k .+ (1:k)]

    [(I + Q) * ũ - (u + v);
     u - Pc(ũ - v);
     ũ - u]
end

function Ju(z)
    ũ = z[1:k]
    u = z[k .+ (1:k)]
    v = z[2k .+ (1:k)]

    JP = zeros(m, m)
    dif = ũ[p .+ (1:m)] - v[p .+ (1:m)]
    for i = 1:m
        if dif[i] >= 0.0
            JP[i, i] = 1.0
        else
            JP[i, i] = 0.0
        end
    end

    if z[k] - z[3k] >= 0.0
        ℓ = 1.0
    else
        ℓ = 0.0
    end

    [-I zeros(p, m) zeros(p) I zeros(p, m) zeros(p) I zeros(p, m) zeros(p);
     zeros(m, p) -JP zeros(m) zeros(m, p) I zeros(m) zeros(m, p) JP zeros(m);
     [zeros(1, p) zeros(1, m) -ℓ zeros(1, p) zeros(1, m) 1.0 zeros(1, p) zeros(1, m) ℓ]]
end

function J(z)
    [(I + Q) -I -I
     Ju(z);
     I -I zeros(k, k)]
end

ũ = zeros(k)
u = zeros(k)
v = zeros(k)
ũ[end] = 1.0
u[end] = 1.0
v[end] = 1.0

z = [ũ; u; v]
z = rand(3k)
F(z)
rank(Ju(z))
J(z)
rank(ForwardDiff.jacobian(F, z))
norm(J(z) - ForwardDiff.jacobian(F, z), Inf)
gmres(J(z), -F(z))
# (J(z)' * J(z)) \ (J(z)' * F(z))
# Δ = zero(z)
# Δ[k] = 1.0
# Δ[2k] = 1.0
# Δ[3k] = 1.0
# gmres!(Δ, J(z), -1.0 * F(z), abstol = 1.0)

function solve()
    ũ = zeros(k)
    u = zeros(k)
    v = zeros(k)
    ũ[end] = 1.0
    u[end] = 1.0
    v[end] = 1.0

    z = [ũ; u; v]
    # z = rand(3k)
    # Δ = zero(z)

    extra_iters = 0

    for i = 1:500
        _F = F(z)
        _J = J(z)
        Δ = gmres(_J, 1.0 * _F, abstol = 1.0e-12, maxiter = i + extra_iters)

        iter = 0
        α = 1.0
        while norm(F(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(_F)^2.0 && α > 1.0e-4
            α = 0.5 * α
            # println("   α = $α")
            iter += 1
            if iter > 100
                @error "line search fail"
                # return z
                x = z[k .+ (1:p)]
                y = z[k + p .+ (1:m)]
                τ = z[k + k]
                κ = z[3k]
                println("τ = $τ")
                println("κ = $κ")
                return x ./ τ, z
            end
        end

        if α <= 1.0e-4
            extra_iters += 1
        end

        println("iter ($i) - norm: $(norm(F(z)))")

        z .-= α * Δ
    end

    # return z
    x = z[k .+ (1:p)]
    y = z[k + p .+ (1:m)]
    τ = z[k + k]
    κ = z[3k]
    println("τ = $τ")
    println("κ = $κ")
    return x ./ τ, z
end

x_sol, z_sol = solve()

x_sol
b - A * x_sol
norm(x_sol - x.value)

(Q - I) * Pc(z_sol[1:k]) + z_sol[1:k]


using SMACOF
using Plots
using Random
gr()
ENV["GKSwstype"] = "100"
Random.seed!(2021)

include("../example_template.jl")

gifdir = "examples/wda_smacof/"
filename = "noiseless"

generate_gif((X, Δ) -> wda_smacof(Δ,
                    η=0.9,
                    Xinit=rand(size(X)...),
                    verbose=true,
                    anchors=X,
                    ε=1e-4,
                    DA_itmax=30,
                    return_history=true)[2],
            gifdir, filename, n=500)


using Test
using QuantumLattices.Prerequisites.Combinatorics

@testset "Combinations" begin
    @test Combinations{0}("abc")|>collect == [()]
    @test Combinations{1}("abc")|>collect == [('a',), ('b',), ('c',)]
    @test Combinations{2}("abc")|>collect == [('a', 'b'), ('a', 'c'), ('b', 'c')]
    @test Combinations{3}("abc")|>collect == [('a', 'b', 'c')]
end

@testset "DulCombinations" begin
    @test DulCombinations{0}("abc")|>collect == [()]
    @test DulCombinations{1}("abc")|>collect == [('a',), ('b',), ('c',)]
    @test DulCombinations{2}("abc")|>collect == [('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'b'), ('b', 'c'), ('c', 'c')]
    @test DulCombinations{3}("abc")|>collect == [('a', 'a', 'a'), ('a', 'a', 'b'), ('a', 'a', 'c'), ('a', 'b', 'b'), ('a', 'b', 'c'), ('a', 'c', 'c'), ('b', 'b', 'b'), ('b', 'b', 'c'), ('b', 'c', 'c'), ('c', 'c', 'c')]
end

@testset "Permutations" begin
    @test Permutations{0}("abc")|>collect == [()]
    @test Permutations{1}("abc")|>collect == [('a',), ('b',), ('c',)]
    @test Permutations{2}("abc")|>collect == [('a', 'b'), ('a', 'c'), ('b', 'a'), ('b', 'c'), ('c', 'a'), ('c', 'b')]
    @test Permutations{3}("abc")|>collect == [('a', 'b', 'c'), ('a', 'c', 'b'), ('b', 'a', 'c'), ('b', 'c', 'a'), ('c', 'a', 'b'), ('c', 'b', 'a')]
end

@testset "DulPermutations" begin
    @test DulPermutations{0}("abc")|>collect == [()]
    @test DulPermutations{1}("abc")|>collect == [('a',), ('b',), ('c',)]
    @test DulPermutations{2}("abc")|>collect == [('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'a'), ('b', 'b'), ('b', 'c'), ('c', 'a'), ('c', 'b'), ('c', 'c')]
    @test DulPermutations{3}("abc")|>collect == [(i, j, k) for i∈"abc" for j∈"abc" for k∈"abc"]
end


include("utility.jl")

base = "ATGC"
prob1 = [0.5, 0.25, 0.125, 0.125]
prob2 = [0.25, 0.25, 0.25, 0.25]
prob3 = [0.125, 0.125, 0.25, 0.5]

function generate_modelA(filename, length, base, prob)
    for i = 2:size(prob,1)
        prob[i] += prob[i-1]
    end
    # print(prob)

    f = open("data/$filename", "w")
    write(f, "$length\n")
    for i = 1:length
        cur = rand()
        # println(cur)
        j = 1
        while cur > prob[j]
            j += 1
        end
        write(f, "$(base[j])")
    end
end

function generate_modelB(filename, length, base)

end

generate_modelA("modelC.txt", 1000, base, prob2)

println("loss for modelA with compressor 1: ", kl_sdistance(prob2, prob1))
println("loss for modelA with compressor 1: ", kl_sdistance(prob2, prob2))
println("loss for modelA with compressor 1: ", kl_sdistance(prob2, prob3))


using DataFrames

"""
translate the element (row in the dataset) to index in the
histogram of the domain, according to the columns dimension sizes.
params:
    element:    iterable that contain the values of the its columns.
    domains:    the domain sizes of the columns.
"""
function get_flatten_index(element, domains)
    counter = 0
    pow = 1
    for (i,s) in zip(element,domains)
        counter += i*pow
        pow *= s
    end
    counter
end

"""
translate the index in the histogram of the domain to the corresponding
element in the dataset (row in the dataset), according to the columns dimension sizes.
params:
    index:    index in the histogram (integer).
    domains:    the domain sizes of the columns.
"""
function reverse_flatten_index(index, dim)

    new_index = zeros(Int64, length(dim))

    for (i,s) in enumerate(dim)
        new_index[i] += index % s
        index -= index % s
        index /= s
    end

    new_index

end

"""
get a list of values, and return how many times the each value is found.
return: dictionary{value : counter}
"""
function countMembers(itr)
    d = Dict{eltype(itr), Int64}()
    for val in itr
        if isa(val, Number) && isnan(val)
            continue
        end
        d[val] = get!(d, val, 0) + 1
    end
    return d
end

"""
create a mapping for the elements of a domain, into integers
    params:
        domain_elements:    the elements of the domain
    return:
        A tuple of two mappings: (index => elemnt, elemnt => index)
"""
function getMapping(domain_elements::Array)

    direct_mapping = Dict{Int64, Any}()
    reverse_mapping2 = Dict{Any, Int64}()

    for (i,element) in enumerate(domain_elements)
        direct_mapping[i-1] = element
        reverse_mapping2[element] = i-1
    end

    direct_mapping, reverse_mapping2

end

"""
Get a dataset and make it into a histogram of its domain.
    params:
        data:               A dataframe containing the dataset
        domain_elements:    list of the domain elements of every column (a list of lists)
    return:
        histogram:  the histogram of the dataset over the domain.
        mapping:    a list of tuples of mappings (index => elemnt, elemnt => index) for every column domain.
"""
function createHistogram(data::DataFrames.DataFrame, domains_elements::Array)

    n, d = size(data)

    domain_sizes = [length(domain)  for domain in domains_elements]
    histogram_size = prod(domain_sizes)

    histogram = zeros(Float64, histogram_size)

    mappings = []

    for domain_elements in domains_elements
        push!(mappings, getMapping(domain_elements))
    end

    for i = 1:n
        x = [data[i,j] for j=1:d]
        index = [mappings[j][2][x[j]] for j=1:d]
        histogram_index = get_flatten_index(index, domain_sizes) + 1
        histogram[histogram_index] += 1
    end

    histogram, mappings

end

""""
create histogram from a matrix of data.
"""
function createHistogram(data::AbstractArray, domains_elements::Array)

    createHistogram(DataFrame(samples), domains_elements)

end

"""
Recreate a database from histogram according to the mapping of the original elements of the dataset domain.
(the mapping from the creating of the histogram. see "createHistogram" function)
params:
    histogram:  a histogram of a dataset.
    mapping:    a list of tuples of mappings (index => elemnt, elemnt => index) for every column domain of the dataset.
    domain_elements:    list of the domain elements of every column (a list of lists)

return:
    df: A dataframe containing the recreated dataset
"""
function recreateDataset(histogram::Array{Float64, 1}, mappings::Array, domains_elements::Array)

    n = sum(histogram)
    d = length(domains_elements)

    domain_sizes = [length(domain)  for domain in domains_elements]

    df = DataFrame()
    for i=1:d
        df[i] = []
    end

    for index = 1:length(histogram)
        domain_element_encoded = reverse_flatten_index(index-1, domain_sizes)
        domain_element = [mappings[j][1][domain_element_encoded[j]] for j=1:d]
        for i=1:histogram[index]
            push!(df, copy(domain_element))
        end
    end

    df

end

function readContingencyTable(data::DataFrames.DataFrame, dimension::Int64)

    @assert length(size(data)) == 2

    n,temp = size(data)

    @assert temp == 2
    @assert n > 0

    hist = zeros(Float64, 2^dimension)

    for i in 1:n
        dig = digits(data[[:1]][1][i], base = 10)
        index = sum([dig[k]*2^(k-1) for k=1:length(dig)])
        hist[index] += data[[:2]][1][i]
    end

    hist
end


"""
    ILP(optimizer, max_iter)

ILP iteratively solves a linearized primal optimization to compute maximum allowable disturbance.
It iteratively adds the linear constraint to the problem.

# Problem requirement
1. Network: any depth, ReLU activation
2. Input: hyperrectangle
3. Output: PolytopeComplement

# Return
`AdversarialResult`

# Method
Iteratively solve a linear encoding of the problem.
It only considers the linear piece of the network that has the same activation pattern as the reference input.
Default `optimizer` is `GLPKSolverMIP()`.
We provide both iterative method and non-iterative method to solve the LP problem.
Default `iterative` is `true`.

# Property
Sound but not complete.

# Reference
[O. Bastani, Y. Ioannou, L. Lampropoulos, D. Vytiniotis, A. Nori, and A. Criminisi,
"Measuring Neural Net Robustness with Constraints,"
in *Advances in Neural Information Processing Systems*, 2016.](https://arxiv.org/abs/1605.07262)
"""
@with_kw struct ILP <: Solver
    optimizer = GLPK.Optimizer
    iterative::Bool = true
end

function solve(solver::ILP, problem::Problem)
    nnet = problem.network
    x = problem.input.center
    model = Model(solver)
    model[:δ] = δ = get_activation(nnet, x)
    z = init_vars(model, nnet, :z, with_input=true)
    add_complementary_set_constraint!(model, problem.output, last(z))
    o = max_disturbance!(model, first(z) - problem.input.center)

    if !solver.iterative
        encode_network!(model, nnet, StandardLP())
        optimize!(model)
        termination_status(model) != OPTIMAL && return AdversarialResult(:unknown)
        x = value(first(z))
        return interpret_result(solver, x, problem.input)
    end

    encode_network!(model, nnet, LinearRelaxedLP())
    while true
        optimize!(model)
        termination_status(model) != OPTIMAL && return AdversarialResult(:unknown)
        x = value(first(z))
        matched, index = match_activation(nnet, x, δ)
        if matched
            return interpret_result(solver, x, problem.input)
        end
        add_constraint!(model, nnet, z, δ, index)
    end
end

function interpret_result(solver::ILP, x, input)
    radius = abs.(x .- center(input))

    if all(radius .>= radius_hyperrectangle(input))
        return AdversarialResult(:holds, minimum(radius))
    else
        return AdversarialResult(:violated, minimum(radius))
    end
end

function add_constraint!(model::Model,
                         nnet::Network,
                         z::Vector{Vector{VariableRef}},
                         δ::Vector{Vector{Bool}},
                         (i, j)::Tuple{Int64, Int64})
    layer = nnet.layers[i]
    val = layer.weights[j, :]' * z[i] + layer.bias[j]
    if δ[i][j]
        @constraint(model, val >= 0.0)
    else
        @constraint(model, val <= 0.0)
    end
end

function match_activation(nnet::Network, x::Vector{Float64}, δ::Vector{Vector{Bool}})
    curr_value = x
    for (i, layer) in enumerate(nnet.layers)
        curr_value = layer.weights * curr_value + layer.bias
        matched, j = match_activation(layer, δ[i], curr_value)

        !matched && return false, (i,j)

        curr_value = layer.activation(curr_value)
    end
    return (true, nothing)
end

match_activation(L::Layer{Id}, args...) = (true, nothing)
function match_activation(L::Layer{ReLU}, δᵢ, val, tol = 1e-4)
    for (j, val) in enumerate(val)
        act = δᵢ[j]
        if act
            val < -tol && return false, j
        else
            val > tol  && return false, j
        end

    end
    return (true, nothing)
end


import Clustering
import PyPlot as Plt
import Seaborn as Sns

function clustermap(mtx::T where T <: AbstractMatrix{Float64}, gene_names::Vector{String}; gene_ord::Union{<:AbstractVector{<:Integer}, Bool}=true,
        cell_ord::Union{<:AbstractVector{<:Integer},Bool}=true, diag_genes::Bool=false, transpose::Bool=false, kwargs...)
    if typeof(cell_ord) === Bool
        if cell_ord
            cell_dists = 1 .- cor(mtx);
            cell_ord = Clustering.hclust(cell_dists, linkage=:ward).order;
        else
            cell_ord = 1:size(mtx, 2)
        end
    end

    if typeof(gene_ord) === Bool
        if gene_ord
            if diag_genes
                gene_ord = sortperm(vec(mapslices(x -> findmax(x)[2], mtx[:, cell_ord], dims=2)), rev=true);
            else
                gene_dists = 1 .- cor(mtx');
                gene_ord = Clustering.hclust(gene_dists, linkage=:ward).order;
            end
        else
            gene_ord = 1:size(mtx, 1)
        end
    end

    mtx = mtx[gene_ord, cell_ord]
    tick_fun = Plt.yticks
    if !transpose
        mtx = copy(mtx')
        tick_fun = Plt.xticks
    end

    plt = Sns.heatmap(mtx; kwargs...)
    tick_fun(1:length(gene_ord), gene_names[gene_ord], rotation=(transpose ? 0 : 90))
    return plt, cell_ord, gene_ord
end

function label_axis!(ax, label::String; x=-0.2, y=1.05, fontsize=24)
    ax.text(x, y, label, transform=ax.transAxes, fontsize=fontsize, va="top", ha="right", fontfamily="sans-serif", fontname="Helvetica")
end

function plot_embedding(embedding::Matrix{<:Real}, groups::Vector; mark_groups::Bool=true, ax=Plt.gca(),
        rasterized::Bool=true, fontsize::Int=10, palette::Union{Nothing, Vector{String}}=nothing, kwargs...)
    annot_labels = sort(unique(groups))
    ids_per_clust = B.split_ids(denserank(groups))
    kwargs = Dict{Symbol, Any}(kwargs...)
    for (i,ids) in enumerate(ids_per_clust)
        if palette !== nothing
            kwargs[:color] = palette[i]
        end

        ax.scatter(embedding[1,ids], embedding[2,ids]; label=annot_labels[i], rasterized=rasterized, kwargs...)
    end

    if mark_groups
        clust_centers = hcat([median(embedding[:, ids], dims=2) for ids in ids_per_clust]...);
        bbox_style = Dict(:facecolor=>"white", :alpha=>0.75, :edgecolor=>"none", :pad=>0.0)
        for (t, x, y) in zip(annot_labels, clust_centers[1,:], clust_centers[2,:])
            ax.annotate(t, (x,y), horizontalalignment="center", fontsize=fontsize, bbox=bbox_style)
        end
    end

    return ax
end


module TestCall

using SpecialFunctions

Base.@ccallable erfv(x::Vector{Float64})::Vector{Float64} = SpecialFunctions.erf.(x)

Base.@ccallable function erfv2(out::Vector{Float64}, x::Vector{Float64})::Cvoid
    # @assert length(out) == length(x)
    @inbounds for i in eachindex(out)
        out[i] = SpecialFunctions.erf(x[i])
    end
end

export erfv, erfv2

end


include("BorderedArrays.jl")

####################################################################################################
"""
	dt = DotTheta( (x,y) -> dot(x,y) / length(x) )

This parametric type allows to define a new dot product from the one saved in `dt::dot`. More precisely:

	dt(u1, u2, p1::T, p2::T, theta::T) where {T <: Real}

computes, the weigthed dot product ``\\langle (u_1,p_1), (u_2,p_2)\\rangle_\\theta = \\theta \\Re \\langle u_1,u_2\\rangle  +(1-\\theta)p_1p_2`` where ``u_i\\in\\mathbb R^N``. The ``\\Re`` factor is put to ensure a real valued result despite possible complex valued arguments.

	normtheta(u, p::T, theta::T)

Compute, the norm associated to weigthed dot product ``\\langle (u_1,p_1), (u_2,p_2)\\rangle_\\theta``.

!!! info "Info"
    This is used in the pseudo-arclength constraint with the dot product ``\\frac{1}{N} \\langle u_1,u_2\\rangle,\\quad u_i\\in\\mathbb R^N``
"""
struct DotTheta{Tdot}
	dot::Tdot		# defaults to (x,y) -> dot(x,y) / length(x)
end

DotTheta() = DotTheta( (x,y) -> dot(x,y) / length(x))

# Implementation of the dot product associated to DotTheta
function (dt::DotTheta)(u1, u2, p1::T, p2::T, theta::T) where {T <: Real}
	# we restrict the type of the parameters because for complex problems, we still want the parameter to be real
	return real(dt.dot(u1, u2) * theta + p1 * p2 * (one(T) - theta))
end

# Implementation of the norm associated to DotTheta
function (dt::DotTheta)(u, p::T, theta::T) where T
	# we restrict the type of the parameters because for complex problems, we still want the parameter to be real
	return sqrt(dt(u, u, p, p, theta))
end

(dt::DotTheta)(a::BorderedArray{vec, T}, b::BorderedArray{vec, T}, theta::T) where {vec, T} = dt(a.u, b.u, a.p, b.p, theta)

(dt::DotTheta)(a::BorderedArray{vec, T}, theta::T) where {vec, T} = dt(a.u, a.p, theta)

####################################################################################################
# equation of the arc length constraint
function arcLengthEq(dt::DotTheta, u, p, du, dp, theta, ds)
	return dt(u, du, p, dp, theta) - ds
end

####################################################################################################
abstract type AbstractTangentPredictor end
abstract type AbstractSecantPredictor <: AbstractTangentPredictor end

"""
	Natural predictor / corrector
"""
struct NaturalPred <: AbstractSecantPredictor end

"""
	Secant tangent predictor
"""
struct SecantPred <: AbstractSecantPredictor end

"""
	Bordered Tangent predictor
"""
struct BorderedPred <: AbstractTangentPredictor end

function getPredictor!(z_pred::M, z_old::M, tau::M, ds, algo::Talgo) where {T, vectype, M <: BorderedArray{vectype, T}, Talgo <: AbstractTangentPredictor}
	# we perform z_pred = z_old + ds * tau
	copyto!(z_pred, z_old)
	axpy!(ds, tau, z_pred)
end

# generic corrector based on Bordered formulation
function corrector(Fhandle, Jhandle, z_old::M, tau_old::M, z_pred::M,
			ds, theta, contparams, dottheta::DotTheta,
			algo::Talgo, linearalgo = MatrixFreeLBS();
			normC = norm,
			callback = (x, f, J, res, iteration; kwargs...) -> true, kwargs...) where {T, vectype, M<:BorderedArray{vectype, T}, Talgo <: AbstractTangentPredictor}
	return newtonPALC(Fhandle, Jhandle,
			z_old, tau_old, z_pred,
			ds, theta,
			contparams, dottheta; linearbdalgo = linearalgo, normN = normC, callback = callback, kwargs...)
end

# corrector based on natural formulation
function corrector(Fhandle, Jhandle, z_old::M, tau_old::M, z_pred::M,
			ds, theta, contparams, dottheta::DotTheta,
			algo::NaturalPred, linearalgo = MatrixFreeLBS();
			normC = norm,
			callback = (x, f, J, res, iteration; kwargs...) -> true, kwargs...) where {T, vectype, M <: BorderedArray{vectype, T}}
	res = newton(u -> Fhandle(u, z_pred.p),
				 u -> Jhandle(u, z_pred.p),
				 z_pred.u, contparams.newtonOptions;
				 normN = normC, callback = callback, kwargs...)
	return BorderedArray(res[1], z_pred.p), res[2], res[3], res[4]
end

# tangent computation using Natural / Secant predictor
# tau_new is the tangent prediction
function getTangent!(tau_new::M, z_new::M, z_old::M, tau_old::M, F, J,
	ds, theta, contparams, normtheta::DotTheta,
	algo::Talgo, verbosity, linearalgo) where {T, vectype, M <: BorderedArray{vectype, T}, Talgo <: AbstractSecantPredictor}
	(verbosity > 0) && println("--> predictor = ", algo)
	# secant predictor: tau = z_new - z_old; tau *= sign(ds) / normtheta(tau)
	copyto!(tau_new, z_new)
	minus!(tau_new, z_old)
	if algo isa SecantPred
		α = sign(ds) / normtheta(tau_new, theta)
	else
		α = sign(ds) / abs(tau_new.p)
	end
	rmul!(tau_new, α)
end

# tangent computation using Bordered system
# tau_new is the tangent prediction
function getTangent!(tau_new::M, z_new::M, z_old::M, tau_old::M, F, J,
	ds, theta, contparams, dottheta::DotTheta,
	algo::BorderedPred, verbosity, linearbdalgo) where {T, vectype, M <: BorderedArray{vectype, T}}
	(verbosity > 0) && println("--> predictor = Bordered")
	# tangent predictor
	ϵ = contparams.finDiffEps
	# dFdl = (F(z_old.u, z_old.p + ϵ) - F(z_old.u, z_old.p)) / ϵ
	dFdl = F(z_old.u, z_old.p + ϵ)
	minus!(dFdl, F(z_old.u, z_old.p))
	rmul!(dFdl, 1/ϵ)

	# tau = getTangent(J(z_old.u, z_old.p), dFdl, tau_old, theta, contparams.newtonOptions.linsolve)
	tau_normed = copyto!(similar(tau_old), tau_old) #copy(tau_old)
	rmul!(tau_normed, theta / length(tau_old.u), 1 - theta)
	# extract tangent as solution of bordered linear system, using zero(z_old.u)
	tauu, taup, flag, it = linearbdalgo( J(z_old.u, z_old.p), dFdl,
			tau_normed, rmul!(similar(z_old.u), false), T(1), theta)

	# the new tangent vector must preserve the direction along the curve
	α = T(1) / dottheta(tauu, tau_old.u, taup, tau_old.p, theta)

	# tau_new = α * tau
	copyto!(tau_new.u, tauu)
	tau_new.p = taup
	rmul!(tau_new, α)
end
################################################################################################
function arcLengthScaling(theta, contparams, tau::M, verbosity) where {M <: BorderedArray}
	# the arclength scaling algorithm is based on Salinger, Andrew G, Nawaf M Bou-Rabee, Elizabeth A Burroughs, Roger P Pawlowski, Richard B Lehoucq, Louis Romero, and Edward D Wilkes. “LOCA 1.0 Library of Continuation Algorithms: Theory and Implementation Manual,” March 1, 2002. https://doi.org/10.2172/800778.

	thetanew = theta
	g = abs(tau.p * theta)
	(verbosity > 0) && print("Theta changes from $(theta) to ")
	if (g > contparams.gMax)
		thetanew = contparams.gGoal / tau.p * sqrt( abs(1.0 - g^2) / abs(1.0 - tau.p^2) )
		if (thetanew < contparams.thetaMin)
		  thetanew = contparams.thetaMin;
		end
	end
	(verbosity > 0) && print("$(thetanew)\n")
	return thetanew
end
################################################################################################
function stepSizeControl(ds, theta, contparams, converged::Bool, it_newton_number::Int64, tau::M, verbosity) where {T, vectype, M<:BorderedArray{vectype, T}}
	if converged == false
		if  abs(ds) <= contparams.dsmin
			(verbosity > 0) && printstyled("*"^80*"\nFailure to converge with given tolerances\n"*"*"^80, color=:red)
			# we stop the continuation
			return ds, theta, true
		end
		dsnew = sign(ds) * max(abs(ds) / 2, contparams.dsmin);
		(verbosity > 0) && printstyled("Halving continuation step, ds=$(dsnew)\n", color=:red)
	else
		# control to have the same number of Newton iterations
		Nmax = contparams.newtonOptions.maxIter
		factor = (Nmax - it_newton_number) / Nmax
		dsnew = ds * (1 + contparams.a * factor^2)
		(verbosity > 0) && @show 1 + contparams.a * factor^2
	end

	# control step to stay between bounds
	if abs(dsnew) < contparams.dsmin
		dsnew = sign(dsnew) * contparams.dsmin
	end

	if abs(dsnew) > contparams.dsmax
		dsnew = sign(dsnew) * contparams.dsmax
	end

	if contparams.doArcLengthScaling
		thetanew = arcLengthScaling(theta, contparams, tau, verbosity)
	else
		thetanew = theta
	end
	@assert abs(dsnew) >= contparams.dsmin "Error with ds value"
	return dsnew, thetanew, false
end
####################################################################################################
"""
This is the classical matrix-free Newton Solver used to solve `F(x, p) = 0` together
with the scalar condition `n(x, p) = (x - x0) * xp + (p - p0) * lp - n0`
"""
function newtonPALC(F, Jh,
						z0::BorderedArray{vectype, T},
						tau0::BorderedArray{vectype, T},
						z_pred::BorderedArray{vectype, T},
						ds::T, theta::T,
						contparams::ContinuationPar{T},
						dottheta::DotTheta;
						linearbdalgo = BorderingBLS(),
						normN = norm,
						callback = (x, f, J, res, iteration, optionsN; kwargs...) ->  true, kwargs...) where {T, vectype}
	# Extract parameters
	newtonOpts = contparams.newtonOptions
	@unpack tol, maxIter, verbose, alpha, almin, linesearch = newtonOpts
	@unpack finDiffEps = contparams

	N = (x, p) -> arcLengthEq(dottheta, minus(x, z0.u), p - z0.p, tau0.u, tau0.p, theta, ds)
	normAC = (resf, resn) -> max(normN(resf), abs(resn))

	# Initialise iterations
	x = copyto!(similar(z_pred.u), z_pred.u) # copy(z_pred.u)
	p = z_pred.p
	x_pred = copyto!(similar(x), x) # copy(x)

	# Initialise residuals
	res_f = F(x, p);  res_n = N(x, p)

	dX   = copyto!(similar(res_f), res_f) # copy(res_f)
	dp   = T(0)
	# dFdp = (F(x, p + finDiffEps) - res_f) / finDiffEps
	dFdp = copyto!(similar(res_f), F(x, p + finDiffEps)) # copy(F(x, p + finDiffEps))
	minus!(dFdp, res_f)	# dFdp = dFdp - res_f
	rmul!(dFdp, T(1) / finDiffEps)

	res     = normAC(res_f, res_n)
	resHist = [res]
	it = 0

	# Displaying results
	verbose && displayIteration(it, 1, res)
	step_ok = true

	# Main loop
	while (res > tol) & (it < maxIter) & step_ok
		# copyto!(dFdp, (F(x, p + epsi) - F(x, p)) / epsi)
		copyto!(dFdp, F(x, p + finDiffEps)); minus!(dFdp, res_f); rmul!(dFdp, T(1) / finDiffEps)

		J = Jh(x, p)
		u, up, flag, liniter = linearbdalgo(J, dFdp, tau0, res_f, res_n, theta)

		if linesearch
			step_ok = false
			while !step_ok & (alpha > almin)
				# x_pred = x - alpha * u
				copyto!(x_pred, x)
				axpy!(-alpha, u, x_pred)

				p_pred = p - alpha * up
				copyto!(res_f, F(x_pred, p_pred))

				res_n  = N(x_pred, p_pred)
				res = normAC(res_f, res_n)

				if res < resHist[end]
					if (res < resHist[end] / 2) & (alpha < 1)
						alpha *= 2
					end
					step_ok = true
					copyto!(x, x_pred)
					p  = p_pred
				else
					alpha /= 2
				end
			end
		else
			minus!(x, u) 	# x .= x .- u
			p = p - up

			copyto!(res_f, F(x, p))

			res_n  = N(x, p)
			res = normAC(res_f, res_n)
		end

		push!(resHist, res)
		it += 1

		callback(x, res_f, J, res, it, contparams; kwargs...) == false && (it = maxIter)
		verbose && displayIteration(it, 1, res, liniter)

	end
	return BorderedArray(x, p), resHist, resHist[end] < tol, it
end
####################################################################################################


#
# A single execution of this file generates a single pair (a(r),r).
# The pair is then stored in "LS_pureIOU_2000_narrow_saved.jld2", together with
# all previously computed pairs.
#      If running for the first time (no "LS_pureIOU_2000_narrow_saved.jld2" exists)
#      uncomment lines 24-29. Put the #'s back for all subsequent runs.
#
include("LS-method-7s-init.jl");

# Assume the bounds on the wealth generated by the dual method and create the grid.
#
w_bounds=(-1.6292720455478487, 7.18257249498523);
begin
    Lg=2000 #the only place to change the grid size
    grda=[w_bounds[1]+i*(w_bounds[end]-w_bounds[1])/(Lg-1) for i=0:(Lg-1)]
    β=exp(-ρ)
end;

## Uncomment and execute the next block the very *first* time,
## *only* if "LS_pureIOU_2000_narrow_saved.jld2" does not exist in the current path.
## Put the #'s back for all subsequent runs.

#begin
#    a_list=Array{Float64}(undef, 0)
#    r_list=Array{Float64}(undef, 0)
#    LS_pureIOU_2000_narrow_saved=[r_list,a_list]
#    @save "LS_pureIOU_2000_narrow_saved.jld2" LS_pureIOU_2000_narrow_saved
#end;

##
## The next block guesses the next candidate for r.
## It can be bypassed by setting r=<value> explicitly. See line 68.
##

begin
    @load "LS_pureIOU_2000_narrow_saved.jld2" LS_pureIOU_2000_narrow_saved;
    r_list=LS_pureIOU_2000_narrow_saved[1]
    a_list=LS_pureIOU_2000_narrow_saved[2]
    if length(r_list)==0
        r=0.03
        elseif length(r_list)==1
        r=0.02
        else
        r=0.5*(r_list[findall(x->x<0,a_list)[end]]+r_list[findall(x->x>0,a_list)[end]])
    end
    push!(r_list,r)
end;

# the next line could take up to an hour
# arguments 2 and 3are the requested digits of convergence
#
@time Λ1,tbl2=LS_method(r,5,5,nos,hours,wage,β,grda,Lg,CPROB,R);

## Two important checkpoints!
sum(Λ1) # should print 1
[sum(Λ1[i,:]) for i=1:nos]-BSP # should print a list of zeros

begin
    policy=[grda[Int64(tbl2[kk,ii])] for kk=1:nos, ii=1:Lg]
    a=sum(Λ1.*policy) # the aggregate demand a(r)
end;

##
## If running with an explicit r (see line 32)
## be sure that a_list and r_list exist before running the next block.
##

begin
    a_list=push!(a_list,a)
    LS_pureIOU_2000_narrow_saved=[r_list,a_list,Λ1,tbl2]
    @save "LS_pureIOU_2000_narrow_saved.jld2" LS_pureIOU_2000_narrow_saved
end;