import time
import tensorflow as tf

from model import evaluate
from model import srgan

from tensorflow.keras.applications.vgg19 import preprocess_input
from tensorflow.keras.losses import BinaryCrossentropy
from tensorflow.keras.losses import MeanAbsoluteError
from tensorflow.keras.losses import MeanSquaredError
from tensorflow.keras.metrics import Mean
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.optimizers.schedules import PiecewiseConstantDecay


class Trainer:
    def __init__(self,
                 model,
                 loss,
                 learning_rate,
                 checkpoint_dir='./ckpt/edsr'):

        self.now = None
        self.loss = loss
        self.checkpoint = tf.train.Checkpoint(step=tf.Variable(0),
                                              psnr=tf.Variable(-1.0),
                                              optimizer=Adam(learning_rate),
                                              model=model)
        self.checkpoint_manager = tf.train.CheckpointManager(checkpoint=self.checkpoint,
                                                             directory=checkpoint_dir,
                                                             max_to_keep=3)

        self.restore()

    @property
    def model(self):
        return self.checkpoint.model

    def train(self, train_dataset, valid_dataset, steps, evaluate_every=1000, save_best_only=False):
        loss_mean = Mean()

        ckpt_mgr = self.checkpoint_manager
        ckpt = self.checkpoint

        self.now = time.perf_counter()

        for lr, hr in train_dataset.take(steps - ckpt.step.numpy()):
            ckpt.step.assign_add(1)
            step = ckpt.step.numpy()

            loss = self.train_step(lr, hr)
            loss_mean(loss)

            if step % evaluate_every == 0:
                loss_value = loss_mean.result()
                loss_mean.reset_states()

                # Compute PSNR on validation dataset
                psnr_value = self.evaluate(valid_dataset)

                duration = time.perf_counter() - self.now
                print(f'{step}/{steps}: loss = {loss_value.numpy():.3f}, PSNR = {psnr_value.numpy():3f} ({duration:.2f}s)')

                if save_best_only and psnr_value <= ckpt.psnr:
                    self.now = time.perf_counter()
                    # skip saving checkpoint, no PSNR improvement
                    continue

                ckpt.psnr = psnr_value
                ckpt_mgr.save()

                self.now = time.perf_counter()

    @tf.function
    def train_step(self, lr, hr):
        with tf.GradientTape() as tape:
            lr = tf.cast(lr, tf.float32)
            hr = tf.cast(hr, tf.float32)

            sr = self.checkpoint.model(lr, training=True)
            loss_value = self.loss(hr, sr)

        gradients = tape.gradient(loss_value, self.checkpoint.model.trainable_variables)
        self.checkpoint.optimizer.apply_gradients(zip(gradients, self.checkpoint.model.trainable_variables))

        return loss_value

    def evaluate(self, dataset):
        return evaluate(self.checkpoint.model, dataset)

    def restore(self):
        if self.checkpoint_manager.latest_checkpoint:
            self.checkpoint.restore(self.checkpoint_manager.latest_checkpoint)
            print(f'Model restored from checkpoint at step {self.checkpoint.step.numpy()}.')


class EdsrTrainer(Trainer):
    def __init__(self,
                 model,
                 checkpoint_dir,
                 learning_rate=PiecewiseConstantDecay(boundaries=[200000], values=[1e-4, 5e-5])):
        super().__init__(model, loss=MeanAbsoluteError(), learning_rate=learning_rate, checkpoint_dir=checkpoint_dir)

    def train(self, train_dataset, valid_dataset, steps=300000, evaluate_every=1000, save_best_only=True):
        super().train(train_dataset, valid_dataset, steps, evaluate_every, save_best_only)


class WdsrTrainer(Trainer):
    def __init__(self,
                 model,
                 checkpoint_dir,
                 learning_rate=PiecewiseConstantDecay(boundaries=[200000], values=[1e-3, 5e-4])):
        super().__init__(model, loss=MeanAbsoluteError(), learning_rate=learning_rate, checkpoint_dir=checkpoint_dir)

    def train(self, train_dataset, valid_dataset, steps=300000, evaluate_every=1000, save_best_only=True):
        super().train(train_dataset, valid_dataset, steps, evaluate_every, save_best_only)




tensorflow.train.CheckpointManager tensorflow.Variable tensorflow.cast tensorflow.keras.optimizers.Adam tensorflow.keras.losses.MeanAbsoluteError tensorflow.keras.optimizers.schedules.PiecewiseConstantDecay tensorflow.keras.metrics.Mean tensorflow.GradientTape
[(29, 'tensorflow.train.CheckpointManager', 'tf.train.CheckpointManager', ([], {'checkpoint': 'self.checkpoint', 'directory': 'checkpoint_dir', 'max_to_keep': '(3)'}), True, 'import tensorflow as tf\n'), (40, 'tensorflow.keras.metrics.Mean', 'Mean', ([], {}), False, 'from tensorflow.keras.metrics import Mean\n'), (45, 'time.perf_counter', 'time.perf_counter', ([], {}), False, 'import time\n'), (89, 'model.evaluate', 'evaluate', (['self.checkpoint.model', 'dataset'], {}), False, 'from model import evaluate\n'), (101, 'tensorflow.keras.optimizers.schedules.PiecewiseConstantDecay', 'PiecewiseConstantDecay', ([], {'boundaries': '[200000]', 'values': '[0.0001, 5e-05]'}), False, 'from tensorflow.keras.optimizers.schedules import PiecewiseConstantDecay\n'), (112, 'tensorflow.keras.optimizers.schedules.PiecewiseConstantDecay', 'PiecewiseConstantDecay', ([], {'boundaries': '[200000]', 'values': '[0.001, 0.0005]'}), False, 'from tensorflow.keras.optimizers.schedules import PiecewiseConstantDecay\n'), (76, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (77, 'tensorflow.cast', 'tf.cast', (['lr', 'tf.float32'], {}), True, 'import tensorflow as tf\n'), (78, 'tensorflow.cast', 'tf.cast', (['hr', 'tf.float32'], {}), True, 'import tensorflow as tf\n'), (25, 'tensorflow.Variable', 'tf.Variable', (['(0)'], {}), True, 'import tensorflow as tf\n'), (26, 'tensorflow.Variable', 'tf.Variable', (['(-1.0)'], {}), True, 'import tensorflow as tf\n'), (27, 'tensorflow.keras.optimizers.Adam', 'Adam', (['learning_rate'], {}), False, 'from tensorflow.keras.optimizers import Adam\n'), (72, 'time.perf_counter', 'time.perf_counter', ([], {}), False, 'import time\n'), (102, 'tensorflow.keras.losses.MeanAbsoluteError', 'MeanAbsoluteError', ([], {}), False, 'from tensorflow.keras.losses import MeanAbsoluteError\n'), (113, 'tensorflow.keras.losses.MeanAbsoluteError', 'MeanAbsoluteError', ([], {}), False, 'from tensorflow.keras.losses import MeanAbsoluteError\n'), (61, 'time.perf_counter', 'time.perf_counter', ([], {}), False, 'import time\n'), (65, 'time.perf_counter', 'time.perf_counter', ([], {}), False, 'import time\n')]