#
# Copyright 2021 Quantum Benchmark Inc.
#

"""
State Preparation and Measurement Noise
=======================================
"""

#%%

import trueq as tq
import matplotlib.pyplot as plt
import numpy as np

#%%
# Adding Measurement Bitflip Noise
# --------------------------------

#%%
# The most convenient way to add measurement noise is through the
# :py:meth:`~trueq.Simulator.add_readout_error` noise source. There are
# several ways to specify readout error with this method.
#
# The most basic method is to provide a single number. In the following example, each
# qubit will get a (symmetric) 1% bitflip error.
sim = tq.Simulator().add_readout_error(0.01)
circuit = tq.Circuit({range(5): tq.Meas()})
sim.run(circuit, n_shots=1000)
circuit.results

#%%
# We can use a dictionary to specify the readout error of particular qubits. Qubits
# that are not explicitly assigned a readout error get the value attached to `None`
# which is equal to 0 by default. In this example, all qubits get 1% readout error,
# except qubit 1 gets a 50% readout error.
sim = tq.Simulator().add_readout_error(0.01, {1: 0.5})
circuit = tq.Circuit({range(5): tq.Meas()})
sim.run(circuit, n_shots=1000)
circuit.results

#%%
# In both of the above examples, any probability can be replaced with a pair of
# probabilities [p10, p01], where p10 is the probability of flipping a '0' to a '1'
# right before measurement, and p01 is the probability of flipping a '1' to a '0' right
# before measurement. In the following example, qubits get 1% readout error, except
# qubit 1 gets 5% readout error, and qubit 3 get asymmetric readout error of 1% on
# :math:`|0\rangle` and 7% on :math:`|1\rangle`.
sim = tq.Simulator().add_readout_error(0.01, {1: 0.05, 3: [0.01, 0.07]})
circuit = tq.Circuit({range(5): tq.Meas()})
sim.run(circuit, n_shots=1000)
circuit.results

#%%
# Adding Measurement POVM Noise
# -----------------------------

#%%
# Measurement noise is applied only when a simulator's
# :py:meth:`~trueq.Simulator.run` method is called. Behind the scenes,
# measurement noise is impelemented by applying a positive-operator valued measurement
# (POVM) channel to the quantum state directly before measurement, thereby modifying the
# probability distribution from which ditstrings are sampled.
#
# As a quick reminder (look elsewhere for a full description, e.g. `Wikipedia
# <https://en.wikipedia.org/wiki/POVM>`_ or the `lecture notes of John Watrous
# <https://cs.uwaterloo.ca/~watrous/TQI/>`_), a POVM is a set of positive-semi-definite
# matrices :math:`\{P_i\}_{i=1}^N` that sum to the identity matrix, :math:`\sum_{i=1}^N
# P_i = \mathbb{I}`. The probability of observing the outcome :math:`i`  when measuring
# a state :math:`\rho` is equal to :math:`p_i:=\operatorname{Tr}(\rho P_i)`. One can
# also take a POVM defined on a single qubit, and tensor it together to get a POVM on a
# collection of qubits. For example, if the POVM above describes measurement on a single
# qubit, the probability of measuring the outcomes :math:`(i_0, i_1, i_2)` on three
# qubits is simply :math:`p_{i_0}p_{i_1}p_{i_2}`. This is why we use the
# :py:class:`~trueq.math.tensor.Tensor` object as a primitive for specifying POVMs; its
# purpose is to store tensor project structures sparsely without actually taking
# kronecker products between subsystems unless necessary.
#
# The :py:meth:`~trueq.Simulator.add_povm` method can be used to directly add
# a POVM measurement channel. A POVM is specified as a
# :py:class:`~trueq.math.tensor.Tensor` object, where the input dimension is the number
# of POVM elements (per subsystem) and the output dimension pair is the dimension of the
# system.

#%%
# In this example, we construct a coherent measurement error on qubit 3.

# Define a set of ideal POVM operators for each subsystem
proj0 = np.array([[1, 0], [0, 0]])  # project onto this one to get a "0"
proj1 = np.array([[0, 0], [0, 1]])  # project onto this one to get a "1"
ideal = np.array([proj0, proj1])

# Define a unitary rotation about X by 20 degrees
u = tq.Gate.from_generators("X", 20).mat

# Define noisy POVM operators by rotating the ideal POVM operators by U, which
# coherently changes the basis of the measurement
twisted = [u @ x @ u.conj().T for x in ideal]

# The spawn is the default POVM, and value customizes the POVM of specific qubit labels
povm = tq.math.Tensor(
    2,  # number of POVMs per subsystem
    (2, 2),  # number of POVMs, (dimensions of qubit subsystems)
    spawn=ideal,  # default measurement operation
    value={(3,): twisted},  # measurement operation on qubit 3
    dtype=np.complex128,  # tensor is real by default
)

# initialize a simulator with the above POVM specifications and test it
sim = tq.Simulator().add_povm(povm)
circuit = tq.Circuit({range(5): tq.Meas()})
sim.run(circuit, n_shots=1000)
circuit.results

#%%
# In this example, we add a correlated readout error on qubits (0,1).

# Define a confusion matrix for a 2-qubit system where the readout of 00, 01, 10 is
# ideal, but 11 is flipped to 01 5% of the time
confusion = [[1, 0, 0, 0], [0, 1, 0, 0.05], [0, 0, 1, 0], [0, 0, 0, 0.95]]

# Add the readout error and test the simulator
sim = tq.Simulator().add_readout_error(errors={(0, 1): confusion})
circuit = tq.Circuit([{(0, 1): tq.Gate.x}, {(0, 1): tq.Meas()}])
sim.run(circuit, n_shots=10000)
circuit.results

#%%
# In this example, we add a third measurement label even though the simulation is on
# qubits, so that the outcomes '0', '1', and '2' are all possible.

# Here we specify that the state |0> results in the outcome '0' 99% of the time, but
# also results in the outcome '2' 1% of the time. Similarly, the state |1> results in
# the outcome '1' 70% of the time, but also results in '2' 25% of the time and '0' 5%
# of the time.
povm = tq.math.Tensor(
    3,
    (2, 2),
    spawn=[np.diag([0.99, 0.05]), np.diag([0, 0.7]), np.diag([0.01, 0.25])],
    dtype=np.complex128,
)
sim = tq.Simulator().add_povm(povm)
circuit = tq.Circuit([{1: tq.Gate.h}, {(0, 1): tq.Meas()}])
sim.run(circuit, n_shots=10000)
circuit.results

#%%
# Adding State Preparation Noise
# ------------------------------

#%%
# State preparation noise can be added using the
# :py:meth:`~trueq.Simulator.add_prep` noise source.
#
# .. note::
#
#    The :py:meth:`~trueq.Simulator.add_prep` noise source only takes place when a
#    :py:class:`~trueq.Prep()` object is encountered.
#
# We can either enter the noise as the probability of a bitflip during preparation of
# :math:`|0\rangle`:
sim = tq.Simulator().add_prep(0.01)
circuit = tq.Circuit([{0: tq.Prep()}])
tq.visualization.plot_mat(sim.state(circuit).mat())

#%%
# Or we can specify the density matrix we want to prepare with:
sim = tq.Simulator().add_prep([[0.75, 0], [0, 0.25]])
circuit = tq.Circuit([{0: tq.Prep()}])
tq.visualization.plot_mat(sim.state(circuit).mat())

#%%
# We can place different preparations on different qubits. Here, we have a 1%
# preparation bitflip error by default, but a 20 degree rotation error on qubit 3.
u = tq.Gate.from_generators("Y", 20).mat
sim = tq.Simulator().add_prep({None: 0.01, 3: u @ np.diag([1, 0]) @ u.conj().T})
circuit = tq.Circuit([{0: tq.Prep(), 3: tq.Prep()}])
tq.visualization.plot_mat(sim.state(circuit).mat())

#%%
# We can specify preparation states of larger dimension to add leakage levels.
sim = tq.Simulator().add_prep(np.diag([0.97, 0.02, 0.01]))
circuit = tq.Circuit([{0: tq.Prep()}, {0: tq.Gate.x}])
tq.visualization.plot_mat(sim.state(circuit).mat())

#%%
# Note that if a preparation takes place mid-circuit, then that system is
# traced out (with no post-selection) and replaced by the specified preparation. This
# is a non-unitary action and will generally result in a mixed state.
#
# In the following example we prepare the bell state :math:`|00\rangle+|11\rangle`
# and then prepare the second qubit in :math:`|1\rangle`. Thus, we end up with the
# final state :math:`\frac{\mathbb{I}}{2}\otimes|1\rangle\langle 1|`.

sim = tq.Simulator().add_prep(np.diag([0, 1]))
circuit = tq.Circuit([{0: tq.Gate.h}, {(0, 1): tq.Gate.cnot}, {1: tq.Prep()}])
tq.visualization.plot_mat(sim.state(circuit).mat())

#%%
# When the simulator is asked to find the operator of a circuit containing a preparation
# object, the preparation object is simulated as a projection step.
plt.figure(figsize=(10, 5))
sim = tq.Simulator().add_prep(np.diag([1, 0]))

op1 = sim.operator(tq.Circuit([{0: tq.Gate.id}])).upgrade().mat()
tq.visualization.plot_mat(op1, ax=plt.subplot(121))
plt.title("Superoperator without Prep() projection")

op2 = sim.operator(tq.Circuit([{0: tq.Prep()}, {0: tq.Gate.id}])).mat()
tq.visualization.plot_mat(op2, ax=plt.subplot(122))
plt.title("Superoperator with Prep() projection")