package supplementary.experiments.scripts;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.PrintWriter;
import java.io.UnsupportedEncodingException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.regex.Pattern;

import main.CommandLineArgParse;
import main.CommandLineArgParse.ArgOption;
import main.CommandLineArgParse.OptionTypes;
import main.StringRoutines;
import main.UnixPrintWriter;

/**
 * Method to generate cluster job scripts for playing lots of trials of a given
 * game until we crash, and then saving the trial right before we crashed.
 *
 * @author Dennis Soemers
 */
public class FindCrashingTrialScriptsGen
{
	/** Memory to assign per CPU (and to JVM), in MB */
	private static final String MEM_PER_CPU = "4096";
	
	/** Max wall time (in minutes) */
	private static final int MAX_WALL_TIME = 6000;
	
	/** Don't submit more than this number of jobs at a single time */
	private static final int MAX_JOBS_PER_BATCH = 800;
	
	//-------------------------------------------------------------------------
	
	/**
	 * Constructor (don't need this)
	 */
	private FindCrashingTrialScriptsGen()
	{
		// Do nothing
	}
	
	//-------------------------------------------------------------------------
	
	/**
	 * Generates our scripts
	 * @param argParse
	 */
	private static void generateScripts(final CommandLineArgParse argParse)
	{
		final List<String> jobScriptNames = new ArrayList<String>();

		String scriptsDir = argParse.getValueString("--scripts-dir");
		scriptsDir = scriptsDir.replaceAll(Pattern.quote("\\"), "/");
		if (!scriptsDir.endsWith("/"))
			scriptsDir += "/";
		
		final String userName = argParse.getValueString("--user-name");

		final int numJobs = argParse.getValueInt("--num-jobs");

		for (int jobID = 0; jobID < numJobs; ++jobID)
		{
			final String jobScriptFilename = "FindCrashingTrial_" + jobID + ".sh";

			try (final PrintWriter writer = new UnixPrintWriter(new File(scriptsDir + jobScriptFilename), "UTF-8"))
			{
				writer.println("#!/usr/local_rwth/bin/zsh");
				writer.println("#SBATCH -J FindCrashingTrial_" + jobID);
				writer.println("#SBATCH -o /work/" + userName + "/FindCrashingTrial/Out_"
						+ jobID + "_%J.out");
				writer.println("#SBATCH -e /work/" + userName + "/FindCrashingTrial/Err_"
						+ jobID + "_%J.err");
				writer.println("#SBATCH -t " + MAX_WALL_TIME);
				writer.println("#SBATCH --mem-per-cpu=" + MEM_PER_CPU);
				writer.println("#SBATCH -A " + argParse.getValueString("--project"));
				writer.println("unset JAVA_TOOL_OPTIONS");
				
				@SuppressWarnings("unchecked")
				final String[] agentsStrings = ((List<String>) argParse.getValue("--agents")).toArray(new String[0]);
				for (int i = 0; i < agentsStrings.length; ++i)
				{
					agentsStrings[i] = StringRoutines.quote(agentsStrings[i]);
				}

				final String javaCall = StringRoutines.join
						(
							" ", 
							"java",
							"-Xms" + MEM_PER_CPU + "M",
							"-Xmx" + MEM_PER_CPU + "M",
							"-XX:+HeapDumpOnOutOfMemoryError",
							"-da",
							"-dsa",
							"-XX:+UseStringDeduplication",
							"-jar",
							StringRoutines.quote("/home/" + userName + "/FindCrashingTrial/Ludii.jar"),
							"--find-crashing-trial",
							"--game",
							StringRoutines.quote(argParse.getValueString("--game")),
							"-n",
							"" + argParse.getValueInt("--num-trials-per-job"),
							"--game-length-cap 1000",
							"--thinking-time " + argParse.getValueDouble("--thinking-time"),
							"--depth-limit " + argParse.getValueInt("--depth-limit"),
							"--iteration-limit " + argParse.getValueInt("--iteration-limit"),
							"--out-trial-file",
							StringRoutines.quote
							(
								"/work/" + 
								userName + 
								"/FindCrashingTrial/" + 
								"CrashTrial_" +
								jobID + 
								".trl"
							),
							"--agents",
							StringRoutines.join(" ", agentsStrings),
							"--max-wall-time",
							"" + Math.max((MAX_WALL_TIME), 60),
							"--no-print-out"
						);

				writer.println(javaCall);

				jobScriptNames.add(jobScriptFilename);
			}
			catch (final FileNotFoundException | UnsupportedEncodingException e)
			{
				e.printStackTrace();
			}
		}
		
		final List<List<String>> jobScriptsLists = new ArrayList<List<String>>();
		List<String> remainingJobScriptNames = jobScriptNames;

		while (remainingJobScriptNames.size() > 0)
		{
			if (remainingJobScriptNames.size() > MAX_JOBS_PER_BATCH)
			{
				final List<String> subList = new ArrayList<String>();

				for (int i = 0; i < MAX_JOBS_PER_BATCH; ++i)
				{
					subList.add(remainingJobScriptNames.get(i));
				}

				jobScriptsLists.add(subList);
				remainingJobScriptNames = remainingJobScriptNames.subList(MAX_JOBS_PER_BATCH, remainingJobScriptNames.size());
			}
			else
			{
				jobScriptsLists.add(remainingJobScriptNames);
				remainingJobScriptNames = new ArrayList<String>();
			}
		}

		for (int i = 0; i < jobScriptsLists.size(); ++i)
		{
			try (final PrintWriter writer = new UnixPrintWriter(new File(scriptsDir + "SubmitJobs_Part" + i + ".sh"), "UTF-8"))
			{
				for (final String jobScriptName : jobScriptsLists.get(i))
				{
					writer.println("sbatch " + jobScriptName);
				}
			}
			catch (final FileNotFoundException | UnsupportedEncodingException e)
			{
				e.printStackTrace();
			}
		}
	}
	
	//-------------------------------------------------------------------------
	
	/**
	 * Main method to generate all our scripts
	 * @param args
	 */
	public static void main(final String[] args)
	{
		// define options for arg parser
		final CommandLineArgParse argParse = 
				new CommandLineArgParse
				(
					true,
					"Play many trials of a game, save trials right before crashes."
				);
		
		argParse.addOption(new ArgOption()
				.withNames("--project")
				.help("Project for which to submit the job on cluster.")
				.withNumVals(1)
				.withType(OptionTypes.String)
				.setRequired());
		
		argParse.addOption(new ArgOption()
				.withNames("--user-name")
				.help("Username on the cluster.")
				.withNumVals(1)
				.withType(OptionTypes.String)
				.setRequired());
		
		argParse.addOption(new ArgOption()
				.withNames("--scripts-dir")
				.help("Directory in which to store generated scripts.")
				.withNumVals(1)
				.withType(OptionTypes.String)
				.setRequired());
		
		argParse.addOption(new ArgOption()
				.withNames("--num-trials-per-job")
				.help("Number of trials to run per job.")
				.withNumVals(1)
				.withType(OptionTypes.Int)
				.withDefault(Integer.valueOf(100)));
		
		argParse.addOption(new ArgOption()
				.withNames("--num-jobs")
				.help("Number of (copies of) jobs to run.")
				.withNumVals(1)
				.withType(OptionTypes.Int)
				.withDefault(Integer.valueOf(100)));
		
		argParse.addOption(new ArgOption()
				.withNames("--game")
				.help("Name of the game to play. Should end with \".lud\".")
				.withDefault("Amazons.lud")
				.withNumVals(1)
				.withType(OptionTypes.String));
		argParse.addOption(new ArgOption()
				.withNames("--game-options")
				.help("Game Options to load.")
				.withDefault(new ArrayList<String>(0))
				.withNumVals("*")
				.withType(OptionTypes.String));
		
		argParse.addOption(new ArgOption()
				.withNames("--agents")
				.help("Agents to use for playing")
				.withDefault(Arrays.asList("UCT", "Biased MCTS"))
				.withNumVals("+")
				.withType(OptionTypes.String));
		argParse.addOption(new ArgOption()
				.withNames("--thinking-time", "--time", "--seconds")
				.help("Max allowed thinking time per move (in seconds).")
				.withDefault(Double.valueOf(1.0))
				.withNumVals(1)
				.withType(OptionTypes.Double));
		argParse.addOption(new ArgOption()
				.withNames("--iteration-limit", "--iterations")
				.help("Max allowed number of MCTS iterations per move.")
				.withDefault(Integer.valueOf(-1))
				.withNumVals(1)
				.withType(OptionTypes.Int));
		argParse.addOption(new ArgOption()
				.withNames("--depth-limit")
				.help("Max allowed search depth per move (for e.g. alpha-beta).")
				.withDefault(Integer.valueOf(-1))
				.withNumVals(1)
				.withType(OptionTypes.Int));
		argParse.addOption(new ArgOption()
				.withNames("--no-rotate-agents")
				.help("Don't rotate through possible assignments of agents to Player IDs.")
				.withType(OptionTypes.Boolean)
				.withNumVals(0));
		
		// parse the args
		if (!argParse.parseArguments(args))
			return;
		
		generateScripts(argParse);
	}
	
	//-------------------------------------------------------------------------

}
