from .linalg import Matrix, Vector, identity, invert_matrix, matmul, matvec, np, scale_matrix, transpose


def _empty_matrix(matrix: Matrix) -> bool:
    if np is not None and hasattr(matrix, "size"):
        return int(matrix.size) == 0
    return not matrix


def ridge_regression_readout(
    states: list[Vector],
    targets: list[Vector],
    *,
    regularization: float,
) -> Matrix:
    if not states or not targets:
        raise ValueError("States and targets must be non-empty for ridge readout.")
    if np is not None:
        state_matrix = np.asarray(states, dtype=np.float64).T
        target_matrix = np.asarray(targets, dtype=np.float64).T
        gram = state_matrix @ state_matrix.T
        regularized = gram + (regularization * np.eye(gram.shape[0], dtype=np.float64))
        cross_covariance = target_matrix @ state_matrix.T
        return np.linalg.solve(regularized.T, cross_covariance.T).T.tolist()

    state_matrix = transpose(states)
    target_matrix = transpose(targets)
    gram = matmul(state_matrix, transpose(state_matrix))
    regularized = [
        [
            gram[row][col] + (regularization if row == col else 0.0)
            for col in range(len(gram[row]))
        ]
        for row in range(len(gram))
    ]
    inverse = invert_matrix(regularized)
    cross_covariance = matmul(target_matrix, transpose(state_matrix))
    return matmul(cross_covariance, inverse)


def ridge_regression_readout_from_moments(
    gram: Matrix,
    cross_covariance: Matrix,
    *,
    regularization: float,
) -> Matrix:
    if _empty_matrix(gram) or _empty_matrix(cross_covariance):
        raise ValueError("Gram and cross-covariance moments must be non-empty for ridge readout.")
    if np is not None:
        gram_array = np.asarray(gram, dtype=np.float64)
        regularized = gram_array + (regularization * np.eye(gram_array.shape[0], dtype=np.float64))
        cross_covariance_array = np.asarray(cross_covariance, dtype=np.float64)
        return np.linalg.solve(regularized.T, cross_covariance_array.T).T

    regularized = [
        [
            gram[row][col] + (regularization if row == col else 0.0)
            for col in range(len(gram[row]))
        ]
        for row in range(len(gram))
    ]
    inverse = invert_matrix(regularized)
    return matmul(cross_covariance, inverse)


def ridge_regression_readout_from_diagonal_moments(
    feature_second_moment: Vector,
    cross_covariance: Matrix,
    *,
    regularization: float,
) -> Matrix:
    if _empty_matrix(feature_second_moment) or _empty_matrix(cross_covariance):
        raise ValueError("Diagonal moments and cross-covariance must be non-empty for ridge readout.")
    if np is not None:
        denominator = np.asarray(feature_second_moment, dtype=np.float64) + regularization
        denominator = np.where(np.abs(denominator) > 1e-12, denominator, regularization)
        cross_covariance_array = np.asarray(cross_covariance, dtype=np.float64)
        return cross_covariance_array / denominator[None, :]

    denominator = [
        value + regularization if abs(value + regularization) > 1e-12 else regularization
        for value in feature_second_moment
    ]
    return [
        [
            value / denominator[col]
            for col, value in enumerate(row)
        ]
        for row in cross_covariance
    ]


def apply_readout(weights: Matrix, state: Vector) -> Vector:
    return matvec(weights, state)