import os
from concurrent.futures import ThreadPoolExecutor
import gradio as gr
from heartBPM_modified_copy import heart
from eyebrow_detection_modified_copy import stress
from age_estimator.mivolo.demo_copy import main as age_estimation_main

def process_video(video_file):
    # Validate the input file path
    if not video_file or not os.path.isfile(video_file):
        return {'error': 'Invalid video path'}

    # Run functions in parallel
    with ThreadPoolExecutor() as executor:
        heart_future = executor.submit(heart, video_file)
        stress_future = executor.submit(stress, video_file, duration=30)
        
        # Define parameters for age estimation
        output_folder = 'output'
        detector_weights = 'age_estimator/mivolo/models/yolov8x_person_face.pt'
        checkpoint = 'age_estimator/mivolo/models/model_imdb_cross_person_4.22_99.46.pth.tar'
        device = 'cpu'
        with_persons = True
        disable_faces = False
        draw = True
        
        age_future = executor.submit(
            age_estimation_main, video_file, output_folder, detector_weights, checkpoint, device, with_persons, disable_faces, draw
        )
        
        # Retrieve results
        avg_bpm, frames_processed = heart_future.result()
        stressed_count, not_stressed_count, most_frequent_label = stress_future.result()
        absolute_age, lower_bound, upper_bound = age_future.result()

    # Compile results
    results = {
        'Average BPM': avg_bpm,
        'Most Frequent State': most_frequent_label,
        'Age Range': f"{lower_bound} - {upper_bound}"
    }

    return results

# Define Gradio interface
gr.Interface(
    fn=process_video,
    inputs=gr.Video(label="Upload a video file"),
    outputs="json",
    title="Parallel Video Processing for Heart Rate, Stress, and Age Estimation"
).launch()