"""Test namedtuple unpickling memory usage issues."""

from __future__ import print_function

import collections
import gc
import pickletools
import random
import string
import subprocess
import sys

try:
    import cPickle as pickle
except ImportError:
    import pickle  # cPickle is just pickle on Python 3

if sys.platform == 'win32':
    # From my recipe at:
    # http://code.activestate.com/recipes/578513-get-memory-usage-of-windows-processes-using-getpro/
    import ctypes
    from ctypes import wintypes

    GetCurrentProcess = ctypes.windll.kernel32.GetCurrentProcess
    GetCurrentProcess.argtypes = []
    GetCurrentProcess.restype = wintypes.HANDLE

    SIZE_T = ctypes.c_size_t

    class PROCESS_MEMORY_COUNTERS_EX(ctypes.Structure):
        _fields_ = [
            ('cb', wintypes.DWORD),
            ('PageFaultCount', wintypes.DWORD),
            ('PeakWorkingSetSize', SIZE_T),
            ('WorkingSetSize', SIZE_T),
            ('QuotaPeakPagedPoolUsage', SIZE_T),
            ('QuotaPagedPoolUsage', SIZE_T),
            ('QuotaPeakNonPagedPoolUsage', SIZE_T),
            ('QuotaNonPagedPoolUsage', SIZE_T),
            ('PagefileUsage', SIZE_T),
            ('PeakPagefileUsage', SIZE_T),
            ('PrivateUsage', SIZE_T),
        ]

    GetProcessMemoryInfo = ctypes.windll.psapi.GetProcessMemoryInfo
    GetProcessMemoryInfo.argtypes = [
        wintypes.HANDLE,
        ctypes.POINTER(PROCESS_MEMORY_COUNTERS_EX),
        wintypes.DWORD,
    ]
    GetProcessMemoryInfo.restype = wintypes.BOOL

    def get_current_process():
        """Return handle to current process."""
        return GetCurrentProcess()

    def get_memory_info(process=None):
        """Return Win32 process memory counters structure as a dict."""
        if process is None:
            process = get_current_process()
        counters = PROCESS_MEMORY_COUNTERS_EX()
        ret = GetProcessMemoryInfo(process, ctypes.byref(counters),
                                   ctypes.sizeof(counters))
        if not ret:
            raise ctypes.WinError()
        info = dict((name, getattr(counters, name))
                    for name, _ in counters._fields_)
        return info

    def get_memory_usage(process=None):
        """Return this process's memory usage in bytes."""
        info = get_memory_info(process=process)
        return info['PrivateUsage']
else:
    def get_memory_usage():
        """Dummy version of get_memory_usage() for non-Windows systems."""
        return int(raw_input('Enter memory process is using in KB:')) * 1024

try:
    xrange = xrange
except NameError:
    xrange = range

Field = collections.namedtuple('Field', 'name value')

def random_string(length, chars=string.digits + string.ascii_lowercase):
    return ''.join(random.choice(chars) for i in xrange(length))

def make_thing():
    return [Field(random_string(10), random_string(10)) for i in xrange(100)]

def make_things(n=1000):
    return dict((i, make_thing()) for i in xrange(n))

def print_mem_usage(prefix):
    gc.collect()
    print(prefix, int(round(get_memory_usage() / 1024.0)), 'KB')

def dump(filename, obj):
    with open(filename, 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

def load(filename):
    with open(filename, 'rb') as f:
        return pickle.load(f)

def optimize(in_filename, out_filename):
    with open(in_filename, 'rb') as f:
        orig = f.read()
    opt = pickletools.optimize(orig)
    with open(out_filename, 'wb') as f:
        f.write(opt)
    return len(orig), len(opt)

def do_dump():
    print(sys.version)
    dump('unpickletest.dat', make_things())

def do_load():
    things = load('unpickletest.dat')
    print_mem_usage('Memory usage after loading normal:')

def do_opt():
    before, after = optimize('unpickletest.dat', 'unpickletest_opt.dat')
    print('File size before and after optimization:',
          int(round(before / 1024.0)), 'KB ->', int(round(after / 1024.0)), 'KB')

def do_load_opt():
    things = load('unpickletest_opt.dat')
    print_mem_usage('Memory usage after loading optimized:')

def do_all():
    subprocess.call([sys.executable, 'unpickletest.py', 'dump'])
    subprocess.call([sys.executable, 'unpickletest.py', 'opt'])
    subprocess.call([sys.executable, 'unpickletest.py', 'load'])
    subprocess.call([sys.executable, 'unpickletest.py', 'load_opt'])

if __name__ == '__main__':
    random.seed('Seed the random number generator with a constant')
    command = sys.argv[1] if len(sys.argv) > 1 else 'all'
    func = globals()['do_' + command]
    func()