import gradio as gr
import numpy as np
import torch
import os
from PIL import Image
from datasets import load_dataset
from transformers import CLIPProcessor, CLIPModel

# ── 1. Configuration ───────────────────────────────────────
MODEL_ID = "openai/clip-vit-base-patch32"
DATASET_ID = "prithivMLmods/Shoe-Net-10K"
TOP_K = 3

# ── 2. Load model ──────────────────────────────────────────
print("Loading CLIP model...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CLIPModel.from_pretrained(MODEL_ID).to(device)
processor = CLIPProcessor.from_pretrained(MODEL_ID)
model.eval()
print(f"Model loaded on {device}")

# ── 3. Load dataset ────────────────────────────────────────
print("Loading shoe dataset...")
dataset = load_dataset(DATASET_ID, split="train")
print(f"Dataset loaded: {len(dataset)} images")

# ── 4. Load pre-computed embeddings (.npy arrays) ──────────
print("Loading pre-computed embeddings from npy artifacts...")
embeddings_matrix = np.load("shoe_embeddings.npy").astype(np.float32)
image_indices = np.load("dataset_image_indices.npy")

# Ensure image indices are fully flattened and cast to python integers
image_indices = np.atleast_1d(image_indices).flatten().astype(int)

cluster_names_map = {
    "0": "Pure White & Ultra Light Sneaker Profile",
    "1": "Jet Black & Deep Dark Footwear Profile",
    "2": "Warm Tones, Earthy Tan & Luxury Browns",
    "3": "Muted Grey, Pastels & Mid-Tone Casuals",
    "4": "Vibrant, Chromatic & Colorful Selection"
}
print(f"Loaded {len(embeddings_matrix)} embeddings of dim {embeddings_matrix.shape[1]}")


# ── 5. Embedding functions ─────────────────────────────────
def embed_image(pil_image: Image.Image) -> np.ndarray:
    img = pil_image.convert('RGB')
    inputs = processor(images=img, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model.get_image_features(**inputs)
        features = outputs if isinstance(outputs, torch.Tensor) else outputs[0]

    features = features / features.norm(dim=-1, keepdim=True)
    return features.cpu().numpy().flatten()


def embed_text(text: str) -> np.ndarray:
    inputs = processor(text=[text], return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        outputs = model.get_text_features(**inputs)
        features = outputs if isinstance(outputs, torch.Tensor) else outputs[0]

    features = features / features.norm(dim=-1, keepdim=True)
    return features.cpu().numpy().flatten()


# ── 6. Recommendation engine ───────────────────────────────
def recommend(query_emb: np.ndarray) -> list:
    sims = np.dot(embeddings_matrix, query_emb.reshape(-1, 1)).flatten()
    sorted_idx = np.argsort(sims)[::-1]

    results = []
    for i in sorted_idx:
        if sims[i] > 0.9999:
            continue

        try:
            dataset_idx = int(image_indices[i])

            if dataset_idx < 0 or dataset_idx >= len(dataset):
                continue

            img_data = dataset[dataset_idx]['image']

            if not isinstance(img_data, Image.Image):
                img = Image.open(img_data)
            else:
                img = img_data

            cluster_id = int(i % 5)
            cluster_name = cluster_names_map.get(str(cluster_id), f"Style Profile {cluster_id}")

            results.append({
                "image": img,
                "score": float(sims[i]),
                "rank": len(results) + 1,
                "index": dataset_idx,
                "cluster": cluster_name
            })
        except Exception:
            continue

        if len(results) == TOP_K:
            break

    return results


# ── 7. Gradio handler functions ────────────────────────────
def recommend_by_image(query_image):
    if query_image is None:
        return (None, "Upload an image first",
                None, "Upload an image first",
                None, "Upload an image first")
    try:
        query_emb = embed_image(query_image)
        results = recommend(query_emb)
        outputs = []
        for r in results:
            outputs.append(r["image"])
            outputs.append(
                f"Rank #{r['rank']}\n"
                f"Similarity: {r['score']:.4f}\n"
                f"Style: {r['cluster']}\n"
                f"Dataset Index: {r['index']}"
            )
        while len(outputs) < 6:
            outputs.extend([None, "No additional matches found"])
        return tuple(outputs[:6])
    except Exception as e:
        err = f"Error during image recommendation: {str(e)}"
        return None, err, None, err, None, err


def recommend_by_text(query_text):
    if not query_text or not query_text.strip():
        return (None, "Enter a text description first",
                None, "Enter a text description first",
                None, "Enter a text description first")
    try:
        query_emb = embed_text(query_text.strip())
        results = recommend(query_emb)
        outputs = []
        for r in results:
            outputs.append(r["image"])
            outputs.append(
                f"Rank #{r['rank']}\n"
                f"Similarity: {r['score']:.4f}\n"
                f"Style: {r['cluster']}\n"
                f"Dataset Index: {r['index']}"
            )
        while len(outputs) < 6:
            outputs.extend([None, "No additional matches found"])
        return tuple(outputs[:6])
    except Exception as e:
        err = f"Error during text recommendation: {str(e)}"
        return None, err, None, err, None, err


# ── Pre-build example images as files (avoids loading PIL objects directly into gr.Examples) ──
EXAMPLES_DIR = "example_images"
os.makedirs(EXAMPLES_DIR, exist_ok=True)
example_image_paths = []
for idx in [0, 100, 500, 1000, 2500]:
    try:
        img = dataset[idx]['image']
        if not isinstance(img, Image.Image):
            img = Image.open(img)
        path = os.path.join(EXAMPLES_DIR, f"example_{idx}.jpg")
        img.convert("RGB").save(path)
        example_image_paths.append(path)
    except Exception as e:
        print(f"Could not prepare example image {idx}: {e}")


# ── 8. Gradio UI ───────────────────────────────────────────
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="blue"),
    title="Shoe Visual Recommender",
    css="""
    .title-text { text-align: center; }
    .result-box { background: #f0f7ff; border-radius: 8px; padding: 8px; }
    """
) as demo:

    gr.Markdown("""
    <div class="title-text">

    # Shoe Visual Recommender
    ### AI-powered shoe recommendation using CLIP embeddings + Cosine Similarity

    </div>
    """)

    gr.Markdown("""
    **How it works:**
    1. Upload a shoe image **or** type a text description
    2. CLIP converts your input to a 512-dimensional embedding vector
    3. Cosine similarity finds the 3 most similar shoes from 10,000 options
    4. Results are ranked by similarity score
    ---
    """)

    with gr.Tabs():

        # ── Tab 1: Image Input ──
        with gr.TabItem("Search by Image"):
            gr.Markdown("### Upload a shoe image to find visually similar shoes")

            with gr.Row():
                with gr.Column(scale=1):
                    img_input = gr.Image(
                        type="pil",
                        label="Upload Shoe Image",
                        height=280
                    )
                    img_btn = gr.Button(
                        "Find Similar Shoes",
                        variant="primary",
                        size="lg"
                    )
                    if example_image_paths:
                        gr.Examples(
                            examples=[[p] for p in example_image_paths],
                            inputs=[img_input],
                            label="Example Shoes from Dataset"
                        )

                with gr.Column(scale=3):
                    gr.Markdown("### Top 3 Most Similar Shoes")
                    with gr.Row():
                        img_out1 = gr.Image(label="Best Match", height=200)
                        img_out2 = gr.Image(label="2nd Match", height=200)
                        img_out3 = gr.Image(label="3rd Match", height=200)
                    with gr.Row():
                        img_info1 = gr.Textbox(label="Match #1 Details", lines=4, interactive=False)
                        img_info2 = gr.Textbox(label="Match #2 Details", lines=4, interactive=False)
                        img_info3 = gr.Textbox(label="Match #3 Details", lines=4, interactive=False)

            img_btn.click(
                fn=recommend_by_image,
                inputs=[img_input],
                outputs=[img_out1, img_info1, img_out2, img_info2, img_out3, img_info3]
            )

        # ── Tab 2: Text Input ──
        with gr.TabItem("Search by Text"):
            gr.Markdown("### Describe the shoe you're looking for")

            with gr.Row():
                with gr.Column(scale=1):
                    text_input = gr.Textbox(
                        label="Describe a Shoe",
                        placeholder="e.g. 'black elegant flat ballet shoe'",
                        lines=3
                    )
                    text_btn = gr.Button(
                        "Find Matching Shoes",
                        variant="primary",
                        size="lg"
                    )
                    gr.Examples(
                        examples=[
                            ["black elegant flat ballet shoe"],
                            ["casual white sneaker with bow"],
                            ["brown leather oxford formal shoe"],
                            ["colorful summer sandal"],
                            ["dark pointed toe heels"]
                        ],
                        inputs=[text_input],
                        label="Example Queries"
                    )

                with gr.Column(scale=3):
                    gr.Markdown("### Top 3 Matching Shoes")
                    with gr.Row():
                        txt_out1 = gr.Image(label="Best Match", height=200)
                        txt_out2 = gr.Image(label="2nd Match", height=200)
                        txt_out3 = gr.Image(label="3rd Match", height=200)
                    with gr.Row():
                        txt_info1 = gr.Textbox(label="Match #1 Details", lines=4, interactive=False)
                        txt_info2 = gr.Textbox(label="Match #2 Details", lines=4, interactive=False)
                        txt_info3 = gr.Textbox(label="Match #3 Details", lines=4, interactive=False)

            text_btn.click(
                fn=recommend_by_text,
                inputs=[text_input],
                outputs=[txt_out1, txt_info1, txt_out2, txt_info2, txt_out3, txt_info3]
            )

        # ── Tab 3: About ──
        with gr.TabItem("About This Project"):
            gr.Markdown("""
            ## About This Shoe Recommender

            ### Dataset
            - **Source:** [Shoe-Net-10K](https://huggingface.co/datasets/prithivMLmods/Shoe-Net-10K)
            - **Size:** 10,000 shoe images across 5 categories
            - **Balance:** Perfectly balanced — 2,000 images per category

            ### Model
            - **Model:** [CLIP ViT-B/32](https://huggingface.co/openai/clip-vit-base-patch32) by OpenAI
            - **Embedding dim:** 512
            - **Why CLIP?** Supports both image AND text queries in the same embedding space

            ### Pipeline
            1. All 10,000 shoe images → CLIP → 512-dim normalized embeddings
            2. Augmented with pixel features (brightness, RGB) for clustering
            3. K-Means (K=5) clustering on augmented 516-dim features
            4. User query → CLIP → cosine similarity → Top-3 results

            ### Clustering Results
            - Algorithm: K-Means (K=5)
            - Evaluation: Silhouette Score = 0.29 (K=5)
            - Visualization: t-SNE + UMAP
            """)

    gr.Markdown("""
    ---
    **Model:** CLIP ViT-B/32 | **Dataset:** Shoe-Net-10K (10K images) | **Method:** Cosine Similarity
    """)

demo.launch()