maining.py

import os
import re
import json
import pickle
import random
import time
import numpy as np
from collections import defaultdict, Counter
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.decomposition import TruncatedSVD
import sqlite3
from datetime import datetime

class AdvancedTrainingAI:
    def __init__(self, data_folder="dataset", model_name="advanced_ai_model"):
        """
        Initialize Advanced Training AI with learning capabilities
        """
        self.data_folder = data_folder
        self.db_path = f"{model_name}_data.db"
        self.model_path = f"{model_name}.pkl"
        self.embedding_dim = 100  # Dimension for semantic embeddings
        
        # Model components
        self.vectorizer = None
        self.svd_model = None
        self.document_embeddings = None
        self.documents = []
        self.word_vectors = {}
        self.word_clusters = {}
        self.response_templates = {}
        self.semantic_index = {}
        
        # Initialize database
        self.init_database()
        
        print("🧠 Initializing Advanced Training AI...")
        print("📁 Loading data from dataset folder...")
        
        # Load and process data
        self.load_data()
        self.train_model()
    
    def init_database(self):
        """
        Initialize SQLite database for data storage
        """
        self.conn = sqlite3.connect(self.db_path)
        self.cursor = self.conn.cursor()
        
        # Create tables
        self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS documents (
                id INTEGER PRIMARY KEY,
                content TEXT,
                source_file TEXT,
                embedding BLOB,
                processed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS word_vectors (
                word TEXT PRIMARY KEY,
                vector BLOB,
                frequency INTEGER
            )
        ''')
        
        self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS response_templates (
                id INTEGER PRIMARY KEY,
                template_type TEXT,
                template TEXT,
                context_words TEXT
            )
        ''')
        
        self.conn.commit()
    
    def load_data(self):
        """
        Load data from text files and process for training
        """
        txt_files = self.find_txt_files()
        
        if not txt_files:
            print("❌ No .txt files found in dataset folder!")
            return
        
        print(f"📊 Found {len(txt_files)} .txt files")
        
        total_documents = 0
        batch_size = 1000
        
        for file_path in txt_files:
            print(f"📖 Processing: {os.path.basename(file_path)}")
            
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                
                # Split into chunks for processing
                chunks = self.split_into_chunks(content, file_path)
                total_documents += len(chunks)
                
                # Process in batches
                for i in range(0, len(chunks), batch_size):
                    batch = chunks[i:i+batch_size]
                    self.process_document_batch(batch)
                    
                    if i % (batch_size * 5) == 0 and i > 0:
                        print(f"   ✓ Processed {i:,} chunks...")
                
                print(f"   ✅ Extracted {len(chunks):,} document chunks")
                
            except Exception as e:
                print(f"   ❌ Error: {e}")
        
        print(f"🎉 Total processed: {total_documents:,} document chunks")
        
        # Extract response templates
        self.extract_response_templates()
    
    def find_txt_files(self):
        """
        Find all .txt files in dataset folder
        """
        txt_files = []
        
        if os.path.exists(self.data_folder):
            for root, dirs, files in os.walk(self.data_folder):
                for file in files:
                    if file.endswith('.txt'):
                        txt_files.append(os.path.join(root, file))
        else:
            print(f"⚠️  Dataset folder '{self.data_folder}' not found!")
        
        return txt_files
    
    def split_into_chunks(self, content, source_file, chunk_size=500):
        """
        Split content into meaningful chunks for training
        """
        # Clean content
        content = self.clean_text(content)
        
        # Split by paragraphs first
        paragraphs = re.split(r'\n\s*\n', content)
        
        chunks = []
        current_chunk = ""
        
        for para in paragraphs:
            para = para.strip()
            if not para:
                continue
                
            # If paragraph is very long, split by sentences
            if len(para) > chunk_size * 2:
                sentences = re.split(r'(?<=[.!?])\s+', para)
                for sentence in sentences:
                    if len(current_chunk) + len(sentence) < chunk_size:
                        current_chunk += sentence + " "
                    else:
                        if current_chunk:
                            chunks.append({
                                'content': current_chunk.strip(),
                                'source': source_file
                            })
                        current_chunk = sentence + " "
            else:
                # Add paragraph as is or start new chunk
                if len(current_chunk) + len(para) < chunk_size:
                    current_chunk += para + " "
                else:
                    if current_chunk:
                        chunks.append({
                            'content': current_chunk.strip(),
                            'source': source_file
                        })
                    current_chunk = para + " "
        
        # Add the last chunk
        if current_chunk:
            chunks.append({
                'content': current_chunk.strip(),
                'source': source_file
            })
        
        return chunks
    
    def clean_text(self, text):
        """
        Clean and normalize text
        """
        # Remove special characters but keep punctuation
        text = re.sub(r'[^\w\s\.\?\!,\'\-]', ' ', text)
        # Remove extra whitespace
        text = re.sub(r'\s+', ' ', text)
        return text.strip()
    
    def process_document_batch(self, documents):
        """
        Process a batch of documents
        """
        for doc in documents:
            self.documents.append(doc)
            
            # Extract words for word vectors
            words = re.findall(r'\b[a-zA-Z]{3,}\b', doc['content'].lower())
            for word in words:
                if word not in self.word_vectors:
                    self.word_vectors[word] = 1
                else:
                    self.word_vectors[word] += 1
    
    def extract_response_templates(self):
        """
        Extract response templates from documents
        """
        print("🔍 Extracting response templates...")
        
        # Define template types and their patterns
        template_patterns = {
            'definition': [
                r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+is\s+(?:defined\s+as\s+)?([^.]+)',
                r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+refers\s+to\s+([^.]+)',
                r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+means\s+([^.]+)'
            ],
            'explanation': [
                r'(?:To|In order to)\s+([^,]+),\s+([^.]+)',
                r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+works\s+by\s+([^.]+)',
                r'The\s+(?:process|method)\s+of\s+([^,]+),\s+([^.]+)'
            ],
            'comparison': [
                r'(?:Compared|In comparison)\s+to\s+([^,]+),\s+([^.]+)',
                r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+(?:differs|is different)\s+from\s+([^.]+)',
                r'The\s+(?:difference|distinction)\s+between\s+([^,]+),\s+([^.]+)'
            ],
            'list': [
                r'(?:The|Some)\s+(?:key|main|important)\s+(?:aspects|features|components)\s+of\s+([^:]+):\s*([^.]+)',
                r'(?:Several|Many|Various)\s+(?:types|kinds|forms)\s+of\s+([^:]+):\s*([^.]+)',
                r'(?:Steps|Stages|Phases)\s+(?:in|of|for)\s+([^:]+):\s*([^.]+)'
            ]
        }
        
        templates_found = defaultdict(list)
        
        # Extract templates from documents
        for doc in self.documents:
            content = doc['content']
            
            for template_type, patterns in template_patterns.items():
                for pattern in patterns:
                    matches = re.findall(pattern, content)
                    for match in matches:
                        if isinstance(match, tuple) and len(match) >= 2:
                            subject, description = match[0], match[1]
                            
                            # Create template by replacing specific parts with placeholders
                            if template_type == 'definition':
                                template = "{subject} is {description}."
                            elif template_type == 'explanation':
                                template = "To {subject}, {description}."
                            elif template_type == 'comparison':
                                template = "Compared to {subject}, {description}."
                            elif template_type == 'list':
                                template = "The key aspects of {subject}: {description}."
                            
                            # Extract context words
                            context_words = re.findall(r'\b[a-zA-Z]{3,}\b', subject.lower() + " " + description.lower())
                            
                            templates_found[template_type].append({
                                'template': template,
                                'subject': subject,
                                'description': description,
                                'context_words': context_words
                            })
        
        # Save templates to database
        for template_type, templates in templates_found.items():
            for template_data in templates:
                self.cursor.execute(
                    'INSERT INTO response_templates (template_type, template, context_words) VALUES (?, ?, ?)',
                    (template_type, template_data['template'], json.dumps(template_data['context_words']))
                )
        
        self.conn.commit()
        
        # Store in memory
        self.response_templates = templates_found
        
        print(f"✅ Extracted templates: {sum(len(t) for t in templates_found.values()):,}")
        for t_type, templates in templates_found.items():
            print(f"   • {t_type}: {len(templates):,}")
    
    def train_model(self):
        """
        Train the AI model
        """
        print("🧠 Training AI model...")
        
        if not self.documents:
            print("❌ No documents to train on!")
            return
        
        # Step 1: Create TF-IDF vectorizer
        print("   ⚙️ Creating TF-IDF vectors...")
        self.vectorizer = TfidfVectorizer(
            max_features=10000,
            ngram_range=(1, 2),
            stop_words='english'
        )
        
        # Extract document content
        doc_contents = [doc['content'] for doc in self.documents]
        
        # Fit and transform documents
        tfidf_matrix = self.vectorizer.fit_transform(doc_contents)
        
        # Step 2: Apply dimensionality reduction for semantic understanding
        print("   ⚙️ Creating semantic embeddings...")
        self.svd_model = TruncatedSVD(n_components=self.embedding_dim)
        self.document_embeddings = self.svd_model.fit_transform(tfidf_matrix)
        
        # Step 3: Create semantic clusters
        print("   ⚙️ Creating semantic clusters...")
        self.create_semantic_clusters()
        
        # Step 4: Build semantic index
        print("   ⚙️ Building semantic index...")
        self.build_semantic_index()
        
        # Step 5: Save model
        self.save_model()
        
        print("✅ Model training complete!")
    
    def create_semantic_clusters(self):
        """
        Create semantic clusters of words
        """
        # Extract vocabulary
        vocabulary = self.vectorizer.get_feature_names_out()
        
        # Create word vectors using SVD components
        components = self.svd_model.components_
        word_vectors = {}
        
        for i, word in enumerate(vocabulary):
            # Get the word's values across all SVD components
            word_vec = components[:, i]
            word_vectors[word] = word_vec
        
        # Simple clustering by dominant component
        word_clusters = defaultdict(list)
        for word, vec in word_vectors.items():
            # Find the component where this word has highest value
            dominant_component = np.argmax(np.abs(vec))
            word_clusters[dominant_component].append((word, vec[dominant_component]))
        
        # Sort words in each cluster by importance
        for component in word_clusters:
            word_clusters[component] = sorted(word_clusters[component], key=lambda x: abs(x[1]), reverse=True)
        
        self.word_vectors = word_vectors
        self.word_clusters = dict(word_clusters)
    
    def build_semantic_index(self):
        """
        Build semantic index for faster retrieval
        """
        semantic_index = defaultdict(list)
        
        # For each document, find its dominant topics
        for i, doc_embedding in enumerate(self.document_embeddings):
            # Get top 3 components for this document
            top_components = np.argsort(np.abs(doc_embedding))[-3:]
            
            for component in top_components:
                semantic_index[component].append((i, doc_embedding[component]))
        
        # Sort documents in each component by relevance
        for component in semantic_index:
            semantic_index[component] = sorted(semantic_index[component], key=lambda x: abs(x[1]), reverse=True)
        
        self.semantic_index = dict(semantic_index)
    
    def save_model(self):
        """
        Save trained model to disk
        """
        model_data = {
            'vectorizer': self.vectorizer,
            'svd_model': self.svd_model,
            'document_embeddings': self.document_embeddings,
            'word_clusters': self.word_clusters,
            'semantic_index': self.semantic_index,
            'documents': self.documents
        }
        
        with open(self.model_path, 'wb') as f:
            pickle.dump(model_data, f)
        
        print(f"💾 Model saved to {self.model_path}")
    
    def load_model(self):
        """
        Load trained model from disk
        """
        if os.path.exists(self.model_path):
            with open(self.model_path, 'rb') as f:
                model_data = pickle.load(f)
                
                self.vectorizer = model_data['vectorizer']
                self.svd_model = model_data['svd_model']
                self.document_embeddings = model_data['document_embeddings']
                self.word_clusters = model_data['word_clusters']
                self.semantic_index = model_data['semantic_index']
                self.documents = model_data['documents']
            
            return True
        return False
    
    def answer_question(self, question):
        """
        Generate an answer to a question using the trained model
        """
        if not self.load_model():
            return {
                'answer': "Model not trained yet. Please run training first.",
                'confidence': 0,
                'sources': []
            }
        
        # Step 1: Analyze the question
        question_type, question_focus = self.analyze_question(question)
        
        # Step 2: Find relevant documents
        relevant_docs = self.find_relevant_documents(question, top_k=5)
        
        if not relevant_docs:
            return {
                'answer': "I don't have enough information to answer this question accurately.",
                'confidence': 0,
                'sources': []
            }
        
        # Step 3: Generate answer based on question type and relevant docs
        answer = self.generate_answer(question, question_type, question_focus, relevant_docs)
        
        # Step 4: Calculate confidence
        confidence = self.calculate_confidence(question, answer, relevant_docs)
        
        # Step 5: Format sources
        sources = [doc['source'] for doc in relevant_docs]
        
        return {
            'answer': answer,
            'confidence': confidence,
            'sources': sources
        }
    
    def analyze_question(self, question):
        """
        Analyze question to determine type and focus
        """
        question = question.lower()
        
        # Determine question type
        if re.search(r'\b(what|who|which)\s+(is|are|was|were)\b', question):
            question_type = 'definition'
        elif re.search(r'\b(how)\s+(does|do|can|could|would|should)\b', question):
            question_type = 'explanation'
        elif re.search(r'\b(what|which)\s+(are|is)\s+the\s+(differences|difference|distinctions|distinction)\b', question):
            question_type = 'comparison'
        elif re.search(r'\b(what|which)\s+(are|is)\s+the\s+(steps|stages|phases|components|features|aspects)\b', question):
            question_type = 'list'
        else:
            question_type = 'general'
        
        # Extract focus words
        focus_words = []
        
        # Remove question words and common verbs
        cleaned_question = re.sub(r'\b(what|who|how|why|when|where|which|is|are|do|does|can|could|would|should|the|a|an)\b', '', question)
        
        # Extract nouns and important words (simple approach)
        focus_words = [word for word in re.findall(r'\b[a-z]{3,}\b', cleaned_question) 
                      if word not in ['this', 'that', 'these', 'those', 'there', 'here', 'from', 'with', 'about']]
        
        return question_type, focus_words
    
    def find_relevant_documents(self, question, top_k=5):
        """
        Find documents relevant to the question using semantic search
        """
        # Convert question to TF-IDF vector
        question_vector = self.vectorizer.transform([question])
        
        # Convert to semantic embedding
        question_embedding = self.svd_model.transform(question_vector)[0]
        
        # Find dominant components in question
        top_components = np.argsort(np.abs(question_embedding))[-3:]
        
        # Collect candidate documents from semantic index
        candidate_docs = []
        for component in top_components:
            if component in self.semantic_index:
                # Get top documents for this component
                component_docs = self.semantic_index[component][:20]  # Get top 20 per component
                candidate_docs.extend([doc_idx for doc_idx, _ in component_docs])
        
        # Remove duplicates
        candidate_docs = list(set(candidate_docs))
        
        if not candidate_docs:
            # Fallback to direct similarity search
            similarities = cosine_similarity(question_vector, 
                                           self.vectorizer.transform([doc['content'] for doc in self.documents]))
            candidate_docs = np.argsort(similarities[0])[-20:].tolist()
        
        # Re-rank candidates by direct similarity
        candidate_contents = [self.documents[idx]['content'] for idx in candidate_docs]
        candidate_vectors = self.vectorizer.transform(candidate_contents)
        question_vector = self.vectorizer.transform([question])
        
        similarities = cosine_similarity(question_vector, candidate_vectors)[0]
        
        # Sort by similarity
        ranked_indices = np.argsort(similarities)[::-1]
        
        # Get top_k documents
        relevant_docs = []
        for i in range(min(top_k, len(ranked_indices))):
            doc_idx = candidate_docs[ranked_indices[i]]
            relevant_docs.append({
                'content': self.documents[doc_idx]['content'],
                'source': self.documents[doc_idx]['source'],
                'similarity': similarities[ranked_indices[i]]
            })
        
        return relevant_docs
    
    def generate_answer(self, question, question_type, focus_words, relevant_docs):
        """
        Generate an answer based on question type and relevant documents
        """
        # Extract relevant content
        relevant_content = [doc['content'] for doc in relevant_docs]
        
        # Different strategies based on question type
        if question_type == 'definition':
            return self.generate_definition_answer(question, focus_words, relevant_content)
        elif question_type == 'explanation':
            return self.generate_explanation_answer(question, focus_words, relevant_content)
        elif question_type == 'comparison':
            return self.generate_comparison_answer(question, focus_words, relevant_content)
        elif question_type == 'list':
            return self.generate_list_answer(question, focus_words, relevant_content)
        else:
            return self.generate_general_answer(question, focus_words, relevant_content)
    
    def generate_definition_answer(self, question, focus_words, relevant_content):
        """
        Generate definition-type answer
        """
        # Look for definition patterns in relevant content
        definition_patterns = [
            r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+is\s+(?:defined\s+as\s+)?([^.]+)',
            r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+refers\s+to\s+([^.]+)',
            r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+means\s+([^.]+)'
        ]
        
        definitions = []
        for content in relevant_content:
            for pattern in definition_patterns:
                matches = re.findall(pattern, content)
                for match in matches:
                    if isinstance(match, tuple) and len(match) >= 2:
                        subject, description = match[0], match[1]
                        # Check if subject is related to focus words
                        if any(focus_word in subject.lower() for focus_word in focus_words):
                            definitions.append((subject, description))
        
        if definitions:
            # Use the most relevant definition
            subject, description = definitions[0]
            answer = f"{subject} is {description}."
            
            # Add additional information if available
            if len(definitions) > 1:
                answer += f" Additionally, it can also be described as {definitions[1][1]}."
            
            # Add context from relevant content
            for content in relevant_content:
                # Find sentences containing the subject
                sentences = re.split(r'(?<=[.!?])\s+', content)
                for sentence in sentences:
                    if subject.lower() in sentence.lower() and sentence not in answer:
                        if len(answer) + len(sentence) < 500:  # Limit answer length
                            answer += f" {sentence}"
                            break
            
            return answer
        
        # Fallback: synthesize from relevant content
        return self.synthesize_answer(question, focus_words, relevant_content)
    
    def generate_explanation_answer(self, question, focus_words, relevant_content):
        """
        Generate explanation-type answer
        """
        # Extract process/method descriptions
        explanation_patterns = [
            r'(?:To|In order to)\s+([^,]+),\s+([^.]+)',
            r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+works\s+by\s+([^.]+)',
            r'The\s+(?:process|method)\s+of\s+([^,]+),\s+([^.]+)'
        ]
        
        explanations = []
        for content in relevant_content:
            for pattern in explanation_patterns:
                matches = re.findall(pattern, content)
                for match in matches:
                    if isinstance(match, tuple) and len(match) >= 2:
                        subject, description = match[0], match[1]
                        # Check if subject is related to focus words
                        if any(focus_word in subject.lower() for focus_word in focus_words):
                            explanations.append((subject, description))
        
        if explanations:
            # Construct explanation answer
            subject, description = explanations[0]
            answer = f"To {subject}, {description}."
            
            # Add steps or additional details if available
            steps = []
            for content in relevant_content:
                step_matches = re.findall(r'(?:First|Second|Third|Next|Finally|Then|Step \d+)[,:]?\s+([^.]+)', content)
                steps.extend(step_matches)
            
            if steps:
                answer += " The process involves these steps:"
                for i, step in enumerate(steps[:5]):  # Limit to 5 steps
                    answer += f"\n{i+1}. {step}"
            
            return answer
        
        # Fallback: synthesize from relevant content
        return self.synthesize_answer(question, focus_words, relevant_content)
    
    def generate_comparison_answer(self, question, focus_words, relevant_content):
        """
        Generate comparison-type answer
        """
        # Look for comparison patterns
        comparison_patterns = [
            r'(?:Compared|In comparison)\s+to\s+([^,]+),\s+([^.]+)',
            r'([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\s+(?:differs|is different)\s+from\s+([^.]+)',
            r'The\s+(?:difference|distinction)\s+between\s+([^,]+),\s+([^.]+)'
        ]
        
        comparisons = []
        for content in relevant_content:
            for pattern in comparison_patterns:
                matches = re.findall(pattern, content)
                for match in matches:
                    if isinstance(match, tuple) and len(match) >= 2:
                        subject, description = match[0], match[1]
                        comparisons.append((subject, description))
        
        if comparisons:
            # Construct comparison answer
            answer = "The key differences are: "
            
            for i, (subject, description) in enumerate(comparisons[:3]):  # Limit to 3 comparisons
                if i > 0:
                    answer += " Additionally, "
                answer += f"compared to {subject}, {description}"
            
            return answer
        
        # Fallback: synthesize from relevant content
        return self.synthesize_answer(question, focus_words, relevant_content)
    
    def generate_list_answer(self, question, focus_words, relevant_content):
        """
        Generate list-type answer
        """
        # Look for list patterns
        list_items = []
        
        # First try to find explicit lists
        for content in relevant_content:
            # Look for numbered lists
            numbered_items = re.findall(r'(?:^|\n)(?:\d+\.|\*|\-)\s+([^\n]+)', content)
            if numbered_items:
                list_items.extend(numbered_items[:5])  # Limit to 5 items
            
            # Look for "key aspects", "features", etc.
            list_intro = re.search(r'(?:key|main|important)\s+(?:aspects|features|components|elements)\s+(?:of|include|are)[:\s]+([^.]+)', content)
            if list_intro:
                items = list_intro.group(1).split(',')
                list_items.extend([item.strip() for item in items])
        
        if list_items:
            # Construct list answer
            topic = ' '.join(focus_words).title()
            answer = f"The key points about {topic} include:"
            
            for i, item in enumerate(list_items[:5]):  # Limit to 5 items
                answer += f"\n{i+1}. {item}"
            
            return answer
        
        # Fallback: extract key sentences
        key_sentences = []
        for content in relevant_content:
            sentences = re.split(r'(?<=[.!?])\s+', content)
            for sentence in sentences:
                # Check if sentence contains focus words
                if any(focus_word in sentence.lower() for focus_word in focus_words):
                    key_sentences.append(sentence)
        
        if key_sentences:
            topic = ' '.join(focus_words).title()
            answer = f"Key information about {topic}:"
            
            for i, sentence in enumerate(key_sentences[:5]):  # Limit to 5 sentences
                answer += f"\n{i+1}. {sentence}"
            
            return answer
        
        # Final fallback
        return self.synthesize_answer(question, focus_words, relevant_content)
    
    def generate_general_answer(self, question, focus_words, relevant_content):
        """
        Generate general answer for other question types
        """
        return self.synthesize_answer(question, focus_words, relevant_content)
    
    def synthesize_answer(self, question, focus_words, relevant_content):
        """
        Synthesize an answer from relevant content
        """
        # Extract relevant sentences
        relevant_sentences = []
        for content in relevant_content:
            sentences = re.split(r'(?<=[.!?])\s+', content)
            for sentence in sentences:
                # Check if sentence contains focus words
                if any(focus_word in sentence.lower() for focus_word in focus_words):
                    relevant_sentences.append(sentence)
        
        if not relevant_sentences:
            # If no relevant sentences found, use first few sentences from relevant content
            for content in relevant_content:
                sentences = re.split(r'(?<=[.!?])\s+', content)
                relevant_sentences.extend(sentences[:2])  # Take first 2 sentences from each content
        
        # Limit number of sentences
        if len(relevant_sentences) > 5:
            relevant_sentences = relevant_sentences[:5]
        
        # Combine sentences into coherent answer
        answer = " ".join(relevant_sentences)
        
        # Add introduction based on question
        if question.lower().startswith('what'):
            topic = ' '.join(focus_words).title()
            answer = f"Regarding {topic}: {answer}"
        elif question.lower().startswith('how'):
            answer = f"Here's how it works: {answer}"
        elif question.lower().startswith('why'):
            answer = f"The reason is: {answer}"
        
        return answer
    
    def calculate_confidence(self, question, answer, relevant_docs):
        """
        Calculate confidence score for the answer
        """
        if not relevant_docs:
            return 0.0
        
        # Base confidence on document similarity
        base_confidence = max([doc['similarity'] for doc in relevant_docs])
        
        # Adjust based on answer quality
        quality_score = 0.0
        
        # Length factor (too short or too long answers may be less reliable)
        answer_length = len(answer)
        if 100 <= answer_length <= 500:
            quality_score += 0.2
        elif 50 <= answer_length < 100:
            quality_score += 0.1
        
        # Focus words coverage
        question_focus = self.analyze_question(question)[1]
        focus_coverage = sum(1 for word in question_focus if word.lower() in answer.lower()) / max(1, len(question_focus))
        quality_score += focus_coverage * 0.3
        
        # Answer structure (sentences, paragraphs)
        sentences = len(re.split(r'(?<=[.!?])\s+', answer))
        if 2 <= sentences <= 5:
            quality_score += 0.2
        elif sentences > 5:
            quality_score += 0.1
        
        # Final confidence score (base + quality adjustments)
        confidence = base_confidence * 0.7 + quality_score
        
        # Cap at 1.0
        return min(confidence, 1.0)
    
    def get_stats(self):
        """
        Get statistics about the model
        """
        if not self.load_model():
            return {
                'status': 'Model not trained',
                'documents': 0,
                'vocabulary': 0
            }
        
        stats = {
            'status': 'Model trained',
            'documents': len(self.documents),
            'vocabulary': len(self.vectorizer.get_feature_names_out()) if self.vectorizer else 0,
            'semantic_dimensions': self.embedding_dim,
            'semantic_clusters': len(self.word_clusters),
            'response_templates': sum(len(templates) for templates in self.response_templates.values()) if hasattr(self, 'response_templates') else 0
        }
        
        return stats

def main():
    """
    Main function to run the Advanced Training AI
    """
    print("🧠 ADVANCED TRAINING AI")
    print("="*50)
    
    try:
        # Initialize AI
        ai = AdvancedTrainingAI(data_folder="dataset")
        
        # Show statistics
        stats = ai.get_stats()
        print(f"\n📊 MODEL STATISTICS:")
        for key, value in stats.items():
            print(f"• {key}: {value}")
        
        print("\n" + "="*50)
        print("🎯 AI ready to answer questions!")
        print("Type 'quit' to exit")
        print("Type 'stats' to see statistics")
        print("="*50)
        
        while True:
            question = input("\n❓ Your question: ").strip()
            
            if question.lower() == 'quit':
                print("👋 Goodbye!")
                break
            elif question.lower() == 'stats':
                stats = ai.get_stats()
                print(f"\n📊 MODEL STATISTICS:")
                for key, value in stats.items():
                    print(f"• {key}: {value}")
                continue
            
            if not question:
                continue
            
            print("🤔 Thinking...")
            
            # Answer question
            start_time = time.time()
            result = ai.answer_question(question)
            processing_time = time.time() - start_time
            
            print(f"\n🧠 **Answer:** (generated in {processing_time:.2f}s)")
            print(result['answer'])
            print(f"\n🎯 **Confidence:** {result['confidence']:.2f}")
            
            if result['sources']:
                print(f"📚 **Sources:** {len(result['sources'])} documents")
                for i, source in enumerate(result['sources'][:3], 1):
                    print(f"   {i}. {os.path.basename(source)}")
    
    except Exception as e:
        print(f"❌ Error: {e}")
        print("\nPlease check:")
        print("1. 'dataset' folder exists and contains .txt files")
        print("2. Required libraries are installed")
        print("3. Sufficient memory available")

if __name__ == "__main__":
    main()