# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/function_approximation.ipynb (unless otherwise specified).

__all__ = ['KerasPipeline', 'HighLevelKerasPipeline']

# Cell
from abc import ABC, abstractmethod
from fastcore.test import *
import jovsatools
from jovsatools import data_generator
from nbdev.showdoc import *
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers.experimental.preprocessing import Rescaling

test_eq(tf.__version__, "2.2.0")

# Cell

class KerasPipeline(ABC):
    """Scafolding for a model building pipeline.

    This base class will contain scafolding for all pipelines
    used. The goal of this base class is mainly to add safeguard
    and assumption enforcement.
    """
    def __init__(self):
        # safe guard: to ensure that various session don't collide
        keras.backend.clear_session()
        self.train_x, self.train_y, self.test_x, self.test_y = None, None, None, None
        self.model = None
        self.evaluation = None
        self.train_history = None

    @abstractmethod
    def get_data(self):
        pass

    @abstractmethod
    def build_model(self):
        pass

    @abstractmethod
    def run_training(self):
        pass

    @abstractmethod
    def evaluate_pipeline(self):
        pass

    def keras_model_summary(self):
        assert self.model is not None
        return self.model.summary()

    @abstractmethod
    def __call__(self):
        # checking that self.get_data() is implimented as expected
        assert list(
            map(lambda x: x is not None, [self.train_x, self.train_y, self.test_x, self.test_y])
            )  == [True] * 4

        # checking that self.model() is implimented as expected
        assert self.model is not None

        # checking that self.train() is implimented as expected
        assert self.train_history is not None
        assert self.model._is_compiled == True

        # checking that self.evaluate() is implimented as expected
        assert self.evaluation is not None

        # TODO (jovsa): run tests that ensure pipeline trains model above a benchmark


# Cell
class HighLevelKerasPipeline(KerasPipeline):
    def __init__(self):
        super().__init__()

    def get_data(self):
        additional_y, train_n = 0, 1000
        data = data_generator.MNISTDataGenerator(additional_y)
        datasets  = data(train_n)
        self.train_x, self.train_y, self.test_x, self.test_y = datasets

    def build_model(self):
        inputs = keras.Input(shape=(784, 1))
        # Rescale images to [0, 1]
        x = Rescaling(scale=1./255)(inputs)
        x = keras.layers.Flatten()(x)
        x = keras.layers.Dense(units=128, activation='relu')(x)
        x = keras.layers.Dense(units=128, activation='relu')(x)
        outputs = keras.layers.Dense(units=10, activation='softmax')(x)

        self.model = keras.Model(inputs=inputs, outputs=outputs)

    def run_training(self):
        self.model.compile(
            optimizer=keras.optimizers.RMSprop(learning_rate=1e-3),
            loss=keras.losses.SparseCategoricalCrossentropy(),
            metrics=[keras.metrics.SparseCategoricalAccuracy()])

        self.train_history = self.model.fit(
            x=self.train_x, y=self.train_y, epochs=10, batch_size=64,
            verbose=1, validation_split=0.1)

    def evaluate_pipeline(self):
        self.evaluation = self.model.evaluate(x=self.test_x, y=self.test_y)
        print(self.evaluation)

    def __call__(self):
        self.get_data()
        self.build_model()
        self.run_training()
        print(self.keras_model_summary())
        self.evaluate_pipeline()
        super().__call__()