import ast
import textwrap
import inspect
from typing import Optional

DTYPE_MAP = {
  "float16": "half",
  "float32": "float",
  "float64": "double",
  "int8": "int8_t",
  "int16": "int16_t",
  "int32": "int32_t",
  "int64": "int64_t",
  "uint8": "uint8_t",
  "uint16": "uint16_t",
  "uint32": "uint32_t",
  "uint64": "uint64_t",
  "bool": "bool",
}

OP_MAP = {
  ast.Add: "+",
  ast.Sub: "-",
  ast.Mult: "*",
  ast.Div: "/",
  ast.FloorDiv: "/",
  ast.Mod: "%",
  ast.BitAnd: "&",
  ast.BitOr: "|",
  ast.BitXor: "^",
  ast.LShift: "<<",
  ast.RShift: ">>",
  ast.Eq: "==",
  ast.NotEq: "!=",
  ast.Lt: "<",
  ast.LtE: "<=",
  ast.Gt: ">",
  ast.GtE: ">=",
  ast.And: "&&",
  ast.Or: "||",
  ast.USub: "-",
  ast.UAdd: "+",
  ast.Not: "!",
  ast.Invert: "~",
}

TRITON_MATH = {
  "exp": "expf",
  "exp2": "exp2f",
  "log": "logf",
  "log2": "log2f",
  "sqrt": "sqrtf",
  "abs": "fabsf",
  "cos": "cosf",
  "sin": "sinf",
  "sigmoid": "__sigmoid",
  "tanh": "tanhf",
  "ceil": "ceilf",
  "floor": "floorf",
  "rint": "rintf",
  "rsqrt": "rsqrtf",
  "erf": "erff",
}


class TranslationError(Exception):
  pass


class TileVar:
  def __init__(self, name: str, shape: tuple, dtype: str = "float"):
    self.name = name
    self.shape = shape
    self.dtype = dtype

  @property
  def numel(self):
    n = 1
    for s in self.shape:
      n *= s
    return n


class CudaEmitter(ast.NodeVisitor):
  def __init__(self, func, constexprs: dict, meta: dict):
    self.func = func
    self.constexprs = constexprs
    self.meta = meta
    self.lines = []
    self.indent = 2
    self.tile_vars: dict[str, TileVar] = {}
    self.scalar_vars: set = set()
    self.temp_counter = 0
    self.range_vars: dict[str, tuple] = {}
    self.pid_dims_used: set = set()
    self._block_size: Optional[tuple] = None
    self._infer_block_size()

  def _infer_block_size(self):
    bs = []
    for k in ["BLOCK_SIZE", "BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]:
      if k in self.meta:
        bs.append(self.meta[k])
    if bs:
      self._block_size = tuple(bs)

  def emit(self, line: str):
    self.lines.append(" " * self.indent + line)

  def fresh_temp(self, prefix="__tmp"):
    self.temp_counter += 1
    return f"{prefix}_{self.temp_counter}"

  def translate(self) -> str:
    src = inspect.getsource(self.func)
    src = textwrap.dedent(src)
    tree = ast.parse(src)
    func_def = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef))
    self._translate_function(func_def)
    return "\n".join(self.lines)

  def _translate_function(self, node: ast.FunctionDef):
    params = self._build_params(node.args)
    self.lines.append(f"__global__ void {node.name}({params}) {{")
    self._emit_pid_setup()
    for stmt in node.body:
      if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant):
        continue
      self._translate_stmt(stmt)
    self.lines.append("}")

  def _build_params(self, args: ast.arguments) -> str:
    parts = []
    for a in args.args:
      if a.annotation:
        ann_str = ast.unparse(a.annotation)
        if "constexpr" in ann_str:
          self.constexprs[a.arg] = self.meta.get(a.arg)
          parts.append(f"const int {a.arg}")
        elif ann_str.startswith("tl.") or "." in ann_str:
          c_type = self._triton_type_to_ptr(ann_str)
          parts.append(f"{c_type} {a.arg}")
        else:
          self.scalar_vars.add(a.arg)
          parts.append(f"int64_t {a.arg}")
      else:
        self.scalar_vars.add(a.arg)
        parts.append(f"int64_t {a.arg}")
    return ", ".join(parts)

  def _triton_type_to_ptr(self, ann: str) -> str:
    mapping = {
      "tl.float16": "half*",
      "tl.float32": "float*",
      "tl.float64": "double*",
      "tl.int8": "int8_t*",
      "tl.int16": "int16_t*",
      "tl.int32": "int32_t*",
      "tl.int64": "int64_t*",
      "tl.uint8": "uint8_t*",
      "tl.uint16": "uint16_t*",
      "tl.uint32": "uint32_t*",
      "tl.uint64": "uint64_t*",
    }
    for k, v in mapping.items():
      if k in ann:
        return v
    return "float*"

  def _emit_pid_setup(self):
    self.emit("int __pid0 = blockIdx.x;")
    self.emit("int __pid1 = blockIdx.y;")
    self.emit("int __pid2 = blockIdx.z;")

  def _translate_stmt(self, node):
    if isinstance(node, ast.Assign):
      self._translate_assign(node)
    elif isinstance(node, ast.AugAssign):
      self._translate_augassign(node)
    elif isinstance(node, ast.For):
      self._translate_for(node)
    elif isinstance(node, ast.If):
      self._translate_if(node)
    elif isinstance(node, ast.Return):
      pass
    elif isinstance(node, ast.Expr):
      val = self._translate_expr(node.value)
      if val:
        self.emit(f"{val};")
    else:
      self.emit(f"/* unhandled: {ast.dump(node)} */")

  def _translate_assign(self, node: ast.Assign):
    rhs = self._translate_expr(node.value)
    for target in node.targets:
      self._get_assign_target(target, rhs)

  def _get_assign_target(self, target, rhs_expr: str) -> Optional[str]:
    if isinstance(target, ast.Name):
      name = target.id
      if name in self.tile_vars:
        tv = self.tile_vars[name]
        self.emit(f"for (int __i = 0; __i < {tv.numel}; ++__i) {{ {name}[__i] = {rhs_expr}[__i]; }}")
        return None
      if name in self.scalar_vars:
        self.emit(f"{name} = {rhs_expr};")
        return None
      if rhs_expr in self.tile_vars:
        tv = self.tile_vars[rhs_expr]
        new_tv = TileVar(name, tv.shape, tv.dtype)
        self.tile_vars[name] = new_tv
        shape_str = " * ".join(str(s) for s in tv.shape)
        self.emit(f"{tv.dtype} {name}[{shape_str}];")
        self.emit(f"for (int __i = 0; __i < {tv.numel}; ++__i) {{ {name}[__i] = {rhs_expr}[__i]; }}")
        return None
      cdtype = self._infer_scalar_decl_type(rhs_expr)
      self.scalar_vars.add(name)
      self.emit(f"{cdtype} {name} = {rhs_expr};")
      return None
    if isinstance(target, ast.Subscript):
      lhs = self._translate_subscript_store(target, rhs_expr)
      if lhs:
        self.emit(f"{lhs} = {rhs_expr};")
      return None
    return None

  def _infer_scalar_decl_type(self, rhs: str) -> str:
    if rhs in self.tile_vars:
      return self.tile_vars[rhs].dtype
    if rhs.startswith("blockIdx"):
      return "int64_t"
    if rhs in self.scalar_vars or rhs in self.constexprs:
      return "int64_t"
    try:
      v = float(rhs)
      if "." in rhs or "f" in rhs.lower() or "e" in rhs.lower():
        return "float"
      return "int64_t"
    except ValueError:
      pass
    return "int64_t"

  def _translate_augassign(self, node: ast.AugAssign):
    op = OP_MAP.get(type(node.op), "?")
    lhs = self._translate_expr(node.target)
    rhs = self._translate_expr(node.value)
    lhs_tv = self.tile_vars.get(lhs)
    rhs_tv = self.tile_vars.get(rhs)
    if lhs_tv:
      numel = lhs_tv.numel
      if rhs_tv:
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {lhs}[__i] {op}= {rhs}[__i]; }}")
      else:
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {lhs}[__i] {op}= {rhs}; }}")
    else:
      self.emit(f"{lhs} {op}= {rhs};")

  def _translate_for(self, node: ast.For):
    if isinstance(node.target, ast.Name) and isinstance(node.iter, ast.Call):
      var = node.target.id
      call = node.iter
      fname = self._get_call_name(call)
      if fname in ("range", "tl.range"):
        args = [self._translate_expr(a) for a in call.args]
        start = args[0] if len(args) > 1 else "0"
        stop = args[1] if len(args) > 1 else args[0]
        step = args[2] if len(args) > 2 else "1"
        self.range_vars[var] = (start, stop, step)
        self.emit(f"for (int {var} = {start}; {var} < {stop}; {var} += {step}) {{")
        self.indent += 2
        for stmt in node.body:
          self._translate_stmt(stmt)
        self.indent -= 2
        self.emit("}")
        return
    self.emit(f"/* unhandled for */")

  def _translate_if(self, node: ast.If):
    cond = self._translate_expr(node.test)
    self.emit(f"if ({cond}) {{")
    self.indent += 2
    for stmt in node.body:
      self._translate_stmt(stmt)
    self.indent -= 2
    if node.orelse:
      self.emit("} else {")
      self.indent += 2
      for stmt in node.orelse:
        self._translate_stmt(stmt)
      self.indent -= 2
    self.emit("}")

  def _translate_expr(self, node) -> str:
    if isinstance(node, ast.Constant):
      return str(node.value)
    if isinstance(node, ast.Name):
      return self._map_name(node.id)
    if isinstance(node, ast.BinOp):
      return self._translate_binop(node)
    if isinstance(node, ast.UnaryOp):
      op = OP_MAP.get(type(node.op), "?")
      return f"({op}{self._translate_expr(node.operand)})"
    if isinstance(node, ast.BoolOp):
      return self._translate_boolop(node)
    if isinstance(node, ast.Compare):
      return self._translate_compare(node)
    if isinstance(node, ast.Call):
      return self._translate_call(node)
    if isinstance(node, ast.Attribute):
      return self._translate_attr(node)
    if isinstance(node, ast.Subscript):
      return self._translate_subscript_load(node)
    if isinstance(node, ast.IfExp):
      t = self._translate_expr(node.test)
      b = self._translate_expr(node.body)
      o = self._translate_expr(node.orelse)
      return f"({t} ? {b} : {o})"
    if isinstance(node, ast.Tuple):
      elts = [self._translate_expr(e) for e in node.elts]
      return "make_tuple(" + ", ".join(elts) + ")"
    return f"/* expr:{ast.dump(node)} */"

  def _map_name(self, name: str) -> str:
    pid_map = {
      "pid": "__pid0",
      "pid_m": "__pid0",
      "pid_n": "__pid1",
      "pid_k": "__pid2",
    }
    if name in pid_map:
      return pid_map[name]
    if name in self.constexprs:
      return name
    return name

  def _translate_binop(self, node: ast.BinOp) -> str:
    op = OP_MAP.get(type(node.op), "?")
    lhs = self._translate_expr(node.left)
    rhs = self._translate_expr(node.right)
    lhs_tv = self.tile_vars.get(lhs)
    rhs_tv = self.tile_vars.get(rhs)
    if lhs_tv or rhs_tv:
      tv = lhs_tv or rhs_tv
      numel = tv.numel
      shape_str = " * ".join(str(s) for s in tv.shape)
      out_dtype = self._binop_dtype(lhs, rhs, lhs_tv, rhs_tv, op)
      tmp = self.fresh_temp("__bop")
      self.tile_vars[tmp] = TileVar(tmp, tv.shape, out_dtype)
      self.emit(f"{out_dtype} {tmp}[{shape_str}];")
      if lhs_tv and rhs_tv:
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {tmp}[__i] = ({out_dtype})({lhs}[__i]) {op} ({out_dtype})({rhs}[__i]); }}")
      elif lhs_tv:
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {tmp}[__i] = ({out_dtype})({lhs}[__i]) {op} ({out_dtype})({rhs}); }}")
      else:
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {tmp}[__i] = ({out_dtype})({lhs}) {op} ({out_dtype})({rhs}[__i]); }}")
      return tmp
    return f"({lhs} {op} {rhs})"

  def _binop_dtype(self, lhs: str, rhs: str, lhs_tv, rhs_tv, op: str) -> str:
    lhs_dtype = lhs_tv.dtype if lhs_tv else self._scalar_dtype(lhs)
    rhs_dtype = rhs_tv.dtype if rhs_tv else self._scalar_dtype(rhs)
    if "float" in lhs_dtype or "half" in lhs_dtype or "double" in lhs_dtype:
      return lhs_dtype
    if "float" in rhs_dtype or "half" in rhs_dtype or "double" in rhs_dtype:
      return rhs_dtype
    return "int64_t"

  def _scalar_dtype(self, name: str) -> str:
    if name in self.scalar_vars or name in self.constexprs:
      return "int64_t"
    if name.startswith("blockIdx"):
      return "int64_t"
    try:
      float(name)
      return "float" if "." in name or "f" in name.lower() else "int64_t"
    except ValueError:
      pass
    return "int64_t"

  def _translate_boolop(self, node: ast.BoolOp) -> str:
    op = OP_MAP.get(type(node.op), "?")
    vals = [self._translate_expr(v) for v in node.values]
    return f"({f' {op} '.join(vals)})"

  def _translate_compare(self, node: ast.Compare) -> str:
    lhs = self._translate_expr(node.left)
    lhs_tv = self.tile_vars.get(lhs)
    results = []
    cur = lhs
    cur_tv = lhs_tv
    for op, comp_node in zip(node.ops, node.comparators):
      op_str = OP_MAP.get(type(op), "?")
      rhs = self._translate_expr(comp_node)
      rhs_tv = self.tile_vars.get(rhs)
      if cur_tv or rhs_tv:
        tv = cur_tv or rhs_tv
        numel = tv.numel
        shape_str = " * ".join(str(s) for s in tv.shape)
        tmp = self.fresh_temp("__cmp")
        self.tile_vars[tmp] = TileVar(tmp, tv.shape, "bool")
        self.emit(f"bool {tmp}[{shape_str}];")
        if cur_tv and rhs_tv:
          self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {tmp}[__i] = {cur}[__i] {op_str} {rhs}[__i]; }}")
        elif cur_tv:
          self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {tmp}[__i] = {cur}[__i] {op_str} {rhs}; }}")
        else:
          self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {tmp}[__i] = {cur} {op_str} {rhs}[__i]; }}")
        results.append(tmp)
        cur = tmp
        cur_tv = self.tile_vars[tmp]
      else:
        results.append(f"({cur} {op_str} {rhs})")
        cur = rhs
        cur_tv = rhs_tv
    return results[-1] if results else f"/* cmp */"

  def _translate_call(self, node: ast.Call) -> str:
    fname = self._get_call_name(node)
    args = [self._translate_expr(a) for a in node.args]
    kwargs = {kw.arg: self._translate_expr(kw.value) for kw in node.keywords}

    if fname in ("tl.zeros", "tl.full"):
      shape_node = node.args[0] if node.args else None
      if shape_node:
        shape = self._eval_shape_from_node(shape_node)
        fill = "0.0f"
        if fname == "tl.full" and len(node.args) > 1:
          fill = self._translate_expr(node.args[1])
        dtype_node = None
        if fname == "tl.zeros" and len(node.args) > 1:
          dtype_node = ast.unparse(node.args[1])
        elif fname == "tl.full" and len(node.args) > 2:
          dtype_node = ast.unparse(node.args[2])
        for kw in node.keywords:
          if kw.arg == "dtype":
            dtype_node = ast.unparse(kw.value)
        cdtype = self._triton_dtype_to_c(dtype_node or "float32")
        numel = 1
        for s in shape:
          numel *= s
        tmp = self.fresh_temp("__tile")
        self.tile_vars[tmp] = TileVar(tmp, shape, cdtype)
        self.emit(f"{cdtype} {tmp}[{numel}];")
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {tmp}[__i] = {fill}; }}")
        return tmp

    return self._dispatch_triton_call(fname, args, kwargs, node)

  def _eval_shape_from_node(self, node) -> tuple:
    if isinstance(node, ast.Tuple):
      dims = []
      for e in node.elts:
        dims.append(self._eval_const_node(e))
      return tuple(dims)
    if isinstance(node, ast.Constant):
      return (int(node.value),)
    if isinstance(node, ast.Name):
      v = self.meta.get(node.id)
      if v is not None:
        return (int(v),)
    return (1,)

  def _dispatch_triton_call(self, fname: str, args: list, kwargs: dict, node) -> str:
    if fname == "tl.program_id":
      axis = int(args[0]) if args else int(kwargs.get("axis", "0"))
      return f"blockIdx.{'xyz'[axis]}"

    if fname == "tl.arange":
      start, end = args[0], args[1]
      size = self._eval_const(end) - self._eval_const(start)
      tmp = self.fresh_temp("__arange")
      self.tile_vars[tmp] = TileVar(tmp, (size,), "int32_t")
      self.emit(f"int32_t {tmp}[{size}];")
      self.emit(f"for (int __i = 0; __i < {size}; ++__i) {{ {tmp}[__i] = {start} + __i; }}")
      return tmp

    if fname == "tl.load":
      return self._emit_load(args, kwargs)

    if fname == "tl.store":
      self._emit_store(args, kwargs)
      return ""

    if fname == "tl.dot":
      return self._emit_dot(args, kwargs)

    if fname in ("tl.zeros", "tl.full"):
      return self._emit_zeros_or_full(fname, args, kwargs)

    if fname == "tl.where":
      return self._emit_where(args, kwargs)

    if fname in ("tl.sum", "tl.max", "tl.min"):
      return self._emit_reduce(fname, args, kwargs)

    if fname == "tl.reshape":
      return self._emit_reshape(args, kwargs)

    if fname == "tl.trans":
      return self._emit_trans(args)

    if fname == "tl.expand_dims":
      return args[0]

    if fname == "tl.broadcast_to":
      return self._emit_broadcast(args, kwargs)

    if fname == "tl.sqrt":
      return self._emit_elementwise("sqrtf", args)

    if fname == "tl.exp":
      return self._emit_elementwise("expf", args)

    if fname == "tl.log":
      return self._emit_elementwise("logf", args)

    if fname in ("tl.cast", "tl.view"):
      dtype = args[1] if len(args) > 1 else kwargs.get("dtype", "float")
      cdtype = self._triton_dtype_to_c(dtype)
      tmp = self.fresh_temp("__cast")
      src_tv = self.tile_vars.get(args[0])
      if src_tv:
        shape_str = " * ".join(str(s) for s in src_tv.shape)
        numel = src_tv.numel
        self.tile_vars[tmp] = TileVar(tmp, src_tv.shape, cdtype)
        self.emit(f"{cdtype} {tmp}[{shape_str}];")
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {tmp}[__i] = ({cdtype}){args[0]}[__i]; }}")
        return tmp
      return f"(({cdtype}){args[0]})"

    if fname == "tl.constexpr":
      return args[0] if args else "0"

    if fname == "tl.multiple_of":
      return args[0] if args else "0"

    if fname == "tl.max_contiguous":
      return args[0] if args else "0"

    if fname in ("tl.atomic_add", "tl.atomic_max", "tl.atomic_min",
                 "tl.atomic_and", "tl.atomic_or", "tl.atomic_xor"):
      return self._emit_atomic(fname, args, kwargs)

    if fname == "tl.debug_barrier":
      self.emit("__syncthreads();")
      return ""

    if fname in TRITON_MATH:
      return self._emit_elementwise(TRITON_MATH[fname], args)

    base = fname.split(".")[-1]
    if base in TRITON_MATH:
      return self._emit_elementwise(TRITON_MATH[base], args)

    return f"{fname}({', '.join(args)})"

  def _emit_load(self, args: list, kwargs: dict) -> str:
    ptr = args[0]
    mask = args[1] if len(args) > 1 else kwargs.get("mask")
    other = args[2] if len(args) > 2 else kwargs.get("other", "0.0f")
    tmp = self.fresh_temp("__load")
    ptr_tv = self.tile_vars.get(ptr)
    mask_tv = self.tile_vars.get(mask) if mask else None

    if ptr_tv:
      numel = ptr_tv.numel
      shape_str = " * ".join(str(s) for s in ptr_tv.shape)
      self.tile_vars[tmp] = TileVar(tmp, ptr_tv.shape, "float")
      self.emit(f"float {tmp}[{shape_str}];")
      if mask_tv:
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{")
        self.indent += 2
        self.emit(f"{tmp}[__i] = {mask}[__i] ? *(float*)({ptr}[__i]) : (float)({other});")
        self.indent -= 2
        self.emit("}")
      elif mask:
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{")
        self.indent += 2
        self.emit(f"{tmp}[__i] = {mask} ? *(float*)({ptr}[__i]) : (float)({other});")
        self.indent -= 2
        self.emit("}")
      else:
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {tmp}[__i] = *(float*)({ptr}[__i]); }}")
      return tmp

    tmp2 = self.fresh_temp("__scalar_load")
    self.scalar_vars.add(tmp2)
    if mask:
      self.emit(f"float {tmp2} = {mask} ? *(float*)({ptr}) : (float)({other});")
    else:
      self.emit(f"float {tmp2} = *(float*)({ptr});")
    return tmp2

  def _emit_store(self, args: list, kwargs: dict):
    ptr = args[0]
    val = args[1]
    mask = args[2] if len(args) > 2 else kwargs.get("mask")
    ptr_tv = self.tile_vars.get(ptr)
    val_tv = self.tile_vars.get(val)
    mask_tv = self.tile_vars.get(mask) if mask else None
    if ptr_tv and val_tv:
      numel = ptr_tv.numel
      if mask_tv:
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{")
        self.indent += 2
        self.emit(f"if ({mask}[__i]) *(float*)({ptr}[__i]) = {val}[__i];")
        self.indent -= 2
        self.emit("}")
      elif mask:
        self.emit(f"if ({mask}) {{")
        self.indent += 2
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ *(float*)({ptr}[__i]) = {val}[__i]; }}")
        self.indent -= 2
        self.emit("}")
      else:
        self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ *(float*)({ptr}[__i]) = {val}[__i]; }}")
    elif val_tv:
      self.emit(f"*(float*)({ptr}) = {val}[0];")
    else:
      self.emit(f"*(float*)({ptr}) = {val};")

  def _emit_dot(self, args: list, kwargs: dict) -> str:
    a, b = args[0], args[1]
    acc = args[2] if len(args) > 2 else kwargs.get("acc")
    a_tv = self.tile_vars.get(a)
    b_tv = self.tile_vars.get(b)
    tmp = self.fresh_temp("__dot")

    M = self.meta.get("BLOCK_SIZE_M", self.meta.get("BLOCK_SIZE", 1))
    N = self.meta.get("BLOCK_SIZE_N", self.meta.get("BLOCK_SIZE", 1))
    K = self.meta.get("BLOCK_SIZE_K", self.meta.get("BLOCK_SIZE", 1))

    if a_tv:
      if len(a_tv.shape) == 2:
        M, K = a_tv.shape
      else:
        K = a_tv.numel // M if M else a_tv.numel
    if b_tv:
      if len(b_tv.shape) == 2:
        K, N = b_tv.shape
      else:
        N = b_tv.numel // K if K else b_tv.numel

    M, N, K = int(M), int(N), int(K)

    self.tile_vars[tmp] = TileVar(tmp, (M, N), "float")
    self.emit(f"float {tmp}[{M * N}];")
    if acc and acc in self.tile_vars:
      self.emit(f"for (int __i = 0; __i < {M * N}; ++__i) {{ {tmp}[__i] = {acc}[__i]; }}")
    else:
      self.emit(f"for (int __i = 0; __i < {M * N}; ++__i) {{ {tmp}[__i] = 0.0f; }}")
    self.emit(f"for (int __m = 0; __m < {M}; ++__m) {{")
    self.indent += 2
    self.emit(f"for (int __n = 0; __n < {N}; ++__n) {{")
    self.indent += 2
    self.emit(f"for (int __k = 0; __k < {K}; ++__k) {{")
    self.indent += 2
    self.emit(f"{tmp}[__m * {N} + __n] += {a}[__m * {K} + __k] * {b}[__k * {N} + __n];")
    self.indent -= 2
    self.emit("}")
    self.indent -= 2
    self.emit("}")
    self.indent -= 2
    self.emit("}")
    return tmp

  def _emit_zeros_or_full(self, fname: str, args: list, kwargs: dict) -> str:
    shape_node = None
    fill = "0.0f"
    dtype_str = "float32"

    orig_args = getattr(self, "_current_call_args", [])
    if fname == "tl.zeros":
      shape_node_raw = orig_args[0] if orig_args else None
      dtype_str = orig_args[1] if len(orig_args) > 1 else "float32"
    else:
      shape_node_raw = orig_args[0] if orig_args else None
      fill_raw = orig_args[1] if len(orig_args) > 1 else "0"
      dtype_str = orig_args[2] if len(orig_args) > 2 else "float32"
      fill = str(fill_raw)

    shape = self._parse_shape_from_str(args[0] if args else "(1,)")
    cdtype = self._triton_dtype_to_c(str(dtype_str))
    numel = 1
    for s in shape:
      numel *= s
    shape_str = str(numel)
    tmp = self.fresh_temp("__tile")
    self.tile_vars[tmp] = TileVar(tmp, shape, cdtype)
    self.emit(f"{cdtype} {tmp}[{shape_str}];")
    self.emit(f"for (int __i = 0; __i < {shape_str}; ++__i) {{ {tmp}[__i] = {fill}; }}")
    return tmp

  def _parse_shape_from_str(self, shape_arg: str) -> tuple:
    shape_arg = shape_arg.strip()
    if shape_arg.startswith("(") and shape_arg.endswith(")"):
      inner = shape_arg[1:-1]
      parts = [p.strip().rstrip(",") for p in inner.split(",") if p.strip()]
      dims = []
      for p in parts:
        dims.append(self._eval_const(p))
      return tuple(dims) if dims else (1,)
    try:
      return (int(shape_arg),)
    except Exception:
      return (self._eval_const(shape_arg),)

  def _emit_where(self, args: list, kwargs: dict) -> str:
    cond, x, y = args[0], args[1], args[2]
    tmp = self.fresh_temp("__where")
    cond_tv = self.tile_vars.get(cond)
    x_tv = self.tile_vars.get(x)
    if cond_tv and x_tv:
      numel = x_tv.numel
      shape_str = " * ".join(str(s) for s in x_tv.shape)
      self.tile_vars[tmp] = TileVar(tmp, x_tv.shape, x_tv.dtype)
      self.emit(f"{x_tv.dtype} {tmp}[{shape_str}];")
      self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{")
      self.indent += 2
      self.emit(f"  {tmp}[__i] = {cond}[__i] ? {x}[__i] : {y};")
      self.indent -= 2
      self.emit("}")
      return tmp
    return f"({cond} ? {x} : {y})"

  def _emit_reduce(self, fname: str, args: list, kwargs: dict) -> str:
    inp = args[0]
    axis = int(args[1]) if len(args) > 1 else int(kwargs.get("axis", "0"))
    tmp = self.fresh_temp("__reduce")
    src_tv = self.tile_vars.get(inp)
    if src_tv:
      numel = src_tv.numel
      if fname == "tl.sum":
        init = "0.0f"
        op = "+="
      elif fname == "tl.max":
        init = "-__FLT_MAX__"
        op = None
      else:
        init = "__FLT_MAX__"
        op = None
      self.scalar_vars.add(tmp)
      self.emit(f"float {tmp} = {init};")
      self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{")
      self.indent += 2
      if op:
        self.emit(f"  {tmp} {op} {inp}[__i];")
      elif fname == "tl.max":
        self.emit(f"  {tmp} = ({inp}[__i] > {tmp}) ? {inp}[__i] : {tmp};")
      else:
        self.emit(f"  {tmp} = ({inp}[__i] < {tmp}) ? {inp}[__i] : {tmp};")
      self.indent -= 2
      self.emit("}")
      return tmp
    return f"/* reduce({inp}) */"

  def _emit_reshape(self, args: list, kwargs: dict) -> str:
    inp = args[0]
    return inp

  def _emit_trans(self, args: list) -> str:
    inp = args[0]
    src_tv = self.tile_vars.get(inp)
    if src_tv and len(src_tv.shape) == 2:
      M, N = src_tv.shape
      tmp = self.fresh_temp("__trans")
      self.tile_vars[tmp] = TileVar(tmp, (N, M), src_tv.dtype)
      self.emit(f"{src_tv.dtype} {tmp}[{N}][{M}];")
      self.emit(f"for (int __m = 0; __m < {M}; ++__m) {{")
      self.indent += 2
      self.emit(f"for (int __n = 0; __n < {N}; ++__n) {{")
      self.indent += 2
      self.emit(f"{tmp}[__n][__m] = {inp}[__m * {N} + __n];")
      self.indent -= 2
      self.emit("}")
      self.indent -= 2
      self.emit("}")
      return f"reinterpret_cast<{src_tv.dtype}*>({tmp})"
    return inp

  def _emit_broadcast(self, args: list, kwargs: dict) -> str:
    return args[0]

  def _emit_atomic(self, fname: str, args: list, kwargs: dict) -> str:
    ptr = args[0]
    val = args[1] if len(args) > 1 else kwargs.get("val", "0")
    op_map = {
      "tl.atomic_add": "atomicAdd",
      "tl.atomic_max": "atomicMax",
      "tl.atomic_min": "atomicMin",
      "tl.atomic_and": "atomicAnd",
      "tl.atomic_or": "atomicOr",
      "tl.atomic_xor": "atomicXor",
    }
    cuda_op = op_map.get(fname, "atomicAdd")
    ptr_tv = self.tile_vars.get(ptr)
    if ptr_tv:
      numel = ptr_tv.numel
      self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{")
      self.indent += 2
      self.emit(f"  {cuda_op}(&{ptr}[__i], {val});")
      self.indent -= 2
      self.emit("}")
      return ""
    return f"{cuda_op}({ptr}, {val})"

  def _emit_elementwise(self, cfunc: str, args: list) -> str:
    inp = args[0]
    src_tv = self.tile_vars.get(inp)
    if src_tv:
      numel = src_tv.numel
      shape_str = " * ".join(str(s) for s in src_tv.shape)
      tmp = self.fresh_temp("__ew")
      self.tile_vars[tmp] = TileVar(tmp, src_tv.shape, src_tv.dtype)
      self.emit(f"{src_tv.dtype} {tmp}[{shape_str}];")
      self.emit(f"for (int __i = 0; __i < {numel}; ++__i) {{ {tmp}[__i] = {cfunc}({inp}[__i]); }}")
      return tmp
    return f"{cfunc}({inp})"

  def _translate_attr(self, node: ast.Attribute) -> str:
    obj = self._translate_expr(node.value)
    attr = node.attr
    if obj == "tl":
      return f"tl.{attr}"
    return f"{obj}.{attr}"

  def _translate_subscript_load(self, node: ast.Subscript) -> str:
    val = self._translate_expr(node.value)
    if isinstance(node.slice, ast.Tuple):
      idxs = []
      for e in node.slice.elts:
        if isinstance(e, ast.Constant) and e.value is None:
          idxs.append(None)
        elif isinstance(e, ast.Slice):
          idxs.append(":")
        else:
          idxs.append(self._translate_expr(e))
      valid = [i for i in idxs if i not in (None, ":")]
      if valid:
        return f"{val}[{']['.join(valid)}]"
      return val
    if isinstance(node.slice, ast.Slice):
      return val
    if isinstance(node.slice, ast.Constant) and node.slice.value is None:
      return val
    idx = self._translate_expr(node.slice)
    return f"{val}[{idx}]"

  def _translate_subscript_store(self, node: ast.Subscript, rhs: str) -> Optional[str]:
    val = self._translate_expr(node.value)
    if isinstance(node.slice, ast.Tuple):
      idxs = [self._translate_expr(e) for e in node.slice.elts]
      return f"{val}[{']['.join(idxs)}]"
    idx = self._translate_expr(node.slice)
    return f"{val}[{idx}]"

  def _get_call_name(self, node: ast.Call) -> str:
    if isinstance(node.func, ast.Name):
      return node.func.id
    if isinstance(node.func, ast.Attribute):
      parts = []
      cur = node.func
      while isinstance(cur, ast.Attribute):
        parts.append(cur.attr)
        cur = cur.value
      if isinstance(cur, ast.Name):
        parts.append(cur.id)
      return ".".join(reversed(parts))
    return "unknown"

  def _is_tile_expr(self, expr: str) -> bool:
    return expr.startswith("__") or expr in self.tile_vars

  def _infer_tile_shape(self, name: str) -> TileVar:
    if name in self.tile_vars:
      return self.tile_vars[name]
    return TileVar(name, (1,), "float")

  def _eval_shape_arg(self, shape_arg) -> tuple:
    if isinstance(shape_arg, ast.Tuple):
      dims = []
      for e in shape_arg.elts:
        dims.append(self._eval_const_node(e))
      return tuple(dims)
    if isinstance(shape_arg, str):
      try:
        return (int(shape_arg),)
      except Exception:
        return (1,)
    return (1,)

  def _eval_const_node(self, node) -> int:
    if isinstance(node, ast.Constant):
      return int(node.value)
    if isinstance(node, ast.Name) and node.id in self.meta:
      return int(self.meta[node.id])
    return 1

  def _eval_const(self, expr: str) -> int:
    try:
      return int(expr)
    except Exception:
      return self.meta.get(expr, 1)

  def _triton_dtype_to_c(self, dtype: str) -> str:
    clean = str(dtype).replace("tl.", "").replace("triton.", "")
    if "float32" in clean or clean == "float":
      return "float"
    if "float16" in clean:
      return "half"
    if "float64" in clean:
      return "double"
    if "int64" in clean:
      return "int64_t"
    if "int32" in clean:
      return "int32_t"
    if "int16" in clean:
      return "int16_t"
    if "int8" in clean:
      return "int8_t"
    if "bool" in clean:
      return "bool"
    return "float"
