from __future__ import annotations
"""
# https://github.com/naist-nlp/mbrs/blob/main/LICENSE
MIT License

Copyright (c) 2024 Hiroyuki Deguchi

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

"""
Timer module including some useful global objects.

Example:
    >>> from mbrs import timer
    >>> for hyps, src in zip(hypotheses, sources):
    ...     with timer.measure("encode/hypotheses") as t:
    ...         h = metric.encode(hyps)
    ...     t.set_delta_ncalls(len(hyps))
    ...     with timer.measure("encode/source"):
    ...         s = metric.encode([src])
    ...     with timer.measure("score"):
    ...         scores = metric.score(h, s)
    >>> res = timer.aggregate().result()  # return the result table
"""

import contextlib
import time
from collections import defaultdict
from dataclasses import dataclass, field


class Stopwatch:
    """Stopwatch class to measure the elapsed time.

    Example:
        >>> timer = Stopwatch()
        >>> for i in range(10):
        ...     with timer():
        ...         time.sleep(1)
        >>> print(f"{timer.elapsed_time:.3f}")
        10.000

        >>> timer = Stopwatch()
        >>> for i in range(10):
        ...     with timer() as t:
        ...         time.sleep(2)
        ...         t.set_delta_ncalls(2)
        >>> print(f"{timer.elapsed_time:.3f}")
        20.000
        >>> print(f"{timer.ncalls}")
        20
    """

    def __init__(self) -> None:
        self.reset()

    def reset(self) -> None:
        """Reset the stopwatch."""
        self._acc_time = 0.0
        self._acc_ncalls = 0
        self._delta_ncalls = 1

    @contextlib.contextmanager
    def __call__(self):
        """Measure the time."""
        _acc_time = self._acc_time
        _acc_ncalls = self._acc_ncalls
        start = time.perf_counter()
        try:
            self.set_delta_ncalls(1)  # Set to 1 in the default
            yield self
        finally:
            # Treat nest calling
            if _acc_time == self._acc_time:
                self._acc_time = _acc_time + time.perf_counter() - start
            if _acc_ncalls == self._acc_ncalls:
                self._acc_ncalls = _acc_ncalls + self._delta_ncalls

    def set_delta_ncalls(self, delta: int = 1):
        """Set delta for counting the number of calls."""
        self._delta_ncalls = delta

    @property
    def elpased_time(self) -> float:
        """Return the total elapsed time."""
        return self._acc_time

    @property
    def ncalls(self) -> int:
        """Return the number of calls."""
        return self._acc_ncalls


class StopwatchDict(defaultdict[str, Stopwatch]):
    """A dictionary of the :class:`Stopwatch` class.

    Example:
        >>> timers = StopwatchDict()
        >>> for i in range(10):
        ...     with timers("A"):
        ...         time.sleep(1)
        >>> for i in range(3):
        ...     with timers("B"):
        ...         time.sleep(1)
        >>> print(f"{timers.total}")
        {"A": 10.000, "B": 3.000}
    """

    def __init__(self) -> None:
        super().__init__(Stopwatch)

    def reset(self) -> None:
        """Reset all stopwatches."""
        for t in self.values():
            t.reset()

    @contextlib.contextmanager
    def __call__(self, name: str):
        """Measure the time."""
        with self[name]() as timer:
            try:
                yield timer
            finally:
                pass

    @property
    def elapsed_time(self) -> dict[str, float]:
        """Return the total elapsed time."""
        return {k: v.elpased_time for k, v in self.items()}

    @property
    def ncalls(self) -> dict[str, int]:
        """Return the number of calls."""
        return {k: v.ncalls for k, v in self.items()}


measure = StopwatchDict()


@dataclass
class ProfileTree:
    elapsed_time: float = -1.0
    ncalls: int = -1
    children: dict[str, ProfileTree] = field(default_factory=dict)

    @property
    def is_leaf(self) -> bool:
        """Return whether the node is leaf or not."""
        return len(self.children) == 0

    def aggregate(self):
        if self.is_leaf:
            if self.elapsed_time < 0.0:
                raise RuntimeError("Missing elapsed_time")
            if self.ncalls < 0:
                raise RuntimeError("Missing ncalls")
        else:
            for child in self.children.values():
                child.aggregate()
            if self.elapsed_time < 0.0:
                self.elapsed_time = 0.0
                for child in self.children.values():
                    self.elapsed_time += child.elapsed_time
            if self.ncalls < 0:
                self.ncalls = 0
                for child in self.children.values():
                    self.ncalls += child.ncalls

    @classmethod
    def build(cls, timers: StopwatchDict, separetor: str = "/"):
        root = cls()
        for name, timer in timers.items():
            prefix = name.split(separetor)
            node = root
            for path in prefix:
                if path not in node.children:
                    node.children[path] = ProfileTree()
                node = node.children[path]
            node.elapsed_time = timer.elpased_time
            node.ncalls = timer.ncalls
        root.aggregate()
        return root

    def result(self, nsentences: int = -1) -> list[dict[str, str | int | float]]:
        def _result(name: str, node: ProfileTree) -> list[dict[str, str | int | float]]:
            stat = {
                "name": name.strip("/"),
                "total time (sec)": node.elapsed_time,
                "citation": node.ncalls,
                "msec/citation": node.elapsed_time * 1000 / node.ncalls,
            }
            if nsentences > 0:
                stat["msec/paper"] = node.elapsed_time * 1000 / nsentences
                stat["citations/paper"] = node.ncalls / nsentences
            res = [stat]
            for path, child in node.children.items():
                res += _result(name + "/" + path, child)
            return res

        res = _result("", self)
        return res[1:]  # Remove the root node


def aggregate() -> ProfileTree:
    """Aggregate the timers.

    Returns:
        ProfileTree: The root of the profile tree.
    """
    return ProfileTree.build(measure)
