from PIL import Image
import torch
import gradio as gr
import spaces

# Load models on CPU initially
model2 = torch.hub.load(
    "AK391/animegan2-pytorch:main",
    "generator",
    pretrained=True,
    device="cpu",
    progress=False
)

model1 = torch.hub.load(
    "AK391/animegan2-pytorch:main", 
    "generator", 
    pretrained="face_paint_512_v1",  
    device="cpu"
)

face2paint = torch.hub.load(
    'AK391/animegan2-pytorch:main', 
    'face2paint', 
    size=512, 
    device="cpu",
    side_by_side=False
)

@spaces.GPU
def inference(img, ver):
    """Convert portrait to anime style"""
    if ver == 'Version 2':
        model = model2
    else:
        model = model1
    
    model = model.to('cuda')
    out = face2paint(model, img, device='cuda')
    
    return out

# Apple-inspired CSS with glassmorphism and smooth animations
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=SF+Pro+Display:wght@300;400;500;600;700&display=swap');

/* Global styling with Apple aesthetics */
.gradio-container {
    max-width: 1400px !important;
    margin: auto !important;
    font-family: -apple-system, BlinkMacSystemFont, 'SF Pro Display', 'Segoe UI', Roboto, sans-serif !important;
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%) !important;
}

/* Glassmorphism header */
#header {
    text-align: center;
    padding: 3rem 2rem 2rem 2rem;
    background: rgba(255, 255, 255, 0.7);
    backdrop-filter: blur(20px) saturate(180%);
    -webkit-backdrop-filter: blur(20px) saturate(180%);
    border-radius: 24px;
    margin-bottom: 2rem;
    border: 1px solid rgba(255, 255, 255, 0.8);
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.08);
}

#header h1 {
    font-size: 3.5rem;
    font-weight: 700;
    margin-bottom: 0.75rem;
    letter-spacing: -1.5px;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
}

#header p {
    font-size: 1.2rem;
    color: #6b7280;
    max-width: 700px;
    margin: 0 auto;
    line-height: 1.7;
    font-weight: 400;
}

/* Main content card with Apple-style elevation */
#main-content {
    background: rgba(255, 255, 255, 0.85);
    backdrop-filter: blur(20px) saturate(180%);
    -webkit-backdrop-filter: blur(20px) saturate(180%);
    border-radius: 24px;
    padding: 2.5rem;
    margin-bottom: 2rem;
    border: 1px solid rgba(255, 255, 255, 0.9);
    box-shadow: 0 12px 48px rgba(0, 0, 0, 0.1);
}

/* Image containers with smooth corners */
.image-container {
    border-radius: 16px;
    overflow: hidden;
    box-shadow: 0 4px 24px rgba(0, 0, 0, 0.08);
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1);
}

.image-container:hover {
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.12);
    transform: translateY(-2px);
}

/* Apple-style button */
.primary-btn button {
    background: linear-gradient(180deg, #667eea 0%, #5a67d8 100%) !important;
    border: none !important;
    font-weight: 600 !important;
    font-size: 1.05rem !important;
    padding: 0.875rem 2.5rem !important;
    border-radius: 12px !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow: 0 4px 16px rgba(102, 126, 234, 0.3) !important;
    letter-spacing: -0.2px !important;
}

.primary-btn button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 24px rgba(102, 126, 234, 0.4) !important;
    background: linear-gradient(180deg, #5a67d8 0%, #4c51bf 100%) !important;
}

.primary-btn button:active {
    transform: translateY(0px) !important;
    box-shadow: 0 2px 8px rgba(102, 126, 234, 0.3) !important;
}

/* Radio button styling - Apple segmented control inspired */
.version-selector label {
    font-weight: 600 !important;
    font-size: 0.95rem !important;
    margin-bottom: 0.75rem !important;
    color: #374151 !important;
    letter-spacing: -0.2px !important;
}

.version-selector .wrap {
    background: rgba(243, 244, 246, 0.8) !important;
    border-radius: 12px !important;
    padding: 0.5rem !important;
}

/* Examples section */
#examples {
    margin-top: 2.5rem;
    background: rgba(255, 255, 255, 0.7);
    backdrop-filter: blur(20px) saturate(180%);
    -webkit-backdrop-filter: blur(20px) saturate(180%);
    border-radius: 20px;
    padding: 2rem;
    border: 1px solid rgba(255, 255, 255, 0.8);
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.06);
}

#examples h3 {
    font-weight: 600;
    letter-spacing: -0.5px;
    color: #1f2937;
    margin-bottom: 1.5rem;
}

/* Example items */
.gradio-examples {
    gap: 1rem !important;
}

.gradio-examples button {
    border-radius: 12px !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    border: 1px solid rgba(229, 231, 235, 0.8) !important;
}

.gradio-examples button:hover {
    transform: scale(1.02) !important;
    box-shadow: 0 4px 16px rgba(0, 0, 0, 0.1) !important;
}

/* Footer with subtle styling */
#footer {
    text-align: center;
    padding: 2.5rem 1rem;
    margin-top: 2rem;
    background: rgba(255, 255, 255, 0.5);
    backdrop-filter: blur(20px) saturate(180%);
    -webkit-backdrop-filter: blur(20px) saturate(180%);
    border-radius: 20px;
    border: 1px solid rgba(255, 255, 255, 0.7);
}

#footer a {
    color: #667eea;
    text-decoration: none;
    font-weight: 600;
    transition: all 0.2s ease;
    letter-spacing: -0.2px;
}

#footer a:hover {
    color: #5a67d8;
    text-decoration: underline;
}

/* Input/Output labels */
.block-label {
    font-weight: 600 !important;
    font-size: 0.95rem !important;
    color: #374151 !important;
    letter-spacing: -0.2px !important;
    margin-bottom: 0.75rem !important;
}

/* Smooth scrollbar */
::-webkit-scrollbar {
    width: 10px;
}

::-webkit-scrollbar-track {
    background: rgba(243, 244, 246, 0.5);
    border-radius: 10px;
}

::-webkit-scrollbar-thumb {
    background: rgba(156, 163, 175, 0.5);
    border-radius: 10px;
    transition: all 0.3s ease;
}

::-webkit-scrollbar-thumb:hover {
    background: rgba(107, 114, 128, 0.7);
}

/* Mobile responsiveness */
@media (max-width: 768px) {
    #header h1 {
        font-size: 2.5rem;
    }
    
    #header p {
        font-size: 1rem;
    }
    
    #main-content {
        padding: 1.5rem;
    }
    
    .primary-btn button {
        width: 100% !important;
    }
}

/* Smooth transitions for all interactive elements */
* {
    transition: background-color 0.3s cubic-bezier(0.4, 0, 0.2, 1),
                border-color 0.3s cubic-bezier(0.4, 0, 0.2, 1),
                color 0.3s cubic-bezier(0.4, 0, 0.2, 1);
}
"""

# Create the Gradio interface
with gr.Blocks(fill_height=True) as demo:
    
    # Header
    with gr.Column(elem_id="header"):
        gr.Markdown("# ⚡ AnimeGAN v2")
        gr.Markdown("Transform your portraits into stunning anime-style artwork with AI")
    
    # Main content
    with gr.Column(elem_id="main-content"):
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="📸 Upload Your Portrait",
                    type="pil",
                    sources=["upload", "clipboard"],
                    height=400,
                    elem_classes="image-container"
                )
                
                version_selector = gr.Radio(
                    choices=[
                        'Version 1',
                        'Version 2'
                    ],
                    value='Version 2',
                    label="🎨 Style Version",
                    info="Version 1: More stylized | Version 2: More robust",
                    elem_classes="version-selector"
                )
                
                with gr.Row():
                    submit_btn = gr.Button(
                        "✨ Generate Anime Style",
                        variant="primary",
                        size="lg",
                        elem_classes="primary-btn"
                    )
            
            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="🎨 Anime Result",
                    type="pil",
                    height=400,
                    interactive=False,
                    elem_classes="image-container"
                )
    
    # Examples section
    with gr.Column(elem_id="examples"):
        gr.Markdown("### 💫 Try These Examples")
        gr.Examples(
            examples=[
                ['groot.jpeg', 'Version 2'],
                ['bill.png', 'Version 1'],
                ['tony.png', 'Version 1'],
                ['elon.png', 'Version 2'],
                ['IU.png', 'Version 1'],
                ['billie.png', 'Version 2'],
                ['will.png', 'Version 2'],
                ['beyonce.png', 'Version 1'],
                ['gongyoo.jpeg', 'Version 1']
            ],
            inputs=[input_image, version_selector],
            outputs=output_image,
            fn=inference,
            cache_examples=False,
            examples_per_page=9
        )
    
    # Footer
    with gr.Column(elem_id="footer"):
        gr.Markdown(
            """
            <div style='display: flex; justify-content: center; align-items: center; gap: 1.5rem; flex-wrap: wrap;'>
                <a href='https://github.com/bryandlee/animegan2-pytorch' target='_blank'>
                    📚 GitHub Repository
                </a>
                <span style='color: rgba(107, 114, 128, 0.3);'>•</span>
                <a href='https://huggingface.co/spaces/akhaliq/anycoder' target='_blank'>
                    🚀 Built with anycoder
                </a>
            </div>
            
            <p style='margin-top: 1.5rem; font-size: 0.95rem; color: #6b7280;'>
                💡 <strong>Tip:</strong> Use a cropped portrait photo for best results
            </p>
            <p style='margin-top: 0.75rem; font-size: 0.9rem; color: #9ca3af;'>
                ⚡ Powered by Zero GPU - efficient GPU usage on Hugging Face Spaces
            </p>
            """
        )
    
    # Event handlers
    submit_btn.click(
        fn=inference,
        inputs=[input_image, version_selector],
        outputs=output_image,
        api_name="generate"
    )
    
    input_image.change(
        fn=inference,
        inputs=[input_image, version_selector],
        outputs=output_image,
        show_progress="minimal"
    )

# Launch with Apple-inspired clean theme
demo.launch(
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="slate",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        text_size="md",
        spacing_size="lg",
        radius_size="lg"
    ).set(
        button_primary_background_fill="linear-gradient(180deg, #667eea 0%, #5a67d8 100%)",
        button_primary_background_fill_hover="linear-gradient(180deg, #5a67d8 0%, #4c51bf 100%)",
        block_title_text_weight="600",
        block_label_text_weight="600",
        block_background_fill="rgba(255, 255, 255, 0.7)",
        input_background_fill="rgba(255, 255, 255, 0.9)",
    ),
    css=custom_css,
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
    ]
)