"""This module implements the UnitaryBuilder class."""
from __future__ import annotations

import logging
from typing import cast
from typing import Sequence
from typing import TYPE_CHECKING

import numpy as np
import numpy.typing as npt

from bqskit.qis.unitary.unitary import RealVector
from bqskit.qis.unitary.unitary import Unitary
from bqskit.qis.unitary.unitarymatrix import UnitaryMatrix
from bqskit.utils.typing import is_integer
from bqskit.utils.typing import is_valid_radixes

if TYPE_CHECKING:
    from bqskit.ir.location import CircuitLocationLike

logger = logging.getLogger(__name__)


class UnitaryBuilder(Unitary):
    """
    An object for fast unitary accumulation using tensor networks.

    A UnitaryBuilder is similar to a StringBuilder in the sense that it is an
    efficient way to string together or accumulate :class:`Unitary` objects.
    This class uses concepts from tensor networks to efficiently multiply
    unitary matrices.
    """

    def __init__(self, num_qudits: int, radixes: Sequence[int] = []) -> None:
        """
        UnitaryBuilder constructor.

        Args:
            num_qudits (int): The number of qudits to build a Unitary for.

            radixes (Sequence[int]): A sequence with its length equal
                to `num_qudits`. Each element specifies the base of a
                qudit. Defaults to qubits.

        Raises:
            ValueError: If `num_qudits` is nonpositive.

            ValueError: If the length of `radixes` is not equal to
                `num_qudits`.

        Examples:
            >>> builder = UnitaryBuilder(4)  # Creates a 4-qubit builder.
        """

        if not is_integer(num_qudits):
            raise TypeError(
                'Expected int for num_qudits, got %s.' %
                type(num_qudits),
            )

        if num_qudits <= 0:
            raise ValueError(
                'Expected positive number for num_qudits, got %d.' %
                num_qudits,
            )

        self._num_qudits = num_qudits
        self._radixes = tuple(radixes if len(radixes) > 0 else [2] * num_qudits)

        if not is_valid_radixes(self.radixes):
            raise TypeError('Invalid qudit radixes.')

        if len(self.radixes) != self.num_qudits:
            raise ValueError(
                'Expected length of radixes to be equal to num_qudits:'
                ' %d != %d' % (len(self.radixes), self.num_qudits),
            )

        self._num_params = 0
        self._dim = int(np.prod(self.radixes))
        self.tensor = np.identity(self.dim, dtype=np.complex128)
        self.tensor = self.tensor.reshape(self.radixes * 2)

    def get_unitary(self, params: RealVector = []) -> UnitaryMatrix:
        """Build the unitary, see :func:`Unitary.get_unitary` for more."""
        utry = self.tensor.reshape((self.dim, self.dim))
        return UnitaryMatrix(utry, self.radixes, False)

    def apply_right(
        self,
        utry: UnitaryMatrix,
        location: CircuitLocationLike,
        inverse: bool = False,
        check_arguments: bool = True,
    ) -> None:
        """
        Apply the specified unitary on the right of this UnitaryBuilder.

        ..
                 .-----.   .------.
              0 -|     |---|      |-
              1 -|     |---| utry |-
                 .     .   '------'
                 .     .
                 .     .
            n-1 -|     |------------
                 '-----'

        Args:
            utry (UnitaryMatrix): The unitary to apply.

            location (CircuitLocationLike): The qudits to apply the unitary on.

            inverse (bool): If true, apply the inverse of the unitary.

            check_arguments (bool): If true, check the inputs for type and
                value errors.

        Raises:
            ValueError: If `utry`'s size does not match the given location.

            ValueError: if `utry`'s radixes does not match the given location.

        Notes:
            - Applying the unitary on the right is equivalent to multiplying
              the unitary on the left of the tensor. The notation comes
              from the quantum circuit perspective.

            - This operation is performed using tensor contraction.
        """
        from bqskit.ir.location import CircuitLocation
        if check_arguments:
            if not isinstance(utry, UnitaryMatrix):
                raise TypeError('Expected UnitaryMatrix, got %s', type(utry))

            if not CircuitLocation.is_location(location, self.num_qudits):
                raise TypeError('Invalid location.')

            location = CircuitLocation(location)

            if len(location) != utry.num_qudits:
                raise ValueError('Unitary and location size mismatch.')

            for utry_radix, bldr_radix_idx in zip(utry.radixes, location):
                if utry_radix != self.radixes[bldr_radix_idx]:
                    raise ValueError('Unitary and location radix mismatch.')

        left_perm = list(cast(CircuitLocation, location))
        mid_perm = [x for x in range(self.num_qudits) if x not in left_perm]
        right_perm = [x + self.num_qudits for x in range(self.num_qudits)]

        left_dim = int(np.prod([self.radixes[x] for x in left_perm]))

        utry = utry.dagger if inverse else utry

        perm = left_perm + mid_perm + right_perm
        self.tensor = self.tensor.transpose(perm)
        self.tensor = self.tensor.reshape((left_dim, -1))
        self.tensor = utry @ self.tensor

        shape = list(self.radixes) * 2
        shape = [shape[p] for p in perm]
        self.tensor = self.tensor.reshape(shape)
        inv_perm = list(np.argsort(perm))
        self.tensor = self.tensor.transpose(inv_perm)

    def apply_left(
        self,
        utry: UnitaryMatrix,
        location: CircuitLocationLike,
        inverse: bool = False,
        check_arguments: bool = True,
    ) -> None:
        """
        Apply the specified unitary on the left of this UnitaryBuilder.

        ..
                 .------.   .-----.
              0 -|      |---|     |-
              1 -| gate |---|     |-
                 '------'   .     .
                            .     .
                            .     .
            n-1 ------------|     |-
                            '-----'

        Args:
            utry (UnitaryMatrix): The unitary to apply.

            location (CircuitLocationLike): The qudits to apply the unitary on.

            inverse (bool): If true, apply the inverse of the unitary.

            check_arguments (bool): If true, check the inputs for type and
                value errors.

        Raises:
            ValueError: If `utry`'s size does not match the given location.

            ValueError: if `utry`'s radixes does not match the given location.

        Notes:
            - Applying the unitary on the left is equivalent to multiplying
              the unitary on the right of the tensor. The notation comes
              from the quantum circuit perspective.

            - This operation is performed using tensor contraction.
        """
        from bqskit.ir.location import CircuitLocation
        if check_arguments:
            if not isinstance(utry, UnitaryMatrix):
                raise TypeError('Expected UnitaryMatrix, got %s', type(utry))

            if not CircuitLocation.is_location(location, self.num_qudits):
                raise TypeError('Invalid location.')

            location = CircuitLocation(location)

            if len(location) != utry.num_qudits:
                raise ValueError('Unitary and location size mismatch.')

            for utry_radix, bldr_radix_idx in zip(utry.radixes, location):
                if utry_radix != self.radixes[bldr_radix_idx]:
                    raise ValueError('Unitary and location radix mismatch.')

        location = cast(CircuitLocation, location)
        left_perm = list(range(self.num_qudits))
        mid_perm = [
            x + self.num_qudits
            for x in left_perm
            if x not in location
        ]
        right_perm = [x + self.num_qudits for x in location]

        right_dim = int(
            np.prod([
                self.radixes[x - self.num_qudits]
                for x in right_perm
            ]),
        )

        utry = utry.dagger if inverse else utry

        perm = left_perm + mid_perm + right_perm
        self.tensor = self.tensor.transpose(perm)
        self.tensor = self.tensor.reshape((-1, right_dim))
        self.tensor = self.tensor @ utry

        shape = list(self.radixes) * 2
        shape = [shape[p] for p in perm]
        self.tensor = self.tensor.reshape(shape)
        inv_perm = list(np.argsort(perm))
        self.tensor = self.tensor.transpose(inv_perm)

    def eval_apply_right(
        self,
        M: npt.NDArray[np.complex128],
        location: CircuitLocationLike,
    ) -> npt.NDArray[np.complex128]:
        """
        Evaluate the application of `M` on the right of this UnitaryBuilder.

        See :func:`apply_right` for more info.
        """
        from bqskit.ir.location import CircuitLocation
        left_perm = list(cast(CircuitLocation, location))
        mid_perm = [x for x in range(self.num_qudits) if x not in left_perm]
        right_perm = [x + self.num_qudits for x in range(self.num_qudits)]

        left_dim = int(np.prod([self.radixes[x] for x in left_perm]))

        perm = left_perm + mid_perm + right_perm
        tensor_copy = self.tensor.copy()
        tensor_copy = tensor_copy.transpose(perm)
        tensor_copy = tensor_copy.reshape((left_dim, -1))
        tensor_copy = M @ tensor_copy  # TODO: Require out matrix to avoid copy

        shape = list(self.radixes) * 2
        shape = [shape[p] for p in perm]
        tensor_copy = tensor_copy.reshape(shape)
        inv_perm = list(np.argsort(perm))
        tensor_copy = tensor_copy.transpose(inv_perm)
        out_M = tensor_copy.reshape((self.dim, self.dim))
        return out_M

    def eval_apply_left(
        self,
        M: npt.NDArray[np.complex128],
        location: CircuitLocationLike,
    ) -> npt.NDArray[np.complex128]:
        """
        Evaluate the application of `M` on the left of this UnitaryBuilder.

        See :func:`apply_left` for more info.
        """
        from bqskit.ir.location import CircuitLocation
        location = cast(CircuitLocation, location)
        left_perm = list(range(self.num_qudits))
        mid_perm = [
            x + self.num_qudits
            for x in left_perm
            if x not in location
        ]
        right_perm = [x + self.num_qudits for x in location]

        right_dim = int(
            np.prod([
                self.radixes[x - self.num_qudits]
                for x in right_perm
            ]),
        )

        perm = left_perm + mid_perm + right_perm
        tensor_copy = self.tensor.copy()
        tensor_copy = tensor_copy.transpose(perm)
        tensor_copy = tensor_copy.reshape((-1, right_dim))
        tensor_copy = tensor_copy @ M

        shape = list(self.radixes) * 2
        shape = [shape[p] for p in perm]
        tensor_copy = tensor_copy.reshape(shape)
        inv_perm = list(np.argsort(perm))
        tensor_copy = tensor_copy.transpose(inv_perm)
        out_M = tensor_copy.reshape((self.dim, self.dim))
        return out_M

    def calc_env_matrix(
            self, location: Sequence[int],
    ) -> npt.NDArray[np.complex128]:
        """
        Calculates the environment matrix w.r.t. the specified location.

        Args:
            location (Sequence[int]): Calculate the environment matrix with
                respect to the qudit indices in location.

        Returns:
            np.ndarray: The environmental matrix.
        """

        left_perm = list(range(self.num_qudits))
        left_perm = [x for x in left_perm if x not in location]
        left_perm = left_perm + [x + self.num_qudits for x in left_perm]
        right_perm = list(location) + [x + self.num_qudits for x in location]

        perm = left_perm + right_perm
        a = np.transpose(self.tensor, perm)
        a = np.reshape(
            a, (
                2 ** (self.num_qudits - len(location)),
                2 ** (self.num_qudits - len(location)),
                2 ** len(location),
                2 ** len(location),
            ),
        )
        return np.trace(a)
