# --- COMPATIBILITY PATCH FOR HUGGINGFACE_HUB ---
import huggingface_hub
if not hasattr(huggingface_hub, "cached_download"):
    huggingface_hub.cached_download = huggingface_hub.hf_hub_download
# -----------------------------------------------

import gradio as gr
import numpy as np
import cv2
import torch
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
from diffusers import UNet2DConditionModel

# --- 1. Sizing Logic (Metric Geometry) ---
def get_person_size(image_np):
    try:
        # Ensure pose_landmarker.task is uploaded to your Space!
        base_options = python.BaseOptions(model_asset_path='pose_landmarker.task')
        options = vision.PoseLandmarkerOptions(base_options=base_options)
        detector = vision.PoseLandmarker.create_from_options(options)
        
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image_np)
        detection_result = detector.detect(mp_image)
        
        if not detection_result.pose_landmarks:
            return "Detection Failed: No body detected."

        landmarks = detection_result.pose_landmarks[0]
        # Euclidean Distance between shoulders (Points 11 and 12)
        ls, rs = landmarks[11], landmarks[12]
        dist = np.sqrt((ls.x - rs.x)**2 + (ls.y - rs.y)**2)

        # Mathematical thresholding for sizing
        if dist > 0.28: return "Extra Large (XL)"
        elif dist > 0.23: return "Large (L)"
        elif dist > 0.18: return "Medium (M)"
        else: return "Small (S)"
    except Exception as e:
        return f"Size Estimation Error: {str(e)}"

# --- 2. Main Processing Function ---
def run_tryon(person_img, garment_img):
    if person_img is None:
        return None, "Please upload a person image."
    
    # Calculate size using our Geometric module
    size_label = get_person_size(person_img)
    
    # Returning the original image as a placeholder for the VTON output
    # This proves the frontend and sizing logic are working!
    return person_img, f"Recommended Size: {size_label}"

# --- 3. Gradio Interface (The UI) ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 👔 MSc Math FYP: Virtual Try-On & Sizing")
    gr.Markdown("Using Euclidean distance for sizing and Diffusion Models for Try-On.")
    
    with gr.Row():
        with gr.Column():
            person_input = gr.Image(label="1. Upload Person Image", type="numpy")
            garment_input = gr.Image(label="2. Upload Cloth Image", type="numpy")
            submit_btn = gr.Button("Run Analysis", variant="primary")
        
        with gr.Column():
            output_img = gr.Image(label="Virtual Result")
            output_text = gr.Textbox(label="Body Size Estimation")

    submit_btn.click(
        fn=run_tryon,
        inputs=[person_input, garment_input],
        outputs=[output_img, output_text]
    )

if __name__ == "__main__":
    demo.launch()