import streamlit as st
import cv2
import mediapipe as mp
import numpy as np
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf
from tensorflow.keras.models import load_model
from PIL import Image

# --- Setup ---
BASE_PATH = r"C:\Users\MANII\Desktop\AI_Hairstyle_Project"
MODEL_PATH = os.path.join(BASE_PATH, "face_shape_model_v2.h5")

# MediaPipe
BaseOptions = mp.tasks.BaseOptions
FaceDetector = mp.tasks.vision.FaceDetector
FaceDetectorOptions = mp.tasks.vision.FaceDetectorOptions
VisionRunningMode = mp.tasks.vision.RunningMode
TFLITE_PATH = os.path.join(BASE_PATH, "blaze_face_short_range.tflite")

# Class Labels
CLASS_NAMES = {0: 'Heart', 1: 'Oblong', 2: 'Oval', 3: 'Round', 4: 'Square'}

# Hairstyle Recommendations
RECOMMENDATIONS = {
    'Heart': {
        'styles': ['Side Part', 'Quiff', 'Fringe'],
        'avoid': 'Volume on top',
        'reason': 'Chin area balanced ho jata hai'
    },
    'Oblong': {
        'styles': ['Buzz Cut', 'Crop Top', 'Side Swept'],
        'avoid': 'Long straight styles',
        'reason': 'Face width add hoti hai'
    },
    'Oval': {
        'styles': ['Any Style', 'Undercut', 'Pompadour'],
        'avoid': 'Kuch bhi avoid nahi',
        'reason': 'Oval face sab styles suit karta hai'
    },
    'Round': {
        'styles': ['Fade', 'Mohawk', 'Textured Top'],
        'avoid': 'Bowl cut',
        'reason': 'Face elongated dikhta hai'
    },
    'Square': {
        'styles': ['Buzz Cut', 'Crew Cut', 'Short Sides'],
        'avoid': 'Flat top',
        'reason': 'Strong jawline complement hoti hai'
    }
}

# Load ML Model
@st.cache_resource
def load_face_model():
    return load_model(MODEL_PATH)

ml_model = load_face_model()

# --- UI ---
st.set_page_config(page_title="AI Men's Hairstyle", layout="centered")
st.title("✂️ Men's AI Virtual Hairstyle Try-On")
st.markdown("Photo upload karo — AI face shape detect karega aur best hairstyle suggest karega")

uploaded_file = st.file_uploader("Apni Photo Upload Karein", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    img = cv2.imdecode(file_bytes, 1)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    st.image(img_rgb, caption="Uploaded Photo", width=300)

    with st.spinner("AI analyze kar raha hai..."):

        # --- Face Shape Prediction ---
        img_resized = cv2.resize(img_rgb, (224, 224))
        img_array = np.expand_dims(img_resized / 255.0, axis=0)
        predictions = ml_model.predict(img_array)
        predicted_class = np.argmax(predictions[0])
        confidence = predictions[0][predicted_class] * 100
        face_shape = CLASS_NAMES[predicted_class]

    # --- Results ---
    st.success(f" Face Shape Detected: **{face_shape}** ({confidence:.1f}% confidence)")

    rec = RECOMMENDATIONS[face_shape]

    col1, col2 = st.columns(2)
    with col1:
        st.subheader(" Recommended Styles")
        for style in rec['styles']:
            st.write(f"• {style}")
        st.caption(f"Why: {rec['reason']}")

    with col2:
        st.subheader(" Avoid")
        st.write(rec['avoid'])

    # --- Virtual Try-On ---
    st.subheader("🎭 Virtual Try-On")
    style_choice = st.selectbox("Hairstyle choose karo:", rec['styles'] + ['Buzz Cut', 'Second Style'])

    hair_file = "buzz_cut.png" if "Buzz" in style_choice else "style.png"
    hair_path = os.path.join(BASE_PATH, hair_file)
    hair = cv2.imread(hair_path, cv2.IMREAD_UNCHANGED)

    if hair is not None:
        options = FaceDetectorOptions(
            base_options=BaseOptions(model_asset_path=TFLITE_PATH),
            running_mode=VisionRunningMode.IMAGE
        )
        with FaceDetector.create_from_options(options) as detector:
            mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img_rgb)
            result = detector.detect(mp_image)

        if result.detections:
            detection = result.detections[0]
            bbox = detection.bounding_box
            h, w, _ = img_rgb.shape

            face_w = int(bbox.width * 1.1)
            face_h = int(bbox.height * 0.6)
            hair_resized = cv2.resize(hair, (face_w, face_h))

            x1 = max(0, bbox.origin_x - int(face_w * 0.1))
            y1 = max(0, bbox.origin_y - int(face_h * 0.7))
            x2 = min(w, x1 + face_w)
            y2 = min(h, y1 + face_h)

            output = img_rgb.copy()
            hair_crop = hair_resized[0:(y2-y1), 0:(x2-x1)]

            if hair_crop.shape[2] == 4:
                alpha = hair_crop[:,:,3] / 255.0
                for c in range(3):
                    output[y1:y2, x1:x2, c] = (
                        hair_crop[:,:,c] * alpha +
                        output[y1:y2, x1:x2, c] * (1 - alpha)
                    )
            else:
                output[y1:y2, x1:x2] = hair_crop[:,:,:3]

            st.image(output, caption=f"Try-On: {style_choice}", width=300)
        else:
            st.warning("Face detect nahi hua try-on ke liye.")
    else:
        st.error(f"Hair image nahi mili: {hair_path}")