import cv2
import numpy as np
import glob
import os, operator

def save_landmarks(coords, img_path):
    base = os.path.splitext(img_path)[0]
    print(base)
    path = img_path.replace('.png', '.pts')
    print(path)
    with open(path, "w") as file:
        for coord in coords:
            file.write("%i," % coord[0])
            file.write("%i\n" % coord[1])

def pos_of_landmarks(img):
    red = img[:,:,2]
    blue = img[:,:,0]
    green = img[:,:,1]

    indices = np.argwhere((red==255)&(blue==0)&(green==0))

    coords = []
    for index in range(len(indices)):
        i, j = indices[index]
        if index == 0:
            coords.append([i, j])
        for apex in range(len(coords)):
            n, k = coords[apex]
            if (abs(i - n) > 5 or abs(j - k) > 5) and index != 0:
                cond = True
            else:
                cond = False
                break
        if cond:
            coords.append([i, j])

    return img, coords

def find_face(cascade_path, frame):
    face_cascade = cv2.CascadeClassifier(cascade_path)
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)
    for (x, y, w, h) in faces:
        cv2.rectangle(frame, (x, y), (x+w, y+h+50), (255, 0, 0), 2)
    return frame, faces

def resampling(frame, faces):
    for (x, y, w, h) in faces:
        ROI = frame[y:y+h+50, x:x+w]
        resize = cv2.resize(ROI, (300,300), fx=2, fy=2, interpolation = cv2.INTER_CUBIC)
    return resize

def average_of_landmarks(img, coords, faces):
    totalCoords = [(0, 0)] * 51
    totalBBw = 0
    totalBBh = 0
    bbX = 0
    bbY = 0
    for (x, y, w, h) in faces:
        bbX = x
        bbY = y

    newCoords = []
    for index in range(len(coords)):
        y, x = coords[index]
        newCoords.append([x-bbX, y-bbY])
        x1, y1 = totalCoords[index]
        totalCoordX = x1 + abs(x-bbX)
        totalCoordY = y1 + abs(y-bbY)
        totalCoords[index] = (totalCoordX, totalCoordY)

    avgCoords = []
    for index in range(len(totalCoords)):
        x, y = totalCoords[index]
        newX = int(round(x/5))
        newY = int(round(y/5))
        avgCoords.append((newX, newY))

    print(avgCoords)
    for (i,j) in avgCoords:
        cv2.circle(img,(j,i), 5, (255,0,0), -1)

    return img, avgCoords

def draw_landmarks(img, coords):
    for (i,j) in coords:
        cv2.circle(img,(j,i), 5, (255,225,255), 1)
    return img

def average_face(avg_coords, coords):
    if len(avg_coords) < len(coords):
        return coords
    for i in range(len(coords)):
        x,y = coords[i]
        j,k = avg_coords[i]
        avg_coords[i] = (x+j, y+k)
    return avg_coords

cascade_path = 'frontal_face_features.xml'

paths_annotated = glob.glob('Ibrahim emotions/sad_annotated/*.png')
paths_original  = glob.glob('Ibrahim emotions/Sad/*.png')

avg_coords = []
count = 0
for path in paths_annotated:
    count += 1
    frame = cv2.imread(path)
    img, faces = find_face(cascade_path, frame)
    resized_img = resampling(img, faces)
    img, coords = pos_of_landmarks(resized_img)
    avg_coords = average_face(avg_coords, coords)

avg_coords[:] = [ (int(i/count), int(j/count)) for (i,j) in avg_coords]
print(avg_coords)
print(len(avg_coords))

for path in paths_original:
    frame = cv2.imread(path)
    img, faces = find_face(cascade_path, frame)
    resized_img = resampling(img, faces)
    img = draw_landmarks(resized_img, avg_coords)
    #img, avg_coords = average_of_landmarks(img, coords, faces)

    #save_landmarks(coords, path)

    cv2.imshow('image',img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()