import psutil
import time
import sys

import tensorflow as tf


class CostBenchmarkSpyCallback(tf.keras.callbacks.Callback):
    def __init__(self, prefix='[COST BENCHMARK]', suffix=''):
        self.log_options = dict(
            prefix=prefix,
            suffix=suffix,
        )
        self.time = dict(
            created=time.time(),
            boot=psutil.boot_time(),
            train={},
            epoch={},
            batch={},
        )
        self.delta = dict(
            boot=self.time['created'] - self.time['boot'],
            train=None,
            sync=None,
            epoch={},
            batch={},
        )

        self._log_time('Boot', self.delta.get("boot"))

    def on_train_begin(self, logs=None):
        self.time.get('train').update({'start': time.time()})
        self.delta.update({'sync': self.time.get('train').get('start') - self.time.get('created')})

        self._log_time('Initialization', self.delta.get('sync'))

    def on_train_end(self, logs=None):
        self.time.get('train').update({'end': time.time()})
        self.delta.update({'train': self.time.get('train').get('end') - self.time.get('train').get('start')})

        self._log_time('Train', self.delta.get('train'))

    def on_epoch_begin(self, epoch, logs=None):
        self.time.get('epoch').update({epoch: dict(start=time.time())})

    def on_epoch_end(self, epoch, logs=None):
        self.time.get('epoch').get(epoch).update({'end': time.time()})

        epoch_delta = self.time.get('epoch').get(epoch).get('end') - \
                      self.time.get('epoch').get(epoch).get('start')
        self.delta.get('epoch').update({ epoch: epoch_delta })
        self._log_time('Epoch ', epoch_delta, epoch + 1)

    def on_train_batch_begin(self, batch, logs=None):
        self.time.get('batch').update({batch: dict(start=time.time())})

    def on_train_batch_end(self, batch, logs=None):
        self.time.get('batch').get(batch).update({'end': time.time()})

        batch_delta = self.time.get('batch').get(batch).get('end') - \
                      self.time.get('batch').get(batch).get('start')
        self.delta.get('batch').update({ batch: batch_delta })
        self._log_time('Batch', batch_delta, batch + 1)

    def _log_time(self, type, time, iterator_number=None):
        prefix = self.log_options.get('prefix')
        suffix = self.log_options.get('suffix')
        if (iterator_number):
            suffix = suffix + f' - ITERATION: {iterator_number}'

        print(f'{prefix} {type}: {time} seconds {suffix}', file=sys.stdout)

import psutil
import time
import sys

import tensorflow as tf


class CostBenchmarkSpyCallback(tf.keras.callbacks.Callback):
    def __init__(self, prefix='[COST BENCHMARK]', suffix=''):
        self.log_options = dict(
            prefix=prefix,
            suffix=suffix,
        )
        self.time = dict(
            created=time.time(),
            boot=psutil.boot_time(),
            train={},
            epoch={},
            batch={},
        )
        self.delta = dict(
            boot=self.time['created'] - self.time['boot'],
            train=None,
            sync=None,
            epoch={},
            batch={},
        )

        self._log_time('Boot', self.delta.get("boot"))

    def on_train_begin(self, logs=None):
        self.time.get('train').update({'start': time.time()})
        self.delta.update({'sync': self.time.get('train').get('start') - self.time.get('created')})

        self._log_time('Initialization', self.delta.get('sync'))

    def on_train_end(self, logs=None):
        self.time.get('train').update({'end': time.time()})
        self.delta.update({'train': self.time.get('train').get('end') - self.time.get('train').get('start')})

        self._log_time('Train', self.delta.get('train'))

    def on_epoch_begin(self, epoch, logs=None):
        self.time.get('epoch').update({epoch: dict(start=time.time())})

    def on_epoch_end(self, epoch, logs=None):
        self.time.get('epoch').get(epoch).update({'end': time.time()})

        epoch_delta = self.time.get('epoch').get(epoch).get('end') - \
                      self.time.get('epoch').get(epoch).get('start')
        self.delta.get('epoch').update({ epoch: epoch_delta })
        self._log_time('Epoch ', epoch_delta, epoch + 1)

    def on_train_batch_begin(self, batch, logs=None):
        self.time.get('batch').update({batch: dict(start=time.time())})

    def on_train_batch_end(self, batch, logs=None):
        self.time.get('batch').get(batch).update({'end': time.time()})

        batch_delta = self.time.get('batch').get(batch).get('end') - \
                      self.time.get('batch').get(batch).get('start')
        self.delta.get('batch').update({ batch: batch_delta })
        self._log_time('Batch', batch_delta, batch + 1)

    def _log_time(self, type, time, iterator_number=None):
        prefix = self.log_options.get('prefix')
        suffix = self.log_options.get('suffix')
        if (iterator_number):
            suffix = suffix + f' - ITERATION: {iterator_number}'

        print(f'{prefix} {type}: {time} seconds {suffix}', file=sys.stdout)
