#
# Copyright 2022 Keysight Technologies Inc.
#

"""
Writing Custom Noise Sources
============================
"""

#%%
# This example will demonstrate how to write a custom noise source for the simulator.
# See also :doc:`this guide<../../guides/run/simulator>` for more general information.

import numpy as np
import trueq as tq
import trueq.simulation as tqs

#%%
# A custom noise source can be implemented by subclassing
# :py:class:`~trueq.simulation.noise_source.NoiseSource` and overriding its
# :py:meth:`~trueq.simulation.noise_source.NoiseSource.apply` method.
#
# The :py:meth:`~trueq.simulation.noise_source.NoiseSource.apply` method will be passed
# three arguments: ``cycle_wrapper``, ``backend``, and ``circuit_cache``. The job of
# apply is to give instructions to the ``backend`` based on the contents of
# ``cycle_wrapper``. The ``backend`` is an instance of
# :py:class:`~trueq.simulation.backend.SimulationBackend` whose runner contains methods
# like :py:meth:`~trueq.simulation.backend.SimulationBackend.Runner.process_gate` and
# :py:meth:`~trueq.simulation.backend.SimulationBackend.Runner.process_superop` which
# the noise source's :py:meth:`~trueq.simulation.noise_source.NoiseSource.apply` can
# use.
#
# Note that ``cycle_wrapper`` is not a :py:class:`~trueq.Cycle` instance as one
# might expect, but instead a thin wrapper called
# :py:class:`~trueq.simulation.match.CycleWrapper` which in turn wraps each cycle
# operation as :py:class:`~trueq.simulation.match.OpWrapper`\.
#
# These wrapper types are required to store metadata on each operation (also discussed
# further below). However, noise sources will rarely need to deal with these wrapper
# objects directly. This is because noise sources (by strongly encouraged convention,
# though not by contract) own a :py:class:`~trueq.simulation.match.Match` instance which
# takes ownership of the mechanics of these wrappers. This is seen in the example below,
# where the :py:class:`~trueq.simulation.match.Match.iter_gates` method is used to yield
# gate objects with their corresponding gate labels.

# predefine noise matrices to be applied to qubits following ideal gates
# note that rowstack_subsys is the superoperator convention required by the simulator
p = 0.01
s1 = tq.math.Superop.from_kraus(
    [np.sqrt(1 - p) * np.eye(2), np.sqrt(p) * tq.Gate.x.mat]
)

p = 0.02
s2 = tq.math.Superop.from_kraus(
    [np.sqrt(1 - p) * np.eye(4), np.sqrt(p) * tq.math.random_unitary(4)]
)


class ExampleNoise(tqs.NoiseSource):
    def __init__(self, match=None):
        # this simulator hardcodes the subsystem dimension to 2
        super().__init__(dim=2, match=match)

    def make_circuit_cache(self, circuit):
        # this return will be made available to apply() as circuit_cache for every
        # cycle in the circuit
        return set(circuit.labels)

    def apply(self, cycle_wrappers, backend, circuit_cache):
        # in this method we ask the backend to process certain operations

        used_labels = set()
        # loop through the gates in the cycle and ask the backend to process them
        # note that we must set noise_only=False since we are simulating each gate
        for labels, gate in self.match.iter_gates(cycle_wrappers, noise_only=False):
            used_labels.update(labels)
            # first do ideal simulation of this gate
            backend.process_gate(labels, gate)

            # now add some superoperator noise to each qubit the gate acts on
            if len(labels) == 1:
                backend.process_superop(labels, s1)
            elif len(labels) == 2:
                backend.process_superop(labels, s2)

        # apply a 20 degree Z rotation to every qubit without a gate in this cycle
        for label in circuit_cache.difference(used_labels):
            backend.process_gate((label,), tq.Gate.rp("Z", 20))


# %%
# Now we can instantiate a simulator that uses this noise source and do any simulator
# calculation. Here, we use the simulator to predict the output of a hypothetical
# KNR experiment.
sim = tq.Simulator().append_noise_source(ExampleNoise())

cycle = {(0, 1): tq.Gate.cnot, 2: tq.Gate.x}
fit = sim.predict_knr(cycle, twirl=tq.Twirl("P", range(5)))
fit.plot.knr_heatmap()

# %%
# Caching
# -------
#
# There are two distinct caches available to a noise source. The first is a private
# noise source attribute that is called ``_cache`` (by convention, not by contract) in
# built-in noise models which persists throughout the lifetime of the noise source. For
# example, a gate-dependent noise model may consider caching gate noise if it is
# expensive to recompute everytime the same gate is encountered.
#
# The second cache comes as the third argument to the
# :py:meth:`~trueq.simulation.noise_source.NoiseSource.apply` method. This cache is
# instantiated by
# :py:meth:`~trueq.simulation.noise_source.NoiseSource.make_circuit_cache` just before a
# new circuit is simulated and is presented as the third argument to the apply method
# for every cycle within that circuit. It typically stores information about the circuit
# that is not available in every cycle, such as all of the labels the circuit acts on,
# or which measurements it performs at the end.

# %%
# Wrappers and metadata
# ---------------------
#
# The variable ``cycle_wrapper`` in the example above has type
# :py:class:`~trueq.simulation.match.CycleWrapper` which is a thin wrapper for a cycle
# and also wraps every cycle operation with
# :py:class:`~trueq.simulation.match.OpWrapper`\.  During simulation, everytime a new
# cycle is encountered it is wrapped once, and the same wrapped cycle instance is passed
# to all noise sources.  The cycle wrapper exists to enable efficient looping logic used
# by :py:class:`~trueq.simulation.match.Match` and to persist operation wrappers for
# multiple noise sources\. The operation wrapper exists to store
# two important pieces of information:
#
#  1. Whether some noise source has previously simulated the same operation of the
#     cycle. This is stored as a boolean called  ``has_been_simulated`` and is set to
#     true whenever a match method is called with the  argument ``noise_only=False``, as
#     it is in the above example.
#  2. Whether no noise sources (except the ideal final simulation if relevant) should
#     touch the operation again. This is stored as a boolean called ``no_more_noise``
#     and is set to true whenever it is yielded by a match method whose ``exclusive``
#     value is true.