import unittest
from tinygrad import Tensor, Context, Device
from tinygrad.engine.realize import get_program
from tinygrad.renderer import Opt, OptOps
from tinygrad.uop.ops import KernelInfo

class TestLinearizerRewrite(unittest.TestCase):
  def test_reduction(self):
    t = Tensor.ones((64,64), device="NULL").contiguous().realize()
    out = (t*2).sum(axis=1)
    with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0):
      si = out.schedule()[-1]
      opts_to_apply = []
      opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
      opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4))
      ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
      prg = get_program(ast, Device["CPU"].renderer)
      print(prg.src)

  def test_arange(self):
    out = Tensor.arange(32, device="NULL")
    with Context(SPLIT_REDUCEOP=0, DEVECTORIZE=0):
      si = out.schedule()[-1]
      opts_to_apply = []
      opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
      opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4))
      ast = si.ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
      prg = get_program(ast, Device["CPU"].renderer)
      print(prg.src)

  def test_kernel_info(self):
    out = Tensor.arange(4, device="NULL")
    si = out.schedule()[-1]

    ast = si.ast.replace(arg=KernelInfo(opts_to_apply=()))
    prg = get_program(ast, Device["CPU"].renderer)
    assert prg.applied_opts == (), f"expected no opts, got {prg}"

    prg = get_program(ast.replace(arg=None), Device["CPU"].renderer)
    assert prg.applied_opts != (), f"expected opts to apply, got {prg.applied_opts}"

    prg = get_program(ast.replace(arg=KernelInfo(name="custom")), Device["CPU"].renderer)
    self.assertEqual(prg.name, "custom")

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