import itertools, functools
from collections import defaultdict
from dataclasses import replace
from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType
from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, graph_rewrite
from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate, sym
from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape
from tinygrad.renderer import Renderer

# ***** image load valid simplification *****

@functools.cache
def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
  # can drop valid if idx is out of bound when valid is False
  drop_stmt = []
  for i,stmt in enumerate(valid.split_uop(Ops.AND)):
    if (res:=parse_valid(stmt)) is None: continue
    X, is_upper_bound, c = res

    # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
    if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in X.split_uop(Ops.ADD)):
      testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx)
      if testidx.index(0).vmax < 0 or testidx.index(1).vmax < 0:
        drop_stmt.append(stmt)
        continue

    # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1]
    lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1)
    if lo <= hi:
      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype, param=True)
      subs = [{X: fake}]
      # idx may not have X itself, so also substitute a term of X: v -> fake - (X - v)
      terms = list(X.split_uop(Ops.ADD))
      v = next((u for u in terms if u.op in GroupOp.Irreducible and u.op is not Ops.CONST), None)
      if v is not None and (rest:=[u for u in terms if u is not v]): subs.append({v: fake - UOp.usum(*rest)})
      if any((testidx:=graph_rewrite(coord.substitute(sub), sym)).vmin >= b or testidx.vmax < 0
             for sub in subs for coord,b in zip(idx.src, (width, height))):
        drop_stmt.append(stmt)
  return drop_stmt

def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
  idx = uop_given_valid(valid, start_idx)
  return None if idx is start_idx or idx is start_idx.simplify() else buf.index(idx.valid(valid))

def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None:
  if not is_image_shape(buf._shape): return None
  if idx_x.dtype != idx_y.dtype: idx_x, idx_y = idx_x.cast(dtypes.int), idx_y.cast(dtypes.int)
  start_idx = idx_x.stack(idx_y)
  idx = uop_given_valid(valid, start_idx)
  drop_stmt = _drop_valid_stmts(valid, idx, buf._shape[0], buf._shape[1])

  if not drop_stmt and idx is start_idx: return None
  new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
  idx_y, idx_x = idx.index(1), idx.index(0)
  if new_valid is not None: return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid))
  return buf.index(idx_y, idx_x)

indexing_simplify = PatternMatcher([
  # image load valid idx simplification
  (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
  (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)),
                                         UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load),
])

# get list of (height, width) that do not require pitch padding
def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]:
  if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return []
  MAXW, pxls = 16384, size // 4
  if base not in (dtypes.half, dtypes.float) or size > 4*MAXW*MAXW: return []
  # height=1 images just need to abide by alignment requirements in bytes, not pixels!
  if size % (ALIGN * 4) != 0: return [] if (base.itemsize * size) % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)]
  return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]

def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
  shapes, ren = ctx
  if not IMAGE or ren.target.device not in {"QCOM", "CL", "PYTHON", "NULL"}: return None
  valid, x = x.get_valid(), x.get_idx()
  # search for dims that drop the most valid statements
  best_drop, cands = -1, []
  for ch, cw in [shapes[buf.arg.slot]] if buf.arg.slot in shapes else image_valid_dims(buf.dtype, buf.max_numel(), ren.target.arch):
    cidx = uop_given_valid(valid, ((x//4)%cw).stack(x//(4*cw)))
    dropped = len(_drop_valid_stmts(valid, cidx, ch, cw))
    if dropped > best_drop: best_drop, cands = dropped, [(ch, cw, cidx)]
    elif dropped == best_drop: cands.append((ch, cw, cidx))
  # if no candidates, we don't rewrite
  if len(cands) == 0: return None
  # and tiebreak with indexing complexity (ie. number of nodes)
  h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice))
  # the image dims are stored in the param's arg, the size stays the flat buffer len
  buf = buf.replace(arg=replace(buf.arg, image=(h, w)))
  shapes[buf.arg.slot] = (h, w)
  if valid.op is not Ops.CONST or valid.val is not True:
    return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid))
  else:
    return buf.index(cidx.src[1], cidx.src[0])

pm_simplify_add_image = PatternMatcher([
  (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image),
  # image load/store is always float
  (UPat(Ops.INDEX, dtype=dtypes.float, name="x").store(UPat(name="d", dtype=dtypes.half)), lambda x,d: x.store(d.cast(dtypes.float))),
  (UPat.var("x", dtype=dtypes.float).cast(dtypes.half).cast(dtypes.float), lambda x: x),
])

def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
  if getenv("DMC"): return sink

  # collect
  memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp, object], dict[int, list[UOp]]] = defaultdict(dict)
  for u in sink.toposort():
    # TODO: this should handle images too, it's just memory coalescing
    if u.op in {Ops.LOAD, Ops.STORE}:
      assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores"
      assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}"
      buf, idx_u = u.src[0].src
      if buf.addrspace == AddrSpace.REG: continue
      idx, valid = idx_u.get_idx(), idx_u.get_valid()
      root_src: UOp|str
      if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].val
      elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].val
      elif idx.op is Ops.CONST and idx.val is Invalid: root_src, arg = "INVALID", 0
      elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.val
      else: root_src, arg = idx, 0
      # loads/stores only coalesce with others carrying the same arg (e.g. the nontemporal flag)
      memory[(u.op, buf, root_src, valid, u.arg)].setdefault(arg, []).append(u)

  # build replacements
  replacements = {}
  for (op,buf,base,valid,ld_arg),offsets in memory.items():
    # allowed lengths (copied in)
    lengths = []
    must_divide = True
    if ctx is not None and ctx.target.device == "DSP":
      lengths = [128,64,32,16,8,4]
      must_divide = False
    elif buf.dtype not in (dtypes.float, dtypes.half, dtypes.int, dtypes.uint, *dtypes.fp8s) and not is_image_shape(buf._shape):
      pass
    elif buf.addrspace == AddrSpace.REG:
      pass
    elif is_image_shape(buf._shape):
      lengths = [4]
    elif ctx is not None and ctx.supports_float4:
      # TODO: a better way to get this than ctx
      lengths = [8,4,2] if buf.dtype == dtypes.half and getenv("ALLOW_HALF8") else [4,2]
    lengths.append(1)  # worst case, it's not folded
    # do the grouping
    grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
    for full_grp in grouped_offsets:
      while len(full_grp):
        offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(full_grp[0])
        length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.divides(l) is not None)][0]
        grp = full_grp[:length]
        # NOTE: we apply the valid again after we determine the length
        offset = offset.valid(valid) if valid is not None else offset
        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset)
        if op == Ops.STORE:
          datas = []
          for i,g in enumerate(grp):
            assert len(offsets[g]) == 1, f"attempting multiple stores: {len(offsets[g])}"
            datas.append(offsets[g][0].src[1])
          store = idx.store(UOp.stack(*datas) if len(datas) > 1 else datas[0])
          for i,g in enumerate(grp): replacements[offsets[g][0]] = store
        else:
          ld = idx.load(arg=ld_arg)
          for i,g in enumerate(grp):
            for oo in offsets[g]:
              replacements[oo] = ld.index(i) if len(grp) > 1 else ld
        full_grp = full_grp[length:]

  # apply
  return sink.substitute(replacements, name="memory coalescing")
