import numpy as np
import matplotlib.pyplot as plt

# householder function from question 2
def householder_qr(A: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """
    Perform Householder QR factorization on matrix A to obtain orthogonal matrix Q.
    
    Parameters:
    A (np.ndarray): The input matrix to be orthogonalized.
    
    Returns:
    tuple[np.ndarray, np.ndarray]: The orthogonal matrix Q and upper triangular matrix R such that A = QR.
    """
    m, n = A.shape
    Q = np.eye(m, dtype=A.dtype)
    R = A.copy()
    
    for k in range(n):
        # Extract the vector x
        # Target: x -> [||x||, 0, 0, ..., 0]^T
        x = R[k:, k] # x = [a_k, a_(k+1), ..., a_m]^T
        s = np.sign(x[0]) if x[0] != 0 else 1
        e = np.zeros_like(x) # e = [||x||, 0, 0, ..., 0]^T
        e[0] = s * np.linalg.norm(x)
        # Compute the Householder vector
        # x - e is the reflection plane normal
        u = x - e
        if np.linalg.norm(u) < 1e-14:
            continue  # Skip the transformation if u is too small
        v = u / np.linalg.norm(u)
        # Apply the Householder transformation to R[k:, k:]
        R[k:, k:] -= 2 * np.outer(v, v.conj().T @ R[k:, k:])
        # Update Q
        # Q_k = I - 2 v v^*
        Q_k = np.eye(m, dtype=A.dtype)
        Q_k[k:, k:] -= 2 * np.outer(v, v.conj().T)
        Q = Q @ Q_k

    return Q, R

def householder_lstsq(A: np.ndarray, b: np.ndarray) -> np.ndarray:
    """
    Solve the least squares problem Ax = b using Householder reflections.
    
    Parameters:
    A (np.ndarray): The input matrix A.
    b (np.ndarray): The right-hand side vector b.
    
    Returns:
    np.ndarray: The solution vector x.
    """
    """
    Ax = b, A = QR
    Rx = Q^*b
    x = R^{-1}Q^*b
    R is upper triangular
    Q is orthogonal
    """
    m, n = A.shape # m is N + 1, n is n + 1
    Q, R = householder_qr(A)
    # Compute Q^*b
    Qt_b = Q.conj().T @ b
    # Solve Rx = Q^*b using back substitution
    x = np.zeros(n, dtype=A.dtype)
    for i in range(n-1, -1, -1):
        x[i] = (Qt_b[i] - R[i, i+1:] @ x[i+1:]) / R[i, i]
    return x

def f1(t: float) -> float:
    """
    f1(t) = 1 / (1 + t^2)
    """
    return 1.0 / (1.0 + t**2)

def build_A(N: int, n: int, f: callable) -> tuple[np.ndarray, np.ndarray]:
    """
    Build the matrix A for the least squares problem.
    
    Parameters:
    n + 1 (int): The number of columns in the matrix A.
    
    Returns:
    np.ndarray: The constructed matrix A.
    """
    """
    {t_j}j=0->N, t_0=-5, t_N=5, t_j divides [-5, 5] equally
    A = [[1, t_0, t_0^2, t_0^3, ..., t_0^n],
        [1, t_1, t_1^2, t_1^3, ..., t_1^n],
        ...
        [1, t_N, t_N^2, t_N^3, ..., t_N^n]] (N + 1) x (n + 1)

    b = [f(t_0), f(t_1), ..., f(t_N)]^T (N + 1) x 1
    x = [a_0, a_1, a_2, ..., a_n]^T (n + 1) x 1
    Ax = b, x is the coeffs of the polynomial
    """
    A = np.zeros((N + 1, n + 1), dtype=float)
    t = np.linspace(-5, 5, N + 1)
    for i in range(N + 1):
        A[i, :] = [t[i]**j for j in range(n + 1)]
    b = np.array([f(t_i) for t_i in t])
    return A, b

def question_3c():
    Ns = [30, 30, 30]
    ns = [5, 15, 30]
    # p(t) = a_n t^n + a_(n-1) t^(n-1) + ... + a_1 t + a_0
    # plot f(t) and p(t) on the same figure for each (N, n)
    t = np.linspace(-5, 5, 1000)
    f_t = f1(t)
    plt.figure(figsize=(15, 5))
    for i in range(len(Ns)):
        N = Ns[i]
        n = ns[i]
        A = np.zeros((N + 1, n + 1), dtype=float)
        A, b = build_A(N, n, f1)
        x = householder_lstsq(A, b)
        p_t = sum([x[j] * t**j for j in range(n + 1)])
        plt.subplot(1, 3, i + 1)
        plt.plot(t, f_t, label='f(t)', color='blue')
        plt.plot(t, p_t, label='p(t)', color='red', linestyle='--')
        plt.title(f'N={N}, n={n}')
        plt.xlabel('t')
        plt.ylabel('y')
        plt.legend()
        plt.grid(True, alpha=0.3)
    plt.suptitle('Least Squares Polynomial Approximation using Householder QR')
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    plt.savefig('Project1_A3c_least_squares_approximation.png')
    #plt.show()

def f2(n:int, t: float) -> float:
    """
    f2(t) = t^n + t^(n-1) + ... + t^2 + t + 1
    """
    return sum([t**i for i in range(n + 1)])

def question_3e():
    # n = 4, 6, 8, 10, ..., 20
    ns = list(range(4, 21, 2))
    # N = 2n
    # quantity is log10||(a_0-1,...,a_n-1)||_2)
    losses = []
    for n in ns:
        N = 2 * n
        A, b = build_A(N, n, lambda t: f2(n, t))
        x = householder_lstsq(A, b)
        true_coeffs = np.array([1.0 for _ in range(n + 1)])
        loss = np.linalg.norm(x - true_coeffs)
        losses.append(np.log10(loss))
    plt.figure(figsize=(10, 6))
    plt.plot(ns, losses, marker='o')
    plt.xlabel('Polynomial Degree n')
    plt.ylabel('log10(||a - true_coeffs||_2)')
    plt.title('Coefficient Error vs Polynomial Degree using Householder QR')
    plt.xticks(ns)
    plt.grid(True, alpha=0.3)
    plt.savefig('Project1_A3e_coefficient_error.png')
    #plt.show()


if __name__ == "__main__":
    question_3c()
    question_3e()