from deepface import DeepFace
import glob
import os
import json
import numpy as np

baseFolderPath = "videos"
output_file = "embeddings.json"

embeddings = []

backends = [
  'opencv', 
  'ssd', 
  'dlib', 
  'mtcnn', 
  'fastmtcnn',
  'retinaface', 
  'mediapipe',
  'yolov8',
  'yunet',
  'centerface',
]

def main():
    video_folders = setup()
    for folder_info in video_folders:
        foldername, video_path, image_paths = folder_info
        run(image_paths, foldername, video_path)
    saveToJson()

def processImage(image, idx, foldername, video_path):
    try:

        # Initialize variables to accumulate sums
        sums = {
            'angry': 0,
            'disgust': 0,
            'fear': 0,
            'happy': 0,
            'sad': 0,
            'surprise': 0,
            'neutral': 0
        }

        anaylse_obj = DeepFace.analyze(
            img_path=image,
            actions=['emotion'],
            enforce_detection=True,
            detector_backend=backends[3]
        )

        print(anaylse_obj)

        emotions = anaylse_obj[0]['emotion']

        num_of_faces = len(anaylse_obj)

        # Get average emotion
        for face in anaylse_obj:
            emotions = face['emotion']
            for emotion in sums:
                sums[emotion] += emotions.get(emotion, 0)

        averages = {emotion: sums[emotion] / num_of_faces for emotion in sums}

        print(f"averages: {averages}")

        raw_emotions = [value for value in averages.values()]  

        normalized_emotions = [value / 100 for value in raw_emotions]

        #normalised_emotions = l2_normalize(np.array(raw_emotions))

        # normalised_emotions = normalised_emotions.tolist()

        entry = {
            "folder": foldername,
            "frame": image,  # Extract the filename from the path
            "video": video_path,
            "faces": num_of_faces,
            "vector": normalized_emotions
        }
        
        # Add the entry to the JSON data list
        embeddings.append(entry)
    except:
        pass
        #print("Could not locate face in image. Skipping.")


def setup():
    video_folders = []

    subfolders = [f.path for f in os.scandir(baseFolderPath) if f.is_dir()]

    for subfolder in subfolders:
        frames_folder_path = os.path.join(subfolder, "frames")
        
        # Find the video file in the subfolder
        video_files = [f for f in os.listdir(subfolder) if f.endswith(('.mp4', '.avi', '.mov'))]
        if video_files:
            video_path = os.path.join(subfolder, video_files[0])

        # List of all image files in the frames folder
        image_files = glob.glob(os.path.join(frames_folder_path, '*'))

        # Filter by specific image extensions (e.g., .jpg, .png)
        image_paths = [img for img in image_files if img.endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.gif'))]

        # Sort the image paths by filename (assuming the filenames are numerically ordered)
        image_paths.sort(key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))

        # Add the information to the list
        video_folders.append((os.path.basename(subfolder), video_path, image_paths))

    return video_folders
        

def run(image_paths, foldername, video_path):
    for idx, image in enumerate(image_paths):
        processImage(image, idx, foldername, video_path)

def l2_normalize(vector):
    norm = np.linalg.norm(vector)
    if norm == 0:
        return vector
    return vector / norm

def saveToJson():
    with open(output_file, 'w') as f:
        json.dump(embeddings, f, indent=4)
    print("Successfully created JSON file.")

if __name__ == "__main__":
    main()