#!/usr/bin/env python
__author__ = 'vmaidel'
import pandas as pd
import numpy as np
import json
#from dateutil import parser
#import sys
#import os
from sklearn.cluster import MeanShift, estimate_bandwidth
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from itertools import cycle
from sklearn.cluster import DBSCAN
from sklearn import metrics
from sklearn.preprocessing import StandardScaler

def extractSubject_id(subject_json):
    return subject_json.keys()[0]

def createCenters(row):
    num_of_boxes = row['num_of_boxes']
    centers = []
    for i in range(1,int(num_of_boxes)+1):
        name_of_column = "T1_Box"
        name_of_column = name_of_column+str(i)
        #calculate the coordinates of the clusters
        x_center = float(row[name_of_column+"_x"])+float(row[name_of_column+"_width"])/2
        y_center = float(row[name_of_column+"_y"])+float(row[name_of_column+"_height"])/2
        #append the coordinate tuples into a list of centers for that classification
        centers.append(tuple((x_center,y_center)))
    return centers

def plotBoxes(aggregated_df,df,cluster_centers):
    print "Saving the plots to the output folder..."

    #get the subject_id - there will only be one subject in this dataframe, so just take the first id
    subject = df.subject_id.values[0]

    # Plot just the coordinates
    x=df[['center_x']].values.tolist()
    y=df[['center_y']].values.tolist()
    plt.scatter(x,y)

    plt.savefig("output/scatterplot.svg")

    plt.figure(figsize=(16,12))
    plt.clf()
    cluster_labels = df.cluster_label.values

    good_clusters = np.unique(aggregated_df.cluster_label.values)
  
    #plot each cluster in different color
    ax = plt.figure(figsize=(16,12)).add_subplot(111, aspect='equal')
    colors = cycle('bgrcmybgrcmybgrcmybgrcmy')
    for k, col in zip(range(len(good_clusters)), colors):
        my_members = cluster_labels == k
        cluster_center = cluster_centers[k]
        plt.plot(df.ix[my_members, 'center_x'], df.ix[my_members, 'center_y'], col + '.')
        plt.plot(cluster_center[0], cluster_center[1], 'x', markerfacecolor=col,markeredgecolor='k', markersize=7)
        for index,row in df.ix[my_members,:].iterrows():
              ax.add_patch(
              patches.Rectangle(
                (float(df.ix[index,'tl_x']), float(df.ix[index,'tl_y'])),
                float(df.ix[index,'width']),
                float(df.ix[index,'height']),color=col,
                fill=False      # remove background
                )
              )
              #plot the aggregated boxes
              ax.add_patch(
              patches.Rectangle(
                (float(aggregated_df.ix[aggregated_df.cluster_label==k,'agg_tl_x'].values), float(aggregated_df.ix[aggregated_df.cluster_label==k,'agg_tl_y'].values)),
                float(aggregated_df.ix[aggregated_df.cluster_label==k,'agg_width']),
                float(aggregated_df.ix[aggregated_df.cluster_label==k,'agg_height']),color="black",linewidth=3,
                fill=False      # remove background
                )
              )
    plt.title('Estimated number of clusters: %d' % len(good_clusters))
    plt.savefig("output/scatterplot_meanshift_with_boxes_"+str(subject)+".svg")

def meanShiftClustering(centers_df,subject):
    #estimate the bandwidth to use with the mean shift algorithm. The quantile represents the distance used between the box centers to define the cluster. Smaller quantile, means smaller distance between points that would end up in the same cluster
    centers_df=centers_df.reset_index()
    bandwidth=estimate_bandwidth(centers_df[['center_x','center_y']].as_matrix(), quantile=0.0055)
    #instantiate the mean shift algorithm
    ms = MeanShift(bandwidth=bandwidth, bin_seeding=True)
    #fit the algorithm on the box center coordinates
    ms.fit(centers_df[['center_x','center_y']])
    #get the resulting clustesr labels
    labels = ms.labels_
    #get the resulting centers of each *cluster*
    cluster_centers = ms.cluster_centers_

    labels_unique = np.unique(labels)
    #calculate the number of clusters by using the length of the list that contains all the unique labels
    n_clusters_ = len(labels_unique)

    #concatenate the centers data frame (which contains all the box coordinates, their dimensions, and their centers) with the clustering labels generated by the clustering
    boxes_df = pd.concat([centers_df,pd.DataFrame(labels,columns=['cluster_label'])],axis=1)

    #the aggregate function in the groupby, includes two functions: count and median
    f = {'Number of boxes in a cluster': ['count'],'Median': ['median']}
    #group by the label of each cluster and aggregate the boxes' top left coordinates and dimensions by applying the median
    aggregated_df = boxes_df.groupby('cluster_label')['cluster_label','tl_x','tl_y','width','height'].agg(f).reset_index()
    #change column names for a more descriptive name
    aggregated_df.columns = ['cluster_label','median_cluster_label','agg_tl_x','agg_tl_y','agg_width','agg_height','boxes_in_cluster','count_tl_x','count_tl_y','count_width','count_height']
    #leave out the unnecessary columns
    aggregated_df = aggregated_df[['cluster_label','agg_tl_x','agg_tl_y','agg_width','agg_height','boxes_in_cluster']]
    
    #Look at the output of the plotBoxes function (svg file) and determine at which THRESHOLD value there is a desired number of clusters (appears at the top of the plot) and that it visually matches the actual grid
    THRESHOLD = 5

    #filter out all the clusters that have less than a certain number of boxes in each cluster
    #use the old-weather-aggregator-with-plot.py script to check what the best threshold is
    aggregated_df = aggregated_df.loc[aggregated_df.boxes_in_cluster>THRESHOLD,:]
    good_clusters = np.unique(aggregated_df.cluster_label.values)

    print "for subject_id:"+str(subject)

    print "number of estimated clusters overall: %d" % n_clusters_

    print "number of estimated clusters, after small clusters were filtered out: %d" % len(good_clusters)

    print "clusters with more than %d boxes per cluster:" % THRESHOLD
    print aggregated_df.columns
    print aggregated_df.head()

    #save the aggregated boxes and their clusters into a csv file, separate file for each subject
    print "Saving the output/aggregated_df_%s.csv file..." % str(subject)
    aggregated_df.to_csv("output/aggregated_df_"+str(subject)+".csv",index=False)

    #make sure that only the boxes that belong to the good_clusters (have more boxes than the threshhold) remain in the boxes_df dataframe and then save the dataframe
    boxes_df = boxes_df.loc[boxes_df['cluster_label'].isin(good_clusters),:]
    print "Saving the output/clustered_df_%s.csv file..." % str(subject)
    boxes_df.to_csv("output/clustered_df_"+str(subject)+".csv",index=False)

    plotBoxes(aggregated_df,boxes_df,cluster_centers)

def preprocessing():
    # read the csv files
    print "Reading classifications csv file for old weather..."
    #change the name of the input files if needed:
    classifications_df=pd.read_csv("old-weather-grid-testing-classifications.csv")

    #apply a json.loads function on the whole annotations column
    classifications_df['annotation_json']=classifications_df['annotations'].map(lambda x: json.loads(x))
    #apply a json.loads function on the subject_data column
    classifications_df['subject_json']=classifications_df['subject_data'].map(lambda x: json.loads(x))
    #extract the subject id from the subject json
    classifications_df['subject_id']=classifications_df['subject_json'].apply(extractSubject_id)
    #convert the 'created at' column to a pandas date time format 
    classifications_df['created_at']=pd.to_datetime(classifications_df.created_at)

    #filter out irrelevant old rows
    classifications_df=classifications_df.loc[classifications_df.created_at>=pd.to_datetime('07/25/2016'),:]
    
    #extract the elements from the annotation json
    for index, row in classifications_df.iterrows():
        for i in row['annotation_json']:
            if (type(i['value']) is list) and (i['task']=="T1"):
                #create a column for each box drawn
                box = 0
                for rectangle in i['value']:
                    box+=1
                    classifications_df.loc[index,"T1_Box"+str(box)+"_x"]=str(rectangle['x'])
                    classifications_df.loc[index,"T1_Box"+str(box)+"_y"]=str(rectangle['y'])
                    classifications_df.loc[index,"T1_Box"+str(box)+"_width"]=str(rectangle['width'])
                    classifications_df.loc[index,"T1_Box"+str(box)+"_height"]=str(rectangle['height'])
                classifications_df.loc[index,'num_of_boxes']=box

    #delete the unnecessary columns
    #classifications_df.drop(['annotation_json','subject_id','locations','metadata_json','subject_json'], axis=1, inplace=True)

    #delete classifications that did not contain any boxes
    classifications_df=classifications_df.loc[~classifications_df.num_of_boxes.isnull(),:]

    #calculate box centers and create a new column which contains a list of all centers in that classification
    classifications_df['centers']=classifications_df.iloc[:,17:len(classifications_df.columns)].apply(createCenters,axis=1)
    #save the preprocessed classifications file to a csv
    classifications_df.to_csv('expanded-old-weather.csv',sep=',',index = False,encoding='utf-8')

    #unpack the centers into "one center per row", into a centers_df dataframe
    centers_df=pd.DataFrame(columns=['subject_id','box_location','center_x','center_y','tl_x','tl_y','width','height'])
    for index, row in classifications_df.iterrows():
        for idx, center in enumerate(row['centers']):
            to_append = pd.DataFrame({'subject_id':row['subject_id'],'box_location': ["T1_Box"+str(idx+1)+"_"+str(row['classification_id'])], 'center_x': [center[0]],'center_y':[center[1]],'tl_x':float(classifications_df.ix[index,"T1_Box"+str(idx+1)+"_x"]),'tl_y':float(classifications_df.ix[index,"T1_Box"+str(idx+1)+"_y"]),'width':float(classifications_df.ix[index,"T1_Box"+str(idx+1)+"_width"]),'height':float(classifications_df.ix[index,"T1_Box"+str(idx+1)+"_height"])})
            centers_df = pd.concat([centers_df,to_append],axis=0,ignore_index=True)        
    return centers_df

def clustering(centers_df):
    subjects = np.unique(centers_df[['subject_id']].values.tolist())
    for subject in subjects:
        #limit the dataframe only to the subject we are currently running the clustering on
        this_subject_centers_df=centers_df.loc[centers_df.subject_id==subject,:]
        #apply the mean shift clustering algorithm on the centers, can try other algorithms, just make sure that the algorithm does not require the number of clusters as a paramter
        meanShiftClustering(this_subject_centers_df,subject)
        
centers_df = preprocessing()
#perform clustering
clustering(centers_df)
