app.py

import os
import sys
import json
import logging
import warnings
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
import hashlib
import pickle
from datetime import datetime
import time
import asyncio
from concurrent.futures import ThreadPoolExecutor

from transformers import AutoModelForSeq2SeqLM  # ✅ Needed for T5 and FLAN models


# Suppress warnings for cleaner output
warnings.filterwarnings("ignore")

# Core dependencies
import gradio as gr
import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer
import faiss
import torch
from transformers import (
    AutoTokenizer, 
    AutoModelForSeq2SeqLM,  
    BitsAndBytesConfig,
    pipeline
)

# Medical knowledge validation
import re

import subprocess
import gradio as gr

def train_model():
    result = subprocess.run(["python", "finetune_flan_t5.py"], capture_output=True, text=True)
    if result.returncode == 0:
        return "✅ Model training complete!"
    else:
        return f"❌ Training failed:\n{result.stderr}"


# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class MedicalFactChecker:
    """Enhanced medical fact checker with faster validation"""
    
    def __init__(self):
        self.medical_facts = self._load_medical_facts()
        self.contraindications = self._load_contraindications()
        self.dosage_patterns = self._compile_dosage_patterns()
        self.definitive_patterns = [
            re.compile(r, re.IGNORECASE) for r in [
                r'always\s+(?:use|take|apply)',
                r'never\s+(?:use|take|apply)',
                r'will\s+(?:cure|heal|fix)',
                r'guaranteed\s+to',
                r'completely\s+(?:safe|effective)'
            ]
        ]
        
    def _load_medical_facts(self) -> Dict[str, Any]:
        """Pre-loaded medical facts for Gaza context"""
        return {
            "burn_treatment": {
                "cool_water": "Use clean, cool (not ice-cold) water for 10-20 minutes",
                "no_ice": "Never apply ice directly to burns",
                "clean_cloth": "Cover with clean, dry cloth if available"
            },
            "wound_care": {
                "pressure": "Apply direct pressure to control bleeding",
                "elevation": "Elevate injured limb if possible",
                "clean_hands": "Clean hands before treating wounds when possible"
            },
            "infection_signs": {
                "redness": "Increasing redness around wound",
                "warmth": "Increased warmth at wound site",
                "pus": "Yellow or green discharge",
                "fever": "Fever may indicate systemic infection"
            }
        }
    
    def _load_contraindications(self) -> Dict[str, List[str]]:
        """Pre-loaded contraindications for common treatments"""
        return {
            "aspirin": ["children under 16", "bleeding disorders", "stomach ulcers"],
            "ibuprofen": ["kidney disease", "heart failure", "stomach bleeding"],
            "hydrogen_peroxide": ["deep wounds", "closed wounds", "eyes"],
            "tourniquets": ["non-life-threatening bleeding", "without proper training"]
        }
    
    def _compile_dosage_patterns(self) -> List[re.Pattern]:
        """Pre-compiled dosage patterns"""
        patterns = [
            r'\d+\s*mg\b',  # milligrams
            r'\d+\s*g\b',   # grams
            r'\d+\s*ml\b',  # milliliters
            r'\d+\s*tablets?\b',  # tablets
            r'\d+\s*times?\s+(?:per\s+)?day\b',  # frequency
            r'every\s+\d+\s+hours?\b'  # intervals
        ]
        return [re.compile(pattern, re.IGNORECASE) for pattern in patterns]
    
    def check_medical_accuracy(self, response: str, context: str) -> Dict[str, Any]:
        """Enhanced medical accuracy check with Gaza-specific considerations"""
        issues = []
        warnings = []
        accuracy_score = 0.0
        
        # Check for contraindications (faster keyword matching)
        response_lower = response.lower()
        for medication, contra_list in self.contraindications.items():
            if medication in response_lower:
                for contra in contra_list:
                    if any(word in response_lower for word in contra.split()):
                        issues.append(f"Potential contraindication: {medication} with {contra}")
                        accuracy_score -= 0.3
                        break
        
        # Context alignment using Jaccard similarity
        if context:
            resp_words = set(response_lower.split())
            ctx_words = set(context.lower().split())
            context_similarity = len(resp_words & ctx_words) / len(resp_words | ctx_words) if ctx_words else 0.0
            if context_similarity < 0.5:  # Lowered threshold for Gaza context
                warnings.append(f"Low context similarity: {context_similarity:.2f}")
                accuracy_score -= 0.1
        else:
            context_similarity = 0.0
        
        # Gaza-specific resource checks
        gaza_resources = ["clean water", "sterile", "hospital", "ambulance", "electricity"]
        if any(resource in response_lower for resource in gaza_resources):
            warnings.append("Consider resource limitations in Gaza context")
            accuracy_score -= 0.05
        
        # Unsupported claims check
        for pattern in self.definitive_patterns:
            if pattern.search(response):
                issues.append(f"Unsupported definitive claim detected")
                accuracy_score -= 0.4
                break
        
        # Dosage validation
        for pattern in self.dosage_patterns:
            if pattern.search(response):
                warnings.append("Dosage detected - verify with professional")
                accuracy_score -= 0.1
                break
        
        confidence_score = max(0.0, min(1.0, 0.8 + accuracy_score))
        
        return {
            "confidence_score": confidence_score,
            "issues": issues,
            "warnings": warnings,
            "context_similarity": context_similarity,
            "is_safe": len(issues) == 0 and confidence_score > 0.5
        }

class OptimizedGazaKnowledgeBase:
    """Optimized knowledge base that loads pre-made FAISS index and assets"""
    
    def __init__(self, vector_store_dir: str = "./vector_store"):
        self.vector_store_dir = Path(vector_store_dir)
        self.faiss_index = None
        self.embedding_model = None
        self.chunks = []
        self.metadata = []
        self.is_initialized = False
        
    def initialize(self):
        """Load pre-made FAISS index and associated data"""
        try:
            logger.info("🔄 Loading pre-made FAISS index and assets...")
            
            # 1. Load FAISS index
            index_path = self.vector_store_dir / "index.faiss"
            if not index_path.exists():
                raise FileNotFoundError(f"FAISS index not found at {index_path}")
            
            self.faiss_index = faiss.read_index(str(index_path))
            logger.info(f"✅ Loaded FAISS index: {self.faiss_index.ntotal} vectors, {self.faiss_index.d} dimensions")
            
            # 2. Load chunks
            chunks_path = self.vector_store_dir / "chunks.txt"
            if not chunks_path.exists():
                raise FileNotFoundError(f"Chunks file not found at {chunks_path}")
            
            with open(chunks_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # Parse chunks from the formatted file
            current_chunk = ""
            for line in lines:
                line = line.strip()
                if line.startswith("=== Chunk") and current_chunk:
                    self.chunks.append(current_chunk.strip())
                    current_chunk = ""
                elif not line.startswith("===") and not line.startswith("Source:") and not line.startswith("Length:"):
                    current_chunk += line + " "
            
            # Add the last chunk
            if current_chunk:
                self.chunks.append(current_chunk.strip())
            
            logger.info(f"✅ Loaded {len(self.chunks)} text chunks")
            
            # 3. Load metadata
            metadata_path = self.vector_store_dir / "metadata.pkl"
            if metadata_path.exists():
                with open(metadata_path, 'rb') as f:
                    metadata_dict = pickle.load(f)
                
                if isinstance(metadata_dict, dict) and 'metadata' in metadata_dict:
                    self.metadata = metadata_dict['metadata']
                    logger.info(f"✅ Loaded {len(self.metadata)} metadata entries")
                else:
                    logger.warning("⚠️ Metadata format not recognized, using empty metadata")
                    self.metadata = [{}] * len(self.chunks)
            else:
                logger.warning("⚠️ No metadata file found, using empty metadata")
                self.metadata = [{}] * len(self.chunks)
            
            # 4. Initialize embedding model for query encoding
            logger.info("🔄 Loading embedding model for queries...")
            self.embedding_model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
            logger.info("✅ Embedding model loaded")
            
            # 5. Verify data consistency
            if len(self.chunks) != self.faiss_index.ntotal:
                logger.warning(f"⚠️ Mismatch: {len(self.chunks)} chunks vs {self.faiss_index.ntotal} vectors")
                # Trim chunks to match index size
                self.chunks = self.chunks[:self.faiss_index.ntotal]
                self.metadata = self.metadata[:self.faiss_index.ntotal]
                logger.info(f"✅ Trimmed to {len(self.chunks)} chunks to match index")
            
            self.is_initialized = True
            logger.info("🎉 Knowledge base initialization complete!")
            
        except Exception as e:
            logger.error(f"❌ Failed to initialize knowledge base: {e}")
            raise
    
    def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """Search using pre-made FAISS index"""
        if not self.is_initialized:
            raise RuntimeError("Knowledge base not initialized")
        
        try:
            # 1. Encode query
            query_embedding = self.embedding_model.encode([query])
            query_vector = np.array(query_embedding, dtype=np.float32)
            
            # 2. Search FAISS index
            distances, indices = self.faiss_index.search(query_vector, k)
            
            # 3. Prepare results
            results = []
            for i, (distance, idx) in enumerate(zip(distances[0], indices[0])):
                if idx >= 0 and idx < len(self.chunks):  # Valid index
                    chunk_metadata = self.metadata[idx] if idx < len(self.metadata) else {}
                    
                    result = {
                        "text": self.chunks[idx],
                        "score": float(1.0 / (1.0 + distance)),  # Convert distance to similarity score
                        "source": chunk_metadata.get("source", "unknown"),
                        "chunk_index": int(idx),
                        "distance": float(distance),
                        "metadata": chunk_metadata
                    }
                    results.append(result)
            
            logger.info(f"🔍 Search for '{query[:50]}...' returned {len(results)} results")
            return results
            
        except Exception as e:
            logger.error(f"❌ Search error: {e}")
            return []
    
    def get_stats(self) -> Dict[str, Any]:
        """Get knowledge base statistics"""
        if not self.is_initialized:
            return {"status": "not_initialized"}
        
        return {
            "status": "initialized",
            "total_chunks": len(self.chunks),
            "total_vectors": self.faiss_index.ntotal,
            "embedding_dimension": self.faiss_index.d,
            "index_type": type(self.faiss_index).__name__,
            "sources": list(set(meta.get("source", "unknown") for meta in self.metadata))
        }

class OptimizedGazaRAGSystem:
    """Optimized RAG system using pre-made assets"""
    
    def __init__(self, vector_store_dir: str = "./vector_store"):
        self.knowledge_base = OptimizedGazaKnowledgeBase(vector_store_dir)
        self.fact_checker = MedicalFactChecker()
        self.llm = None
        self.tokenizer = None
        self.system_prompt = self._create_system_prompt()
        self.generation_pipeline = None
        self.response_cache = {}
        self.executor = ThreadPoolExecutor(max_workers=2)
        
    def initialize(self):
        """Initialize the optimized RAG system"""
        logger.info("🚀 Initializing Optimized Gaza RAG System...")
        self.knowledge_base.initialize()
        logger.info("✅ Optimized Gaza RAG System ready!")
        

    
    def _initialize_llm(self):
        """Load flan-t5-base for CPU fallback"""
        model_name = "google/flan-t5-base"
        try:
            logger.info(f"🔄 Loading fallback CPU model: {model_name}")
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            self.llm = AutoModelForSeq2SeqLM.from_pretrained(model_name)
            self.generation_pipeline = pipeline(
                "text2text-generation",  # ✅ correct pipeline for T5
                model=self.llm,
                tokenizer=self.tokenizer
            )
            logger.info("✅ FLAN-T5 model loaded successfully")
        except Exception as e:
            logger.error(f"❌ Error loading FLAN-T5 model: {e}")
            self.llm = None


    
    def _initialize_fallback_llm(self):
        """Enhanced fallback model with better error handling"""
        try:
            logger.info("🔄 Loading fallback model...")
            
            fallback_model = "microsoft/DialoGPT-small"
            self.tokenizer = AutoTokenizer.from_pretrained(fallback_model)
            self.llm = AutoModelForCausalLM.from_pretrained(
                fallback_model,
                torch_dtype=torch.float32,
                low_cpu_mem_usage=True
            )
            
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token
            
            self.generation_pipeline = pipeline(
                "text-generation",
                model=self.llm,
                tokenizer=self.tokenizer,
                return_full_text=False
            )
            
            logger.info("✅ Fallback model loaded successfully")
            
        except Exception as e:
            logger.error(f"❌ Fallback model failed: {e}")
            self.llm = None
            self.generation_pipeline = None
    
    def _create_system_prompt(self) -> str:
        """Enhanced system prompt for Gaza context"""
        return """You are a medical AI assistant specifically designed for Gaza healthcare workers operating under siege conditions. 

CRITICAL GUIDELINES:
- Provide practical first aid guidance considering limited resources (water, electricity, medical supplies)
- Always prioritize patient safety and recommend professional medical help when available
- Consider Gaza's specific challenges: blockade, limited hospitals, frequent power outages
- Suggest alternative treatments when standard medical supplies are unavailable
- Never provide definitive diagnoses - only supportive care guidance
- Be culturally sensitive and aware of the humanitarian crisis context

RESOURCE CONSTRAINTS TO CONSIDER:
- Limited clean water availability
- Frequent electricity outages
- Restricted medical supply access
- Overwhelmed healthcare facilities
- Limited transportation for medical emergencies

Provide clear, actionable advice while emphasizing the need for professional medical care when possible."""
    
    async def generate_response_async(self, query: str, progress_callback=None) -> Dict[str, Any]:
        """Async response generation with progress tracking"""
        start_time = time.time()
        
        if progress_callback:
            progress_callback(0.1, "🔍 Checking cache...")
        
        # Check cache first
        query_hash = hashlib.md5(query.encode()).hexdigest()
        if query_hash in self.response_cache:
            cached_response = self.response_cache[query_hash]
            cached_response["cached"] = True
            cached_response["response_time"] = 0.1
            if progress_callback:
                progress_callback(1.0, "💾 Retrieved from cache!")
            return cached_response
        
        try:
            if progress_callback:
                progress_callback(0.2, "🤖 Initializing LLM...")
            
            # Initialize LLM only when needed
            if self.llm is None:
                await asyncio.get_event_loop().run_in_executor(
                    self.executor, self._initialize_llm
                )
            
            if progress_callback:
                progress_callback(0.4, "🔍 Searching knowledge base...")
                
            # Enhanced knowledge retrieval using pre-made index
            search_results = await asyncio.get_event_loop().run_in_executor(
                self.executor, self.knowledge_base.search, query, 5
            )
            
            if progress_callback:
                progress_callback(0.6, "📝 Preparing context...")
            
            context = self._prepare_context(search_results)
            
            if progress_callback:
                progress_callback(0.8, "🧠 Generating response...")
            
            # Generate response
            response = await asyncio.get_event_loop().run_in_executor(
                self.executor, self._generate_response, query, context
            )
            
            if progress_callback:
                progress_callback(0.9, "🛡️ Validating safety...")
            
            # Enhanced safety check
            safety_check = self.fact_checker.check_medical_accuracy(response, context)
            
            # Prepare final response
            final_response = self._prepare_final_response(
                response, 
                search_results, 
                safety_check,
                time.time() - start_time
            )
            
            # Cache the response (limit cache size)
            if len(self.response_cache) < 100:
                self.response_cache[query_hash] = final_response
            
            if progress_callback:
                progress_callback(1.0, "✅ Complete!")
            
            return final_response
            
        except Exception as e:
            logger.error(f"❌ Error generating response: {e}")
            if progress_callback:
                progress_callback(1.0, f"❌ Error: {str(e)}")
            return self._create_error_response(str(e))
    
    def _prepare_context(self, search_results: List[Dict[str, Any]]) -> str:
        """Enhanced context preparation with better formatting"""
        if not search_results:
            return "No specific medical guidance found in knowledge base. Provide general first aid principles."
        
        context_parts = []
        for i, result in enumerate(search_results, 1):
            source = result.get('source', 'unknown')
            text = result.get('text', '')
            score = result.get('score', 0.0)
            
            # Truncate long text but preserve important information
            if len(text) > 400:
                text = text[:400] + "..."
            
            context_parts.append(f"[Source {i}: {source} - Relevance: {score:.2f}]\n{text}")
        
        return "\n\n".join(context_parts)
    
    def _generate_response(self, query: str, context: str) -> str:
        """Generate response using T5-style seq2seq model with Gaza-specific context"""
        if self.llm is None or self.tokenizer is None:
            return self._generate_fallback_response(query, context)
            
        prompt = f"""{self.system_prompt}   
MEDICAL KNOWLEDGE CONTEXT:
{context}

PATIENT QUESTION: {query}

RESPONSE (provide practical, Gaza-appropriate medical guidance):"""
        try:
            inputs = self.tokenizer(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=512,
                padding="max_length"
            )
            input_ids = inputs["input_ids"]
            attention_mask = inputs["attention_mask"]
            device = self.llm.device if hasattr(self.llm, "device") else "cpu"
            input_ids = input_ids.to(device)
            attention_mask = attention_mask.to(device)
            
            with torch.no_grad():
                outputs = self.llm.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=256,
                temperature=0.3,
                pad_token_id=self.tokenizer.eos_token_id,
                do_sample=True,
                repetition_penalty=1.15,
                no_repeat_ngram_size=3
                )
            response_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            lines = response_text.split('\n')
            unique_lines = []
            for line in lines:
                line = line.strip()
                if line and line not in unique_lines and len(line) > 10:
                    unique_lines.append(line)
            
            final_response = '\n'.join(unique_lines)
            logger.info(f"🧪 Final cleaned response:\n{final_response}")
            return final_response
        except Exception as e:
            logger.error(f"❌ Error in LLM generate(): {e}")
            return self._generate_fallback_response(query, context)


    
    def _generate_fallback_response(self, query: str, context: str) -> str:
        """Enhanced fallback response with Gaza-specific guidance"""
        gaza_guidance = {
            "burn": "For burns: Use clean, cool water if available. If water is scarce, use clean cloth. Avoid ice. Seek medical help urgently.",
            "bleeding": "For bleeding: Apply direct pressure with clean cloth. Elevate if possible. If severe, seek immediate medical attention.",
            "wound": "For wounds: Clean hands if possible. Apply pressure to stop bleeding. Cover with clean material. Watch for infection signs.",
            "infection": "Signs of infection: Redness, warmth, swelling, pus, fever. Seek medical care immediately if available.",
            "pain": "For pain management: Rest, elevation, cold/warm compress as appropriate. Avoid aspirin in children."
        }
        
        query_lower = query.lower()
        for condition, guidance in gaza_guidance.items():
            if condition in query_lower:
                return f"{guidance}\n\nContext from medical sources:\n{context[:200]}..."
        
        return f"Medical guidance for: {query}\n\nGeneral advice: Prioritize safety, seek professional help when available, consider resource limitations in Gaza.\n\nRelevant information:\n{context[:600]}..."
    
    def _prepare_final_response(
        self, 
        response: str, 
        search_results: List[Dict[str, Any]], 
        safety_check: Dict[str, Any],
        response_time: float
    ) -> Dict[str, Any]:
        """Enhanced final response preparation with more metadata"""
        
        # Add safety warnings if needed
        if not safety_check["is_safe"]:
            response = f"⚠️ MEDICAL CAUTION: {response}\n\n🚨 Please verify this guidance with a medical professional when possible."
        
        # Add Gaza-specific disclaimer
        response += "\n\n📍 Gaza Context: This guidance considers resource limitations. Adapt based on available supplies and seek professional medical care when accessible."
        
        # Extract unique sources
        sources = list(set(res.get("source", "unknown") for res in search_results)) if search_results else []
        
        # Calculate confidence based on multiple factors
        base_confidence = safety_check.get("confidence_score", 0.5)
        context_bonus = 0.1 if search_results else 0.0
        safety_penalty = 0.2 if not safety_check.get("is_safe", True) else 0.0
        
        final_confidence = max(0.0, min(1.0, base_confidence + context_bonus - safety_penalty))
        
        return {
            "response": response,
            "confidence": final_confidence,
            "sources": sources,
            "search_results_count": len(search_results),
            "safety_issues": safety_check.get("issues", []),
            "safety_warnings": safety_check.get("warnings", []),
            "response_time": round(response_time, 2),
            "timestamp": datetime.now().isoformat()[:19],
            "cached": False
        }
    
    def _create_error_response(self, error_msg: str) -> Dict[str, Any]:
        """Enhanced error response with helpful information"""
        return {
            "response": f"⚠️ System Error: Unable to process your medical query at this time.\n\nError: {error_msg}\n\n🚨 For immediate medical emergencies, seek professional help directly.\n\n📞 Gaza Emergency Numbers:\n- Palestinian Red Crescent: 101\n- Civil Defense: 102",
            "confidence": 0.0,
            "sources": [],
            "search_results_count": 0,
            "safety_issues": ["System error occurred"],
            "safety_warnings": ["Unable to validate medical accuracy"],
            "response_time": 0.0,
            "timestamp": datetime.now().isoformat()[:19],
            "cached": False,
            "error": True
        }

# Global system instance
optimized_rag_system = None

def initialize_optimized_system(vector_store_dir: str = "./vector_store"):
    """Initialize optimized system with pre-made assets"""
    global optimized_rag_system
    if optimized_rag_system is None:
        try:
            optimized_rag_system = OptimizedGazaRAGSystem(vector_store_dir)
            optimized_rag_system.initialize()
            logger.info("✅ Optimized Gaza RAG System initialized successfully")
        except Exception as e:
            logger.error(f"❌ Failed to initialize optimized system: {e}")
            raise
    return optimized_rag_system

def process_medical_query_with_progress(query: str, progress=gr.Progress()) -> Tuple[str, str, str]:
    """Enhanced query processing with detailed progress tracking and status updates"""
    if not query.strip():
        return "Please enter a medical question.", "", "⚠️ No query provided"
    
    try:
        # Initialize system with progress
        progress(0.05, desc="🔧 Initializing optimized system...")
        system = initialize_optimized_system()
        
        # Create async event loop for progress tracking
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        def progress_callback(value, desc):
            progress(value, desc=desc)
        
        try:
            # Run async generation with progress
            result = loop.run_until_complete(
                system.generate_response_async(query, progress_callback)
            )
        finally:
            loop.close()
        
        # Prepare response with enhanced metadata
        response = result["response"]
        
        # Prepare detailed metadata
        metadata_parts = [
            f"🎯 Confidence: {result['confidence']:.1%}",
            f"⏱️ Response: {result['response_time']}s",
            f"📚 Sources: {result['search_results_count']} found"
        ]
        
        if result.get('cached'):
            metadata_parts.append("💾 Cached")
        
        if result.get('sources'):
            metadata_parts.append(f"📖 Refs: {', '.join(result['sources'][:2])}")
        
        metadata = " | ".join(metadata_parts)
        
        # Prepare status with warnings/issues
        status_parts = []
        if result.get('safety_warnings'):
            status_parts.append(f"⚠️ {len(result['safety_warnings'])} warnings")
        if result.get('safety_issues'):
            status_parts.append(f"🚨 {len(result['safety_issues'])} issues")
        if not status_parts:
            status_parts.append("✅ Safe response")
        
        status = " | ".join(status_parts)
        
        return response, metadata, status
        
    except Exception as e:
        logger.error(f"❌ Error processing query: {e}")
        error_response = f"⚠️ Error processing your query: {str(e)}\n\n🚨 For medical emergencies, seek immediate professional help."
        error_metadata = f"❌ Error at {datetime.now().strftime('%H:%M:%S')}"
        error_status = "🚨 System error occurred"
        return error_response, error_metadata, error_status

def get_system_stats() -> str:
    """Get system statistics for display"""
    try:
        system = initialize_optimized_system()
        stats = system.knowledge_base.get_stats()
        
        if stats["status"] == "initialized":
            return f"""
📊 **System Statistics:**
- Status: ✅ Initialized
- Total Chunks: {stats['total_chunks']:,}
- Vector Dimension: {stats['embedding_dimension']}
- Index Type: {stats['index_type']}
- Sources: {len(stats['sources'])} documents
- Available Sources: {', '.join(stats['sources'][:5])}{'...' if len(stats['sources']) > 5 else ''}
"""
        else:
            return "📊 System Status: ❌ Not Initialized"
    except Exception as e:
        return f"📊 System Status: ❌ Error - {str(e)}"

def create_optimized_gradio_interface():
    """Create optimized Gradio interface with enhanced features"""
    
    # Enhanced CSS with medical theme
    css = """
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

    * {
        font-family: 'Inter', sans-serif !important;
    }

    .gradio-container {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        min-height: 100vh;
    }
    
    .main-container {
        background: rgba(255, 255, 255, 0.95);
        backdrop-filter: blur(10px);
        border-radius: 20px;
        padding: 30px;
        margin: 20px;
        box-shadow: 0 20px 40px rgba(0,0,0,0.1);
        border: 1px solid rgba(255,255,255,0.2);
    }
    
    .header-section {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        border-radius: 15px;
        padding: 25px;
        margin-bottom: 25px;
        text-align: center;
        box-shadow: 0 10px 30px rgba(102, 126, 234, 0.3);
    }
    
    .query-container {
        background: linear-gradient(135deg, #f8f9ff 0%, #e8f2ff 100%);
        border-radius: 15px;
        padding: 20px;
        margin: 15px 0;
        border: 2px solid #667eea;
        transition: all 0.3s ease;
    }
    
    .response-container {
        background: linear-gradient(135deg, #fff 0%, #f8f9ff 100%);
        border-radius: 15px;
        padding: 20px;
        margin: 15px 0;
        border: 2px solid #4CAF50;
        min-height: 300px;
    }
    
    .submit-btn {
        background: linear-gradient(135deg, #4CAF50 0%, #45a049 100%) !important;
        color: white !important;
        border: none !important;
        border-radius: 12px !important;
        padding: 15px 30px !important;
        font-size: 16px !important;
        font-weight: 600 !important;
        cursor: pointer !important;
        transition: all 0.3s ease !important;
        box-shadow: 0 6px 20px rgba(76, 175, 80, 0.3) !important;
    }
    
    .submit-btn:hover {
        transform: translateY(-3px) !important;
        box-shadow: 0 10px 30px rgba(76, 175, 80, 0.4) !important;
    }
    
    .stats-container {
        background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%);
        border-radius: 12px;
        padding: 15px;
        margin: 10px 0;
        border-left: 5px solid #2196F3;
        font-size: 14px;
    }
    """

    with gr.Blocks(
        css=css, 
        title="🏥 Optimized Gaza First Aid Assistant",
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="green",
            neutral_hue="slate"
        )
    ) as interface:
        
        # Header Section
        with gr.Row(elem_classes=["main-container"]):
            gr.HTML("""
            <div class="header-section">
                <h1 style="margin: 0; font-size: 2.5em; font-weight: 700;">
                    🏥 Optimized Gaza First Aid Assistant
                </h1>
                <h2 style="margin: 10px 0 0 0; font-size: 1.2em; font-weight: 400; opacity: 0.9;">
                    Powered by Pre-computed FAISS Index & 768-dim Embeddings
                </h2>
                <p style="margin: 15px 0 0 0; font-size: 1em; opacity: 0.8;">
                    Lightning-fast medical guidance using pre-processed knowledge base
                </p>
            </div>
            """)
        
        # System Stats
        with gr.Row(elem_classes=["main-container"]):
            with gr.Group(elem_classes=["stats-container"]):
                stats_display = gr.Markdown(
                    value=get_system_stats(),
                    label="📊 System Status"
                )
        
        # Main Interface
        with gr.Row(elem_classes=["main-container"]):
            with gr.Column(scale=2):
                # Query Input Section
                with gr.Group(elem_classes=["query-container"]):
                    gr.Markdown("### 🩺 Medical Query Input")
                    query_input = gr.Textbox(
                        label="Describe your medical situation",
                        placeholder="Enter your first aid question or describe the medical emergency...",
                        lines=4
                    )
                    
                    with gr.Row():
                        submit_btn = gr.Button(
                            "🔍 Get Medical Guidance", 
                            variant="primary",
                            elem_classes=["submit-btn"],
                            scale=3
                        )
                        clear_btn = gr.Button(
                            "🗑️ Clear", 
                            variant="secondary",
                            scale=1
                        )
                        train_btn = gr.Button("Train")
                        train_output = gr.Textbox(label="Training Status", lines=10)
                        train_btn.click(train_model, outputs=train_output)
            
            with gr.Column(scale=1):
                # Quick Access
                gr.Markdown("""
                ### ⚡ Optimized Features
                
                **🚀 Performance:**
                - Pre-computed FAISS index
                - 768-dimensional embeddings
                - Lightning-fast search
                - Optimized for Gaza context
                
                **📚 Knowledge Base:**
                - WHO medical protocols
                - ICRC war surgery guides
                - MSF field manuals
                - Gaza-specific adaptations
                
                **🛡️ Safety Features:**
                - Real-time fact checking
                - Contraindication detection
                - Gaza resource warnings
                - Professional disclaimers
                """)
        
        # Response Section
        with gr.Row(elem_classes=["main-container"]):
            with gr.Column():
                # Main Response
                with gr.Group(elem_classes=["response-container"]):
                    gr.Markdown("### 🩹 Medical Guidance Response")
                    response_output = gr.Textbox(
                        label="AI Medical Guidance",
                        lines=15,
                        interactive=False,
                        placeholder="Your medical guidance will appear here..."
                    )
                
                # Metadata and Status
                with gr.Row():
                    with gr.Column(scale=1):
                        metadata_output = gr.Textbox(
                            label="📊 Response Metadata",
                            lines=2,
                            interactive=False,
                            placeholder="Response metadata will appear here..."
                        )
                    
                    with gr.Column(scale=1):
                        status_output = gr.Textbox(
                            label="🛡️ Safety Status",
                            lines=2,
                            interactive=False,
                            placeholder="Safety validation status will appear here..."
                        )
        
        # Examples Section
        with gr.Row(elem_classes=["main-container"]):
            gr.Markdown("### 💡 Example Medical Scenarios")
            
            example_queries = [
                "How to treat severe burns when clean water is extremely limited?",
                "Managing gunshot wounds with only basic household supplies",
                "Recognizing and treating infection in wounds without antibiotics",
                "Emergency care for children during extended power outages",
                "Treating compound fractures without proper medical equipment"
            ]
            
            gr.Examples(
                examples=example_queries,
                inputs=query_input,
                label="Click any example to try it:",
                examples_per_page=5
            )
        
        # Event Handlers
        submit_btn.click(
            process_medical_query_with_progress,
            inputs=query_input,
            outputs=[response_output, metadata_output, status_output],
            show_progress=True
        )
        
        query_input.submit(
            process_medical_query_with_progress,
            inputs=query_input,
            outputs=[response_output, metadata_output, status_output],
            show_progress=True
        )
        
        clear_btn.click(
            lambda: ("", "", "", ""),
            outputs=[query_input, response_output, metadata_output, status_output]
        )
        
        # Refresh stats button
        refresh_stats_btn = gr.Button("🔄 Refresh System Stats", variant="secondary")
        refresh_stats_btn.click(
            lambda: get_system_stats(),
            outputs=stats_display
        )
    
    return interface

def main():
    """Enhanced main function with optimized system initialization"""
    logger.info("🚀 Starting Optimized Gaza First Aid Assistant")
    
    try:
        # Check for vector store directory
        vector_store_dir = "./vector_store"
        if not Path(vector_store_dir).exists():
            # Try alternative paths
            alt_paths = ["./results/vector_store", "./results/vector_store_extracted"]
            for alt_path in alt_paths:
                if Path(alt_path).exists():
                    vector_store_dir = alt_path
                    logger.info(f"📁 Found vector store at: {vector_store_dir}")
                    break
            else:
                raise FileNotFoundError("Vector store directory not found. Please ensure pre-made assets are available.")
        
        # System initialization with detailed logging
        logger.info(f"🔧 Loading optimized system from: {vector_store_dir}")
        system = initialize_optimized_system(vector_store_dir)
        
        # Verify system components
        stats = system.knowledge_base.get_stats()
        logger.info(f"✅ Knowledge base loaded: {stats['total_chunks']} chunks, {stats['embedding_dimension']}D")
        logger.info(f"✅ Sources: {len(stats['sources'])} documents")
        logger.info("✅ Medical fact checker ready")
        logger.info("✅ Optimized FAISS indexing active")
        
        # Create and launch optimized interface
        logger.info("🎨 Creating optimized Gradio interface...")
        interface = create_optimized_gradio_interface()
        
        logger.info("🌐 Launching optimized interface...")
        interface.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False,
            max_threads=6,
            show_error=True,
            quiet=False
        )
        
    except Exception as e:
        logger.error(f"❌ Failed to start Optimized Gaza First Aid Assistant: {e}")
        print(f"\n🚨 STARTUP ERROR: {e}")
        print("\n🔧 Troubleshooting Steps:")
        print("1. Ensure vector_store directory exists with index.faiss, chunks.txt, and metadata.pkl")
        print("2. Check if all dependencies are installed: pip install -r requirements.txt")
        print("3. Verify sufficient memory is available (minimum 4GB RAM recommended)")
        print("4. Check system logs for detailed error information")
        print("\n📞 For technical support, check the application logs above.")
        sys.exit(1)

if __name__ == "__main__":
    main()