# %%
#
# Copyright 2021 Quantum Benchmark Inc.
#

"""
Phase Tracking with the Compiler
================================
"""

# %%
# Phase tracking is a technique that exploits the ability of quantum control hardware to
# perform accurate and arbitrary phase rotations of pulse shapes in order to reduce the
# number of pulse shapes required to form a device's native quantum gate set. Moreover,
# it causes the native gate set to be continuously parametric which greatly reduces the
# synthesis cost of arbitrary unitaries. For example, if an :math:`R_x(90)` gate is
# tuned up, then any single qubit unitary can be synthesized using two of these pulses
# and phase tracking.
#
# Below, we call these phase rotation gates *virtual gates* because they are not
# directly used to update the quantum state. Instead, virtual gates are better thought
# of as pulse compiler directives, where the pulse compiler is a program that prepares a
# sequence of waveforms for the control electronics given a circuit of native gates. In
# a multi-qubit device that employs phase tracking, the pulse compiler accumulates a
# phase on each qubit in the circuit whenever a virtual gate is encountered. Any time a
# non-virtual gate is encountered on a given subset of qubits, the corresponding
# cummulative phases are used to modify the phases of the pulse shape.
#
# We illustrate how phase tracking works using the following single-qubit circuit
# expressed as rotations in matrix multiplication order. This circuit is assumed to
# be the output of a pre-compiler that has already converted some original circuit into
# :math:`R_x(90)` and :math:`R_z(\phi)` rotations.
#
# .. math::
#
#   R_z(d)\cdot R_x(90) \cdot R_z(c)\cdot R_x(90)
#       \cdot R_z(b) \cdot R_x(90) \cdot R_z(a)
#
# This can be rewritten as
#
# .. math::
#
#     R_z(a+b+c+d) \cdot R_z(-(a+b+c)) \cdot R_x(90) \cdot R_z(a+b+c)
#         \\ \quad
#         \cdot R_z(-(a+b)) \cdot R_x(90) \cdot R_z(a+b) \cdot R_z(-a)
#         \cdot R_x(90) \cdot R_z(a)
#
# and further simplified to
#
# .. math::
#
#     R_z(a+b+c+d)R_{a+b+c}(90)R_{a+b}(90) R_a(90)
#
# where we have defined a nutation about some vector in the x-y plane as
#
# .. math ::
#
#     R_\phi(\theta)=R_z(-\phi)R_x(90)R_z(\phi)
#         =\operatorname{exp}(-i \theta (\cos(\phi)X-\sin(\phi)Y)/2).
#
# Thus if a pulse shape for the :math:`R_x(90)` gate is tuned up, and our control
# electronics are able to rotate it in quadrature by arbitrary angles thereby producing
# operations :math:`R_\phi(90)`, then we have sufficient control to perform the
# original circuit. Note further that any qubit unitary can be decomposed into an
# alternating sequence :math:`R_z(c)\cdot R_x(90) \cdot R_z(b) \cdot R_x(90) \cdot
# R_z(a)` where :math:`a`, :math:`b`, :math:`c` are its ZXZ Euler angles. It follows
# that this is the *only* pulse shape that is required to perform single qubit gates.
#
# For most systems, the final accumulated phase in the z-axis, :math:`R_z(a+b+c+d)`,
# does not actually need to be implemented on the qubit. This is because the final
# operation will be a measurement of the qubit along the z-axis, whose outcome will not
# be affected by a z-rotation.
#
# The process of phase tracking on multi-qubit gates is similar to the single qubit
# case, which will be seen in the examples below. As previously mentioned, a different
# phase needs to be tracked for each qubit, and pulse shapes may require a parameter for
# every qubit they act on.

import trueq as tq
import numpy as np

# %%
# Using the Controlled-Z Gate
# ---------------------------
# A device whose only native two-qubit gate is the CZ gate is the simplest case because
# this commutes with the virtual Z gate on both qubits. This means we only need to apply
# accumulated phases to the single qubit gates.
#
# First, we define a config object that contains all of the necessary gate factories.
# Since we will be doing phase tracking, we require not only the parametric native gate
# factories and the virtual gate factory, but also a factory for each of the static
# gates that correspond to the parametric gates at parameter value 0.

# factories for our parametric gates. the first is a 90 degree rotation about an
# axis in the x-y plane, equivalent to z(-a)*x(90)*z(a). the second is our cz gate
# which happens to need no parameters because it commutes with virtual gates
r90 = tq.config.GateFactory.from_hamiltonian(
    "r90", [["Z", "-theta"], ["X", "90"], ["Z", "theta"]]
)
cz = tq.config.GateFactory.from_matrix("cz", np.diag([1, 1, 1, -1]))

# factory for our virtual gate
z = tq.config.GateFactory.from_hamiltonian("z", [["Z", "theta"]])

# factory for static gates, equal to our parametric gates at 0
x90 = tq.config.GateFactory.from_matrix("x90", r90(0))
factories = [x90, z, cz, r90]

# %%
# Next, we initialize our phase-tracking compiler pattern. Note that it chooses the
# virtual gate factory by looking for a rotation about the Z-axis in units of degrees.
# If this is not the case for you, specify any other single qubit rotation manually
# using ``virtual=factory``.
phase_tracker = tq.compilation.PhaseTrack(factories=factories)

# %%
# Next we set up our compiler and the rest of its passes. There are two essential steps:
#
#  1. Convert all abstract gates into the static gates and the virtual gate
#     (:py:class:`~trueq.compilation.Native1Q` and
#     :py:class:`~trueq.compilation.Native2Q`).
#  2. Use the :py:class:`~trueq.compilation.PhaseTrack` pattern to accumulate phases and
#     compile them into the parametric gate parameters.
#
# There are additionally justification, merge, and identity removal book-keeping passes
# that can be adjusted to one's needs.
compiler = tq.Compiler(
    [
        tq.compilation.Native2Q(factories),
        tq.compilation.Justify(),
        tq.compilation.Merge(),
        tq.compilation.RemoveId(),
        tq.compilation.Native1Q(factories),
        phase_tracker,
        tq.compilation.RemoveEmptyCycle(),
    ]
)

# %%
# We illustrate the output of this compiler using a circuit that generates the GHZ state
# on :math:`n` qubits. Before compilation, our abstract circuit on 4 qubits is as
# follows:


def ghz(n):
    circuit = tq.Circuit([{0: tq.Gate.h}])
    for idx in range(n - 1):
        circuit.append({(idx, idx + 1): tq.Gate.cnot})
    circuit.measure_all()
    return circuit


ghz(4)

# %%
# The output bitstring distribution is as expected:
tq.Simulator().state(ghz(4)).to_results().plot()

# %%
# Following compilation, our circuit contains only gates that were generated using the
# ``r90`` or ``cz`` factory. The parameters the factories used to construct the gates
# are stored in the gate objects, along with the gates' matrix representations. Single
# qubit gates, under this compiler, are always decomposed into two 90 degree nutations.
circuit = compiler.compile(ghz(4))
circuit

# %%
# Finally, our output bitstring distribution is unchanged, as expected.
tq.Simulator().state(circuit).probabilities().to_results().plot()

# %%
# Using the Cross-Resonance Gate
# ------------------------------
# The maximally entangling cross-resonance gate is equal to
#
# .. math::
#     CR=\begin{pmatrix}
#         1 & -i & 0 & 0 \\ -i & 1 & 0 & 0 \\ 0 & 0 & 1 & i \\ 0 & 0 & i & 1
#     \end{pmatrix}/\sqrt{2}
#     =\frac{I\otimes I-iZ\otimes X}{\sqrt{2}}.
#
# By adjusting the phase of the microwave control used to generate this gate, we can
# rotate action of the second qubit about the z-axis, which results in the parametric
# gate
#
# .. math::
#
#     CR(\phi)
#     = \frac{I\otimes I-i\cos(\phi)Z\otimes X+i\sin(\phi)Z\otimes Y}{\sqrt{2}}.
#
#
# This parametric gate is defined below as a gate factory, along with other factories
# necessary for this example.

# define parametric gate factories
r90 = tq.config.GateFactory.from_hamiltonian(
    "r90", [["Z", "-theta"], ["X", "90"], ["Z", "theta"]]
)
cr = tq.config.GateFactory.from_hamiltonian(
    "cr", [["IZ", "-theta"], ["ZX", "90"], ["IZ", "theta"]]
)

# factory for our virtual gate
z = tq.config.GateFactory.from_hamiltonian("z", [["Z", "theta"]])

# factory for static gates, equal to our parametric gates at 0
x90 = tq.config.GateFactory.from_matrix("x90", r90(0))
cr0 = tq.config.GateFactory.from_matrix("cr0", cr(0))

# put everything into a factory list. static gates should go first
factories = [x90, z, cr0, r90, cr]

# %%
# As in the CZ section, we define a compiler that performs phase-tracking using our
# native gates, and test it on a circuit that generates a 4-qubit GHZ state.
compiler = tq.Compiler(
    [
        tq.compilation.Native2Q(factories),
        tq.compilation.Justify(),
        tq.compilation.Merge(),
        tq.compilation.RemoveId(),
        tq.compilation.Native1Q(factories),
        tq.compilation.PhaseTrack(factories),
        tq.compilation.Justify(),
    ]
)

circuit = compiler.compile(ghz(4))
circuit

# %%
# Finally, we verify that our bitstring output is as expected.
tq.Simulator().state(circuit).probabilities().to_results().plot()

# %%
# Using the Molmer-Sorensen Gate
# ------------------------------
# A maximally entangling MS gate has the form
#
# .. math::
#
#     MS(90)=\begin{pmatrix}
#         1 & 0 & 0 & -i \\ 0 & 1 & -i & 0 \\ 0 & -i & 1 & 0 \\ -i & 0 & 0 & 1
#     \end{pmatrix}/\sqrt{2}
#     =\frac{I\otimes I-iX\otimes X}{/\sqrt{2}}
#
# which, for a trapped ion, is generated by lasing on both qubits simultaneously.
# However, by changing the phase of the frequency sources sent to the optical modulators
# on each qubit, the azimuthal angle of the MS gate can be changed independently on both
# qubits qubits. This allows the implementation of the parametric gate
#
# .. math::
#
#     MS(90, \phi_1, \phi_2) = \frac{II-iR_{\phi_1}(90)\otimes R_{\phi_2}(90)}{\sqrt{2}}
#
# where :math:`R_\phi(\theta)` is defined above.
#
# This parametric gate is defined below as a gate factory, along with other factories
# necessary for this example.

# factories for our parametric gates. the first is a 90 degree rotation about an axis in
# the x-y plane, equivalent to z(a)*x(90)*z(-a). the second is a molmer-sorensen gate
# with phase rotation a on the first qubit and b on the second qubit, i.e. equivalent to
# (z(a) & z(b))*xx(90)*(z(-a) & z(-b))
r90 = tq.config.GateFactory.from_hamiltonian(
    "r90", [["Z", "-theta"], ["X", "90"], ["Z", "theta"]]
)
ms = tq.config.GateFactory.from_hamiltonian(
    "ms", [["IZ", "-b"], ["ZI", "-a"], ["XX", 90], ["ZI", "a"], ["IZ", "b"]]
)

# factory for our virtual gate
z = tq.config.GateFactory.from_hamiltonian("z", [["Z", "theta"]])

# factory for static gates, equal to our parametric gates at 0
x90 = tq.config.GateFactory.from_matrix("x90", r90(0))
ms0 = tq.config.GateFactory.from_matrix("ms0", ms(0, 0))

# put everything into a factory list
factories = [x90, z, ms0, r90, ms]

# %%
phase_tracker = tq.compilation.PhaseTrack(factories)
phase_tracker._rules

# %%
# As in the CZ section, we define a compiler that performs phase-tracking using our
# native gates, and test it on a circuit that generates a 4-qubit GHZ state.
compiler = tq.Compiler(
    [
        tq.compilation.Native2Q(factories),
        tq.compilation.Justify(),
        tq.compilation.Merge(),
        tq.compilation.RemoveId(),
        tq.compilation.Native1Q(factories),
        tq.compilation.PhaseTrack(factories),
        tq.compilation.Justify(),
    ]
)

circuit = compiler.compile(ghz(4))
circuit

# %%
# Finally, we verify that our bitstring output is as expected.
tq.Simulator().state(circuit).probabilities().to_results().plot()