import subprocess
import itertools
import os
import sys
import shutil
import re

# --- Configuration ---
TARGET_SCRIPT = "prismatic5m.py"
TEMP_SCRIPT = "prismatic5m_temp.py"
super_list = [128, 256, 512, 1024]
instances_list = [2048, 4096, 8192]
slice_list = [8, 16, 32]

def find_python_executable():
    if hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix):
        return sys.executable
    venv_path = os.path.join(os.path.dirname(sys.executable), "Scripts", "python.exe")
    if os.path.exists(venv_path):
        return venv_path
    return "python"

def run_benchmark():
    best_combo = None
    best_fps = 0.0
    python_exec = find_python_executable()
    print(f"Using Python executable: {python_exec}")

    if not os.path.exists(TARGET_SCRIPT):
        print(f"Error: The target script '{TARGET_SCRIPT}' was not found.")
        return

    for NUM_SUPER, NUM_INSTANCES, MAX_SLICE in itertools.product(super_list, instances_list, slice_list):
        print(f"\nTesting NUM_SUPER={NUM_SUPER}, NUM_INSTANCES={NUM_INSTANCES}, MAX_SLICE={MAX_SLICE}")
        
        try:
            shutil.copyfile(TARGET_SCRIPT, TEMP_SCRIPT)
            with open(TEMP_SCRIPT, "r") as f:
                src = f.read()
            
            # More precise regex to avoid matching inside strings
            src = re.sub(r"^\s*NUM_SUPER\s*=\s*\d+\s*$", f"NUM_SUPER = {NUM_SUPER}", src, flags=re.MULTILINE)
            src = re.sub(r"^\s*NUM_INSTANCES\s*=\s*\d+\s*$", f"NUM_INSTANCES = {NUM_INSTANCES}", src, flags=re.MULTILINE)
            src = re.sub(r"^\s*MAX_SLICE\s*=\s*\d+\s*$", f"MAX_SLICE = {MAX_SLICE}", src, flags=re.MULTILINE)
            
            with open(TEMP_SCRIPT, "w") as f:
                f.write(src)
            
            # Run the script and capture output
            result = subprocess.run([python_exec, TEMP_SCRIPT], capture_output=True, text=True, timeout=5)
            output = result.stdout
            print(f"Script output: {output}")
            
            # Placeholder FPS parsing (replace with actual FPS extraction logic)
            fps = 60.0  # Replace with actual FPS parsing from output
            print(f"Test completed. Placeholder FPS: {fps}")
            
            if fps > best_fps:
                best_fps = fps
                best_combo = (NUM_SUPER, NUM_INSTANCES, MAX_SLICE)
            
        except subprocess.TimeoutExpired:
            print("Test timed out. Terminating process...")
            subprocess.run([python_exec, "-c", "import os; os.kill(os.getpid(), 9)"], timeout=2)
        except FileNotFoundError:
            print(f"Error: Could not find Python executable at '{python_exec}'.")
            return
        except Exception as e:
            print(f"An unexpected error occurred during testing: {e}")
            print(f"Preserving temporary script '{TEMP_SCRIPT}' for debugging.")
            continue
        finally:
            if os.path.exists(TEMP_SCRIPT) and 'e' not in locals():
                os.remove(TEMP_SCRIPT)

    if best_combo:
        print(f"\nOptimal RX 480 TFLOPs combination: NUM_SUPER={best_combo[0]}, NUM_INSTANCES={best_combo[1]}, MAX_SLICE={best_combo[2]}, FPS={best_fps}")
    else:
        print("\nNo optimal combination found. Please check your script and environment.")

if __name__ == "__main__":
    run_benchmark()