import os
import re
import numpy as np

def generate_genuine_comparison(L):
    idx_1 = np.random.randint(0, len(L))
    while True:
        idx_2 = np.random.randint(0, len(L))
        if ((L[idx_1][1] == L[idx_2][1]) and (idx_1 != idx_2)):
            break
    return [L[idx_1][0], L[idx_2][0]]


def generate_impostor_comparison(L):
    idx_1 = np.random.randint(0, len(L))
    while True:
        idx_2 = np.random.randint(0, len(L))
        if (L[idx_1][1] != L[idx_2][1]):
            break
    return [L[idx_1][0], L[idx_2][0]]



input_folder = '/Users/hugomcp/Desktop/both_eyes'

tot_comparisons = 1000

proportions_LVT = [0.7, 0.1, 0.2]

files = os.listdir(input_folder)

pattern = r"C(\d+)_S(\d+)_I(\d+)\.jpg"

LV = []
T = []
for file in files:
    if not file.endswith('jpg'):
        continue

    match = re.match(pattern, file)

    if 'S1' in file:
        LV.append([file, int(match.group(1)), int(match.group(2)), int(match.group(3))])
    else:
        T.append([file, int(match.group(1)), int(match.group(2)), int(match.group(3))])



comparisons = []


for i in range(int(tot_comparisons*proportions_LVT[0]/2)):
    print(i)
    comparisons.append([generate_genuine_comparison(LV), 1, 0])
    comparisons.append([generate_impostor_comparison(LV), 0, 0])

for i in range(int(tot_comparisons*proportions_LVT[1]/2)):
    comparisons.append([generate_genuine_comparison(LV), 1, 1])
    comparisons.append([generate_impostor_comparison(LV), 0, 1])

for i in range(int(tot_comparisons*proportions_LVT[2]/2)):
    comparisons.append([generate_genuine_comparison(T), 1, 2])
    comparisons.append([generate_impostor_comparison(T), 0, 2])


with open("/Users/hugomcp/Desktop/output.csv", "w") as file:
    for item in comparisons:
        file.write("%s, %s, %d, %d \n" % (item[0][0], item[0][1], item[1], item[2]))