persistence.py

"""
Persistence layer for saving/loading evolved DSL functions
"""

import json
import pickle
import os
from typing import Dict, Any, List
from dataclasses import asdict
import importlib.util

from dsl import DSL, DSLFunction, DSLType

class DSLPersistence:
    """Handle saving and loading DSL state"""
    
    def __init__(self, storage_dir: str = "dsl_storage"):
        self.storage_dir = storage_dir
        os.makedirs(storage_dir, exist_ok=True)
    
    def save_dsl(self, dsl: DSL, name: str = "current_dsl") -> str:
        """Save DSL to disk"""
        filepath = os.path.join(self.storage_dir, f"{name}.json")
        
        # Serialize DSL functions (without implementation objects)
        dsl_data = {
            "functions": {},
            "metadata": {
                "total_functions": len(dsl.functions),
                "primitive_count": 0,
                "evolved_count": 0
            }
        }
        
        for name, func in dsl.functions.items():
            # Check if it's a primitive (has implementation but no body)
            is_primitive = func.implementation is not None and func.body is None
            
            dsl_data["functions"][name] = {
                "name": func.name,
                "params": func.params,
                "param_types": [t.value for t in func.param_types],
                "return_type": func.return_type.value,
                "body": func.body,
                "fitness_score": func.fitness_score,
                "usage_count": func.usage_count,
                "is_primitive": is_primitive
            }
            
            if is_primitive:
                dsl_data["metadata"]["primitive_count"] += 1
            else:
                dsl_data["metadata"]["evolved_count"] += 1
        
        with open(filepath, 'w') as f:
            json.dump(dsl_data, f, indent=2)
        
        print(f"💾 DSL saved to {filepath}")
        return filepath
    
    def load_dsl(self, name: str = "current_dsl") -> DSL:
        """Load DSL from disk"""
        filepath = os.path.join(self.storage_dir, f"{name}.json")
        
        if not os.path.exists(filepath):
            print(f"❌ DSL file {filepath} not found, creating new DSL")
            return DSL()
        
        with open(filepath, 'r') as f:
            dsl_data = json.load(f)
        
        # Create new DSL (this initializes primitives)
        dsl = DSL()
        
        # Add evolved functions (non-primitives)
        for func_name, func_data in dsl_data["functions"].items():
            if not func_data.get("is_primitive", False):
                # Reconstruct DSL function
                param_types = [DSLType(t) for t in func_data["param_types"]]
                return_type = DSLType(func_data["return_type"])
                
                evolved_func = DSLFunction(
                    name=func_data["name"],
                    params=func_data["params"],
                    param_types=param_types,
                    return_type=return_type,
                    body=func_data["body"],
                    fitness_score=func_data["fitness_score"],
                    usage_count=func_data["usage_count"]
                )
                
                # Try to reconstruct implementation from body
                if func_data["body"]:
                    try:
                        evolved_func.implementation = self._create_implementation(
                            evolved_func, dsl
                        )
                    except Exception as e:
                        print(f"⚠️  Could not reconstruct implementation for {func_name}: {e}")
                
                dsl.add_function(evolved_func)
        
        print(f"📂 DSL loaded from {filepath}")
        print(f"   Primitives: {dsl_data['metadata']['primitive_count']}")
        print(f"   Evolved: {dsl_data['metadata']['evolved_count']}")
        
        return dsl
    
    def _create_implementation(self, func: DSLFunction, dsl: DSL):
        """Recreate function implementation from code body"""
        if not func.body:
            return None
        
        # Create namespace with DSL functions
        namespace = {"__builtins__": {}}
        namespace.update({name: f for name, f in dsl.functions.items()})
        
        # Execute the function definition
        exec(func.body, namespace)
        
        # Return the function object
        return namespace[func.name]
    
    def export_to_python(self, dsl: DSL, filename: str = None) -> str:
        """Export evolved functions as standalone Python code"""
        if filename is None:
            filename = os.path.join(self.storage_dir, "evolved_functions.py")
        
        code_lines = [
            "# Automatically generated evolved functions",
            "# Generated by MCTS + Evolution Coding Agent",
            "",
            "# Primitive functions (manually implemented)",
        ]
        
        # Add primitive function implementations
        primitives = {
            "add": "def add(x, y): return x + y",
            "sub": "def sub(x, y): return x - y", 
            "mul": "def mul(x, y): return x * y",
            "div": "def div(x, y): return x // y if y != 0 else 0",
            "eq": "def eq(x, y): return x == y",
            "lt": "def lt(x, y): return x < y",
            "gt": "def gt(x, y): return x > y",
            "if_then_else": "def if_then_else(cond, then_val, else_val): return then_val if cond else else_val",
            "identity": "def identity(x): return x"
        }
        
        for name, implementation in primitives.items():
            code_lines.append(implementation)
        
        code_lines.extend(["", "# Evolved functions"])
        
        # Add evolved functions
        for name, func in dsl.functions.items():
            if func.body and name not in primitives:
                code_lines.append("")
                code_lines.append(f"# Fitness: {func.fitness_score:.3f}")
                code_lines.append(func.body)
        
        # Write to file
        with open(filename, 'w') as f:
            f.write('\n'.join(code_lines))
        
        print(f"🐍 Python code exported to {filename}")
        return filename
    
    def list_saved_dsls(self) -> List[str]:
        """List all saved DSL files"""
        json_files = [f[:-5] for f in os.listdir(self.storage_dir) 
                     if f.endswith('.json')]
        return json_files
    
    def backup_dsl(self, dsl: DSL, cycle_id: int) -> str:
        """Create a backup of DSL after each cycle"""
        backup_name = f"dsl_cycle_{cycle_id}"
        return self.save_dsl(dsl, backup_name)

class SessionManager:
    """Manage DSL evolution sessions"""
    
    def __init__(self, session_name: str = "default"):
        self.session_name = session_name
        self.persistence = DSLPersistence(f"sessions/{session_name}")
        self.session_log = []
    
    def start_session(self, resume: bool = True) -> DSL:
        """Start or resume a session"""
        if resume:
            dsl = self.persistence.load_dsl("latest")
            print(f"🔄 Resumed session '{self.session_name}'")
        else:
            dsl = DSL()
            print(f"🆕 Started new session '{self.session_name}'")
        
        return dsl
    
    def save_cycle(self, dsl: DSL, cycle_id: int, summary: Dict[str, Any]):
        """Save DSL state and cycle summary"""
        # Save DSL
        self.persistence.save_dsl(dsl, "latest")
        self.persistence.backup_dsl(dsl, cycle_id)
        
        # Log cycle summary
        self.session_log.append({
            "cycle_id": cycle_id,
            "timestamp": summary.get("timestamp"),
            "functions_added": summary.get("new_functions", 0),
            "total_functions": len(dsl.functions),
            "best_fitness": summary.get("best_fitness", 0.0)
        })
        
        # Save session log
        log_file = os.path.join(self.persistence.storage_dir, "session_log.json")
        with open(log_file, 'w') as f:
            json.dump(self.session_log, f, indent=2)
        
        print(f"💾 Cycle {cycle_id} saved to session '{self.session_name}'")
    
    def get_session_summary(self) -> Dict[str, Any]:
        """Get summary of current session"""
        if not self.session_log:
            return {"cycles": 0, "total_functions": 0}
        
        latest = self.session_log[-1]
        return {
            "session_name": self.session_name,
            "cycles_completed": len(self.session_log),
            "total_functions": latest["total_functions"],
            "best_fitness_achieved": max(log["best_fitness"] for log in self.session_log),
            "functions_discovered": sum(log["functions_added"] for log in self.session_log)
        }

# Example usage integration
def create_persistent_system(session_name: str = "my_evolution"):
    """Create a bootstrap system with persistence"""
    
    # Create session manager
    session = SessionManager(session_name)
    
    # Load or create DSL
    dsl = session.start_session(resume=True)
    
    return session, dsl

# Usage in main system
def persistent_bootstrap_cycle(session: SessionManager, dsl: DSL, cycle_id: int):
    """Run bootstrap cycle with automatic saving"""
    from mcts_gpt4o import GPT4BootstrapSystem
    
    # Run cycle
    system = GPT4BootstrapSystem(initial_dsl=dsl)
    summary = system.run_bootstrap_cycle_async(cycle_id=cycle_id)
    
    # Save results
    session.save_cycle(dsl, cycle_id, summary)
    
    return summary