import torch
import trme_torch
import time

def test_integration():
    print("🔧 Testing PyTorch Integration...")
    print("-" * 50)

    N, K, M = 128, 128, 128

    # Random Inputs
    input = torch.randn(N, K, requires_grad=True)
    weight = torch.randn(M, K, requires_grad=True) # Linear layer is M x K

    # 1. TRME Forward
    start = time.perf_counter()
    out_trme = trme_torch.matmul(input, weight)
    end = time.perf_counter()
    print(f"TRME Forward: {end - start:.6f} s")
    print(f"Output Mean: {out_trme.mean().item():.4f}")

    # 2. PyTorch Standard
    start = time.perf_counter()
    out_ref = torch.nn.functional.linear(input, weight)
    end = time.perf_counter()
    print(f"Torch Forward: {end - start:.6f} s")

    # Note: Outputs won't match exactly due to aggressive quantization (2-bit vs Float32)
    # We just check shapes and execution.
    assert out_trme.shape == out_ref.shape
    print("✅ Shapes match.")

    # 3. Backward
    loss = out_trme.sum()
    loss.backward()
    print("✅ Backward pass successful.")
    print("-" * 50)

if __name__ == "__main__":
    test_integration()
