# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/017_Benchmarking.ipynb.

# %% auto 0
__all__ = ['logger']

# %% ../../nbs/017_Benchmarking.ipynb 1
from collections import deque
from datetime import datetime, timedelta
from functools import wraps
from statistics import mean, stdev
from typing import *

from .logger import get_logger

# %% ../../nbs/017_Benchmarking.ipynb 4
logger = get_logger("fastkafka.benchmark")

# %% ../../nbs/017_Benchmarking.ipynb 5
def _benchmark(
    interval: Union[int, timedelta] = 1,
    *,
    sliding_window_size: Optional[int] = None,
    func_name: str,
    benchmark_results: Dict[str, Dict[str, Any]],
) -> None:
    """Used to record the benchmark results(throughput, average throughput, standard deviation) of a given function

    Args:
        interval: the time interval after which the benchmark results are logged.
        sliding_window_size: the maximum number of benchmark results to use to calculate average throughput and standard deviation.
        func_name: the name of the function to be benchmarked.
        benchmark_results: a dictionary containing the benchmark results of all functions.
    """
    if isinstance(interval, int):
        interval = timedelta(seconds=interval)
    if func_name not in benchmark_results:
        benchmark_results[func_name] = {
            "count": 0,
            "last_count": 0,
            "start": None,
            "last_start": None,
            "history": [],
        }
        if sliding_window_size is not None:
            benchmark_results[func_name]["history"] = deque(maxlen=sliding_window_size)

    benchmark_results[func_name]["count"] += 1

    if benchmark_results[func_name]["count"] == 1:
        benchmark_results[func_name]["start"] = benchmark_results[func_name][
            "last_start"
        ] = datetime.utcnow()

    diff = datetime.utcnow() - benchmark_results[func_name]["last_start"]
    if diff >= interval:
        throughput = (
            benchmark_results[func_name]["count"]
            - benchmark_results[func_name]["last_count"]
        ) / (diff / timedelta(seconds=1))
        log_msg = f"Throughput = {throughput:5,.0f}"

        if sliding_window_size is not None:
            benchmark_results[func_name]["history"].append(throughput)

            log_msg += f", Avg throughput = {mean(benchmark_results[func_name]['history']):5,.0f}"
        #             if len(benchmark_results[func_name]["history"]) > 1:
        #                 log_msg += f", Standard deviation of throughput is {stdev(benchmark_results[func_name]['history']):5,.0f}"
        log_msg = (
            log_msg
            + f" - For {func_name}(interval={interval.seconds},{sliding_window_size=})"
        )
        logger.info(log_msg)

        benchmark_results[func_name]["last_start"] = datetime.utcnow()
        benchmark_results[func_name]["last_count"] = benchmark_results[func_name][
            "count"
        ]
