import gradio as gr
import torch
from diffusers import StableDiffusionInpaintPipeline, StableDiffusionXLInpaintPipeline
from PIL import Image
import numpy as np
from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
import cv2

class RafayyVirtualTryOn:
    def __init__(self):
        # Initialize SDXL for better quality
        self.inpaint_model = StableDiffusionXLInpaintPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0",
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            variant="fp16"
        )
        if torch.cuda.is_available():
            self.inpaint_model.to("cuda")
            
        # Initialize enhanced segmentation model
        self.segmenter = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
        self.processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
        
    def enhance_mask(self, mask):
        """Enhance the clothing mask for better results"""
        kernel = np.ones((5,5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.GaussianBlur(mask, (5,5), 0)
        return mask

    def get_clothing_mask(self, image):
        """Extract and enhance clothing mask"""
        inputs = self.processor(images=image, return_tensors="pt")
        outputs = self.segmenter(**inputs)
        logits = outputs.logits.squeeze()
        
        clothing_mask = (logits.argmax(0) == 5).float().numpy()
        clothing_mask = (clothing_mask * 255).astype(np.uint8)
        clothing_mask = self.enhance_mask(clothing_mask)
        
        return Image.fromarray(clothing_mask)

    def try_on(self, original_image, prompt, style_strength=0.7, progress=gr.Progress()):
        """Enhanced virtual try-on with style control"""
        try:
            progress(0, desc="Initializing...")
            
            # Image preprocessing
            original_image = Image.fromarray(original_image)
            if original_image.mode != "RGB":
                original_image = original_image.convert("RGB")
            
            progress(0.2, desc="Analyzing clothing...")
            mask = self.get_clothing_mask(original_image)
            
            # Enhanced prompt engineering
            style_prompts = {
                "quality": "ultra detailed, 8k uhd, high quality, professional photo",
                "lighting": "perfect lighting, studio lighting, professional photography",
                "realism": "hyperrealistic, photorealistic, highly detailed"
            }
            
            full_prompt = f"""
            A person wearing {prompt}, {style_prompts['quality']}, 
            {style_prompts['lighting']}, {style_prompts['realism']}
            """
            
            negative_prompt = """
            low quality, blurry, distorted, deformed, unrealistic, 
            bad proportions, bad lighting, oversaturated, undersaturated
            """
            
            progress(0.4, desc="Generating new clothing...")
            result = self.inpaint_model(
                prompt=full_prompt,
                negative_prompt=negative_prompt,
                image=original_image,
                mask_image=mask,
                num_inference_steps=50,
                guidance_scale=7.5 * style_strength
            ).images[0]
            
            progress(1.0, desc="Finalizing...")
            return result
            
        except Exception as e:
            raise gr.Error(f"Processing Error: {str(e)}")

# Initialize the model
model = RafayyVirtualTryOn()

# Custom CSS with professional styling
custom_css = """
.gradio-container {
    font-family: 'Poppins', sans-serif;
    max-width: 1200px !important;
    margin: auto !important;
}
#component-0 {
    max-width: 100% !important;
    margin-bottom: 20px !important;
}
.gr-button {
    background: linear-gradient(90deg, #2193b0, #6dd5ed) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
}
.gr-button:hover {
    background: linear-gradient(90deg, #6dd5ed, #2193b0) !important;
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(33, 147, 176, 0.3) !important;
    transition: all 0.3s ease;
}
.gr-input {
    border: 2px solid #e0e0e0 !important;
    border-radius: 8px !important;
    padding: 12px !important;
}
.gr-input:focus {
    border-color: #2193b0 !important;
    box-shadow: 0 0 0 2px rgba(33, 147, 176, 0.2) !important;
}
.gr-panel {
    border-radius: 12px !important;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.1) !important;
}
.footer {
    background: linear-gradient(to right, #f8f9fa, #e9ecef);
    padding: 20px;
    border-radius: 10px;
    margin-top: 30px;
}
"""

# Create enhanced Gradio interface
demo = gr.Interface(
    fn=model.try_on,
    inputs=[
        gr.Image(label="📸 Upload Your Photo", type="numpy"),
        gr.Textbox(
            label="🎨 Describe New Clothing",
            placeholder="e.g., 'elegant black suit with silk lapels', 'designer red dress with gold accents'",
            lines=2
        ),
        gr.Slider(
            label="Style Strength",
            minimum=0.1,
            maximum=1.0,
            value=0.7,
            step=0.1
        )
    ],
    outputs=gr.Image(label="✨ Your New Look", type="pil"),
    title="🌟 Rafayy's Professional Virtual Try-On Studio 🌟",
    description="""
    <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
        <h3 style="color: #2193b0; font-size: 24px; margin-bottom: 10px;">Transform Your Style with AI</h3>
        <p style="color: #666; font-size: 16px;">Experience the future of fashion with our advanced virtual try-on technology.</p>
        <div style="margin-top: 20px; padding: 15px; background: rgba(33, 147, 176, 0.1); border-radius: 10px;">
            <p style="font-weight: bold; color: #2193b0;">Premium Features:</p>
            <p>✓ High-Resolution Output</p>
            <p>✓ Advanced Clothing Detection</p>
            <p>✓ Professional Style Enhancement</p>
        </div>
    </div>
    """,
    examples=[
        ["example1.jpg", "luxury black suit with silk details", 0.8],
        ["example2.jpg", "designer white dress with lace accents", 0.7],
        ["example3.jpg", "professional navy blazer with gold buttons", 0.9],
        ["example4.jpg", "casual denim jacket with vintage wash", 0.6]
    ],
    css=custom_css
)

# Enhanced footer with professional information
demo.footer = """
<div class="footer">
    <div style="text-align: center;">
        <h4 style="color: #2193b0; margin-bottom: 15px;">🎯 Professional Tips for Best Results</h4>
        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 20px;">
            <div>
                <p style="font-weight: bold;">📸 Photo Guidelines</p>
                <p>• Use well-lit, front-facing photos</p>
                <p>• Ensure clear visibility of clothing</p>
                <p>• Avoid complex backgrounds</p>
            </div>
            <div>
                <p style="font-weight: bold;">✍️ Description Tips</p>
                <p>• Be specific about style details</p>
                <p>• Include color and material</p>
                <p>• Mention design elements</p>
            </div>
            <div>
                <p style="font-weight: bold;">⚙️ Settings</p>
                <p>• Adjust style strength as needed</p>
                <p>• Higher values for bold changes</p>
                <p>• Lower values for subtle effects</p>
            </div>
        </div>
        <p style="margin-top: 20px; color: #666;">© 2024 Rafayy AI Studio | Professional Virtual Try-On Service</p>
    </div>
</div>
"""

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