import unittest
import numpy as np
from tinygrad import Tensor, GlobalCounters, Context, Device
from tinygrad.dtype import DTypeLike, dtypes
from tinygrad.engine.realize import run_linear
from tinygrad.helpers import DEBUG, get_single_element
from test.helpers import check_schedule

def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
  # only support axis =-1
  x = x_in.reshape(-1, x_in.shape[-1])
  nr_dim, r_dim = x.shape

  inp = x.reshape(nr_dim, 1, 1, r_dim).expand(nr_dim, r_dim, 1, r_dim)
  imx = x.reshape(nr_dim, 1, r_dim, 1).expand(nr_dim, r_dim, r_dim, r_dim).max(axis=-2, keepdim=True)
  m = inp - imx.detach()
  if dtype is not None: m = m.cast(dtype)
  e = m.exp()
  ss = e.sum(axis=-1, keepdim=True)

  inp = x.reshape(nr_dim, r_dim, 1, 1)
  imx = x.reshape(nr_dim, 1, r_dim, 1).expand(nr_dim, r_dim, r_dim, 1).max(axis=-2, keepdim=True)
  m = inp - imx.detach()
  if dtype is not None: m = m.cast(dtype)
  e = m.exp()

  out = e.div(ss).reshape(x_in.shape)
  return out

def run_one_schedule_item(out):
  linear = out.schedule_linear()
  get_single_element(linear.src)
  run_linear(linear)

class TestFuse(unittest.TestCase):
  def _test_fuse(self, fxn, *args, atol=1e-6, allow_multiple=False, **kwargs):
    GlobalCounters.reset()
    out_single = fxn(*args, **kwargs)
    if not allow_multiple: run_one_schedule_item(out_single)
    np_single = out_single.numpy()
    GlobalCounters.reset()
    np_multi = fxn(*args, **kwargs).numpy()
    np.testing.assert_allclose(np_single, np_multi, atol=atol)

  @unittest.skip("needs RANGEIFY>1")
  def test_fuse_norm(self):
    a = Tensor.rand(50,50).realize()
    self._test_fuse(lambda a: a / a.mean(axis=1), a)

  @unittest.skip("needs RANGEIFY>1")
  def test_fuse_argmax(self):
    a = Tensor.rand(50,50).realize()
    self._test_fuse(lambda a: a.argmax(axis=-1), a)

  @unittest.skip("needs RANGEIFY>1")
  def test_fuse_softmax(self):
    a = Tensor.rand(50,50).realize()
    self._test_fuse(lambda a: a.softmax(axis=-1), a)

  def test_fuse_gemm_softmax(self):
    a = Tensor.rand(50,50).realize()
    b = Tensor.rand(50,50).realize()
    self._test_fuse(lambda a,b: ((a@b).relu()+a).contiguous().softmax(axis=-1), a,b, allow_multiple=True)

  @unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), f"no float16 on {Device.DEFAULT}")
  @unittest.skip("needs RANGEIFY>1")
  def test_fuse_softmax_dtype(self):
    a = Tensor.rand(50,50).realize()
    self._test_fuse(lambda a: a.softmax(axis=-1, dtype='half'), a, atol=3e-4)

  def test_fuse_arange_eye(self):
    self._test_fuse(lambda: (Tensor.arange(10).reshape(10,1).expand(10,10) == Tensor.arange(10).reshape(1,10).expand(10,10)).clone())

  @unittest.skip("needs RANGEIFY>1")
  def test_double_gemm(self):
    N = 32
    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
      a = (Tensor.rand(N,N)-0.5).realize()
      b = (Tensor.rand(N,N)-0.5).realize()
      c = (Tensor.rand(N,N)-0.5).realize()
    self._test_fuse(lambda a,b,c: a@b@c, a, b, c, atol=1e-5)

  def test_embedding(self):
    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
      vocab_sz = 123
      embed_sz = 16
      weight = (Tensor.rand(vocab_sz, embed_sz)-0.5).realize()
      a = Tensor([1, 1, 2, 3]).realize()
    def embedding(idx:Tensor):
      arange = Tensor.arange(vocab_sz).unsqueeze(-1)
      big_shp = idx.shape + (vocab_sz, embed_sz)
      arange, vals = arange.expand(big_shp), weight.expand(big_shp)
      idx = idx.reshape(idx.shape+(1, 1)).expand(big_shp)
      return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
    self._test_fuse(embedding, a, atol=1e-5)

  @unittest.skip("needs RANGEIFY>1")
  def test_attention_kernel_count(self):
    wq = Tensor.empty(32, 32)
    wk = Tensor.empty(32, 32)
    wv = Tensor.empty(32, 32)
    x = Tensor.empty(2, 100, 32)
    q = (x @ wq).contiguous()
    k = (x @ wk).contiguous()
    v = (x @ wv).contiguous()
    attn = q.scaled_dot_product_attention(k, v)
    check_schedule(attn, 4) # 3 matmul and 1 attention

  @unittest.skip("needs RANGEIFY>1")
  def test_flash_attention(self):
    BS = 4
    HEADS = 2
    MATDIM = 16
    EMB = 8
    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
      q = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
      k = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
      v = Tensor.randn(BS, HEADS, MATDIM, EMB).realize()
    # TODO: OPT is breaking things. NOOPT isn't linearizing
    with Context(NOOPT=1):
      self._test_fuse(Tensor.scaled_dot_product_attention, q, k, v, atol=1e-5)

  def test_mismatch_reduce(self):
    a = Tensor.ones(16, 10).contiguous().realize()
    b = Tensor.ones(16, 20).contiguous().realize()
    c = (a.sum(axis=1) + b.sum(axis=1))
    self.assertListEqual(c.tolist(), [30]*16)

  @unittest.skipUnless(Device.DEFAULT == "METAL", "METAL TC")
  def test_fuse_and_tc_opt(self):
    A = Tensor.randn(8, 8).realize()
    B = Tensor.randn(8, 8).realize()
    C = Tensor.ones(1, 8, 8).pad(((1,1), None, None),).sum(0)
    out = (C + (A @ B))
    out.realize()

class TestSoftmaxFusion(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    with Context(TRACK_MATCH_STATS=0): cls.test = Tensor.rand(32, 10).contiguous().realize()

  def setUp(self):
    GlobalCounters.reset()

  def test_norm(self):
    print("*** norm ***")
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      # NOTE: there's an implied expand on the mean here
      sout = self.test / self.test.mean(-1, keepdim=True)
      sout.realize()

    print("*** single kernel norm ***")
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      inp = self.test.reshape(32, 10, 1)
      div = self.test.reshape(32, 1, 10).expand(32, 10, 10).mean(axis=-1, keepdim=True)
      out = (inp / div).reshape(32, 10)
      out.realize()

    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)

  def test_softmax(self):
    # this is the softmax from scaled_dot_product_attention
    # it becomes 3 kernels
    print("*** softmax ***")
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      sout = self.test.softmax(-1)
      sout.realize()

    print("*** single kernel softmax ***")
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      out = single_kernel_softmax(self.test)
      out.realize()

    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)

  @unittest.skip("needs RANGEIFY>1")
  def test_auto_softmax(self):
    print("*** softmax ***")
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      sout = self.test.softmax(-1)
      sout.realize()

    print("*** auto single kernel softmax ***")
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      out = self.test.contiguous().softmax(-1)
      run_one_schedule_item(out)

    np.testing.assert_allclose(sout.numpy(), out.numpy(), atol=3e-7)

  def test_softmax_bw(self):
    print("*** softmax bw ***")
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      self.test.softmax(-1).sum().backward()
      sg = self.test.grad.realize()

    self.test.grad = None

    print("*** single kernel softmax bw ***")
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      single_kernel_softmax(self.test).sum().backward()
      g = self.test.grad.realize()

    np.testing.assert_allclose(sg.numpy(), g.numpy(), atol=1e-7)

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