pandora_lib_promethion/src/io/bed.rs

//! BED file I/O, overlap queries, and genomic region annotation.
//!
//! All coordinates follow **BED convention: 0-based, half-open `[start, end)`**.
//! This matches [`GenomeRange`]'s internal `Range<u32>` representation, so BED
//! values are stored as-is without conversion.
//!
//! # Main types and functions
//!
//! | Item | Purpose |
//! |------|---------|
//! | [`BedRow`] | One parsed BED line (up to BED6) |
//! | [`read_bed`] | Load a BED file into memory, skipping headers and blank lines |
//! | [`fetch_bed`] | Random-access region fetch from a Tabix-indexed `.bed.gz` |
//! | [`annotate_with_bed`] | Annotate a [`Variants`] collection from a BED region set |
//! | [`bedrow_overlaps_par`] | Parallel overlap query: which BED rows hit a set of query ranges |
//! | [`GenesBedIndex`] | Per-contig indexed structure for fast gene-name lookup |
//! | [`convert_bgz_with_tabix`] | Compress a BED file to BGZF + Tabix index via [`BgzTabixWriter`] |
//! | [`parse_centromere_intervals`] | Extract centromeric intervals from a UCSC cytoband file |

use std::{
    io::{BufRead, BufReader},
    path::Path,
    str::FromStr,
    sync::Arc,
};

use anyhow::Context;
use log::warn;
use rayon::prelude::*;

use crate::{
    annotation::Annotation,
    io::{readers::TabixReader, writers::{BgzTabixWriter, IndexFormat}},
    positions::{
        GenomePosition, GenomeRange, GetGenomeRange, extract_contig_indices, find_contig_indices, overlaps_par
    },
    variant::variant_collection::Variants,
};

use super::readers::get_reader;

/// One parsed row from a BED file (up to BED6).
///
/// Coordinates in `range` are **0-based, half-open** as stored in the BED file.
/// Optional fields (`name`, `score`, `strand`) are `None` when the column is absent
/// or unparseable.
#[derive(Debug, Clone)]
pub struct BedRow {
    /// Genomic interval — 0-based half-open `[start, end)`
    pub range: GenomeRange,
    /// BED column 4: feature name
    pub name: Option<String>,
    /// BED column 5: score (0–1000 per spec, stored as `u16`)
    pub score: Option<u16>,
    /// BED column 6: strand — `true` = `+`, `false` = `-`
    pub strand: Option<bool>,
}

/// Parse a tab-separated BED line into a [`BedRow`].
///
/// Expects at least 3 tab-separated columns: `chrom`, `start`, `end`.
/// Columns 4–6 (`name`, `score`, `strand`) are optional; missing or unparseable
/// values are stored as `None`. Coordinates are stored as-is (0-based half-open).
///
/// # Errors
///
/// Returns an error if fewer than 3 columns are present or if `start`/`end`
/// cannot be parsed as `u32`.
impl FromStr for BedRow {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> anyhow::Result<Self> {
        let v: Vec<&str> = s.split('\t').collect();
        let range: GenomeRange = (
            *v.first()
                .ok_or(anyhow::anyhow!("Can't get contig from: {s}"))?,
            *v.get(1)
                .ok_or(anyhow::anyhow!("Can't get position from: {s}"))?,
            *v.get(2)
                .ok_or(anyhow::anyhow!("Can't get position from: {s}"))?,
        )
            .try_into()
            .context(format!("Can't parse range from: {s}"))?;

        Ok(Self {
            range,
            name: v.get(3).map(|v| v.to_string()),
            score: v.get(4).and_then(|v| v.parse().ok()),
            strand: v.get(5).and_then(|&v| match v {
                "+" => Some(true),
                "-" => Some(false),
                _ => None,
            }),
        })
    }
}

impl GetGenomeRange for BedRow {
    fn range(&self) -> &GenomeRange {
        &self.range
    }
}

fn is_bed_metadata_line(line: &str) -> bool {
    let line = line.trim_start();
    line.is_empty()
        || line.starts_with('#')
        || line.starts_with("track")
        || line.starts_with("browser")
}

fn bed_row_overlaps_region(row: &BedRow, region: &GenomeRange) -> bool {
    if row.range.contig != region.contig {
        return false;
    }

    let row_start = row.range.range.start;
    let row_end = row.range.range.end;

    if row_start == row_end {
        region.range.contains(&row_start)
    } else {
        row_end > region.range.start && row_start < region.range.end
    }
}

/// Load a BED file into memory as a vector of [`BedRow`]s.
///
/// Skips blank lines and lines starting with `#`, `track`, or `browser`
/// (standard UCSC header conventions). All other lines are parsed via
/// [`BedRow::from_str`].
///
/// # Arguments
///
/// * `path` - Path to a BED file (plain text or gzip-compressed)
///
/// # Returns
///
/// A vector of parsed rows in file order.
///
/// # Errors
///
/// Returns an error if the file cannot be opened or if any data line fails to parse.
/// I/O errors on individual lines are logged as warnings and skipped.
pub fn read_bed(path: &str) -> anyhow::Result<Vec<BedRow>> {
    let reader = BufReader::new(get_reader(path)?);

    let mut res = Vec::new();
    for (i, line) in reader.lines().enumerate() {
        let line_no = i + 1;
        match line {
            Ok(line) => {
                if is_bed_metadata_line(&line) {
                    continue;
                }
                res.push(line.parse().with_context(|| {
                    format!("failed parsing BED record at {path}:{line_no}: {line}")
                })?);
            }
            Err(e) => warn!("Can't read {path}:{line_no}: {e}"),
        }
    }

    Ok(res)
}

/// Annotate variants that overlap a BED region set and return coverage statistics.
///
/// Loads the BED file, finds all variants whose position overlaps any BED region
/// (via [`overlaps_par`]), pushes `annotation` onto each matching variant, and
/// returns the total covered base-pair count alongside the overlap count.
///
/// # Arguments
///
/// * `variants` - Variant collection to annotate in place
/// * `bed_path` - Path to the BED file defining the region class
/// * `annotation` - Annotation tag to attach to each overlapping variant
///
/// # Returns
///
/// `(total_bp, overlap_count)` where `total_bp` is the sum of all BED interval
/// lengths and `overlap_count` is the number of variants that received the annotation.
///
/// # Errors
///
/// Returns an error if the BED file cannot be read.
pub fn annotate_with_bed(
    variants: &mut Variants,
    bed_path: &str,
    annotation: Annotation,
) -> anyhow::Result<(usize, usize)> {
    let bed_rows = read_bed(bed_path)?;
    let ranges: Vec<&GenomeRange> = bed_rows.iter().map(|b| &b.range).collect();

    let total_bp: usize = ranges.iter().map(|r| r.length() as usize).sum();
    let positions: Vec<&GenomePosition> = variants.data.iter().map(|v| &v.position).collect();
    let overlaps = overlaps_par(&positions, &ranges);

    for &idx in &overlaps {
        variants.data[idx].annotations.push(annotation.clone());
    }

    Ok((total_bp, overlaps.len()))
}

/// Return every BED row that overlaps at least one query range, parallelised per contig.
///
/// Both inputs must be sorted by `(contig, start)`. Each matching row appears exactly
/// once in the output regardless of how many queries it overlaps.
///
/// Uses an anchored-scan algorithm (O(n + m)) that is correct even when rows or
/// queries overlap each other (e.g. nested genes). The anchor `j0` advances only
/// when a query definitively ends before the current row starts, so no hit can be
/// missed.
///
/// # Arguments
///
/// * `rows` - BED rows sorted by `(contig, start)`
/// * `queries` - Query ranges sorted by `(contig, start)`
///
/// # Returns
///
/// Matching BED rows in per-contig order. Cross-contig order is unspecified (parallel
/// execution).
pub fn bedrow_overlaps_par(rows: &[BedRow], queries: &[&GenomeRange]) -> Vec<BedRow>
where
    BedRow: Clone + Send + Sync,
{
    let row_ranges: Vec<&GenomeRange> = rows.iter().map(|r| &r.range).collect();
    let (row_contigs, query_contigs) = rayon::join(
        || extract_contig_indices(&row_ranges),
        || extract_contig_indices(queries),
    );

    row_contigs
        .into_par_iter()
        .filter_map(|(contig, r_start, r_end)| {
            let (q_start, q_end) = find_contig_indices(&query_contigs, contig)?;

            let r_slice = &rows[r_start..r_end];
            let q_slice = &queries[q_start..q_end];

            let mut hits = Vec::new();
            let mut j0 = 0usize;

            for row in r_slice {
                let r = row.range();

                while j0 < q_slice.len() && q_slice[j0].range.end <= r.range.start {
                    j0 += 1;
                }

                if j0 < q_slice.len() && q_slice[j0].range.start < r.range.end {
                    hits.push(row.clone());
                }
            }

            Some(hits)
        })
        .flatten()
        .collect()
}

/// Per-contig indexed BED structure for fast gene-name lookup.
///
/// Rows are sorted by `(contig, start, end)` at construction time and stored in
/// 256 contig slots (one per encoded `u8` contig). Lookups are O(rows before `end`)
/// with early termination — correct even for very long host genes.
#[derive(Clone)]
pub struct GenesBedIndex {
    by_contig: Vec<Arc<[BedRow]>>,
}

impl GenesBedIndex {
    /// Build a [`GenesBedIndex`] from an unsorted vector of [`BedRow`]s.
    ///
    /// Rows are sorted by `(contig, start, end)` internally.
    pub fn new(mut rows: Vec<BedRow>) -> Self {
        rows.sort_unstable_by(|a, b| {
            a.range
                .contig
                .cmp(&b.range.contig)
                .then_with(|| a.range.range.start.cmp(&b.range.range.start))
                .then_with(|| a.range.range.end.cmp(&b.range.range.end))
        });

        let mut tmp: Vec<Vec<BedRow>> = vec![Vec::new(); 256];
        for r in rows {
            tmp[r.range.contig as usize].push(r);
        }

        Self {
            by_contig: tmp.into_iter().map(Arc::<[BedRow]>::from).collect(),
        }
    }

    /// Return the names of all genes overlapping `[start, end)` on `contig`.
    ///
    /// Coordinates are 0-based half-open. If `start > end` the arguments are
    /// swapped defensively. Genes without a name field are silently skipped.
    ///
    /// # Arguments
    ///
    /// * `contig` - Encoded contig index (see `contig_to_num`)
    /// * `start` - Query start, 0-based inclusive
    /// * `end` - Query end, 0-based exclusive
    ///
    /// # Returns
    ///
    /// Names of overlapping genes in start-sorted order.
    pub fn query_genes(&self, contig: u8, start: u32, end: u32) -> Vec<String> {
        let (s, e) = if start <= end {
            (start, end)
        } else {
            (end, start)
        };
        let rows = &self.by_contig[contig as usize];

        rows.iter()
            .take_while(|r| r.range.range.start < e)
            .filter(|r| r.range.range.end > s)
            .filter_map(|r| r.name.clone())
            .collect()
    }
}

/// Compress a BED file to BGZF and build a Tabix index (`.tbi`) in a single pass.
///
/// Delegates all BGZF and Tabix mechanics to [`BgzTabixWriter`]. BED-specific work
/// here is limited to column parsing and coordinate conversion.
///
/// Output files are `<input>.gz` and `<input>.gz.tbi`. Both are written atomically.
/// The input must be tab-delimited, contig-grouped, and coordinate-sorted.
///
/// # Coordinate conversion
///
/// BED coordinates are **0-based half-open `[start, end)`**. [`BgzTabixWriter::write_bed_record`]
/// converts them to 1-based Tabix positions automatically:
/// - `tabix_start = bed_start + 1`
/// - `tabix_end   = bed_end`  (0-based exclusive = 1-based inclusive numerically)
///
/// # Arguments
///
/// * `input` - Path to the input BED file (plain text or BGZF)
/// * `force` - Overwrite existing output files if present
///
/// # Errors
///
/// Returns an error if the input cannot be read, any data line is malformed, or
/// the output cannot be written.
pub fn convert_bgz_with_tabix(input: impl AsRef<Path>, force: bool) -> anyhow::Result<()> {
    let input = input.as_ref();
    let out_bgz = format!("{}.gz", input.display());

    let reader = get_reader(&input.to_string_lossy())?;
    let mut reader = BufReader::new(reader);
    let mut writer = BgzTabixWriter::new(&out_bgz, IndexFormat::Tbi, force)?;

    let mut line = String::new();
    loop {
        line.clear();
        if reader.read_line(&mut line)? == 0 {
            break;
        }

        if is_bed_metadata_line(&line) {
            writer.write_header(line.as_bytes())?;
            continue;
        }

        let mut it = line.split('\t');
        let rname = it.next().context("BED: missing contig")?;
        let start0: u32 = it
            .next()
            .context("BED: missing start")?
            .parse()
            .context("BED: invalid start")?;
        let end0: u32 = it
            .next()
            .context("BED: missing end")?
            .trim()
            .parse()
            .context("BED: invalid end")?;

        writer.write_bed_record(line.as_bytes(), rname, start0, end0)?;
    }

    writer.finish()
}

/// Fetch BED rows overlapping `region` from a Tabix-indexed BGZF file.
///
/// Delegates the BGZF seek and line extraction to [`fetch_tabix_lines`], then
/// parses each returned line as a [`BedRow`] and applies a secondary overlap
/// filter to discard any false positives from the Tabix bin query.
///
/// # Coordinate system
///
/// `region` is **0-based half-open** `[start, end)`. The secondary filter uses
/// the same convention: a row `[rs, re)` is kept iff `rs < region.end && re > region.start`.
///
/// # Arguments
///
/// * `bgz_path` - Path to the `.bed.gz` file; index expected at `<bgz_path>.tbi`
/// * `region` - Genomic region to query (0-based half-open)
///
/// # Returns
///
/// Parsed [`BedRow`]s that genuinely overlap `region`, in file order.
///
/// # Errors
///
/// Returns an error if the index or file cannot be read, or if a fetched line
/// fails to parse as a BED row.
pub fn fetch_bed_with_reader<R: std::io::Read + std::io::Seek>(
    reader: &mut TabixReader<R>,
    region: &GenomeRange,
) -> anyhow::Result<Vec<BedRow>> {
    let mut rows = Vec::new();
    reader.fetch_with(region, |line| {
        let row: BedRow = line
            .parse()
            .with_context(|| format!("failed to parse BED line: {line}"))?;
        // Secondary overlap filter: discard tabix-bin false positives.
        if bed_row_overlaps_region(&row, region) {
            rows.push(row);
        }
        Ok(())
    })?;
    Ok(rows)
}

/// Parse centromeric intervals from a UCSC cytoband BED file.
///
/// Identifies rows whose stain field (column 5) equals `acen` and returns their
/// coordinates as `(chrom, start, end)`. Coordinates are 0-based half-open as
/// stored in the file.
///
/// Expected column layout: `chrom\tstart\tend\tband_name\tstain`
///
/// # Arguments
///
/// * `path` - Path to the cytoband BED file
///
/// # Returns
///
/// A vector of `(chrom, start, end)` tuples for all `acen` bands, in file order.
///
/// # Errors
///
/// Returns an error if the file cannot be opened, a line has fewer than 5 columns,
/// or a coordinate field cannot be parsed as `u64`.
pub fn parse_centromere_intervals(path: &str) -> anyhow::Result<Vec<(String, u64, u64)>> {
    let reader = BufReader::new(
        get_reader(path).with_context(|| format!("Cannot open cytobands BED: {path}"))?,
    );

    let mut intervals = Vec::new();
    for (i, line) in reader.lines().enumerate() {
        let line_no = i + 1;
        let line = line.with_context(|| format!("failed reading {path}:{line_no}"))?;
        if is_bed_metadata_line(&line) {
            continue;
        }
        let f: Vec<&str> = line.split('\t').collect();
        anyhow::ensure!(
            f.len() >= 5,
            "Expected 5 columns in cytoband BED at {path}:{line_no}, got {}: {line}",
            f.len()
        );
        if f[4].trim() == "acen" {
            intervals.push((
                f[0].to_string(),
                f[1].parse::<u64>()
                    .with_context(|| format!("bad cytoband start at {path}:{line_no}: {}", f[1]))?,
                f[2].parse::<u64>()
                    .with_context(|| format!("bad cytoband end at {path}:{line_no}: {}", f[2]))?,
            ));
        }
    }
    Ok(intervals)
}

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

    #[test]
    fn bed_to_bgz() -> anyhow::Result<()> {
        test_init();

        let a = "/mnt/beegfs02/scratch/t_steimle/data/wgs/CHAHA/norm/CHAHA_norm_modkit_pileup.bed";
        convert_bgz_with_tabix(a, true)?;
        Ok(())
    }

    #[test]
    fn zero_width_bed_row_overlaps_containing_base_only() {
        let row: BedRow = "chr1\t10\t10\tins".parse().unwrap();

        assert!(!bed_row_overlaps_region(
            &row,
            &GenomeRange::new("chr1", 9, 10)
        ));
        assert!(bed_row_overlaps_region(
            &row,
            &GenomeRange::new("chr1", 10, 11)
        ));
        assert!(!bed_row_overlaps_region(
            &row,
            &GenomeRange::new("chr1", 11, 12)
        ));
        assert!(!bed_row_overlaps_region(
            &row,
            &GenomeRange::new("chr2", 10, 11)
        ));
    }
}