litsea-cli 0.6.0

use std::error::Error;
use std::io::{self, BufRead, Write};
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use clap::{Args, Parser, Subcommand};

use litsea::version;
use litsea::{AdaBoost, AveragedPerceptron, Extractor, Language, PosTrainer, Segmenter, Trainer};

/// Arguments for the extract command.
#[derive(Debug, Args)]
#[command(about = "Extract features from a corpus")]
struct ExtractArgs {
    #[arg(short, long, default_value = "japanese", value_parser = Language::from_str)]
    language: Language,

    /// Extract features from a POS-tagged corpus (format: "word/POS word/POS ...")
    #[arg(long)]
    pos: bool,

    corpus_file: PathBuf,
    features_file: PathBuf,
}

/// Arguments for the train command.
#[derive(Debug, Args)]
#[command(about = "Train a segmenter")]
struct TrainArgs {
    #[arg(short, long, default_value = "0.01")]
    threshold: f64,

    #[arg(short = 'i', long, default_value = "100")]
    num_iterations: usize,

    #[arg(short = 'm', long)]
    load_model_uri: Option<String>,

    /// Train a POS tagging model with the Averaged Perceptron
    #[arg(long)]
    pos: bool,

    /// Number of training epochs for the POS model
    #[arg(long, default_value = "10")]
    num_epochs: usize,

    features_file: PathBuf,
    model_file: PathBuf,
}

/// Arguments for the segment command.
#[derive(Debug, Args)]
#[command(about = "Segment a sentence")]
struct SegmentArgs {
    #[arg(short, long, default_value = "japanese", value_parser = Language::from_str)]
    language: Language,

    /// Segment with POS tagging (requires an Averaged Perceptron model)
    #[arg(long)]
    pos: bool,

    model_uri: String,
}

/// Subcommands for litsea CLI.
#[derive(Debug, Subcommand)]
enum Commands {
    Extract(ExtractArgs),
    Train(TrainArgs),
    Segment(SegmentArgs),
}

/// Arguments for the litsea command.
#[derive(Debug, Parser)]
#[command(
    name = "litsea",
    author,
    about = "A morphological analysis command line interface",
    version = version(),
    propagate_version = true,
)]
struct CommandArgs {
    #[command(subcommand)]
    command: Commands,
}

/// Extract features from a corpus file and write them to a specified output file.
/// This function reads sentences from the corpus file, segments them into words,
/// and writes the extracted features to the output file.
///
/// # Arguments
/// * `args` - The arguments for the extract command [`ExtractArgs`].
///
/// # Returns
/// Returns a Result indicating success or failure.
fn extract(args: ExtractArgs) -> Result<(), Box<dyn Error>> {
    let extractor = Extractor::new(args.language);

    if args.pos {
        extractor.extract_with_pos(args.corpus_file.as_path(), args.features_file.as_path())?;
    } else {
        extractor.extract(args.corpus_file.as_path(), args.features_file.as_path())?;
    }

    eprintln!("Feature extraction completed successfully.");
    Ok(())
}

/// Train a segmenter using the provided arguments.
/// This function initializes a Trainer with the specified parameters,
/// loads a model if specified, and trains the model using the features file.
///
/// # Arguments
/// * `args` - The arguments for the train command [`TrainArgs`].
///
/// # Returns
/// Returns a Result indicating success or failure.
async fn train(args: TrainArgs) -> Result<(), Box<dyn Error>> {
    let running = Arc::new(AtomicBool::new(true));
    let r = running.clone();

    ctrlc::set_handler(move || {
        if r.load(Ordering::SeqCst) {
            r.store(false, Ordering::SeqCst);
        } else {
            std::process::exit(0);
        }
    })?;

    if args.pos {
        // Train the POS tagging model with the Averaged Perceptron
        let mut trainer = PosTrainer::new(args.num_epochs, args.features_file.as_path())?;

        if let Some(model_uri) = &args.load_model_uri {
            trainer.load_model(model_uri).await?;
        }

        let metrics = trainer.train(&running, args.model_file.as_path())?;

        eprintln!("Result Metrics (POS):");
        eprintln!("  Accuracy: {:.2}% ( {} )", metrics.accuracy, metrics.num_instances);
        eprintln!("  Macro Precision: {:.2}%", metrics.macro_precision);
        eprintln!("  Macro Recall: {:.2}%", metrics.macro_recall);
    } else {
        // Train the word segmentation model with AdaBoost
        let mut trainer =
            Trainer::new(args.threshold, args.num_iterations, args.features_file.as_path())?;

        if let Some(model_uri) = &args.load_model_uri {
            trainer.load_model(model_uri).await?;
        }

        let metrics = trainer.train(&running, args.model_file.as_path())?;

        eprintln!("Result Metrics:");
        eprintln!(
            "  Accuracy: {:.2}% ( {} / {} )",
            metrics.accuracy,
            metrics.true_positives + metrics.true_negatives,
            metrics.num_instances
        );
        eprintln!(
            "  Precision: {:.2}% ( {} / {} )",
            metrics.precision,
            metrics.true_positives,
            metrics.true_positives + metrics.false_positives
        );
        eprintln!(
            "  Recall: {:.2}% ( {} / {} )",
            metrics.recall,
            metrics.true_positives,
            metrics.true_positives + metrics.false_negatives
        );
        eprintln!(
            "  Confusion Matrix:\n    True Positives: {}\n    False Positives: {}\n    False Negatives: {}\n    True Negatives: {}",
            metrics.true_positives,
            metrics.false_positives,
            metrics.false_negatives,
            metrics.true_negatives
        );
    }

    Ok(())
}

/// Writes one output line, treating a closed downstream pipe as normal
/// termination.
///
/// # Arguments
/// * `writer` - The output writer.
/// * `line` - The line to write (a newline is appended).
///
/// # Returns
/// `Ok(true)` to continue writing, `Ok(false)` when the downstream consumer
/// closed the pipe (e.g. `litsea segment model | head -1`), or the original
/// error for any other I/O failure.
fn write_output_line<W: Write>(writer: &mut W, line: &str) -> io::Result<bool> {
    match writeln!(writer, "{}", line) {
        Ok(()) => Ok(true),
        Err(e) if e.kind() == io::ErrorKind::BrokenPipe => Ok(false),
        Err(e) => Err(e),
    }
}

/// Flushes the output writer, treating a closed downstream pipe as normal
/// termination.
///
/// # Arguments
/// * `writer` - The output writer to flush.
///
/// # Returns
/// `Ok(())` on success or broken pipe; any other I/O error is returned so it
/// is surfaced instead of being lost in the writer's drop.
fn flush_output<W: Write>(writer: &mut W) -> io::Result<()> {
    match writer.flush() {
        Ok(()) => Ok(()),
        Err(e) if e.kind() == io::ErrorKind::BrokenPipe => Ok(()),
        Err(e) => Err(e),
    }
}

/// Segment a sentence using the trained model.
/// This function loads the AdaBoost model from the specified file,
/// reads sentences from standard input, segments them into words,
/// and writes the segmented sentences to standard output.
///
/// A downstream consumer closing stdout early (broken pipe) terminates the
/// command successfully instead of reporting an error.
///
/// # Arguments
/// * `args` - The arguments for the segment command [`SegmentArgs`].
///
/// # Returns
/// Returns a Result indicating success or failure.
async fn segment(args: SegmentArgs) -> Result<(), Box<dyn Error>> {
    let language = args.language;

    let stdin = io::stdin();
    let stdout = io::stdout();
    let mut writer = io::BufWriter::new(stdout.lock());

    if args.pos {
        // Joint segmentation + POS tagging with an Averaged Perceptron model
        let mut pos_learner = AveragedPerceptron::new();
        pos_learner.load_model(args.model_uri.as_str()).await?;

        let segmenter = Segmenter::with_pos_learner(language, pos_learner);

        for line in stdin.lock().lines() {
            let line = line?;
            let line = line.trim();
            if line.is_empty() {
                continue;
            }
            let tokens = segmenter.segment_with_pos(line)?;
            let formatted: Vec<String> =
                tokens.iter().map(|(word, pos)| format!("{}/{}", word, pos)).collect();
            if !write_output_line(&mut writer, &formatted.join(" "))? {
                return Ok(());
            }
        }
    } else {
        // Word segmentation only, with an AdaBoost model
        let mut learner = AdaBoost::new(0.01, 100);
        learner.load_model(args.model_uri.as_str()).await?;

        let segmenter = Segmenter::with_learner(language, learner);

        for line in stdin.lock().lines() {
            let line = line?;
            let line = line.trim();
            if line.is_empty() {
                continue;
            }
            let tokens = segmenter.segment(line);
            if !write_output_line(&mut writer, &tokens.join(" "))? {
                return Ok(());
            }
        }
    }

    flush_output(&mut writer)?;
    Ok(())
}

async fn run() -> Result<(), Box<dyn std::error::Error>> {
    let args = CommandArgs::parse();

    match args.command {
        Commands::Extract(args) => extract(args),
        Commands::Train(args) => train(args).await,
        Commands::Segment(args) => segment(args).await,
    }
}

#[tokio::main]
async fn main() {
    if let Err(e) = run().await {
        eprintln!("Error: {}", e);
        std::process::exit(1);
    }
}