import spaces
import gradio as gr
import torch
from transformers import CLIPModel, CLIPProcessor
from PIL import Image

# ─── Model Loading (module-level for ZeroGPU) ───────────────────────────
MODEL_ID = "kwanY/styleid"
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()

# ─── Inference ───────────────────────────────────────────────────────────

@spaces.GPU(duration=30)
def compare_faces(img1: Image.Image, img2: Image.Image):
    """Compare face identity between two images using StyleID."""
    if img1 is None or img2 is None:
        return "⚠️ Please upload both images.", ""

    def get_embedding(img):
        inputs = processor(images=img.convert("RGB"), return_tensors="pt").to(device)
        with torch.no_grad():
            emb = model.get_image_features(**inputs)
            emb = emb / emb.norm(dim=-1, keepdim=True)
        return emb

    e1 = get_embedding(img1)
    e2 = get_embedding(img2)
    score = (e1 * e2).sum(dim=-1).item()

    # Interpretation based on paper thresholds
    if score >= 0.55:
        verdict = "✅ Confidently Same Person"
        color = "#22c55e"
    elif score >= 0.45:
        verdict = "✅ Very Likely Same Person"
        color = "#4ade80"
    elif score >= 0.35:
        verdict = "🟡 Probably Same Person"
        color = "#facc15"
    elif score >= 0.25:
        verdict = "🟠 Uncertain — Borderline"
        color = "#fb923c"
    else:
        verdict = "❌ Likely Different People"
        color = "#ef4444"

    # Create a visual gauge
    score_pct = max(0, min(100, int(score * 100)))
    gauge_html = f"""
    <div style="text-align:center; padding: 20px;">
        <div style="font-size: 48px; font-weight: bold; color: {color};">{score:.3f}</div>
        <div style="font-size: 20px; margin-top: 8px; color: {color}; font-weight: 600;">{verdict}</div>
        <div style="margin-top: 16px; background: #1f2937; border-radius: 12px; height: 24px; overflow: hidden; position: relative;">
            <div style="background: linear-gradient(90deg, #ef4444, #fb923c, #facc15, #4ade80, #22c55e); height: 100%; width: {score_pct}%; border-radius: 12px; transition: width 0.5s;"></div>
        </div>
        <div style="display: flex; justify-content: space-between; font-size: 11px; color: #9ca3af; margin-top: 4px;">
            <span>Different</span>
            <span>Uncertain</span>
            <span>Same</span>
        </div>
    </div>
    """

    details = f"""### How to Read the Score
| Range | Meaning |
|-------|---------|
| **> 0.55** | ✅ Confidently same person |
| **0.45 – 0.55** | ✅ Very likely same person |
| **0.35 – 0.45** | 🟡 Probably same person |
| **0.25 – 0.35** | 🟠 Uncertain / borderline |
| **< 0.25** | ❌ Likely different people |

**Your Score: {score:.4f}**

> **Note:** StyleID is optimized for comparing a real photo against its stylized version (anime, cartoon, painting, sketch, etc.). Works best with a single prominent face per image.
"""

    return gauge_html, details

# ─── UI ──────────────────────────────────────────────────────────────────

DESCRIPTION = """
# 🎭 StyleID — Does AI Preserve Your Face?

Upload a **real photo** and a **stylized portrait** (anime, cartoon, painting, sketch, etc.)  
→ StyleID checks if the identity is preserved across art styles.

**Powered by [StyleID](https://huggingface.co/kwanY/styleid)** — a perception-aware face identity model trained to match human judgments across stylizations.

### 🔬 How it works
StyleID uses a fine-tuned CLIP ViT-L/14 encoder calibrated against human psychometric experiments.
Unlike standard face recognition (ArcFace, AdaFace), StyleID is specifically designed for **cross-style** identity matching — 
real photo ↔ anime, cartoon, painting, caricature, 3D render, etc.

### ⚡ Use Cases
- Check if your AI-stylized portrait still looks like you
- Evaluate face-swap / face-transfer quality
- Compare identity preservation across different AI art tools
- Research: benchmark stylization models for ID preservation
"""

TIPS = """
### 💡 Tips for Best Results
1. **One face per image** — StyleID doesn't do face detection
2. **Face should be prominent** — at least 25% of the image
3. **Frontal or 3/4 view works best** — extreme profiles may reduce scores
4. **Works across styles** — anime, cartoon, oil painting, watercolor, pencil sketch, 3D render, caricature
5. **NOT a security tool** — for research/evaluation only
"""

with gr.Blocks(
    title="StyleID — Face Identity Across Art Styles",
    theme=gr.themes.Soft(
        primary_hue="violet",
        secondary_hue="blue",
    ),
    css="footer { display: none !important; }"
) as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column():
            img1 = gr.Image(type="pil", label="📸 Original Photo (Real Face)", height=400)
        with gr.Column():
            img2 = gr.Image(type="pil", label="🎨 Stylized Portrait (Anime, Cartoon, etc.)", height=400)

    btn = gr.Button("🔍 Compare Identity", variant="primary", size="lg")

    with gr.Row():
        with gr.Column():
            gauge = gr.HTML(label="Identity Score")
        with gr.Column():
            details = gr.Markdown(label="Details")

    btn.click(
        fn=compare_faces,
        inputs=[img1, img2],
        outputs=[gauge, details],
    )

    with gr.Accordion("💡 Tips & Info", open=False):
        gr.Markdown(TIPS)

    gr.Markdown("""
---
**Model:** [kwanY/styleid](https://huggingface.co/kwanY/styleid) | **Paper:** [StyleID (arXiv 2604.21689)](https://arxiv.org/abs/2604.21689) | **License:** Research use only — not for biometric authentication.
""")

demo.launch()