import os
from enum import Enum

from tensorflow.compat.v1 import variable_scope, placeholder
from tensorflow.compat.v1.summary import image, FileWriter, scalar


class Phase(Enum):
    TRAIN = 'TRAIN'
    VAL = 'VAL'
    TEST = 'TEST'


class Logger:
    def __init__(self, sess, summary_dir):
        self.sess = sess
        self.summary_placeholders = {}
        self.summary_ops = {}
        self.train_summary_writer = FileWriter(os.path.join(summary_dir, Phase.TRAIN.value), self.sess.graph)
        self.val_summary_writer = FileWriter(os.path.join(summary_dir, Phase.VAL.value))
        self.test_summary_writer = FileWriter(os.path.join(summary_dir, Phase.TEST.value))

    # it can summarize scalars and images.
    def summarize(self, step, phase: Phase = Phase.TRAIN, scope="", summaries_dict=None):
        """
        :param step: the step of the summary
        :param phase: use the train summary writer or the test one
        :param scope: variable scope
        :param summaries_dict: the dict of the summaries values (tag,value)
        :return:
        """
        if phase == Phase.TRAIN:
            summary_writer = self.train_summary_writer
        elif phase == Phase.VAL:
            summary_writer = self.val_summary_writer
        elif phase == Phase.TEST:
            summary_writer = self.test_summary_writer
        else:
            raise ValueError(f'Illegal Argument for summarizer: {phase.value}')

        with variable_scope(scope):

            if summaries_dict is not None:
                summary_list = []
                for tag, value in summaries_dict.items():
                    if value is None:
                        continue
                    if tag not in self.summary_ops:
                        if len(value.shape) <= 1:
                            self.summary_placeholders[tag] = placeholder('float32', value.shape, name=tag)
                            self.summary_ops[tag] = scalar(tag, self.summary_placeholders[tag])
                        else:
                            self.summary_placeholders[tag] = placeholder('float32', [None] + list(value.shape[1:]), name=tag)
                            self.summary_ops[tag] = image(tag, self.summary_placeholders[tag], max_outputs=100)

                    summary_list.append(self.sess.run(self.summary_ops[tag], {self.summary_placeholders[tag]: value}))

                for summary in summary_list:
                    summary_writer.add_summary(summary, step)
                summary_writer.flush()