import numpy as np

from practical3.gradient_descent import stochastic_gradient_descent


class Perceptron:
    def __init__(self, epochs=10, weights = None):
        """Initialize Perceptron model

        :param dims: number of dimensions of data
        :param epochs: number of iterations over whole data
        """
        
        self.epochs = epochs
        self.weights = weights

    @staticmethod
    def _generate_initial_weights(dims):
        # Fill with random initial values
        np.random.seed(0)
        return np.random.rand(dims,)

    def fit(self, data, labels):
        """Fit the model and fix weight vector

        :param data: [N, rows, cols] dimensional numpy array of floats
        :param labels: [N] dimensional numpy array of 1s and -1s denoting class
        :yield: the function will yield weight vector after each update
        """
        data = self.transform_data(data, fit=True)
        
        if self.weights is None:
            self.w = self._generate_initial_weights(data.shape[1])
        else:
            self.w = np.loadtxt(self.weights, delimiter = ',')
            
        yield self.w
        for num_epoch in range(self.epochs):
            print(f"epoch N{num_epoch+1}:", end='\r', flush=True)
            # FIXME: Won't work correctly for windows, sorry :/
            for dw in stochastic_gradient_descent(
                    data, labels, self.gradloss):
                
                self.w -= dw
            yield self.w

    def _loss(self, vec, label):
        """Calculate loss on single data point

        :param vec: The input datapoint, numpy array of [dims] shape
        :param label: 1 or -1, which denotes the class of the input vector vec
        :return: loss generated by the given datapoint
        Note, that the result should be non_negative.
        >>> model = Perceptron(1)
        >>> model.w = np.array([1])
        >>> model._loss(np.array([3]), 1)
        0
        >>> model._loss(np.array([3]), -1)
        3
        >>> model._loss(np.array([-3]), 1)
        3
        >>> model._loss(np.array([-3]), -1)
        0
        >>> model._loss(np.array([0]), 1) # the boundary
        0
        >>> model._loss(np.array([0]), -1) # the boundary
        0
        """
        vec = Perceptron.transform_datapoint(self, vec)
        pred = 1 if np.dot(vec, self.w) >= 0.0 else -1
        return -1 * (label != pred)*(label * np.dot(vec, self.w))
    
    def loss(self, data, labels):
        return sum(self._loss(vec, label)
                   for vec, label in zip(data, labels))

    def _gradloss(self, vec, label):
        """Calculate the gradient of _loss on single data point
        >>> model = Perceptron(2)
        >>> model.w = np.array([1, 2])
        >>> np.all(model._gradloss(np.array([2, 1]), 1) == np.array([0, 0]))
        True
        >>> np.all(model._gradloss(np.array([2, 1]), -1) == np.array([-2, -1]))
        True
        """

        
        pred = 1 if np.dot(vec, self.w) >= 0.0 else -1
        return -1 * (label != pred)*(label * vec)
    
    def gradloss(self, data, labels):
    
        return sum(self._gradloss(vec, label)
                   for vec, label in zip(data, labels))
        
    def predict(self, data):
        """Calculate labels for each datapoing of the given data

        :param data: [N, rows, cols] dimensional numpy array to predict classes
        :return: numpy array of 1s and -1s,
                 where return_i denotes data_i's class
        >>> model = Perceptron(2)
        >>> model.w = np.array([1, 2])
        >>> np.all(model.predict(np.array([[2, 1], [1, 0], [0, -1]]))
        ...        == np.array([1, 1, -1]))
        True
        """
        
        data = self.transform_data(data)
        return [1 if np.dot(datapoint, self.w) >= 0.0 else -1 for datapoint in data]
    
    def accuracy(self, data, labels):
    
        res = self.predict(data) == labels
        return np.sum(res) *100 / len(res)
        
    def transform_data(self, data, fit=False):
        """Transform the data adding features, scaling, etc.
        Don't forget to save scaling parameters on `self` if fit is True

        :param data: [N, rows, cols] shaped numpy array of images
        :param fit: if True, scaling parameters should be calculated,
                    if False, scaling should be done with already calculated    parameters
                    Note: any processing is not required
        :return: [N, dims] shaped numpy array
        """

        data = [datapoint.flatten() / 255.0 for datapoint in data]
        data = np.array(data)
        
        if fit:
            self.mask = []
            
            for i in range(data.shape[1]):
                if set(data[:, i]) == {0}:
                    self.mask.append(False)
                else:
                    self.mask.append(True)
            
        data = np.c_[np.ones(len(data)), data[:, self.mask]]
        
        return np.array(data)
        
    def transform_datapoint(self, vec):
        vec = vec.flatten() / 255.0
        if vec.shape[0] != 784:
            print(vec)
        vec  = vec[self.mask]
        vec = np.insert(vec, 0, 1.0)

        return np.array(vec)