import numpy as np
import time
import matplotlib.pyplot as plt

def cgs_q(A: np.ndarray) -> np.ndarray:
    """
    Perform Classical Gram-Schmidt orthogonalization on matrix A.
    
    Parameters:
    A (np.ndarray): The input matrix to be orthogonalized.
    
    Returns:
    np.ndarray: The orthogonalized matrix Q.
    """
    m, n = A.shape
    Q = np.zeros((m, n), dtype=A.dtype)
    
    for j in range(n):
        v = A[:, j]
        for i in range(j):
            # r_ij = q_i^* a_j
            r_ij = np.dot(Q[:, i].conj(), A[:, j])
            # v = a_j - r_ij * q_i
            v = v - r_ij * Q[:, i]
        # v = a_j - sum_{i=1}^{j-1} r_ij * q_i
        # r_jj = ||v||
        r_jj = np.linalg.norm(v)
        #if r_jj < 1e-14:
        #    print(f"Warning: Vector {j} is close to zero after orthogonalization.")
        Q[:, j] = v / r_jj
    
    return Q

def mgs_q(A: np.ndarray) -> np.ndarray:
    """
    Perform Modified Gram-Schmidt orthogonalization on matrix A.
    
    Parameters:
    A (np.ndarray): The input matrix to be orthogonalized.
    
    Returns:
    np.ndarray: The orthogonalized matrix Q.
    """
    m, n = A.shape
    Q = np.zeros((m, n), dtype=A.dtype)
    V = A.copy()
    
    for j in range(n):
        for i in range(j):
            # r_ij = q_i^* a_j
            r_ij = np.dot(Q[:, i].conj(), V[:, j])
            # a_j = a_j - r_ij * q_i
            V[:, j] = V[:, j] - r_ij * Q[:, i]
        # r_jj = ||a_j||
        r_jj = np.linalg.norm(V[:, j])
        #if r_jj < 1e-14:
        #    print(f"Warning: Vector {j} is close to zero after orthogonalization.")
        Q[:, j] = V[:, j] / r_jj
    return Q

def householder_q(A: 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:
    np.ndarray: The orthogonalized matrix Q.
    """
    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

def cgs_twice_q(A: np.ndarray) -> np.ndarray:
    """
    Perform Classical Gram-Schmidt orthogonalization twice on matrix A.
    
    Parameters:
    A (np.ndarray): The input matrix to be orthogonalized.
    
    Returns:
    np.ndarray: The orthogonalized matrix Q.
    """
    return cgs_q(cgs_q(A))

def ortho_loss(Q: np.ndarray) -> float:
    """
    Compute the orthogonality loss of matrix Q.
    
    Parameters:
    Q (np.ndarray): The orthogonal matrix to be evaluated.
    
    Returns:
    float: The orthogonality loss defined as log10(||I - Q^* Q||_F).
    """
    m, n = Q.shape
    I = np.eye(n, dtype=Q.dtype)
    E = I - Q.conj().T @ Q
    loss = np.linalg.norm(E, ord='fro')
    return np.log10(loss)

def build_hilbert(m: int, n: int) -> np.ndarray:
    """
    Build an m x n Hilbert matrix.
    
    Parameters:
    m (int): Number of rows.
    n (int): Number of columns.
    
    Returns:
    np.ndarray: The m x n Hilbert matrix.
    """
    H = np.zeros((m, n), dtype=float)
    for i in range(m):
        for j in range(n):
            H[i, j] = 1.0 / (i + j + 1)
    return H


def question_2bcd(n_min=2, n_max=20):
    # Question 2(b)
    ns = list(range(n_min, n_max + 1))
    methods = {
        "Classical Gram-Schmidt": cgs_q,
        "Modified Gram-Schmidt": mgs_q,
        "Householder": householder_q,
        "Classical Gram-Schmidt Twice": cgs_twice_q
    }
    losses = {name: [] for name in methods.keys()}
    times = {name: [] for name in methods.keys()}

    for n in ns:
        H = build_hilbert(n, n)
        # run ten times and take the average
        times_temp = {name: [] for name in methods.keys()}
        losses_temp = {name: [] for name in methods.keys()}
        for _ in range(10):
            for name, method in methods.items():
                start_time = time.time()
                Q = method(H)
                end_time = time.time()
                loss = ortho_loss(Q)
                times_temp[name].append(end_time - start_time)
                losses_temp[name].append(loss)
        for name in methods.keys():
            times[name].append(np.mean(times_temp[name]))
            losses[name].append(np.log(np.mean(np.exp(losses_temp[name]))))
            #print(f"n={n}, Method={name}, Loss={loss}, Time={end_time - start_time}s")
    
    # Plot the orthogonality loss
    plt.figure(figsize=(10, 6))
    for name, loss in losses.items():
        plt.plot(ns, loss, marker='o', label=name)
    plt.xlabel('Matrix Size n')
    plt.ylabel('Orthogonality Loss (log10 scale)')
    plt.title('Orthogonality Loss vs Matrix Size for Different QR Methods')
    # set x ticks to be integers
    plt.xticks(ns)
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig('Project1_A2d_orthogonality_loss.png')
    #plt.show()
    # Plot the time taken
    plt.figure(figsize=(10, 6))
    for name, time_list in times.items():
        plt.plot(ns, time_list, marker='o', label=name)
    plt.xlabel('Matrix Size n')
    plt.ylabel('Time Taken (seconds)')
    plt.title('Time Taken vs Matrix Size for Different QR Methods')
    # set x ticks to be integers
    plt.xticks(ns)
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig('Project1_A2d_time_taken.png')
    #plt.show()

    # Table of results (n=4, 9, 11)
    print("\nTable of Orthogonality Loss and Time for n=4, 9, 11:")
    print(f"{'n':<5} {'Method':<30} {'Loss':<20} {'Time (s)':<10}")
    for n in [4, 9, 11]:
        for name in methods.keys():
            idx = ns.index(n)
            print(f"{n:<5} {name:<30} {losses[name][idx]:<20} {times[name][idx]:<10}")
            
if __name__ == "__main__":
    question_2bcd()