1 files changed, 189 insertions, 44 deletions
diff --git a/src/main.rs b/src/main.rs
index 1404a9f..b6a4ad0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,8 @@
-use std::{path::Path, time::Duration};
+use std::{path::PathBuf, time::Duration};
 
-use crate::{board::Turn, claude_player::ClaudePlayer, game::GameHandler, local_player::LocalPlayer, network::Network, neural_player::NeuralPlayer, player::Player};
+use clap::{Args, Parser, Subcommand, ValueEnum};
+
+use crate::{claude_player::ClaudePlayer, game::GameHandler, local_player::LocalPlayer, network::Network, neural_player::NeuralPlayer, player::Player};
 
 mod board;
 mod player;
@@ -17,63 +19,206 @@ mod network;
 mod neural_player;
 mod train;
 
-fn main() -> std::io::Result<()> {
-	let mut args: Vec<String> = std::env::args().collect();
-	// `--from <weights>` can go anywhere after `train`
-	let from = args.iter().position(|a| a == "--from").map(|i|
+/// Tigers and Sheep, a hunt game on a 5x5 board. With no command, plays a game in the terminal.
+#[derive(Parser)]
+#[command(version, args_conflicts_with_subcommands = true)]
+struct Cli
+{
+	#[command(subcommand)]
+	command: Option<Command>,
+	#[command(flatten)]
+	play: PlayArgs,
+}
+
+#[derive(Subcommand)]
+enum Command
+{
+	/// Play a game in the terminal (the default)
+	Play(PlayArgs),
+	/// Play against a trained network
+	Neural
 	{
-		let path = args.get(i + 1).cloned().expect("usage: --from <weights file>");
-		args.drain(i..=i + 1);
-		path
-	});
-	let arg = |i: usize| args.get(i).map(String::as_str);
-	let number = |i: usize, default: usize| arg(i).map_or(default, |n| n.parse().expect("expected a number"));
-	match arg(1)
+		/// Weights file written by `train`
+		weights: PathBuf,
+		/// The side you play
+		#[arg(value_enum, default_value_t = Side::Sheep)]
+		side: Side,
+		/// Seconds the network may think per move
+		#[arg(short, long, default_value_t = 2.0)]
+		time: f64,
+	},
+	/// Time the computer player's search on fixed positions, without the TUI
+	Bench
 	{
-		// `bench [depth]` times the computer player's search without the TUI
-		Some("bench") =>
+		/// Search depth
+		#[arg(default_value_t = 6)]
+		depth: usize,
+	},
+	/// Train a network by self-play
+	///
+	/// Starts from `--from`, or else the newest gen-NNN.bin in the out dir, or else random weights.
+	Train(TrainArgs),
+}
+
+#[derive(Args)]
+struct PlayArgs
+{
+	/// Who plays the sheep: `human`, `claude`, or a weights file for a neural player
+	#[arg(short = 'S', long, default_value = "claude", value_parser = parse_player)]
+	sheep: PlayerKind,
+	/// Who plays the tigers: `human`, `claude`, or a weights file for a neural player
+	#[arg(short = 'T', long, default_value = "claude", value_parser = parse_player)]
+	tiger: PlayerKind,
+	/// Seconds each computer player may think per move
+	#[arg(short, long, default_value_t = 2.0)]
+	time: f64,
+}
+
+#[derive(Clone)]
+enum PlayerKind
+{
+	Human,
+	Claude,
+	Neural(PathBuf),
+}
+
+fn parse_player(s: &str) -> Result<PlayerKind, String>
+{
+	match s
+	{
+		"human" => Ok(PlayerKind::Human),
+		"claude" => Ok(PlayerKind::Claude),
+		path if PathBuf::from(path).is_file() => Ok(PlayerKind::Neural(path.into())),
+		_ => Err(format!("expected `human`, `claude` or a weights file, but `{s}` is none of these")),
+	}
+}
+
+#[derive(Clone, Copy, ValueEnum)]
+enum Side
+{
+	Sheep,
+	Tiger,
+}
+
+#[derive(Args)]
+struct TrainArgs
+{
+	/// Generations to train
+	#[arg(default_value_t = train::Config::default().generations)]
+	generations: usize,
+	/// Self-play games per generation
+	#[arg(default_value_t = train::Config::default().games_per_generation)]
+	games: usize,
+	/// Where each generation's weights are saved
+	#[arg(default_value = "weights")]
+	out_dir: PathBuf,
+	/// Start from this weights file
+	#[arg(long)]
+	from: Option<PathBuf>,
+	/// Search depth for every move in self-play and matches
+	#[arg(long, default_value_t = train::Config::default().depth)]
+	depth: usize,
+	/// A game this many plies long is a draw
+	#[arg(long, default_value_t = train::Config::default().ply_cap)]
+	ply_cap: usize,
+	/// Plies of random moves at the start of each self-play game
+	#[arg(long, default_value_t = train::Config::default().random_plies)]
+	random_plies: usize,
+	/// Chance of a random move after that
+	#[arg(long, default_value_t = train::Config::default().epsilon)]
+	epsilon: f32,
+	/// TD(λ): 1 trains on each game's result alone, smaller trusts later search scores more
+	#[arg(long, default_value_t = train::Config::default().lambda)]
+	lambda: f32,
+	/// Train on the samples of this many of the latest generations
+	#[arg(long, default_value_t = train::Config::default().buffer_generations)]
+	buffer_generations: usize,
+	/// Don't balance positions from sheep-win and tiger-win games
+	#[arg(long)]
+	no_balance: bool,
+	/// Passes over the samples per generation
+	#[arg(long, default_value_t = train::Config::default().epochs)]
+	epochs: usize,
+	/// Positions per optimiser step
+	#[arg(long, default_value_t = train::Config::default().batch_size)]
+	batch_size: usize,
+	/// Adam's learning rate
+	#[arg(long, default_value_t = train::Config::default().learning_rate)]
+	learning_rate: f32,
+	/// Games against ClaudePlayer, and against the previous generation, after each generation
+	#[arg(long, default_value_t = train::Config::default().match_games)]
+	match_games: usize,
+}
+
+impl From<TrainArgs> for train::Config
+{
+	fn from(args: TrainArgs) -> Self
+	{
+		Self
 		{
-			bench::run(number(2, 6));
-			return Ok(());
+			generations: args.generations,
+			games_per_generation: args.games,
+			out_dir: args.out_dir,
+			from: args.from,
+			depth: args.depth,
+			ply_cap: args.ply_cap,
+			random_plies: args.random_plies,
+			epsilon: args.epsilon,
+			lambda: args.lambda,
+			buffer_generations: args.buffer_generations,
+			balance_outcomes: !args.no_balance,
+			epochs: args.epochs,
+			batch_size: args.batch_size,
+			learning_rate: args.learning_rate,
+			match_games: args.match_games,
 		}
-		// `train [generations] [games per generation] [out dir] [--from <weights>]` trains a network
-		// by self-play, starting from `--from`, or the newest generation in the out dir, or random weights
-		Some("train") =>
+	}
+}
+
+fn main() -> std::io::Result<()>
+{
+	let cli = Cli::parse();
+	match cli.command
+	{
+		Some(Command::Bench { depth }) =>
 		{
-			let defaults = train::Config::default();
-			let cfg = train::Config
-			{
-				generations: number(2, defaults.generations),
-				games_per_generation: number(3, defaults.games_per_generation),
-				out_dir: arg(4).map_or(defaults.out_dir.clone(), Into::into),
-				from: from.map(Into::into),
-				..defaults
-			};
-			return train::run(&cfg);
+			bench::run(depth);
+			Ok(())
 		}
-		// `neural <weights> [sheep|tiger]`: you play the given side against a trained network
-		Some("neural") =>
+		Some(Command::Train(args)) => train::run(&args.into()),
+		Some(Command::Neural { weights, side, time }) =>
 		{
-			let net = Network::load(Path::new(arg(2).expect("usage: neural <weights file> [sheep|tiger]")))?;
-			let side = match arg(3) { Some("tiger") => Turn::Tiger, _ => Turn::Sheep };
-			let mut computer = NeuralPlayer::new(net, Duration::from_secs(2), 20);
-			let mut human = LocalPlayer::default();
-			let (sheep, tiger): (&mut dyn Player, &mut dyn Player) = match side
+			let human = PlayerKind::Human;
+			let computer = PlayerKind::Neural(weights);
+			let (sheep, tiger) = match side
 			{
-				Turn::Sheep => (&mut human, &mut computer),
-				Turn::Tiger => (&mut computer, &mut human),
+				Side::Sheep => (human, computer),
+				Side::Tiger => (computer, human),
 			};
-			return play(sheep, tiger);
+			play(PlayArgs { sheep, tiger, time })
 		}
-		_ => (),
+		Some(Command::Play(args)) => play(args),
+		None => play(cli.play),
 	}
-	play(&mut ClaudePlayer::new(Duration::from_secs(2)), &mut ClaudePlayer::new(Duration::from_secs(2)))
 }
 
-fn play(sheep: &mut dyn Player, tiger: &mut dyn Player) -> std::io::Result<()>
+fn make_player(kind: PlayerKind, time_limit: Duration) -> std::io::Result<Box<dyn Player>>
+{
+	Ok(match kind
+	{
+		PlayerKind::Human => Box::new(LocalPlayer::default()),
+		PlayerKind::Claude => Box::new(ClaudePlayer::new(time_limit)),
+		PlayerKind::Neural(path) => Box::new(NeuralPlayer::new(Network::load(&path)?, time_limit, 20)),
+	})
+}
+
+fn play(args: PlayArgs) -> std::io::Result<()>
 {
+	let time_limit = Duration::from_secs_f64(args.time);
+	let mut sheep = make_player(args.sheep, time_limit)?;
+	let mut tiger = make_player(args.tiger, time_limit)?;
 	let mut game = GameHandler::default();
-	if let Some(winner) = game.play_game(sheep, tiger)?
+	if let Some(winner) = game.play_game(sheep.as_mut(), tiger.as_mut())?
 	{
 		println!("{winner:?} wins!");
 	}