# Short-Term (Working) Memory Example
class ShortTermMemory:
    def __init__(self, max_length=10):
        self.memory = []
        self.max_length = max_length
    
    def add(self, item):
        if len(self.memory) >= self.max_length:
            self.memory.pop(0)  # Remove oldest
        self.memory.append(item)
    
    def get_context(self):
        return self.memory

# Usage
stm = ShortTermMemory(max_length=5)
for i in range(10):
    stm.add(f"observation {i}")

print("Current Context:", stm.get_context())

# Episodic Memory Example
class EpisodicMemory:
    def __init__(self):
        self.episodes = []
    
    def add_episode(self, interaction):
        self.episodes.append(interaction)
    
    def recall_latest(self, n=1):
        return self.episodes[-n:]

# Usage
episodic = EpisodicMemory()
episodic.add_episode({"user": "Hi", "agent": "Hello!"})
episodic.add_episode({"user": "What's the weather?", "agent": "Sunny."})

print("Latest Episode:", episodic.recall_latest(1))

# Semantic Memory Example
class SemanticMemory:
    def __init__(self):
        self.knowledge_base = {}
    
    def add_fact(self, key, value):
        self.knowledge_base[key] = value
    
    def query_fact(self, key):
        return self.knowledge_base.get(key, "Unknown")

# Usage
semantic = SemanticMemory()
semantic.add_fact("Paris", "Capital of France")
semantic.add_fact("Sun", "A star at the center of the solar system")

print("Query Paris:", semantic.query_fact("Paris"))

# Procedural Memory Example
class ProceduralMemory:
    def __init__(self):
        self.skills = {}
    
    def add_skill(self, name, func):
        self.skills[name] = func
    
    def execute_skill(self, name, *args, **kwargs):
        if name in self.skills:
            return self.skills[name](*args, **kwargs)
        else:
            return "Skill not found."

# Example skill
def add_numbers(a, b):
    return a + b

# Usage
procedural = ProceduralMemory()
procedural.add_skill("addition", add_numbers)

print("Execute Skill:", procedural.execute_skill("addition", 5, 7))

# Reflective Memory Example
import time

class ReflectiveMemory:
    def __init__(self):
        self.memory_log = []
    
    def add_memory(self, content):
        entry = {"content": content, "timestamp": time.time(), "validated": False}
        self.memory_log.append(entry)
    
    def validate_memory(self, index):
        self.memory_log[index]["validated"] = True
    
    def fetch_validated(self):
        return [m for m in self.memory_log if m["validated"]]

# Usage
reflective = ReflectiveMemory()
reflective.add_memory("Saw a red car.")
reflective.validate_memory(0)

print("Validated Memories:", reflective.fetch_validated())

class ContextWindow:
    def __init__(self, max_tokens=1024):
        self.history = []
        self.max_tokens = max_tokens
        self.tokenizer = lambda x: x.split()  # Simplistic tokenizer
    
    def add(self, text):
        self.history.append(text)
        # Truncate if exceeds max tokens
        while self.total_tokens() > self.max_tokens:
            self.history.pop(0)
    
    def total_tokens(self):
        return sum(len(self.tokenizer(h)) for h in self.history)
    
    def get_context(self):
        return "\n".join(self.history)

# Usage
context = ContextWindow(max_tokens=50)
context.add("User: Hi")
context.add("Agent: Hello, how can I help?")
context.add("User: Tell me about Paris.")

print("Context Window:\n", context.get_context())

import faiss
import numpy as np

class VectorMemory:
    def __init__(self, dim=128):
        self.index = faiss.IndexFlatL2(dim)
        self.embeddings = []
        self.data = []
    
    def add_memory(self, embedding, content):
        self.index.add(np.array([embedding]).astype(np.float32))
        self.embeddings.append(embedding)
        self.data.append(content)
    
    def search(self, query_embedding, k=1):
        distances, indices = self.index.search(np.array([query_embedding]).astype(np.float32), k)
        return [self.data[idx] for idx in indices[0]]

# Usage
np.random.seed(0)
memory = VectorMemory(dim=128)
memory.add_memory(np.random.rand(128), "Memory 1: About Paris")
memory.add_memory(np.random.rand(128), "Memory 2: About Rome")

query = np.random.rand(128)
print("Nearest Memories:", memory.search(query, k=1))

import torch
import torch.nn as nn

class DifferentiableMemory(nn.Module):
    def __init__(self, memory_size, key_dim, value_dim):
        super().__init__()
        self.keys = nn.Parameter(torch.randn(memory_size, key_dim))
        self.values = nn.Parameter(torch.randn(memory_size, value_dim))
    
    def forward(self, query):
        similarities = torch.matmul(query, self.keys.T)  # [batch, memory_size]
        attn_weights = torch.softmax(similarities, dim=-1)
        readout = torch.matmul(attn_weights, self.values)
        return readout

# Usage
dmemory = DifferentiableMemory(memory_size=100, key_dim=64, value_dim=256)
query = torch.randn(1, 64)
output = dmemory(query)

print("Readout from Differentiable Memory:", output.shape)

class HybridMemoryAgent:
    def __init__(self, memory_system, generator_model):
        self.memory = memory_system
        self.generator = generator_model  # Any generative model like a language model

    def respond(self, query_embedding):
        relevant_memories = self.memory.search(query_embedding, k=5)
        prompt = " ".join(relevant_memories)
        response = self.generator.generate(prompt)
        return response

# Mock generator
class SimpleGenerator:
    def generate(self, text):
        return f"Generated based on: {text}"

# Usage
hybrid_agent = HybridMemoryAgent(memory_system=memory, generator_model=SimpleGenerator())
print("Agent Response:", hybrid_agent.respond(np.random.rand(128)))

import openai

class OpenAIGenerator:
    def __init__(self, model_name="gpt-4", temperature=0.3):
        self.model_name = model_name
        self.temperature = temperature

    def generate(self, prompt):
        response = openai.ChatCompletion.create(
            model=self.model_name,
            messages=[{"role": "system", "content": "You are a helpful, intelligent AI assistant."},
                      {"role": "user", "content": prompt}],
            temperature=self.temperature,
            max_tokens=500,
        )
        return response['choices'][0]['message']['content']

import numpy as np
import time

class MemoryAugmentedAgent:
    def __init__(self):
        self.short_term = ShortTermMemory(max_length=10)
        self.episodic = EpisodicMemory()
        self.semantic = SemanticMemory()
        self.procedural = ProceduralMemory()
        self.reflective = ReflectiveMemory()
        self.vector_memory = VectorMemory(dim=128)
        self.generator = OpenAIGenerator(model_name="gpt-4")

    def get_embedding(self, text):
        """Fake deterministic embedding for simplicity."""
        np.random.seed(hash(text) % 10000)
        return np.random.rand(128)

    def process_query(self, user_input):
        # Step 1: Update working memory
        self.short_term.add(user_input)

        # Step 2: Recall past episodes
        recent_interactions = self.episodic.recall_latest(3)

        # Step 3: Semantic retrieval
        query_emb = self.get_embedding(user_input)
        retrieved_facts = self.vector_memory.search(query_emb, k=2)

        # Step 4: Attempt skill execution
        skill_response = None
        if "add" in user_input:
            numbers = [int(s) for s in user_input.split() if s.isdigit()]
            if len(numbers) >= 2:
                skill_response = self.procedural.execute_skill("addition", numbers[0], numbers[1])

        # Step 5: Reflective memory lookup
        validated_memories = self.reflective.fetch_validated()

        # Step 6: Compose prompt
        prompt = f"""
        Context:
        - Recent conversations: {recent_interactions}
        - Retrieved facts: {retrieved_facts}
        - Validated memories: {validated_memories}
        - Skill execution output: {skill_response}
        
        Current user query:
        {user_input}
        """

        # Step 7: Get model response
        final_response = self.generator.generate(prompt)

        # Step 8: Update memories
        self.episodic.add_episode({"user": user_input, "agent": final_response})
        self.reflective.add_memory(f"Interacted about: {user_input}")

        return final_response

    def consolidate_memory(self):
        """Summarizes old episodic memories into a single entry."""
        if len(self.episodic.episodes) < 5:
            return  # Consolidate only if needed

        summary_prompt = "Summarize the following conversation history:\n"
        for ep in self.episodic.episodes:
            summary_prompt += f"User: {ep['user']}\nAgent: {ep['agent']}\n"

        summary = self.generator.generate(summary_prompt)
        self.episodic.episodes = [{"user": "Summary", "agent": summary}]
        print("\n[Memory Consolidated] New episodic summary created.")

    def learn_skill(self, skill_name, skill_definition):
        """
        Dynamically add a new skill at runtime.
        skill_definition should be a lambda expression in string form.
        """
        try:
            new_skill = eval(skill_definition)
            self.procedural.add_skill(skill_name, new_skill)
            return f"Skill '{skill_name}' learned successfully."
        except Exception as e:
            return f"Skill learning failed: {str(e)}"

# Initialize agent
agent = MemoryAugmentedAgent()

# Seed basic knowledge
agent.semantic.add_fact("Eiffel Tower", "A famous monument in Paris.")
agent.vector_memory.add_memory(agent.get_embedding("Eiffel Tower"), "It is located in Paris.")
agent.procedural.add_skill("addition", lambda x, y: f"The sum is {x + y}")
agent.reflective.add_memory("Initialized memory systems.")
agent.reflective.validate_memory(0)

# Example conversation
queries = [
    "Tell me about the Eiffel Tower.",
    "What is 5 plus 6?",
    "Teach you a new skill: multiply two numbers. Skill: lambda x, y: f'Multiplication is {x * y}'",
    "What's 7 times 8?",
    "Summarize everything we discussed."
]

for q in queries:
    # If query teaches a new skill
    if "Teach you a new skill" in q:
        parts = q.split("Skill:")
        skill_name = "multiplication"
        skill_code = parts[1].strip()
        print(agent.learn_skill(skill_name, skill_code))
    else:
        response = agent.process_query(q)
        print(f"\nUser: {q}\nAgent: {response}")

# Consolidate memory after interaction
agent.consolidate_memory()