import subprocess
import os
import sys

ds = 57
perturbation_names = {
 	"variable_renaming": 1, 
 	"neutral_operator": 2, 
 	"arithmetic_commutativity": 3, 
 	"comparison_commutativity": 4, 
 	"neutral_assignment": 5, 
 	"neutral_expression": 6, 
}

perturbation_modes = {
	"all": 1,
	"one": 2,
	"proba": 3,
}

sample_instance = "2"
perturbation_name = "variable_renaming"
perturbation_mode = "all"

start_checkpoint = 1
end_checkpoint = 8

env = os.environ.copy()
procs = []
files = []
try:

	# Creating the eval folders and creating the clones
	print("Creating the eval folders and creating the clones")
	for checkpoint in range(start_checkpoint, end_checkpoint + 1):
		checkpoint_dir = f"eval-checkpoint-{checkpoint}-{perturbation_name+sample_instance}"
		os.makedirs(checkpoint_dir, exist_ok=True)
		files.append(open(f"{checkpoint_dir}/cloning.out", "w"))
		proc = subprocess.Popen(
			[
				"python",
				os.path.abspath(f"indistribution_perturbator.py"),
				str(perturbation_names[perturbation_name]),
				str(perturbation_modes[perturbation_mode]),
				os.path.abspath(f"eval-checkpoint-{checkpoint}/infers/all-results.csv"),
				os.path.abspath(f"{checkpoint_dir}/robustness_snippets.txt"),
				"1.0",
				"1",
				"11",

			],
			cwd = checkpoint_dir,
			stdout = files[-1],
			stderr = files[-1],
		)
		procs.append(proc)
	
	# Wait for all children
	for proc in procs:
		proc.wait()
	
	for file in files: file.close()
	files = []
	procs = []
	
	# Tracing the clones
	print("Tracing the clones")
	for checkpoint in range(start_checkpoint, end_checkpoint + 1):
		checkpoint_dir = f"eval-checkpoint-{checkpoint}-{perturbation_name+sample_instance}"
		files.append(open(f"{checkpoint_dir}/tracing.out", "w"))
		proc = subprocess.Popen(
			[
				"python",
				os.path.abspath("./tinypy_code_tracer.py"),
				os.path.abspath(f"{checkpoint_dir}/robustness_snippets.txt"),
				os.path.abspath(f"{checkpoint_dir}/traced_robustness_snippets.txt"),
			],
			cwd = checkpoint_dir,
			stdout = files[-1],
			stderr = files[-1],
			env = env,
		)
		procs.append(proc)
	
	# Wait for all children
	for proc in procs:
		proc.wait()
	
	for file in files: file.close()
	procs = []
	files = []
	
	# Evaluating on the clones
	print("Evaluating on the clones")
	for checkpoint in range(start_checkpoint, end_checkpoint + 1):
		checkpoint_dir = f"eval-checkpoint-{checkpoint}-{perturbation_name+sample_instance}"
		files.append(open(f"{checkpoint_dir}/evaluation.out", "w"))
		env["CUDA_VISIBLE_DEVICES"] = str((checkpoint-1) % 8)
		
		p = subprocess.Popen(
			[
				"python",
				os.path.abspath("optimus_eval.py"),
				os.path.abspath(f"/data/yb2618/Tiny-Language-Models-Framework/research_project_4/robustness_evolution_study/train-14/on-ds-{ds}/checkpoints/checkpoint_{checkpoint}.00.pth"),
				os.path.abspath(f"{checkpoint_dir}/traced_robustness_snippets.txt"),
				"",
			],
			cwd = checkpoint_dir,
			stdout = files[-1],
			stderr = files[-1],
			env = env,
		)
		procs.append(p)
	
	# Wait for all children
	for proc in procs:
		proc.wait()
	
	for file in files: file.close()


except KeyboardInterrupt:
	print("\nCtrl+C pressed! Terminating all child processes...")
	for proc in procs:
		proc.terminate()  # send SIGTERM
	for proc in procs:
		proc.wait()
	for f in files:
		f.close()
	sys.exit(1)
