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

# --- Configuration ---
# Set the name of the Python script to be tested
TARGET_SCRIPT = "prismatic5m.py"
# Set the name of the temporary file for testing
TEMP_SCRIPT = "prismatic5m_temp.py"

# Define the ranges for the benchmark parameters
super_list = [128, 256, 512, 1024]
instances_list = [2048, 4096, 8192]
slice_list = [8, 16, 32]

# --- Main Script Logic ---

def find_python_executable():
    """
    Attempts to find the Python executable, prioritizing the one in the
    current virtual environment. This helps resolve ModuleNotFoundError issues.
    """
    # Check if a virtual environment is active
    if hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix):
        # We are inside a virtual environment, so use the current executable
        return sys.executable
    
    # Fallback to a common location for a virtual environment's python.exe
    venv_path = os.path.join(os.path.dirname(sys.executable), "Scripts", "python.exe")
    if os.path.exists(venv_path):
        return venv_path
    
    # Fallback to the default system python
    return "python"

def run_benchmark():
    """
    Iterates through parameter combinations, modifies the target script,
    and runs it as a subprocess to test performance.
    """
    best_combo = None
    best_tflops = 0.0

    # Get the correct Python executable path
    python_exec = find_python_executable()
    print(f"Using Python executable: {python_exec}")

    # Ensure the target script exists
    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:
            # Create a temporary copy of the script for modification
            shutil.copyfile(TARGET_SCRIPT, TEMP_SCRIPT)
            
            # Read the temporary script's content
            with open(TEMP_SCRIPT, "r") as f:
                src = f.read()
            
            # Replace the parameters dynamically using regular expressions
            # This makes the script more resilient to changes in initial values
            src = re.sub(r"NUM_SUPER\s*=\s*\d+", f"NUM_SUPER = {NUM_SUPER}", src)
            src = re.sub(r"NUM_INSTANCES\s*=\s*\d+", f"NUM_INSTANCES = {NUM_INSTANCES}", src)
            src = re.sub(r"MAX_SLICE\s*=\s*\d+", f"MAX_SLICE = {MAX_SLICE}", src)
            
            # Write the modified content back to the temporary file
            with open(TEMP_SCRIPT, "w") as f:
                f.write(src)
            
            # Launch the subprocess using the correct Python executable
            proc = subprocess.Popen([python_exec, TEMP_SCRIPT])
            
            # In a real-world scenario, you would monitor the GPU here.
            # The current script waits for 5 seconds and then terminates.
            try:
                proc.wait(timeout=5)
                
                # Placeholder for a real TFLOPS calculation. 
                # This assumes a specific calculation based on the parameters.
                # In a real application, you would measure this directly from the GPU.
                tflops = (NUM_SUPER * NUM_INSTANCES * MAX_SLICE) / 1e12 
                print(f"Test completed. Estimated TFLOPS: {tflops:.2f}")
                
                if tflops > best_tflops:
                    best_tflops = tflops
                    best_combo = (NUM_SUPER, NUM_INSTANCES, MAX_SLICE)
            
            except subprocess.TimeoutExpired:
                print("Test timed out. Terminating process...")
                proc.terminate()
                proc.wait() # Ensure the process is fully terminated
            
        except FileNotFoundError:
            print(f"Error: Could not find Python executable at '{python_exec}'. "
                  f"Please ensure it's in your system PATH or that your virtual environment is activated.")
            return # Exit on critical error
        except Exception as e:
            print(f"An unexpected error occurred during testing: {e}")
        finally:
            # Clean up the temporary file
            if os.path.exists(TEMP_SCRIPT):
                os.remove(TEMP_SCRIPT)

    # Print the final results
    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]}, TFLOPS={best_tflops:.2f}")
    else:
        print("\nNo optimal combination found. Please check your script and environment.")

if __name__ == "__main__":
    run_benchmark()
