mapanything/utils/timing.py

"""
Utility functions for timing code blocks
"""

import time
from contextlib import ContextDecorator

import numpy as np


class BlockTimeManager:
    """
    Manages a collection of timers and their formatting options.

    This class serves as a central registry for Timer objects, allowing them to be
    accessed by name and maintaining their formatting preferences.

    Attributes:
        timers (dict): Dictionary mapping timer names to Timer objects
        timer_fmts (dict): Dictionary mapping timer names to their display formats
        window_size (int): Default window size for calculating windowed averages
        buf_size (int): Default buffer size for storing timing measurements
    """

    def __init__(self, window_size=10, buf_size=100000):
        self.timers = dict()
        self.timer_fmts = dict()
        self.window_size = window_size
        self.buf_size = buf_size


btm = BlockTimeManager(window_size=100000)


class Timer:
    """
    Core timing class that tracks execution times.

    This class provides the fundamental timing functionality, storing timing measurements
    and calculating various statistics.

    Attributes:
        name (str): Identifier for this timer
        buf_size (int): Maximum number of timing measurements to store
        window_size (int): Number of most recent measurements to use for windowed statistics
        measures_arr (numpy.ndarray): Array storing start and end times of measurements
        current_start (float or None): Start time of current measurement
        current_end (float or None): End time of current measurement
    """

    def __init__(self, name, window_size, buf_size=100000):
        self.name = name
        self.buf_size = buf_size
        self.window_size = window_size
        self.init()

    def init(self):
        """Initialize or reset the timer's state."""
        self.measures_arr = np.empty((0, 2))  # LIFO
        self.current_start = None
        self.current_end = None

    def reset(self):
        """Reset the timer to its initial state."""
        self.init()

    def tic(self):
        """Start a new timing measurement."""
        if self.current_start is not None:
            # another tic executed before a toc
            self.toc()
        self.current_start = time.perf_counter()

    def toc(self):
        """End the current timing measurement."""
        self.current_end = time.perf_counter()
        self._add_current_measure()

    def _add_current_measure(self):
        """Add the current timing measurement to the measurements array."""
        self.measures_arr = np.concatenate(
            [
                np.array([[self.current_start, self.current_end]]),
                self.measures_arr[: self.buf_size],
            ]
        )
        self.current_start = None
        self.current_end = None

    @property
    def avg(self) -> float:
        """Calculate the average execution time across all measurements."""
        return np.mean(self.measures_arr[:, 1] - self.measures_arr[:, 0])

    @property
    def wavg(self) -> float:
        """Calculate the windowed average execution time using the most recent measurements."""
        return np.mean(
            self.measures_arr[: self.window_size, 1]
            - self.measures_arr[: self.window_size, 0]
        )

    @property
    def max(self) -> float:
        """Return the maximum execution time."""
        return np.max(self.measures_arr[:, 1] - self.measures_arr[:, 0])

    @property
    def min(self) -> float:
        """Return the minimum execution time."""
        return np.min(self.measures_arr[:, 1] - self.measures_arr[:, 0])

    @property
    def total(self) -> float:
        """Return the total execution time across all measurements."""
        return np.sum(self.measures_arr[:, 1] - self.measures_arr[:, 0])

    @property
    def latest(self) -> float:
        """Return the most recent execution time."""
        return self.measures_arr[0, 1] - self.measures_arr[0, 0]

    @property
    def median(self) -> float:
        """Return the median execution time."""
        return np.median(self.measures_arr[:, 1] - self.measures_arr[:, 0])

    @property
    def var(self) -> float:
        """Return the variance of execution times."""
        return np.var(self.measures_arr[:, 1] - self.measures_arr[:, 0])


class BlockTimer(ContextDecorator):
    """
    A context manager and decorator for timing code blocks.

    This class provides a convenient interface for timing code execution, either as a
    context manager (with statement) or as a decorator. It uses the Timer class for
    the actual timing functionality.

    Attributes:
        name (str): Identifier for this timer
        fmt (str or None): Format string for displaying timing information
        timer (Timer): The underlying Timer object
        num_calls (int): Number of times this timer has been called
    """

    @staticmethod
    def timers():
        """Return a list of all registered timer names."""
        return list(btm.timers.keys())

    def __init__(self, name, fmt=None, window_size=100):
        self.name = name
        if name in btm.timers:
            self.timer = btm.timers[name]
            # restore format
            self.fmt = fmt if fmt is not None else btm.timer_fmts[name]
        else:
            self.timer = Timer(name, btm.window_size, btm.buf_size)
            btm.timers[name] = self.timer
            btm.timer_fmts[name] = fmt
        self.timer.window_size = window_size
        self._default_fmt = "[{name}] num: {num} latest: {latest:.4f} --wind_avg: {wavg:.4f} -- avg: {avg:.4f} --var: {var:.4f} -- total: {total:.4f}"
        if fmt == "default":
            self.fmt = self._default_fmt
        # extend here for new formats
        else:
            self.fmt = None

        self.num_calls = 0

    def __enter__(self) -> "Timer":
        """Start timing when entering a context."""
        self.tic()
        return self

    def __exit__(self, *args):
        """End timing when exiting a context and optionally display results."""
        self.toc()
        if self.fmt is not None:
            print(str(self))

    def __str__(self) -> str:
        """Return a string representation of the timer."""
        return self.display()

    def reset(self):
        """Reset the timer and call counter."""
        self.timer.reset()
        self.num_calls = 0

    def display(self, fmt=None):
        """
        Format and return timing information.

        Args:
            fmt (str, optional): Format string to use. If None, uses the timer's format.

        Returns:
            str: Formatted timing information
        """
        if fmt is None:
            if self.fmt is not None:
                fmt = self.fmt
            else:
                fmt = self._default_fmt
        return fmt.format(
            name=self.name,
            num=self.num_calls,
            latest=self.latest,
            wavg=self.wavg,
            avg=self.avg,
            var=self.var,
            total=self.total,
        )

    def tic(self):
        """Start a new timing measurement and increment the call counter."""
        self.timer.tic()
        self.num_calls += 1

    def toc(self, display=False):
        """
        End the current timing measurement.

        Args:
            display (bool): Whether to return a formatted display string

        Returns:
            str or None: Formatted timing information if display is True
        """
        self.timer.toc()
        if display:
            return self.display()

    @property
    def latest(self) -> float:
        """Return the most recent execution time."""
        return self.timer.latest

    @property
    def avg(self) -> float:
        """Return the average execution time."""
        return self.timer.avg

    @property
    def wavg(self) -> float:
        """Return the windowed average execution time."""
        return self.timer.wavg

    @property
    def max(self) -> float:
        """Return the maximum execution time."""
        return self.timer.max

    @property
    def min(self) -> float:
        """Return the minimum execution time."""
        return self.timer.min

    @property
    def total(self) -> float:
        """Return the total execution time."""
        return self.timer.total

    @property
    def median(self) -> float:
        """Return the median execution time."""
        return self.timer.median

    @property
    def var(self) -> float:
        """Return the variance of execution times."""
        return self.timer.var


if __name__ == "__main__":

    @BlockTimer("fct", "default")
    def fct(bobo):
        time.sleep(0.5)

    fct(2)

    for i in range(10):
        with BlockTimer("affe", "default"):
            time.sleep(0.1)
    for i in range(1000):
        with BlockTimer("test", None):
            time.sleep(0.001)

        # BlockTimer("test").display = f"""avg: {BlockTimer("test").avg}  total: {BlockTimer("test").total}"""
        # print(str(BlockTimer("test")))

    print(BlockTimer("test"))
    BlockTimer("test").tic()
    BlockTimer("t2", "default").tic()
    time.sleep(0.4)
    print(BlockTimer("t2").toc(True))

    time.sleep(0.4)
    print(BlockTimer("test").toc(True))