import numpy as np
import time

# Reproducible setup
np.random.seed(42)
X = np.random.randn(100, 784).astype(np.float32)
W = np.random.randn(784, 100).astype(np.float32)

def benchmark(func, name="Method"):
    start = time.perf_counter()
    result = func()
    elapsed = time.perf_counter() - start
    print(f"{name:40} {elapsed*1000:8.2f} ms")
    return result, elapsed

def method_numpy():
    return np.dot(X, W)

result_baseline, _ = benchmark(method_numpy, "NumPy dot (baseline)")

def method_nikhilam_quantized():
    """
    Use Nikhilam's deficiency concept.
    Quantize to int8 for speed, then scale back.
    """
    # Find base and deficiency for quantization
    scale = 127 / max(abs(X).max(), abs(W).max())
    
    X_q = np.round(X * scale).astype(np.int8)
    W_q = np.round(W * scale).astype(np.int8)
    
    # int8 dot product is much faster
    result_q = np.dot(X_q, W_q).astype(np.float32)
    
    # Scale back to float
    return result_q / (scale ** 2)

result_nikhilam, _ = benchmark(method_nikhilam_quantized, "Nikhilam quantized (int8)")

def method_vedic_chunking():
    """
    Inspired by Vedic block multiplication.
    Break 784 into 28 chunks of 28 for better cache locality.
    """
    chunk_size = 28
    n_chunks = 784 // chunk_size
    result = np.zeros((100, 100), dtype=np.float32)
    
    for i in range(n_chunks):
        start = i * chunk_size
        end = start + chunk_size
        X_chunk = X[:, start:end]
        W_chunk = W[start:end, :]
        result += np.dot(X_chunk, W_chunk)
    
    return result

result_vedic, _ = benchmark(method_vedic_chunking, "Vedic chunking (28-block)")

def method_paravartya_transpose():
    """
    Paravartya: Transpose and apply.
    Layout W in contiguous memory for better access patterns.
    """
    W_T = W.T  # 100 x 784 (contiguous)
    
    # Process row by row with pre-transposed weight
    result = np.zeros((100, 100), dtype=np.float32)
    
    for i in range(100):
        # Each row of X dotted with transposed W
        result[i] = np.dot(X[:, i], W_T)  # Wait, that doesn't work
    
    # Correct approach:
    result = np.zeros((100, 100), dtype=np.float32)
    for j in range(100):
        result[:, j] = X @ W[:, j]  # Each column of W
    
    return result

# Actually, simpler version:
def method_paravartya_fast():
    """Transpose for memory contiguity, then standard dot."""
    W_T = W.T.copy()  # 100 x 784, now row-major access
    return X @ W_T.T  # Back to original multiplication with warm cache

result_paravartya, _ = benchmark(method_paravartya_fast, "Paravartya transpose")


from concurrent.futures import ThreadPoolExecutor

def method_dhwajank_parallel():
    """
    Dhwajank (Flag) method: Parallel computation of independent rows.
    Each thread computes one "flag" section.
    """
    def compute_row_slice(start, end):
        return X[:, start:end] @ W[start:end, :]
    
    n_workers = 4
    row_chunks = np.array_split(range(784), n_workers)
    
    with ThreadPoolExecutor(max_workers=n_workers) as executor:
        futures = [executor.submit(compute_row_slice, c[0], c[-1]+1) 
                   for c in row_chunks if len(c) > 0]
        partials = [f.result() for f in futures]
    
    return sum(partials)

result_dhwajank, _ = benchmark(method_dhwajank_parallel, "Dhwajank parallel (4 threads)")

def method_anurupyena():
    """
    Anurupyena: Proportional scaling to simpler numbers.
    Scale both matrices to have smaller magnitude ranges.
    """
    # Normalize to [-1, 1] range
    X_norm = X / np.abs(X).max()
    W_norm = W / np.abs(W).max()
    
    # Compute in normalized space
    result = X_norm @ W_norm
    
    # Scale back
    scale_factor = np.abs(X).max() * np.abs(W).max()
    return result * scale_factor

result_anurupyena, _ = benchmark(method_anurupyena, "Anurupyena normalized")