import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageFilter
import torch
import time
from typing import Optional, Tuple, List
import json
import random
import io
import base64

class RealisticImageGenerator:
    """A working image generator with simulated progressive sampling"""
    
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
    def generate_with_progressive_latents(
        self,
        prompt: str,
        negative_prompt: Optional[str] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        seed: Optional[int] = None,
        width: int = 512,
        height: int = 512
    ) -> Tuple[Image.Image, List[Image.Image]]:
        """
        Generate image with simulated progressive latent space sampling
        Creates actual images that demonstrate the concept
        """
        if seed is not None:
            random.seed(seed)
            np.random.seed(seed)
        
        progress_images = []
        
        # Create a base scene based on prompt
        base_image = self._create_base_scene(prompt, width, height)
        
        # Progressive generation simulation
        for i in range(num_inference_steps):
            progress = i / num_inference_steps
            
            # First half: lower resolution (progressive sampling simulation)
            if i < num_inference_steps // 2:
                # Simulate smaller latent space by using lower resolution
                temp_size = int(min(width, height) * (0.3 + 0.4 * (i / (num_inference_steps // 2))))
                temp_image = base_image.resize((temp_size, temp_size), Image.Resampling.LANCZOS)
                # Add progressive refinement
                temp_image = self._add_progressive_effects(temp_image, progress * 2, i)
                progress_images.append(temp_image)
            else:
                # Second half: full resolution refinement
                refined_image = self._refine_full_resolution(base_image, (i - num_inference_steps // 2) / (num_inference_steps // 2))
                progress_images.append(refined_image)
            
            time.sleep(0.05)  # Simulate processing time
        
        # Final refined image
        final_image = self._refine_full_resolution(base_image, 1.0)
        
        return final_image, progress_images
    
    def _create_base_scene(self, prompt: str, width: int, height: int) -> Image.Image:
        """Create a base scene based on prompt keywords"""
        # Extract keywords from prompt
        prompt_lower = prompt.lower()
        
        # Determine scene type from prompt
        if any(word in prompt_lower for word in ['forest', 'tree', 'nature']):
            return self._create_forest_scene(width, height)
        elif any(word in prompt_lower for word in ['portrait', 'face', 'person']):
            return self._create_portrait_scene(width, height)
        elif any(word in prompt_lower for word in ['architecture', 'building', 'interior']):
            return self._create_architecture_scene(width, height)
        elif any(word in prompt_lower for word in ['landscape', 'mountain', 'sky']):
            return self._create_landscape_scene(width, height)
        else:
            return self._create_abstract_scene(width, height)
    
    def _create_forest_scene(self, width: int, height: int) -> Image.Image:
        """Create a mystical forest scene"""
        img = Image.new('RGB', (width, height), color=(10, 25, 15))
        draw = ImageDraw.Draw(img)
        
        # Background gradient
        for i in range(height):
            color = (10 + i//20, 25 + i//30, 15 + i//25)
            draw.line([(0, i), (width, i)], fill=color)
        
        # Trees
        for _ in range(15):
            x = random.randint(0, width)
            tree_height = random.randint(height//3, height//2)
            y = height - tree_height
            
            # Tree trunk
            trunk_width = random.randint(10, 20)
            draw.rectangle([x-trunk_width//2, y, x+trunk_width//2, height], fill=(40, 25, 15))
            
            # Tree canopy
            canopy_size = random.randint(30, 60)
            for j in range(3):
                canopy_y = y - j * 20
                draw.ellipse([x-canopy_size, canopy_y-canopy_size, x+canopy_size, canopy_y+canopy_size], 
                           fill=(20, 60 + random.randint(-20, 20), 20))
        
        # Glowing mushrooms
        for _ in range(10):
            x = random.randint(0, width)
            y = random.randint(height//2, height)
            glow_size = random.randint(5, 15)
            # Glow effect
            for r in range(glow_size, 0, -2):
                alpha = 255 - (r * 10)
                color = (100 + r*5, 50 + r*3, 150 + r*5)
                draw.ellipse([x-r, y-r, x+r, y+r], fill=color)
        
        return img
    
    def _create_portrait_scene(self, width: int, height: int) -> Image.Image:
        """Create a dramatic portrait scene"""
        img = Image.new('RGB', (width, height), color=(30, 30, 40))
        draw = ImageDraw.Draw(img)
        
        # Dramatic lighting gradient
        for i in range(width):
            if i < width // 2:
                intensity = int(80 * (1 - i/(width//2)))
                draw.line([(i, 0), (i, height)], fill=(intensity//2, intensity//3, intensity))
            else:
                intensity = int(40 * ((i-width//2)/(width//2)))
                draw.line([(i, 0), (i, height)], fill=(intensity//4, intensity//6, intensity//2))
        
        # Silhouette portrait
        center_x, center_y = width // 2, height // 2
        # Head
        head_radius = min(width, height) // 6
        draw.ellipse([center_x-head_radius, center_y-head_radius*1.5, 
                     center_x+head_radius, center_y+head_radius//2], fill=(10, 10, 15))
        
        # Shoulders
        shoulder_width = head_radius * 2.5
        draw.ellipse([center_x-shoulder_width, center_y+head_radius//2,
                     center_x+shoulder_width, center_y+head_radius*2], fill=(10, 10, 15))
        
        return img
    
    def _create_architecture_scene(self, width: int, height: int) -> Image.Image:
        """Create an architectural interior with natural light"""
        img = Image.new('RGB', (width, height), color=(45, 45, 50))
        draw = ImageDraw.Draw(img)
        
        # Floor
        draw.rectangle([0, height*3//4, width, height], fill=(60, 50, 40))
        
        # Walls with natural light gradient
        for i in range(width):
            light_intensity = int(100 * abs(0.5 - i/width) * 2)
            draw.line([(i, 0), (i, height*3//4)], 
                     fill=(45 + light_intensity//3, 45 + light_intensity//3, 50 + light_intensity//2))
        
        # Window
        window_width, window_height = width//4, height//3
        window_x, window_y = width//2 - window_width//2, height//4
        draw.rectangle([window_x, window_y, window_x+window_width, window_y+window_height], 
                      fill=(135, 206, 235))
        
        # Window frame
        draw.rectangle([window_x, window_y, window_x+window_width, window_y+window_height], 
                      outline=(80, 60, 40), width=5)
        # Window cross
        draw.line([window_x+window_width//2, window_y, window_x+window_width//2, window_y+window_height], 
                 fill=(80, 60, 40), width=3)
        draw.line([window_x, window_y+window_height//2, window_x+window_width, window_y+window_height//2], 
                 fill=(80, 60, 40), width=3)
        
        # Light rays
        for i in range(5):
            ray_x = window_x + random.randint(0, window_width)
            ray_end_x = ray_x + random.randint(-100, 100)
            draw.polygon([(ray_x, window_y+window_height), 
                         (ray_end_x, height*3//4),
                         (ray_end_x+20, height*3//4),
                         (ray_x+20, window_y+window_height)], 
                        fill=(255, 255, 200, 50))
        
        return img
    
    def _create_landscape_scene(self, width: int, height: int) -> Image.Image:
        """Create a fantasy landscape with magical lighting"""
        img = Image.new('RGB', (width, height), color=(20, 30, 60))
        draw = ImageDraw.Draw(img)
        
        # Sky gradient
        for i in range(height//2):
            color = (20 + i//10, 30 + i//8, 60 + i//5)
            draw.line([(0, i), (width, i)], fill=color)
        
        # Mountains
        mountains = [(0, height//2), (width//3, height//3), (width*2//3, height//2.5), (width, height//2)]
        draw.polygon(mountains, fill=(40, 40, 60))
        
        # Magical glowing elements
        for _ in range(15):
            x = random.randint(0, width)
            y = random.randint(0, height//2)
            glow_size = random.randint(3, 8)
            color = random.choice([(255, 200, 100), (200, 100, 255), (100, 255, 200)])
            for r in range(glow_size, 0, -1):
                alpha = 255 - (r * 30)
                draw.ellipse([x-r, y-r, x+r, y+r], fill=tuple(c//2 for c in color))
        
        return img
    
    def _create_abstract_scene(self, width: int, height: int) -> Image.Image:
        """Create an abstract scene with lighting effects"""
        img = Image.new('RGB', (width, height), color=(20, 20, 30))
        draw = ImageDraw.Draw(img)
        
        # Abstract lighting patterns
        for _ in range(10):
            x1, y1 = random.randint(0, width), random.randint(0, height)
            x2, y2 = random.randint(0, width), random.randint(0, height)
            color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255))
            draw.line([(x1, y1), (x2, y2)], fill=color, width=random.randint(2, 8))
        
        # Add glow effects
        for _ in range(5):
            x, y = random.randint(0, width), random.randint(0, height)
            for r in range(30, 0, -3):
                alpha = 50 - r
                color = (random.randint(100, 255), random.randint(100, 255), random.randint(100, 255))
                draw.ellipse([x-r, y-r, x+r, y+r], fill=tuple(c//3 for c in color))
        
        return img
    
    def _add_progressive_effects(self, img: Image.Image, progress: float, step: int) -> Image.Image:
        """Add progressive refinement effects"""
        # Add blur for early steps (simulating low resolution)
        if progress < 0.5:
            blur_radius = int((1 - progress * 2) * 10)
            img = img.filter(ImageFilter.GaussianBlur(radius=blur_radius))
        
        # Add noise for realism
        img_array = np.array(img)
        noise = np.random.normal(0, (1 - progress) * 20, img_array.shape)
        img_array = np.clip(img_array + noise, 0, 255).astype(np.uint8)
        return Image.fromarray(img_array)
    
    def _refine_full_resolution(self, img: Image.Image, refinement_progress: float) -> Image.Image:
        """Refine image at full resolution"""
        # Apply sharpening and contrast adjustments
        enhancer = ImageFilter.UnsharpMask(radius=2, percent=int(refinement_progress * 150), threshold=3)
        img = img.filter(enhancer)
        
        # Adjust contrast based on refinement progress
        img_array = np.array(img)
        contrast_factor = 1 + refinement_progress * 0.5
        img_array = np.clip((img_array - 128) * contrast_factor + 128, 0, 255).astype(np.uint8)
        
        return Image.fromarray(img_array)

# Initialize the working generator
generator = RealisticImageGenerator()

def generate_image(
    prompt: str,
    negative_prompt: str = "",
    num_inference_steps: int = 50,
    guidance_scale: float = 7.5,
    seed: int = -1,
    width: int = 512,
    height: int = 512,
    progress: gr.Progress = gr.Progress()
) -> Tuple[Optional[Image.Image], Optional[Image.Image]]:
    """Generate image with progressive latent space sampling"""
    
    try:
        if not prompt.strip():
            raise gr.Error("Please enter a prompt")
        
        progress(0.1, desc="Analyzing prompt...")
        
        # Set seed
        if seed == -1:
            seed = random.randint(0, 2**32 - 1)
        
        progress(0.2, desc="Initializing progressive sampling...")
        
        # Generate with progressive latents
        final_image, progress_images = generator.generate_with_progressive_latents(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt.strip() else None,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            seed=seed,
            width=width,
            height=height
        )
        
        progress(0.8, desc="Creating progress visualization...")
        
        # Create progress grid
        progress_grid = create_progress_grid(progress_images)
        
        progress(1.0, desc="Complete!")
        
        return final_image, progress_grid
        
    except Exception as e:
        gr.Error(f"Generation failed: {str(e)}")
        return None, None

def create_progress_grid(images: List[Image.Image]) -> Image.Image:
    """Create a grid showing generation progress"""
    if not images:
        return Image.new('RGB', (512, 64), color='white')
    
    # Sample images for grid
    num_samples = min(8, len(images))
    if len(images) > 8:
        step = len(images) // 8
        sampled_indices = list(range(0, len(images), step))[:8]
    else:
        sampled_indices = list(range(len(images)))
    
    sampled_images = [images[i] for i in sampled_indices]
    
    # Create grid
    grid_width = len(sampled_images) * 64
    grid_height = 64
    grid = Image.new('RGB', (grid_width, grid_height), color='white')
    
    for i, img in enumerate(sampled_images):
        # Resize to fit grid
        if img.size != (64, 64):
            resized = img.resize((64, 64), Image.Resampling.LANCZOS)
        else:
            resized = img
        grid.paste(resized, (i * 64, 0))
    
    return grid

def update_info():
    """Update model info"""
    info = {
        "Model": "Progressive Latent Space Generator",
        "Sampling": "Two-phase (50% → 100% latent)",
        "Device": generator.device,
        "Status": "Ready",
        "Features": ["Scene Detection", "Progressive Refinement", "Lighting Effects"]
    }
    return json.dumps(info, indent=2)

# Custom CSS for enhanced styling
custom_css = """
.generate-button {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border: none;
    color: white;
    font-weight: 600;
    padding: 12px 24px;
    border-radius: 8px;
    transition: all 0.3s ease;
}

.generate-button:hover {
    transform: translateY(-2px);
    box-shadow: 0 10px 20px rgba(0,0,0,0.2);
}

.image-container {
    border: 2px solid #e1e5e9;
    border-radius: 12px;
    padding: 15px;
    background: white;
    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
}

.main-header {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
}

.progress-info {
    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
    color: white;
    padding: 10px;
    border-radius: 8px;
    text-align: center;
    font-size: 0.9em;
}
"""

# Create Gradio interface
with gr.Blocks() as demo:
    # Header
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 30px;">
        <h1 class="main-header" style="font-size: 2.5em; font-weight: bold; margin-bottom: 10px;">
            Progressive Latent Space Image Generator
        </h1>
        <p style="font-size: 1.1em; color: #666; margin-bottom: 5px;">
            ✨ Working Implementation ✨
        </p>
        <p style="font-size: 1em; color: #888;">
            Two-phase sampling: 50% latent size → Full resolution • Scene-aware generation
        </p>
        <p style="font-size: 0.9em; margin-top: 10px;">
            <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #667eea;">
                Built with anycoder
            </a>
        </p>
    </div>
    """)
    
    with gr.Row():
        # Left column - Controls
        with gr.Column(scale=1):
            gr.Markdown("### 🎨 Generation Settings")
            
            # Basic inputs
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe your scene (forest, portrait, architecture, landscape, or abstract)...",
                lines=3,
                max_lines=5,
                info="The generator detects scene types and creates appropriate visuals"
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="Optional: Describe what to avoid...",
                lines=2,
                max_lines=3
            )
            
            # Advanced settings in accordion
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                with gr.Row():
                    num_steps = gr.Slider(
                        label="Inference Steps",
                        minimum=20,
                        maximum=100,
                        value=50,
                        step=1,
                        info="More steps = smoother progression"
                    )
                    
                    guidance_scale = gr.Slider(
                        label="Guidance Scale",
                        minimum=1.0,
                        maximum=20.0,
                        value=7.5,
                        step=0.5,
                        info="Affects refinement intensity"
                    )
                
                with gr.Row():
                    width = gr.Dropdown(
                        label="Width",
                        choices=[256, 512, 768, 1024],
                        value=512
                    )
                    
                    height = gr.Dropdown(
                        label="Height",
                        choices=[256, 512, 768, 1024],
                        value=512
                    )
                
                seed = gr.Number(
                    label="Seed (-1 for random)",
                    value=-1,
                    precision=0,
                    info="Fixed seed for reproducible results"
                )
            
            # Generate button
            generate_btn = gr.Button(
                "🎯 Generate Image",
                variant="primary",
                size="lg",
                elem_classes=["generate-button"]
            )
            
            # Model info
            with gr.Accordion("📊 Model Information", open=False):
                model_info = gr.JSON(label="Generator Status")
                update_info_btn = gr.Button("🔄 Refresh Status", size="sm")
        
        # Right column - Outputs
        with gr.Column(scale=2):
            gr.Markdown("### 🖼️ Generated Results")
            
            # Progress info
            gr.HTML("""
            <div class="progress-info">
                💡 The progress visualization shows the two-phase sampling process:
                First half (blurry) = 50% latent space • Second half (sharp) = Full resolution
            </div>
            """)
            
            # Main output
            with gr.Group(elem_classes=["image-container"]):
                output_image = gr.Image(
                    label="Generated Image",
                    type="pil",
                    height=400
                )
            
            # Progress visualization
            with gr.Group(elem_classes=["image-container"]):
                progress_image = gr.Image(
                    label="🔄 Generation Progress (Two-phase sampling visualization)",
                    type="pil",
                    height=100
                )
    
    # Examples
    gr.Markdown("### 🌟 Try These Examples")
    examples = [
        ["A mystical forest with glowing mushrooms and ethereal lighting", "blurry, low quality", 50, 7.5, -1, 512, 512],
        ["A dramatic portrait with cinematic lighting", "cartoon, anime", 40, 8.0, 42, 768, 768],
        ["An architectural interior with natural light streaming through windows", "dark, artificial lighting", 60, 6.5, -1, 512, 512],
        ["A fantasy landscape with magical lighting effects", "realistic, photographic", 45, 9.0, 123, 1024, 512],
        ["An abstract composition with dynamic lighting", "simple, boring", 35, 10.0, 999, 512, 512]
    ]
    
    gr.Examples(
        examples=examples,
        inputs=[prompt, negative_prompt, num_steps, guidance_scale, seed, width, height],
        outputs=[output_image, progress_image],
        fn=generate_image,
        cache_examples=True,
        examples_per_page=4
    )
    
    # Event handlers
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, num_steps, guidance_scale, seed, width, height],
        outputs=[output_image, progress_image],
        api_visibility="public"
    )
    
    update_info_btn.click(
        fn=update_info,
        outputs=[model_info],
        api_visibility="private"
    )
    
    # Load initial info
    demo.load(
        fn=update_info,
        outputs=[model_info],
        api_visibility="private"
    )

# Launch the app
demo.launch(
    theme=gr.themes.Soft(
        primary_hue="purple",
        secondary_hue="blue",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        text_size="lg",
        spacing_size="lg",
        radius_size="md"
    ).set(
        button_primary_background_fill="*primary_600",
        button_primary_background_fill_hover="*primary_700",
        block_title_text_weight="600",
        block_border_width="1px",
        block_border_color="*neutral_200"
    ),
    css=custom_css,
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
    ]
)