import numpy as np
from PIL import Image, ImageDraw
import gradio as gr
import base64
from io import BytesIO
from diffusers import StableDiffusionPipeline
import torch

# Define style names and their corresponding prompts
style_configs = {
    "Watercolor": {
        "prompt": "a serene mountain landscape with lake and sunset, watercolor painting style, yellow sun, artistic",
        "negative_prompt": "digital art, photorealistic, sketch"
    },
    "Cyberpunk": {
        "prompt": "cyberpunk city at night with neon signs in yellow, futuristic buildings, blade runner style",
        "negative_prompt": "daytime, natural, watercolor"
    },
    "Anime": {
        "prompt": "anime character portrait, Studio Ghibli style, yellow hair, bright colors",
        "negative_prompt": "photorealistic, western art"
    },
    "Oil Painting": {
        "prompt": "still life with yellow sunflowers in vase, oil painting style, Van Gogh inspired",
        "negative_prompt": "watercolor, digital art, photograph"
    },
    "Sketch": {
        "prompt": "pencil sketch of a landscape with yellow highlights, detailed drawing",
        "negative_prompt": "color, painting, digital"
    }
}

# Initialize Stable Diffusion pipeline
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
)
pipe = pipe.to(device)

# Cache for generated images
image_cache = {}

def generate_styled_image(style):
    """Generate an image using Stable Diffusion based on style"""
    if style in image_cache:
        return image_cache[style]
    
    config = style_configs[style]
    image = pipe(
        prompt=config["prompt"],
        negative_prompt=config["negative_prompt"],
        num_inference_steps=30,
        guidance_scale=7.5
    ).images[0]
    
    # Cache the generated image
    image_cache[style] = image
    return image

def yellow_loss(image, strength=0.8):
    """Reduces yellow colors in the image."""
    try:
        img_array = np.array(image).astype(np.float32) / 255.0
        r, g, b = img_array[:, :, 0], img_array[:, :, 1], img_array[:, :, 2]
        
        yellow_mask = np.logical_and(np.logical_and(r > 0.5, g > 0.5), b < 0.4)
        
        if np.any(yellow_mask):
            r[yellow_mask] *= (1 - strength * 0.7)
            g[yellow_mask] *= (1 - strength)
            b[yellow_mask] += (1 - b[yellow_mask]) * strength
            
            img_array[:, :, 0] = r
            img_array[:, :, 1] = g
            img_array[:, :, 2] = b
        
        return Image.fromarray((img_array * 255).astype(np.uint8))
    except Exception as e:
        print(f"Error in yellow_loss: {e}")
        return image

def apply_color_loss(style, strength, image_input=None):
    """Apply yellow loss to an image."""
    try:
        if image_input is not None:
            try:
                image = Image.fromarray(image_input) if isinstance(image_input, np.ndarray) else image_input
                image.thumbnail((512, 512), Image.LANCZOS)
            except Exception as e:
                print(f"Error processing input image: {e}")
                image = generate_styled_image(style)
        else:
            image = generate_styled_image(style)
        
        # Apply yellow loss
        result = yellow_loss(image, strength)
        
        # Create side-by-side comparison
        comparison = Image.new('RGB', (image.width * 2 + 10, image.height), (240, 240, 240))
        comparison.paste(image, (0, 0))
        comparison.paste(result, (image.width + 10, 0))
        
        # Add labels
        draw = ImageDraw.Draw(comparison)
        draw.text((10, 10), f"Original ({style})", fill=(255, 255, 255), stroke_fill=(0, 0, 0), stroke_width=2)
        draw.text((image.width + 20, 10), f"Yellow Loss: {strength:.1f}", fill=(255, 255, 255), stroke_fill=(0, 0, 0), stroke_width=2)
        
        return comparison
    
    except Exception as e:
        print(f"Error in apply_color_loss: {e}")
        return Image.new('RGB', (512, 256), (200, 200, 200))

# Create Gradio interface
demo = gr.Interface(
    fn=apply_color_loss,
    inputs=[
        gr.Dropdown(choices=list(style_configs.keys()), value="Watercolor", label="Style"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.8, step=0.1, label="Yellow Loss Strength"),
        gr.Image(label="Upload an image (optional)", type="pil")
    ],
    outputs=gr.Image(label="Result (Before and After)"),
    title="Yellow Loss Demo",
    description="This demo shows how yellow loss affects different artistic styles. Each style is generated using Stable Diffusion.",
    examples=[
        ["Watercolor", 0.8, None],
        ["Cyberpunk", 0.5, None],
        ["Anime", 0.9, None],
        ["Oil Painting", 0.7, None],
        ["Sketch", 0.6, None]
    ],
    cache_examples=True
)

# Launch the app
if __name__ == "__main__":
    demo.launch()