"""
    
    Mini profiling tool for your projects
    https://github.com/kadir014/miniprofiler

"""

from typing import Callable, Any

from dataclasses import dataclass
from contextlib import contextmanager
from time import perf_counter
from statistics import stdev


__version__ = "0.0.2"


def percentile(samples: list[float], percent: float) -> float:
    """
    Calculate the Nth percentile given timing samples.
    
    Parameters
    ----------
    samples
        Array of timing calculations
    percent
        Percentile in range [0, 100]
    """
    
    n = len(samples)
    if n == 0:
        return 0.0

    samples = sorted(samples)
    k = (n - 1.0) * (percent / 100.0)
    f = int(k)
    c = int(min(f + 1, n - 1))

    if f == c:
        return samples[int(k)]

    return samples[f] + (samples[c] - samples[f]) * (k - f)


@dataclass
class MetricSummary:
    """
    Profiling summary over a period of samples for a single metric.

    Attributes
    ----------
    avg
        Average of timings in seconds.
    min
        Minimum of timings in seconds.
    max
        Maximum of timings in seconds.
    p95
        95th percentile timing in seconds.
    p99
        99th percentile timing in seconds.
    p01
        1st percentile timing in seconds.
    p05
        5th percentile timing in seconds.
    stdev
        Standard deviation of timings in seconds.
    last
        Latest timing sample recorded in seconds.
    samples
        Number of timing samples aggregated to calculate summary statistics.
    total
        Sum of all timing samples aggregated in seconds.
    name
        Name of the metric.
    """

    avg: float
    min: float
    max: float
    p95: float
    p99: float
    p01: float
    p05: float
    stdev: float
    last: float
    samples: int
    total: float
    name: str


class Profiler:
    """
    A lightweight execution-time profiler.

    Records repeated timing samples across multiple runs and computes summary
    statistics such as average time, percentiles and standard deviation.

    Sample usage
    ```
    profiler = Profiler()

    while app_running:
        with profiler.profile("some_metric"):
            # profiled, heavy code here

    print("Summary:")
    print(profiler["some_metric"])
    ```
    """
    
    def __init__(self, max_samples: int | None = None) -> None:
        """
        Parameters
        ----------
        max_samples
            Maximum number of recent timing samples to retain per metric.
            If None, samples accumulate indefinitely.
        """
        self.__timings = dict()

        self.max_samples = max_samples

    def __getitem__(self, metric: str) -> MetricSummary:
        """ Return a profiled metric summary. """

        return MetricSummary(
            avg = self.__timings[metric]["avg"],
            min = self.__timings[metric]["min"],
            max = self.__timings[metric]["max"],
            p95 = self.__timings[metric]["p95"],
            p99 = self.__timings[metric]["p99"],
            p01 = self.__timings[metric]["p01"],
            p05 = self.__timings[metric]["p05"],
            stdev = self.__timings[metric]["stdev"],
            last = self.__timings[metric]["last"],
            samples = len(self.__timings[metric]["acc"]),
            total = self.__timings[metric]["total"],
            name = metric
        )

    @contextmanager
    def profile(self, metric: str):
        """
        Profile piece of code.
        
        Parameters
        ----------
        metric
            Metric name to profile
        """

        if metric not in self.__timings:
            self.__timings[metric] = {
                "acc": [],
                "avg": 0.0,
                "min": 0.0,
                "max": 0.0,
                "p95": 0.0,
                "p99": 0.0,
                "last": 0.0,
                "stdev": 0.0
            }

        start = perf_counter()
        
        try:
            yield None

        finally:
            elapsed = perf_counter() - start
            self._accumulate(metric, elapsed)

    def _accumulate(self, metric: str, value: float) -> None:
        """
        Accumulate metric timing.
        
        Parameters
        ----------
        metric
            Metric name to accumulate.
        value
            Timing value to accumulate.
        """

        acc = self.__timings[metric]["acc"]
        acc.append(value)

        if self.max_samples is not None:
            # The reason for the while loop is here that, if max_samples is
            # changed during profiling, there is a chance that the accumulated
            # values will never catch up with a single if check.
            while len(acc) > self.max_samples:
                acc.pop(0)

        s = sum(acc)
        n = len(acc)

        if n < 2:
            s_stdev = 0.0
        else:
            s_stdev = stdev(acc)

        self.__timings[metric]["total"] = s
        self.__timings[metric]["avg"] = s / n
        self.__timings[metric]["min"] = min(acc)
        self.__timings[metric]["max"] = max(acc)
        self.__timings[metric]["p95"] = percentile(acc, 95.0)
        self.__timings[metric]["p99"] = percentile(acc, 99.0)
        self.__timings[metric]["p01"] = percentile(acc, 1.0)
        self.__timings[metric]["p05"] = percentile(acc, 5.0)
        self.__timings[metric]["last"] = value
        self.__timings[metric]["stdev"] = s_stdev


def bench(
        *funcs: Callable[..., Any],
        runs: int = 10,
        samples: int = 1000,
        unit: str = "ms",
        digits: int = 2,
        table_name_gap: int = 16,
        table_col_gap: int = 10
        ) -> None:
    """
    Small benchmarking & comparison utility for a number of functions.

    Parameters
    ----------
    funcs
        Functions to benchmark and compare
    runs
        Number of total complete runs
    samples
        Number of iterations for each run
    unit
        Time unit used to display benchmarked values (e.g "ms", "s", "m", ...)
    digits
        Number of digits to round the values to when printing
    table_name_gap
        Number of characters allowed for function names in printed table
    table_col_gap
        Number of characters allowed for each value column in printed table
    """

    profiler = Profiler()

    unit_s = {
        "h": 1.0 / 3600.0,
        "m": 1.0 / 60.0,
        "s": 1.0,
        "cs": 100.0,
        "ms": 1000.0,
        "us": 1000000.0,
        "ns": 1000000000.0
    }[unit]
    def _avg(metric: str) -> str: return str(round(profiler[metric].avg * unit_s, digits))
    def _p01(metric: str) -> str: return str(round(profiler[metric].p01 * unit_s, digits))
    def _p99(metric: str) -> str: return str(round(profiler[metric].p95 * unit_s, digits))
    def _stdev(metric: str) -> str: return str(round(profiler[metric].stdev * unit_s, digits))

    with profiler.profile("please_dont_overwrite_benchmark"):
        for func in funcs:
            name = func.__name__
            print(f"Benching '{name}'...")

            for run_i in range(runs):
                with profiler.profile(name):
                    for _ in range(samples):
                        func()

                print(f"Run #{run_i}: {profiler[name].last}s")
            
            print()

    final = profiler["please_dont_overwrite_benchmark"]
    print(f"The benchmark is completed in {round(final.last, 3)}s ({runs * samples * len(funcs)} calls)")
    print(f"Displaying results for {runs} runs and {samples} samples per 1 function")
    print()

    # Print timings table
    c = table_col_gap
    print(f"{unit:<{table_name_gap}}{'avg':<{c}}{'1st %':<{c}}{'99th %':<{c}}{'stdev':<{c}}")
    print("─" * table_name_gap + "─" * (c * 4))
    for func in funcs:
        name = func.__name__

        display = name
        if len(display) >= table_name_gap:
            display = display[:table_name_gap-4] + "..."

        print(f"{display:<{table_name_gap}}{_avg(name):<{c}}{_p01(name):<{c}}{_p99(name):<{c}}{_stdev(name):<{c}}")
    print()

    # Compare functions
    func_names = [func.__name__ for func in funcs]
    sorted_avg = sorted(func_names, key=lambda name: profiler[name].avg)
    sorted_stdev = sorted(func_names, key=lambda name: profiler[name].stdev)

    print("Fastest to slowest by average")
    print("─────────────────────────────")
    print(" < ".join(sorted_avg))
    print()
    print("Lowest to highest standard deviation")
    print("────────────────────────────────────")
    print(" < ".join(sorted_stdev))


__all__ = (
    "__version__",
    "percentile",
    "MetricSummary",
    "Profiler",
    "bench"
)