knowledge_engine.py

"""
Knowledge Engine - Layer 2: Knowledge management and retrieval
"""

import json
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime
import hashlib
from dataclasses import dataclass
from enum import Enum


class KnowledgeType(Enum):
    """Types of knowledge in the system"""
    CONCEPT = "concept"
    FORMULA = "formula"
    EXAMPLE = "example"
    PRACTICE_PROBLEM = "practice_problem"
    STUDY_GUIDE = "study_guide"
    EXAM_TIP = "exam_tip"


@dataclass
class KnowledgeNode:
    """Knowledge node in the knowledge graph"""
    id: str
    type: KnowledgeType
    content: Dict[str, Any]
    topics: List[str]
    difficulty: str
    prerequisites: List[str]
    related_nodes: List[str]
    metadata: Dict[str, Any]
    created_at: datetime
    last_accessed: datetime
    access_count: int


class KnowledgeEngine:
    """Knowledge management and retrieval engine"""
    
    def __init__(self, knowledge_base_path: str = "data/cfa_knowledge_base.json"):
        self.knowledge_graph: Dict[str, KnowledgeNode] = {}
        self.topic_index: Dict[str, List[str]] = {}
        self.concept_index: Dict[str, List[str]] = {}
        self.load_knowledge_base(knowledge_base_path)
        self.access_log = []
        
    def load_knowledge_base(self, path: str):
        """Load and index knowledge base"""
        try:
            with open(path, 'r', encoding='utf-8') as f:
                knowledge_data = json.load(f)
                self._build_knowledge_graph(knowledge_data)
                print(f"✅ Knowledge base loaded: {len(self.knowledge_graph)} nodes")
        except FileNotFoundError:
            print(f"⚠️ Knowledge base not found at {path}, creating minimal graph")
            self._create_minimal_graph()
    
    def _build_knowledge_graph(self, knowledge_data: Dict[str, Any]):
        """Build knowledge graph from loaded data"""
        # Process topics
        for topic_name, topic_data in knowledge_data.get('topics', {}).items():
            # Create concept nodes
            concepts = topic_data.get('key_concepts', [])
            for concept in concepts:
                node_id = self._generate_node_id(concept['name'], KnowledgeType.CONCEPT)
                node = KnowledgeNode(
                    id=node_id,
                    type=KnowledgeType.CONCEPT,
                    content={
                        'name': concept['name'],
                        'definition': concept.get('definition', ''),
                        'formula': concept.get('formula', ''),
                        'example': concept.get('example', '')
                    },
                    topics=[topic_name],
                    difficulty=topic_data.get('difficulty', 'medium'),
                    prerequisites=[],
                    related_nodes=[],
                    metadata={
                        'cfa_level': topic_data.get('cfa_level', 'I'),
                        'weight': topic_data.get('weight', ''),
                        'applications': topic_data.get('common_applications', [])
                    },
                    created_at=datetime.now(),
                    last_accessed=datetime.now(),
                    access_count=0
                )
                self.knowledge_graph[node_id] = node
                
                # Update indices
                self._add_to_index(node_id, topic_name, concept['name'])
        
        # Process formulas
        for category, formulas in knowledge_data.get('formulas', {}).items():
            for formula in formulas:
                node_id = self._generate_node_id(formula['name'], KnowledgeType.FORMULA)
                node = KnowledgeNode(
                    id=node_id,
                    type=KnowledgeType.FORMULA,
                    content={
                        'name': formula['name'],
                        'formula': formula.get('formula', ''),
                        'description': formula.get('description', ''),
                        'variables': formula.get('variables', {})
                    },
                    topics=[category],
                    difficulty='medium',
                    prerequisites=[],
                    related_nodes=[],
                    metadata={'category': category},
                    created_at=datetime.now(),
                    last_accessed=datetime.now(),
                    access_count=0
                )
                self.knowledge_graph[node_id] = node
                self._add_to_index(node_id, category, formula['name'])
        
        # Build relationships
        self._build_relationships()
    
    def _create_minimal_graph(self):
        """Create minimal knowledge graph if no data file"""
        # Core TVM concepts
        tvm_concepts = [
            {
                'name': 'Present Value',
                'definition': 'Current worth of future cash flows',
                'formula': 'PV = FV / (1 + r)^n'
            },
            {
                'name': 'Future Value',
                'definition': 'Future worth of current investment',
                'formula': 'FV = PV × (1 + r)^n'
            }
        ]
        
        for concept in tvm_concepts:
            node_id = self._generate_node_id(concept['name'], KnowledgeType.CONCEPT)
            node = KnowledgeNode(
                id=node_id,
                type=KnowledgeType.CONCEPT,
                content=concept,
                topics=['Time Value of Money'],
                difficulty='medium',
                prerequisites=[],
                related_nodes=[],
                metadata={'cfa_level': 'I', 'weight': '8-12%'},
                created_at=datetime.now(),
                last_accessed=datetime.now(),
                access_count=0
            )
            self.knowledge_graph[node_id] = node
            self._add_to_index(node_id, 'Time Value of Money', concept['name'])
    
    def _generate_node_id(self, name: str, ktype: KnowledgeType) -> str:
        """Generate unique node ID"""
        hash_input = f"{name}_{ktype.value}_{datetime.now().timestamp()}"
        return hashlib.md5(hash_input.encode()).hexdigest()[:12]
    
    def _add_to_index(self, node_id: str, topic: str, concept: str):
        """Add node to search indices"""
        # Topic index
        if topic not in self.topic_index:
            self.topic_index[topic] = []
        self.topic_index[topic].append(node_id)
        
        # Concept index (by keywords)
        keywords = concept.lower().split()
        for keyword in keywords:
            if len(keyword) > 3:  # Ignore short words
                if keyword not in self.concept_index:
                    self.concept_index[keyword] = []
                self.concept_index[keyword].append(node_id)
    
    def _build_relationships(self):
        """Build relationships between knowledge nodes"""
        # Simple relationship building - in production, use NLP
        node_ids = list(self.knowledge_graph.keys())
        
        for i, node_id in enumerate(node_ids):
            node = self.knowledge_graph[node_id]
            related = []
            
            # Find related nodes by topic
            for other_id in node_ids:
                if other_id != node_id:
                    other_node = self.knowledge_graph[other_id]
                    # Same topic
                    if any(topic in other_node.topics for topic in node.topics):
                        related.append(other_id)
                    # Related concepts (simple keyword matching)
                    node_content = str(node.content).lower()
                    other_content = str(other_node.content).lower()
                    common_words = set(node_content.split()) & set(other_content.split())
                    if len(common_words) >= 2:
                        related.append(other_id)
            
            # Limit relationships
            node.related_nodes = list(set(related))[:10]
            self.knowledge_graph[node_id] = node
    
    def retrieve_knowledge(self, 
                          query: str, 
                          user_profile: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        Retrieve knowledge based on query
        
        Args:
            query: Knowledge query
            user_profile: User's learning profile
        
        Returns:
            Retrieved knowledge
        """
        # Log access
        self.access_log.append({
            'query': query,
            'timestamp': datetime.now().isoformat(),
            'user_profile': user_profile
        })
        
        # Search for relevant nodes
        search_results = self._search_knowledge(query)
        
        # Rank results
        ranked_results = self._rank_results(search_results, query, user_profile)
        
        # Prepare response
        response = self._prepare_knowledge_response(ranked_results, query, user_profile)
        
        # Update access counts
        for node_id, _ in ranked_results[:5]:
            if node_id in self.knowledge_graph:
                node = self.knowledge_graph[node_id]
                node.last_accessed = datetime.now()
                node.access_count += 1
                self.knowledge_graph[node_id] = node
        
        return response
    
    def _search_knowledge(self, query: str) -> List[str]:
        """Search knowledge graph for relevant nodes"""
        query_lower = query.lower()
        found_nodes = set()
        
        # Search by topic
        for topic in self.topic_index.keys():
            if topic.lower() in query_lower:
                found_nodes.update(self.topic_index[topic])
        
        # Search by keywords
        query_words = query_lower.split()
        for word in query_words:
            if len(word) > 3 and word in self.concept_index:
                found_nodes.update(self.concept_index[word])
        
        # If no results, try broader search
        if not found_nodes:
            for node_id, node in self.knowledge_graph.items():
                node_text = str(node.content).lower()
                if any(word in node_text for word in query_words[:3]):
                    found_nodes.add(node_id)
        
        return list(found_nodes)
    
    def _rank_results(self, 
                     node_ids: List[str], 
                     query: str,
                     user_profile: Optional[Dict[str, Any]]) -> List[Tuple[str, float]]:
        """Rank search results by relevance"""
        ranked = []
        
        for node_id in node_ids:
            node = self.knowledge_graph.get(node_id)
            if not node:
                continue
            
            relevance_score = self._calculate_relevance(node, query, user_profile)
            ranked.append((node_id, relevance_score))
        
        # Sort by relevance
        ranked.sort(key=lambda x: x[1], reverse=True)
        return ranked
    
    def _calculate_relevance(self, 
                           node: KnowledgeNode, 
                           query: str,
                           user_profile: Optional[Dict[str, Any]]) -> float:
        """Calculate relevance score for a node"""
        score = 0.0
        
        # Text similarity (simplified)
        query_lower = query.lower()
        node_text = str(node.content).lower()
        
        # Keyword matches
        query_words = set(query_lower.split())
        node_words = set(node_text.split())
        keyword_overlap = len(query_words & node_words) / max(len(query_words), 1)
        score += keyword_overlap * 0.4
        
        # Topic relevance
        if user_profile:
            study_focus = user_profile.get('study_focus', [])
            topic_overlap = len(set(node.topics) & set(study_focus))
            score += topic_overlap * 0.2
        
        # Popularity/recency
        recency_factor = 1.0 - ((datetime.now() - node.last_accessed).days / 365)
        popularity_factor = min(node.access_count / 100, 1.0)
        score += (recency_factor * 0.2 + popularity_factor * 0.2)
        
        # Difficulty match
        if user_profile:
            experience = user_profile.get('experience_level', 'beginner')
            difficulty_score = self._match_difficulty(node.difficulty, experience)
            score += difficulty_score * 0.1
        
        return min(score, 1.0)
    
    def _match_difficulty(self, node_difficulty: str, user_experience: str) -> float:
        """Match node difficulty to user experience"""
        difficulty_map = {'easy': 0, 'medium': 1, 'hard': 2}
        experience_map = {'beginner': 0, 'intermediate': 1, 'advanced': 2}
        
        node_level = difficulty_map.get(node_difficulty, 1)
        user_level = experience_map.get(user_experience, 1)
        
        # Prefer nodes at or slightly above user level
        diff = abs(node_level - user_level)
        if diff == 0:
            return 1.0
        elif diff == 1:
            return 0.7
        else:
            return 0.3
    
    def _prepare_knowledge_response(self, 
                                   ranked_results: List[Tuple[str, float]], 
                                   query: str,
                                   user_profile: Optional[Dict[str, Any]]) -> Dict[str, Any]:
        """Prepare structured knowledge response"""
        response = {
            'query': query,
            'results': [],
            'topics': set(),
            'total_results': len(ranked_results),
            'timestamp': datetime.now().isoformat()
        }
        
        # Get top results
        for node_id, relevance in ranked_results[:10]:  # Top 10 results
            node = self.knowledge_graph.get(node_id)
            if node:
                node_data = {
                    'id': node.id,
                    'type': node.type.value,
                    'content': node.content,
                    'topics': node.topics,
                    'difficulty': node.difficulty,
                    'relevance': relevance,
                    'metadata': node.metadata,
                    'access_count': node.access_count
                }
                response['results'].append(node_data)
                response['topics'].update(node.topics)
        
        response['topics'] = list(response['topics'])
        
        # Add learning path if user profile provided
        if user_profile and response['results']:
            response['learning_path'] = self._generate_learning_path(
                response['results'], user_profile
            )
        
        return response
    
    def _generate_learning_path(self, 
                               results: List[Dict[str, Any]], 
                               user_profile: Dict[str, Any]) -> Dict[str, Any]:
        """Generate learning path based on results and user profile"""
        # Group by topic
        topics = {}
        for result in results:
            for topic in result['topics']:
                if topic not in topics:
                    topics[topic] = []
                topics[topic].append(result)
        
        # Select primary topic
        primary_topic = max(topics.items(), key=lambda x: len(x[1]))[0] if topics else None
        
        if not primary_topic:
            return {'suggested_path': [], 'estimated_time': 0}
        
        # Generate learning path
        learning_path = []
        
        # Start with concepts
        concepts = [r for r in results if r['type'] == 'concept']
        for concept in concepts[:3]:
            learning_path.append({
                'type': 'concept',
                'title': concept['content'].get('name', 'Concept'),
                'time': 15,  # minutes
                'description': f"Learn: {concept['content'].get('definition', '')[:100]}..."
            })
        
        # Add formulas
        formulas = [r for r in results if r['type'] == 'formula']
        for formula in formulas[:2]:
            learning_path.append({
                'type': 'formula',
                'title': formula['content'].get('name', 'Formula'),
                'time': 10,
                'description': f"Master: {formula['content'].get('formula', '')}"
            })
        
        # Add practice
        learning_path.append({
            'type': 'practice',
            'title': f'Practice {primary_topic} Problems',
            'time': 20,
            'description': 'Apply concepts through practice problems'
        })
        
        estimated_time = sum(step['time'] for step in learning_path)
        
        return {
            'primary_topic': primary_topic,
            'suggested_path': learning_path,
            'estimated_time': estimated_time,
            'suitable_for': user_profile.get('experience_level', 'beginner')
        }
    
    def add_knowledge(self, 
                     knowledge_type: KnowledgeType, 
                     content: Dict[str, Any], 
                     topics: List[str],
                     difficulty: str = 'medium',
                     metadata: Optional[Dict[str, Any]] = None) -> str:
        """
        Add new knowledge to the graph
        
        Args:
            knowledge_type: Type of knowledge
            content: Knowledge content
            topics: Related topics
            difficulty: Difficulty level
            metadata: Additional metadata
        
        Returns:
            Node ID
        """
        # Generate node
        node_name = content.get('name', 'Unnamed Knowledge')
        node_id = self._generate_node_id(node_name, knowledge_type)
        
        node = KnowledgeNode(
            id=node_id,
            type=knowledge_type,
            content=content,
            topics=topics,
            difficulty=difficulty,
            prerequisites=[],
            related_nodes=[],
            metadata=metadata or {},
            created_at=datetime.now(),
            last_accessed=datetime.now(),
            access_count=0
        )
        
        # Add to graph
        self.knowledge_graph[node_id] = node
        
        # Update indices
        for topic in topics:
            if topic not in self.topic_index:
                self.topic_index[topic] = []
            self.topic_index[topic].append(node_id)
        
        # Index content keywords
        content_str = str(content).lower()
        for word in content_str.split():
            if len(word) > 3:
                if word not in self.concept_index:
                    self.concept_index[word] = []
                self.concept_index[word].append(node_id)
        
        print(f"✅ Added new knowledge: {node_name} ({node_id})")
        return node_id
    
    def get_knowledge_stats(self) -> Dict[str, Any]:
        """Get knowledge engine statistics"""
        stats = {
            'total_nodes': len(self.knowledge_graph),
            'by_type': {},
            'by_topic': {},
            'by_difficulty': {},
            'access_stats': {
                'total_accesses': sum(node.access_count for node in self.knowledge_graph.values()),
                'most_accessed': []
            }
        }
        
        # Count by type
        for node in self.knowledge_graph.values():
            ktype = node.type.value
            stats['by_type'][ktype] = stats['by_type'].get(ktype, 0) + 1
            
            # Count by topic
            for topic in node.topics:
                stats['by_topic'][topic] = stats['by_topic'].get(topic, 0) + 1
            
            # Count by difficulty
            stats['by_difficulty'][node.difficulty] = stats['by_difficulty'].get(node.difficulty, 0) + 1
        
        # Most accessed nodes
        sorted_nodes = sorted(self.knowledge_graph.values(), 
                            key=lambda x: x.access_count, 
                            reverse=True)[:5]
        stats['access_stats']['most_accessed'] = [
            {
                'id': node.id,
                'name': node.content.get('name', 'Unnamed'),
                'access_count': node.access_count,
                'type': node.type.value
            }
            for node in sorted_nodes
        ]
        
        return stats
    
    def export_knowledge_base(self, path: str):
        """Export knowledge base to file"""
        export_data = {
            'metadata': {
                'export_date': datetime.now().isoformat(),
                'total_nodes': len(self.knowledge_graph),
                'version': '1.0'
            },
            'nodes': []
        }
        
        for node in self.knowledge_graph.values():
            node_data = {
                'id': node.id,
                'type': node.type.value,
                'content': node.content,
                'topics': node.topics,
                'difficulty': node.difficulty,
                'metadata': node.metadata,
                'access_count': node.access_count,
                'created_at': node.created_at.isoformat(),
                'last_accessed': node.last_accessed.isoformat()
            }
            export_data['nodes'].append(node_data)
        
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(export_data, f, indent=2, ensure_ascii=False)
        
        print(f"✅ Knowledge base exported to {path}")