import numpy as np
import random
from matplotlib import pyplot as plt

def main(extend_attribute):
    vectors, answers = loadData('hw3_dataset.txt', extend_attribute)
    row, col = vectors.shape
    weights = np.array([random.random() for i in range(col + 1)])
    # weights = np.array([0.3334, 0.3334, 0.3334, 0.3334, 0.3334, -1.0001])

    epoch = 0

    while(True):
        # print(epoch)
        # print(weights)
        predicts = classification(vectors, weights)

        # print(predicts == answers)
        if not np.all(predicts == answers):
            # print(predicts)
            # print(answers)
            weights = update(predicts, answers, weights, vectors)
            epoch += 1

        else:
            epoch += 1
            break

        # input()
    # print(epoch)
    # print(weights)
    return epoch


def classification(vectors, weights):
    row, col = vectors.shape
    ones = np.ones((row, 1), dtype=np.int64)
    vectors = np.append(vectors, ones, axis=-1)

    weights = np.tile(weights, (row, 1))
    weighted = np.multiply(vectors, weights)
    output = np.array([w.sum() for w in weighted])

    # print(output)

    return output > 0

def update(predicts, labels, weights, vectors):
    row, col = vectors.shape
    ones = np.ones((row, 1), dtype=np.int64)
    vectors = np.append(vectors, ones, axis=-1)


    predicts = np.where(predicts == True, 1, 0)
    labels = np.where(labels == True, 1, 0)

    # print('loss = %d' % np.abs(predicts - labels).sum())

    for i in range(row):
        # print(learning_rate * (labels[i] - predicts[i]) * weights)
        weights += learning_rate * (labels[i] - predicts[i]) * vectors[i]

    return weights


def loadData(filename, N):
    data = np.genfromtxt(filename, dtype='str')
    vectors = np.array([list(bstring) for bstring in data[:, 1]]).astype(np.int64)
    row, col = vectors.shape
    extended = np.random.randint(2, size=(row, N))
    vectors = np.concatenate((vectors, extended), axis=1)
    answers = np.array([vector.sum() >= 3 for vector in vectors])


    return vectors, answers


if __name__ == '__main__':

    learning_rate = 0.002
    epochs = []
    extend_attributes = []
    for i in range(20):
        epochnum = main(i + 1)
        extend_attributes.append(i + 1)
        epochs.append(epochnum)
        # print('learning_rate = %.3f, epoch = %d' %(learning_rate, epochnum))
    print(epochs)

    plt.plot(extend_attributes, epochs)
    plt.xlabel('extended attributes')
    plt.ylabel('epoch')
    plt.savefig('hw3_2.jpg')



    # learning_rate = 0.001
    # print(main())