batch_executor.py

"""Batch and parallel execution strategies for simulations.

Separates:

A) SINGLE SIMULATION MODE
   One mission at a time, sequential execution
   Use for: Single trajectory, detailed debugging, small runs
   
B) BATCH/PARALLEL MODE
   Multiple independent simulations
   Use for: Monte Carlo, parameter sweeps, optimization
   Backends: multiprocessing, Ray, MPI

Architecture:
  Batch Manager ─→ Task Queue ─→ Worker Pool
      ↓                               ↓
   Config       [Task 1] ────→ [Worker 1]
   Generator    [Task 2] ────→ [Worker 2]
                [Task 3] ────→ [Worker 3]
                     ↓
                  Results Database
"""

from dataclasses import dataclass
from typing import Callable, List, Dict, Any, Optional, Iterable
from pathlib import Path
import json
import logging


logger = logging.getLogger(__name__)


# ============================================================================
# Batch Task Definitions
# ============================================================================

@dataclass
class SimulationTask:
    """One independent simulation task."""
    
    task_id: int
    mission_name: str
    config: Dict[str, Any]
    solver_config: Optional[Dict[str, Any]] = None
    random_seed: Optional[int] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Serialize to dict."""
        return {
            "task_id": self.task_id,
            "mission_name": self.mission_name,
            "config": self.config,
            "solver_config": self.solver_config,
            "random_seed": self.random_seed,
        }


# ============================================================================
# Parameter Space Generators
# ============================================================================

class ParameterSwGenerator:
    """Generate parameter sweep tasks (Cartesian product)."""
    
    @staticmethod
    def generate(
        base_config: Dict[str, Any],
        parameter_ranges: Dict[str, List[Any]],
        mission_name: str = "param_sweep",
    ) -> Iterable[SimulationTask]:
        """Generate all combinations of parameter values.
        
        Args:
            base_config: Base configuration to modify
            parameter_ranges: Dict mapping param_name -> list of values
            mission_name: Base name for missions
        
        Yields:
            SimulationTask for each combination
        """
        from itertools import product
        
        param_names = list(parameter_ranges.keys())
        param_values = [parameter_ranges[name] for name in param_names]
        
        for task_id, values in enumerate(product(*param_values)):
            config = base_config.copy()
            
            # Set parameters
            for param_name, value in zip(param_names, values):
                config[param_name] = value
            
            name = mission_name + f"_task_{task_id}"
            
            yield SimulationTask(
                task_id=task_id,
                mission_name=name,
                config=config,
            )


class MonteCarloGenerator:
    """Generate random parameter Monte Carlo tasks."""
    
    @staticmethod
    def generate(
        base_config: Dict[str, Any],
        parameter_distributions: Dict[str, tuple],
        num_samples: int,
        mission_name: str = "monte_carlo",
        seed: Optional[int] = None,
    ) -> Iterable[SimulationTask]:
        """Generate random parameter samples.
        
        Args:
            base_config: Base configuration
            parameter_distributions: Dict mapping param_name -> (dist_type, min, max)
                                    dist_type: "uniform", "normal", "lognormal"
            num_samples: Number of random samples
            seed: Random seed
        
        Yields:
            SimulationTask for each sample
        """
        import random
        import numpy as np
        
        if seed is not None:
            random.seed(seed)
            np.random.seed(seed)
        
        for task_id in range(num_samples):
            config = base_config.copy()
            sample_seed = seed + task_id if seed else None
            
            # Generate random parameters
            for param_name, (dist_type, min_val, max_val) in parameter_distributions.items():
                if dist_type == "uniform":
                    config[param_name] = random.uniform(min_val, max_val)
                elif dist_type == "normal":
                    config[param_name] = np.random.normal(min_val, max_val)
                elif dist_type == "lognormal":
                    config[param_name] = np.random.lognormal(min_val, max_val)
            
            name = mission_name + f"_sample_{task_id}"
            
            yield SimulationTask(
                task_id=task_id,
                mission_name=name,
                config=config,
                random_seed=sample_seed,
            )


class SensitivityAnalysisGenerator:
    """Generate tasks for local sensitivity analysis (one-at-a-time)."""
    
    @staticmethod
    def generate(
        base_config: Dict[str, Any],
        parameter_perturbations: Dict[str, float],
        mission_name: str = "sensitivity",
    ) -> Iterable[SimulationTask]:
        """Generate baseline + perturbed variants.
        
        Args:
            base_config: Baseline configuration
            parameter_perturbations: Dict mapping param_name -> perturbation_factor
                                    (e.g., 0.9 = -10%, 1.1 = +10%)
            mission_name: Base name
        
        Yields:
            SimulationTask: baseline + one task per perturbed parameter
        """
        task_id = 0
        
        # Baseline task
        yield SimulationTask(
            task_id=task_id,
            mission_name=mission_name + "_baseline",
            config=base_config.copy(),
        )
        task_id += 1
        
        # Perturbed tasks
        for param_name, factor in parameter_perturbations.items():
            if param_name not in base_config:
                logger.warning(f"Parameter {param_name} not in base config")
                continue
            
            perturbed_config = base_config.copy()
            baseline_val = base_config[param_name]
            perturbed_config[param_name] = baseline_val * factor
            
            yield SimulationTask(
                task_id=task_id,
                mission_name=mission_name + f"_{param_name}_{factor:.2f}",
                config=perturbed_config,
            )
            task_id += 1


# ============================================================================
# Single Simulation Mode
# ============================================================================

class SingleSimulationExecutor:
    """Execute one mission at a time (sequential)."""
    
    @staticmethod
    def run(
        config: Dict[str, Any],
        mission_name: str = "single_mission",
    ) -> Dict[str, Any]:
        """Run single mission synchronously.
        
        Args:
            config: SimulationConfigJSON as dict
            mission_name: Mission name
        
        Returns:
            Dict with keys: mission_id, status, results
        """
        from api.mission import simple_mission
        
        logger.info(f"[SingleSim] Starting: {mission_name}")
        
        # Create and run mission
        mission = simple_mission(mission_name, t_end=config.get("simulation_duration_days", 1.0) * 86400)
        
        try:
            result = mission.run()
            
            return {
                "status": "success",
                "mission_name": mission_name,
                "result": result,
                "summary": result.summary(),
            }
        except Exception as e:
            logger.error(f"[SingleSim] Error: {e}")
            return {
                "status": "failed",
                "mission_name": mission_name,
                "error": str(e),
            }


# ============================================================================
# Batch/Parallel Mode - Multiprocessing Backend
# ============================================================================

class MultiprocessingBatchExecutor:
    """Batch executor using Python multiprocessing."""
    
    @staticmethod
    def worker(task: SimulationTask) -> Dict[str, Any]:
        """Worker function: execute one task (runs in separate process)."""
        from api.mission import simple_mission
        
        logger.info(f"[Worker {task.task_id}] Starting: {task.mission_name}")
        
        try:
            mission = simple_mission(
                task.mission_name,
                t_end=task.config.get("simulation_duration_days", 1.0) * 86400,
            )
            result = mission.run()
            
            return {
                "task_id": task.task_id,
                "status": "success",
                "mission_name": task.mission_name,
                "summary": result.summary(),
            }
        except Exception as e:
            logger.error(f"[Worker {task.task_id}] Error: {e}")
            return {
                "task_id": task.task_id,
                "status": "failed",
                "mission_name": task.mission_name,
                "error": str(e),
            }
    
    @staticmethod
    def run_batch(
        tasks: List[SimulationTask],
        num_workers: Optional[int] = None,
    ) -> List[Dict[str, Any]]:
        """Execute batch of tasks in parallel using multiprocessing.
        
        Args:
            tasks: List of SimulationTask objects
            num_workers: Number of worker processes (default: CPU count)
        
        Returns:
            List of result dicts (one per task)
        """
        from multiprocessing import Pool
        import os
        
        if num_workers is None:
            num_workers = os.cpu_count() or 1
        
        logger.info(f"[MultiprocessingBatch] Spawning {num_workers} workers for {len(tasks)} tasks")
        
        results = []
        with Pool(num_workers) as pool:
            for result in pool.imap_unordered(
                MultiprocessingBatchExecutor.worker,
                tasks,
            ):
                results.append(result)
                logger.info(f"[MultiprocessingBatch] Task {result['task_id']} complete: {result['status']}")
        
        return results


# ============================================================================
# Batch/Parallel Mode - Ray Backend (Optional)
# ============================================================================

class RayBatchExecutor:
    """Batch executor using Ray distributed framework.
    
    Requires: pip install ray
    """
    
    @staticmethod
    def run_batch(
        tasks: List[SimulationTask],
        num_workers: Optional[int] = None,
        ray_address: str = "auto",
    ) -> List[Dict[str, Any]]:
        """Execute batch using Ray.
        
        Args:
            tasks: List of SimulationTask objects
            num_workers: Number of parallel workers
            ray_address: Ray cluster address ("auto" for local)
        
        Returns:
            List of result dicts
        """
        try:
            import ray
        except ImportError:
            raise ImportError("Ray not installed. Install with: pip install ray")
        
        # Initialize Ray
        if not ray.is_initialized():
            ray.init(address=ray_address)
        
        logger.info(f"[Ray] Ray initialized with {num_workers} workers")
        
        # Define remote task function
        @ray.remote
        def ray_worker(task: SimulationTask) -> Dict[str, Any]:
            from api.mission import simple_mission
            
            logger.info(f"[Ray Worker {task.task_id}] Starting: {task.mission_name}")
            
            try:
                mission = simple_mission(
                    task.mission_name,
                    t_end=task.config.get("simulation_duration_days", 1.0) * 86400,
                )
                result = mission.run()
                
                return {
                    "task_id": task.task_id,
                    "status": "success",
                    "mission_name": task.mission_name,
                    "summary": result.summary(),
                }
            except Exception as e:
                logger.error(f"[Ray Worker {task.task_id}] Error: {e}")
                return {
                    "task_id": task.task_id,
                    "status": "failed",
                    "mission_name": task.mission_name,
                    "error": str(e),
                }
        
        # Submit all tasks
        futures = [ray_worker.remote(task) for task in tasks]
        
        # Collect results (in order of completion)
        results = []
        for future in ray.get(futures):
            results.append(future)
            logger.info(f"[Ray] Task {future['task_id']} complete: {future['status']}")
        
        return results


# ============================================================================
# Batch Coordinator
# ============================================================================

class BatchCoordinator:
    """High-level interface for batch simulation execution."""
    
    def __init__(self, backend: str = "multiprocessing", num_workers: Optional[int] = None):
        """Initialize batch coordinator.
        
        Args:
            backend: "multiprocessing", "ray", or "sequential"
            num_workers: Number of parallel workers
        """
        self.backend = backend
        self.num_workers = num_workers
    
    def execute(
        self,
        tasks: Iterable[SimulationTask],
        progress_callback: Optional[Callable[[int, int], None]] = None,
    ) -> List[Dict[str, Any]]:
        """Execute batch of tasks.
        
        Args:
            tasks: Iterable of SimulationTask objects
            progress_callback: Called with (completed, total)
        
        Returns:
            List of result dicts
        """
        # Convert to list if needed
        task_list = list(tasks)
        logger.info(f"[BatchCoordinator] Executing {len(task_list)} tasks using {self.backend}")
        
        if self.backend == "sequential":
            results = self._execute_sequential(task_list, progress_callback)
        elif self.backend == "multiprocessing":
            results = MultiprocessingBatchExecutor.run_batch(task_list, self.num_workers)
        elif self.backend == "ray":
            results = RayBatchExecutor.run_batch(task_list, self.num_workers)
        else:
            raise ValueError(f"Unknown backend: {self.backend}")
        
        return results
    
    @staticmethod
    def _execute_sequential(
        tasks: List[SimulationTask],
        progress_callback: Optional[Callable[[int, int], None]] = None,
    ) -> List[Dict[str, Any]]:
        """Execute tasks sequentially (for testing)."""
        results = []
        for i, task in enumerate(tasks):
            from api.mission import simple_mission
            
            logger.info(f"[Sequential] Task {task.task_id}/{len(tasks)}: {task.mission_name}")
            
            try:
                mission = simple_mission(
                    task.mission_name,
                    t_end=task.config.get("simulation_duration_days", 1.0) * 86400,
                )
                result = mission.run()
                
                results.append({
                    "task_id": task.task_id,
                    "status": "success",
                    "mission_name": task.mission_name,
                    "summary": result.summary(),
                })
            except Exception as e:
                results.append({
                    "task_id": task.task_id,
                    "status": "failed",
                    "mission_name": task.mission_name,
                    "error": str(e),
                })
            
            if progress_callback:
                progress_callback(i + 1, len(tasks))
        
        return results


# ============================================================================
# Batch Results Analysis
# ============================================================================

class BatchResultsAnalyzer:
    """Analyze results from batch execution."""
    
    @staticmethod
    def summary(results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Generate batch execution summary."""
        successful = [r for r in results if r.get("status") == "success"]
        failed = [r for r in results if r.get("status") == "failed"]
        
        return {
            "total_tasks": len(results),
            "successful": len(successful),
            "failed": len(failed),
            "success_rate": len(successful) / len(results) if results else 0.0,
            "results": results,
        }
    
    @staticmethod
    def pareto_front(
        results: List[Dict[str, Any]],
        objectives: Dict[str, str],  # {obj_name: "minimize"|"maximize"}
    ) -> List[Dict[str, Any]]:
        """Extract Pareto front from multi-objective results.
        
        Args:
            results: Batch results
            objectives: Dict mapping objective_name -> direction
        
        Returns:
            List of non-dominated solutions
        """
        successful = [r for r in results if r.get("status") == "success"]
        
        # Simple Pareto front extraction (O(n^2))
        pareto = []
        for candidate in successful:
            summary = candidate.get("summary", {})
            
            # Check if candidate is dominated by any other solution
            is_dominated = False
            for other in successful:
                if candidate == other:
                    continue
                
                other_summary = other.get("summary", {})
                
                # Check if 'other' dominates 'candidate'
                # 'other' dominates if it's better in all objectives
                other_dominates = True
                for obj_name, direction in objectives.items():
                    if obj_name not in summary or obj_name not in other_summary:
                        # Skip objectives not in results
                        continue
                    
                    cand_val = summary[obj_name]
                    other_val = other_summary[obj_name]
                    
                    if direction == "minimize":
                        # For minimization, lower is better
                        # If other_val > cand_val, then candidate is better, so other does NOT dominate
                        if other_val > cand_val:
                            other_dominates = False
                            break
                    else:  # maximize
                        # For maximization, higher is better
                        # If other_val < cand_val, then candidate is better, so other does NOT dominate
                        if other_val < cand_val:
                            other_dominates = False
                            break
                
                if other_dominates:
                    is_dominated = True
                    break
            
            if not is_dominated:
                pareto.append(candidate)
        
        return pareto


# ============================================================================
# Example Usage
# ============================================================================

if __name__ == "__main__":
    # Example: Parameter sweep
    print("\n" + "="*70)
    print("PARAMETER SWEEP EXAMPLE")
    print("="*70)
    
    base_config = {
        "sail_area": 1000.0,
        "reflectivity": 0.90,
        "simulation_duration_days": 1.0,
    }
    
    param_ranges = {
        "sail_area": [500, 1000, 1500, 2000],
        "reflectivity": [0.85, 0.90, 0.95],
    }
    
    tasks = list(ParameterSwGenerator.generate(base_config, param_ranges))
    print(f"Generated {len(tasks)} tasks")
    print(f"First task: {tasks[0]}")
    
    # Example: Monte Carlo
    print("\n" + "="*70)
    print("MONTE CARLO EXAMPLE")
    print("="*70)
    
    distributions = {
        "sail_area": ("uniform", 500, 2000),
        "reflectivity": ("uniform", 0.80, 0.95),
    }
    
    mc_tasks = list(MonteCarloGenerator.generate(base_config, distributions, 5, seed=42))
    print(f"Generated {len(mc_tasks)} Monte Carlo samples")
    print(f"First sample: {mc_tasks[0]}")