import tensorflow as tf
import numpy as np
import scipy.misc
from tensorboardX.src.summary_pb2 import Summary
from tensorboardX import SummaryWriter

try:
    from StringIO import StringIO
except ImportError:
    from io import BytesIO


class Logger(object):

    def __init__(self, log_dir):
        '''
        Create a summary writer logging to log_dir
        '''
        self.writer = tf.summary.FileWriter(log_dir)
        self.writerX = SummaryWriter(log_dir=log_dir)

    def scalar_summary(self, tag, value, step):
        '''
        Log scalar value
        '''
        summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)])
        self.writer.add_summary(summary, step)

    def add_image(self, tag, img, step):
        '''
        Log img
        '''
        summary = Summary(value=[Summary.Value(tag=tag, image=imgs)])
        self.writerX.add_summary(summary, step)

    def histo_summary(self, tag, values, step, bins=1000):
        '''
        Log a histogram of the tensor of values.
        '''

        # Create histogram:
        counts, bin_edges = np.histogram(values, bins=bins)

        # Fill the fields of the histogram proto
        hist = tf.HistogramProto()
        his.min = float(np.min(values))
        his.max = float(np.max(values))
        hist.num = int(np.prod(values.shape))
        hist.sum = float(np.sum(values))
        hist.sum_squares = float(np.sum(values**2))

        # Drop the start of the first bin
        bin_edges = bin_edges[1:]

        # Add bin edges and counts:
        for edge in bin_edges:
            hist.bucket_limit.append(edge)
        for c in counts:
            hist.bucket.append(c)

        # Create and write Summary
        summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)])
        self.writer.add_summary(summary, step)
        self.writer.flush()
