import numpy as np
from gradient_descent import stochastic_gradient_descent


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

        :param dims: number of dimensions of data
        :param epochs: number of iterations over whole data
        """
        self.dims = dims
        self.epochs = epochs
        self.w = self._generate_initial_weights(dims)

    @staticmethod
    def _generate_initial_weights(dims):
        # FIXME: Fill with random initial values
        return np.ones(dims)

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

        :param data: [N X dims] 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
        """
        yield self.w
        for num_epoch in range(self.epochs):
            print(f"epoch N{num_epoch}:", 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
        """
        raise NotImplementedError()
    
    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])
        >>> model._gradloss(np.array([2, 1]), 1)
        array([0, 0])
        >>> model._gradloss(np.array([2, 1]), -1)
        array([-2, -1])
        """
        raise NotImplementedError()
    
    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 X dims] 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])
        >>> model.predict(np.array([[2, 1], [1, 0], [0, -1]]))
        array([1, 1, -1])
        """
        raise NotImplementedError()
