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

def softmax(X_train,W,k):
#    print(W)
#    print(X_train.dot(W))
#    print(np.sum(X_train.dot(W),axis=1))
    return np.exp(X_train.dot(W[:,k]))/(np.sum(np.exp(X_train.dot(W))))

class LogisticRegression():
    def __init__(self,epsilon=10**-2,lamb=10**-2):
        self.epsilon=epsilon
        self.lamb=lamb
    def train(self,X_train,y_train):
        train_no,feature_no=X_train.shape
        self.n_cat=len(set(y_train))
        W=np.random.randn(feature_no+1,self.n_cat)
        W[:,-1]=np.zeros((feature_no+1))
        X_cat=np.ones((train_no,feature_no+1))
        X_cat[:,:feature_no]=X_train
        dw=np.zeros((feature_no+1,self.n_cat))
        
        old_W=np.full((W.shape),np.inf)
#        print(old_W-W)
#        print(abs(old_W-W)<self.epsilon)
        while np.any(abs(old_W-W)>self.epsilon):
#            print('pass')
            old_W=W
            
            for k in range(self.n_cat-1):
                summ=np.zeros((feature_no+1,1))
                for i in range(train_no):
#                    print(X_cat[i].T.shape)
#                    print(y_train[i]==k)
#                    print(softmax(X_cat[i],W,k,self.n_cat))
                    summ+=X_cat[i].reshape(feature_no+1,1)*(int(y_train[i]==k)-softmax(X_cat[i],W,k))
                   
                dw[:,k]=summ.ravel()
#            print(dw)
#            print(W)
            
            W=W+self.lamb*dw
        self.W=W
        print(W)
        
    def predict(self,X_test):
        test_no,feature_no=X_test.shape
        X_cat=np.ones((test_no,feature_no+1))
        X_cat[:,:feature_no]=X_test
        result=np.zeros((test_no,self.n_cat))
        for i in range(test_no):
            for k in range(self.n_cat):
                result[i,k]=softmax(X_cat[i],self.W,k)
        y_pred=np.argmax(result,axis=1)
        
        return y_pred
        
if __name__=='__main__':
    iris=load_iris()
    X_train=iris.data
#    print(iris.data)
    y_train=iris.target
    index=np.random.permutation(np.arange(len(X_train)))
    X_train=X_train[index]
    y_train=y_train[index]
    X_test=X_train
    
    model=LogisticRegression()
    model.train(X_train,y_train)
    y_pred=model.predict(X_test)
    print('Logistic Regression',y_pred)
    print(sum(y_train==y_pred))
    
    model=linear_model.LogisticRegression(solver='sag',multi_class='multinomial')
    model.fit(X_train,y_train)
    y_pred=model.predict(X_test)
    print('sklearn',y_pred)
    print(sum(y_train==y_pred))  
    print(model.coef_)
    
    
    
    