import json
import time
from math import log10, sqrt
from typing import Callable, Dict, List

from .report import print_table
from .statstream import StatStream


def chrono(func: Callable):
    def chrono_decorator(*args, **kwargs):
        start = time.time()
        value = func(*args, **kwargs)
        t = time.time() - start
        print(f"{func.__name__:>30} ran in {t:10.4f} s")
        return value

    return chrono_decorator


class _DummyContext:
    def __init__(self, **args):
        pass

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass


class ChronoContext:
    """
    sync is a function that can be set to make the timer wait before ending.
    This is useful when timing async calls like cuda calls
    """

    def __init__(
        self,
        name,
        stream: StatStream,
        sync: Callable,
        parent,
        verbose=False,
        endline="\n",
    ):
        self.name = name
        self.stream = stream
        self.start = 0
        self.sync = sync
        self.parent = parent
        self.verbose = verbose
        self.newline = endline

    def __enter__(self):
        # Sync before starting timer to make sure previous work is not timed as well
        self.depth = self.parent.depth
        self.parent.depth += 1
        self.sync()

        # if self.verbose:
        #    print(f'{" " * self.depth * 2} [{self.depth:3d}] >  {self.name}', end='')

        self.start = time.time()
        return self.stream

    def __exit__(self, exception_type, exc_val, traceback):
        # Sync before ending timer to make sure all the work is accounted for
        self.sync()
        self.end = time.time()

        self.parent.depth -= 1
        if exception_type is None:
            self.stream.update(self.end - self.start)

        if self.verbose:
            print(
                f'{self.newline}{" " * self.depth * 2} [{self.depth:3d}] <  {self.name:>30}: (obs: {self.stream.val:8.4f} s, '
                f"avg: {self.stream.avg:8.4f}, "
                f"cnt: {self.stream.count})",
                end="",
            )

    @property
    def count(self):
        return self.stream.current_count


class MultiStageChrono:
    def __init__(self, skip_obs=10, sync=None, disabled=False, name=None):
        self.chronos = {}
        self.skip_obs = skip_obs
        self.sync = sync
        self.name = name
        self.disabled = disabled
        self.depth = 0
        if sync is None:
            self.sync = lambda: None

    def time(self, name, skip_obs=None, **kwargs):
        if self.disabled:
            return _DummyContext()

        # if self.name is not None:
        #    name = '{}.{}'.format(self.name, name)

        val = self.chronos.get(name)

        if val is None:
            val = StatStream(self.skip_obs)
            if skip_obs is not None:
                val = StatStream(skip_obs)
            self.chronos[name] = val

        # inherit sync from parent
        if kwargs.get("sync") is None:
            kwargs["sync"] = self.sync

        return ChronoContext(name, val, parent=self, **kwargs)

    def make_table(self, common: List = None, transform=None):
        common = common or []
        table = []

        for i, (name, stream) in enumerate(self.chronos.items()):
            table.append([name] + stream.to_array(transform) + common)

        return table

    def report(self, *args, format="csv", **kwargs):
        if format == "csv":
            return self.report_csv(*args, **kwargs)
        print(self.to_json(*args, **kwargs))

    def report_csv(
        self,
        speed=False,
        size=1,
        file_name=None,
        common: Dict[str, str] = None,
        skip_header=False,
    ):
        if self.disabled:
            return

        common = common or {}

        # split map in two
        items = list(common.items())

        common_header = list(map(lambda item: item[0], items))
        common = list(map(lambda item: item[1], items))

        header = ["Stage", "Average", "Deviation", "Min", "Max", "count"]
        header.extend(common_header)

        table = (
            self.make_table(common, lambda x: size / x)
            if speed
            else self.make_table(common)
        )
        print_table(header, table, file_name, skip_header)

    def to_dict(self, base=None):
        items = base
        if items is None:
            items = {}

        if self.name is not None:
            items["name"] = self.name

        for key, stream in self.chronos.items():
            items[key] = stream.to_dict()

        return items

    def to_json(self, base=None, *args, **kwargs):
        if "indent" not in kwargs:
            kwargs["indent"] = "  "
        return json.dumps(self.to_dict(base), *args, **kwargs)


def time_this(chrono, *cargs, **ckwargs):
    def toplevel_decorator(fun):
        def wrapper(*args, **kwargs):
            with chrono.time(fun.__name__, *cargs, **ckwargs):
                return fun(*args, **kwargs)

        return wrapper

    return toplevel_decorator


def estimated_time_to_arrival(i: int, n: int, timer: StatStream):
    # return ETA and +/- offset
    avg = timer.avg
    if avg == 0:
        avg = timer.val

    eta = (n - i - 1) * avg
    return eta, 2.95 * timer.sd * sqrt(n - i - 1)


def get_div_fmt(val: float):
    div, fmt = 60, "min"
    if val < div:
        div = 1
        fmt = "s"
    return div, fmt


def show_eta(i: int, n: int, timer: StatStream, end="\n"):
    eta, offset = estimated_time_to_arrival(i, n, timer)
    size = int(log10(n) + 1)

    div, fmt = get_div_fmt(eta)

    eta = f"{eta / div:6.2f} {fmt}"

    div, fmt = get_div_fmt(offset)
    conf = f"{offset / sqrt(div):6.2f} {fmt}"

    total = ""
    if timer.avg != 0:
        t = timer.avg * (timer.count + timer.drop_obs)
        div, fmt = get_div_fmt(t)
        total = f" | Total {t / div:6.2f} {fmt}"

    print(f"[{i:{size}d}/{n:{size}d}] Remaining {eta} +/- {conf}{total}", end=end)


if __name__ == "__main__":

    timer = MultiStageChrono(0, disabled=False)

    with timer.time("all", verbose=True):
        for i in range(0, 10):

            with timer.time("forward_back", verbose=True) as timer:
                with timer.time("forward", verbose=True):
                    time.sleep(1)

                    if i % 2 == 0:
                        time.sleep(0.25)

                with timer.time("backward", verbose=True, skip_obs=3):
                    time.sleep(1)

                    if i % 2 == 0:
                        time.sleep(0.25)

        show_eta(i, 10, timer)

    print()
    timer.report()
    print(chrono.to_json(base={"main": 1}, indent="   "))

    print()
    timer.report(format="json")

    @time_this(timer, verbose=True)
    def test():
        time.sleep(1)

    test()
