using SecondQuantizedAlgebra
using QuantumOpticsBase
using Latexify
using Symbolics: @variables, Num
using Test
import SecondQuantizedAlgebra: simplify, QAdd, QSym, HilbertSpace

@testset "phase_space" begin
    @testset "Construction — product space" begin
        h = FockSpace(:c) ⊗ PhaseSpace(:q)
        x = Position(h, :x, 2)
        p = Momentum(h, :p, 2)
        @test x.space_index == 2
        @test p.space_index == 2
        @test_throws ArgumentError Position(h, :x, 1)  # FockSpace, not PhaseSpace
        @test_throws ArgumentError Momentum(h, :p, 1)
        @test_throws ArgumentError Position(h, :x, 3)  # out of range
    end

    @testset "Adjoint — Hermitian" begin
        h = PhaseSpace(:q)
        x = Position(h, :x)
        p = Momentum(h, :p)
        @test x' == x
        @test p' == p
    end

    @testset "Lazy multiplication" begin
        h = PhaseSpace(:q)
        x = Position(h, :x)
        p = Momentum(h, :p)

        m = x * p
        @test m isa QAdd

        s = x + p
        @test s isa QAdd
    end

    @testset "Simplify: [X, P] = i" begin
        h = PhaseSpace(:q)
        x = Position(h, :x)
        p = Momentum(h, :p)

        # p*x out of order → x*p - i
        result = simplify(p * x)
        @test result isa QAdd
        @test length(result) == 2  # x·p + (-i)

        # x*p already ordered → x·p (1 term)
        result2 = simplify(x * p)
        @test result2 isa QAdd
        @test length(result2) == 1
    end

    @testset "Simplify: mixed spaces don't interact" begin
        h = PhaseSpace(:q1) ⊗ PhaseSpace(:q2)
        x1 = Position(h, :x, 1)
        p2 = Momentum(h, :p, 2)

        # Different spaces — no commutation
        result = simplify(p2 * x1)
        @test length(result) == 1
    end

    @testset "Higher-order algebra" begin
        h = PhaseSpace(:motion)
        x = Position(h, :x)
        p = Momentum(h, :p)

        @test isequal(
            simplify(normal_order(p * x * x)), simplify(x * x * p - 2im * x)
        )
    end

    @testset "Multi-space algebra" begin
        hps1 = PhaseSpace(:motion1)
        hps2 = PhaseSpace(:motion2)
        h = hps1 ⊗ hps2

        x1 = Position(h, :x_1, 1)
        p1 = Momentum(h, :p_1, 1)
        x2 = Position(h, :x_2, 2)

        @test isequal(p1 * x2, x2 * p1)
        @test isequal(
            simplify(normal_order(p1 * x2 * x1)),
            simplify((x1 * p1 - 1im) * x2),
        )
    end

    @testset "Hamiltonian construction and commutators" begin
        hps1 = PhaseSpace(:motion1)
        hps2 = PhaseSpace(:motion2)
        h = hps1 ⊗ hps2

        x1 = Position(h, :x_1, 1)
        p1 = Momentum(h, :p_1, 1)
        x2 = Position(h, :x_2, 2)
        p2 = Momentum(h, :p_2, 2)

        @variables ω::Real m::Real miv::Real c::Real
        H = 0.5 * Num(miv) * p1^2 + 0.5Num(m) * Num(ω)^2 * x1^2 + Num(c) * x1^4

        @test iszero(simplify(commutator(H, x2)))
        @test iszero(simplify(commutator(H, p2)))
        @test isequal(simplify(1im * commutator(H, x1)), simplify(p1 * Num(miv)))
        @test isequal(
            simplify(1im * commutator(H, p1)),
            simplify(-x1 * Num(m) * Num(ω)^2 - 4 * Num(c) * x1^3),
        )
    end

    @testset "Numeric conversion" begin
        h = PhaseSpace(:q)
        x = Position(h, :x)
        p = Momentum(h, :p)
        b = FockBasis(10)

        x_num = to_numeric(x, b)
        p_num = to_numeric(p, b)
        @test x_num ≈ (destroy(b) + create(b)) / sqrt(2)
        @test p_num ≈ im * (create(b) - destroy(b)) / sqrt(2)
    end

    @testset "Numeric: commutator [X, P] = i" begin
        h = PhaseSpace(:q)
        x = Position(h, :x)
        p = Momentum(h, :p)
        b = FockBasis(30)

        x_num = to_numeric(x, b)
        p_num = to_numeric(p, b)
        comm = x_num * p_num - p_num * x_num
        # Should be approximately i*I (exact in infinite basis, approximate in truncated)
        @test comm.data[1, 1] ≈ im atol = 1.0e-10
    end

    @testset "Numeric: composite basis" begin
        h = FockSpace(:c) ⊗ PhaseSpace(:q)
        x = Position(h, :x, 2)
        bf = FockBasis(3)
        bq = FockBasis(5)
        bc = bf ⊗ bq
        @test to_numeric(x, bc; op_type = identity) isa LazySum
        @test to_numeric(x, bc) isa Operator
    end

    @testset "Printing" begin
        h = PhaseSpace(:q)
        @test sprint(show, h) == "ℋ(q)"

        x = Position(h, :x)
        p = Momentum(h, :p)
        @test sprint(show, x) == "x"
        @test sprint(show, p) == "p"
    end

    @testset "LaTeX" begin
        h = PhaseSpace(:q)
        x = Position(h, :x)
        p = Momentum(h, :p)

        sx = latexify(x)
        sp = latexify(p)
        @test occursin("hat", sx)
        @test occursin("x", sx)
        @test occursin("hat", sp)
        @test occursin("p", sp)

        # MIME dispatch
        @test repr(MIME"text/latex"(), x) == sx
        @test repr(MIME"text/latex"(), p) == sp
    end

    @testset "@qnumbers" begin
        h = PhaseSpace(:q)
        @qnumbers x::Position(h)
        @qnumbers p::Momentum(h)
        @test is_position(x)
        @test operator_name(x) == :x
        @test is_momentum(p)
        @test operator_name(p) == :p
    end

    @testset "Type stability" begin
        h = PhaseSpace(:q)
        x = Position(h, :x)
        p = Momentum(h, :p)

        @inferred Position(:x, 1)
        @inferred Momentum(:p, 1)
        @inferred adjoint(x)
        @inferred adjoint(p)
        @inferred isequal(x, x)
        @inferred isequal(p, p)
        @inferred hash(x, UInt(0))
        @inferred hash(p, UInt(0))
        @inferred x * p
        @inferred p * x
        @inferred x + p
        @inferred commutator(x, p)
    end

    @static if VERSION >= v"1.12"
        @testset "Allocations" begin
            h = PhaseSpace(:q)
            x = Position(h, :x)
            p = Momentum(h, :p)
            x2 = Position(h, :x)
            p2 = Momentum(h, :p)

            @test @allocations(Position(:x, 1)) == 0
            @test @allocations(Momentum(:p, 1)) == 0
            @test @allocations(adjoint(x)) == 0
            @test @allocations(adjoint(p)) == 0
            @test @allocations(isequal(x, x2)) == 0
            @test @allocations(isequal(p, p2)) == 0
            @test @allocations(hash(x, UInt(0))) == 0
            @test @allocations(hash(p, UInt(0))) == 0
        end
    end
end
