batch_runner.py

# ====================
# batch_runner.py
# ====================
# !/usr/bin/env python3
"""
Batch runner for BeDDeM parameter sensitivity analysis
"""

import os
import pandas as pd
import numpy as np
from pathlib import Path
import subprocess
import itertools
import sys
import json


class BeddemBatchRunner:
    """Batch runner for parameter sensitivity analysis"""

    def __init__(self, base_data_dir="data"):
        self.base_data_dir = Path(base_data_dir)
        self.results_dir = Path("batch_results")
        self.results_dir.mkdir(exist_ok=True)

    def generate_parameter_combinations(self):
        """Generate parameter combinations for sensitivity analysis with proper TIB weights"""
        # TIB Level 1 weight ranges for sensitivity analysis
        time_weights = [1.0, 3.0, 5.0]
        price_weights = [1.0, 3.0, 5.0]
        role_weights = [1.0, 3.0, 5.0]  # Environmental role weight
        norm_weights = [1.0, 3.0, 5.0]  # Social norm weight

        combinations = []
        for tw, pw, rw, nw in itertools.product(time_weights, price_weights, role_weights, norm_weights):
            combinations.append({
                'time_weight': tw,
                'price_weight': pw,
                'role_weight': rw,
                'norm_weight': nw,
                'experiment_id': f"tw_{tw}_pw_{pw}_rw_{rw}_nw_{nw}"
            })

        return combinations

    def create_experiment_data(self, params):
        """Create data files for a specific experiment with proper TIB structure"""
        exp_dir = self.results_dir / params['experiment_id']
        exp_dir.mkdir(exist_ok=True)

        data_dir = exp_dir / "data"
        data_dir.mkdir(exist_ok=True)

        # Create modified agent.csv with proper TIB weights structure
        agent_data = f"""agent_id,price_weight,time_weight,norm_weight,role_weight,self_concept_weight,emotion_weight,consequence_weight,social_weight,affect_weight,intention_weight,habit_weight,facilitating_weight
1,{params['price_weight']},{params['time_weight']},{params['norm_weight']},{params['role_weight']},2.0,1.0,3.0,2.0,1.0,3.0,2.0,1.0
2,{params['price_weight']},{params['time_weight']},{params['norm_weight']},{params['role_weight']},3.0,2.0,4.0,3.0,2.0,4.0,3.0,2.0
3,{params['price_weight']},{params['time_weight']},{params['norm_weight']},{params['role_weight']},1.0,1.0,2.0,1.0,1.0,2.0,1.0,1.0"""

        with open(data_dir / "agent.csv", "w") as f:
            f.write(agent_data)

        # Create basic vehicle.csv if it doesn't exist
        if not (self.base_data_dir / "vehicle.csv").exists():
            vehicle_data = """vehicle_id,cost_per_km,speed_kmh,comfort_level,availability
walk,0.0,5.0,0.3,1.0
car,0.30,50.0,0.9,1.0
bus,0.15,25.0,0.6,1.0
bike,0.02,15.0,0.7,1.0
train,0.20,80.0,0.8,1.0
tram,0.18,30.0,0.7,1.0"""
            with open(data_dir / "vehicle.csv", "w") as f:
                f.write(vehicle_data)
        else:
            # Copy existing vehicle.csv
            dst_file = data_dir / "vehicle.csv"
            dst_file.write_text((self.base_data_dir / "vehicle.csv").read_text())

        # Create basic location.csv if it doesn't exist
        if not (self.base_data_dir / "location.csv").exists():
            location_data = """loc_id,train,bus,tram
1,1,1,0
2,1,1,1
3,0,1,0
4,1,0,0"""
            with open(data_dir / "location.csv", "w") as f:
                f.write(location_data)
        else:
            # Copy existing location.csv
            dst_file = data_dir / "location.csv"
            dst_file.write_text((self.base_data_dir / "location.csv").read_text())

        # Create basic schedule.csv if it doesn't exist
        if not (self.base_data_dir / "schedule.csv").exists():
            schedule_data = """task_id,agent_id,start_time,end_time,origin,destination,distance,required
1,1,8.0,9.0,1,2,15.0,true
2,1,17.0,18.0,2,1,15.0,true
3,2,9.0,10.0,1,3,10.0,true
4,2,16.0,17.0,3,1,10.0,true
5,3,7.0,8.0,1,4,25.0,true
6,3,18.0,19.0,4,1,25.0,true"""
            with open(data_dir / "schedule.csv", "w") as f:
                f.write(schedule_data)
        else:
            # Copy existing schedule.csv
            dst_file = data_dir / "schedule.csv"
            dst_file.write_text((self.base_data_dir / "schedule.csv").read_text())

        return str(data_dir)

    def create_config_file(self, data_dir, exp_dir):
        """Create a configuration file for the experiment"""
        config = {
            'data_path': data_dir,
            'max_steps': 24,
            'output_file': str(exp_dir / 'beddem_output.csv')
        }

        config_file = exp_dir / 'config.json'
        with open(config_file, 'w') as f:
            json.dump(config, f, indent=2)

        return str(config_file)

    def run_experiment(self, params, num_processes=1):
        """Run a single experiment"""
        print(f"Running experiment: {params['experiment_id']}")

        exp_dir = self.results_dir / params['experiment_id']
        data_dir = self.create_experiment_data(params)
        config_file = self.create_config_file(data_dir, exp_dir)

        # Create a simple test script for this experiment
        test_script = exp_dir / "run_experiment.py"
        test_script_content = f'''#!/usr/bin/env python3
import sys
import os
sys.path.insert(0, "{os.getcwd()}")

from beddem_repast4py import BeddemSimulation

# Run simulation with specific data path
simulation = BeddemSimulation("{config_file}")
try:
    simulation.run()
    print(f"Experiment {params['experiment_id']} completed successfully")
except Exception as e:
    print(f"Error in experiment {params['experiment_id']}: {{e}}")
    import traceback
    traceback.print_exc()
'''

        with open(test_script, "w") as f:
            f.write(test_script_content)

        # Run the experiment
        try:
            if num_processes > 1:
                cmd = ["mpirun", "-n", str(num_processes), "python", str(test_script)]
            else:
                cmd = ["python", str(test_script)]

            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                cwd=os.getcwd()
            )

            # Save output
            with open(exp_dir / "stdout.log", "w") as f:
                f.write(result.stdout)

            with open(exp_dir / "stderr.log", "w") as f:
                f.write(result.stderr)

            # Check if output file was created
            output_files = list(exp_dir.glob("beddem_output*.csv"))
            success = result.returncode == 0 and len(output_files) > 0

            if not success:
                print(f"  Failed: return code {result.returncode}, output files: {len(output_files)}")
                if result.stderr:
                    print(f"  Error: {result.stderr[:200]}...")
            else:
                print(f"  Success: {len(output_files)} output file(s) created")

            return success

        except Exception as e:
            print(f"Error running experiment {params['experiment_id']}: {e}")
            # Save error info
            with open(exp_dir / "error.log", "w") as f:
                f.write(str(e))
            return False

    def run_batch(self, num_processes=1, max_experiments=None):
        """Run all experiments in batch"""
        combinations = self.generate_parameter_combinations()

        if max_experiments:
            combinations = combinations[:max_experiments]

        print(f"Running {len(combinations)} experiments...")

        results = []
        successful = 0

        for i, params in enumerate(combinations):
            print(f"\n--- Experiment {i + 1}/{len(combinations)} ---")
            success = self.run_experiment(params, num_processes)

            if success:
                successful += 1

            results.append({
                'experiment_id': params['experiment_id'],
                'time_weight': params['time_weight'],
                'price_weight': params['price_weight'],
                'role_weight': params['role_weight'],
                'norm_weight': params['norm_weight'],
                'success': success
            })

        # Save batch results summary
        df = pd.DataFrame(results)
        df.to_csv(self.results_dir / "batch_summary.csv", index=False)

        print(f"\nBatch run completed: {successful}/{len(combinations)} successful")
        print(f"Results saved to {self.results_dir}")
        return results

    def analyze_results(self):
        """Analyze batch results"""
        print("Analyzing batch results...")

        # Check what directories exist
        exp_dirs = [d for d in self.results_dir.iterdir() if d.is_dir() and d.name.startswith("tw_")]
        print(f"Found {len(exp_dirs)} experiment directories")

        # Collect all results
        all_results = []
        successful_experiments = 0

        for exp_dir in exp_dirs:
            print(f"Checking {exp_dir.name}...")

            # Look for output files
            output_files = list(exp_dir.glob("beddem_output*.csv"))

            if output_files:
                try:
                    # Read the output file
                    df = pd.read_csv(output_files[0])

                    # Extract parameters from directory name
                    # Format: tw_X.X_pw_X.X_rw_X.X_nw_X.X
                    parts = exp_dir.name.split("_")
                    if len(parts) >= 8:
                        time_weight = float(parts[1])
                        price_weight = float(parts[3])
                        role_weight = float(parts[5])
                        norm_weight = float(parts[7])

                        # Add parameter columns
                        df['time_weight'] = time_weight
                        df['price_weight'] = price_weight
                        df['role_weight'] = role_weight
                        df['norm_weight'] = norm_weight
                        df['experiment_id'] = exp_dir.name

                        all_results.append(df)
                        successful_experiments += 1
                        print(f"  ✓ Loaded {len(df)} decisions")
                    else:
                        print(f"  ✗ Cannot parse experiment name: {exp_dir.name}")

                except Exception as e:
                    print(f"  ✗ Error reading output: {e}")
            else:
                # Check for error logs
                error_files = list(exp_dir.glob("*.log"))
                if error_files:
                    print(f"  ✗ No output files, but found logs: {[f.name for f in error_files]}")
                else:
                    print(f"  ✗ No output files or logs found")

        print(f"\nSuccessfully loaded {successful_experiments} experiments")

        if all_results:
            # Combine all results
            combined_df = pd.concat(all_results, ignore_index=True)
            combined_df.to_csv(self.results_dir / "combined_results.csv", index=False)

            # Basic analysis
            print(f"\n=== Batch Analysis Results ===")
            print(f"Total decisions recorded: {len(combined_df)}")
            print(f"Unique experiments: {combined_df['experiment_id'].nunique()}")

            # Mode choice distribution
            if 'mode' in combined_df.columns:
                print(f"\nOverall mode distribution:")
                mode_counts = combined_df['mode'].value_counts()
                for mode, count in mode_counts.items():
                    percentage = (count / len(combined_df)) * 100
                    print(f"  {mode.upper():8s}: {count:4d} ({percentage:.1f}%)")

                # Mode choice by parameters
                print(f"\nMode choice by time and price weights:")
                mode_analysis = combined_df.groupby(['time_weight', 'price_weight', 'mode']).size().unstack(
                    fill_value=0)
                print(mode_analysis)

                # Save detailed analysis
                mode_analysis.to_csv(self.results_dir / "mode_analysis.csv")

                # Parameter sensitivity analysis
                print(f"\nParameter sensitivity (mode diversity by weight):")
                for param in ['time_weight', 'price_weight', 'role_weight', 'norm_weight']:
                    param_diversity = combined_df.groupby(param)['mode'].nunique()
                    print(f"  {param:15s}: {param_diversity.to_dict()}")

            return combined_df
        else:
            print("No results found to analyze.")
            print("\nTroubleshooting tips:")
            print("1. Check if experiments ran successfully: --run-batch")
            print("2. Look at individual experiment logs in batch_results/")
            print("3. Verify the BeDDeM simulation is working independently")
            return None


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Run BeDDeM Batch Analysis")
    parser.add_argument("--run-batch", action="store_true",
                        help="Run batch experiments")
    parser.add_argument("--analyze", action="store_true",
                        help="Analyze existing results")
    parser.add_argument("-n", "--processes", type=int, default=1,
                        help="Number of MPI processes per experiment")
    parser.add_argument("--max-experiments", type=int, default=None,
                        help="Limit number of experiments (for testing)")

    args = parser.parse_args()

    runner = BeddemBatchRunner()

    if args.run_batch:
        runner.run_batch(args.processes, args.max_experiments)

    if args.analyze:
        runner.analyze_results()

    if not args.run_batch and not args.analyze:
        print("Use --run-batch to run experiments or --analyze to analyze results")
        print("Example: python batch_runner.py --run-batch --max-experiments 4")