litsea-cli 0.5.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(
    author,
    about = "Extract features from a corpus",
    version = version(),
)]
struct ExtractArgs {
    #[arg(short, long, default_value = "japanese", value_parser = Language::from_str)]
    language: Language,

    /// 品詞付きコーパスから特徴量を抽出する（コーパスフォーマット: "単語/品詞 単語/品詞 ..."）
    #[arg(long, default_value = "false")]
    pos: bool,

    corpus_file: PathBuf,
    features_file: PathBuf,
}

/// Arguments for the train command.
#[derive(Debug, Args)]
#[command(author,
    about = "Train a segmenter",
    version = version(),
)]
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>,

    /// Averaged Perceptronで品詞推定モデルを学習する
    #[arg(long, default_value = "false")]
    pos: bool,

    /// 品詞モデル学習時のエポック数
    #[arg(long, default_value = "10")]
    num_epochs: usize,

    features_file: PathBuf,
    model_file: PathBuf,
}

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

    /// 品詞推定付きで分割する（Averaged Perceptronモデルを使用）
    #[arg(long, default_value = "false")]
    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(),
)]
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 mut 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 {
        // 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 {
        // 既存の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(())
}

/// 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.
///
/// # 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 {
        // Averaged Perceptronモデルで品詞推定付き分割
        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();
            writeln!(writer, "{}", formatted.join(" "))?;
        }
    } else {
        // 既存のAdaBoostモデルで単語分割のみ
        let mut learner = AdaBoost::new(0.01, 100);
        learner.load_model(args.model_uri.as_str()).await?;

        let segmenter = Segmenter::new(language, Some(learner));

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

    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);
    }
}