#!/bin/python

### Same Size Clustering ###
###
### This is a variation of the k-means clustering that produces equally sized clusters.
### The algorithm consists of two phases:
### 1. Initialization:
###     Compute the desired cluster size: n/k
###     Initialize means with k-means
###     Order points by the biggest benefit (delta distance) of best over worst assignment
###     Assign points to their prefered cluster until the cluster is full, then resort remaining
###     objects by excluding the clusters that are full.
### 2. Refinement of clustering:
###     This is done in an iterative fashion, until there is no change in clustering or the max
###     number of iterations have been reached.
###     Interation:
###         Compute current cluster means.
###         For each object, compute the distance to the cluster means.
###         Sort elements based on the delta of the current assignment and the alternative best 
###         possible assignment, sort in descending order.
###         For each element by priority:
###             For each cluster to whom it doesn't belong, by element gain:
###                 if there is an element on the transfer list of the cluster, and swapping the 
###                 two element yields improvement, swap the two elements;
###             End For
###             If the element is not swapped, add to outgoing transfer list
###         End For
###         If no transfers were done, or max iteration was reached, terminate.
###     Since any transfer must decrease variance, thus the clustering will converge.
### 
### Author: Xin Cindy Yan
### Credit: The algorithm is adapted from the java script by Erich Schubert as on the ELKI webpage:
###         http://elki.dbs.ifi.lmu.de/wiki/Releases
### Date:   03/27/2015

import sys
import numpy as np
import random
#import mykmns
#import readdata
import extract
from sklearn import cluster
from scipy.spatial import distance_matrix

#------------Input---------------
# n = 2000 # total number of points
# cluster_size = 4 # size of each cluster
# k = (n+cluster_size-1) // cluster_size # total number of clusters
infl = 't4p2k.pdb' # the input pdb file with all dataset
outfl = 'cluster2k.pdb' # the output pdb file with the clustering information
boxl = np.array([39.0571])#, 39.0571])
max_iter = 100


def kmeans_missing(X, n_clusters, max_iter=500):
    """Perform K-Means clustering on data with missing values.

    Args:
      X: An [n_samples, n_features] array of data to cluster.
      n_clusters: Number of clusters to form.
      max_iter: Maximum number of EM iterations to perform.

    Returns:
      labels: An [n_samples] vector of integer labels.
      centroids: An [n_clusters, n_features] array of cluster centroids.
      X_hat: Copy of X with the missing values filled in.
    """

    # Initialize missing values to their column means
    missing = ~np.isfinite(X)
    mu = np.nanmean(X, 0)
    X_hat = np.where(missing, mu, X)

    for i in xrange(max_iter):
        if i > 0:
            # initialize KMeans with the previous set of centroids. this is much
            # faster and makes it easier to check convergence (since labels
            # won't be permuted on every iteration), but might be more prone to
            # getting stuck in local minima.
            cls = cluster.KMeans(n_clusters, init=prev_centroids)
        else:
            # do multiple random initializations in parallel
            cls = cluster.KMeans(n_clusters, n_jobs=-1)

        # perform clustering on the filled-in data
        labels = cls.fit_predict(X_hat)
        centroids = cls.cluster_centers_
        distance = cls.fit_transform(X_hat)

        # fill in the missing values based on their cluster centroids
        X_hat[missing] = centroids[labels][missing]

        # when the labels have stopped changing then we have converged
        if i > 0 and np.all(labels == prev_labels):
            break

        prev_labels = labels
        prev_centroids = cls.cluster_centers_
        prev_distance = cls.inertia_

        
    return distance, labels, centroids, X_hat


def filter_by(df, constraints):
    """Filter MultiIndex by sublevels."""
    indexer = [constraints[name] if name in constraints else slice(None)
               for name in df.index.names]
    return df.loc[tuple(indexer)] if len(df.shape) == 1 else df.loc[tuple(indexer),]


# Choose initial means with K-means method
def ChooseInitialMeans(X, k):
    #kmeans_missing(X, n_clusters, max_iter=500):
    distance, labels, centroids, X_hat = kmeans_missing(X, k, max_iter=500)
    return centroids

# Returns the distance between an object x and a cluster mean c
def GetDist(x, c):
    dist = np.linalg.norm(x-c-boxl*np.around((x-c)/boxl))
    return dist

# Returns the distance matrix of each object to each cluster mean,
# and each row (i'th obj) of the matrix is ordered so that the dist are in ascending order
def SetDistMat(X, n, k, means):
    distmat_dtype = [('key',int), ('dist',float)]
    distmat = np.empty((n,k),dtype=distmat_dtype)
    for i in range(n):
        print X[i]
        distmat[i,:] = [(c[0], GetDist(X[i], c[1])) for c in enumerate(means)]
        distmat[i,:] = np.sort(distmat[i,:], order='dist')
    return distmat


# Returns the priority list in which objects are ordered by ascending benefit 
# of best over worst assignment
def Get_plst(assigned, distmat, full, n, k):
    plst = []
    for i in range(n):
        if (i not in assigned):
            j = 0
            while j<k:
                if (not full[distmat[i,j][0]]):
                    bestkey = distmat[i,j][0]
                    mindist = distmat[i,j][1]
                    break
                else:
                    j += 1
            for j in range(k-1,-1,-1):
                if (not full[distmat[i,j][0]]):
                    maxdist = distmat[i,j][1]
                    break
            plst.append((i, bestkey, maxdist-mindist))
    plst.sort(key=lambda t:t[2])
    return plst

def InitialAssignment(distmat, cluster_size, n, k):
    clusters = {}
    full = np.zeros(k,dtype=bool) # a boolean array that records which clusters are full
    assigned = [] # a list of objects who has been assigned to a cluster
    plst = Get_plst(assigned, distmat, full, n, k)
    while (len(plst)):
        temp = plst.pop()
        try:
            if (len(clusters[temp[1]])<cluster_size):
                clusters[temp[1]].append(temp[0])
                assigned.append(temp[0])
            else:
                full[temp[1]] = True
                plst = Get_plst(assigned, distmat, full, n, k)
        except KeyError:
            clusters[temp[1]] = [temp[0]]
            assigned.append(temp[0])
    return clusters

def CalcMeans(X, oldmeans, clusters, k):
    means = np.zeros((k,3)) # should be three originally
    keys = sorted(clusters.keys())
    for key in keys:
        for i in clusters[key]:
            means[key] += X[i]-boxl*np.around((X[i]-oldmeans[key])/boxl)
        means[key] /= len(clusters[key])
        means[key] -= boxl*np.around(means[key]/boxl)
    return means

def SortObj(X, clusters, means, distmat):
    objlst = [] # list of objects ordered in asceding delta of the current 
                # assignment and the best possible alternate assignment
    keys = sorted(clusters.keys())
    for key in keys:
        for i in clusters[key]:
            currdist = GetDist(X[i],means[key])
            mindist = distmat[i,0][1]
            objlst.append((i, key, currdist-mindist))
    objlst.sort(key=lambda t:t[2], reverse=True)
    return objlst

"""Mostly useless in comparison to the rest of the code, but important"""
def Transfer(obj, clufrom, cluto, clusters):
    clusters[clufrom].remove(obj)
    clusters[cluto].append(obj)
    return clusters

def run_equal_size_kmeans(X, n, k, cluster_size):
    # Set up the database of objects
    #X = readdata.read_data(infl)
    # Choose initial means with K-means
    means = ChooseInitialMeans(X, k)
    # Set up initial clusters
    distmat = SetDistMat(X, n, k, means) 
    print distmat
    clusters = InitialAssignment(distmat, n, k, cluster_size)
    ## debug code
    #keys = sorted(clusters.keys())
    #for key in keys:
    #    print("cluster %i:"%key)
    #    print(clusters[key])
    ## end of debug
    # Iteration step
    for iter in range(max_iter):
        active = 0 # indicate the number of transfers in the current iteration
        tranlst = (-1)*np.ones(k, dtype='int') # set up transfer list for each cluster
        # Compute the cluster means
        oldmeans = means.copy()
        means = CalcMeans(X, oldmeans, clusters)
        # Get statistics about the clustering
        #ClusterStat(X, means, clusters)
        ## debug code
        #print("old means:")
        #print(oldmeans)
        #print("new means:")
        #print(means)
        ## end of debug
        # For each object, compute the distances to the cluster means
        distmat = SetDistMat(X, n, k, means)
        # Sort objects based on the delta of the current assignment and the best 
        # possible alternate assignment
        objlst = SortObj(X, clusters, means, distmat)
        ##debug code
        #print(objlst)
        ##return
        #end of debug
        # For each element by prioty:
        while (len(objlst)):
            (i, key, temp) = objlst.pop()
            obj2key = GetDist(X[i], means[key])
            transferred = False #record if any transfering has occured to i 
            if (key == distmat[i,0][0]):
                ##debug
                #print("%i is already the opt cluster for obj %i. no transfer"%(clu, i))
                ##end of debug
                continue
            # For each other clusters by element gain:
            else:
                for j in range(k):
                    clu = distmat[i,j][0] # the key of another cluster
                    objgain = obj2key - distmat[i,j][1] # gain by transfering i from cluster key to clu
                    if (clu==key): # already in the cluster
                        continue
                    if (len(clusters[clu]) < cluster_size):
                        active += 1
                        transferred = True
                        clusters = Transfer(i, key, clu, clusters)
                        ##debug
                        #print("cluster %i not full. transfer obj %i from cluster %i to it."%(clu, i, key))
                        ##end of debug
                        break
                    elif (tranlst[clu] != -1): # if the tranlst of another cluster is not empty
                        # distance between the obj in the tranlst and the current cluster
                        tran2key = GetDist(X[tranlst[clu]], means[key])
                        tran2clu = GetDist(X[tranlst[clu]], means[clu])
                        # gain by transfering the obj in tranlst from cluster clu to key
                        trangain = tran2clu - tran2key
                        if (objgain + trangain > 0): # transfer if the sum of gains are positive, ie net gain
                            active += 2
                            transferred = True
                            clusters = Transfer(i, key, clu, clusters)
                            clusters = Transfer(tranlst[clu], clu, key, clusters)
                            ##debug
                            #print("obj %i is transfered from cluster %i to %i"%(i, key, clu))
                            #print("obj %i is transfered from cluster %i to %i"%(tranlst[clu], clu, key))
                            #print("objgain: %f, trangain: %f"%(objgain, trangain))
                            ##end of debug
                            tranlst[clu] = -1 # reset the tranlst to empty
                            break
                if (not transferred):
                    tranlst[key] = i
                    ##debug
                    #print("add obj %i in cluster %i to the transfer list"%(i, key))
                    ##end of debug
        # nothing is transferred during this iteration, return the clustering result
        if (not active):
                break
        #debug code
        print("number of transfers in iter %i: %i\n"%(iter+1, active))
        #end of debug
    print("K-means clustering converged in %d iterations!\n"%(iter+1))
    # Output the clustering results
    WriteResult(outfl, X, means, clusters, k)
    ClusterStat(X, means, clusters, k)
