table.rs

use std::collections::HashMap;
use std::path::Path;
use anyhow::Result;
use std::fs::File;
use std::io::{BufWriter, Write};
use crate::merge::MergedRecord;

/// Represents a multi-sample ASV abundance table (R's makeSequenceTable output).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SequenceTable {
    /// List of sample names (representing rows).
    pub samples: Vec<String>,
    /// List of global ASV sequences (representing columns, sorted by total abundance descending).
    pub sequences: Vec<Vec<u8>>,
    /// 2D abundance grid: matrix[sample_idx][asv_idx]
    pub matrix: Vec<Vec<u32>>,
}

/// Constructs a unified sequence table from a list of samples and their merged records.
pub fn make_sequence_table(
    samples_data: &[(String, Vec<MergedRecord>)]
) -> SequenceTable {
    // 1. Compile total global abundance for each unique sequence
    let mut global_abundance: HashMap<Vec<u8>, u64> = HashMap::new();
    for (_sample_name, records) in samples_data {
        for record in records {
            let seq = &record.sequence;
            let entry = global_abundance.entry(seq.clone()).or_insert(0);
            *entry += record.abundance as u64;
        }
    }

    // 2. Sort sequences by global abundance in descending order
    let mut sorted_seqs: Vec<(Vec<u8>, u64)> = global_abundance.into_iter().collect();
    sorted_seqs.sort_by(|a, b| b.1.cmp(&a.1));

    let sequences: Vec<Vec<u8>> = sorted_seqs.into_iter().map(|(seq, _)| seq).collect();

    // 3. Build the 2D abundance matrix
    let n_samples = samples_data.len();
    let n_seqs = sequences.len();
    let mut matrix = vec![vec![0u32; n_seqs]; n_samples];
    let mut samples = Vec::with_capacity(n_samples);

    for (sample_idx, (sample_name, records)) in samples_data.iter().enumerate() {
        samples.push(sample_name.clone());

        // Map sequence to its abundance in this sample
        let mut sample_map: HashMap<&[u8], u32> = HashMap::with_capacity(records.len());
        for record in records {
            sample_map.insert(&record.sequence, record.abundance);
        }

        for (seq_idx, seq) in sequences.iter().enumerate() {
            if let Some(&abundance) = sample_map.get(seq.as_slice()) {
                matrix[sample_idx][seq_idx] = abundance;
            }
        }
    }

    SequenceTable {
        samples,
        sequences,
        matrix,
    }
}

impl SequenceTable {
    /// Writes the sequence table to a standard tab-separated value (TSV) file.
    pub fn write_tsv<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        let file = File::create(path.as_ref())?;
        let mut writer = BufWriter::new(file);

        // Header: "SampleID" \t Sequence1 \t Sequence2 ...
        writer.write_all(b"SampleID")?;
        for seq in &self.sequences {
            writer.write_all(b"\t")?;
            writer.write_all(seq)?;
        }
        writer.write_all(b"\n")?;

        // Rows: SampleName \t Abund1 \t Abund2 ...
        for (sample_idx, sample_name) in self.samples.iter().enumerate() {
            writer.write_all(sample_name.as_bytes())?;
            for seq_idx in 0..self.sequences.len() {
                writer.write_all(b"\t")?;
                let count_str = self.matrix[sample_idx][seq_idx].to_string();
                writer.write_all(count_str.as_bytes())?;
            }
            writer.write_all(b"\n")?;
        }

        writer.flush()?;
        Ok(())
    }
}

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

    #[test]
    fn test_make_sequence_table() {
        let sample1 = (
            "Sample1".to_string(),
            vec![
                MergedRecord {
                    sequence: b"ACTG".to_vec(),
                    abundance: 10,
                    forward_center_idx: 0,
                    reverse_center_idx: 0,
                    overlap_len: 10,
                    mismatch_count: 0,
                },
                MergedRecord {
                    sequence: b"AAAA".to_vec(),
                    abundance: 2,
                    forward_center_idx: 1,
                    reverse_center_idx: 1,
                    overlap_len: 10,
                    mismatch_count: 0,
                },
            ],
        );

        let sample2 = (
            "Sample2".to_string(),
            vec![
                MergedRecord {
                    sequence: b"GGGG".to_vec(),
                    abundance: 20,
                    forward_center_idx: 2,
                    reverse_center_idx: 2,
                    overlap_len: 10,
                    mismatch_count: 0,
                },
                MergedRecord {
                    sequence: b"ACTG".to_vec(),
                    abundance: 5,
                    forward_center_idx: 0,
                    reverse_center_idx: 0,
                    overlap_len: 10,
                    mismatch_count: 0,
                },
            ],
        );

        let table = make_sequence_table(&[sample1, sample2]);

        // Global abundance:
        // GGGG: 20 (Sample 2)
        // ACTG: 15 (Sample 1: 10, Sample 2: 5)
        // AAAA: 2 (Sample 1)
        // So global ranking order of sequences must be GGGG, ACTG, AAAA!
        assert_eq!(table.sequences.len(), 3);
        assert_eq!(table.sequences[0], b"GGGG");
        assert_eq!(table.sequences[1], b"ACTG");
        assert_eq!(table.sequences[2], b"AAAA");

        assert_eq!(table.samples, vec!["Sample1".to_string(), "Sample2".to_string()]);

        // Row 1 (Sample1) counts: GGGG=0, ACTG=10, AAAA=2
        assert_eq!(table.matrix[0], vec![0, 10, 2]);

        // Row 2 (Sample2) counts: GGGG=20, ACTG=5, AAAA=0
        assert_eq!(table.matrix[1], vec![20, 5, 0]);
    }
}