import numpy as np
from sklearn import linear_model
from sklearn import datasets

class Perceptron():
    def __init__(self,model='prime',max_iter=1000,lamb=10**-2):
        self.model=model
        self.max_iter=max_iter
        self.lamb=lamb
    
    def train(self,X_train,y_train):
        self.X_train=X_train
        self.y_train=y_train
        train_no,feature_no=X_train.shape
        W=np.zeros((feature_no,1))
        alpha=np.zeros((train_no,1))
        b=0
        flag=True
        iter_no=0
        while flag:
            iter_no+=1
            flag=False
            for i in range(train_no):
                if self.model=='dual':
                    W=np.zeros((feature_no,1))
                    for j in range(train_no):
                        W+=alpha[j]*y_train[j]*X_train[j].reshape(feature_no,1)
                    if y_train[i]*(X_train[i].dot(W)+b)<=0:
                        alpha[i]=alpha[i]+self.lamb
                        b=b+self.lamb*y_train[i]
                        flag=True
                        break                
                if self.model=='prime':
                    if y_train[i]*(X_train[i].dot(W)+b)<=0:
                        W=W+self.lamb*y_train[i]*X_train[i].reshape(feature_no,1)
                        b=b+self.lamb*y_train[i]
                        flag=True
                        break
                    
            if iter_no>self.max_iter:
                break
        self.W=W
        self.b=b
        self.alpha=alpha
        
    def predict(self,X_test):
        train_no,feature_no=self.X_train.shape
        if self.model=='prime':
            y_pred=np.sign(X_test.dot(self.W)+self.b)
        if self.model=='dual':
            W=np.zeros((feature_no,1))
            for j in range(train_no):
                W+=self.alpha[j]*self.y_train[j]*self.X_train[j].reshape(feature_no,1)
            y_pred=np.sign(X_test.dot(W)+self.b)
        return y_pred.ravel()
        
if __name__=='__main__':
#    X_train=np.array([[3,3],[4,3],[1,1]])
#    y_train=np.array([1,1,-1])
#    X_test=np.array([[100,100],[99,99]])
    
    index=np.random.permutation(150)
    iris=datasets.load_iris()
    X_train=iris['data'][:,(2,3)]
    y_train=(iris['target']==2).astype(np.float64)
    y_train[y_train==0]=-1
    y_train_1=y_train
    X_test=X_train
    X_train=X_train[index]
    y_train=y_train[index]
    
    model=Perceptron(model='prime',max_iter=10**3,lamb=10**-4)
    model.train(X_train,y_train)
    y_pred=model.predict(X_test)
    print('Perceptron',y_pred)
    print(sum(y_pred==y_train_1))
    
    model=linear_model.Perceptron()
    model.fit(X_train,y_train)
    y_pred=model.predict(X_test)
    print('sklearn',y_pred)
    print(sum(y_pred==y_train_1))
    
                