purpose_agent/orchestrator.py

"""
Orchestrator — The main loop tying Actor, Purpose Function, Experience Replay,
and Heuristic Optimizer together.

Implements the self-improvement loop:

    ┌─────────────────────────────────────────────────────────────────┐
    │                     ORCHESTRATOR LOOP                          │
    │                                                                 │
    │  ┌──────────┐   action   ┌─────────────┐   s_new              │
    │  │  ACTOR   │ ────────►  │ ENVIRONMENT │ ──────────┐          │
    │  │(+memory) │            │ (your code) │           │          │
    │  └────▲─────┘            └─────────────┘           │          │
    │       │                                             ▼          │
    │       │  heuristics    ┌────────────────┐   (s, a, s')        │
    │       │◄───────────────│   OPTIMIZER    │◄─────────┐          │
    │       │                │ (distillation) │          │          │
    │       │                └────────────────┘          │          │
    │       │                                             │          │
    │       │                ┌────────────────┐   Φ(s)→Φ(s')       │
    │       │                │   PURPOSE FN   │──────────┤          │
    │       │                │ (state critic) │          │          │
    │       │                └────────────────┘          │          │
    │       │                                             │          │
    │       │                ┌────────────────┐          │          │
    │       └────────────────│ EXPERIENCE     │◄─────────┘          │
    │                        │ REPLAY BUFFER  │                      │
    │                        └────────────────┘                      │
    └─────────────────────────────────────────────────────────────────┘

Usage:
    from purpose_agent import Orchestrator, MockLLMBackend

    # 1. Define your environment
    class MyEnv(Environment):
        def execute(self, action, current_state):
            # ... do something ...
            return new_state

    # 2. Create orchestrator
    orch = Orchestrator(
        llm=MockLLMBackend(),  # or HFInferenceBackend(), OpenAICompatibleBackend()
        environment=MyEnv(),
        available_actions={"search": "Search for items", "move": "Move to location"},
    )

    # 3. Run a task
    result = orch.run_task(
        purpose="Find the hidden treasure in the maze",
        initial_state=State(data={"position": [0, 0], "inventory": []}),
        max_steps=20,
    )

    # 4. The agent self-improves — run more tasks and it gets better
    result2 = orch.run_task(purpose="Find the second treasure", ...)
"""

from __future__ import annotations

import json
import logging
import time
from abc import ABC, abstractmethod
from typing import Any, Callable

from purpose_agent.types import (
    Action,
    Heuristic,
    MemoryTier,
    PurposeScore,
    State,
    Trajectory,
    TrajectoryStep,
)
from purpose_agent.actor import Actor
from purpose_agent.purpose_function import PurposeFunction
from purpose_agent.experience_replay import ExperienceReplay
from purpose_agent.optimizer import HeuristicOptimizer
from purpose_agent.llm_backend import LLMBackend

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Environment Interface
# ---------------------------------------------------------------------------

class Environment(ABC):
    """
    Abstract environment that the Agent acts in.
    
    Implement this for your specific use case:
    - Web navigation: wrap a browser automation tool
    - Code generation: wrap a code executor
    - Game: wrap a game API
    - Simulated: mock environment for testing
    
    The Orchestrator calls execute() with the agent's action and current state,
    and expects a new state back.
    """

    @abstractmethod
    def execute(self, action: Action, current_state: State) -> State:
        """
        Execute an action in the environment and return the resulting state.
        
        Args:
            action: The action to execute
            current_state: The state before the action
            
        Returns:
            The new state after the action
        """
        ...

    def reset(self) -> State:
        """
        Reset the environment and return the initial state.
        Override if your environment needs resetting between tasks.
        """
        return State(data={})

    def is_terminal(self, state: State) -> bool:
        """
        Check if the state is terminal (task complete or impossible to continue).
        Override for environments with natural termination conditions.
        """
        return False


class SimpleEnvironment(Environment):
    """
    A simple environment backed by a user-provided execute function.
    
    Usage:
        env = SimpleEnvironment(
            execute_fn=lambda action, state: new_state,
            initial_state=State(data={"x": 0}),
        )
    """

    def __init__(
        self,
        execute_fn: Callable[[Action, State], State],
        initial_state: State | None = None,
        terminal_fn: Callable[[State], bool] | None = None,
    ):
        self._execute_fn = execute_fn
        self._initial_state = initial_state or State(data={})
        self._terminal_fn = terminal_fn

    def execute(self, action: Action, current_state: State) -> State:
        return self._execute_fn(action, current_state)

    def reset(self) -> State:
        return self._initial_state

    def is_terminal(self, state: State) -> bool:
        if self._terminal_fn:
            return self._terminal_fn(state)
        return False


# ---------------------------------------------------------------------------
# Task Result
# ---------------------------------------------------------------------------

class TaskResult:
    """Result of running a task through the Orchestrator."""

    def __init__(self, trajectory: Trajectory, final_state: State):
        self.trajectory = trajectory
        self.final_state = final_state

    @property
    def success(self) -> bool:
        """Was the task successful? (final Φ > 7.0)"""
        phi = self.trajectory.final_phi
        return phi is not None and phi > 7.0

    @property
    def total_steps(self) -> int:
        return len(self.trajectory.steps)

    @property
    def cumulative_reward(self) -> float:
        return self.trajectory.cumulative_reward

    @property
    def final_phi(self) -> float | None:
        return self.trajectory.final_phi

    def summary(self) -> str:
        lines = [
            f"Task: {self.trajectory.task_description}",
            f"Purpose: {self.trajectory.purpose}",
            f"Steps: {self.total_steps}",
            f"Success Rate: {self.trajectory.success_rate:.1%}",
            f"Cumulative Reward: {self.cumulative_reward:.2f}",
            f"Net Delta: {self.trajectory.total_delta:.2f}",
            f"Final Φ: {self.final_phi:.2f}" if self.final_phi is not None else "Final Φ: N/A",
            f"Task Success: {'✓' if self.success else '✗'}",
        ]
        return "\n".join(lines)


# ---------------------------------------------------------------------------
# Orchestrator
# ---------------------------------------------------------------------------

class Orchestrator:
    """
    Main orchestration loop for the self-improving agent.
    
    Ties together all modules:
    - Actor: Decides actions based on state + memory
    - Purpose Function: Scores state transitions (Φ improvement)
    - Experience Replay: Stores trajectories for future retrieval
    - Heuristic Optimizer: Extracts winning strategies from good trajectories
    
    Self-improvement happens via the memory feedback loop:
    1. Actor uses heuristics from memory to decide actions
    2. Purpose Function scores each transition
    3. Experience Replay stores the full trajectory
    4. Optimizer distills high-reward trajectories into new heuristics
    5. Actor's memory is updated with new heuristics → better next time
    
    Args:
        llm: Default LLM backend (used for all modules unless overridden)
        critic_llm: Optional separate LLM for the Purpose Function
        optimizer_llm: Optional separate LLM for the Optimizer
        environment: The environment the agent acts in
        available_actions: Dict of {action_name: description}
        experience_buffer_size: Max trajectories in experience replay
        persistence_dir: Directory for persistent storage (experience replay, heuristics)
        on_step: Optional callback called after each step (for monitoring)
    """

    def __init__(
        self,
        llm: LLMBackend,
        environment: Environment,
        available_actions: dict[str, str] | None = None,
        critic_llm: LLMBackend | None = None,
        optimizer_llm: LLMBackend | None = None,
        experience_buffer_size: int = 500,
        persistence_dir: str | None = None,
        on_step: Callable[[TrajectoryStep], None] | None = None,
        optimize_every_n_tasks: int = 1,
    ):
        self.environment = environment
        self.on_step = on_step
        self.optimize_every_n_tasks = optimize_every_n_tasks
        self._tasks_since_optimize = 0

        # Persistence
        replay_path = None
        if persistence_dir:
            import os
            os.makedirs(persistence_dir, exist_ok=True)
            replay_path = f"{persistence_dir}/experience_replay.json"

        # Initialize modules
        self.actor = Actor(
            llm=llm,
            available_actions=available_actions,
        )
        self.purpose_fn = PurposeFunction(
            llm=critic_llm or llm,
        )
        self.experience_replay = ExperienceReplay(
            capacity=experience_buffer_size,
            persistence_path=replay_path,
        )
        self.optimizer = HeuristicOptimizer(
            llm=optimizer_llm or llm,
        )

        # Load existing heuristics into Actor memory
        self.sync_memory()

    # ------------------------------------------------------------------
    # Main Task Loop
    # ------------------------------------------------------------------

    def run_task(
        self,
        purpose: str,
        initial_state: State | None = None,
        max_steps: int = 20,
        early_stop_phi: float = 9.0,
        task_description: str | None = None,
    ) -> TaskResult:
        """
        Run a complete task through the agent loop.
        
        The loop for each step:
        1. Actor decides an action (with thought + prediction)
        2. Environment executes the action → new state
        3. Purpose Function evaluates: Φ(s_new) vs Φ(s_old)
        4. Trajectory step is recorded
        5. Check termination conditions
        
        After the task:
        - Trajectory is added to Experience Replay
        - If enough tasks have run, Optimizer extracts new heuristics
        - Actor's memory is updated
        
        Args:
            purpose: The goal description
            initial_state: Starting state (or environment.reset() if None)
            max_steps: Maximum steps before forced termination
            early_stop_phi: Stop if Φ exceeds this value (goal ~achieved)
            task_description: Optional description (defaults to purpose)
        """
        task_desc = task_description or purpose
        current_state = initial_state or self.environment.reset()

        # Reset Purpose Function per-trajectory stats
        self.purpose_fn.reset_trajectory_stats()

        # Retrieve relevant past experiences for context
        relevant_experiences = self.experience_replay.retrieve(task_desc, top_k=3)
        self._inject_experience_context(relevant_experiences)

        # Create trajectory
        trajectory = Trajectory(
            task_description=task_desc,
            purpose=purpose,
        )

        # History for Actor context
        history: list[dict[str, Any]] = []

        logger.info(f"═══ Starting task: {task_desc} (max {max_steps} steps) ═══")

        for step_idx in range(max_steps):
            step_start = time.time()

            # Step 1: Actor decides
            action = self.actor.decide(
                purpose=purpose,
                current_state=current_state,
                history=history,
            )

            logger.info(
                f"Step {step_idx + 1}: Action={action.name}, "
                f"Thought={action.thought[:100]}..."
            )

            # Check for DONE action
            if action.name.upper() == "DONE":
                logger.info("Agent signaled DONE — ending task")
                # Still score the final state to record final Φ
                final_score = self.purpose_fn.evaluate(
                    state_before=current_state,
                    action=action,
                    state_after=current_state,
                    purpose=purpose,
                )
                trajectory.steps.append(TrajectoryStep(
                    state_before=current_state,
                    action=action,
                    state_after=current_state,
                    score=final_score,
                    step_index=step_idx + 1,
                    wall_time_s=time.time() - step_start,
                ))
                break

            # Step 2: Environment executes
            try:
                new_state = self.environment.execute(action, current_state)
            except Exception as e:
                logger.error(f"Environment execution failed: {e}")
                new_state = State(
                    data={**current_state.data, "_error": str(e)},
                    summary=f"Error: {e}",
                )

            # Step 3: Purpose Function evaluates
            score = self.purpose_fn.evaluate(
                state_before=current_state,
                action=action,
                state_after=new_state,
                purpose=purpose,
            )

            # Step 4: Record step
            step = TrajectoryStep(
                state_before=current_state,
                action=action,
                state_after=new_state,
                score=score,
                step_index=step_idx + 1,
                wall_time_s=time.time() - step_start,
            )
            trajectory.steps.append(step)

            # Update history for Actor context
            history.append({
                "action": f"{action.name}({json.dumps(action.params, default=str)})",
                "result": new_state.describe()[:200],
                "score": f"Δ={score.delta:+.2f}" if score else "N/A",
            })

            # Callback
            if self.on_step:
                self.on_step(step)

            logger.info(
                f"  → Φ: {score.phi_before:.1f} → {score.phi_after:.1f} "
                f"(Δ={score.delta:+.2f}, conf={score.confidence:.2f})"
            )

            # Step 5: Check termination
            if score.phi_after >= early_stop_phi:
                logger.info(f"Early stop: Φ={score.phi_after:.1f} ≥ {early_stop_phi}")
                break

            if self.environment.is_terminal(new_state):
                logger.info("Environment signaled terminal state")
                break

            current_state = new_state

        # Post-task processing
        result = TaskResult(trajectory=trajectory, final_state=current_state)
        self.post_task(trajectory, relevant_experiences)

        logger.info(f"═══ Task complete ═══\n{result.summary()}")
        return result

    # ------------------------------------------------------------------
    # Post-Task: Experience Storage + Optimization
    # ------------------------------------------------------------------

    def post_task(
        self,
        trajectory: Trajectory,
        used_experiences: list[Any] | None = None,
    ) -> None:
        """Post-task processing: store trajectory, maybe optimize, sync memory.
        
        Public API — called by HITLOrchestrator, AsyncOrchestrator, and
        any custom orchestration wrapper after a task completes.
        """
        used_experiences = used_experiences or []

        # Store in experience replay
        record = self.experience_replay.add(trajectory)

        # Update Q-values for retrieved experiences that were used
        task_success = trajectory.success_rate > 0.5
        for exp in used_experiences:
            self.experience_replay.update_q_value(
                exp.id, reward=1.0 if task_success else 0.0
            )

        # Update heuristic usage stats
        for h in self.actor.strategic_memory + self.actor.procedural_memory:
            self.optimizer.update_heuristic_usage(h.id, was_successful=task_success)

        # Periodic optimization
        self._tasks_since_optimize += 1
        if self._tasks_since_optimize >= self.optimize_every_n_tasks:
            self._run_optimization()
            self._tasks_since_optimize = 0

    def _run_optimization(self) -> None:
        """Run the heuristic optimization cycle."""
        logger.info("Running optimization cycle...")

        # Get best trajectories
        top_trajectories = self.experience_replay.get_top_trajectories(
            n=5, min_success_rate=0.3
        )

        if not top_trajectories:
            logger.info("No qualifying trajectories for optimization")
            return

        # Run optimizer
        self.optimizer.optimize(top_trajectories)

        # Sync updated heuristics to Actor memory
        self.sync_memory()

    def sync_memory(self) -> None:
        """Push current heuristic library to Actor's memory tiers.
        
        Public API — call after manually modifying the heuristic library
        (e.g., human-injected heuristics via HITL).
        """
        self.actor.update_strategic_memory(
            self.optimizer.get_heuristics_by_tier(MemoryTier.STRATEGIC)
        )
        self.actor.update_procedural_memory(
            self.optimizer.get_heuristics_by_tier(MemoryTier.PROCEDURAL)
        )

        # Tool memory from heuristics
        tool_heuristics = self.optimizer.get_heuristics_by_tier(MemoryTier.TOOL)
        tool_tips = {h.pattern: h.strategy for h in tool_heuristics}
        if tool_tips:
            self.actor.update_tool_memory(tool_tips)

    def _inject_experience_context(self, experiences: list[Any]) -> None:
        """
        Inject retrieved experience context into Actor's procedural memory.
        
        This is the CER (arxiv:2506.06698) retrieval injection pattern:
        relevant past trajectories → distilled into SOPs → added to Actor context.
        """
        injected = []
        for exp in experiences:
            for h in exp.heuristics:
                if h.tier == MemoryTier.PROCEDURAL:
                    injected.append(h)

        if injected:
            current = self.actor.procedural_memory or []
            self.actor.procedural_memory = current + injected
            logger.debug(f"Injected {len(injected)} experience-based SOPs")

    # ------------------------------------------------------------------
    # Inspection / Monitoring
    # ------------------------------------------------------------------

    @property
    def stats(self) -> dict[str, Any]:
        """Get current framework statistics."""
        return {
            "experience_replay": self.experience_replay.stats,
            "heuristic_library_size": len(self.optimizer.heuristic_library),
            "heuristics_by_tier": {
                tier.value: len(self.optimizer.get_heuristics_by_tier(tier))
                for tier in MemoryTier
            },
            "tasks_since_optimize": self._tasks_since_optimize,
        }

    def get_heuristic_report(self) -> str:
        """Human-readable report of all learned heuristics."""
        lines = ["═══ Learned Heuristics Report ═══\n"]

        for tier in MemoryTier:
            heuristics = self.optimizer.get_heuristics_by_tier(tier)
            lines.append(f"\n{'─' * 40}")
            lines.append(f"  {tier.value.upper()} ({len(heuristics)} heuristics)")
            lines.append(f"{'─' * 40}")

            for h in heuristics:
                lines.append(f"\n  [{h.id}] Q={h.q_value:.3f} (used {h.times_used}x, "
                             f"{h.times_succeeded} successes)")
                lines.append(f"  Pattern:  {h.pattern}")
                lines.append(f"  Strategy: {h.strategy}")
                if h.steps:
                    for i, step in enumerate(h.steps, 1):
                        lines.append(f"    {i}. {step}")

        return "\n".join(lines)