# -*- coding: utf-8 -*-#
'''
# Name:         Level3_ColorAndShape_CNN
# Description:  
# Author:       super
# Date:         2020/6/15
'''

import time
import matplotlib.pyplot as plt

from MiniFramework.NeuralNet_4_2 import *
from ExtendedDataReader.GeometryDataReader import *

train_data_name = "../data/ch17.train_shape_color.npz"
test_data_name = "../data/ch17.test_shape_color.npz"

name = ["red-circle", "red-rect", "red-tri", "green-circle", "green-rect", "green-tri", "blue-circle", "blue-rect",
        "blue-tri", ]


def LoadData(mode):
    print("reading data...")
    dr = GeometryDataReader(train_data_name, test_data_name, mode)
    dr.ReadData()
    dr.NormalizeX()
    dr.NormalizeY(NetType.MultipleClassifier, base=0)
    dr.Shuffle()
    dr.GenerateValidationSet(k=10)
    return dr


def cnn_model():
    num_output = 9
    max_epoch = 20
    batch_size = 16
    learning_rate = 0.1
    params = HyperParameters_4_2(
        learning_rate, max_epoch, batch_size,
        net_type=NetType.MultipleClassifier,
        init_method=InitialMethod.MSRA,
        optimizer_name=OptimizerName.SGD)

    net = NeuralNet_4_2(params, "color_shape_cnn")

    c1 = ConvLayer((3, 28, 28), (8, 3, 3), (1, 1), params)
    net.add_layer(c1, "c1")
    r1 = ActivationLayer(Relu())
    net.add_layer(r1, "relu1")
    p1 = PoolingLayer(c1.output_shape, (2, 2), 2, PoolingTypes.MAX)
    net.add_layer(p1, "p1")

    c2 = ConvLayer(p1.output_shape, (16, 3, 3), (1, 0), params)
    net.add_layer(c2, "c2")
    r2 = ActivationLayer(Relu())
    net.add_layer(r2, "relu2")
    p2 = PoolingLayer(c2.output_shape, (2, 2), 2, PoolingTypes.MAX)
    net.add_layer(p2, "p2")

    params.learning_rate = 0.1

    f3 = FcLayer_2_0(p2.output_size, 32, params)
    net.add_layer(f3, "f3")
    bn3 = BnLayer(f3.output_size)
    net.add_layer(bn3, "bn3")
    r3 = ActivationLayer(Relu())
    net.add_layer(r3, "relu3")

    f4 = FcLayer_2_0(f3.output_size, num_output, params)
    net.add_layer(f4, "f4")
    s4 = ClassificationLayer(Softmax())
    net.add_layer(s4, "s4")

    return net


def show_samples(x, y, title):
    fig, ax = plt.subplots(nrows=6, ncols=6, figsize=(9, 9))
    for i in range(36):
        ax[i // 6, i % 6].imshow(x[i].transpose(1, 2, 0))
        ax[i // 6, i % 6].set_title(name[np.argmax(y[i])])
        ax[i // 6, i % 6].axis('off')
    # endfor
    plt.suptitle(title)
    plt.show()


def show_result(x, y,y_raw, title):
    fig, ax = plt.subplots(nrows=6, ncols=6, figsize=(9, 9))
    for i in range(36):
        ax[i // 6, i % 6].imshow(x[i].transpose(1, 2, 0))
        if np.argmax(y[i]) == np.argmax([y_raw[i]]):
            ax[i // 6, i % 6].set_title(name[np.argmax(y[i])])
        else:
            ax[i // 6, i % 6].set_title(name[np.argmax(y[i])], fontdict={'color':'r'})
        ax[i // 6, i % 6].axis('off')
    # endfor
    plt.suptitle(title)
    plt.show()


def train_cnn():
    dataReader = LoadData("image")
    net = cnn_model()
    x, y = dataReader.GetBatchTrainSamples(36, 0)
    print(x.shape)
    print(y.shape)
    show_samples(x, y, "sample")
    net.train(dataReader, checkpoint=0.5, need_test=True)
    net.ShowLossHistory(XCoordinate.Iteration)
    X_test, Y_test = dataReader.GetTestSet()
    X = dataReader.XTest[0:36]
    Z = net.inference(X_test[0:36])
    Y = Y_test[0:36]
    print(X.shape)
    print(Y.shape)
    show_result(X, Z, Y, "cnn:predication")


if __name__ == '__main__':
    train_cnn()