'''
steps

1. build, lots of random data, but on simple models (ie. narrow and short dataframes)
2. use NN to test on a simple model
3. cyclical learning rate

# run experiment on kernel size -> train time/train quality

'''
import numpy as np
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import Dense,Dropout,Flatten
from keras.constraints import maxnorm
from keras.optimizers import SGD, Adam
from keras.layers.convolutional import Conv2D,MaxPooling2D
from keras.callbacks import EarlyStopping
from keras.utils import np_utils
from keras import backend as K, regularizers
from keras.layers.advanced_activations import ELU
from sklearn.model_selection import train_test_split
import keras

   
class CNN_2_BBN_Trainer():
    
    def __init__(self,*,complexity = "low",compression = True,compressFactor = 30):
        self.complexity = "low"
        self.compression = True
        self.compressFactor = 30
        self.momentum = 0.9
        self.epochs = 30
        self.holdout = True



    def createModel(self,*,complexity="low",compressFactor=30,dimensions=6,num_classes=2):
        ''' Creates the Convolutional Neural Network Model 
        
        Arguments:     
        complexity     -- the complexity (predominantly depth) of the conv neural net (default = "low")
        compressFactor -- the cmopression factor, equivalent to the length/height of the image
        dimensions     -- the dimensions of the data, equivalent to the width of the image
        num_classes    -- the number of classes, i.e. 2 or 3 (i.e. no edge + undirected edge or 
                          2 x directed edges)
       
        Returns:
        model          -- the convolutional neural network
        ''' 
        K.set_image_dim_ordering('th')
#        kernel_depth = 10
#        kernel_breadth = dimensions
        if complexity == "low":
            model = Sequential()
            model.add(Conv2D(filters=32, # dimensionality of output
                             kernel_size=(10, dimensions), 
                             input_shape=(1, compressFactor, dimensions), 
                             padding='same',
                             #activation=elu,#'relu',
                             bias_regularizer=regularizers.l1(0.01),
                             kernel_initializer='random_uniform',
                             kernel_constraint=maxnorm(3)))
            model.add(keras.layers.advanced_activations.ELU(alpha=1.0))
            model.add(Dropout(0.2)) #0.5
            model.add(Conv2D(filters=32,
                             kernel_size=(10, 6),
                             #activation=elu,#'relu',
                             padding='same',
                             bias_regularizer=regularizers.l1(0.01),
                             kernel_initializer='random_uniform',
                             kernel_constraint=maxnorm(3)))
            model.add(keras.layers.advanced_activations.ELU(alpha=1.0))
            model.add(Dropout(0.4)) #0.5
            model.add(Conv2D(filters=32,
                             kernel_size=(10, 6),
                             #activation=elu,#'relu',
                             padding='same',
                             kernel_initializer='random_uniform',
                             bias_regularizer=regularizers.l1(0.01),
                             kernel_constraint=maxnorm(3)))
            model.add(keras.layers.advanced_activations.ELU(alpha=1.0))
            model.add(MaxPooling2D(pool_size=(2, 2)))
            model.add(Flatten())
            model.add(Dense(150,#512, 200 best so far
                            #activation=elu,#'relu',
                            bias_regularizer=regularizers.l1(0.01),
                            kernel_initializer='random_uniform',
                            kernel_constraint=maxnorm(3)))
            model.add(keras.layers.advanced_activations.ELU(alpha=1.0))
            model.add(Dropout(0.5))
            model.add(Dense(num_classes, activation='softmax'))

        self.model = model
        
    def compileMod(self,opti = "Adam",lrate=0.01):
        self.lrate = lrate
        self.decay = self.lrate/self.epochs
        if opti == "SGD":
            optimizer = SGD(lr=self.lrate,momentum=self.momentum, decay=self.decay, nesterov=True)
        else:
            optimizer = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=self.decay)
        self.model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])

    def fitMod(self,*,epochs=30,batch_size=32):
        self.history = self.model.fit(self.X_train, self.y_train,
                                      validation_data=(self.X_val, self.y_val),
                                      epochs=epochs,batch_size=32)

    def importData(self,X,y,*,hardCodeSplit=False,holdout= True):
        self.holdout = holdout
        self.X = X
        self.y = y
        self.baseRate = max(np.mean(y, axis=0))
        if hardCodeSplit == True:
            Xyrange = len(X)//3 * 2
            self.X_train = self.X[:Xyrange]    
            X_valtest    = self.X[Xyrange:]  
            self.y_train = self.y[:Xyrange]   
            y_valtest    = self.y[Xyrange:]    
            if holdout == True:
                Xyrange = len(X_valtest)//2
                self.X_val,self.X_test,self.y_val,self.y_test = X[:Xyrange], X[Xyrange:], y[:Xyrange], y[Xyrange:]  
            else:
                self.X_val = X_valtest
                self.y_val = y_valtest
        else:
            self.X_train,X_valtest,self.y_train,y_valtest=train_test_split(self.X,self.y,
                                                                           test_size=0.5,
                                                                           random_state=99)
            if holdout == True:
                self.X_val,self.X_test,self.y_val,self.y_test=train_test_split(X_valtest,y_valtest,
                                                                               test_size=0.5,
                                                                               random_state=42)
            else:
                self.X_val = X_valtest
                self.y_val = y_valtest
        
    def dataSummary(self):        
        print ("------- Data Inputs ---------")
        print ("X shape -", self.X.shape)
        print ("y shape -", self.y.shape)
        print ("X val shape -", self.X_val.shape)
        print ("y val shape -", self.y_val.shape)
        print ("Holdout...? ", self.holdout)
        if self.holdout == True:
            print ("X_test shape -", self.X_test.shape)
            print ("y_test, shape -", self.y_test.shape)
        
    def printRes(self):
        print ("------- Results ---------")
        print ("Base Rate", self.baseRate,"/",abs(1-self.baseRate))
        valScores = self.model.evaluate(self.X_train,self.y_train, verbose=0)
        print("Train Accuracy: %.2f%%" % (valScores[1]*100))
        print("Train Loss: %.2f%%" % (valScores[0]*100))
        valScores = self.model.evaluate(self.X_val, self.y_val, verbose=0)
        print("Val Accuracy: %.2f%%" % (valScores[1]*100))
        print("Val Loss: %.2f%%" % (valScores[0]*100))
        
        if self.holdout == True:
            testScores = self.model.evaluate(self.X_test, self.y_test, verbose=0)
            print("Holdout Accuracy: %.2f%%" % (testScores[1]*100))
            print("Holdout Loss: %.2f%%" % (testScores[0]*100))

    def printGraphs(self):
        # list all data in history
        print(self.history.history.keys())
        # summarize history for accuracy
        plt.plot(self.history.history['acc'])
        plt.plot(self.history.history['val_acc'])
        plt.ylabel('accuracy')
        plt.xlabel('epoch')
        plt.legend(['train', 'test'], loc='upper left')
        plt.show()
        # summarize history for loss
        plt.plot(self.history.history['loss'])
        plt.plot(self.history.history['val_loss'])
        plt.title('model loss')
        plt.ylabel('loss')
        plt.xlabel('epoch')
        plt.legend(['train', 'test'], loc='lower left')
        plt.show()
        
    def runTest(self,X,y):
        cnnCont = CNNTrainerContainer()
        cnnCont.importData(X,y)      
        cnnCont.dataSummary()
        cnnCont.createModel()
        cnnCont.createModel(dimensions=6,num_classes=2)      
        cnnCont.model.summary()
        print("Base Rate -", cnnCont.baseRate)
        cnnCont.compileMod()
        cnnCont.fitMod()
        cnnCont.printRes()
        cnnCont.printGraphs()
        
    def saveModel(self,*,name):
        self.model.save(name +".h5")
        

