# ruff: noqa: E501 E712 F401
from dataclasses import replace
from tinygrad import dtypes, Device
from tinygrad.uop.ops import UOp, AxisType, Ops, KernelInfo
from tinygrad.codegen.opt import Opt, OptOps # pylint: disable=unused-import
from tinygrad.engine.realize import get_runtime
from tinygrad.codegen import to_program
from tinygrad.helpers import dedup, getenv
from tinygrad.device import Buffer
from tinygrad.dtype import Invalid

# PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx

def vision_conv_143():
  c0 = UOp.param(0, dtypes.half, shape=(16, 1024, 4))
  c2 = UOp.range(32, 3)
  c5 = UOp.range(128, 4)
  c8 = UOp.range(16, 2)
  c16 = UOp.range(7, 0, AxisType.REDUCE)
  c17 = c8*2+c16
  c24 = ((c17<3)!=True)&(c17<35)
  c26 = UOp.range(7, 1, AxisType.REDUCE)
  c27 = c2*2+c26
  c32 = ((c27<3)!=True)&(c27<67)
  c34 = UOp.param(1, dtypes.half, shape=(32, 1024, 4))
  c38 = c5//2
  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(Invalid, dtypes.weakint))
  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
  c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
  c63 = UOp.param(3, dtypes.float, 128)
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*128+c5+c8*4096)).store(c65).end(c8, c2, c5)

  opts = None
  # JITBEAM=2
  # (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.SWAP, axis=1, arg=2))
  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))

def vision_conv_153():
  c0 = UOp.param(0, dtypes.half, shape=(8, 1024, 4))
  c2 = UOp.range(16, 3)
  c5 = UOp.range(256, 4)
  c8 = UOp.range(8, 2)
  c16 = UOp.range(7, 0, AxisType.REDUCE)
  c17 = c8*2+c16
  c24 = ((c17<3)!=True)&(c17<19)
  c26 = UOp.range(7, 1, AxisType.REDUCE)
  c27 = c2*2+c26
  c32 = ((c27<3)!=True)&(c27<35)
  c34 = UOp.param(1, dtypes.half, shape=(16, 1024, 4))
  c38 = c5//2
  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(Invalid, dtypes.weakint))
  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
  c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
  c63 = UOp.param(3, dtypes.float, 256)
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*256+c5+c8*4096)).store(c65).end(c8, c2, c5)

  opts = None
  # JITBEAM=2
  # (Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.NOLOCALS, axis=None, arg=None), Opt(op=OptOps.UPCAST, axis=2, arg=2), Opt(op=OptOps.SWAP, axis=1, arg=2))
  return c67.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))

def dm_conv_172():
  c0 = UOp.param(0, dtypes.half, shape=(1, 240, 4))
  c2 = UOp.range(960, 4)
  c5 = UOp.param(1, dtypes.half, shape=(8, 384, 4))
  c7 = UOp.range(32, 0, AxisType.REDUCE)
  c10 = UOp.range(4, 1, AxisType.REDUCE)
  c13 = UOp.range(12, 3, AxisType.REDUCE)
  c18 = UOp.range(8, 2, AxisType.REDUCE)
  c23 = UOp.param(2, dtypes.half, shape=(240, 128, 4))
  c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
  c37 = UOp.param(3, dtypes.float, 960)
  c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
  c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
  c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
  c55 = c0.index(c2).store(c53).end(c2)

  opts = None
  # JITBEAM=2
  # (Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.GROUPTOP, axis=1, arg=32), Opt(op=OptOps.UNROLL, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.GROUP, axis=1, arg=0))
  return c55.sink(arg=KernelInfo(name="conv", opts_to_apply=opts))

ast = {143: vision_conv_143, 153: vision_conv_153, 172: dm_conv_172}[getenv("NUM", 143)]()

renderer = Device.default.renderer
allocator = Device.default.allocator

ps = to_program(ast, renderer)
rt = get_runtime(Device.DEFAULT, ps)

gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda u: u.arg)
# print(len(gs))
# print([g.dtype for g in gs])
bufs = [Buffer(ps.arg.device, g.max_numel(), g.dtype).ensure_allocated() for g in gs]

gsize, lsize = ps.arg.launch_dims({})
t = rt(*[b._buf for b in bufs], global_size=gsize, local_size=lsize, vals=ps.arg.vals({}), wait=True)
print(f"{t*1e6:.2f} us")
