text:
module DegreesOfFreedom

using MacroTools: @capture
using Printf: @printf, @sprintf
using StaticArrays: SVector
using ..QuantumOperators: OperatorUnit, ID, OperatorPack, OperatorSum, Operator, Operators, valuetolatextext, idtype
using ..Spatials: AbstractPID, AbstractBond, Point, Bonds, pidtype
using ...Essentials: dtype
using ...Interfaces: decompose, dimension, add!
using ...Prerequisites: Float, atol, rtol, decimaltostr, concatenate
using ...Prerequisites.Traits: rawtype, fulltype, efficientoperations, commontype, parametertype
using ...Prerequisites.CompositeStructures: CompositeDict, CompositeTuple, NamedContainer
using ...Prerequisites.VectorSpaces: VectorSpace, VectorSpaceStyle, VectorSpaceCartesian, VectorSpaceDirectSummed, VectorSpaceDirectProducted

import LinearAlgebra: ishermitian
import ..QuantumOperators: script, optype
import ..Spatials: pidtype, rcoord, icoord
import ...Essentials: kind, reset!, update!
import ...Interfaces: id, value, rank, expand, expand!, ⊗, ⊕
import ...Prerequisites.Traits: parameternames, isparameterbound, contentnames, getcontent

export IID, SimpleIID, CompositeIID, Internal, SimpleInternal, CompositeInternal
export AbstractOID, Index, CompositeOID, OID, statistics, iidtype, ishermitian, indextype, oidtype, Hilbert
export IIDSpace, Subscript, Subscripts, @subscript_str, subscriptexpr, wildcard, diagonal, noconstrain
export AbstractCoupling, Coupling, Couplings, couplingcenters, couplingpoints, couplinginternals, @couplings
export Metric, OIDToTuple, Table
export TermFunction, TermAmplitude, TermCouplings, TermModulate, Term, ismodulatable, abbr

# IID and Internal
"""
    IID <: OperatorUnit

The id of an internal degree of freedom.
"""
abstract type IID <: OperatorUnit end

"""
    SimpleIID <: IID

The id of a simple internal degree of freedom.
"""
abstract type SimpleIID <: IID end
@inline statistics(iid::SimpleIID) = statistics(typeof(iid))

"""
    CompositeIID{T<:Tuple{Vararg{SimpleIID}}} <: IID

The composition of several single internal ids.
"""
struct CompositeIID{T<:Tuple{Vararg{SimpleIID}}} <: IID
    contents::T
end
Base.show(io::IO, ciid::CompositeIID) = @printf io "%s" join((string(ciid[i]) for i = 1:rank(ciid)), " ⊗ ")
@inline Base.length(ciid::CompositeIID) = length(typeof(ciid))
@inline Base.length(::Type{<:CompositeIID{T}}) where {T<:Tuple{Vararg{SimpleIID}}} = fieldcount(T)
@inline Base.getindex(ciid::CompositeIID, i::Int) = ciid.contents[i]
@inline Base.getproperty(ciid::CompositeIID, name::Symbol) = ciidgetproperty(ciid, Val(name))
@inline ciidgetproperty(ciid::CompositeIID, ::Val{:contents}) = getfield(ciid, :contents)
@inline ciidgetproperty(ciid::CompositeIID, ::Val{name}) where name = getproperty(getfield(ciid, :contents), name)

"""
    CompositeIID(contents::SimpleIID...)

Construct a composite iid from a set of simple iids.
"""
@inline CompositeIID(contents::SimpleIID...) = CompositeIID(contents)

"""
    rank(ciid::CompositeIID) -> Int
    rank(::Type{<:CompositeIID{T}}) where {T<:Tuple{Vararg{SimpleIID}}} -> Int

Get the number of simple iids in a composite iid.
"""
@inline rank(ciid::CompositeIID) = rank(typeof(ciid))
@inline rank(::Type{<:CompositeIID{T}}) where {T<:Tuple{Vararg{SimpleIID}}} = fieldcount(T)

"""
    iidtype(ciid::CompositeIID, i::Integer)
    iidtype(::Type{<:CompositeIID{T}}, i::Integer) where {T<:Tuple{Vararg{SimpleIID}}}

Get the type of the ith simple iid in a composite iid.
"""
@inline iidtype(ciid::CompositeIID, i::Integer) = iidtype(typeof(ciid), i)
@inline iidtype(::Type{<:CompositeIID{T}}, i::Integer) where {T<:Tuple{Vararg{SimpleIID}}} = fieldtype(T, i)

"""
    ⊗(iid₁::SimpleIID, iid₂::SimpleIID) -> CompositeIID
    ⊗(iid::SimpleIID, ciid::CompositeIID) -> CompositeIID
    ⊗(ciid::CompositeIID, iid::SimpleIID) -> CompositeIID
    ⊗(ciid₁::CompositeIID, ciid₂::CompositeIID) -> CompositeIID

Direct product between simple iids and composite iids.
"""
@inline ⊗(iid₁::SimpleIID, iid₂::SimpleIID) = CompositeIID(iid₁, iid₂)
@inline ⊗(iid::SimpleIID, ciid::CompositeIID) = CompositeIID(iid, ciid.contents...)
@inline ⊗(ciid::CompositeIID, iid::SimpleIID) = CompositeIID(ciid.contents..., iid)
@inline ⊗(ciid₁::CompositeIID, ciid₂::CompositeIID) = CompositeIID(ciid₁.contents..., ciid₂.contents...)

"""
    Internal{I<:IID} <: VectorSpace{I}

The whole internal degrees of freedom at a single point.
"""
abstract type Internal{I<:IID} <: VectorSpace{I} end

"""
    SimpleInternal{I<:SimpleIID} <: Internal{I}

The simple internal degrees of freedom at a single point.
"""
abstract type SimpleInternal{I<:SimpleIID} <: Internal{I} end
@inline VectorSpaceStyle(::Type{<:SimpleInternal}) = VectorSpaceCartesian()
@inline Base.show(io::IO, i::SimpleInternal) = @printf io "%s(%s)" i|>typeof|>nameof join(("$name=$(getfield(i, name))" for name in i|>typeof|>fieldnames), ", ")

"""
    statistics(i::SimpleInternal) -> Symbol
    statistics(::Type{<:SimpleInternal{I}}) where {I<:SimpleIID} -> Symbol

Get the statistics of a simple internal space.
"""
@inline statistics(i::SimpleInternal) = statistics(typeof(i))
@inline statistics(::Type{<:SimpleInternal{I}}) where {I<:SimpleIID} = statistics(I)

"""
    match(iid::SimpleIID, i::SimpleInternal) -> Bool
    match(::Type{I}, ::Type{SI}) where {I<:SimpleIID, SI<:SimpleInternal}

Judge whether a simple iid or a simple iid type matches a simple internal space or a simple internal space type.

Here, "match" means that the eltype of the simple internal space has the same type name with the simple iid.
"""
@inline Base.match(iid::SimpleIID, i::SimpleInternal) = match(typeof(iid), typeof(i))
@inline Base.match(::Type{I}, ::Type{SI}) where {I<:SimpleIID, SI<:SimpleInternal} = nameof(I)==nameof(eltype(SI))

"""
    filter(iid::SimpleIID, i::SimpleInternal) -> Union{Nothing, typeof(i)}
    filter(::Type{I}, i::SimpleInternal) where {I<:SimpleIID} -> Union{Nothing, typeof(i)}

Filter a simple internal space with respect to the input `iid` or type `I`.
"""
@inline Base.filter(iid::SimpleIID, i::SimpleInternal) = filter(typeof(iid), i)
@inline Base.filter(::Type{I}, i::SimpleInternal) where {I<:SimpleIID} = match(I, typeof(i)) ? i : nothing

"""
    filter(iid::SimpleIID, ::Type{T}) where {T<:SimpleInternal}
    filter(::Type{I}, ::Type{T}) where {I<:SimpleIID, T<:SimpleInternal}

Filter the type of a simple internal space with respect to the input `iid` or type `I`.
"""
@inline Base.filter(iid::SimpleIID, ::Type{T}) where {T<:SimpleInternal} = filter(typeof(iid), T)
@inline Base.filter(::Type{I}, ::Type{T}) where {I<:SimpleIID, T<:SimpleInternal} = match(I, T) ? T : nothing

"""
    CompositeInternal{K, T<:Tuple{Vararg{SimpleInternal}}} <: Internal{IID}

The composition of several single internal spaces.
"""
struct CompositeInternal{K, T<:Tuple{Vararg{SimpleInternal}}} <: Internal{IID}
    contents::T
    function CompositeInternal{K}(contents::Tuple{Vararg{SimpleInternal}}) where K
        @assert K∈(:⊕, :⊗) "CompositeInternal error: kind must be either :⊕ for direct sum or :⊗ for direct product."
        new{K, typeof(contents)}(contents)
    end
end
@inline kind(ci::CompositeInternal) = kind(typeof(ci))
@inline kind(::Type{<:CompositeInternal{K}}) where K = K
@inline Base.show(io::IO, ci::CompositeInternal) = @printf io "%s" join((string(ci.contents[i]) for i = 1:rank(ci)), @sprintf(" %s ", kind(ci)))
@inline VectorSpaceStyle(::Type{<:CompositeInternal{:⊕}}) = VectorSpaceDirectSummed()
@inline VectorSpaceStyle(::Type{<:CompositeInternal{:⊗}}) = VectorSpaceDirectProducted()
@inline @generated Base.eltype(::Type{<:CompositeInternal{:⊕, T}}) where {T<:Tuple{Vararg{SimpleInternal}}} = mapreduce(eltype, typejoin, fieldtypes(T), init=Union{})
@inline @generated Base.eltype(::Type{<:CompositeInternal{:⊗, T}}) where {T<:Tuple{Vararg{SimpleInternal}}} = CompositeIID{Tuple{map(eltype, fieldtypes(T))...}}
@inline SimpleIID(iid::SimpleIID, ::CompositeInternal) = iid
@inline CompositeIID(iid::Tuple, ::CompositeInternal) = CompositeIID(iid)

"""
    CompositeInternal{K}(contents::SimpleInternal...) where K

Construct a composite internal space from a set of simple internal spaces.
"""
@inline CompositeInternal{K}(contents::SimpleInternal...) where K = CompositeInternal{K}(contents)

"""
    rank(ci::CompositeInternal) -> Int
    rank(::Type{<:CompositeInternal{K, T}}) where {K, T<:Tuple{Vararg{SimpleInternal}}} -> Int

Get the number of simple internal spaces in a composite internal space.
"""
@inline rank(ci::CompositeInternal) = rank(typeof(ci))
@inline rank(::Type{<:CompositeInternal{K, T}}) where {K, T<:Tuple{Vararg{SimpleInternal}}} = fieldcount(T)

"""
    ⊗(i₁::SimpleInternal, i₂::SimpleInternal) -> CompositeInternal{:⊗}
    ⊗(i::SimpleInternal, ci::CompositeInternal{:⊗}) -> CompositeInternal{:⊗}
    ⊗(ci::CompositeInternal{:⊗}, i::SimpleInternal) -> CompositeInternal{:⊗}
    ⊗(ci₁::CompositeInternal{:⊗}, ci₂::CompositeInternal{:⊗}) -> CompositeInternal{:⊗}

Direct product between simple internal spaces and composite internal spaces.
"""
@inline ⊗(i₁::SimpleInternal, i₂::SimpleInternal) = CompositeInternal{:⊗}(i₁, i₂)
@inline ⊗(i::SimpleInternal, ci::CompositeInternal{:⊗}) = CompositeInternal{:⊗}(i, ci.contents...)
@inline ⊗(ci::CompositeInternal{:⊗}, i::SimpleInternal) = CompositeInternal{:⊗}(ci.contents..., i)
@inline ⊗(ci₁::CompositeInternal{:⊗}, ci₂::CompositeInternal{:⊗}) = CompositeInternal{:⊗}(ci₁.contents..., ci₂.contents...)

"""
    ⊕(i₁::SimpleInternal, i₂::SimpleInternal) -> CompositeInternal{:⊕}
    ⊕(i::SimpleInternal, ci::CompositeInternal{:⊕}) -> CompositeInternal{:⊕}
    ⊕(ci::CompositeInternal{:⊕}, i::SimpleInternal) -> CompositeInternal{:⊕}
    ⊕(ci₁::CompositeInternal{:⊕}, ci₂::CompositeInternal{:⊕}) -> CompositeInternal{:⊕}

Direct product between simple internal spaces and composite internal spaces.
"""
@inline ⊕(i₁::SimpleInternal, i₂::SimpleInternal) = CompositeInternal{:⊕}(i₁, i₂)
@inline ⊕(i::SimpleInternal, ci::CompositeInternal{:⊕}) = CompositeInternal{:⊕}(i, ci.contents...)
@inline ⊕(ci::CompositeInternal{:⊕}, i::SimpleInternal) = CompositeInternal{:⊕}(ci.contents..., i)
@inline ⊕(ci₁::CompositeInternal{:⊕}, ci₂::CompositeInternal{:⊕}) = CompositeInternal{:⊕}(ci₁.contents..., ci₂.contents...)

"""
    filter(iid::SimpleIID, ci::CompositeInternal) -> Union{Nothing, SimpleInternal, CompositeInternal}
    filter(::Type{I}, ci::CompositeInternal) where {I<:SimpleIID} -> Union{Nothing, SimpleInternal, CompositeInternal}

Filter the composite internal space and select those that matches `I` or the type of `iid`.
"""
@inline Base.filter(iid::SimpleIID, ci::CompositeInternal) = filter(typeof(iid), ci)
@inline Base.filter(::Type{I}, ci::CompositeInternal) where {I<:SimpleIID} = filterhelper₁(I, ci, filtermatches(I, typeof(ci))|>Val)
@inline @generated function filtermatches(::Type{I}, ::Type{CompositeInternal{K, C}}) where {I<:SimpleIID, K, C<:Tuple{Vararg{SimpleInternal}}}
    exprs = []
    for i = 1:fieldcount(C)
        T = fieldtype(C, i)
        push!(exprs, :(match(I, $T)))
    end
    return Expr(:tuple, exprs...)
end
@inline @generated function filterhelper₁(::Type{I}, ci::CompositeInternal, ::Val{BS}) where {I<:SimpleIID, BS}
    exprs = []
    for (i, B) in enumerate(BS)
        B && push!(exprs, :(ci.contents[$i]))
    end
    length(exprs)==0 && return
    length(exprs)==1 && return first(exprs)
    return Expr(:call, Expr(:curly, :CompositeInternal, QuoteNode(kind(ci))), exprs...)
end

"""
    filter(iid::SimpleIID, ::Type{C}) where {C<:CompositeInternal}
    filter(::Type{I}, ::Type{C}) where {I<:SimpleIID, C<:CompositeInternal}

Filter the type of a composite internal space and select those that matches `I` or the type of `iid`.
"""
@inline Base.filter(iid::SimpleIID, ::Type{C}) where {C<:CompositeInternal} = filter(typeof(iid), C)
@inline Base.filter(::Type{I}, ::Type{C}) where {I<:SimpleIID, C<:CompositeInternal} = filterhelper₂(I, C, filtermatches(I, C)|>Val)
@inline @generated function filterhelper₂(::Type{I}, ::Type{CompositeInternal{K, C}}, ::Val{BS}) where {I<:SimpleIID, K, C<:Tuple{Vararg{SimpleInternal}}, BS}
    exprs = []
    for (i, B) in enumerate(BS)
        B && push!(exprs, :(fieldtype(C, $i)))
    end
    length(exprs)==0 && return
    length(exprs)==1 && return first(exprs)
    return Expr(:curly, :CompositeInternal, QuoteNode(K), Expr(:curly, :Tuple, exprs...))
end

# OID
"""
    AbstractOID <: OperatorUnit

Abstract type of operator id.
"""
abstract type AbstractOID <: OperatorUnit end

"""
    Index{P<:AbstractPID, I<:SimpleIID} <: AbstractOID

The index of a degree of freedom, which consist of the spatial part and the internal part.
"""
struct Index{P<:AbstractPID, I<:SimpleIID} <: AbstractOID
    pid::P
    iid::I
end
@inline parameternames(::Type{<:Index}) = (:pid, :iid)
@inline isparameterbound(::Type{<:Index}, ::Val{:pid}, ::Type{P}) where {P<:AbstractPID} = !isconcretetype(P)
@inline isparameterbound(::Type{<:Index}, ::Val{:iid}, ::Type{I}) where {I<:SimpleIID} = !isconcretetype(I)

"""
    pidtype(index::Index)
    pidtype(::Type{I}) where {I<:Index}

Get the type of the spatial part of an index.
"""
@inline pidtype(index::Index) = pidtype(typeof(index))
@inline pidtype(::Type{I}) where {I<:Index} = parametertype(I, 1)

"""
    iidtype(index::Index)
    iidtype(::Type{I}) where {I<:Index}

Get the type of the internal part of an index.
"""
@inline iidtype(index::Index) = iidtype(typeof(index))
@inline iidtype(::Type{I}) where {I<:Index} = parametertype(I, 2)

"""
    statistics(index::Index) -> Symbol
    statistics(::Type{<:Index{<:AbstractPID, I}}) where {I<:SimpleIID} -> Symbol

Get the statistics of an index.
"""
@inline statistics(index::Index) = statistics(typeof(index))
@inline statistics(::Type{<:Index{<:AbstractPID, I}}) where {I<:SimpleIID} = statistics(I)

"""
    adjoint(index::Index) -> typeof(index)

Get the adjoint of an index.
"""
@inline Base.adjoint(index::Index) = rawtype(typeof(index))(index.pid, adjoint(index.iid))

"""
    CompositeOID{I<:Index} <: AbstractOID

The abstract type of composite operator id.
"""
abstract type CompositeOID{I<:Index} <: AbstractOID end
@inline contentnames(::Type{<:CompositeOID}) = (:index,)
@inline parameternames(::Type{<:CompositeOID}) = (:index,)
@inline isparameterbound(::Type{<:CompositeOID}, ::Val{:index}, ::Type{I}) where {I<:Index} = !isconcretetype(I)

"""
    indextype(::CompositeOID)
    indextype(::Type{<:CompositeOID})

Get the index type of a composite operator id.
"""
@inline indextype(oid::CompositeOID) = indextype(typeof(oid))
@inline @generated indextype(::Type{I}) where {I<:CompositeOID} = parametertype(supertype(I, :CompositeOID), :index)

"""
    statistics(oid::CompositeOID) -> Symbol
    statistics(::Type{<:CompositeOID{I}}) where {I<:Index} -> Symbol

Get the statistics of a composite operator id.
"""
@inline statistics(oid::CompositeOID) = statistics(typeof(oid))
@inline statistics(::Type{<:CompositeOID{I}}) where {I<:Index} = statistics(I)

"""
    OID{I<:Index, V<:SVector} <: CompositeOID{I}

Operator id.
"""
struct OID{I<:Index, V<:SVector} <: CompositeOID{I}
    index::I
    rcoord::V
    icoord::V
    OID(index::Index, rcoord::V, icoord::V) where {V<:SVector} = new{typeof(index), V}(index, oidcoord(rcoord), oidcoord(icoord))
end
@inline contentnames(::Type{<:OID}) = (:index, :rcoord, :icoord)
@inline parameternames(::Type{<:OID}) = (:index, :coord)
@inline isparameterbound(::Type{<:OID}, ::Val{:coord}, ::Type{V}) where {V<:SVector} = !isconcretetype(V)
@inline Base.hash(oid::OID, h::UInt) = hash((oid.index, Tuple(oid.rcoord)), h)
@inline Base.propertynames(::ID{OID}) = (:indexes, :rcoords, :icoords)
@inline Base.show(io::IO, oid::OID) = @printf io "OID(%s, %s, %s)" oid.index oid.rcoord oid.icoord
@inline oidcoord(vector::SVector) = vector
@inline oidcoord(vector::SVector{N, Float}) where N = SVector(ntuple(i->vector[i]===-0.0 ? 0.0 : vector[i], Val(N)))

"""
    OID(index::Index, rcoord, icoord)
    OID(index::Index; rcoord, icoord)

Construct an operator id.
"""
@inline OID(index::Index, rcoord, icoord) = OID(index, SVector{length(rcoord)}(rcoord), SVector{length(icoord)}(icoord))
@inline OID(index::Index; rcoord, icoord) = OID(index, rcoord, icoord)

"""
    adjoint(oid::OID) -> typeof(oid)

Get the adjoint of an operator id.
"""
@inline Base.adjoint(oid::OID) = OID(oid.index', oid.rcoord, oid.icoord)

"""
    oidtype(I::Type{<:SimpleInternal}, P::Type{<:Point}, ::Val)

Get the compatible oid type from the combination of the internal part and the spatial part.
"""
@inline function oidtype(I::Type{<:SimpleInternal}, P::Type{<:Point}, ::Val)
    fulltype(OID, NamedTuple{(:index, :coord), Tuple{fulltype(Index, NamedTuple{(:pid, :iid), Tuple{pidtype(P), eltype(I)}}), SVector{dimension(P), dtype(P)}}})
end

"""
    rcoord(opt::Operator{<:Number, <:ID{OID}}) -> SVector

Get the whole rcoord of an operator.
"""
@inline function rcoord(opt::Operator{<:Number, <:ID{OID}})
    rank(opt)==1 && return id(opt)[1].rcoord
    rank(opt)==2 && return id(opt)[1].rcoord-id(opt)[2].rcoord
    error("rcoord error: not supported rank($(rank(opt))) of $(nameof(opt)).")
end

"""
    icoord(opt::Operator{<:Number, <:ID{OID}}) -> SVector

Get the whole icoord of an operator.
"""
@inline function icoord(opt::Operator{<:Number, <:ID{OID}})
    rank(opt)==1 && return id(opt)[1].icoord
    rank(opt)==2 && return id(opt)[1].icoord-id(opt)[2].icoord
    error("icoord error: not supported rank($(rank(opt))) of $(nameof(opt)).")
end

"""
    script(::Val{:rcoord}, oid::OID; kwargs...) -> String
    script(::Val{:icoord}, oid::OID; kwargs...) -> String

Get the `:rcoord/:icoord` script of an oid.
"""
@inline script(::Val{:rcoord}, oid::OID; kwargs...) = @sprintf "[%s]" join(valuetolatextext.(oid.rcoord), ", ")
@inline script(::Val{:icoord}, oid::OID; kwargs...) = @sprintf "[%s]" join(valuetolatextext.(oid.icoord), ", ")

"""
    script(::Val{:integralicoord}, oid::OID; vectors, kwargs...)

Get the integral script of the icoord of an oid.
"""
function script(::Val{:integralicoord}, oid::OID; vectors, kwargs...)
    rcoeff = decompose(oid.icoord, vectors...)
    icoeff = Int.(round.(rcoeff))
    @assert isapprox(efficientoperations, rcoeff, icoeff) "script error: mismatched icoord of oid and input vectors."
    return @sprintf "[%s]" join(icoeff, ", ")
end

"""
    script(::Val{attr}, oid::OID; kwargs...) where attr

Get the `attr` script of an oid, which is contained in its index.
"""
@inline script(::Val{attr}, oid::OID; kwargs...) where attr = script(Val(attr), oid.index; kwargs...)

# Hilbert
"""
    Hilbert{I<:Internal, P<:AbstractPID} <: CompositeDict{P, I}

Hilbert space at a lattice.
"""
struct Hilbert{I<:Internal, P<:AbstractPID} <: CompositeDict{P, I}
    map::Function
    contents::Dict{P, I}
end
@inline contentnames(::Type{<:Hilbert}) = (:map, :contents)

"""
    Hilbert(ps::Pair...)
    Hilbert(kv)

Construct a Hilbert space the same way as a Dict.
"""
@inline Hilbert(ps::Pair...) = Hilbert(ps)
function Hilbert(kv)
    contents = Dict(kv)
    map = pid -> contents[pid]
    return Hilbert(map, contents)
end

"""
    Hilbert(map::Function, pids::AbstractVector{<:AbstractPID})

Construct a Hilbert space from a function and a set of point ids.

Here, `map` maps a `AbstractPID` to an `Internal`.
"""
@inline function Hilbert(map::Function, pids::AbstractVector{<:AbstractPID})
    I = commontype(map, Tuple{Vararg{Any}}, Internal)
    return Hilbert{I}(map, pids)
end

"""
    Hilbert{I}(map::Function, pids::AbstractVector{<:AbstractPID}) where {I<:Internal}

Construct a Hilbert space from a function and a set of point ids.

Here, `map` maps a `AbstractPID` to an `Internal`.
"""
function Hilbert{I}(map::Function, pids::AbstractVector{<:AbstractPID}) where {I<:Internal}
    contents = Dict{pids|>eltype, I}()
    for pid in pids
        contents[pid] = map(pid)
    end
    return Hilbert(map, contents)
end

"""
    reset!(hilbert::Hilbert, pids) -> Hilbert

Reset the Hilbert space with new pids.
"""
function reset!(hilbert::Hilbert, pids)
    empty!(hilbert)
    for pid in pids
        hilbert[pid] = hilbert.map(pid)::valtype(hilbert)
    end
    hilbert
end

# Coupling and Couplings
"""
    IIDSpace{I<:IID, V<:Internal, Kind} <: VectorSpace{IID}

The space expanded by a "labeled" iid.
"""
struct IIDSpace{I<:IID, V<:Internal, Kind} <: VectorSpace{IID}
    iid::I
    internal::V
    IIDSpace(iid::IID, internal::Internal, ::Val{Kind}=Val(:info)) where Kind = new{typeof(iid), typeof(internal), Kind}(iid, internal)
end
@inline kind(iidspace::IIDSpace) = kind(typeof(iidspace))
@inline kind(::Type{<:IIDSpace{<:IID, <:Internal, Kind}}) where Kind = Kind
@inline Base.eltype(::Type{<:IIDSpace{<:IID, V}}) where {V<:Internal} = eltype(V)

@inline VectorSpaceStyle(::Type{<:IIDSpace{<:SimpleIID, <:SimpleInternal}}) = VectorSpaceCartesian()
@inline Base.CartesianIndex(iid::SimpleIID, iidspace::IIDSpace{<:SimpleIID, <:SimpleInternal}) = CartesianIndex(iid, iidspace.internal)
@inline function Base.getindex(::VectorSpaceCartesian, iidspace::IIDSpace{<:SimpleIID, <:SimpleInternal}, index::CartesianIndex)
    return rawtype(eltype(iidspace))(index, iidspace.internal)
end

@inline VectorSpaceStyle(::Type{<:IIDSpace{<:CompositeIID, <:CompositeInternal}}) = VectorSpaceDirectProducted()
@inline function getcontent(iidspace::IIDSpace{<:CompositeIID, <:CompositeInternal}, ::Val{:contents})
    return map((iid, internal)->IIDSpace(iid, internal, Val(kind(iidspace))), CompositeIID(iidspace).contents, CompositeInternal(iidspace).contents)
end
@inline CompositeIID(iids::Tuple{Vararg{SimpleIID}}, ::IIDSpace{<:CompositeIID, <:CompositeInternal}) = CompositeIID(iids)
@inline CompositeIID(iidspace::IIDSpace{<:CompositeIID, <:CompositeInternal}) = iidspace.iid
@inline CompositeInternal(iidspace::IIDSpace{<:CompositeIID, <:CompositeInternal}) = iidspace.internal

"""
    expand(iid::SimpleIID, internal::SimpleInternal) -> IIDSpace
    expand(iids::NTuple{N, SimpleIID}, internals::NTuple{N, SimpleInternal}) where N -> IIDSpace

Get the space expanded by a set of "labeled" iids.
"""
@inline expand(iid::SimpleIID, internal::SimpleInternal) = expand((iid,), (internal,))
@inline expand(iids::NTuple{N, SimpleIID}, internals::NTuple{N, SimpleInternal}) where N = IIDSpace(CompositeIID(iids), CompositeInternal{:⊗}(internals))

@inline diagonal(xs...) = length(xs)<2 ? true : all(map(==(xs[1]), xs))
@inline noconstrain(_...) = true
const wildcard = Symbol("*")
"""
    Subscript{P<:Tuple} <: CompositeTuple{P}

The subscript representative of a certain internal degree of freedom.
"""
struct Subscript{P<:Tuple} <: CompositeTuple{P}
    pattern::P
    rep::String
    constraint::Function
end
@inline contentnames(::Type{<:Subscript}) = (:contents, :rep, :constraint)
@inline getcontent(subscript::Subscript, ::Val{:contents}) = subscript.pattern
@inline Base.:(==)(subs₁::Subscript, subs₂::Subscript) = subs₁.pattern==subs₂.pattern && subs₁.rep==subs₂.rep
@inline Base.:isequal(subs₁::Subscript, subs₂::Subscript) = isequal(subs₁.pattern, subs₂.pattern) && isequal(subs₁.rep, subs₂.rep)
function Base.show(io::IO, subscript::Subscript)
    if subscript.rep ∈ ("diagonal", "noconstrain", "constant")
        @printf io "[%s]" join(subscript.pattern, " ")
    else
        @printf io "%s" subscript.rep
    end
end

"""
    Subscript(N::Int)
    Subscript(pattern::Tuple, check_constant::Bool=false)

Construct the subscript representative of a certain internal degree of freedom.
"""
@inline Subscript(N::Int) = Subscript(Val(N))
@inline Subscript(::Val{N}) where N = Subscript(ntuple(i->wildcard, Val(N)), "diagonal", diagonal)
@inline Subscript(pattern::Tuple, check_constant::Bool=false) = Subscript(pattern, Val(check_constant))
@inline function Subscript(pattern::Tuple, ::Val{false})
    any(map(p->isa(p, Symbol), pattern)) && error("Subscript error: wrong constant pattern.")
    return Subscript(pattern, "noconstrain", noconstrain)
end
@inline function Subscript(pattern::Tuple, ::Val{true})
    any(map(p->isa(p, Symbol), pattern)) && error("Subscript error: wrong constant pattern.")
    return Subscript(pattern, "constant", (xs...)->xs==pattern)
end

"""
    rank(subscript::Subscript) -> Int
    rank(::Type{<:Subscript}) -> Int

Get the number of the whole variables of the subscript.
"""
@inline rank(subscript::Subscript) = rank(typeof(subscript))
@inline rank(::Type{T}) where {T<:Subscript} = length(T)

"""
    match(subscript::Subscript, values::Tuple) -> Bool

Judge whether a set of values matches the pattern specified by the subscript.
"""
@inline function Base.match(subscript::Subscript, values::Tuple)
    @assert length(subscript)==length(values) "match error: mismatched length of values."
    return subscript.constraint(values...)::Bool
end

"""
    subscript"..." -> Subscript

Construct the subscript from a literal string.
"""
macro subscript_str(str)
    expr = Meta.parse(str)
    expr.head==:toplevel || return subscriptexpr(expr)
    @assert length(expr.args)==2 && isa(expr.args[2], Bool) "@subscript_str error: wrong pattern."
    return subscriptexpr(expr.args[1], expr.args[2])
end
function subscriptexpr(expr::Expr, check_constant::Bool=false)
    if @capture(expr, op_(cp_))
        @assert op.head∈(:hcat, :vect) "subscriptexpr error: wrong pattern."
        pattern, condition = Tuple(op.args), cp
        rep = @sprintf "[%s](%s)" join(pattern, " ") condition
    else
        @assert expr.head∈(:hcat, :vect) "subscriptexpr error: wrong pattern."
        pattern, condition = Tuple(expr.args), true
        rep = @sprintf "[%s]" join(pattern, " ")
    end
    if !any(map(p->isa(p, Symbol), pattern))
        check_constant && return :(Subscript($pattern, Val(true)))
        return :(Subscript($pattern, Val(false)))
    end
    paramargs, groups = Symbol[], Dict{Symbol, Vector{Symbol}}()
    for sub in pattern
        isa(sub, Symbol) || begin
            paramarg = gensym("paramarg")
            push!(paramargs, paramarg)
            condition = Expr(Symbol("&&"), condition, Expr(:call, :(==), paramarg, sub))
            continue
        end
        if sub∉paramargs
            push!(paramargs, sub)
            groups[sub]=[sub]
        else
            paramarg = gensym("paramarg")
            push!(paramargs, paramarg)
            push!(groups[sub], :(==))
            push!(groups[sub], paramarg)
        end
    end
    for group in values(groups)
        length(group)==1 && continue
        condition = Expr(Symbol("&&"), condition, Expr(:comparison, group...))
    end
    name = gensym("subconstraint")
    constraint = :($name($(paramargs...)) = $condition)
    return Expr(:block, constraint, :(Subscript($pattern, $rep, $name)))
end

"""
    Subscripts{T<:Tuple{Vararg{NamedContainer{Subscript}}}, R<:Tuple{Vararg{Pair{UnitRange{Int}, <:Tuple{Vararg{String}}}}}} <: OperatorUnit

The complete set of subscripts of the internal degrees of freedom.
"""
struct Subscripts{T<:Tuple{Vararg{NamedContainer{Subscript}}}, R<:Tuple{Vararg{Pair{UnitRange{Int}, <:Tuple{Vararg{String}}}}}} <: OperatorUnit
    contents::T
    rep::R
end
@inline Base.:(==)(subs₁::Subscripts, subs₂::Subscripts) = subs₁.contents==subs₂.contents
@inline Base.hash(subscripts::Subscripts, h::UInt) = hash(subscripts.rep, h)
@inline Base.length(subscripts::Subscripts) = length(typeof(subscripts))
@inline Base.length(::Type{<:Subscripts{T}}) where {T<:Tuple{Vararg{NamedContainer{Subscript}}}} = fieldcount(T)
@inline Base.getindex(subscripts::Subscripts, i::Integer) = subscripts.contents[i]
@inline Base.iterate(subscripts::Subscripts, state=1) = state>length(subscripts) ? nothing : (subscripts[state], state+1)
function Base.show(io::IO, subscripts::Subscripts)
    for (i, segment) in enumerate(subscripts.contents)
        i>1 && @printf io "%s" " × "
        for (j, (field, subscript)) in enumerate(pairs(segment))
            j>1 && @printf io "%s" " ⊗ "
            @printf io "%s%s" field subscript
        end
    end
end
function Base.repr(subscripts::Subscripts, slice, field::Symbol)
    result = []
    for (i, segment) in enumerate(slice)
        i>1 && push!(result, "×")
        push!(result, @sprintf "%s" getfield(subscripts.contents[segment], field))
    end
    return join(result)
end

"""
    Subscripts(contents::NamedContainer{Subscript}...)

Construct the complete set of subscripts.
"""
function Subscripts(contents::NamedContainer{Subscript}...)
    for segment in contents
        length(segment)>1 && @assert mapreduce(length, ==, values(segment)) "Subscripts error: mismatched ranks."
    end
    reps = map(content->map(subscript->subscript.rep, values(content)), contents)
    counts = (0, cumsum(map(content->rank(first(content)), contents))...)
    return Subscripts(contents, map((i, rep)->(counts[i]+1:counts[i+1])=>rep, ntuple(i->i, Val(fieldcount(typeof(reps)))), reps))
end

"""
    rank(subscripts::Subscripts) -> Int
    rank(::Type{<:Subscripts{T}}) where {T<:Tuple{Vararg{NamedContainer{Subscript}}}} -> Int

Get the rank of the subscripts.
"""
@inline rank(subscripts::Subscripts) = rank(typeof(subscripts))
@inline @generated function rank(::Type{<:Subscripts{T}}) where {T<:Tuple{Vararg{NamedContainer{Subscript}}}}
    sum(rank(fieldtype(fieldtype(T, i), 1)) for i = 1:fieldcount(T))
end

"""
    rank(subscripts::Subscripts, i::Integer) -> Int
    rank(::Type{<:Subscripts{T}}, i::Integer) where {T<:Tuple{Vararg{NamedContainer{Subscript}}}} -> Int

Get the rank of the ith homogenous segment of the subscripts.
"""
@inline rank(subscripts::Subscripts, i::Integer) = rank(typeof(subscripts), i)
@inline rank(::Type{<:Subscripts{T}}, i::Integer) where {T<:Tuple{Vararg{NamedContainer{Subscript}}}} = rank(fieldtype(fieldtype(T, i), 1))

"""
    match(subscripts::Subscripts, iids::Tuple{Vararg{SimpleIID}}) -> Bool
    match(subscripts::Subscripts, ciid::CompositeIID) -> Bool

Judge whether a composite iid matches the patterns specified by the subscripts.
"""
@inline Base.match(subscripts::Subscripts, ciid::CompositeIID) = match(subscripts, ciid.contents)
@generated function Base.match(subscripts::Subscripts, iids::Tuple{Vararg{SimpleIID}})
    length(subscripts)==0 && return true
    @assert rank(subscripts)==fieldcount(iids) "match error: mismatched rank of iids and subscripts."
    exprs, count = [], 1
    for i = 1:length(subscripts)
        start, stop = count, count+rank(subscripts, i)-1
        for field in fieldnames(fieldtype(fieldtype(subscripts, :contents), i))
            field = QuoteNode(field)
            paramvalue = Expr(:tuple, [:(getfield(iids[$j], $field)) for j = start:stop]...)
            push!(exprs, :(match(getfield(subscripts[$i], $field), $paramvalue)))
        end
        count = stop+1
    end
    return Expr(:call, :all, Expr(:tuple, exprs...))
end

"""
    *(subscripts₁::Subscripts, subscripts₂::Subscripts) -> Subscripts

Get the combination of two sets of subscripts.
"""
@inline Base.:*(subscripts₁::Subscripts, subscripts₂::Subscripts) = Subscripts((subscripts₁.contents..., subscripts₂.contents...)...)

"""
    AbstractCoupling{V, I<:ID{OperatorUnit}} <: OperatorPack{V, I}

The abstract coupling intra/inter internal degrees of freedom at different lattice points.
"""
abstract type AbstractCoupling{V, I<:ID{OperatorUnit}} <: OperatorPack{V, I} end
@inline ID{SimpleIID}(coupling::AbstractCoupling) = id(coupling)
@inline Subscripts(coupling::AbstractCoupling) = Subscripts()

"""
    couplingcenters(coupling::AbstractCoupling, bond::AbstractBond, info::Val) -> NTuple{rank(coupling), Int}

Get the acting centers of the coupling on a bond.
"""
@inline couplingcenters(coupling::AbstractCoupling, point::Point, ::Val) = ntuple(i->1, Val(rank(coupling)))

"""
    couplingpoints(coupling::AbstractCoupling, bond::AbstractBond, info::Val) -> NTuple{rank(coupling), eltype(bond)}

Get the points where each order of the coupling acts on.
"""
@inline function couplingpoints(coupling::AbstractCoupling, bond::AbstractBond, info::Val)
    centers = couplingcenters(coupling, bond, info)
    return NTuple{rank(coupling), eltype(bond)}(bond[centers[i]] for i = 1:rank(coupling))
end

"""
    couplinginternals(coupling::AbstractCoupling, bond::AbstractBond, hilbert::Hilbert, info::Val) -> NTuple{rank(coupling), SimpleInternal}

Get the internal spaces where each order of the coupling acts on.
"""
@inline function couplinginternals(coupling::AbstractCoupling, bond::AbstractBond, hilbert::Hilbert, info::Val)
    centers = couplingcenters(coupling, bond, info)
    internals = ntuple(i->hilbert[bond[centers[i]].pid], Val(rank(coupling)))
    return map((iid, internal)->filter(iid, internal), ID{SimpleIID}(coupling), internals)
end

"""
    expand(coupling::AbstractCoupling, bond::AbstractBond, hilbert::Hilbert, info::Val)

Expand a coupling with the given bond and Hilbert space.
"""
function expand(coupling::AbstractCoupling, bond::AbstractBond, hilbert::Hilbert, info::Val)
    points = couplingpoints(coupling, bond, info)
    internals = couplinginternals(coupling, bond, hilbert, info)
    @assert rank(coupling)==length(points)==length(internals) "expand error: mismatched rank."
    return CExpand(value(coupling), points, IIDSpace(CompositeIID(ID{SimpleIID}(coupling)), CompositeInternal{:⊗}(internals), info), Subscripts(coupling))
end
struct CExpand{V, N, P<:AbstractPID, SV<:SVector, S<:IIDSpace, C<:Subscripts}
    value::V
    pids::NTuple{N, P}
    rcoords::NTuple{N, SV}
    icoords::NTuple{N, SV}
    iidspace::S
    subscripts::C
end
function CExpand(value, points::NTuple{N, P}, iidspace::IIDSpace, subscripts::Subscripts) where {N, P<:Point}
    pids = NTuple{N, pidtype(P)}(points[i].pid for i = 1:N)
    rcoords = NTuple{N, SVector{dimension(P), dtype(P)}}(points[i].rcoord for i = 1:N)
    icoords = NTuple{N, SVector{dimension(P), dtype(P)}}(points[i].icoord for i = 1:N)
    return CExpand(value, pids, rcoords, icoords, iidspace, subscripts)
end
@inline Base.eltype(ex::CExpand) = eltype(typeof(ex))
@inline @generated function Base.eltype(::Type{<:CExpand{V, N, P, SV, S}}) where {V, N, P<:AbstractPID, SV<:SVector, S<:IIDSpace}
    return Operator{V, Tuple{map(I->OID{Index{P, I}, SV}, fieldtypes(fieldtype(eltype(S), :contents)))...}}
end
@inline Base.IteratorSize(::Type{<:CExpand}) = Base.SizeUnknown()
function Base.iterate(ex::CExpand, state=iterate(ex.iidspace))
    result = nothing
    while !isnothing(state)
        ciid, state = state
        if match(ex.subscripts, ciid)
            result = Operator(ex.value, ID(OID, ID(Index, ex.pids, ciid.contents), ex.rcoords, ex.icoords)), iterate(ex.iidspace, state)
            break
        else
            state = iterate(ex.iidspace, state)
        end
    end
    return result
end

"""
    Coupling{V, I<:ID{SimpleIID}, C<:Subscripts} <: AbstractCoupling{V, Tuple{CompositeIID{I}, C}}

The coupling intra/inter internal degrees of freedom at different lattice points.
"""
struct Coupling{V, I<:ID{SimpleIID}, C<:Subscripts} <: AbstractCoupling{V, Tuple{CompositeIID{I}, C}}
    value::V
    cid::I
    subscripts::C
    function Coupling(value::Number, cid::ID{SimpleIID}, subscripts::Subscripts=Subscripts())
        new{typeof(value), typeof(cid), typeof(subscripts)}(value, cid, subscripts)
    end
end
@inline parameternames(::Type{<:Coupling}) = (:value, :cid, :subscripts)
@inline isparameterbound(::Type{<:Coupling}, ::Val{:cid}, ::Type{I}) where {I<:ID{SimpleIID}} = !isconcretetype(I)
@inline isparameterbound(::Type{<:Coupling}, ::Val{:subscripts}, ::Type{C}) where {C<:Subscripts} = !isconcretetype(C)
@inline getcontent(coupling::Coupling, ::Val{:id}) = ID(CompositeIID(coupling.cid), coupling.subscripts)
@inline rank(::Type{<:Coupling{V, I} where V}) where {I<:ID{SimpleIID}} = fieldcount(I)
@inline Coupling(value::Number, id::Tuple{CompositeIID, Subscripts}) = Coupling(value, id[1].contents, id[2])
@inline ID{SimpleIID}(coupling::Coupling) = coupling.cid
@inline Subscripts(coupling::Coupling) = coupling.subscripts

"""
    *(cp₁::Coupling, cp₂::Coupling) -> Coupling

Get the multiplication between two couplings.
"""
@inline Base.:*(cp₁::Coupling, cp₂::Coupling) = Coupling(cp₁.value*cp₂.value, ID(cp₁.cid, cp₂.cid), cp₁.subscripts*cp₂.subscripts)

"""
    Couplings(cps::AbstractCoupling...)

A pack of couplings intra/inter internal degrees of freedom at different lattice points.

Alias for `OperatorSum{<:AbstractCoupling, <:ID{OperatorUnit}}`.
"""
const Couplings{C<:AbstractCoupling, I<:ID{OperatorUnit}} = OperatorSum{C, I}
@inline Couplings(cps::AbstractCoupling...) = OperatorSum(cps)
@inline Couplings(cps::Couplings) = cps

"""
    @couplings cps -> Couplings

Convert an expression/literal to a set of couplings.
"""
macro couplings(cps) :(Couplings($(esc(cps)))) end

# Metric
"""
    Metric <: Function

The rules for measuring a concrete oid so that oids can be compared.

As a function, every instance should accept only one positional argument, i.e. the concrete oid to be measured.
"""
abstract type Metric <: Function end
@inline Base.:(==)(m₁::T, m₂::T) where {T<:Metric} = ==(efficientoperations, m₁, m₂)
@inline Base.isequal(m₁::T, m₂::T) where {T<:Metric} = isequal(efficientoperations, m₁, m₂)
@inline (M::Type{<:Metric})(::Type{I}) where {I<:CompositeOID} = M(indextype(I))
@inline (metric::Metric)(oid::CompositeOID) = metric(getcontent(oid, :index))
@inline Base.valtype(::Type{M}, ::Type{I}) where {M<:Metric, I<:CompositeOID} = valtype(M, indextype(I))
@inline (M::Type{<:Metric})(::Type{H}) where {H<:Hilbert} = M(Index{H|>keytype, H|>valtype|>eltype})

"""
    OIDToTuple{Fields} <: Metric

A rule that converts an oid to a tuple by iterating over a set of selected fields in a specific order.
"""
struct OIDToTuple{Fields} <: Metric
    OIDToTuple(fields::Tuple{Vararg{Symbol}}) = new{fields}()
end
@inline OIDToTuple(fields::Symbol...) = OIDToTuple(fields)

"""
    keys(::OIDToTuple{Fields}) where Fields -> Fields
    keys(::Type{<:OIDToTuple{Fields}}) where Fields -> Fields

Get the names of the selected fields.
"""
@inline Base.keys(::OIDToTuple{Fields}) where Fields = Fields
@inline Base.keys(::Type{<:OIDToTuple{Fields}}) where Fields = Fields

"""
    filter(f::Function, oidtotuple::OIDToTuple) -> OIDToTuple

Filter the selected fields.
"""
@inline Base.filter(f::Function, oidtotuple::OIDToTuple) = OIDToTuple(Tuple(field for field in keys(oidtotuple) if f(field)))

"""
    OIDToTuple(::Type{I}) where {I<:Index}

Construct the conversion rule from the information of subtypes of `AbstractOID`.
"""
@inline OIDToTuple(::Type{I}) where {I<:Index} = OIDToTuple(fieldnames(pidtype(I))..., (fieldnames(iidtype(I)))...)

"""
    valtype(::Type{<:OIDToTuple}, ::Type{<:Index})

Get the valtype of applying an `OIDToTuple` rule to a subtype of `AbstractOID`.
"""
@inline @generated function Base.valtype(::Type{M}, ::Type{I}) where {M<:OIDToTuple, I<:Index}
    types = []
    for field in keys(M)
        if hasfield(pidtype(I), field)
            push!(types, fieldtype(pidtype(I), field))
        elseif hasfield(iidtype(I), field)
            push!(types, fieldtype(iidtype(I), field))
        end
    end
    return  Expr(:curly, :Tuple, types...)
end

"""
    (oidtotuple::OIDToTuple)(index::Index) -> Tuple

Convert a concrete oid to a tuple.
"""
@inline @generated function (oidtotuple::OIDToTuple)(index::Index)
    exprs = []
    for name in keys(oidtotuple)
        field = QuoteNode(name)
        if hasfield(pidtype(index), name)
            push!(exprs, :(getfield(index.pid, $field)))
        elseif hasfield(iidtype(index), name)
            push!(exprs, :(getfield(index.iid, $field)))
        end
    end
    return Expr(:tuple, exprs...)
end

"""
    Table{I, B<:Metric} <: CompositeDict{I, Int}

The table of oid-sequence pairs.
"""
struct Table{I, B<:Metric} <: CompositeDict{I, Int}
    by::B
    contents::Dict{I, Int}
end
@inline contentnames(::Type{<:Table}) = (:by, :contents)
@inline Table{I}(by::Metric) where {I<:AbstractOID} = Table(by, Dict{valtype(typeof(by), I), Int}())
@inline vec2dict(vs::AbstractVector) = Dict{eltype(vs), Int}(v=>i for (i, v) in enumerate(vs))

"""
    getindex(table::Table, oid::AbstractOID) -> Int

Inquiry the sequence of an oid.
"""
@inline Base.getindex(table::Table, oid::AbstractOID) = table[table.by(oid)]

"""
    haskey(table::Table, oid::AbstractOID) -> Bool
    haskey(table::Table, id::ID{AbstractOID}) -> Tuple{Vararg{Bool}}

Judge whether a single oid or a set of oids have been assigned with sequences in table.
"""
@inline Base.haskey(table::Table, oid::AbstractOID) = haskey(table, table.by(oid))
@inline Base.haskey(table::Table, id::ID{AbstractOID}) = map(oid->haskey(table, oid), id)

"""
    Table(oids::AbstractVector{<:AbstractOID}, by::Metric=OIDToTuple(eltype(oids)))

Convert a set of concrete oids to the corresponding table of oid-sequence pairs.

The input oids are measured by the input `by` function with the duplicates removed. The resulting unique values are sorted, which determines the sequence of the input `oids`. Note that two oids have the same sequence if their converted values are equal to each other.
"""
@inline Table(oids::AbstractVector{<:AbstractOID}, by::Metric=OIDToTuple(eltype(oids))) = Table(by, [by(oid) for oid in oids]|>unique!|>sort!|>vec2dict)

"""
    Table(hilbert::Hilbert, by::Metric=OIDToTuple(typeof(hilbert))) -> Table

Get the oid-sequence table of a Hilbert space.
"""
function Table(hilbert::Hilbert, by::Metric=OIDToTuple(typeof(hilbert)))
    result = Index{hilbert|>keytype, hilbert|>valtype|>eltype}[]
    for (pid, internal) in hilbert
        for iid in internal
            push!(result, (result|>eltype)(pid, iid))
        end
    end
    return Table(result, by)
end

"""
    union(tables::Table...) -> Table

Unite several oid-sequence tables.
"""
function Base.union(tables::Table...)
    @assert mapreduce(table->table.by, ==, tables) "union error: all input tables should have the same `by` attribute."
    indices = (tables|>eltype|>keytype)[]
    for table in tables
        for index in keys(table)
            push!(indices, index)
        end
    end
    return Table(tables[1].by, indices|>unique!|>sort!|>vec2dict)
end

"""
    reset!(table::Table, oids::AbstractVector{<:AbstractOID}) -> Table

Reset a table by a new set of oids.
"""
function reset!(table::Table, oids::AbstractVector{<:AbstractOID})
    empty!(table)
    for (i, id) in enumerate([table.by(oid) for oid in oids]|>unique!|>sort!)
        table[id] = i
    end
    return table
end

"""
    reset!(table::Table, hilbert::Hilbert) -> Table

Reset a table by a Hilbert space.
"""
function reset!(table::Table, hilbert::Hilbert)
    indices = Index{hilbert|>keytype, hilbert|>valtype|>eltype}[]
    for (pid, internal) in hilbert
        for iid in internal
            push!(indices, (indices|>eltype)(pid, iid))
        end
    end
    reset!(table, indices)
end

# Term
"""
    TermFunction <: Function

Abstract type for concrete term functions.
"""
abstract type TermFunction <: Function end
@inline Base.:(==)(tf1::TermFunction, tf2::TermFunction) = ==(efficientoperations, tf1, tf2)
@inline Base.isequal(tf1::TermFunction, tf2::TermFunction) = isequal(efficientoperations, tf1, tf2)

"""
    TermAmplitude(amplitude::Union{Function, Nothing}=nothing)

The function for the amplitude of a term.
"""
struct TermAmplitude{A<:Union{Function, Nothing}} <: TermFunction
    amplitude::A
    TermAmplitude(amplitude::Union{Function, Nothing}=nothing) = new{typeof(amplitude)}(amplitude)
end
@inline (termamplitude::TermAmplitude{Nothing})(args...; kwargs...) = 1
@inline (termamplitude::TermAmplitude{<:Function})(args...; kwargs...) = termamplitude.amplitude(args...; kwargs...)

"""
    TermCouplings(couplings::Union{Couplings, Function})

The function for the couplings of a term.
"""
struct TermCouplings{C₁<:Union{Function, Couplings}, C₂<:Union{Couplings, Nothing}} <: TermFunction
    couplings::C₁
    TermCouplings(couplings::Couplings) = new{typeof(couplings), Nothing}(couplings)
    TermCouplings{C}(couplings::Function) where {C<:Couplings} = new{typeof(couplings), C}(couplings)
    TermCouplings(couplings::Function) = new{typeof(couplings), commontype(couplings, Tuple{Vararg{Any}}, Couplings)}(couplings)
end
@inline Base.valtype(termcouplings::TermCouplings) = valtype(typeof(termcouplings))
@inline Base.valtype(::Type{<:TermCouplings{C}}) where {C<:Couplings} = C
@inline Base.valtype(::Type{<:TermCouplings{<:Function, C}}) where {C<:Couplings} = C
@inline (termcouplings::TermCouplings{<:Couplings})(args...; kwargs...) = termcouplings.couplings
@inline (termcouplings::TermCouplings{<:Function})(args...; kwargs...) = termcouplings.couplings(args...; kwargs...)

"""
    TermModulate(id::Symbol, modulate::Function)
    TermModulate(id::Symbol, modulate::Bool)

The function for the modulation of a term.
"""
struct TermModulate{M<:Union{Function, Val{true}, Val{false}}, id} <: TermFunction
    modulate::M
    TermModulate(id::Symbol, modulate::Function) = new{typeof(modulate), id}(modulate)
    TermModulate(id::Symbol, modulate::Bool=true) = new{Val{modulate}, id}(modulate|>Val)
end
@inline (termmodulate::TermModulate{Val{true}, id})(args...; kwargs...) where id = get(kwargs, id, nothing)
@inline (termmodulate::TermModulate{<:Function})(args...; kwargs...) = termmodulate.modulate(args...; kwargs...)
@inline ismodulatable(termmodulate::TermModulate) = ismodulatable(typeof(termmodulate))
@inline ismodulatable(::Type{<:TermModulate{Val{B}}}) where B = B
@inline ismodulatable(::Type{<:TermModulate{<:Function}}) = true

"""
    Term{K, I, V, B, C<:TermCouplings, A<:TermAmplitude, M<:TermModulate}

A term of a quantum lattice system.
"""
mutable struct Term{K, I, V, B, C<:TermCouplings, A<:TermAmplitude, M<:TermModulate}
    value::V
    bondkind::B
    couplings::C
    amplitude::A
    modulate::M
    factor::V
    function Term{K, I}(value, bondkind, couplings::TermCouplings, amplitude::TermAmplitude, modulate::TermModulate, factor) where {K, I}
        @assert isa(K, Symbol) "Term error: kind must be a Symbol."
        @assert isa(I, Symbol) "Term error: id must be a Symbol."
        @assert value==value' "Term error: only real values are allowed. Complex values should be specified by the amplitude function."
        V, B, C, A, M = typeof(value), typeof(bondkind), typeof(couplings), typeof(amplitude), typeof(modulate)
        new{K, I, V, B, C, A, M}(value, bondkind, couplings, amplitude, modulate, factor)
    end
end
@inline Base.:(==)(term1::Term, term2::Term) = ==(efficientoperations, term1, term2)
@inline Base.isequal(term1::Term, term2::Term) = isequal(efficientoperations, term1, term2)
@inline function Base.show(io::IO, term::Term)
    @printf io "%s{%s}(id=%s, value=%s, bondkind=%s, factor=%s)" kind(term) rank(term) id(term) decimaltostr(term.value) term.bondkind decimaltostr(term.factor)
end

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

Construct a term.
"""
@inline function Term{K}(id::Symbol, value, bondkind;
        couplings::Union{Function, Couplings},
        amplitude::Union{Function, Nothing}=nothing,
        modulate::Union{Function, Bool}=false
        ) where K
    Term{K, id}(value, bondkind, TermCouplings(couplings), TermAmplitude(amplitude), TermModulate(id, modulate), 1)
end

"""
    kind(term::Term) -> Symbol
    kind(::Type{<:Term) -> Symbol

Get the kind of a term.
"""
@inline kind(term::Term) = kind(typeof(term))
@inline kind(::Type{<:Term{K}}) where K = K

"""
    id(term::Term) -> Symbol
    id(::Type{<:Term) -> Symbol

Get the id of a term.
"""
@inline id(term::Term) = id(typeof(term))
@inline id(::Type{<:Term{K, I} where K}) where I = I

"""
    valtype(term::Term)
    valtype(::Type{<:Term)

Get the value type of a term.
"""
@inline Base.valtype(term::Term) = valtype(typeof(term))
@inline Base.valtype(::Type{<:Term{K, I, V} where {K, I}}) where V = V

"""
    rank(term::Term) -> Int
    rank(::Type{<:Term) -> Int

Get the rank of a term.
"""
@inline rank(term::Term) = rank(typeof(term))
@inline rank(::Type{<:Term{K, I, V, B, C} where {K, I, V, B}}) where {C<:TermCouplings} = rank(eltype(valtype(C)))

"""
    abbr(term::Term) -> Symbol
    abbr(::Type{<:Term}) -> Symbol

Get the abbreviation of the kind of a term.
"""
@inline abbr(term::Term) = abbr(typeof(term))
@inline abbr(::Type{<:Term}) = :tm

"""
    ismodulatable(term::Term) -> Bool
    ismodulatable(::Type{<:Term}) -> Bool

Judge whether a term could be modulated by its modulate function.
"""
@inline ismodulatable(term::Term) = ismodulatable(typeof(term))
@inline ismodulatable(::Type{<:Term{K, I, V, B, <:TermCouplings, <:TermAmplitude, M} where {K, I, V, B}}) where M = ismodulatable(M)

"""
    ishermitian(term::Term) -> Bool
    ishermitian(::Type{<:Term}) -> Bool
"""
@inline ishermitian(term::Term) = ishermitian(typeof(term))
@inline ishermitian(::Type{<:Term}) = error("ishermitian error: not implemented.")

"""
    repr(term::Term, bond::AbstractBond, hilbert::Hilbert) -> String

Get the repr representation of a term on a bond with a given Hilbert space.
"""
function Base.repr(term::Term, bond::AbstractBond, hilbert::Hilbert)
    cache = String[]
    if term.bondkind == bond|>kind
        value = term.value * term.amplitude(bond) * term.factor
        if abs(value) ≠ 0
            for coupling in term.couplings(bond)
                isnothing(iterate(expand(coupling, bond, hilbert, term|>kind|>Val))) || push!(cache, @sprintf "%s: %s" abbr(term) repr(value*coupling))
            end
        end
    end
    return join(cache, "\n")
end

"""
    replace(term::Term; kwargs...) -> Term

Replace some attributes of a term with key word arguments.
"""
@inline @generated function Base.replace(term::Term; kwargs...)
    exprs = [:(get(kwargs, $name, getfield(term, $name))) for name in QuoteNode.(term|>fieldnames)]
    return :(Term{kind(term), id(term)}($(exprs...)))
end

"""
    +(term::Term) -> Term
    -(term::Term) -> Term
    *(term::Term, factor) -> Term
    *(factor, term::Term) -> Term
    /(term::Term, factor) -> Term

Allowed arithmetic operations for a term.
"""
@inline Base.:+(term::Term) = term
@inline Base.:-(term::Term) = term * (-1)
@inline Base.:*(term::Term, factor) = factor * term
@inline Base.:*(factor, term::Term) = replace(term, factor=factor*term.factor)
@inline Base.:/(term::Term, factor) = term * (one(term.value)/factor)

"""
    one(term::Term) -> Term

Get a unit term.
"""
@inline Base.one(term::Term) = replace(term, value=one(term.value))

"""
    zero(term::Term) -> Term

Get a zero term.
"""
@inline Base.zero(term::Term) = replace(term, value=zero(term.value))

"""
    update!(term::Term, args...; kwargs...) -> Term

Update the value of a term by its `modulate` function.
"""
function update!(term::Term, args...; kwargs...)
    @assert ismodulatable(term) "update! error: not modulatable term."
    value = term.modulate(args...; kwargs...)
    isnothing(value) || (term.value = value)
    return term
end

"""
    optype(::Type{T}, ::Type{H}, ::Type{B}) where {T<:Term, H<:Hilbert, B<:AbstractBond}

Get the compatible `Operator` type from the type of a term, a Hilbert space and a bond.
"""
@inline function optype(::Type{T}, ::Type{H}, ::Type{B}) where {T<:Term, H<:Hilbert, B<:AbstractBond}
    C = eltype(valtype(fieldtype(T, :couplings)))
    @assert C<:Coupling "optype error: not supported."
    oidtypes = ntuple(i->oidtype(filter(fieldtype(parametertype(C, :cid), i) , valtype(H)), eltype(B), Val(kind(T))), Val(rank(C)))
    return fulltype(Operator, NamedTuple{(:value, :id), Tuple{valtype(T), Tuple{oidtypes...}}})
end

"""
    expand!(operators::Operators, term::Term, bond::AbstractBond, hilbert::Hilbert; half::Bool=false, table::Union{Nothing, Table}=nothing) -> Operators
    expand!(operators::Operators, term::Term, bonds, hilbert::Hilbert; half::Bool=false, table::Union{Nothing, Table}=nothing) -> Operators

Expand the operators of a term on a bond/set-of-bonds with a given Hilbert space.

The `half` parameter determines the behavior of generating operators, which falls into the following two categories
* `true`: "Hermitian half" of the generated operators
* `false`: "Hermitian whole" of the generated operators
"""
function expand!(operators::Operators, term::Term, bond::AbstractBond, hilbert::Hilbert; half::Bool=false, table::Union{Nothing, Table}=nothing)
    if term.bondkind == bond|>kind
        value = term.value * term.amplitude(bond) * term.factor
        if abs(value) ≠ 0
            hermitian = ishermitian(term)
            record = (isnothing(hermitian) && length(operators)>0) ? Set{idtype(eltype(operators))}() : nothing
            for coupling in term.couplings(bond)
                for opt in expand(coupling, bond, hilbert, term|>kind|>Val)
                    isapprox(opt.value, 0.0, atol=atol, rtol=rtol) && continue
                    !isnothing(table) && !all(haskey(table, opt.id)) && continue
                    if hermitian == true
                        add!(operators, opt*valtype(eltype(operators))(value/(half ? 2 : 1)))
                    elseif hermitian == false
                        opt = opt * valtype(eltype(operators))(value)
                        add!(operators, opt)
                        half || add!(operators, opt')
                    else
                        if !(isnothing(record) ? haskey(operators, opt.id') : opt.id'∈record)
                            isnothing(record) || push!(record, opt.id)
                            opt = opt * valtype(eltype(operators))(value/termfactor(opt.id, term|>kind|>Val))
                            add!(operators, opt)
                            half || add!(operators, opt')
                        end
                    end
                end
            end
        end
    end
    return operators
end
@inline function expand!(operators::Operators, term::Term, bonds, hilbert::Hilbert; half::Bool=false, table::Union{Nothing, Table}=nothing)
    for bond in bonds
        expand!(operators, term, bond, hilbert; half=half, table=table)
    end
    return operators
end
@inline function expand!(operators::Operators, term::Term, bonds::Bonds, hilbert::Hilbert; half::Bool=false, table::Union{Nothing, Table}=nothing)
    map(bs->expand!(operators, term, bs, hilbert; half=half, table=table), bonds.bonds)
    return operators
end
@inline termfactor(id::ID{AbstractOID}, ::Val) = ishermitian(id) ? 2 : 1

"""
    expand(term::Term, bond::AbstractBond, hilbert::Hilbert; half::Bool=false, table::Union{Nothing, Table}=nothing) -> Operators
    expand(term::Term, bonds::Bonds, hilbert::Hilbert; half::Bool=false, table::Union{Nothing, Table}=nothing) -> Operators

Expand the operators of a term on a bond/set-of-bonds with a given Hilbert space.
"""
@inline function expand(term::Term, bond::AbstractBond, hilbert::Hilbert; half::Bool=false, table::Union{Nothing, Table}=nothing)
    M = optype(term|>typeof, hilbert|>typeof, bond|>typeof)
    expand!(Operators{M}(), term, bond, hilbert; half=half, table=table)
end
@inline function expand(term::Term, bonds::Bonds, hilbert::Hilbert; half::Bool=false, table::Union{Nothing, Table}=nothing)
    M = optype(term|>typeof, hilbert|>typeof, bonds|>eltype)
    expand!(Operators{M}(), term, bonds, hilbert; half=half, table=table)
end

end #module


text:
# This file was generated, do not modify it. # hide
cm = machine(CompositeModel(), X, y)
res = evaluate!(cm, resampling=Holdout(fraction_train=0.8),
                measure=rms)
round(res.measurement[1], sigdigits=3)

ext: jl


text:
using Test
@testset "Tests for triangulation" begin
# import Triangle
@testset "Test function is_left_of" begin
    points = [
         0  0
         1  0
         1  1
         1 -1
        -1  1
        -1 -1
    ]
    @test is_left_of(points, 1, 2, 3) == true
    @test is_left_of(points, 1, 2, 4) == false
    @test is_left_of(points, 1, 2, 5) == true
    @test is_left_of(points, 1, 2, 6) == false
end
@testset "Test function is_right_of" begin
    points = [
         0  0
         1  0
         1  1
         1 -1
        -1  1
        -1 -1
    ]
    @test is_right_of(points, 1, 2, 3) == false
    @test is_right_of(points, 1, 2, 4) == true
    @test is_right_of(points, 1, 2, 5) == false
    @test is_right_of(points, 1, 2, 6) == true
end

@testset "Test of Triangle constructors" begin
    points = [
         0  0
         1  0
         1  1
    ]
    @testset "For two vertices" begin
        t = Triangle(1, 2)
        @test t.v3 == 0
        @test t.t1 == t.t2 == t.t3
        next = t.t3
        @test next.v3 == 0
        @test next.v1 == t.v2
        @test next.v2 == t.v1
        @test next.t1 == next.t2 == next.t3 == t
    end
    @testset "For three vertices" begin
        t = Triangle(points, 1, 2, 3)
        # v1 -> v2 -> v3 should be clockwise
        @test t.v1 == 1
        @test t.v2 == 3
        @test t.v3 == 2
        @test is_left_of(points, t.v1, t.v3, t.v2)
        @test t.v1 == 1

        @test t.t1.t3 == t.t2.t3 == t.t3.t3 == t
        @test t.t1.t1 == t.t3.t2 == t.t2
        @test t.t1.t2 == t.t2.t1 == t.t3
        @test t.t2.t2 == t.t3.t1 == t.t1

        @test t.t2.v2 == t.t3.v1 == t.v1
        @test t.t1.v1 == t.t3.v2 == t.v2
        @test t.t1.v2 == t.t2.v1 == t.v3
    end
end
@testset "Test is_hull" begin
    points = [
         0  0
         1  0
         1  1
    ]
    t = Triangle(points, 1, 2, 3)
    @test is_hull(t.t1)
    @test !is_hull(t)
end
@testset "Test hull navigation function" begin
    points = [
         0  0
         1  0
         1  1
    ]
    t = Triangle(points, 1, 2, 3)
    @test next(t.t1) == t.t2
    @test prev(t.t1) == t.t3
end
@testset "Test find_base_tangente" begin
    points = [
         0  0
         1  0
         1  1
         2  0
         3  0
         3  1
    ]
    t1 = Triangle(points, 1, 2, 3)
    t2 = Triangle(points, 4, 5, 6)
    t3 = Triangle(2, 3)
    t4 = Triangle(4, 5)

    t_left, t_right = find_base_tangente(points, t1.t1, t2.t2)
    @test (t_left.v1, t_right.v2) == (1, 4)
end
@testset "Test _is_above_base" begin
    @testset "Test _is_above_base_left" begin
        @testset "For two edges triangles" begin
            points = [
                0 0
                0 1
                1 0
            ]
            t = Triangle(1, 2)
            @test _is_above_base_left(t, points, 1, 3) == true
            points = [
                0 -1
                0 0
                1 0
            ]
            t = Triangle(2, 1)
            @test _is_above_base_left(t, points, 2, 3) == false
        end
        @testset "For classic triangles" begin
            points = [
                0 0
                1 2
                2 1
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 1, 4) == true
            points = [
                0 0
                1 -1
                1 -2
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 1, 4) == false
            points = [
                0 1
                1 0
                2 0
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 3, 4) == true
            points = [
                0 -1
                1 0
                2 0
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 3, 4) == false
            points = [
                0 1
                1 1
                2 0
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 3, 4) == true
            points = [
                0 -1
                1 1
                2 0
                3 0
            ]
            t = Triangle(points, 1, 2, 3)
            @test _is_above_base_left(t, points, 3, 4) == false
        end
    end
    @testset "Test _is_above_base_right" begin
        @testset "For two edges triangles" begin
            points = [
                0 0
                1 0
                1 1
            ]
            t = Triangle(2, 3)
            @test _is_above_base_right(t, points, 1, 2) == true
            points = [
                0 0
                1 0
                1 -1
            ]
            t = Triangle(3, 2)
            @test _is_above_base_right(t, points, 1, 2) == false
        end
        @testset "For classic triangles" begin
            points = [
                0 0
                1 3
                2 1
                3 0
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 4) == true
            points = [
                0 0
                1 -3
                2 -2
                3 0
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 4) == false
            points = [
                0 0
                1 1
                2 0
                3 1
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 3) == true
            points = [
                0 0
                1 -1
                2 0
                3 -1
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 3) == false
            points = [
                0 0
                1 0
                2 1
                3 1
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 2) == true
            points = [
                0 0
                1 0
                2 -1
                3 -1
            ]
            t = Triangle(points, 2, 3, 4)
            @test _is_above_base_right(t, points, 1, 2) == false
        end
    end
end
@testset "Test _next" begin
    @testset "For two edges triangles" begin
        t = Triangle(1, 2)
        @test _next_left(t, 1) == t.t3
        @test _next_left(_next_left(t, 1), 1) == t
        @test _next_right(t, 1) == t.t3
        @test _next_right(_next_right(t, 1), 1) == t
    end
    @testset "For classic triangles" begin
        points = [
            0 0
            1 0
            1 1
        ]
        t = Triangle(points, 1, 2, 3)
        @test _next_left(t, 1) == t.t3
        @test _next_left(t, 3) == t.t1
        @test _next_left(t, 2) == t.t2
        @test _next_right(t, 1) == t.t2
        @test _next_right(t, 3) == t.t3
        @test _next_right(t, 2) == t.t1
    end
end
end

max_stars_repo_name: Klafyvel/Land-of-Julia

max_issues_repo_name: Klafyvel/Land-of-Julia

max_forks_repo_name: Klafyvel/Land-of-Julia


text:
using PyCall
@pyimport matplotlib.pyplot as plt

using Gen

#import Gen: Distribution, logpdf

# Example from Section 4 of Reversible jump Markov chain Monte Carlo
# computation and Bayesian model determination

########################
# custom distributions #
########################

# minimum of k draws from uniform_continuous(lower, upper)

# we can sequentially sample the order statistics of a collection of K uniform
# continuous samples on the interval [a, b], by:
# x1 ~ min_uniform_continuous(a, b, K)
# x2 | x1 ~ min_uniform_continuous(x1, b, K-1)
# ..
# xK | x1 .. x_{K-1} ~ min_uniform_continuous(x_{K-1}, b, 1)

struct MinUniformContinuous <: Distribution{Float64} end
const min_uniform_continuous = MinUniformContinuous()

function Gen.logpdf(::MinUniformContinuous, x::Float64, lower::T, upper::U, k::Int) where {T<:Real,U<:Real}
    if x > lower && x < upper
        (k-1) * log(upper - x) + log(k) - k * log(upper - lower)
    else
        -Inf
    end
end

function Gen.random(::MinUniformContinuous, lower::T, upper::U, k::Int) where {T<:Real,U<:Real}
    # inverse CDF method
    p = rand()
    upper - (upper - lower) * (1. - p)^(1. / k)
end

# piecewise homogenous Poisson process

# n intervals - n + 1 bounds
# (b_1, b_2]
# (b_2, b_3]
# ..
# (b_n, b_{n+1}]

function compute_total(bounds, rates)
    num_intervals = length(rates)
    if length(bounds) != num_intervals + 1
        error("Number of bounds does not match number of rates")
    end
    total = 0.
    bounds_ascending = true
    for i=1:num_intervals
        lower = bounds[i]
        upper = bounds[i+1]
        rate = rates[i]
        len = upper - lower
        if len <= 0
            bounds_ascending = false
        end
        total += len * rate
    end
    (total, bounds_ascending)
end

struct PiecewiseHomogenousPoissonProcess <: Distribution{Vector{Float64}} end
const piecewise_poisson_process = PiecewiseHomogenousPoissonProcess()

function Gen.logpdf(::PiecewiseHomogenousPoissonProcess, x::Vector{Float64}, bounds::Vector{Float64}, rates::Vector{Float64})
    cur = 1
    upper = bounds[cur+1]
    lpdf = 0.
    for xi in sort(x)
        if xi < bounds[1] || xi > bounds[end]
            error("x ($xi) lies outside of interval")
        end
        while xi > upper
            cur += 1
            upper = bounds[cur+1]
        end
        lpdf += log(rates[cur])
    end
    (total, bounds_ascending) = compute_total(bounds, rates)
    if bounds_ascending
        lpdf - total
    else
        -Inf
    end
end

function Gen.random(::PiecewiseHomogenousPoissonProcess, bounds::Vector{Float64}, rates::Vector{Float64})
    x = Vector{Float64}()
    num_intervals = length(rates)
    for i=1:num_intervals
        lower = bounds[i]
        upper = bounds[i+1]
        rate = (upper - lower) * rates[i]
        n = random(poisson, rate)
        for j=1:n
            push!(x, random(uniform_continuous, lower, upper))
        end
    end
    x
end

#########
# model #
#########

@gen function model(T::Float64)

    # prior on number of change points
    k = @addr(poisson(3.), "k")

    # prior on the location of (sorted) change points
    change_pts = Vector{Float64}(undef, k)
    lower = 0.
    for i=1:k
        cp = @addr(min_uniform_continuous(lower, T, k-i+1), "cp$i")
        change_pts[i] = cp
        lower = cp
    end

    # k + 1 rate values
    # h$i is the rate for cp$(i-1) to cp$i where cp0 := 0 and where cp$(k+1) := T
    alpha = 1.
    beta = 200.
    rates = Float64[@addr(Gen.gamma(alpha, 1. / beta), "h$i") for i=1:k+1]

    # poisson process
    bounds = vcat([0.], change_pts, [T])
    @addr(piecewise_poisson_process(bounds, rates), "points")
end

function render(trace; ymax=0.02)
    T = get_call_record(trace).args[1]
    assignment = get_assignment(trace)
    k = assignment["k"]
    bounds = vcat([0.], sort([assignment["cp$i"] for i=1:k]), [T])
    rates = [assignment["h$i"] for i=1:k+1]
    for i=1:length(rates)
        lower = bounds[i]
        upper = bounds[i+1]
        rate = rates[i]
        plt.plot([lower, upper], [rate, rate], color="black", linewidth=2)
    end
    points = assignment["points"]
    plt.scatter(points, -rand(length(points)) * (ymax/5.), color="black", s=5)
    ax = plt.gca()
    xlim = [0., T]
    plt.plot(xlim, [0., 0.], "--")
    ax[:set_xlim](xlim)
    ax[:set_ylim](-ymax/5., ymax)
end

function show_prior_samples()
    plt.figure(figsize=(16,16))
    T = 40000.
    for i=1:16
        plt.subplot(4, 4, i)
        trace = simulate(model, (T,))
        render(trace; ymax=0.015)
    end
    plt.tight_layout(pad=0)
    plt.savefig("prior_samples.pdf")
end

#############################
# height and position moves #
#############################

@gen function height_proposal(prev, i::Int)
    prev_assignment = get_assignment(prev)
    height = prev_assignment["h$i"]
    @addr(uniform_continuous(height/2., height*2.), "h$i")
end

@gen function position_proposal(prev, i::Int)
    prev_assignment = get_assignment(prev)
    k = prev_assignment["k"]
    lower = (i == 1) ? 0. : prev_assignment["cp$(i-1)"]
    upper = (i == k) ? T : prev_assignment["cp$(i+1)"]
    @addr(uniform_continuous(lower, upper), "cp$i")
end

function height_move(trace)
    k = get_assignment(trace)["k"]
    i = random(uniform_discrete, 1, k+1)
    mh(model, height_proposal, (i,), trace)
end

function position_move(trace)
    k = get_assignment(trace)["k"]
    i = random(uniform_discrete, 1, k)
    mh(model, position_proposal, (i,), trace)
end

######################
# birth / death move #
######################

# insert a new change point at i, where 1 <= i <= k+1
# the current change point at i, and all after, will be shifted right
# the new change point will be placed between cp$(i-1) and the current cp$i
@gen function birth_proposal(prev, T, i::Int)
    prev_assignment = get_assignment(prev)
    k = prev_assignment["k"]
    lower = (i == 1) ? 0. : prev_assignment["cp$(i-1)"]
    upper = (i == k+1) ? T : prev_assignment["cp$i"]
    @addr(uniform_continuous(lower, upper), "new-cp")
    @addr(uniform_continuous(0., 1.), "u")
end

@gen function death_proposal(prev) end

function birth_move_new_heights(cur_height, new_cp, prev_cp, next_cp, u)
    d_prev = new_cp - prev_cp
    d_next = next_cp - new_cp
    @assert d_prev > 0
    @assert d_next > 0
    d_total = d_prev + d_next
    log_cur_height = log(cur_height)
    log_ratio = log(1 - u) - log(u)
    new_h_prev = exp(log_cur_height - (d_prev / d_total) * log_ratio)
    new_h_next = exp(log_cur_height + (d_next / d_total) * log_ratio)
    @assert new_h_prev > 0.
    @assert new_h_next > 0.
    (new_h_prev, new_h_next)
end

const MODEL = :model
const PROPOSAL = :proposal

@inj function birth_injection(T, i::Int)

    # increment k
    k = @read(MODEL => "k")
    @write(k+1, MODEL => "k")

    # changepoints
    for j=1:i-1
        @copy(MODEL => "cp$j", MODEL => "cp$j")
    end
    @copy(PROPOSAL => "new-cp", MODEL => "cp$i")
    for j=i:k
        @copy(MODEL => "cp$j", MODEL => "cp$(j+1)")
    end

    # compute new heights
    cur_height = @read(MODEL => "h$i")
    prev_cp = (i == 1) ? 0. : @read(MODEL => "cp$(i-1)")
    next_cp = (i == k+1) ? T : @read(MODEL => "cp$i")
    new_cp = @read(PROPOSAL => "new-cp")
    u = @read(PROPOSAL => "u")
    (new_h_prev, new_h_next) = birth_move_new_heights(cur_height, new_cp, prev_cp, next_cp, u)

    # heights
    for j=1:i-1
        @copy(MODEL => "h$j", MODEL => "h$j")
    end
    @write(new_h_prev, MODEL => "h$i")
    @write(new_h_next, MODEL => "h$(i+1)")
    for j=i+1:k+1
        @copy(MODEL => "h$j", MODEL => "h$(j+1)")
    end

    @copy(MODEL => "points", MODEL => "points")
end

function death_move_u_new_height(prev_height, next_height, cur_cp, prev_cp, next_cp)
    d_prev = cur_cp - prev_cp
    d_next = next_cp - cur_cp
    @assert d_prev > 0
    @assert d_next > 0
    d_total = d_prev + d_next
    log_prev_height = log(prev_height)
    log_next_height = log(next_height)
    new_height = exp((d_prev / d_total) * log_prev_height + (d_next / d_total) * log_next_height)
    u = prev_height / (prev_height + next_height)
    @assert new_height > 0.
    (new_height, u)
end

@inj function death_injection(T, i::Int)

    # decrement k
    k = @read(MODEL => "k")
    @assert k > 0
    @write(k-1, MODEL => "k")

    # change points
    for j=1:i-1
        @copy(MODEL => "cp$j", MODEL => "cp$j")
    end
    for j=i+1:k
        @copy(MODEL => "cp$j", MODEL => "cp$(j-1)")
    end

    # compute new height
    cur_cp = @read(MODEL => "cp$i")
    prev_cp = (i == 1) ? 0. : @read(MODEL => "cp$(i-1)")
    next_cp = (i == k) ? T : @read(MODEL => "cp$(i+1)")
    prev_height = @read(MODEL => "h$i")
    next_height = @read(MODEL => "h$(i+1)")
    (new_height, u) = death_move_u_new_height(prev_height, next_height, cur_cp, prev_cp, next_cp)

    # heights
    for j=1:i-1
        @copy(MODEL => "h$j", MODEL => "h$j")
    end
    @write(new_height, MODEL => "h$i")
    for j=i+2:k+1
        @copy(MODEL => "h$j", MODEL => "h$(j-1)")
    end

    @copy(MODEL => "points", MODEL => "points")

    @copy(MODEL => "cp$i", PROPOSAL => "new-cp")
    @write(u, PROPOSAL => "u")
end

function birth_move(trace)
    # if k > 0, then prob_b = 0.25 and prob_d = 0.25
    # the probability that the one we introduce will get deleted, given that
    # we choose a death move, is: 1/(k+1); which is also the probability that
    # we choose the i that we choose here
    k = get_assignment(trace)["k"]
    i = random(uniform_discrete, 1, k+1)
    # prob_b = 1, but prob_d = 0.25, so we correct by -log(4)
    correction = (new_trace) -> (k == 0 ? -log(4) : 0)
    T = get_call_record(trace).args[1]
    rjmcmc(model,
        birth_proposal, (T, i),
        death_proposal, (),
        birth_injection, (T, i),
        trace, correction)
end

function death_move(trace)
    k = get_assignment(trace)["k"]
    @assert k > 0
    i = random(uniform_discrete, 1, k)
    correction = (new_trace) -> (k == 0 ? log(4) : 0)
    T = get_call_record(trace).args[1]
    rjmcmc(model,
        death_proposal, (),
        birth_proposal, (T, i),
        death_injection, (T, i),
        trace, correction)
end

##########################
# Generic MCMC inference #
##########################

function resimulation_mh(selection, trace)
    model_args = get_call_record(trace).args
    (new_trace, weight) = regenerate(model, model_args, NoChange(), trace, selection)
    if log(rand()) < weight
        # accept
        return new_trace
    else
        # reject
        return trace
    end
end

k_selection = DynamicAddressSet()
push_leaf_node!(k_selection, "k")

function generic_mcmc_step(trace)
    k = get_assignment(trace)["k"]
    if k > 0
        prob_h = 1./3
        prob_p = 1./3
        prob_change_k = 1./3
    else
        prob_h = 0.
        prob_p = 0.
        prob_change_k = 1
    end
    move_type = random(categorical, [prob_h, prob_p, prob_change_k])
    if move_type == 1
        height_move(trace)
    elseif move_type == 2
        position_move(trace)
    else
        resimulation_mh(k_selection, trace)
    end
end

#########################
# RJMCMC MCMC inference #
#########################0

function mcmc_step(trace)
    k = get_assignment(trace)["k"]
    if k > 0
        prob_h = 0.25
        prob_p = 0.25
        prob_b = 0.25
        prob_d = 0.25
    else
        prob_h = 0.
        prob_p = 0.
        prob_b = 1.
        prob_d = 0.
    end
    move_type = random(categorical, [prob_h, prob_p, prob_b, prob_d])
    if move_type == 1
        height_move(trace)
    elseif move_type == 2
        position_move(trace)
    elseif move_type == 3
        birth_move(trace)
    elseif move_type == 4
        death_move(trace)
    else
        error("Unknown move type $move_type")
    end
end

function do_mcmc(T, num_steps::Int)
    (trace, _) = generate(model, (T,), observations)
    for iter=1:num_steps
        if iter % 1000 == 0
            println("iter $iter of $num_steps, k: $(get_assignment(trace)["k"])")
        end
        #trace = mcmc_step(trace)
        trace = generic_mcmc_step(trace)
    end
    trace
end

########################
# inference experiment #
########################

Gen.load_generated_functions()

import Random
Random.seed!(1)

# load data set
import CSV
function load_data_set()
    df = CSV.read("coal.csv")
    dates = df[1]
    dates = dates .- minimum(dates)
    dates * 365.25 # convert years to days
end

const points = load_data_set()
const T = maximum(points)
const observations = DynamicAssignment()
observations["points"] = points

function show_posterior_samples()
    plt.figure(figsize=(16,16))
    for i=1:16
        println("replicate $i")
        tic()
        plt.subplot(4, 4, i)
        trace = do_mcmc(T, 5000)#10000)
        toc()
        render(trace; ymax=0.015)
    end
    plt.tight_layout(pad=0)
    plt.savefig("posterior_samples.pdf")
end

function get_rate_vector(trace, test_points)
    assignment = get_assignment(trace)
    k = assignment["k"]
    cps = [assignment["cp$i"] for i=1:k]
    hs = [assignment["h$i"] for i=1:k+1]
    rate = Vector{Float64}()
    cur_h_idx = 1
    cur_h = hs[cur_h_idx]
    next_cp_idx = 1
    upper = (next_cp_idx == k + 1) ? T : cps[next_cp_idx]
    for x in test_points
        while x > upper
            next_cp_idx += 1
            upper = (next_cp_idx == k + 1) ? T : cps[next_cp_idx]
            cur_h_idx += 1
            cur_h = hs[cur_h_idx]
        end
        push!(rate, cur_h)
    end
    rate
end

# compute posterior mean rate curve

function plot_posterior_mean_rate()
    test_points = collect(1.0:10.0:T)
    rates = Vector{Vector{Float64}}()
    num_samples = 0
    num_steps = 5000 # 20000
    for reps=1:10
        (trace, _) = generate(model, (T,), observations)
        for iter=1:num_steps
            if iter % 1000 == 0
                println("iter $iter of $num_steps, k: $(get_assignment(trace)["k"])")
            end
            trace = mcmc_step(trace)
            if iter > 4000
                num_samples += 1
                rate_vector = get_rate_vector(trace, test_points)
                @assert length(rate_vector) == length(test_points)
                push!(rates, rate_vector)
            end
        end
    end
    posterior_mean_rate = zeros(length(test_points))
    for rate in rates
        posterior_mean_rate += rate / Float64(num_samples)
    end
    ymax = 0.010
    plt.figure()
    plt.plot(test_points, posterior_mean_rate, color="black")
    plt.scatter(points, -rand(length(points)) * (ymax/6.), color="black", s=5)
    ax = plt.gca()
    xlim = [0., T]
    plt.plot(xlim, [0., 0.], "--")
    ax[:set_xlim](xlim)
    ax[:set_ylim](-ymax/5., ymax)
    plt.savefig("posterior_mean_rate.pdf")
end

function plot_trace_plot()
    # show the number of clusters
    (trace, _) = generate(model, (T,), observations)
    num_clusters_vec = Int[]
    burn_in = 20000
    for iter=1:burn_in + 5000
        (trace, accept) = mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, get_assignment(trace)["k"])
        end
    end
    plt.figure()
    plt.plot(num_clusters_vec)
    ax = plt.gca()
    plt.savefig("trace_plot_rjmcmc.pdf")
end

function plot_trace_plot()
    plt.figure(figsize=(8, 4))

    # generic
    (trace, _) = generate(model, (T,), observations)
    num_clusters_vec = Int[]
    burn_in = 20000
    for iter=1:burn_in + 5000
        trace = generic_mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, get_assignment(trace)["k"])
        end
    end
    plt.subplot(2, 1, 1)
    plt.plot(num_clusters_vec, "r")

    # reversible jump
    (trace, _) = generate(model, (T,), observations)
    height1 = Float64[]
    num_clusters_vec = Int[]
    burn_in = 20000
    for iter=1:burn_in + 5000
        trace = mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, get_assignment(trace)["k"])
        end
    end
    plt.subplot(2, 1, 2)
    plt.plot(num_clusters_vec, "b")

    ax = plt.gca()
    plt.savefig("trace_plot.pdf")
end

println("showing prior samples...")
show_prior_samples()

println("showing posterior samples...")
show_posterior_samples()

println("estimating posterior mean rate...")
plot_posterior_mean_rate()

println("making trace plot...")
plot_trace_plot()

max_stars_repo_name: BoltzmannBrain/Gen

max_issues_repo_name: BoltzmannBrain/Gen

max_forks_repo_name: BoltzmannBrain/Gen


text:
println(Int('a'))
println(Char(97))

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
import Base.==
## custom json formatting for types
struct B
    id::Int
    name::String
    B(id::Int) = new(id)
    B(id::Int, name::String) = new(id, name)
end
B(id::Int, ::Nothing) = B(id)

b1 = B(1, "harry")
b2 = B(2, "hermione")
b3 = B(3, "ron")

partialB = B(4)

@test JSON2.write(b1) == "{\"id\":1,\"name\":\"harry\"}"
@test JSON2.read(JSON2.write(b1), B) == b1
@test JSON2.write(b2) == "{\"id\":2,\"name\":\"hermione\"}"
@test JSON2.read(JSON2.write(b2), B) == b2
@test JSON2.write(b3) == "{\"id\":3,\"name\":\"ron\"}"
@test JSON2.read(JSON2.write(b3), B) == b3

struct B1
    id::Int
    name::String
    B1(id::Int) = new(id)
    B1(id::Int, nm::String) = new(id, nm)
end

JSON2.@format B1 begin
    name => (omitempty=true,)
end

@test JSON2.write(B1(4)) == "{\"id\":4}"
@test JSON2.read(JSON2.write(B1(4)), B1) == B1(4)

struct B2
    id::Int
    name::String
    B2(id::Int) = new(id)
    B2(id::String, ::Nothing) = new(parse(Int, id))
end

JSON2.@format B2 begin
    id => (jsontype=String,)
    name => (omitempty=true,)
end

@test JSON2.write(B2(4)) == "{\"id\":\"4\"}"
@test JSON2.read(JSON2.write(B2(4)), B2) == B2(4)

struct B3
    id::Int
    name::Union{Nothing, String}
    B3(id::Int) = new(id)
    B3(::Nothing) = new()
end

JSON2.@format B3 begin
    id => (exclude=true,)
end

@test JSON2.write(B3(4)) == "{\"name\":null}"
bb3 = JSON2.read(JSON2.write(B3(4)), B3)
@test !isdefined(bb3, :name)

struct B4
    id::Int
    name::String
    B4(id::Int) = new(id)
    B4() = new()
end

JSON2.@format B4 begin
    id => (exclude=true,)
    name => (omitempty=true, jsontype=String)
end

@test JSON2.write(B4(4)) == "{}"
bb4 = JSON2.read(JSON2.write(B4(4)), B4)
@test !isdefined(bb4, :name)

struct B5
    id::Int
    name::String
    B5(; id=0, name="", kwargs...) = new(id, name)
end

JSON2.@format B5 keywordargs

@test JSON2.read("{\"id\":1}", B5) == B5(; id=1)
@test JSON2.read("{}", B5) == B5()
@test JSON2.read("{\"id\":1,\"name\":\"hey\"}", B5) == B5(; id=1, name="hey")
@test JSON2.read("{\"id\":1,\"name2\":\"hey\"}", B5) == B5(; id=1)

const jsons = [
    """{"a": 1, "b": 2, "c": 3, "d": 4}""",
    """{"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}""",
    """{"a": 1, "d": 4}""",
    """{"a": 1, "e": 5, "d": 4}""",
    """{"c": 3, "a": 1, "d": 4, "b": 2}""",
    """{"c": 3, "a": 1, "e": 5, "d": 4, "b": 2}""",
    """{"d": 4, "a": 1}""",
    """{"d": 4, "e": 5, "a": 1}""",
]

# default read
struct AJ1
    a::Int
    b::Int
    c::Int
    d::Int
end

@test JSON2.read(jsons[1], AJ1) == AJ1(1, 2, 3, 4)
@test JSON2.read(jsons[2], AJ1) == AJ1(1, 2, 3, 4)

# providing field defaults
struct AJ2
    a::Int
    b::Int
    c::Int
    d::Int
end

JSON2.@format AJ2 begin
    b => (default=2,)
    c => (default=3,)
end

@test JSON2.read(jsons[3], AJ2) == AJ2(1, 2, 3, 4)
@test JSON2.read(jsons[4], AJ2) == AJ2(1, 2, 3, 4)

# noargs
mutable struct AJ3
    a::Int
    b::Int
    c::Int
    d::Int
    AJ3() = new(0, 0, 0, 0)
end
==(a::AJ3, b::AJ3) = a.a == b.a && a.b == b.b && a.c == b.c && a.d == b.d
JSON2.@format AJ3 noargs begin
    b => (default=2,)
    c => (default=3,)
end

const aj3 = AJ3()
aj3.a = 1
aj3.b = 2
aj3.c = 3
aj3.d = 4
@test JSON2.read(jsons[5], AJ3) == aj3
@test JSON2.read(jsons[6], AJ3) == aj3
@test JSON2.read(jsons[7], AJ3) == aj3
@test JSON2.read(jsons[8], AJ3) == aj3

# keywordargs
struct AJ4
    a::Int
    b::Int
    c::Int
    d::Int
    AJ4(; a=1, b=0, c=0, d=4, kwargs...) = new(a, b, c, d)
end
==(a::AJ4, b::AJ4) = a.a == b.a && a.b == b.b && a.c == b.c && a.d == b.d
JSON2.@format AJ4 keywordargs begin
    b => (default=2,)
    c => (default=3,)
end

const aj4 = AJ4(b=2, c=3)
@test JSON2.read(jsons[5], AJ4) == aj4
@test JSON2.read(jsons[6], AJ4) == aj4
@test JSON2.read(jsons[7], AJ4) == aj4
@test JSON2.read(jsons[8], AJ4) == aj4


text:
scores = [44, 42, 42, 41, 41, 41, 39]
names = ["Solomon", "Jason", "Errol", "Garry",
         "Bernard", "Barry", "Stephen"]

srank = rankstandard(scores)
mrank = rankmodified(scores)
drank = rankdense(scores)
orank = rankordinal(scores)
frank = rankfractional(scores)

println("    Name    Score  Std  Mod  Den  Ord  Frac")
for i in 1:length(scores)
    print(@sprintf("   %-7s", names[i]))
    print(@sprintf("%5d ", scores[i]))
    print(@sprintf("%5d", srank[i]))
    print(@sprintf("%5d", mrank[i]))
    print(@sprintf("%5d", drank[i]))
    print(@sprintf("%5d", orank[i]))
    print(@sprintf("%7.2f", frank[i]))
    println()
end


text:
"""
	setNorm!{T}(simulator::Quantum2dFractalSimulator{T}, norm::T, which::Int)
	setNorm!{T}(simulator::Quantum2dFractalSimulator{T}, norm::T, normname::AbstractString)
set the given normalization factor to the normalizationFactor Array
to the current iteration and given index information cell.

ref: `simulator_fractal.jl`
"""
function setNorm!{T}(simulator::Quantum2dFractalSimulator{T}, norm::T, which::Int)
	iteration = getCount(simulator)
	setNorm!(simulator, norm, which, iteration)
end

function setNorm!{T}(simulator::Quantum2dFractalSimulator{T}, norm::T, normname::AbstractString)
	setNorm!(simulator, norm, getIndexOf(simulator, normname))
end


text:
@doc raw"""
Holds the tableau of a stochastic implicit Runge-Kutta method.

`qdrift` holds the RK coefficients for the drift part,
`qdiff` holds the RK coefficients for the diffusion part of the SDE.

Order of the tableau is not included, because unlike in the deterministic
setting, it depends on the properties of the noise (e.g., the dimension of
the Wiener process and the commutativity properties of the diffusion matrix).

Orders stored in `qdrift` and `qdiff` are understood as the classical orders of these methods.
"""
struct TableauSIRK{T} <: AbstractTableauIRK{T}
    name::Symbol
    s::Int
    qdrift::CoefficientsRK{T}
    qdiff::CoefficientsRK{T}

    function TableauSIRK{T}(name, s, qdrift, qdiff) where {T}
        @assert s == qdrift.s == qdiff.s
        new(name, s, qdrift, qdiff)
    end
end

function TableauSIRK(name::Symbol, qdrift::CoefficientsRK{T}, qdiff::CoefficientsRK{T}) where {T}
    TableauSIRK{T}(name, qdrift.s, qdrift, qdiff)
end

function TableauSIRK(name::Symbol, order_drift::Int, a_drift::Matrix{T}, b_drift::Vector{T}, c_drift::Vector{T}, order_diff::Int, a_diff::Matrix{T}, b_diff::Vector{T}, c_diff::Vector{T}) where {T}
    TableauSIRK{T}(name, length(c_drift), CoefficientsRK(name, order_drift, a_drift, b_drift, c_drift), CoefficientsRK(name, order_diff, a_diff, b_diff, c_diff))
end

# TODO function readTableauSIRKFromFile(dir::AbstractString, name::AbstractString)

@doc raw"""
Parameters for right-hand side function of implicit Runge-Kutta methods.

`A` - if positive, the upper bound of the Wiener process increments; if `A=0.0`, no truncation
"""
mutable struct ParametersSIRK{DT, TT, D, M, S, ET <: NamedTuple} <: Parameters{DT,TT}
    equ::ET
    tab::TableauSIRK{TT}
    Δt::TT

    ΔW::Vector{DT}
    ΔZ::Vector{DT}
    A::DT

    t::TT
    q::Vector{DT}

    function ParametersSIRK{DT,D,M}(equ::ET, tab::TableauSIRK{TT}, Δt::TT, ΔW::Vector{DT}, ΔZ::Vector{DT}, A::DT) where {DT, TT, D, M, ET <: NamedTuple}
        @assert M == length(ΔW) == length(ΔZ)
        new{DT, TT, D, M, tab.s, ET}(equ, tab, Δt, ΔW, ΔZ, A, zero(TT), zeros(DT,D))
    end
end

@doc raw"""
Structure for holding the internal stages `Q`, the values of the drift vector
and the diffusion matrix evaluated at the internal stages `V=v(Q)`, `B=B(Q)`,
and the increments `Y = a_drift*v(Q)*Δt + a_diff*B(Q)*ΔW`.
"""
struct IntegratorCacheSIRK{DT,D,M,S} <: SDEIntegratorCache{DT,D,M}
    Q::Vector{Vector{DT}}
    V::Vector{Vector{DT}}
    B::Vector{Matrix{DT}}
    Y::Vector{Vector{DT}}

    v::Vector{DT}
    b::Matrix{DT}
    y::Vector{DT}

    ΔQ::Vector{DT}
    Y1::Vector{DT}
    Y2::Vector{DT}
    V1::Vector{DT}
    V2::Vector{DT}
    B1::Matrix{DT}
    B2::Matrix{DT}
    Δw::Vector{DT}
    Δy::Vector{DT}

    function IntegratorCacheSIRK{DT,D,M,S}() where {DT,D,M,S}
        # create internal stage vectors
        Q = create_internal_stage_vector(DT, D, S)
        V = create_internal_stage_vector(DT, D, S)
        B = create_internal_stage_matrix(DT, D, M, S)
        Y = create_internal_stage_vector(DT, D, S)

        # create velocity and update vector
        v = zeros(DT,D)
        b = zeros(DT,D,M)
        y = zeros(DT,D)

        # create temporary vectors
        ΔQ = zeros(DT,D)
        Y1 = zeros(DT,D)
        Y2 = zeros(DT,D)
        V1 = zeros(DT,D)
        V2 = zeros(DT,D)
        B1 = zeros(DT,D,M)
        B2 = zeros(DT,D,M)
        Δw = zeros(DT,M)
        Δy = zeros(DT,M)

        new(Q, V, B, Y, v, b, y, ΔQ, Y1, Y2, V1, V2, B1, B2, Δw, Δy)
    end
end

function Integrators.IntegratorCache{ST}(params::ParametersSIRK{DT,TT,D,M,S}; kwargs...) where {ST,DT,TT,D,M,S}
    IntegratorCacheSIRK{ST,D,M,S}(; kwargs...)
end

@inline Integrators.CacheType(ST, params::ParametersSIRK{DT,TT,D,M,S}) where {DT,TT,D,M,S} = IntegratorCacheSIRK{ST,D,M,S}

"Stochastic implicit Runge-Kutta integrator."
struct IntegratorSIRK{DT, TT, D, M, S,
                PT <: ParametersSIRK{DT,TT},
                ST <: NonlinearSolver{DT}} <: StochasticIntegratorRK{DT,TT,D,M,S}
    params::PT
    solver::ST
    caches::CacheDict{PT}

    function IntegratorSIRK(params::ParametersSIRK{DT,TT,D,M,S}, solver::ST, caches) where {DT,TT,D,M,S,ST}
        new{DT, TT, D, M, S, typeof(params), ST}(params, solver, caches)
    end

    function IntegratorSIRK{DT,D,M}(equations::NamedTuple, tableau::TableauSIRK{TT}, Δt::TT; K::Int=0) where {DT,TT,D,M}
        # get number of stages
        S = tableau.s

        # K - the integer in the bound A = √(2 K Δt |log Δt|) due to Milstein & Tretyakov; K=0 no truncation
        K==0 ? A = 0.0 : A = sqrt( 2*K*Δt*abs(log(Δt)) )

        # create params
        params = ParametersSIRK{DT,D,M}(equations, tableau, Δt, zeros(DT,M), zeros(DT,M), A)

        # create cache dict
        caches = CacheDict(params)

        # create solver
        solver = create_nonlinear_solver(DT, D*S, params, caches)

        # create integrator
        IntegratorSIRK(params, solver, caches)
    end

    function IntegratorSIRK(equation::SDE{DT,TT}, tableau::TableauSIRK{TT}, Δt::TT; kwargs...) where {DT,TT}
        IntegratorSIRK{DT, ndims(equation), equation.m}(get_function_tuple(equation), tableau, Δt; kwargs...)
    end
end

@doc raw"""
This function computes initial guesses for Y and assigns them to `int.solver.x`.
The prediction is calculated using an explicit integrator.

SIMPLE SOLUTION
The simplest initial guess for `Y` is 0:
`int.solver.x .= zeros(eltype(int), tableau(int).s*ndims(int))`

USING AN EXPLICIT INTEGRATOR TO COMPUTE AN INITIAL GUESS
Below we use the R2 method of Burrage & Burrage to calculate
the internal stages at the times `c[1]...c[s]`.
This approach seems to give very good approximations if the time step
and magnitude of noise are not too large. If the noise intensity is too big,
one may have to perform a few iterations of the explicit method with a smaller
time step, use a higher-order explicit method (e.g. CL or G5), or use
the simple solution above.

When calling this function, int.params should contain the data:
`int.params.q`  - the solution at the previous time step
`int.params.t`  - the time of the previous step
`int.params.ΔW` - the increment of the Brownian motion for the current step
"""
function initial_guess!(int::IntegratorSIRK{DT,TT}, sol::AtomicSolutionSDE{DT,TT},
                        cache::IntegratorCacheSIRK{DT}=int.caches[DT]) where {DT,TT}
    local t2::TT
    local Δt_local::TT

    # Evaluating the functions v and B at t,q - same for all stages
    int.params.equ[:v](int.params.t, int.params.q, cache.V1)
    int.params.equ[:B](int.params.t, int.params.q, cache.B1)

    for i in eachstage(int)
        Δt_local = tableau(int).qdrift.c[i]  * int.params.Δt
        cache.Δw .= tableau(int).qdrift.c[i] .* int.params.ΔW

        mul!(cache.ΔQ, cache.B1, cache.Δw)
        @. cache.Q[i] = int.params.q + 2. / 3. * Δt_local * cache.V1 + 2. / 3. * cache.ΔQ

        t2 = int.params.t + 2. / 3. * Δt_local

        int.params.equ[:v](t2, cache.Q[i], cache.V2)
        int.params.equ[:B](t2, cache.Q[i], cache.B2)

        # Calculating the Y's and assigning them to the array int.solver.x as initial guesses
        mul!(cache.Y1, cache.B1, cache.Δw)
        mul!(cache.Y2, cache.B2, cache.Δw)

        for j in eachdim(int)
            int.solver.x[(i-1)*ndims(int)+j] = Δt_local*(1. / 4. * cache.V1[j] + 3. / 4. * cache.V2[j]) + 1. / 4. * cache.Y1[j] + 3. / 4. * cache.Y2[j]
        end
    end
end

@doc raw"""
Unpacks the data stored in `x = (Y[1][1], Y[1][2], ... Y[1][D], Y[2][1], ...)`
into `Y`, calculates the internal stages `Q`, the values of the RHS
of the SDE ( v(Q) and B(Q) ), and assigns them to `V` and `B`.
Unlike for FIRK, here `Y = a v(Q) Δt + â B(Q) ΔW`
"""
function compute_stages!(x::Vector{ST}, Q::Vector{Vector{ST}}, V::Vector{Vector{ST}},
                         B::Vector{Matrix{ST}}, Y::Vector{Vector{ST}},
                         params::ParametersSIRK{DT,TT,D}) where {ST,DT,TT,D}

    local tᵢ::TT

    @assert length(Y) == length(Q) == length(V) == length(B)

    # copy x to Y and calculate Q
    for i in eachindex(Q,Y)
        @assert size(B[i],1) == length(Q[i]) == length(V[i])
        for k in eachindex(Q[i])
            Y[i][k] = x[D*(i-1)+k]
        end
        Q[i] .= params.q .+ Y[i]
    end

    # compute V = v(Q) and B=B(Q)
    for i in eachindex(Q,V,B)
        tᵢ = params.t + params.Δt * params.tab.qdrift.c[i]
        # calculates v(t,Q[i]) and assigns to the i-th column of V
        params.equ[:v](tᵢ, Q[i], V[i])
        # calculates B(t,Q[i]) and assigns to the matrix B[i]
        params.equ[:B](tᵢ, Q[i], B[i])
    end
end

"Compute stages of stochastic implicit Runge-Kutta methods."
function Integrators.function_stages!(x::Vector{ST}, b::Vector{ST},
                params::ParametersSIRK{DT,TT,D,M,S},
                caches::CacheDict) where {ST,DT,TT,D,M,S}

    # get cache for internal stages
    cache = caches[ST]

    compute_stages!(x, cache.Q, cache.V, cache.B, cache.Y, params)

    local y1::ST
    local y2::ST
    local y3::ST
    local y4::ST

    # compute b = - (Y-AV)
    for i in 1:S
        for k in 1:D
            y1 = 0
            y2 = 0
            y3 = 0
            y4 = 0
            for j in 1:S
                y1 += params.tab.qdrift.a[i,j] * cache.V[j][k] * params.Δt
                y2 += params.tab.qdrift.â[i,j] * cache.V[j][k] * params.Δt
                for l in 1:M
                    y3 += params.tab.qdiff.a[i,j] * cache.B[j][k,l] * params.ΔW[l]
                    y4 += params.tab.qdiff.â[i,j] * cache.B[j][k,l] * params.ΔW[l]
                end
            end
            b[D*(i-1)+k] = - cache.Y[i][k] + (y1 + y2) + (y3 + y4)
        end
    end
end

"Integrate SDE with a stochastic implicit Runge-Kutta integrator."
function Integrators.integrate_step!(int::IntegratorSIRK{DT,TT}, sol::AtomicSolutionSDE{DT,TT},
                                     cache::IntegratorCacheSIRK{DT}=int.caches[DT]) where {DT,TT}
    # update nonlinear solver parameters from cache
    update_params!(int, sol)

    # compute initial guess and assign to int.solver.x
    initial_guess!(int, sol, cache)

    # reset cache
    reset!(sol, timestep(int))

    # call nonlinear solver
    solve!(int.solver)

    # print solver status
    print_solver_status(int.solver.status, int.solver.params)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.solver.status, int.solver.params)

    # compute the drift vector field and the diffusion matrix at internal stages
    compute_stages!(int.solver.x, cache.Q, cache.V, cache.B, cache.Y, int.params)

    # compute final update
    update_solution!(sol, cache.V, cache.B, tableau(int).qdrift.b, tableau(int).qdiff.b, int.params.Δt, int.params.ΔW, cache.Δy)
    update_solution!(sol, cache.V, cache.B, tableau(int).qdrift.b̂, tableau(int).qdiff.b̂, int.params.Δt, int.params.ΔW, cache.Δy)
end