"""A `dowel.logger.LogOutput` for tensorboard.

It receives the input data stream from `dowel.logger`, then add them to
tensorboard summary operations through tensorboardX.

Note:
Neither TensorboardX nor TensorBoard supports log parametric
distributions. We add this feature by sampling data from a
`tfp.distributions.Distribution` object.
"""
import functools
import warnings
import os
import wandb

import matplotlib.pyplot as plt
import numpy as np
import scipy.stats
import tensorboardX as tbX
try:
    import tensorflow as tf
except ImportError:
    tf = None

from dowel import Histogram
from dowel import LoggerWarning
from dowel import LogOutput
from dowel import TabularInput
from dowel.utils import colorize


class TensorBoardOutput(LogOutput):
    """TensorBoard output for logger.

    Args:
        log_dir(str): The save location of the tensorboard event files.
        x_axis(str): The name of data used as x-axis for scalar tabular.
            If None, x-axis will be the number of dump() is called.
        additional_x_axes(list[str]): Names of data to used be as additional
            x-axes.
        flush_secs(int): How often, in seconds, to flush the added summaries
            and events to disk.
        histogram_samples(int): Number of samples to generate when logging
            random distribution.
    """

    def __init__(self,
                 log_dir,
                 x_axis=None,
                 additional_x_axes=None,
                 flush_secs=120,
                 histogram_samples=1e3):
        # Also log to wandb

        if x_axis is None:
            assert not additional_x_axes, (
                'You have to specify an x_axis if you want additional axes.')

        additional_x_axes = additional_x_axes or []

        self._writer = tbX.SummaryWriter(log_dir, flush_secs=flush_secs)
        self._x_axis = x_axis
        self._additional_x_axes = additional_x_axes
        self._default_step = 0
        self._histogram_samples = int(histogram_samples)
        self._added_graph = False
        self._waiting_for_dump = []
        # Used in tests to emulate Tensorflow not being installed.
        self._tf = tf

        self.wandb = False
        
        # if 'WANDB_API_KEY' in os.environ:
        #     self.wandb = True
        # else:
        #     self.wandb = False

        self._warned_once = set()
        self._disable_warnings = False

    @property
    def types_accepted(self):
        """Return the types that the logger may pass to this output."""
        if self._tf is None:
            return (TabularInput, )
        else:
            return (TabularInput, self._tf.Graph)

    def record(self, data, prefix=''):
        """Add data to tensorboard summary.

        Args:
            data: The data to be logged by the output.
            prefix(str): A prefix placed before a log entry in text outputs.

        """
        if isinstance(data, TabularInput):
            self._waiting_for_dump.append(
                functools.partial(self._record_tabular, data))
        elif self._tf is not None and isinstance(data, self._tf.Graph):
            self._record_graph(data)
        else:
            raise ValueError('Unacceptable type.')

    def _record_tabular(self, data, step):
        if self._x_axis:
            nonexist_axes = []
            for axis in [self._x_axis] + self._additional_x_axes:
                if axis not in data.as_dict:
                    nonexist_axes.append(axis)
            if nonexist_axes:
                self._warn('{} {} exist in the tabular data.'.format(
                    ', '.join(nonexist_axes),
                    'do not' if len(nonexist_axes) > 1 else 'does not'))

        for key, value in data.as_dict.items():
            if isinstance(value,
                          np.ScalarType) and self._x_axis in data.as_dict:
                if self._x_axis is not key:
                    x = data.as_dict[self._x_axis]
                    self._record_kv(key, value, x)

                for axis in self._additional_x_axes:
                    if key is not axis and key in data.as_dict:
                        x = data.as_dict[axis]
                        self._record_kv('{}/{}'.format(key, axis), value, x)
            else:
                self._record_kv(key, value, step)
            data.mark(key)

    def _record_kv(self, key, value, step):
        if isinstance(value, str):
            self._writer.add_text(key, value, step)
        elif isinstance(value, np.ScalarType):
            self._writer.add_scalar(key, value, step)
            if self.wandb:
                wandb.log({key: value}, step=step)
        elif isinstance(value, plt.Figure):
            self._writer.add_figure(key, value, step)
            if self.wandb:
                wandb.log({key: wandb.Image(value)}, step=step)
        elif isinstance(value, np.ndarray) and value.ndim == 5:
            self._writer.add_video(key, value, step, fps=15)
            if self.wandb:
                wandb.log({key: wandb.Video(value)}, step=step)
        elif isinstance(value, scipy.stats._distn_infrastructure.rv_frozen):
            shape = (self._histogram_samples, ) + value.mean().shape
            self._writer.add_histogram(key, value.rvs(shape), step)
        elif isinstance(value, scipy.stats._multivariate.multi_rv_frozen):
            self._writer.add_histogram(key, value.rvs(self._histogram_samples),
                                       step)
        elif isinstance(value, Histogram):
            self._writer.add_histogram(key, np.asarray(value), step)

    def _record_graph(self, graph):
        graph_def = graph.as_graph_def(add_shapes=True)
        event = tbX.proto.event_pb2.Event(
            graph_def=graph_def.SerializeToString())
        self._writer.file_writer.add_event(event)

    def dump(self, step=None):
        """Flush summary writer to disk."""
        # Log the tabular inputs, now that we have a step
        for p in self._waiting_for_dump:
            p(step or self._default_step)
        self._waiting_for_dump.clear()

        # Flush output files
        for w in self._writer.all_writers.values():
            w.flush()

        self._default_step += 1

    def close(self):
        """Flush all the events to disk and close the file."""
        self._writer.close()
        if self.wandb:
            wandb.finish()

    def _warn(self, msg):
        """Warns the user using warnings.warn.

        The stacklevel parameter needs to be 3 to ensure the call to logger.log
        is the one printed.
        """
        if not self._disable_warnings and msg not in self._warned_once:
            warnings.warn(
                colorize(msg, 'yellow'), NonexistentAxesWarning, stacklevel=3)
        self._warned_once.add(msg)
        return msg


class NonexistentAxesWarning(LoggerWarning):
    """Raise when the specified x axes do not exist in the tabular."""
