# Copyright 2018-2021 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
This module contains the qml.purity measurement.
"""

from typing import Sequence
import pennylane as qml
from pennylane.wires import Wires

from .measurements import StateMeasurement


def purity(wires):
    r"""The purity of the system prior to measurement.

    .. math::
        \gamma = \text{Tr}(\rho^2)

    where :math:`\rho` is the density matrix. The purity of a normalized quantum state satisfies
    :math:`\frac{1}{d} \leq \gamma \leq 1`, where :math:`d` is the dimension of the Hilbert space.
    A pure state has a purity of 1.

    Args:
        wires (Sequence[int] or int): The wires of the subsystem

    Returns:
        PurityMP: Measurement process instance

    **Example**

    .. code-block:: python3

        dev = qml.device("default.mixed", wires=2)

        @qml.qnode(dev)
        def circuit_purity(p):
            qml.Hadamard(wires=0)
            qml.CNOT(wires=[0, 1])
            qml.BitFlip(p, wires=0)
            qml.BitFlip(p, wires=1)
            return qml.purity(wires=[0,1])

    >>> circuit_purity(0.1)
    array(0.7048)

    .. seealso:: :func:`pennylane.qinfo.transforms.purity` and :func:`pennylane.math.purity`
    """
    wires = Wires(wires)
    return PurityMP(wires=wires)


class PurityMP(StateMeasurement):
    """Measurement process that computes the purity of the system prior to measurement.

    Please refer to :func:`purity` for detailed documentation.

    Args:
        obs (.Observable): The observable that is to be measured as part of the
            measurement process. Not all measurement processes require observables (for
            example ``Probability``); this argument is optional.
        wires (.Wires): The wires the measurement process applies to.
            This can only be specified if an observable was not provided.
        eigvals (array): A flat array representing the eigenvalues of the measurement.
            This can only be specified if an observable was not provided.
        id (str): custom label given to a measurement instance, can be useful for some
            applications where the instance has to be identified
    """

    @property
    def numeric_type(self):
        return float

    def shape(self, device=None):
        if qml.active_return():
            return self._shape_new(device)
        if device is None or device.shot_vector is None:
            return (1,)
        num_shot_elements = sum(s.copies for s in device.shot_vector)
        return (num_shot_elements,)

    def _shape_new(self, device=None):
        if device is None or device.shot_vector is None:
            return ()
        num_shot_elements = sum(s.copies for s in device.shot_vector)
        return tuple(() for _ in range(num_shot_elements))

    def process_state(self, state: Sequence[complex], wire_order: Wires):
        wire_map = dict(zip(wire_order, list(range(len(wire_order)))))
        indices = [wire_map[w] for w in self.wires]
        return qml.math.purity(state, indices=indices, c_dtype=state.dtype)
