rules.jl

##
# This file defines manual simplification and expansion rules for specific operations of quantum objects.
##

##
# Predicate functions
##
function hasscalings(xs)
    any(xs) do x
        operation(x) == *
    end
end
_isa(T) = x->isa(x,T)
_vecisa(T) = x->all(_isa(T), x)

##
# Simplification rules
## 

# Pauli identities
RULES_PAULI = [
    @rule(~o1::_isa(XGate)*~o2::_isa(XGate) => I),
    @rule(~o1::_isa(YGate)*~o2::_isa(YGate) => I),
    @rule(~o1::_isa(ZGate)*~o2::_isa(ZGate) => I),
    @rule(~o1::_isa(XGate)*~o2::_isa(YGate) => im*Z),
    @rule(~o1::_isa(YGate)*~o2::_isa(ZGate) => im*X),
    @rule(~o1::_isa(ZGate)*~o2::_isa(XGate) => im*Y),
    @rule(~o1::_isa(YGate)*~o2::_isa(XGate) => -im*Z),
    @rule(~o1::_isa(ZGate)*~o2::_isa(YGate) => -im*X),
    @rule(~o1::_isa(XGate)*~o2::_isa(ZGate) => -im*Y),
    @rule(~o1::_isa(HGate)*~o2::_isa(XGate)*~o3::_isa(HGate) => Z),
    @rule(~o1::_isa(HGate)*~o2::_isa(YGate)*~o3::_isa(HGate) => -Y),
    @rule(~o1::_isa(HGate)*~o2::_isa(ZGate)*~o3::_isa(HGate) => X),

    @rule(~o::_isa(XGate)*~k::isequal(X1) => X1),
    @rule(~o::_isa(YGate)*~k::isequal(X1) => -im*X2),
    @rule(~o::_isa(ZGate)*~k::isequal(X1) => X2),

    @rule(~o::_isa(XGate)*~k::isequal(X2) => -X2),
    @rule(~o::_isa(YGate)*~k::isequal(X2) => im*X1),
    @rule(~o::_isa(ZGate)*~k::isequal(X2) => X1),

    @rule(~o::_isa(XGate)*~k::isequal(Y1) => im*Y2),
    @rule(~o::_isa(YGate)*~k::isequal(Y1) => Y1),
    @rule(~o::_isa(ZGate)*~k::isequal(Y1) => Y2),

    @rule(~o::_isa(XGate)*~k::isequal(Y2) => -im*Y1),
    @rule(~o::_isa(YGate)*~k::isequal(Y2) => -Y2),
    @rule(~o::_isa(ZGate)*~k::isequal(Y2) => Y1),

    @rule(~o::_isa(XGate)*~k::isequal(Z1) => Z2),
    @rule(~o::_isa(YGate)*~k::isequal(Z1) => im*Z2),
    @rule(~o::_isa(ZGate)*~k::isequal(Z1) => Z1),

    @rule(~o::_isa(XGate)*~k::isequal(Z2) => Z1),
    @rule(~o::_isa(YGate)*~k::isequal(Z2) => -im*Z1),
    @rule(~o::_isa(ZGate)*~k::isequal(Z2) => -Z2),

    @rule(~o::_isa(HGate)*~k::isequal(X1) => Z1),
    @rule(~o::_isa(HGate)*~k::isequal(X2) => Z2),
    @rule(~o::_isa(HGate)*~k::isequal(Y1) => (X1+im*X2)/sqrt(2)),
    @rule(~o::_isa(HGate)*~k::isequal(Y2) => (X1-im*X2)/sqrt(2)),
    @rule(~o::_isa(HGate)*~k::isequal(Z1) => X1),
    @rule(~o::_isa(HGate)*~k::isequal(Z2) => X2)
]

# Commutator identities
RULES_COMMUTATOR = [
    @rule(commutator(~o1::_isa(XGate), ~o2::_isa(YGate)) => 2*im*Z),
    @rule(commutator(~o1::_isa(YGate), ~o2::_isa(ZGate)) => 2*im*X),
    @rule(commutator(~o1::_isa(ZGate), ~o2::_isa(XGate)) => 2*im*Y),
    @rule(commutator(~o1::_isa(YGate), ~o2::_isa(XGate)) => -2*im*Z),
    @rule(commutator(~o1::_isa(ZGate), ~o2::_isa(YGate)) => -2*im*X),
    @rule(commutator(~o1::_isa(XGate), ~o2::_isa(ZGate)) => -2*im*Y)
]

# Anticommutator identities
RULES_ANTICOMMUTATOR = [
    @rule(anticommutator(~o1::_isa(XGate), ~o2::_isa(XGate)) => 2*I),
    @rule(anticommutator(~o1::_isa(YGate), ~o2::_isa(YGate)) => 2*I),
    @rule(anticommutator(~o1::_isa(ZGate), ~o2::_isa(ZGate)) => 2*I),
    @rule(anticommutator(~o1::_isa(XGate), ~o2::_isa(YGate))=> 0),
    @rule(anticommutator(~o1::_isa(YGate), ~o2::_isa(ZGate)) => 0),
    @rule(anticommutator(~o1::_isa(ZGate), ~o2::_isa(XGate)) => 0),
    @rule(anticommutator(~o1::_isa(YGate), ~o2::_isa(XGate)) => 0),
    @rule(anticommutator(~o1::_isa(ZGate), ~o2::_isa(YGate)) => 0),
    @rule(anticommutator(~o1::_isa(XGate), ~o2::_isa(ZGate)) => 0),
    @rule(commutator(~o1::_isa(DestroyOp), ~o2::_isa(CreateOp)) => IdentityOp(basis(~o1))),
    @rule(commutator(~o1::_isa(CreateOp), ~o2::_isa(DestroyOp)) => -IdentityOp(basis(~o1))),
    @rule(commutator(~o1::_isa(NumberOp), ~o2::_isa(DestroyOp)) => -(~o2)),
    @rule(commutator(~o1::_isa(DestroyOp), ~o2::_isa(NumberOp)) => (~o1)),
    @rule(commutator(~o1::_isa(NumberOp), ~o2::_isa(CreateOp)) => (~o2)),
    @rule(commutator(~o1::_isa(CreateOp), ~o2::_isa(NumberOp)) => -(~o1))
]

RULES_FOCK = [
    @rule(~o::_isa(DestroyOp) * ~k::isequal(vac) => SZeroKet()),
    @rule(~o::_isa(CreateOp) * ~k::_isa(FockState) => Term(sqrt,[(~k).idx+1])*FockState((~k).idx+1)),
    @rule(~o::_isa(DestroyOp) * ~k::_isa(FockState) => Term(sqrt,[(~k).idx])*FockState((~k).idx-1)),
    @rule(~o::_isa(NumberOp) * ~k::_isa(FockState) => (~k).idx*(~k)),
    @rule(~o::_isa(DestroyOp) * ~k::_isa(CoherentState) => (~k).alpha*(~k)),
    @rule(~o::_isa(PhaseShiftOp) * ~k::_isa(CoherentState) => CoherentState((~k).alpha * exp(-im*(~o).phase))),
    @rule(dagger(~o1::_isa(PhaseShiftOp)) * ~o2::_isa(DestroyOp) * ~o1 => ~o2*exp(-im*((~o1).phase))),
    @rule(~o1::_isa(PhaseShiftOp) * ~o2::_isa(DestroyOp) * dagger(~o1) => ~o2*exp(im*((~o1).phase))),
    @rule(dagger(~o1::_isa(DisplaceOp)) * ~o2::_isa(DestroyOp) * ~o1 => (~o2) + (~o1).alpha*IdentityOp(basis(~o2))),
    @rule(dagger(~o1::_isa(DisplaceOp)) * ~o2::_isa(CreateOp) * ~o1 => (~o2) + conj((~o1).alpha)*IdentityOp(basis(~o2))),
    @rule(~o::_isa(DisplaceOp) * ~k::((x->(isa(x,FockState) && x.idx == 0))) => CoherentState((~o).alpha)),
    @rule(~o::_isa(SqueezeOp) * ~k::isequal(vac) => SqueezedState((~o).z)),
    @rule(~o::_isa(TwoSqueezeOp) * ~k::isequal(vac ⊗ vac) => TwoSqueezedState((~o).z))
]

RULES_ROT = [
    @rule(~r::_isa(RotXGate) => (-im*X) where (~r).θ in [π, 1π]),
    @rule(~r::_isa(RotYGate) => (-im*Y) where (~r).θ in [π, 1π]),
    @rule(~r::_isa(RotZGate) => (-im*Z) where (~r).θ in [π, 1π]),
    @rule(~r::_isa(RotXGate) => try Rx(mod((~r).θ, 4π)) catch end),
    @rule(~r::_isa(RotYGate) => try Ry(mod((~r).θ, 4π)) catch end),
    @rule(~r::_isa(RotZGate) => try Rz(mod((~r).θ, 4π)) catch end),
    @rule(~r::_isa(RotXGate) => try (~r).θ ≥ 2π ? -Rx((~r).θ - 2π) : nothing catch end),
    @rule(~r::_isa(RotYGate) => try (~r).θ ≥ 2π ? -Ry((~r).θ - 2π) : nothing catch end),
    @rule(~r::_isa(RotZGate) => try (~r).θ ≥ 2π ? -Rz((~r).θ - 2π) : nothing catch end),
    @rule(~r1::_isa(RotXGate) * ~r2::_isa(RotXGate) => try Rx((~r1).θ + (~r2).θ) catch end),
    @rule(~r1::_isa(RotYGate) * ~r2::_isa(RotYGate) => try Ry((~r1).θ + (~r2).θ) catch end),
    @rule(~r1::_isa(RotZGate) * ~r2::_isa(RotZGate) => try Rz((~r1).θ + (~r2).θ) catch end),
    @rule(exp(~α * ~x::_isa(XGate)) => try if real(~α) == 0 Rx(-2imag(~α)) end catch end),
    @rule(exp(~α * ~x::_isa(YGate)) => try if real(~α) == 0 Ry(-2imag(~α)) end catch end),
    @rule(exp(~α * ~x::_isa(ZGate)) => try if real(~α) == 0 Rz(-2imag(~α)) end catch end)
]

RULES_SIMPLIFY = [RULES_PAULI; RULES_COMMUTATOR; RULES_ANTICOMMUTATOR; RULES_FOCK; RULES_ROT]

##
# Simplification rewriters
##

qsimplify_pauli = Chain(RULES_PAULI)
qsimplify_commutator = Chain(RULES_COMMUTATOR)
qsimplify_anticommutator = Chain(RULES_ANTICOMMUTATOR)
qsimplify_fock = Chain(RULES_FOCK)
qsimplify_rot = Chain(RULES_ROT)

"""
    qsimplify(s; rewriter=nothing)

Manually simplify a symbolic expression of quantum objects. 

If the keyword `rewriter` is not specified, then `qsimplify` will apply every defined rule to the expression. 
For performance or single-purpose motivations, the user has the option to define a specific rewriter for `qsimplify` to apply to the expression.
The defined rewriters for simplification are the following objects:
    - `qsimplify_pauli`
    - `qsimplify_commutator`
    - `qsimplify_anticommutator`
    - `qsimplify_fock`

```jldoctest
julia> qsimplify(σʸ*commutator(σˣ*σᶻ, σᶻ))
(0 - 2im)Z

julia> qsimplify(anticommutator(σˣ, σˣ), rewriter=qsimplify_anticommutator)
2𝕀
```
"""
function qsimplify(s; rewriter=nothing)
    if QuantumSymbolics.isexpr(s)
        if isnothing(rewriter)
            Fixpoint(Prewalk(Chain(RULES_SIMPLIFY)))(s)
        else
            Fixpoint(Prewalk(rewriter))(s)
        end
    else
        error("Object $(s) of type $(typeof(s)) is not an expression.")
    end
end

##
# Expansion rules
## 

RULES_EXPAND = [
    @rule(commutator(~o1, ~o2) => (~o1)*(~o2) - (~o2)*(~o1)),
    @rule(anticommutator(~o1, ~o2) => (~o1)*(~o2) + (~o2)*(~o1)),
    @rule(~o1 ⊗ +(~~ops) => +(map(op -> ~o1 ⊗ op, ~~ops)...)),
    @rule(+(~~ops) ⊗ ~o1 => +(map(op -> op ⊗ ~o1, ~~ops)...)),
    @rule(~o1 * +(~~ops) => +(map(op -> ~o1 * op, ~~ops)...)),
    @rule(+(~~ops) * ~o1 => +(map(op -> op * ~o1, ~~ops)...)),
    @rule(⊗(~~ops1::_vecisa(Symbolic{AbstractBra})) * ⊗(~~ops2::_vecisa(Symbolic{AbstractKet})) => *(map(*, ~~ops1, ~~ops2)...)),
    @rule(⊗(~~ops1::_vecisa(Symbolic{AbstractOperator})) * ⊗(~~ops2::_vecisa(Symbolic{AbstractOperator})) => ⊗(map(*, ~~ops1, ~~ops2)...)),
]

# 

##
# Expansion rewriter
##

"""
    qexpand(s)

Manually expand a symbolic expression of quantum objects. 

```jldoctest
julia> @op A; @op B; @op C;

julia> qexpand(commutator(A, B))
-1BA+AB

julia> qexpand(A⊗(B+C))
(A⊗B)+(A⊗C)

julia> @ket k₁; @ket k₂;

julia> qexpand(A*(k₁+k₂))
A|k₁⟩+A|k₂⟩
```
"""
function qexpand(s)
    if QuantumSymbolics.isexpr(s)
        Fixpoint(Prewalk(Chain(RULES_EXPAND)))(s)
    else
        error("Object $(s) of type $(typeof(s)) is not an expression.")
    end
end