#
# Copyright 2021 Quantum Benchmark Inc.
#

"""
Streamlined Randomized Benchmarking (SRB)
=========================================
"""

#%%
# These examples demonstrate how to use SRB, the most basic QCVV protocol, to estimate
# the process infidelity of a single-qubit or two-qubit gateset [#f1]_ . See
# :tqdoc:`SRB` for more information about this protocol.

#%%
# Hello World
# -----------
# The simplest thing we can do is run SRB on a single-qubit. Qubits are labeled by
# non-negative integers, and here we target qubit 0. We choose to use sequence lengths
# ``[4, 30, 50]``, and, by default, 30 random circuits are generated per sequence
# length.  Below, we display the first circuit where the 4 random Clifford gates can be
# seen, plus a fifth gate which is the product of their inverses and a random Pauli
# matrix, where the additional Pauli matrix is introduced to make the fit more robust.
# The empty cycles with non-zero markers delimit the random Cliffords, and more
# generally, in other protocols like :meth:`~trueq.make_cb()` and
# :meth:`~trueq.make_irb()`, contain interleaved cycles of interest.


import trueq as tq
import trueq.simulation as tqs

circuits = tq.make_srb(0, [4, 30, 50])
circuits[0]

#%%
# Next, we simulate each of the circuits with the built-in simulator. We can see the
# bitstring counts of the first circuit below.

tq.Simulator().add_overrotation(0.04).add_depolarizing(0.01).run(circuits)
circuits[0].results.plot()

#%%
# It is important to look at the decay curve to verify that appropriate sequence lengths
# were chosen. The most informative sequence length to measure occurs around
# :math:`m=1/(1-p)` where the decay is :math:`y=Ap^m`. If :math:`A\approx 1` and
# :math:`p\approx 1`, then this corresponds to :math:`y\approx1/e\approx0.37`.
# Therefore, in this example, our sequence lengths are a bit too short, and ``[4, 70]``
# would, in retrospect, have been more appropriate.
circuits.plot.raw()

#%%
# Finally, we can look at the values of the estimate:

circuits.fit()

#%%
# Isolated vs. Simultaneous SRB
# -----------------------------

# %%
# The core function :meth:`~trueq.make_srb()` generates a circuit collection that
# implements simultaneous SRB, where gates are placed on all specified qubits of a
# quantum device within each circuit. Different random gates are independently chosen
# during each cycle. Qubit labels are entered into its first argument and specify which
# qubits are acted on simultaneously. You can choose which qubits get single-qubit
# twirls, and which qubits get two-qubit twirls, by nesting the qubit labels
# appropriately; see the following examples.

# simultaneous single-qubit RB on qubits 0 and 1; 60 circuits total
circuits = tq.make_srb([0, 1], [4, 100], 30)

# isolated single-qubit RB on qubits 0 and 1; 120 circuits total
circuits = tq.make_srb([0], [4, 100], 30).append(tq.make_srb([1], [4, 100], 30))

# two-qubit RB on 0 and 1, simultaneous with single-qubit RB on 2; 60 circuits total
circuits = tq.make_srb([[0, 1], 2], [4, 100], 30)

#%%
# The invocations above all use two sequence lengths (the number of random gate cycles
# in a circuit), 4 and 100, and produce 30 random circuits per sequence length. Note
# that specifying qubit labels ``[0, 1, 2]`` is a convenient short-hand for
# ``[[0], [1], [2]]``. This imples, for example, that ``[[0,1]]`` has a much different
# meaning than ``[0,1]``.

#%%
# Isolated SRB
# ------------
# In isolated SRB, we characterize qubits or qubit-pairs one at a time with disjoint
# experiments. This is done to assess gate quality on a small subregister in the context
# of an idling environment. In this example, we build a single circuit collection to
# estimate the average gate infidelity of three pairs on a 6 qubit device.

# generate a circuit collection
circuits = tq.CircuitCollection()
for pair in [[0, 1], [2, 3], [4, 5]]:
    circuits += tq.make_srb([pair], [4, 10, 20])

#%%
# Next, we transpile the circuits into native gates and simulate them.

# Initialize a simulator
sim = tq.Simulator().add_overrotation(0.01, 0.02)
sim.add_stochastic_pauli(px=0.01, match=tqs.LabelMatch((0, 1)))
sim.add_stochastic_pauli(pz=0.005, match=tqs.LabelMatch((2, 3, 4, 5)))

# initialize a transpiler into native gates (U3 and CNOT)
cnot_factory = tq.config.GateFactory.from_matrix("cnot", tq.Gate.cnot.mat)
config = tq.Config(factories=[tq.config.u3_factory, cnot_factory])
t = tq.Compiler.from_config(config)

# transpile and simulate
circuits = t.compile(circuits)
sim.run(circuits)

#%%
# Plot the results.
circuits.plot.raw([[0, 1], [2, 3], [4, 5]])
circuits.plot.compare_rb()

#%%
# Simultaneous SRB
# ----------------
# In simultaneous SRB, we characterize qubits or qubit-pairs in a single set of
# experiments. This captures crosstalk between these subsystems which is usually not
# present when other qubits are idle. In this example, we add more circuits to the
# circuits from the previous section.

# add simultaneous circuits to our old collection
simul_circuits = t.compile(tq.make_srb([[0, 1], [2, 3], [4, 5]], [4, 10, 20]))
circuits += simul_circuits
sim.run(simul_circuits)

# plot the results
circuits.plot.raw([[0, 1], [2, 3], [4, 5]])
circuits.plot.compare_rb()

#%%
# Changing the Twirling Group Gateset
# -----------------------------------
# The above examples characterize the average gate fidelity of the Clifford group, which
# is historically by far the most common group to use. However, any unitary 2-design
# will suffice, and so for SRB, the other notable choice is the set of all unitaries
# sampled using the Haar measure. We make this alternate choice below.

su_circuits = tq.make_srb([[0, 1], [2, 3], [4, 5]], [4, 6, 12], twirl="U")
su_circuits = t.compile(su_circuits)
sim.run(su_circuits)

#%%
# Our native gates here are still CNOT and U3. It takes on average 1.5 CNOTs to
# synthesize a two-qubit Clifford gate, but it almost always takes 3 CNOTs to
# synthesize a Haar-random two-qubit unitary group. Thus we see higher error rates
# than in the section above.

#%%
# Simulate and plot the results.
su_circuits.plot.raw([[0, 1], [2, 3], [4, 5]])
su_circuits.plot.compare_rb()


#%%
# .. rubric:: Footnotes
#
# .. [#f1] "Streamlined" refers to the method we use to eliminate the constant offset
#    (the :math:`B` in :math:`Ap^m+B`) that typically appears in standard RB.