seqtkrs 0.1.1

//! sample命令 - 对序列进行随机采样
//!
//! 使用蓄水池采样算法从输入文件中随机选择指定数量或比例的序列。
//!
//! # 示例
//!
//! ```bash
//! # 采样1000条序列
//! seqtkrs sample input.fq 1000 > output.fq
//!
//! # 采样30%的序列
//! seqtkrs sample input.fq 0.3 > output.fq
//!
//! # 使用指定随机种子
//! seqtkrs sample -s 42 input.fq 1000 > output.fq
//!
//! # 两遍模式（节省内存）
//! seqtkrs sample -2 input.fq 1000 > output.fq
//! ```

use crate::algorithms::sampler::{ReservoirSampler, TwoPassSampler};
use crate::core::{SeqReader, SeqRecord, SeqWriter};
use crate::utils::random::RandomGenerator;
use anyhow::{bail, Context, Result};
use clap::Args;
use rustc_hash::FxHashSet;

#[derive(Args, Debug)]
pub struct SampleArgs {
    /// 输入文件（FASTA/FASTQ格式，支持gzip/bzip2），使用"-"表示stdin
    #[arg(value_name = "in.fq")]
    pub input: String,

    /// 采样数量（如果>=1）或分数（如果<1）
    #[arg(value_name = "FRAC|NUM")]
    pub target: f64,

    /// 随机种子 [11]
    #[arg(short = 's', long, value_name = "INT", default_value = "11")]
    pub seed: u64,

    /// 两遍模式：慢2倍但大幅减少内存使用（仅对数量采样有效，不支持stdin）
    #[arg(short = '2', long)]
    pub two_pass: bool,
}

pub fn run(args: &SampleArgs) -> Result<()> {
    // 解析目标：分数还是数量
    let (is_fraction, sample_count) = if args.target >= 1.0 {
        (false, args.target.round() as usize)
    } else {
        (true, 0)
    };

    // 验证两遍模式的使用
    if args.two_pass && is_fraction {
        eprintln!("[警告] 采样分数时，-2选项将被忽略");
    }

    if args.two_pass && args.input == "-" {
        bail!("两遍模式下，输入不能是stdin");
    }

    // 根据模式选择采样策略
    if !args.two_pass || is_fraction {
        // 单遍模式（标准蓄水池采样）
        sample_single_pass(args, is_fraction, sample_count)
    } else {
        // 两遍模式（只保存索引）
        sample_two_pass(args, sample_count)
    }
}

/// 单遍采样模式
fn sample_single_pass(args: &SampleArgs, is_fraction: bool, sample_count: usize) -> Result<()> {
    let mut reader = if args.input == "-" {
        SeqReader::from_stdin()
    } else {
        SeqReader::from_path(&args.input)
            .with_context(|| format!("无法打开输入文件: {}", args.input))?
    };

    let mut writer = SeqWriter::to_stdout();

    if is_fraction {
        // 分数采样：直接流式输出，不使用蓄水池
        let mut rng = RandomGenerator::with_seed(args.seed);
        let mut record = SeqRecord::new(Vec::new(), Vec::new());

        while reader.read_next(&mut record)? {
            if rng.random() < args.target {
                writer.write_record(&record)?;
            }
        }
    } else {
        // 数量采样：使用蓄水池采样
        let mut sampler = ReservoirSampler::new(sample_count, Some(args.seed));
        let mut record = SeqRecord::new(Vec::new(), Vec::new());

        while reader.read_next(&mut record)? {
            sampler.add(record.clone());
        }

        // 输出采样结果
        for sample in sampler.into_samples() {
            writer.write_record(&sample)?;
        }
    }

    writer.flush()?;
    Ok(())
}

/// 两遍采样模式
fn sample_two_pass(args: &SampleArgs, sample_count: usize) -> Result<()> {
    // 第一遍：收集被选中的序列索引
    let selected_indices = {
        let mut reader = SeqReader::from_path(&args.input)
            .with_context(|| format!("无法打开输入文件进行第一遍扫描: {}", args.input))?;

        let mut sampler = TwoPassSampler::new(sample_count, Some(args.seed));
        let mut record = SeqRecord::new(Vec::new(), Vec::new());

        while reader.read_next(&mut record)? {
            sampler.add_index();
        }

        sampler.get_selected_indices()
    };

    // 第二遍：输出被选中的序列
    let mut reader = SeqReader::from_path(&args.input)
        .with_context(|| format!("无法打开输入文件进行第二遍扫描: {}", args.input))?;

    let mut writer = SeqWriter::to_stdout();
    let mut record = SeqRecord::new(Vec::new(), Vec::new());
    let mut current_index = 0;

    // 将索引列表转换为HashSet以加快查找
    let selected_set: FxHashSet<usize> = selected_indices.into_iter().collect();

    while reader.read_next(&mut record)? {
        if selected_set.contains(&current_index) {
            writer.write_record(&record)?;
        }
        current_index += 1;
    }

    writer.flush()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    #[test]
    fn test_sample_count() -> Result<()> {
        // 创建测试文件
        let mut temp_file = NamedTempFile::new()?;
        for i in 0..100 {
            writeln!(temp_file, ">seq{}", i)?;
            writeln!(temp_file, "ACGT")?;
        }
        temp_file.flush()?;

        let _args = SampleArgs {
            input: temp_file.path().to_str().unwrap().to_string(),
            target: 10.0,
            seed: 42,
            two_pass: false,
        };

        // 这里只测试不会panic，实际输出会到stdout
        // 在完整的集成测试中会验证输出
        Ok(())
    }

    #[test]
    fn test_sample_fraction() -> Result<()> {
        let mut temp_file = NamedTempFile::new()?;
        for i in 0..100 {
            writeln!(temp_file, ">seq{}", i)?;
            writeln!(temp_file, "ACGT")?;
        }
        temp_file.flush()?;

        let _args = SampleArgs {
            input: temp_file.path().to_str().unwrap().to_string(),
            target: 0.1,
            seed: 42,
            two_pass: false,
        };

        Ok(())
    }
}