# Source code for labml.tracker

    from typing import Dict, overload, Optional

    from labml.internal.tracker import tracker_singleton as _internal

    def set_global_step(global_step: Optional[int]):
        """
        Set the current step for tracking

        Arguments:
            global_step (int): Global step
        """
        _internal().set_global_step(global_step)

    def add_global_step(increment_global_step: int = 1):
        """
        Increment the current step for tracking

        Arguments:
            increment_global_step (int, optional): By how much to increment the global step.
                Defaults to ``1`` if not provided.
        """
        _internal().add_global_step(int(increment_global_step))

    def get_global_step() -> int:
        """
        Returns current step
        """
        return _internal().global_step

    def set_histogram(name: str, is_print: bool = False):
        """
        Set indicator type to be a histogram.
        It will log the tracked values as a histogram.

        Arguments:
            name (str): Name of the indicator
            is_print: (bool, optional): Whether the indicator should be printed in console.
                Defaults to ``False``.
        """
        from labml.internal.tracker.indicators.numeric import Histogram
        _internal().add_indicator(Histogram(name, is_print))

    def set_scalar(name: str, is_print: bool = False):
        """
        Set indicator type to be a scalar.
        It will log a scalar of the tracked values.
        If there are multiple values it will log the mean.

        Arguments:
            name (str): Name of the indicator
            is_print: (bool, optional): Whether the indicator should be printed in console.
                Defaults to ``False``.
        """
        from labml.internal.tracker.indicators.numeric import Scalar
        _internal().add_indicator(Scalar(name, is_print))

    def _add_dict(values: Dict[str, any]):
        for k, v in values.items():
            _internal().store(k, v)

    @overload
    def add(values: Dict[str, any]):
        ...

    @overload
    def add(name: str, value: any):
        ...

    @overload
    def add(**kwargs: any):
        ...

    def add(*args, **kwargs):
        """
        This has multiple overloads

        .. function:: add(values: Dict[str, any])
            :noindex:

        .. function:: add(name: str, value: any)
            :noindex:

        .. function:: add(**kwargs: any)
            :noindex:

        This add tracking information to a temporary queue.
        These will be saved when :func:`labml.tracker.save` is called.

        You should use :func:`labml.tracker.add` to improve performance since
        saving tracking information consumes time.
        Although saving takes negligible amount of time it can add up if called very frequently.

        Arguments:
            values (Dict[str, any]): A dictionary of key-value pairs to track
            name (str): The name of the value to be tracked
            value (any): The value to be tracked
            kwargs: Key-value pairs to track
        """
        if len(args) > 2:
            raise TypeError('tracker.add should be called as add(name, value), add(dictionary) or add(k=v,k2=v2...)')

        if len(args) == 0:
            _add_dict(kwargs)
        elif len(args) == 1:
            if kwargs:
                raise TypeError('tracker.add should be called as add(name, value), add(dictionary) or add(k=v,k2=v2...)')
            if not isinstance(args[0], dict):
                raise TypeError('tracker.add should be called as add(name, value), add(dictionary) or add(k=v,k2=v2...)')
            _add_dict(args[0])
        elif len(args) == 2:
            if kwargs:
                raise TypeError('tracker.add should be called as add(name, value), add(dictionary) or add(k=v,k2=v2...)')
            if not isinstance(args[0], str):
                raise TypeError('tracker.add should be called as add(name, value), add(dictionary) or add(k=v,k2=v2...)')
            _internal().store(args[0], args[1])

    @overload
    def save():
        ...

    @overload
    def save(global_step: int):
        ...

    @overload
    def save(values: Dict[str, any]):
        ...

    @overload
    def save(name: str, value: any):
        ...

    @overload
    def save(**kwargs: any):
        ...

    @overload
    def save(global_step: int, values: Dict[str, any]):
        ...

    @overload
    def save(global_step: int, name: str, value: any):
        ...

    @overload
    def save(global_step: int, **kwargs: any):
        ...

    def save(*args, **kwargs):
        r"""
        This has multiple overloads

        .. function:: save()
            :noindex:

        .. function:: save(global_step: int)
            :noindex:

        .. function:: save(values: Dict[str, any])
            :noindex:

        .. function:: save(name: str, value: any)
            :noindex:

        .. function:: save(**kwargs: any)
            :noindex:

        .. function:: save(global_step: int, values: Dict[str, any])
            :noindex:

        .. function:: save(global_step: int, name: str, value: any)
            :noindex:

        .. function:: save(global_step: int, **kwargs: any)
            :noindex:

        This saves the tracking information.

        Arguments:
            global_step (int): The current step
            values (Dict[str, any]): A dictionary of key-value pairs to track
            name (str): The name of the value to be tracked
            value (any): The value to be tracked
            kwargs: Key-value pairs to track
        """
        if len(args) > 0 and type(args[0]) == int:
            _internal().set_global_step(args[0])
            args = args[1:]

        if len(args) > 0 or len(kwargs.keys()) > 0:
            add(*args, **kwargs)

        _internal().write()

    def new_line():
        r"""
        Prints a new line.

        Equivalent to ``logger.log``, but handles distributed training where only the rank=0
        process is tracking data.
        """
        _internal().new_line()

    def namespace(name: str):
        r"""
        Set a namespace for tracking
        """
        return _internal().namespace(name)

    def reset():
        r"""
        Reset indicators, for a new experiment
        """
        _internal().reset_store()