from typing import Callable
import functools
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher
from tinygrad.renderer import Renderer

# *** integer division ***

@functools.lru_cache(None)
def magicgu(vmax:int, d:int) -> tuple[int,int]:
  # calculate m,s such that x//d == (x*m) >> s for all 0 <= x <= vmax, d>0; adapted from Hacker's Delight, Chapter 10
  nc = (vmax+1)//(d) * d - 1
  nbits = vmax.bit_length()
  for s in range(0, 2*nbits + 1):
    if 2**s > nc*(d - 1 - (2**s - 1) % d):
      m = (2**s + d - 1 - (2**s - 1) % d)//d
      return m, s
  assert False

def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None:
  if d <= 0 or x.vmin < 0: return None
  if (vmax:=min(x.vmax, x.dtype.max)) < d: return x.const_like(0)
  m,s = magicgu(vmax, d)
  if m*vmax <= x.dtype.max: return (x*m) >> s
  # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
  if (k := (d & -d).bit_length()-1) > 0:
    if (ret:=fast_idiv(ren, x >> k, d >> k, dont_cast=True)) is not None: return ret
  if dont_cast: return None
  # the next integer width that holds x*m
  widen = {dtypes.int8:dtypes.int16, dtypes.int16:dtypes.int32, dtypes.int32:dtypes.int64, dtypes.int64:dtypes.uint64,
           dtypes.uint8:dtypes.uint16, dtypes.uint16:dtypes.uint32, dtypes.uint32:dtypes.uint64}
  if (next_dtype := widen.get(x.dtype)) is not None and next_dtype in ren.supported_dtypes():
    if m*vmax <= next_dtype.max: return ((x.cast(next_dtype)*m) >> s).cast(x.dtype)
  return None

# ***** threefry *****

def threefry2x32(x: UOp, key: UOp):
  # split x and key from uint64 to two uint32
  x0, x1 = x.cast(dtypes.uint32), (x >> 32).cast(dtypes.uint32)
  key0, key1 = key.cast(dtypes.uint32), (key >> 32).cast(dtypes.uint32)

  rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
  ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0]
  xr:list[UOp] = [x0 + ks[-1], x1 + ks[0]]
  for i in range(5):
    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] << r) + (xr[1] >> (32 - r)))
    xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)]

  return (xr[1].cast(dtypes.uint64) << 32) | xr[0].cast(dtypes.uint64)

# ***** decomposition patterns *****

def floordiv_to_idiv(a:UOp, b:UOp) -> UOp:
  if (a.vmin >= 0 and b.vmin >= 0) or (a.vmax <= 0 and b.vmax <= 0): return a.alu(Ops.CDIV, b)
  return a.alu(Ops.CDIV, b) - (a.alu(Ops.CMOD, b).ne(0) & (a<0).ne(b<0))

def floormod_to_mod(a:UOp, b:UOp) -> UOp:
  if (a.vmin >= 0 and b.vmin >= 0) or (a.vmax <= 0 and b.vmax <= 0): return a.alu(Ops.CMOD, b)
  r = a.alu(Ops.CMOD, b)
  # use where instead of mul to avoid being fused into MULACC (which int64 long-decomp doesn't handle)
  return r + (r.ne(0) & (a<0).ne(b<0)).where(b, b.const_like(0))

powers_of_two: dict[int, int] = {2**i:i for i in range(64)}
@functools.cache
def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher:
  # these are rewrites that make things simpler
  pat: list[tuple[UPat, Callable]] = []
  # FLOORDIV by 2**y -> x >> y (an arithmetic shift is exactly floor division for any sign); fires before floordiv_to_idiv
  if Ops.SHR in ops: pat.append((UPat.var("x", dtypes.ints)//UPat.cvar("c"),
    lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None))
  pat.append((UPat.var("a")//UPat.var("b"), floordiv_to_idiv))
  # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
  if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None))
  pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))
  # no real hardware supports THREEFRY, but NullRenderer does
  if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, src=(UPat.var("x"), UPat.var("key"))), threefry2x32))
  return PatternMatcher(pat)

@functools.cache
def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> PatternMatcher:
  pat: list[tuple[UPat, Callable]] = []
  # MAX can be rewritten as CMPLT + WHERE (max function is annoying on many cstyle backends). late: WHERE loses the bound MAX carries
  if Ops.MAX not in ops and Ops.CMPLT in ops: pat += [(UPat(Ops.MAX, name="m"), lambda m: (m.src[0] < m.src[1]).where(m.src[1], m.src[0]))]
  if Ops.OR in ops: pat += [(UPat.var("x", dtypes.bool).logical_not()&UPat.var("y", dtypes.bool).logical_not(),
    lambda x,y: (x | y).logical_not())]
  # rewrite MUL/CDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
  if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.val, 0)) else None)]
  if Ops.SHR in ops:
    # uint CDIV by 2**v -> x >> v (FLOORDIV is lowered to CDIV by the rule above before reaching here)
    pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.uints), UPat.cvar("c"))),
      lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None)]
    # signed CDIV (trunc) by 2**v -> (x + (x<0 ? c-1 : 0)) >> v
    pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("c"))),
      lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v
        if (v:=powers_of_two.get(c.val, 0)) else None)]
    if not disable_fast_idiv:
      # fast_idiv handles non-pow2 divisors on non-negative inputs
      pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))), lambda ctx, x, d: fast_idiv(ctx, x, d.val))]
      # rewrite raw CMOD -> x - d*fast_idiv(x,d), only when fast_idiv can actually divide;
      # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail)
      pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
        lambda ctx, x, d: x - d * q if (q:=fast_idiv(ctx, x, d.val)) is not None else None)]
  if Ops.NEG in ops:
    pat += [(UPat.var('x')*-1, lambda ctx,x: x.alu(Ops.NEG))]
    if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda ctx,x,y: x.alu(Ops.SUB, y))]
  if Ops.CMPLT in ops:
    # These are late rewrites because simplex expects equalities to be a certain format
    pat += [
      ((UPat.var("x", dtypes.sints) < UPat.cvar("c")).logical_not(), lambda x,c: c-1<x),
      ((UPat.cvar("c") < UPat.var("x", dtypes.sints)).logical_not(), lambda x,c: x<c+1),
      (UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*UPat.cvar("c"), lambda x,y,c: y*(-c)<x),
      (UPat.var("x", dtypes.sints)*-1 < UPat.cvar("c"), lambda x,c:-c<x),
      ((UPat.cvar("c1")<UPat.var("x", dtypes.sints)) & (UPat.var("x", dtypes.sints)<UPat.cvar("c2")),
        lambda x,c1,c2: x.eq(c1+1) if c1.val+1==c2.val-1 else None),  # (c-1)<x & x<(c+1) -> x==c
    ]
  if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))]
  if Ops.MULACC in ops:
    pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
    # also fuse (x << n) + c → MULACC(x, 2^n, c) since MUL→SHL may run first
    if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1<<n.val), c))]
  # some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
  if Ops.FDIV in ops:
    pat += [(UPat.var("x").reciprocal(), lambda x: UOp.const(1.0).alu(Ops.FDIV, x))]
    pat += [(UPat.var("a") * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
  return PatternMatcher(pat)
