# -*- coding: utf-8 -*-
# @Author: Aliaksei Halachkin
# @Date:   2017-04-19 14:32:45
# @Last Modified by:   Aliaksei Halachkin
# @Last Modified time: 2017-04-21 23:02:15

import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.preprocessing import MinMaxScaler
from sklearn.neighbors import KNeighborsClassifier
import weight


class LocalModel(BaseEstimator, ClassifierMixin):
    def __init__(self, weights, 
                       local_model,
                       n_neighbors,
                       pca = None,
                       dist_tf=None, #function that transform distance to sample weight
                       log=False):
        self.weights = weights
        self.weight = weight.Weight(weights)
        self.minmax = MinMaxScaler()
        self.pca = pca
        self.knn =  KNeighborsClassifier(n_neighbors=n_neighbors, 
                                         weights = 'uniform', 
                                         algorithm = 'brute',
                                         metric = 'minkowski', p=2)
        self.local_model = local_model
        self.n_neighbors = n_neighbors
        self.dist_tf = dist_tf
        self.log = log
        if log:
            self.sample_weights_ = []
            self.neighbor_labels = []

    def fit(self, X, y):
        Xw = self.minmax.fit_transform(X)
        Xw = self.weight.fit_transform(Xw)

        self.knn.fit(Xw, y)
        if not self.pca is None:
            self.pca.fit(Xw, y)
        self.X = X
        self.y = y

    def predict(self, X):
        Xw = self.minmax.transform(X)
        Xw = self.weight.transform(Xw)
        ypred = []
        c = 0
        for xw in Xw:
            xw = xw.reshape(1,-1)
            dist, neighbors_ind = self.knn.kneighbors(xw,
                                                      n_neighbors=self.n_neighbors,
                                                      return_distance=True)
            Xlocal, ylocal = self.X[neighbors_ind[0], :], self.y[neighbors_ind[0]]
            if self.log:
                self.neighbor_labels.append(ylocal)
            if len(set(ylocal)) > 1:

                if not self.pca is None:
                    Xlocal = self.pca.transform(Xlocal)

                if(self.dist_tf is None): 
                    self.local_model.fit(Xlocal, ylocal)
                else: 
                    sample_weights = self.dist_tf(dist[0])
                    if self.log:
                        self.sample_weights_.append(sample_weights)
                    self.local_model.fit(Xlocal, ylocal, sample_weights)

                x = X[c].reshape(1, -1)
                ypred.append(self.local_model.predict(x)[0])
            else:
                ypred.append(ylocal[0])

            c += 1

        return np.array(ypred)
