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)
       to simulate realistic highly-redundant data or physical attractors."""
    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

def alc_factor(X, lambda_reg=1e-6):
    m, n = X.shape
    # Oversampling parameter as per specification (Phase 1)
    s = max(min(m, n) // 4, 2)  
    
    # Randomized range finder [cite: 12]
    Omega = np.random.randn(n, s) [cite: 13]
    Y = X @ Omega [cite: 14]
    
    # FIXED: Use 'reduced' mode instead of 'economic' to guarantee a 2-tuple (Q, R) return
    Q, _ = np.linalg.qr(Y, mode='reduced') [cite: 14]
    
    B = Q.T @ X [cite: 15]
    U, sigma, Vt = np.linalg.svd(B, full_matrices=False) [cite: 15, 16]
    
    # Cumulative energy
    cum_energy = np.cumsum(sigma**2) / np.sum(sigma**2) [cite: 38]
    
    # Cost function optimization (Phase 2)
    best_score = -1
    best_r = 1
    for r in range(1, len(sigma) + 1): [cite: 38]
        storage_saved = (m * n - r * (m + n)) / (m * n) [cite: 39]
        error = 1.0 - cum_energy[r - 1] [cite: 39]
        score = storage_saved / (error + lambda_reg) [cite: 39]
        if score > best_score: [cite: 39]
            best_score = score [cite: 39]
            best_r = r [cite: 39]
            
    # Low-rank factors (Aligned Q with SVD subspace U)
    L = (Q @ U[:, :best_r]) @ np.diag(sigma[:best_r]) [cite: 39]
    R = Vt[:best_r, :].T [cite: 39]
    return L, R
    
def alc_matmul(A, B, lambda_reg=1e-6):
    """ALC Compressed Multiplication (Phase 3)"""
    LA, RA = alc_factor(A, lambda_reg)
    LB, RB = alc_factor(B, lambda_reg)
    
    # Inner small core multiplication
    core = RA.T @ LB 
    # Two-step final multiplication
    C = LA @ (core @ RB.T)   
    return C

def run_benchmark(matrix_size=1500, effective_rank=30):
    print(f"=== BENCHMARK CONFIGURATION ===")
    print(f"Matrix Dimensions: {matrix_size} x {matrix_size}")
    print(f"Effective Rank:    {effective_rank}")
    print(f"===============================\n")
    
    # 1. Generate matrices
    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)
    
    # 2. Ordinary MatMul
    print("[2/3] Running Ordinary Matrix Multiplication...")
    start_time = time.time()
    C_dense = A @ B
    dense_time = time.time() - start_time
    
    # 3. ALC MatMul
    print("[3/3] Running Adaptive Linear Compression MatMul...")
    start_time = time.time()
    C_alc = alc_matmul(A, B, lambda_reg=1e-5)
    alc_time = time.time() - start_time
    
    # 4. Accuracy Assessment
    frob_error = np.linalg.norm(C_dense - C_alc) / np.linalg.norm(C_dense)
    
    # Display Summary Table
    print("\n### Benchmark Performance Results")
    print("-" * 65)
    print(f"{'Method':<20} | {'Execution Time (s)':<20} | {'Relative Error':<15}")
    print("-" * 65)
    print(f"{'Ordinary (Dense)':<20} | {dense_time:<20.5f} | {'0.00000 (Ref)':<15}")
    print(f"{'ALC (Compressed)':<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 Frobenius error of **{frob_error:.2e}**.")

if __name__ == "__main__":
    # You can scale up matrix_size (e.g., 2000, 3000) to see a wider gap in performance
    run_benchmark(matrix_size=1600, effective_rank=40)