#!/usr/bin/python 

""" 
    Skeleton code for k-means clustering mini-project.
"""




import pickle
import numpy
import matplotlib.pyplot as plt
import sys
sys.path.append("../tools/")
from feature_format import featureFormat, targetFeatureSplit
import pprint




def Draw(pred, features, poi, mark_poi=False, name="image.png", f1_name="feature 1", f2_name="feature 2"):
    """ some plotting code designed to help you visualize your clusters """

    ### plot each cluster with a different color--add more colors for
    ### drawing more than five clusters
    colors = ["b", "c", "k", "m", "g"]
    for ii, pp in enumerate(pred):
        plt.scatter(features[ii][0], features[ii][1], color = colors[pred[ii]])

    ### if you like, place red stars over points that are POIs (just for funsies)
    if mark_poi:
        for ii, pp in enumerate(pred):
            if poi[ii]:
                plt.scatter(features[ii][0], features[ii][1], color="r", marker="*")
    plt.xlabel(f1_name)
    plt.ylabel(f2_name)
    plt.savefig(name)
    plt.show()



### load in the dict of dicts containing all the data on each person in the dataset
data_dict = pickle.load( open("../final_project/final_project_dataset.pkl", "r") )
### there's an outlier--remove it! 
data_dict.pop("TOTAL", 0)
# pprint.pprint (data_dict)
# for element in data_dict:
#     print data_dict[element]['exercised_stock_options']
exercised = featureFormat(data_dict, ["exercised_stock_options"])
# print (type(exercised))
# print max(exercised)
# print min(exercised)

salary = featureFormat(data_dict, ["salary"])
# print (type(exercised))
# print max(salary)
# print min(salary)


### the input features we want to use 
### can be any key in the person-level dictionary (salary, director_fees, etc.) 
feature_1 = "salary"
feature_2 = "exercised_stock_options"
# feature_3 = "total_payments"

poi  = "poi"
features_list = [poi, feature_1, feature_2]
# print (features_list)
data = featureFormat(data_dict, features_list )
# print (data)
poi, finance_features = targetFeatureSplit( data )
# print poi
# print (finance_features)

from sklearn.preprocessing import MinMaxScaler

# data = [[-1, 2], [-0.5, 6], [0, 10], [1, 18]]
scaler = MinMaxScaler(feature_range=(0, 1), copy=True)
scaler.fit(finance_features)

# print(scaler.data_max_)
#
print finance_features
finance_features = scaler.transform(finance_features)
print (scaler.transform([200000, 1000000]))


### in the "clustering with 3 features" part of the mini-project,
### you'll want to change this line to 
### for f1, f2, _ in finance_features:
### (as it's currently written, the line below assumes 2 features)
for f1, f2 in finance_features:
    plt.scatter( f1, f2 )
plt.show()

### cluster here; create predictions of the cluster labels
### for the data and store them to a list called pred

# import numpy as np
# from sklearn.cluster import KMeans
# kmeans_model = KMeans(n_clusters=3, random_state=1).fit(X)
# labels = kmeans_model.labels_
# metrics.calinski_harabaz_score(X, labels)
#
#
from sklearn.cluster import KMeans
import numpy as np
# X = np.array([[1, 2], [1, 4], [1, 0],
#               [4, 2], [4, 4], [4, 0]])
kmeans = KMeans(n_clusters=2, random_state=0).fit(finance_features)
# kmeans.labels_
#
pred = kmeans.predict(finance_features)
#
# kmeans.cluster_centers_

### rename the "name" parameter when you change the number of features
### so that the figure gets saved to a different file
try:
    Draw(pred, finance_features, poi,
         mark_poi=False,
         name="clusters.pdf",
         f1_name=feature_1,
         f2_name=feature_2)
         # f3_name=feature_3)
except NameError:
    print "no predictions object named pred found, no clusters to plot"
