# -*- coding: utf-8 -*-
"""Implementation of the perceptron algorithm.

Attributes:
    W (int):
        Number of columns in our input feature vectors
    alpha (float, optional)
        Learning rate
"""
import numpy as np


class Perceptron:
    """Implementation of the perceptron algorithm.
    """

    def __init__(self, N, alpha=0.1):
        """Initialize perceptron object.

        Arguments:
            N {int} -- number of columns in our input feature vectors

        Keyword Arguments:
            alpha {float} -- learning rate (default: {0.1})
        """
        # initialize the weight matrix and store the learning rate
        self.W = np.random.randn(N + 1) / np.sqrt(N)  # +1 for bias
        self.alpha = alpha

    def step(self, x):
        """ apply the step function

        Arguments:
            x {float} -- input data

        Returns:
           int -- if x is positive we return 1, otherwise we return 0
        """
        return 1 if x > 0 else 0

    def fit(self, X, y, epochs=10):
        """Train the perceptron.

        Arguments:
            x {float} -- Training data
            y {float} -- target output class labels

        Keyword Arguments:
            epochs {int} -- number of training epochs (default: {10})
        """
        # insert a column of 1's as the last entry in the feature matrix -- this little trick
        # allows us to treat the bias as a trainable parameter within the weight matrix
        X = np.c_[X, np.ones((X.shape[0]))]
        # loop over the desired number of epochs
        for _ in np.arange(0, epochs):
            # loop over each individual data point
            for (x, target) in zip(X, y):
                # take the dot product between the input features and the weight matrix, then pass
                # this value through the step function to obtain the prediction
                prediction = self.step(np.dot(x, self.W))
                # only perform a weight update if our prediction does not match the target
                if prediction != target:
                    # determine the error
                    error = prediction - target
                    # update the weight matrix
                    self.W += -self.alpha * error * x

    def predict(self, X, add_bias=True):
        """[summary]

        Arguments:
            X {array} -- The data points we’ll be predicting class labels for.

        Keyword Arguments:
            add_bias {bool} -- A boolean indicating whether we need to add a column of 1’s to X to
                               perform the bias trick. (default: {True})

        Returns:
            int -- if dot product is positive we return 1, otherwise we return 0
        """
        # ensure our input is a matrix
        X = np.atleast_2d(X)
        # check to see if the bias column should be added
        if add_bias:
            # insert a column of 1's as the last entry in the feature matrix (bias)
            X = np.c_[X, np.ones((X.shape[0]))]
        # take the dot product between the input features and the weight matrix, then pass the
        # value through the step function
        return self.step(np.dot(X, self.W))
