import numpy as np
import time
import warnings

# Suppress any unexpected legacy matrix warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

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
    Omega = np.random.randn(n, s)
    Y = X @ Omega
    
    # FIXED: Using 'reduced' mode ensures a reliable 2-tuple (Q, R) unpacking across all NumPy versions
    Q, _ = np.linalg.qr(Y, mode='reduced')
    
    B = Q.T @ X
    U, sigma, Vt = np.linalg.svd(B, full_matrices=False)
    
    # Cumulative energy
    cum_energy = np.cumsum(sigma**2) / np.sum(sigma**2)
    
    # Cost function optimization (Phase 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
            
    # Low-rank factors (Aligned Q with SVD subspace U)
    L = (Q @ U[:, :best_r]) @ np.diag(sigma[:best_r])
    R = Vt[:best_r, :].T
    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__":
    run_benchmark(matrix_size=1600, effective_rank=40)