seqtkrs 0.1.1

// BED区域解析模块
// 支持解析BED文件和名称列表

use anyhow::{Context, Result};
use rustc_hash::FxHashMap;
use std::fs::File;
use std::io::{BufRead, BufReader};

/// 单个区域
#[derive(Debug, Clone, Copy)]
pub struct Region {
    /// 起始位置（0-based，包含）
    pub start: usize,
    /// 结束位置（0-based，不包含）
    pub end: usize,
    /// 链方向：Some(b'+')=正链, Some(b'-')=负链, None=未指定
    pub strand: Option<u8>,
}

impl Region {
    /// 创建新区域
    pub fn new(start: usize, end: usize, strand: Option<u8>) -> Self {
        Self { start, end, strand }
    }
}

/// 区域列表
#[derive(Debug, Clone)]
pub struct RegionList {
    /// (start, end) 对，start是0-based，end是exclusive
    pub regions: Vec<(u64, u64)>,
    /// 链方向: 1=正链, -1=负链, 0=未指定
    pub strands: Vec<i8>,
}

impl RegionList {
    /// 创建新的区域列表
    pub fn new() -> Self {
        Self {
            regions: Vec::new(),
            strands: Vec::new(),
        }
    }

    /// 添加一个区域
    pub fn add_region(&mut self, start: u64, end: u64, strand: i8) {
        self.regions.push((start, end));
        self.strands.push(strand);
    }

    /// 获取区域迭代器
    pub fn iter(&self) -> impl Iterator<Item = Region> + '_ {
        self.regions
            .iter()
            .zip(&self.strands)
            .map(|(&(start, end), &strand)| Region {
                start: start as usize,
                end: end as usize,
                strand: match strand {
                    1 => Some(b'+'),
                    -1 => Some(b'-'),
                    _ => None,
                },
            })
    }

    /// 检查位置是否在任何区域内
    pub fn contains(&self, pos: u64) -> bool {
        self.regions
            .iter()
            .any(|&(start, end)| pos >= start && pos < end)
    }

    /// 获取区域数量
    pub fn len(&self) -> usize {
        self.regions.len()
    }

    /// 判断是否为空
    pub fn is_empty(&self) -> bool {
        self.regions.is_empty()
    }
}

impl Default for RegionList {
    fn default() -> Self {
        Self::new()
    }
}

/// BED文件解析结果
/// Key: 序列名称，Value: 该序列的区域列表
pub type RegionHash = FxHashMap<Vec<u8>, RegionList>;

/// 解析BED文件或名称列表
///
/// 支持两种格式：
/// 1. 名称列表（每行一个序列名）
/// 2. BED格式（tab分隔：chr start end [name] [score] [strand]）
pub fn parse_regions(path: &str) -> Result<RegionHash> {
    let file = File::open(path).with_context(|| format!("无法打开区域文件: {}", path))?;
    let reader = BufReader::new(file);

    let mut hash: RegionHash = FxHashMap::default();

    for (line_num, line) in reader.lines().enumerate() {
        let line = line.with_context(|| format!("读取文件{}第{}行失败", path, line_num + 1))?;

        // 跳过空行和注释行
        if line.is_empty() || line.starts_with('#') {
            continue;
        }

        let fields: Vec<&str> = line.split('\t').collect();
        if fields.is_empty() {
            continue;
        }

        let name = fields[0].as_bytes().to_vec();
        let entry = hash.entry(name).or_default();

        // 只有名称（名称列表格式）
        if fields.len() == 1 {
            entry.add_region(0, u64::MAX, 0);
            continue;
        }

        // BED格式: chr start end [name] [score] [strand]
        let start: u64 = fields.get(1).and_then(|s| s.parse().ok()).unwrap_or(0);

        let end: u64 = fields
            .get(2)
            .and_then(|s| s.parse().ok())
            .unwrap_or(u64::MAX);

        let strand: i8 = fields
            .get(5)
            .map(|s| match *s {
                "+" => 1,
                "-" => -1,
                _ => 0,
            })
            .unwrap_or(0);

        entry.add_region(start, end, strand);
    }

    Ok(hash)
}

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

    #[test]
    fn test_region_list_basic() {
        let mut list = RegionList::new();
        assert!(list.is_empty());
        assert_eq!(list.len(), 0);

        list.add_region(100, 200, 1);
        assert!(!list.is_empty());
        assert_eq!(list.len(), 1);

        assert!(list.contains(100));
        assert!(list.contains(150));
        assert!(list.contains(199));
        assert!(!list.contains(200));
        assert!(!list.contains(99));
    }

    #[test]
    fn test_parse_name_list() -> Result<()> {
        let mut temp_file = NamedTempFile::new()?;
        writeln!(temp_file, "seq1")?;
        writeln!(temp_file, "seq2")?;
        writeln!(temp_file, "# 注释行")?;
        writeln!(temp_file, "")?; // 空行
        writeln!(temp_file, "seq3")?;

        let regions = parse_regions(temp_file.path().to_str().unwrap())?;

        assert_eq!(regions.len(), 3);
        assert!(regions.contains_key(b"seq1".as_slice()));
        assert!(regions.contains_key(b"seq2".as_slice()));
        assert!(regions.contains_key(b"seq3".as_slice()));

        // 名称列表格式应该添加整个序列
        let seq1_regions = &regions[b"seq1".as_slice()];
        assert_eq!(seq1_regions.len(), 1);
        assert_eq!(seq1_regions.regions[0], (0, u64::MAX));

        Ok(())
    }

    #[test]
    fn test_parse_bed_format() -> Result<()> {
        let mut temp_file = NamedTempFile::new()?;
        writeln!(temp_file, "chr1\t100\t200\tregion1\t0\t+")?;
        writeln!(temp_file, "chr1\t300\t400\tregion2\t0\t-")?;
        writeln!(temp_file, "chr2\t500\t600")?;

        let regions = parse_regions(temp_file.path().to_str().unwrap())?;

        assert_eq!(regions.len(), 2);

        let chr1_regions = &regions[b"chr1".as_slice()];
        assert_eq!(chr1_regions.len(), 2);
        assert_eq!(chr1_regions.regions[0], (100, 200));
        assert_eq!(chr1_regions.strands[0], 1); // +
        assert_eq!(chr1_regions.regions[1], (300, 400));
        assert_eq!(chr1_regions.strands[1], -1); // -

        let chr2_regions = &regions[b"chr2".as_slice()];
        assert_eq!(chr2_regions.len(), 1);
        assert_eq!(chr2_regions.regions[0], (500, 600));
        assert_eq!(chr2_regions.strands[0], 0); // 未指定

        Ok(())
    }

    #[test]
    fn test_multiple_regions_same_chr() -> Result<()> {
        let mut temp_file = NamedTempFile::new()?;
        writeln!(temp_file, "chr1\t100\t200")?;
        writeln!(temp_file, "chr1\t300\t400")?;
        writeln!(temp_file, "chr1\t500\t600")?;

        let regions = parse_regions(temp_file.path().to_str().unwrap())?;

        let chr1_regions = &regions[b"chr1".as_slice()];
        assert_eq!(chr1_regions.len(), 3);

        assert!(chr1_regions.contains(150));
        assert!(chr1_regions.contains(350));
        assert!(chr1_regions.contains(550));
        assert!(!chr1_regions.contains(250));

        Ok(())
    }
}