run_experiments_multithread_python.py

import os
import subprocess
import csv
import time
import re
import argparse
from concurrent.futures import ProcessPoolExecutor, as_completed
from functools import partial
from tqdm import tqdm
import multiprocessing

# ============================================================
#region CONFIGURATION
# ============================================================

# ASP / Clingo
CLINGO_BIN = r"C:\Program Files\clingo\clingo.exe"
ASP_PROGRAM = "political_districting.lp"
ASP_BENCHMARK_DIR = "benchmarks"
TIMEOUT_SECONDS_CLINGO = 300
PYTHON_TIMEOUT_SAFETY_MARGIN = 30

CLINGO_CONFIGS = {
    "default": [],
}

CLINGO_MODEL_COUNT = "0"
CLINGO_QUIET_OPTION = ["--quiet=1"]

# MiniZinc
MINIZINC_BIN = r"C:\Program Files\MiniZinc\minizinc.exe"
MINIZINC_MODEL = "political_districting.mzn"
MINIZINC_IMPROVED_MODEL = "political_districting_improved.mzn"
MZN_BENCHMARK_DIR = "benchmarks_dzn"
MINIZINC_TIMEOUT = 300

# Results
RESULTS_FILE = "results_python.csv"

# ============================================================
#region PARSING HELPERS
# ============================================================

def parse_clingo_summary(output_text, clingo_return_code, python_timed_out):
    status = "ERROR"
    max_reps = "NA"
    num_models = "NA"
    raw_opt = "NA"

    if python_timed_out:
        return "TIMEOUT", "NA", "NA", "NA"

    # Optimization Value
    opt_match = re.search(r"Optimization\s*:\s*(-?\d+)", output_text)
    if opt_match:
        raw_opt = int(opt_match.group(1))
    
    reps_match = re.search(r"representatives_for_optimized_party\((\d+)\)", output_text)
    if reps_match:
        max_reps = int(reps_match.group(1))
    else:
        max_reps = abs(raw_opt) if opt_match else "NA"

    # Model count
    models_match = re.search(r"Models\s*:\s*(\d+)", output_text)
    if models_match:
        num_models = models_match.group(1)

    # Status Mapping
    if "OPTIMUM FOUND" in output_text:
        status = "OPTIMAL"
    elif "SATISFIABLE" in output_text:
        status = "SATISFIABLE"
    elif "UNSATISFIABLE" in output_text:
        status = "UNSATISFIABLE"
    elif "INTERRUPTED" in output_text or clingo_return_code in [1, 30]: 
        status = "TIMEOUT"

    return status, max_reps, num_models, raw_opt

def parse_minizinc_output(output_text, exit_code, python_timed_out):
    status = "ERROR" 
    max_reps = "NA"
    num_models = "NA"
    raw_opt = "NA"

    if python_timed_out:
        return "TIMEOUT", "NA", "NA", "NA"

    matches = re.findall(r"Representatives for party\s+\d+\s*:\s*(\d+)", output_text)
    
    if matches:
        max_reps = int(matches[-1])
        num_models = len(matches)
        status = "OPTIMAL" if "==========" in output_text else "SATISFIABLE"
    elif "UNSATISFIABLE" in output_text:
        status = "UNSATISFIABLE"
    elif exit_code != 0 or "time limit exceeded" in output_text.lower():
        status = "TIMEOUT"

    return status, max_reps, num_models, raw_opt

# ============================================================
#region SOLVER RUNNERS
# ============================================================

def check_solver_available(name, command):
    """
    Checks whether a solver is available by trying to run '<command> --version'.
    Returns True if available, False otherwise.
    """
    try:
        process = subprocess.run(
            [command, "--version"],
            capture_output=True,
            text=True,
            timeout=5
        )
        if process.returncode == 0:
            print(f"[OK] {name} found: {command}")
            first_line = process.stdout.strip().splitlines()[0] if process.stdout else ""
            print(f"     Version: {first_line}")
            return True
        else:
            print(f"[WARN] {name} returned non-zero exit code when checking version.")
            print(f"       Command: {command}")
            print(f"       Output: {process.stdout.strip()}")
            return False

    except FileNotFoundError:
        print(f"[ERROR] {name} not found: {command}")
        print("        The executable does not exist or is not on PATH.")
        return False

    except Exception as e:
        print(f"[ERROR] Unexpected error while checking {name}: {e}")
        return False

# ============================================================
#region TASK RUNNERS
# ============================================================

def run_clingo_task(clingo_bin, asp_program, instance_path, party, strategy_name, strategy_options_list, timeout_seconds, python_margin, dry_run=False):
    """
    Run a single clingo task. Returns a result dict.
    """
    command = [
        clingo_bin,
        asp_program,
        instance_path,
        f"-c", f"party_to_optimize={party}",
        f"--time-limit={timeout_seconds}",
    ]
    command.extend(CLINGO_QUIET_OPTION)
    command.append(CLINGO_MODEL_COUNT)
    command.extend(strategy_options_list)

    if dry_run:     
        return {
            "solver": "clingo",
            "Instance": os.path.basename(instance_path),
            "PartyOptimized": party,
            "Strategy": strategy_name,
            "ModelFile": os.path.basename(asp_program),
            "max_reps": "DRY",
            "time_sec": 0.0,
            "status": "DRY_RUN",
            "models": "NA",
            "exit_code": "DRY",
            "raw_opt": "NA",
            "cmd": " ".join(command)
        }

    start_time = time.perf_counter()
    python_timeout = False
    output_text = ""
    exit_code = "NA"

    try:
        process = subprocess.run(
            command,
            capture_output=True,
            text=True,
            timeout=timeout_seconds + python_margin,
            check=False
        )
        output_text = (process.stdout or "") + "\n---\n" + (process.stderr or "")
        exit_code = process.returncode

    except subprocess.TimeoutExpired:
        python_timeout = True
        output_text = "Execution timed out by Python script."
    except FileNotFoundError:
        output_text = "Clingo executable not found."
        exit_code = "FNF_ERROR"
    except Exception as e:
        output_text = f"Python error: {str(e)}"
        exit_code = "PYTHON_SUBPROC_ERROR"

    elapsed = time.perf_counter() - start_time

    status, max_reps, num_models, raw_opt = parse_clingo_summary(output_text, exit_code, python_timeout)

    return {
        "solver": "clingo",
        "Instance": os.path.basename(instance_path),
        "PartyOptimized": party,
        "Strategy": strategy_name,
        "ModelFile": os.path.basename(asp_program),
        "max_reps": max_reps,
        "time_sec": elapsed,
        "status": status,
        "models": num_models,
        "exit_code": exit_code,
        "raw_opt": raw_opt,
        "cmd": None
    }


def run_minizinc_task(minizinc_bin, minizinc_model, instance_path, party, timeout_seconds, python_margin, dry_run=False):
    """
    Run a single MiniZinc task. Returns a result dict.
    """
    command = [
        minizinc_bin,
        "--solver", "chuffed",
        "-D", f"party_to_optimize={party}",
        minizinc_model,
        instance_path,
        "--time-limit", str(timeout_seconds * 1000),
    ]

    solver_name = "minizinc" if os.path.basename(minizinc_model) == os.path.basename(MINIZINC_MODEL) else "improved_minizinc"

    if dry_run:
        return {
            "solver": solver_name,
            "Instance": os.path.basename(instance_path),
            "PartyOptimized": party,
            "Strategy": "chuffed",
            "ModelFile": os.path.basename(minizinc_model),
            "max_reps": "DRY",
            "time_sec": 0.0,
            "status": "DRY_RUN",
            "models": "NA",
            "exit_code": "DRY",
            "raw_opt": "NA",
            "cmd": " ".join(command)
        }

    start_time = time.perf_counter()
    python_timeout = False
    output_text = ""
    exit_code = "NA"

    try:
        process = subprocess.run(
            command,
            capture_output=True,
            text=True,
            timeout=timeout_seconds + python_margin,
            check=False
        )
        output_text = (process.stdout or "") + "\n---\n" + (process.stderr or "")
        exit_code = process.returncode

        if exit_code != 0:
            print(f"DEBUG: MiniZinc failed on {instance_path}. Error:\n{output_text}")

    except subprocess.TimeoutExpired:
        python_timeout = True
        output_text = "Execution timed out by Python script."
    except FileNotFoundError:
        output_text = "MiniZinc executable not found."
        exit_code = "FNF_ERROR"
    except Exception as e:
        output_text = f"Python error: {str(e)}"
        exit_code = "PYTHON_SUBPROC_ERROR"

    elapsed = time.perf_counter() - start_time

    status, reps, models, opt = parse_minizinc_output(output_text, exit_code, python_timeout)

    return {
        "solver": solver_name,
        "Instance": os.path.basename(instance_path),
        "PartyOptimized": party,
        "Strategy": "chuffed",
        "ModelFile": os.path.basename(minizinc_model),
        "max_reps": reps,
        "time_sec": elapsed,
        "status": status,
        "models": models,
        "exit_code": exit_code,
        "raw_opt": "NA",
        "cmd": None
    }

# ============================================================
#region SUMMARY REPORT
# ============================================================

def print_summary(summary):
    active_solvers = {k: v for k, v in summary.items() if v['count'] > 0}

    if not active_solvers:
        print("\n" + "!" * 95)
        print(f"{'!!! NO EXPERIMENTS EXECUTED !!!':^95}")
        print("!" * 95 + "\n")
        return

    print("\n" + "=" * 95)
    print(f"{'EXPERIMENT PERFORMANCE SUMMARY':^95}")
    print("=" * 95)
    
    header = f"{'Solver':<20} | {'Runs':>6} | {'Opt':>6} | {'SAT':>6} | {'UNSAT':>6} | {'Timeout':>8} | {'Err':>5} | {'Avg Time'}"
    print(header)
    print("-" * 95)

    for solver, stats in active_solvers.items():
        avg = stats['total_time'] / stats['count'] if stats['count'] > 0 else 0.0
        display_name = solver.replace("_", " ").title()
        
        row = (f"{display_name:<20} | "
               f"{stats['count']:>6} | "
               f"{stats['optimal']:>6} | "
               f"{stats['sat']:>6} | "
               f"{stats['unsat']:>6} | "
               f"{stats['timeout']:>8} | "
               f"{stats.get('error',0):>5} | "
               f"{avg:>8.2f}s")
        print(row)
    print("=" * 95 + "\n")

# ============================================================
#region MAIN
# ============================================================

def main():
    parser = argparse.ArgumentParser(
        description="Gerrymandering Experiment Runner\nRuns ASP (Clingo) and MiniZinc models in parallel to compare efficiency.",
        formatter_class=argparse.RawTextHelpFormatter,
        epilog="""
            Modes:
            asp                Run only the ASP (Clingo) model.
            minizinc           Run only the base MiniZinc model.
            improved_minizinc  Run only the optimized MiniZinc model (Fixed Roots).
            base_comparison    Run 'asp' and 'minizinc' (Standard comparison).
            all                Run all three strategies for full benchmarking.
                    """
    )

    parser.add_argument(
        "--solver",
        choices=["asp", "minizinc", "improved_minizinc", "base_comparison", "all"],
        default="base_comparison",
        help="Target solver strategy. Default is 'base_comparison'."
    )

    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Validation mode: Print generated commands without starting execution."
    )

    parser.add_argument(
        "--workers",
        type=int,
        default=max(1, (multiprocessing.cpu_count() // 2)),
        help="Number of parallel threads. Default: half of available CPU cores."
    )

    parser.add_argument(
        "--limit",
        type=int,
        default=0,
        help="Subset limit: Run only the first N instances (Useful for debugging)."
    )

    args = parser.parse_args()

    # Determine which solvers to run
    run_asp = args.solver in ("asp", "base_comparison", "all")
    run_mzn = args.solver in ("minizinc", "base_comparison", "all")
    run_improved = args.solver in ("improved_minizinc", "all")

    print("Checking solver availability...\n")
    clingo_ok = check_solver_available("Clingo", CLINGO_BIN if os.path.isabs(CLINGO_BIN) else "clingo")
    minizinc_ok = check_solver_available("MiniZinc", MINIZINC_BIN)
    print()

    if run_asp and not clingo_ok:
        print("Clingo is required but not available. Aborting.")
        return

    if (run_mzn or run_improved) and not minizinc_ok:
        print("MiniZinc is required but not available. Aborting.")
        return

    write_header = not os.path.isfile(RESULTS_FILE)

    summary = {
    "clingo": {"count": 0, "optimal": 0, "sat": 0, "unsat": 0, "timeout": 0, "total_time": 0.0},
    "minizinc": {"count": 0, "optimal": 0, "sat": 0, "unsat": 0, "timeout": 0, "total_time": 0.0},
    "improved_minizinc": {"count": 0, "optimal": 0, "sat": 0, "unsat": 0, "timeout": 0, "total_time": 0.0}
    }

    limit = args.limit if args.limit and args.limit > 0 else None

    asp_instances = sorted(
        f for f in os.listdir(ASP_BENCHMARK_DIR) if f.endswith(".lp")
    ) if run_asp else []
    if limit:
        asp_instances = asp_instances[:limit]

    mzn_instances = sorted(
        f for f in os.listdir(MZN_BENCHMARK_DIR) if f.endswith(".dzn")
    ) if (run_mzn or run_improved) else []
    if limit:
        mzn_instances = mzn_instances[:limit]

    # Build task list
    tasks = []
    if run_asp:
        for instance_filename in asp_instances:
            instance_path = os.path.join(ASP_BENCHMARK_DIR, instance_filename)
            for party in [0, 1]:
                for strategy_name, strategy_opts in CLINGO_CONFIGS.items():
                    tasks.append(("clingo", instance_path, party, strategy_name, strategy_opts, instance_filename))

    if run_mzn:
        for instance_filename in mzn_instances:
            instance_path = os.path.join(MZN_BENCHMARK_DIR, instance_filename)
            for party in [0, 1]:
                tasks.append(("minizinc", instance_path, party, "minizinc", [], instance_filename))
    
    
    if run_improved:
        for instance_filename in mzn_instances:
            instance_path = os.path.join(MZN_BENCHMARK_DIR, instance_filename)
            for party in [0, 1]:
                tasks.append(("improved_minizinc", instance_path, party, "improved_minizinc", [], instance_filename))

    total_runs = len(tasks)
    if total_runs == 0:
        print("No tasks to run (check benchmark directories and --solver option).")
        return

    # Prepare CSV writer in main process
    csvfile = open(RESULTS_FILE, 'a', newline='', encoding='utf-8')
    csv_writer = csv.writer(csvfile)
    if write_header:
        csv_writer.writerow([
            "Instance", "PartyOptimized","ModelFile", "Strategy",
            "MaxReps", "TimeSec", "Status",
            "Models", "SolverReturnCode", "RawOptValue"
        ])
        csvfile.flush()

    # Create partial functions for pool workers
    clingo_partial = partial(
        run_clingo_task,
        CLINGO_BIN if os.path.isabs(CLINGO_BIN) else "clingo",
        ASP_PROGRAM,
        timeout_seconds=TIMEOUT_SECONDS_CLINGO,
        python_margin=PYTHON_TIMEOUT_SAFETY_MARGIN,
        dry_run=args.dry_run
    )

    minizinc_partial = partial(
        run_minizinc_task,
        MINIZINC_BIN,
        MINIZINC_MODEL,
        timeout_seconds=MINIZINC_TIMEOUT,
        python_margin=PYTHON_TIMEOUT_SAFETY_MARGIN,
        dry_run=args.dry_run
    )

    improved_partial = partial(
        run_minizinc_task,
        MINIZINC_BIN,
        MINIZINC_IMPROVED_MODEL,
        timeout_seconds=MINIZINC_TIMEOUT,
        python_margin=PYTHON_TIMEOUT_SAFETY_MARGIN,
        dry_run=args.dry_run
    )

    # Use ProcessPoolExecutor to run tasks in parallel
    workers = max(1, args.workers)
    print(f"Starting {workers} worker(s). Total tasks: {total_runs}\n")

    futures = []
    try:
        with ProcessPoolExecutor(max_workers=workers) as executor:
            for t in tasks:
                solver_type, instance_path, party, strategy_name, strategy_opts, instance_filename = t
                if solver_type == "clingo":
                    futures.append((executor.submit(clingo_partial, instance_path, party, strategy_name, strategy_opts), instance_filename))
                elif solver_type == "minizinc":
                    futures.append((executor.submit(minizinc_partial, instance_path, party), instance_filename))
                else:
                    futures.append((executor.submit(improved_partial, instance_path, party), instance_filename))

            future_to_instance = {fut: inst for fut, inst in futures}

            with tqdm(total=total_runs, desc="Running experiments") as pbar:
                for completed in as_completed(future_to_instance):
                    instance_filename = future_to_instance[completed]
                    try:
                        res = completed.result()
                    except Exception as e:
                        res = {
                            "solver": "unknown",
                            "strategy": "error",
                            "party": "NA",
                            "max_reps": "NA",
                            "time_sec": 0.0,
                            "status": f"WORKER_EXCEPTION: {e}",
                            "models": "NA",
                            "exit_code": "WORKER_EXCEPTION",
                            "raw_opt": "NA"
                        }

                    csv_writer.writerow([
                        res.get("Instance"),
                        res.get("PartyOptimized"),
                        res.get("ModelFile"),
                        res.get("Strategy"),
                        res.get("max_reps"),
                        f"{res.get('time_sec', 0.0):.3f}",
                        res.get("status"),
                        res.get("models"),
                        res.get("exit_code"),
                        res.get("raw_opt")
                    ])
                    csvfile.flush()

                    solver_key = res.get("solver", "").lower()
                    if solver_key in summary:
                        summary[solver_key]["count"] += 1
                        summary[solver_key]["total_time"] += float(res.get("time_sec", 0.0) or 0.0)
                        
                    status = str(res.get("status", ""))

                    if status == "OPTIMAL":
                        summary[solver_key]["optimal"] += 1
                    elif status == "SATISFIABLE":
                        summary[solver_key]["sat"] += 1
                    elif status == "UNSATISFIABLE":
                        summary[solver_key]["unsat"] += 1
                    elif status == "TIMEOUT":
                        summary[solver_key]["timeout"] += 1
                    else:
                        # This catches ERROR or UNKNOWN
                        summary[solver_key].setdefault("error", 0)
                        summary[solver_key]["error"] += 1

                    pbar.update(1)

    except KeyboardInterrupt:
        print("\nInterrupted by user. Attempting to shut down workers...")
    finally:
        csvfile.close()

    print_summary(summary)
    print(f"Results saved to {RESULTS_FILE}")


if __name__ == "__main__":
    main()