import numpy as np

def proj(A:np.ndarray, b:np.ndarray) -> np.ndarray:
    """
    Project vector b onto the column space of matrix A using the formula:
    proj_A(b) = A(A^* A)^(-1) A^* b

    Parameters:
    A (np.ndarray): The matrix whose column space we are projecting onto.
    b (np.ndarray): The vector to be projected.

    Returns:
    np.ndarray: The projection of b onto the column space of A.
    """
    # Check if A is full rank
    if np.linalg.matrix_rank(A) < min(A.shape):
        raise ValueError("Matrix A must be full rank.")
    # Compute the projection
    AhA_inv = np.linalg.inv(A.conj().T @ A)
    projection = A @ AhA_inv @ A.conj().T @ b
    return projection

def proj_SVD(A:np.ndarray, b:np.ndarray) -> np.ndarray:
    """
    Project vector b onto the column space of matrix A using SVD.

    Parameters:
    A (np.ndarray): The matrix whose column space we are projecting onto.
    b (np.ndarray): The vector to be projected.

    Returns:
    np.ndarray: The projection of b onto the column space of A.
    """
    # A = U S V^*
    U, S, Vh = np.linalg.svd(A, full_matrices=False)
    """
    proj_A(b) = A(A^* A)^(-1) A^* b
    = U S V^* (V S^2 V^*)^(-1) V S U^* b
    = U S V^* V S^(-2) V^* V S U^* b
    = U U^* b
    If A = U S V^*, then the projection onto the column space of A is:
        proj_A(b) = U_r U_r^* b
    where U_r are the left singular vectors corresponding to nonzero singular values.
    """
    U_r = U[:, :np.linalg.matrix_rank(A)]  # Take only the relevant columns
    # Compute the projection using the SVD components
    projection = U_r @ U_r.conj().T @ b
    return projection

def build_A(eps: float) -> np.ndarray:
    """
    Build the matrix A(ε) as specified in the problem statement.
    """
    A = np.array([
        [1.0, 1.0, 0.0, (1.0 - eps)],
        [1.0, 0.0, 1.0, (eps + (1.0 - eps))],
        [0.0, 1.0, 0.0, eps],
        [1.0, 0.0, 0.0, (1.0 - eps)],
        [1.0, 0.0, 0.0, (eps + (1.0 - eps))]
    ], dtype=float)
    return A

def build_b() -> np.ndarray:
    """
    Build the vector b as specified in the problem statement.
    """
    b = np.array([2.0, 1.0, 3.0, 1.0, 4.0])
    return b

def question_1a(A: np.ndarray, b: np.ndarray):
    # Question 1(a)
    print("Question 1(a):")
    projection = proj(A, b)
    print("Projection of b onto the column space of A(1): (Using proj function)")
    print(projection)
    projection_svd = proj_SVD(A, b)
    print("Projection of b onto the column space of A(1): (Using proj_SVD function)")
    print(projection_svd)
    print("Difference between the two methods:")
    print(projection - projection_svd)

def question_1b(A: np.ndarray, b: np.ndarray):
    # Question 1(b)
    print("\nQuestion 1(b):")
    for eps in [1.0, 1e-1, 1e-2, 1e-4, 1e-8, 1e-16, 1e-32]:
        print(f"\nFor ε = {eps}:")
        A_eps = build_A(eps)
        try:
            projection_eps = proj(A_eps, b)
            print(f"Projection of b onto the column space of A({eps}):")
            print(projection_eps)
        except np.linalg.LinAlgError as e:
            print(f"LinAlgError for eps={eps}: {e}")
        except ValueError as ve:
            print(f"ValueError for eps={eps}: {ve}")
        projection_svd_eps = proj_SVD(A_eps, b)
        print(f"Projection of b onto the column space of A({eps}) using SVD:")
        print(projection_svd_eps)
        print("Difference between the two methods:")
        try:
            print(projection_eps - projection_svd_eps)
        except:
            print("Could not compute difference due to previous error.")
        projection_eps = None  # Reset for next iteration
    
if __name__ == "__main__":
    A = build_A(eps=1.0)
    b = build_b()
    question_1a(A, b)
    question_1b(A, b)
    
