from PIL import Image
import torch
import gradio as gr

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

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

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

def inference(img, ver):
    """Convert portrait to anime style"""
    if ver == 'Version 2':
        out = face2paint(model2, img)
    else:
        out = face2paint(model1, img)
    return out

# Custom CSS for modern, mobile-friendly design
custom_css = """
/* Container styling */
.gradio-container {
    max-width: 1200px !important;
    margin: auto !important;
}

/* Header styling */
#header {
    text-align: center;
    padding: 2rem 1rem 1rem 1rem;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 16px;
    margin-bottom: 2rem;
    color: white;
}

#header h1 {
    font-size: 2.5rem;
    font-weight: 800;
    margin-bottom: 0.5rem;
    letter-spacing: -0.5px;
}

#header p {
    font-size: 1.1rem;
    opacity: 0.95;
    max-width: 600px;
    margin: 0 auto;
    line-height: 1.6;
}

/* Main content area */
#main-content {
    background: rgba(255, 255, 255, 0.05);
    border-radius: 16px;
    padding: 2rem;
    margin-bottom: 2rem;
}

/* Image containers */
.image-container {
    border-radius: 12px;
    overflow: hidden;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}

/* Button styling */
.primary-btn button {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    font-weight: 600 !important;
    font-size: 1.1rem !important;
    padding: 0.75rem 2rem !important;
    border-radius: 12px !important;
    transition: transform 0.2s ease !important;
}

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

/* Radio button styling */
.version-selector label {
    font-weight: 600 !important;
    font-size: 1rem !important;
    margin-bottom: 0.75rem !important;
}

/* Examples styling */
#examples {
    margin-top: 2rem;
}

/* Footer */
#footer {
    text-align: center;
    padding: 2rem 1rem;
    margin-top: 2rem;
    border-top: 1px solid rgba(255, 255, 255, 0.1);
}

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

#footer a:hover {
    color: #764ba2;
}

/* Mobile responsiveness */
@media (max-width: 768px) {
    #header h1 {
        font-size: 2rem;
    }
    
    #header p {
        font-size: 1rem;
    }
    
    #main-content {
        padding: 1rem;
    }
}
"""

# 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")
    
    # 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: 1rem; flex-wrap: wrap;'>
                <a href='https://github.com/bryandlee/animegan2-pytorch' target='_blank'>
                    📚 GitHub Repository
                </a>
                <span style='color: rgba(255,255,255,0.3);'>•</span>
                <a href='https://huggingface.co/spaces/akhaliq/anycoder' target='_blank'>
                    🚀 Built with anycoder
                </a>
            </div>
            
            <p style='margin-top: 1rem; font-size: 0.9rem; color: rgba(255,255,255,0.6);'>
                💡 <strong>Tip:</strong> Use a cropped portrait photo for best results
            </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 modern theme
demo.launch(
    theme=gr.themes.Soft(
        primary_hue="indigo",
        secondary_hue="purple",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        text_size="lg",
        spacing_size="lg",
        radius_size="lg"
    ).set(
        button_primary_background_fill="*primary_600",
        button_primary_background_fill_hover="*primary_700",
        block_title_text_weight="600",
        block_label_text_weight="600",
    ),
    css=custom_css,
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
    ]
)