import cv2
import numpy as np
import matplotlib.pyplot as plt
import mediapipe as mp
import gradio as gr
import mediapipe as mp
import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model

# classification head to MobileNetV2
base_model = MobileNetV2(weights="imagenet", include_top=False, input_shape=(128, 128, 3))
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(1, activation='sigmoid')(x)  # Sigmoid activation for binary classification (0 or 1)
model = Model(inputs=base_model.input, outputs=x)

def classify_gender(image):
    # Resize image to 128x128 for MobileNetV2
    img_resized = cv2.resize(image, (128, 128))
    img_array = img_to_array(img_resized)
    img_array = preprocess_input(img_array)
    img_array = np.expand_dims(img_array, axis=0)

    # Predict gender (assuming 0: Male, 1: Female)
    prediction = model.predict(img_array)
    gender = "Male" if prediction[0][0] < 0.5 else "Female"
    return gender

# Initialize Mediapipe Pose model
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=True)

# Function to estimate height using dynamic calibration based on shoulder width
def estimate_real_world_height_and_area(image):
    height, width, _ = image.shape
    results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    if not results.pose_landmarks:
        print("No landmarks detected.")
        return None, None, None  # Return None if no landmarks detected

    landmarks = results.pose_landmarks.landmark

    # Reference average shoulder width (in cm)
    AVERAGE_SHOULDER_WIDTH_CM = 45

    # Points for height (nose to ankle) and shoulder width
    nose = landmarks[mp_pose.PoseLandmark.NOSE]
    left_shoulder = landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER]
    right_shoulder = landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER]
    left_ankle = landmarks[mp_pose.PoseLandmark.LEFT_ANKLE]
    right_ankle = landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE]

    # Calculate pixel height between nose (head) and ankle
    head_y = int(nose.y * height)
    ankle_y = int((left_ankle.y + right_ankle.y) / 2 * height)
    pixel_height = abs(ankle_y - head_y)

    # Calculate shoulder width in pixels
    shoulder_width_pixels = abs(left_shoulder.x - right_shoulder.x) * width

    # Determine the pixel-to-cm ratio using shoulder width
    pixel_to_cm_ratio = AVERAGE_SHOULDER_WIDTH_CM / shoulder_width_pixels if shoulder_width_pixels != 0 else 0

    # Calculate real-world height using this ratio
    real_world_height_cm = pixel_height * pixel_to_cm_ratio

    # Estimate torso area in pixels
    torso_width = abs(left_shoulder.x - right_shoulder.x) * width
    torso_height = abs(left_shoulder.y - left_ankle.y) * height
    torso_area_pixels = torso_width * torso_height
    real_torso_area = torso_area_pixels * (pixel_to_cm_ratio ** 2)

    return real_world_height_cm, real_torso_area, landmarks

# Main Gradio prediction function with scaling and plotting
def predict(image):
    try:
        # Classify Gender
        gender = classify_gender(image)  # Assuming this function is defined

        # Estimate height and area
        real_height, real_torso_area, landmarks = estimate_real_world_height_and_area(image)

        if real_height is None or real_torso_area is None:
            return "No landmarks detected; please use a clearer image.", None, None, None

        # Plot with scale
        fig, ax = plt.subplots(figsize=(10, 10))
        ax.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

        # Plot landmarks if they exist
        if landmarks:
            for landmark in landmarks:
                ax.scatter(landmark.x * image.shape[1], landmark.y * image.shape[0], color="red", s=10)

        # Add height and area text
        ax.text(10, 30, f"Predicted Gender: {gender}", color="yellow", fontsize=12, backgroundcolor='black')
        ax.text(10, 60, f"Estimated Height: {real_height:.2f} cm", color="yellow", fontsize=12, backgroundcolor='black')
        ax.text(10, 90, f"Estimated Torso Area: {real_torso_area:.2f} cm^2", color="yellow", fontsize=12, backgroundcolor='black')

        # Add grid scale to background
        ax.set_xticks(np.arange(0, image.shape[1], 50))
        ax.set_yticks(np.arange(0, image.shape[0], 50))
        ax.grid(color='white', linestyle='--', linewidth=0.5)

        # Save annotated plot for output
        plt.savefig("output_plot.png")
        plt.close(fig)

        return gender, real_height, real_torso_area, "output_plot.png"
    except Exception as e:
        print(f"Error encountered: {e}")
        return f"An error occurred: {e}", None, None, None

# Define Gradio interface
image_input = gr.Image(label="Input Image")
text_output = gr.Textbox(label="Predicted Gender")
height_output = gr.Textbox(label="Estimated Height (cm)")
area_output = gr.Textbox(label="Estimated Torso Area (cm^2)")
image_output = gr.Image(label="Annotated Image with Measurements")

# Set up Gradio interface
gr_interface = gr.Interface(
    fn=predict,
    inputs=image_input,
    outputs=[text_output, height_output, area_output, image_output],
    title="Gender, Height, and Body Measurement Estimation",
    description="Upload an image to predict gender, estimate height, and calculate 2D torso area with a scale overlay."
)

# Launch the Gradio app
gr_interface.launch()