seqtkrs 0.1.1

// comp 命令：计算序列的碱基组成统计
//
// 该命令统计每条序列的详细碱基组成，包括：
// - A/C/G/T 碱基计数
// - 模糊碱基（2-way, 3-way, 4-way）计数
// - CpG 二核苷酸数量
// - 转换（transitions）和颠换（transversions）
// - CpG 上下文中的转换

use anyhow::Result;
use clap::Args;
use rustc_hash::FxHashMap;

use crate::core::tables::LOOKUP_TABLES;
use crate::core::{SeqReader, SeqRecord};
use crate::utils::region::{parse_regions, RegionList};

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

    /// 只统计大写碱基
    #[arg(short = 'u')]
    pub upper_only: bool,

    /// BED 文件路径，指定要统计的区域
    #[arg(short = 'r', value_name = "FILE")]
    pub region_file: Option<String>,
}

/// 碱基组成统计结果
#[derive(Debug, Default)]
struct CompStats {
    /// A/C/G/T 计数
    bases: [u64; 4],
    /// 2-way 模糊碱基计数
    ambig_2: u64,
    /// 3-way 模糊碱基计数
    ambig_3: u64,
    /// 4-way 模糊碱基计数
    ambig_4: u64,
    /// CpG 二核苷酸数量
    cpg: u64,
    /// 颠换（transversions）数量
    tv: u64,
    /// 转换（transitions）数量
    ts: u64,
    /// CpG 上下文中的转换数量
    cpg_ts: u64,
}

pub fn run(args: &CompArgs) -> Result<()> {
    // 解析区域文件（如果提供）
    let regions = if let Some(ref region_file) = args.region_file {
        Some(parse_regions(region_file)?)
    } else {
        None
    };

    // 创建输入读取器
    let input = args.input.as_deref().unwrap_or("-");
    let mut reader = if input == "-" {
        SeqReader::from_stdin()
    } else {
        SeqReader::from_path(input)?
    };

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

    // 处理每条序列
    while reader.read_next(&mut record)? {
        if let Some(ref region_map) = regions {
            // 如果提供了区域文件，按区域统计
            process_with_regions(&record, region_map, args.upper_only)?;
        } else {
            // 否则统计整条序列
            let stats = compute_composition(&record.seq, 0, record.seq.len(), args.upper_only);
            print_stats(&record.name, None, record.seq.len(), &stats);
        }
    }

    Ok(())
}

/// 按区域处理序列
fn process_with_regions(
    record: &SeqRecord,
    region_map: &FxHashMap<Vec<u8>, RegionList>,
    upper_only: bool,
) -> Result<()> {
    // 查找该序列的区域
    if let Some(region_list) = region_map.get(&record.name) {
        for &(start, end) in &region_list.regions {
            let start = start as usize;
            let end = end as usize;

            // 确保区域在序列范围内
            if start >= record.seq.len() {
                continue;
            }
            let end = end.min(record.seq.len());

            // 计算该区域的组成
            let stats = compute_composition(&record.seq, start, end, upper_only);

            // 输出结果（包含区域坐标）
            print_stats(&record.name, Some((start, end)), end - start, &stats);
        }
    }

    Ok(())
}

/// nt16到4碱基编码的转换表
const NT16_TO_4: [u8; 16] = [4, 0, 1, 4, 2, 4, 4, 4, 3, 4, 4, 4, 4, 4, 4, 4];

/// 计算序列的碱基组成
fn compute_composition(seq: &[u8], start: usize, end: usize, upper_only: bool) -> CompStats {
    let mut stats = CompStats::default();

    if start >= end || end > seq.len() {
        return stats;
    }

    // 用于跟踪前一个碱基（用于 CpG 检测）
    let prev_base = if start > 0 {
        seq[start - 1]
    } else {
        b'a' // 小写字母，表示不存在
    };
    let mut prev_code = LOOKUP_TABLES.nt16[prev_base as usize];

    // 遍历序列
    for i in start..end {
        let base = seq[i];
        let code = LOOKUP_TABLES.nt16[base as usize];
        let bitcnt = LOOKUP_TABLES.bitcnt[code as usize];

        // 判断是否为 CpG
        let mut is_cpg = false;
        if code == 2 || code == 10 {
            // C 或 Y (Y = C/T)
            // 检查下一个碱基是否为 G 或 R (R = A/G)
            if i + 1 < end {
                let next_code = LOOKUP_TABLES.nt16[seq[i + 1] as usize];
                if next_code == 4 || next_code == 5 {
                    is_cpg = true;
                }
            }
        } else if code == 4 || code == 5 {
            // G 或 R (R = A/G)
            // 检查前一个碱基是否为 C 或 Y (Y = C/T)
            if prev_code == 2 || prev_code == 10 {
                is_cpg = true;
            }
        }

        // 如果只统计大写碱基，检查是否为大写
        if !upper_only || base.is_ascii_uppercase() {
            // 统计碱基类型
            if bitcnt > 1 {
                // 模糊碱基（2-way, 3-way, 4-way）
                stats.ambig_2 += (bitcnt == 2) as u64;
                stats.ambig_3 += (bitcnt == 3) as u64;
                stats.ambig_4 += (bitcnt == 4) as u64;
            } else if bitcnt == 1 {
                // 精确碱基 (A/C/G/T)
                let nt4 = NT16_TO_4[code as usize];
                if nt4 < 4 {
                    stats.bases[nt4 as usize] += 1;
                }
            }

            // 统计转换和颠换
            if code == 10 || code == 5 {
                // R (A/G) 或 Y (C/T) - 转换
                stats.ts += 1;
            } else if bitcnt == 2 {
                // 其他 2-way 模糊碱基 - 颠换
                stats.tv += 1;
            }

            // 统计 CpG
            if is_cpg {
                stats.cpg += 1;
                if code == 10 || code == 5 {
                    stats.cpg_ts += 1;
                }
            }
        }

        prev_code = code;
    }

    stats
}

/// 打印统计结果
fn print_stats(name: &[u8], region: Option<(usize, usize)>, length: usize, stats: &CompStats) {
    // 输出序列名称
    print!("{}", String::from_utf8_lossy(name));

    // 如果有区域坐标，输出区域
    if let Some((start, end)) = region {
        print!("\t{}\t{}", start, end);
    } else {
        print!("\t{}", length);
    }

    // 输出统计结果
    println!(
        "\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}",
        stats.bases[0], // A
        stats.bases[1], // C
        stats.bases[2], // G
        stats.bases[3], // T
        stats.ambig_2,  // 2-way
        stats.ambig_3,  // 3-way
        stats.ambig_4,  // 4-way
        stats.cpg,      // CpG
        stats.tv,       // transversions
        stats.ts,       // transitions
        stats.cpg_ts    // CpG transitions
    );
}

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

    #[test]
    fn test_compute_composition_simple() {
        let seq = b"ACGT";
        let stats = compute_composition(seq, 0, seq.len(), false);

        assert_eq!(stats.bases[0], 1); // A
        assert_eq!(stats.bases[1], 1); // C
        assert_eq!(stats.bases[2], 1); // G
        assert_eq!(stats.bases[3], 1); // T
    }

    #[test]
    fn test_compute_composition_cpg() {
        let seq = b"ACGT";
        let stats = compute_composition(seq, 0, seq.len(), false);

        // CG 应该被识别为 CpG
        assert_eq!(stats.cpg, 2); // C 和 G 各算一次
    }

    #[test]
    fn test_compute_composition_upper_only() {
        let seq = b"ACGTacgt";
        let stats = compute_composition(seq, 0, seq.len(), true);

        // 只统计大写碱基
        assert_eq!(stats.bases[0], 1); // A (大写)
        assert_eq!(stats.bases[1], 1); // C (大写)
        assert_eq!(stats.bases[2], 1); // G (大写)
        assert_eq!(stats.bases[3], 1); // T (大写)
    }

    #[test]
    fn test_compute_composition_ambiguous() {
        let seq = b"ACGTRYN"; // R=A/G, Y=C/T, N=any
        let stats = compute_composition(seq, 0, seq.len(), false);

        // R 和 Y 是 2-way 模糊碱基，同时也是转换
        assert_eq!(stats.ambig_2, 2);
        assert_eq!(stats.ts, 2);
    }
}