# Your existing imports remain the same
import gradio as gr
import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt
import json
import time
import os
from functools import partial
import datetime
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# AI Snipper Custom CSS
ai_snipper_css = """
/* AI Snipper Color Variables */
:root {
    --ai-cyan: #06b6d4;
    --ai-blue: #3b82f6;
    --ai-purple: #8b5cf6;
    --ai-teal: #14b8a6;
    --ai-indigo: #6366f1;
    
    --bg-primary: #0a0e16;
    --bg-secondary: #1a2332;
    --bg-tertiary: #2a3441;
    --bg-card: #1e293b;
    --bg-card-hover: #334155;
    
    --text-primary: #ffffff;
    --text-secondary: #e2e8f0;
    --text-muted: #94a3b8;
    --text-accent: #06b6d4;
    
    --border-primary: #475569;
    --border-accent: #06b6d4;
    
    --gradient-primary: linear-gradient(135deg, #06b6d4, #3b82f6, #8b5cf6);
    --gradient-secondary: linear-gradient(135deg, #0a0e16 0%, #1a2332 50%, #2a3441 100%);
    --gradient-button: linear-gradient(135deg, #06b6d4, #3b82f6);
    --gradient-card: linear-gradient(135deg, #1e293b 0%, #334155 100%);
}

/* Main container styling */
.gradio-container {
    background: var(--gradient-secondary) !important;
    color: var(--text-primary) !important;
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
    min-height: 100vh !important;
}

/* Header styling */
.gradio-container h1 {
    background: var(--gradient-primary) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    background-clip: text !important;
    text-align: center !important;
    font-size: 3rem !important;
    font-weight: 800 !important;
    margin-bottom: 1rem !important;
    text-shadow: 0 0 20px rgba(6, 182, 212, 0.3) !important;
}

.gradio-container h2, .gradio-container h3 {
    color: var(--text-primary) !important;
    font-weight: 600 !important;
}

/* Card/Panel styling */
.gr-box, .gr-form, .gr-panel {
    background: var(--gradient-card) !important;
    border: 1px solid var(--border-primary) !important;
    border-radius: 16px !important;
    box-shadow: 0 10px 15px -3px rgb(0 0 0 / 0.1) !important;
    backdrop-filter: blur(10px) !important;
}

/* Group styling */
.gr-group {
    background: var(--gradient-card) !important;
    border: 1px solid var(--border-primary) !important;
    border-radius: 12px !important;
    padding: 1.5rem !important;
}

/* Text inputs */
.gr-textbox textarea, .gr-textbox input {
    background: var(--bg-secondary) !important;
    border: 1px solid var(--border-primary) !important;
    border-radius: 12px !important;
    color: var(--text-primary) !important;
    padding: 1rem !important;
    font-family: inherit !important;
}

.gr-textbox textarea:focus, .gr-textbox input:focus {
    border-color: var(--border-accent) !important;
    box-shadow: 0 0 0 3px rgba(6, 182, 212, 0.1) !important;
    outline: none !important;
}

.gr-textbox textarea::placeholder, .gr-textbox input::placeholder {
    color: var(--text-muted) !important;
}

/* Labels */
.gr-textbox label, .gr-slider label, .gr-radio label, .gr-checkbox label {
    color: var(--text-secondary) !important;
    font-weight: 500 !important;
    margin-bottom: 0.5rem !important;
}

/* Buttons */
.gr-button {
    background: var(--gradient-button) !important;
    border: none !important;
    border-radius: 12px !important;
    color: var(--text-primary) !important;
    font-weight: 600 !important;
    padding: 1rem 2rem !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 6px -1px rgb(0 0 0 / 0.1) !important;
    font-family: inherit !important;
}

.gr-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 10px 20px rgba(6, 182, 212, 0.3) !important;
    filter: brightness(1.1) !important;
}

.gr-button.secondary {
    background: var(--bg-card) !important;
    color: var(--text-secondary) !important;
    border: 1px solid var(--border-primary) !important;
}

.gr-button.secondary:hover {
    background: var(--bg-card-hover) !important;
    border-color: var(--border-accent) !important;
    color: var(--text-primary) !important;
}

/* File upload areas */
.gr-file-upload {
    background: var(--bg-card) !important;
    border: 2px dashed var(--border-accent) !important;
    border-radius: 16px !important;
    color: var(--text-secondary) !important;
    transition: all 0.3s ease !important;
}

.gr-file-upload:hover {
    border-color: var(--ai-cyan) !important;
    background: var(--bg-card-hover) !important;
}

/* Audio input */
.gr-audio {
    background: var(--gradient-card) !important;
    border: 1px solid var(--border-primary) !important;
    border-radius: 12px !important;
}

/* Sliders */
.gr-slider input[type="range"] {
    background: var(--bg-secondary) !important;
}

.gr-slider input[type="range"]::-webkit-slider-track {
    background: var(--bg-secondary) !important;
    border-radius: 6px !important;
}

.gr-slider input[type="range"]::-webkit-slider-thumb {
    background: var(--gradient-button) !important;
    border: none !important;
    border-radius: 50% !important;
    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.2) !important;
}

/* Radio buttons and checkboxes */
.gr-radio input[type="radio"] {
    accent-color: var(--ai-cyan) !important;
}

.gr-checkbox input[type="checkbox"] {
    accent-color: var(--ai-cyan) !important;
}

/* Tabs */
.gr-tab-nav {
    background: var(--gradient-card) !important;
    border-radius: 12px !important;
    padding: 0.5rem !important;
}

.gr-tab-nav button {
    background: transparent !important;
    color: var(--text-secondary) !important;
    border: none !important;
    border-radius: 8px !important;
    margin: 0 4px !important;
    padding: 0.75rem 1.5rem !important;
    transition: all 0.3s ease !important;
    font-weight: 500 !important;
}

.gr-tab-nav button.selected {
    background: var(--gradient-button) !important;
    color: var(--text-primary) !important;
    box-shadow: 0 2px 4px rgba(6, 182, 212, 0.3) !important;
}

.gr-tab-nav button:hover:not(.selected) {
    background: var(--bg-card-hover) !important;
    color: var(--text-primary) !important;
}

/* Tab content */
.gr-tabitem {
    background: var(--gradient-card) !important;
    border: 1px solid var(--border-primary) !important;
    border-radius: 12px !important;
    padding: 1.5rem !important;
    margin-top: 1rem !important;
}

/* Progress bars */
.gr-progress {
    background: var(--bg-secondary) !important;
    border-radius: 6px !important;
}

.gr-progress-bar {
    background: var(--gradient-button) !important;
    border-radius: 6px !important;
}

/* Accordion */
.gr-accordion {
    background: var(--gradient-card) !important;
    border: 1px solid var(--border-primary) !important;
    border-radius: 12px !important;
}

.gr-accordion summary {
    background: var(--bg-card) !important;
    color: var(--text-primary) !important;
    padding: 1rem !important;
    border-radius: 12px !important;
    cursor: pointer !important;
    font-weight: 600 !important;
}

.gr-accordion[open] summary {
    border-bottom: 1px solid var(--border-primary) !important;
    border-radius: 12px 12px 0 0 !important;
}

/* JSON output */
.gr-json {
    background: var(--bg-secondary) !important;
    border: 1px solid var(--border-primary) !important;
    border-radius: 12px !important;
    color: var(--text-primary) !important;
}

/* HTML output areas */
.gr-html {
    background: var(--gradient-card) !important;
    border: 1px solid var(--border-primary) !important;
    border-radius: 12px !important;
    padding: 1rem !important;
}

/* Plot containers */
.gr-plot {
    background: var(--gradient-card) !important;
    border: 1px solid var(--border-primary) !important;
    border-radius: 12px !important;
    padding: 1rem !important;
}

/* Rows and columns */
.gr-row {
    gap: 1.5rem !important;
}

.gr-column {
    gap: 1rem !important;
}

/* Scrollbars */
::-webkit-scrollbar {
    width: 8px;
    height: 8px;
}

::-webkit-scrollbar-track {
    background: var(--bg-secondary);
    border-radius: 4px;
}

::-webkit-scrollbar-thumb {
    background: var(--gradient-button);
    border-radius: 4px;
}

::-webkit-scrollbar-thumb:hover {
    background: var(--ai-cyan);
}

/* Custom DNA-themed elements */
.dna-header {
    position: relative;
    text-align: center;
    padding: 2rem 0;
    margin-bottom: 2rem;
}

.dna-header::before {
    content: '';
    position: absolute;
    top: 0;
    left: 50%;
    transform: translateX(-50%);
    width: 100px;
    height: 4px;
    background: var(--gradient-primary);
    border-radius: 2px;
}

.dna-subtitle {
    color: var(--text-muted) !important;
    font-size: 1.2rem !important;
    margin-top: 1rem !important;
    font-weight: 400 !important;
}

/* Example button styling */
.example-buttons .gr-button {
    background: var(--bg-card) !important;
    color: var(--text-accent) !important;
    border: 1px solid var(--border-accent) !important;
    font-size: 0.875rem !important;
    padding: 0.5rem 1rem !important;
}

.example-buttons .gr-button:hover {
    background: var(--gradient-button) !important;
    color: var(--text-primary) !important;
    border-color: transparent !important;
}

/* Status messages */
.status-message {
    text-align: center !important;
    padding: 1rem !important;
    border-radius: 8px !important;
    margin: 1rem 0 !important;
    font-weight: 500 !important;
}

.status-loading {
    background: rgba(6, 182, 212, 0.1) !important;
    border: 1px solid var(--border-accent) !important;
    color: var(--text-accent) !important;
}

.status-success {
    background: rgba(20, 184, 166, 0.1) !important;
    border: 1px solid var(--ai-teal) !important;
    color: var(--ai-teal) !important;
}

.status-error {
    background: rgba(239, 68, 68, 0.1) !important;
    border: 1px solid #ef4444 !important;
    color: #ef4444 !important;
}

/* Footer hiding */
footer {
    visibility: hidden !important;
}

/* Mobile responsiveness */
@media (max-width: 768px) {
    .gradio-container h1 {
        font-size: 2rem !important;
    }
    
    .gr-button {
        width: 100% !important;
        justify-content: center !important;
    }
    
    .gr-row {
        flex-direction: column !important;
    }
}
"""

# Keep all your existing function code exactly the same
# [Your existing global variables and all functions remain unchanged]

# Global variables to store models
tokenizer = None
ner_pipeline = None
pos_pipeline = None
intent_classifier = None
semantic_model = None
stt_model = None  # Speech-to-text model
models_loaded = False

# Database to store keyword ranking history (in-memory database for this example)
ranking_history = {}

# [Keep all your existing functions - load_models, speech_to_text, etc.]
# I'm not repeating them here to save space, but they should remain exactly the same

# Updated Gradio interface with AI Snipper styling
with gr.Blocks(
    css=ai_snipper_css,
    title="🧬 AI Snipper Keyword DNA Analyzer",
    theme=gr.themes.Base(
        primary_hue="cyan",
        secondary_hue="blue", 
        neutral_hue="slate",
        font=[gr.themes.GoogleFont("Inter"), "sans-serif"]
    )
) as demo:
    
    # Custom header with DNA theme
    gr.HTML("""
    <div class="dna-header">
        <h1>🧬 Keyword DNA Analyzer</h1>
        <p class="dna-subtitle">
            Decode the genetic structure of your keywords with AI-powered analysis
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Voice search capabilities with improved styling
            with gr.Group():
                gr.Markdown("### 🎯 Enter Keyword")
                with gr.Row():
                    input_text = gr.Textbox(
                        label="Keyword to analyze", 
                        placeholder="e.g. artificial intelligence, machine learning, SEO strategy...",
                        lines=2
                    )
                
                with gr.Row():
                    audio_input = gr.Audio(
                        type="filepath", 
                        label="🎤 Or use voice search"
                    )
                    voice_submit_btn = gr.Button(
                        "🔄 Convert Voice to Text", 
                        variant="secondary"
                    )
            
            # SERP settings with better organization
            with gr.Accordion("⚙️ Analysis Settings", open=False):
                with gr.Row():
                    forecast_months = gr.Slider(
                        minimum=3, 
                        maximum=12, 
                        value=6, 
                        step=1, 
                        label="📈 Forecast Months"
                    )
                    include_serp = gr.Checkbox(
                        label="🔍 Include SERP Analysis", 
                        value=True
                    )
                
                growth_scenario = gr.Radio(
                    ["Conservative", "Moderate", "Aggressive"], 
                    value="Moderate", 
                    label="📊 Growth Scenario"
                )
            
            # Status indicator with custom styling
            status_html = gr.HTML(
                '<div class="status-message">🚀 Enter a keyword and click "Analyze DNA" to begin</div>'
            )
            
            # Main analyze button
            analyze_btn = gr.Button(
                "🧬 Analyze DNA", 
                variant="primary",
                size="lg"
            )
            
            # Example buttons with custom styling
            gr.Markdown("### 💡 Try These Examples")
            with gr.Row(elem_classes="example-buttons"):
                example_btns = []
                examples = [
                    "preprocessing", 
                    "breakdown", 
                    "artificial intelligence", 
                    "transformer model", 
                    "machine learning"
                ]
                for example in examples:
                    example_btns.append(gr.Button(f"✨ {example}"))
            
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("🔬 Token Visualization"):
                    token_viz_html = gr.HTML()
                
                with gr.Tab("📊 Full Analysis"):
                    analysis_html = gr.HTML()
                
                with gr.Tab("📈 Evolution Chart"):
                    evolution_chart = gr.Plot(label="Keyword Evolution Forecast")
                
                with gr.Tab("🌐 SERP Results"):
                    serp_html = gr.HTML()
                
                with gr.Tab("📉 Ranking History"):
                    ranking_chart = gr.Plot(label="Keyword Ranking History")
                
                with gr.Tab("💾 Raw Data"):
                    json_output = gr.JSON()
    
    # Event handlers remain the same but with updated status messages
    voice_submit_btn.click(
        handle_voice_input,
        inputs=[audio_input],
        outputs=[input_text]
    )
    
    # Updated status messages with custom styling
    analyze_btn.click(
        lambda: '<div class="status-message status-loading">🔄 Loading models and analyzing... This may take a moment.</div>',
        outputs=status_html
    ).then(
        analyze_keyword,
        inputs=[input_text, forecast_months, growth_scenario, include_serp],
        outputs=[token_viz_html, analysis_html, json_output, evolution_chart, serp_html, ranking_chart, input_text]
    ).then(
        lambda: '<div class="status-message status-success">✅ Analysis complete! Check the results above.</div>',
        outputs=status_html
    )
    
    # Example buttons with enhanced interaction
    for btn in example_btns:
        def set_example(btn_label):
            # Remove the emoji prefix for the actual keyword
            return btn_label.replace("✨ ", "")
        
        btn.click(
            set_example,
            inputs=[btn],
            outputs=[input_text]
        ).then(
            lambda: '<div class="status-message status-loading">🔄 Loading models and analyzing... This may take a moment.</div>',
            outputs=status_html
        ).then(
            analyze_keyword,
            inputs=[input_text, forecast_months, growth_scenario, include_serp],
            outputs=[token_viz_html, analysis_html, json_output, evolution_chart, serp_html, ranking_chart, input_text]
        ).then(
            lambda: '<div class="status-message status-success">✅ Analysis complete! Check the results above.</div>',
            outputs=status_html
        )

# Launch configuration
if __name__ == "__main__":
    demo.launch(
        share=True,
        show_error=True,
        debug=True
    )