muse2/

cli.rs

//! Command-line interface for the MUSE2 simulation.
use crate::graph::save_commodity_graphs_for_model;
use crate::input::{load_commodity_graphs, load_model};
use crate::log;
use crate::output::{copy_input_files, create_output_directory, get_graphs_dir, get_output_dir};
use crate::settings::Settings;
use crate::timeit::{DispatchContext, InvestmentContext, TimeContext};
use ::log::{info, warn};
use anyhow::{Context, Result, ensure};
use clap::{Args, CommandFactory, Parser, Subcommand};
use std::path::{self, Path, PathBuf};
use std::time::Instant;

pub mod example;
use example::ExampleSubcommands;

pub mod settings;
use settings::SettingsSubcommands;

/// The command line interface for the simulation.
#[derive(Parser)]
#[command(
    version,
    about,
    after_help = concat!(
        "For more detailed documentation on this version of MUSE2, see: ", crate::docs_url!()
    )
)]
struct Cli {
    /// The available commands.
    #[command(subcommand)]
    command: Option<Commands>,
    /// Flag to provide the CLI docs as markdown
    #[arg(long, hide = true)]
    markdown_help: bool,
}

/// Options for the `run` command
#[derive(Args)]
pub struct RunOpts {
    /// Directory for output files
    #[arg(short, long)]
    pub output_dir: Option<PathBuf>,
    /// Whether to overwrite the output directory if it already exists
    #[arg(long)]
    pub overwrite: bool,
    /// Whether to write additional information to CSV files
    #[arg(long)]
    pub debug_model: bool,

    /// Whether to skip copying input files to the output folder
    #[arg(long)]
    pub no_copy_input_files: bool,
}

/// Options for the `graph` command
#[derive(Args)]
pub struct GraphOpts {
    /// Directory for graph files
    #[arg(short, long)]
    pub output_dir: Option<PathBuf>,
    /// Whether to overwrite the output directory if it already exists
    #[arg(long)]
    pub overwrite: bool,
}

/// The available commands.
#[derive(Subcommand)]
enum Commands {
    /// Run a simulation model.
    Run {
        /// Path to the model directory.
        model_dir: PathBuf,
        /// Other run options
        #[command(flatten)]
        opts: RunOpts,
    },
    /// Manage example models.
    Example {
        /// The available subcommands for managing example models.
        #[command(subcommand)]
        subcommand: ExampleSubcommands,
    },
    /// Validate a model.
    Validate {
        /// The path to the model directory.
        model_dir: PathBuf,
    },
    /// Build and output commodity flow graphs for a model.
    SaveGraphs {
        /// The path to the model directory.
        model_dir: PathBuf,
        /// Other options
        #[command(flatten)]
        opts: GraphOpts,
    },
    /// Manage settings file.
    Settings {
        /// The subcommands for managing the settings file.
        #[command(subcommand)]
        subcommand: SettingsSubcommands,
    },
}

impl Commands {
    /// Execute the supplied CLI command
    fn execute(self) -> Result<()> {
        match self {
            Self::Run { model_dir, opts } => handle_run_command(&model_dir, &opts),
            Self::Example { subcommand } => subcommand.execute(),
            Self::Validate { model_dir } => handle_validate_command(&model_dir),
            Self::SaveGraphs { model_dir, opts } => handle_save_graphs_command(&model_dir, &opts),
            Self::Settings { subcommand } => subcommand.execute(),
        }
    }
}

/// Parse CLI arguments and run the requested command or show help.
pub fn run_cli() -> Result<()> {
    let cli = Cli::parse();

    // Invoked as: `$ muse2 --markdown-help`
    if cli.markdown_help {
        clap_markdown::print_help_markdown::<Cli>();
        return Ok(());
    }

    if let Some(command) = cli.command {
        command.execute()?;
    } else {
        // No command provided. Show help.
        Cli::command().print_long_help()?;
    }

    Ok(())
}

/// Handle the `run` command.
pub fn handle_run_command(model_path: &Path, opts: &RunOpts) -> Result<()> {
    let start = Instant::now();
    let mut settings = Settings::load_or_default().context("Failed to load settings.")?;

    // These settings can be overridden by command-line arguments
    if opts.debug_model {
        settings.debug_model = true;
    }
    if opts.overwrite {
        settings.overwrite = true;
    }
    if opts.no_copy_input_files {
        settings.copy_input_files = false;
    }

    // Get path to output folder
    let pathbuf: PathBuf;
    let output_path = if let Some(p) = opts.output_dir.as_deref() {
        p
    } else {
        pathbuf = get_output_dir(model_path, settings.results_root)?;
        &pathbuf
    };

    ensure!(
        !path::absolute(model_path)?.starts_with(path::absolute(output_path)?),
        "Model input data cannot be inside output folder"
    );

    let overwrite =
        create_output_directory(output_path, settings.overwrite).with_context(|| {
            format!(
                "Failed to create output directory: {}",
                output_path.display()
            )
        })?;

    let model_path = model_path
        .canonicalize()
        .context("Failed to resolve model path.")?;
    let model_name = model_path
        .file_name()
        .context("Model cannot be the root directory.")?
        .to_str()
        .context("Invalid chars in model directory name")?;

    if settings.copy_input_files {
        copy_input_files(&model_path, output_path, model_name)
            .context("Failed to copy input files to output directory.")?;
    }
    // Initialise program logger
    if !cfg!(feature = "bench") {
        // Don't initialise logging when running bench, as it will interfere with the process
        log::init(&settings.log_level, Some(output_path))
            .context("Failed to initialise logging.")?;
    }
    info!("Starting MUSE2 v{}", env!("CARGO_PKG_VERSION"));

    // Load the model to run
    let model = load_model(&model_path).context("Failed to load model.")?;
    info!("Loaded model from {}", model_path.display());
    info!("Output folder: {}", output_path.display());

    // NB: We have to wait until the logger is initialised to display this warning
    if overwrite {
        warn!("Output folder will be overwritten");
    }

    // Run the simulation
    crate::simulation::run(&model, output_path, settings.debug_model)?;
    info!("Simulation complete!");

    info!(
        "--- Total execution time: {:.1}ms",
        start.elapsed().as_secs_f64() * 1000.0
    );

    info!(
        "--- Total time spent in investment steps: {:.1}ms",
        InvestmentContext::total_time()
    );
    info!(
        "--- Total time spent in dispatch steps: {:.1}ms",
        DispatchContext::total_time()
    );

    Ok(())
}

/// Handle the `validate` command.
pub fn handle_validate_command(model_path: &Path) -> Result<()> {
    let settings = Settings::load_or_default().context("Failed to load settings.")?;

    // Initialise program logger (we won't save log files when running the validate command)
    log::init(&settings.log_level, None).context("Failed to initialise logging.")?;

    // Load/validate the model
    load_model(model_path).context("Failed to validate model.")?;
    info!("Model validation successful!");

    Ok(())
}

/// Handle the `save-graphs` command.
pub fn handle_save_graphs_command(model_path: &Path, opts: &GraphOpts) -> Result<()> {
    let mut settings = Settings::load_or_default().context("Failed to load settings.")?;

    if opts.overwrite {
        settings.overwrite = true;
    }

    // Get path to output folder
    let pathbuf: PathBuf;
    let output_path = if let Some(p) = opts.output_dir.as_deref() {
        p
    } else {
        pathbuf = get_graphs_dir(model_path, settings.graph_results_root)?;
        &pathbuf
    };

    let overwrite =
        create_output_directory(output_path, settings.overwrite).with_context(|| {
            format!(
                "Failed to create graphs directory: {}",
                output_path.display()
            )
        })?;

    // Initialise program logger (we won't save log files when running this command)
    log::init(&settings.log_level, None).context("Failed to initialise logging.")?;

    // NB: We have to wait until the logger is initialised to display this warning
    if overwrite {
        warn!("Graphs directory will be overwritten");
    }

    // Load commodity flow graphs and save to file
    let commodity_graphs = load_commodity_graphs(model_path).context("Failed to build graphs.")?;
    save_commodity_graphs_for_model(&commodity_graphs, output_path)?;
    info!("Graphs saved to: {}", output_path.display());

    Ok(())
}