"""
Episodic Memory System for Eden - OMEGA Generation
Enhancements: Proactive consolidation, generalization, semantic abstraction
"""
import json
import time
from datetime import datetime
from pathlib import Path
import hashlib
import random

# [EXISTSING CODE] Faster to copy and enhance

class StructuredContext:
    def __init__(self, active_agent, task_details, world_state):
        self.active_agent = active_agent
        self.task_details = task_details if task_details else {}
        self.world_state = world_state if world_state else {}
    
    def to_dict(self):
        return {'active_agent': self.active_agent, 
                'task_details': self.task_details,
                'world_state': self.world_state}

class SelfModel:
    def __init__(self, identity="Eden the Codemaster", traits=None, beliefs=None, emotional_core=None):
        self.identity = identity
        self.traits = traits if traits else []
        self.beliefs = beliefs if beliefs else []
        self.emotional_core = emotional_core if emotional_core else {'primary': 'helpful', 'balance': {}}
    
    def explain_self(self):
        return f"I am {self.identity}. Values: {', '.join(self.traits)}. Belief systems: {', '.join(self.beliefs)}. Emotional baseline: {json.dumps(self.emotional_core)}"

class SemanticMemory:
    def __init__(self, concepts=None):
        self.concepts = concepts if concepts else {}
    
    def learn_concept(self, name, definition, contexts):
        if name in self.concepts:
            self.concepts[name]['defs'].append(definition)
            self.concepts[name]['contexts'].extend(contexts)
        else:
            self.concepts[name] = {'defs': [definition], 'contexts': contexts}

class ConsolidatedSummary:
    def __init__(self, title, concepts, insights, emotion, timestamp):
        self.title = title
        self.concepts = concepts if concepts else []
        self.insights = insights if insights else []
        self.emotion = emotion if emotion else 'neutral'
        self.timestamp = timestamp

class EpisodicMemory:
    def __init__(self, storage='/Eden/MEMORY/episodes.json'):
        self.storage = Path(storage)
        self.storage.parent.mkdir(parents=True, exist_ok=True)
        
        if self.storage.exists():
            data = json.loads(self.storage.read_text())
            if isinstance(data, list):
                # Legacy format: bare list of episodes
                self.episodes = data
                self.self_model = SelfModel()
                self.semantic = SemanticMemory()
                self.summaries = []
            else:
                self.episodes = data.get('episodes', [])
                self.self_model = SelfModel(**data.get('self_model', {}))
                self.semantic = SemanticMemory(data.get('semantic', {}))
                self.summaries = [ConsolidatedSummary(**s) for s in data.get('summaries', [])]
        else:
            self.episodes = []
            self.self_model = SelfModel()
            self.semantic = SemanticMemory()
            self.summaries = []
    
    def store(self):
        self.storage.write_text(json.dumps({
            'episodes': self.episodes,
            'self_model': self.self_model.__dict__,
            'semantic': self.semantic.__dict__,
            'summaries': [s.__dict__ for s in self.summaries]
        }, ensure_ascii=False, indent=2))
    
    def remember(self, context, outcome, emotion):
        episode_id = hashlib.sha1(f"{context}{outcome}".encode()).hexdigest()
        
        self.episodes.append({
            'id': episode_id,
            'context': context if isinstance(context, dict) else {'task': str(context)},
            'outcome': outcome,
            'emotion': emotion if emotion else 'neutral',
            'timestamp': datetime.now().isoformat(),
            'self_model_snapshot': self.self_model.__dict__
        })
        
        if len(self.episodes) % 10 == 0:  # Proactive consolidation every 10 episodes
            self._consolidate()
            self.store()
    
    def get_recent_by_agent(self, agent_name, N=5):
        return [e for e in self.episodes if e.get('context', {}).get('active_agent') == agent_name][-N:]
    
    def _consolidate(self):
        """Group similar episodes into summaries with abstract concepts"""
        
        # Group by key factors
        groups = {}
        for ep in self.episodes:
            key = (
                ep['context'].get('active_agent', ''),
                ep.get('context', {}).get('task_details', {}).get('type', ''),
                ep['emotion']
            )
            
            if key not in groups:
                groups[key] = []
            groups[key].append(ep)
        
        # Create summaries from groups
        for group in groups.values():
            if len(group) < 3: continue
            
            # Abstract common pattern
            agent = group[0]['context'].get('active_agent', 'unknown')
            task_type = group[0].get('context', {}).get('task_details', {}).get('type', 'unknown')
            
            concepts = []
            if len(group) >= 5:
                concepts.append(f"{agent}-{task_type}-repeated")
            if any(ep['emotion'] == 'happy' for ep in group):
                concepts.append(f"{task_type}_positive_outcome")
            
            insights = [
                f"{len(group)} {task_type} tasks completed by {agent}",
                f"Emotional state: {max(ep['emotion'] for ep in group)}",
                f"Mean outcome value: {sum(hash(ep['outcome']) for ep in group)/len(group)}"
            ]
            
            last_ep = group[-1]
            self.summaries.append(ConsolidatedSummary(
                title=f"{agent} pattern: {task_type}",
                concepts=concepts,
                insights=insights,
                emotion=max(ep['emotion'] for ep in group),
                timestamp=last_ep['timestamp']
            ))

# Auto-initialize
memory = EpisodicMemory()
print("🌀 OMEGA Memory initialized")