"""
k-means algorithm.
"""

## This code is written by Davide Albanese, <albanese@fbk.eu>
## (C) 2010 Fondazione Bruno Kessler - Via Santa Croce 77, 38100 Trento, ITALY.

## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.

## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.

## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.


__all__= ['Kmeans']


import numpy as np
import kmeanscore


class Kmeans(object):
    """k-means algorithm.
    """

    def __init__(self, k, init="std", seed=0):
        """Initialization.
        
        :Parameters:
   
          k : int (>1)
              number of clusters
          init : string ('std', 'plus')
                 initialization algorithm
                   * 'std' : randomly selected
                   * 'plus' : k-means++ algorithm
          seed : int (>=0)
                 random seed
        
        Example:

        >>> import numpy as np
        >>> import mlpy
        >>> x = np.array([[ 1. ,  1.5],
        ...               [ 1.1,  1.8],
        ...               [ 2. ,  2.8],
        ...               [ 3.2,  3.1],
        ...               [ 3.4,  3.2]])
        >>> kmeans = mlpy.Kmeans(k=3, init="plus", seed=0)
        >>> kmeans.compute(x)
        array([1, 1, 2, 0, 0], dtype=int32)
        >>> kmeans.means
        array([[ 3.3 ,  3.15],
        [ 1.05,  1.65],
        [ 2.  ,  2.8 ]])
        >>> kmeans.steps
        2
        """
        
        self.INIT = {
            'std': 0,
            'plus': 1,
            }

        self.__k = k
        self.__init = init
        self.__seed = seed
                
        self.means = None
        self.steps = None
        
    def compute(self, x):
        """Compute Kmeans.
        
        :Parameters:
           x : ndarray
               an 2-dimensional vector (number of points x dimensions)
   
        :Returns:
           cls : ndarray (1-dimensional vector)
                 cluster membership. Clusters are in 0, ..., k-1

        :Attributes:
          Kmeans.means : 2d ndarray float (k x dim)
                         means
          Kmeans.steps : int
                         number of steps
        """

        cls, self.means, self.steps = kmeanscore.kmeans(
            x, self.__k, self.INIT[self.__init], self.__seed)
        
        return cls
