1use anyhow::{Context, Result};
27use rustc_hash::{FxHashMap, FxHashSet};
28use std::fs::File;
29use std::io::{BufRead, BufReader, BufWriter, Read, Write};
30use std::path::{Path, PathBuf};
31use std::process::Command;
32use std::sync::atomic::{AtomicUsize, AtomicBool, Ordering};
33use std::sync::{Arc, Mutex, mpsc};
34use std::time::Duration;
35use std::thread;
36
37const NCBI_FTP_BASE: &str = "https://ftp.ncbi.nlm.nih.gov/genomes";
38const NCBI_TAXDUMP_URL: &str = "https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/taxdump.tar.gz";
39const NCBI_DATASETS_API: &str = "https://api.ncbi.nlm.nih.gov/datasets/v2";
40const PLSDB_META_URL: &str = "https://ccb-microbe.cs.uni-saarland.de/plsdb2025/download_meta.tar.gz";
41const PLSDB_FASTA_URL: &str = "https://ccb-microbe.cs.uni-saarland.de/plsdb2025/download_fasta";
42const API_BATCH_SIZE: usize = 1000; #[derive(Clone, Debug, Default)]
47pub struct PlsdbOptions {
48 pub dir: Option<PathBuf>,
50 pub skip: bool,
52}
53
54#[derive(Clone, Debug)]
56pub struct FlankBuildConfig {
57 pub flanking_length: usize,
59 pub queue_buffer_gb: u32,
61 pub plsdb: PlsdbOptions,
63}
64
65impl Default for FlankBuildConfig {
66 fn default() -> Self {
67 Self {
68 flanking_length: 1000,
69 queue_buffer_gb: 30,
70 plsdb: PlsdbOptions::default(),
71 }
72 }
73}
74
75#[derive(Clone, Debug)]
77pub struct AssemblyInfo {
78 pub accession: String,
79 pub taxid: String,
80 pub species_taxid: String,
81 pub organism_name: String,
82}
83
84#[derive(Clone, Debug)]
86pub struct PlasmidInfo {
87 pub accession: String,
88 pub taxonomy_uid: String,
89 pub genus: String,
90 pub species: String,
91}
92
93#[derive(Clone, Debug)]
95struct PafHit {
96 query_name: String,
97 query_start: usize,
98 query_end: usize,
99 gene_name: String,
100 gene_length: usize,
101 score: i64,
102 mapq: u8,
103 divergence: f32,
104 gap_count: usize,
105 raw_line: String,
106}
107
108impl PafHit {
109 fn from_paf_line(line: &str) -> Option<Self> {
110 let fields: Vec<&str> = line.split('\t').collect();
111 if fields.len() < 12 {
112 return None;
113 }
114
115 let query_name = fields[0].to_string();
116 let query_start: usize = fields[2].parse().ok()?;
117 let query_end: usize = fields[3].parse().ok()?;
118 let gene_name = fields[5].to_string();
119 let gene_length: usize = fields[6].parse().ok()?;
120 let mapq: u8 = fields[11].parse().unwrap_or(0);
121
122 let matches: usize = fields[9].parse().unwrap_or(0);
123 let block_len: usize = fields[10].parse().unwrap_or(0);
124
125 let mut score: i64 = 0;
126 let mut divergence: f32 = 1.0;
127
128 for field in &fields[12..] {
129 if let Some(val) = field.strip_prefix("AS:i:") {
130 score = val.parse().unwrap_or(0);
131 } else if let Some(val) = field.strip_prefix("de:f:") {
132 divergence = val.parse().unwrap_or(1.0);
133 }
134 }
135
136 let gap_count = block_len.saturating_sub(matches);
137
138 Some(PafHit {
139 query_name,
140 query_start,
141 query_end,
142 gene_name,
143 gene_length,
144 score,
145 mapq,
146 divergence,
147 gap_count,
148 raw_line: line.to_string(),
149 })
150 }
151
152 fn overlaps(&self, other: &PafHit) -> bool {
153 if self.query_name != other.query_name {
154 return false;
155 }
156 let start = self.query_start.max(other.query_start);
157 let end = self.query_end.min(other.query_end);
158 if start >= end {
159 return false;
160 }
161 let overlap = end - start;
162 let self_len = self.query_end - self.query_start;
163 let other_len = other.query_end - other.query_start;
164 overlap * 2 > self_len || overlap * 2 > other_len
165 }
166}
167
168fn compare_paf_hits(a: &PafHit, b: &PafHit) -> std::cmp::Ordering {
170 use std::cmp::Ordering;
171 b.score.cmp(&a.score)
172 .then_with(|| b.gene_length.cmp(&a.gene_length))
173 .then_with(|| b.mapq.cmp(&a.mapq))
174 .then_with(|| a.divergence.partial_cmp(&b.divergence).unwrap_or(Ordering::Equal))
175 .then_with(|| a.gap_count.cmp(&b.gap_count))
176 .then_with(|| a.gene_name.cmp(&b.gene_name))
177}
178
179fn deduplicate_paf_hits(hits: Vec<PafHit>) -> Vec<PafHit> {
181 if hits.is_empty() {
182 return hits;
183 }
184
185 let mut groups: Vec<Vec<PafHit>> = Vec::new();
186 for hit in hits {
187 let mut found = false;
188 for group in &mut groups {
189 if group.iter().any(|h| h.overlaps(&hit)) {
190 group.push(hit.clone());
191 found = true;
192 break;
193 }
194 }
195 if !found {
196 groups.push(vec![hit]);
197 }
198 }
199
200 groups
201 .into_iter()
202 .map(|mut g| {
203 g.sort_by(compare_paf_hits);
204 g.into_iter().next().unwrap()
205 })
206 .collect()
207}
208
209#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
216pub enum GenomeSizeCategory {
217 Small,
219 Medium,
221 Large,
223}
224
225const SMALL_GENOME_THRESHOLD: u64 = 5 * 1024 * 1024; const LARGE_GENOME_THRESHOLD: u64 = 20 * 1024 * 1024; impl GenomeSizeCategory {
230 pub fn from_size(size_bytes: u64) -> Self {
232 if size_bytes < SMALL_GENOME_THRESHOLD {
233 GenomeSizeCategory::Small
234 } else if size_bytes < LARGE_GENOME_THRESHOLD {
235 GenomeSizeCategory::Medium
236 } else {
237 GenomeSizeCategory::Large
238 }
239 }
240
241 pub fn thread_count(&self, max_threads: usize) -> usize {
243 let recommended = match self {
244 GenomeSizeCategory::Small => 4,
245 GenomeSizeCategory::Medium => 6,
246 GenomeSizeCategory::Large => 8,
247 };
248 recommended.min(max_threads)
250 }
251
252 pub fn name(&self) -> &'static str {
254 match self {
255 GenomeSizeCategory::Small => "small (<5MB)",
256 GenomeSizeCategory::Medium => "medium (5-20MB)",
257 GenomeSizeCategory::Large => "large (>20MB)",
258 }
259 }
260}
261
262fn bucket_genomes_by_size(genome_files: &[PathBuf]) -> FxHashMap<GenomeSizeCategory, Vec<PathBuf>> {
265 let mut buckets: FxHashMap<GenomeSizeCategory, Vec<PathBuf>> = FxHashMap::default();
266
267 for path in genome_files {
268 let size = std::fs::metadata(path)
269 .map(|m| m.len())
270 .unwrap_or(0);
271 let category = GenomeSizeCategory::from_size(size);
272 buckets.entry(category).or_default().push(path.clone());
273 }
274
275 buckets
276}
277
278#[derive(Clone, Debug)]
285pub struct CatalogEntry {
286 pub accession: String, pub taxid: String, pub species_taxid: String, pub genus: String, pub species: String, pub organism_name: String, pub source: String, }
294
295pub struct GenomeCatalog {
297 entries: FxHashMap<String, CatalogEntry>,
298}
299
300impl Default for GenomeCatalog {
301 fn default() -> Self {
302 Self::new()
303 }
304}
305
306impl GenomeCatalog {
307 pub fn new() -> Self {
308 Self {
309 entries: FxHashMap::default(),
310 }
311 }
312
313 pub fn add_assembly_with_taxonomy(
315 &mut self,
316 asm: &AssemblyInfo,
317 source: &str,
318 taxonomy: Option<&TaxonomyDB>,
319 ) {
320 let (parsed_genus, species) = parse_organism_name(&asm.organism_name);
321
322 let genus = resolve_genus(&asm.organism_name, &asm.taxid, &parsed_genus, taxonomy);
324
325 let entry = CatalogEntry {
326 accession: asm.accession.clone(),
327 taxid: asm.taxid.clone(),
328 species_taxid: asm.species_taxid.clone(),
329 genus,
330 species,
331 organism_name: asm.organism_name.clone(),
332 source: source.to_string(),
333 };
334
335 self.entries.insert(asm.accession.clone(), entry.clone());
337 if let Some(base) = asm.accession.split('.').next() {
338 self.entries.insert(base.to_string(), entry);
339 }
340 }
341
342 pub fn add_plasmid_with_taxonomy(
344 &mut self,
345 plasmid: &PlasmidInfo,
346 taxonomy: Option<&TaxonomyDB>,
347 ) {
348 let organism_name = format!("{} {}", plasmid.genus, plasmid.species);
350
351 let genus = resolve_genus(&organism_name, &plasmid.taxonomy_uid, &plasmid.genus, taxonomy);
353
354 let entry = CatalogEntry {
355 accession: plasmid.accession.clone(),
356 taxid: plasmid.taxonomy_uid.clone(),
357 species_taxid: String::new(),
358 genus,
359 species: plasmid.species.clone(),
360 organism_name,
361 source: "plsdb".to_string(),
362 };
363
364 self.entries.insert(plasmid.accession.clone(), entry.clone());
366
367 let base = plasmid.accession
369 .strip_prefix("NZ_")
370 .unwrap_or(&plasmid.accession);
371 let no_ver = base.split('.').next().unwrap_or(base);
372 self.entries.insert(no_ver.to_string(), entry.clone());
373
374 let underscore_key = plasmid.accession.replace('.', "_");
377 self.entries.insert(underscore_key, entry);
378 }
379
380 pub fn get_genus(&self, accession: &str) -> Option<&str> {
382 if let Some(entry) = self.entries.get(accession) {
384 return Some(&entry.genus);
385 }
386
387 if let Some(base) = accession.split('.').next() {
389 if let Some(entry) = self.entries.get(base) {
390 return Some(&entry.genus);
391 }
392 }
393
394 let stripped = accession.strip_prefix("NZ_").unwrap_or(accession);
396 if let Some(entry) = self.entries.get(stripped) {
397 return Some(&entry.genus);
398 }
399 let stripped_base = stripped.split('.').next().unwrap_or(stripped);
400 if let Some(entry) = self.entries.get(stripped_base) {
401 return Some(&entry.genus);
402 }
403
404 if let Some(last_underscore) = accession.rfind('_') {
407 let suffix = &accession[last_underscore + 1..];
408 if !suffix.is_empty() && suffix.chars().all(|c| c.is_ascii_digit()) {
410 let with_dot = format!("{}.{}", &accession[..last_underscore], suffix);
411 if let Some(entry) = self.entries.get(&with_dot) {
412 return Some(&entry.genus);
413 }
414 let stripped_with_dot = with_dot.strip_prefix("NZ_").unwrap_or(&with_dot);
416 if let Some(entry) = self.entries.get(stripped_with_dot) {
417 return Some(&entry.genus);
418 }
419 }
420 }
421
422 None
423 }
424
425 pub fn save(&self, path: &Path) -> Result<()> {
427 let file = File::create(path)?;
428 let mut writer = BufWriter::new(file);
429
430 writeln!(writer, "accession\ttaxid\tspecies_taxid\tgenus\tspecies\torganism_name\tsource")?;
431
432 let mut seen: FxHashSet<&str> = FxHashSet::default();
434 for entry in self.entries.values() {
435 if seen.contains(entry.accession.as_str()) {
436 continue;
437 }
438 seen.insert(&entry.accession);
439
440 writeln!(
441 writer,
442 "{}\t{}\t{}\t{}\t{}\t{}\t{}",
443 entry.accession,
444 entry.taxid,
445 entry.species_taxid,
446 entry.genus,
447 entry.species,
448 entry.organism_name.replace('\t', " "),
449 entry.source
450 )?;
451 }
452
453 Ok(())
454 }
455
456 pub fn len(&self) -> usize {
458 let mut seen: FxHashSet<&str> = FxHashSet::default();
459 for entry in self.entries.values() {
460 seen.insert(&entry.accession);
461 }
462 seen.len()
463 }
464
465 #[allow(dead_code)]
467 pub fn is_empty(&self) -> bool {
468 self.entries.is_empty()
469 }
470}
471
472fn parse_organism_name(organism: &str) -> (String, String) {
474 let parts: Vec<&str> = organism.split_whitespace().collect();
475
476 let genus = parts.first()
477 .map(|s| s.to_string())
478 .unwrap_or_else(|| "Unknown".to_string());
479
480 let species = if parts.len() >= 2 {
481 parts[1..].join(" ")
483 } else {
484 "Unknown".to_string()
485 };
486
487 (clean_genus(&genus), species)
488}
489
490fn clean_genus(genus: &str) -> String {
492 let cleaned = if let Some(idx) = genus.find('(') {
495 genus[..idx].trim()
496 } else {
497 genus.trim()
498 };
499
500 cleaned.to_string()
501}
502
503fn resolve_genus(
509 organism_name: &str,
510 taxid_str: &str,
511 _parsed_genus: &str,
512 taxonomy: Option<&TaxonomyDB>,
513) -> String {
514 let organism_lower = organism_name.to_lowercase();
516 if organism_lower.contains("uncultured") {
517 return "uncultured".to_string();
518 }
519
520 if let Some(tax_db) = taxonomy {
522 if let Ok(taxid) = taxid_str.parse::<u32>() {
523 if let Some(genus) = tax_db.get_genus(taxid) {
525 return clean_genus(&genus);
526 }
527
528 if let Some((name, _rank)) = tax_db.get_genus_or_higher(taxid) {
530 return clean_genus(&name);
531 }
532 }
533 }
534
535 let cleaned = clean_genus(organism_name);
538 if cleaned.is_empty() {
539 "unknown".to_string()
540 } else {
541 cleaned
542 }
543}
544
545pub struct TaxonomyDB {
548 names: FxHashMap<u32, String>,
550 nodes: FxHashMap<u32, (u32, String)>,
552}
553
554impl Default for TaxonomyDB {
555 fn default() -> Self {
556 Self::new()
557 }
558}
559
560impl TaxonomyDB {
561 pub fn new() -> Self {
562 Self {
563 names: FxHashMap::default(),
564 nodes: FxHashMap::default(),
565 }
566 }
567
568 pub fn load(taxdump_dir: &Path) -> Result<Self> {
570 let mut db = Self::new();
571
572 let names_path = taxdump_dir.join("names.dmp");
573 let nodes_path = taxdump_dir.join("nodes.dmp");
574
575 if !names_path.exists() || !nodes_path.exists() {
576 anyhow::bail!("Taxonomy files not found in {}", taxdump_dir.display());
577 }
578
579 eprintln!(" Loading names.dmp...");
582 let file = File::open(&names_path)?;
583 let reader = BufReader::new(file);
584
585 for line in reader.lines() {
586 let line = line?;
587 let fields: Vec<&str> = line.split("\t|\t").collect();
588 if fields.len() < 4 {
589 continue;
590 }
591
592 let name_class = fields[3].trim_end_matches("\t|");
594 if name_class != "scientific name" {
595 continue;
596 }
597
598 let taxid: u32 = match fields[0].parse() {
599 Ok(id) => id,
600 Err(_) => continue,
601 };
602 let name = fields[1].to_string();
603
604 db.names.insert(taxid, name);
605 }
606
607 eprintln!(" Loading nodes.dmp...");
610 let file = File::open(&nodes_path)?;
611 let reader = BufReader::new(file);
612
613 for line in reader.lines() {
614 let line = line?;
615 let fields: Vec<&str> = line.split("\t|\t").collect();
616 if fields.len() < 3 {
617 continue;
618 }
619
620 let taxid: u32 = match fields[0].parse() {
621 Ok(id) => id,
622 Err(_) => continue,
623 };
624 let parent_taxid: u32 = match fields[1].parse() {
625 Ok(id) => id,
626 Err(_) => continue,
627 };
628 let rank = fields[2].to_string();
629
630 db.nodes.insert(taxid, (parent_taxid, rank));
631 }
632
633 eprintln!(" Loaded {} taxa, {} nodes", db.names.len(), db.nodes.len());
634 Ok(db)
635 }
636
637 pub fn get_genus(&self, taxid: u32) -> Option<String> {
639 let mut current = taxid;
640 let mut visited = 0;
641
642 while visited < 50 {
644 if let Some((parent, rank)) = self.nodes.get(¤t) {
645 if rank == "genus" {
646 return self.names.get(¤t).cloned();
647 }
648
649 if *parent == current {
651 return None;
652 }
653
654 current = *parent;
655 visited += 1;
656 } else {
657 return None;
658 }
659 }
660
661 None
662 }
663
664 pub fn get_genus_or_higher(&self, taxid: u32) -> Option<(String, String)> {
667 let mut current = taxid;
668 let mut visited = 0;
669
670 let target_ranks = ["genus", "family", "order", "class", "phylum"];
672 let mut best_match: Option<(String, String, usize)> = None; while visited < 50 {
676 if let Some((parent, rank)) = self.nodes.get(¤t) {
677 if let Some(priority) = target_ranks.iter().position(|r| r == rank) {
679 if let Some(name) = self.names.get(¤t) {
680 if rank == "genus" {
682 return Some((name.clone(), rank.clone()));
683 }
684 if best_match.is_none() || priority < best_match.as_ref().unwrap().2 {
686 best_match = Some((name.clone(), rank.clone(), priority));
687 }
688 }
689 }
690
691 if *parent == current {
693 break;
694 }
695
696 current = *parent;
697 visited += 1;
698 } else {
699 break;
700 }
701 }
702
703 best_match.map(|(name, rank, _)| (name, rank))
704 }
705}
706
707#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
709pub struct BuildState {
710 pub started_at: String,
711 pub last_updated: String,
712 pub current_step: String,
713 pub completed_steps: Vec<String>,
714 pub genomes_downloaded: usize,
715 pub plsdb_extracted: usize,
716 pub alignments_done: bool,
717 pub flanking_extracted: bool,
718}
719
720impl Default for BuildState {
721 fn default() -> Self {
722 Self::new()
723 }
724}
725
726impl BuildState {
727 pub fn new() -> Self {
728 let now = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
729 Self {
730 started_at: now.clone(),
731 last_updated: now,
732 current_step: String::new(),
733 completed_steps: Vec::new(),
734 genomes_downloaded: 0,
735 plsdb_extracted: 0,
736 alignments_done: false,
737 flanking_extracted: false,
738 }
739 }
740
741 pub fn load(path: &Path) -> Option<Self> {
742 std::fs::read_to_string(path)
743 .ok()
744 .and_then(|s| serde_json::from_str(&s).ok())
745 }
746
747 pub fn save(&self, path: &Path) -> Result<()> {
748 let json = serde_json::to_string_pretty(self)?;
749 std::fs::write(path, json)?;
750 Ok(())
751 }
752
753 pub fn update_step(&mut self, step: &str) {
754 self.current_step = step.to_string();
755 self.last_updated = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
756 }
757
758 pub fn complete_step(&mut self, step: &str) {
759 if !self.completed_steps.contains(&step.to_string()) {
760 self.completed_steps.push(step.to_string());
761 }
762 self.last_updated = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
763 }
764
765 pub fn is_completed(&self, step: &str) -> bool {
766 self.completed_steps.contains(&step.to_string())
767 }
768}
769
770pub struct FlankingDbBuilder {
772 output_dir: PathBuf,
773 threads: usize,
774 amr_db_path: PathBuf,
775 email: String,
777 config: FlankBuildConfig,
779}
780
781impl FlankingDbBuilder {
782 pub fn new(amr_db: &Path, output_dir: &Path, threads: usize, email: &str, config: FlankBuildConfig) -> Self {
783 Self {
784 output_dir: output_dir.to_path_buf(),
785 threads,
786 amr_db_path: amr_db.to_path_buf(),
787 email: email.to_string(),
788 config,
789 }
790 }
791
792 pub fn build(&self) -> Result<()> {
794 eprintln!("\n============================================================");
795 eprintln!(" ARGenus Flanking Database Builder (Rust)");
796 eprintln!("============================================================");
797 eprintln!("Output directory: {}", self.output_dir.display());
798 eprintln!("NCBI Email: {}", self.email);
799 eprintln!("Download method: NCBI Datasets API (batch)");
800 eprintln!("Threads: {}", self.threads);
801 eprintln!("Flanking length: {} bp", self.config.flanking_length);
802 eprintln!("Queue buffer: {} GB", self.config.queue_buffer_gb);
803 eprintln!("Alignment: minimap2 (AMR indexed)");
804 eprintln!();
805
806 std::fs::create_dir_all(&self.output_dir)?;
807 let genomes_dir = self.output_dir.join("genomes");
808 std::fs::create_dir_all(&genomes_dir)?;
809 let plsdb_dir = self.config.plsdb.dir.clone().unwrap_or_else(|| self.output_dir.join("plsdb"));
811 if !self.config.plsdb.skip {
812 std::fs::create_dir_all(&plsdb_dir)?;
813 }
814 let taxonomy_dir = self.output_dir.join("taxonomy");
815 let temp_dir = self.output_dir.join("temp");
816 std::fs::create_dir_all(&temp_dir)?;
817 let state_path = self.output_dir.join("build_state.json");
818
819 let mut state = BuildState::load(&state_path).unwrap_or_else(|| {
821 eprintln!("Starting fresh build...");
822 BuildState::new()
823 });
824
825 if !state.completed_steps.is_empty() {
826 eprintln!("Resuming from previous state:");
827 eprintln!(" Started: {}", state.started_at);
828 eprintln!(" Last updated: {}", state.last_updated);
829 eprintln!(" Completed steps: {}", state.completed_steps.join(", "));
830 eprintln!(" Genomes downloaded: {}", state.genomes_downloaded);
831 eprintln!(" PLSDB extracted: {}", state.plsdb_extracted);
832 eprintln!();
833 }
834
835 let mut catalog = GenomeCatalog::new();
837 let catalog_path = self.output_dir.join("genome_catalog.tsv");
838
839 if !state.is_completed("taxonomy_download") {
841 state.update_step("taxonomy_download");
842 state.save(&state_path)?;
843 eprintln!("[1/9] Downloading NCBI taxonomy database...");
844 self.download_taxdump(&taxonomy_dir)?;
845 state.complete_step("taxonomy_download");
846 state.save(&state_path)?;
847 } else {
848 eprintln!("[1/9] Taxonomy database already downloaded, skipping...");
849 }
850
851 eprintln!("[2/9] Loading taxonomy database...");
853 let taxonomy = match TaxonomyDB::load(&taxonomy_dir) {
854 Ok(db) => {
855 eprintln!(" Taxonomy database loaded successfully");
856 Some(db)
857 }
858 Err(e) => {
859 eprintln!(" Warning: Failed to load taxonomy: {}", e);
860 eprintln!(" Genus will be parsed from organism names instead");
861 None
862 }
863 };
864
865 eprintln!("[3/9] Downloading NCBI assembly summaries (GenBank)...");
867 let assemblies = self.download_assembly_summaries()?;
868 eprintln!(" GenBank assemblies: {} (Complete Genome + Chromosome)", assemblies.len());
869
870 for asm in &assemblies {
872 catalog.add_assembly_with_taxonomy(asm, "genbank", taxonomy.as_ref());
873 }
874
875 let standalone_plasmids = if self.config.plsdb.skip {
877 eprintln!("[4/9] PLSDB skipped (--skip-plsdb)");
878 eprintln!("[5/9] PLSDB metadata skipped");
879 Vec::new()
880 } else if self.config.plsdb.dir.is_some() {
881 eprintln!("[4/9] Using pre-downloaded PLSDB: {}", plsdb_dir.display());
883 let nuccore_csv = plsdb_dir.join("nuccore.csv");
885 let fasta_path = plsdb_dir.join("sequences.fasta");
886 if !nuccore_csv.exists() {
887 let meta_tar = plsdb_dir.join("meta.tar.gz");
889 if meta_tar.exists() {
890 eprintln!(" Extracting meta.tar.gz...");
891 std::process::Command::new("tar")
892 .args(["-xzf", meta_tar.to_str().unwrap(), "-C", plsdb_dir.to_str().unwrap()])
893 .status()
894 .with_context(|| "Failed to extract PLSDB metadata")?;
895 } else {
896 anyhow::bail!("PLSDB directory missing nuccore.csv and meta.tar.gz: {}", plsdb_dir.display());
897 }
898 }
899 if !fasta_path.exists() {
900 anyhow::bail!("PLSDB directory missing sequences.fasta: {}", plsdb_dir.display());
901 }
902 state.complete_step("plsdb_download");
903 state.save(&state_path)?;
904
905 eprintln!("[5/9] Loading PLSDB metadata...");
907 let plasmids = self.load_plsdb_plasmids(&plsdb_dir)?;
908 eprintln!(" Standalone circular+complete plasmids: {} (not in any assembly)",
909 plasmids.len());
910 plasmids
911 } else {
912 if !state.is_completed("plsdb_download") {
914 state.update_step("plsdb_download");
915 state.save(&state_path)?;
916 eprintln!("[4/9] Downloading PLSDB database...");
917 self.download_plsdb(&plsdb_dir)?;
918 state.complete_step("plsdb_download");
919 state.save(&state_path)?;
920 } else {
921 eprintln!("[4/9] PLSDB database already downloaded, skipping...");
922 }
923
924 eprintln!("[5/9] Loading PLSDB metadata...");
926 let plasmids = self.load_plsdb_plasmids(&plsdb_dir)?;
927 eprintln!(" Standalone circular+complete plasmids: {} (not in any assembly)",
928 plasmids.len());
929 plasmids
930 };
931
932 for plasmid in &standalone_plasmids {
934 catalog.add_plasmid_with_taxonomy(plasmid, taxonomy.as_ref());
935 }
936
937 eprintln!("[6/9] Saving unified genome catalog...");
939 catalog.save(&catalog_path)?;
940 eprintln!(" Catalog entries: {}", catalog.len());
941 eprintln!(" Saved to: {}", catalog_path.display());
942
943 let paf_output = self.output_dir.join("all_alignments.paf");
946 let merged_hits = self.output_dir.join("merged_alignment_hits.tsv");
947 let output_tsv = self.output_dir.join("all_flanking_sequences.tsv");
948
949 if !state.is_completed("alignment_flanking") {
950 state.update_step("genome_download");
952 state.save(&state_path)?;
953
954 eprintln!("[7/9] Batch processing: Download + Align NCBI genomes...");
955 eprintln!(" Producer-consumer pattern with {} GB queue buffer", self.config.queue_buffer_gb);
956
957 let downloaded = self.download_and_align_batches(
958 &assemblies,
959 &genomes_dir,
960 &temp_dir,
961 &paf_output,
962 &catalog,
963 )?;
964 state.genomes_downloaded = downloaded;
965 state.complete_step("genome_download");
966 state.save(&state_path)?;
967 eprintln!(" Processed {} genomes", downloaded);
968
969 if self.config.plsdb.skip {
971 eprintln!("[8/9] PLSDB processing skipped (--skip-plsdb)");
972 state.plsdb_extracted = 0;
973 } else {
974 state.update_step("plsdb_extract");
975 state.save(&state_path)?;
976 eprintln!("[8/9] Processing PLSDB sequences...");
977 let plsdb_count = self.align_plsdb(
978 &standalone_plasmids,
979 &plsdb_dir,
980 &temp_dir,
981 &paf_output,
982 )?;
983 state.plsdb_extracted = plsdb_count;
984 state.complete_step("plsdb_extract");
985 state.save(&state_path)?;
986 }
987
988 state.update_step("alignment_flanking");
990 state.save(&state_path)?;
991
992 eprintln!("[9/9] Processing alignment results...");
994 self.convert_paf_to_merged(&paf_output, &merged_hits)?;
995 state.alignments_done = true;
996 state.save(&state_path)?;
997
998 eprintln!(" Extracting flanking sequences (flanking_length: {} bp)...",
1000 self.config.flanking_length);
1001 self.extract_flanking_sequences(
1002 &merged_hits,
1003 &genomes_dir,
1004 &output_tsv,
1005 &catalog,
1006 )?;
1007 if !self.config.plsdb.skip {
1009 self.extract_flanking_from_plsdb(
1010 &merged_hits,
1011 &plsdb_dir,
1012 &output_tsv,
1013 &catalog,
1014 )?;
1015 }
1016 state.flanking_extracted = true;
1017 state.complete_step("alignment_flanking");
1018 state.save(&state_path)?;
1019
1020 eprintln!(" Cleaning up PAF file...");
1022 std::fs::remove_file(&paf_output).ok();
1023 } else {
1024 eprintln!("[7-9] Batch pipeline already completed, skipping...");
1025 }
1026
1027 if temp_dir.exists() {
1029 eprintln!("\nCleaning up temp files...");
1030 std::fs::remove_dir_all(&temp_dir).ok();
1031 }
1032
1033 let fdb_path = self.output_dir.join("flanking.fdb");
1035 if !state.is_completed("fdb_build") && output_tsv.exists() {
1036 state.update_step("fdb_build");
1037 state.save(&state_path)?;
1038 eprintln!("\n[10/10] Building FDB from TSV (external sort + zstd)...");
1039
1040 let buffer_size_mb = 1024;
1042 crate::fdb::build(&output_tsv, &fdb_path, buffer_size_mb, self.threads)?;
1043
1044 state.complete_step("fdb_build");
1045 state.save(&state_path)?;
1046 } else if fdb_path.exists() {
1047 eprintln!("[10/10] FDB already built, skipping...");
1048 }
1049
1050 eprintln!("\n============================================================");
1051 eprintln!(" Build Complete");
1052 eprintln!("============================================================");
1053 eprintln!("Output files:");
1054 eprintln!(" - {}", output_tsv.display());
1055 if fdb_path.exists() {
1056 let fdb_size = std::fs::metadata(&fdb_path).map(|m| m.len()).unwrap_or(0);
1057 eprintln!(" - {} ({:.1} MB)", fdb_path.display(), fdb_size as f64 / 1024.0 / 1024.0);
1058 }
1059 eprintln!(" - {}", catalog_path.display());
1060
1061 Ok(())
1062 }
1063
1064 fn download_assembly_summaries(&self) -> Result<Vec<AssemblyInfo>> {
1067 let mut assemblies = Vec::new();
1068
1069 for kingdom in ["bacteria", "archaea"] {
1071 let url = format!(
1072 "{}/genbank/{}/assembly_summary.txt",
1073 NCBI_FTP_BASE, kingdom
1074 );
1075 eprintln!(" Downloading genbank/{}...", kingdom);
1076
1077 match self.download_and_parse_assembly_summary(&url, None) {
1078 Ok((mut asm, _)) => {
1079 eprintln!(" Found {} assemblies (Complete Genome + Chromosome)", asm.len());
1080 assemblies.append(&mut asm);
1081 }
1082 Err(e) => {
1083 eprintln!(" Warning: Failed to download {}: {}", url, e);
1084 }
1085 }
1086 }
1087
1088 eprintln!(" Total assemblies: {}", assemblies.len());
1089 Ok(assemblies)
1090 }
1091
1092 fn download_plsdb(&self, plsdb_dir: &Path) -> Result<()> {
1094 let meta_tar = plsdb_dir.join("meta.tar.gz");
1095 let fasta_path = plsdb_dir.join("sequences.fasta");
1096
1097 if !plsdb_dir.join("nuccore.csv").exists() {
1099 eprintln!(" Downloading PLSDB metadata...");
1100 self.download_file(PLSDB_META_URL, &meta_tar)?;
1101
1102 eprintln!(" Extracting metadata...");
1104 let status = Command::new("tar")
1105 .args(["-xzf", meta_tar.to_str().unwrap(), "-C", plsdb_dir.to_str().unwrap()])
1106 .status()
1107 .with_context(|| "Failed to extract PLSDB metadata")?;
1108
1109 if !status.success() {
1110 anyhow::bail!("tar extraction failed");
1111 }
1112
1113 std::fs::remove_file(&meta_tar).ok();
1115 } else {
1116 eprintln!(" PLSDB metadata already exists, skipping download...");
1117 }
1118
1119 if !fasta_path.exists() {
1121 eprintln!(" Downloading PLSDB sequences (~7GB)...");
1122 self.download_file(PLSDB_FASTA_URL, &fasta_path)?;
1123 } else {
1124 eprintln!(" PLSDB sequences already exist, skipping download...");
1125 }
1126
1127 Ok(())
1128 }
1129
1130 fn download_taxdump(&self, taxonomy_dir: &Path) -> Result<()> {
1132 let tar_path = taxonomy_dir.join("taxdump.tar.gz");
1133 let names_path = taxonomy_dir.join("names.dmp");
1134
1135 if names_path.exists() {
1137 eprintln!(" Taxdump already exists, skipping download...");
1138 return Ok(());
1139 }
1140
1141 std::fs::create_dir_all(taxonomy_dir)?;
1142
1143 eprintln!(" Downloading NCBI taxdump (~60MB)...");
1145 self.download_file(NCBI_TAXDUMP_URL, &tar_path)?;
1146
1147 eprintln!(" Extracting taxdump...");
1149 let status = Command::new("tar")
1150 .args(["-xzf", tar_path.to_str().unwrap(), "-C", taxonomy_dir.to_str().unwrap()])
1151 .status()
1152 .with_context(|| "Failed to extract taxdump")?;
1153
1154 if !status.success() {
1155 anyhow::bail!("tar extraction failed");
1156 }
1157
1158 std::fs::remove_file(&tar_path).ok();
1160
1161 Ok(())
1162 }
1163
1164 fn download_file(&self, url: &str, output_path: &Path) -> Result<()> {
1166 for attempt in 0..3 {
1167 match self.download_file_once(url, output_path) {
1168 Ok(_) => return Ok(()),
1169 Err(e) if attempt < 2 => {
1170 eprintln!(" Download failed (attempt {}): {}", attempt + 1, e);
1171 eprintln!(" Retrying in 5 seconds...");
1172 std::thread::sleep(Duration::from_secs(5));
1173 continue;
1174 }
1175 Err(e) => return Err(e),
1176 }
1177 }
1178 Ok(())
1179 }
1180
1181 fn download_file_once(&self, url: &str, output_path: &Path) -> Result<()> {
1183 let response = ureq::get(url)
1184 .set("User-Agent", "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36")
1185 .timeout(Duration::from_secs(7200)) .call()
1187 .with_context(|| format!("Failed to download {}", url))?;
1188
1189 let mut file = File::create(output_path)?;
1190 let mut reader = response.into_reader();
1191
1192 let mut buffer = [0u8; 65536];
1194 let mut total = 0usize;
1195 loop {
1196 match reader.read(&mut buffer) {
1197 Ok(0) => break,
1198 Ok(n) => {
1199 file.write_all(&buffer[..n])?;
1200 total += n;
1201 if total.is_multiple_of(100 * 1024 * 1024) {
1202 eprintln!(" Downloaded {} MB...", total / (1024 * 1024));
1203 }
1204 }
1205 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
1206 Err(e) => return Err(e.into()),
1207 }
1208 }
1209 eprintln!(" Total: {} MB", total / (1024 * 1024));
1210
1211 Ok(())
1212 }
1213
1214 fn download_and_parse_assembly_summary(
1218 &self,
1219 url: &str,
1220 exclude_set: Option<&FxHashSet<String>>,
1221 ) -> Result<(Vec<AssemblyInfo>, Vec<String>)> {
1222 let response = ureq::get(url)
1223 .timeout(Duration::from_secs(600))
1224 .call()
1225 .with_context(|| format!("Failed to download {}", url))?;
1226
1227 let reader = BufReader::new(response.into_reader());
1228 let mut assemblies = Vec::new();
1229 let mut paired_gca = Vec::new();
1230 let mut skipped_duplicates = 0usize;
1231
1232 for line in reader.lines() {
1233 let line = line?;
1234 if line.starts_with('#') {
1235 continue;
1236 }
1237
1238 let fields: Vec<&str> = line.split('\t').collect();
1239 if fields.len() < 20 {
1240 continue;
1241 }
1242
1243 let accession = fields[0];
1244 let assembly_level = fields[11];
1245
1246 if assembly_level != "Complete Genome" && assembly_level != "Chromosome" {
1248 continue;
1249 }
1250
1251 let ftp_path = fields[19];
1252 if ftp_path == "na" || ftp_path.is_empty() {
1253 continue;
1254 }
1255
1256 if let Some(exclude) = exclude_set {
1258 let acc_base = accession.split('.').next().unwrap_or(accession);
1260 if exclude.contains(accession) || exclude.contains(acc_base) {
1261 skipped_duplicates += 1;
1262 continue;
1263 }
1264 }
1265
1266 if fields.len() > 18 && accession.starts_with("GCF_") {
1270 let paired_asm = fields[17];
1271 let paired_comp = fields[18];
1272 if paired_asm != "na" && !paired_asm.is_empty() && paired_comp == "identical" {
1273 paired_gca.push(paired_asm.to_string());
1274 if let Some(base) = paired_asm.split('.').next() {
1276 paired_gca.push(base.to_string());
1277 }
1278 }
1279 }
1280
1281 assemblies.push(AssemblyInfo {
1282 accession: accession.to_string(),
1283 taxid: fields[5].to_string(),
1284 species_taxid: fields[6].to_string(),
1285 organism_name: fields[7].to_string(),
1286 });
1287 }
1288
1289 if skipped_duplicates > 0 {
1290 eprintln!(" Skipped {} duplicates (have identical RefSeq)", skipped_duplicates);
1291 }
1292
1293 Ok((assemblies, paired_gca))
1294 }
1295
1296 fn load_plsdb_plasmids(&self, plsdb_dir: &Path) -> Result<Vec<PlasmidInfo>> {
1298 let nuccore_path = plsdb_dir.join("nuccore.csv");
1299 let taxonomy_path = plsdb_dir.join("taxonomy.csv");
1300
1301 let mut taxonomy: FxHashMap<String, (String, String)> = FxHashMap::default();
1303 if taxonomy_path.exists() {
1304 let tax_file = File::open(&taxonomy_path)?;
1305 let tax_reader = BufReader::new(tax_file);
1306
1307 for (idx, line) in tax_reader.lines().enumerate() {
1308 let line = line?;
1309 if idx == 0 {
1310 continue; }
1312
1313 let fields: Vec<&str> = line.split(',').collect();
1314 if fields.len() >= 10 {
1315 let uid = fields[0].to_string();
1316 let genus = fields[8].to_string();
1317 let species = fields[9].to_string();
1318 taxonomy.insert(uid, (genus, species));
1319 }
1320 }
1321 }
1322
1323 let mut plasmids = Vec::new();
1325 let nuc_file = File::open(&nuccore_path)?;
1326 let nuc_reader = BufReader::new(nuc_file);
1327
1328 for (idx, line) in nuc_reader.lines().enumerate() {
1329 let line = line?;
1330 if idx == 0 {
1331 continue; }
1333
1334 let fields = parse_csv_line(&line);
1336 if fields.len() < 15 {
1337 continue;
1338 }
1339
1340 let completeness = &fields[4];
1341 let topology = &fields[14];
1342 let assembly_uid = &fields[9]; if completeness != "complete" {
1346 continue;
1347 }
1348 if topology != "circular" {
1349 continue;
1350 }
1351
1352 if assembly_uid != "-1" {
1355 continue;
1356 }
1357
1358 let taxonomy_uid = &fields[12];
1359 let (genus, species) = taxonomy.get(taxonomy_uid)
1360 .cloned()
1361 .unwrap_or_else(|| ("Unknown".to_string(), "Unknown".to_string()));
1362
1363 plasmids.push(PlasmidInfo {
1364 accession: fields[1].clone(),
1365 taxonomy_uid: taxonomy_uid.clone(),
1366 genus,
1367 species,
1368 });
1369 }
1370
1371 Ok(plasmids)
1372 }
1373
1374 fn download_and_align_batches(
1378 &self,
1379 assemblies: &[AssemblyInfo],
1380 genomes_dir: &Path,
1381 temp_dir: &Path,
1382 paf_output: &Path,
1383 _catalog: &GenomeCatalog,
1384 ) -> Result<usize> {
1385 let all_accessions_path = temp_dir.join("accessions_all.txt");
1386 let done_accessions_path = temp_dir.join("accessions_done.txt");
1387
1388 let accessions: Vec<String> = assemblies.iter()
1390 .map(|a| a.accession.clone())
1391 .collect();
1392
1393 if !all_accessions_path.exists() {
1394 eprintln!(" Saving {} accessions to temp file...", accessions.len());
1395 let mut file = File::create(&all_accessions_path)?;
1396 for acc in &accessions {
1397 writeln!(file, "{}", acc)?;
1398 }
1399 }
1400
1401 let mut done_set: FxHashSet<String> = FxHashSet::default();
1403 if done_accessions_path.exists() {
1404 let file = File::open(&done_accessions_path)?;
1405 let reader = BufReader::new(file);
1406 for acc in reader.lines().map_while(Result::ok) {
1407 done_set.insert(acc.trim().to_string());
1408 }
1409
1410 if !done_set.is_empty() {
1412 if !paf_output.exists() {
1413 eprintln!(" WARNING: PAF file missing but {} genomes marked as done", done_set.len());
1414 eprintln!(" Clearing done list and reprocessing all genomes...");
1415 done_set.clear();
1416 std::fs::remove_file(&done_accessions_path).ok();
1417 } else {
1418 let paf_size = std::fs::metadata(paf_output)?.len();
1420 if paf_size == 0 {
1421 eprintln!(" WARNING: PAF file is empty but {} genomes marked as done", done_set.len());
1422 eprintln!(" Clearing done list and reprocessing all genomes...");
1423 done_set.clear();
1424 std::fs::remove_file(&done_accessions_path).ok();
1425 } else {
1426 let mut missing_genomes = Vec::new();
1428 for acc in &done_set {
1429 let genome_path = genomes_dir.join(format!("{}.fna", acc));
1430 if !genome_path.exists() {
1431 missing_genomes.push(acc.clone());
1432 }
1433 }
1434 if !missing_genomes.is_empty() {
1435 eprintln!(" WARNING: {} genome files missing for done accessions", missing_genomes.len());
1436 for acc in &missing_genomes {
1437 done_set.remove(acc);
1438 }
1439 eprintln!(" Removed missing genomes from done list, will re-download");
1440 }
1441
1442 if !done_set.is_empty() {
1443 eprintln!(" Resume: {} already processed (PAF: {} MB)",
1444 done_set.len(), paf_size / 1024 / 1024);
1445 }
1446 }
1447 }
1448 }
1449 }
1450
1451 let remaining: Vec<&String> = accessions.iter()
1453 .filter(|a| !done_set.contains(*a))
1454 .collect();
1455
1456 if remaining.is_empty() {
1457 eprintln!(" All {} genomes already processed", accessions.len());
1458 return Ok(accessions.len());
1459 }
1460
1461 eprintln!(" Remaining: {}/{} genomes to process", remaining.len(), accessions.len());
1462
1463 let bytes_per_batch = API_BATCH_SIZE * 3 * 1024 * 1024; let queue_capacity = std::cmp::max(
1467 1,
1468 (self.config.queue_buffer_gb as usize * 1024 * 1024 * 1024) / bytes_per_batch
1469 );
1470 eprintln!(" Queue capacity: {} batches (based on {} GB buffer)", queue_capacity, self.config.queue_buffer_gb);
1471
1472 let batches: Vec<Vec<String>> = remaining.chunks(API_BATCH_SIZE)
1474 .map(|chunk| chunk.iter().map(|s| (*s).clone()).collect())
1475 .collect();
1476 let total_batches = batches.len();
1477
1478 let (tx, rx) = mpsc::sync_channel::<(usize, Vec<PathBuf>)>(queue_capacity);
1480
1481 let downloaded = Arc::new(AtomicUsize::new(done_set.len()));
1483 let aligned = Arc::new(AtomicUsize::new(done_set.len()));
1484 let download_error = Arc::new(AtomicBool::new(false));
1485
1486 let genomes_dir = genomes_dir.to_path_buf();
1488 let temp_dir_producer = temp_dir.to_path_buf();
1489 let temp_dir_consumer = temp_dir.to_path_buf();
1490 let email = self.email.clone();
1491 let downloaded_clone = Arc::clone(&downloaded);
1492 let download_error_clone = Arc::clone(&download_error);
1493 let total_accessions = accessions.len();
1494
1495 let producer = thread::spawn(move || -> Result<()> {
1497 for (batch_idx, batch) in batches.into_iter().enumerate() {
1498 if download_error_clone.load(Ordering::Relaxed) {
1499 break;
1500 }
1501
1502 let zip_path = temp_dir_producer.join(format!("batch_{:04}.zip", batch_idx));
1503
1504 let url = format!("{}/genome/download", NCBI_DATASETS_API);
1506 let acc_list: Vec<&str> = batch.iter().map(|s| s.as_str()).collect();
1507 let request_body = serde_json::json!({
1508 "accessions": acc_list,
1509 "include_annotation_type": ["GENOME_FASTA"]
1510 });
1511
1512 let response = ureq::post(&url)
1513 .set("Content-Type", "application/json")
1514 .set("ncbi-client-id", &email)
1515 .timeout(Duration::from_secs(3600))
1516 .send_json(&request_body);
1517
1518 match response {
1519 Ok(resp) => {
1520 let mut zip_file = File::create(&zip_path)?;
1522 let mut reader = resp.into_reader();
1523 std::io::copy(&mut reader, &mut zip_file)?;
1524 drop(zip_file);
1525
1526 let genome_files = Self::extract_batch_to_files_static(&zip_path, &genomes_dir)?;
1528
1529 std::fs::remove_file(&zip_path).ok();
1531
1532 downloaded_clone.fetch_add(batch.len(), Ordering::Relaxed);
1533 eprintln!(" [Download] Batch {}/{}: {} genomes ({}/{})",
1534 batch_idx + 1, total_batches, batch.len(),
1535 downloaded_clone.load(Ordering::Relaxed), total_accessions);
1536
1537 if tx.send((batch_idx, genome_files)).is_err() {
1539 break; }
1541 }
1542 Err(e) => {
1543 eprintln!(" [Download] Batch {} failed: {}", batch_idx + 1, e);
1544 download_error_clone.store(true, Ordering::Relaxed);
1545 break;
1546 }
1547 }
1548
1549 std::thread::sleep(Duration::from_millis(200));
1551 }
1552 Ok(())
1553 });
1554
1555 let paf_file = Arc::new(Mutex::new(
1557 std::fs::OpenOptions::new()
1558 .create(true)
1559 .append(true)
1560 .open(paf_output)?
1561 ));
1562 let done_file = Arc::new(Mutex::new(
1563 std::fs::OpenOptions::new()
1564 .create(true)
1565 .append(true)
1566 .open(&done_accessions_path)?
1567 ));
1568
1569 let amr_db = self.amr_db_path.clone();
1570 let max_threads = self.threads;
1571
1572 while let Ok((batch_idx, genome_files)) = rx.recv() {
1573 if genome_files.is_empty() {
1574 continue;
1575 }
1576
1577 let buckets = bucket_genomes_by_size(&genome_files);
1579 let mut batch_hits: Vec<PafHit> = Vec::new();
1580
1581 for (category, bucket_files) in &buckets {
1583 if bucket_files.is_empty() {
1584 continue;
1585 }
1586
1587 let bucket_threads = category.thread_count(max_threads);
1588 let bucket_suffix = format!("batch_{:04}_{:?}", batch_idx, category);
1589 let bucket_fasta = temp_dir_consumer.join(format!("{}.fas", bucket_suffix));
1590
1591 {
1593 let mut fasta_writer = BufWriter::new(File::create(&bucket_fasta)?);
1594 for genome_path in bucket_files {
1595 let filename = genome_path.file_name()
1596 .and_then(|n| n.to_str())
1597 .unwrap_or("unknown.fna");
1598
1599 let file = File::open(genome_path)?;
1600 let reader = BufReader::new(file);
1601 for line in reader.lines() {
1602 let line = line?;
1603 if let Some(stripped) = line.strip_prefix('>') {
1604 let contig_id = stripped.split_whitespace().next().unwrap_or(stripped);
1605 writeln!(fasta_writer, ">{}|{}", contig_id, filename)?;
1606 } else {
1607 writeln!(fasta_writer, "{}", line)?;
1608 }
1609 }
1610 }
1611 fasta_writer.flush()?;
1612 }
1613
1614 let bucket_paf = temp_dir_consumer.join(format!("{}.paf", bucket_suffix));
1616 let status = Command::new("minimap2")
1617 .args([
1618 "-cx", "asm20",
1619 "-t", &bucket_threads.to_string(),
1620 amr_db.to_str().unwrap(),
1621 bucket_fasta.to_str().unwrap(),
1622 "-o", bucket_paf.to_str().unwrap(),
1623 ])
1624 .stdout(std::process::Stdio::null())
1625 .stderr(std::process::Stdio::null())
1626 .status();
1627
1628 if let Ok(s) = status {
1629 if s.success() && bucket_paf.exists() {
1630 let paf_content = std::fs::read_to_string(&bucket_paf)?;
1632 let hits: Vec<PafHit> = paf_content
1633 .lines()
1634 .filter_map(PafHit::from_paf_line)
1635 .collect();
1636 batch_hits.extend(hits);
1637 }
1638 }
1639
1640 std::fs::remove_file(&bucket_fasta).ok();
1642 std::fs::remove_file(&bucket_paf).ok();
1643 }
1644
1645 let dedup_hits = deduplicate_paf_hits(batch_hits);
1647
1648 {
1650 let mut paf = paf_file.lock().unwrap();
1651 for hit in &dedup_hits {
1652 writeln!(paf, "{}", hit.raw_line)?;
1653 }
1654 paf.flush()?; }
1656
1657 {
1659 let mut done = done_file.lock().unwrap();
1660 for genome_path in &genome_files {
1661 if let Some(stem) = genome_path.file_stem().and_then(|s| s.to_str()) {
1662 writeln!(done, "{}", stem)?;
1663 }
1664 }
1665 done.flush()?; }
1667
1668 let count = genome_files.len();
1669 aligned.fetch_add(count, Ordering::Relaxed);
1670
1671 let bucket_info: Vec<String> = buckets.iter()
1673 .map(|(cat, files)| format!("{}={}", cat.name().split_whitespace().next().unwrap_or("?"), files.len()))
1674 .collect();
1675 eprintln!(" [Align] Batch {}/{}: {} genomes [{}] ({}/{})",
1676 batch_idx + 1, total_batches, count,
1677 bucket_info.join(", "),
1678 aligned.load(Ordering::Relaxed), total_accessions);
1679 }
1680
1681 let _ = producer.join();
1683
1684 Ok(aligned.load(Ordering::Relaxed))
1685 }
1686
1687 fn extract_batch_to_files_static(zip_path: &Path, genomes_dir: &Path) -> Result<Vec<PathBuf>> {
1689 let zip_file = File::open(zip_path)?;
1690 let mut archive = zip::ZipArchive::new(zip_file)?;
1691 let mut genome_files = Vec::new();
1692
1693 for i in 0..archive.len() {
1694 let mut file = archive.by_index(i)?;
1695 let name = file.name().to_string();
1696
1697 if name.ends_with("_genomic.fna") || name.ends_with("_genomic.fasta") {
1698 let parts: Vec<&str> = name.split('/').collect();
1699 if let Some(acc_dir) = parts.iter().find(|p| p.starts_with("GCA_") || p.starts_with("GCF_")) {
1700 let filename = format!("{}.fna", acc_dir);
1701 let output_path = genomes_dir.join(&filename);
1702
1703 let mut content = Vec::new();
1705 file.read_to_end(&mut content)?;
1706 std::fs::write(&output_path, &content)?;
1707
1708 genome_files.push(output_path);
1709 }
1710 }
1711 }
1712
1713 Ok(genome_files)
1714 }
1715
1716 fn align_plsdb(
1718 &self,
1719 plasmids: &[PlasmidInfo],
1720 plsdb_dir: &Path,
1721 temp_dir: &Path,
1722 paf_output: &Path,
1723 ) -> Result<usize> {
1724 let fasta_path = plsdb_dir.join("sequences.fasta");
1725
1726 if !fasta_path.exists() {
1727 eprintln!(" Warning: sequences.fasta not found in PLSDB directory");
1728 return Ok(0);
1729 }
1730
1731 let target_accs: FxHashSet<_> = plasmids.iter()
1733 .map(|p| p.accession.clone())
1734 .collect();
1735
1736 if target_accs.is_empty() {
1737 eprintln!(" No PLSDB plasmids to process");
1738 return Ok(0);
1739 }
1740
1741 let acc_list: Vec<_> = target_accs.iter().collect();
1743 let batches: Vec<_> = acc_list.chunks(API_BATCH_SIZE).collect();
1744 let total_batches = batches.len();
1745 let mut total_processed = 0usize;
1746
1747 eprintln!(" Processing {} PLSDB plasmids in {} batches...",
1748 target_accs.len(), total_batches);
1749
1750 let file = File::open(&fasta_path)?;
1752 let reader = BufReader::new(file);
1753
1754 let mut seq_map: FxHashMap<String, String> = FxHashMap::default();
1756 let mut current_acc: Option<String> = None;
1757 let mut current_seq = String::new();
1758
1759 for line in reader.lines() {
1760 let line = line?;
1761 if line.starts_with('>') {
1762 if let Some(acc) = current_acc.take() {
1763 if target_accs.contains(&acc) {
1764 seq_map.insert(acc, std::mem::take(&mut current_seq));
1765 }
1766 }
1767 let header = line.trim_start_matches('>');
1768 let acc = header.split_whitespace().next().unwrap_or("").to_string();
1769 current_acc = Some(acc);
1770 current_seq.clear();
1771 } else {
1772 current_seq.push_str(line.trim());
1773 }
1774 }
1775 if let Some(acc) = current_acc {
1776 if target_accs.contains(&acc) {
1777 seq_map.insert(acc, current_seq);
1778 }
1779 }
1780
1781 eprintln!(" Loaded {} target sequences from PLSDB", seq_map.len());
1782
1783 for (batch_idx, batch) in batches.iter().enumerate() {
1785 let mut size_buckets: FxHashMap<GenomeSizeCategory, Vec<(&str, &str)>> = FxHashMap::default();
1787
1788 for acc in *batch {
1789 if let Some(seq) = seq_map.get(*acc) {
1790 let category = GenomeSizeCategory::from_size(seq.len() as u64);
1792 size_buckets.entry(category)
1793 .or_default()
1794 .push((acc.as_str(), seq.as_str()));
1795 }
1796 }
1797
1798 let mut batch_hits: Vec<PafHit> = Vec::new();
1799
1800 for (category, bucket_seqs) in &size_buckets {
1802 if bucket_seqs.is_empty() {
1803 continue;
1804 }
1805
1806 let bucket_threads = category.thread_count(self.threads);
1807 let bucket_suffix = format!("plsdb_batch_{:04}_{:?}", batch_idx, category);
1808 let bucket_fasta = temp_dir.join(format!("{}.fas", bucket_suffix));
1809
1810 {
1812 let mut fasta_writer = BufWriter::new(File::create(&bucket_fasta)?);
1813 for (acc, seq) in bucket_seqs {
1814 let filename = format!("{}.fna", acc.replace('.', "_"));
1815 writeln!(fasta_writer, ">{}|{}", acc, filename)?;
1816 writeln!(fasta_writer, "{}", seq)?;
1817 }
1818 fasta_writer.flush()?;
1819 }
1820
1821 let bucket_paf = temp_dir.join(format!("{}.paf", bucket_suffix));
1823 let status = Command::new("minimap2")
1824 .args([
1825 "-cx", "asm20",
1826 "-t", &bucket_threads.to_string(),
1827 self.amr_db_path.to_str().unwrap(),
1828 bucket_fasta.to_str().unwrap(),
1829 "-o", bucket_paf.to_str().unwrap(),
1830 ])
1831 .stdout(std::process::Stdio::null())
1832 .stderr(std::process::Stdio::null())
1833 .status();
1834
1835 if let Ok(s) = status {
1836 if s.success() && bucket_paf.exists() {
1837 let paf_content = std::fs::read_to_string(&bucket_paf)?;
1839 let hits: Vec<PafHit> = paf_content
1840 .lines()
1841 .filter_map(PafHit::from_paf_line)
1842 .collect();
1843 batch_hits.extend(hits);
1844 }
1845 }
1846
1847 std::fs::remove_file(&bucket_fasta).ok();
1849 std::fs::remove_file(&bucket_paf).ok();
1850 }
1851
1852 let dedup_hits = deduplicate_paf_hits(batch_hits);
1854
1855 {
1857 let mut paf_file = std::fs::OpenOptions::new()
1858 .create(true)
1859 .append(true)
1860 .open(paf_output)?;
1861 for hit in &dedup_hits {
1862 writeln!(paf_file, "{}", hit.raw_line)?;
1863 }
1864 }
1865
1866 total_processed += batch.len();
1867
1868 let bucket_info: Vec<String> = size_buckets.iter()
1870 .map(|(cat, seqs)| format!("{}={}", cat.name().split_whitespace().next().unwrap_or("?"), seqs.len()))
1871 .collect();
1872 eprintln!(" [PLSDB] Batch {}/{}: {} plasmids [{}] ({}/{})",
1873 batch_idx + 1, total_batches, batch.len(),
1874 bucket_info.join(", "),
1875 total_processed, target_accs.len());
1876 }
1877
1878 for (acc, seq) in &seq_map {
1880 let out_path = plsdb_dir.join(format!("{}.fna", acc.replace('.', "_")));
1881 if !out_path.exists() {
1882 let mut out_file = File::create(&out_path)?;
1883 writeln!(out_file, ">{}", acc)?;
1884 writeln!(out_file, "{}", seq)?;
1885 }
1886 }
1887
1888 eprintln!(" Aligned {} PLSDB sequences", total_processed);
1889 Ok(total_processed)
1890 }
1891
1892 fn extract_flanking_from_plsdb(
1894 &self,
1895 hits_path: &Path,
1896 plsdb_dir: &Path,
1897 output_path: &Path,
1898 catalog: &GenomeCatalog,
1899 ) -> Result<()> {
1900 let hits_file = File::open(hits_path)?;
1901 let reader = BufReader::new(hits_file);
1902
1903 let mut plsdb_hits: FxHashMap<String, Vec<(String, String, usize, usize)>> = FxHashMap::default();
1905
1906 for (i, line) in reader.lines().enumerate() {
1907 let line = line?;
1908 if i == 0 && line.starts_with("gene") {
1910 continue;
1911 }
1912
1913 let fields: Vec<&str> = line.split('\t').collect();
1914 if fields.len() < 4 {
1915 continue;
1916 }
1917
1918 let gene = fields[0];
1919 let contig_file = fields[1];
1920 let start: usize = fields[2].parse().unwrap_or(0);
1921 let end: usize = fields[3].parse().unwrap_or(0);
1922
1923 let (contig_id, genome_file) = if let Some(pipe_pos) = contig_file.rfind('|') {
1925 (contig_file[..pipe_pos].to_string(), contig_file[pipe_pos + 1..].to_string())
1926 } else {
1927 continue;
1928 };
1929
1930 if genome_file.starts_with("NZ_") || genome_file.starts_with("CP") ||
1932 genome_file.starts_with("AP") || genome_file.starts_with("NC_") {
1933 plsdb_hits.entry(genome_file)
1934 .or_default()
1935 .push((gene.to_string(), contig_id, start, end));
1936 }
1937 }
1938
1939 if plsdb_hits.is_empty() {
1940 return Ok(());
1941 }
1942
1943 let mut writer = std::fs::OpenOptions::new()
1945 .create(true)
1946 .append(true)
1947 .open(output_path)?;
1948
1949 let mut extracted = 0usize;
1950
1951 for (genome_file, hits) in &plsdb_hits {
1952 let genome_path = plsdb_dir.join(genome_file);
1953 if !genome_path.exists() {
1954 continue;
1955 }
1956
1957 let sequences = match self.load_genome_sequences(&genome_path) {
1959 Ok(seqs) => seqs,
1960 Err(_) => continue,
1961 };
1962
1963 let seq_map: FxHashMap<&str, &str> = sequences.iter()
1964 .map(|(header, seq)| {
1965 let key = header.split_whitespace().next().unwrap_or(header.as_str());
1966 (key, seq.as_str())
1967 })
1968 .collect();
1969
1970 let genome_acc = genome_file.trim_end_matches(".fna");
1971 let base_genus = catalog.get_genus(genome_acc).map(|s| s.to_string());
1972
1973 for (gene, contig_id, start, end) in hits {
1974 let contig_seq = match seq_map.get(contig_id.as_str()) {
1975 Some(seq) => *seq,
1976 None => continue,
1977 };
1978
1979 let contig_len = contig_seq.len();
1980 if *start >= contig_len || *end > contig_len || start >= end {
1981 continue;
1982 }
1983
1984 let genus = base_genus.clone().unwrap_or_else(|| "Unknown".to_string());
1986
1987 let upstream_start = start.saturating_sub(self.config.flanking_length);
1989 let upstream = &contig_seq[upstream_start..*start];
1990
1991 let downstream_end = std::cmp::min(*end + self.config.flanking_length, contig_len);
1992 let downstream = &contig_seq[*end..downstream_end];
1993
1994 writeln!(writer, "{}\t{}\t{}\t{}\t{}\t{}\t{}",
1995 gene, contig_id, genus, start, end, upstream, downstream)?;
1996 extracted += 1;
1997 }
1998 }
1999
2000 if extracted > 0 {
2001 eprintln!(" Extracted {} PLSDB flanking sequences", extracted);
2002 }
2003 Ok(())
2004 }
2005
2006 fn convert_paf_to_merged(
2008 &self,
2009 paf_path: &Path,
2010 output_path: &Path,
2011 ) -> Result<()> {
2012 let mut hits: FxHashSet<(String, String, usize, usize)> = FxHashSet::default();
2013
2014 if paf_path.exists() {
2017 let file = File::open(paf_path)?;
2018 let reader = BufReader::new(file);
2019 for line in reader.lines() {
2020 let line = line?;
2021 let fields: Vec<&str> = line.split('\t').collect();
2022 if fields.len() < 12 {
2023 continue;
2024 }
2025 let contig_file = fields[0].to_string(); let gene = fields[5].to_string(); let start: usize = fields[2].parse().unwrap_or(0);
2028 let end: usize = fields[3].parse().unwrap_or(0);
2029 hits.insert((gene, contig_file, start, end));
2030 }
2031 }
2032
2033 let mut writer = BufWriter::new(File::create(output_path)?);
2035 for (gene, contig_file, start, end) in &hits {
2036 writeln!(writer, "{}\t{}\t{}\t{}", gene, contig_file, start, end)?;
2037 }
2038
2039 eprintln!(" Found {} unique hits", hits.len());
2040 Ok(())
2041 }
2042
2043 fn extract_flanking_sequences(
2047 &self,
2048 hits_path: &Path,
2049 genomes_dir: &Path,
2050 output_path: &Path,
2051 catalog: &GenomeCatalog,
2052 ) -> Result<()> {
2053 if output_path.exists() {
2055 let metadata = std::fs::metadata(output_path)?;
2056 if metadata.len() > 0 {
2057 eprintln!(" Flanking sequences file already exists ({} MB), skipping extraction...",
2058 metadata.len() / (1024 * 1024));
2059 return Ok(());
2060 }
2061 }
2062
2063 let hits_file = File::open(hits_path)?;
2064 let reader = BufReader::new(hits_file);
2065
2066 let mut genome_hits: FxHashMap<String, Vec<(String, String, usize, usize)>> = FxHashMap::default();
2069
2070 for (i, line) in reader.lines().enumerate() {
2071 let line = line?;
2072 if i == 0 && line.starts_with("gene") {
2074 continue;
2075 }
2076
2077 let fields: Vec<&str> = line.split('\t').collect();
2078 if fields.len() < 4 {
2079 continue;
2080 }
2081
2082 let gene = fields[0];
2083 let contig_file = fields[1]; let start: usize = fields[2].parse().unwrap_or(0);
2085 let end: usize = fields[3].parse().unwrap_or(0);
2086
2087 let (contig_id, genome_file) = if let Some(pipe_pos) = contig_file.rfind('|') {
2089 (contig_file[..pipe_pos].to_string(), contig_file[pipe_pos + 1..].to_string())
2090 } else {
2091 (contig_file.to_string(), extract_genome_file(contig_file))
2093 };
2094
2095 genome_hits.entry(genome_file)
2096 .or_default()
2097 .push((gene.to_string(), contig_id, start, end));
2098 }
2099
2100 let output_file = File::create(output_path)?;
2102 let mut writer = BufWriter::new(output_file);
2103
2104 writeln!(writer, "Gene\tContig\tGenus\tStart\tEnd\tUpstream\tDownstream")?;
2106
2107 let total_genomes = genome_hits.len();
2108 let mut processed = 0;
2109 let mut genus_found = 0usize;
2110 let mut genus_unknown = 0usize;
2111 let mut sequences_extracted = 0usize;
2112
2113 for (genome_file, hits) in &genome_hits {
2114 let genome_path = genomes_dir.join(genome_file);
2115 if !genome_path.exists() {
2116 continue;
2117 }
2118
2119 let sequences = match self.load_genome_sequences(&genome_path) {
2121 Ok(seqs) => seqs,
2122 Err(_) => continue,
2123 };
2124
2125 let seq_map: FxHashMap<&str, &str> = sequences.iter()
2127 .map(|(header, seq)| {
2128 let key = header.split_whitespace().next().unwrap_or(header.as_str());
2129 (key, seq.as_str())
2130 })
2131 .collect();
2132
2133 let genome_acc = genome_file.trim_end_matches(".fna");
2135 let base_genus = catalog.get_genus(genome_acc).map(|s| s.to_string());
2136
2137 for (gene, contig_id, start, end) in hits {
2139 let contig_seq = match seq_map.get(contig_id.as_str()) {
2141 Some(seq) => *seq,
2142 None => continue,
2143 };
2144
2145 let contig_len = contig_seq.len();
2146
2147 if *start >= contig_len || *end > contig_len || start >= end {
2149 continue;
2150 }
2151
2152 let upstream_start = start.saturating_sub(self.config.flanking_length);
2153 let downstream_end = std::cmp::min(end + self.config.flanking_length, contig_len);
2154
2155 let upstream = if *start > 0 && upstream_start < *start {
2156 &contig_seq[upstream_start..*start]
2157 } else {
2158 ""
2159 };
2160
2161 let downstream = if *end < contig_len && *end < downstream_end {
2162 &contig_seq[*end..downstream_end]
2163 } else {
2164 ""
2165 };
2166
2167 let genus = if let Some(ref g) = base_genus {
2171 genus_found += 1;
2172 g.clone()
2173 } else if let Some(g) = catalog.get_genus(contig_id) {
2174 genus_found += 1;
2175 g.to_string()
2176 } else {
2177 genus_unknown += 1;
2178 "Unknown".to_string()
2179 };
2180
2181 writeln!(
2182 writer,
2183 "{}\t{}\t{}\t{}\t{}\t{}\t{}",
2184 gene, contig_id, genus, start, end, upstream, downstream
2185 )?;
2186 sequences_extracted += 1;
2187 }
2188
2189 processed += 1;
2190 if processed % 10000 == 0 || processed == total_genomes {
2191 eprintln!(" Processed {}/{} genomes ({} sequences)",
2192 processed, total_genomes, sequences_extracted);
2193 }
2194 }
2195
2196 eprintln!(" Extracted {} flanking sequences from {} genomes",
2197 sequences_extracted, processed);
2198 eprintln!(" Genus resolved: {}, Unknown: {}", genus_found, genus_unknown);
2199 Ok(())
2200 }
2201
2202 fn load_genome_sequences(&self, genome_path: &Path) -> Result<Vec<(String, String)>> {
2204 let file = File::open(genome_path)?;
2205 let reader = BufReader::new(file);
2206
2207 let mut sequences = Vec::new();
2208 let mut current_name: Option<String> = None;
2209 let mut current_seq = String::new();
2210
2211 for line in reader.lines() {
2212 let line = line?;
2213
2214 if let Some(stripped) = line.strip_prefix('>') {
2215 if let Some(name) = current_name.take() {
2216 sequences.push((name, std::mem::take(&mut current_seq)));
2217 }
2218 current_name = Some(stripped.to_string());
2219 current_seq.clear();
2220 } else {
2221 current_seq.push_str(line.trim());
2222 }
2223 }
2224
2225 if let Some(name) = current_name {
2226 sequences.push((name, current_seq));
2227 }
2228
2229 Ok(sequences)
2230 }
2231}
2232
2233fn parse_csv_line(line: &str) -> Vec<String> {
2235 let mut fields = Vec::new();
2236 let mut current = String::new();
2237 let mut in_quotes = false;
2238
2239 for ch in line.chars() {
2240 match ch {
2241 '"' => in_quotes = !in_quotes,
2242 ',' if !in_quotes => {
2243 fields.push(std::mem::take(&mut current));
2244 }
2245 _ => current.push(ch),
2246 }
2247 }
2248 fields.push(current);
2249 fields
2250}
2251
2252fn extract_genome_file(contig: &str) -> String {
2254 if let Some(start) = contig.find("GC") {
2256 let remainder = &contig[start..];
2257 if let Some(end) = remainder.find(|c: char| !c.is_alphanumeric() && c != '_' && c != '.') {
2258 return format!("{}.fna", &remainder[..end]);
2259 }
2260 return format!("{}.fna", remainder);
2261 }
2262
2263 let acc = contig.split_whitespace().next().unwrap_or(contig);
2265 format!("{}.fna", acc.replace('.', "_"))
2266}
2267
2268pub fn build(output_dir: &Path, arg_db: &Path, threads: usize, email: &str, config: FlankBuildConfig) -> Result<()> {
2285 let builder = FlankingDbBuilder::new(arg_db, output_dir, threads, email, config);
2286 builder.build()
2287}