using Test
using StaticArrays: SVector
using QuantumLattices.Essentials.SpinPackage
using QuantumLattices.Essentials.Spatials: PID,Point,Bond
using QuantumLattices.Essentials.DegreesOfFreedom: Table,OID,isHermitian,IDFConfig,Operators,oidtype,otype,script,iid,latexformat
using QuantumLattices.Essentials.Terms: Couplings,@subscript,statistics,abbr
using QuantumLattices.Interfaces: dims,inds,expand,matrix,permute,rank
using QuantumLattices.Prerequisites: Float
using QuantumLattices.Mathematics.Combinatorics: Permutations
using QuantumLattices.Mathematics.AlgebraOverFields: ID
using QuantumLattices.Mathematics.VectorSpaces: IsMultiIndexable,MultiIndexOrderStyle

@testset "SID" begin
    @test SID|>fieldnames==(:orbital,:spin,:tag)
    @test SID(orbital=1,spin=0.5,tag='z')'==SID(orbital=1,spin=0.5,tag='z')
    @test SID(orbital=1,spin=0.5,tag='x')'==SID(orbital=1,spin=0.5,tag='x')
    @test SID(orbital=1,spin=0.5,tag='y')'==SID(orbital=1,spin=0.5,tag='y')
    @test SID(orbital=1,spin=0.5,tag='+')'==SID(orbital=1,spin=0.5,tag='-')
    @test SID(orbital=1,spin=0.5,tag='-')'==SID(orbital=1,spin=0.5,tag='+')
end

@testset "matrix" begin
    @test isapprox(matrix(SID(1,0.5,'z')),[[-0.5,0.0] [0.0,0.5]])
    @test isapprox(matrix(SID(1,0.5,'x')),[[0.0,0.5] [0.5,0.0]])
    @test isapprox(matrix(SID(1,0.5,'y')),[[0.0,-0.5im] [0.5im,0.0]])
    @test isapprox(matrix(SID(1,0.5,'+')),[[0.0,1.0] [0.0,0.0]])
    @test isapprox(matrix(SID(1,0.5,'-')),[[0.0,0.0] [1.0,0.0]])

    @test isapprox(matrix(SID(1,1.0,'z')),[[-1.0,0.0,0.0] [0.0,0.0,0.0] [0.0,0.0,1.0]])
    @test isapprox(matrix(SID(1,1.0,'x')),[[0.0,√2/2,0.0] [√2/2,0.0,√2/2] [0.0,√2/2,0.0]])
    @test isapprox(matrix(SID(1,1.0,'y')),[[0.0,-√2im/2,0.0] [√2im/2,0.0,-√2im/2] [0.0,√2im/2,0.0]])
    @test isapprox(matrix(SID(1,1.0,'+')),[[0.0,√2,0.0] [0.0,0.0,√2] [0.0,0.0,0.0]])
    @test isapprox(matrix(SID(1,1.0,'-')),[[0.0,0.0,0.0] [√2,0.0,0.0] [0.0,√2,0.0]])
end

@testset "Spin" begin
    spin=Spin(atom=1,norbital=2,spin=1.0)
    @test IsMultiIndexable(Spin)==IsMultiIndexable(true)
    @test MultiIndexOrderStyle(Spin)==MultiIndexOrderStyle('C')
    @test dims(spin)==(2,5)
    @test inds(SID(1,1.0,'z'),spin)==(1,3)
    @test SID((1,1),spin)==SID(1,1.0,'x')
    @test spin|>collect==[  SID(1,1.0,'x'),SID(1,1.0,'y'),SID(1,1.0,'z'),SID(1,1.0,'+'),SID(1,1.0,'-'),
                            SID(2,1.0,'x'),SID(2,1.0,'y'),SID(2,1.0,'z'),SID(2,1.0,'+'),SID(2,1.0,'-')
                            ]
end

@testset "SIndex" begin
    @test SIndex|>fieldnames==(:scope,:site,:orbital,:spin,:tag)
    @test union(PID{Char},SID)==SIndex{Char}
end

@testset "script" begin
    @test script(OID(SIndex('S',1,2,0.5,'z')),Val(:site))==1
    @test script(OID(SIndex('S',1,2,0.5,'z')),Val(:orbital))==2
    @test script(OID(SIndex('S',1,2,0.5,'z')),Val(:spin))==0.5
    @test script(OID(SIndex('S',1,2,0.5,'z')),Val(:tag))=='z'
end

@testset "oidtype" begin
    @test oidtype(SID,Point{2,PID{Int}},Nothing,Val(true))==OID{SIndex{Int},SVector{2,Float},SVector{2,Float},Nothing}
    @test oidtype(SID,Point{2,PID{Int}},Table,Val(true))==OID{SIndex{Int},SVector{2,Float},SVector{2,Float},Int}
    @test oidtype(SID,Point{2,PID{Int}},Nothing,Val(false))==OID{SIndex{Int},Nothing,Nothing,Nothing}
    @test oidtype(SID,Point{2,PID{Int}},Table,Val(false))==OID{SIndex{Int},Nothing,Nothing,Int}
end

@testset "SOperator" begin
    @test latexformat(SOperator)==soptdefaultlatex
    opt=SOperator(1.0,(SIndex('a',1,1,0.5,'+'),SIndex('a',1,1,0.5,'-')))
    @test opt|>statistics==opt|>typeof|>statistics=='B'
    @test opt'==SOperator(1.0,(SIndex('a',1,1,0.5,'+'),SIndex('a',1,1,0.5,'-')))
    @test isHermitian(opt)
    @test repr(opt)=="S^{+}_{1,1}S^{-}_{1,1}"
end

@testset "permute" begin
    soptrep(opt::SOperator)=opt.value*prod([matrix(opt.id[i].index|>iid) for i=1:rank(opt)])
    for S in (0.5,1.0,1.5)
        oids=[OID(SIndex('S',1,2,S,tag),seq=i) for (i,tag) in enumerate(('x','y','z','+','-'))]
        table=Dict(oid.index=>oid.seq for oid in oids)
        for (id1,id2) in Permutations{2}(oids)
            left=soptrep(SOperator(1,ID(id1,id2)))
            right=sum([soptrep(opt) for opt in permute(SOperator,id1,id2,table)])
            @test isapprox(left,right)
        end
    end
    id1=OID(SIndex('S',1,2,0.5,'z'))
    id2=OID(SIndex('S',2,2,0.5,'z'))
    @test permute(SOperator,id1,id2,nothing)==(SOperator(1,ID(id2,id1)),)
end

@testset "SCID" begin
    @test SCID(center=1,atom=1,orbital=1,tag='z')|>typeof|>fieldnames==(:center,:atom,:orbital,:tag,:subscript)
end

@testset "SpinCoupling" begin
    @test SpinCoupling{2}(1.0,tags=('+','-'))|>string=="SpinCoupling(value=1.0,tags=(+,-))"
    @test SpinCoupling{2}(1.0,atoms=(1,1),tags=('z','z'))|>string=="SpinCoupling(value=1.0,atoms=(1,1),tags=(z,z))"
    @test SpinCoupling{2}(1.0,atoms=(1,1),orbitals=(1,2),tags=('-','+'))|>string=="SpinCoupling(value=1.0,atoms=(1,1),orbitals=(1,2),tags=(-,+))"
    @test SpinCoupling{2}(2.0,tags=('x','y'))|>repr=="2.0 SxSy"

    sc1=SpinCoupling{2}(1.5,tags=('+','-'),atoms=(1,2),orbitals=(@subscript (x,y) with x>y),centers=(1,2))
    sc2=SpinCoupling{2}(2.0,tags=('+','-'),atoms=(1,2),orbitals=(@subscript (x,y) with x<y),centers=(1,2))
    @test sc1|>repr=="1.5 S+S- sl(1,2)⊗ob(x,y)@(1,2) with $(sc1.id[1].subscript) && $(sc1.id[2].subscript)"
    @test sc2|>repr=="2.0 S+S- sl(1,2)⊗ob(x,y)@(1,2) with $(sc2.id[1].subscript) && $(sc2.id[2].subscript)"

    sc=sc1*sc2
    @test sc|>repr=="3.0 S+S-S+S- sl(1,2,1,2)⊗ob(x,y,x,y)@(1,2,1,2) with $(sc1.id[1].subscript) && $(sc1.id[2].subscript) && $(sc2.id[1].subscript) && $(sc2.id[2].subscript)"

    ex=expand(SpinCoupling{2}(2.0,tags=('+','-'),atoms=(1,1)),PID(1,1),Spin(atom=2,norbital=2,spin=1.0))
    @test collect(ex)==[]

    ex=expand(SpinCoupling{2}(2.0,tags=('+','-'),atoms=(1,2),orbitals=(1,2)),(PID(1,1),PID(1,2)),(Spin(atom=1,norbital=2,spin=1.0),Spin(atom=2,norbital=2,spin=1.0)))
    @test IsMultiIndexable(typeof(ex))==IsMultiIndexable(true)
    @test MultiIndexOrderStyle(typeof(ex))==MultiIndexOrderStyle('C')
    @test dims(ex)==(1,)
    @test collect(ex)==[(2.0,(SIndex(1,1,1,1.0,'+'),SIndex(1,2,2,1.0,'-')))]

    ex=expand(SpinCoupling{4}(2.0,tags=('+','-','+','-'),orbitals=(@subscript (α,α,β,β) with α<β)),PID(1,1),Spin(norbital=3,spin=1.0))
    @test dims(ex)==(3,)
    @test collect(ex)==[(2.0,(SIndex(1,1,1,1.0,'+'),SIndex(1,1,1,1.0,'-'),SIndex(1,1,2,1.0,'+'),SIndex(1,1,2,1.0,'-'))),
                        (2.0,(SIndex(1,1,1,1.0,'+'),SIndex(1,1,1,1.0,'-'),SIndex(1,1,3,1.0,'+'),SIndex(1,1,3,1.0,'-'))),
                        (2.0,(SIndex(1,1,2,1.0,'+'),SIndex(1,1,2,1.0,'-'),SIndex(1,1,3,1.0,'+'),SIndex(1,1,3,1.0,'-')))
    ]
end

@testset "Heisenberg" begin
    @test Heisenberg(orbitals=(1,2))==Couplings(SpinCoupling{2}(1//1,tags=('z','z'),orbitals=(1,2)),
                                                SpinCoupling{2}(1//2,tags=('+','-'),orbitals=(1,2)),
                                                SpinCoupling{2}(1//2,tags=('-','+'),orbitals=(1,2))
                                                )
    @test Heisenberg("xyz")==Couplings( SpinCoupling{2}(1,tags=('x','x')),
                                        SpinCoupling{2}(1,tags=('y','y')),
                                        SpinCoupling{2}(1,tags=('z','z'))
                                        )
end

@testset "Ising" begin
    @test Ising('x',atoms=(1,2))==Couplings(SpinCoupling{2}(1,tags=('x','x'),atoms=(1,2)))
    @test Ising('y',atoms=(1,2))==Couplings(SpinCoupling{2}(1,tags=('y','y'),atoms=(1,2)))
    @test Ising('z',atoms=(1,2))==Couplings(SpinCoupling{2}(1,tags=('z','z'),atoms=(1,2)))
end

@testset "Gamma" begin
    @test Gamma('x',orbitals=(1,1))==SpinCoupling{2}(1,tags=('y','z'),orbitals=(1,1))+SpinCoupling{2}(1,tags=('z','y'),orbitals=(1,1))
    @test Gamma('y',atoms=(1,2))==SpinCoupling{2}(1,tags=('z','x'),atoms=(1,2))+SpinCoupling{2}(1,tags=('x','z'),atoms=(1,2))
    @test Gamma('z')==SpinCoupling{2}(1,tags=('x','y'))+SpinCoupling{2}(1,tags=('y','x'))
end

@testset "DM" begin
    @test DM('x',orbitals=(1,1))==SpinCoupling{2}(1,tags=('y','z'),orbitals=(1,1))-SpinCoupling{2}(1,tags=('z','y'),orbitals=(1,1))
    @test DM('y',atoms=(1,2))==SpinCoupling{2}(1,tags=('z','x'),atoms=(1,2))-SpinCoupling{2}(1,tags=('x','z'),atoms=(1,2))
    @test DM('z')==SpinCoupling{2}(1,tags=('x','y'))-SpinCoupling{2}(1,tags=('y','x'))
end

@testset "Sᵅ" begin
    @test Sˣ(atom=1,orbital=1)==Couplings(SpinCoupling{1}(1,tags=('x',),atoms=(1,),orbitals=(1,)))
    @test Sʸ(atom=1)==Couplings(SpinCoupling{1}(1,tags=('y',),atoms=(1,)))
    @test Sᶻ(orbital=1)==Couplings(SpinCoupling{1}(1,tags=('z',),orbitals=(1,)))
end

@testset "spincoupling" begin
    sc=sc"1.0 S+S- sl(1,1)⊗ob(α,β)@(1,2) with α<β"
    ob=sc.subscripts[1].identifier
    @test repr(sc)=="1.0 S+S- sl(1,1)⊗ob(α,β)@(1,2) with $ob && $ob"

    sc=sc"1.0 S+S- sl(1,1)⊗ob(α,β)@(1,2)"
    ob=sc.subscripts[1].identifier
    @test repr(sc)=="1.0 S+S- sl(1,1)⊗ob(α,β)@(1,2) with $ob && $ob"

    sc=sc"1.0 S+S- sl(1,1)⊗ob(α,β)"
    ob=sc.subscripts[1].identifier
    @test repr(sc)=="1.0 S+S- sl(1,1)⊗ob(α,β) with $ob && $ob"

    sc=sc"1.0 S+S- sl(1,1)⊗ob(1,2)@(1,2)"
    @test repr(sc)=="1.0 S+S- sl(1,1)⊗ob(1,2)@(1,2)"

    sc=sc"1.0 S+S- sl(1,1)@(1,2)"
    @test repr(sc)=="1.0 S+S- sl(1,1)@(1,2)"

    sc=sc"1.0 S+S- @(1,2)"
    @test repr(sc)=="1.0 S+S- @(1,2)"

    sc=sc"1.0 S+S-"
    @test repr(sc)=="1.0 S+S-"
end

@testset "spincouplings" begin
    @test heisenberg"sl(1,1)⊗ob(1,3)@(1,2)"==Heisenberg(centers=(1,2),atoms=(1,1),orbitals=(1,3))
    @test heisenberg"sl(1,1)⊗ob(1,3)"==Heisenberg(atoms=(1,1),orbitals=(1,3))
    @test heisenberg"@(1,2)"==Heisenberg(centers=(1,2))
    @test heisenberg"ob(1,3)"==Heisenberg(orbitals=(1,3))
    @test heisenberg"sl(1,1)"==Heisenberg(atoms=(1,1))
    @test heisenberg""==heisenberg"+-z"==Heisenberg()
    @test heisenberg"xyz"==Heisenberg("xyz")

    @test ising"x"==Ising('x') && ising"x sl(1,1)⊗ob(1,3)@(1,2)"==Ising('x',centers=(1,2),atoms=(1,1),orbitals=(1,3))
    @test ising"y"==Ising('y') && ising"y sl(1,1)⊗ob(1,3)@(1,2)"==Ising('y',centers=(1,2),atoms=(1,1),orbitals=(1,3))
    @test ising"z"==Ising('z') && ising"z sl(1,1)⊗ob(1,3)@(1,2)"==Ising('z',centers=(1,2),atoms=(1,1),orbitals=(1,3))

    @test gamma"x"==Gamma('x') && gamma"x sl(1,1)⊗ob(1,3)@(1,2)"==Gamma('x',centers=(1,2),atoms=(1,1),orbitals=(1,3))
    @test gamma"y"==Gamma('y') && gamma"y sl(1,1)⊗ob(1,3)@(1,2)"==Gamma('y',centers=(1,2),atoms=(1,1),orbitals=(1,3))
    @test gamma"z"==Gamma('z') && gamma"z sl(1,1)⊗ob(1,3)@(1,2)"==Gamma('z',centers=(1,2),atoms=(1,1),orbitals=(1,3))

    @test dm"x"==DM('x') && dm"x sl(1,1)⊗ob(1,3)@(1,2)"==DM('x',centers=(1,2),atoms=(1,1),orbitals=(1,3))
    @test dm"y"==DM('y') && dm"y sl(1,1)⊗ob(1,3)@(1,2)"==DM('y',centers=(1,2),atoms=(1,1),orbitals=(1,3))
    @test dm"z"==DM('z') && dm"z sl(1,1)⊗ob(1,3)@(1,2)"==DM('z',centers=(1,2),atoms=(1,1),orbitals=(1,3))

    @test sˣ""==Sˣ() && sˣ"sl(1)⊗ob(2)"==Sˣ(atom=1,orbital=2)
    @test sʸ""==Sʸ() && sʸ"sl(1)⊗ob(2)"==Sʸ(atom=1,orbital=2)
    @test sᶻ""==Sᶻ() && sᶻ"sl(1)⊗ob(2)"==Sᶻ(atom=1,orbital=2)
end

@testset "SpinTerm" begin
    term=SpinTerm{1}(:h,1.5,0,couplings=Sᶻ())
    @test term|>abbr==:sp
    @test otype(term|>typeof,OID{SIndex{Int},SVector{2,Float},SVector{2,Float},Nothing})==SOperator{Float,ID{OID{SIndex{Int},SVector{2,Float},SVector{2,Float},Nothing},1}}
    @test otype(term|>typeof,OID{SIndex{Int},SVector{2,Float},SVector{2,Float},Int})==SOperator{Float,ID{OID{SIndex{Int},SVector{2,Float},SVector{2,Float},Int},1}}

    point=Point(PID('a',1),(0.5,0.5),(0.0,0.0))
    config=IDFConfig{Spin}(pid->Spin(atom=pid.site%2,norbital=2,spin=0.5),[point.pid])
    table=Table(config,usualspinindextotuple)
    term=SpinTerm{1}(:h,1.5,0,couplings=Sᶻ())
    operators=Operators(SOperator(1.5,ID(OID(SIndex('a',1,1,0.5,'z'),[0.5,0.5],[0.0,0.0],1))),
                        SOperator(1.5,ID(OID(SIndex('a',1,2,0.5,'z'),[0.5,0.5],[0.0,0.0],2)))
    )
    @test expand(term,point,config,table)==operators

    bond=Bond(1,Point(PID('a',1),(0.0,0.0),(0.0,0.0)),Point(PID('b',1),(0.5,0.5),(0.0,0.0)))
    config=IDFConfig{Spin}(pid->Spin(atom=pid.site%2,norbital=2,spin=0.5),[bond.spoint.pid,bond.epoint.pid])
    table=Table(config,usualspinindextotuple)
    term=SpinTerm{2}(:J,1.5,1,couplings=Heisenberg())
    operators=Operators(SOperator(1.50,ID(OID(SIndex('b',1,2,0.5,'z'),[0.5,0.5],[0.0,0.0],4),OID(SIndex('a',1,2,0.5,'z'),[0.0,0.0],[0.0,0.0],2))),
                        SOperator(0.75,ID(OID(SIndex('b',1,2,0.5,'-'),[0.5,0.5],[0.0,0.0],4),OID(SIndex('a',1,2,0.5,'+'),[0.0,0.0],[0.0,0.0],2))),
                        SOperator(0.75,ID(OID(SIndex('b',1,1,0.5,'-'),[0.5,0.5],[0.0,0.0],3),OID(SIndex('a',1,1,0.5,'+'),[0.0,0.0],[0.0,0.0],1))),
                        SOperator(0.75,ID(OID(SIndex('b',1,1,0.5,'+'),[0.5,0.5],[0.0,0.0],3),OID(SIndex('a',1,1,0.5,'-'),[0.0,0.0],[0.0,0.0],1))),
                        SOperator(1.50,ID(OID(SIndex('b',1,1,0.5,'z'),[0.5,0.5],[0.0,0.0],3),OID(SIndex('a',1,1,0.5,'z'),[0.0,0.0],[0.0,0.0],1))),
                        SOperator(0.75,ID(OID(SIndex('b',1,2,0.5,'+'),[0.5,0.5],[0.0,0.0],4),OID(SIndex('a',1,2,0.5,'-'),[0.0,0.0],[0.0,0.0],2)))
    )
    @test expand(term,bond,config,table)==operators
end
