#!/usr/bin/env python
# coding: utf-8

# In[21]:


import torch


# In[22]:


import weakref, traceback


# In[23]:


torch.backends.mps.is_available()


# In[24]:


t = { 1,2,3}
import gc
len (gc.get_referrers(t))


# In[25]:


gc.get_referrers(t)[0].keys()


# In[26]:


[(o['__package__'],o['__name__']) for o in gc.get_referrers(t)]


# In[27]:


import sys


# In[28]:


import gc
def tensors():
    objs = gc.get_objects()
    for v in objs:
        if isinstance(v, torch.Tensor):
            yield  weakref.ref(v)

def print_tensor_size(message=""):
    gc.collect()
    if message: print(message)
    for v in tensors():
        if isinstance(v(), torch.Tensor):
            print( id(v()), v().shape, len(gc.get_referrers(v())))
            #show_holders(v)
            
    gc.collect()

print_tensor_size()


# In[ ]:





# In[29]:


# [type(o) for l in gc.get_referrers(t) for o in l]


# In[30]:


def test_exception_raise(should_raise=False):
    try:
        allocate_tensor(should_raise)
    finally:
        print_tensor_size('Finally block')
   
def allocate_tensor(should_raise=False):  
    x = torch.rand(10000, 10000)
    print_tensor_size('after allocation') 
    x *= 2
    if should_raise:
        raise Exception('I am a bad function')
    return x


# In[31]:


def clean_up():
    import sys, gc, traceback
    if 'last_traceback' in sys.__dict__:
        print(sys.last_traceback)
        print(get_ipython().InteractiveTB.tb)
        traceback.clear_frames(sys.last_traceback)
        del sys.last_traceback 
    if 'last_type' in sys.__dict__: del sys.last_type
    if 'last_value' in sys.__dict__: del sys.last_value
    gc.collect()
# above code translated to use hasttr and delattr

# def clean_up():
#     import sys, gc, traceback
#     if hasattr(sys, 'last_traceback'):
#         traceback.clear_frames(sys.last_traceback)
#         delattr(sys, 'last_traceback')
#     if hasattr(sys, 'last_type'): delattr(sys, 'last_type')
#     if hasattr(sys, 'last_value'): delattr(sys, 'last_value')
#     gc.collect()    


# In[32]:


import types
class hiddenref:
    def __new__(cls, r):
        try:
            return weakref.ref(r)
        except:
            self = super().__new__(cls)
            self.hiddenref=r
    def __call__(self):
        return self.hiddenref


# In[33]:


def dump_holders(v, t, f=repr):
    for r in gc.get_referrers(v()):
        if str(t) in str(type(r)):
            v = f(r)
            if v:
                print(v[:400])


# In[34]:


def get_holders(v, t):
    for r in gc.get_referrers(v()):
        if str(t) in str(type(r)):
            if isinstance(r, dict) and 'hiddenref' in r: continue
            yield hiddenref(r)


# In[35]:


test_exception_raise()
print_tensor_size('Outside')


# In[36]:


test_exception_raise(should_raise=True)


# In[37]:


import types

def find_type(variables, type, along_type=None, debug=False):
    seen = set()
    objs = []
    
    q = [v() if isinstance(v, (weakref.ref,hiddenref)) else v for v in variables]
    while q:
        v = q.pop()
        gc.collect()
        for r in gc.get_referrers(v):
            if id(r) in seen: continue
            seen.add(id(r))
            
            if debug: print(type(r), id(r))
            if isinstance(r, dict) and 'hiddenref' in r: continue
            if isinstance(r, type):
                objs.append(r)
            else:
                if debug: print(len(q), len(seen))
                if along_type and isinstance(r, along_type):
                    q.append(r)
        if debug: print('q',len(q), [id(v) for v in q])
    return objs

def find_tracebacks(*variables):
    frames = find_type(variables, types.FrameType, debug=False)
    tracebacks = find_type(frames, types.TracebackType,  types.FrameType, debug=False)
    return list(set(tracebacks))


# In[38]:


def print_ref_graph(obj):
    """Returns a list of references keep this object alive"""
    orig_obj = obj
    seen = set()
 
    gc.collect()
    level = 0
    branches = []
    while obj is not None:
        nobj = None
        path_printed = False
        for nobj in gc.get_referrers(obj):
            if path_printed and id(nobj) not in seen and isinstance(nobj, (list, dict, types.FrameType, types.TracebackType)):
                branches.append((level-1, nobj))
                continue
            print(' '*level, end='')
            if id(nobj) in seen: 
                print(f'{id(nobj)} cycle, skipping')
                nobj = None
                continue
            else: 
                seen.add(id(nobj))

            if isinstance(nobj, dict) and 'hiddenref' in nobj: 
                nobj = None
                continue
            elif isinstance(nobj, types.FrameType): 
                print('frame', nobj.f_code.co_name)
            elif isinstance(nobj, types.TracebackType): 
                print('traceback', nobj.tb_frame.f_code.co_name, id(nobj))
            elif isinstance(nobj, (tuple,list)):
                for i in range(len(nobj)):
                    if nobj[i] is obj:
                        print(type(nobj), i)
                        break
                else:
                    print('notfound', type(nobj))
            elif isinstance(nobj, dict):
                for k in nobj:
                    if nobj[k] is obj:
                        print(type(nobj), k)
                        break
                else:
                    print('notfound', type(nobj))
            elif isinstance(nobj, types.ModuleType):
                print('module', nobj.__name__)
            else:
                print(type(nobj), nobj.__name__ if hasattr(nobj, '__name__') else id(nobj))
            
            path_printed = True  # effectively break
            obj = nobj
            level += 1
             
            
        if not path_printed:
            if branches: 
               
                level, obj = branches.pop()
                print("-"*(level-1),"Starting over")
            else:
                print('done')
                obj = None
        else: 
           pass
       
for tb in find_tracebacks(*list(tensors())):
    print_ref_graph(tb)
    


# In[ ]:





# In[42]:


test_exception_raise(should_raise=True)


# In[160]:


def make_empty_traceback():
    try: 
        raise Exception('Empty traceback')
    except Exception as e:
        return e.__traceback__
make_empty_traceback()


# In[43]:


print_tensor_size('Before clean up')
agressive_clear_traceback()
clean_up()
print_tensor_size('After clean up')


# In[46]:


print_tensor_size('After clean up')


# In[ ]: