import numpy as np
from numpy.typing import NDArray

from qoptcraft.math import trace_product


def projection_subspace(matrix: NDArray, basis: list[NDArray]) -> NDArray:
    """Projection of a matrix onto the basis of a subspace."""
    projection = np.zeros_like(matrix)
    matrix_conj = matrix.conj()
    for bi in basis:
        projection += trace_product(matrix_conj, bi) * bi
    return projection


def get_coordinates(matrix: NDArray, basis: list[NDArray]) -> NDArray:
    """Coefficients of the projection of a matrix onto a subspace."""
    matrix_conj = matrix.conj()
    coefs = []
    for bi in basis:
        coefs.append(trace_product(matrix_conj, bi))
    return np.array(coefs)


def matrix_from_coordinates(coords, basis):
    matrix = np.zeros_like(basis[0])
    for coord, basis_matrix in zip(coords, basis, strict=True):
        matrix += coord * basis_matrix
    return matrix
