import numpy as np

from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Dropout
from keras.layers import Flatten
from keras.layers.convolutional import Conv1D
from keras.layers.convolutional import MaxPooling1D
from keras.utils import np_utils
from keras import backend as K
import warnings
warnings.simplefilter('ignore')

from keras.optimizers import Optimizer

from keras import backend as K

if K.backend() == 'tensorflow':
    import tensorflow as tf
from six.moves import zip

if K.backend() == 'tensorflow':
    import tensorflow as tf
from tensorflow.keras.optimizers import Optimizer
from tensorflow.keras import backend as kb
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.framework import ops


class WAME(Optimizer):

    def __init__(
            self,
            learning_rate=0.001,
            alpha=0.9,
            eta_plus=1.2,
            eta_minus=0.1,
            zeta_min=0.01,
            zeta_max=100,
            name='WAME',
            **kwargs
    ):
       
        super(WAME, self).__init__(name, **kwargs)

        self._set_hyper('learning_rate', learning_rate)
        self._set_hyper('alpha', alpha)
        self._set_hyper('eta_plus', eta_plus)
        self._set_hyper('eta_minus', eta_minus)
        self._set_hyper('zeta_min', zeta_min)
        self._set_hyper('zeta_max', zeta_max)
        self.epsilon = kb.epsilon()

    def _prepare_local(self, var_device, var_dtype, apply_state):
        """
        """
        super(WAME, self)._prepare_local(var_device, var_dtype, apply_state)

        apply_state[(var_device, var_dtype)].update(
            {
                'alpha': array_ops.identity(self._get_hyper('alpha', var_dtype)),
                'eta_plus': array_ops.identity(self._get_hyper('eta_plus', var_dtype)),
                'eta_minus': array_ops.identity(self._get_hyper('eta_minus', var_dtype)),
                'zeta_min': array_ops.identity(self._get_hyper('zeta_min', var_dtype)),
                'zeta_max': array_ops.identity(self._get_hyper('zeta_max', var_dtype)),
                'epsilon': ops.convert_to_tensor_v2(self.epsilon, var_dtype)
            }
        )

    def _resource_apply_dense(self, grad, var, apply_state=None):
       
        var_device, var_dtype = var.device, var.dtype.base_dtype

        coefficients = (
                (apply_state or {}).get((var_device, var_dtype))
                or self._fallback_apply_state(var_device, var_dtype)
        )

        # Static Coefficients
        learning_rate = coefficients['lr_t']
        alpha = coefficients['alpha']
        eta_plus = coefficients['eta_plus']
        eta_minus = coefficients['eta_minus']
        zeta_min = coefficients['zeta_min']
        zeta_max = coefficients['zeta_max']
        epsilon = coefficients['epsilon']

        # Variables
        previous_gradient = self.get_slot(var, 'previous_gradient')
        z = self.get_slot(var, 'z')
        theta = self.get_slot(var, 'theta')
        zeta = self.get_slot(var, 'zeta')

        # Calculate weight delta
        zeta_t = kb.switch(
            math_ops.less(grad * previous_gradient, 0),
            theta * eta_plus,
            kb.switch(
                math_ops.greater(grad * previous_gradient, 0),
                theta * eta_minus,
                zeta
            )
        )
        zeta_t = kb.clip(zeta_t, zeta_min, zeta_max)
        z_t = (alpha * z) + ((1 - alpha) * zeta_t)
        theta_t = (alpha * theta) + ((1 - alpha) * math_ops.square(grad))
        w_delta = - (learning_rate * (z_t * grad * (1 / (theta_t + epsilon))))
        var_t = var + w_delta

        # Update variables
        state_ops.assign(previous_gradient, grad, use_locking=self._use_locking)
        state_ops.assign(zeta, zeta_t, use_locking=self._use_locking)
        state_ops.assign(z, z_t, use_locking=self._use_locking)
        state_ops.assign(theta, theta_t, use_locking=self._use_locking)

        return state_ops.assign(var, var_t, use_locking=self._use_locking).op

    def _create_slots(self, var_list):
        """ Defines trainable TensorFlow variables.
        Args:
            var_list: list of `Variable` objects that will be minimized using this optimizer.
        """
        for var in var_list:
            self.add_slot(var, 'z')
            self.add_slot(var, 'theta')
            self.add_slot(var, 'previous_gradient')
            self.add_slot(var, 'zeta')

    def get_config(self):
        """ Returns the config of the optimizer.
        An optimizer config is a Python dictionary (serializable)
        containing the configuration of an optimizer.
        The same optimizer can be reinstantiated later
        (without any saved state) from this configuration.
        It is good practice to define the get_config and from_config
        methods when writing a custom model or layer class.
        This allows you to easily update the computation later if needed.
        Returns:
            dict: Model configuration dictionary.
        """
        config = super(WAME, self).get_config()

        config.update({
            'learning_rate': self._serialize_hyperparameter('learning_rate'),
            'alpha': self._serialize_hyperparameter('alpha'),
            'eta_plus': self._serialize_hyperparameter('eta_plus'),
            'eta_minus': self._serialize_hyperparameter("eta_minus"),
            'zeta_min': self._serialize_hyperparameter("zeta_min"),
            'zeta_max': self._serialize_hyperparameter("zeta_max"),
        })

        return config

"""Baseline Model"""

# Implement baseline model with Keras Sequential
# the model with some pre assigned hypeparameter

def deep_cnn_model(shape):
    
    # Model Layers
    model = Sequential()
    model.add(Conv1D(64, 5, padding='same',
                     input_shape=shape,
                     activation='relu'))
    model.add(Conv1D(64, 5,  activation='relu',padding="same"))
    model.add(MaxPooling1D(pool_size=2,padding='same'))
    model.add(Dropout(0.2))
    model.add(Flatten())
    model.add(Dense(3, activation='softmax'))
    
    # Compile Model
    model.compile(loss='categorical_crossentropy',
                  optimizer= WAME(),
                  metrics=['accuracy'])
    
    return model

# import dataset
import pandas as pd
white_wine = pd.read_csv("https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-white.csv",
                sep=";")
red_wine = pd.read_csv("https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv",
                       sep=";")

white_wine.head()

white_wine.shape

red_wine.head()

red_wine.shape

# combine both of the red and white wine dataset
df  = pd.concat([white_wine,red_wine])

df["quality"].unique() # we divide the wine quality in three major groups on the conditions

df['quality_label']=df['quality'].apply(lambda x: 'low' if x<=5 else 'medium' if x<=7 else 'high')

print(df.quality_label.value_counts())
# there is a class imbalance problem we need to fix so that the model can learn equally in train set
df.drop("quality",axis=1,inplace=True)

# divide data into train and test split
X=df.drop("quality_label",axis=1).values
y=df["quality_label"].values

X.shape

y.shape

# convert strngs to nums
from sklearn.preprocessing import LabelEncoder
lbl= LabelEncoder()
y=lbl.fit_transform(y)
from tensorflow.keras.utils import to_categorical
# y= to_categorical(y,num_classes=3)

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.20, random_state=101)

y_train= to_categorical(y_train,3)
y_test = to_categorical(y_test,3)

# now train the model using default parameters
x_train=X_train.reshape(-1,1,X_train.shape[1])

x_train.shape

y_train.shape

model=deep_cnn_model(shape=(x_train.shape[1],x_train.shape[2]))

model.summary()

from tensorflow.keras.callbacks import ReduceLROnPlateau
model.fit(x_train, 
          y_train,validation_split=0.2,
          
          epochs=100,batch_size=128)

# predict on test set
x_test= X_test.reshape(-1,1,X_test.shape[1])
y_pred =  np.argmax(model.predict(x_test),axis=1)

y_act = np.argmax(y_test,axis=1)
y_act_train =  np.argmax(y_train,axis=1)
y_pred_train= np.argmax(model.predict(x_train),axis=1)
from sklearn import metrics
print(metrics.accuracy_score(y_pred,y_act))
print(metrics.accuracy_score(y_pred_train,y_act_train))

print(metrics.f1_score(y_pred,y_act,average="weighted"))
print(metrics.f1_score(y_pred_train,y_act_train,average="weighted"))

print(metrics.classification_report(y_pred,y_act))

# next do hyperparameter tuning for the Neural Network
# we tune the activation function , number of neurons and drop_out rate of the network
def deep_cnn_model_tune(shape=(1,11),neurons1=64,activation1="relu",neurons2=64,activation2="relu",drop_rate=0.2):
    
    # Model Layers
    model = Sequential()
    model.add(Conv1D(neurons1, 5, padding='same',
                     input_shape=shape,
                     activation=activation1))
    model.add(Conv1D(neurons2, 5,  activation=activation2,padding="same"))
    model.add(MaxPooling1D(pool_size=2,padding='same'))
    model.add(Dropout(drop_rate))
    model.add(Flatten())
    model.add(Dense(3, activation='softmax'))
    
    # Compile Model
    model.compile(loss='categorical_crossentropy',
                  optimizer= WAME(),
                  metrics=['accuracy'])
    
    return model

from sklearn.model_selection import RandomizedSearchCV
from tensorflow.keras.wrappers.scikit_learn import KerasClassifier

tune_params={
    "neurons1":[64,128,256],
    "activation1":["relu","sigmoid"],
    "neurons2":[64,128,256],
    "activation2":["relu","sigmoid"],
    "drop_rate":[0.2,0.4,0.8]
}
tune_params

model = KerasClassifier(build_fn=deep_cnn_model_tune,epochs=100)

np.random.seed(101) # set seed for reproducibility


# Run random search model. Iterations over 5 epochs.
randomCV = RandomizedSearchCV(model, param_distributions=tune_params, cv=2,     n_iter=5,
                              verbose=0)

randomCV.fit(x_train,y_train)

randomCV.best_params_

y_pred_cv = randomCV.predict(x_test)
y_pred_cv_train = randomCV.predict(x_train)

print(metrics.accuracy_score(y_act,y_pred_cv))
print(metrics.accuracy_score(y_act_train,y_pred_cv_train))

print(metrics.f1_score(y_act,y_pred_cv))
print(metrics.f1_score(y_act_train,y_pred_cv_train))

print(metrics.classification_report(y_act,y_pred_cv))

print(metrics.confusion_matrix(y_act,y_pred_cv))