import pandas as pd
import torch
import torch.nn.functional as F
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import gradio as gr
from datasets import load_dataset

# ============================================================
# 1. INITIALIZATION (Synced with Research Environment)
# ============================================================
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the CLIP model and processor using the Transformers library
# Ensuring 100% parity with the embeddings generated in Colab
MODEL_ID = "openai/clip-vit-base-patch32"
processor = CLIPProcessor.from_pretrained(MODEL_ID)
model = CLIPModel.from_pretrained(MODEL_ID).to(device)

# Load metadata and the reference dataset
df = pd.read_parquet('hm_style_data.parquet')
ds = load_dataset("tomytjandra/h-and-m-fashion-caption", split="train")

# ============================================================
# 2. ENGINE (Standardized Vector Generation with Object Fix)
# ============================================================

def get_embedding(user_input, input_type):
    """
    Generates a normalized feature vector. 
    Includes a safety check to extract the tensor from model output objects.
    """
    model.eval()
    with torch.no_grad():
        if input_type == "image":
            # Process image and extract features
            image = Image.open(user_input).convert("RGB")
            inputs = processor(images=image, return_tensors="pt").to(device)
            outputs = model.get_image_features(**inputs)
        else:
            # Process text and extract features
            inputs = processor(text=[user_input], return_tensors="pt", padding=True).to(device)
            outputs = model.get_text_features(**inputs)
            
        # --- THE COLAB FIX: Ensure we are working with a Tensor, not an Object ---
        if not isinstance(outputs, torch.Tensor):
            embedding = getattr(outputs, "pooler_output", outputs)
        else:
            embedding = outputs
            
        # Apply L2 Normalization using the functional interface
        embedding = F.normalize(embedding, p=2, dim=-1)
        
    return embedding.cpu().numpy().flatten()

def search_styles(text_input, image_input):
    try:
        if image_input:
            user_vector = get_embedding(image_input, "image")
        elif text_input:
            user_vector = get_embedding(text_input, "text")
        else:
            return None, "Upload inspiration to begin..."

        stored_embeddings = np.stack(df['embedding'].values)
        scores = cosine_similarity(user_vector.reshape(1, -1), stored_embeddings).flatten()
        df['similarity_score'] = scores
        top_matches = df.sort_values(by='similarity_score', ascending=False).head(4)

        return [ds[int(row['item_id'])]['image'] for _, row in top_matches.iterrows()], ""
    except Exception as e:
        return None, f"Status: {str(e)}"

# ============================================================
# 3. UI DESIGN
# ============================================================

BG_URL = "https://huggingface.co/spaces/lia-prop13/HM-Style-Scout/resolve/main/background.png"

custom_css = f"""
/* 1. Global Background & Footer Removal */
.gradio-container, body, #component-0 {{
    background-image: url('{BG_URL}') !important;
    background-repeat: repeat !important;
    background-size: 550px !important;
    background-attachment: fixed !important;
    background-color: #f3f3f3 !important;
}}
footer {{ display: none !important; }}

/* 2. Glass Card UI */
.main-card {{
    background: rgba(255, 255, 255, 0.96) !important;
    backdrop-filter: blur(20px);
    border-radius: 40px !important;
    padding: 30px 50px !important;
    margin: 15px auto !important;
    box-shadow: 0 30px 60px rgba(0,0,0,0.1) !important;
    max-width: 950px !important;
}}

/* 3. Typography */
.hm-title {{ font-family: 'Inter', sans-serif; font-weight: 900; font-size: 36px; text-align: center; color: #000; margin-bottom: 5px !important; }}
.hm-subtitle {{ font-family: 'Inter', sans-serif; font-weight: 400; font-size: 15px; text-align: center; color: #444; margin-bottom: 5px; }}
.hm-pro-tip {{ font-family: 'Inter', sans-serif; font-size: 13px; text-align: center; color: #777; font-style: italic; margin-bottom: 25px; }}

/* 4. Image Upload - Neutralizing Buttons & Hiding Toolbar */
.image-upload {{ border: 1px solid rgba(0,0,0,0.05) !important; border-radius: 25px !important; background: white !important; overflow: hidden; }}

/* Hiding the bar and labels */
.image-upload label, .image-upload .image-footer, .image-upload .icon-buttons, .image-upload .icon-button {{ 
    display: none !important; 
}}

/* If buttons still appear, make them neutral (Not Orange) */
.image-upload button, .image-upload .image-footer button {{
    color: #000 !important;
    background: transparent !important;
    border: none !important;
}}
.image-upload .upload-text {{ visibility: hidden !important; }}
.image-upload .upload-text::after {{
    content: "Upload or drag image";
    visibility: visible !important;
    display: block !important;
    font-size: 15px !important;
    color: #888 !important;
}}

/* 5. Gallery & Interaction (No Orange Focus) */
.gallery-container, .gallery-container * {{ border: none !important; outline: none !important; box-shadow: none !important; background: transparent !important; }}
.gallery-container img {{ transition: transform 0.4s ease !important; border-radius: 12px !important; }}
.gallery-container img:hover {{ transform: scale(1.05) !important; }}

/* 6. Buttons & Inputs */
.search-button {{ background: #000 !important; color: #fff !important; border-radius: 30px !important; height: 50px !important; font-weight: 700 !important; margin-top: 15px !important; border: none !important; }}
.text-input {{ border: 1px solid rgba(0,0,0,0.05) !important; border-radius: 15px !important; margin-top: 10px !important; }}

/* Heights Symmetry */
.image-upload {{ height: 300px !important; }}
.gallery-container {{ height: 435px !important; }}
"""

with gr.Blocks(css=custom_css) as demo:
    with gr.Column(elem_classes="main-card"):
        gr.HTML("""
            <div class='hm-title'>H&M Style Matcher</div>
            <div class='hm-subtitle'>Upload a photo of a clothing item or describe your dream fit. We'll find your next H&M favorite in seconds.</div>
            <div class='hm-pro-tip'> Pro tip: Search using either text OR an image. Remember to clear one before trying the other!</div>
        """)
        
        with gr.Row(equal_height=True):
            with gr.Column(scale=2):
                input_img = gr.Image(
                    label=None, 
                    type="filepath", 
                    show_label=False, 
                    container=False, 
                    elem_classes="image-upload"
                )
                input_txt = gr.Textbox(
                    placeholder="Type your style inspiration here...",
                    show_label=False,
                    container=False, 
                    elem_classes="text-input"
                )
                search_btn = gr.Button("FIND MATCHES", elem_classes="search-button")
                status = gr.Markdown("")

            with gr.Column(scale=3):
                style_gallery = gr.Gallery(
                    show_label=False, 
                    columns=2, 
                    object_fit="contain",
                    elem_classes="gallery-container"
                )

    search_btn.click(
        fn=search_styles, 
        inputs=[input_txt, input_img], 
        outputs=[style_gallery, status]
    )

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