##
# 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) => SymbolicUtils.term(sqrt, (~k).idx+1)*FockState((~k).idx+1)),
    @rule(~o::_isa(DestroyOp) * ~k::_isa(FockState) => SymbolicUtils.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_SIMPLIFY = [RULES_PAULI; RULES_COMMUTATOR; RULES_ANTICOMMUTATOR; RULES_FOCK]

##
# Custom tree walker for QuantumSymbolics types
# SymbolicUtils v4's Prewalk/Postwalk only walk BasicSymbolic trees.
# Our types are not BasicSymbolic, so we need our own walker.
##

"""Apply a rewriter in pre-order to a QuantumSymbolics expression tree."""
function _prewalk(rw, x)
    y = rw(x)
    if !isnothing(y)
        x = y
    end
    if isexpr(x) && iscall(x)
        args = arguments(x)
        newargs = map(a -> _prewalk(rw, a), args)
        # Always reconstruct through maketerm to ensure proper normalization
        # (e.g., flattening of associative operations like * and ⊗).
        return maketerm(typeof(x), operation(x), newargs, metadata(x))
    end
    return x
end

"""Apply a rewriter repeatedly in pre-order until no changes occur."""
function _fixpoint_prewalk(rw, x)
    while true
        y = _prewalk(rw, x)
        isequal(x, y) && return x
        x = y
    end
end

"""Recursively substitute subexpressions in a QuantumSymbolics expression tree."""
function _substitute(expr, dict)
    haskey(dict, expr) && return dict[expr]
    if isexpr(expr) && iscall(expr)
        args = arguments(expr)
        newargs = map(a -> _substitute(a, dict), args)
        return maketerm(typeof(expr), operation(expr), newargs, metadata(expr))
    end
    return expr
end

# Extend SymbolicUtils.substitute for our Symbolic types so that
# `Symbolics.substitute` (which re-exports it) works on quantum expressions.
SymbolicUtils.substitute(expr::Symbolic, dict; kwargs...) = _substitute(expr, dict)

##
# Simplification rewriters
##

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

"""
    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