import streamlit as st
from PIL import Image
import time
from utils.levels import complete_level, render_page, initialize_level
from utils.login import get_login, initialize_login
from utils.database import get_database
import os
import time
import face_recognition
import json
import numpy as np
from collections import defaultdict
import pickle as pkl

initialize_login()
initialize_level()

LEVEL = 3

PKL_PATH = 'dataset/database.pkl'

def step3_page():
    st.header("Training the Model")
    st.subheader("Face encoding")
    st.markdown(
        """
    ### What is Face Encoding?
    In face recognition, face encodings are numerical representations of facial features that are used to uniquely identify individuals. 
    These encodings are obtained by extracting relevant facial information from an input image or video frame.
    
    Face encodings are typically computed using deep learning models, such as Convolutional Neural Networks (CNNs), 
    that are trained on large datasets of labeled faces. During the training process, these models learn to recognize patterns and extract discriminative features from facial images.
    """
    )
    st.image(
        "https://miro.medium.com/v2/resize:fit:720/format:webp/1*V_wNVR0QvLQ7JZyUwMTv8w.jpeg",
        use_column_width=True,
    )

    st.markdown(
        """
        Once the face encodings are obtained, they can be stored in a database or used for face recognition tasks. 
        During face recognition, the encodings of input faces are compared to the stored encodings to determine if a match exists. 
        Various similarity metrics, such as Euclidean distance or cosine similarity, can be utilized to measure the similarity between 
        face encodings and determine potential matches.
    """
    )

    st.info(
        "Now it's your turn to train your model to create face encodings to each of the faces in the known-face database that you have created in the previous step!"
    )

    img_dir = os.path.join(".sessions", get_login()["username"], "known_faces")
    images = os.listdir(img_dir)
    if len(images) > 0:
        st.info("First, let's see your saved faces in your known-face database.")
        cols = st.columns(len(images))
        for i, img in enumerate(images):
            face_name = img.split("_")[0]
            cols[i].image(os.path.join(img_dir, img), use_column_width=True)
            cols[i].write(face_name)
    st.info("Now it's your turn to train the model! Click on the button below to train the model with your data to generate face encodings!")
    if st.button("Train Model"):
        my_bar = st.progress(0, text="Training....")
        if len(images) > 0:
            database = get_database(PKL_PATH)
            for i in range(100):
                time.sleep(0.1)
                my_bar.progress(i, text="Training....")
            my_bar.progress(100, text="Successfully Trained!")
            st.success("Model trained successfully!")
            st.info("Now, lets generate face encodings for each face in known-face database using the model you just trained!")
            my_bar = st.progress(0, text="Generating face encodings...")
            for i, img in enumerate(images):
                face_image = face_recognition.load_image_file(os.path.join(img_dir, img))
                my_face_encoding = face_recognition.face_encodings(face_image)[0]
                face_name = img.split("_")[0]
                face_id = img.split(".")[0]

                # check if id already exists
                existing_id = [database[i]["face_id"] for i in database.keys()]
                if face_id in existing_id:
                    st.error(f"Encoding already created for : {face_id}")
                else:
                    database[i] = {'face_id': face_id,
                                   'name': face_name,
                                   'encoding': my_face_encoding}
                    with open(PKL_PATH, 'wb') as f:
                        pkl.dump(database, f)
                time.sleep(0.5)
                my_bar.progress(int((i + 1) / len(images) * 100), text="Generating face encodings...")
            my_bar.progress(100, text="Successfully encoded all the known faces!")
            st.success("Face encoding completed successfully!")

        else:
            my_bar.empty()
            st.error("You have not taken any images yet! Do the previous steps first!")
    if st.button("Complete"):
        complete_level(LEVEL)


render_page(step3_page, LEVEL)