import numpy as np
from collections import deque
from ditk import logging
from time import time

from ding.framework import task
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from ding.framework.context import Context


def epoch_timer(print_per: int = 1, smooth_window: int = 10):
    """
    Overview:
        Print time cost of each epoch.
    Arguments:
        - print_per (:obj:`int`): Print each N epoch.
        - smooth_window (:obj:`int`): The window size to smooth the mean.
    """
    records = deque(maxlen=print_per * smooth_window)

    def _epoch_timer(ctx: "Context"):
        start = time()
        yield
        time_cost = time() - start
        records.append(time_cost)
        if ctx.total_step % print_per == 0:
            logging.info(
                "[Epoch Timer][Node:{:>2}]: Cost: {:.2f}ms, Mean: {:.2f}ms".format(
                    task.router.node_id or 0, time_cost * 1000,
                    np.mean(records) * 1000
                )
            )

    return _epoch_timer