"""The perceptron from perceptron.html, as plain Python.

predict() and train() are the same two functions as on the web page.
Run it:  python3 perceptron.py
"""


def predict(picture, points):
    total = 0
    for i in range(len(picture)):
        if picture[i] == 1:            # a black square
            total = total + points[i]  # add its points
    return total > 0                   # more than 0: YES


def train(picture, answer, points):
    guess = predict(picture, points)
    if guess == answer:
        return                 # right: change nothing
    if answer == True:
        change = 1             # it is YES: add points
    else:
        change = -1            # it is NO: take points away
    for i in range(len(picture)):
        if picture[i] == 1:    # only the black squares
            points[i] = points[i] + change


# ---------- the same examples as the page ----------

def picture_of(*rows):
    """picture_of('010', '111', '010') -> [0, 1, 0, 1, 1, 1, 0, 1, 0]"""
    return [int(square) for row in rows for square in row]


LINES_IN_THE_MIDDLE = [
    (picture_of('000', '111', '000'), True),     # flat line
    (picture_of('010', '010', '010'), False),    # standing line
]

LINES_ANYWHERE = [
    (picture_of('111', '000', '000'), True), (picture_of('100', '100', '100'), False),
    (picture_of('000', '111', '000'), True), (picture_of('010', '010', '010'), False),
    (picture_of('000', '000', '111'), True), (picture_of('001', '001', '001'), False),
]

SEVEN_OR_ONE = [
    (picture_of('11111', '00010', '00100', '01000', '01000'), True),
    (picture_of('00100', '00100', '00100', '00100', '00100'), False),
    (picture_of('11110', '00010', '00100', '00100', '00100'), True),
    (picture_of('00100', '01100', '00100', '00100', '01110'), False),
    (picture_of('01111', '00001', '00010', '00100', '00100'), True),
    (picture_of('00010', '00010', '00010', '00010', '00010'), False),
    (picture_of('11111', '00001', '00001', '00001', '00001'), True),
    (picture_of('01000', '01000', '01000', '01000', '01000'), False),
]


# ---------- training on all the examples, round after round ----------

def train_until_right(examples, max_rounds=30):
    """Like the page's "Train until all right" button. Gives back (rounds, points)."""
    points = [0] * len(examples[0][0])
    for round_number in range(1, max_rounds + 1):
        mistakes = 0
        for picture, answer in examples:
            if predict(picture, points) != answer:
                mistakes = mistakes + 1
            train(picture, answer, points)
        if mistakes == 0:
            return round_number, points
    return None, points


def show(points):
    """Print points as rows, like the grid on the page."""
    size = int(len(points) ** 0.5)
    for row in range(size):
        print('   ', ' '.join(f'{p:3}' for p in points[row * size:(row + 1) * size]))


if __name__ == '__main__':
    for name, examples in [('Lines in the middle', LINES_IN_THE_MIDDLE),
                           ('Lines anywhere', LINES_ANYWHERE),
                           ('7 or 1', SEVEN_OR_ONE)]:
        rounds, points = train_until_right(examples)
        if rounds:
            print(f'{name}: all right after {rounds} rounds. points =')
        else:
            right = sum(predict(p, points) == a for p, a in examples)
            print(f'{name}: still wrong after 30 rounds ({right}/{len(examples)} right). points =')
        show(points)