#!/usr/bin/python
# Written by Boudewijn Schoon
# see LICENSE.txt for license information

"""
Use the garbage collector to monitor memory usage
"""

from types import (IntType, FloatType, StringType, UnicodeType, TupleType, ListType, DictType, FunctionType,
                   ModuleType, FrameType)
import gc
import thread
import time
import logging

logger = logging.getLogger(__name__)


def _get_default_footprint(obj, depth):
    return 4


def _get_int_footprint(obj, depth):
    return 4


def _get_float_footprint(obj, depth):
    return 8


def _get_string_footprint(obj, depth):
    return len(obj)


def _get_unicode_footprint(obj, depth):
    return 2 * len(obj)


def _get_tuple_footprint(obj, depth):
    if depth == 0:
        return 4 + 4 * len(obj)
    else:
        return 4 + 4 * len(obj) + sum(map(lambda o: get_memory_footprint(o, depth), obj))


def _get_list_footprint(obj, depth):
    if depth == 0:
        return 8 + 4 * len(obj)
    else:
        if len(obj) in (2, 3):
            logger.info("Len: %s %s", type(obj[0]), type(obj[1]))
            logger.info(repr(obj))
            return 42
        logger.info("Len: %d", len(obj))
        return 8 + 4 * len(obj) + sum(map(lambda o: get_memory_footprint(o, depth), obj))


def _get_dict_footprint(obj, depth):
    if depth == 0:
        return 32 + 8 * len(obj)
    else:
        return 32 + 8 * len(obj) \
            + sum(map(lambda o: get_memory_footprint(o, depth), obj.iterkeys())) \
            + sum(map(lambda o: get_memory_footprint(o, depth), obj.itervalues()))

memory_footprint_map = {IntType: _get_int_footprint,
                        FloatType: _get_float_footprint,
                        StringType: _get_float_footprint,
                        UnicodeType: _get_unicode_footprint,
                        TupleType: _get_tuple_footprint,
                        ListType: _get_list_footprint,
                        DictType: _get_dict_footprint}


def get_memory_footprint(obj, depth=100):
    return memory_footprint_map.get(type(obj), _get_default_footprint)(obj, depth - 1)


def _get_default_description(obj):
    return type(obj)


def _get_function_description(obj):
    return "<function '%s' from '%s'>" % (obj.__name__, obj.__module__)


def _get_module_description(obj):
    return str(obj)


def _get_frame_description(obj):
    return "<frame for '%s' from %s:%d >" % (obj.f_code.co_name, obj.f_code.co_filename, obj.f_code.co_firstlineno)

description_map = {FunctionType: _get_function_description,
                   ModuleType: _get_module_description,
                   FrameType: _get_frame_description}


def get_description(obj):
    return description_map.get(type(obj), _get_default_description)(obj)


def get_datetime():
    return time.strftime("%Y/%m/%d %H:%M:%S")


def byte_uint_to_human(i, fmt="%(value).1f%(unit)s"):
    """Convert a number into a formatted string.

    format: %(value)d%(unit)s
    1           --> 1B
    1024        --> 1KB
    1048576     --> 1MB
    1073741824  --> 1GB

    format: %(value).1f %(unit-long)s
    1           --> 1.0 byte
    2           --> 2.0 bytes

    todo:
    - uint_to_human(1025, format="%(value)d %(unit-long)s") --> '1 kilobytes'
      however, this should result in '1 kilobyte'

    """
    assert type(i) in (int, long)
    assert i >= 0
    assert isinstance(fmt, str)
    dic = {}
    if i < 1024:
        dic["value"] = i
        dic["unit"] = "B"
        dic["unit-long"] = (i == 1 and "byte" or "bytes")
    elif i < 1048576:
        dic["value"] = i / 1024.0
        dic["unit"] = "KB"
        dic["unit-long"] = (i == 1024 and "kilobyte" or "kilobytes")
    elif i < 1073741824:
        dic["value"] = i / 1048576.0
        dic["unit"] = "MB"
        dic["unit-long"] = (i == 1048576 and "megabyte" or "megabytes")
    else:
        dic["value"] = i / 1073741824.0
        dic["unit"] = "GB"
        dic["unit-long"] = (i == 1073741824 and "gigabyte" or "gigabytes")

    return fmt % dic


def monitor(delay=10.0, interval=60.0, min_footprint=100000):
    def parallel():
        time.sleep(delay)

        history = [min_footprint]
        while True:
            high_foot = 0
            history = history[-2:]
            low_foot = min(history)
            datetime = get_datetime()
            logger.info("Memory: %s using minimal footprint: %s", datetime, byte_uint_to_human(low_foot))

            gc.collect()
            for obj in gc.get_objects():
                if type(obj) in (TupleType, ListType, DictType, StringType, UnicodeType):
                    try:
                        footprint = get_memory_footprint(obj)
                    except:
                        logger.error("Memory: %s unable to get footprint for %s", datetime, get_description(obj))
                    else:
                        if footprint > high_foot:
                            high_foot = footprint
                        if footprint >= low_foot:
                            logger.info("Memory: %s, %s footprint: %s", datetime, get_description(obj),
                                        byte_uint_to_human(footprint))
                            for referrer in gc.get_referrers(obj):
                                logger.info("Memory: %s REF %s", datetime, get_description(referrer))
                            logger.info("Memory")

            history.append(high_foot)
            time.sleep(interval)

    thread.start_new_thread(parallel, ())


def main():
    """
    Test the memory monitor
    """
    monitor(1.0)
    time.sleep(10)

if __name__ == "__main__":
    main()
