from collections import defaultdict
from contextlib import contextmanager
from dataclasses import dataclass, field
from math import sqrt
import timeit


@dataclass
class _Frame:
    """Represents one entry on the active timing stack."""

    exclude_inner: bool
    exclude_time: float = field(default=0.0)


class ExecTimeProfiler:
    """Profiles execution time among different stages in a single thread.

    Warnings:
        This class is NOT thread-safe.
        Also, do not use it to profile executions in multiple threads.
    """

    __slots__ = ("_stack", "_stages")

    def __init__(self):
        self._stack: list[_Frame] = []
        self._stages: defaultdict[str, list[float]] = defaultdict(list)

    @property
    def _top(self) -> _Frame | None:
        """Returns the current top-of-stack frame, or None if empty."""
        return self._stack[-1] if self._stack else None

    @contextmanager
    def timeit(self, stage_name: str, exclude_inner: bool = False):
        """Times a section of code.

        Args:
            stage_name: Stage name, should be unique for each distinctive stage.
            exclude_inner: Whether to exclude counting for inner stages.
        """
        start = timeit.default_timer()
        self._stack.append(_Frame(exclude_inner))
        try:
            yield
        finally:
            elapsed = timeit.default_timer() - start
            frame = self._stack.pop()
            self._stages[stage_name].append(elapsed - frame.exclude_time)
            if (parent := self._top) is not None and parent.exclude_inner:
                parent.exclude_time += elapsed

    @property
    def stages(self) -> list[tuple[str, int, float, float]]:
        """Returns for each profiled stage: name, size, sum, std-dev."""
        ret = []
        for name, samples in self._stages.items():
            size = len(samples)
            total = sum(samples)
            mean = total / size
            stdev = 0.0 if size == 1 else sqrt(sum((x - mean) ** 2 for x in samples) / (size - 1))
            ret.append((name, size, total, stdev))
        return ret

    def to_string(self, unit: str = "sec") -> str:
        """Converts the stages to string."""
        assert unit in ("sec", "ms")
        scale = 1000 if unit == "ms" else 1
        parts = [
            f"{name}:{total * scale:.2f}/{size}/{stdev:.2f}"
            for name, size, total, stdev in self.stages
        ]
        return f"[{', '.join(parts)}]"

    def get_stage_times(self, name: str, unit: str = "sec") -> list[float]:
        """Gets all the elapsed times for a stage."""
        assert unit in ("sec", "ms")
        scale = 1000 if unit == "ms" else 1
        return [x * scale for x in self._stages.get(name, [])]

    def clear(self):
        """Clears all profiled data."""
        self._stages.clear()


class NopProfiler(ExecTimeProfiler):
    """No-op profiler used as a default argument value."""

    @contextmanager
    def timeit(self, stage_name: str, exclude_inner: bool = False):
        yield


# Set as default argument value so that a function's structure is kept the same.
nop_profiler = NopProfiler()
