"""im2col_tile on a bf16 raw-pixel crop (0..255 integers, exact in bf16) vs host_im2col; fp32 path unchanged."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels as tk
from hamer.fused import host_im2col
dev = open_dev()
try:
    torch.manual_seed(0)
    L1 = ttnn.L1_MEMORY_CONFIG
    u = torch.randint(0, 256, (1, 3, 256, 192)).float()
    xt = ttnn.from_torch(u.to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=L1)
    o = ttnn.to_torch(tk.im2col_tile(xt)).float()
    print("[i2c] bf16 exact", torch.equal(o, host_im2col(u)), float((o - host_im2col(u)).abs().max()), flush=True)
    x = torch.randn(1, 3, 256, 192)
    xf = ttnn.from_torch(x, dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=L1)
    print("[i2c] fp32 exact", torch.equal(ttnn.to_torch(tk.im2col_tile(xf)), host_im2col(x)), flush=True)
    if "time" in sys.argv:
        print(f"[i2c] bf16 {time_op(dev, lambda: tk.im2col_tile(xt)):.2f} us  fp32 {time_op(dev, lambda: tk.im2col_tile(xf)):.2f} us", flush=True)
finally:
    ttnn.close_device(dev)