## This file is part of mlpy.
## k-Nearest Neighbor (kNN) based on kNN
## C-libraries developed by Stefano Merler.
   
## This code is written by Davide Albanese, <albanese@fbk.eu>.
## (C) 2008 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__ = ['Knn']

from numpy import *
import nncore


class Knn:
    """
    k-Nearest Neighbor (KNN).

    Example:
    
    >>> import numpy as np
    >>> import mlpy
    >>> xtr = np.array([[1.0, 2.0, 3.1, 1.0],  # first sample
    ...                 [1.0, 2.0, 3.0, 2.0],  # second sample
    ...                 [1.0, 2.0, 3.1, 1.0]]) # third sample
    >>> ytr = np.array([1, -1, 1])             # classes
    >>> myknn = mlpy.Knn(k = 1)                # initialize knn class
    >>> myknn.compute(xtr, ytr)              # compute knn
    1
    >>> myknn.predict(xtr)                   # predict knn model on training data
    array([ 1, -1,  1])
    >>> xts = np.array([4.0, 5.0, 6.0, 7.0]) # test point
    >>> myknn.predict(xts)                   # predict knn model on test point
    -1
    >>> myknn.realpred                       # real-valued prediction
    0.0
    """
    
    def __init__(self, k, dist = 'se'):
        """
        Initialize the Knn class.
        
        :Parameters:
          k : int (odd > = 1)
            number of NN
          dist : string ('se' = SQUARED EUCLIDEAN, 'e' = EUCLIDEAN)
            adopted distance
        """
        
        DIST = {'se': 1, # DIST_SQUARED_EUCLIDEAN
                'e':  2  # DIST_EUCLIDEAN
                }
        
        self.__k     = int(k)
        self.__dist  = DIST[dist]

        self.__x       = None
        self.__y       = None
        self.__classes = None

        self.realpred  = None
        
        self.__computed = False

    def compute(self, x, y):
        """
        Store x and y data.

        :Parameters:  
          x : 2d ndarray float (samples x feats)
            training data
          y : 1d ndarray integer (-1 or 1 for binary classification)
            : 1d ndarray integer (1, ..., nclasses for multiclass classificatio) 
            classes

        :Returns:
          1
        
        :Raises:
          ValueError
            if not (1 <= k <= #samples)
          ValueError
            if there aren'e at least 2 classes
          ValueError
            if, in case of 2-classes problems, the lables are not 1 and -1
          ValueError
            if, in case of n-classes problems, the lables are not int from 1 to n
        """
       
        self.__classes = unique(y).astype(int)
              
        
        if self.__k <= 0 or self.__k >= x.shape[0]:
            raise ValueError("k must be in [1, #samples)")

        if self.__classes.shape[0] < 2:
            raise ValueError("Number of classes must be >= 2")

        elif self.__classes.shape[0] == 2:
            if self.__classes[0] != -1 or self.__classes[1] != 1:
                raise ValueError("For binary classification classes must be -1 and 1")

        elif self.__classes.shape[0] > 2:
            if not alltrue(self.__classes == arange(1, self.__classes.shape[0] + 1)):
                raise ValueError("For %d-class classification classes must be 1, ..., %d"
                                 % (self.__classes.shape[0], self.__classes.shape[0]))

        self.__x = x.copy()
        self.__y = y.copy()

        self.__computed = True
        return 1


    def predict(self, p):
        """
        Predict knn model on a test point(s).

        :Parameters:
          p : 1d or 2d ndarray float (sample(s) x feats)
            test sample(s)

        :Returns:
          the predicted value(s) on success:
          integer or 1d numpy array integer (-1 or 1) for binary classification
          integer or 1d numpy array integer (1, ..., nclasses) for multiclass classification
          0 on succes with non unique classification
          -2 otherwise

        :Raises:
          StandardError
            if no Knn method computed
        """

        if self.__computed == False:
            raise StandardError("No Knn method compute() run")

        if p.ndim == 1:
            pred = nncore.predictnn(self.__x, self.__y, p, self.__classes, self.__k, self.__dist)
            self.realpred = 0.0
        elif p.ndim == 2:
            pred = empty(p.shape[0], dtype = int)
            for i in range(p.shape[0]):
                pred[i] = nncore.predictnn(self.__x, self.__y, p[i], self.__classes, self.__k, self.__dist)
            self.realpred = zeros(p.shape[0])

        return pred

