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()
torch.manual_seed(0)
for mem in [ttnn.L1_MEMORY_CONFIG]:
    x = torch.randn(1, 3, 256, 192)
    xt = ttnn.from_torch(x, dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=mem)
    o = ttnn.to_torch(tk.im2col_tile(xt))
    ref = host_im2col(x)
    print("exact", torch.equal(o, ref), float((o - ref).abs().max()))
    x2 = torch.randn(1, 3, 256, 192)
    h = ttnn.from_torch(x2, dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT)
    ttnn.copy_host_to_device_tensor(h, xt)
    print("exact after copy", torch.equal(ttnn.to_torch(tk.im2col_tile(xt)), host_im2col(x2)))
    print("us", time_op(dev, lambda: tk.im2col_tile(xt)))
    print("tilize us", time_op(dev, lambda: ttnn.tilize_with_zero_padding(ttnn.from_torch(host_im2col(x), dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev), use_multicore=True)) if False else "")
ttnn.close_device(dev)