import numpy as np
from numpy.typing import NDArray

from qoptcraft.math import commutator


def is_unitary(matrix: NDArray, tol=1e-10) -> bool:
    """Check if a matrix is unitary."""
    identity = np.eye(matrix.shape[0])
    return np.allclose(matrix.conj().T @ matrix, identity, atol=tol) and np.allclose(
        matrix @ matrix.conj().T, identity, atol=tol
    )


def is_hermitian(matrix: NDArray, rtol=1e-5) -> bool:
    """Check if a matrix is Hermitian."""
    return np.allclose(matrix, matrix.conj().T, rtol=rtol)


class NonHermitianError(Exception):
    """Raised when a matrix is not Hermitian."""

    def __init__(self, matrix: NDArray, atol: float, rtol: float):
        message = (
            f"Matrix is not Hermitian within tolerance {rtol = }, {atol= }.\n"
            f"Matrix:\n{matrix}\n"
            f"Conjugate Transpose:\n{matrix.conj().T}\n"
        )
        super().__init__(message)


class NonSymmetricError(Exception):
    """Raised when a matrix is not Hermitian."""

    def __init__(self, matrix: NDArray, atol: float, rtol: float):
        message = (
            f"Matrix is not Hermitian within tolerance {rtol = }, {atol= }.\n"
            f"Matrix:\n{matrix}\n"
            f"Conjugate Transpose:\n{matrix.T}\n"
        )
        super().__init__(message)


def assert_hermitian(matrix: NDArray, atol=1e-8, rtol=1e-5) -> None:
    """Assert that a matrix is Hermitian.

    Raises:
        NonHermitianError: If the matrix is not Hermitian within the given tolerance.
    """
    if not np.allclose(matrix, matrix.conj().T, atol=atol, rtol=rtol):
        raise NonHermitianError(matrix, atol, rtol)


def assert_symmetric(matrix: NDArray, atol=1e-8, rtol=1e-5) -> None:
    """Assert that a matrix is Hermitian.

    Raises:
        NonHermitianError: If the matrix is not Hermitian within the given tolerance.
    """
    if not np.allclose(matrix, matrix.T, atol=atol, rtol=rtol):
        raise NonSymmetricError(matrix, atol, rtol)


def assert_equivalent_unitaries(unitary_1, unitary_2, atol=1e-10, rtol=1e-7):
    prod = unitary_1.conj().T @ unitary_2
    return np.testing.assert_allclose(
        prod * np.abs(prod[0,0]) / prod[0,0], np.eye(unitary_1.shape[0]), atol=atol, rtol=rtol,
        err_msg="Unitaries are not equivalent up to global phase."
    )


def assert_commute(matrix_1: NDArray, matrix_2: NDArray, tol=1e-10):
    commutator(matrix_1, matrix_2)
    if not np.allclose(commutator(matrix_1, matrix_2), np.zeros_like(matrix_1), rtol=tol):
        raise ValueError("Matrices don't commute.")


def are_same_subspace(A_basis, B_basis, tol=1e-6):
    # Combine columns from both bases
    A_basis = [matrix.flatten() for matrix in A_basis]
    A_basis = np.array(A_basis).T
    B_basis = [matrix.flatten() for matrix in B_basis]
    B_basis = np.array(B_basis).T
    C = np.hstack((A_basis, B_basis))

    # Compute ranks
    rank_A = np.linalg.matrix_rank(A_basis, tol)
    rank_B = np.linalg.matrix_rank(B_basis, tol)
    rank_C = np.linalg.matrix_rank(C, tol)

    print(f"rank(A) = {rank_A}, rank(B) = {rank_B}, rank(C) = {rank_C}")

    # If all ranks are equal, spans are the same
    return rank_A == rank_B == rank_C
