md_analysis 0.1.0

/// 三阶段 worker 流水线：reader → stats workers → writer
///
/// 使用 std::sync primitives（Mutex + Condvar）实现 fan-out channel，
/// 解决标准库 mpsc 不支持多消费者（multiple receivers）的问题。
use rust_htslib::bam::Read;
use std::path::Path;
use std::{format, thread};

use crate::binning::FinalResult;
use crate::config::{NUM_BINS, default_worker_count};
use crate::pbar::{DEFAULT_INTERVAL, get_spin_pb};
use crate::record::RecordData;
use crate::writer::write_final_result_csv;

/// 流水线配置参数
#[derive(Clone)]
pub struct PipelineConfig {
    /// stats worker 数量
    pub worker_count: usize,
}

impl PipelineConfig {
    pub fn new(worker_count: usize) -> Self {
        Self { worker_count }
    }
}

impl Default for PipelineConfig {
    fn default() -> Self {
        Self {
            worker_count: default_worker_count(),
        }
    }
}

/// 运行三阶段并行流水线，返回 (Histogram, PerBaseHistograms)。
pub fn run_pipeline(
    bam_path: &str,
    config: &PipelineConfig,
    slice_n: usize,
    output_dir: impl AsRef<Path>,
) {
    let worker_cnt = config.worker_count.max(1);

    let bam_path = bam_path.to_string();

    let length_min = 0;
    let length_max = usize::MAX;

    let (reader_sender, reader_recv) = crossbeam::channel::bounded(1000);

    thread::scope(|thread_scope| {
        thread_scope.spawn(move || {
            let mut reader = rust_htslib::bam::Reader::from_path(&bam_path)
                .expect(&format!("read {} error", bam_path));
            for record in reader.records() {
                if let Ok(record) = record {
                    if record.seq_len() >= length_min && record.seq_len() < length_max {
                        match reader_sender.send(record) {
                            Ok(_) => {}
                            Err(e) => {
                                tracing::error!("reader send error. {:#?}", e);
                            }
                        }
                    }
                } else {
                    tracing::warn!("error record");
                }
            }
        });

        let (worker_sender, worker_recv) = crossbeam::channel::bounded(1000);
        for _ in 0..worker_cnt {
            thread_scope.spawn({
                let reader_recv = reader_recv.clone();
                let worker_sender = worker_sender.clone();
                move || {
                    let mut result: FinalResult<NUM_BINS> = if slice_n > 0 {
                        FinalResult::new_first_n_last_n()
                    } else {
                        FinalResult::new_all()
                    };

                    for record in reader_recv {
                        let record_data: RecordData = (&record).into();

                        if slice_n > 0 {
                            let first_start = 0;
                            let first_end = (first_start + slice_n).min(record_data.seq_len);

                            let last_start = record_data.seq_len.saturating_sub(slice_n);
                            let last_end = record_data.seq_len;

                            result.get_first_n_total_hist_mut().collect_ar_dw(
                                (&record_data.ar[first_start..first_end]).iter(),
                                (&record_data.dw[first_start..first_end]).iter(),
                            );

                            result.get_first_n_per_base_hist_mut().collect_ar_dw(
                                &record_data.bases.as_str()[first_start..first_end],
                                &record_data.ar[first_start..first_end],
                                &record_data.dw[first_start..first_end],
                            );

                            result.get_last_n_total_hist_mut().collect_ar_dw(
                                (&record_data.ar[last_start..last_end]).iter(),
                                (&record_data.dw[last_start..last_end]).iter(),
                            );

                            result.get_last_n_per_base_hist_mut().collect_ar_dw(
                                &record_data.bases.as_str()[last_start..last_end],
                                &record_data.ar[last_start..last_end],
                                &record_data.dw[last_start..last_end],
                            );
                        } else {
                            result
                                .get_total_hist_mut()
                                .collect_ar_dw(record_data.ar.iter(), record_data.dw.iter());

                            result.get_total_per_base_hist_mut().collect_ar_dw(
                                record_data.bases.as_str(),
                                &record_data.ar,
                                &record_data.dw,
                            );
                        }
                    }

                    worker_sender.send(result)
                }
            });
        }
        drop(reader_recv);
        drop(worker_sender);

        let mut final_result: FinalResult<NUM_BINS> = if slice_n > 0 {
            FinalResult::new_first_n_last_n()
        } else {
            FinalResult::new_all()
        };

        let pb = get_spin_pb(format!("processing"), DEFAULT_INTERVAL);

        for result in worker_recv {
            final_result.merge(&result);
            pb.inc(1);
        }
        pb.finish();

        write_final_result_csv(output_dir.as_ref(), &final_result).unwrap();
    });
}