use crate::kernel::muskingum::MuskingumCungeKernel;
use anyhow::{Context, Result};
use clap::Parser;
use colored::Colorize;
use num_cpus;
use std::path::PathBuf;
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
route_dir: PathBuf,
#[arg(long)]
hf: Option<PathBuf>,
#[arg(short = 'i', long)]
input_dir: Option<PathBuf>,
#[arg(short, long)]
output_dir: Option<PathBuf>,
#[arg(short = 't', long, default_value_t = 300)]
internal_timestep_seconds: usize,
#[arg(short, long, default_value_t = MuskingumCungeKernel::TRouteModernized)]
kernel: MuskingumCungeKernel,
#[arg(short, long, default_value_t = num_cpus::get())]
num_threads: usize,
}
pub fn print_banner(config: &Config) {
eprintln!(" {}", "🌊 Route RS".cyan().bold());
eprintln!(" Kernel: {}", format!("{}", config.kernel).green());
eprintln!(" Timestep: {}s", config.internal_timestep_seconds);
eprintln!(" Threads: {}", config.num_threads);
eprintln!(
" GeoPackage: {}",
config.gpkg_file.display().to_string().dimmed()
);
eprintln!();
}
pub struct Config {
pub config_dir: PathBuf,
pub csv_dir: PathBuf,
pub gpkg_file: PathBuf,
pub internal_timestep_seconds: usize,
pub output_dir: PathBuf,
pub kernel: MuskingumCungeKernel,
pub num_threads: usize,
}
pub fn get_args() -> Result<Config> {
let args = Args::parse();
let root_dir = args.route_dir;
let csv_dir = args.input_dir.unwrap_or_else(|| root_dir.join("outputs").join("ngen"));
let config_dir = root_dir.join("config");
let output_dir = args.output_dir.unwrap_or_else(|| root_dir.join("outputs").join("troute"));
if !root_dir.exists() || !root_dir.is_dir() {
return Err(anyhow::anyhow!(
"Given root directory does not exist or is not a directory: {:?}",
root_dir
))
.with_context(|| format!("Failed to access root directory: {:?}", root_dir));
}
let dirs_to_check: Vec<&PathBuf> = if args.hf.is_some() {
vec![&csv_dir, &output_dir]
} else {
vec![&csv_dir, &config_dir, &output_dir]
};
let mut missing_dirs = Vec::new();
for dir in dirs_to_check {
if !dir.exists() || !dir.is_dir() {
missing_dirs.push(dir);
}
}
if !missing_dirs.is_empty() {
return Err(anyhow::anyhow!(
"Missing required directories: {:?}",
missing_dirs
))
.with_context(|| format!("Failed to access required directories: {:?}", missing_dirs));
}
let gpkg_file = if let Some(hf) = args.hf {
if !hf.exists() {
return Err(anyhow::anyhow!(
"Specified hydrofabric file does not exist: {:?}",
hf
));
}
hf
} else {
config_dir
.read_dir()
.context("Failed to read config directory")?
.filter_map(Result::ok)
.find(|entry| entry.path().extension().map_or(false, |ext| ext == "gpkg"))
.ok_or_else(|| anyhow::anyhow!("No .gpkg file found in config directory"))?
.path()
};
let cfg = Config {
config_dir,
csv_dir,
gpkg_file,
internal_timestep_seconds: args.internal_timestep_seconds,
output_dir,
kernel: args.kernel,
num_threads: args.num_threads,
};
print_banner(&cfg);
Ok(cfg)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_args_parsing_defaults() {
let args = Args::parse_from(["test", "test_route_dir"]);
assert_eq!(args.route_dir, PathBuf::from("test_route_dir"));
assert_eq!(args.internal_timestep_seconds, 300);
match args.kernel {
MuskingumCungeKernel::TRouteModernized => {}
_ => panic!("Expected default kernel to be TRouteModernized"),
}
}
#[test]
fn test_args_parsing_custom() {
let args = Args::parse_from([
"test",
"test_route_dir",
"-t",
"600",
"-k",
"t-route-legacy",
]);
assert_eq!(args.route_dir, PathBuf::from("test_route_dir"));
assert_eq!(args.internal_timestep_seconds, 600);
match args.kernel {
MuskingumCungeKernel::TRouteLegacy => {}
_ => panic!("Expected kernel to be TRouteLegacy"),
}
}
}