app.py

"""
Image Evaluator - Gradio App for HuggingFace Spaces

AI image quality assessment using:
- Soft-TIFA for prompt alignment
- VLM-as-Judge for holistic assessment
- Multi-image comparison
- Technical metrics (sharpness, colorfulness, contrast, CLIP)

Powered by Qwen2.5-VL-7B
"""

import gradio as gr
from PIL import Image
from typing import Optional, List
import time
import spaces

# Global evaluator (loaded on first use)
evaluator = None


def get_evaluator():
    """Lazy load evaluator."""
    global evaluator
    if evaluator is None:
        from evaluator import ImageEvaluator
        evaluator = ImageEvaluator()
    return evaluator


# Dark theme CSS
DARK_CSS = """
/* Base dark theme */
.gradio-container {
    background-color: #09090b !important;
    color: #fafafa !important;
}

/* Main blocks */
.dark {
    --background-fill-primary: #09090b !important;
    --background-fill-secondary: #18181b !important;
    --border-color-primary: #27272a !important;
    --text-color: #fafafa !important;
    --text-color-subdued: #a1a1aa !important;
}

/* Cards and panels */
.panel, .form, .block {
    background-color: #18181b !important;
    border: 1px solid #27272a !important;
    border-radius: 12px !important;
}

/* Input fields */
input, textarea, select {
    background-color: #18181b !important;
    border: 1px solid #27272a !important;
    color: #fafafa !important;
    border-radius: 8px !important;
}

input:focus, textarea:focus {
    border-color: #3b82f6 !important;
    outline: none !important;
}

/* Buttons */
.primary {
    background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%) !important;
    border: none !important;
    color: white !important;
}

.secondary {
    background-color: #27272a !important;
    border: 1px solid #3f3f46 !important;
    color: #fafafa !important;
}

/* Tabs */
.tab-nav {
    background-color: #09090b !important;
    border-bottom: 1px solid #27272a !important;
}

.tab-nav button {
    color: #a1a1aa !important;
    background: transparent !important;
}

.tab-nav button.selected {
    color: #3b82f6 !important;
    border-bottom: 2px solid #3b82f6 !important;
}

/* Image upload */
.image-container {
    background-color: #18181b !important;
    border: 2px dashed #27272a !important;
    border-radius: 12px !important;
}

/* Labels */
label, .label-wrap {
    color: #a1a1aa !important;
    font-weight: 500 !important;
}

/* Checkboxes */
.checkbox-group {
    background-color: #18181b !important;
    padding: 12px !important;
    border-radius: 8px !important;
}

/* Scrollbar */
::-webkit-scrollbar {
    width: 8px;
    height: 8px;
}

::-webkit-scrollbar-track {
    background: #18181b;
}

::-webkit-scrollbar-thumb {
    background: #3f3f46;
    border-radius: 4px;
}

::-webkit-scrollbar-thumb:hover {
    background: #52525b;
}

/* Result cards */
.result-card {
    background: #18181b;
    border: 1px solid #27272a;
    border-radius: 12px;
    padding: 20px;
    margin: 8px 0;
}

/* Score display */
.score-large {
    font-size: 3em;
    font-weight: 700;
    background: linear-gradient(135deg, #3b82f6 0%, #8b5cf6 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}

/* Winner badge */
.winner-badge {
    background: linear-gradient(135deg, #f59e0b 0%, #d97706 100%);
    color: #000;
    padding: 4px 12px;
    border-radius: 20px;
    font-weight: 600;
    font-size: 0.85em;
}

/* Rank badge */
.rank-badge {
    display: inline-flex;
    align-items: center;
    justify-content: center;
    width: 28px;
    height: 28px;
    border-radius: 50%;
    font-weight: 700;
    font-size: 0.9em;
}

.rank-1 { background: linear-gradient(135deg, #fbbf24 0%, #f59e0b 100%); color: #000; }
.rank-2 { background: linear-gradient(135deg, #94a3b8 0%, #64748b 100%); color: #fff; }
.rank-3 { background: linear-gradient(135deg, #c2855a 0%, #a16207 100%); color: #fff; }
.rank-4 { background: #3f3f46; color: #a1a1aa; }

/* Progress bars */
.progress-bar {
    background: #27272a;
    border-radius: 4px;
    height: 6px;
    overflow: hidden;
}

.progress-fill {
    height: 100%;
    border-radius: 4px;
    transition: width 0.3s ease;
}

.progress-green { background: linear-gradient(90deg, #22c55e 0%, #16a34a 100%); }
.progress-yellow { background: linear-gradient(90deg, #eab308 0%, #ca8a04 100%); }
.progress-red { background: linear-gradient(90deg, #ef4444 0%, #dc2626 100%); }
.progress-blue { background: linear-gradient(90deg, #3b82f6 0%, #2563eb 100%); }

/* Skeleton loading animation */
@keyframes skeleton-pulse {
    0%, 100% { opacity: 0.4; }
    50% { opacity: 0.7; }
}

.skeleton {
    background: linear-gradient(90deg, #27272a 0%, #3f3f46 50%, #27272a 100%);
    background-size: 200% 100%;
    animation: skeleton-pulse 1.5s ease-in-out infinite;
    border-radius: 6px;
}

.skeleton-text {
    height: 14px;
    margin: 8px 0;
}

.skeleton-title {
    height: 20px;
    width: 60%;
    margin-bottom: 16px;
}

.skeleton-score {
    height: 48px;
    width: 120px;
    margin: 12px auto;
}

.skeleton-bar {
    height: 6px;
    width: 100%;
    margin: 8px 0;
}

.skeleton-card {
    background: #18181b;
    border: 1px solid #27272a;
    border-radius: 12px;
    padding: 20px;
}
"""


# Skeleton placeholders for loading states
SKELETON_OVERALL = '''
<div class="skeleton-card" style="margin-bottom: 16px;">
    <div class="skeleton skeleton-text" style="width: 40%;"></div>
    <div class="skeleton skeleton-score"></div>
    <div class="skeleton skeleton-text" style="width: 30%;"></div>
    <div class="skeleton skeleton-bar" style="margin-top: 16px;"></div>
    <div class="skeleton skeleton-text" style="width: 70%; margin-top: 12px;"></div>
</div>
'''

SKELETON_BREAKDOWN = '''
<div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr)); gap: 12px;">
    <div style="background: #27272a; padding: 16px; border-radius: 10px; text-align: center;">
        <div class="skeleton skeleton-text" style="width: 80%; margin: 0 auto;"></div>
        <div class="skeleton" style="height: 28px; width: 60px; margin: 8px auto;"></div>
        <div class="skeleton skeleton-bar"></div>
    </div>
    <div style="background: #27272a; padding: 16px; border-radius: 10px; text-align: center;">
        <div class="skeleton skeleton-text" style="width: 80%; margin: 0 auto;"></div>
        <div class="skeleton" style="height: 28px; width: 60px; margin: 8px auto;"></div>
        <div class="skeleton skeleton-bar"></div>
    </div>
    <div style="background: #27272a; padding: 16px; border-radius: 10px; text-align: center;">
        <div class="skeleton skeleton-text" style="width: 80%; margin: 0 auto;"></div>
        <div class="skeleton" style="height: 28px; width: 60px; margin: 8px auto;"></div>
        <div class="skeleton skeleton-bar"></div>
    </div>
</div>
'''

SKELETON_SOFT_TIFA = '''
<div class="skeleton-card" style="margin-top: 16px;">
    <div class="skeleton skeleton-title"></div>
    <div class="skeleton skeleton-text" style="width: 50%;"></div>
    <div style="margin-top: 12px;">
        <div style="display: flex; justify-content: space-between; padding: 8px 0; border-bottom: 1px solid #27272a;">
            <div class="skeleton skeleton-text" style="width: 60%;"></div>
            <div class="skeleton skeleton-text" style="width: 15%;"></div>
        </div>
        <div style="display: flex; justify-content: space-between; padding: 8px 0; border-bottom: 1px solid #27272a;">
            <div class="skeleton skeleton-text" style="width: 50%;"></div>
            <div class="skeleton skeleton-text" style="width: 15%;"></div>
        </div>
        <div style="display: flex; justify-content: space-between; padding: 8px 0; border-bottom: 1px solid #27272a;">
            <div class="skeleton skeleton-text" style="width: 70%;"></div>
            <div class="skeleton skeleton-text" style="width: 15%;"></div>
        </div>
    </div>
</div>
'''

SKELETON_VLM = '''
<div class="skeleton-card">
    <div class="skeleton skeleton-title"></div>
    <div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 12px;">
        <div class="skeleton skeleton-text"></div>
        <div class="skeleton skeleton-text"></div>
        <div class="skeleton skeleton-text"></div>
        <div class="skeleton skeleton-text"></div>
    </div>
</div>
'''

SKELETON_TECHNICAL = '''
<div class="skeleton-card">
    <div class="skeleton skeleton-title"></div>
    <div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 8px;">
        <div class="skeleton skeleton-text"></div>
        <div class="skeleton skeleton-text"></div>
        <div class="skeleton skeleton-text"></div>
        <div class="skeleton skeleton-text"></div>
    </div>
</div>
'''

SKELETON_WINNER = '''
<div class="skeleton-card" style="border-color: #3f3f46; text-align: center; padding: 32px;">
    <div class="skeleton skeleton-text" style="width: 20%; margin: 0 auto;"></div>
    <div class="skeleton" style="height: 56px; width: 150px; margin: 16px auto;"></div>
    <div class="skeleton" style="height: 32px; width: 120px; margin: 12px auto; border-radius: 20px;"></div>
    <div class="skeleton skeleton-text" style="width: 60%; margin: 16px auto;"></div>
</div>
'''

SKELETON_RANKINGS = '''
<div class="skeleton-card" style="margin-top: 16px;">
    <div class="skeleton skeleton-title"></div>
    <div style="margin-top: 16px;">
        <div style="display: flex; gap: 16px; padding: 12px 0; border-bottom: 1px solid #27272a;">
            <div class="skeleton skeleton-text" style="width: 25%;"></div>
            <div class="skeleton" style="height: 28px; width: 28px; border-radius: 50%;"></div>
            <div class="skeleton" style="height: 28px; width: 28px; border-radius: 50%;"></div>
        </div>
        <div style="display: flex; gap: 16px; padding: 12px 0; border-bottom: 1px solid #27272a;">
            <div class="skeleton skeleton-text" style="width: 25%;"></div>
            <div class="skeleton" style="height: 28px; width: 28px; border-radius: 50%;"></div>
            <div class="skeleton" style="height: 28px; width: 28px; border-radius: 50%;"></div>
        </div>
        <div style="display: flex; gap: 16px; padding: 12px 0; border-bottom: 1px solid #27272a;">
            <div class="skeleton skeleton-text" style="width: 25%;"></div>
            <div class="skeleton" style="height: 28px; width: 28px; border-radius: 50%;"></div>
            <div class="skeleton" style="height: 28px; width: 28px; border-radius: 50%;"></div>
        </div>
    </div>
</div>
'''

SKELETON_INDIVIDUAL = '''
<div class="skeleton-card" style="margin-top: 16px;">
    <div class="skeleton skeleton-title"></div>
    <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 16px; margin-top: 16px;">
        <div style="background: #27272a; border-radius: 12px; padding: 16px; text-align: center;">
            <div class="skeleton skeleton-text" style="width: 50%; margin: 0 auto;"></div>
            <div class="skeleton" style="height: 40px; width: 80px; margin: 12px auto;"></div>
            <div class="skeleton" style="height: 24px; width: 40px; margin: 8px auto; border-radius: 4px;"></div>
        </div>
        <div style="background: #27272a; border-radius: 12px; padding: 16px; text-align: center;">
            <div class="skeleton skeleton-text" style="width: 50%; margin: 0 auto;"></div>
            <div class="skeleton" style="height: 40px; width: 80px; margin: 12px auto;"></div>
            <div class="skeleton" style="height: 24px; width: 40px; margin: 8px auto; border-radius: 4px;"></div>
        </div>
    </div>
</div>
'''

SKELETON_EDIT = '''
<div class="skeleton-card" style="padding: 24px;">
    <div class="skeleton skeleton-text" style="width: 40%;"></div>
    <div class="skeleton skeleton-score"></div>
    <div class="skeleton skeleton-text" style="width: 25%;"></div>
    <div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 12px; margin-top: 20px;">
        <div style="background: #27272a; padding: 16px; border-radius: 10px; text-align: center;">
            <div class="skeleton skeleton-text" style="width: 80%; margin: 0 auto;"></div>
            <div class="skeleton" style="height: 28px; width: 60px; margin: 8px auto;"></div>
        </div>
        <div style="background: #27272a; padding: 16px; border-radius: 10px; text-align: center;">
            <div class="skeleton skeleton-text" style="width: 80%; margin: 0 auto;"></div>
            <div class="skeleton" style="height: 28px; width: 60px; margin: 8px auto;"></div>
        </div>
    </div>
</div>
'''

SKELETON_DETAILS = '''
<div class="skeleton-card" style="margin-top: 16px;">
    <div class="skeleton skeleton-title"></div>
    <div style="background: #27272a; border-radius: 12px; padding: 16px; margin-bottom: 12px;">
        <div style="display: flex; justify-content: space-between; margin-bottom: 12px;">
            <div class="skeleton skeleton-text" style="width: 20%;"></div>
            <div class="skeleton" style="height: 24px; width: 60px;"></div>
        </div>
        <div style="display: grid; grid-template-columns: repeat(5, 1fr); gap: 8px;">
            <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                <div class="skeleton skeleton-text" style="width: 70%;"></div>
                <div class="skeleton" style="height: 18px; width: 50px; margin-top: 4px;"></div>
                <div class="skeleton skeleton-bar" style="margin-top: 4px;"></div>
            </div>
            <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                <div class="skeleton skeleton-text" style="width: 70%;"></div>
                <div class="skeleton" style="height: 18px; width: 50px; margin-top: 4px;"></div>
                <div class="skeleton skeleton-bar" style="margin-top: 4px;"></div>
            </div>
            <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                <div class="skeleton skeleton-text" style="width: 70%;"></div>
                <div class="skeleton" style="height: 18px; width: 50px; margin-top: 4px;"></div>
                <div class="skeleton skeleton-bar" style="margin-top: 4px;"></div>
            </div>
            <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                <div class="skeleton skeleton-text" style="width: 70%;"></div>
                <div class="skeleton" style="height: 18px; width: 50px; margin-top: 4px;"></div>
                <div class="skeleton skeleton-bar" style="margin-top: 4px;"></div>
            </div>
            <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                <div class="skeleton skeleton-text" style="width: 70%;"></div>
                <div class="skeleton" style="height: 18px; width: 50px; margin-top: 4px;"></div>
                <div class="skeleton skeleton-bar" style="margin-top: 4px;"></div>
            </div>
        </div>
    </div>
    <div style="background: #27272a; border-radius: 12px; padding: 16px;">
        <div style="display: flex; justify-content: space-between; margin-bottom: 12px;">
            <div class="skeleton skeleton-text" style="width: 20%;"></div>
            <div class="skeleton" style="height: 24px; width: 60px;"></div>
        </div>
        <div style="display: grid; grid-template-columns: repeat(5, 1fr); gap: 8px;">
            <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                <div class="skeleton skeleton-text" style="width: 70%;"></div>
                <div class="skeleton" style="height: 18px; width: 50px; margin-top: 4px;"></div>
                <div class="skeleton skeleton-bar" style="margin-top: 4px;"></div>
            </div>
            <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                <div class="skeleton skeleton-text" style="width: 70%;"></div>
                <div class="skeleton" style="height: 18px; width: 50px; margin-top: 4px;"></div>
                <div class="skeleton skeleton-bar" style="margin-top: 4px;"></div>
            </div>
            <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                <div class="skeleton skeleton-text" style="width: 70%;"></div>
                <div class="skeleton" style="height: 18px; width: 50px; margin-top: 4px;"></div>
                <div class="skeleton skeleton-bar" style="margin-top: 4px;"></div>
            </div>
            <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                <div class="skeleton skeleton-text" style="width: 70%;"></div>
                <div class="skeleton" style="height: 18px; width: 50px; margin-top: 4px;"></div>
                <div class="skeleton skeleton-bar" style="margin-top: 4px;"></div>
            </div>
            <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                <div class="skeleton skeleton-text" style="width: 70%;"></div>
                <div class="skeleton" style="height: 18px; width: 50px; margin-top: 4px;"></div>
                <div class="skeleton skeleton-bar" style="margin-top: 4px;"></div>
            </div>
        </div>
    </div>
</div>
'''


def format_score_card(score: float, label: str, show_bar: bool = True) -> str:
    """Format a score as a styled card."""
    if score >= 0.7:
        color_class = "progress-green"
        text_color = "#22c55e"
    elif score >= 0.5:
        color_class = "progress-yellow"
        text_color = "#eab308"
    else:
        color_class = "progress-red"
        text_color = "#ef4444"

    bar_html = f'''
        <div class="progress-bar" style="margin-top: 8px;">
            <div class="progress-fill {color_class}" style="width: {int(score * 100)}%;"></div>
        </div>
    ''' if show_bar else ""

    return f'''
    <div style="background: #27272a; padding: 16px; border-radius: 10px; text-align: center;">
        <div style="color: #a1a1aa; font-size: 0.85em; margin-bottom: 4px;">{label}</div>
        <div style="color: {text_color}; font-size: 1.4em; font-weight: 700;">{score:.3f}</div>
        {bar_html}
    </div>
    '''


def get_grade_colors(grade: str) -> tuple:
    """Get colors for a grade (F=red, D=orange, C=yellow, B=lime, A=green)."""
    grade_colors = {
        "A": ("#22c55e", "#16a34a"),
        "B": ("#84cc16", "#65a30d"),
        "C": ("#eab308", "#ca8a04"),
        "D": ("#f97316", "#ea580c"),
        "F": ("#ef4444", "#dc2626"),
    }
    return grade_colors.get(grade.upper(), ("#71717a", "#52525b"))


def get_grade_status(grade: str) -> str:
    """Get status text for a grade."""
    statuses = {
        "A": "EXCELLENT",
        "B": "GOOD",
        "C": "AVERAGE",
        "D": "BELOW AVG",
        "F": "POOR",
    }
    return statuses.get(grade.upper(), "UNKNOWN")


def format_grade_badge(grade: str, passed: bool) -> str:
    """Format grade as a badge with color gradient from F (red) to A (green)."""
    color1, color2 = get_grade_colors(grade)
    status = get_grade_status(grade)
    bg = f"linear-gradient(135deg, {color1} 0%, {color2} 100%)"

    return f'''
    <div style="display: flex; align-items: center; gap: 12px; margin-top: 12px;">
        <span style="background: {bg}; color: white; padding: 6px 16px; border-radius: 6px; font-weight: 700; font-size: 1.3em;">
            {grade}
        </span>
        <span style="color: {color1}; font-weight: 600;">
            {status}
        </span>
    </div>
    '''


def format_rank_badge(rank: int) -> str:
    """Format ranking position as badge."""
    return f'<span class="rank-badge rank-{min(rank, 4)}">{rank}</span>'


@spaces.GPU(duration=300)
def evaluate_single(
    image: Image.Image,
    prompt: str,
    include_soft_tifa: bool,
    include_vlm: bool,
    include_technical: bool,
    progress=gr.Progress()
) -> tuple:
    """Evaluate a single AI-generated image."""
    if image is None:
        return ("Please upload an image.", "", "", "", "")

    progress(0.1, desc="Loading models...")
    try:
        eval_instance = get_evaluator()
    except Exception as e:
        return (f"Error loading models: {str(e)}", "", "", "", "")

    progress(0.2, desc="Evaluating image...")
    prompt_text = prompt.strip() if prompt else None

    try:
        result = eval_instance.evaluate(
            image=image,
            prompt=prompt_text,
            include_soft_tifa=include_soft_tifa and bool(prompt_text),
            include_vlm=include_vlm,
            include_technical=include_technical,
        )
    except Exception as e:
        return (f"Evaluation error: {str(e)}", "", "", "", "")

    progress(0.9, desc="Formatting results...")
    score = result.score

    # Overall score card
    overall_html = f'''
    <div style="background: #18181b; border: 1px solid #27272a; border-radius: 16px; padding: 24px; margin-bottom: 16px;">
        <div style="color: #a1a1aa; font-size: 0.9em; text-transform: uppercase; letter-spacing: 0.05em;">Overall Score</div>
        <div class="score-large" style="margin: 8px 0;">{score.overall:.3f}</div>
        {format_grade_badge(score.grade, score.passed)}
        <div style="color: #71717a; font-size: 0.85em; margin-top: 16px;">
            Confidence: {score.confidence:.0%} &nbsp;|&nbsp; Time: {result.evaluation_time:.1f}s
        </div>
        <div style="background: #27272a; padding: 12px 16px; border-radius: 8px; margin-top: 12px; color: #d4d4d8; font-size: 0.9em;">
            {score.recommendation}
        </div>
    </div>
    '''

    # Breakdown scores
    breakdown = score.breakdown
    metrics = [
        ("Prompt Alignment", breakdown.prompt_alignment),
        ("Technical Quality", breakdown.technical_quality),
        ("Aesthetic Appeal", breakdown.aesthetic_appeal),
        ("Realism", breakdown.realism),
        ("Artifacts (inv)", breakdown.artifacts),
    ]

    breakdown_html = '<div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr)); gap: 12px;">'
    for name, value in metrics:
        if value is not None:
            breakdown_html += format_score_card(value, name)
    breakdown_html += '</div>'

    # Soft-TIFA results
    soft_tifa_html = ""
    if result.soft_tifa:
        st = result.soft_tifa
        soft_tifa_html = f'''
        <div style="background: #18181b; border: 1px solid #27272a; border-radius: 12px; padding: 20px; margin-top: 16px;">
            <div style="color: #fafafa; font-weight: 600; margin-bottom: 12px;">Soft-TIFA Analysis</div>
            <div style="color: #a1a1aa; font-size: 0.85em; margin-bottom: 12px;">
                Primitives: {st.primitives_count} &nbsp;|&nbsp; Atom: {st.atom_score:.3f} &nbsp;|&nbsp; Prompt: {st.prompt_score:.3f}
            </div>
            <div style="max-height: 200px; overflow-y: auto;">
        '''
        for pr in st.primitive_results[:10]:
            icon = "●" if pr.score >= 0.7 else "○"
            color = "#22c55e" if pr.score >= 0.7 else "#ef4444"
            soft_tifa_html += f'''
                <div style="display: flex; justify-content: space-between; padding: 8px 0; border-bottom: 1px solid #27272a;">
                    <span style="color: {color};">{icon} {pr.content}</span>
                    <span style="color: #71717a;">{pr.score:.2f}</span>
                </div>
            '''
        soft_tifa_html += '</div></div>'

    # VLM Assessment
    vlm_html = ""
    if result.vlm_assessment:
        vlm = result.vlm_assessment
        vlm_html = f'''
        <div style="background: #18181b; border: 1px solid #27272a; border-radius: 12px; padding: 20px;">
            <div style="color: #fafafa; font-weight: 600; margin-bottom: 12px;">VLM Assessment</div>
            <div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 12px; color: #d4d4d8;">
                <div>Technical: <span style="color: #3b82f6;">{vlm.technical_quality:.1f}/10</span></div>
                <div>Aesthetic: <span style="color: #8b5cf6;">{vlm.aesthetic_appeal:.1f}/10</span></div>
                <div>Realism: <span style="color: #06b6d4;">{vlm.realism:.1f}/10</span></div>
                <div>Overall: <span style="color: #22c55e;">{vlm.overall:.1f}/10</span></div>
            </div>
        '''
        if vlm.artifacts_detected:
            vlm_html += f'''
            <div style="background: #450a0a; border: 1px solid #7f1d1d; padding: 10px 14px; border-radius: 8px; margin-top: 12px; color: #fca5a5; font-size: 0.85em;">
                <strong>Artifacts ({vlm.artifacts_severity}):</strong> {', '.join(vlm.artifacts_detected[:5])}
            </div>
            '''
        vlm_html += '</div>'

    # Technical metrics
    tech_html = ""
    if result.technical_metrics:
        tm = result.technical_metrics
        tech_html = f'''
        <div style="background: #18181b; border: 1px solid #27272a; border-radius: 12px; padding: 20px;">
            <div style="color: #fafafa; font-weight: 600; margin-bottom: 12px;">Technical Metrics</div>
            <div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 8px; color: #d4d4d8; font-size: 0.9em;">
        '''
        if tm.clip_score is not None:
            tech_html += f'<div>CLIP: <span style="color: #3b82f6;">{tm.clip_score:.3f}</span></div>'
        if tm.sharpness is not None:
            tech_html += f'<div>Sharpness: <span style="color: #22c55e;">{tm.sharpness:.3f}</span></div>'
        if tm.colorfulness is not None:
            tech_html += f'<div>Colorfulness: <span style="color: #f59e0b;">{tm.colorfulness:.3f}</span></div>'
        if tm.contrast is not None:
            tech_html += f'<div>Contrast: <span style="color: #8b5cf6;">{tm.contrast:.3f}</span></div>'
        tech_html += '</div></div>'

    return (overall_html, breakdown_html, soft_tifa_html, vlm_html, tech_html)


@spaces.GPU(duration=600)
def compare_images(
    img1: Image.Image,
    img2: Image.Image,
    img3: Image.Image,
    img4: Image.Image,
    prompt: str,
    progress=gr.Progress()
) -> tuple:
    """Compare 2-4 images against a prompt."""
    images = [img for img in [img1, img2, img3, img4] if img is not None]

    if len(images) < 2:
        return ("Please upload at least 2 images to compare.", "", "", "")

    if not prompt.strip():
        return ("Please enter a prompt to compare images against.", "", "", "")

    progress(0.1, desc="Loading models...")
    try:
        eval_instance = get_evaluator()
    except Exception as e:
        return (f"Error loading models: {str(e)}", "", "", "")

    progress(0.2, desc="Comparing images...")
    try:
        result = eval_instance.compare_images(images, prompt.strip())
    except Exception as e:
        return (f"Comparison error: {str(e)}", "", "", "")

    progress(0.9, desc="Formatting results...")

    # Winner announcement
    winner_html = f'''
    <div style="background: linear-gradient(135deg, #18181b 0%, #1f1f23 100%); border: 1px solid #f59e0b; border-radius: 16px; padding: 24px; text-align: center;">
        <div style="color: #f59e0b; font-size: 0.9em; text-transform: uppercase; letter-spacing: 0.1em; margin-bottom: 8px;">Winner</div>
        <div style="font-size: 3em; font-weight: 700; color: #fafafa; margin-bottom: 8px;">
            Image {result.winner_index + 1}
        </div>
        <div style="display: inline-block; background: linear-gradient(135deg, #f59e0b 0%, #d97706 100%); color: #000; padding: 6px 20px; border-radius: 20px; font-weight: 600;">
            BEST OVERALL
        </div>
        <div style="color: #a1a1aa; font-size: 0.9em; margin-top: 16px; max-width: 500px; margin-left: auto; margin-right: auto;">
            {result.winner_reasoning}
        </div>
        <div style="color: #71717a; font-size: 0.8em; margin-top: 12px;">
            Evaluation time: {result.evaluation_time:.1f}s
        </div>
    </div>
    '''

    # Rankings table
    rankings_html = '''
    <div style="background: #18181b; border: 1px solid #27272a; border-radius: 12px; padding: 20px; margin-top: 16px;">
        <div style="color: #fafafa; font-weight: 600; margin-bottom: 16px;">Rankings by Criterion</div>
        <div style="overflow-x: auto;">
            <table style="width: 100%; border-collapse: collapse; color: #d4d4d8; font-size: 0.9em;">
                <thead>
                    <tr style="border-bottom: 1px solid #27272a;">
                        <th style="text-align: left; padding: 12px 8px; color: #a1a1aa;">Criterion</th>
    '''
    for i in range(result.num_images):
        rankings_html += f'<th style="text-align: center; padding: 12px 8px; color: #a1a1aa;">Image {i+1}</th>'
    rankings_html += '</tr></thead><tbody>'

    criteria_labels = {
        "prompt_alignment": "Prompt Alignment (45%)",
        "technical_quality": "Technical Quality (20%)",
        "aesthetic_appeal": "Aesthetic Appeal (15%)",
        "realism": "Realism (10%)"
    }

    for criterion, label in criteria_labels.items():
        if criterion in result.rankings_by_criterion:
            ranking_data = result.rankings_by_criterion[criterion]
            rankings_html += f'<tr style="border-bottom: 1px solid #27272a;"><td style="padding: 12px 8px;">{label}</td>'
            for i in range(result.num_images):
                rank = ranking_data.ranking.index(i + 1) + 1 if (i + 1) in ranking_data.ranking else i + 1
                score = ranking_data.scores[i] if i < len(ranking_data.scores) else 0.5
                rankings_html += f'''
                    <td style="text-align: center; padding: 12px 8px;">
                        {format_rank_badge(rank)}
                        <div style="color: #71717a; font-size: 0.8em; margin-top: 4px;">{score:.3f}</div>
                    </td>
                '''
            rankings_html += '</tr>'

    # Overall row
    rankings_html += '<tr style="background: #27272a;"><td style="padding: 12px 8px; font-weight: 600;">Overall</td>'
    for i in range(result.num_images):
        rank = result.overall_ranking.index(i + 1) + 1 if (i + 1) in result.overall_ranking else i + 1
        score = result.overall_scores[i] if i < len(result.overall_scores) else 0.5
        is_winner = i == result.winner_index
        rankings_html += f'''
            <td style="text-align: center; padding: 12px 8px;">
                {format_rank_badge(rank)}
                {'<span class="winner-badge" style="display: inline-block; margin-left: 8px;">WINNER</span>' if is_winner else ''}
                <div style="color: #71717a; font-size: 0.8em; margin-top: 4px;">{score:.3f}</div>
            </td>
        '''
    rankings_html += '</tr></tbody></table></div></div>'

    # Individual scores summary
    individual_html = '''
    <div style="background: #18181b; border: 1px solid #27272a; border-radius: 12px; padding: 20px; margin-top: 16px;">
        <div style="color: #fafafa; font-weight: 600; margin-bottom: 16px;">Individual Scores Summary</div>
        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 16px;">
    '''

    for i, score in enumerate(result.individual_scores):
        is_winner = i == result.winner_index
        border_color = "#f59e0b" if is_winner else "#27272a"
        grade_color1, grade_color2 = get_grade_colors(score.grade)
        grade_bg = f"linear-gradient(135deg, {grade_color1} 0%, {grade_color2} 100%)"

        individual_html += f'''
        <div style="background: #27272a; border: 2px solid {border_color}; border-radius: 12px; padding: 16px; text-align: center;">
            <div style="color: #a1a1aa; font-size: 0.85em;">Image {i+1}</div>
            <div style="font-size: 2em; font-weight: 700; color: {'#f59e0b' if is_winner else grade_color1}; margin: 8px 0;">
                {score.overall:.3f}
            </div>
            <div style="background: {grade_bg}; color: white; padding: 4px 12px; border-radius: 4px; font-weight: 600; font-size: 0.9em;">
                {score.grade}
            </div>
            {'<div style="margin-top: 8px;"><span class="winner-badge">WINNER</span></div>' if is_winner else ''}
        </div>
        '''

    individual_html += '</div></div>'

    # Detailed breakdown per image
    details_html = '''
    <div style="background: #18181b; border: 1px solid #27272a; border-radius: 12px; padding: 20px; margin-top: 16px;">
        <div style="color: #fafafa; font-weight: 600; margin-bottom: 16px;">Detailed Breakdown</div>
    '''

    for i, (score, eval_result) in enumerate(zip(result.individual_scores, result.individual_results)):
        is_winner = i == result.winner_index
        border_color = "#f59e0b" if is_winner else "#3f3f46"

        details_html += f'''
        <div style="background: #27272a; border: 1px solid {border_color}; border-radius: 12px; padding: 16px; margin-bottom: 12px;">
            <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 12px;">
                <div style="color: #fafafa; font-weight: 600;">
                    Image {i+1} {'<span class="winner-badge" style="margin-left: 8px;">WINNER</span>' if is_winner else ''}
                </div>
                <div style="color: {'#f59e0b' if is_winner else '#3b82f6'}; font-weight: 700; font-size: 1.2em;">
                    {score.overall:.3f}
                </div>
            </div>
            <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(120px, 1fr)); gap: 8px;">
        '''

        # Breakdown metrics
        metrics = [
            ("Prompt", score.breakdown.prompt_alignment, "#3b82f6"),
            ("Technical", score.breakdown.technical_quality, "#22c55e"),
            ("Aesthetic", score.breakdown.aesthetic_appeal, "#8b5cf6"),
            ("Realism", score.breakdown.realism, "#06b6d4"),
            ("Artifacts", score.breakdown.artifacts, "#f59e0b"),
        ]

        for name, value, color in metrics:
            if value is not None:
                bar_width = int(value * 100)
                details_html += f'''
                <div style="background: #18181b; padding: 8px; border-radius: 6px;">
                    <div style="color: #a1a1aa; font-size: 0.75em; margin-bottom: 4px;">{name}</div>
                    <div style="color: {color}; font-weight: 600; font-size: 0.9em;">{value:.3f}</div>
                    <div style="background: #3f3f46; height: 4px; border-radius: 2px; margin-top: 4px;">
                        <div style="background: {color}; height: 100%; width: {bar_width}%; border-radius: 2px;"></div>
                    </div>
                </div>
                '''

        details_html += '</div>'

        # VLM Assessment if available
        if eval_result.vlm_assessment:
            vlm = eval_result.vlm_assessment
            details_html += f'''
            <div style="margin-top: 12px; padding-top: 12px; border-top: 1px solid #3f3f46;">
                <div style="color: #a1a1aa; font-size: 0.8em; margin-bottom: 8px;">VLM Assessment</div>
                <div style="display: flex; gap: 12px; flex-wrap: wrap; font-size: 0.85em;">
                    <span style="color: #d4d4d8;">Tech: <span style="color: #3b82f6;">{vlm.technical_quality:.1f}</span></span>
                    <span style="color: #d4d4d8;">Aesthetic: <span style="color: #8b5cf6;">{vlm.aesthetic_appeal:.1f}</span></span>
                    <span style="color: #d4d4d8;">Realism: <span style="color: #06b6d4;">{vlm.realism:.1f}</span></span>
                    <span style="color: #d4d4d8;">Overall: <span style="color: #22c55e;">{vlm.overall:.1f}</span></span>
                </div>
            '''
            if vlm.artifacts_detected:
                details_html += f'''
                <div style="background: #450a0a; padding: 6px 10px; border-radius: 4px; margin-top: 8px; color: #fca5a5; font-size: 0.8em;">
                    Artifacts ({vlm.artifacts_severity}): {', '.join(vlm.artifacts_detected[:3])}
                </div>
                '''
            details_html += '</div>'

        # Soft-TIFA summary if available
        if eval_result.soft_tifa:
            st = eval_result.soft_tifa
            details_html += f'''
            <div style="margin-top: 12px; padding-top: 12px; border-top: 1px solid #3f3f46;">
                <div style="color: #a1a1aa; font-size: 0.8em; margin-bottom: 4px;">Soft-TIFA</div>
                <div style="font-size: 0.85em; color: #d4d4d8;">
                    Primitives: {st.primitives_count} | Atom: <span style="color: #3b82f6;">{st.atom_score:.3f}</span> | Prompt: <span style="color: #22c55e;">{st.prompt_score:.3f}</span>
                </div>
            </div>
            '''

        details_html += '</div>'

    details_html += '</div>'

    return (winner_html, rankings_html, individual_html, details_html)


# Build Gradio interface
with gr.Blocks(title="Image Evaluator", css=DARK_CSS, theme=gr.themes.Base()) as demo:

    gr.HTML('''
    <div style="text-align: center; padding: 20px 0 30px 0;">
        <h1 style="color: #fafafa; font-size: 2.2em; font-weight: 700; margin: 0;">Image Evaluator</h1>
        <p style="color: #71717a; font-size: 1em; margin-top: 8px;">
            AI image quality assessment powered by <span style="color: #3b82f6;">Qwen2.5-VL-7B</span>
        </p>
    </div>
    ''')

    with gr.Tabs():
        # Single Image Evaluation Tab
        with gr.TabItem("Single Evaluation"):
            with gr.Row():
                with gr.Column(scale=1):
                    image_input = gr.Image(label="Upload Image", type="pil", height=400)
                    prompt_input = gr.Textbox(
                        label="Generation Prompt",
                        placeholder="Enter the prompt used to generate this image...",
                        lines=3
                    )
                    with gr.Row():
                        soft_tifa_check = gr.Checkbox(label="Soft-TIFA", value=True, info="Prompt alignment")
                        vlm_check = gr.Checkbox(label="VLM Judge", value=True, info="Holistic assessment")
                        technical_check = gr.Checkbox(label="Technical", value=True, info="CLIP, sharpness")

                    evaluate_btn = gr.Button("Evaluate Image", variant="primary", size="lg")

                with gr.Column(scale=1):
                    overall_output = gr.HTML(value=SKELETON_OVERALL)
                    breakdown_output = gr.HTML(value=SKELETON_BREAKDOWN)
                    soft_tifa_output = gr.HTML(value=SKELETON_SOFT_TIFA)

            with gr.Row():
                vlm_output = gr.HTML(value=SKELETON_VLM)
                technical_output = gr.HTML(value=SKELETON_TECHNICAL)

            evaluate_btn.click(
                fn=evaluate_single,
                inputs=[image_input, prompt_input, soft_tifa_check, vlm_check, technical_check],
                outputs=[overall_output, breakdown_output, soft_tifa_output, vlm_output, technical_output],
            )

        # Compare Images Tab
        with gr.TabItem("Compare Images"):
            gr.HTML('''
            <div style="color: #a1a1aa; font-size: 0.9em; margin-bottom: 20px; text-align: center;">
                Upload 2-4 images to compare them against a prompt. The AI will rank them across multiple criteria.
            </div>
            ''')

            with gr.Row():
                img1 = gr.Image(label="Image 1", type="pil", height=250)
                img2 = gr.Image(label="Image 2", type="pil", height=250)
                img3 = gr.Image(label="Image 3 (optional)", type="pil", height=250)
                img4 = gr.Image(label="Image 4 (optional)", type="pil", height=250)

            compare_prompt = gr.Textbox(
                label="Comparison Prompt",
                placeholder="Enter the prompt to compare images against...",
                lines=2
            )

            compare_btn = gr.Button("Compare Images", variant="primary", size="lg")

            winner_output = gr.HTML(value=SKELETON_WINNER)
            rankings_output = gr.HTML(value=SKELETON_RANKINGS)
            individual_output = gr.HTML(value=SKELETON_INDIVIDUAL)
            details_output = gr.HTML(value=SKELETON_DETAILS)

            compare_btn.click(
                fn=compare_images,
                inputs=[img1, img2, img3, img4, compare_prompt],
                outputs=[winner_output, rankings_output, individual_output, details_output],
            )

        # Edit Evaluation Tab
        with gr.TabItem("Edit Evaluation"):
            gr.HTML('''
            <div style="color: #a1a1aa; font-size: 0.9em; margin-bottom: 20px; text-align: center;">
                Evaluate image editing quality by comparing source and edited images.
            </div>
            ''')

            with gr.Row():
                with gr.Column():
                    source_input = gr.Image(label="Source Image (Before)", type="pil", height=300)
                with gr.Column():
                    edited_input = gr.Image(label="Edited Image (After)", type="pil", height=300)

            edit_instruction = gr.Textbox(
                label="Edit Instruction",
                placeholder="Enter the editing instruction that was applied...",
                lines=2
            )

            edit_btn = gr.Button("Evaluate Edit", variant="primary", size="lg")
            edit_output = gr.HTML(value=SKELETON_EDIT)

            @spaces.GPU(duration=300)
            def evaluate_edit_handler(source, edited, instruction, progress=gr.Progress()):
                if source is None or edited is None:
                    return "Please upload both source and edited images."
                if not instruction.strip():
                    return "Please enter the edit instruction."

                progress(0.1, desc="Loading models...")
                try:
                    from evaluator import EditEvaluator
                    edit_eval = EditEvaluator()
                except Exception as e:
                    return f"Error loading models: {str(e)}"

                progress(0.3, desc="Evaluating edit...")
                try:
                    result = edit_eval.evaluate(source, edited, instruction.strip())
                except Exception as e:
                    return f"Evaluation error: {str(e)}"

                score = result.score
                return f'''
                <div style="background: #18181b; border: 1px solid #27272a; border-radius: 16px; padding: 24px;">
                    <div style="color: #a1a1aa; font-size: 0.9em; text-transform: uppercase;">Edit Quality Score</div>
                    <div class="score-large" style="margin: 8px 0;">{score.overall:.3f}</div>
                    {format_grade_badge(score.grade, score.passed)}
                    <div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 12px; margin-top: 20px;">
                        {format_score_card(score.breakdown.instruction_following or 0, "Instruction Following")}
                        {format_score_card(score.breakdown.preservation or 0, "Preservation")}
                        {format_score_card(score.breakdown.edit_quality or 0, "Edit Quality")}
                        {format_score_card(score.breakdown.artifacts or 0, "Artifacts (inv)")}
                    </div>
                    <div style="background: #27272a; padding: 12px 16px; border-radius: 8px; margin-top: 16px; color: #d4d4d8; font-size: 0.9em;">
                        {score.recommendation}
                    </div>
                    <div style="color: #71717a; font-size: 0.8em; margin-top: 12px;">
                        Evaluation time: {result.evaluation_time:.1f}s
                    </div>
                </div>
                '''

            edit_btn.click(
                fn=evaluate_edit_handler,
                inputs=[source_input, edited_input, edit_instruction],
                outputs=[edit_output],
            )

    gr.HTML('''
    <div style="text-align: center; padding: 30px 0 20px 0; color: #52525b; font-size: 0.85em;">
        Powered by Qwen2.5-VL-7B &nbsp;|&nbsp; Soft-TIFA &nbsp;|&nbsp; CLIP &nbsp;|&nbsp; LPIPS
    </div>
    ''')


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