import numpy as np
import pandas as pd
from sklearn.metrics import r2_score
import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = (12, 7)

def printModelEvaluations(model, x, y, x_test, y_test, plot=True):
    print(f'Best model: {model}')
    
    y_pred, rmse = predict_error(model,x,y)
    print('Training error:')
    print(f'RMSE: {rmse}')
    r2 = r2_score(y,y_pred)
    print(f'Rsq: {r2}')  

    y_pred_test, rmse_test = predict_error(model,x_test,y_test)
    print('Test error:')
    print(f'RMSE: {rmse_test}')
    r2_test = r2_score(y_test,y_pred_test)
    print(f'Rsq: {r2_test}')  

    y = np.squeeze(y)
    # y = np.squeeze(y.values)
    y_test = np.squeeze(y_test)
    # y_test = np.squeeze(y_test.values)
    x = x
    # x = x.values
    if plot:

        plt.subplot(2,2,1)
        chartResidualVsYsFor(y_pred, y)
        plt.subplot(2,2,2)
        chartPredictedVsActual(y_pred, y)
        # chartPredictionAndActualVsX(x,y,y_pred)

        plt.subplot(2,2,3)
        chartResidualVsYsFor(y_pred_test, y_test, type = 'validation')
        plt.subplot(2,2,4)
        chartPredictedVsActual(y_pred_test, y_test, type = 'validation')
        # chartPredictionAndActualVsX(x_test,y_test,y_pred_test, type = 'validation')

        plt.suptitle(f'Train/Valid RMSE: {rmse:0.2f}/{rmse_test:0.2f}. Rsq: {r2:0.2f}/{r2_test:0.2f}. N: {len(y)}/{len(y_test)}')

    return(model)

def chartResidualVsYsFor(y_pred, y, type = 'Train'):
    plt.scatter(y_pred, y_pred - y)
    plt.hlines(y = 0, xmin = min(y_pred), xmax = max(y_pred))
    plt.title(f'Residual Errors ({type})')

def chartPredictedVsActual(y_pred, y, type = 'Train'):
    plt.scatter(y_pred, y)
    plt.title(f'Predicted vs Actual ({type})')

def chartPredictionAndActualVsX(x,y,y_pred,type = 'Train'):
    plt.scatter(x,y)
    lim = plt.axis()
    plt.scatter(x,y_pred, c = 'red')
    plt.title(f'Prediction versus actual ({type})')


def predict_error(model,x,y):
    y_pred = model.predict(x)
    #rmse = np.sqrt(np.sum((y_pred - y)**2))[0]
    rmse = np.sqrt(np.sum((y_pred - y)**2))
    return (y_pred, rmse)