#
# Copyright 2021 Quantum Benchmark Inc.
#

"""
Circuit Batching
================
"""

#%%
# This example demonstrates how to use :py:meth:`~trueq.CircuitCollection.batch` to
# group a collection of circuits into numbered batches of arbitrary size. This has a
# number of useful applications, such as ensuring that you are always sending an
# optimal number of circuits to your device.

import trueq as tq
from time import sleep

# initialize a noisy simulator
sim = tq.Simulator().add_stochastic_pauli(px=0.02)

# generate 90 SRB circuits
circuits = tq.make_srb([[0], [1, 2]], [4, 32, 64], 30)

#%%
# We call the batch method to group the circuits into three batches with 30 circuits
# per batch. Note that the existing order of ``circuits`` is used in the batching; the
# argument ``sequencer`` can be used to customize this behaviour. For example, it can
# be used to riffle circuit depths.

batches = circuits.batch(30)

# execute each batch on the simulator
for batch in batches:
    print(f"executing batch {batch.keys().batch.pop()} with {len(batch)} circuits")
    sim.run(batch)
    sleep(1)

# display all unique batches of SRB circuits
circuits.keys().batch

#%%
# When the circuits were executed on the simulator, a one second delay was added
# between each batch, which can be visualized using
# :py:meth:`~trueq.visualization.plotters.timestamps.timestamps`.

# plot the timestamps of results for all circuits
circuits.plot.timestamps()

#%%
# The :py:meth:`~trueq.CircuitCollection.subset` method can also be used to further
# manipulate the batches as follows.

# plot the timestamps of results for batches 0 and 2
circuits.subset(batch={0, 2}).plot.timestamps()

#%%
# Extra Circuits
# --------------
#
# This section illustrates how to use the ``extra_circuits`` argument of
# :py:meth:`~trueq.CircuitCollection.batch` to append a set of identical circuits to
# each batch. For example, it can be used to include :tqdoc:`RCAL` circuits in each
# batch to correct the readout errors caused by the drift in your device.

# generate 90 SRB circuits
circuits = tq.make_srb([[0], [1, 2]], [4, 32, 64], 30)

# generate 2 RCAL circuits
rcal_circuits = tq.make_rcal([0, 1, 2])

# group the circuits into three batches with 30 SRB and 2 RCAL circuits per batch
batches = circuits.batch(32, extra_circuits=rcal_circuits)

# execute each batch on the simulator
for batch in batches:
    sim.run(batch)

# verify the contents of each batch
for batch in circuits.keys().batch:
    srb = circuits.subset(batch=batch, protocol="SRB")
    rcal = circuits.subset(batch=batch, protocol="RCAL")
    print(f"batch {batch} contains {len(srb)} SRB and {len(rcal)} RCAL circuits")