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

class TestPickle(unittest.TestCase):
  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).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).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_no_storage_aliasing(self):
    # loading the same pickle twice gives fully independent storage: the buffers (and their BUFFER uops) are never shared
    t = Tensor([1,2,3,4]).realize()
    st = pickle.dumps(t)
    t1, t2 = pickle.loads(st), pickle.loads(st)
    self.assertIsNot(t1.uop, t2.uop)
    self.assertIsNot(t1.uop.base.buffer, t2.uop.base.buffer)
    t1.assign(Tensor([9,9,9,9])).realize()
    self.assertListEqual(t1.tolist(), [9,9,9,9])
    self.assertListEqual(t2.tolist(), [1,2,3,4])

  def test_pickle_view_is_self_contained(self):
    # a pickled graph carries its own buffer: data from earlier loads of related graphs must not leak into it
    t = Tensor([1,2,3,4]).realize()
    t1 = pickle.loads(pickle.dumps(t))
    t1.assign(Tensor([9,9,9,9])).realize()
    # loading a view of the original tensor must give the pickled values ([2,3]), not the mutated values from the other load
    v2 = pickle.loads(pickle.dumps(t[1:3]))
    self.assertListEqual(v2.realize().tolist(), [2,3])

  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])

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()