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

baseFolderPath = "main_dataset"
lostFolderPath = "lost-images"
output_file = "sv_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

        print(f"Processing {foldername}\n")

        run(image_paths, foldername, video_path)
    saveToJson()

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

        print(f"Processing {image}", end="\r")

        # 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[5],
            silent=True
        )

        emotions = anaylse_obj[0]['emotion']

        selected_emotions = [
            emotions['happy'],
            emotions['sad'],
            emotions['angry'],
            emotions['neutral']
        ]

        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}

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

        #normalized_emotions = [value / 100 for value in raw_emotions]

        normalized_emotions = [value / 100 for value in selected_emotions]


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

        # normalised_emotions = normalised_emotions.tolist()

        entry = {
            "folder": foldername,
            "image": image,  # Extract the filename from the path
            "error": 0,
            "vector": normalized_emotions
        }
        
        # Add the entry to the JSON data list
        embeddings.append(entry)
    except:
        # if model does not detect any emotions
        # set lost -> 1 & generate a random emotion vector
        random_vector = np.random.rand(4)
        random_vector_list = random_vector.tolist()

        entry = {
            "folder": foldername,
            "image": image,  # Extract the filename from the path
            "error": 1,
            "vector": random_vector_list
        }
        
        # Add the entry to the JSON data list
        embeddings.append(entry)
        pass

    print(entry["vector"])
    


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, "segmented_images")
        
        # 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]))

        video_path = "null"

        # 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()