"""Mutually unbiased basis states.

If a system prepared in an eigenstate of one of the bases gives an equal probability of (1/d) when measured with respect
to the other bases, mutually unbiased basis states are orthonormal bases in the Hilbert space Cᵈ.
"""

import math

import numpy as np

from toqito.matrices import gen_pauli


def mutually_unbiased_basis(dim: int) -> list[np.ndarray]:
    r"""Generate list of MUBs for a given dimension [@wikipediamub].

    Note that this function only works if the dimension provided is prime or a power of a prime. Otherwise, we don't
    know how to generate general MUBs.

    Examples:
        For the case of dimension 2, the three mutually unbiased bases are provided by:

        \[
            M_0 = \left\{|0\rangle, |1\rangle \right\}, \\
            M_1 = \left\{\frac{|0\rangle + |1\rangle}{\sqrt{2}}, \frac{|0\rangle - |1\rangle}{\sqrt{2}}\right\}
            M_2 = \left\{\frac{|0\rangle + i|1\rangle}{\sqrt{2}}, \frac{|0\rangle - i|1\rangle}{\sqrt{2}}\right\}
        \]

        The six vectors above are obtained accordingly:

        ```python exec="1" source="above" session="mubs"
        from toqito.states import mutually_unbiased_basis
        mubs = mutually_unbiased_basis(2)
        print(len(mubs))
        ```

        ```python exec="1" source="above" session="mubs"
        lst =[vec.shape for vec in mubs]
        print(lst)
        ```
    Args:
        dim: The dimension of the mutually unbiased bases to produce.

    Returns:
        The set of mutually unbiased bases of dimension `dim` (if known).

    """
    # The first basis will always be the standard basis:
    mats = [np.eye(dim)]

    pauli_x = gen_pauli(1, 0, dim)
    if _is_prime(dim):
        pauli_z = gen_pauli(0, 1, dim)

        for j in range(dim, 0, -1):
            _, eigen_vec = np.linalg.eig(pauli_x @ pauli_z ** (j))
            mats.append(eigen_vec)
    elif _is_prime_power(dim) and not _is_prime(dim):
        raise ValueError(f"Dimension {dim} is a prime power but not prime (more complicated no support at the moment).")
    else:
        raise ValueError(f"No general construction of MUBs is known for dimension: {dim}.")

    mubs: list[np.ndarray] = []
    for mat in mats:
        nrows, ncols = mat.shape[0], mat.shape[1]
        for row in range(nrows):
            mub: list[np.ndarray] = []
            for col in range(ncols):
                mub.append(mat[col][row])
            mubs.append(np.array(mub))
    return mubs


def _is_prime(n: int) -> bool:
    """Check if n is a prime number."""
    return n >= 2 and all(n % i != 0 for i in range(2, math.isqrt(n) + 1))


def _is_prime_power(n: int) -> bool:
    """Determine if a number is a prime power.

    A number is a prime power if it can be written as p^k, where p is a prime number and k is an integer greater than 0.

    Args:
        n: An integer to check for being a prime power.

    Returns:
        True if n is a prime power, False otherwise.

    """
    if n < 2:
        return False
    for p in range(2, math.isqrt(n) + 1):
        if n % p == 0:
            if not _is_prime(p):
                continue
            while n % p == 0:
                n //= p
            return n == 1
    return True
