import functools, math, pathlib
from tinygrad import Tensor, dtypes
from tinygrad.helpers import getenv
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from tinygrad.renderer import Estimates
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler

BLOCK_ROW = 256

@functools.cache
def _router_mfma_fwd(out:UOp, x:UOp, weight:UOp, bias:UOp, *, dname:str) -> UOp:
  *lead, K = x.shape
  M = math.prod(lead)
  E = weight.shape[0]
  threads = UOp.special(256, "lidx0")
  workgroups = UOp.special((M + 63) // 64, "gidx0")
  sink = UOp.sink(out.base, x.base, weight.base, bias.base, threads, workgroups,
                  arg=KernelInfo(f"moe_router_mfma_{M}_{K}_{E}", estimates=Estimates(ops=2*M*E*K, mem=(M*K+E*K+E)*2+M*E*4)))
  amd = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
  src = (amd/"moe_router_mfma.cpp").read_text()
  lib = HIPCCCompiler("gfx950", [f"-I{(amd/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS",
                                 f"-DROUTER_M={M}", f"-DROUTER_K={K}", f"-DROUTER_E={E}"]).compile_cached(src)
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))

def _router_mfma_bwd(gradient:UOp, kernel:UOp) -> tuple:
  _, x_u, weight_u, bias_u = kernel.src[1:5]
  x, weight, bias = (Tensor(u, device=u.device) for u in (x_u, weight_u, bias_u))
  reference = x.float() @ weight.float().T + bias.float()
  grad_x, grad_weight, grad_bias = reference.gradient(x, weight, bias, gradient=Tensor(gradient, device=x_u.device))
  return None, grad_x.uop, grad_weight.uop, grad_bias.uop

def router_mfma(x:Tensor, weight:Tensor, bias:Tensor) -> Tensor:
  assert x.ndim >= 2 and weight.ndim == 2 and bias.ndim == 1
  K = x.shape[-1]
  E = weight.shape[0]
  assert weight.shape == (E, K) and bias.shape == (E,)
  assert x.dtype == weight.dtype == bias.dtype == dtypes.bfloat16
  assert E == 32 and K % 64 == 0
  if isinstance(x.device, tuple):
    assert x.uop.axis == 0, f"router MFMA requires axis-0 sharding, got axis={x.uop.axis}"
    local_shape = x.uop.shard_shape
    assert local_shape[-1] == K and math.prod(local_shape[:-1]) % 64 == 0, f"unsupported local router shape {local_shape}"
  else:
    assert math.prod(x.shape[:-1]) % 64 == 0
  x, weight, bias = x.contiguous(), weight.contiguous(), bias.contiguous()
  out = _sharded_invalids((*x.shape[:-1], E), dtypes.float32, x.device)
  out, *_ = Tensor.custom_kernel(out, x, weight, bias,
    fxn=functools.partial(_router_mfma_fwd, dname=str(x.device)), grad_fxn=_router_mfma_bwd)
  return out

def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
  if isinstance(device, tuple):
    per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device)
    return Tensor(per.uop.unshard(0), device=device)
  return Tensor.invalids(*shape, dtype=dtype, device=device)

def _atomic_add(device:str) -> str:
  return "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);" if device == "AMD" \
    else "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"

def _blk_for(D:int) -> int:
  blk = 64
  while D % blk: blk //= 2
  return blk

def _kv_ranges(G, N, D, BLK):
  g = UOp.range(G, 0)
  m = UOp.range(N, 1)
  jo = UOp.range(D // BLK, 2)
  ji = UOp.range(BLK, 3, AxisType.LOCAL)
  return g, m, jo * BLK + ji, jo, ji

def _ggather_fwd_kernel(out:UOp, table:UOp, idx:UOp) -> UOp:
  G, M, D = out.shape
  g, m, j, jo, ji = _kv_ranges(G, M, D, _blk_for(D))
  row = idx.index(g, m).cast(dtypes.weakint)
  val = table.index(g, row, j).load()
  return out.index(g, m, j).store(val).end(g, m, jo, ji).sink(
    arg=KernelInfo(name=f"ggather_fwd_{M}_{D}", opts_to_apply=()))

def _ggather_zero_kernel(out:UOp) -> UOp:
  i = UOp.range(out.numel(), 0)
  return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).end(i).sink(arg=KernelInfo(name="ggather_zero"))

def _sharded_zeros(shape:tuple[int, ...], dtype, device) -> Tensor:
  return Tensor.custom_kernel(_sharded_invalids(shape, dtype, device), fxn=_ggather_zero_kernel)[0]

def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
  _, table_u, idx_u = kernel.src[1:4]
  dev = table_u.device
  device = (dev[0] if isinstance(dev, tuple) else dev).split(":")[0]
  G, R, D = table_u.shape
  gt = _sharded_zeros((G, R, D), dtypes.float32, dev)
  go = Tensor(gradient, device=dev)
  atomic_str = _atomic_add(device)
  def _bwd_kernel(gtab:UOp, gout:UOp, idx:UOp) -> UOp:
    Gk, M, Dk = gout.shape
    g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
    row = idx.index(g, m).cast(dtypes.weakint)
    val = gout.index(g, m, j).load().cast(dtypes.float32)
    atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=(atomic_str, dtypes.void))
    return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
  grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
  return (None, grad_table.cast(table_u.dtype).uop, None)

def grouped_gather_rows(table:Tensor, idx:Tensor, n_groups:int) -> Tensor:
  G, R, D = table.shape
  M = idx.shape[1]
  out = _sharded_invalids((G, M, D), table.dtype, table.device)
  return Tensor.custom_kernel(out, table, idx, fxn=_ggather_fwd_kernel, grad_fxn=_ggather_bwd)[0]

def _gscatter_fwd_kernel(out:UOp, src:UOp, idx:UOp) -> UOp:
  G, M, D = out.shape
  k = idx.shape[1] // src.shape[1]
  g, m, j, jo, ji = _kv_ranges(G, idx.shape[1], D, _blk_for(D))
  row = idx.index(g, m).cast(dtypes.weakint)
  val = src.index(g, (m // k).cast(dtypes.weakint), j).load()
  return out.index(g, row, j).store(val).end(g, m, jo, ji).sink(
    arg=KernelInfo(name=f"gscatter_fwd_{idx.shape[1]}_{D}", opts_to_apply=()))

def _gscatter_bwd(gradient:UOp, kernel:UOp) -> tuple:
  _, src_u, idx_u = kernel.src[1:4]
  dev = src_u.device
  G, T_l, D = src_u.shape
  k = idx_u.shape[1] // T_l
  if getenv("GGATHER_SUM_HIP", 0):
    from extra.gptoss_kernels.gather_sum import gather_sum
    assert k == 4
    return None, gather_sum(Tensor(gradient), Tensor(idx_u)).uop, None
  sel = grouped_gather_rows(Tensor(gradient, device=dev), Tensor(idx_u, device=dev), G)
  # FP8 dispatch carries BF16 gradients back to the quantizer, not FP8-rounded gradients.
  grad_dtype = dtypes.bfloat16 if src_u.dtype == dtypes.fp8e4m3 else src_u.dtype
  return (None, sel.reshape(G, T_l, k, D).sum(2).cast(grad_dtype).uop, None)

def grouped_scatter_rows(src:Tensor, idx:Tensor, m_l:int) -> Tensor:
  G, T_l, D = src.shape
  zero = _sharded_zeros((G, m_l, D), src.dtype, src.device)
  return Tensor.custom_kernel(zero, src, idx, fxn=_gscatter_fwd_kernel, grad_fxn=_gscatter_bwd)[0]

def m_max_for(t_local:int, experts_per_tok:int, n_experts:int) -> int:
  return (-(-t_local * experts_per_tok // BLOCK_ROW) + n_experts) * BLOCK_ROW

class Routing:
  def __init__(self, weights:Tensor, dest_row:Tensor, off:Tensor, m_l:int, n_groups:int, t_local:int,
               topi:Tensor|None=None, counts:Tensor|None=None):
    self.weights, self.dest_row = weights, dest_row
    self.off = off
    self.topi = topi
    self.counts = counts
    self.m_l, self.n_groups, self.t_local = m_l, n_groups, t_local

  @property
  def tile_e(self) -> Tensor:
    G, E = self.off.shape[0], self.off.shape[1] - 1
    tr = Tensor.arange(self.m_l // BLOCK_ROW, dtype=dtypes.int32).reshape(1, -1, 1) * BLOCK_ROW
    tr = tr.shard(self.off.device) if isinstance(self.off.device, tuple) else tr.to(self.off.device)
    return ((tr >= self.off[:, :E].reshape(G, 1, E)).sum(-1) - 1).cast(dtypes.int32).reshape(-1)

  @property
  def rows_e(self) -> Tensor:
    return self.tile_e.reshape(-1, 1).expand(-1, BLOCK_ROW).reshape(-1)

def n_groups_of(t:Tensor) -> int:
  return len(t.device) if isinstance(t.device, tuple) else 1

def route(logits:Tensor, experts_per_tok:int, n_experts:int) -> Routing:
  T, E = logits.shape
  k, G = experts_per_tok, n_groups_of(logits)
  assert T % G == 0, f"tokens {T} must split across {G} devices"
  T_l = T // G

  if getenv("FUSED_ROUTER_TOPK", 0):
    from extra.gptoss_kernels.router_topk import fused_router_topk
    weights, topi = fused_router_topk(logits.reshape(G, T_l, E))
  else:
    topv, topi = logits.reshape(G, T_l, E).topk(k)
    weights = topv.softmax(-1)
  return route_topk(weights, topi, n_experts)

def route_topk(weights:Tensor, topi:Tensor, n_experts:int) -> Routing:
  G, T_l, k = weights.shape
  E, m_l = n_experts, m_max_for(T_l, k, n_experts)
  m = topi.reshape(G, T_l * k).cast(dtypes.int32).one_hot(E).cast(dtypes.int32)

  counts = m.sum(1)
  pad = ((counts + (BLOCK_ROW - 1)) // BLOCK_ROW) * BLOCK_ROW
  off = pad.cumsum(1).pad(((0, 0), (1, 0)))
  dest_row = ((m.cumsum(1) + off[:, :E].reshape(G, 1, E)) * m).sum(-1).sub(1).cast(dtypes.int32)
  return Routing(weights, dest_row, off, m_l, G, T_l, topi=topi, counts=counts)

def dispatch(x:Tensor, r:Routing) -> Tensor:
  G, D = r.n_groups, x.shape[-1]
  return grouped_scatter_rows(x.reshape(G, r.t_local, D), r.dest_row, r.m_l).reshape(G * r.m_l, D)

def dispatch_fp8(x:Tensor|tuple[Tensor, Tensor], r:Routing) -> tuple[Tensor, Tensor]:
  from extra.gptoss_kernels.quantize_mxfp8 import quantize_mxfp8_fused_qe8
  q, e8 = x if isinstance(x, tuple) else quantize_mxfp8_fused_qe8(x)
  if getenv("DISPATCH_GATHER", 0):
    from extra.gptoss_kernels.dispatch import inverse_rows, dispatch_gather
    assert r.counts is not None
    src_row = inverse_rows(r.dest_row, r.counts, r.off, r.m_l)
    return tuple(dispatch_gather(t.reshape(r.n_groups, r.t_local, -1), r.dest_row, src_row) for t in (q, e8))
  return dispatch(q, r), dispatch(e8, r)

def combine(y:Tensor, r:Routing, n_tokens:int, experts_per_tok:int) -> Tensor:
  G, D, k = r.n_groups, y.shape[-1], experts_per_tok
  if getenv("FUSED_COMBINE", 0):
    from extra.gptoss_kernels.combine import fused_combine
    assert k == 4
    return fused_combine(y.reshape(G, r.m_l, D), r.dest_row, r.weights.reshape(G, r.t_local, k)).reshape(n_tokens, D)
  sel = grouped_gather_rows(y.reshape(G, r.m_l, D), r.dest_row, G).reshape(G, r.t_local, k, D)
  return (sel * r.weights.reshape(G, r.t_local, k, 1).cast(sel.dtype)).sum(2).reshape(n_tokens, D).cast(y.dtype)
