1use anyhow::{Context, Result};
19use rustc_hash::FxHashMap;
20use std::fs::File;
21use std::io::{BufRead, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
22use std::path::Path;
23use std::process::Command;
24
25use crate::snp::{self, SnpStatus};
26
27const FDB_MAGIC: &[u8; 8] = b"FLANKDB\0";
28
29#[derive(Debug, Clone)]
37pub struct ArgPosition {
38 pub arg_name: String,
40 pub contig_name: String,
42 pub contig_seq: String,
44 pub contig_len: usize,
46 pub arg_start: usize,
48 pub arg_end: usize,
50 pub strand: char,
52}
53
54pub const GENUS_TIE_PCT: f64 = 1.0;
58
59#[derive(Debug, Clone)]
60pub struct GenusResult {
61 pub arg_name: String,
63 pub contig_name: String,
65 pub genus: Option<String>,
67 pub confidence: f64,
69 pub specificity: f64,
71 pub upstream_len: usize,
73 pub downstream_len: usize,
75 pub top_matches: Vec<(String, f64)>,
77 pub snp_status: SnpStatus,
79 pub upstream_seq: String,
82 pub downstream_seq: String,
84 pub context: String,
87 pub n_genera_tied: usize,
90 pub species: Option<String>,
93 pub species_top_matches: Vec<(String, f64)>,
95 pub n_species_tied: usize,
97}
98
99impl Default for GenusResult {
100 fn default() -> Self {
101 Self {
102 arg_name: String::new(),
103 contig_name: String::new(),
104 genus: None,
105 confidence: 0.0,
106 specificity: 0.0,
107 upstream_len: 0,
108 downstream_len: 0,
109 top_matches: vec![],
110 snp_status: SnpStatus::NotApplicable,
111 upstream_seq: String::new(),
112 downstream_seq: String::new(),
113 context: "NA".to_string(),
114 n_genera_tied: 0,
115 species: None,
116 species_top_matches: vec![],
117 n_species_tied: 0,
118 }
119 }
120}
121
122#[derive(Debug, Clone)]
126pub struct FlankingRecord {
127 pub contig: String,
129 pub genus: String,
131 pub upstream: String,
133 pub downstream: String,
135}
136
137#[derive(Debug, Clone)]
143struct FdbIndexEntry {
144 offset: u64,
145 compressed_len: u32,
146 record_count: u32,
147}
148
149pub struct FlankingDatabase {
165 file: File,
166 index: FxHashMap<String, FdbIndexEntry>,
167 gene_name_to_key: FxHashMap<String, String>,
170}
171
172impl FlankingDatabase {
173 pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
180 let mut file = File::open(path.as_ref())
181 .with_context(|| format!("Failed to open fdb: {}", path.as_ref().display()))?;
182
183 let mut magic = [0u8; 8];
185 file.read_exact(&mut magic)?;
186 if &magic != FDB_MAGIC {
187 anyhow::bail!("Invalid fdb magic");
188 }
189
190 let mut buf4 = [0u8; 4];
191 let mut buf8 = [0u8; 8];
192
193 file.read_exact(&mut buf4)?;
194 let _version = u32::from_le_bytes(buf4);
195
196 file.read_exact(&mut buf4)?;
197 let gene_count = u32::from_le_bytes(buf4);
198
199 file.read_exact(&mut buf8)?;
200 let index_offset = u64::from_le_bytes(buf8);
201
202 file.seek(SeekFrom::Start(index_offset))?;
204 let mut index = FxHashMap::default();
205
206 for _ in 0..gene_count {
207 let mut buf2 = [0u8; 2];
208 file.read_exact(&mut buf2)?;
209 let name_len = u16::from_le_bytes(buf2) as usize;
210
211 let mut name_buf = vec![0u8; name_len];
212 file.read_exact(&mut name_buf)?;
213 let gene = String::from_utf8(name_buf)?;
214
215 file.read_exact(&mut buf8)?;
216 let offset = u64::from_le_bytes(buf8);
217
218 file.read_exact(&mut buf4)?;
219 let compressed_len = u32::from_le_bytes(buf4);
220
221 file.read_exact(&mut buf4)?;
222 let record_count = u32::from_le_bytes(buf4);
223
224 index.insert(gene, FdbIndexEntry {
225 offset,
226 compressed_len,
227 record_count,
228 });
229 }
230
231 let mut gene_name_to_key = FxHashMap::default();
234 for full_key in index.keys() {
235 let gene_name = full_key.split('|').next().unwrap_or(full_key);
237 if !gene_name_to_key.contains_key(gene_name) {
239 gene_name_to_key.insert(gene_name.to_string(), full_key.clone());
240 }
241 }
242
243 Ok(Self { file, index, gene_name_to_key })
244 }
245
246 pub fn has_gene(&self, gene: &str) -> bool {
249 if self.index.contains_key(gene) {
251 return true;
252 }
253 self.gene_name_to_key.contains_key(gene)
255 }
256
257 pub fn get_gene_records(&mut self, gene: &str) -> Result<Vec<FlankingRecord>> {
262 let lookup_key = if self.index.contains_key(gene) {
264 gene.to_string()
265 } else if let Some(full_key) = self.gene_name_to_key.get(gene) {
266 full_key.clone()
267 } else {
268 anyhow::bail!("Gene not found: {}", gene);
269 };
270
271 let entry = self.index.get(&lookup_key)
272 .ok_or_else(|| anyhow::anyhow!("Gene not found in index: {}", lookup_key))?
273 .clone();
274
275 self.file.seek(SeekFrom::Start(entry.offset))?;
277 let mut compressed = vec![0u8; entry.compressed_len as usize];
278 self.file.read_exact(&mut compressed)?;
279
280 let decompressed = zstd::decode_all(&compressed[..])?;
282 let content = String::from_utf8(decompressed)?;
283
284 let mut records = Vec::with_capacity(entry.record_count as usize);
286 let mut lines = content.lines();
287
288 let _header = lines.next();
290
291 for line in lines {
292 if line.is_empty() {
293 continue;
294 }
295 let fields: Vec<&str> = line.split('\t').collect();
296 if fields.len() < 7 {
298 continue;
299 }
300
301 records.push(FlankingRecord {
302 contig: fields[1].to_string(),
303 genus: fields[2].to_string(),
304 upstream: fields[5].to_string(),
305 downstream: fields[6].to_string(),
306 });
307 }
308
309 Ok(records)
310 }
311
312 pub fn get_genus_distribution(&mut self, gene: &str) -> Result<FxHashMap<String, usize>> {
316 let records = self.get_gene_records(gene)?;
317 let mut dist: FxHashMap<String, usize> = FxHashMap::default();
318
319 for rec in records {
320 *dist.entry(rec.genus).or_default() += 1;
321 }
322
323 Ok(dist)
324 }
325}
326
327pub struct GenusClassifier {
336 db: FlankingDatabase,
337 minimap2_path: String,
338 min_identity: f64,
339 min_align_len: usize,
340 max_flanking: usize,
341 plasmid_contigs: rustc_hash::FxHashSet<String>,
344 species_identity: f64,
346 species_map: FxHashMap<String, String>,
348 plasmid_hi: f64,
351 plasmid_lo: f64,
352}
353
354impl GenusClassifier {
355 pub fn new<P: AsRef<Path>>(
364 db_path: P,
365 minimap2_path: &str,
366 min_identity: f64,
367 min_align_len: usize,
368 max_flanking: usize,
369 plasmid_contigs_path: Option<&Path>,
370 species_identity: f64,
371 species_map_path: Option<&Path>,
372 plasmid_hi: f64,
373 plasmid_lo: f64,
374 ) -> Result<Self> {
375 let db = FlankingDatabase::open(db_path)?;
376 let mut plasmid_contigs = rustc_hash::FxHashSet::default();
377 if let Some(p) = plasmid_contigs_path {
378 let f = File::open(p).with_context(|| format!("open plasmid list {:?}", p))?;
379 for line in BufReader::new(f).lines() {
380 let acc = line?.trim().to_string();
381 if !acc.is_empty() { plasmid_contigs.insert(acc); }
382 }
383 }
384 let mut species_map = FxHashMap::default();
385 if let Some(p) = species_map_path {
386 let f = File::open(p).with_context(|| format!("open species map {:?}", p))?;
387 for line in BufReader::new(f).lines() {
388 let line = line?;
389 if let Some((c, sp)) = line.split_once('\t') {
390 let (c, sp) = (c.trim(), sp.trim());
391 if !c.is_empty() && !sp.is_empty() { species_map.insert(c.to_string(), sp.to_string()); }
392 }
393 }
394 }
395 Ok(Self {
396 db,
397 minimap2_path: minimap2_path.to_string(),
398 min_identity,
399 min_align_len,
400 max_flanking,
401 plasmid_contigs,
402 species_identity,
403 species_map,
404 plasmid_hi,
405 plasmid_lo,
406 })
407 }
408
409 pub fn classify_batch(&mut self, positions: &[ArgPosition], threads: usize) -> Result<Vec<GenusResult>> {
413 let mut results = Vec::with_capacity(positions.len());
414
415 for pos in positions {
416 let result = self.classify_single(pos, threads)?;
417 results.push(result);
418 }
419
420 Ok(results)
421 }
422
423 pub fn classify_single(&mut self, pos: &ArgPosition, threads: usize) -> Result<GenusResult> {
432 let (upstream, downstream) = self.extract_flanking_regions(pos);
434
435 let upstream_len = upstream.len();
436 let downstream_len = downstream.len();
437
438 let snp_status = snp::verify_snp(
440 &pos.contig_seq,
441 &pos.arg_name,
442 0,
443 pos.arg_end - pos.arg_start,
444 pos.arg_start,
445 pos.arg_end,
446 pos.strand,
447 );
448
449 if upstream_len < 50 && downstream_len < 50 {
451 return Ok(GenusResult {
452 arg_name: pos.arg_name.clone(),
453 contig_name: pos.contig_name.clone(),
454 genus: None,
455 confidence: 0.0,
456 specificity: 0.0,
457 upstream_len,
458 downstream_len,
459 top_matches: vec![],
460 snp_status,
461 upstream_seq: upstream.clone(),
462 downstream_seq: downstream.clone(),
463 context: "NA".to_string(),
464 n_genera_tied: 0,
465 species: None,
466 species_top_matches: vec![],
467 n_species_tied: 0,
468 });
469 }
470
471 if !self.db.has_gene(&pos.arg_name) {
473 return Ok(GenusResult {
474 arg_name: pos.arg_name.clone(),
475 contig_name: pos.contig_name.clone(),
476 genus: None,
477 confidence: 0.0,
478 specificity: 0.0,
479 upstream_len,
480 downstream_len,
481 top_matches: vec![("gene_not_in_db".to_string(), 0.0)],
482 snp_status,
483 upstream_seq: upstream.clone(),
484 downstream_seq: downstream.clone(),
485 context: "NA".to_string(),
486 n_genera_tied: 0,
487 species: None,
488 species_top_matches: vec![],
489 n_species_tied: 0,
490 });
491 }
492
493 let ref_records = self.db.get_gene_records(&pos.arg_name)?;
495 if ref_records.is_empty() {
496 return Ok(GenusResult {
497 arg_name: pos.arg_name.clone(),
498 contig_name: pos.contig_name.clone(),
499 genus: None,
500 confidence: 0.0,
501 specificity: 0.0,
502 upstream_len,
503 downstream_len,
504 top_matches: vec![("no_ref_records".to_string(), 0.0)],
505 snp_status,
506 upstream_seq: upstream.clone(),
507 downstream_seq: downstream.clone(),
508 context: "NA".to_string(),
509 n_genera_tied: 0,
510 species: None,
511 species_top_matches: vec![],
512 n_species_tied: 0,
513 });
514 }
515
516 let temp_dir = std::env::temp_dir();
518 let pid = std::process::id();
519 let query_path = temp_dir.join(format!("argenus_query_{}.fas", pid));
520 let ref_path = temp_dir.join(format!("argenus_ref_{}.fas", pid));
521 let paf_path = temp_dir.join(format!("argenus_align_{}.paf", pid));
522
523 {
525 let mut query_file = BufWriter::new(File::create(&query_path)?);
526 if !upstream.is_empty() {
527 writeln!(query_file, ">upstream")?;
528 writeln!(query_file, "{}", upstream)?;
529 }
530 if !downstream.is_empty() {
531 writeln!(query_file, ">downstream")?;
532 writeln!(query_file, "{}", downstream)?;
533 }
534 }
535
536 {
538 let mut ref_file = BufWriter::new(File::create(&ref_path)?);
539 for (i, rec) in ref_records.iter().enumerate() {
540 if !rec.upstream.is_empty() {
541 writeln!(ref_file, ">{}|{}|up_{}", rec.genus, rec.contig, i)?;
542 writeln!(ref_file, "{}", rec.upstream)?;
543 }
544 if !rec.downstream.is_empty() {
545 writeln!(ref_file, ">{}|{}|down_{}", rec.genus, rec.contig, i)?;
546 writeln!(ref_file, "{}", rec.downstream)?;
547 }
548 }
549 }
550
551 let output = Command::new(&self.minimap2_path)
553 .args(["-x", "sr", "-t", &threads.to_string(), "-c", "--secondary=yes", "-N", "100", "-k", "15", "-w", "5"])
554 .arg(&ref_path)
555 .arg(&query_path)
556 .arg("-o").arg(&paf_path)
557 .stderr(std::process::Stdio::null())
558 .output()
559 .context("Failed to run minimap2")?;
560
561 if !output.status.success() {
562 let _ = std::fs::remove_file(&query_path);
564 let _ = std::fs::remove_file(&ref_path);
565 let _ = std::fs::remove_file(&paf_path);
566
567 return Ok(GenusResult {
568 arg_name: pos.arg_name.clone(),
569 contig_name: pos.contig_name.clone(),
570 genus: None,
571 confidence: 0.0,
572 specificity: 0.0,
573 upstream_len,
574 downstream_len,
575 top_matches: vec![("minimap2_failed".to_string(), 0.0)],
576 snp_status,
577 upstream_seq: upstream.clone(),
578 downstream_seq: downstream.clone(),
579 context: "NA".to_string(),
580 n_genera_tied: 0,
581 species: None,
582 species_top_matches: vec![],
583 n_species_tied: 0,
584 });
585 }
586
587 let (genus_scores, plasmid_frac, species_scores) = self.calculate_genus_scores(&paf_path)?;
589
590 let _ = std::fs::remove_file(&query_path);
592 let _ = std::fs::remove_file(&ref_path);
593 let _ = std::fs::remove_file(&paf_path);
594
595 let context = if self.plasmid_contigs.is_empty() || genus_scores.is_empty() {
600 "NA".to_string()
601 } else if plasmid_frac >= self.plasmid_hi {
602 "plasmid".to_string()
603 } else if plasmid_frac <= self.plasmid_lo {
604 "chromosome".to_string()
605 } else {
606 "ambiguous".to_string()
607 };
608
609 let genus_dist = self.db.get_genus_distribution(&pos.arg_name)?;
611 let total_in_db: usize = genus_dist.values().sum();
612
613 let mut sorted_scores: Vec<(String, f64)> = genus_scores.into_iter().collect();
615 sorted_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
616
617 let (genus, confidence, specificity) = if let Some((best_genus, best_score)) = sorted_scores.first() {
618 let genus_count = genus_dist.get(best_genus).copied().unwrap_or(0);
619 let specificity = if total_in_db > 0 {
620 (genus_count as f64 / total_in_db as f64) * 100.0
621 } else {
622 0.0
623 };
624
625 (Some(best_genus.clone()), *best_score, specificity)
626 } else {
627 (None, 0.0, 0.0)
628 };
629
630 let n_genera_tied = match sorted_scores.first() {
633 Some((_, best)) => sorted_scores.iter()
634 .filter(|(g, s)| !g.is_empty() && *s >= best - GENUS_TIE_PCT)
635 .count(),
636 None => 0,
637 };
638
639 let top_matches: Vec<(String, f64)> = sorted_scores.into_iter().take(5).collect();
640
641 let mut sorted_species: Vec<(String, f64)> = species_scores.into_iter().collect();
643 sorted_species.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
644 let species = sorted_species.first().map(|(s, _)| s.clone());
645 let n_species_tied = match sorted_species.first() {
646 Some((_, best)) => sorted_species.iter()
647 .filter(|(s, sc)| !s.is_empty() && *sc >= best - GENUS_TIE_PCT)
648 .count(),
649 None => 0,
650 };
651 let species_top_matches: Vec<(String, f64)> = sorted_species.into_iter().take(5).collect();
652
653 Ok(GenusResult {
654 arg_name: pos.arg_name.clone(),
655 contig_name: pos.contig_name.clone(),
656 genus,
657 confidence,
658 specificity,
659 upstream_len,
660 downstream_len,
661 top_matches,
662 snp_status,
663 upstream_seq: upstream.clone(),
664 downstream_seq: downstream.clone(),
665 context,
666 n_genera_tied,
667 species,
668 species_top_matches,
669 n_species_tied,
670 })
671 }
672
673 fn extract_flanking_regions(&self, pos: &ArgPosition) -> (String, String) {
677 let seq = &pos.contig_seq;
678
679 let upstream_end = pos.arg_start;
681 let upstream_start = upstream_end.saturating_sub(self.max_flanking);
682 let upstream = if upstream_end > upstream_start {
683 seq[upstream_start..upstream_end].to_string()
684 } else {
685 String::new()
686 };
687
688 let downstream_start = pos.arg_end;
690 let downstream_end = (downstream_start + self.max_flanking).min(seq.len());
691 let downstream = if downstream_end > downstream_start {
692 seq[downstream_start..downstream_end].to_string()
693 } else {
694 String::new()
695 };
696
697 if pos.strand == '-' {
699 (reverse_complement(&downstream), reverse_complement(&upstream))
700 } else {
701 (upstream, downstream)
702 }
703 }
704
705 fn calculate_genus_scores(&self, paf_path: &Path) -> Result<(FxHashMap<String, f64>, f64, FxHashMap<String, f64>)> {
707 let file = File::open(paf_path)?;
708 let reader = BufReader::new(file);
709
710 let mut genus_matches: FxHashMap<String, Vec<f64>> = FxHashMap::default();
711 let min_identity_pct = self.min_identity * 100.0;
712 let (mut plasmid_hits, mut total_hits) = (0usize, 0usize);
714 let mut species_matches: FxHashMap<String, Vec<f64>> = FxHashMap::default();
717 let species_pct = self.species_identity * 100.0;
718 let do_species = self.species_identity > 0.0 && !self.species_map.is_empty();
719
720 for line in reader.lines() {
721 let line = line?;
722 let fields: Vec<&str> = line.split('\t').collect();
723 if fields.len() < 12 {
724 continue;
725 }
726
727 let block_len: usize = fields[10].parse().unwrap_or(0);
728 let matches: usize = fields[9].parse().unwrap_or(0);
729
730 if block_len < self.min_align_len {
731 continue;
732 }
733
734 let identity = if block_len > 0 {
735 (matches as f64 / block_len as f64) * 100.0
736 } else {
737 0.0
738 };
739
740 if identity < min_identity_pct {
741 continue;
742 }
743
744 let target_name = fields[5];
746 let mut toks = target_name.split('|');
747 let genus = toks.next().unwrap_or("");
748 let contig = toks.next().unwrap_or("");
749 if genus.is_empty() {
752 continue;
753 }
754 total_hits += 1;
755 if !contig.is_empty() && self.plasmid_contigs.contains(contig) {
756 plasmid_hits += 1;
757 }
758 genus_matches.entry(genus.to_string()).or_default().push(identity);
759
760 if do_species && identity >= species_pct {
762 if let Some(sp) = self.species_map.get(contig) {
763 species_matches.entry(sp.clone()).or_default().push(identity);
764 }
765 }
766 }
767
768 let mut genus_scores: FxHashMap<String, f64> = FxHashMap::default();
770 for (genus, scores) in genus_matches {
771 if scores.is_empty() {
772 continue;
773 }
774 let avg_identity = scores.iter().sum::<f64>() / scores.len() as f64;
776 let count_bonus = (scores.len() as f64).ln().max(1.0);
777 genus_scores.insert(genus, avg_identity * count_bonus / count_bonus.max(1.0));
778 }
779
780 let mut species_scores: FxHashMap<String, f64> = FxHashMap::default();
781 for (sp, scores) in species_matches {
782 if scores.is_empty() { continue; }
783 let avg = scores.iter().sum::<f64>() / scores.len() as f64;
784 species_scores.insert(sp, avg);
785 }
786
787 let plasmid_frac = if total_hits > 0 { plasmid_hits as f64 / total_hits as f64 } else { 0.0 };
788 Ok((genus_scores, plasmid_frac, species_scores))
789 }
790
791}
792
793fn reverse_complement(seq: &str) -> String {
799 seq.chars()
800 .rev()
801 .map(|c| match c.to_ascii_uppercase() {
802 'A' => 'T',
803 'T' => 'A',
804 'G' => 'C',
805 'C' => 'G',
806 _ => 'N',
807 })
808 .collect()
809}
810
811#[cfg(test)]
816mod tests {
817 use super::*;
818
819 #[test]
820 fn test_reverse_complement() {
821 assert_eq!(reverse_complement("ATGC"), "GCAT");
822 assert_eq!(reverse_complement("AAAA"), "TTTT");
823 assert_eq!(reverse_complement(""), "");
824 }
825
826 #[test]
827 fn test_genus_result_default() {
828 let result = GenusResult::default();
829 assert!(result.genus.is_none());
830 assert_eq!(result.confidence, 0.0);
831 }
832}