import random
import numpy
import matplotlib.pyplot as plt
import pickle

from outlier_cleaner import outlierCleaner


# Load up some practice data with outliers in it
ages = pickle.load( open("practice_outliers_ages_UNIX.pkl", "rb") )
net_worths = pickle.load( open("practice_outliers_net_worths_UNIX.pkl", "rb") )


# Reshape ages and net_worths into 2D numpy arrays
ages       = numpy.reshape( numpy.array(ages), (len(ages), 1))
net_worths = numpy.reshape( numpy.array(net_worths), (len(net_worths), 1))


# Split the data for train and test sets
from sklearn.model_selection import train_test_split
ages_train, ages_test, net_worths_train, net_worths_test = train_test_split(ages, net_worths, test_size=0.1, random_state=42)


### Fit the data using Linear Regression
from sklearn import linear_model

reg = linear_model.LinearRegression()
reg.fit(ages_train, net_worths_train)

print("\nTraining Set Slope: ", reg.coef_)
print("Training Set Score", reg.score(ages_test, net_worths_test))

try:
    plt.plot(ages, reg.predict(ages), color="blue")
except NameError:
    pass
plt.scatter(ages, net_worths)
plt.show()


# Identify and remove the most outlier-y points
cleaned_data = []

predictions = reg.predict(ages_train)
cleaned_data = outlierCleaner( predictions, ages_train, net_worths_train )
# except NameError:
#     print("your regression object doesn't exist, or isn't name reg")
#     print("can't make predictions to use in identifying outliers")


# Only run this code if cleaned_data is returning data
if len(cleaned_data) > 0:
    ages, net_worths, errors = zip(*cleaned_data)
    ages       = numpy.reshape( numpy.array(ages), (len(ages), 1))
    net_worths = numpy.reshape( numpy.array(net_worths), (len(net_worths), 1))

    # Refit the cleaned data
    try:
        reg.fit(ages, net_worths)
        plt.plot(ages, reg.predict(ages), color="blue")
    except NameError:
        print("you don't seem to have regression imported/created,")
        print("   or else your regression object isn't named reg")
        print("   either way, only draw the scatter plot of the cleaned data")
    plt.scatter(ages, net_worths)
    plt.xlabel("ages")
    plt.ylabel("net worths")
    plt.show()
    
    print("\nCleaned Set Slope: ", reg.coef_)
    print("Cleaned Set Score", reg.score(ages_test, net_worths_test))

else:
    print("outlierCleaner() is returning an empty list, no refitting to be done")

