from PIL import Image
import numpy as np
import random
import sys


def arraymutate(array, frequency, numpixels):
    """" Mutates given 2D-array with a specific frequency. In case of mutation, numpixels pixels get changed"""
    if random.random() <= frequency:
        # mutate numpixel times by flipping 0 to 1
        k = 0
        while k < numpixels:
            # generate random index and change 0 to 1 in that spot
            index1, index2 = random.randint(0, len(array)-1), random.randint(0, len(array[0])-1)
            if array[index1, index2] != 1:
                array[index1, index2] = 1
                k += 1
    return array

# We want to read in some images and modify them to create our training data. These are the labels for the pictures and the number of images of that type.
shapes = [("back", 11), ("circ", 8), ("dash", 37), ("slash", 11), ("underscore", 9), ("vert", 32), ("blank", 1)]
desiredoutput = {"back":"100000", "circ":"010000", "dash":"001000", "slash":"000100", "underscore":"000010", "vert":"000001", "blank":"000000"}
dataUnchanged = []

# read in pre-defined images, convert them to 0/1 matrices and store them in a list
for tupel in shapes:
    for i in range(0, tupel[1]):
        image = Image.open("{0}{1}.bmp".format(tupel[0], str(i))).convert("1")
        matrix = (np.reshape(list(image.getdata()), (8, 8)) / 255).astype(int)
        dataUnchanged.append((1 - matrix, desiredoutput[tupel[0]]))

trainingdata = []
# desired number of examples in the training data set
numberexamples = sys.argv[1]
# desired frequency of mutation
mutationfrequency = float(sys.argv[2])
# number of pixels that get changed randomly
pixelmutation = [1, 2, 3, 4]

# create a list of tupels with input patterns and desired output. The input patterns may get changed slighty via arraymutate function

for n in range(int(numberexamples)):
    example, pixelchanges = random.choice(dataUnchanged), random.choice(pixelmutation)
    trainingdata.append((arraymutate(example[0], mutationfrequency, pixelchanges), example[1]))
    # print statement is just for debugging
    #print("currently doing training object number" + str(n))

with open("recog_input", "wb") as fp:
    np.save(fp, trainingdata)