import gradio as gr
import numpy as np
from PIL import Image
import random
import warnings
warnings.filterwarnings("ignore")

# Try to import AI dependencies
try:
    import torch
    from transformers import pipeline, BlipProcessor, BlipForConditionalGeneration
    TRANSFORMERS_AVAILABLE = True
    print("✅ AI models available")
except ImportError:
    TRANSFORMERS_AVAILABLE = False
    print("⚠️ Using lightweight mode")

class FashionAI:
    def __init__(self):
        self.setup_ai_models()
        self.load_fashion_data()
        
    def setup_ai_models(self):
        """Load AI models if available"""
        if TRANSFORMERS_AVAILABLE:
            try:
                print("🔄 Loading AI models...")
                # Use smaller models for HF Spaces
                self.caption_model = pipeline(
                    "image-to-text", 
                    model="nlpconnect/vit-gpt2-image-captioning",
                    device=0 if torch.cuda.is_available() else -1
                )
                
                self.classifier = pipeline(
                    "zero-shot-classification",
                    model="facebook/bart-large-mnli",
                    device=0 if torch.cuda.is_available() else -1
                )
                
                self.use_ai = True
                print("✅ AI models loaded successfully!")
                
            except Exception as e:
                print(f"❌ AI model error: {e}")
                self.use_ai = False
        else:
            self.use_ai = False
            print("📱 Running in lightweight mode")
    
    def load_fashion_data(self):
        """Load fashion knowledge base"""
        self.color_seasons = {
            "spring": {
                "colors": ["coral", "peach", "bright yellow", "warm pink", "kelly green"],
                "avoid": ["black", "burgundy", "navy", "cool blues"]
            },
            "summer": {
                "colors": ["soft blue", "lavender", "rose pink", "powder blue", "mint"],
                "avoid": ["orange", "bright yellow", "warm browns"]
            },
            "autumn": {
                "colors": ["rust", "golden yellow", "olive green", "burnt orange", "chocolate"],
                "avoid": ["bright pink", "cool blues", "pure white"]
            },
            "winter": {
                "colors": ["true red", "bright white", "royal blue", "emerald", "black"],
                "avoid": ["beige", "orange", "golden yellow"]
            }
        }
        
        self.body_types = {
            "pear": {
                "focus": "shoulders and upper body",
                "recommend": ["A-line tops", "boat necks", "structured jackets"],
                "avoid": ["tight bottoms", "hip-emphasizing details"]
            },
            "apple": {
                "focus": "legs and neckline", 
                "recommend": ["A-line dresses", "V-necks", "empire waist"],
                "avoid": ["tight around waist", "horizontal stripes"]
            },
            "hourglass": {
                "focus": "natural waist",
                "recommend": ["fitted styles", "wrap dresses", "belted outfits"],
                "avoid": ["loose shapeless clothes", "hiding waist"]
            },
            "rectangle": {
                "focus": "creating curves",
                "recommend": ["peplum tops", "layered looks", "ruffles"],
                "avoid": ["straight cuts", "baggy clothes"]
            }
        }

    def analyze_image(self, image):
        """Analyze fashion image with AI or fallback"""
        if image is None:
            return "❌ Please upload an image first!"
        
        try:
            if self.use_ai:
                return self.ai_image_analysis(image)
            else:
                return self.fallback_analysis(image)
        except Exception as e:
            return f"⚠️ Analysis error: {str(e)}\n\nUsing basic analysis...\n\n{self.fallback_analysis(image)}"

    def ai_image_analysis(self, image):
        """AI-powered image analysis"""
        try:
            # Generate caption
            caption_result = self.caption_model(image)
            caption = caption_result[0]['generated_text'] if caption_result else "fashion item"
            
            # Classify style
            style_labels = ["casual", "formal", "elegant", "sporty", "trendy", "professional"]
            style_result = self.classifier(caption, style_labels)
            style = style_result['labels'][0]
            
            # Generate comprehensive analysis
            analysis = f"""# 🔍 **AI Fashion Analysis**

👀 **What I See**: {caption.capitalize()}

✨ **Style Category**: {style.title()}

🎨 **Color Recommendations**: 
Based on the {style} style, I recommend:
• {random.choice(['Navy & White', 'Black & Gold', 'Coral & Cream', 'Emerald & Silver'])}
• {random.choice(['Soft pastels', 'Bold jewel tones', 'Neutral earth tones', 'Classic monochromes'])}

💡 **Styling Tips**:
• Perfect for {random.choice(['professional settings', 'casual outings', 'special occasions', 'everyday wear'])}
• Pair with {random.choice(['statement accessories', 'classic heels', 'comfortable flats', 'a structured bag'])}
• Consider {random.choice(['layering with a blazer', 'adding a belt', 'mixing textures', 'playing with proportions'])}

🎯 **Best Occasions**: {', '.join(random.sample(['Work meetings', 'Dinner dates', 'Weekend brunches', 'Shopping trips', 'Social events'], 3))}

✨ *AI-powered analysis complete!*"""

            return analysis
            
        except Exception as e:
            return f"AI analysis failed: {e}"

    def fallback_analysis(self, image):
        """Rule-based fallback analysis"""
        try:
            # Basic color analysis
            img_array = np.array(image)
            avg_color = np.mean(img_array)
            
            if avg_color > 180:
                color_desc = "light and bright"
                season = "spring"
            elif avg_color < 80:
                color_desc = "dark and dramatic"  
                season = "winter"
            else:
                color_desc = "balanced tones"
                season = "autumn"
            
            season_info = self.color_seasons[season]
            
            return f"""# 🔍 **Fashion Analysis Results**

🎨 **Color Profile**: {color_desc.title()}

🌈 **Your Season**: {season.title()}
• **Perfect Colors**: {', '.join(season_info['colors'][:3])}
• **Avoid**: {', '.join(season_info['avoid'][:2])}

💡 **Styling Suggestions**:
• This piece has {color_desc} that work beautifully for {season} color palettes
• Consider pairing with complementary {season} colors
• Perfect for creating sophisticated, coordinated looks

🎯 **Versatile Styling**:
• Dress it up with heels and jewelry for evening
• Keep it casual with flats and minimal accessories  
• Layer with complementary pieces for different occasions

✨ *Analysis complete - you're ready to style!*"""
            
        except Exception as e:
            return f"Basic analysis: This appears to be a fashion item with styling potential! Consider the colors and silhouette when creating outfits."

    def chat_response(self, message, history):
        """Generate chat responses"""
        msg_lower = message.lower()
        
        responses = {
            'color': """🌈 **Color Magic!**

**Find Your Season:**
• **Spring**: Warm, bright colors (coral, peach, bright yellow)
• **Summer**: Cool, soft colors (lavender, powder blue, rose pink)  
• **Autumn**: Warm, rich colors (rust, golden yellow, olive green)
• **Winter**: Cool, dramatic colors (true red, royal blue, black)

**Quick Test**: Look at your wrist veins:
• Green veins = Warm undertones (Spring/Autumn)
• Blue veins = Cool undertones (Summer/Winter)

What colors do you gravitate toward naturally?""",

            'body': """👗 **Body Type Styling Guide**

**Pear Shape** (smaller shoulders, fuller hips):
✅ A-line tops, boat necks, structured jackets
❌ Tight bottoms, hip details

**Apple Shape** (fuller midsection):
✅ A-line dresses, V-necks, empire waist
❌ Tight waistlines, horizontal stripes

**Hourglass** (balanced curves):
✅ Fitted styles, wrap dresses, belts
❌ Loose, shapeless clothing

**Rectangle** (straight up and down):
✅ Peplum tops, layers, ruffles, belts
❌ Straight cuts, baggy styles

Which shape sounds most like you?""",

            'work': """👔 **Professional Power Dressing**

**Essential Pieces:**
• Well-fitted blazer (navy, black, gray)
• Tailored pants or pencil skirt
• Classic button-down shirts
• Closed-toe shoes (modest heel)
• Quality, minimal jewelry

**Color Strategy:**
• Neutrals as base (black, navy, gray, white)
• Add one accent color per outfit
• Avoid overly bright or distracting patterns

**Pro Tips:**
• Fit is EVERYTHING in professional wear
• Invest in quality basics over trendy pieces
• Keep makeup and accessories understated

Ready to build your power wardrobe?""",

            'date': """💕 **Date Night Perfection**

**The Golden Rules:**
• Wear something that makes YOU feel amazing
• Comfort + confidence = irresistible combination
• Match the vibe (casual coffee vs fancy dinner)

**Go-To Options:**
• **Casual**: Great jeans + silk blouse + cute flats
• **Dinner**: Little black dress + statement jewelry + heels
• **Activity**: Cute sundress + comfortable wedges

**Final Touch:**
• Subtle, flattering makeup
• Signature scent (not overpowering!)
• Genuine smile and positive energy

What kind of date are you planning?""",

            'default': """✨ **Your Fashion AI Assistant**

I'm here to help with all things style! I can assist with:

🎨 **Color Analysis** - Find your perfect palette
👗 **Body Type Styling** - Flattering fits for your shape  
💼 **Professional Wardrobe** - Power dressing tips
💕 **Special Occasions** - Perfect outfits for events
🛍️ **Wardrobe Building** - Smart shopping strategies
📸 **Image Analysis** - Upload photos for personalized advice

**Popular Questions:**
"What colors suit me?" | "How to dress for my body type?" | "Professional outfit ideas" | "Date night styling"

What fashion challenge can I solve for you today?"""
        }
        
        # Match user intent
        for keyword, response in responses.items():
            if keyword != 'default' and keyword in msg_lower:
                return response
        
        return responses['default']

    def personal_style_guide(self, skin_tone, body_type, style_prefs, occasion):
        """Generate personalized style recommendations"""
        guide = ["# 🌟 **Your Personal Style Guide**\n"]
        
        # Color recommendations
        if skin_tone != "Not sure":
            color_map = {
                "Warm": ("Spring/Autumn", ["coral", "peach", "golden yellow", "rust", "olive green"]),
                "Cool": ("Summer/Winter", ["soft blue", "lavender", "true red", "royal blue", "emerald"]),
                "Neutral": ("Flexible", ["navy", "black", "white", "gray", "most colors work"])
            }
            
            season, colors = color_map[skin_tone]
            guide.append(f"## 🎨 Perfect Colors for {skin_tone} Undertones ({season})")
            guide.append(f"**Your Palette**: {', '.join(colors)}")
            guide.append("")
        
        # Body type styling
        if body_type != "Not sure":
            body_key = body_type.lower().replace(' ', '_').replace('inverted_triangle', 'rectangle')
            if body_key in self.body_types:
                body_info = self.body_types[body_key]
                guide.append(f"## 👗 Styling for {body_type} Shape")
                guide.append(f"**Focus on**: {body_info['focus']}")
                guide.append(f"**Recommended**: {', '.join(body_info['recommend'])}")
                guide.append(f"**Avoid**: {', '.join(body_info['avoid'])}")
                guide.append("")
        
        # Style preferences
        if style_prefs:
            guide.append(f"## ✨ Your Style DNA: {', '.join(style_prefs)}")
            
            style_tips = {
                "Casual": "Comfortable, versatile pieces that mix and match",
                "Professional": "Tailored, classic pieces in quality fabrics",
                "Elegant": "Refined silhouettes with luxurious details",
                "Trendy": "Current styles with fashion-forward elements",
                "Minimalist": "Clean lines, neutral colors, capsule wardrobe",
                "Bohemian": "Flowing fabrics, artistic prints, layered accessories"
            }
            
            for style in style_prefs[:3]:  # Limit to top 3
                if style in style_tips:
                    guide.append(f"**{style}**: {style_tips[style]}")
            guide.append("")
        
        # Occasion-specific advice
        occasion_guide = {
            "Work/Professional": "Sharp blazers, tailored fits, neutral colors with subtle personality",
            "Casual Day": "Comfortable yet put-together, versatile pieces that transition well",
            "Date Night": "Something that makes you feel confident and authentic to your style",
            "Party/Event": "Statement pieces, bold colors, interesting textures and details",
            "Wedding Guest": "Elegant without upstaging, avoid white, consider the venue",
            "Travel": "Comfortable layers, wrinkle-resistant fabrics, versatile pieces"
        }
        
        guide.append(f"## 🎯 Perfect for {occasion}")
        guide.append(f"{occasion_guide.get(occasion, 'Versatile styling for any occasion')}")
        guide.append("")
        
        # Final styling tips
        guide.append("## 💡 **Your Style Action Plan**")
        guide.append("• **Start with fit** - well-fitted basics are your foundation")
        guide.append("• **Build gradually** - invest in quality pieces over time") 
        guide.append("• **Mix and match** - create multiple looks with fewer pieces")
        guide.append("• **Accessorize strategically** - transform outfits with small changes")
        guide.append("• **Stay true to you** - confidence is your best accessory!")
        
        return "\n".join(guide)

# Create the Gradio interface
def create_fashion_interface():
    fashion_ai = FashionAI()
    
    # Modern styling
    css = """
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
    
    .gradio-container {
        font-family: 'Inter', sans-serif !important;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    }
    
    .main {
        background: rgba(255, 255, 255, 0.95) !important;
        border-radius: 25px !important;
        backdrop-filter: blur(10px) !important;
        box-shadow: 0 25px 50px rgba(0,0,0,0.15) !important;
        margin: 20px !important;
    }
    
    button {
        background: linear-gradient(135deg, #667eea, #764ba2) !important;
        border: none !important;
        border-radius: 25px !important;
        color: white !important;
        font-weight: 600 !important;
        transition: all 0.3s ease !important;
        box-shadow: 0 5px 15px rgba(102, 126, 234, 0.3) !important;
    }
    
    button:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 10px 25px rgba(102, 126, 234, 0.4) !important;
    }
    
    .tab-nav button.selected {
        background: linear-gradient(135deg, #667eea, #764ba2) !important;
        transform: translateY(-2px) !important;
    }
    
    img {
        border-radius: 15px !important;
        box-shadow: 0 10px 30px rgba(0,0,0,0.2) !important;
        transition: all 0.3s ease !important;
    }
    
    .markdown {
        background: white !important;
        border-radius: 15px !important;
        padding: 25px !important;
        box-shadow: 0 5px 15px rgba(0,0,0,0.1) !important;
        border-left: 4px solid #667eea !important;
    }
    """

    with gr.Blocks(
        title="🤖✨ Advanced Fashion AI Stylist",
        theme=gr.themes.Soft(primary_hue="purple", secondary_hue="pink"),
        css=css
    ) as demo:
        
        # Header
        gr.HTML("""
        <div style="text-align: center; padding: 40px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 25px 25px 0 0; margin-bottom: 0;">
            <h1 style="font-size: 3.5rem; margin: 0; text-shadow: 2px 2px 10px rgba(0,0,0,0.3);">🤖✨ Fashion AI Stylist</h1>
            <p style="font-size: 1.4rem; margin: 15px 0; opacity: 0.9;">Advanced AI-Powered Fashion Analysis & Personal Styling</p>
            
            <div style="display: flex; justify-content: center; gap: 30px; margin-top: 30px; flex-wrap: wrap;">
                <div style="text-align: center; padding: 20px; background: rgba(255, 255, 255, 0.15); border-radius: 15px; min-width: 120px; backdrop-filter: blur(10px);">
                    <div style="font-size: 2.5rem; margin-bottom: 10px;">👁️</div>
                    <div style="font-weight: 600;">AI Vision</div>
                </div>
                <div style="text-align: center; padding: 20px; background: rgba(255, 255, 255, 0.15); border-radius: 15px; min-width: 120px; backdrop-filter: blur(10px);">
                    <div style="font-size: 2.5rem; margin-bottom: 10px;">🎨</div>
                    <div style="font-weight: 600;">Color Analysis</div>
                </div>
                <div style="text-align: center; padding: 20px; background: rgba(255, 255, 255, 0.15); border-radius: 15px; min-width: 120px; backdrop-filter: blur(10px);">
                    <div style="font-size: 2.5rem; margin-bottom: 10px;">💬</div>
                    <div style="font-weight: 600;">Smart Chat</div>
                </div>
                <div style="text-align: center; padding: 20px; background: rgba(255, 255, 255, 0.15); border-radius: 15px; min-width: 120px; backdrop-filter: blur(10px);">
                    <div style="font-size: 2.5rem; margin-bottom: 10px;">✨</div>
                    <div style="font-weight: 600;">Personal Style</div>
                </div>
            </div>
        </div>
        """)
        
        with gr.Tabs():
            # AI Image Analysis
            with gr.TabItem("📸 AI Image Analysis"):
                gr.Markdown("## 🔍 Upload & Analyze Your Fashion Images")
                
                with gr.Row():
                    with gr.Column(scale=1):
                        image_input = gr.Image(
                            label="📷 Drop your fashion image here",
                            type="pil",
                            height=400
                        )
                        analyze_btn = gr.Button(
                            "✨ Analyze with AI", 
                            variant="primary", 
                            size="lg"
                        )
                        
                        gr.Markdown("""
                        ### 🎯 **What I Can Analyze:**
                        - **Clothing items** and complete outfits
                        - **Color palettes** and seasonal recommendations  
                        - **Style categories** and fashion themes
                        - **Styling suggestions** for different occasions
                        - **Mix & match** ideas with your existing wardrobe
                        
                        *Upload any fashion image for instant AI-powered insights!*
                        """)
                    
                    with gr.Column(scale=1):
                        analysis_output = gr.Markdown(
                            value="🎭 **Ready for Analysis!**\n\nUpload a fashion image and click '✨ Analyze with AI' to discover styling insights, color recommendations, and personalized fashion advice!\n\n*Your personal fashion consultant is just one click away...*"
                        )
                
                analyze_btn.click(
                    fashion_ai.analyze_image,
                    inputs=[image_input],
                    outputs=[analysis_output]
                )
            
            # Smart Fashion Chat
            with gr.TabItem("💬 Fashion Chat"):
                gr.Markdown("## 🤖 Chat with Your Personal Fashion Stylist")
                
                chatbot = gr.Chatbot(
                    value=[["", "👋 Hello gorgeous! I'm your personal AI fashion stylist. I'm here to help you discover your perfect style, find amazing color combinations, and create stunning outfits for any occasion!\n\nWhat fashion adventure shall we embark on today? ✨"]],
                    height=500
                )
                
                with gr.Row():
                    msg = gr.Textbox(
                        placeholder="Ask me anything! Colors, styling, body types, occasions, trends...",
                        show_label=False,
                        scale=4
                    )
                    send_btn = gr.Button("Send ✨", scale=1)
                
                # Quick suggestion buttons
                with gr.Row():
                    gr.Button("🌈 What colors suit me?", size="sm").click(
                        lambda: "What colors suit me best?", outputs=[msg]
                    )
                    gr.Button("💕 Date night outfit ideas", size="sm").click(
                        lambda: "I need the perfect date night outfit!", outputs=[msg]
                    )
                    gr.Button("👔 Professional wardrobe help", size="sm").click(
                        lambda: "Help me build a professional wardrobe", outputs=[msg]
                    )
                    gr.Button("👗 Body type styling tips", size="sm").click(
                        lambda: "How should I dress for my body type?", outputs=[msg]
                    )
                
                def respond(message, chat_history):
                    if message.strip():
                        bot_response = fashion_ai.chat_response(message, chat_history)
                        chat_history.append([message, bot_response])
                    return chat_history, ""
                
                msg.submit(respond, [msg, chatbot], [chatbot, msg])
                send_btn.click(respond, [msg, chatbot], [chatbot, msg])
            
            # Personal Style Assistant
            with gr.TabItem("✨ Personal Style"):
                gr.Markdown("## 🌟 Create Your Personal Style Profile")
                
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### 👤 Tell me about yourself")
                        
                        skin_tone = gr.Radio(
                            choices=["Warm", "Cool", "Neutral", "Not sure"],
                            label="🌈 Skin Undertone (look at your wrist veins: green=warm, blue=cool)",
                            value="Not sure"
                        )
                        
                        body_type = gr.Radio(
                            choices=["Pear", "Apple", "Hourglass", "Rectangle", "Inverted Triangle", "Not sure"],
                            label="👗 Body Type",
                            value="Not sure"
                        )
                        
                        style_prefs = gr.CheckboxGroup(
                            choices=["Casual", "Professional", "Elegant", "Trendy", "Minimalist", "Bohemian"],
                            label="✨ Style Preferences (select all that resonate)",
                            value=[]
                        )
                        
                        occasion = gr.Dropdown(
                            choices=["Work/Professional", "Casual Day", "Date Night", "Party/Event", "Wedding Guest", "Travel"],
                            label="🎯 Current Styling Need",
                            value="Casual Day"
                        )
                        
                        style_btn = gr.Button("🌟 Create My Style Guide", variant="primary", size="lg")
                    
                    with gr.Column():
                        personal_results = gr.Markdown(
                            value="✨ **Your Personal Style Journey Starts Here**\n\nFill out your preferences on the left to receive a comprehensive, personalized style guide tailored specifically for you!\n\n🎯 *Get ready to discover your perfect style formula...*"
                        )
                
                style_btn.click(
                    fashion_ai.personal_style_guide,
                    inputs=[skin_tone, body_type, style_prefs, occasion],
                    outputs=[personal_results]
                )
        
        # Footer
        gr.HTML(f"""
        <div style="text-align: center; padding: 25px; margin-top: 20px; background: linear-gradient(135deg, #f8f9fa, #e9ecef); border-radius: 0 0 25px 25px; border-top: 1px solid #dee2e6;">
            <div style="margin-bottom: 15px;">
                <span style="font-size: 1.2rem; margin: 0 15px;">🤖</span>
                <span style="font-size: 1.2rem; margin: 0 15px;">✨</span>
                <span style="font-size: 1.2rem; margin: 0 15px;">👗</span>
                <span style="font-size: 1.2rem; margin: 0 15px;">🎨</span>
                <span style="font-size: 1.2rem; margin: 0 15px;">💫</span>
            </div>
            <p style="margin: 0; color: #666; font-size: 1rem; font-weight: 500;">
                <strong>🚀 AI Mode:</strong> {'Advanced AI Enhanced' if TRANSFORMERS_AVAILABLE else 'Lightweight & Fast'} • 
                <strong>⚡ Status:</strong> Ready to Style • 
                <strong>✨ Your Fashion Journey Awaits!</strong>
            </p>
        </div>
        """)

    return demo

# Launch the app
if __name__ == "__main__":
    print("🚀 Starting Advanced Fashion AI Stylist...")
    demo = create_fashion_interface()
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )