#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
    Testing the Gap Statistic to find the k in k-means

    http://datasciencelab.wordpress.com/2013/12/27/finding-the-k-in-k-means-clustering/
"""
from __future__ import division
import sys
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt 


################################################################################ 
#                       ARRAY  MANIPULATION CODE
################################################################################

def norm(a, axis=-1):
    """NumPy 1.8 style norm(). Needed for the NumPy 1.7.1 I am using.
        
        a: A NumPy ndarray
        axis: Axis to calculate norm along. Whole ndarray is normed if axis is None
        
        Returns: norm along axis
    """
    if axis is None:
        return np.linalg.norm(a)
    return np.apply_along_axis(np.linalg.norm, axis, a)


def subtract_outer(a, b):
    """The outer difference of a and b where
        a_b = subtract_outer(a, b) => a_b[i, j, :] = a[i, :] - b[j, :] 

        a: A NumPy ndarray
        b: A NumPy ndarray 
        
        Returns: outer difference of a and b
    """
    assert a.shape[1] == b.shape[1]
    assert len(a.shape) == 2 and len(b.shape) == 2 
    n = a.shape[1]
    a_b = np.empty((a.shape[0], b.shape[0], n))
    for i in xrange(n):
        a_b[:, :, i] = np.subtract.outer(a[:, i], b[:, i])
    return a_b        


################################################################################ 
#                      CODE TO FIND K IN K-MEANS
################################################################################    
    
def find_centers(X, k):
    """Divide the points in X into k clusters
        
        X: points 
        k: number of clusters to separate points into
        
        Returns: mu, labels
            mu: centers of clusters
            labels: indexes of points in X belonging to each cluster
    """
    estimator = KMeans(init='k-means++', n_clusters=k, n_init=10)
    estimator.fit(X)
    mu = estimator.cluster_centers_
    labels = estimator.labels_
    return mu, labels


def Wk(X, mu, labels):
    """Compute the intra-cluster distances for the k clusters in X described by mu and labels.

        X: points
        mu: centers of clusters
        labels: indexes of points in X belonging to each cluster

        Returns: Normalized intra-cluster distance as defined in 
            http://datasciencelab.wordpress.com/2013/12/27/finding-the-k-in-k-means-clustering/
    """
    k = mu.shape[0]
    clusters = [X[np.where(labels == i)] for i in xrange(k)]
    n = [x.shape[0] for x in clusters]
    return sum(norm(clusters[i] - mu[i], None)**2/(2 * n[i]) for i in xrange(k))


def bounding_box(X):
    """Compute the bounding box for the points in X. This is the highest and lowest 
            x and y coordinates of all the points.
    
        X: points
        
        Returns: (xmin, xmax), (ymin, ymax)
            xmin, xmax: min and max of x coordinates of X
            ymin, ymax: min and max of y coordinates of X
    """
    x, y = X[:, 0], X[:, 1]
    return (x.min(), x.max()), (y.min(), y.max()) 


def gap_statistic(X, min_k, max_k, b):
    """Calculate gap statistic for X for k = min_k through k = max_k
        using b reference data sets
    
        X: points
        min_k: lowest k to test
        max_k: highest k to test
        b: number of reference data sets to test against    
 
        Returns: Generator yielding k, logWk, logWkb, sk for min_k <= k <= max_k
            k: This k
            Wks: log(intra-cluster distance) for k
            Wkbs: average reference log(intra-cluster distance) for k
            sk: Normalized std dev log(intra-cluster distance) for k
    """
    N = X.shape[0]
    (xmin, xmax), (ymin, ymax) = bounding_box(X)

    def reference_results(k):
        # Create b reference data sets
        BWkbs = np.zeros(b)
        for i in xrange(b):
            Xb = np.vstack([np.random.uniform(xmin, xmax, N),
                            np.random.uniform(ymin, ymax, N)]).T
            mu, labels = find_centers(Xb, k)
            BWkbs[i] = np.log(Wk(Xb, mu, labels))
        logWkb = np.sum(BWkbs)/b
        sk = np.sqrt(np.sum((BWkbs - logWkb)**2)/b) * np.sqrt(1 + 1/b)
        return logWkb, sk

    for k in xrange(min_k, max_k + 1):
        mu, labels = find_centers(X, k)
        logWk = np.log(Wk(X, mu, labels))
        logWkb, sk = reference_results(k) 
        yield k, logWk, logWkb, sk 


# Parameters for gap statistic calculation
B = 20      # Number of reference data sets
MIN_K = 1   # Lowest k to test
MAX_K = 10  # Highest k to test        

def find_k(X, verbose=1):
    """Find the best k for k-means gap for X using the Gap Statistic
    
        X: points
        verbose: verbosity level: 0, 1 or 2    
 
        Returns: best k if found, otherwise -1
    """
    for i, (k1, logWk1, logWkb1, sk1) in enumerate(gap_statistic(X, MIN_K, MAX_K + 1, B)):
        gap1 = logWkb1 - logWk1
        if i > 0: 
            if verbose >= 2: 
                print('%5d %5.2f %5.2f %5.2f : %5.2f' % (k, logWk, logWkb, sk, gap))
            if gap > gap1 - sk1:
                return k
        k, logWk, logWkb, sk, gap = k1, logWk1, logWkb1, sk1, gap1

    return -1    
    
 
################################################################################ 
#                       TESTING     CODE
################################################################################
 
# GRID_NUMBER is a square number close to GRID_NUMBER_TARGET  
GRID_NUMBER_TARGET = 1000
GRID_WIDTH = int(np.sqrt(GRID_NUMBER_TARGET))
GRID_NUMBER = GRID_WIDTH**2 
 
# UNIFORM_GRID an array of GRID_WIDTH x GRID_WIDTH evenly spaced points on [-1, 1] x [-1, 1]
xv, yv = np.meshgrid(np.linspace(-1, 1, GRID_WIDTH), np.linspace(-1, 1, GRID_WIDTH))
UNIFORM_GRID = np.vstack([xv.ravel(), yv.ravel()]).T

 
def maximally_spaced_points(k, r):
    """Return maximally spaced points in square of radius 1 around origin
        (i.e. square containing x, y such that -1 <= x <= 1, -1 <= y <= 1)
        Try to keep points at least distance r from edges of square 

        k: number of points
        r: desired minimum distance from point to edge of square    
 
        Returns: ndarray of N 2-d points
    """

    if k == 1:
        return np.random.uniform(-min(r, 0.5), min(r, 0.5), size=(k, 2))

    scale = 1.0 - min(r, 0.5)    

    # Start by randomly distributing points over unit radius square
    x0 = np.random.uniform(-1.0, 1.0, size=(k, 2))

    # Maximize minimum distance between centroids
 
    for m in xrange(10):
        changed = False
        for i in xrange(k):
            # Test replacing ith element in x0 all with elements in UNIFORM_GRID to find 
            #  the one that maximizes the minimum distance to elements other than ith in x0 
            # If this minimum distance is greater than current_min then make it the ith element 
            #  in x0    
            x1 = np.vstack((x0[:i, :], x0[i+1:, :]))

            # minimum distance between ith element in x0 and all other elements in x0
            current_min = norm(x1 - x0[i], 1).min()

            # diffs[j] = minimum distance between jth element in UNIFORM_GRID and all elements 
            # in x0 other than ith
            diffs = norm(subtract_outer(UNIFORM_GRID, x1)).min(axis=-1)

            #  max_j_min = index of element in UNIFORM_GRID that maximizes 
            #      minimum distance between ith element in x0 and all other elements in x0
            max_j_min = np.argmax(diffs)

            if diffs[max_j_min] > current_min:
                x0[i] = UNIFORM_GRID[max_j_min]
                changed = True
      
        if not changed and m > 1:
            break

    # Shrink square to get points r-ish distant from edges of unit radius square
    return x0 * scale      


def init_board_gauss(N, k, r):
    """Initialize board of N points with k clusters

        Board is square of radius 1 around origin
        (i.e. square containing x, y such that -1 <= x <= 1, -1 <= y <= 1)
 
        Try to space cluster centers as far apart as possible while keeping them at least distance
        r from edges of unit radius square. This is done in an approximate way by generating 
        random points around maximally spaced nuclei.
  
        N: number of points
        k: number of cluster
        r: desired std dev of points in cluster from cluster center    
 
        Returns: X, centroids, labels 
            X: points
            centroids: centroids of clusters
            labels: cluster index for each point
    """

    def add_cluster(X, j0, j1, cx, cy, s):
        """Add a cluster of normally distributed points to x for indexes [j0,j1) 
            around center cx, cy and std dev s.
            
            X: points