seqtkrs 0.1.1

// fqchk 命令：FASTQ 文件质量检查
//
// 该命令统计 FASTQ 文件的详细质量信息，包括：
// - 每个位置的碱基分布（%A, %C, %G, %T, %N）
// - 每个位置的质量值分布
// - 平均质量和错误率质量
// - 低质量/高质量百分比

use anyhow::Result;
use clap::Args;

use crate::core::tables::LOOKUP_TABLES;
use crate::core::{SeqReader, SeqRecord};

#[derive(Args, Debug)]
pub struct FqchkArgs {
    /// 输入的 FASTQ 文件（"-" 表示标准输入）
    #[arg(value_name = "in.fq")]
    pub input: String,

    /// 质量阈值，用于区分低质量和高质量碱基（设为0显示所有质量值分布）
    #[arg(short = 'q', default_value = "20")]
    pub quality_threshold: usize,
}

/// 位置统计结构
#[derive(Debug, Clone)]
struct PosStats {
    /// 质量值分布 (0-93)
    quality: [u64; 94],
    /// 碱基分布 (A, C, G, T, N)
    bases: [u64; 5],
}

impl PosStats {
    fn new() -> Self {
        Self {
            quality: [0; 94],
            bases: [0; 5],
        }
    }

    /// 添加一个位置的统计
    fn add(&mut self, base: u8, qual: u8, offset: u8) {
        // 统计质量值
        let q = (qual.saturating_sub(offset) as usize).min(93);
        self.quality[q] += 1;

        // 统计碱基
        let nt6 = LOOKUP_TABLES.nt6[base as usize];
        let base_idx = if (1..=4).contains(&nt6) {
            (nt6 - 1) as usize
        } else {
            4 // N 或其他
        };
        self.bases[base_idx] += 1;
    }

    /// 合并另一个统计
    fn merge(&mut self, other: &PosStats) {
        for i in 0..94 {
            self.quality[i] += other.quality[i];
        }
        for i in 0..5 {
            self.bases[i] += other.bases[i];
        }
    }

    /// 获取碱基总数
    fn total_bases(&self) -> u64 {
        self.bases.iter().sum()
    }

    /// 计算平均质量值
    fn avg_quality(&self) -> f64 {
        let total: u64 = self.quality.iter().sum();
        if total == 0 {
            return 0.0;
        }

        let sum: u64 = self
            .quality
            .iter()
            .enumerate()
            .map(|(q, &count)| q as u64 * count)
            .sum();

        sum as f64 / total as f64
    }

    /// 计算错误率质量（error rate quality）
    fn error_quality(&self, perr: &[f64; 94]) -> f64 {
        let total: u64 = self.quality.iter().sum();
        if total == 0 {
            return 0.0;
        }

        let psum: f64 = self
            .quality
            .iter()
            .enumerate()
            .map(|(q, &count)| count as f64 * perr[q])
            .sum();

        -4.343 * ((psum + 1e-6) / (total as f64 + 1e-6)).ln()
    }
}

pub fn run(args: &FqchkArgs) -> Result<()> {
    // 创建输入读取器
    let mut reader = if args.input == "-" {
        SeqReader::from_stdin()
    } else {
        SeqReader::from_path(&args.input)?
    };

    // 统计变量
    let mut n_reads = 0u64;
    let mut tot_len = 0u64;
    let mut min_len = usize::MAX;
    let mut max_len = 0usize;

    // 位置统计数组（动态增长）
    let mut pos_stats: Vec<PosStats> = Vec::new();

    // 错误概率表
    let mut perr = [0.0f64; 94];
    for (k, item) in perr.iter_mut().enumerate() {
        *item = 10.0f64.powf(-0.1 * k as f64);
    }
    // 前几个质量值的错误概率设为0.5
    perr[0] = 0.5;
    perr[1] = 0.5;
    perr[2] = 0.5;
    perr[3] = 0.5;

    // 质量值偏移（Phred+33）
    let offset = 33u8;

    // 创建记录缓冲区
    let mut record = SeqRecord::new(Vec::new(), Vec::new());

    // 第一遍：读取所有序列并统计
    while reader.read_next(&mut record)? {
        // 跳过没有质量值的序列
        if !record.is_fastq() {
            continue;
        }

        n_reads += 1;
        let seq_len = record.seq.len();
        tot_len += seq_len as u64;
        min_len = min_len.min(seq_len);
        max_len = max_len.max(seq_len);

        // 扩展位置统计数组
        while pos_stats.len() < seq_len {
            pos_stats.push(PosStats::new());
        }

        // 统计每个位置的碱基和质量值
        if let Some(ref qual) = record.qual {
            for i in 0..seq_len {
                pos_stats[i].add(record.seq[i], qual[i], offset);
            }
        }
    }

    // 如果没有读取到任何序列，直接返回
    if n_reads == 0 {
        eprintln!("警告：未找到任何 FASTQ 序列");
        return Ok(());
    }

    // 计算所有位置的统计（用于显示不同质量值的数量）
    let mut all_stats = PosStats::new();
    for pos in &pos_stats {
        all_stats.merge(pos);
    }

    // 统计有多少种不同的质量值
    let n_diff_q = all_stats.quality.iter().filter(|&&q| q > 0).count();

    // 输出汇总信息
    println!(
        "min_len: {}; max_len: {}; avg_len: {:.2}; {} distinct quality values",
        min_len,
        max_len,
        tot_len as f64 / n_reads as f64,
        n_diff_q
    );

    // 输出表头
    print!("POS\t#bases\t%A\t%C\t%G\t%T\t%N\tavgQ\terrQ");
    if args.quality_threshold == 0 {
        // 显示所有质量值的分布
        for k in 0..=93 {
            if all_stats.quality[k] > 0 {
                print!("\t%Q{}", k);
            }
        }
    } else {
        // 显示低质量和高质量的百分比
        print!("\t%low\t%high");
    }
    println!();

    // 输出 ALL 行（所有位置的统计）
    print_pos_stats(&all_stats, 0, &perr, args.quality_threshold, &all_stats);

    // 输出每个位置的统计
    for (i, pos) in pos_stats.iter().enumerate() {
        print_pos_stats(pos, i + 1, &perr, args.quality_threshold, &all_stats);
    }

    Ok(())
}

/// 打印位置统计
fn print_pos_stats(
    stats: &PosStats,
    pos: usize,
    perr: &[f64; 94],
    qthres: usize,
    all_stats: &PosStats,
) {
    let total = stats.total_bases();
    if total == 0 {
        return;
    }

    // 位置（0表示ALL）
    if pos == 0 {
        print!("ALL");
    } else {
        print!("{}", pos);
    }

    // 碱基总数
    print!("\t{}", total);

    // 碱基百分比
    for i in 0..5 {
        let pct = 100.0 * stats.bases[i] as f64 / total as f64;
        print!("\t{:.1}", pct);
    }

    // 平均质量
    print!("\t{:.1}", stats.avg_quality());

    // 错误率质量
    print!("\t{:.1}", stats.error_quality(perr));

    // 质量值分布或低/高质量百分比
    if qthres == 0 {
        // 显示所有质量值的分布
        for k in 0..=93 {
            if all_stats.quality[k] > 0 {
                let pct = 100.0 * stats.quality[k] as f64 / total as f64;
                print!("\t{:.2}", pct);
            }
        }
    } else {
        // 计算低质量和高质量的数量
        let sum_low: u64 = stats.quality[..qthres].iter().sum();
        let sum_high = total - sum_low;

        let pct_low = 100.0 * sum_low as f64 / total as f64;
        let pct_high = 100.0 * sum_high as f64 / total as f64;

        print!("\t{:.1}\t{:.1}", pct_low, pct_high);
    }

    println!();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_pos_stats_add() {
        let mut stats = PosStats::new();

        // 添加一个 A 碱基，质量值为 40 (Phred+33 = 73 = 'I')
        stats.add(b'A', b'I', 33);

        assert_eq!(stats.bases[0], 1); // A
        assert_eq!(stats.quality[40], 1); // Q40
    }

    #[test]
    fn test_pos_stats_merge() {
        let mut stats1 = PosStats::new();
        stats1.add(b'A', b'I', 33);

        let mut stats2 = PosStats::new();
        stats2.add(b'C', b'I', 33);

        stats1.merge(&stats2);

        assert_eq!(stats1.bases[0], 1); // A
        assert_eq!(stats1.bases[1], 1); // C
        assert_eq!(stats1.quality[40], 2); // Q40 x2
    }

    #[test]
    fn test_pos_stats_avg_quality() {
        let mut stats = PosStats::new();

        // 添加 Q30 和 Q40
        stats.quality[30] = 1;
        stats.quality[40] = 1;

        let avg = stats.avg_quality();
        assert!((avg - 35.0).abs() < 0.01);
    }

    #[test]
    fn test_pos_stats_total_bases() {
        let mut stats = PosStats::new();
        stats.bases[0] = 10; // A
        stats.bases[1] = 20; // C
        stats.bases[2] = 30; // G
        stats.bases[3] = 40; // T

        assert_eq!(stats.total_bases(), 100);
    }
}