from models import AbstractModel
from pretreat import *
from keras.callbacks import TensorBoard, ModelCheckpoint
from keras.layers import Conv2D, MaxPooling2D, Dropout, Flatten, Dense
from keras.models import Input, Model, load_model
from keras.optimizers import SGD
import numpy as np
import config as C


class CNNModel(AbstractModel):

    def __init__(self):
        super().__init__()

        self.best_weights_path = os.path.join(C.weights_path, 'model01best.h5')
        self.weights_path = os.path.join(C.weights_path, 'model01.h5')

        self.label_handler = LabelHandler(C.NUM)

        self.x_train_path = os.path.join(C.train_path, 'inputs.npy')
        self.y_train_path = os.path.join(C.train_path, 'labels.npy')

        self.x_test_path = os.path.join(C.test_path, 'inputs.npy')
        self.y_test_path = os.path.join(C.test_path, 'labels.npy')

    def build(self):
        h, w, c = C.Height, C.Width, C.Depth
        inputs = Input((h, w, c))
        conv1_1 = Conv2D(32, 3, activation='relu', padding='same', kernel_initializer='he_normal')(inputs)
        conv1_2 = Conv2D(32, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1_1)
        pool1 = MaxPooling2D((2, 2))(conv1_2)
        conv2_1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool1)
        conv2_2 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2_1)
        pool2 = MaxPooling2D((2, 2))(conv2_2)
        conv3_1 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool2)
        conv3_2 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3_1)
        pool3 = MaxPooling2D((2, 2))(conv3_2)
        conv4_1 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool3)
        conv4_2 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4_1)
        conv4_3 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4_2)
        pool4 = MaxPooling2D((2, 2))(conv4_3)
        # conv5_1 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool4)
        # conv5_2 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv5_1)
        # conv5_3 = Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv5_2)
        # pool5 = MaxPooling2D((2,2))(conv5_3)
        flat = Flatten()(pool4)
        drop = Dropout(0.25)(flat)
        # dense1 = Dense(512,activation='relu')(flat)
        # drop1 = Dropout(0.5)(dense1)
        # outputs = Dense(40,activation='relu',kernel_regularizer=regularizers.l2(0.01))(drop1)
        outputs = [Dense(LabelHandler.classes, activation='softmax', name='c%d' % (i + 1))(drop) for i in range(C.NUM)]
        # outputs = Reshape((4,10))(outputs)
        model = Model(inputs, outputs)
        # sgd = SGD(0.1, decay=1e-3, momentum=0.9, nesterov=True)
        model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
        model.summary()

        return model

    @staticmethod
    def decode(y):
        n = len(y)
        y = y.reshape((C.NUM, n, LabelHandler.classes))
        return list(y)

    def train(self, batch_size=256, epochs=200):
        # self.model.load_weights(self.best_model_path)
        x, y = self.x_train, self.y_train
        print(x.shape, y.shape)
        y = self.decode(y)
        save_best = ModelCheckpoint(self.best_weights_path, monitor='val_loss', verbose=1, save_best_only=True,
                                    save_weights_only=True)
        tensor_board = TensorBoard(log_dir='log/log_cnn')
        callbacks = [save_best, tensor_board]
        with self.graph.as_default():
            with self.sess.as_default():
                self.model.fit(x, y, batch_size=batch_size, epochs=epochs, verbose=1, validation_split=0.1, callbacks=callbacks,
                               shuffle=True)
                self.model.save_weights(self.weights_path)

    def test(self):
        y = self.decode(self.y_test)
        score = self.model.evaluate(self.x_test, y, batch_size=64, verbose=1)
        print(score)

    def inference(self, x: np.ndarray):
        if len(x.shape) == 3:
            x = np.expand_dims(x, axis=0)
        with self.graph.as_default():
            with self.sess.as_default():
                results = self.model.predict(x, batch_size=64, verbose=1)
        return [self.label_handler.decode(res) for res in results]

    def pre_process(self, x):
        return img_pretreat(x)

    def post_process(self, x):
        return self.label_handler.decode(x[0])

