bam_tide 1.2.0

// src/bin/bam-quant.rs
//
// Quantify a 10x-style BAM against a splice index (GTF-derived) into scdata.
//
// Pipeline:
// 1) Stream BAM -> collect jobs (cell_u64, umi_u64, spliced_read, span)
// 2) Rayon over chunks -> local Scdata + MappingInfo
// 3) Merge locals into one Scdata + MappingInfo
//
// Assumptions (adjust if your crate differs):
// - Tags: CB (cell barcode), UB (UMI). CB may end with "-1" etc.
// - record_to_blocks(&Record) exists and returns Vec<RefBlock> 0-based half-open.
// - SpliceIndex has:
//     - transcripts: Vec<Transcript>
//     - chrom_names or equivalent to build chr_name -> chr_id mapping
//     - candidates_for_span_union(chr_id, start0, end0) -> Vec<TranscriptId>
// - Transcript has: gene_id: usize and match_spliced_read(&SplicedRead, MatchOptions) -> MatchHit
// - Strand enum has Plus/Minus (or similar)
// - MappingInfo has fields ok_reads, pcr_duplicates, local_dup, and ideally a merge method.
//   If no merge method exists, see `merge_mapping_info_fallback()` below.

use anyhow::{Context, Result};
use clap::{Parser, ValueEnum};
use rayon::prelude::*;
use rust_htslib::bam::{Read, Reader, Record};
use std::path::PathBuf;

use scdata::cell_data::GeneUmiHash;
use scdata::{IndexedGenes, MatrixValueType, Scdata};

use mapping_info::MappingInfo;

use int_to_str::int_to_str::IntToStr;

// for the log file...
use std::fs::File;
use std::io::Write;

const CHUNK: usize = 2_000_000;

// ---- Your splice index / transcript-matching crate ----
// Adjust these paths to your actual crate/module names.
use gtf_splice_index::types::RefBlock;
use gtf_splice_index::{MatchClass, MatchOptions, SpliceIndex, SplicedRead, Strand};

// ---- Your ref-block conversion ----
// Adjust these paths to where RefBlock + record_to_blocks live in bam_tide.
use bam_tide::core::ref_block::record_to_blocks;
use bam_tide::compute_io_threads;

#[derive(Debug, Clone, Copy, ValueEnum)]
pub enum QuantMode {
    Gene,
    Transcript,
}

#[derive(Parser, Debug, Clone)]
#[command(
    name = "bam-quant",
    about = "Quantify 10x BAM against splice index into scdata"
)]


pub struct QuantCli {
    /// Input BAM (10x / CellRanger BAM)
    #[arg(long, short)]
    pub bam: std::path::PathBuf,

    /// Splice index path (built from GTF beforehand)
    #[arg(long, short)]
    pub index: std::path::PathBuf,

    /// Outpath for the 10x mmx formated outfiles
    #[arg(long, short)]
    pub outpath: std::path::PathBuf,



    /// Split Intronic from rest
    ///
    /// This is currently not recommended as exon intron detection seams to be too strict from normal sequencing data
    #[arg(long, short, default_value_t = false)]
    pub split_intronic: bool,

    /// Minimum MAPQ
    #[arg(long, default_value_t = 0)]
    pub min_mapq: u8,

    /// Use only read1 (recommended for 10x; reduces duplicate mate counting noise)
    #[arg(long, default_value_t = false)]
    pub read1_only: bool,

    /// threads for bam read process (default 4)
    #[arg(long, default_value_t = 4)]
    pub threads: usize,

    /// Collect Gene or Transcript names
    #[arg(long, value_enum, default_value_t = QuantMode::Gene)]
    pub quant_mode: QuantMode,

    /// Max reads to process (debug/dev)
    #[arg(long)]
    pub max_reads: Option<usize>,

    /// Min read counts per reported cell (debug/dev)
    #[arg(long, default_value_t = 400)]
    pub min_cell_counts: usize,

    // ------------------------------
    // MatchOptions exposed to user
    // ------------------------------
    /// If true, require read blocks to be on a compatible strand.
    #[arg(long, default_value_t = false)]
    pub require_strand: bool,

    /// If true, require the read to have the exact same splice junction chain as the transcript.
    #[arg(long, default_value_t = false)]
    pub require_exact_junction_chain: bool,

    /// Maximum allowed 5′ overhang (bp). If exceeded -> OverhangTooLarge.
    #[arg(long, default_value_t = 100)]
    pub max_5p_overhang_bp: u32,

    /// Maximum allowed 3′ overhang (bp). If exceeded -> OverhangTooLarge.
    #[arg(long, default_value_t = 100)]
    pub max_3p_overhang_bp: u32,

    /// Allowed sequencing error gap. If exceeded -> JunctionMismatch.
    #[arg(long, default_value_t = 5)]
    pub allowed_intronic_gap_size: u32,
}

#[derive(Clone)]
struct Job {
    cell: u64,
    umi: u64,
    //chr_id: usize,
    //start0: u32,
    //end0: u32,
    spliced: SplicedRead,
}

fn aux_tag_str<'a>(rec: &'a Record, tag: [u8; 2]) -> Option<&'a str> {
    use rust_htslib::bam::record::Aux;
    match rec.aux(&tag).ok()? {
        Aux::String(s) => Some(s),
        _ => None,
    }
}

/// Many 10x barcodes look like "AAAC...-1". Strip the suffix before encoding.
fn normalize_10x_barcode(cb: &str) -> &str {
    match cb.split_once('-') {
        Some((core, _)) => core,
        None => cb,
    }
}

/// Convert a DNA string (ACGT only) into a packed u64 via 2-bit encoding.
fn dna_to_u64(seq: &str) -> Option<u64> {
    if !seq.bytes().all(|b| matches!(b, b'A' | b'C' | b'G' | b'T')) {
        return None;
    }
    let tool = IntToStr::new(seq.as_bytes());
    Some(tool.into_u64())
}

/// Build SplicedRead + union span [start0, end0) from a BAM record.
fn record_to_spliced_read(rec: &Record, chr_id: usize) -> Option<(SplicedRead, u32, u32)> {
    let blocks: Vec<RefBlock> = record_to_blocks(rec);
    if blocks.is_empty() {
        return None;
    }

    let strand = if rec.is_reverse() {
        Strand::Minus
    } else {
        Strand::Plus
    };

    let mut spliced = SplicedRead::new(chr_id, strand, blocks);
    let (start0, end0) = spliced.finalize();
    Some((spliced, start0, end0))
}


/*
/// Build a chr_name -> chr_id lookup.
///
/// Adjust this if your SpliceIndex already provides a method, e.g.
/// `idx.chr_name_to_id(&str) -> Option<usize>`.
fn build_chr_map(idx: &SpliceIndex) -> std::collections::HashMap<String, usize> {
    // Common patterns:
    // - idx.chr_names: Vec<String>
    // - idx.chrom_names: Vec<String>
    //
    // Change the field name here to match your SpliceIndex.
    let names: &Vec<String> = &idx.chr_names;

    let mut map = std::collections::HashMap::with_capacity(names.len());
    for (i, n) in names.iter().enumerate() {
        map.insert(n.clone(), i);
    }
    map
}*/

fn main() -> Result<()> {
    let args = QuantCli::parse();

    if args.threads > 0 {
        // If build_global fails (already initialized), ignore.
        let _ = rayon::ThreadPoolBuilder::new()
            .num_threads(args.threads)
            .build_global();
    }

    let idx = SpliceIndex::load(&args.index)
        .with_context(|| format!("reading index {}", args.index.display()))?;
    println!("{idx}");

    //let chr_map = build_chr_map(&idx);
    let chr_map = build_chr_map_fuzzy(&idx);

    let match_opts = MatchOptions {
        require_strand: args.require_strand,
        require_exact_junction_chain: args.require_exact_junction_chain,
        max_5p_overhang_bp: args.max_5p_overhang_bp,
        max_3p_overhang_bp: args.max_3p_overhang_bp,
        allowed_intronic_gap_size: args.allowed_intronic_gap_size,
    };

    let mut reader = Reader::from_path(&args.bam)
        .with_context(|| format!("bam file could not be read: {}", args.bam.display()))?;

    let hts_threads = compute_io_threads(args.threads);
    if let Err(e) = reader.set_threads(hts_threads) {
        eprintln!(
            "Warning: failed to enable HTSlib threading ({}). Continuing single-threaded.",
            e
        )
    };
    reader.set_threads(hts_threads)?;
    let header = reader.header().clone();

    // ----------------------------
    // Stage 1: stream BAM -> jobs
    // ----------------------------
    let mut jobs: Vec<Job> = Vec::new();
    let mut n_seen: usize = 0;

    // ----------------------------
    //    merge partial results
    // ----------------------------
    let mut merged = Scdata::new(1, MatrixValueType::Real);
    let mut merged_intron = Scdata::new(1, MatrixValueType::Real);
    let mut merged_report = MappingInfo::new(
        None,
        args.min_mapq as f32,
        args.max_reads.unwrap_or(usize::MAX),
    );
    merged_report.start_counter();

    #[cfg(debug_assertions)]
    let mut i = 0;
    for r in reader.records() {
        #[cfg(debug_assertions)]
        {
            i += 1;
            if i - n_seen > 1_000_0000 {
                panic!(
                    "We have found more than 1 mio reads that did not pass initial filters.\n read {i}, processed {n_seen},  Wrong gene model?\n{merged_report}"
                );
            }
        }
        let rec = r.context("BAM read error")?;

        if rec.is_unmapped() {
            merged_report.report("unmapped");
            continue;
        }
        if rec.mapq() < args.min_mapq {
            merged_report.report("mapq failed");
            continue;
        }
        if rec.is_secondary() || rec.is_supplementary() {
            merged_report.report("secondary or supplemantary");
            continue;
        }
        if args.read1_only && !rec.is_first_in_template() {
            merged_report.report("read!=1");
            continue;
        }

        let cb_raw = match aux_tag_str(&rec, *b"CB") {
            Some(v) => v,
            None => {
                merged_report.report("no CB tag");
                continue;
            }
        };
        let ub = match aux_tag_str(&rec, *b"UB") {
            Some(v) => v,
            None => {
                merged_report.report("no UB tag");
                continue;
            }
        };

        let cb = normalize_10x_barcode(cb_raw);

        let cell = match dna_to_u64(cb) {
            Some(v) => v,
            None => continue,
        };
        let umi = match dna_to_u64(ub) {
            Some(v) => v,
            None => continue,
        };

        let tid = rec.tid();
        if tid < 0 {
            merged_report.report("tid below 1");
            continue;
        }
        let chr_name = std::str::from_utf8(header.tid2name(tid as u32))
            .context("Invalid chromosome name in BAM header")?;

        let chr_id = match chr_map.get(chr_name) {
            Some(&id) => id,
            None => {
                merged_report.report("contig not in index");
                continue; // contig not in splice index
            }
        };

        let (spliced, _start0, _end0) = match record_to_spliced_read(&rec, chr_id) {
            Some(v) => v,
            None => continue,
        };

        jobs.push(Job {
            cell,
            umi,
            //chr_id,
            //start0,
            //end0,
            spliced,
        });

        n_seen += 1;
        if let Some(maxr) = args.max_reads {
            if n_seen >= maxr {
                break;
            }
        }

        if jobs.len() >= CHUNK {
            println!("Processing chunk of size {}", CHUNK);
            merged_report.stop_file_io_time();

            match args.quant_mode {
                QuantMode::Gene => process_chunk_gene(
                    &jobs,
                    &idx,
                    match_opts,
                    &mut merged,
                    &mut merged_intron,
                    &mut merged_report,
                    args.min_mapq,
                )?,
                QuantMode::Transcript => process_chunk_transcript(
                    &jobs,
                    &idx,
                    match_opts,
                    &mut merged,
                    &mut merged_intron,
                    &mut merged_report,
                    args.min_mapq,
                )?,
            }

            jobs.clear();
        }
    }

    // ----------------------------------------
    // Stage 2: parallel chunks -> partial scdata
    // ----------------------------------------

    if jobs.len() > 0 {
        println!("Processing final chunk of size {}", jobs.len());
        merged_report.stop_file_io_time();
        match args.quant_mode {
            QuantMode::Gene => process_chunk_gene(
                &jobs,
                &idx,
                match_opts,
                &mut merged,
                &mut merged_intron,
                &mut merged_report,
                args.min_mapq,
            )?,
            QuantMode::Transcript => process_chunk_transcript(
                &jobs,
                &idx,
                match_opts,
                &mut merged,
                &mut merged_intron,
                &mut merged_report,
                args.min_mapq,
            )?,
        }
        jobs.clear();
    }

    println!("Writing outfiles");

    let features = match args.quant_mode {
        QuantMode::Gene => {
            //let names = idx.gene_names();
            IndexedGenes::from_names(&idx.gene_names())
        },
        QuantMode::Transcript => { 
            //let names = idx.transcript_names();
            IndexedGenes::from_names(&idx.transcript_names())
        },
    };
    merged_report.stop_single_processor_time();

    

    if args.split_intronic{
        // 1) Compute passing cells ONCE from the real data (merged)
        let pass = merged.passing_cell_set_by_umi(args.min_cell_counts);
        // 2) Apply to BOTH datasets
        merged.restrict_to_cells(&pass);
        merged_intron.restrict_to_cells(&pass);
        merged_report.stop_multi_processor_time();
        
        println!("Writing matrix files");
        let _ = merged.write_sparse(&args.outpath, &features, 0);
        println!("Writing intronic matrix files");
        let _ = merged_intron.write_sparse(
            &add_suffix(&args.outpath, "_intronic"),
            &features,
            0,
        );
    }else {
        merged.merge( &merged_intron );
        let pass = merged.passing_cell_set_by_umi(args.min_cell_counts);
        merged.restrict_to_cells(&pass);

        merged_report.stop_multi_processor_time();

        println!("Writing matrix files");
        let _ = merged.write_sparse(&args.outpath, &features, 0);

    }

    merged_report.stop_file_io_time();

    println!("{merged_report}"); 

    let log_str = format!("{merged_report}");
    let log_path = args.outpath.with_extension("log");
    let mut file = File::create(&log_path)
        .expect("failed to create log file");
    file.write_all(log_str.as_bytes())
        .expect("failed to write log file");

    Ok(())
}

fn add_suffix(path: &PathBuf, suffix: &str) -> PathBuf {
    let parent = path.parent().unwrap_or_else(|| std::path::Path::new(""));
    let stem = path.file_stem().unwrap().to_string_lossy();
    let ext = path.extension().map(|e| e.to_string_lossy());

    let new_name = match ext {
        Some(e) => format!("{stem}{suffix}.{e}"),
        None => format!("{stem}{suffix}"),
    };

    parent.join(new_name)
}

fn process_chunk_gene(
    jobs: &[Job],
    idx: &SpliceIndex,
    match_opts: MatchOptions,
    merged: &mut Scdata,
    merged_intron: &mut Scdata,
    merged_report: &mut MappingInfo,
    min_mapq: u8,
) -> Result<()> {
    let threads = rayon::current_num_threads().max(1);
    let chunk_size = (jobs.len() / threads).max(10_000);

    let partials: Vec<(Scdata, Scdata, MappingInfo)> = jobs
        .par_chunks(chunk_size)
        .map(|chunk| {
            let mut sc = Scdata::new(1, MatrixValueType::Real);
            let mut sc_intron = Scdata::new(1, MatrixValueType::Real);
            let mut rep = MappingInfo::new(None, min_mapq as f32, usize::MAX);

            for job in chunk { 
                let gene_hits = idx.match_genes(&job.spliced, match_opts);
                if gene_hits.is_empty() {
                    rep.report("no hit");
                    continue;
                }

                let g = &gene_hits[0];
                rep.report(g.best_hit.class.to_string());

                let gid = g.gene_id; // (GeneId is usize in your code)
                //println!("cell: {}",job.cell);

                if g.best_hit.class == MatchClass::Intronic {
                    sc_intron.try_insert(&job.cell, GeneUmiHash(gid, job.umi), 1.0, &mut rep);
                } else {
                    sc.try_insert(&job.cell, GeneUmiHash(gid, job.umi), 1.0, &mut rep);
                }
            }

            (sc, sc_intron, rep)
        })
        .collect();

    merged_report.stop_multi_processor_time();

    for (sc, sc_intron, rep) in partials {
        merged.merge(&sc);
        merged_intron.merge(&sc_intron);
        merged_report.merge(&rep);
    }
    merged_report.stop_single_processor_time();

    Ok(())
}


fn process_chunk_transcript(
    jobs: &[Job],
    idx: &SpliceIndex,
    match_opts: MatchOptions,
    merged: &mut Scdata,
    merged_intron: &mut Scdata,
    merged_report: &mut MappingInfo,
    min_mapq: u8,
) -> Result<()> {
    let threads = rayon::current_num_threads().max(1);
    let chunk_size = (jobs.len() / threads).max(10_000);

    let partials: Vec<(Scdata, Scdata, MappingInfo)> = jobs
        .par_chunks(chunk_size)
        .map(|chunk| {
            let mut sc = Scdata::new(1, MatrixValueType::Real);
            let mut sc_intron = Scdata::new(1, MatrixValueType::Real);
            let mut rep = MappingInfo::new(None, min_mapq as f32, usize::MAX);

            for job in chunk {
                let tx_hits = idx.match_transcripts(&job.spliced, match_opts);
                if tx_hits.is_empty() {
                    rep.report("no hit");
                    continue;
                }

                let h = &tx_hits[0];
                rep.report(h.hit.class.to_string());

                let tid = h.transcript_id;

                if h.hit.class == MatchClass::Intronic {
                    sc_intron.try_insert(&job.cell, GeneUmiHash(tid, job.umi), 1.0, &mut rep);
                } else {
                    sc.try_insert(&job.cell, GeneUmiHash(tid, job.umi), 1.0, &mut rep);
                }
            }

            (sc, sc_intron, rep)
        })
        .collect();
    merged_report.stop_multi_processor_time();

    for (sc, sc_intron, rep) in partials {
        merged.merge(&sc);
        merged_intron.merge(&sc_intron);
        merged_report.merge(&rep);
    }
    merged_report.stop_single_processor_time();

    Ok(())
}

fn build_chr_map_fuzzy(idx: &SpliceIndex) -> std::collections::HashMap<String, usize> {
    let mut map = std::collections::HashMap::new();

    for (i, n) in idx.chr_names.iter().enumerate() {
        // exact
        map.entry(n.clone()).or_insert(i);

        // without chr prefix
        let no_chr = n.strip_prefix("chr").unwrap_or(n).to_string();
        map.entry(no_chr).or_insert(i);

        // with chr prefix
        let with_chr = if n.starts_with("chr") {
            n.clone()
        } else {
            format!("chr{n}")
        };
        map.entry(with_chr).or_insert(i);

        // mito aliases
        if n == "MT" {
            map.entry("chrM".to_string()).or_insert(i);
            map.entry("M".to_string()).or_insert(i);
        }
        if n == "chrM" {
            map.entry("MT".to_string()).or_insert(i);
            map.entry("M".to_string()).or_insert(i);
        }
    }

    map
}