import unittest
from tinygrad.helpers import DEBUG, Context
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import UPat, rewrite_group, GroupOp, Ops
from tinygrad.uop.upat import _get_code, upat_compile
import dis

@rewrite_group()
def do_compile(up):
  print("\n***** COMPILE", up)
  match_code = _get_code(up, False)
  match = upat_compile(up, lambda **kwargs: None)
  print(match_code[0])
  if DEBUG >= 2: dis.dis(match)
  return match_code[0]

@Context(SPEC=0)
class TestUPatCompile(unittest.TestCase):
  def test_double(self):
    up = UPat.var("x") * UPat.cvar("c0") + UPat.var("x") * UPat.cvar("c1")
    do_compile(up)

  def test_single(self):
    up = UPat.var("x") + UPat.var("y")
    do_compile(up)

  def test_xpx(self):
    up = UPat.var("x") + UPat.var("x")
    do_compile(up)

  def test_xp0(self):
    up = UPat.var("x") + 0
    do_compile(up)

  def test_bool(self):
    up = UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool)
    do_compile(up)

  def test_single_c(self):
    up = (UPat.var("x") + UPat.var("y")) * UPat.var("c")
    do_compile(up)

  def test_const_folding(self):
    up = UPat(GroupOp.ALU-{Ops.THREEFRY}, name="a", src=UPat((Ops.CONST, Ops.STACK)))
    do_compile(up)

  @unittest.skip("fix this")
  def test_range_named(self):
    # this should be one src, but this should also still work
    up = UPat(Ops.CAST, dtypes.float, UPat.var("x", dtypes.bfloat16))
    do_compile(up)

if __name__ == "__main__":
  unittest.main()