from keras.callbacks import TensorBoard, ModelCheckpoint
from keras.layers import Conv2D, MaxPooling2D, Dropout, Dense, Reshape, GRU, Lambda
from keras.models import Input, Model
from keras.layers.wrappers import Bidirectional
from keras import backend as K

from models import AbstractModel
from pretreat import *
import config as C
from keras.optimizers import Adam

RNN_SIZE = 128  # 循环神经网络尺寸
N_CLASS = LabelHandlerCTC.classes  # 类别个数

CONV_SHAPE = (None, 13, 5120)  # 最后一层卷积层输出形状


def ctc_lambda_func(args):
    """
    ctc 回调函数
    :param args:
    :return:
    """
    y_pred, labels, input_length, label_length = args
    return K.ctc_batch_cost(labels, y_pred, input_length, label_length)


def get_acc(base_model, batch_num=64):
    batch_acc = 0
    x_test = np.load(os.path.join(C.test_path, 'inputs.npy'))
    y_test = np.load(os.path.join(C.test_path, 'labels_ctc.npy'))
    y_pred = base_model.predict(x_test, batch_size=batch_num)
    shape = y_pred.shape
    ctc_decode = K.ctc_decode(y_pred, input_length=np.ones(shape[0]) * CONV_SHAPE[1])[0][0]
    out = K.get_value(ctc_decode)
    if out.shape[1] == 4:
        batch_acc += ((y_test == out).sum(axis=1) == 4).mean()
    # batch_acc += ((y_test == out).sum(axis=1) == 4).mean()
    return batch_acc


class Evaluate(ModelCheckpoint):
    def __init__(self, base_model, file_path):
        super(Evaluate, self).__init__(file_path)
        self.accs = []
        self.monitor = 'val_loss'
        self.verbose = 1

        self.monitor_op = np.less
        self.best = np.Inf
        self.base_model = base_model

    def on_epoch_end(self, epoch, logs=None):
        acc = get_acc(self.base_model) * 100
        self.accs.append(acc)
        print('acc: %f%%' % acc)

        current = logs.get(self.monitor)
        if self.monitor_op(current, self.best):
            if self.verbose > 0:
                print('\nEpoch %05d: %s improved from %0.5f to %0.5f,'
                      ' saving model to %s'
                      % (epoch + 1, self.monitor, self.best,
                         current, self.filepath))
            self.best = current
            self.base_model.save(self.filepath, overwrite=True)
        else:
            if self.verbose > 0:
                print('\nEpoch %05d: %s did not improve from %0.5f' %
                      (epoch + 1, self.monitor, self.best))


class CNNCTCModel(AbstractModel):

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

        self.best_weights_path = os.path.join(C.weights_path, 'model02best.h5')
        self.weights_path = os.path.join(C.weights_path, 'model02.h5')

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

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

        self.label_handler = LabelHandlerCTC(C.NUM)

    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)
        shape = conv3_2.get_shape()
        rsp = Reshape((int(shape[1]), int(shape[2] * shape[3])))(conv3_2)

        des = Dense(32)(rsp)

        gru_1 = Bidirectional(GRU(RNN_SIZE, return_sequences=True, kernel_initializer='he_normal'), merge_mode='sum')(
            des)
        gru_2 = Bidirectional(GRU(RNN_SIZE, return_sequences=True, kernel_initializer='he_normal'),
                              merge_mode='concat')(gru_1)

        x = Dropout(0.25)(gru_2)
        x = Dense(N_CLASS, kernel_initializer='he_normal', activation='softmax')(x)
        self.base_model = Model(inputs=inputs, outputs=x)
        # base_model.summary()

        labels = Input(name='the_labels', shape=[C.NUM], dtype='float32')
        input_length = Input(name='input_length', shape=[1], dtype='int64')
        label_length = Input(name='label_length', shape=[1], dtype='int64')
        loss_out = Lambda(ctc_lambda_func, output_shape=(1,),
                          name='ctc')([x, labels, input_length, label_length])

        model = Model(inputs=[inputs, labels, input_length, label_length], outputs=[loss_out])
        model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer=Adam(1e-5))
        model.summary()

        return model

    def train(self, batch_size=256, epochs=500):
        evaluator = Evaluate(self.base_model, self.best_weights_path)
        X, Y = self.x_train, self.y_train
        # save_best = ModelCheckpoint(self.best_weights_path, monitor='val_loss', verbose=1, save_best_only=True)
        tensor_board = TensorBoard(log_dir='log/cnn_ctc')
        callback_list = [tensor_board, evaluator]
        with self.graph.as_default():
            with self.sess.as_default():
                self.model.fit([X, Y, np.array(np.ones(len(X)) * int(CONV_SHAPE[1])), np.array(np.ones(len(X)) * C.NUM)], Y,
                               batch_size=batch_size, epochs=epochs, verbose=1, validation_split=0.1, callbacks=callback_list)
                self.base_model.save(self.weights_path)

    def inference(self, x):
        if len(x.shape) == 3:
            x = np.expand_dims(x, axis=0)
        with self.graph.as_default():
            with self.sess.as_default():
                y_pred = self.base_model.predict(x, batch_size=64)
                out = K.get_value(K.ctc_decode(y_pred, input_length=[y_pred.shape[1]] * y_pred.shape[0], )[0][0][:, :C.NUM])
        return out

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

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

    def test(self):
        x_test, y_test = self.x_test, self.y_test

        print(x_test.shape)
        out = self.inference(x_test)

        error = 0
        for i, m in enumerate(out):
            if not all(out[i] == y_test[i]):
                print("pred", self.label_handler.decode(m), end="\t")
                print("truth", self.label_handler.decode(y_test[i]))
                error += 1
        acc = np.all((out == y_test), axis=1).mean()
        w_acc = (out == y_test).sum() / (y_test.shape[0] * y_test.shape[1])

        print("总数%d 预测错误数%d 正确率%.2f%% " % (len(out), error, acc * 100))
        # print("单词正确率%.2f%%" % (w_acc*100))
        # acc = 1
        # for y_,y in zip(Y_test,out):
        #     cmp = (y_ == y)
        #     if not False in cmp:
        #         acc+=1
        #     else:
        #         print(y_,y)
        # ratio = acc/len(Y_test)
        # print(acc,ratio)
        # out = ''.join([characters[x] for x in out[0]])
        # y_true = ''.join([characters[x] for x in y_test[0]])
        #
        # plt.imshow(X_test[0].transpose(1, 0, 2))
        # plt.title('pred:' + str(out) + '\ntrue: ' + str(y_true))
        #
        # argmax = np.argmax(y_pred, axis=2)[0]
        # list(zip(argmax, ''.join([characters2[x] for x in argmax])))
