import numpy as np
import pandas as pd
from carbontracker.tracker import CarbonTracker
from tensorflow.keras.callbacks import Callback
from utime.utils import get_memory_usage
from mpunet.utils import highlighted
from mpunet.logging import ScreenLogger
from collections import defaultdict
from datetime import timedelta


class Validation(Callback):
    """
    Validation computation callback.

    Samples a number of validation batches from a deepsleep
    ValidationMultiSequence object
    and computes for all tasks:
        - Batch-wise validation loss
        - Batch-wise metrics as specified in model.metrics_tensors
        - Epoch-wise pr-class and average precision
        - Epoch-wise pr-class and average recall
        - Epoch-wise pr-class and average dice coefficients
    ... and adds all results to the log dict

    Note: The purpose of this callback over the default tf.keras evaluation
    mechanism is to calculate certain metrics over the entire epoch of data as
    opposed to averaged batch-wise computations.
    """
    def __init__(self,
                 val_sequence,
                 max_val_studies_per_dataset=20,
                 logger=None, verbose=True):
        """
        Args:
            val_sequence: A deepsleep ValidationMultiSequence object
            logger:       An instance of a MultiPlanar Logger that prints to
                          screen and/or file
            verbose:      Print progress to screen - OBS does not use Logger
        """
        super().__init__()
        self.logger = logger or ScreenLogger()
        self.sequences = val_sequence.sequences
        self.verbose = verbose
        self.max_studies = max_val_studies_per_dataset
        self.n_classes = val_sequence.n_classes
        self.IDs = val_sequence.IDs
        self.print_round = 3
        self.log_round = 4

    def _compute_counts(self, pred, true, ignore_class=None):
        # Argmax and CM elements
        pred = pred.argmax(-1).ravel()
        true = true.ravel()

        if ignore_class:
            mask = np.where(true != ignore_class)
            true = true[mask]
            pred = pred[mask]

        # Compute relevant CM elements
        # We select the number following the largest class integer when
        # y != pred, then bincount and remove the added dummy class
        tps = np.bincount(np.where(true == pred, true, self.n_classes),
                          minlength=self.n_classes+1)[:-1].astype(np.uint64)
        rel = np.bincount(true, minlength=self.n_classes).astype(np.uint64)
        sel = np.bincount(pred, minlength=self.n_classes).astype(np.uint64)
        return tps, rel, sel

    def predict(self):
        # Get tensors to run and their names
        metrics = self.model.loss_functions + self.model.metrics
        metrics_names = self.model.metrics_names
        self.model.reset_metrics()
        assert len(metrics_names) == len(metrics)

        # Prepare arrays for CM summary stats
        true_pos, relevant, selected, metrics_results = {}, {}, {}, {}
        for id_, sequence in zip(self.IDs, self.sequences):
            # Add count arrays to the result dictionaries
            true_pos[id_] = np.zeros(shape=(self.n_classes,), dtype=np.uint64)
            relevant[id_] = np.zeros(shape=(self.n_classes,), dtype=np.uint64)
            selected[id_] = np.zeros(shape=(self.n_classes,), dtype=np.uint64)

            # Get validation sleep study loader
            n_val = min(len(sequence.dataset_queue), self.max_studies)
            study_iterator = sequence.dataset_queue.get_study_iterator(n_val)

            # Predict and evaluate on all studies
            per_study_metrics = defaultdict(list)
            for i, sleep_study_context in enumerate(study_iterator):
                if self.verbose:
                    s = "   {}Validation subject: {}/{}".format(f"[{id_}] "
                                                                if id_ else "",
                                                                i+1,
                                                                n_val)
                    print(s, end="\r", flush=True)

                with sleep_study_context as ss:
                    x, y = sequence.get_single_study_full_seq(ss.identifier,
                                                              reshape=True)
                    pred = self.model.predict_on_batch(x)

                # Compute counts
                tps, rel, sel = self._compute_counts(pred=pred.numpy(),
                                                     true=y,
                                                     ignore_class=5)
                true_pos[id_] += tps
                relevant[id_] += rel
                selected[id_] += sel

                # Run all metrics
                for metric, name in zip(metrics, metrics_names):
                    per_study_metrics[name].append(metric(y, pred).numpy())

            # Compute mean metrics for the dataset
            metrics_results[id_] = {}
            for metric, name in zip(metrics, metrics_names):
                metrics_results[id_][name] = np.mean(per_study_metrics[name])
            self.model.reset_metrics()
            self.logger("")
        self.logger("")
        return true_pos, relevant, selected, metrics_results

    @staticmethod
    def _compute_dice(tp, rel, sel):
        # Get data masks (to avoid div. by zero warnings)
        # We set precision, recall, dice to 0 in for those particular cls.
        sel_mask = sel > 0
        rel_mask = rel > 0

        # prepare arrays
        precisions = np.zeros(shape=tp.shape, dtype=np.float32)
        recalls = np.zeros_like(precisions)
        dices = np.zeros_like(precisions)

        # Compute precisions, recalls
        precisions[sel_mask] = tp[sel_mask] / sel[sel_mask]
        recalls[rel_mask] = tp[rel_mask] / rel[rel_mask]

        # Compute dice
        intrs = (2 * precisions * recalls)
        union = (precisions + recalls)
        dice_mask = union > 0
        dices[dice_mask] = intrs[dice_mask] / union[dice_mask]

        return precisions, recalls, dices

    def _print_val_results(self, precisions, recalls, dices, metrics, epoch,
                           name, classes):
        # Log the results
        # We add them to a pd dataframe just for the pretty print output
        index = ["cls %i" % i for i in classes]
        metric_keys, metric_vals = map(list, list(zip(*metrics.items())))
        col_order = metric_keys + ["precision", "recall", "dice"]
        nan_arr = np.empty(shape=len(precisions))
        nan_arr[:] = np.nan
        value_dict = {"precision": precisions,
                      "recall": recalls,
                      "dice": dices}
        value_dict.update({key: nan_arr for key in metrics})
        val_results = pd.DataFrame(value_dict,
                                   index=index).loc[:, col_order]  # ensure order
        # Transpose the results to have metrics in rows
        val_results = val_results.T
        # Add mean and set in first row
        means = metric_vals + [precisions.mean(), recalls.mean(), dices.mean()]
        val_results["mean"] = means
        cols = list(val_results.columns)
        cols.insert(0, cols.pop(cols.index('mean')))
        val_results = val_results.loc[:, cols]

        # Print the df to screen
        self.logger(highlighted(("[%s] Validation Results for "
                            "Epoch %i" % (name, epoch)).lstrip(" ")))
        print_string = val_results.round(self.print_round).to_string()
        self.logger(print_string.replace("NaN", "---") + "\n")

    def on_epoch_end(self, epoch, logs={}):
        self.logger("\n")
        # Predict and get CM
        TPs, relevant, selected, metrics = self.predict()
        for id_ in self.IDs:
            tp, rel, sel = TPs[id_], relevant[id_], selected[id_]
            precisions, recalls, dices = self._compute_dice(tp=tp, sel=sel, rel=rel)
            classes = np.arange(len(dices))

            # Add to log
            n = (id_ + "_") if len(self.IDs) > 1 else ""
            logs[f"{n}val_dice"] = dices.mean().round(self.log_round)
            logs[f"{n}val_precision"] = precisions.mean().round(self.log_round)
            logs[f"{n}val_recall"] = recalls.mean().round(self.log_round)
            for m_name, value in metrics[id_].items():
                logs[f"{n}val_{m_name}"] = value.round(self.log_round)

            if self.verbose:
                self._print_val_results(precisions=precisions,
                                        recalls=recalls,
                                        dices=dices,
                                        metrics=metrics[id_],
                                        epoch=epoch,
                                        name=id_,
                                        classes=classes)

        if len(self.IDs) > 1:
            # Print cross-dataset mean values
            if self.verbose:
                self.logger(highlighted(f"[ALL DATASETS] Means Across Classes"
                                        f" for Epoch {epoch}"))
            fetch = ("val_dice", "val_precision", "val_recall")
            m_fetch = tuple(["val_" + s for s in self.model.metrics_names])
            to_print = {}
            for f in fetch + m_fetch:
                scores = [logs["%s_%s" % (name, f)] for name in self.IDs]
                res = np.mean(scores)
                logs[f] = res.round(self.log_round)  # Add to log file
                to_print[f.split("_")[-1]] = list(scores) + [res]
            if self.verbose:
                df = pd.DataFrame(to_print)
                df.index = self.IDs + ["mean"]
                self.logger(df.round(self.print_round))
            self.logger("")


class MemoryConsumption(Callback):
    def __init__(self, max_gib=None, round_=2, logger=None, set_limit=False):
        self.max_gib = max_gib
        self.logger = logger
        self.round_ = round_
        if set_limit:
            import resource
            _, hard = resource.getrlimit(resource.RLIMIT_AS)
            resource.setrlimit(resource.RLIMIT_AS,
                               (self._gib_to_bytes(max_gib), hard))
            self.logger("Setting memory limit to {} GiB".format(max_gib))

    @staticmethod
    def _gib_to_bytes(gib):
        return gib * (1024 ** 3)

    @staticmethod
    def _bytes_to_gib(bytes):
        return bytes / (1024 ** 3)

    def on_epoch_end(self, epoch, logs={}):
        mem_bytes = get_memory_usage()
        mem_gib = round(self._bytes_to_gib(mem_bytes), self.round_)
        logs['memory_usage_gib'] = mem_gib
        if self.max_gib and mem_gib >= self.max_gib:
            self.logger.warn("Stopping training from callback 'MemoryConsumption'! "
                             "Total memory consumption of {} GiB exceeds limitation"
                             " (self.max_gib = {}) ".format(mem_gib, self.max_gib))
            self.model.stop_training = True


class MaxTrainingTime(Callback):
    def __init__(self, max_minutes, log_name='train_time_total', logger=None):
        """
        TODO
        Args:
        """
        super().__init__()
        self.max_minutes = int(max_minutes)
        self.log_name = log_name
        self.logger = logger or ScreenLogger()

    def on_epoch_end(self, epochs, logs={}):
        """
        TODO

        Args:
            epochs:
            logs:

        Returns:

        """
        train_time_str = logs.get(self.log_name, None)
        if not train_time_str:
            self.logger.warn("Did not find log entry '{}' (needed in callback "
                             "'MaxTrainingTime')".format(self.log_name))
            return
        train_time_m = timedelta(
            days=int(train_time_str[:2]),
            hours=int(train_time_str[4:6]),
            minutes=int(train_time_str[8:10]),
            seconds=int(train_time_str[12:14])
        ).total_seconds() / 60
        if train_time_m >= self.max_minutes:
            # Stop training
            self.warn("Stopping training from callback 'MaxTrainingTime'! "
                      "Total training length of {} minutes exceeded (now {}) "
                      "".format(self.max_minutes, train_time_m))
            self.model.stop_training = True


class CarbonUsageTracking(Callback):
    """
    tf.keras Callback for the Carbontracker package.
    See https://github.com/lfwa/carbontracker.
    """
    def __init__(self, epochs, add_to_logs=True, monitor_epochs=-1,
                 epochs_before_pred=-1, devices_by_pid=True, **additional_tracker_kwargs):
        """
        Accepts parameters as per CarbonTracker.__init__
        Sets other default values for key parameters.

        Args:
            add_to_logs: bool, Add total_energy_kwh and total_co2_g to the keras logs after each epoch
            For other arguments, please refer to CarbonTracker.__init__
        """
        super().__init__()
        self.tracker = None
        self.add_to_logs = bool(add_to_logs)
        self.parameters = {"epochs": epochs,
                           "monitor_epochs": monitor_epochs,
                           "epochs_before_pred": epochs_before_pred,
                           "devices_by_pid": devices_by_pid}
        self.parameters.update(additional_tracker_kwargs)

    def on_train_end(self, logs={}):
        """ Ensure actual consumption is reported """
        self.tracker.stop()

    def on_epoch_begin(self, epoch, logs={}):
        """ Start tracking this epoch """
        if self.tracker is None:
            # At this point all CPUs should be discoverable
            self.tracker = CarbonTracker(**self.parameters)
        self.tracker.epoch_start()

    def on_epoch_end(self, epoch, logs={}):
        """ End tracking this epoch """
        self.tracker.epoch_end()
        if self.add_to_logs:
            energy_kwh = self.tracker.tracker.total_energy_per_epoch().sum()
            co2eq_g = self.tracker._co2eq(energy_kwh)
            logs["total_energy_kwh"] = round(energy_kwh, 6)
            logs["total_co2_g"] = round(co2eq_g, 6)

pandas.DataFrame