from datetime import datetime
from threading import RLock
from typing import Any, Dict, List, Optional

import pandas as pd
from triad.utils.convert import to_timedelta
from tune import Monitor, TrialReport


class PrintBest(Monitor):
    def __init__(self):
        super().__init__()
        self._lock = RLock()
        self._bins: Dict[str, "_ReportBin"] = {}

    def on_report(self, report: TrialReport) -> None:
        with self._lock:
            key = str(report.trial.keys)
            if key not in self._bins:
                self._bins[key] = _ReportBin(new_best_only=True)
            rbin = self._bins[key]
        if rbin.on_report(report):
            print(report.trial.keys, report.metric, report.jsondict)


class NotebookSimpleChart(Monitor):
    def __init__(self, interval: Any = "1sec"):
        super().__init__()
        self._lock = RLock()
        self._last: Any = None
        self._bins: Dict[str, "_ReportBin"] = {}
        self._interval = to_timedelta(interval)

    def on_report(self, report: TrialReport) -> None:
        now = datetime.now()
        with self._lock:
            key = str(report.trial.keys)
            if key not in self._bins:
                self._bins[key] = _ReportBin(new_best_only=True)
            rbin = self._bins[key]

        rbin.on_report(report)

        with self._lock:
            if self._last is None or now - self._last > self._interval:
                import matplotlib.pyplot as plt
                from IPython.display import clear_output

                clear_output()
                df = pd.concat(
                    [
                        pd.DataFrame(
                            x.records,
                            columns=[
                                "partition",
                                "rung",
                                "time",
                                "id",
                                "metric",
                                "best_metric",
                            ],
                        )
                        for x in self._bins.values()
                    ]
                )
                self.plot(df)
                plt.show()
                self._last = datetime.now()
                for best in [x.best for x in self._bins.values() if x.best is not None]:
                    if best is not None:
                        print(best.trial.keys, best.metric, best.jsondict)

    def plot(self, df: pd.DataFrame) -> None:
        return  # pragma: no cover


class NotebookSimpleRungs(NotebookSimpleChart):
    def __init__(self, interval: Any = "1sec"):
        super().__init__(interval)

    def plot(self, df: pd.DataFrame) -> None:
        import seaborn as sns

        sns.lineplot(data=df, x="rung", y="metric", hue="id", marker="o", legend=False)


class NotebookSimpleTimeSeries(NotebookSimpleChart):
    def __init__(self, interval: Any = "1sec"):
        super().__init__(interval)

    def plot(self, df: pd.DataFrame) -> None:
        import seaborn as sns

        sns.lineplot(data=df, x="time", y="best_metric", hue="partition", marker="o")


class NotebookSimpleHist(NotebookSimpleChart):
    def __init__(self, interval: Any = "1sec"):
        super().__init__(interval)

    def plot(self, df: pd.DataFrame) -> None:
        import seaborn as sns

        sns.histplot(data=df, x="metric", hue="partition")


class _ReportBin:
    def __init__(self, new_best_only: bool = False):
        super().__init__()
        self._values: List[List[Any]] = []
        self._lock = RLock()
        self._best_report: Optional[TrialReport] = None
        self._new_best_only = new_best_only

    def on_report(self, report: TrialReport) -> bool:
        with self._lock:
            now = datetime.now()
            updated = False
            if (
                self._best_report is None
                or report.sort_metric < self._best_report.sort_metric
            ):
                self._best_report = report
                updated = True
            if updated or not self._new_best_only:
                self._values.append(
                    [
                        str(report.trial.keys),
                        report.rung,
                        now,
                        report.trial_id,
                        report.metric,
                        self._best_report.metric,
                    ]
                )
            return updated

    @property
    def best(self) -> Optional[TrialReport]:
        with self._lock:
            return self._best_report

    @property
    def records(self) -> List[List[Any]]:
        with self._lock:
            return list(self._values)
