import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import tensorflow as tf
import tf_bodypix
from tf_bodypix.api import download_model, load_model, BodyPixModelPaths
from tf_bodypix.draw import draw_poses
from tensorflow.keras import preprocessing
import cv2
import json
from matplotlib import pyplot as plt
import numpy as np
from calculations import measure_body_sizes
import gradio as gr
import pandas as pd

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

rainbow = [
  [110, 64, 170], [143, 61, 178], [178, 60, 178], [210, 62, 167],
  [238, 67, 149], [255, 78, 125], [255, 94, 99],  [255, 115, 75],
  [255, 140, 56], [239, 167, 47], [217, 194, 49], [194, 219, 64],
  [175, 240, 91], [135, 245, 87], [96, 247, 96],  [64, 243, 115],
  [40, 234, 141], [28, 219, 169], [26, 199, 194], [33, 176, 213],
  [47, 150, 224], [65, 125, 224], [84, 101, 214], [99, 81, 195]
]

def process_images(front_img, side_img, real_height_cm):
    fimage_array = preprocessing.image.img_to_array(front_img)
    simage_array = preprocessing.image.img_to_array(side_img)
    
    # BodyPix prediction
    frontresult = bodypix_model.predict_single(fimage_array)
    sideresult = bodypix_model.predict_single(simage_array)
    
    front_mask = frontresult.get_mask(threshold=0.75)
    side_mask = sideresult.get_mask(threshold=0.75)
    
    front_colored_mask = frontresult.get_colored_part_mask(front_mask, rainbow)
    side_colored_mask = sideresult.get_colored_part_mask(side_mask, rainbow)
    
    frontposes = frontresult.get_poses()
    sideposes = sideresult.get_poses()
    
    # Calculate body sizes
    body_sizes = measure_body_sizes(side_colored_mask, front_colored_mask, sideposes, frontposes, real_height_cm, rainbow)
    print(body_sizes)
    measurements_df = pd.DataFrame([body_sizes]) if isinstance(body_sizes, dict) else pd.DataFrame(body_sizes)
    
    # Save measurements to CSV
    csv_file = "Body-measurement.csv"
    if not os.path.exists(csv_file):
        measurements_df.to_csv(csv_file, index=False)
    else:
        measurements_df.to_csv(csv_file, mode='a', header=False, index=False)
    
    # Prepare measurements for display
    measurement_display = measurements_df.to_html(index=False, justify="center", border=1)
    print(measurement_display)
    return f"""
    <h3 style="text-align: center;">Body Measurements</h3>
    <div style="text-align: center;">
    </div>
    {measurement_display}
    """

# Create the Gradio interface
interface = gr.Interface(
    fn=process_images,
    inputs=[
        gr.Image(sources=["webcam", "upload"], type="numpy", label="Front Pose"),
        gr.Image(sources=["webcam", "upload"], type="numpy", label="Side Pose"),
        gr.Number(label="Enter Your Height (cm)")
    ],
    outputs=gr.HTML(label="Measurement Results"),  # Use HTML output to display the measurements
    title="Body Sizing System Demo",
    description="Capture two webcam images: Front View and Side View, and input your height in cm."
)

# Launch the app
interface.launch(share=True)