import numpy as np
import random

dataset = {}
dataset['target_name'] = np.array(['More Needed', 'Correct', 'Less needed'])
dataset['feature_name'] = np.array(['digit1', 'digit2', 'digit3', 'digit4'])
#dataset['features'] = np.array()

for i in range(20000):
    digit1 = random.randint(0,9)
    digit2 = random.randint(0,9)
    added = digit1 + digit2
    use_correct = bool(random.getrandbits(1))
    if use_correct:
        digit3, digit4 = f'{added:02}'[:1], f'{added:02}'[1:]
        target = 1 #correct
    else:
        other = random.randint(0,19)
        digit3, digit4 = f'{other:02}'[:1], f'{other:02}'[1:]
        #'correct' if added==other else {True: 'More', False: 'Less'}.get(added-other<0, 'h')
        target = 1 if added==other else {True: 0, False: 2}.get(added-other<0, 'h')


    new_feature = np.array([digit1, digit2, digit3, digit4]).astype('int')
    if not 'features' in dataset:
        dataset['features'] = new_feature
        dataset['target'] = np.array([]).astype('int')
    else:
        dataset['features'] = np.vstack([dataset['features'], new_feature])

    dataset['target'] = np.append (dataset['target'], [target])

#print(dataset['features'])
#print(dataset['target'])

from sklearn.neighbors import KNeighborsClassifier
knn = KNeighborsClassifier(n_neighbors=1)

knn.fit(dataset['features'], dataset['target'])

###############
#Now lets predict
#4 + 6 = 12
#Gave 11 as result
what_is = np.array([[4, 6, 1, 1]])

prediction = knn.predict(what_is)

print(dataset['target_name'][prediction])
#Answer: need 'More Needed' which is correct

#This actually just remembered the exact digits but a good way to illustrate how data is stored and 'predicted'