seqwish 0.1.3

use crate::intervaltree::AdaptiveTree;
use crate::intervaltree::IntervalTree;
use crate::paf::PafRow;
use crate::pos::{decr_pos, incr_pos, incr_pos_by, is_rev, make_pos_t, offset, PosT};
use crate::seqindex::SeqIndex;
use flate2::read::MultiGzDecoder;
use std::fs::File;
use std::io::{self, BufRead, BufReader};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;

/// Hash function for match parameters
/// Uses a Wang hash-like mixing algorithm
pub fn match_hash(q: u64, t: u64, l: u64) -> u64 {
    let mut seed = q | t | l;
    seed ^= q
        .wrapping_add(0x9e3779b97f4a7c15)
        .wrapping_add(seed << 17)
        .wrapping_add(seed >> 9);
    seed ^= t
        .wrapping_add(0x9e3779b97f4a7c15)
        .wrapping_add(seed << 7)
        .wrapping_add(seed >> 23);
    seed ^= l
        .wrapping_add(0x9e3779b97f4a7c15)
        .wrapping_add(seed << 9)
        .wrapping_add(seed >> 2);
    seed
}

/// Determine if a match should be kept based on sparsification factor
/// Returns true if the match hash is below the threshold determined by f
pub fn keep_sparse(q: u64, t: u64, l: u64, f: f32) -> bool {
    // Match C++ behavior: compare hash as f64 against max * f
    (match_hash(q, t, l) as f64) < (u64::MAX as f64 * f as f64)
}

/// Worker thread that processes PAF alignments
fn paf_worker(
    reader: Arc<Mutex<Box<dyn BufRead + Send>>>,
    more: Arc<AtomicBool>,
    aln_iitree: Arc<Mutex<AdaptiveTree<u64, PosT>>>,
    seqidx: Arc<SeqIndex>,
    min_match_len: u64,
    sparsification_factor: f32,
) -> io::Result<()> {
    while more.load(Ordering::Relaxed) {
        // Read one line with mutex protection
        let line = {
            let mut reader_guard = reader.lock().unwrap();
            let mut line = String::new();
            match reader_guard.read_line(&mut line) {
                Ok(0) => {
                    // EOF
                    more.store(false, Ordering::Relaxed);
                    return Ok(());
                }
                Ok(_) => {
                    // Strip newline
                    if line.ends_with('\n') {
                        line.pop();
                        if line.ends_with('\r') {
                            line.pop();
                        }
                    }
                    line
                }
                Err(e) => {
                    more.store(false, Ordering::Relaxed);
                    return Err(e);
                }
            }
        };

        if line.is_empty() {
            continue;
        }

        // Parse PAF row
        let paf = match PafRow::from_line(&line) {
            Some(p) => p,
            None => continue, // Skip malformed lines
        };

        // Check if coordinates are reasonable
        if paf.query_sequence_length == 0
            || paf.target_sequence_length == 0
            || paf.query_start >= paf.query_sequence_length
            || paf.query_end > paf.query_sequence_length
            || paf.query_start >= paf.query_end
            || paf.target_start >= paf.target_sequence_length
            || paf.target_end > paf.target_sequence_length
            || paf.target_start >= paf.target_end
        {
            continue;
        }

        // Look up sequence indices
        let query_idx = match seqidx.rank_of_seq_named(&paf.query_sequence_name) {
            Some(idx) => idx,
            None => continue,
        };

        let target_idx = match seqidx.rank_of_seq_named(&paf.target_sequence_name) {
            Some(idx) => idx,
            None => continue,
        };

        // Determine orientations and starting positions
        let q_rev = !paf.query_target_same_strand;
        let q_all_pos = if q_rev {
            seqidx
                .pos_in_all_seqs_by_id(query_idx, paf.query_end, false)
                .unwrap()
                - 1
        } else {
            seqidx
                .pos_in_all_seqs_by_id(query_idx, paf.query_start, false)
                .unwrap()
        };

        let t_all_pos = seqidx
            .pos_in_all_seqs_by_id(target_idx, paf.target_start, false)
            .unwrap();

        let mut q_pos = make_pos_t(q_all_pos, q_rev);
        let mut t_pos = make_pos_t(t_all_pos, false);

        // Process CIGAR operations
        for c in &paf.cigar {
            match c.op {
                b'M' | b'=' | b'X' => {
                    let mut q_pos_match_start = q_pos;
                    let mut t_pos_match_start = t_pos;
                    let mut match_len = 0u64;

                    // Helper to add a match to the iitree
                    let add_match =
                        |q_start: PosT, q_end: PosT, t_start: PosT, t_end: PosT, len: u64| {
                            if len >= min_match_len
                                && (sparsification_factor == 0.0
                                    || keep_sparse(
                                        offset(q_start),
                                        offset(t_start),
                                        len,
                                        sparsification_factor,
                                    ))
                            {
                                let mut tree = aln_iitree.lock().unwrap();
                                if is_rev(q_end) {
                                    // Reverse query
                                    let mut x_pos = q_end;
                                    decr_pos(&mut x_pos);
                                    tree.add(
                                        offset(x_pos),
                                        offset(q_start) + 1,
                                        make_pos_t(offset(t_end) - 1, true),
                                    )
                                    .ok();
                                    tree.add(
                                        offset(t_start),
                                        offset(t_end),
                                        make_pos_t(offset(q_start), true),
                                    )
                                    .ok();
                                } else {
                                    // Forward query
                                    tree.add(offset(q_start), offset(q_end), t_start).ok();
                                    tree.add(offset(t_start), offset(t_end), q_start).ok();
                                }
                            }
                        };

                    // Process each base in the match
                    for _ in 0..c.len {
                        let query_base = seqidx.at_pos(q_pos).ok_or_else(|| {
                            io::Error::new(io::ErrorKind::Other, "Invalid query position")
                        })?;
                        let target_base = seqidx.at_pos(t_pos).ok_or_else(|| {
                            io::Error::new(io::ErrorKind::Other, "Invalid target position")
                        })?;

                        if query_base == target_base
                            && query_base != 'N'
                            && offset(q_pos) != offset(t_pos)
                        {
                            // Start or extend match
                            if match_len == 0 {
                                q_pos_match_start = q_pos;
                                t_pos_match_start = t_pos;
                            }
                            match_len += 1;
                            incr_pos(&mut q_pos);
                            incr_pos(&mut t_pos);
                        } else {
                            // Mismatch or end of match
                            if match_len > 0 {
                                add_match(
                                    q_pos_match_start,
                                    q_pos,
                                    t_pos_match_start,
                                    t_pos,
                                    match_len,
                                );
                                match_len = 0;
                            }
                            incr_pos(&mut q_pos);
                            incr_pos(&mut t_pos);
                        }
                    }

                    // Handle any final match
                    if match_len > 0 {
                        add_match(
                            q_pos_match_start,
                            q_pos,
                            t_pos_match_start,
                            t_pos,
                            match_len,
                        );
                    }
                }
                b'I' => {
                    // Insertion in query (skip bases in query)
                    incr_pos_by(&mut q_pos, c.len as usize);
                }
                b'D' => {
                    // Deletion in query (skip bases in target)
                    incr_pos_by(&mut t_pos, c.len as usize);
                }
                _ => {
                    // Unknown operation, skip
                }
            }
        }
    }

    Ok(())
}

/// Unpack PAF alignments into an interval tree
///
/// Reads a PAF file (gzipped or plain text), spawns worker threads to process alignments,
/// and populates the interval tree with match intervals.
pub fn unpack_paf_alignments(
    paf_file: &str,
    aln_iitree: Arc<Mutex<AdaptiveTree<u64, PosT>>>,
    seqidx: Arc<SeqIndex>,
    min_match_len: u64,
    sparsification_factor: f32,
    num_threads: usize,
) -> io::Result<()> {
    // Open PAF file (auto-detect gzipped or plain text)
    let file = File::open(paf_file)?;
    let reader: Box<dyn BufRead + Send> = if paf_file.ends_with(".gz") {
        Box::new(BufReader::new(MultiGzDecoder::new(file)))
    } else {
        Box::new(BufReader::new(file))
    };
    let reader = Arc::new(Mutex::new(reader));

    let more = Arc::new(AtomicBool::new(true));

    // Spawn worker threads
    let mut handles = Vec::new();
    for _ in 0..num_threads {
        let reader_clone = Arc::clone(&reader);
        let more_clone = Arc::clone(&more);
        let tree_clone = Arc::clone(&aln_iitree);
        let seqidx_clone = Arc::clone(&seqidx);

        let handle = thread::spawn(move || {
            paf_worker(
                reader_clone,
                more_clone,
                tree_clone,
                seqidx_clone,
                min_match_len,
                sparsification_factor,
            )
        });

        handles.push(handle);
    }

    // Wait for all workers to finish
    for handle in handles {
        match handle.join() {
            Ok(result) => result?,
            Err(_) => {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "Worker thread panicked",
                ))
            }
        }
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::seqindex::SeqIndex;
    use flate2::write::GzEncoder;
    use flate2::Compression;
    use std::fs::File;
    use std::io::Write;

    #[test]
    fn test_match_hash_deterministic() {
        // Same inputs should give same hash
        let hash1 = match_hash(100, 200, 50);
        let hash2 = match_hash(100, 200, 50);
        assert_eq!(hash1, hash2);
    }

    #[test]
    fn test_match_hash_different_inputs() {
        // Different inputs should give different hashes (usually)
        let hash1 = match_hash(100, 200, 50);
        let hash2 = match_hash(100, 200, 51);
        assert_ne!(hash1, hash2);

        let hash3 = match_hash(100, 201, 50);
        assert_ne!(hash1, hash3);

        let hash4 = match_hash(101, 200, 50);
        assert_ne!(hash1, hash4);
    }

    #[test]
    fn test_match_hash_zero_inputs() {
        let hash = match_hash(0, 0, 0);
        // Verify it matches the C++ implementation
        assert_eq!(hash, 668627425924508912);
    }

    #[test]
    fn test_keep_sparse_always_keep() {
        // f = 1.0 should keep everything
        assert!(keep_sparse(100, 200, 50, 1.0));
        assert!(keep_sparse(0, 0, 0, 1.0));
        assert!(keep_sparse(u64::MAX, u64::MAX, u64::MAX, 1.0));
    }

    #[test]
    fn test_keep_sparse_never_keep() {
        // f = 0.0 should keep nothing
        assert!(!keep_sparse(100, 200, 50, 0.0));
        assert!(!keep_sparse(0, 0, 0, 0.0));
        assert!(!keep_sparse(1, 1, 1, 0.0));
    }

    #[test]
    fn test_keep_sparse_threshold_behavior() {
        // The hash function produces values in a clustered range
        // Test that the threshold comparison works correctly
        let hash_val = match_hash(100, 200, 50);

        // With f = 1.0, everything should pass
        assert!((hash_val as f64) < (u64::MAX as f64 * 1.0));

        // With f = 0.0, nothing should pass
        assert!(!((hash_val as f64) < (u64::MAX as f64 * 0.0)));

        // With intermediate values, it depends on the hash
        // Just verify the comparison logic works
        let small_f = 0.0001;
        let large_f = 0.9999;
        let _small_result = keep_sparse(100, 200, 50, small_f);
        let large_result = keep_sparse(100, 200, 50, large_f);
        assert!(large_result); // Large f should keep most things
    }

    #[test]
    fn test_keep_sparse_consistency() {
        // Same inputs should give same result
        let result1 = keep_sparse(12345, 67890, 111, 0.3);
        let result2 = keep_sparse(12345, 67890, 111, 0.3);
        assert_eq!(result1, result2);
    }

    #[test]
    fn test_unpack_paf_alignments_integration() -> std::io::Result<()> {
        use std::time::{SystemTime, UNIX_EPOCH};

        // Create unique temporary paths using timestamp and thread ID
        let unique_id = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let fasta_path = format!("/tmp/test_alignments_{}.fasta", unique_id);
        let paf_path = format!("/tmp/test_alignments_{}.paf.gz", unique_id);
        let iitree_path = format!("/tmp/test_alignments_{}.iitree", unique_id);

        // Create temporary FASTA file with two sequences
        let mut fasta_file = File::create(&fasta_path)?;
        writeln!(fasta_file, ">seq1")?;
        writeln!(fasta_file, "ACGTACGTACGT")?; // 12 bases
        writeln!(fasta_file, ">seq2")?;
        writeln!(fasta_file, "ACGTACGTACGT")?; // 12 bases (identical)
        fasta_file.flush()?;

        // Build SeqIndex
        let mut seqidx = SeqIndex::new();
        seqidx
            .build_index(&fasta_path)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
        let seqidx = Arc::new(seqidx);

        // Create temporary gzipped PAF file
        // PAF format: query_name, query_len, query_start, query_end, strand, target_name, target_len, target_start, target_end, matches, block_len, mapq, [tags]
        let paf_file = File::create(&paf_path)?;
        let mut gz = GzEncoder::new(paf_file, Compression::default());

        // Simple alignment: seq1[0..12] aligns to seq2[0..12] with perfect match
        writeln!(
            gz,
            "seq1\t12\t0\t12\t+\tseq2\t12\t0\t12\t12\t12\t60\tcg:Z:12M"
        )?;
        gz.finish()?;

        // Create iitree
        let mut tree = AdaptiveTree::new_disk(&iitree_path)?;
        tree.open_writer()?;
        let tree = Arc::new(Mutex::new(tree));

        // Process PAF file
        unpack_paf_alignments(
            &paf_path,
            Arc::clone(&tree),
            Arc::clone(&seqidx),
            1,   // min_match_len
            0.0, // sparsification_factor (keep all)
            1,   // num_threads
        )?;

        // Index the tree
        {
            let mut tree_guard = tree.lock().unwrap();
            tree_guard.index()?;
        }

        // Verify intervals were added
        let tree_guard = tree.lock().unwrap();
        let num_intervals = tree_guard.len();

        // We expect at least 2 intervals (one for each direction of the alignment)
        assert!(
            num_intervals >= 2,
            "Expected at least 2 intervals, got {}",
            num_intervals
        );

        // Cleanup
        std::fs::remove_file(&fasta_path).ok();
        std::fs::remove_file(&paf_path).ok();
        std::fs::remove_file(&iitree_path).ok();

        Ok(())
    }
}