import gradio as gr
import os
import json
import time
from datetime import datetime
from groq import Groq
import pandas as pd
from typing import Dict, List, Tuple, Optional
import re
import uuid

# Initialize Groq client
def get_groq_client():
    api_key = os.getenv("GROQ_API_KEY")
    if not api_key:
        raise ValueError("GROQ_API_KEY environment variable not set")
    return Groq(api_key=api_key)

class PowerSystemsConsultant:
    def __init__(self):
        self.groq_client = get_groq_client()
        self.knowledge_base = self.load_knowledge_base()
        self.conversation_history = []
        
    def load_knowledge_base(self):
        """Load the power systems knowledge base"""
        try:
            with open('data/knowledge_base.json', 'r', encoding='utf-8') as f:
                return json.load(f)
        except FileNotFoundError:
            return self.get_default_knowledge_base()
    
    def get_default_knowledge_base(self):
        """Default knowledge base with power systems content"""
        return {
            "faults": {
                "symmetrical_faults": {
                    "description": "Three-phase faults where all phases are equally affected",
                    "characteristics": "Balanced conditions, easiest to analyze",
                    "analysis_method": "Single-phase equivalent circuit"
                },
                "unsymmetrical_faults": {
                    "line_to_ground": "Most common fault (70-80% of all faults)",
                    "line_to_line": "Second most common fault (15-20%)",
                    "double_line_to_ground": "Less common but severe"
                }
            },
            "protection": {
                "overcurrent_protection": "Time-current characteristic curves, coordination",
                "differential_protection": "Used for transformers, generators, buses",
                "distance_protection": "Impedance-based protection for transmission lines"
            },
            "standards": {
                "IEEE_standards": ["IEEE C37", "IEEE 1547", "IEEE 519"],
                "IEC_standards": ["IEC 61850", "IEC 60909", "IEC 61131"]
            },
            "diagrams": {
                "single_line_diagram": "Shows electrical connections and components",
                "protection_coordination": "Time-current curves for protective devices",
                "fault_analysis": "Sequence networks for fault calculations"
            }
        }
    
    def retrieve_relevant_context(self, query: str) -> str:
        """Simple keyword-based retrieval from knowledge base"""
        query_lower = query.lower()
        relevant_info = []
        
        # Search through knowledge base
        for category, content in self.knowledge_base.items():
            if any(keyword in query_lower for keyword in [category]):
                if isinstance(content, dict):
                    for key, value in content.items():
                        if any(keyword in query_lower for keyword in key.split('_')):
                            relevant_info.append(f"{category.title()} - {key}: {value}")
                else:
                    relevant_info.append(f"{category.title()}: {content}")
        
        return "\n".join(relevant_info[:5])  # Limit to top 5 matches
    
    def generate_response(self, user_query: str, chat_history: List[Tuple[str, str]]) -> str:
        """Generate response using Groq LLM with RAG context"""
        try:
            # Retrieve relevant context
            context = self.retrieve_relevant_context(user_query)
            
            # Prepare system prompt
            system_prompt = f"""You are a Power Systems Mini-Consultant AI assistant specializing in electrical power systems. 
            You help with:
            1. Fault analysis and calculations
            2. Protection system design
            3. Standards interpretation
            4. Study materials and exam preparation
            5. Practice problem generation
            
            Use the following context from the knowledge base:
            {context}
            
            Provide clear, technical explanations with practical examples. Include relevant formulas, standards references, and safety considerations when appropriate."""
            
            # Prepare conversation context
            messages = [{"role": "system", "content": system_prompt}]
            
            # Add chat history
            for human_msg, ai_msg in chat_history[-5:]:  # Last 5 exchanges
                messages.append({"role": "user", "content": human_msg})
                messages.append({"role": "assistant", "content": ai_msg})
            
            # Add current query
            messages.append({"role": "user", "content": user_query})
            
            # Generate response using Groq
            response = self.groq_client.chat.completions.create(
                model="mixtral-8x7b-32768",
                messages=messages,
                max_tokens=1500,
                temperature=0.7
            )
            
            return response.choices[0].message.content
            
        except Exception as e:
            return f"Error generating response: {str(e)}. Please check your GROQ_API_KEY and try again."

    def generate_practice_pack(self, topic: str, difficulty: str, num_questions: int) -> str:
        """Generate practice questions for power systems topics"""
        try:
            prompt = f"""Generate {num_questions} practice questions about {topic} in power systems 
            with {difficulty} difficulty level. Include:
            
            1. Multiple choice questions with 4 options
            2. Short answer questions
            3. Calculation problems with given data
            4. Provide answers and brief explanations
            
            Format as a structured practice pack suitable for exam preparation."""
            
            response = self.groq_client.chat.completions.create(
                model="mixtral-8x7b-32768",
                messages=[{"role": "user", "content": prompt}],
                max_tokens=2000,
                temperature=0.8
            )
            
            return response.choices[0].message.content
        except Exception as e:
            return f"Error generating practice pack: {str(e)}"

    def explain_standard(self, standard_name: str) -> str:
        """Explain power system standards"""
        try:
            prompt = f"""Provide a comprehensive explanation of the {standard_name} standard in power systems.
            Include:
            1. Purpose and scope
            2. Key requirements
            3. Applications
            4. Related standards
            5. Practical implementation notes
            
            Keep it technically accurate but accessible."""
            
            response = self.groq_client.chat.completions.create(
                model="mixtral-8x7b-32768",
                messages=[{"role": "user", "content": prompt}],
                max_tokens=1500,
                temperature=0.6
            )
            
            return response.choices[0].message.content
        except Exception as e:
            return f"Error explaining standard: {str(e)}"

# Initialize the consultant
try:
    consultant = PowerSystemsConsultant()
    initialization_status = "✅ Power Systems Consultant initialized successfully!"
except Exception as e:
    consultant = None
    initialization_status = f"❌ Initialization failed: {str(e)}"

def chat_interface(message, history):
    """Main chat interface function"""
    if consultant is None:
        return "❌ System not initialized. Please check GROQ_API_KEY."
    
    if not message.strip():
        return "Please enter a question about power systems."
    
    try:
        response = consultant.generate_response(message, history)
        return response
    except Exception as e:
        return f"Error: {str(e)}"

def generate_practice_questions(topic, difficulty, num_questions):
    """Generate practice pack interface"""
    if consultant is None:
        return "❌ System not initialized. Please check GROQ_API_KEY."
    
    try:
        practice_pack = consultant.generate_practice_pack(topic, difficulty, int(num_questions))
        return practice_pack
    except Exception as e:
        return f"Error generating practice pack: {str(e)}"

def explain_standard_interface(standard):
    """Standards explanation interface"""
    if consultant is None:
        return "❌ System not initialized. Please check GROQ_API_KEY."
    
    try:
        explanation = consultant.explain_standard(standard)
        return explanation
    except Exception as e:
        return f"Error explaining standard: {str(e)}"

# Create Gradio interface
with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="cyan",
        neutral_hue="slate"
    ),
    title="Power Systems Mini-Consultant",
    css="""
        .header-text { 
            text-align: center; 
            font-size: 2.5em; 
            font-weight: bold; 
            color: #1f4e79; 
            margin-bottom: 0.5em;
        }
        .subtitle-text { 
            text-align: center; 
            font-size: 1.2em; 
            color: #666; 
            margin-bottom: 2em;
        }
        .status-box {
            padding: 1em;
            border-radius: 8px;
            margin: 1em 0;
        }
        .feature-box {
            border: 2px solid #e0e7ff;
            border-radius: 12px;
            padding: 1.5em;
            margin: 1em 0;
            background: linear-gradient(135deg, #f8fafc 0%, #e2e8f0 100%);
        }
    """
) as app:
    
    # Header
    gr.HTML("""
        <div class="header-text">⚡ Power Systems Mini-Consultant</div>
        <div class="subtitle-text">
            Advanced AI Assistant for Fault Analysis, Protection Systems & Standards
        </div>
    """)
    
    # Status display
    gr.HTML(f"""
        <div class="status-box" style="background-color: {'#d4edda' if '✅' in initialization_status else '#f8d7da'}; 
                                      color: {'#155724' if '✅' in initialization_status else '#721c24'};">
            <strong>System Status:</strong> {initialization_status}
        </div>
    """)
    
    with gr.Tabs():
        # Main Chat Tab
        with gr.TabItem("💬 AI Consultant Chat"):
            gr.HTML("""
                <div class="feature-box">
                    <h3>🤖 Intelligent Power Systems Consultant</h3>
                    <p>Ask questions about fault analysis, protection systems, standards, calculations, and more!</p>
                    <ul>
                        <li>Fault analysis and calculations</li>
                        <li>Protection system design</li>
                        <li>Standards interpretation (IEEE, IEC)</li>
                        <li>Study guidance and explanations</li>
                    </ul>
                </div>
            """)
            
            chatbot = gr.Chatbot(
                height=400,
                placeholder="Your Power Systems AI Consultant will appear here...",
                show_label=False
            )
            
            msg = gr.Textbox(
                placeholder="Ask about faults, protection, standards, calculations...",
                show_label=False,
                container=False
            )
            
            with gr.Row():
                submit_btn = gr.Button("Send 📤", variant="primary")
                clear_btn = gr.Button("Clear Chat 🗑️", variant="secondary")
            
            # Example questions
            gr.HTML("""
                <div class="feature-box">
                    <h4>💡 Example Questions:</h4>
                    <ul>
                        <li>"Explain symmetrical vs unsymmetrical faults"</li>
                        <li>"How to calculate short circuit current?"</li>
                        <li>"What is IEEE C37 standard about?"</li>
                        <li>"Design overcurrent protection for a 33kV feeder"</li>
                    </ul>
                </div>
            """)
            
            submit_btn.click(
                chat_interface,
                inputs=[msg, chatbot],
                outputs=[chatbot]
            ).then(
                lambda: "",
                outputs=[msg]
            )
            
            msg.submit(
                chat_interface,
                inputs=[msg, chatbot],
                outputs=[chatbot]
            ).then(
                lambda: "",
                outputs=[msg]
            )
            
            clear_btn.click(lambda: None, outputs=[chatbot])
        
        # Practice Pack Generator
        with gr.TabItem("📚 Practice Pack Generator"):
            gr.HTML("""
                <div class="feature-box">
                    <h3>📝 AI-Generated Practice Questions</h3>
                    <p>Generate customized practice packs for power systems topics with varying difficulty levels.</p>
                </div>
            """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    topic_input = gr.Dropdown(
                        choices=[
                            "Fault Analysis", "Protection Systems", "Power Flow",
                            "Stability Studies", "Relay Coordination", "Transformers",
                            "Transmission Lines", "Load Flow", "Short Circuit Analysis",
                            "Harmonics", "Power Quality", "SCADA Systems"
                        ],
                        label="📋 Select Topic",
                        value="Fault Analysis"
                    )
                    
                    difficulty_input = gr.Radio(
                        choices=["Beginner", "Intermediate", "Advanced"],
                        label="🎯 Difficulty Level",
                        value="Intermediate"
                    )
                    
                    num_questions_input = gr.Slider(
                        minimum=5,
                        maximum=20,
                        step=1,
                        value=10,
                        label="🔢 Number of Questions"
                    )
                    
                    generate_btn = gr.Button("Generate Practice Pack 🎓", variant="primary")
                
                with gr.Column(scale=2):
                    practice_output = gr.Textbox(
                        label="📖 Generated Practice Pack",
                        lines=20,
                        placeholder="Your practice questions will appear here..."
                    )
            
            generate_btn.click(
                generate_practice_questions,
                inputs=[topic_input, difficulty_input, num_questions_input],
                outputs=[practice_output]
            )
        
        # Standards Explorer
        with gr.TabItem("📋 Standards Explorer"):
            gr.HTML("""
                <div class="feature-box">
                    <h3>📚 Power Systems Standards Guide</h3>
                    <p>Get detailed explanations of IEEE, IEC, and other international standards.</p>
                </div>
            """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    standard_input = gr.Dropdown(
                        choices=[
                            "IEEE C37.2", "IEEE 1547", "IEEE 519", "IEEE C57.12.00",
                            "IEC 61850", "IEC 60909", "IEC 61131", "IEC 60255",
                            "ANSI C84.1", "IEEE 242", "IEEE 399", "IEEE C37.010"
                        ],
                        label="📑 Select Standard",
                        value="IEEE C37.2"
                    )
                    
                    explain_btn = gr.Button("Explain Standard 📖", variant="primary")
                
                with gr.Column(scale=2):
                    standard_output = gr.Textbox(
                        label="📋 Standard Explanation",
                        lines=15,
                        placeholder="Standard explanation will appear here..."
                    )
            
            explain_btn.click(
                explain_standard_interface,
                inputs=[standard_input],
                outputs=[standard_output]
            )
        
        # Study Resources
        with gr.TabItem("📚 Study Resources"):
            gr.HTML("""
                <div class="feature-box">
                    <h3>📖 Power Systems Study Guide</h3>
                    <p>Essential formulas, concepts, and quick reference materials.</p>
                </div>
            """)
            
            with gr.Accordion("⚡ Fault Analysis Formulas", open=False):
                gr.HTML("""
                    <h4>Three-Phase Fault Current:</h4>
                    <p><code>I_fault = V_nominal / Z_total</code></p>
                    
                    <h4>Line-to-Ground Fault:</h4>
                    <p><code>I_fault = 3 * V_nominal / (Z1 + Z2 + Z0)</code></p>
                    
                    <h4>Line-to-Line Fault:</h4>
                    <p><code>I_fault = √3 * V_nominal / (Z1 + Z2)</code></p>
                """)
            
            with gr.Accordion("🛡️ Protection Principles", open=False):
                gr.HTML("""
                    <h4>Overcurrent Protection:</h4>
                    <ul>
                        <li>Inverse Time Characteristics</li>
                        <li>Coordination Time Intervals</li>
                        <li>Pickup Current Settings</li>
                    </ul>
                    
                    <h4>Distance Protection:</h4>
                    <ul>
                        <li>Zone 1: 80-90% of line length</li>
                        <li>Zone 2: 120% of protected line + 50% of next line</li>
                        <li>Zone 3: Backup protection</li>
                    </ul>
                """)
            
            with gr.Accordion("📊 Key Standards Summary", open=False):
                gr.HTML("""
                    <h4>IEEE Standards:</h4>
                    <ul>
                        <li><strong>IEEE C37:</strong> Circuit breakers and switchgear</li>
                        <li><strong>IEEE 1547:</strong> Distributed energy resources</li>
                        <li><strong>IEEE 519:</strong> Harmonic control</li>
                        <li><strong>IEEE 242:</strong> Protection and coordination</li>
                    </ul>
                    
                    <h4>IEC Standards:</h4>
                    <ul>
                        <li><strong>IEC 61850:</strong> Substation automation</li>
                        <li><strong>IEC 60909:</strong> Short-circuit calculations</li>
                        <li><strong>IEC 60255:</strong> Protective relays</li>
                    </ul>
                """)
    
    # Footer
    gr.HTML("""
        <div style="text-align: center; margin-top: 2em; padding: 1em; background-color: #f8f9fa; border-radius: 8px;">
            <p><strong>⚡ Power Systems Mini-Consultant</strong></p>
            <p>Built with Groq AI • Deployed on Hugging Face Spaces</p>
            <p><em>For educational and professional use in power systems engineering</em></p>
        </div>
    """)

if __name__ == "__main__":
    app.launch()