import unittest, pickle, types, tracemalloc
import numpy as np
from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
from tinygrad.uop.ops import PatternMatcher, UPat, UOp, deconstruct_function
from test.helpers import KernelCountException

class TestPickle(unittest.TestCase):
  def test_pickle_code_object(self):
    y = lambda x: x*2  # noqa: E731
    code_str = pickle.dumps(y.__code__)
    fxn = types.FunctionType(pickle.loads(code_str), globals())
    self.assertEqual(fxn(2), 4)

  def test_deconstruct_function_nested_comprehension(self):
    # pre PEP 709, each comprehension is its own code object, so dtypes here is referenced two code objects deep
    def fxn(): return [[dtypes.int for _ in range(2)] for _ in range(2)]
    self.assertEqual(types.FunctionType(*deconstruct_function(fxn))(), fxn())

  def test_pickle_pattern_matcher(self):
    pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
    sink = UOp.const(2)
    tt = pm.rewrite(sink)
    pm_str = pickle.dumps(pm)
    pm2 = pickle.loads(pm_str)
    self.assertEqual(pm2.rewrite(sink).key, tt.key)

  def test_pickle_main_pattern_matcher(self):
    from tinygrad.uop.symbolic import sym
    ssym = pickle.dumps(sym)
    dsym = pickle.loads(ssym)
    self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location)

  def test_pickle_realized_tensor(self):
    print("** init")
    t = Tensor.rand(10, 10).realize()
    st = pickle.dumps(t)
    t_values = t.numpy()
    del t # free buffers
    print("** post pickle")
    GlobalCounters.reset()
    t2:Tensor = pickle.loads(st)
    np.testing.assert_equal(t_values, t2.numpy())
    # expect at most one COPY kernel
    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)

  def test_pickle_realized_tensor_alt(self):
    print("** init")
    t = Tensor.rand(10, 10).to("CPU").realize()
    st = pickle.dumps(t)
    t_values = t.numpy()
    del t # free buffers
    print("** post pickle")
    t2:Tensor = pickle.loads(st)
    assert t2.uop.is_realized
    np.testing.assert_equal(t_values, t2.numpy())

  def test_pickle_realized_tensor_alt2(self):
    print("** init")
    t = Tensor.rand(10, 10).to("CPU").realize()
    tensor_uop = t.uop
    assert tensor_uop.is_realized, f"expected {tensor_uop} to be realized"
    t_values = t.numpy()
    # pickle
    st = pickle.dumps(t)
    # free buffers
    del t
    del tensor_uop
    print("** post pickle")
    t2:Tensor = pickle.loads(st)
    assert t2.uop.is_realized, f"expected {t2.uop} to be realized"
    np.testing.assert_equal(t_values, t2.numpy())

  # NOTE: currently Buffer exists on the uop, not tensor
  def test_pickle_buffer_uop(self):
    t = Tensor.arange(4).clone().realize()
    a = t.uop
    assert a.is_realized
    self.assertIsNotNone(buffer:=a.base.realized)
    s = pickle.dumps(a)
    # free buffers
    del a
    del buffer
    a2:UOp = pickle.loads(s)
    self.assertListEqual(a2.base.realized.as_memoryview().cast("I").tolist(), [0, 1, 2, 3])

  @unittest.skipIf(DEV.interface.startswith("MOCK"), "mock device buffers live in host RAM, not VRAM")
  def test_pickle_oob_ram(self):
    N, M = 8, 10**6
    ts = [Tensor.rand(M, dtype='float32').realize() for _ in range(N)]
    tracemalloc.start()
    st = pickle.dumps(ts, protocol=5, buffer_callback=lambda pb: pb.release())
    self.assertLess(tracemalloc.get_traced_memory()[1], N*M*4)
    tracemalloc.reset_peak()
    def make_fake_buffers():
      for _ in range(N):
        Device[Device.DEFAULT].synchronize()
        yield pickle.PickleBuffer(bytearray(M*4))
    pickle.loads(st, buffers=make_fake_buffers())
    self.assertLess(tracemalloc.get_traced_memory()[1], N*M*4)
    tracemalloc.stop()

  def test_pickle_unrealized_tensor(self):
    t = Tensor.ones(10, 10)
    st = pickle.dumps(t)
    t2:Tensor = pickle.loads(st)
    np.testing.assert_equal(t.numpy(), t2.numpy())

  def test_pickle_variable(self):
    v = Variable("i", 1, 20).bind(10)
    t1 = Tensor.ones(10, v).contiguous()
    t2 = Tensor.ones(10, v).contiguous()
    ret = (t1+t2).sum(1)
    st = pickle.dumps(ret)
    del ret
    vt2 = pickle.loads(st)
    np.testing.assert_equal(vt2.numpy(), 20)

  def test_pickle_buffer_view(self):
    t = Tensor.arange(10).clone(device="CPU").realize()
    vt = t[3:5].contiguous().realize()
    assert hasattr(vt.uop.buffer, 'base')
    ref_value = vt.tolist()
    st = pickle.dumps(vt)
    del t, vt
    vt2 = pickle.loads(st)
    assert hasattr(vt2.uop.buffer, 'base')
    assert ref_value == vt2.tolist()

  def test_pickle_numpy(self):
    t = Tensor(np.array([1,2,3,4.]), dtype=dtypes.float32)
    st = pickle.dumps(t)
    t2:Tensor = pickle.loads(st)
    np.testing.assert_equal(t.numpy(), t2.numpy())

  def test_pickle_jit(self):
    @TinyJit
    def add(a, b): return a.sum()+b+1
    for _ in range(3): add(Tensor.rand(10, 10), Tensor.rand(10, 10))
    st = pickle.dumps(add)
    del add

    add_fxn = pickle.loads(st)
    x = Tensor.ones(10, 10).contiguous().realize()
    y = Tensor.ones(10, 10).contiguous().realize()
    print("post jit")
    out = add_fxn(x, y)
    np.testing.assert_equal(out.numpy(), 102)

  def test_pickle_jit_no_del(self):
    @TinyJit
    def fn(x): return x + 1.0
    for _ in range(3): fn(Tensor.randn(4))
    loaded = pickle.loads(pickle.dumps(fn))
    self.assertEqual(loaded(Tensor([1.0,2.0,3.0,4.0])).tolist(), [2.0,3.0,4.0,5.0])

  def test_pickle_context_var(self):
    v = ContextVar("test_var", 0)
    with Context(test_var=1):
      vs = pickle.dumps(v)
    v2 = pickle.loads(vs)
    self.assertEqual(v2.value, 1)

  def test_pickle_schedule(self):
    a = Tensor([1,2])
    out = a + 2
    sched = out.schedule_linear()
    pk = pickle.dumps(sched)
    sched_pk = pickle.loads(pk)
    self.assertEqual(sched_pk.src[-1].src[0], sched.src[-1].src[0])

  def test_pickle_renderer(self):
    from tinygrad.device import Device
    pk = pickle.dumps(Device.default.renderer)
    pickle.loads(pk)

class TestPickleJIT(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    N = 10
    @TinyJit
    def add(a, b): return a.sum()+b+1
    for _ in range(3): add(Tensor.rand(N, N), Tensor.rand(N, N))
    cls.st = pickle.dumps(add)
    del add

  def test_inspect(self):
    import io
    class FakeClass:
      def __init__(self, *args, **kwargs):
        print(self.module, self.name)
    class InspectUnpickler(pickle.Unpickler):
      def find_class(self, module, name): return type("SpecializedFakeClass", (FakeClass,), {"name": name, "module": module})
    InspectUnpickler(io.BytesIO(self.st)).load()

  @unittest.skip("we are still saving intermediate buffers")
  def test_size(self):
    # confirm no intermediate buffers are saved
    self.assertLess(len(self.st), 1_000_000)

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