import collections
import sys
import typing as tp  # NOQA

from chainer.backends import cuda
from chainer import function_hook


try:
    MemoryHook = cuda.cupy.cuda.memory_hook.MemoryHook  # type: tp.Any # to handle https://github.com/python/mypy/issues/2477 # NOQA
    memory_hook_available = True
except Exception as e:
    _resolution_error = e
    MemoryHook = object
    memory_hook_available = False


class CupyMemoryProfileHook(function_hook.FunctionHook):
    """Function hook for measuring memory usage of functions in cupy memory pool.

    Example:
        Code example::

            from chainer.function_hooks import CupyMemoryProfileHook
            hook = CupyMemoryProfileHook()
            with hook:
                trainer.run()
            hook.print_report()

        Output example::

                   FunctionName  UsedBytes  AcquiredBytes  Occurrence
                 LinearFunction     5.16GB       179.98MB        3900
                           ReLU     0.99GB       458.97MB        2600
            SoftmaxCrossEntropy     0.01GB         5.08MB        1300
                       Accuracy     0.00GB         0.35MB         700

        where *FunctionName* is the name of function that calls the hook, and
        *UsedBytes* is the memory bytes the function used from cupy memory
        pool, and *AcquiredBytes* is the actual memory bytes the cupy memory
        pool acquired from GPU device on the function call, and *Occurrence*
        is the number of calls.
    Attributes:
        call_history: List of measurement results. It consists of the name of