#%%
import os
import numpy as np
import dlib
from skimage import io
from scipy.spatial import distance
from scipy.special import softmax
import webbrowser
from matplotlib import pyplot as plt
from PIL import Image
from multiprocessing import Pool
#%%


def make_description(name):
    global sp, facerec, detector, direct
    img = io.imread(direct + name)
    dets_webcam = detector(img, 1)
    for k, d in enumerate(dets_webcam):
        shape = sp(img, d)
    return facerec.compute_face_descriptor(img, shape)


#%%
print('Введите 1, если Вы будете искать мужчину;')
print('Введите 2, если Вы будете искать женщину.')
sex = int(input()) - 1
#%%
direct = './data/photos/women/' if sex else './data/photos/men/'
fileNames = os.listdir(direct)
amount = len(fileNames)
#%%
vect = np.zeros(amount)
similarityMatrix = np.zeros((amount, amount))
descripions = np.zeros((amount, 128))
#%%
sp = dlib.shape_predictor('models/shape_predictor_68_face_landmarks.dat')
facerec = dlib.face_recognition_model_v1('models/dlib_face_recognition_resnet_model_v1.dat')
detector = dlib.get_frontal_face_detector()
#%%
with Pool() as p:
    descripions = np.array(p.map(make_description, fileNames))
#%%
for i in range(0, amount - 1):
    for j in range(i + 1, amount):
        similarityMatrix[i, j] = similarityMatrix[j, i] = distance.euclidean(descripions[i], descripions[j])
#%%
print(similarityMatrix)
# %%
answers = []


def ask_image(i: int) -> bool:
    img = Image.open(f'{direct}/{fileNames[i]}')
    plt.imshow(img)
    plt.show()
    ans = ''
    while ans not in ('+', '-'):
        print("+/-/stop?")
        ans = input()
        if ans == 'stop':
            raise ValueError("Stop")
    answers.append(ans)
    return ans == '+'


# %%
k = 1


def run_iter(visited):
    global vect

    vc = vect.copy() + 1e-4
    vc[visited] = 1e9
    p = softmax(1 / vc)
    i = np.random.choice(amount, p=p)
    while visited[i]:
        i = np.random.choice(amount, p=p)
    visited[i] = True
    res = ask_image(i)
    if res:
        vect += k * (similarityMatrix[i] ** 2)
    else:
        vect += k / ((similarityMatrix[i] ** 2) + 1e-4)  # To avoid division by zero


def run_cycle():
    visited = np.zeros(amount, dtype=bool)
    while True:
        try:
            run_iter(visited)
        except ValueError:
            break


# %%
ans_num = np.array(list(map(lambda x: x == '+', answers)), dtype=int)
# %%
plt.scatter(np.arange(amount), ans_num)
plt.show()
# %%
ans_cumsum = np.cumsum(list(map(lambda x: 1 if x == 1 else -1, ans_num)))
# %%
plt.plot(np.arange(amount), ans_cumsum)
plt.show()