import os
import cv2
import time
import numpy as np
import gradio as gr
import threading
from tensorflow.keras import preprocessing
import pandas as pd
from calculations import measure_body_sizes
import tf_bodypix
from tf_bodypix.api import download_model, load_model, BodyPixModelPaths

# Load BodyPix model
bodypix_model = load_model(download_model(BodyPixModelPaths.MOBILENET_FLOAT_50_STRIDE_16))

# Create output directory
output_dir = "output_images"
os.makedirs(output_dir, exist_ok=True)

def save_image(image, filename):
    filepath = os.path.join(output_dir, filename)
    cv2.imwrite(filepath, image)
    print(f"Image saved as {filepath}")

def process_webcam_image(webcam_image, label, real_height_cm):
    """Process a webcam image"""
    if webcam_image is None:
        raise gr.Error(f"No {label} image captured")
    
    # Convert Gradio image to NumPy array
    frame = webcam_image

    filename = f"{label}_image.jpg"
    save_image(frame, filename)
    return frame

def process_images(front_img, side_img, real_height_cm):
    """Process front and side images to extract body measurements"""
    # Validate inputs
    if front_img is None or side_img is None:
        raise gr.Error("Both front and side images are required")

    # Convert images to array
    front_image_array = preprocessing.image.img_to_array(front_img)
    side_image_array = preprocessing.image.img_to_array(side_img)

    # Predict body parts
    front_result = bodypix_model.predict_single(front_image_array)
    side_result = bodypix_model.predict_single(side_image_array)

    # Create masks
    front_mask = front_result.get_mask(threshold=0.75)
    side_mask = side_result.get_mask(threshold=0.75)

    # Color the masks
    front_colored_mask = front_result.get_colored_part_mask(front_mask)
    side_colored_mask = side_result.get_colored_part_mask(side_mask)

    # Get poses
    front_poses = front_result.get_poses()
    side_poses = side_result.get_poses()

    # Calculate body sizes
    body_sizes = measure_body_sizes(
        side_colored_mask, 
        front_colored_mask, 
        side_poses, 
        front_poses, 
        real_height_cm, 
        rainbow=None
    )
    
    # Convert to DataFrame
    measurements_df = pd.DataFrame(body_sizes)
    return measurements_df

def create_gradio_interface():
    """Create Gradio interface for body measurements"""
    with gr.Blocks() as demo:
        gr.Markdown("# Body Measurement Tool")
        
        with gr.Row():
            height_input = gr.Number(label="Your Height (cm)")
            
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Front View")
                front_webcam = gr.Image(source="webcam", type="numpy", label="Front View")
                front_capture_button = gr.Button("Capture Front View")
                
            with gr.Column():
                gr.Markdown("### Side View")
                side_webcam = gr.Image(source="webcam", type="numpy", label="Side View")
                side_capture_button = gr.Button("Capture Side View")
        
        # Capture buttons functionality
        current_front_image = gr.State(None)
        current_side_image = gr.State(None)
        
        def capture_front_image(webcam_input):
            if webcam_input is None:
                raise gr.Error("No image in webcam")
            return webcam_input, webcam_input
        
        def capture_side_image(webcam_input):
            if webcam_input is None:
                raise gr.Error("No image in webcam")
            return webcam_input, webcam_input
        
        front_capture_button.click(
            fn=capture_front_image, 
            inputs=front_webcam, 
            outputs=[current_front_image, front_webcam]
        )
        
        side_capture_button.click(
            fn=capture_side_image, 
            inputs=side_webcam, 
            outputs=[current_side_image, side_webcam]
        )
        
        # Measurement calculation button
        calculate_button = gr.Button("Calculate Measurements")
        output_df = gr.DataFrame(label="Body Measurements")
        
        calculate_button.click(
            fn=lambda height, front, side: process_images(front, side, height),
            inputs=[height_input, current_front_image, current_side_image],
            outputs=output_df
        )
    
    return demo

def main():
    demo = create_gradio_interface()
    demo.launch(share=True, debug=True)

if __name__ == "__main__":
    main()