from sklearn.cluster import KMeans
import numpy as np
import pandas as pd

class Clustering(object):

    '''

    IN: tdidf, df, n = trained tdidf model, urls (list), 
                titles (list) #clusters (int)

    WRITE OUT: 3 txt files with clustering results
    OUT: None
    A class for performing clustering on the provided 
    urls in the using a given tdidf model.

    '''

    def __init__(self, tdidf, urls, titles, n):
        self.tdidf=tdidf

        self.urls=urls
        self.urls=[self._ret_str(a) for a in self.urls]

        self.titles=titles

        self.n=n
        
    def _ret_str(self, obj):
        
        '''
        IN: either list or string
        OUT: string

        Ensures each item of list is a string 
        and not a list. 
        '''

        if type(obj)==list:
            return obj[0]
        else:
            return obj

    def run_clustering(self):
        
        '''runs Kmeans model with data contained
        in class instance'''

        print 'running KMeans with cluster #:', self.n

        Xs=self.tdidf.transform(self.urls)
        self.X=Xs.toarray()

        self.kmeansmodel=KMeans(n_clusters=self.n, 
            max_iter=300, n_jobs=-1, init='k-means++')

        self.kmeansmodel.fit_predict(self.X)
        
    def centroids(self):

        '''Print out the centroids.'''

        print "cluster centers:"
        print self.kmeansmodel.cluster_centers_

        self.features=self.kmeansmodel.get_feature_names()
        
    def top_feats_eclust(self):

        '''Find the top 10 features for each cluster and
            print them.'''

        top_centroids = self.kmeansmodel.cluster_centers_.argsort()[:,-1:-11:-1]

        print "top features for each cluster:"
        for num, centroid in enumerate(top_centroids):
            print "%d: %s" % (num, ", ".join(features[i] for i in centroid))


    def url_sample(self,n_per_cl=20):

        #Print out the titles of a random sample of the articles assigned to each
        # cluster to get a sense of the topic.

        assigned_cluster = self.kmeansmodel.transform(self.X).argmin(axis=1)

        f=open('cluster_output/clusteroutput'+str(self.n)+'.txt', 'wb')
        for i in range(self.kmeansmodel.n_clusters):
            
            cluster = np.arange(0, self.X.shape[0])[assigned_cluster==i]
            sample_urls = np.random.choice(cluster, n_per_cl, replace=True)
            
            f.write('<textarea rows="23" cols="100">')
            f.write('                              ++++++++++++++++++++++CLUSTER ' + str(i+1) + \
                '++++++++++++++++++++++ '+'\n\n\n')

            #iterate through indexes corresponding to URLs of different 
            #clusters
            for url_n in sample_urls:
                
                txttit=self.titles[url_n]
                txttit=''.join([a for a in txttit if (ord(a) < 127) & (ord(a) > 0)])
                f.write( txttit)
                
                f.write('\n\n')
                    
            f.write('</textarea>') 
            f.write('\n')
        f.close()
