"""Two-hop activation broadcast (kernels/mcbench/bcast2.cpp), round 7.
usage: bcast_test.py mode=rc|cr [cols=12] [rows=10] [sb=49152] [noc=0|1] [verify=1] [time=1] [nodata=0]
  rc: shard r starts at (0, r); hop A = row multicast, hop B = chain of column multicasts (12 columns in parallel)
  cr: shard c starts at (c, 0) (c < NSH); hop A = column multicast, hop B = chain of row multicasts (rows in parallel)"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *

K = str(Path(__file__).resolve().parents[2] / "hamer/kernels/mcbench/bcast2.cpp")
args = dict(a.split("=") for a in sys.argv[1:])
mode = args.get("mode", "rc")
C, R = int(args.get("cols", 12)), int(args.get("rows", 10))
SB = int(args.get("sb", 49152))
noc = int(args.get("noc", 0))
verify = args.get("verify", "1") == "1"
NSH = R if mode == "rc" else int(args.get("nsh", 10))
assert NSH <= 14

dev = open_dev()
try:
    phys = {(x, y): dev.worker_core_from_logical_core(ttnn.CoreCoord(x, y)) for x in range(C) for y in range(R)}
    cores = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(C - 1, R - 1))])

    def rect(x0, y0, x1, y1):
        p0, p1 = phys[(x0, y0)], phys[(x1, y1)]
        return [min(p0.x, p1.x), min(p0.y, p1.y), max(p0.x, p1.x), max(p0.y, p1.y)]

    out = ttnn.from_torch(torch.zeros(C * R, 8, dtype=torch.int32), dtype=ttnn.uint32, layout=ttnn.ROW_MAJOR_LAYOUT,
                          device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)

    def build(empty=False):
        rt = ttnn.RuntimeArgs()
        for x in range(C):
            for y in range(R):
                a = [0xff, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0, 0, 0, 0, 0, 0, 0, 0, 0, SB, out.buffer_address(),
                     y * C + x]
                if not empty:
                    if mode == "rc":
                        # hop A: row y, head (0, y) holds shard y; hop B: column x, chain y = 0..R-1
                        p = phys[(0, y)]
                        a[0] = y
                        a[1] = 1 if x == 0 else 0
                        a[2:6] = rect(0, y, C - 1, y); a[6] = C
                        a[7:11] = rect(x, 0, x, R - 1); a[11] = R
                        a[12] = y; a[13] = y - 1 if y else 0
                        a[14] = 0 if x == 0 else 1; a[15:17] = [p.x, p.y]
                        q = phys[(x, 0)]
                        a[17] = 0 if y == 0 else 1; a[18:20] = [q.x, q.y]; a[20] = 1 if y == 0 else 0
                    else:
                        # hop A: column x (x < NSH), head (x, 0) holds shard x; hop B: row y, chain x = 0..NSH-1
                        if x < NSH:
                            p = phys[(x, 0)]
                            a[0] = x
                            a[1] = 1 if y == 0 else 0
                            a[2:6] = rect(x, 0, x, R - 1); a[6] = R
                            a[12] = x; a[13] = x - 1 if x else 0
                            a[14] = 0 if y == 0 else 1; a[15:17] = [p.x, p.y]
                        a[7:11] = rect(0, y, C - 1, y); a[11] = C
                        q = phys[(0, y)]
                        a[17] = 0 if x == 0 else 1; a[18:20] = [q.x, q.y]; a[20] = 1 if x == 0 else 0
                    a[21] = NSH
                rt[x][y] = a
        cfg = ttnn.ReaderConfigDescriptor() if noc == 0 else ttnn.WriterConfigDescriptor()
        defs = ([("BC_VERIFY", "1")] if verify else []) + ([("BC_NO_DATA", "1")] if args.get("nodata") == "1" else []) \
            + [(d, "1") for d in filter(None, args.get("defs", "").split(","))]
        kd = ttnn.KernelDescriptor(kernel_source=K, source_type=ttnn.KernelDescriptor.SourceType.FILE_PATH, core_ranges=cores,
                                   compile_time_args=list(ttnn.TensorAccessorArgs(out).get_compile_time_args()),
                                   defines=defs, runtime_args=rt, config=cfg)
        cbs = [ttnn.CBDescriptor(total_size=NSH * SB if not empty else 2048, core_ranges=cores, format_descriptors=[
                   ttnn.CBFormatDescriptor(buffer_index=0, data_format=ttnn.bfloat16, page_size=2048)]),
               ttnn.CBDescriptor(total_size=256, core_ranges=cores, format_descriptors=[
                   ttnn.CBFormatDescriptor(buffer_index=1, data_format=ttnn.uint32, page_size=256)])]
        sems = [ttnn.SemaphoreDescriptor(id=i, core_ranges=cores, initial_value=0) for i in range(2 + NSH)]
        return ttnn.ProgramDescriptor(kernels=[kd], semaphores=sems, cbs=cbs)

    prog = build()
    ttnn.generic_op([out, out], prog)
    ttnn.synchronize_device(dev)
    if verify:
        r = ttnn.to_torch(out).reshape(C * R, 8)
        bad = int(r[:, 0].sum()); tags = int(((r[:, 1].to(torch.int64) & 0xFFFF0000) == 0xC0DE0000).sum())
        print(f"[bc] {mode} verify: cores reporting {tags}/{C*R}, mismatched words {bad}", flush=True)
    if args.get("prof") == "1":
        pe = build(empty=True)
        prof_run(dev, "empty", lambda: (ttnn.generic_op([out, out], pe), None)[1])
        prof_run(dev, "bcast", lambda: (ttnn.generic_op([out, out], prog), None)[1])
    if args.get("time", "1") == "1":
        reps = int(args.get("reps", 20))
        pe = build(empty=True)
        te = time_op(dev, lambda: (ttnn.generic_op([out, out], pe), None)[1], reps=reps)
        tb = time_op(dev, lambda: (ttnn.generic_op([out, out], prog), None)[1], reps=reps)
        tot = NSH * SB
        print(f"[bc] {mode} noc{noc} {C}x{R} NSH={NSH} SB={SB}: empty {te:.2f} us, bcast {tb:.2f} us -> {tb - te:.2f} us"
              f" ({tot / max(tb - te, 1e-3) / 1e3:.1f} GB/s per receiver)", flush=True)
finally:
    from open_device_12x10 import close_device
    close_device(dev)
if args.get("prof") == "1":
    prof_report()