# Copyright 2018 The Cirq Developers
#
# 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
#
#     https://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.

from types import NotImplementedType
from typing import Any, Optional, TypeVar, Union

import numpy as np
from typing_extensions import Protocol

from cirq._doc import doc_private
from cirq.protocols import qid_shape_protocol
from cirq.protocols.apply_unitary_protocol import apply_unitaries, ApplyUnitaryArgs
from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits

# This is a special indicator value used by the unitary method to determine
# whether or not the caller provided a 'default' argument. It must be of type
# np.ndarray to ensure the method has the correct type signature in that case.
# It is checked for using `is`, so it won't have a false positive if the user
# provides a different np.array([]) value.
RaiseTypeErrorIfNotProvided: np.ndarray = np.array([])

TDefault = TypeVar('TDefault')


class SupportsUnitary(Protocol):
    """An object that may be describable by a unitary matrix."""

    @doc_private
    def _unitary_(self) -> Union[np.ndarray, NotImplementedType]:
        """A unitary matrix describing this value, e.g. the matrix of a gate.

        This method is used by the global `cirq.unitary` method. If this method
        is not present, or returns NotImplemented, it is assumed that the
        receiving object doesn't have a unitary matrix (resulting in a TypeError
        or default result when calling `cirq.unitary` on it). (The ability to
        return NotImplemented is useful when a class cannot know if it has a
        matrix until runtime, e.g. cirq.X**c normally has a matrix but
        cirq.X**sympy.Symbol('a') doesn't.)

        The order of cells in the matrix is always implicit with respect to the
        object being called. For example, for gates the matrix must be ordered
        with respect to the list of qubits that the gate is applied to. For
        operations, the matrix is ordered to match the list returned by its
        `qubits` attribute. The qubit-to-amplitude order mapping matches the
        ordering of numpy.kron(A, B), where A is a qubit earlier in the list
        than the qubit B.

        Returns:
            A unitary matrix describing this value, or NotImplemented if there
            is no such matrix.
        """

    @doc_private
    def _has_unitary_(self) -> bool:
        """Whether this value has a unitary matrix representation.

        This method is used by the global `cirq.has_unitary` method.  If this
        method is not present, or returns NotImplemented, it will fallback
        to using _unitary_ with a default value, or False if neither exist.

        Returns:
            True if the value has a unitary matrix representation, False
            otherwise.
        """


def unitary(
    val: Any, default: Union[np.ndarray, TDefault] = RaiseTypeErrorIfNotProvided
) -> Union[np.ndarray, TDefault]:
    """Returns a unitary matrix describing the given value.

    The matrix is determined by any one of the following techniques:

    - The value has a `_unitary_` method that returns something besides None or
        NotImplemented. The matrix is whatever the method returned.
    - The value has a `_decompose_` method that returns a list of operations,
        and each operation in the list has a unitary effect. The matrix is
        created by aggregating the sub-operations' unitary effects.
    - The value has an `_apply_unitary_` method, and it returns something
        besides None or NotImplemented. The matrix is created by applying
        `_apply_unitary_` to an identity matrix.

    If none of these techniques succeeds, it is assumed that `val` doesn't have
    a unitary effect. The order in which techniques are attempted is
    unspecified.

    Args:
        val: The value to describe with a unitary matrix.
        default: Determines the fallback behavior when `val` doesn't have
            a unitary effect. If `default` is not set, a TypeError is raised. If
            `default` is set to a value, that value is returned.

    Returns:
        If `val` has a unitary effect, the corresponding unitary matrix.
        Otherwise, if `default` is specified, it is returned.

    Raises:
        TypeError: `val` doesn't have a unitary effect and no default value was
            specified.
    """
    strats = [
        _strat_unitary_from_unitary,
        _strat_unitary_from_apply_unitary,
        _strat_unitary_from_decompose,
    ]
    for strat in strats:
        result = strat(val)
        if result is None:
            break
        if result is not NotImplemented:
            return result

    if default is not RaiseTypeErrorIfNotProvided:
        return default
    raise TypeError(
        "cirq.unitary failed. "
        "Value doesn't have a (non-parameterized) unitary effect.\n"
        "\n"
        f"type: {type(val)}\n"
        f"value: {val!r}\n"
        "\n"
        "The value failed to satisfy any of the following criteria:\n"
        "- A `_unitary_(self)` method that returned a value "
        "besides None or NotImplemented.\n"
        "- A `_decompose_(self)` method that returned a "
        "list of unitary operations.\n"
        "- An `_apply_unitary_(self, args) method that returned a value "
        "besides None or NotImplemented."
    )


def _strat_unitary_from_unitary(val: Any) -> Optional[np.ndarray]:
    """Attempts to compute a value's unitary via its _unitary_ method."""
    getter = getattr(val, '_unitary_', None)
    if getter is None:
        return NotImplemented
    return getter()


def _strat_unitary_from_apply_unitary(val: Any) -> Optional[np.ndarray]:
    """Attempts to compute a value's unitary via its _apply_unitary_ method."""
    # Check for the magic method.
    method = getattr(val, '_apply_unitary_', None)
    if method is None:
        return NotImplemented

    # Get the qid_shape.
    val_qid_shape = qid_shape_protocol.qid_shape(val, None)
    if val_qid_shape is None:
        return NotImplemented

    # Apply unitary effect to an identity matrix.
    result = method(ApplyUnitaryArgs.for_unitary(qid_shape=val_qid_shape))

    if result is NotImplemented or result is None:
        return result
    state_len = np.prod(val_qid_shape, dtype=np.int64)
    return result.reshape((state_len, state_len))


def _strat_unitary_from_decompose(val: Any) -> Optional[np.ndarray]:
    """Attempts to compute a value's unitary via its _decompose_ method."""
    # Check if there's a decomposition.
    operations, qubits, val_qid_shape = _try_decompose_into_operations_and_qubits(val)
    if operations is None:
        return NotImplemented

    all_qubits = frozenset(q for op in operations for q in op.qubits)
    work_qubits = frozenset(qubits)
    ancillas = tuple(sorted(all_qubits.difference(work_qubits)))

    ordered_qubits = ancillas + tuple(qubits)
    val_qid_shape = qid_shape_protocol.qid_shape(ancillas) + val_qid_shape

    # Apply sub-operations' unitary effects to an identity matrix.
    result = apply_unitaries(
        operations, ordered_qubits, ApplyUnitaryArgs.for_unitary(qid_shape=val_qid_shape), None
    )

    # Package result.
    if result is None:
        return None

    state_len = np.prod(val_qid_shape, dtype=np.int64)
    result = result.reshape((state_len, state_len))
    # Assuming borrowable qubits are restored to their original state and
    # clean qubits restord to the zero state then the desired unitary is
    # the upper left square.
    work_state_len = np.prod(val_qid_shape[len(ancillas) :], dtype=np.int64)
    return result[:work_state_len, :work_state_len]