agents/knowledge_agent.py

"""
Knowledge Agent - Retrieves and explains CFA concepts using local embeddings
Free version using sentence-transformers and ChromaDB
"""

import json
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime
import re
import numpy as np
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.config import Settings
import os


class KnowledgeAgent:
    """Retrieves and explains CFA knowledge concepts using local embeddings"""
    
    def __init__(self, knowledge_base_path: str = "data/cfa_knowledge_base.json"):
        self.knowledge_base_path = knowledge_base_path
        self.knowledge_base = self._load_knowledge_base()
        self.search_history = []
        
        # Initialize embedding model (free, runs locally)
        try:
            self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
            print("✅ Sentence transformer model loaded")
        except Exception as e:
            print(f"⚠️ Could not load sentence transformer: {e}")
            self.embedding_model = None
        
        # Initialize vector database
        self.chroma_client = None
        self.collection = None
        self._initialize_vector_db()
        
    def _load_knowledge_base(self) -> Dict[str, Any]:
        """Load CFA knowledge base from JSON file"""
        try:
            # Try multiple possible paths
            possible_paths = [
                self.knowledge_base_path,
                os.path.join(os.path.dirname(__file__), '..', self.knowledge_base_path),
                os.path.join(os.getcwd(), self.knowledge_base_path)
            ]
            
            for path in possible_paths:
                if os.path.exists(path):
                    with open(path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        print(f"✅ Knowledge base loaded from {path}")
                        return data
            
            print("⚠️ Knowledge base file not found, creating default")
            return self._create_default_knowledge_base()
            
        except Exception as e:
            print(f"❌ Error loading knowledge base: {e}")
            return self._create_default_knowledge_base()
    
    def _create_default_knowledge_base(self) -> Dict[str, Any]:
        """Create a default knowledge base if file not found"""
        return {
            "topics": {
                "Time Value of Money": {
                    "description": "Core concept of finance dealing with the value of money over time",
                    "key_concepts": [
                        {
                            "name": "Present Value",
                            "definition": "Current worth of a future sum of money given a specified rate of return",
                            "formula": "PV = FV / (1 + r)^n",
                            "example": "The present value of $110 received one year from now at 10% discount rate is $100"
                        },
                        {
                            "name": "Future Value",
                            "definition": "Value of a current asset at a specified date in the future",
                            "formula": "FV = PV × (1 + r)^n",
                            "example": "$100 invested today at 10% annual interest will grow to $110 in one year"
                        }
                    ],
                    "common_applications": [
                        "Valuing bonds and stocks",
                        "Capital budgeting decisions",
                        "Retirement planning",
                        "Loan amortization"
                    ],
                    "cfa_level": "I",
                    "weight": "8-12%",
                    "difficulty": "Medium"
                },
                "Statistics": {
                    "description": "Collection, analysis, interpretation, and presentation of data",
                    "key_concepts": [
                        {
                            "name": "Mean",
                            "definition": "Average of a set of numbers",
                            "formula": "μ = Σ x_i / n",
                            "example": "Mean of [2, 4, 6, 8] is (2+4+6+8)/4 = 5"
                        },
                        {
                            "name": "Standard Deviation",
                            "definition": "Measure of the amount of variation or dispersion",
                            "formula": "σ = √[Σ(x_i - μ)² / (n-1)]",
                            "example": "Standard deviation measures how spread out numbers are"
                        }
                    ],
                    "common_applications": [
                        "Portfolio analysis",
                        "Risk assessment",
                        "Performance measurement"
                    ],
                    "cfa_level": "I",
                    "weight": "5-8%",
                    "difficulty": "Medium"
                }
            },
            "formulas": {
                "TVM": [
                    {"name": "Future Value", "formula": "FV = PV(1+r)^n"},
                    {"name": "Present Value", "formula": "PV = FV/(1+r)^n"}
                ]
            },
            "glossary": {
                "NPV": "Net Present Value - difference between present value of cash inflows and outflows",
                "IRR": "Internal Rate of Return - discount rate that makes NPV zero",
                "Beta": "Measure of a stock's volatility relative to the market",
                "Alpha": "Measure of performance on a risk-adjusted basis"
            }
        }
    
    def _initialize_vector_db(self):
        """Initialize ChromaDB vector database"""
        try:
            # Create persistent directory for ChromaDB
            chroma_dir = os.path.join(os.path.dirname(__file__), '..', 'chroma_db')
            os.makedirs(chroma_dir, exist_ok=True)
            
            self.chroma_client = chromadb.PersistentClient(
                path=chroma_dir,
                settings=Settings(anonymized_telemetry=False)
            )
            
            # Create or get collection
            self.collection = self.chroma_client.get_or_create_collection(
                name="cfa_knowledge",
                metadata={"description": "CFA Knowledge Base Embeddings"}
            )
            
            print("✅ ChromaDB initialized")
            
            # Index knowledge if collection is empty
            if self.collection.count() == 0:
                self._index_knowledge_base()
                
        except Exception as e:
            print(f"⚠️ Could not initialize ChromaDB: {e}")
            self.chroma_client = None
            self.collection = None
    
    def _index_knowledge_base(self):
        """Index knowledge base in vector database"""
        if not self.embedding_model or not self.collection:
            return
        
        print("📚 Indexing knowledge base...")
        
        documents = []
        metadatas = []
        ids = []
        
        # Index topics
        for topic_name, topic_data in self.knowledge_base.get('topics', {}).items():
            # Topic description
            doc_text = f"Topic: {topic_name}. Description: {topic_data.get('description', '')}"
            documents.append(doc_text)
            metadatas.append({
                "type": "topic",
                "topic": topic_name,
                "cfa_level": topic_data.get('cfa_level', 'I'),
                "difficulty": topic_data.get('difficulty', 'Medium')
            })
            ids.append(f"topic_{topic_name}")
            
            # Key concepts
            for concept in topic_data.get('key_concepts', []):
                doc_text = f"Concept: {concept['name']}. Definition: {concept.get('definition', '')}"
                if 'formula' in concept:
                    doc_text += f" Formula: {concept['formula']}"
                if 'example' in concept:
                    doc_text += f" Example: {concept['example']}"
                
                documents.append(doc_text)
                metadatas.append({
                    "type": "concept",
                    "topic": topic_name,
                    "concept": concept['name'],
                    "has_formula": 'formula' in concept
                })
                ids.append(f"concept_{topic_name}_{concept['name']}")
        
        # Index formulas
        for category, formulas in self.knowledge_base.get('formulas', {}).items():
            for formula in formulas:
                doc_text = f"Formula: {formula['name']}. Formula: {formula.get('formula', '')}"
                if 'description' in formula:
                    doc_text += f" Description: {formula['description']}"
                
                documents.append(doc_text)
                metadatas.append({
                    "type": "formula",
                    "category": category,
                    "formula_name": formula['name']
                })
                ids.append(f"formula_{category}_{formula['name']}")
        
        # Index glossary terms
        for term, definition in self.knowledge_base.get('glossary', {}).items():
            doc_text = f"Term: {term}. Definition: {definition}"
            documents.append(doc_text)
            metadatas.append({
                "type": "glossary",
                "term": term
            })
            ids.append(f"glossary_{term}")
        
        # Generate embeddings and add to collection
        if documents:
            try:
                embeddings = self.embedding_model.encode(documents).tolist()
                self.collection.add(
                    documents=documents,
                    metadatas=metadatas,
                    ids=ids,
                    embeddings=embeddings
                )
                print(f"✅ Indexed {len(documents)} knowledge items")
            except Exception as e:
                print(f"❌ Error indexing documents: {e}")
    
    def process(self,
                query: str,
                user_profile: Dict[str, Any],
                conversation_history: List[Dict] = None) -> Dict[str, Any]:
        """
        Retrieve and explain CFA knowledge concepts
        
        Args:
            query: Knowledge query
            user_profile: User's learning profile
            conversation_history: Previous conversation
        
        Returns:
            Knowledge explanation
        """
        # Analyze query to determine knowledge needs
        knowledge_type = self._analyze_knowledge_query(query)
        
        # Search knowledge base
        if self.collection and self.embedding_model:
            # Use vector search
            search_results = self._vector_search(query, knowledge_type, user_profile)
        else:
            # Fallback to keyword search
            search_results = self._keyword_search(query, knowledge_type)
        
        # Format response
        response = self._format_knowledge_response(query, search_results, knowledge_type, user_profile)
        
        # Log search
        self._log_search(query, knowledge_type, search_results)
        
        return {
            'response': response,
            'knowledge_type': knowledge_type,
            'search_results': search_results,
            'confidence': self._calculate_confidence(search_results),
            'metadata': {
                'search_method': 'vector' if self.collection else 'keyword',
                'result_count': len(search_results.get('items', []))
            }
        }
    
    def _analyze_knowledge_query(self, query: str) -> str:
        """Analyze query to determine type of knowledge needed"""
        query_lower = query.lower()
        
        if any(term in query_lower for term in ['what is', 'define', 'definition of', 'meaning of']):
            return 'definition'
        elif any(term in query_lower for term in ['explain', 'how does', 'how to', 'understand']):
            return 'explanation'
        elif any(term in query_lower for term in ['formula', 'calculate', 'equation', 'solve']):
            return 'formula'
        elif any(term in query_lower for term in ['example', 'case study', 'application', 'real world']):
            return 'example'
        elif any(term in query_lower for term in ['compare', 'difference between', 'vs', 'versus']):
            return 'comparison'
        elif any(term in query_lower for term in ['list', 'what are the', 'types of', 'categories']):
            return 'list'
        else:
            return 'general'
    
    def _vector_search(self, query: str, knowledge_type: str, user_profile: Dict[str, Any]) -> Dict[str, Any]:
        """Search using vector embeddings"""
        results = {
            'items': [],
            'topics': set(),
            'method': 'vector'
        }
        
        try:
            # Generate query embedding
            query_embedding = self.embedding_model.encode(query).tolist()
            
            # Search collection
            search_results = self.collection.query(
                query_embeddings=[query_embedding],
                n_results=10,
                include=["documents", "metadatas", "distances"]
            )
            
            # Process results
            if search_results['documents']:
                for i, (doc, metadata, distance) in enumerate(zip(
                    search_results['documents'][0],
                    search_results['metadatas'][0],
                    search_results['distances'][0]
                )):
                    # Convert distance to similarity score (0-1)
                    similarity = 1 - (distance / 2)  # Rough approximation
                    
                    item = {
                        'content': doc,
                        'metadata': metadata,
                        'similarity': similarity,
                        'type': metadata.get('type', 'unknown')
                    }
                    
                    results['items'].append(item)
                    
                    # Extract topics
                    if 'topic' in metadata:
                        results['topics'].add(metadata['topic'])
            
            # Filter by knowledge type if needed
            if knowledge_type != 'general':
                filtered_items = []
                for item in results['items']:
                    if self._matches_knowledge_type(item, knowledge_type):
                        filtered_items.append(item)
                results['items'] = filtered_items[:5]  # Limit to 5
            
        except Exception as e:
            print(f"❌ Vector search error: {e}")
            # Fall back to keyword search
            return self._keyword_search(query, knowledge_type)
        
        return results
    
    def _keyword_search(self, query: str, knowledge_type: str) -> Dict[str, Any]:
        """Fallback keyword search"""
        results = {
            'items': [],
            'topics': set(),
            'method': 'keyword'
        }
        
        query_lower = query.lower()
        query_terms = set(query_lower.split())
        
        # Search topics
        for topic_name, topic_data in self.knowledge_base.get('topics', {}).items():
            topic_lower = topic_name.lower()
            
            # Check topic name
            if any(term in topic_lower for term in query_terms):
                results['topics'].add(topic_name)
                results['items'].append({
                    'content': f"Topic: {topic_name}. {topic_data.get('description', '')}",
                    'metadata': {'type': 'topic', 'topic': topic_name},
                    'similarity': 0.8,
                    'type': 'topic'
                })
            
            # Search concepts within topic
            for concept in topic_data.get('key_concepts', []):
                concept_lower = concept['name'].lower()
                if any(term in concept_lower for term in query_terms):
                    results['topics'].add(topic_name)
                    content = f"Concept: {concept['name']}. {concept.get('definition', '')}"
                    if 'formula' in concept:
                        content += f" Formula: {concept['formula']}"
                    results['items'].append({
                        'content': content,
                        'metadata': {'type': 'concept', 'topic': topic_name, 'concept': concept['name']},
                        'similarity': 0.7,
                        'type': 'concept'
                    })
        
        # Search glossary
        for term, definition in self.knowledge_base.get('glossary', {}).items():
            term_lower = term.lower()
            if term_lower in query_lower or any(term_word in query_terms for term_word in term_lower.split()):
                results['items'].append({
                    'content': f"Term: {term}. Definition: {definition}",
                    'metadata': {'type': 'glossary', 'term': term},
                    'similarity': 0.9,
                    'type': 'glossary'
                })
        
        # Sort by similarity (keyword match quality)
        results['items'].sort(key=lambda x: x['similarity'], reverse=True)
        results['items'] = results['items'][:5]  # Limit to 5
        
        return results
    
    def _matches_knowledge_type(self, item: Dict[str, Any], knowledge_type: str) -> bool:
        """Check if item matches the requested knowledge type"""
        item_type = item.get('type')
        
        type_mapping = {
            'definition': ['glossary', 'concept'],
            'explanation': ['topic', 'concept'],
            'formula': ['concept', 'formula'],
            'example': ['concept'],
            'comparison': ['topic', 'concept'],
            'list': ['topic']
        }
        
        if knowledge_type in type_mapping:
            return item_type in type_mapping[knowledge_type]
        
        return True
    
    def _format_knowledge_response(self, 
                                 query: str, 
                                 results: Dict[str, Any], 
                                 knowledge_type: str,
                                 user_profile: Dict[str, Any]) -> str:
        """Format knowledge response based on query type"""
        
        response_parts = []
        
        if not results.get('items'):
            response_parts.append(f"I couldn't find specific information about '{query}'.")
            response_parts.append("\n**Try asking about:**")
            response_parts.append("- Time Value of Money (TVM)")
            response_parts.append("- Statistics and probability")
            response_parts.append("- Portfolio management")
            response_parts.append("- Financial ratios")
            response_parts.append("- CFA ethics")
            return "\n".join(response_parts)
        
        # Add header based on knowledge type
        headers = {
            'definition': "📚 Definitions",
            'explanation': "📖 Explanations",
            'formula': "🧮 Formulas",
            'example': "💡 Examples",
            'comparison': "⚖️ Comparisons",
            'list': "📋 Lists",
            'general': "📚 Knowledge Results"
        }
        
        response_parts.append(f"## {headers.get(knowledge_type, '📚 Knowledge')}")
        
        # Add results
        for i, item in enumerate(results['items'][:3], 1):  # Limit to 3 items
            response_parts.append(f"\n### Result {i}")
            
            content = item['content']
            metadata = item.get('metadata', {})
            
            # Format based on type
            if metadata.get('type') == 'topic':
                response_parts.append(f"**Topic:** {metadata.get('topic', 'Unknown')}")
                response_parts.append(f"{content}")
                
            elif metadata.get('type') == 'concept':
                response_parts.append(f"**Concept:** {metadata.get('concept', 'Unknown')}")
                response_parts.append(f"**Topic:** {metadata.get('topic', 'General')}")
                response_parts.append(f"{content}")
                
            elif metadata.get('type') == 'glossary':
                response_parts.append(f"**Term:** {metadata.get('term', 'Unknown')}")
                response_parts.append(f"{content}")
                
            elif metadata.get('type') == 'formula':
                response_parts.append(f"**Formula:** {metadata.get('formula_name', 'Unknown')}")
                response_parts.append(f"{content}")
            
            else:
                response_parts.append(f"{content}")
        
        # Add study tips if relevant
        if results.get('topics'):
            topics = list(results['topics'])[:2]  # Limit to 2 topics
            if topics:
                response_parts.append("\n---\n")
                response_parts.append("**📝 Study Tips:**")
                
                for topic in topics:
                    tips = self._get_study_tips(topic, user_profile)
                    if tips:
                        response_parts.append(f"\n**For {topic}:**")
                        for tip in tips[:2]:
                            response_parts.append(f"- {tip}")
        
        return "\n".join(response_parts)
    
    def _get_study_tips(self, topic: str, user_profile: Dict[str, Any]) -> List[str]:
        """Get study tips for a topic"""
        tips_db = {
            'Time Value of Money': [
                "Practice with different compounding periods (annual, semi-annual, monthly)",
                "Master the TVM calculator keys: N, I/Y, PV, PMT, FV",
                "Understand the relationship between discount rate and present value"
            ],
            'Statistics': [
                "Focus on interpreting results, not just calculations",
                "Understand when to use population vs sample statistics",
                "Practice hypothesis testing steps systematically"
            ],
            'Portfolio Management': [
                "Memorize the CAPM formula: E(R) = Rf + β(Rm - Rf)",
                "Understand the efficient frontier and capital market line",
                "Practice calculating portfolio variance with different correlations"
            ]
        }
        
        tips = tips_db.get(topic, [
            "Break down complex concepts into smaller parts",
            "Create flashcards for key terms and formulas",
            "Practice with real-world examples"
        ])
        
        # Adjust based on user profile
        experience = user_profile.get('experience_level', 'beginner')
        if experience == 'beginner':
            tips.insert(0, "Start with foundational concepts before advancing")
        elif experience == 'advanced':
            tips.append("Focus on application and analysis questions")
        
        return tips[:3]  # Return top 3 tips
    
    def _calculate_confidence(self, results: Dict[str, Any]) -> float:
        """Calculate confidence in search results"""
        if not results.get('items'):
            return 0.2
        
        # Calculate average similarity
        similarities = [item.get('similarity', 0.5) for item in results['items']]
        avg_similarity = sum(similarities) / len(similarities)
        
        # Adjust based on result count
        count_factor = min(len(results['items']) / 5, 1.0) * 0.2
        
        confidence = avg_similarity + count_factor
        return min(confidence, 0.95)
    
    def _log_search(self, query: str, knowledge_type: str, results: Dict[str, Any]):
        """Log knowledge searches"""
        log_entry = {
            'timestamp': datetime.now().isoformat(),
            'query': query,
            'knowledge_type': knowledge_type,
            'result_count': len(results.get('items', [])),
            'method': results.get('method', 'unknown'),
            'topics': list(results.get('topics', set()))
        }
        self.search_history.append(log_entry)
        
        # Keep only last 100 searches
        if len(self.search_history) > 100:
            self.search_history = self.search_history[-100:]
    
    def get_search_history(self, limit: int = 10) -> List[Dict]:
        """Get recent search history"""
        return self.search_history[-limit:]
    
    def get_metrics(self) -> Dict[str, Any]:
        """Get knowledge agent metrics"""
        if not self.search_history:
            return {
                'total_searches': 0,
                'vector_db_ready': self.collection is not None,
                'embedding_model_ready': self.embedding_model is not None,
                'knowledge_items': self.collection.count() if self.collection else 0
            }
        
        # Calculate success rate (searches with results)
        successful = sum(1 for entry in self.search_history if entry['result_count'] > 0)
        success_rate = successful / len(self.search_history) if self.search_history else 0
        
        # Get most searched topics
        topic_counts = {}
        for entry in self.search_history:
            for topic in entry.get('topics', []):
                topic_counts[topic] = topic_counts.get(topic, 0) + 1
        
        common_topics = sorted(topic_counts.items(), key=lambda x: x[1], reverse=True)[:3]
        
        return {
            'total_searches': len(self.search_history),
            'success_rate': success_rate,
            'common_topics': common_topics,
            'vector_db_ready': self.collection is not None,
            'embedding_model_ready': self.embedding_model is not None,
            'knowledge_items': self.collection.count() if self.collection else 0,
            'recent_queries': [entry['query'][:50] for entry in self.search_history[-5:]]
        }