#!/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 MLRegressor

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 regresor de series temporales')
    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 test / 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=[.01],
                       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('--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 * "*")

    mlr = MLRegressor(capas=[100])
    epoca, ecm = mlr.fit(X_train, y_train, epoch=args.epoch[0],
                                      error=args.ecm[0],
                                      learn_rate=args.learning[0])

    res_train = mlr.evaluar(X_train)
    res_test = mlr.evaluar(X_test)

    ecm_train = np.square(res_train - y_train).mean()
    ecm_test = np.square(res_test - y_test).mean()

    # Prediccion recursiva sobre datos de X_test

    x_recursivo = X_train[-1]

    prediccion_recursiva = np.empty(len(y_test))
    ecm_test = np.square(res_test - y_test).mean()

    for i in range(len(prediccion_recursiva)):
        pred = mlr.evaluar(np.array([x_recursivo]))
        prediccion_recursiva[i] = pred[0,0]
        x_recursivo = np.roll(x_recursivo, -1)
        x_recursivo[-1] = prediccion_recursiva[i]

    #ecm_recursivo = np.square(res_test - y_test).mean()


    print("Error cuadrático medio sobre datos de entrenamiento", ecm_train)
    print("Error cuadrático medio sobre datos de test", ecm_test)

    print("||y_train - ML(X_train)||_inf = ", np.abs(res_train - y_train).max())
    print("||y_test - ML(X_test)||_inf = ", np.abs(res_test - y_test).max())


    if args.plot:
        import matplotlib.pyplot as plt
        plt.plot(ecm)
        plt.title("Error cuadrático medio en entrenamiento")
        plt.xlabel("Época")
        plt.ylabel("ECM")

        for res, name in zip([res_train, res_test], ["train", "test"]):
            plt.figure()
            plt.title("Predicción sobre datos de " + name)

            if res.shape[1] == 1:
                plt.scatter(np.arange(len(res)), res.flatten())
            else:
                l = res.shape[1]
                a = np.arange(1,len(res) + l)
                for i in range(len(res)):
                    plt.plot(a[i:i+l], res[i])

        plt.figure()
        plt.title("Predicción recursiva")
        plt.plot(prediccion_recursiva)
        plt.plot(y_test[:,0], linestyle='--')
        plt.legend(['predicción recursiva', 'serie real'])


        plt.show()
