#!/usr/bin/env python3

import joblib
import numpy as np
from keras import models
import sys
from keras.models import Model

def load_data(filename):
    """
    Loads the data from a saved .npz file.
    ### YOU CAN NOT EDIT THIS FUNCTION ###

    :param filename: string, path to the .npz file storing the data.
    :return: two numpy arrays:
        - x, a Numpy array of shape (n_samples, n_features) with the inputs;
        - y, a Numpy array of shape (n_samples, ) with the targets.
    """
    data = np.load(filename)
    x = data['x']
    y = data['y']

    return x, y


def evaluate_predictions(y_true, y_pred):
    """
    Evaluates the mean squared error between the values in y_true and the values
    in y_pred.
    ### YOU CAN NOT EDIT THIS FUNCTION ###

    :param y_true: Numpy array, the true target values from the test set;
    :param y_pred: Numpy array, the values predicted by your model.
    :return: float, the the mean squared error between the two arrays.
    """
    assert y_true.shape == y_pred.shape
    return ((y_true - y_pred) ** 2).mean()


def load_model(filename):
    """
    Loads a Scikit-learn model saved with joblib.dump.
    This is just an example, you can write your own function to load the model.
    Some examples can be found in src/utils.py.

    :param filename: string, path to the file storing the model.
    :return: the model.
    """
    model = joblib.load(filename)

    return model


if __name__ == '__main__':
    # Load the data
    # This will be replaced with the test data when grading the assignment
    data_path = '../data/data.npz'
    x, y = load_data(data_path)

    ############################################################################
    # EDITABLE SECTION OF THE SCRIPT: if you need to edit the script, do it here
    ############################################################################

    y_pred = None

    X = np.zeros([x.shape[0], 5])
    X[:, 0] = 1
    X[:, 1:3] = x
    X[:, 3] = x[:, 0] * x[:, 1]
    X[:, 4] = np.sin(x[:, 0])

    # Load the linear model
    linear_model_path = './linear_regression.pickle'
    linear_model = load_model(linear_model_path)

    # Predict on the given samples
    y_pred = linear_model.predict(X)
    # Load saved Keras model
    nonlinear_model = models.load_model('./nonlinear_model')
    # Load saved normalizing factors for Xs
    normalizers = load_model('./nonlinear_model_normalizers.pickle')

    print("Linear model:")
    mse = evaluate_predictions(y_pred, y)
    print('MSE: {}'.format(mse))

    # Normalize Xs with normalizing factors
    X = X[:, 1:3]
    X -= normalizers["mean"]
    X /= normalizers["std"]

    # Run model on normalized data
    y_pred = nonlinear_model.predict(X)
    y_pred = y_pred[:, 0]

    print("Non-linear model:")

    ############################################################################
    # STOP EDITABLE SECTION: do not modify anything below this point.
    ############################################################################

    # Evaluate the prediction using MSE
    mse = evaluate_predictions(y_pred, y)
    print('MSE: {}'.format(mse))