import gc
import sys
from chimerax.core.commands import open, close
from chimerax.core.atomic import AtomicStructure

# PDB_ID = '2hmg'      # large protein
PDB_ID = '3rec'      # smallest PDB file
# PDB_ID = 'water.cif'   # just waters


def gc_histogram():
    """Returns per-class counts of existing objects."""
    result = {}
    for o in gc.get_objects():
        t = type(o)
        count = result.get(t, 0)
        result[t] = count + 1
    return result


def diff_hists(h1, h2):
    """Prints differences between two results of gcHistogram()."""
    for k in h2:
        if k in h1:
            continue
        h1[k] = 0
    for k in h1:
        if k not in h2:
            h2[k] = 0
        if h1[k] != h2[k]:
            print("%s: %d -> %d (%s%d)" % (
                k, h1[k], h2[k], h2[k] > h1[k] and "+" or "",
                h2[k] - h1[k]))


# open/close once to populate caches
models = open.open(session, PDB_ID)
close.close(session, models)
del models
session.main_view.check_for_drawing_change()

gc.collect()
if 0:
    gc.set_debug(gc.DEBUG_LEAK)
else:
    gc.set_debug(gc.DEBUG_SAVEALL)
hist1 = gc_histogram()
models = open.open(session, PDB_ID)
for m in models:
    print('After open:', m, sys.getrefcount(m))
del m
close.close(session, models)
del models
session.main_view.check_for_drawing_change()

gc.collect()
hist2 = gc_histogram()
diff_hists(hist1, hist2)

for o in gc.get_objects():
    if isinstance(o, AtomicStructure):
        break
else:
    print('No AtomicStructure references!!')
    raise SystemExit(0)
gc.collect()

del hist1, hist2

print(sys.getrefcount(o), 'references to an AtomicStructure')

if 1:
    referrers = gc.get_referrers(o)
    refs = ['%r from %s' % (type(x), type(x).__module__) for x in referrers
            if id(x) != id(globals())]
    print(len(refs), 'referrers:', ', '.join(refs))

if 0:
    referents = gc.get_referents(o)
    refs = ['%r from %s' % (type(x), type(x).__module__) for x in referents]
    print(len(referents), 'referents', ', '.join(refs))

if 0:
    tmp = ["%r from %s" % (type(g), type(g).__module__) for g in gc.garbage]
    print(len(gc.garbage), "items in garbage:", ', '.join(tmp))

raise SystemExit(1)