import numpy as np
import time

def generate_low_rank_matrix(m, n, effective_rank=20, noise_level=1e-4):
    """Generates a matrix with decaying singular values (low numerical rank)
       using NumPy just for data preparation."""
    U, _ = np.linalg.qr(np.random.randn(m, m))
    V, _ = np.linalg.qr(np.random.randn(n, n))
    
    # Linear decay of singular values
    s = np.zeros(min(m, n))
    s[:effective_rank] = np.linspace(100, 1, effective_rank)
    s[effective_rank:] = noise_level
    
    return (U[:, :min(m, n)] @ np.diag(s) @ V[:, :min(m, n)].T).tolist()

def pure_python_matmul(A, B):
    """A fair, pure Python, triple-loop matrix multiplication.
       Bypasses all underlying C-accelerated BLAS/MKL optimizations."""
    rows_A = len(A)
    cols_A = len(A[0])
    cols_B = len(B[0])
    
    # Initialize the output matrix with zeros
    C = [[0.0 for _ in range(cols_B)] for _ in range(rows_A)]
    
    # Standard O(n^3) triple loop matrix multiplication
    for i in range(rows_A):
        for k in range(cols_A):
            element_A = A[i][k]
            for j in range(cols_B):
                C[i][j] += element_A * B[k][j]
    return C

def alc_factor_python(X_list, lambda_reg=1e-6):
    """Performs ALC factorization using NumPy only for the unavoidable
       SVD/QR steps, but returns lists for pure Python handling."""
    X = np.array(X_list)
    m, n = X.shape
    s = max(min(m, n) // 4, 2)  
    
    Omega = np.random.randn(n, s)
    Y = X @ Omega
    
    Q, _ = np.linalg.qr(Y, mode='reduced')
    B = Q.T @ X
    U, sigma, Vt = np.linalg.svd(B, full_matrices=False)
    
    cum_energy = np.cumsum(sigma**2) / np.sum(sigma**2)
    
    best_score = -1
    best_r = 1
    for r in range(1, len(sigma) + 1):
        storage_saved = (m * n - r * (m + n)) / (m * n)
        error = 1.0 - cum_energy[r - 1]
        score = storage_saved / (error + lambda_reg)
        if score > best_score:
            best_score = score
            best_r = r
            
    L = (Q @ U[:, :best_r]) @ np.diag(sigma[:best_r])
    R = Vt[:best_r, :].T
    return L.tolist(), R.tolist()

def alc_matmul_pure(A, B, lambda_reg=1e-6):
    """ALC Compressed Multiplication using pure Python loops for the actual
       multiplication phases."""
    # Step 1: Compress A and B
    LA, RA = alc_factor_python(A, lambda_reg)
    LB, RB = alc_factor_python(B, lambda_reg)
    
    # Step 2: Perform the inner core multiplications via fair pure Python loops
    # Transpose RA in pure Python to get RA.T
    RA_T = [[RA[j][i] for j in range(len(RA))] for i in range(len(RA[0]))]
    core = pure_python_matmul(RA_T, LB)
    
    # Transpose RB in pure Python to get RB.T
    RB_T = [[RB[j][i] for j in range(len(RB))] for i in range(len(RB[0]))]
    intermediate = pure_python_matmul(core, RB_T)
    
    C = pure_python_matmul(LA, intermediate)
    return C

def run_benchmark(matrix_size=180, effective_rank=10):
    # NOTE: Matrix size is reduced compared to the C version because pure 
    # Python triple loops are exponentially slower.
    print(f"=== FAIR PURE PYTHON BENCHMARK CONFIGURATION ===")
    print(f"Matrix Dimensions: {matrix_size} x {matrix_size}")
    print(f"Effective Rank:    {effective_rank}")
    print(f"================================================\n")
    
    print("[1/3] Generating structured low-rank matrices...")
    A = generate_low_rank_matrix(matrix_size, matrix_size, effective_rank=effective_rank)
    B = generate_low_rank_matrix(matrix_size, matrix_size, effective_rank=effective_rank)
    
    print("[2/3] Running Fair Pure Python Ordinary MatMul...")
    start_time = time.time()
    C_dense = pure_python_matmul(A, B)
    dense_time = time.time() - start_time
    
    print("[3/3] Running Fair Pure Python ALC MatMul...")
    start_time = time.time()
    C_alc = alc_matmul_pure(A, B, lambda_reg=1e-5)
    alc_time = time.time() - start_time
    
    # Accuracy Assessment
    frob_error = np.linalg.norm(np.array(C_dense) - np.array(C_alc)) / np.linalg.norm(np.array(C_dense))
    
    print("\n### Benchmark Performance Results (Pure Python Loops)")
    print("-" * 65)
    print(f"{'Method':<20} | {'Execution Time (s)':<20} | {'Relative Error':<15}")
    print("-" * 65)
    print(f"{'Ordinary (Pure Python)':<20} | {dense_time:<20.5f} | {'0.00000 (Ref)':<15}")
    print(f"{'ALC (Pure Python)':<20} | {alc_time:<20.5f} | {frob_error:<15.2e}")
    print("-" * 65)
    
    speedup = dense_time / alc_time
    print(f"\n> **Result**: ALC achieved a **{speedup:.2f}x speedup** with a relative error of **{frob_error:.2e}**.")

if __name__ == "__main__":
    run_benchmark(matrix_size=200, effective_rank=10)