# To add a new cell, type '# %%'
# To add a new markdown cell, type '# %% [markdown]'

#%%
import tensorflow as tf
from tensorflow.keras import Input, layers
import matplotlib.pyplot as plt
import numpy as np
import time
import os
from datetime import datetime
from sklearn.model_selection import train_test_split as tvsplit
# disable logging warning and error
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tf.config.experimental_run_functions_eagerly(True)
#%%
sample_n = 100
epochs = 50
#%%
meana = np.array([1, 1])
cova = np.array([[0.1, 0],[0, 0.1]])

meanb = np.array([2, 2])
covb = np.array([[0.1, 0],[0, 0.1]])

x_red = np.random.multivariate_normal(mean=meana, cov = cova, size=sample_n)
x_green = np.random.multivariate_normal(mean=meanb, cov = covb, size=sample_n)

y_red = np.array([1] * sample_n)
y_green = np.array([0] * sample_n)

plt.scatter(x_red[:, 0], x_red[:, 1], c = 'red' , marker='.', s = 30)
plt.scatter(x_green[:, 0], x_green[:, 1], c = 'green', marker='.', s = 30)
# plt.show()

X = np.concatenate([x_red, x_green]).astype(np.float32)
y = np.concatenate([y_red, y_green]).astype(np.float32)

#%%
X_train, X_test, y_train, y_test = tvsplit(X, y)

#%%
class Logistic(tf.keras.models.Model):
    def __init__(self, hidden_size = 5, output_size=1, dynamic=False, **kwargs):
        '''
        name: String name of the model.
        dynamic: (Subclassed models only) Set this to `True` if your model should
            only be run eagerly, and should not be used to generate a static
            computation graph. This attribute is automatically set for Functional API
            models.
        trainable: Boolean, whether the model's variables should be trainable.
        dtype: (Subclassed models only) Default dtype of the model's weights (
            default of `None` means use the type of the first input). This attribute
            has no effect on Functional API models, which do not have weights of their
            own.
        '''
        super().__init__(dynamic=dynamic, **kwargs)
        self.inputs_ = tf.keras.Input(shape=(2,), name="hello")
        self._set_input_layer(self.inputs_)
        self.hidden_size = hidden_size
        self.dense = layers.Dense(hidden_size, name = "linear")
        self.outlayer = layers.Dense(output_size, 
                        activation = 'sigmoid', name = "out_layer")
        self.build()
        

    def _set_input_layer(self, inputs):
        """add inputLayer to model and display InputLayers in model.summary()

        Args:
            inputs ([dict]): the result from `tf.keras.Input`
        """
        if isinstance(inputs, dict):
            self.inputs_layer = {n: tf.keras.layers.InputLayer(input_tensor=i, name=n) 
                                    for n, i in inputs.items()}
        elif isinstance(inputs, (list, tuple)):
            self.inputs_layer = [tf.keras.layers.InputLayer(input_tensor=i, name=i.name) 
                                    for i in inputs]
        elif tf.is_tensor(inputs):
            self.inputs_layer = tf.keras.layers.InputLayer(input_tensor=inputs, name=inputs.name)
    
    def build(self):
        super(Logistic, self).build(self.inputs_.shape if tf.is_tensor(self.inputs_) else self.inputs_)
        # super(Logistic, self).build((None, 2))
        _ = self.call(self.inputs_)
    

    def call(self, X):
        X = self.dense(X)
        Y = self.outlayer(X)
        Y = tf.squeeze(Y, axis = 1)
        return Y


#%%
model = Logistic()
model.summary()

# %%
optimizer=tf.keras.optimizers.Adam()
loss=tf.keras.losses.BinaryCrossentropy()
metrics=tf.keras.metrics.AUC()

#%%
# @tf.function # in graph mode
def losses(y_true, y_pred, sample_weight=None, loss_obj=loss):
    return loss_obj(y_true, y_pred, sample_weight)

#%%
# @tf.function # in graph mode
def Metrics(y_true, y_pred, sample_weight=None, metrics=metrics):
    metrics.update_state(y_true, y_pred, sample_weight)
    return metrics.result()
#%%
# @tf.function # in graph mode
def grad(model, inputs, labels):
    
    with tf.GradientTape() as tape:
        pred = model(inputs)
        loss_value = losses(labels, pred)
        labels = tf.expand_dims(labels, axis=1)
        metr = Metrics(labels, pred)
    
    return loss_value, tape.gradient(loss_value, model.trainable_variables), metr

# %%
# @tf.function # in graph mode
def train_on_batch(model, inputs, labels, opt=optimizer):
    closs, cgrad, cmetric = grad(model, inputs, labels)
    opt.apply_gradients(zip(
        cgrad,
        model.trainable_variables
    ))
    return closs, cmetric

#%%
for e in range(epochs):
    for ind in range(len(X_train)):
        loss_value, cmetric = train_on_batch(model, X_train[ind][None, :],
                np.expand_dims(y_train[ind], axis=0))
    # tf.print(f"Epochs {e}: loss {loss_value.numpy()}, metric:{cmetric.numpy()}") # in graph mode
    print(f"Epochs {e}: loss {loss_value.numpy()}, metric:{cmetric.numpy()}")

#%%
# model.save("save/subclass/1")