#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Generador de letras con ruido para entrenar autoencoder
# Neurocomputación - Practica 3
# Jorge Arellano y Pablo Marcos

import argparse
import numpy as np
from neuro import Autoencoder

def parse_data(file_train, output):
    """Lee los ficheros de datos de entrada del programa."""

    with open(file_train) as train_file:
        # Leemos cabecera
        n, m = [int(a) for a in train_file.readline().split()]

    data = np.loadtxt(file_train, skiprows=1)

    X = data[:,:n]
    y = data[:,-m:]

    n_ejemplos = len(X)

    # Caso fichero de test indicado como porcentaje
    if output.isdigit():
        percentage = int(output)

        if percentage < 100:

            corte = (n_ejemplos * percentage) // 100
            idx = np.arange(n_ejemplos)
            np.random.shuffle(idx)
            idx_train = idx[:corte]
            idx_test = idx[corte:]

            X_train = X[idx_train]
            y_train = y[idx_train]
            X_test = X[idx_test]
            y_test = y[idx_test]
        else:

            X_train = X_test = X
            y_train = y_test = y

    else:
        X_train = X
        y_train = y

        data_test = np.loadtxt(output, skiprows=1)

        #data_test = np.loadtxt(output, skiprows=1)
        X_test = data_test[:,:n]
        y_test = data_test[:,-m:]


    return X_train, y_train, X_test, y_test


if __name__ == '__main__':

    # Parseador del script
    parser = argparse.ArgumentParser(description='Entrenamiento de autoencoder')
    parser.add_argument('train', type=str, nargs=1,
                        help='Fichero de entrada autoencoder')

    parser.add_argument('test', type=str, nargs='?', default="100",
                        help='Archivo de salida / porcentage')
    parser.add_argument('-o', '--output', type=str, nargs=1, default=[None],
                       help="Fichero de salida")
    parser.add_argument('-e', '--epoch', type=int, nargs=1, default=[100],
                       help="Maximas épocas de entrenamiento")
    parser.add_argument('--learning', type=float, nargs=1, default=[.1],
                       help="Tasa de aprendizaje")
    parser.add_argument('--ecm', type=float, nargs=1, default=[0],
                       help="Condición de parada con error cuadrático medio")
    parser.add_argument('-l', '--layers', type=int, nargs='+', default=[10],
                       help="Capas ocultas")
    parser.add_argument('--n_rondas', type=int, nargs='?',
                       default=1,
                       help='Numero de letras')
    parser.add_argument('--plot', action='store_true',
                        help="Dibuja estadisticas")


    args = parser.parse_args()

    X_train, y_train, X_test, y_test = parse_data(args.train[0], args.test)

    print("Datos de entrenamiento:")
    print("\t", "Fichero", args.train[0])
    print("\t", "Maximo numero de épocas", args.epoch[0])
    print("\t", "Constante de aprendizaje", args.learning[0])
    print("\t", "Criterio de parada ecm < ", args.ecm[0])
    print("\t", "Capas ocultas", args.layers)
    print("\t", "Numero de datos de entrenamiento", len(X_train))
    print("\t", "Numero de datos de test", len(X_test))
    print(40 * "*")

    ae = Autoencoder(capas=[100])
    epoca, ecm, pe, mpe, lrc = ae.fit(X_train, y_train, epoch=args.epoch[0],
                                      error=args.ecm[0],
                                      learn_rate=args.learning[0])

    res = ae.evaluar(X_train)

    for _ in range(args.n_rondas - 1):
        res = ae.evaluar(res)

    diff = np.abs(res - y_train).sum(axis=1)
    print(" *** TRAIN ***")
    print("Numero de pixeles errados sobre datos de train: ", diff.sum())
    print("Numero de pixeles promedio errados sobre datos de train",
          diff.sum()/len(X_train))
    print("Numero de letras de train recuperadas", (diff == 0).sum())

    res = ae.evaluar(X_test)

    for _ in range(args.n_rondas - 1):
        res = ae.evaluar(res)

    diff = np.abs(res - y_test).sum(axis=1)
    print(" *** TEST ***")
    print("Numero de pixeles errados sobre datos de test: ", diff.sum())
    print("Numero de pixeles promedio errados sobre datos de test",
          diff.sum()/len(X_test))
    print("Numero de letras de test recuperadas", (diff == 0).sum())


    if args.plot:
        # Estadisticas por epoca
        import matplotlib.pyplot as plt

        plt.style.use('seaborn')

        plt.figure(1)
        plt.title("Error cuadrático medio")
        plt.xlabel("Época")
        plt.ylabel("Error cuadrático medio")
        plt.plot(ecm)

        plt.figure(2)
        plt.title("Pixeles Errados")
        plt.xlabel("Época")
        plt.ylabel("PE")
        plt.plot(pe)

        plt.figure(3)
        plt.title("Pixeles Errados Promedio")
        plt.xlabel("Época")
        plt.ylabel("MPE")
        plt.plot(mpe)

        plt.figure(4)
        plt.title("Letras Recuperadas")
        plt.xlabel("Época")
        plt.ylabel("LRC")
        plt.plot(lrc)

        fig, ax = plt.subplots(3, 6)
        ax[0, 0].set_ylabel("Letra Original")
        ax[1, 0].set_ylabel("Letra Ruidosa")
        ax[2, 0].set_ylabel("Letra Recuperada")

        idx = np.arange(len(X_test))
        np.random.shuffle(idx)

        for i in range(6):

            ax[0, i].imshow(y_test[idx[i]].reshape((7,5)), cmap="Greys")
            ax[1, i].imshow(X_test[idx[i]].reshape((7,5)), cmap="Greys")
            ax[2, i].imshow(res[idx[i]].reshape((7,5)), cmap="Greys")


            ax[0, i].set_xticks([])
            ax[1, i].set_yticks([])
            ax[0, i].set_yticks([])
            ax[1, i].set_xticks([])
            ax[2, i].set_yticks([])
            ax[2, i].set_xticks([])

        plt.tight_layout()

        plt.show()
