import os
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

import streamlit as st
import cv2
from tqdm import tqdm
import numpy as np
import tensorflow as tf
import pandas as pd
from tempfile import NamedTemporaryFile
from functions import *
au_to_movements= {
    'au1': 'inner brow raiser',
    'au2': 'outer brow raiser',
    'au4': 'brow lowerer',
    'au5': 'upper lid raiser',
    'au6': 'cheek raiser',
    'au9': 'nose wrinkler',
    'au12': 'lip corner puller',
    'au15': 'lip corner depressor',
    'au17': 'chin raiser',
    'au20': 'lip stretcher',
    'au25': 'lips part',
    'au26': 'jaw drop'
}
au_labels = [
    "au1",
    "au12",
    "au15",
    "au17",
    "au2",
    "au20",
    "au25",
    "au26",
    "au4",
    "au5",
    "au6",
    "au9"
]
col=[au_to_movements[i] for i in au_labels]
def binary_focal_loss(gamma=2.0, alpha=0.25):
    def focal_loss(y_true, y_pred):
        # Define epsilon to avoid log(0)
        epsilon = tf.keras.backend.epsilon()
        # Clip predictions to prevent log(0) and log(1 - 0)
        y_pred = tf.clip_by_value(y_pred, epsilon, 1.0 - epsilon)
        # Compute the focal loss
        fl = - alpha * (y_true * (1 - y_pred)**gamma * tf.math.log(y_pred)
                       + (1 - y_true) * (y_pred**gamma) * tf.math.log(1 - y_pred))
        return tf.reduce_mean(fl, axis=-1)
    return focal_loss

loss = binary_focal_loss(gamma=2.0, alpha=0.25)

# Function to read video frames into a list
def read_video_frames(video_path):
    cap = cv2.VideoCapture(video_path)
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frames.append(frame)

    cap.release()
    return frames

# Function to process frames and make predictions
def process_frames(frames, model):
    frames = [get_face(frame) for frame in tqdm(frames)]
    st.text(f"face shape : {frames[0].shape}")
    frame_array = np.array(frames)
    preds = model.predict(frame_array).round()
    return preds

# Function to save predictions to a CSV file
def save_predictions_to_csv(predictions, filename="predictions.csv"):
    df = pd.DataFrame(predictions,columns=col)
    df.to_csv(filename, index=False)
    return filename

# Load your Keras model
def load_model():
    model = tf.keras.models.load_model('incept_v3_10fps_full_dp0.2.keras',
                                       custom_objects={'binary_focal_loss': binary_focal_loss})
    return model

# Streamlit app
def main():
    st.title("Video Frame Prediction App")
    
    # Upload video file
    uploaded_file = st.file_uploader("Upload a video file", type=["mp4", "avi", "mov"])

    if uploaded_file is not None:
        with NamedTemporaryFile(delete=False) as tmp_file:
            tmp_file.write(uploaded_file.read())
            video_path = tmp_file.name

        # Load the model
        model = load_model()

        # Predict button
        if st.button("Predict"):
            # Read frames from video
            st.text("Reading video frames...")
            frames = read_video_frames(video_path)
            st.text(f"Total frames read: {len(frames)}")

            # Process frames and make predictions
            st.text("Processing frames and making predictions...")
            predictions = process_frames(frames, model)
            st.text("Predictions completed!")

            # Save predictions to CSV
            csv_file_path = save_predictions_to_csv(predictions)
            st.text("Predictions saved to CSV!")

            # Make CSV downloadable
            with open(csv_file_path, "rb") as f:
                st.download_button(
                    label="Download CSV",
                    data=f,
                    file_name="predictions.csv",
                    mime="text/csv"
                )

            # Clean up the temporary file
            os.remove(video_path)

if __name__ == "__main__":
    main()