import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import DiffusionPipeline

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=dtype
).to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

def enhance_prompt_for_pattern(prompt):
    """Add specific terms to ensure seamless, tileable patterns."""
    pattern_terms = [
        "seamless pattern",
        "tileable textile design",
        "repeating pattern",
        "high-quality fabric design",
        "continuous pattern",
    ]
    enhanced_prompt = f"{prompt}, {random.choice(pattern_terms)}, suitable for textile printing, high-quality fabric design, seamless edges"
    return enhanced_prompt

@spaces.GPU()
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, 
          num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    enhanced_prompt = enhance_prompt_for_pattern(prompt)
    generator = torch.Generator().manual_seed(seed)
    image = pipe(
        prompt=enhanced_prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=0.0
    ).images[0]
    
    return image, seed

examples = [
    "geometric Art Deco shapes in gold and navy",
    "delicate floral motifs with small roses and leaves",
    "abstract watercolor spots in pastel colors",
    "traditional paisley design in earth tones",
    "modern minimalist lines and circles",
]

# Enhanced CSS for better visual design and mobile responsiveness
css = """
#col-container {
    margin: 0 auto;
    max-width: 800px !important;
    padding: 20px;
}

.main-title {
    text-align: center;
    color: #2d3748;
    margin-bottom: 1rem;
    font-family: 'Poppins', sans-serif;
}

.subtitle {
    text-align: center;
    color: #4a5568;
    margin-bottom: 2rem;
    font-family: 'Inter', sans-serif;
    font-size: 0.95rem;
    line-height: 1.5;
}

.pattern-input {
    border: 2px solid #e2e8f0;
    border-radius: 10px;
    padding: 12px !important;
    margin-bottom: 1rem !important;
    font-size: 1rem;
    transition: all 0.3s ease;
}

.pattern-input:focus {
    border-color: #4299e1;
    box-shadow: 0 0 0 3px rgba(66, 153, 225, 0.1);
}

.generate-button {
    background-color: #4299e1 !important;
    color: white !important;
    padding: 12px 24px !important;
    border-radius: 8px !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
}

.generate-button:hover {
    background-color: #3182ce !important;
    transform: translateY(-1px);
}

.result-image {
    border-radius: 12px;
    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
    margin-top: 1rem;
}

.advanced-settings {
    margin-top: 1.5rem;
    border: 1px solid #e2e8f0;
    border-radius: 10px;
    padding: 1rem;
}

/* Mobile Responsiveness */
@media (max-width: 768px) {
    #col-container {
        padding: 12px;
    }
    
    .main-title {
        font-size: 1.5rem;
    }
    
    .subtitle {
        font-size: 0.9rem;
    }
    
    .pattern-input {
        font-size: 0.9rem;
    }
}

/* Custom styling for examples section */
.examples-section {
    margin-top: 2rem;
    padding: 1rem;
    background: #f7fafc;
    border-radius: 10px;
}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🎨 Deradh's AI Pattern Master
            """,
            elem_classes=["main-title"]
        )
        
        gr.Markdown(
            """
            Create beautiful, seamless patterns for your textile designs using AI. 
            Simply describe your desired pattern, and watch as AI brings your vision to life with 
            professional-quality, repeatable patterns perfect for fabrics and materials.
            """,
            elem_classes=["subtitle"]
        )
        
        with gr.Row():
            prompt = gr.Text(
                label="Pattern Description",
                show_label=False,
                max_lines=1,
                placeholder="Describe your dream pattern (e.g., 'geometric Art Deco shapes in gold and navy')",
                container=False,
                elem_classes=["pattern-input"]
            )
            run_button = gr.Button(
                "✨ Generate",
                scale=0,
                elem_classes=["generate-button"]
            )
        
        result = gr.Image(
            label="Your Generated Pattern",
            show_label=True,
            elem_classes=["result-image"]
        )
        
        with gr.Accordion("🔧 Advanced Settings", open=False):
            with gr.Group(elem_classes=["advanced-settings"]):
                seed = gr.Slider(
                    label="Pattern Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                )
                randomize_seed = gr.Checkbox(
                    label="Randomize Pattern",
                    value=True
                )
                
                with gr.Row():
                    width = gr.Slider(
                        label="Width",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=1024,
                    )
                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=1024,
                    )
                
                num_inference_steps = gr.Slider(
                    label="Generation Quality (Steps)",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4,
                )
        
        with gr.Group(elem_classes=["examples-section"]):
            gr.Markdown("### 💫 Try These Examples")
            gr.Examples(
                examples=examples,
                fn=infer,
                inputs=[prompt],
                outputs=[result, seed],
                cache_examples="lazy"
            )
        
        gr.on(
            triggers=[run_button.click, prompt.submit],
            fn=infer,
            inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps],
            outputs=[result, seed]
        )

demo.launch()