# coding: utf-8
import pathlib

import cv2
import matplotlib.pyplot as plt
import numpy as np
from keras import backend as K
from keras import layers
from keras import models
from keras.callbacks import ReduceLROnPlateau
from keras.utils import to_categorical


def load_data(fn='texts.npz', to=False):
    data = np.load(fn)
    texts, labels = data['texts'], data['labels']
    texts = texts / 255.0
    _, h, w = texts.shape
    texts.shape = (-1, h, w, 1)
    if to:
        labels = to_categorical(labels)
    n = int(texts.shape[0] * 0.9)   # 90%用于训练，10%用于测试
    return (texts[:n], labels[:n]), (texts[n:], labels[n:])


def savefig(history, fn='loss.jpg', start=2):
    # 忽略起点
    loss = history.history['loss'][start - 1:]
    val_loss = history.history['val_loss'][start - 1:]
    epochs = list(range(start, len(loss) + start))
    plt.plot(epochs, loss, 'bo', label='Training loss')
    plt.plot(epochs, val_loss, 'b', label='Validation loss')
    plt.title('Training and validation loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    plt.savefig(fn)


def main():
    (train_x, train_y), (test_x, test_y) = load_data()
    model = models.Sequential([
        layers.Conv2D(64, (3, 3), padding='same', activation='relu', input_shape=(None, None, 1)),
        layers.MaxPooling2D(),  # 19 -> 9
        layers.Conv2D(64, (3, 3), padding='same', activation='relu'),
        layers.MaxPooling2D(),  # 9 -> 4
        layers.Conv2D(64, (3, 3), padding='same', activation='relu'),
        layers.MaxPooling2D(),  # 4 -> 2
        layers.GlobalAveragePooling2D(),
        layers.Dropout(0.25),
        layers.Dense(64, activation='relu'),
        layers.Dense(80, activation='softmax'),
    ])
    model.summary()
    model.compile(optimizer='rmsprop',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    # 当标准评估停止提升时，降低学习速率
    reduce_lr = ReduceLROnPlateau(verbose=1)
    history = model.fit(train_x, train_y, epochs=100,
                        validation_data=(test_x, test_y),
                        callbacks=[reduce_lr])
    savefig(history, start=10)
    model.save('model.v1.0.h5', include_optimizer=False)


def load_data_v2():
    (train_x, train_y), (test_x, test_y) = load_data(to=True)
    # 这里是统计学数据
    (train_v2_x, train_v2_y), (test_v2_x, test_v2_y) = load_data('texts.v2.npz')
    # 合并
    train_x = np.concatenate((train_x, train_v2_x))
    train_y = np.concatenate((train_y, train_v2_y))
    test_x = np.concatenate((test_x, test_v2_x))
    test_y = np.concatenate((test_y, test_v2_y))
    return (train_x, train_y), (test_x, test_y)


def acc(y_true, y_pred):
    return K.cast(K.equal(K.argmax(y_true + y_pred, axis=-1),
                          K.argmax(y_pred, axis=-1)),
                  K.floatx())


def main_v19():     # 1.9
    (train_x, train_y), (test_x, test_y) = load_data_v2()
    model = models.load_model('model.v1.0.h5')
    model.compile(optimizer='RMSprop',
                  loss='categorical_hinge',
                  metrics=[acc])
    reduce_lr = ReduceLROnPlateau(verbose=1)
    history = model.fit(train_x, train_y, epochs=100,
                        validation_data=(test_x, test_y),
                        callbacks=[reduce_lr])
    savefig(history)
    model.save('model.v1.9.h5', include_optimizer=False)


def main_v20():
    (train_x, train_y), (test_x, test_y) = load_data()
    model = models.Sequential([
        layers.Conv2D(64, (3, 3), activation='relu', padding='same', input_shape=(None, None, 1)),
        layers.MaxPooling2D(),  # 19 -> 9
        layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
        layers.MaxPooling2D(),  # 9 -> 4
        layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
        layers.MaxPooling2D(),  # 4 -> 2
        layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
        layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
        layers.MaxPooling2D(),  # 2 -> 1
        layers.GlobalAveragePooling2D(),
        layers.Dropout(0.25),
        layers.Dense(64, activation='relu'),
        layers.Dense(80, activation='softmax'),
    ])
    model.summary()
    model.compile(optimizer='rmsprop',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    model.fit(train_x, train_y, epochs=10,
              validation_data=(test_x, test_y))
    (train_x, train_y), (test_x, test_y) = load_data_v2()
    model.compile(optimizer='rmsprop',
                  loss='categorical_hinge',
                  metrics=[acc])
    reduce_lr = ReduceLROnPlateau(verbose=1)
    history = model.fit(train_x, train_y, epochs=100,
                        validation_data=(test_x, test_y),
                        callbacks=[reduce_lr])
    savefig(history)
    # 保存，并扔掉优化器
    model.save('model.v2.0.h5', include_optimizer=False)


def predict(texts):
    model = models.load_model('model.h5')
    texts = texts / 255.0
    _, h, w = texts.shape
    texts.shape = (-1, h, w, 1)
    labels = model.predict(texts)
    return labels


def _predict():
    texts = np.load('data.npy')
    labels = predict(texts)
    np.save('labels.npy', labels)


def show():
    texts = np.load('data.npy')
    labels = np.load('labels.npy')
    labels = labels.argmax(axis=1)
    pathlib.Path('classify').mkdir(exist_ok=True)
    for idx, (text, label) in enumerate(zip(texts, labels)):
        # 使用聚类结果命名
        fn = f'classify/{label}.{idx}.jpg'
        cv2.imwrite(fn, text)


if __name__ == '__main__':
    main()
    # main_v2()
    _predict()
    show()