import numpy as np
import torch
import torch.nn.functional as F
from scipy.signal import correlate2d
import time

# Reproducible setup
torch.manual_seed(42)
np.random.seed(42)

# Conv2D: Input (batch, channels, height, width)
# Conv2D: Weights (out_channels, in_channels, kH, kW)
B, C, H, W = 4, 64, 112, 112    # Batch, Channels, Height, Width
K, kH, kW = 128, 3, 3           # Out channels, Kernel size

x = torch.randn(B, C, H, W, dtype=torch.float32)
w = torch.randn(K, C, kH, kW, dtype=torch.float32)

def benchmark(func, name="Method", runs=100):
    # Warmup
    for _ in range(10):
        _ = func()
    
    start = time.perf_counter()
    for _ in range(runs):
        result = func()
    elapsed = (time.perf_counter() - start) / runs * 1000
    print(f"{name:45} {elapsed:8.2f} ms")
    return result

# Baseline
def baseline():
    return F.conv2d(x, w, padding=1)

_ = benchmark(baseline, "PyTorch Conv2d (baseline)", runs=100)


def method_vedic_im2col():
    """
    Urdhva-Tiryagbyham (Crosswise): Convert convolution to matrix multiplication.
    Inspired by ancient grid-based multiplication.
    """
    # im2col: Unfold image patches into columns
    # Each column = flattened kernel-sized patch
    col_x = torch.nn.functional.unfold(x, kernel_size=(kH, kW), padding=1)
    # Shape: (B, kH*kW*C, out_H*out_W)
    
    # Reshape weight for matrix multiply
    w_flat = w.view(K, -1)  # (K, kH*kW*C)
    
    # Matrix multiply: (B, K, out_H*out_W)
    out = col_x.transpose(1, 2) @ w_flat.t()
    
    # Reshape to (B, K, out_H, out_W)
    out_H = (H + 2 - kH) 
    out_W = (W + 2 - kW)
    return out.transpose(1, 2).reshape(B, K, out_H, out_W)

_ = benchmark(method_vedic_im2col, "Vedic im2col (manual unfold)", runs=100)