import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import gc

# Global variables for model and tokenizer
model = None
tokenizer = None

def load_model():
    """Load the model and tokenizer from Hugging Face with CPU optimizations"""
    global model, tokenizer
    
    if model is None or tokenizer is None:
        try:
            model_name = "Harshu0117/Materials_IISC_MRC"
            
            # Load tokenizer
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            
            # Load model with CPU optimizations
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                device_map="cpu",
                trust_remote_code=True,
                low_cpu_mem_usage=True,  # Reduce memory usage
                offload_folder="offload"  # Enable model offloading
            )
            
            # Enable CPU optimizations
            model.eval()
            
            # Set pad token if not set
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token
            
            # Clear GPU cache if any
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            # Force garbage collection
            gc.collect()
            
            return "✅ Model loaded successfully with CPU optimizations!"
        except Exception as e:
            return f"❌ Error loading model: {str(e)}"
    
    return "✅ Model already loaded!"

def generate_response(prompt, max_tokens, temperature, top_p, repetition_penalty):
    """Generate response using the loaded model with CPU optimizations"""
    global model, tokenizer
    
    # Load model if not already loaded
    if model is None or tokenizer is None:
        load_result = load_model()
        if "Error" in load_result:
            return load_result
    
    if not prompt.strip():
        return "⚠️ Please enter a question or topic first!"
    
    try:
        # Tokenize input with truncation for faster processing
        inputs = tokenizer(
            prompt.strip(), 
            return_tensors="pt",
            truncation=True,
            max_length=512,  # Reduced from 1024 for faster processing
            padding=True
        )
        
        # Keep on CPU
        inputs = inputs.to("cpu")
        
        # Generate response with optimized settings
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=int(max_tokens),
                temperature=float(temperature),
                top_p=float(top_p),
                repetition_penalty=float(repetition_penalty),
                do_sample=True,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                use_cache=True,
                num_beams=1,  # Use greedy decoding for speed
                early_stopping=True
            )
        
        # Decode response
        response = tokenizer.decode(
            outputs[0], 
            skip_special_tokens=True
        )
        
        # Remove the original prompt from response
        response = response.replace(prompt.strip(), "").strip()
        
        # Clear memory
        del outputs
        gc.collect()
        
        return response
    
    except Exception as e:
        return f"❌ Error generating response: {str(e)}"

# Create Gradio interface
def create_interface():
    # Custom CSS for styling
    css = """
    .gradio-container {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    }
    .gr-button-primary {
        background: linear-gradient(45deg, #FF6B6B, #4ECDC4) !important;
        border: none !important;
        border-radius: 25px !important;
        color: white !important;
        font-weight: bold !important;
        padding: 12px 24px !important;
        font-size: 16px !important;
        transition: all 0.3s ease !important;
    }
    .gr-button-primary:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 4px 12px rgba(0, 0, 0, 0.2) !important;
    }
    .gr-textbox {
        border-radius: 15px !important;
        border: 2px solid #e0e0e0 !important;
        background: rgba(255, 255, 255, 0.95) !important;
    }
    .gr-textbox:focus {
        border-color: #4ECDC4 !important;
        box-shadow: 0 0 10px rgba(78, 205, 196, 0.3) !important;
    }
    .output-text {
        background: rgba(255, 255, 255, 0.95) !important;
        border-radius: 15px !important;
        padding: 20px !important;
        margin: 10px 0 !important;
        border-left: 4px solid #4ECDC4 !important;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1) !important;
    }
    .gr-accordion {
        background: rgba(255, 255, 255, 0.1) !important;
        border-radius: 15px !important;
        border: 1px solid rgba(255, 255, 255, 0.3) !important;
    }
    """
    
    # Create interface
    with gr.Blocks(
        css=css,
        title="🧪 Materials Science AI Assistant",
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="cyan",
            neutral_hue="slate"
        )
    ) as demo:
        
        # Header
        gr.HTML("""
        <div style="text-align: center; padding: 30px; background: rgba(255, 255, 255, 0.95); border-radius: 20px; margin-bottom: 20px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);">
            <h1 style="color: #2c3e50; font-size: 2.5em; margin: 0; text-shadow: 2px 2px 4px rgba(0,0,0,0.1);">
                🧪 Materials Science AI Assistant
            </h1>
            <p style="color: #7f8c8d; font-size: 1.2em; margin: 10px 0 0 0; font-weight: 500;">
                Powered by Fine-tuned LLaMA 3 8B | Specialized in Materials Research
            </p>
        </div>
        """)
        
        # Main interface
        with gr.Row():
            with gr.Column(scale=2):
                # Input area
                gr.Markdown("### 💬 Ask me anything about Materials Science!")
                
                prompt = gr.Textbox(
                    label="Enter your question or topic:",
                    placeholder="e.g., Crystalline MAX Phases and their 2D derivative MXenes",
                    lines=4,
                    max_lines=8
                )
                
                # Advanced options
                with gr.Accordion("⚙️ Advanced Options", open=False):
                    with gr.Row():
                        max_tokens = gr.Slider(
                            label="Max Tokens (Response Length)",
                            minimum=50,
                            maximum=500,
                            value=200,
                            step=10,
                            info="Maximum number of tokens in the response"
                        )
                        
                        temperature = gr.Slider(
                            label="Temperature (Creativity)",
                            minimum=0.1,
                            maximum=1.0,
                            value=0.7,
                            step=0.1,
                            info="Higher values make responses more creative"
                        )
                    
                    with gr.Row():
                        top_p = gr.Slider(
                            label="Top-p (Diversity)",
                            minimum=0.1,
                            maximum=1.0,
                            value=0.9,
                            step=0.1,
                            info="Controls diversity of word choices"
                        )
                        
                        repetition_penalty = gr.Slider(
                            label="Repetition Penalty",
                            minimum=1.0,
                            maximum=2.0,
                            value=1.2,
                            step=0.1,
                            info="Penalty for repeating words/phrases"
                        )
                
                # Generate button
                generate_btn = gr.Button(
                    "🚀 Generate Response",
                    variant="primary",
                    size="lg"
                )
        
        # Output area
        gr.Markdown("### 🤖 AI Response:")
        output = gr.Textbox(
            label="Generated Response",
            lines=10,
            max_lines=20,
            interactive=False,
            elem_classes=["output-text"]
        )
        
        # Example prompts
        gr.Markdown("### 💡 Example Prompts (Click to use):")
        examples = [
            "Crystalline MAX Phases and their 2D derivative MXenes",
            "Properties of titanium carbide MXenes",
            "Synthesis methods for MAX phases",
            "Applications of MXenes in energy storage",
            "Mechanical properties of ceramic materials"
        ]
        
        gr.Examples(
            examples=examples,
            inputs=prompt,
            label="Click any example to try:"
        )
        
        # Footer
        gr.HTML("""
        <div style="text-align: center; padding: 20px; margin-top: 30px; background: rgba(255, 255, 255, 0.1); border-radius: 15px; border: 1px solid rgba(255, 255, 255, 0.3);">
            <p style="color: white; font-size: 16px; margin: 0;">
                🔬 <strong>Specialized in Materials Science</strong> | 🧪 <strong>MAX Phases & MXenes Expert</strong>
            </p>
            <p style="color: rgba(255, 255, 255, 0.8); font-size: 14px; margin: 5px 0 0 0;">
                Built with ❤️ using Gradio & Hugging Face Spaces
            </p>
        </div>
        """)
        
        # Connect the generate button to the function
        generate_btn.click(
            fn=generate_response,
            inputs=[prompt, max_tokens, temperature, top_p, repetition_penalty],
            outputs=output,
            show_progress=True
        )
    
    return demo

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