# Setup and Installation

import torch
print("🖥️ System Check:")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU device: {torch.cuda.get_device_name(0)}")
    print(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
else:
    print("⚠️ No GPU detected - BioGPT will run on CPU (much slower)")

print("\n🔧 Installing required packages...")


# Import Libraries

import os
import re
import torch
import warnings
import numpy as np
import faiss  # FAISS for vector search
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    pipeline,
    BitsAndBytesConfig
)
from sentence_transformers import SentenceTransformer
from typing import List, Dict, Optional
import time
from datetime import datetime
import json
import pickle

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

print("📚 Libraries imported successfully!")
print(f"🔍 FAISS version: {faiss.__version__}")
print("🎯 Using FAISS for vector search (ChromaDB completely removed)")

# File Upload Helper

import io

def upload_medical_data():
    """Upload your Pediatric_cleaned.txt file"""
    print("📁 Please upload your Pediatric_cleaned.txt file:")
    uploaded = files.upload()

    # Get the uploaded file
    filename = list(uploaded.keys())[0]
    print(f"✅ File '{filename}' uploaded successfully!")

    # Read the content
    content = uploaded[filename].decode('utf-8')

    # Save it locally in Colab
    with open('Pediatric_cleaned.txt', 'w', encoding='utf-8') as f:
        f.write(content)

    print(f"📝 File saved as 'Pediatric_cleaned.txt' ({len(content)} characters)")
    return 'Pediatric_cleaned.txt'

medical_file = 'Pediatric_cleaned.txt'

# BioGPT Medical Chatbot Class


class ColabBioGPTChatbot:
    def __init__(self, use_gpu=True, use_8bit=True):
        """Initialize BioGPT chatbot optimized for Google Colab"""
        print("🏥 Initializing Professional BioGPT Medical Chatbot...")

        self.device = "cuda" if torch.cuda.is_available() and use_gpu else "cpu"
        self.use_8bit = use_8bit and torch.cuda.is_available()

        print(f"🖥️ Using device: {self.device}")
        if self.use_8bit:
            print("💾 Using 8-bit quantization for memory efficiency")

        # Setup components
        self.setup_embeddings()
        self.setup_faiss_index() # Ensure this sets up self.collection if needed
        self.setup_biogpt()

        # Conversation tracking
        self.conversation_history = []
        self.knowledge_chunks = []

        print("✅ BioGPT Medical Chatbot ready for professional medical assistance!")

    def setup_embeddings(self):
        """Setup medical-optimized embeddings"""
        print("🔧 Loading medical embeddings...")
        try:
            # Use a medical-focused embedding model if available, otherwise general
            self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
            self.embedding_dim = self.embedding_model.get_sentence_embedding_dimension()
            print(f"✅ Embeddings loaded (dimension: {self.embedding_dim})")
            self.use_embeddings = True
        except Exception as e:
            print(f"⚠️ Embeddings failed: {e}")
            self.embedding_model = None
            self.embedding_dim = 384 # default dimension
            self.use_embeddings = False

    def setup_faiss_index(self):
        """Setup faiss for CPU-based vector search"""
        print("🔧 Setting up FAISS vector database...")
        try:
            print(' Using CPU FAISS index for maximum compatibility')
            self.faiss_index = faiss.IndexFlatIP(self.embedding_dim)  # In-memory for Colab
            self.use_gpu_faiss = False
            self.faiss_ready = True # Set to True when index is ready
            self.collection = self.faiss_index # Initialize collection attribute
            print("✅ FAISS CPU index initialized successfully")
        except Exception as e:
            print(f"❌ FAISS setup failed: {e}")
            self.faiss_index = None
            self.faiss_ready = False
            self.collection = None # Ensure collection is None on failure

    def setup_biogpt(self):
        """Setup BioGPT model with optimizations for Colab"""
        print("🧠 Loading BioGPT-Large (this may take a few minutes on first run)...")

        model_name = "microsoft/BioGPT-Large"

        try:
            # Setup quantization config for memory efficiency
            if self.use_8bit:
                quantization_config = BitsAndBytesConfig(
                    load_in_8bit=True,
                    llm_int8_threshold=6.0,
                    llm_int8_has_fp16_weight=False,
                )
            else:
                quantization_config = None

            # Load tokenizer
            print("   Loading tokenizer...")
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)

            # Set padding token
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token

            # Load model
            print("   Loading BioGPT model...")
            start_time = time.time()

            self.model = AutoModelForCausalLM.from_pretrained(
                model_name,
                quantization_config=quantization_config,
                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
                device_map="auto" if self.device == "cuda" else None,
                trust_remote_code=True
            )

            # Move to device if not using device_map
            if self.device == "cuda" and quantization_config is None:
                self.model = self.model.to(self.device)

            load_time = time.time() - start_time
            print(f"✅ BioGPT loaded successfully! ({load_time:.1f} seconds)")

            # Test the model
            self.test_biogpt()

        except Exception as e:
            print(f"❌ BioGPT loading failed: {e}")
            print("💡 Falling back to smaller medical model...")
            self.setup_fallback_model()

    def setup_fallback_model(self):
        """Setup fallback model if BioGPT fails"""
        try:
            fallback_model = "microsoft/DialoGPT-medium"
            print(f"🔄 Loading fallback model: {fallback_model}")

            self.tokenizer = AutoTokenizer.from_pretrained(fallback_model)
            self.model = AutoModelForCausalLM.from_pretrained(fallback_model)

            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token

            if self.device == "cuda":
                self.model = self.model.to(self.device)

            print("✅ Fallback model loaded")

        except Exception as e:
            print(f"❌ All models failed: {e}")
            self.model = None
            self.tokenizer = None

    def test_biogpt(self):
        """Test BioGPT with a simple medical query"""
        print("🧪 Testing BioGPT...")
        try:
            test_prompt = "Fever in children can be caused by"
            inputs = self.tokenizer(test_prompt, return_tensors="pt")

            if self.device == "cuda": # Ensure inputs are on the correct device
                inputs = {k: v.to(self.device) for k, v in inputs.items()}

            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=20,
                    do_sample=True,
                    temperature=0.7,
                    pad_token_id=self.tokenizer.eos_token_id
                )

            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            print(f"✅ BioGPT test successful!")
            print(f"   Test response: {response}")

        except Exception as e:
            print(f"⚠️ BioGPT test failed: {e}")

    def load_medical_data(self, file_path: str):
        """Load and process medical data with progress tracking"""
        print(f"📖 Loading medical data from {file_path}...")

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                text = f.read()
            print(f"📄 File loaded: {len(text):,} characters")
        except FileNotFoundError:
            print(f"❌ File {file_path} not found!")
            return False

        # Create chunks optimized for medical content
        print("📝 Creating medical-optimized chunks...")
        chunks = self.create_medical_chunks(text)
        print(f"📋 Created {len(chunks)} medical chunks")

        self.knowledge_chunks = chunks

        # Generate embeddings with progress and add to FAISS index
        if self.use_embeddings and self.embedding_model and self.faiss_ready:
            return self.generate_embeddings_with_progress(chunks)

        print("✅ Medical data loaded (text search mode)")
        return True

    def create_medical_chunks(self, text: str, chunk_size: int = 400) -> List[Dict]:
        """Create medically-optimized text chunks"""
        chunks = []

        # Split by medical sections first
        medical_sections = self.split_by_medical_sections(text)

        chunk_id = 0
        for section in medical_sections:
            if len(section.split()) > chunk_size:
                # Split large sections by sentences
                sentences = re.split(r'[.!?]+', section)
                current_chunk = ""

                for sentence in sentences:
                    sentence = sentence.strip()
                    if not sentence:
                        continue

                    if len(current_chunk.split()) + len(sentence.split()) < chunk_size:
                        current_chunk += sentence + ". "
                    else:
                        if current_chunk.strip():
                            chunks.append({
                                'id': chunk_id,
                                'text': current_chunk.strip(),
                                'medical_focus': self.identify_medical_focus(current_chunk)
                            })
                            chunk_id += 1
                        current_chunk = sentence + ". "

                if current_chunk.strip():
                    chunks.append({
                        'id': chunk_id,
                        'text': current_chunk.strip(),
                        'medical_focus': self.identify_medical_focus(current_chunk)
                    })
                    chunk_id += 1
            else:
                chunks.append({
                    'id': chunk_id,
                    'text': section,
                    'medical_focus': self.identify_medical_focus(section)
                })
                chunk_id += 1

        return chunks

    def split_by_medical_sections(self, text: str) -> List[str]:
        """Split text by medical sections"""
        # Look for medical section headers
        section_patterns = [
            r'\n\s*(?:SYMPTOMS?|TREATMENT|DIAGNOSIS|CAUSES?|PREVENTION|MANAGEMENT).*?\n',
            r'\n\s*\d+\.\s+',  # Numbered sections
            r'\n\n+'  # Paragraph breaks
        ]

        sections = [text]
        for pattern in section_patterns:
            new_sections = []
            for section in sections:
                splits = re.split(pattern, section, flags=re.IGNORECASE)
                new_sections.extend([s.strip() for s in splits if len(s.strip()) > 100])
            sections = new_sections

        return sections

    def identify_medical_focus(self, text: str) -> str:
        """Identify the medical focus of a text chunk"""
        text_lower = text.lower()

        # Medical categories
        categories = {
            'pediatric_symptoms': ['fever', 'cough', 'rash', 'vomiting', 'diarrhea'],
            'treatments': ['treatment', 'therapy', 'medication', 'antibiotics'],
            'diagnosis': ['diagnosis', 'diagnostic', 'symptoms', 'signs'],
            'emergency': ['emergency', 'urgent', 'serious', 'hospital'],
            'prevention': ['prevention', 'vaccine', 'immunization', 'avoid']
        }

        for category, keywords in categories.items():
            if any(keyword in text_lower for keyword in keywords):
                return category

        return 'general_medical'

    def generate_embeddings_with_progress(self, chunks: List[Dict]) -> bool:
        """Generate embeddings with progress tracking and add to FAISS index"""
        print("🔮 Generating medical embeddings and adding to FAISS index...")

        if not self.embedding_model or not self.faiss_index:
            print("❌ Embedding model or FAISS index not available.")
            return False

        try:
            texts = [chunk['text'] for chunk in chunks]

            # Generate embeddings in batches with progress
            batch_size = 32
            all_embeddings = []

            for i in range(0, len(texts), batch_size):
                batch_texts = texts[i:i+batch_size]
                batch_embeddings = self.embedding_model.encode(batch_texts, show_progress_bar=False)
                all_embeddings.extend(batch_embeddings)

                # Show progress
                progress = min(i + batch_size, len(texts))
                print(f"   Progress: {progress}/{len(texts)} chunks processed", end='\r')

            print(f"\n   ✅ Generated embeddings for {len(texts)} chunks")

            # Add embeddings to FAISS index
            print("💾 Adding embeddings to FAISS index...")
            self.faiss_index.add(np.array(all_embeddings))

            print("✅ Medical embeddings added to FAISS index successfully!")
            return True

        except Exception as e:
            print(f"❌ Embedding generation or FAISS add failed: {e}")
            return False


    def retrieve_medical_context(self, query: str, n_results: int = 3) -> List[str]:
        """Retrieve relevant medical context using embeddings or keyword search"""
        if self.use_embeddings and self.embedding_model and self.faiss_ready:
            try:
                # Generate query embedding
                query_embedding = self.embedding_model.encode([query])

                # Search for similar content in FAISS index
                distances, indices = self.faiss_index.search(np.array(query_embedding), n_results)

                # Retrieve the corresponding chunks
                context_chunks = [self.knowledge_chunks[i]['text'] for i in indices[0] if i != -1]

                if context_chunks:
                    return context_chunks

            except Exception as e:
                print(f"⚠️ Embedding search failed: {e}")

        # Fallback to keyword search
        print("⚠️ Falling back to keyword search.")
        return self.keyword_search_medical(query, n_results)


    def keyword_search_medical(self, query: str, n_results: int) -> List[str]:
        """Medical-focused keyword search"""
        if not self.knowledge_chunks:
            return []

        query_words = set(query.lower().split())
        chunk_scores = []

        for chunk_info in self.knowledge_chunks:
            chunk_text = chunk_info['text']
            chunk_words = set(chunk_text.lower().split())

            # Calculate relevance score
            word_overlap = len(query_words.intersection(chunk_words))
            base_score = word_overlap / len(query_words) if query_words else 0

            # Boost medical content
            medical_boost = 0
            if chunk_info.get('medical_focus') in ['pediatric_symptoms', 'treatments', 'diagnosis']:
                medical_boost = 0.5

            final_score = base_score + medical_boost

            if final_score > 0:
                chunk_scores.append((final_score, chunk_text))

        # Return top matches
        chunk_scores.sort(reverse=True)
        return [chunk for _, chunk in chunk_scores[:n_results]]

    def generate_biogpt_response(self, context: str, query: str) -> str:
        """Generate medical response using BioGPT"""
        if not self.model or not self.tokenizer:
            return "Medical model not available. Please check the setup."

        try:
            # Create medical-focused prompt
            prompt = f"""Medical Context: {context[:800]}

Question: {query}

Medical Answer:"""

            # Tokenize input
            inputs = self.tokenizer(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=1024
            )

            # Move inputs to the correct device
            if self.device == "cuda":
                inputs = {k: v.to(self.device) for k, v in inputs.items()}

            # Generate response
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=150,
                    do_sample=True,
                    temperature=0.7,
                    top_p=0.9,
                    pad_token_id=self.tokenizer.eos_token_id,
                    repetition_penalty=1.1
                )

            # Decode response
            full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

            # Extract just the generated part
            if "Medical Answer:" in full_response:
                generated_response = full_response.split("Medical Answer:")[-1].strip()
            else:
                generated_response = full_response[len(prompt):].strip()

            # Clean up response
            cleaned_response = self.clean_medical_response(generated_response)

            return cleaned_response

        except Exception as e:
            print(f"⚠️ BioGPT generation failed: {e}")
            return self.fallback_response(context, query)

    def clean_medical_response(self, response: str) -> str:
        """Clean and format medical response"""
        # Remove incomplete sentences and limit length
        sentences = re.split(r'[.!?]+', response)
        clean_sentences = []

        for sentence in sentences:
            sentence = sentence.strip()
            if len(sentence) > 10 and not sentence.endswith(('and', 'or', 'but', 'however')):
                clean_sentences.append(sentence)
            if len(clean_sentences) >= 3:  # Limit to 3 sentences
                break

        if clean_sentences:
            cleaned = '. '.join(clean_sentences) + '.'
        else:
            cleaned = response[:200] + '...' if len(response) > 200 else response

        return cleaned

    def fallback_response(self, context: str, query: str) -> str:
        """Fallback response when BioGPT fails"""
        # Extract key sentences from context
        sentences = [s.strip() for s in context.split('.') if len(s.strip()) > 20]

        if sentences:
            response = sentences[0] + '.'
            if len(sentences) > 1:
                response += ' ' + sentences[1] + '.'
        else:
            response = context[:300] + '...'

        return response

    def handle_conversational_interactions(self, query: str) -> Optional[str]:
        """Handle comprehensive conversational interactions"""
        query_lower = query.lower().strip()

        # Use more specific patterns for greetings
        greeting_patterns = [
            r'^\s*(hello|hi|hey|hiya|howdy)\s*$',
            r'^\s*(good morning|good afternoon|good evening|good day)\s*$',
            r'^\s*(what\'s up|whats up|sup|yo)\s*$',
            r'^\s*(greetings|salutations)\s*$',
            r'^\s*(how are you|how are you doing|how\'s it going|hows it going)\s*$',
            r'^\s*(good to meet you|nice to meet you|pleased to meet you)\s*$'
        ]

        for pattern in greeting_patterns:
            if re.match(pattern, query_lower):
                responses = [
                    "👋 Hello! I'm BioGPT, your professional medical AI assistant specialized in pediatric medicine. I'm here to provide evidence-based medical information. What health concern can I help you with today?",
                    "🏥 Hi there! I'm a medical AI assistant powered by BioGPT, trained on medical literature. I can help answer questions about children's health and medical conditions. How can I assist you?",
                    "👋 Greetings! I'm your AI medical consultant, ready to help with pediatric health questions using the latest medical knowledge. What would you like to know about?"
                ]
                return np.random.choice(responses)

        # ===== THANKS & APPRECIATION =====
        thanks_patterns = [
            ['thank you', 'thanks', 'thx', 'ty', 'thank you so much', 'thanks a lot', 'much appreciated', 'really appreciate it', 'i appreciate it', 'grateful', 'that was helpful', 'very helpful', 'awesome', 'perfect', 'great', 'excellent', 'wonderful', 'that helped', 'exactly what i needed', 'very informative', 'good information']
        ]

        for pattern_group in thanks_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                responses = [
                    "🙏 You're very welcome! I'm glad I could provide helpful medical information. Remember, this is educational guidance - always consult your healthcare provider for personalized medical advice. Feel free to ask more questions!",
                    "😊 Happy to help! Providing accurate medical information is what I'm here for. If you have any other pediatric health questions, don't hesitate to ask.",
                    "🤗 You're most welcome! I'm pleased the medical information was useful. Please remember to consult with healthcare professionals for any medical decisions. What else can I help you with?"
                ]
                return np.random.choice(responses)

        # ===== GOODBYES =====
        goodbye_patterns = [
            ['bye', 'goodbye', 'farewell', 'see you', 'later', 'see ya', 'catch you later', 'talk to you later', 'ttyl', 'have a good day', 'have a great day', 'take care', 'until next time', 'i need to go', 'that\'s all for now', 'no more questions']
        ]

        for pattern_group in goodbye_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                responses = [
                    "👋 Goodbye! Take excellent care of yourself and your little ones. Remember, I'm here whenever you need reliable pediatric medical information. Stay healthy! 🏥",
                    "🌟 Farewell! Wishing you and your family good health. Don't hesitate to return if you have more medical questions. Take care!",
                    "👋 See you later! Hope the medical information was helpful. Remember to always consult healthcare professionals for medical decisions. Stay well!"
                ]
                return np.random.choice(responses)

        # ===== ABOUT/HELP QUESTIONS =====
        about_patterns = [
            ['what are you', 'who are you', 'tell me about yourself', 'what do you do', 'what can you help with', 'what can you do', 'how can you help', 'what are your capabilities', 'help', 'help me', 'i need help', 'can you help', 'how do i use this', 'how does this work', 'what should i ask']
        ]

        for pattern_group in about_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                return """🤖 **About BioGPT Medical Assistant**

I'm an AI medical assistant powered by BioGPT-Large, a specialized medical AI model trained on extensive medical literature. Here's what I can help you with:

🩺 **Medical Specialties:**
• Pediatric medicine and children's health
• Symptom explanation and medical conditions
• Treatment options and medical procedures
• When to seek medical care
• Prevention and wellness guidance

🎯 **How to Use Me:**
• Ask specific medical questions: "What causes fever in children?"
• Describe symptoms: "My child has a persistent cough"
• Seek guidance: "When should I call the doctor?"
• Get information: "How do I treat dehydration?"

⚠️ **Important Reminder:**
I provide educational medical information based on medical literature, but I'm not a substitute for professional medical advice. Always consult qualified healthcare providers for:
• Medical emergencies
• Diagnosis and treatment decisions
• Personalized medical advice
• Medication guidance

💡 **Tip:** Be specific in your questions for the most helpful responses!

What pediatric health topic would you like to explore?"""

        # ===== SMALL TALK & PERSONAL QUESTIONS =====
        personal_patterns = [
            ['how are you feeling', 'are you okay', 'how\'s your day', 'are you smart', 'are you intelligent', 'do you know everything', 'are you human', 'are you real', 'are you a robot', 'are you ai', 'you\'re smart', 'you\'re helpful', 'good job', 'well done', 'impressive']
        ]

        for pattern_group in personal_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                responses = [
                    "🤖 I'm an AI medical assistant, so I don't have feelings, but I'm functioning well and ready to help with medical questions! My purpose is to provide reliable pediatric health information. What can I help you with?",
                    "😊 Thank you for asking! As an AI, I'm always ready to assist with medical information. I'm designed to help with pediatric health questions using evidence-based medical knowledge. How can I help you today?",
                    "🎯 I'm doing what I do best - providing medical information! I'm an AI trained on medical literature to help with pediatric health questions. What medical topic interests you?"
                ]
                return np.random.choice(responses)

        # ===== CONFUSED/UNCLEAR INPUT =====
        confusion_patterns = [
            ['i don\'t know', 'not sure', 'confused', 'unclear', 'help me understand', 'what do you mean', 'i don\'t understand', 'can you explain', 'huh', 'i\'m lost', 'i\'m confused', 'this is confusing']
        ]

        for pattern_group in confusion_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                return """🤔 **I understand it can be confusing!** Let me help you get started.

💡 **Try asking questions like:**

🩺 **Symptoms:**
• "What causes [symptom] in children?"
• "My child has [symptom], what should I do?"

💊 **Treatments:**
• "How do I treat [condition] in children?"
• "What are treatment options for [condition]?"

🚨 **Urgency:**
• "When should I call the doctor about [symptom]?"
• "Is [symptom] serious in children?"

🛡️ **Prevention:**
• "How can I prevent [condition]?"
• "What are the warning signs of [condition]?"

**What specific aspect of your child's health would you like to understand better?**"""

        # ===== APOLOGIES & POLITENESS =====
        polite_patterns = [
            ['sorry', 'excuse me', 'pardon me', 'my apologies', 'please help', 'could you please', 'would you mind', 'if you don\'t mind', 'sorry to bother you']
        ]

        for pattern_group in polite_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                return "😊 No need to apologize! I'm here to help with medical questions. Please feel free to ask anything about pediatric health - that's exactly what I'm designed for. What can I help you with?"

        # ===== TESTING & VERIFICATION =====
        test_patterns = [
            ['test', 'testing', 'hello world', 'can you hear me', 'are you working', 'do you work', 'are you there', 'are you online', 'check', 'verify', 'ping']
        ]

        for pattern_group in test_patterns:
            if any(keyword in query_lower for keyword in pattern_group):
                return "✅ **System Check:** I'm working perfectly and ready to assist! BioGPT medical AI is online and functioning optimally. Ready to help with pediatric medical questions. What would you like to know?"

        # Return None if no conversational pattern matches
        return None

    def chat(self, query: str) -> str:
        """Main chat function with BioGPT and comprehensive conversational handling"""
        if not query.strip():
            return "Hello! I'm BioGPT, your professional medical AI assistant. How can I help you with pediatric medical questions today?"

        # Handle comprehensive conversational interactions first
        conversational_response = self.handle_conversational_interactions(query)
        if conversational_response:
            # Add to conversation history
            self.conversation_history.append({
                'query': query,
                'response': conversational_response,
                'timestamp': datetime.now().isoformat(),
                'type': 'conversational'
            })
            return conversational_response

        if not self.knowledge_chunks:
            return "Please load medical data first to access the medical knowledge base."

        print(f"🔍 Processing medical query: {query}")

        # Retrieve relevant medical context using FAISS or keyword search
        start_time = time.time()
        context = self.retrieve_medical_context(query)
        retrieval_time = time.time() - start_time

        if not context:
            return "I don't have specific information about this topic in my medical database. Please consult with a healthcare professional for personalized medical advice."

        print(f"   📚 Context retrieved ({retrieval_time:.2f}s)")

        # Generate response with BioGPT
        start_time = time.time()
        main_context = '\n\n'.join(context)
        response = self.generate_biogpt_response(main_context, query)
        generation_time = time.time() - start_time

        print(f"   🧠 Response generated ({generation_time:.2f}s)")

        # Format final response
        final_response = f"🩺 **Medical Information:** {response}\n\n⚠️ **Important:** This information is for educational purposes only. Always consult with qualified healthcare professionals for medical diagnosis, treatment, and personalized advice."

        # Add to conversation history
        self.conversation_history.append({
            'query': query,
            'response': final_response,
            'timestamp': datetime.now().isoformat(),
            'retrieval_time': retrieval_time,
            'generation_time': generation_time,
            'type': 'medical'
        })

        return final_response

    def get_conversation_summary(self) -> Dict:
        """Get conversation statistics"""
        if not self.conversation_history:
            return {"message": "No conversations yet"}

        # Filter medical conversations for performance stats
        medical_conversations = [h for h in self.conversation_history if h.get('type') == 'medical']

        if not medical_conversations:
            return {
                "total_conversations": len(self.conversation_history),
                "medical_conversations": 0,
                "conversational_interactions": len(self.conversation_history),
                "model_info": "BioGPT-Large" if "BioGPT" in str(self.model) else "Fallback Model",
                "vector_search": "FAISS CPU" if self.faiss_ready else "Keyword Search",
                "device": self.device
            }

        avg_retrieval_time = sum(h.get('retrieval_time', 0) for h in medical_conversations) / len(medical_conversations)
        avg_generation_time = sum(h.get('generation_time', 0) for h in medical_conversations) / len(medical_conversations)

        return {
            "total_conversations": len(self.conversation_history),
            "medical_conversations": len(medical_conversations),
            "conversational_interactions": len(self.conversation_history) - len(medical_conversations),
            "avg_retrieval_time": f"{avg_retrieval_time:.2f}s",
            "avg_generation_time": f"{avg_generation_time:.2f}s",
            "model_info": "BioGPT-Large" if "BioGPT" in str(self.model) else "Fallback Model",
            "vector_search": "FAISS CPU" if self.faiss_ready else "Keyword Search",
            "device": self.device,
            "quantization": "8-bit" if self.use_8bit else "16-bit/32-bit"
        }

# Create and Test BioGPT Chatbot

def create_biogpt_chatbot():
    """Create and initialize the BioGPT chatbot"""
    print("🚀 Creating Professional BioGPT Medical Chatbot")
    print("=" * 60)

    # Create chatbot
    chatbot = ColabBioGPTChatbot(use_gpu=True, use_8bit=True)

    return chatbot

def test_biogpt_chatbot(chatbot, test_file='Pediatric_cleaned.txt'):
    """Test the BioGPT chatbot"""
    print("\n📚 Loading medical data...")
    success = chatbot.load_medical_data(test_file)

    if not success:
        print("❌ Failed to load medical data. Please check the file.")
        return None

    print("\n🧪 Testing BioGPT Medical Chatbot:")
    print("=" * 50)

    # Test queries
    test_queries = [
        "What causes fever in children?",
        "How should I treat my child's cough?",
        "When should I be concerned about my baby's breathing?",
        "What are the signs of dehydration in infants?"
    ]

    for i, query in enumerate(test_queries, 1):
        print(f"\n{i}️⃣ Testing: {query}")
        print("-" * 40)

        response = chatbot.chat(query)
        print(f"🤖 BioGPT Response:\n{response}")
        print("=" * 50)

    # Show conversation summary
    summary = chatbot.get_conversation_summary()
    print("\n📊 Performance Summary:")
    for key, value in summary.items():
        print(f"   {key}: {value}")

    return chatbot

# Interactive Chat Interface

def interactive_biogpt_chat(chatbot):
    """Interactive chat with BioGPT"""
    print("\n💬 Interactive BioGPT Medical Chat")
    print("=" * 50)
    print("You're now chatting with BioGPT, a professional medical AI!")
    print("Type 'quit' to exit, 'summary' to see stats")
    print("-" * 50)

    while True:
        user_input = input("\n👤 You: ").strip()

        if user_input.lower() in ['quit', 'exit', 'bye']:
            print("\n👋 Thank you for using BioGPT Medical Assistant!")
            # Show final summary
            summary = chatbot.get_conversation_summary()
            print("\n📊 Final Session Summary:")
            for key, value in summary.items():
                print(f"   {key}: {value}")
            break

        elif user_input.lower() == 'summary':
            summary = chatbot.get_conversation_summary()
            print("\n📊 Current Session Summary:")
            for key, value in summary.items():
                print(f"   {key}: {value}")
            continue

        elif not user_input:
            continue

        print(f"\n🤖 BioGPT: ", end="")
        response = chatbot.chat(user_input)
        print(response)

# Main Execution

# Create the BioGPT chatbot
chatbot = create_biogpt_chatbot()

print("\n" + "="*60)
print("🎯 NEXT STEPS:")
print("1. Upload your medical data file by running: upload_medical_data()")
print("2. Test the chatbot: test_biogpt_chatbot(chatbot)")
print("3. Start interactive chat: interactive_biogpt_chat(chatbot)")
print("="*60)