#
# Copyright 2022 Keysight Technologies Inc.
#

"""
Defining Custom Compilers
=========================
"""

#%%
# Building a general compiler which is aware of all possible simplification rules would
# result in an overly complex and rigid tool, and would be difficult to generalize for
# all hardware implementations. By leaving the compiler itself unadorned and allowing
# for custom rules, very complex compilation instructions can be expressed in a simple
# and readable fashion.

#%%
# With this in mind, |True-Q|\'s :py:class:`~trueq.compilation.Compiler` works by
# applying an ordered list of built-in and/or user-specified
# :py:class:`~trueq.compilation.base.Pass`\es to a circuit in order.
# :py:class:`~trueq.compilation.base.Pass` objects define rules for how to decompose,
# replace, or remove :py:class:`~trueq.Gate`\s (or more generally,
# :py:class:`~trueq.Operation`\s), while possibly also adding or removing
# :py:class:`~trueq.Cycle`\s.
#
# Getting started
# ---------------
#
# Let's start with a simple example that demonstrates the actions of a
# compiler. Consider the circuit below:
import trueq as tq

circuit = tq.Circuit(
    cycles=[
        tq.Cycle({(0, 1): tq.Gate.cz}),
        tq.Cycle({0: tq.Gate.h, 1: tq.Gate.h}),
        tq.Cycle({(2, 1): tq.Gate.cx}),
        tq.Cycle({(0, 1): tq.Gate.cy, 2: tq.Gate.z}),
    ]
)
circuit.draw()

#%%
# We want to define a compiler that can replace certain cycles within this circuit
# with different cycles, and furthermore merge any resulting additional single- and
# two-qudit operations. To do that, we define two types of passes: a
# :py:class:`~trueq.compilation.CycleReplacement` and a
# :py:class:`~trueq.compilation.Merge` pass:

replace_pass = tq.compilation.CycleReplacement(
    target=tq.Cycle({(0, 1): tq.Gate.cz}),
    replacement=3 * [tq.Cycle({(0, 1): tq.Gate.cz})],
)
merge_pass = tq.compilation.Merge(max_sys=2)

#%%
# The :py:class:`~trueq.compilation.CycleReplacement` replaces any occurrence of
# the ``target`` cycle with the ``replacement`` cycle. In this case, our
# ``replace_pass`` will replace ``tq.Gate.cz`` on qubits ``(0, 1)`` with three
# consecutive ``tq.Gate.cz`` on qubits ``(0, 1)``. The
# :py:class:`~trueq.compilation.Merge` pass looks for any neighbouring gates which
# act on ``max_sys`` qudits with the same labels and merge them into a single gate.

#%%
# To see how those passes work in practice, let's define two different compilers:
# one which applies only ``replace_pass`` and one which applies only ``merge_pass``.
replace_compiler = tq.compilation.Compiler(passes=[replace_pass])
merge_compiler = tq.compilation.Compiler(passes=[merge_pass])

#%%
# Here is the output from the first compiler:

compilation1 = replace_compiler.compile(circuit)
compilation1.draw()

#%%
# Notice that the compiler constructed with ``replace_pass`` replaced the
# ``tq.Gate.cz`` with three consecutive ``tq.Gate.cz`` gates on qubits ``(0, 1)``, as
# expected.
#
# Now let's apply our ``merge_compiler`` to this circuit:
compilation2 = merge_compiler.compile(compilation1)
compilation2.draw()

#%%
# Note that the compiler constructed using the ``merge_pass`` merged the first four
# cycles into a single cycle, as intended.

#%%
# Applying first our ``replace_compiler`` and then our ``merge_compiler`` has the effect
# of the cycle replacement pass and the cycle merging pass to be executed in that order.
#
# The same can be achieved through a single compiler that we initialize with a list
# containing both passes:
composite_compiler = tq.compilation.Compiler(passes=[replace_pass, merge_pass])

composite_compilation = composite_compiler.compile(circuit)
composite_compilation.draw()

#%%
# Predefined Pass Lists
# ---------------------
#
# In addition to the passes shown above, there are several other useful
# pass classes available as members of the
# :py:class:`~trueq.compilation.Compiler` class, for example:
# :py:attr:`~trueq.compilation.Compiler.HARDWARE_PASSES`\,
# :py:attr:`~trueq.compilation.Compiler.SIMPLIFY_PASSES`\,
# :py:attr:`~trueq.compilation.Compiler.RC_PASSES`\, and
# :py:attr:`~trueq.compilation.Compiler.NATIVE2Q_PASSES`\. Click on the links to their
# documentation for further details on each.
#
# Note that these pass classes need to be instantiated before they can be given to the
# compiler constructor. For certain types of advanced usage, directly invoking the
# constructor may be the preferred method. However, the compiler also has two static
# covenience methods to automate pass instantiation:
# :py:meth:`~trueq.compilation.Compiler.from_config` and
# :py:meth:`~trueq.compilation.Compiler.basic`\. The first of these uses a
# :py:class:`~trueq.Config` object to pass relevant gate factories to each pass, and the
# second is a further specialization that auto-instantiates these factories based on
# a desired entangling gate and mode (which may not be relevant to all passes).

# define a compiler to simplify circuits
compiler1 = tq.Compiler.basic(passes=tq.Compiler.SIMPLIFY_PASSES)

# define a compiler that does randomized compiling
compiler2 = tq.Compiler.basic(passes=tq.Compiler.RC_PASSES)

# define a compiler that decomposes two-qubit gates into the specified entangling gate,
# then randomly compiles, and then converts all gates into hardware compatable gates.
# for this example, we use the maximally entangling cross-resonance gate.
passes = (
    tq.Compiler.NATIVE2Q_PASSES + tq.Compiler.RC_PASSES + tq.Compiler.HARDWARE_PASSES
)
compiler3 = tq.Compiler.basic(entangler=tq.Gate.rp("ZX", 90), passes=passes)

#%%
# We demonstrate the last of these compilers by defining the following circuit:

circuit = tq.Circuit(
    [{range(4): tq.Gate.h}, {(0, 1): tq.Gate.rp("ZZ", 22), (2, 3): tq.Gate.cz}]
).measure_all()
circuit

#%%
# We see that the compiled circuit contains only CNOT gates, Z rotations, and X90
# rotations. Adjacent single qubit gates have been merged together on each qubit.
compiler3.compile(circuit)

#%%
# Custom Pass Lists
# -----------------
# Custom lists of passes can be used if none of the predifined lists have the desired
# behaviour. We can pass lists of these classes to
# :py:meth:`~trueq.compilation.Compiler.from_config` or
# :py:meth:`~trueq.compilation.Compiler.basic`\, as discussed above, or we can use the
# compiler constructor directly, as in the following example.
#
# First, we define a device configuration. Here, we use the Berkeley gate as the
# entangling operation.
b = tq.Gate.from_generators("XX", 90, "YY", 45)
config = tq.Config(
    factories=[
        tq.config.GateFactory.from_matrix("B", b),
        tq.config.GateFactory.from_hamiltonian("x90", [["X", 90]]),
        tq.config.GateFactory.from_hamiltonian("z", [["Z", "phi"]]),
    ],
    mode="ZXZXZ",
)


#%%
# Next, we define a compiler. Unlike,
# :py:attr:`~trueq.compilation.Compiler.HARDWARE_PASSES`\, this compiler starts by
# merging adjacent two-qubit gates. Also, we know that any two qubit gate can be
# decomposed into two Berkeley gates interleaved with single qubit gates. We use
# :py:class:`tq.compilation.NativeDecomp` fixed at ``depth=2`` to force every
# two-qubit gate to decompose into two Berkeley gates, even if only one is required.
decomposer = tq.compilation.NativeDecomp(depth=2, factories=config.factories)
compiler = tq.Compiler(
    [
        tq.compilation.Merge(max_sys=2),
        tq.compilation.Parallel(decomposer),
        tq.compilation.Merge(),
        tq.compilation.Native1Q(config.factories),
        tq.compilation.RemoveEmptyCycle(),
    ]
)

#%%
# Next, we define a circuit to test this compiler on.
circuit = tq.Circuit(
    [
        {range(4): tq.Gate.h},
        {(0, 1): tq.Gate.rp("ZZ", 22)},
        {(0, 1): tq.Gate.cnot, (2, 3): tq.Gate.swap},
    ]
).measure_all()
circuit

#%%
# The resulting compiled circuit is as follows.
compiler.compile(circuit)