module SpinPackage

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

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

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

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

"""
    SID <: IID

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

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

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

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

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

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

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

"""
    Spin <: Internal{SID}

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

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

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

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

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

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

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

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

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

"""
    usualspinindextotuple

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

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

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

"""
    soptdefaultlatex

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Spin term.

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

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

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

end # module


abstract type Var end   # variable

struct DiscreteVariable{N,D,B,M} <: Var
    name::N
    domain::D
    booleans::B
    varmap::M  # forward dictionary from domain to variables
end

# Base.show(io::IO, x::DiscreteVariable) = print(io, "Variable $(x.name) ∈ $(x.domain)")
Base.show(io::IO, x::DiscreteVariable) = print(io, "variable $(x.name)")

DiscreteVariable(name, D) = DiscreteVariable(name, sort(collect(D)))

# index integers as themselves; otherwise just 1..n:
indices(v::Vector{Int}) = v
indices(v) = collect(eachindex(v))

function DiscreteVariable(name, domain::Vector)
    booleans = [Variable(name, i) for i ∈ indices(domain)]
    varmap = Dict(i => v for (i, v) in zip(domain, booleans))

    return DiscreteVariable(name, domain, booleans, varmap)
end

not_both(x, y) = ¬x ∨ ¬y
implies(x, y) = ¬x ∨ y
and(x, y) = ¬(¬x ∨ ¬y)

"Given a set of boolean variables, gives clauses specifying that exactly one of them is true"
function exactly_one(vars)
    clauses = [∨(vars...)]   # one must be true

    # two cannot be true:
    for i in 1:length(vars)
        for j in i+1:length(vars)
            push!(clauses, not_both(vars[i], vars[j]))
        end
    end

    return clauses
end

clauses(x::DiscreteVariable) = exactly_one(x.booleans)

domain(x::Var) = x.domain
# vars(x::DiscreteVariable) = [x]

function decode(varmap, x::Var)
    values = [varmap[v] for v in x.booleans]

    num_true = count(values)

    if num_true ≠ 1
        error("Variable $x has not been successfully solved: values $(x.vars .=> values)")
    end

    return x.domain[findfirst(values)]
end


module CompoundDistributions

using StatsFuns.RFunctions: betarand, gammarand
using LogProbs

# distribution types
export BetaBern, DirMul, DirCat, UniCat, CatDist, ChineseRest, DirCatVar, DirCatVarArith

# conditional distribution type
export SimpleCond

# functions
export support, sample, logscore, add_obs!, rm_obs!

# helper function
export categorical_sample

abstract type Distribution{T} end

function categorical_sample(tokens, weights)
    x = rand() * sum(weights)
    cum_weights = zero(eltype(weights))
    for (t, w) in zip(tokens, weights)
        cum_weights += w
        if cum_weights > x
            return t
        end
    end
end

categorical_sample(d::Dict) = categorical_sample(keys(d), values(d))
categorical_sample(v::Vector) = categorical_sample(1:length(v), v)

############################
### Beta Bernoulli Class ###
############################

mutable struct BetaBern{T, C} <: Distribution{T}
    heads :: T
    tails :: T
    alpha :: C
    beta  :: C
end

BetaBern(a, b) = BetaBern(true, false, a, b)
BetaBern(support) = BetaBern(support..., 1, 1)

support(bb::BetaBern) = [bb.heads, bb.tails]

function sample(bb::BetaBern)
    p = betarand(bb.alpha, bb.beta)
    categorical_sample((bb.heads, bb.tails), (p, 1-p))
end

function logscore(bb::BetaBern, obs)
    k = obs == bb.heads
    LogProb(lbeta(k+bb.alpha, 1-k+bb.beta) - lbeta(bb.alpha, bb.beta))
end

add_obs!(bb::BetaBern, obs) = obs == bb.heads ? bb.alpha += 1 : bb.beta += 1
rm_obs!(bb::BetaBern, obs)  = obs == bb.heads ? bb.alpha -= 1 : bb.beta -= 1

###################################
### Dirichlet Multinomial Class ###
###################################

mutable struct DirMul{T, C} <: Distribution{T}
    counts :: Dict{T, C}
end

DirMul(support) = DirMul(Dict(x => 1.0 for x in support))
support(dm::DirMul) = keys(dm.counts)

function sample(dm::DirMul, n)
    weights = [gammarand(c, 1) for c in values(dm.counts)]
    d = Dict(k=>0 for k in keys(dm.counts))
    for i in 1:n
        d[categorical_sample(keys(dm.counts), weights)] += 1
    end
    d
end

function logscore(dm::DirMul, obs::AbstractDict)
    n = sum(values(obs))
    LogProb(
        log(n) + lbeta(sum(values(dm.counts)), n) -
        sum(
            log(obs[x]) + lbeta(dm.counts[x], obs[x])
            for x in keys(dm.counts) if obs[x] > 0
        )
    )
end

function add_obs!(dm::DirMul, obs::AbstractDict)
    for x in keys(obs)
        dm.counts[x] += obs[x]
    end
end

function rm_obs!(dm::DirMul, obs::AbstractDict)
    for x in keys(obs)
        dm.counts[x] -= obs[x]
    end
end

###################################
### Dirichlet Categorical Class ###
###################################

mutable struct DirCat{T, C} <: Distribution{T}
    counts :: Dict{T, C}
end

DirCat(support) = DirCat(Dict(x => 1.0 for x in support))
support(dc::DirCat) = keys(dc.counts)

function sample(dc::DirCat)
    weights = [gammarand(c, 1) for c in values(dc.counts)]
    categorical_sample(keys(dc.counts), weights)
end

function logscore(dc::DirCat, obs)
    LogProb(lbeta(sum(values(dc.counts)), 1) - lbeta(dc.counts[obs], 1))
end

function add_obs!(dc::DirCat, obs)
    dc.counts[obs] += 1
end

function rm_obs!(dc::DirCat, obs)
    dc.counts[obs] -= 1
end

#########################################################
### Dirichlet Categorical for Variational Bayes Class ###
#########################################################

mutable struct DirCatVar{T, C} <: Distribution{T}
    counts :: Dict{T, C}
end

DirCatVar(support) = DirCatVar(Dict(x => 1.0 for x in support))
support(dc::DirCatVar) = keys(dc.counts)

function sample(dc::DirCatVar)
    weights = [gammarand(c, 1) for c in values(dc.counts)]
    categorical_sample(keys(dc.counts), weights)
end

function logscore(dc::DirCatVar, obs)
    LogProb(digamma(dc.counts[obs]) - digamma(sum(values(dc.counts))))
end

add_obs!(dc::DirCatVar, obs) = dc.counts[obs] += 1
rm_obs!(dc::DirCatVar, obs) = dc.counts[obs] -= 1

###########################################################################
### Dirichlet Categorical for Variational Bayes Class - Arithmetic mean ###
###########################################################################

mutable struct DirCatVarArith{T, C} <: Distribution{T}
    counts :: Dict{T, C}
end

DirCatVarArith(support) = DirCatVarArith(Dict(x => 1.0 for x in support))
support(dc::DirCatVarArith) = keys(dc.counts)

function sample(dc::DirCatVarArith)
    weights = [gammarand(c, 1) for c in values(dc.counts)]
    categorical_sample(keys(dc.counts), weights)
end

function logscore(dc::DirCatVarArith, obs)
    LogProb((dc.counts[obs])/(sum(values(dc.counts))))
end

add_obs!(dc::DirCatVarArith, obs) = dc.counts[obs] += 1
rm_obs!(dc::DirCatVarArith, obs) = dc.counts[obs] -= 1

######################################
### Categorical Distribution Class ###
######################################

mutable struct CatDist{T} <: Distribution{T}
    probs :: Dict{T, LogProb}
end

support(cd::CatDist) = keys(cd.probs)

logscore(cd::CatDist, x) = cd.probs[x]
sample(cd::CatDist) = categorical_sample(cd.probs)
add_obs!(cd::CatDist, obs, context) = nothing
remove_obs!(cd::CatDist, obs, context) = nothing

#################################
### Uniform Categorical Class ###
#################################

struct UniCat{T} <: Distribution{T}
    support :: Vector{T}
end

UniCat(support) = UniCat(collect(support))
support(uc::UniCat) = uc.support

sample(uc::UniCat) = uc.support[rand(1:length(uc.support))]
logscore(uc::UniCat, obs) = obs in uc.support ? LogProb(1 / length(uc.support)) : zero(LogProb)
add_obs!(uc::UniCat, obs, context) = nothing
rm_obs!(uc::UniCat, obs, context) = nothing

################################
### Chinese Restaurant Class ###
################################

mutable struct ChineseRest{Dish, Dist <: Distribution{Dish}} <: Distribution{Dish}
    a :: Float64 # discount parameter
    b :: Float64 # crp parameter
    basedist :: Dist
    tables :: Dict{Dish, Vector{Int}}
    num_tables :: Int
    num_customers :: Int
end

ChineseRest(a, b, basedist::Distribution{T}) where T =
    ChineseRest(a, b, basedist, Dict{T, Vector{Int}}(), 0, 0)
ChineseRest(support) = ChineseRest(0.0, 0.1, UniCat(support))
support(r::ChineseRest) = support(r.basedist)

function sample(r::ChineseRest)
    if flip((r.num_tables * r.a + r.b) / (r.num_customers + r.b))
        sample(r.basedist)
    else
        categorical_sample(
            Dict(dish => sum(r.tables[dish]) for dish in keys(r.tables))
        )
    end
end

function new_table_logscore(r::ChineseRest, dish)
    new_table_prob = LogProb(
        log(r.num_tables * r.a + r.b) - log(r.num_customers + r.b)
    )
    base_dist_prob = logscore(r.basedist, dish)
    new_table_prob * base_dist_prob
end

function logscore(r::ChineseRest, dish)
    if haskey(r.tables, dish)
        numerator = sum(n - r.a for n in r.tables[dish])
        LogProb(log(numerator) - log(r.num_customers + r.b)) + new_table_logscore(r, dish)
    else
        new_table_logscore(r, dish)
    end
end

function add_obs!(r::ChineseRest, dish)
    r.num_customers += 1
    if !haskey(r.tables, dish)
        r.num_tables += 1
        r.tables[dish] = [1]
    else
        r.tables[dish][1] += 1
    end
    nothing
end

function rm_obs!(r::ChineseRest, dish)
    r.num_customers -= 1
    r.tables[dish][1] -= 1
    if r.tables[dish][1] == 0
        r.num_tables -= 1
        delete!(r.tables, dish)
    end
    nothing
end

################################
### Simple Conditional Class ###
################################

mutable struct SimpleCond{C, D, S} # context, distribution, support
    dists   :: Dict{C, D}
    support :: S
end

SimpleCond(dists::AbstractDict) = SimpleCond(
    dists,
    vcat([collect(support(dist)) for dist in values(dists)]...)
)

SimpleCond(g::Base.Generator) = SimpleCond(Dict(g))

sample(sc::SimpleCond, context, args...) = sample(sc.dists[context], args...)
logscore(sc::SimpleCond, obs, context) = logscore(sc.dists[context], obs)
rm_obs!(sc::SimpleCond, obs, context) = rm_obs!(sc.dists[context], obs)

function add_obs!(cond::SimpleCond{C,D,S}, obs, context) where {C,D,S}
    if !haskey(cond.dists, context)
        cond.dists[context] = D(cond.support)
    end
    add_obs!(cond.dists[context], obs)
end

end # module


using Test
using TimeseriesSurrogates
ENV["GKSwstype"] = "100"

N = 500
ts = cumsum(randn(N))
ts_nan = cumsum(randn(N))
ts_nan[1] = NaN
x = cos.(range(0, 20π, length = N)) .+ randn(N)*0.05

@testset "LombScargle" begin
    t = sort((0:N-1) + rand(N))
    tol = 10
    ls = IrregularLombScargle(t, tol = 10, n_total = 20000, n_acc = 5000)

    s = surrogate(x, ls)

    @test all(sort(s) .== sort(x))
end

@testset "WLS" begin
    wts_norescale = WLS(AAFT(), rescale = false)
    wts_rescale = WLS(AAFT())

    s_norescale = surrogate(x, wts_norescale)
    s_rescale = surrogate(x, wts_rescale)
    @test length(s_norescale) == length(x)
    @test length(s_rescale) == length(x)

    # If rescaling, the surrogate will have the same values as the original
    @test sort(x) ≈ sort(s_rescale)

end

@testset "PartialRandomization" begin
    pr = PartialRandomization(0.2)
    s = surrogate(x, pr)

    @test length(s) == length(x)

    @test_throws AssertionError PartialRandomization(-0.01)
    @test_throws AssertionError PartialRandomization(1.01)
end

@testset "PartialRandomizationAAFT" begin
    praaft = PartialRandomizationAAFT(0.5)
    s = surrogate(x, praaft)

    @test length(s) == length(x)
    @test sort(x) ≈ sort(s)

    @test_throws AssertionError PartialRandomizationAAFT(-0.01)
    @test_throws AssertionError PartialRandomizationAAFT(1.01)
end

@testset "RandomCascade" begin
    randomcascade = RandomCascade()
    s = surrogate(x, randomcascade)

    @test length(s) == length(x)
end

@testset "Periodic" begin
	pp = PseudoPeriodic(3, 25, 0.05)
	s = surrogate(x, pp)
	@test length(s) == length(ts)
	@test all(s[i] ∈ x for i in 1:N)
	# Perhaps a more advanced test, e.g. that both components have Fourier peak at
	# the same frequency, should be considered.

	#TODO: Test for noiseradius
end

@testset "PeriodicTwin" begin
    # A better test based on recurrence plots should be considered.
    d, τ = 2, 6
    δ = 0.2
    ρ = noiseradius(x, d, τ, 0.02:0.02:0.5)
    method = PseudoPeriodicTwin(d, τ, δ, ρ)

    sg = surrogenerator(x, method)
    s = sg()[:, 1]
    @test length(s) == length(ts)
    @test all(s[i] ∈ x for i in 1:N)
end

@testset "BlockShuffle" begin
    bs1 = BlockShuffle()
    bs2 = BlockShuffle(4)
    s1 = surrogate(x, bs1)
    s2 = surrogate(x, bs2)

    @test length(s1) == length(x)
    @test length(s2) == length(x)
    @test all([s1[i] ∈ x for i = 1:N])
    @test all([s2[i] ∈ x for i = 1:N])
end

@testset "RandomShuffle" begin
    rs = RandomShuffle()
    s = surrogate(x, rs)

    @test length(s) == length(x)
    @test all([s[i] ∈ x for i = 1:N])
end

@testset "AutoRegressive" begin
	y = TimeseriesSurrogates.AR1(2000, 0.1, 0.5)
    sg = surrogenerator(y, AutoRegressive(1))
	@test 0.4 ≤ abs(sg.init.φ[1]) ≤ 0.6
	s = sg()
    @test length(s) == length(y)
end

@testset "AAFT" begin
    aaft = AAFT()
    s = surrogate(x, aaft)

    @test length(s) == length(x)
    @test all([s[i] ∈ x for i = 1:N])
end

@testset "IAAFT" begin
    iaaft = IAAFT()
    s = surrogate(x, iaaft)

    @test length(s) == length(x)
    @test all([s[i] ∈ x for i = 1:N])
end

@testset "TFTS" begin
    method_preserve_lofreq = TFTS(0.05)
    method_preserve_hifreq = TFTS(-0.05)

    s = surrogate(x, method_preserve_lofreq)
    @test length(s) == length(x)

    s = surrogate(x, method_preserve_hifreq)
    @test length(s) == length(x)

    @test_throws ArgumentError TFTS(0)
end

@testset "TFTDAAFT" begin
    tftdaaft = TFTDAAFT(0.03)
    s = surrogate(x, tftdaaft)
    @test length(s) == length(x)
    @test sort(x) ≈ sort(s)
    @test_throws ArgumentError TFTD(0)
    @test_throws ArgumentError TFTD(1.2)
end

@testset "TFTDIAAFT" begin
    tftdiaaft = TFTDAAFT(0.05)
    s = surrogate(x, tftdiaaft)
    @test length(s) == length(x)
    @test sort(x) ≈ sort(s)
    @test_throws ArgumentError TFTD(0)
    @test_throws ArgumentError TFTD(1.2)
end

@testset "TAAFT" begin
    method_preserve_lofreq = TAAFT(0.05)
    method_preserve_hifreq = TAAFT(-0.05)

    s = surrogate(x, method_preserve_lofreq)
    @test length(s) == length(x)
    @test all([s[i] ∈ x for i = 1:N])

    s = surrogate(x, method_preserve_hifreq)
    @test length(s) == length(x)
    @test all([s[i] ∈ x for i = 1:N])

    @test_throws ArgumentError TAAFT(0)
end

@testset "RandomFourier" begin
    @testset "random phases" begin
        phases = true
        rf = RandomFourier(phases)
        s = surrogate(x, rf)

        @test length(s) == length(x)
    end

    @testset "random amplitudes" begin
        phases = false
        rf = RandomFourier(phases)
        s = surrogate(x, rf)

        @test length(s) == length(x)
    end
end

@testset "TFTDRandomFourier" begin
    @testset "random phases" begin
        phases = true
        rf = TFTDRandomFourier(phases)
        s = surrogate(x, rf)

        @test length(s) == length(x)
    end

    @testset "random amplitudes" begin
        phases = false
        rf = TFTDRandomFourier(phases)
        s = surrogate(x, rf)

        @test length(s) == length(x)
    end
end

@testset "Circ/Cycle shuffle" begin
	x = random_cycles()
	s = surrogate(x, CycleShuffle())
	for a in s
		@test a ∈ x
	end
	s = surrogate(x, CircShift(1:length(x)))
	for a in s
		@test a ∈ x
	end
end

using DelayEmbeddings
@testset "ShufleDims" begin
	X = Dataset(rand(100, 3))
	Y = surrogate(X, ShuffleDimensions())
	for i in 1:100
		@test sort(X[i]) == sort(Y[i])
	end
end


# Calculate the Sum of Distance between the Cities

function sumLen(city)

    n = size(city)[1]
    len = 0

    for ii = 1: (n - 1)

        len += euclidean(city[ii, :], city[ii + 1, :])

    end

    len += euclidean(city[1, :], city[n, :])

    return len

end


"""
Computes number of connected components and the component ID for every node.

Input:
-----
* G::SparseMatrixCSC

Output:
------
* num_cc ::Int32
* conn_comp::Vector{Int32}

"""
function cc(G::SparseMatrixCSC)
	rows = G.colptr - 1
    cols = G.rowval - 1
    nodes = length(rows) - 1
    edges = length(cols)
    rows = map(Int32, rows)
    cols = map(Int32, cols)
    conn_comp = Array(Int32, nodes)
	num_cc = ccall((:cc, gunrock_lib), Cint, (Ptr{Cint}, Int, Int, Ptr{Cint}, Ptr{Cint}), conn_comp, nodes, edges, rows, cols)
    conn_comp += 1
	num_cc, conn_comp
end


using ApproxFunSingularities, ApproxFunOrthogonalPolynomials, ApproxFunBase, InfiniteArrays, BlockArrays, BlockBandedMatrices, IntervalSets, SpecialFunctions, LinearAlgebra, Random, Test
import ApproxFunBase: HeavisideSpace, PointSpace, ArraySpace, DiracSpace, PiecewiseSegment, testbandedoperator, UnionDomain, testtransforms, testfunctional, testbandedblockbandedoperator,
                        resizedata!, CachedOperator, RaggedMatrix

@testset "Sqrt" begin
    x=Fun(identity);
    @time @test sqrt(cos(π/2*x))(.1) ≈ sqrt(cos(.1π/2))

    x=Fun(identity,-2..2)
    @time u=sqrt(4-x^2)/(2π)

    @test u(.1) ≈ sqrt(4-0.1^2)/(2π)
    @test sum(u) ≈ 1

    #this call threw an error, which we check
    @test length(values(u)) == 1

    f = Fun(x->x*cot(π*x/2))
    x = Fun(identity)
    u = Fun(JacobiWeight(1.,1.,ChebyshevInterval()), (f/(1-x^2)).coefficients)
    @test 1/(0.1*cot(π*.1/2)) ≈ (1/u)(.1)

    @test (x/u)(.1) ≈ tan(π*.1/2)

    f=Fun(x->exp(-x^2),Line(0.,0.,-.5,-.5),400)
    @time @test sum(f) ≈ sqrt(π)

    f=Fun(x->exp(x)/sqrt(1-x.^2),JacobiWeight(-.5,-.5))
    @test f(.1) ≈ (x->exp(x)/sqrt(1-x.^2))(.1)

    @time @test norm(Fun(exp,Legendre(0..1))+sqrt(Fun(0..1))) ≈ 2.491141949903508

    @testset "sampling Chebyshev" begin
        x=Fun(identity)
        f = exp(x)/sqrt(1-x^2)
        @time g = cumsum(f)
        @test abs(g(-1)) ≤ 1E-15
        @test g'(0.1) ≈ f(0.1)
    end

    @testset "Complex domains sqrt" begin
        a=1+10*im; b=2-6*im
        d = IntervalCurve(Fun(x->1+a*x+b*x^2))

        x=Fun(d)
        w=sqrt(abs(leftendpoint(d)-x))*sqrt(abs(rightendpoint(d)-x))

        @test sum(w/(x-2.))/(2π*im) ≈ (-4.722196879007759+2.347910413861846im)
        @test linesum(w*log(abs(x-2.)))/π ≈ (88.5579588360686)

        a=Arc(0.,1.,0.,π/2)
        ζ=Fun(identity,a)
        f=Fun(exp,a)*sqrt(abs((ζ-1)*(ζ-im)))
    end

    @time @test norm(Fun(exp,Legendre(0..1))+sqrt(Fun(0..1))) ≈ 2.491141949903508
end

@testset "JacobiWeight" begin
    @testset "Sub-operator re-view bug" begin
        D = Derivative(Chebyshev())
        S = view(D[:, 2:end], Block.(3:4), Block.(2:4))
        @test parent(S) == D
        @test parentindices(S) == (3:4,2:4)
        @test bandwidths(S)  == (-2,2)

        DS=JacobiWeight(1,1,Jacobi(1,1))
        D=Derivative(DS)[2:end,:]
        @test domainspace(D) == DS | (1:∞)
        testbandedoperator(D)
    end

    @testset "Multiplication functions" begin
        x = Fun()
        M = Multiplication(x, JacobiWeight(0,0,Chebyshev()))
        @test exp(M).f == Multiplication(exp(x), Chebyshev()).f
    end

    @testset "Derivative" begin
        S = JacobiWeight(-1.,-1.,Chebyshev(0..1))

        # Checks bug in Derivative(S)
        @test typeof(ConstantSpace(0..1)) <: Space{ClosedInterval{Int},Float64}

        D=Derivative(S)
        f=Fun(S,Fun(exp,0..1).coefficients)
        x=0.1
        @test f(x) ≈ exp(x)*x^(-1)*(1-x)^(-1)/4
        @test (D*f)(x) ≈ -exp(x)*(1+(x-3)*x)/(4*(x-1)^2*x^2)

        S=JacobiWeight(-1.,0.,Chebyshev(0..1))
        D=Derivative(S)

        f=Fun(S,Fun(exp,0..1).coefficients)
        x=.1
        @test f(x) ≈ exp(x)*x^(-1)/2
        @test (D*f)(x) ≈ exp(x)*(x-1)/(2x^2)
    end

    @testset "Jacobi singularity" begin
        x = Fun(identity)
        f = exp(x)/(1-x.^2)

        @test f(.1) ≈ exp(.1)/(1-.1^2)
        f = exp(x)/(1-x.^2).^1
        @test f(.1) ≈ exp(.1)/(1-.1^2)
        f = exp(x)/(1-x.^2).^1.0
        @test f(.1) ≈ exp(.1)/(1-.1^2)

        ## 1/f with poles
        x=Fun(identity)
        f=sin(10x)
        g=1/f

        @test g(.123) ≈ csc(10*.123)
    end

    @testset "Jacobi conversions" begin
        S1,S2=JacobiWeight(3.,1.,Jacobi(1.,1.)),JacobiWeight(1.,1.,Jacobi(0.,1.))
        f=Fun(S1,[1,2,3.])
        C=Conversion(S1,S2)
        Cf=C*f
        @test Cf(0.1) ≈ f(0.1)

        S1,S2=JacobiWeight(3.,2.,Jacobi(1.,1.)),JacobiWeight(1.,1.,Jacobi(0.,0.))
        f=Fun(S1,[1,2,3.])
        C=Conversion(S1,S2)
        Cf=C*f
        @test Cf(0.1) ≈ f(0.1)

        C=Conversion(JacobiWeight(2,0,Jacobi(2,-2)),Legendre())
        @test C[1,1]==4/3 && C[1,2]==-2 && C[2,1]==2 && C[2,2]==-2.4
    end

    @testset "Array Conversion" begin
        a = ArraySpace(JacobiWeight(1/2,1/2, Chebyshev()), 2)
        b = ArraySpace(JacobiWeight(1/2,1/2, Ultraspherical(1)), 2)
        C = Conversion(a, b)

        f = Fun(a, rand(10))
        @test f(0.1) ≈ (C*f)(0.1)

        a = ArraySpace(JacobiWeight(1/2,1/2, Chebyshev()), 2,3)
        b = ArraySpace(JacobiWeight(1/2,1/2, Ultraspherical(1)), 2,3)
        C = Conversion(a, b)

        f = Fun(a, rand(10))
        @test f(0.1) ≈ (C*f)(0.1)
    end

    @testset "Equivalent spaces" begin
        @test norm(Fun(cos,Chebyshev)-Fun(cos,Jacobi(-0.5,-0.5)))<100eps()
        @test norm(Fun(cos,Chebyshev)-Fun(cos,JacobiWeight(0,0)))<100eps()
        @test norm(Fun(cos,Jacobi(-0.5,-0.5))-Fun(cos,JacobiWeight(0,0))) < 100eps()
        @test norm(Fun(cos,Chebyshev)-Fun(cos,JacobiWeight(0,0,Jacobi(-0.5,-0.5))))<100eps()
        @test norm(Fun(cos,Jacobi(-0.5,-0.5))-Fun(cos,JacobiWeight(0,0,Jacobi(-0.5,-0.5))))<100eps()
    end
end

@testset "Ray and Line" begin
    @testset "broadcast!" begin
        f = Fun()
        @test_throws ArgumentError (f .= Fun(Line()))
    end
    @testset "Ray" begin
        @test Inf in Ray()   # this was a bug
        @test Space(0..Inf) == Chebyshev(Ray())
        f=Fun(x->exp(-x),0..Inf)
        @test f(0.1) ≈ exp(-0.1)
        @test f'(.1) ≈ -f(.1)

        x=Fun(identity,Ray())
        f=exp(-x)
        u=integrate(f)
        @test (u(1.)-u(0)-1) ≈ -f(1)

        x=Fun(identity,Ray())
        f=x^(-0.123)*exp(-x)
        @test integrate(f)'(1.) ≈ f(1.)

        @test ≈(sum(Fun(sech,0..Inf)),sum(Fun(sech,0..40));atol=1000000eps())
        @test Line() ∪ Ray() == Line()
        @test Line() ∩ Ray() == Ray()

        f=Fun(sech,Line())
        @test Fun(f,Ray())(2.0) ≈ sech(2.0)
        @test Fun(f,Ray(0.,π))(-2.0) ≈ sech(-2.0)
        @test Fun(sech,Ray(0.,π))(-2.0) ≈ sech(-2.0)
    end

    @testset "Ei (Exp Integral)" begin
        y=Fun(Ray())
        q=integrate(exp(-y)/y)
        @test (q-last(q))(2.) ≈ (-0.04890051070806113)

        ## Line
        f=Fun(x->exp(-x^2),Line())

        @test f'(0.1) ≈ -2*0.1exp(-0.1^2)
        @test (Derivative()*f)(0.1) ≈ -2*0.1exp(-0.1^2)
    end
end

@testset "LogWeight" begin
    x=Fun(identity,-1..1)
    f=exp(x+1)-1
    @test log(f)(0.1) ≈ log(f(0.1))

    x=Fun(identity,0..1)
    f=exp(x)-1
    @test log(f)(0.1) ≈ log(f(0.1))

    x=Fun(identity,0..1)
    @test Fun(exp(x)/x-1/x,Chebyshev)(0.1) ≈ (exp(0.1)-1)/0.1

    x=Fun(identity,0..1)
    f=1/x
    p=integrate(f)
    @test (p-p(1.))(0.5) ≈ log(0.5)

    f=1/(1-x)
    p=integrate(f)
    @test (p-p(0.))(0.5) ≈ -log(1-0.5)

    @testset "#393" begin
        x = Fun(0..1)
        @time f = exp(x)*sqrt(x)*log(1-x)
        @test f(0.1) ≈ exp(0.1)*sqrt(0.1)*log(1-0.1)
    end
end

@testset "DiracDelta and PointSpace" begin
    a,b=DiracDelta(0.),DiracDelta(1.)
    f=Fun(exp)
    g=a+0.2b+f
    @test components(g)[2](0.) ≈ 1.
    @test g(.1) ≈ exp(.1)
    @test sum(g) ≈ (sum(f)+1.2)

    #Checks prevoius bug
    δ=DiracDelta()
    x=Fun()
    w=sqrt(1-x^2)
    @test (w+δ)(0.1) ≈ w(0.1)
    @test sum(w+δ) ≈ sum(w)+1

    ## PointSpace
    f=Fun(x->(x-0.1),PointSpace([0,0.1,1]))
    g = f + Fun(2..3)
    @test f(0.0) ≈ g(0.0) ≈ -0.1
    @test f(0.1) ≈ g(0.1) ≈ 0.0
    @test f(1.0) ≈ g(1.0) ≈ 0.9

    @test g(2.3) ≈ 2.3

    h = a + Fun(2..3)

    # for some reason this test is broken only on Travis
    @test_skip g/h ≈ f/a + Fun(1,2..3)
end

@testset "Multiple roots" begin
    x=Fun(identity,-1..1)

    @test (1/x^2)(0.1) ≈ 100.
    @test (1/x^2)(-0.1) ≈ 100.

    fc=x*(1+x)^2
    @time @test (1/fc)(0.1) ≈ 1/fc(0.1)

    fc=x*(1-x)^2
    @test (1/fc)(0.1) ≈ 1/fc(0.1)
end

@testset "special function singularities" begin
    x=Fun(0..1)
    @time @test erf(sqrt(x))(0.1) ≈ erf(sqrt(0.1))
    @time @test erfc(sqrt(x))(0.1) ≈ erfc(sqrt(0.1))

    ## roots of log(abs(x-y))
    x=Fun(-2..(-1))
    @test space(abs(x)) == Chebyshev(-2 .. (-1))

    @test roots(abs(x+1.2)) ≈ [-1.2]

    f = sign(x+1.2)
    @test space(f) isa PiecewiseSpace
    @test f(-1.4) == -1
    @test f(-1.1) == 1

    f=abs(x+1.2)
    @test abs(f)(-1.3) ≈ f(-1.3)
    @test abs(f)(-1.1) ≈ f(-1.1)
    @test norm(abs(f)-f)<10eps()

    @test norm(sign(f)-Fun(1,space(f)))<10eps()

    @test log(f)(-1.3) ≈ log(abs(-1.3+1.2))
    @test log(f)(-1.1) ≈ log(abs(-1.1+1.2))
end

@testset "SumSpace" begin
    @testset "SumSpace bug" begin
        dsp=JacobiWeight(1.,0.,Jacobi(1.,0.,0..1))⊕JacobiWeight(0.5,0.,Jacobi(0.5,-0.5,0..1))
        rsp=Legendre(0..1)⊕JacobiWeight(0.5,0.,Jacobi(0.5,0.5,0..1))

        C=Conversion(dsp,rsp)

        f=Fun(dsp,[1.,2.,3.,4.,5.])
        @test f(0.1) ≈ (C*f)(0.1)
    end

    @testset "Triple SumSpace" begin
        x=Fun()
        w=log(1-x)+sqrt(1-x^2)
        @time f=w+x
        @test f(0.1) ≈ (w(0.1)+0.1)
        @test (w+1)(0.1) ≈ (w(0.1)+1)
        @test (w+x+1)(0.1) ≈ (w(0.1)+1.1)
        @test ((w+x)+1)(0.1) ≈ (w(0.1)+1.1)
    end

    @testset "Previoius segfault" begin
        x=Fun(identity,-1..1)
        @time f=x+sin(2x)*sqrt(1-x^2)
        @test f(0.1) ≈ 0.1+sin(2*0.1)*sqrt(1-0.1^2)
    end

    @testset "Multiple piecewisespace" begin
        x=Fun(identity,-3 .. -2)+Fun(identity,2..3)
        w=sqrt(9-x^2)
        f=w+Fun()
        @test (f+w)(2.5) ≈ 2w(2.5)
        @test (f+w)(.5) ≈ f(.5)
    end

    @testset "cancellation conversion" begin
        x=Fun(0..1)
        f=exp(x)-1
        Fun(f,JacobiWeight(1.,0.,0..1))
    end

    @testset "Union of ChebyshevDirichlet" begin
        dom = UnionDomain(0..1, 2..3)
        @test components(union(JacobiWeight.(-0.5,-0.5,ChebyshevDirichlet{1,1}.(components(dom)))...)) ==
            (JacobiWeight.(-0.5,-0.5,ChebyshevDirichlet{1,1}.(components(dom)))...,)
    end

    @testset "Ultraspherical special functions" begin
        x = Fun(Ultraspherical(2,0..1))
        sqrt(x)(0.1) ≈ sqrt(0.1)

        f = Fun(x->x*exp(x),Ultraspherical(1,0..1))
        sqrt(f(0.1)) ≈ sqrt(f)(0.1)
    end

    @testset "one for SumSpace" begin
        S = Jacobi(0,1) ⊕ JacobiWeight(1/3,0,Jacobi(1/3,2/3)) ⊕ JacobiWeight(2/3,0,Jacobi(2/3,1/3))
        o = ones(S)
        @test o(0.5) ≈ 1
    end

    @testset "Complex piecewise" begin
        x = Fun(identity, Segment(im,0) ∪ Segment(0,1))

        @test sqrt(1-x)(0.2im) ≈ sqrt(1-0.2im)
        @test sqrt(1-x)(0.2) ≈ sqrt(1-0.2)

        w=2/(sqrt(1-x)*sqrt(1+im*x))
        @test w(0.2im) ≈ 2/(sqrt(1-0.2im)*sqrt(1+im*(0.2im)))
        @test w(0.2) ≈ 2/(sqrt(1-0.2)*sqrt(1+im*(0.2)))
    end
end

@testset "WeightedJacobi" begin
    m=20
    @time testtransforms(JacobiWeight(0.,m,Jacobi(0.,2m+1)))
    f=Fun(x->((1-x)/2).^m.*exp(x),JacobiWeight(0.,m,Jacobi(0.,2m+1)))
    @test abs(f(.1)-(x->((1-x)/2).^m.*exp(x))(.1))<10eps()

    m=10
    @time f=Fun(x->besselj(m,m*(1-x)),JacobiWeight(0.,m,Jacobi(0.,2m+1)))
    @test f(0.) ≈ besselj(m,m)

    @testset "Conversion" begin
        testtransforms(Jacobi(-0.5,-0.5))
        @test norm(Fun(Fun(exp),Jacobi(-.5,-.5))-Fun(exp,Jacobi(-.5,-.5))) < 300eps()

        x=Fun(identity)
        ri=0.5/(1-x)
        @test ((1-x)/2 .* Fun(exp,JacobiWeight(0.,0.,Jacobi(0.,1.))))(.1) ≈ (1-.1)./2*exp(.1)

        @test ((1-x)/2 .* Fun(exp,JacobiWeight(0.,0.,Jacobi(0.,1.))))(.1) ≈ (1-.1)./2*exp(.1)

        @test (ri*Fun(exp,JacobiWeight(0.,0.,Jacobi(0.,1.))))(.1) ≈ .5/(1-.1)*exp(.1)
    end

    @testset "Derivative" begin
        D=Derivative(Jacobi(0.,1.,Segment(1.,0.)))
        @time testbandedoperator(D)

        S=JacobiWeight(0.,0.,Jacobi(0.,1.,Segment(1.,0.)))
        D=Derivative(S)
        testbandedoperator(D)

        f=Fun(exp,domainspace(D))
        @test (D*f-f).coefficients|>norm < eps(100000.)
        @test (f'-f).coefficients|>norm < eps(100000.)
        @test (D^2*f-f).coefficients|>norm < eps(100000000.)
        @test (D*(D*f)-f).coefficients|>norm < eps(100000000.)

        S=JacobiWeight(1,1,Ultraspherical(1))

        f=Fun(S,[1.,2.,3.])
        @test (Derivative(S,2)*f)(0.1) ≈ f''(0.1)
    end

    @testset "special derivative" begin
        x=Fun()
        f=exp(x)*sqrt(1-x^2)
        D=Derivative(WeightedJacobi(.5,.5))

        testtransforms(WeightedJacobi(.5,.5))
        testbandedoperator(D)

        @time g=(D*Fun(f,domainspace(D)))
        @test f'(0.1) ≈ g(0.1)
    end

    @testset "==" begin
        @test WeightedJacobi(0.1,0.2) == WeightedJacobi(0.1+eps(),0.2)
    end

    @testset "subspace bug" begin
        f=Fun(WeightedJacobi(0.1,0.2),rand(10))  # convert to Legendre expansion

        g=(f|(2:∞))

        @test ApproxFunBase.coefficients(g.coefficients,space(g),ApproxFunBase.canonicalspace(g))[1] ==0.
        @test norm((Fun(g,space(f))|(2:∞)-g).coefficients) < 10eps()
    end

    @testset "JacobiWeight cumsum bug Issue #557" begin
        x = Fun(0.0..1.0)
        ν = 2
        @time f = x^(ν/2-1) * exp(-x/2) # 0.05s
        @test cumsum(f)' ≈ f
        @test cumsum(f)(1.0) ≈ 0.7869386805747332 # Mathematic

        x = Fun(Ray())
        ν = 2
        @time f = x^(ν/2-1) * exp(-x/2) # 0.05s
        @test cumsum(f)' ≈ f
        @test cumsum(f)(1.0) ≈ 0.7869386805747332
    end

    @testset "Definite integral" begin
        for S in (WeightedJacobi(0,0), JacobiWeight(0,0, Legendre(1.1..2.3)), Legendre())
            B = DefiniteIntegral(S)
            testfunctional(B)
            @test ApproxFunBase.rowstop(B,1) == 1
            B[1] == arclength(domain(S))
            f = Fun(exp, S)
            B*f == sum(Fun(exp,domain(S)))
        end
    end
end

@testset "WeightedLaguerre" begin
    @testset "WeightedLaguerre cumsum" begin
        α = 2.7
        f = Fun(WeightedLaguerre(α), [1.0]);
        f = Fun(f, JacobiWeight(α,0,Chebyshev(Ray())));
        g = integrate(f)
        g(3.0) - cumsum(Fun(x -> f(x), 0..6))(3.0)
    end

    @testset "Log with squareroot singularities" begin
        a = 1.0; b = 2.0+im
        d = Segment(a,b)
        z = Fun(d)

        f = real(exp(z) / (sqrt(z-a)*sqrt(b-z)))
        S=space(f)
        x=4.0+2im;
        @test linesum(f*log(abs(x-z))) ≈ 13.740676344264614
    end
end

@testset "Hermite Integration" begin
    @test_throws ArgumentError integrate(Fun(GaussWeight(Hermite(2),1), [0.0,1.0]))

    w = Fun(GaussWeight(Hermite(2), 0), [1.0,2.0,3.0])
    g = integrate(w)
    g̃ = Fun(Hermite(2), [0.0, 0.5, 0.5, 0.5])
    @test g(0.1) == g̃(0.1)

    w = Fun(GaussWeight(), Float64[])
    g = integrate(w)
    @test g(0.1) == 0.0

    w = Fun(GaussWeight(), [1.0])
    g = integrate(w)
    @test_skip w̃ = Fun(w, -7..7)
    w̃ = Fun( x-> w(x), -7..7)
    g̃ = cumsum(w̃)
    @test g(3) - g(-7) ≈ g̃(3)

    w = Fun(GaussWeight(), Float64[1.0])
    g = integrate(w)
    @test_skip w̃ = Fun(w, -7..7)
    w̃ = Fun(x -> w(x), -7..7)
    g̃ = cumsum(w̃)
    @test g(3) - g(-7) ≈ g̃(3)

    w = Fun(GaussWeight(Hermite(2), 2), Float64[1.0])
    g = integrate(w)
    @test_skip w̃ = Fun(w, -7..7)
    w̃ = Fun(x -> w(x), -7..7)
    g̃ = cumsum(w̃)
    @test g(3) - g(-7) ≈ g̃(3)

    w = Fun(GaussWeight(), Float64[0.0, 1.0])
    g = integrate(w)
    @test_skip w̃ = Fun(w, -7..7)
    w̃ = Fun(x -> w(x), -7..7)
    g̃ = cumsum(w̃)
    @test g(3) - g(-7) ≈ g̃(3)

    w = Fun(GaussWeight(Hermite(2), 2), Float64[0.0, 1.0])
    g = integrate(w)
    @test_skip w̃ = Fun(w, -7..7)
    w̃ = Fun(x -> w(x), -7..7)
    g̃ = cumsum(w̃)
    @test g(3) - g(-7) ≈ g̃(3)
end

@testset "Laguerre" begin
    @testset "Integration" begin
        @test_throws ArgumentError integrate(Fun(LaguerreWeight(1, Laguerre(2)), [1.0, 2.0]))

        g = integrate(Fun(WeightedLaguerre(), []))
        @test g(0.1) == 0.0

        α = 2.8
        f = Fun(WeightedLaguerre(α), [1.0])
        g = integrate(f)
        @test g(3.0) - cumsum(Fun(x -> f(x), 0..6))(3.0) ≈ g(0.0)

        α = 2
        f = Fun(WeightedLaguerre(α), [1.0])
        g = integrate(f)
        @test g(3.0) - cumsum(Fun(x -> f(x), 0..6))(3.0) ≈ g(0.0)

        α = 2.8
        f = Fun(WeightedLaguerre(α), [0.0, 1.0])
        g = integrate(f)
        f̃ = Fun(x -> f(x), 0 .. 100)
        g̃ = integrate(f̃)
        g̃ = g̃ - last(g̃)
        @test g(3.0) ≈ g̃(3.0)

        α = 2
        f = Fun(WeightedLaguerre(α), [1.0, 1.0])
        g = integrate(f)
        @test g(3.0) - cumsum(Fun(x -> f(x), 0..6))(3.0) ≈ g(0.0)

        α = 5
        f = Fun(WeightedLaguerre(α), [1.0])
        g = integrate(f)
        @test g(3.0) - cumsum(Fun(x -> f(x), 0..6))(3.0) ≈ g(0.0)

        α = 5
        f = Fun(WeightedLaguerre(α), [0.0, 1.0])
        g = integrate(f)
        f̃ = Fun(x -> f(x), 0 .. 100)
        g̃ = integrate(f̃)
        g̃ = g̃ - last(g̃)
        @test g(3.0) ≈ g̃(3.0)

        α = 5
        f = Fun(WeightedLaguerre(α), [1.0, 1.0])
        g = integrate(f)
        @test g(3.0) - cumsum(Fun(x -> f(x), 0..6))(3.0) ≈ g(0.0)
    end

    @testset "Correct domain" begin
        w = Fun(WeightedLaguerre(0.5),[1.0])
        h = cumsum(w)
        @test domain(h) == Ray()
    end
end

include("IntegralEquationsTest.jl")

@testset "PDE" begin
    @testset "Laplacian" begin
        S=WeightedJacobi(1,1)^2
        L=Laplacian(S)
        testbandedblockbandedoperator(L)
    end
    @testset "Zero Dirichlet" begin
        S = JacobiWeight(1.,1.,Jacobi(1.,1.))^2
        Δ = Laplacian(S)

        testbandedblockbandedoperator(Δ)

        u = Fun((x,y)->sin(π*x)*sin(π*y),S)
        f = -2π^2*u

        F = qr(Δ)
        ApproxFunBase.resizedata!(F,:,1000)
        @time v=F\f
        @test norm((u-v).coefficients)<100eps()

        F=qr(Δ)
        ApproxFunBase.resizedata!(F.R_cache,:,100)
        ApproxFunBase.resizedata!(F.R_cache,:,1000)
        @time v=F \ f
        @test norm((u-v).coefficients)<100eps()

        F=qr(Δ)
        @time v=F\f
        @test norm((u-v).coefficients)<100eps()
    end
    @testset "Gaussian zero Dirichlet Poisson" begin
        S=JacobiWeight(1.,1.,Jacobi(1.,1.))^2
        Δ=Laplacian(S)

        f=Fun((x,y)->exp(-10(x+.2)^2-20(y-.1)^2),rangespace(Δ))  #default is [-1,1]^2
        @time v = \(Δ,f;tolerance=1E-14)
        @test norm((Δ*v-f).coefficients) < 1E-14

    end
    @testset "check we dispatch correctly to get fast build" begin
        S = JacobiWeight(1.,1.,Jacobi(1.,1.))^2
        Δ = Laplacian(S)
        @time S = view(Δ.op.ops[1].ops[1].op,Block.(1:40), Block.(1:40))
        @test typeof(S.parent.domaintensorizer) == ApproxFunBase.Trivial2DTensorizer
    end

    @testset "Operator resize" begin
        S=ChebyshevDirichlet()^2
        B=Dirichlet(S)
        f = Fun((x,y)->exp(x)*sin(y),S)
        @test norm((Fun((x,y)->exp(x)*sin(y),∂(domain(S))) - B*f).coefficients) < 100eps()

        S=JacobiWeight(1.,1.,Jacobi(1.,1.))^2
        Δ=Laplacian(S)

        @test cache(Δ)[1:100,1:100]  ≈ Δ[1:100,1:100]
        @test cache(Δ;padding=true)[1:100,1:100]  ≈ Δ[1:100,1:100]

        @test cache(Δ)[5:100,7:100]  ≈ Δ[5:100,7:100]
        @test cache(Δ;padding=true)[5:100,7:100]  ≈ Δ[5:100,7:100]

        # Check that QR is growing correctly
        for col in (1,2,3,10,11,40)
            QR=qr(Δ)
            resizedata!(QR.R_cache,:,col+100)
            resizedata!(QR,:,col)
            QR2=qr!(CachedOperator(RaggedMatrix,Δ;padding=true))
            resizedata!(QR2.R_cache,:,QR.ncols+100)
            resizedata!(QR2,:,QR.ncols)
            n=min(size(QR.H,1),size(QR2.H,1))
            @test QR.H[1:n,1:col] ≈ QR2.H[1:n,1:col]
            @test QR.R_cache[1:col,1:col] ≈ QR2.R_cache[1:col,1:col]
            @test QR.R_cache[1:col+10,1:col+10] ≈ QR2.R_cache[1:col+10,1:col+10]
        end

        QRR=qr(Δ)
        QR2=qr!(CachedOperator(RaggedMatrix,Δ;padding=true))
        for col in (80,200)
            resizedata!(QRR,:,col)
            resizedata!(QR2,:,QRR.ncols)
            n=min(size(QRR.H,1),size(QR2.H,1))
            @test QRR.H[1:n,1:col] ≈ QR2.H[1:n,1:col]
            @test QRR.R_cache[1:col,1:col] ≈ QR2.R_cache[1:col,1:col]
            @test QRR.R_cache[1:col+10,1:col+10] ≈ QR2.R_cache[1:col+10,1:col+10]
        end

        # this checks a bug
        QRR=qr(Δ)
        resizedata!(QRR,:,548)
        resizedata!(QRR,:,430)

        u=Fun((x,y)->sin(π*x)*sin(π*y),S)
        f=-2π^2*u

        QRR=qr(Δ)
        v=QRR\f
        @test norm((u-v).coefficients)<100eps()

        v=Δ\f
        @test norm((u-v).coefficients)<100eps()

        f=Fun((x,y)->exp(-10(x+.2)^2-20(y-.1)^2),rangespace(Δ))  #default is [-1,1]^2
        @time v=\(Δ,f;tolerance=1E-14)
        @test norm((Δ*v-f).coefficients)<1E-14

        KO=Δ.op.ops[1].ops[1].op

        M=BandedBlockBandedMatrix(view(KO,1:4,1:4))
        @test norm(BandedBlockBandedMatrix(view(KO,1:4,2:4))-M[:,2:4]) < 10eps()
        @test norm(BandedBlockBandedMatrix(view(KO,1:4,3:4))-M[:,3:4]) < 10eps()
    end
end

@testset "Cauchy fun" begin
    f = Fun((x,y)->1/(2π*(x^2+y^2+1)^(3/2)),Line()^2)
    @test f(0.1,0.2) ≈ 1/(2π*(0.1^2+0.2^2+1)^(3/2))

    #TODO: improve tolerance
    f = LowRankFun((x,y)->1/(2π*(x^2+y^2+1)^(3/2)),JacobiWeight(2.,2.,Line())^2)
    @test ≈(f(0.1,0.2),1/(2π*(0.1^2+0.2^2+1)^(3/2));atol=1E-4)
end

@testset "Bessel" begin
    @time for ν in (1.,0.5,2.,3.5)
        println("        ν = $ν")
        S=JacobiWeight(-ν,0.,Chebyshev(0..1))
        D=Derivative(S)
        x=Fun(identity,domain(S))
        L=(x^2)*D^2+x*D+(x^2-ν^2);
        u=\([rdirichlet(S);rneumann(S);L],[bessely(ν,1.),.5*(bessely(ν-1.,1.)-bessely(ν+1.,1.)),0];
                    tolerance=1E-10)
        @test ≈(u(.1),bessely(ν,.1);atol=eps(1000000.)*max(abs(u(.1)),1))
        u=Fun(x->bessely(ν,x),S)
        @test ≈(u(.1),bessely(ν,.1);atol=eps(10000.)*max(abs(u(.1)),1))
        u=Fun(x->besselj(ν,x),S)
        @test ≈(u(.1),besselj(ν,.1);atol=eps(10000.)*max(abs(u(.1)),1))
    end

    @time for ν in (1.,0.5,0.123,3.5)
        println("        ν = $ν")
        S=JacobiWeight(ν,0.,Chebyshev(0..1))
        D=Derivative(S)
        x=Fun(identity,domain(S))
        L=(x^2)*D^2+x*D+(x^2-ν^2);

        u=\([rdirichlet(S);rneumann(S);L],[besselj(ν,1.),.5*(besselj(ν-1.,1.)-besselj(ν+1.,1.)),0];
                    tolerance=1E-10)
        @test ≈(u(.1),besselj(ν,.1);atol=eps(1000000.)*max(abs(u(.1)),1))
        u=Fun(x->besselj(ν,x),S)
        @test ≈(u(.1),besselj(ν,.1);atol=eps(10000.)*max(abs(u(.1)),1))
    end
end

@testset "Speed test" begin
    S = JacobiWeight(1.,1.,Jacobi(1.,1.))^2
    Δ = Laplacian(S)

    f = Fun((x,y)->sin(π*x)*sin(π*y),S)

    QR1=qr(Δ)
    ApproxFunBase.resizedata!(QR1,:,400)
        \(QR1,f; tolerance=1E-10)
    QR1=qr(Δ)
        @time Δ[Block.(1:40), Block.(1:40)]
        @time ApproxFunBase.resizedata!(QR1,:,400)
        @time \(QR1,f; tolerance=1E-10)
    println("Laplace Dirichlet: should be ~0.015, 0.015, 0.001")
end


# * Dispersed fields
#   We calculate the effect of dispersion described by the Sellmeier
#   equations in the frequency domain, to which we transform via an
#   FFT. For this to be possible, we only allow the evaluation of the
#   field for a specified sampling frequency, i.e. we don't provide an
#   implementation for evaluating a =DispersedField= at an arbitrary
#   time point.

struct DispersedField <: AbstractField
    a::AbstractField
    m::Medium
    d::Unitful.Length
end

function show(io::IO, f::DispersedField)
    show(io, f.a)
    write(io, "\n  – dispersed through $(f.d) of $(f.m)")
end

disperse(a::AbstractField, m::Medium, d::Unitful.Length) =
    DispersedField(a, m, d)

(f::DispersedField)(t::Union{Unitful.Time,AbstractVector{<:Unitful.Time}}) =
    error("Dispersed fields can only be evaluated by specifying a sampling frequency")

function (f::DispersedField)(fs::Unitful.Frequency = default_sampling_frequency(f.a))
    F̂ = fft(f.a, fs)
    ifft(F̂.*dispersion(f.m, f.d, fftfreq(f.a, fs), frequency(f.a)))*(u"V"*u"m")
end

max_frequency(f::DispersedField) = max_frequency(f.a)

# This is, strictly speaking, not true, since dispersing a pulse will
# in general stretch it in the time domain. It is up to the user to
# ensure that the time window is large enough to contain even the
# stretched pulse, by setting the σmax (or friends) to a large enough
# value.
span(f::DispersedField) = span(f.a)

export disperse


# 2D Heat Diffusion
# LBM
#
# This could only work on a rectangular region! To generalize to more
# complicated mesh, I need 1D indexing of lattices.
#
# Hongyang Zhou, hyzhou@umich.edu 05/01/2019

using PyPlot
include("Parameters.jl")
include("lattice.jl")
include("BC.jl")

using .Parameters, .BC, .lattice

function solveHeatDiffusion2D!(param::Param,lat::Lattice)

   # Set aliases
   nI,nJ,ω,Δt,TotalTime = param.nI,param.nJ,param.ω,param.Δt,param.TotalTime

   nw,w  = lat.nw, lat.w
   T     = zeros(1,nI,nJ)
   feq   = Array{Float64,3}(undef,nw,nI,nJ)
   ftemp = Array{Float64,3}(undef,nw,nI,nJ)
   t     = 0.0     # Current simulation time

   # Advance
   while t < TotalTime

      # Calculate equilibrium distribution
      for iw = 1:nw
         @. feq[iw,:,:] = w[iw]*T[1,:,:]
      end

      # Collision
      for iw = 1:nw
         @. lat.f[iw,:,:] = (1.0-ω)*lat.f[iw,:,:] + ω.*feq[iw,:,:]
      end

      # Streaming
      for j=1:nJ, i=1:nI, iw = 1:nw
         iNew = i + lat.c[1,iw]
         jNew = j + lat.c[2,iw]
         (iNew > nI || iNew < 1 || jNew > nJ || jNew < 1) && continue
         ftemp[iw,iNew,jNew] = lat.f[iw,i,j]
      end
      #copy!(lat.f,ftemp) # This does not work in Julia 1.0 (but 1.1+)!
      lat.f .= ftemp

      # Remaining BCs
      setBCs!(lat,param)

      # Calculate macroscopic T
      T = sum(lat.f, dims=1)

      t += Δt
   end

end

function init_lattice()

   param = setParameters()

   Δx,Δt,nI,nJ,LT = param.Δx,param.Δt,param.nI,param.nJ,param.LatticeType

   if LT == 1
      lat = LatticeD2Q4(Δx,Δt,nI,nJ)
   elseif LT == 2
      lat = LatticeD2Q5(Δx,Δt,nI,nJ)
   elseif LT == 3
      lat = LatticeD2Q9(Δx,Δt,nI,nJ)
   end

   return param,lat
end

function plotSol(lat::Lattice)
   # Visualization
   x = range(0,stop=lat.nI-1,length=lat.nI)
   y = range(0,stop=lat.nJ-1,length=lat.nJ)
   T = sum(lat.f, dims=1)

   # Python follows row major order
   contourf(x,y,T[1,:,:]'); colorbar()
   title("2D Heat Diffusion in a Plate")
   xlabel("x")
   ylabel("y")
   grid(true)

   # 1D cut plot
   #plot(0:nx-1,T[1,:,50],linestyle="-",marker="o",label="LBM")
   #xlim(0,30)
end

function main()
   param,lat = init_lattice()

   solveHeatDiffusion2D!(param,lat)

   plotSol(lat)

end


using StateSpaceModels, BenchmarkTools, CSV
# Local level series
locallevel_series = CSV.read("./benchmark/locallevel.csv"; header = false)
verbose = 0

locallevel = local_level(Matrix{Float64}(locallevel_series))
locallevel.H[1] = 1
locallevel.Q[1] = 1
@benchmark StateSpaceModels.kalman_filter($locallevel)
# 24th October 2019
# BenchmarkTools.Trial:
#   memory estimate:  12.17 KiB
#   allocs estimate:  93
#   --------------
#   minimum time:     9.258 μs (0.00% GC)
#   median time:      9.832 μs (0.00% GC)
#   mean time:        12.296 μs (11.59% GC)
#   maximum time:     4.382 ms (99.48% GC)
#   --------------
#   samples:          10000
#   evals/sample:     1

locallevel = local_level(Matrix{Float64}(locallevel_series))
locallevel.H[1] = 1
locallevel.Q[1] = 1
@benchmark StateSpaceModels.univariate_kalman_filter($locallevel)
# 24th October 2019
# BenchmarkTools.Trial:
#   memory estimate:  8.33 KiB
#   allocs estimate:  36
#   --------------
#   minimum time:     4.977 μs (0.00% GC)
#   median time:      5.458 μs (0.00% GC)
#   mean time:        6.477 μs (10.23% GC)
#   maximum time:     406.721 μs (97.11% GC)
#   --------------
#   samples:          10000
#   evals/sample:     6

# Benchmark structural model
AP = CSV.read("./examples/AirPassengers.csv")
logAP = log.(Vector{Float64}(AP[:Passengers]))

structural_model = structural(logAP, 12)
structural_model.H[1] = 1e-4
structural_model.Q[[1, 5, 9]] .= 1e-4

@benchmark StateSpaceModels.kalman_filter($structural_model)
# 24th October 2019
# BenchmarkTools.Trial:
#   memory estimate:  740.64 KiB
#   allocs estimate:  1455
#   --------------
#   minimum time:     336.863 μs (0.00% GC)
#   median time:      352.551 μs (0.00% GC)
#   mean time:        473.981 μs (13.59% GC)
#   maximum time:     5.512 ms (77.69% GC)
#   --------------
#   samples:          10000
#   evals/sample:     1

structural_model = structural(logAP, 12)
structural_model.H[1] = 1e-4
structural_model.Q[[1, 5, 9]] .= 1e-4

@benchmark StateSpaceModels.univariate_kalman_filter($structural_model)
# 24th October 2019
# BenchmarkTools.Trial:
#   memory estimate:  658.19 KiB
#   allocs estimate:  446
#   --------------
#   minimum time:     246.795 μs (0.00% GC)
#   median time:      256.044 μs (0.00% GC)
#   mean time:        304.342 μs (12.20% GC)
#   maximum time:     3.262 ms (79.72% GC)
#   --------------
#   samples:          10000
#   evals/sample:     1


"""
Creation and manipulation of ∇² differential operator
"""
module Laplacians

export Laplacian

files = (
    "1D/Laplacians1D.jl",
    "2D/Symmetric/Laplacians2D.jl",
    # "3D/Laplacians3D.jl"
    )

using ..Lattices
using ..Points
using SparseArrays

import ..Symmetric, ..Unsymmetric
import LinearAlgebra: I

struct Laplacian{N,CLASS}
    l0::SparseMatrixCSC{ComplexF64,Int}
    complex_scaling::SparseMatrixCSC{ComplexF64,Int}
end

foreach(include,files)

function Base.conj!(lap::Laplacian{N}) where N
    conj!(lap.l0)
    return lap
end
Base.conj(lap::Laplacian{N,CLASS}) where {N,CLASS} = Laplacian{N,CLASS}(conj(lap.l0),lap.complex_scaling)

################################################################################
# PRETTY PRINTING

import ..PRINTED_COLOR_DARK
import ..PRINTED_COLOR_VARIABLE

function Base.show(io::IO,::Laplacian{N,CLASS}) where {N,CLASS}
    print(io,"$(N)D ")
    CLASS<:Symmetric ? print(io,"Symmetric ") : nothing
    CLASS<:Unsymmetric ? print(io,"Unsymmetric ") : nothing
    printstyled(io,"Laplacian", color=PRINTED_COLOR_DARK)
end

end


#--------------------------------------------------------------------
# Spacetime Discretization methods in Julia
# Soham 06-2018
#--------------------------------------------------------------------

function derivOP_corrected{T<:Int}(Nx::T, Ny::T)::Array{Float64, 4}
	" NOTE: This is the corrected version which evaluates the operator
			correctly when nx != ny. However, this doesn't seem to
			give the correct solution."
	operator = zeros(Nx+1, Ny+1, Nx+1, Ny+1)
    for index in CartesianRange(size(operator))
        (m,n,p,q) = index.I
        operator[m,n,p,q] = 2*chebw(p,Nx)*chebw(n,Ny)*chebd(p,m,Nx)*chebd(n,q,Ny)
    end
    return operator
end

function derivOP{T<:Int}(Nx::T, Ny::T)::Array{Float64, 4}
    @assert Nx == Ny
	operator = zeros(Nx+1, Ny+1, Nx+1, Ny+1)
    for index in CartesianRange(size(operator))
        k = index.I[1]
        i = index.I[2]
        l = index.I[3]
        j = index.I[4]
        operator[index] = 2*chebw(i,Ny)*chebw(k,Nx)*chebd(k,l,Nx)*chebd(i,j,Ny)
	end
	return operator
end

function boundaryOP{T<:Int}(Nx::T, Ny::T)::Array{Float64, 4}
    bnd = zeros(Nx+1, Ny+1, Nx+1, Ny+1)
    for index in CartesianRange(size(bnd))
        k = index.I[1]
        i = index.I[2]
        l = index.I[3]
        j = index.I[4]
        if  i==1 || k==1
            bnd[index] = delta(i,j)*delta(k,l)
        end
    end
    return bnd
end

function RHS{T<:Int}(fn::Function, Nx::T, Ny::T, M::T, loc::Array{Int,1})::Array{Float64,2}
    rhs = projectonPatchbyRestriction(fn, Nx, Ny, M, loc)
    for index in CartesianRange(size(rhs))
        (i,j) = index.I
        rhs[i,j] = chebw(i,Nx)*chebw(j,Ny)*rhs[i,j]
    end
    return rhs
end


"""
Represent any source (incident) wave

Subtypes may have a symmetry (such as [`PlaneSource`](@ref)) and will contain information about physical medium.
"""
abstract type AbstractSource{T,P} end

"""
    PlaneSource(medium::P, amplitude::SVector, direction::SVector)

Is a struct type which describes a plane-wave source that drives/forces the whole system. It has three fields: a physical `medium`, an `amplitude` of the source, and the direction the propagate in `direction`.

For any given angular frequency ω, the PlaneSource has the value ``e^{i ω/c \\mathbf v \\cdot \\mathbf x }`` at the point ``\\mathbf x``, where ``c`` is the medium wavespeed and ``\\mathbf v`` is the direction.
"""
struct PlaneSource{T<:AbstractFloat,Dim,FieldDim,P<:PhysicalMedium} <: AbstractSource{T,P}
    medium::P
    direction::SVector{Dim,T}
    position::SVector{Dim,T}
    amplitude::SVector{FieldDim,Complex{T}}
    # Check that P has same Dim and FieldDim
    function PlaneSource{T,Dim,FieldDim,P}(medium::P, direction::AbstractArray{T}, position::AbstractArray{T}, amplitude::AbstractArray{CT} where CT<:Union{T,Complex{T}}) where {T,Dim,FieldDim,P<:PhysicalMedium{T,Dim,FieldDim}}
        normw = sqrt(sum(direction .^2)) # note if direction is complex this is different from norm(direction)
        if !(normw ≈ one(T))
            @warn "The direction will be normalised so that sum(direction .^2) == 1.0"
        end

        if length(direction) != Dim || length(position) != Dim
            throw(DimensionMismatch("The spatial dimensions of the medium do not match dimensions of the direction or position vector."))
        elseif length(amplitude) != FieldDim
            throw(DimensionMismatch("The dimensions of the field in the medium (i.e. always 1 for acoustics) do not match the dimension of the amplitude vector."))
        else
            new{T,Dim,FieldDim,P}(medium,direction ./ normw, position, amplitude)
        end
    end
end

field(s::PlaneSource{T}) where T = function (x::AbstractArray{T}, ω::T)
        s.amplitude * exp(im * (ω / s.medium.c) * dot(x - s.position,s.direction))
    end

field(s::PlaneSource{T,Dim,1}) where {T, Dim} = function (x::AbstractArray{T}, ω::T)
        s.amplitude[1] * exp(im * (ω / s.medium.c) * dot(x - s.position,s.direction))
    end

field(s::PlaneSource{T}, x::AbstractArray{T}, ω::T) where T = field(s)(x,ω)

# Convenience constructor which does not require explicit types/parameters
function PlaneSource(medium::PhysicalMedium{T,Dim,FieldDim},
        direction::AbstractArray{T},
        position::AbstractArray{T} = zeros(T,Dim),
        amplitude::Union{CT,AbstractArray{CT}} where CT<:Union{T,Complex{T}} = SVector{FieldDim}(ones(Complex{T},FieldDim))
    ) where {T,Dim,FieldDim}

    if !(typeof(amplitude) <: AbstractArray)
        amplitude = [amplitude]
    end

    PlaneSource{T,Dim,FieldDim,typeof(medium)}(medium,direction,position,amplitude)
end

function PlaneSource(medium::PhysicalMedium{T,Dim,FieldDim};
        direction::AbstractArray{T} = [one(T); zeros(T,Dim-1)],
        position::AbstractArray{T} = zeros(T,Dim),
        amplitude::Union{CT,AbstractArray{CT}} where CT<:Union{T,Complex{T}} = SVector{FieldDim}(ones(Complex{T},FieldDim))
    ) where {T,Dim,FieldDim}

    PlaneSource(medium,direction,position,amplitude)
end

Symmetry(s::PlaneSource{T,Dim}) where {T,Dim} = (abs(dot(s.direction,azimuthalnormal(Dim))) == norm(s.direction)) ? PlanarAzimuthalSymmetry{Dim}() : PlanarSymmetry{Dim}()

"""
    RegularSource(medium::P, field::Function, coef::Function)

Is a struct type which describes the source field that drives/forces the whole system. It is also known as an incident wave. It has three fields `RegularSource.medium`, `RegularSource.field`, and `RegularSource.coef`.

The source field at the position 'x' and angular frequency 'ω' is given by
```julia-repl
x = [1.0,0.0]
ω = 1.0
RegularSource.field(x,ω)
```

The field `RegularSource.coef`
regular_basis_function(medium::Acoustic{T,2}, ω::T)
"""
struct RegularSource{T<:AbstractFloat,P<:PhysicalMedium,S<:AbstractSymmetry} <: AbstractSource{T,P}
    medium::P
    "Use: field(x,ω)"
    field::Function
    "Use: coefficients(n,x,ω)"
    coefficients::Function
    # Enforce that the Types are the same
    function RegularSource{T,P,S}(medium::P,field::Function,coefficients::Function) where {T,Dim,FieldDim,P<:PhysicalMedium{T,Dim,FieldDim},S<:AbstractSymmetry{Dim}}
        s = new{T,P,S}(medium,field,coefficients)
        self_test(s)
        return s
    end
end

function RegularSource(medium::P,field::Function,coefficients::Function; symmetry::AbstractSymmetry{Dim} = WithoutSymmetry{Dim}()) where {T,Dim,FieldDim,P<:PhysicalMedium{T,Dim,FieldDim}}
    return RegularSource{T,P,typeof(symmetry)}(medium,field,coefficients)
end

Symmetry(reg::RegularSource{T,P,S}) where {T,P,S} = S()

"""
    regular_spherical_coefficients(source::RegularSource)

return a function which can calculate the coefficients of a regular spherical wave basis.
"""
function regular_spherical_coefficients(source::RegularSource)
    source.coefficients
end

"""
Check that the source functions return the correct types
"""
function self_test(source::RegularSource{T,P}) where {P,T}
    ω = one(T)

    # choose rand postion, hopefully not the source position/origin
    x = SVector(rand(T.(0.1:0.1:1.0),spatial_dimension(P))...)

    # Check that the result of field has same dimension and type as PhysicalProperty field
    if field_dimension(P) == 1
        source.field(x,ω)::Union{Complex{T},SVector{field_dimension(P),Complex{T}}}
    else
        source.field(x,ω)::SVector{field_dimension(P),Complex{T}}
    end

    # Check that the result of field has the same dimension as field dimension of P
    if field_dimension(P) == 1
        source.coefficients(1,x,ω)::Union{Vector{Complex{T}},Vector{SVector{field_dimension(P),Complex{T}}}}
    # else # this else is not necessarily correct..
    #     source.coefficients(1,x,ω)::SVector{field_dimension(P),Complex{T}}
    end

    return true
end

field(s::RegularSource) = s.field
field(s::RegularSource{T},x::AbstractArray{T}, ω::T) where T = field(s)(x,ω)

function constant_source(medium::P, num::Complex{T} = zero(Float64) * im) where {P,T}
    return RegularSource{T,P}(medium, (x,ω) -> num, (order,x,ω) -> [num])
end

"""
    regular_spherical_source(PhysicalMedium,regular_coefficients; amplitude = one(T), position = zeros(T,Dim))

``c_n = ```regular_coefficients[n]`, ``x_o=```position`, and let ``v_n(kx)`` represent the regular spherical basis with wavenumber ``k`` at position ``x``. The above returns a `source` where `source.field(x,ω) =```\\sum_n c_n v_n(kx -k x_0)``
"""
function regular_spherical_source(medium::PhysicalMedium{T,Dim},regular_coefficients::AbstractVector{CT};
        symmetry::AbstractSymmetry{Dim} = WithoutSymmetry{Dim}(),
        amplitude::Union{T,Complex{T}} = one(T),
        position::AbstractArray{T} = SVector(zeros(T,Dim)...)) where {T,Dim,CT<:Union{T,Complex{T}}}

    coeff_order = basislength_to_basisorder(typeof(medium),length(regular_coefficients))

    function source_field(x,ω)
        vs = regular_basis_function(medium, ω)(coeff_order,x-position)
        return amplitude * sum(vs .* regular_coefficients)

    end

    function source_coef(order,centre,ω)
        k = ω / medium.c

        V = regular_translation_matrix(medium, max(order,coeff_order), ω, centre - position)

        len2 = basisorder_to_basislength(typeof(medium),order)
        len1 = length(regular_coefficients)

        return amplitude .* sum(regular_coefficients[n] .* V[n,1:len2] for n in 1:len1)
    end

    return RegularSource(medium, source_field, source_coef; symmetry = symmetry)
end

"""
    source_expand(RegularSource, centre; basis_order = 4)

Returns a function of `(x,ω)` which approximates the value of the source at `(x,ω)`. That is, the source is written in terms of a regular basis expansion centred at `centre`.
"""
function source_expand(source::AbstractSource{T}, centre::AbstractVector{T}; basis_order::Int = 4) where T

    # Convert to SVector for efficiency and consistency
    centre = SVector(centre...)

    return function (x::AbstractVector{T}, ω::T)
        vs = regular_basis_function(source.medium, ω)
        regular_coefficients = regular_spherical_coefficients(source)

        sum(regular_coefficients(basis_order,centre,ω) .* vs(basis_order, x - centre))
    end
end

import Base.(+)
function +(s1::RegularSource{T,P},s2::RegularSource{T,P})::RegularSource{T,P} where {P,T}
    if s1.medium != s2.medium
        error("Can not add sources from different physical mediums.")
    end
    field(x,ω) = s1.field(x,ω) + s2.field(x,ω)
    coef(n,centre,ω) = s1.coefficients(n,centre,ω) + s2.coefficients(n,centre,ω)

    sym1 = Symmetry(s1)
    sym2 = Symmetry(s2)
    S = typeof(Symmetry(sym1,sym2))

    RegularSource{T,P,S}(s1.medium,field,coef)
end

import Base.(*)
function *(a,s::RegularSource{T,P,S})::RegularSource{T,P,S} where {P,T,S}
    if typeof(one(Complex{T})*a) != Complex{T}
        error("Multiplying source by $a would cause source type to change, please explicitly cast $a to same type as RegularSource ($T)")
    end

    field(x,ω) = a * s.field(x,ω)
    coef(n,centre,ω) = a .* s.coefficients(n,centre,ω)

    RegularSource{T,P,S}(s.medium,field,coef)
end

*(s::RegularSource,a) = *(a,s)


# specialized methods
"""
    dotdot(::Vec, ::SymmetricFourthOrderTensor, ::Vec)

Computes a special dot product between two vectors and a symmetric fourth order tensor
such that ``a_k C_{ikjl} b_l``.
"""
@generated function dotdot(v1::Vec{dim}, S::SymmetricTensor{4, dim}, v2::Vec{dim}) where {dim}
    idx(i,j,k,l) = compute_index(SymmetricTensor{4, dim}, i, j, k, l)
    exps = Expr(:tuple)
    for j in 1:dim, i in 1:dim
        ex1, ex2 = Expr[],Expr[]
        for l in 1:dim, k in 1:dim
            push!(ex1, :(get_data(v1)[$k] * get_data(S)[$(idx(i,k,j,l))]))
            push!(ex2, :(get_data(v2)[$l]))
        end
        push!(exps.args, reducer(ex1, ex2, true))
    end
    return quote
        $(Expr(:meta, :inline))
        @inbounds return Tensor{2, dim}($exps)
    end
end